KyanChen/BuildingExtraction
5
1import numpy as np2import cv23import torch4 5class Compose(object):6 """Composes several transforms together.7 8 Args:9 transforms (list of ``Transform`` objects): list of transforms to compose.10 11 Example:12 >>> transforms.Compose([13 >>> transforms.CenterCrop(10),14 >>> transforms.ToTensor(),15 >>> ])16 """17 18 def __init__(self, transforms):19 self.transforms = transforms20 21 def __call__(self, data):22 for t in self.transforms:23 data = t(data)24 return data25 26 def __repr__(self):27 format_string = self.__class__.__name__ + '('28 for t in self.transforms:29 format_string += '\n'30 format_string += ' {0}'.format(t)31 format_string += '\n)'32 return format_string33 34 35class ConvertUcharToFloat(object):36 """37 Convert img form uchar to float3238 """39 40 def __call__(self, data):41 data = [x.astype(np.float32) for x in data]42 return data43 44 45class RandomContrast(object):46 """47 Get random contrast img48 """49 def __init__(self, phase, lower=0.8, upper=1.2, prob=0.5):50 self.phase = phase51 self.lower = lower52 self.upper = upper53 self.prob = prob54 assert self.upper >= self.lower, "contrast upper must be >= lower!"55 assert self.lower > 0, "contrast lower must be non-negative!"56 57 def __call__(self, data):58 if self.phase in ['od', 'seg']:59 img, _ = data60 if torch.rand(1) < self.prob:61 alpha = torch.FloatTensor(1).uniform_(self.lower, self.upper)62 img *= alpha.numpy()63 return_data = img, _64 elif self.phase == 'cd':65 img1, label1, img2, label2 = data66 if torch.rand(1) < self.prob:67 alpha = torch.FloatTensor(1).uniform_(self.lower, self.upper)68 img1 *= alpha.numpy()69 if torch.rand(1) < self.prob:70 alpha = torch.FloatTensor(1).uniform_(self.lower, self.upper)71 img2 *= alpha.numpy()72 return_data = img1, label1, img2, label273 return return_data74 75 76class RandomBrightness(object):77 """78 Get random brightness img79 """80 def __init__(self, phase, delta=10, prob=0.5):81 self.phase = phase82 self.delta = delta83 self.prob = prob84 assert 0. <= self.delta < 255., "brightness delta must between 0 to 255"85 86 def __call__(self, data):87 if self.phase in ['od', 'seg']:88 img, _ = data89 if torch.rand(1) < self.prob:90 delta = torch.FloatTensor(1).uniform_(- self.delta, self.delta)91 img += delta.numpy()92 return_data = img, _93 94 elif self.phase == 'cd':95 img1, label1, img2, label2 = data96 if torch.rand(1) < self.prob:97 delta = torch.FloatTensor(1).uniform_(- self.delta, self.delta)98 img1 += delta.numpy()99 if torch.rand(1) < self.prob:100 delta = torch.FloatTensor(1).uniform_(- self.delta, self.delta)101 img2 += delta.numpy()102 return_data = img1, label1, img2, label2103 104 return return_data105 106 107class ConvertColor(object):108 """109 Convert img color BGR to HSV or HSV to BGR for later img distortion.110 """111 def __init__(self, phase, current='RGB', target='HSV'):112 self.phase = phase113 self.current = current114 self.target = target115 116 def __call__(self, data):117 118 if self.phase in ['od', 'seg']:119 img, _ = data120 if self.current == 'RGB' and self.target == 'HSV':121 img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)122 elif self.current == 'HSV' and self.target == 'RGB':123 img = cv2.cvtColor(img, cv2.COLOR_HSV2RGB)124 else:125 raise NotImplementedError("Convert color fail!")126 return_data = img, _127 128 elif self.phase == 'cd':129 img1, label1, img2, label2 = data130 if self.current == 'RGB' and self.target == 'HSV':131 img1 = cv2.cvtColor(img1, cv2.COLOR_RGB2HSV)132 img2 = cv2.cvtColor(img2, cv2.COLOR_RGB2HSV)133 elif self.current == 'HSV' and self.target == 'RGB':134 img1 = cv2.cvtColor(img1, cv2.COLOR_HSV2RGB)135 img2 = cv2.cvtColor(img2, cv2.COLOR_HSV2RGB)136 else:137 raise NotImplementedError("Convert color fail!")138 return_data = img1, label1, img2, label2139 140 return return_data141 142 143class RandomSaturation(object):144 """145 get random saturation img146 apply the restriction on saturation S147 """148 def __init__(self, phase, lower=0.8, upper=1.2, prob=0.5):149 self.phase = phase150 self.lower = lower151 self.upper = upper152 self.prob = prob153 assert self.upper >= self.lower, "saturation upper must be >= lower!"154 assert self.lower > 0, "saturation lower must be non-negative!"155 156 def __call__(self, data):157 if self.phase in ['od', 'seg']:158 img, _ = data159 if torch.rand(1) < self.prob:160 alpha = torch.FloatTensor(1).uniform_(self.lower, self.upper)161 img[:, :, 1] *= alpha.numpy()162 return_data = img, _163 elif self.phase == 'cd':164 img1, label1, img2, label2 = data165 if torch.rand(1) < self.prob:166 alpha = torch.FloatTensor(1).uniform_(self.lower, self.upper)167 img1[:, :, 1] *= alpha.numpy()168 if torch.rand(1) < self.prob:169 alpha = torch.FloatTensor(1).uniform_(self.lower, self.upper)170 img2[:, :, 1] *= alpha.numpy()171 return_data = img1, label1, img2, label2172 return return_data173 174 175class RandomHue(object):176 """177 get random Hue img178 apply the restriction on Hue H179 """180 def __init__(self, phase, delta=10., prob=0.5):181 self.phase = phase182 self.delta = delta183 self.prob = prob184 assert 0 <= self.delta < 360, "Hue delta must between 0 to 360!"185 186 def __call__(self, data):187 if self.phase in ['od', 'seg']:188 img, _ = data189 if torch.rand(1) < self.prob:190 alpha = torch.FloatTensor(1).uniform_(-self.delta, self.delta)191 img[:, :, 0] += alpha.numpy()192 img[:, :, 0][img[:, :, 0] > 360.0] -= 360.0193 img[:, :, 0][img[:, :, 0] < 0.0] += 360.0194 return_data = img, _195 196 elif self.phase == 'cd':197 img1, label1, img2, label2 = data198 if torch.rand(1) < self.prob:199 alpha = torch.FloatTensor(1).uniform_(-self.delta, self.delta)200 img1[:, :, 0] += alpha.numpy()201 img1[:, :, 0][img1[:, :, 0] > 360.0] -= 360.0202 img1[:, :, 0][img1[:, :, 0] < 0.0] += 360.0203 if torch.rand(1) < self.prob:204 alpha = torch.FloatTensor(1).uniform_(-self.delta, self.delta)205 img2[:, :, 0] += alpha.numpy()206 img2[:, :, 0][img2[:, :, 0] > 360.0] -= 360.0207 img2[:, :, 0][img2[:, :, 0] < 0.0] += 360.0208 209 return_data = img1, label1, img2, label2210 211 return return_data212 213 214class RandomChannelNoise(object):215 """216 Get random shuffle channels217 """218 def __init__(self, phase, prob=0.4):219 self.phase = phase220 self.prob = prob221 self.perms = ((0, 1, 2), (0, 2, 1),222 (1, 0, 2), (1, 2, 0),223 (2, 0, 1), (2, 1, 0))224 225 def __call__(self, data):226 if self.phase in ['od', 'seg']:227 img, _ = data228 if torch.rand(1) < self.prob:229 shuffle_factor = self.perms[torch.randint(0, len(self.perms), size=[])]230 img = img[:, :, shuffle_factor]231 return_data = img, _232 233 elif self.phase == 'cd':234 img1, label1, img2, label2 = data235 if torch.rand(1) < self.prob:236 shuffle_factor = self.perms[torch.randint(0, len(self.perms), size=[])]237 img1 = img1[:, :, shuffle_factor]238 if torch.rand(1) < self.prob:239 shuffle_factor = self.perms[torch.randint(0, len(self.perms), size=[])]240 img2 = img2[:, :, shuffle_factor]241 return_data = img1, label1, img2, label2242 243 return return_data244 245 246class ImgDistortion(object):247 """248 Change img by distortion249 """250 def __init__(self, phase, prob=0.5):251 self.phase = phase252 self.prob = prob253 self.operation = [254 RandomContrast(phase),255 ConvertColor(phase, current='RGB', target='HSV'),256 RandomSaturation(phase),257 RandomHue(phase),258 ConvertColor(phase, current='HSV', target='RGB'),259 RandomContrast(phase)260 ]261 self.random_brightness = RandomBrightness(phase)262 self.random_light_noise = RandomChannelNoise(phase)263 264 def __call__(self, data):265 if torch.rand(1) < self.prob:266 data = self.random_brightness(data)267 if torch.rand(1) < self.prob:268 distort = Compose(self.operation[:-1])269 else:270 distort = Compose(self.operation[1:])271 data = distort(data)272 data = self.random_light_noise(data)273 return data274 275 276class ExpandImg(object):277 """278 Get expand img279 """280 def __init__(self, phase, prior_mean, prob=0.5, expand_ratio=0.2):281 self.phase = phase282 self.prior_mean = np.array(prior_mean) * 255283 self.prob = prob284 self.expand_ratio = expand_ratio285 286 def __call__(self, data):287 if self.phase == 'seg':288 img, label = data289 if torch.rand(1) < self.prob:290 return data291 height, width, channels = img.shape292 ratio_width = self.expand_ratio * torch.rand([])293 ratio_height = self.expand_ratio * torch.rand([])294 left, right = torch.randint(high=int(max(1, width * ratio_width)), size=[2])295 top, bottom = torch.randint(high=int(max(1, width * ratio_height)), size=[2])296 img = cv2.copyMakeBorder(297 img, int(top), int(bottom), int(left), int(right), cv2.BORDER_CONSTANT, value=self.prior_mean)298 label = cv2.copyMakeBorder(299 label, int(top), int(bottom), int(left), int(right), cv2.BORDER_CONSTANT, value=0)300 return img, label301 elif self.phase == 'cd':302 img1, label1, img2, label2 = data303 if torch.rand(1) < self.prob:304 return data305 height, width, channels = img1.shape306 ratio_width = self.expand_ratio * torch.rand([])307 ratio_height = self.expand_ratio * torch.rand([])308 left, right = torch.randint(high=int(max(1, width * ratio_width)), size=[2])309 top, bottom = torch.randint(high=int(max(1, width * ratio_height)), size=[2])310 img1 = cv2.copyMakeBorder(311 img1, int(top), int(bottom), int(left), int(right), cv2.BORDER_CONSTANT, value=self.prior_mean)312 label1 = cv2.copyMakeBorder(313 label1, int(top), int(bottom), int(left), int(right), cv2.BORDER_CONSTANT, value=0)314 img2 = cv2.copyMakeBorder(315 img2, int(top), int(bottom), int(left), int(right), cv2.BORDER_CONSTANT, value=self.prior_mean)316 label2 = cv2.copyMakeBorder(317 label2, int(top), int(bottom), int(left), int(right), cv2.BORDER_CONSTANT, value=0)318 return img1, label1, img2, label2319 320 elif self.phase == 'od':321 if torch.rand(1) < self.prob:322 return data323 img, label = data324 height, width, channels = img.shape325 ratio_width = self.expand_ratio * torch.rand([])326 ratio_height = self.expand_ratio * torch.rand([])327 left, right = torch.randint(high=int(max(1, width * ratio_width)), size=[2])328 top, bottom = torch.randint(high=int(max(1, width * ratio_height)), size=[2])329 left = int(left)330 right = int(right)331 top = int(top)332 bottom = int(bottom)333 img = cv2.copyMakeBorder(334 img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=self.prior_mean)335 336 label[:, 1::2] += left337 label[:, 2::2] += top338 return img, label339 340 341class RandomSampleCrop(object):342 """343 Crop344 Arguments:345 img (Image): the image being input during training346 boxes (Tensor): the original bounding boxes in pt form347 label (Tensor): the class label for each bbox348 mode (float tuple): the min and max jaccard overlaps349 Return:350 (img, boxes, classes)351 img (Image): the cropped image352 boxes (Tensor): the adjusted bounding boxes in pt form353 label (Tensor): the class label for each bbox354 """355 def __init__(self,356 phase,357 original_size=[512, 512],358 prob=0.5,359 crop_scale_ratios_range=[0.8, 1.2],360 aspect_ratio_range=[4./5, 5./4]):361 self.phase = phase362 self.prob = prob363 self.scale_range = crop_scale_ratios_range364 self.original_size = original_size365 self.aspect_ratio_range = aspect_ratio_range # h/w366 self.max_try_times = 500367 368 def __call__(self, data):369 if self.phase == 'seg':370 img, label = data371 w, h, c = img.shape372 if torch.rand(1) < self.prob:373 return data374 else:375 try_times = 0376 while try_times < self.max_try_times:377 crop_w = torch.randint(378 min(w, int(self.scale_range[0] * self.original_size[0])),379 min(w + 1, int(self.scale_range[1] * self.original_size[0])),380 size=[]381 )382 crop_h = torch.randint(383 min(h, int(self.scale_range[0] * self.original_size[1])),384 min(h + 1, int(self.scale_range[1] * self.original_size[1])),385 size=[]386 )387 # aspect ratio constraint388 if self.aspect_ratio_range[0] < crop_h / crop_w < self.aspect_ratio_range[1]:389 break390 else:391 try_times += 1392 if try_times >= self.max_try_times:393 print("try times over max threshold!", flush=True)394 return img, label395 396 left = torch.randint(0, w - crop_w + 1, size=[])397 top = torch.randint(0, h - crop_h + 1, size=[])398 img = img[top:(top + crop_h), left:(left + crop_w), :]399 label = label[top:(top + crop_h), left:(left + crop_w)]400 return img, label401 402 elif self.phase == 'od':403 if torch.rand(1) < self.prob:404 return data405 img, label = data406 w, h, c = img.shape407 408 while True:409 crop_w = torch.randint(410 min(w, int(self.scale_range[0] * self.original_size[0])),411 min(w + 1, int(self.scale_range[1] * self.original_size[0])),412 size=[]413 )414 crop_h = torch.randint(415 min(h, int(self.scale_range[0] * self.original_size[1])),416 min(h + 1, int(self.scale_range[1] * self.original_size[1])),417 size=[]418 )419 420 # aspect ratio constraint421 if self.aspect_ratio_range[0] < crop_h / crop_w < self.aspect_ratio_range[1]:422 break423 424 left = torch.randint(0, w - crop_w + 1, size=[])425 top = torch.randint(0, h - crop_h + 1, size=[])426 left = left.numpy()427 top = top.numpy()428 crop_h = crop_h.numpy()429 crop_w = crop_w.numpy()430 img = img[top:(top + crop_h), left:(left + crop_w), :]431 if len(label):432 # keep overlap with gt box IF center in sampled patch433 centers = (label[:, 1:3] + label[:, 3:]) / 2.0434 # mask in all gt boxes that above and to the left of centers435 m1 = (left <= centers[:, 0]) * (top <= centers[:, 1])436 # mask in all gt boxes that under and to the right of centers437 m2 = ((left + crop_w) >= centers[:, 0]) * ((top + crop_h) > centers[:, 1])438 # mask in that both m1 and m2 are true439 mask = m1 * m2440 441 # take only matching gt boxes442 current_label = label[mask, :]443 444 # adjust to crop (by substracting crop's left,top)445 current_label[:, 1::2] -= left446 current_label[:, 2::2] -= top447 label = current_label448 return img, label449 450 451class RandomMirror(object):452 def __init__(self, phase, prob=0.5):453 self.phase = phase454 self.prob = prob455 456 def __call__(self, data):457 if self.phase == 'seg':458 img, label = data459 if torch.rand(1) < self.prob:460 img = img[:, ::-1]461 label = label[:, ::-1]462 return img, label463 elif self.phase == 'cd':464 img1, label1, img2, label2 = data465 if torch.rand(1) < self.prob:466 img1 = img1[:, ::-1]467 label1 = label1[:, ::-1]468 img2 = img2[:, ::-1]469 label2 = label2[:, ::-1]470 return img1, label1, img2, label2471 elif self.phase == 'od':472 img, label = data473 if torch.rand(1) < self.prob:474 _, width, _ = img.shape475 img = img[:, ::-1]476 label[:, 1::2] = width - label[:, 3::-2]477 return img, label478 479 480class RandomFlipV(object):481 def __init__(self, phase, prob=0.5):482 self.phase = phase483 self.prob = prob484 485 def __call__(self, data):486 if self.phase == 'seg':487 img, label = data488 if torch.rand(1) < self.prob:489 img = img[::-1, :]490 label = label[::-1, :]491 return img, label492 elif self.phase == 'cd':493 img1, label1, img2, label2 = data494 if torch.rand(1) < self.prob:495 img1 = img1[::-1, :]496 label1 = label1[::-1, :]497 img2 = img2[::-1, :]498 label2 = label2[::-1, :]499 return img1, label1, img2, label2500 elif self.phase == 'od':501 img, label = data502 if torch.rand(1) < self.prob:503 height, _, _ = img.shape504 img = img[::-1, :]505 label[:, 2::2] = height - label[:, 4:1:-2]506 return img, label507 508 509class Resize(object):510 def __init__(self, phase, size):511 self.phase = phase512 self.size = size513 514 def __call__(self, data):515 if self.phase == 'seg':516 img, label = data517 img = cv2.resize(img, self.size, interpolation=cv2.INTER_LINEAR)518 # for label519 label = cv2.resize(label, self.size, interpolation=cv2.INTER_NEAREST)520 return img, label521 elif self.phase == 'cd':522 img1, label1, img2, label2 = data523 img1 = cv2.resize(img1, self.size, interpolation=cv2.INTER_LINEAR)524 img2 = cv2.resize(img2, self.size, interpolation=cv2.INTER_LINEAR)525 # for label526 label1 = cv2.resize(label1, self.size, interpolation=cv2.INTER_NEAREST)527 label2 = cv2.resize(label2, self.size, interpolation=cv2.INTER_NEAREST)528 return img1, label1, img2, label2529 elif self.phase == 'od':530 img, label = data531 height, width, _ = img.shape532 img = cv2.resize(img, self.size, interpolation=cv2.INTER_LINEAR)533 label[:, 1::2] = label[:, 1::2] / width * self.size[0]534 label[:, 2::2] = label[:, 2::2] / height * self.size[1]535 return img, label536 537 538class Normalize(object):539 def __init__(self, phase, prior_mean, prior_std):540 self.phase = phase541 self.prior_mean = np.array([[prior_mean]], dtype=np.float32)542 self.prior_std = np.array([[prior_std]], dtype=np.float32)543 544 def __call__(self, data):545 if self.phase in ['od', 'seg']:546 img, _ = data547 img = img / 255.548 img = (img - self.prior_mean) / (self.prior_std + 1e-10)549 550 return img, _551 elif self.phase == 'cd':552 img1, label1, img2, label2 = data553 img1 = img1 / 255.554 img1 = (img1 - self.prior_mean) / (self.prior_std + 1e-10)555 img2 = img2 / 255.556 img2 = (img2 - self.prior_mean) / (self.prior_std + 1e-10)557 558 return img1, label1, img2, label2559 560 561class InvNormalize(object):562 def __init__(self, prior_mean, prior_std):563 self.prior_mean = np.array([[prior_mean]], dtype=np.float32)564 self.prior_std = np.array([[prior_std]], dtype=np.float32)565 566 def __call__(self, img):567 img = img * self.prior_std + self.prior_mean568 img = img * 255.569 img = np.clip(img, a_min=0, a_max=255)570 return img571 572 573class Augmentations(object):574 def __init__(self, size, prior_mean=0, prior_std=1, pattern='train', phase='seg', *args, **kwargs):575 self.size = size576 self.prior_mean = prior_mean577 self.prior_std = prior_std578 self.phase = phase579 580 augments = {581 'train': Compose([582 ConvertUcharToFloat(),583 ImgDistortion(self.phase),584 ExpandImg(self.phase, self.prior_mean),585 RandomSampleCrop(self.phase, original_size=self.size),586 RandomMirror(self.phase),587 RandomFlipV(self.phase),588 Resize(self.phase, self.size),589 Normalize(self.phase, self.prior_mean, self.prior_std),590 ]),591 'val': Compose([592 ConvertUcharToFloat(),593 Resize(self.phase, self.size),594 Normalize(self.phase, self.prior_mean, self.prior_std),595 ]),596 'test': Compose([597 ConvertUcharToFloat(),598 Resize(self.phase, self.size),599 Normalize(self.phase, self.prior_mean, self.prior_std),600 ])601 }602 self.augment = augments[pattern]603 604 def __call__(self, data):605 return self.augment(data)606 607