| import os |
| from pathlib import Path |
|
|
| import cv2 |
| from PIL import Image |
| from torchvision import transforms |
|
|
|
|
| class NormalizeImage(object): |
| """Normalize given tensor into given mean and standard dev |
| |
| Args: |
| mean (float): Desired mean to substract from tensors |
| std (float): Desired std to divide from tensors |
| """ |
|
|
| def __init__(self, mean, std): |
| assert isinstance(mean, (float)) |
| if isinstance(mean, float): |
| self.mean = mean |
|
|
| if isinstance(std, float): |
| self.std = std |
|
|
| self.normalize_1 = transforms.Normalize(self.mean, self.std) |
| self.normalize_3 = transforms.Normalize([self.mean] * 3, [self.std] * 3) |
| self.normalize_18 = transforms.Normalize([self.mean] * 18, [self.std] * 18) |
|
|
| def __call__(self, image_tensor): |
| if image_tensor.shape[0] == 1: |
| return self.normalize_1(image_tensor) |
|
|
| elif image_tensor.shape[0] == 3: |
| return self.normalize_3(image_tensor) |
|
|
| elif image_tensor.shape[0] == 18: |
| return self.normalize_18(image_tensor) |
|
|
| else: |
| assert "Please set proper channels! Normalization implemented only for 1, 3 and 18" |
|
|
|
|
| def naive_cutout(img, mask): |
| empty = Image.new("RGBA", (img.size), 0) |
| cutout = Image.composite(img, empty, mask.resize(img.size, Image.LANCZOS)) |
| return cutout |
|
|
|
|
| def resize_by_bigger_index(crop): |
| |
| crop_shape = crop.shape |
| if crop_shape[0] / crop_shape[1] <= 1.33: |
| resized_crop = image_resize(crop, width=768) |
| else: |
| resized_crop = image_resize(crop, height=1024) |
| return resized_crop |
|
|
|
|
| def image_resize(image, width=None, height=None): |
| dim = None |
| (h, w) = image.shape[:2] |
|
|
| if width is None and height is None: |
| return image |
|
|
| if width is None: |
| r = height / float(h) |
| dim = (int(w * r), height) |
|
|
| else: |
| r = width / float(w) |
| dim = (width, int(h * r)) |
|
|
| resized = cv2.resize(image, dim) |
|
|
| return resized |
|
|
|
|
| def convert_to_jpg(image_path, output_dir, size=None): |
| """ |
| Convert image to jpg format |
| :param image_path: image path |
| :param output_dir: output directory |
| :param size: desired size of the image (w, h) |
| """ |
| img = cv2.imread(image_path) |
| if size is not None: |
| img = image_resize(img, width=size[0], height=size[1]) |
|
|
| filename = Path(image_path).name |
| cv2.imwrite(os.path.join(output_dir, filename.split(".")[0] + ".jpg"), img) |
|
|
