247 lines
9.2 KiB
Python
247 lines
9.2 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import numbers
|
|
import pdb
|
|
from typing import Any, Dict, Union
|
|
|
|
import cv2
|
|
import json
|
|
import numpy as np
|
|
import torch
|
|
from torchvision import transforms
|
|
|
|
from modelscope.metainfo import Preprocessors
|
|
from modelscope.preprocessors import Preprocessor
|
|
from modelscope.preprocessors.builder import PREPROCESSORS
|
|
from modelscope.preprocessors.image import LoadImage
|
|
from modelscope.utils.constant import Fields, ModeKeys
|
|
|
|
_cv2_pad_to_str = {
|
|
'constant': cv2.BORDER_CONSTANT,
|
|
'edge': cv2.BORDER_REPLICATE,
|
|
'reflect': cv2.BORDER_REFLECT_101,
|
|
'symmetric': cv2.BORDER_REFLECT,
|
|
}
|
|
|
|
|
|
@PREPROCESSORS.register_module(
|
|
Fields.cv, module_name=Preprocessors.image_sky_change_preprocessor)
|
|
class ImageSkyChangePreprocessor(Preprocessor):
|
|
|
|
def __init__(self,
|
|
model_dir: str = None,
|
|
mode: str = ModeKeys.INFERENCE,
|
|
coarse_model_width=640,
|
|
coarse_model_height=640,
|
|
refine_model_width=1280,
|
|
refine_model_height=1280,
|
|
mean_vec=[0.485, 0.456, 0.406],
|
|
std_vec=[0.229, 0.224, 0.225],
|
|
*args,
|
|
**kwargs):
|
|
"""
|
|
Args:
|
|
model_dir (str): model directory to initialize some resource.
|
|
mode: The mode for the preprocessor.
|
|
coarse_model_width: required width of input tensor of coarse model.
|
|
coarse_model_height: required height of input tensor of coarse model.
|
|
refine_model_width: required width of input tensor of refine model.
|
|
refine_model_height: required height of input tensor of refine model.
|
|
mean_vec: mean of dataset(for transforms.Normalize), default is mean of Imagenet dataset.
|
|
std_vec: standard deviation of dataset(for transforms.Normalize), default is std of Imagenet dataset.
|
|
"""
|
|
super().__init__(mode)
|
|
|
|
# set preprocessor info
|
|
self.coarse_input_size = [coarse_model_width, coarse_model_height]
|
|
self.refine_input_size = [refine_model_width, refine_model_height]
|
|
self.normalize = transforms.Normalize(mean=mean_vec, std=std_vec)
|
|
|
|
def __call__(self, data: Union[str, Dict], **kwargs) -> Dict[str, Any]:
|
|
"""process the raw input data
|
|
Args:
|
|
data (dict): data dict containing following info:
|
|
sky_image, scene_image
|
|
|
|
Example:
|
|
>>> {
|
|
>>> "sky_image": "xxx.jpg" # sky_image path(str)
|
|
>>> "scene_image": "xxx.jpg", # scene_image path(str)
|
|
>>> }
|
|
|
|
Returns:
|
|
Dict[str, Any]: the preprocessed data
|
|
{
|
|
"sky_image": the preprocessed sky image(origin size)
|
|
"sky_image_refine": the preprocessed resized sky image
|
|
"scene_image": the preprocessed scene image(origin size)
|
|
"scene_image_refine": the preprocessed resized scene image
|
|
"img_metas": informations of preprocessed images, e.g. origin shape, pad information, resized shape.
|
|
}
|
|
"""
|
|
if 'sky_image' not in data.keys():
|
|
raise Exception('sky_image not in input data')
|
|
if 'scene_image' not in data.keys():
|
|
raise Exception('scene_image not in input data')
|
|
if isinstance(data['sky_image'], str):
|
|
sky_image = LoadImage.convert_to_ndarray(data['sky_image'])
|
|
sky_image = sky_image.astype(np.uint8) # RGB
|
|
sky_image = cv2.cvtColor(sky_image, cv2.COLOR_RGB2BGR) # BGR
|
|
if sky_image is not None:
|
|
sky_image = self.check_image(sky_image)
|
|
else:
|
|
raise Exception('sky_image is None')
|
|
else:
|
|
raise Exception('sky_image(path of sky image) is not valid')
|
|
if isinstance(data['scene_image'], str):
|
|
scene_image = LoadImage.convert_to_ndarray(data['scene_image'])
|
|
scene_image = scene_image.astype(np.uint8) # RGB
|
|
scene_image = cv2.cvtColor(scene_image, cv2.COLOR_RGB2BGR) # BGR
|
|
if scene_image is not None:
|
|
scene_image = self.check_image(scene_image)
|
|
else:
|
|
raise Exception('scene_image is None')
|
|
else:
|
|
raise Exception('scene_image(path of scene image) is not valid')
|
|
data = {}
|
|
sky_image_refine, sky_img_metas = self.process_single_img(sky_image)
|
|
scene_image_refine, scene_img_metas = self.process_single_img(
|
|
scene_image)
|
|
data['sky_image'] = sky_image
|
|
data['sky_image_refine'] = sky_image_refine
|
|
data['scene_image'] = scene_image
|
|
data['scene_image_refine'] = scene_image_refine
|
|
data['img_metas'] = {
|
|
'sky_img_metas': sky_img_metas,
|
|
'scene_img_metas': scene_img_metas,
|
|
'input_size': {
|
|
'coarse_input_size': self.coarse_input_size,
|
|
'refine_input_size': self.refine_input_size
|
|
}
|
|
}
|
|
return data
|
|
|
|
def process_single_img(self, img):
|
|
img_metas = {}
|
|
img_metas['ori_shape'] = img.shape[0:2] # img: (origin_h, origin_w, 3)
|
|
img, pad_direction = get_refine_input(img, self.refine_input_size)
|
|
img = image_transform(
|
|
img, self.normalize) # torch.Size([3, refine_net_h, refine_net_w])
|
|
img = img.unsqueeze(0)
|
|
img_metas['pad_direction'] = pad_direction
|
|
img_metas['refine_shape'] = img.shape[
|
|
2:] # torch.Size([1, 3, refine_net_h, refine_net_w])
|
|
return img, img_metas
|
|
|
|
def check_image(self, input_img):
|
|
whole_temp_shape = input_img.shape
|
|
if len(whole_temp_shape) == 2:
|
|
input_img = np.stack([input_img, input_img, input_img], axis=2)
|
|
elif whole_temp_shape[2] == 1:
|
|
input_img = np.concatenate([input_img, input_img, input_img],
|
|
axis=2)
|
|
elif whole_temp_shape[2] == 4:
|
|
input_img = input_img[:, :,
|
|
0:3] * 1.0 * input_img[:, :,
|
|
3:4] * 1.0 / 255.0
|
|
return input_img
|
|
|
|
|
|
def get_refine_input(mat, refine_input_size):
|
|
# maxDimMatch: resize
|
|
mat = max_dim_match(mat, refine_input_size)
|
|
# pad image to refine net input size
|
|
mat, pad_direction = center_pad_image_withwh(mat, refine_input_size, 0)
|
|
return mat, pad_direction
|
|
|
|
|
|
def max_dim_match(image, refine_model_size):
|
|
h, w, c = np.shape(image)
|
|
resize_w, resize_h = refine_model_size
|
|
if h != resize_h or w != resize_w:
|
|
h_scale = float(resize_h) / h
|
|
w_scale = float(resize_w) / w
|
|
resize_scale = min(w_scale, h_scale)
|
|
new_h = int(h * resize_scale + 0.5)
|
|
new_w = int(w * resize_scale + 0.5)
|
|
image = cv2.resize(
|
|
image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
|
|
return image
|
|
|
|
|
|
def center_pad_image_withwh(image,
|
|
crop_size,
|
|
padvalue,
|
|
padding_mode='constant'):
|
|
pad_image = image
|
|
h, w = image.shape[0], image.shape[1]
|
|
pad_h = max(crop_size[1] - h, 0)
|
|
pad_w = max(crop_size[0] - w, 0)
|
|
pad_direction = (0, 0, 0, 0)
|
|
if pad_h > 0 or pad_w > 0:
|
|
half_w = int(pad_w / 2 + 0.5)
|
|
half_h = int(pad_h / 2 + 0.5)
|
|
pad_direction = (half_w, half_h, pad_w - half_w, pad_h - half_h)
|
|
pad_image = pad(
|
|
image, pad_direction, padvalue, padding_mode=padding_mode)
|
|
return pad_image, pad_direction
|
|
|
|
|
|
def pad(img, padding, fill=0, padding_mode='constant'):
|
|
if not is_numpy_image(img):
|
|
raise TypeError('img should be numpy ndarray. Got {}'.format(
|
|
type(img)))
|
|
if not isinstance(padding,
|
|
(numbers.Number, tuple, list)) or len(padding) != 4:
|
|
raise TypeError('Got inappropriate padding arg')
|
|
|
|
pad_left = padding[0]
|
|
pad_top = padding[1]
|
|
pad_right = padding[2]
|
|
pad_bottom = padding[3]
|
|
|
|
shape_len = len(img.shape)
|
|
if shape_len == 2:
|
|
return cv2.copyMakeBorder(
|
|
img,
|
|
top=pad_top,
|
|
bottom=pad_bottom,
|
|
left=pad_left,
|
|
right=pad_right,
|
|
borderType=_cv2_pad_to_str[padding_mode],
|
|
value=fill,
|
|
)
|
|
elif shape_len == 3 and img.shape[2] == 1:
|
|
return cv2.copyMakeBorder(
|
|
img,
|
|
top=pad_top,
|
|
bottom=pad_bottom,
|
|
left=pad_left,
|
|
right=pad_right,
|
|
borderType=_cv2_pad_to_str[padding_mode],
|
|
value=fill,
|
|
)[:, :, np.newaxis]
|
|
else:
|
|
return cv2.copyMakeBorder(
|
|
img,
|
|
top=pad_top,
|
|
bottom=pad_bottom,
|
|
left=pad_left,
|
|
right=pad_right,
|
|
borderType=_cv2_pad_to_str[padding_mode],
|
|
value=fill,
|
|
)
|
|
|
|
|
|
def image_transform(img, normalize):
|
|
img = img[:, :, ::-1] # BGR-->RGB to pil format
|
|
img = img.transpose((2, 0, 1)) # h,w,c --> c,h,w
|
|
img = img.astype(np.float32) / 255
|
|
img = normalize(torch.from_numpy(img.copy()))
|
|
return img
|
|
|
|
|
|
def is_numpy_image(img):
|
|
return isinstance(img, np.ndarray) and (img.ndim in {2, 3})
|