def _augment(self, image):
# OpenCV channels ordering
image = image.transpose(1, 2, 0) # W x H x C
if self.resize:
image = cv.resize(image, (self.resize, self.resize))
if self.crop:
assert isinstance(self.crop,
int), 'If crop provided, it should be integer'
assert self.crop < self.resize, 'Crop should be smaller than image size'
shift = np.random.randint(0, self.resize - self.crop - 1)
image = image[shift:(shift + self.crop),
shift:(shift + self.crop), :]
if self.flip_v:
if bernoulli(self.flip_v) == 1:
image = cv.flip(image, 0)
if self.flip_h:
if bernoulli(self.flip_h) == 1:
image = cv.flip(image, 1)
if self.hat:
kernel = np.ones(self.hat, np.uint8)
image = cv.morphologyEx(image, cv.MORPH_TOPHAT, self.hat)
image = image.astype(np.float32)
if self.rescale:
image /= 255
if self.norm:
mean, std = image.mean(axis=(0, 1)), image.std(axis=(0, 1))
image = (image - mean) / (std + 1e-8)
if self.sigma_clip:
image = np.clip(image, -self.sigma_clip, self.sigma_clip)
image = image.transpose(2, 0, 1) # C x W x H
return image
def __getitem__(self, index):
s1, s2 = self.ds1[index], self.ds2[index]
if self.swap and bernoulli(self.swap) == 1:
s1, s2 = s2, s1
return {'site1': s1, 'site2': s2}