39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import torchvision.transforms as TT
|
|
|
|
from scepter.studio.preprocess.processors.base_processor import \
|
|
BaseImageProcessor
|
|
|
|
__all__ = ['CenterCrop', 'PaddingCrop']
|
|
|
|
|
|
class CenterCrop(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def __call__(self, image, **kwargs):
|
|
w, h = image.size
|
|
height_ratio = kwargs.get('height_ratio', 1)
|
|
width_ratio = kwargs.get('width_ratio', 1)
|
|
|
|
output_h_align_height, output_h_align_width = h, int(h / height_ratio *
|
|
width_ratio)
|
|
if output_h_align_height * output_h_align_width <= w * h:
|
|
output_height, output_width = output_h_align_height, output_h_align_width
|
|
else:
|
|
output_height, output_width = int(w / width_ratio *
|
|
height_ratio), w
|
|
image = TT.Resize(max(output_height, output_width))(image)
|
|
image = TT.CenterCrop((output_height, output_width))(image)
|
|
return image
|
|
|
|
|
|
class PaddingCrop(BaseImageProcessor):
|
|
def __init__(self, cfg, language='en'):
|
|
super().__init__(cfg, language=language)
|
|
|
|
def get_caption(self, image, **kwargs):
|
|
return image
|