File size: 2,536 Bytes
bcf59c3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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):
# function resizes and keeps the aspect ratio same
crop_shape = crop.shape # hxwxc
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)
|