| from PIL import Image |
| import torch |
| import numpy as np |
|
|
|
|
| class PreparedSrcImg: |
| def __init__(self, src_rgb, crop_trans_m, x_s_info, f_s_user, x_s_user, mask_ori): |
| self.src_rgb = src_rgb |
| self.crop_trans_m = crop_trans_m |
| self.x_s_info = x_s_info |
| self.f_s_user = f_s_user |
| self.x_s_user = x_s_user |
| self.mask_ori = mask_ori |
|
|
|
|
| def tensor2pil(image): |
| return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8)) |
|
|
|
|
| def pil2tensor(image): |
| return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) |
|
|
|
|
| def rgb_crop(rgb, region): |
| return rgb[region[1]:region[3], region[0]:region[2]] |
|
|
|
|
| def rgb_crop_batch(rgbs, region): |
| return rgbs[:, region[1]:region[3], region[0]:region[2]] |
|
|
|
|
| def get_rgb_size(rgb): |
| return rgb.shape[1], rgb.shape[0] |
|
|
|
|
| def create_transform_matrix(x, y, s_x, s_y): |
| return np.float32([[s_x, 0, x], [0, s_y, y]]) |
|
|
|
|
| def calc_crop_limit(center, img_size, crop_size): |
| pos = center - crop_size / 2 |
| if pos < 0: |
| crop_size += pos * 2 |
| pos = 0 |
|
|
| pos2 = pos + crop_size |
|
|
| if img_size < pos2: |
| crop_size -= (pos2 - img_size) * 2 |
| pos2 = img_size |
| pos = pos2 - crop_size |
|
|
| return pos, pos2, crop_size |
|
|
|
|
| def save_image(numpy_array: np.ndarray, output_path: str): |
| out = Image.fromarray(numpy_array) |
| out.save(output_path, compress_level=1, format="png") |
| return output_path |
|
|
|
|
| def image_path_to_array(image_path: str) -> np.ndarray: |
| image = Image.open(image_path) |
| image_array = np.array(image) |
| if len(image_array.shape) <= 3: |
| image_array = image_array[np.newaxis, ...] |
|
|
| return image_array |
|
|