File size: 1,531 Bytes
07e1105 |
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 |
import torch
import numpy as np
class Normalize(object):
def __init__(self, mean, var):
self.mean = mean
self.var = var
def __call__(self, sample):
if isinstance(sample, dict):
img = sample['img']
gt = sample['gt']
img = (img - self.mean) / self.var
sample = {'img': img, 'gt': gt}
else:
sample = (sample - self.mean) / self.var
return sample
class RandHorizontalFlip(object):
def __init__(self, prob_aug):
self.prob_aug = prob_aug
def __call__(self, sample):
p_aug = np.array([self.prob_aug, 1 - self.prob_aug])
prob_lr = np.random.choice([1, 0], p=p_aug.ravel())
if isinstance(sample, dict):
img = sample['img']
gt = sample['gt']
if prob_lr > 0.5:
img = np.fliplr(img).copy()
sample = {'img': img, 'gt': gt}
else:
if prob_lr > 0.5:
sample = np.fliplr(sample).copy()
return sample
class ToTensor(object):
def __init__(self):
pass
def __call__(self, sample):
if isinstance(sample, dict):
img = sample['img']
gt = sample['gt']
img = torch.from_numpy(img).type(torch.FloatTensor)
gt = torch.from_numpy(gt).type(torch.FloatTensor)
sample = {'img': img, 'gt': gt}
else:
sample = torch.from_numpy(sample).type(torch.FloatTensor)
return sample
|