File size: 986 Bytes
d8afa61
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import cv2
import numpy as np
import torch
from torchvision import transforms

class Resize(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image):
        image = cv2.resize(image, (self.size, self.size))
        return image

class NormalizeImage(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image):
        image = image.astype(np.float32) / 255.0
        image -= np.array(self.mean)
        image /= np.array(self.std)
        return image

class PrepareForNet(object):
    def __call__(self, image):
        image = torch.from_numpy(image)
        if len(image.shape) == 3:
            image = image.permute(2, 0, 1)
        image = image.unsqueeze(0)
        return image

class Compose:
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, img):
        for t in self.transforms:
            img = t(img)
        return img