Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model
0
1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3"""4Implement many useful :class:`Augmentation`.5"""6import numpy as np7import sys8from numpy import random9from typing import Tuple10import torch11from fvcore.transforms.transform import (12 BlendTransform,13 CropTransform,14 HFlipTransform,15 NoOpTransform,16 PadTransform,17 Transform,18 TransformList,19 VFlipTransform,20)21from PIL import Image22 23from detectron2.structures import Boxes, pairwise_iou24 25from .augmentation import Augmentation, _transform_to_aug26from .transform import ExtentTransform, ResizeTransform, RotationTransform27 28__all__ = [29 "FixedSizeCrop",30 "RandomApply",31 "RandomBrightness",32 "RandomContrast",33 "RandomCrop",34 "RandomExtent",35 "RandomFlip",36 "RandomSaturation",37 "RandomLighting",38 "RandomRotation",39 "Resize",40 "ResizeScale",41 "ResizeShortestEdge",42 "RandomCrop_CategoryAreaConstraint",43 "RandomResize",44 "MinIoURandomCrop",45]46 47 48class RandomApply(Augmentation):49 """50 Randomly apply an augmentation with a given probability.51 """52 53 def __init__(self, tfm_or_aug, prob=0.5):54 """55 Args:56 tfm_or_aug (Transform, Augmentation): the transform or augmentation57 to be applied. It can either be a `Transform` or `Augmentation`58 instance.59 prob (float): probability between 0.0 and 1.0 that60 the wrapper transformation is applied61 """62 super().__init__()63 self.aug = _transform_to_aug(tfm_or_aug)64 assert 0.0 <= prob <= 1.0, f"Probablity must be between 0.0 and 1.0 (given: {prob})"65 self.prob = prob66 67 def get_transform(self, *args):68 do = self._rand_range() < self.prob69 if do:70 return self.aug.get_transform(*args)71 else:72 return NoOpTransform()73 74 def __call__(self, aug_input):75 do = self._rand_range() < self.prob76 if do:77 return self.aug(aug_input)78 else:79 return NoOpTransform()80 81 82class RandomFlip(Augmentation):83 """84 Flip the image horizontally or vertically with the given probability.85 """86 87 def __init__(self, prob=0.5, *, horizontal=True, vertical=False):88 """89 Args:90 prob (float): probability of flip.91 horizontal (boolean): whether to apply horizontal flipping92 vertical (boolean): whether to apply vertical flipping93 """94 super().__init__()95 96 if horizontal and vertical:97 raise ValueError("Cannot do both horiz and vert. Please use two Flip instead.")98 if not horizontal and not vertical:99 raise ValueError("At least one of horiz or vert has to be True!")100 self._init(locals())101 102 def get_transform(self, image):103 h, w = image.shape[:2]104 do = self._rand_range() < self.prob105 if do:106 if self.horizontal:107 return HFlipTransform(w)108 elif self.vertical:109 return VFlipTransform(h)110 else:111 return NoOpTransform()112 113 114class Resize(Augmentation):115 """Resize image to a fixed target size"""116 117 def __init__(self, shape, interp=Image.BILINEAR):118 """119 Args:120 shape: (h, w) tuple or a int121 interp: PIL interpolation method122 """123 if isinstance(shape, int):124 shape = (shape, shape)125 shape = tuple(shape)126 self._init(locals())127 128 def get_transform(self, image):129 return ResizeTransform(130 image.shape[0], image.shape[1], self.shape[0], self.shape[1], self.interp131 )132 133 134class ResizeShortestEdge(Augmentation):135 """136 Resize the image while keeping the aspect ratio unchanged.137 It attempts to scale the shorter edge to the given `short_edge_length`,138 as long as the longer edge does not exceed `max_size`.139 If `max_size` is reached, then downscale so that the longer edge does not exceed max_size.140 """141 142 @torch.jit.unused143 def __init__(144 self, short_edge_length, max_size=sys.maxsize, sample_style="range", interp=Image.BILINEAR145 ):146 """147 Args:148 short_edge_length (list[int]): If ``sample_style=="range"``,149 a [min, max] interval from which to sample the shortest edge length.150 If ``sample_style=="choice"``, a list of shortest edge lengths to sample from.151 max_size (int): maximum allowed longest edge length.152 sample_style (str): either "range" or "choice".153 """154 super().__init__()155 assert sample_style in ["range", "choice"], sample_style156 157 self.is_range = sample_style == "range"158 if isinstance(short_edge_length, int):159 short_edge_length = (short_edge_length, short_edge_length)160 if self.is_range:161 assert len(short_edge_length) == 2, (162 "short_edge_length must be two values using 'range' sample style."163 f" Got {short_edge_length}!"164 )165 self._init(locals())166 167 @torch.jit.unused168 def get_transform(self, image):169 h, w = image.shape[:2]170 if self.is_range:171 size = np.random.randint(self.short_edge_length[0], self.short_edge_length[1] + 1)172 else:173 size = np.random.choice(self.short_edge_length)174 if size == 0:175 return NoOpTransform()176 177 newh, neww = ResizeShortestEdge.get_output_shape(h, w, size, self.max_size)178 return ResizeTransform(h, w, newh, neww, self.interp)179 180 @staticmethod181 def get_output_shape(182 oldh: int, oldw: int, short_edge_length: int, max_size: int183 ) -> Tuple[int, int]:184 """185 Compute the output size given input size and target short edge length.186 """187 h, w = oldh, oldw188 size = short_edge_length * 1.0189 scale = size / min(h, w)190 if h < w:191 newh, neww = size, scale * w192 else:193 newh, neww = scale * h, size194 if max(newh, neww) > max_size:195 scale = max_size * 1.0 / max(newh, neww)196 newh = newh * scale197 neww = neww * scale198 neww = int(neww + 0.5)199 newh = int(newh + 0.5)200 return (newh, neww)201 202 203class ResizeScale(Augmentation):204 """205 Takes target size as input and randomly scales the given target size between `min_scale`206 and `max_scale`. It then scales the input image such that it fits inside the scaled target207 box, keeping the aspect ratio constant.208 This implements the resize part of the Google's 'resize_and_crop' data augmentation:209 https://github.com/tensorflow/tpu/blob/master/models/official/detection/utils/input_utils.py#L127210 """211 212 def __init__(213 self,214 min_scale: float,215 max_scale: float,216 target_height: int,217 target_width: int,218 interp: int = Image.BILINEAR,219 ):220 """221 Args:222 min_scale: minimum image scale range.223 max_scale: maximum image scale range.224 target_height: target image height.225 target_width: target image width.226 interp: image interpolation method.227 """228 super().__init__()229 self._init(locals())230 231 def _get_resize(self, image: np.ndarray, scale: float) -> Transform:232 input_size = image.shape[:2]233 234 # Compute new target size given a scale.235 target_size = (self.target_height, self.target_width)236 target_scale_size = np.multiply(target_size, scale)237 238 # Compute actual rescaling applied to input image and output size.239 output_scale = np.minimum(240 target_scale_size[0] / input_size[0], target_scale_size[1] / input_size[1]241 )242 output_size = np.round(np.multiply(input_size, output_scale)).astype(int)243 244 return ResizeTransform(245 input_size[0], input_size[1], output_size[0], output_size[1], self.interp246 )247 248 def get_transform(self, image: np.ndarray) -> Transform:249 random_scale = np.random.uniform(self.min_scale, self.max_scale)250 return self._get_resize(image, random_scale)251 252 253class RandomRotation(Augmentation):254 """255 This method returns a copy of this image, rotated the given256 number of degrees counter clockwise around the given center.257 """258 259 def __init__(self, angle, expand=True, center=None, sample_style="range", interp=None):260 """261 Args:262 angle (list[float]): If ``sample_style=="range"``,263 a [min, max] interval from which to sample the angle (in degrees).264 If ``sample_style=="choice"``, a list of angles to sample from265 expand (bool): choose if the image should be resized to fit the whole266 rotated image (default), or simply cropped267 center (list[[float, float]]): If ``sample_style=="range"``,268 a [[minx, miny], [maxx, maxy]] relative interval from which to sample the center,269 [0, 0] being the top left of the image and [1, 1] the bottom right.270 If ``sample_style=="choice"``, a list of centers to sample from271 Default: None, which means that the center of rotation is the center of the image272 center has no effect if expand=True because it only affects shifting273 """274 super().__init__()275 assert sample_style in ["range", "choice"], sample_style276 self.is_range = sample_style == "range"277 if isinstance(angle, (float, int)):278 angle = (angle, angle)279 if center is not None and isinstance(center[0], (float, int)):280 center = (center, center)281 self._init(locals())282 283 def get_transform(self, image):284 h, w = image.shape[:2]285 center = None286 if self.is_range:287 angle = np.random.uniform(self.angle[0], self.angle[1])288 if self.center is not None:289 center = (290 np.random.uniform(self.center[0][0], self.center[1][0]),291 np.random.uniform(self.center[0][1], self.center[1][1]),292 )293 else:294 angle = np.random.choice(self.angle)295 if self.center is not None:296 center = np.random.choice(self.center)297 298 if center is not None:299 center = (w * center[0], h * center[1]) # Convert to absolute coordinates300 301 if angle % 360 == 0:302 return NoOpTransform()303 304 return RotationTransform(h, w, angle, expand=self.expand, center=center, interp=self.interp)305 306 307class FixedSizeCrop(Augmentation):308 """309 If `crop_size` is smaller than the input image size, then it uses a random crop of310 the crop size. If `crop_size` is larger than the input image size, then it pads311 the right and the bottom of the image to the crop size if `pad` is True, otherwise312 it returns the smaller image.313 """314 315 def __init__(316 self,317 crop_size: Tuple[int],318 pad: bool = True,319 pad_value: float = 128.0,320 seg_pad_value: int = 255,321 ):322 """323 Args:324 crop_size: target image (height, width).325 pad: if True, will pad images smaller than `crop_size` up to `crop_size`326 pad_value: the padding value to the image.327 seg_pad_value: the padding value to the segmentation mask.328 """329 super().__init__()330 self._init(locals())331 332 def _get_crop(self, image: np.ndarray) -> Transform:333 # Compute the image scale and scaled size.334 input_size = image.shape[:2]335 output_size = self.crop_size336 337 # Add random crop if the image is scaled up.338 max_offset = np.subtract(input_size, output_size)339 max_offset = np.maximum(max_offset, 0)340 offset = np.multiply(max_offset, np.random.uniform(0.0, 1.0))341 offset = np.round(offset).astype(int)342 return CropTransform(343 offset[1], offset[0], output_size[1], output_size[0], input_size[1], input_size[0]344 )345 346 def _get_pad(self, image: np.ndarray) -> Transform:347 # Compute the image scale and scaled size.348 input_size = image.shape[:2]349 output_size = self.crop_size350 351 # Add padding if the image is scaled down.352 pad_size = np.subtract(output_size, input_size)353 pad_size = np.maximum(pad_size, 0)354 original_size = np.minimum(input_size, output_size)355 return PadTransform(356 0,357 0,358 pad_size[1],359 pad_size[0],360 original_size[1],361 original_size[0],362 self.pad_value,363 self.seg_pad_value,364 )365 366 def get_transform(self, image: np.ndarray) -> TransformList:367 transforms = [self._get_crop(image)]368 if self.pad:369 transforms.append(self._get_pad(image))370 return TransformList(transforms)371 372 373class RandomCrop(Augmentation):374 """375 Randomly crop a rectangle region out of an image.376 """377 378 def __init__(self, crop_type: str, crop_size):379 """380 Args:381 crop_type (str): one of "relative_range", "relative", "absolute", "absolute_range".382 crop_size (tuple[float, float]): two floats, explained below.383 384 - "relative": crop a (H * crop_size[0], W * crop_size[1]) region from an input image of385 size (H, W). crop size should be in (0, 1]386 - "relative_range": uniformly sample two values from [crop_size[0], 1]387 and [crop_size[1]], 1], and use them as in "relative" crop type.388 - "absolute" crop a (crop_size[0], crop_size[1]) region from input image.389 crop_size must be smaller than the input image size.390 - "absolute_range", for an input of size (H, W), uniformly sample H_crop in391 [crop_size[0], min(H, crop_size[1])] and W_crop in [crop_size[0], min(W, crop_size[1])].392 Then crop a region (H_crop, W_crop).393 """394 # TODO style of relative_range and absolute_range are not consistent:395 # one takes (h, w) but another takes (min, max)396 super().__init__()397 assert crop_type in ["relative_range", "relative", "absolute", "absolute_range"]398 self._init(locals())399 400 def get_transform(self, image):401 h, w = image.shape[:2]402 croph, cropw = self.get_crop_size((h, w))403 assert h >= croph and w >= cropw, "Shape computation in {} has bugs.".format(self)404 h0 = np.random.randint(h - croph + 1)405 w0 = np.random.randint(w - cropw + 1)406 return CropTransform(w0, h0, cropw, croph)407 408 def get_crop_size(self, image_size):409 """410 Args:411 image_size (tuple): height, width412 413 Returns:414 crop_size (tuple): height, width in absolute pixels415 """416 h, w = image_size417 if self.crop_type == "relative":418 ch, cw = self.crop_size419 return int(h * ch + 0.5), int(w * cw + 0.5)420 elif self.crop_type == "relative_range":421 crop_size = np.asarray(self.crop_size, dtype=np.float32)422 ch, cw = crop_size + np.random.rand(2) * (1 - crop_size)423 return int(h * ch + 0.5), int(w * cw + 0.5)424 elif self.crop_type == "absolute":425 return (min(self.crop_size[0], h), min(self.crop_size[1], w))426 elif self.crop_type == "absolute_range":427 assert self.crop_size[0] <= self.crop_size[1]428 ch = np.random.randint(min(h, self.crop_size[0]), min(h, self.crop_size[1]) + 1)429 cw = np.random.randint(min(w, self.crop_size[0]), min(w, self.crop_size[1]) + 1)430 return ch, cw431 else:432 raise NotImplementedError("Unknown crop type {}".format(self.crop_type))433 434 435class RandomCrop_CategoryAreaConstraint(Augmentation):436 """437 Similar to :class:`RandomCrop`, but find a cropping window such that no single category438 occupies a ratio of more than `single_category_max_area` in semantic segmentation ground439 truth, which can cause unstability in training. The function attempts to find such a valid440 cropping window for at most 10 times.441 """442 443 def __init__(444 self,445 crop_type: str,446 crop_size,447 single_category_max_area: float = 1.0,448 ignored_category: int = None,449 ):450 """451 Args:452 crop_type, crop_size: same as in :class:`RandomCrop`453 single_category_max_area: the maximum allowed area ratio of a454 category. Set to 1.0 to disable455 ignored_category: allow this category in the semantic segmentation456 ground truth to exceed the area ratio. Usually set to the category457 that's ignored in training.458 """459 self.crop_aug = RandomCrop(crop_type, crop_size)460 self._init(locals())461 462 def get_transform(self, image, sem_seg):463 if self.single_category_max_area >= 1.0:464 return self.crop_aug.get_transform(image)465 else:466 h, w = sem_seg.shape467 for _ in range(10):468 crop_size = self.crop_aug.get_crop_size((h, w))469 y0 = np.random.randint(h - crop_size[0] + 1)470 x0 = np.random.randint(w - crop_size[1] + 1)471 sem_seg_temp = sem_seg[y0 : y0 + crop_size[0], x0 : x0 + crop_size[1]]472 labels, cnt = np.unique(sem_seg_temp, return_counts=True)473 if self.ignored_category is not None:474 cnt = cnt[labels != self.ignored_category]475 if len(cnt) > 1 and np.max(cnt) < np.sum(cnt) * self.single_category_max_area:476 break477 crop_tfm = CropTransform(x0, y0, crop_size[1], crop_size[0])478 return crop_tfm479 480 481class RandomExtent(Augmentation):482 """483 Outputs an image by cropping a random "subrect" of the source image.484 485 The subrect can be parameterized to include pixels outside the source image,486 in which case they will be set to zeros (i.e. black). The size of the output487 image will vary with the size of the random subrect.488 """489 490 def __init__(self, scale_range, shift_range):491 """492 Args:493 output_size (h, w): Dimensions of output image494 scale_range (l, h): Range of input-to-output size scaling factor495 shift_range (x, y): Range of shifts of the cropped subrect. The rect496 is shifted by [w / 2 * Uniform(-x, x), h / 2 * Uniform(-y, y)],497 where (w, h) is the (width, height) of the input image. Set each498 component to zero to crop at the image's center.499 """500 super().__init__()501 self._init(locals())502 503 def get_transform(self, image):504 img_h, img_w = image.shape[:2]505 506 # Initialize src_rect to fit the input image.507 src_rect = np.array([-0.5 * img_w, -0.5 * img_h, 0.5 * img_w, 0.5 * img_h])508 509 # Apply a random scaling to the src_rect.510 src_rect *= np.random.uniform(self.scale_range[0], self.scale_range[1])511 512 # Apply a random shift to the coordinates origin.513 src_rect[0::2] += self.shift_range[0] * img_w * (np.random.rand() - 0.5)514 src_rect[1::2] += self.shift_range[1] * img_h * (np.random.rand() - 0.5)515 516 # Map src_rect coordinates into image coordinates (center at corner).517 src_rect[0::2] += 0.5 * img_w518 src_rect[1::2] += 0.5 * img_h519 520 return ExtentTransform(521 src_rect=(src_rect[0], src_rect[1], src_rect[2], src_rect[3]),522 output_size=(int(src_rect[3] - src_rect[1]), int(src_rect[2] - src_rect[0])),523 )524 525 526class RandomContrast(Augmentation):527 """528 Randomly transforms image contrast.529 530 Contrast intensity is uniformly sampled in (intensity_min, intensity_max).531 - intensity < 1 will reduce contrast532 - intensity = 1 will preserve the input image533 - intensity > 1 will increase contrast534 535 See: https://pillow.readthedocs.io/en/3.0.x/reference/ImageEnhance.html536 """537 538 def __init__(self, intensity_min, intensity_max):539 """540 Args:541 intensity_min (float): Minimum augmentation542 intensity_max (float): Maximum augmentation543 """544 super().__init__()545 self._init(locals())546 547 def get_transform(self, image):548 w = np.random.uniform(self.intensity_min, self.intensity_max)549 return BlendTransform(src_image=image.mean(), src_weight=1 - w, dst_weight=w)550 551 552class RandomBrightness(Augmentation):553 """554 Randomly transforms image brightness.555 556 Brightness intensity is uniformly sampled in (intensity_min, intensity_max).557 - intensity < 1 will reduce brightness558 - intensity = 1 will preserve the input image559 - intensity > 1 will increase brightness560 561 See: https://pillow.readthedocs.io/en/3.0.x/reference/ImageEnhance.html562 """563 564 def __init__(self, intensity_min, intensity_max):565 """566 Args:567 intensity_min (float): Minimum augmentation568 intensity_max (float): Maximum augmentation569 """570 super().__init__()571 self._init(locals())572 573 def get_transform(self, image):574 w = np.random.uniform(self.intensity_min, self.intensity_max)575 return BlendTransform(src_image=0, src_weight=1 - w, dst_weight=w)576 577 578class RandomSaturation(Augmentation):579 """580 Randomly transforms saturation of an RGB image.581 Input images are assumed to have 'RGB' channel order.582 583 Saturation intensity is uniformly sampled in (intensity_min, intensity_max).584 - intensity < 1 will reduce saturation (make the image more grayscale)585 - intensity = 1 will preserve the input image586 - intensity > 1 will increase saturation587 588 See: https://pillow.readthedocs.io/en/3.0.x/reference/ImageEnhance.html589 """590 591 def __init__(self, intensity_min, intensity_max):592 """593 Args:594 intensity_min (float): Minimum augmentation (1 preserves input).595 intensity_max (float): Maximum augmentation (1 preserves input).596 """597 super().__init__()598 self._init(locals())599 600 def get_transform(self, image):601 assert image.shape[-1] == 3, "RandomSaturation only works on RGB images"602 w = np.random.uniform(self.intensity_min, self.intensity_max)603 grayscale = image.dot([0.299, 0.587, 0.114])[:, :, np.newaxis]604 return BlendTransform(src_image=grayscale, src_weight=1 - w, dst_weight=w)605 606 607class RandomLighting(Augmentation):608 """609 The "lighting" augmentation described in AlexNet, using fixed PCA over ImageNet.610 Input images are assumed to have 'RGB' channel order.611 612 The degree of color jittering is randomly sampled via a normal distribution,613 with standard deviation given by the scale parameter.614 """615 616 def __init__(self, scale):617 """618 Args:619 scale (float): Standard deviation of principal component weighting.620 """621 super().__init__()622 self._init(locals())623 self.eigen_vecs = np.array(624 [[-0.5675, 0.7192, 0.4009], [-0.5808, -0.0045, -0.8140], [-0.5836, -0.6948, 0.4203]]625 )626 self.eigen_vals = np.array([0.2175, 0.0188, 0.0045])627 628 def get_transform(self, image):629 assert image.shape[-1] == 3, "RandomLighting only works on RGB images"630 weights = np.random.normal(scale=self.scale, size=3)631 return BlendTransform(632 src_image=self.eigen_vecs.dot(weights * self.eigen_vals), src_weight=1.0, dst_weight=1.0633 )634 635 636class RandomResize(Augmentation):637 """Randomly resize image to a target size in shape_list"""638 639 def __init__(self, shape_list, interp=Image.BILINEAR):640 """641 Args:642 shape_list: a list of shapes in (h, w)643 interp: PIL interpolation method644 """645 self.shape_list = shape_list646 self._init(locals())647 648 def get_transform(self, image):649 shape_idx = np.random.randint(low=0, high=len(self.shape_list))650 h, w = self.shape_list[shape_idx]651 return ResizeTransform(image.shape[0], image.shape[1], h, w, self.interp)652 653 654class MinIoURandomCrop(Augmentation):655 """Random crop the image & bboxes, the cropped patches have minimum IoU656 requirement with original image & bboxes, the IoU threshold is randomly657 selected from min_ious.658 659 Args:660 min_ious (tuple): minimum IoU threshold for all intersections with661 bounding boxes662 min_crop_size (float): minimum crop's size (i.e. h,w := a*h, a*w,663 where a >= min_crop_size)664 mode_trials: number of trials for sampling min_ious threshold665 crop_trials: number of trials for sampling crop_size after cropping666 """667 668 def __init__(669 self,670 min_ious=(0.1, 0.3, 0.5, 0.7, 0.9),671 min_crop_size=0.3,672 mode_trials=1000,673 crop_trials=50,674 ):675 self.min_ious = min_ious676 self.sample_mode = (1, *min_ious, 0)677 self.min_crop_size = min_crop_size678 self.mode_trials = mode_trials679 self.crop_trials = crop_trials680 681 def get_transform(self, image, boxes):682 """Call function to crop images and bounding boxes with minimum IoU683 constraint.684 685 Args:686 boxes: ground truth boxes in (x1, y1, x2, y2) format687 """688 if boxes is None:689 return NoOpTransform()690 h, w, c = image.shape691 for _ in range(self.mode_trials):692 mode = random.choice(self.sample_mode)693 self.mode = mode694 if mode == 1:695 return NoOpTransform()696 697 min_iou = mode698 for _ in range(self.crop_trials):699 new_w = random.uniform(self.min_crop_size * w, w)700 new_h = random.uniform(self.min_crop_size * h, h)701 702 # h / w in [0.5, 2]703 if new_h / new_w < 0.5 or new_h / new_w > 2:704 continue705 706 left = random.uniform(w - new_w)707 top = random.uniform(h - new_h)708 709 patch = np.array((int(left), int(top), int(left + new_w), int(top + new_h)))710 # Line or point crop is not allowed711 if patch[2] == patch[0] or patch[3] == patch[1]:712 continue713 overlaps = pairwise_iou(714 Boxes(patch.reshape(-1, 4)), Boxes(boxes.reshape(-1, 4))715 ).reshape(-1)716 if len(overlaps) > 0 and overlaps.min() < min_iou:717 continue718 719 # center of boxes should inside the crop img720 # only adjust boxes and instance masks when the gt is not empty721 if len(overlaps) > 0:722 # adjust boxes723 def is_center_of_bboxes_in_patch(boxes, patch):724 center = (boxes[:, :2] + boxes[:, 2:]) / 2725 mask = (726 (center[:, 0] > patch[0])727 * (center[:, 1] > patch[1])728 * (center[:, 0] < patch[2])729 * (center[:, 1] < patch[3])730 )731 return mask732 733 mask = is_center_of_bboxes_in_patch(boxes, patch)734 if not mask.any():735 continue736 return CropTransform(int(left), int(top), int(new_w), int(new_h))737