File size: 4,415 Bytes
0b32e3c |
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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 |
import numpy as np
from PIL import Image
import random
import torch
from torchvision import transforms as T
from torchvision.transforms import functional as F
class transform(object):
def __init__(self, args):
self.Resize = Resize(args.img_size, args.img_size)
self.ToTensor = ToTensor()
self.Norm = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
def __call__(self, image, target):
image, target = self.Resize(image, target)
image, target = self.ToTensor(image, target)
image, target = self.Norm(image, target)
return image, target
def pad_if_smaller(img, size, fill=0):
min_size = min(img.size)
if min_size < size:
ow, oh = img.size
padh = size - oh if oh < size else 0
padw = size - ow if ow < size else 0
img = F.pad(img, (0, 0, padw, padh), fill=fill)
return img
class Compose(object):
def __init__(self, transforms):
self.transforms = transforms
def __call__(self, image, target):
for t in self.transforms:
image, target = t(image, target)
return image, target
class Resize(object):
def __init__(self, h, w):
self.h = h
self.w = w
def __call__(self, image, target):
image = F.resize(image, (self.h, self.w))
# If size is a sequence like (h, w), the output size will be matched to this.
# If size is an int, the smaller edge of the image will be matched to this number maintaining the aspect ratio
target = F.resize(target, (self.h, self.w), interpolation=Image.NEAREST)
return image, target
class RandomResize(object):
def __init__(self, min_size, max_size=None):
self.min_size = min_size
if max_size is None:
max_size = min_size
self.max_size = max_size
def __call__(self, image, target):
size = random.randint(self.min_size, self.max_size) # Return a random integer N such that a <= N <= b. Alias for randrange(a, b+1)
image = F.resize(image, size)
# If size is a sequence like (h, w), the output size will be matched to this.
# If size is an int, the smaller edge of the image will be matched to this number maintaining the aspect ratio
target = F.resize(target, size, interpolation=Image.NEAREST)
return image, target
class RandomHorizontalFlip(object):
def __init__(self, flip_prob):
self.flip_prob = flip_prob
def __call__(self, image, target):
if random.random() < self.flip_prob:
image = F.hflip(image)
target = F.hflip(target)
return image, target
class RandomCrop(object):
def __init__(self, size):
self.size = size
def __call__(self, image, target):
image = pad_if_smaller(image, self.size)
target = pad_if_smaller(target, self.size, fill=0)
crop_params = T.RandomCrop.get_params(image, (self.size, self.size))
image = F.crop(image, *crop_params)
target = F.crop(target, *crop_params)
# print(image.size, target.size)
return image, target
class CenterCrop(object):
def __init__(self, size):
self.size = size
def __call__(self, image, target):
image = F.center_crop(image, self.size)
target = F.center_crop(target, self.size)
return image, target
class ToTensor(object):
def __call__(self, image, target):
image = F.to_tensor(image)
target = torch.as_tensor(np.asarray(target).copy(), dtype=torch.int64)
return image, target
class RandomAffine(object):
def __init__(self, angle, translate, scale, shear, resample=0, fillcolor=None):
self.angle = angle
self.translate = translate
self.scale = scale
self.shear = shear
self.resample = resample
self.fillcolor = fillcolor
def __call__(self, image, target):
affine_params = T.RandomAffine.get_params(self.angle, self.translate, self.scale, self.shear, image.size)
image = F.affine(image, *affine_params)
target = F.affine(target, *affine_params)
return image, target
class Normalize(object):
def __init__(self, mean, std):
self.mean = mean
self.std = std
def __call__(self, image, target):
image = F.normalize(image, mean=self.mean, std=self.std)
return image, target
|