jawahar-konathala/Tryon2
0
1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3 4"""5See "Data Augmentation" tutorial for an overview of the system:6https://detectron2.readthedocs.io/tutorials/augmentation.html7"""8 9import numpy as np10import torch11import torch.nn.functional as F12from fvcore.transforms.transform import (13 CropTransform,14 HFlipTransform,15 NoOpTransform,16 Transform,17 TransformList,18)19from PIL import Image20 21try:22 import cv2 # noqa23except ImportError:24 # OpenCV is an optional dependency at the moment25 pass26 27__all__ = [28 "ExtentTransform",29 "ResizeTransform",30 "RotationTransform",31 "ColorTransform",32 "PILColorTransform",33]34 35 36class ExtentTransform(Transform):37 """38 Extracts a subregion from the source image and scales it to the output size.39 40 The fill color is used to map pixels from the source rect that fall outside41 the source image.42 43 See: https://pillow.readthedocs.io/en/latest/PIL.html#PIL.ImageTransform.ExtentTransform44 """45 46 def __init__(self, src_rect, output_size, interp=Image.BILINEAR, fill=0):47 """48 Args:49 src_rect (x0, y0, x1, y1): src coordinates50 output_size (h, w): dst image size51 interp: PIL interpolation methods52 fill: Fill color used when src_rect extends outside image53 """54 super().__init__()55 self._set_attributes(locals())56 57 def apply_image(self, img, interp=None):58 h, w = self.output_size59 if len(img.shape) > 2 and img.shape[2] == 1:60 pil_image = Image.fromarray(img[:, :, 0], mode="L")61 else:62 pil_image = Image.fromarray(img)63 pil_image = pil_image.transform(64 size=(w, h),65 method=Image.EXTENT,66 data=self.src_rect,67 resample=interp if interp else self.interp,68 fill=self.fill,69 )70 ret = np.asarray(pil_image)71 if len(img.shape) > 2 and img.shape[2] == 1:72 ret = np.expand_dims(ret, -1)73 return ret74 75 def apply_coords(self, coords):76 # Transform image center from source coordinates into output coordinates77 # and then map the new origin to the corner of the output image.78 h, w = self.output_size79 x0, y0, x1, y1 = self.src_rect80 new_coords = coords.astype(np.float32)81 new_coords[:, 0] -= 0.5 * (x0 + x1)82 new_coords[:, 1] -= 0.5 * (y0 + y1)83 new_coords[:, 0] *= w / (x1 - x0)84 new_coords[:, 1] *= h / (y1 - y0)85 new_coords[:, 0] += 0.5 * w86 new_coords[:, 1] += 0.5 * h87 return new_coords88 89 def apply_segmentation(self, segmentation):90 segmentation = self.apply_image(segmentation, interp=Image.NEAREST)91 return segmentation92 93 94class ResizeTransform(Transform):95 """96 Resize the image to a target size.97 """98 99 def __init__(self, h, w, new_h, new_w, interp=None):100 """101 Args:102 h, w (int): original image size103 new_h, new_w (int): new image size104 interp: PIL interpolation methods, defaults to bilinear.105 """106 # TODO decide on PIL vs opencv107 super().__init__()108 if interp is None:109 interp = Image.BILINEAR110 self._set_attributes(locals())111 112 def apply_image(self, img, interp=None):113 assert img.shape[:2] == (self.h, self.w)114 assert len(img.shape) <= 4115 interp_method = interp if interp is not None else self.interp116 117 if img.dtype == np.uint8:118 if len(img.shape) > 2 and img.shape[2] == 1:119 pil_image = Image.fromarray(img[:, :, 0], mode="L")120 else:121 pil_image = Image.fromarray(img)122 pil_image = pil_image.resize((self.new_w, self.new_h), interp_method)123 ret = np.asarray(pil_image)124 if len(img.shape) > 2 and img.shape[2] == 1:125 ret = np.expand_dims(ret, -1)126 else:127 # PIL only supports uint8128 if any(x < 0 for x in img.strides):129 img = np.ascontiguousarray(img)130 img = torch.from_numpy(img)131 shape = list(img.shape)132 shape_4d = shape[:2] + [1] * (4 - len(shape)) + shape[2:]133 img = img.view(shape_4d).permute(2, 3, 0, 1) # hw(c) -> nchw134 _PIL_RESIZE_TO_INTERPOLATE_MODE = {135 Image.NEAREST: "nearest",136 Image.BILINEAR: "bilinear",137 Image.BICUBIC: "bicubic",138 }139 mode = _PIL_RESIZE_TO_INTERPOLATE_MODE[interp_method]140 align_corners = None if mode == "nearest" else False141 img = F.interpolate(142 img, (self.new_h, self.new_w), mode=mode, align_corners=align_corners143 )144 shape[:2] = (self.new_h, self.new_w)145 ret = img.permute(2, 3, 0, 1).view(shape).numpy() # nchw -> hw(c)146 147 return ret148 149 def apply_coords(self, coords):150 coords[:, 0] = coords[:, 0] * (self.new_w * 1.0 / self.w)151 coords[:, 1] = coords[:, 1] * (self.new_h * 1.0 / self.h)152 return coords153 154 def apply_segmentation(self, segmentation):155 segmentation = self.apply_image(segmentation, interp=Image.NEAREST)156 return segmentation157 158 def inverse(self):159 return ResizeTransform(self.new_h, self.new_w, self.h, self.w, self.interp)160 161 162class RotationTransform(Transform):163 """164 This method returns a copy of this image, rotated the given165 number of degrees counter clockwise around its center.166 """167 168 def __init__(self, h, w, angle, expand=True, center=None, interp=None):169 """170 Args:171 h, w (int): original image size172 angle (float): degrees for rotation173 expand (bool): choose if the image should be resized to fit the whole174 rotated image (default), or simply cropped175 center (tuple (width, height)): coordinates of the rotation center176 if left to None, the center will be fit to the center of each image177 center has no effect if expand=True because it only affects shifting178 interp: cv2 interpolation method, default cv2.INTER_LINEAR179 """180 super().__init__()181 image_center = np.array((w / 2, h / 2))182 if center is None:183 center = image_center184 if interp is None:185 interp = cv2.INTER_LINEAR186 abs_cos, abs_sin = (abs(np.cos(np.deg2rad(angle))), abs(np.sin(np.deg2rad(angle))))187 if expand:188 # find the new width and height bounds189 bound_w, bound_h = np.rint(190 [h * abs_sin + w * abs_cos, h * abs_cos + w * abs_sin]191 ).astype(int)192 else:193 bound_w, bound_h = w, h194 195 self._set_attributes(locals())196 self.rm_coords = self.create_rotation_matrix()197 # Needed because of this problem https://github.com/opencv/opencv/issues/11784198 self.rm_image = self.create_rotation_matrix(offset=-0.5)199 200 def apply_image(self, img, interp=None):201 """202 img should be a numpy array, formatted as Height * Width * Nchannels203 """204 if len(img) == 0 or self.angle % 360 == 0:205 return img206 assert img.shape[:2] == (self.h, self.w)207 interp = interp if interp is not None else self.interp208 return cv2.warpAffine(img, self.rm_image, (self.bound_w, self.bound_h), flags=interp)209 210 def apply_coords(self, coords):211 """212 coords should be a N * 2 array-like, containing N couples of (x, y) points213 """214 coords = np.asarray(coords, dtype=float)215 if len(coords) == 0 or self.angle % 360 == 0:216 return coords217 return cv2.transform(coords[:, np.newaxis, :], self.rm_coords)[:, 0, :]218 219 def apply_segmentation(self, segmentation):220 segmentation = self.apply_image(segmentation, interp=cv2.INTER_NEAREST)221 return segmentation222 223 def create_rotation_matrix(self, offset=0):224 center = (self.center[0] + offset, self.center[1] + offset)225 rm = cv2.getRotationMatrix2D(tuple(center), self.angle, 1)226 if self.expand:227 # Find the coordinates of the center of rotation in the new image228 # The only point for which we know the future coordinates is the center of the image229 rot_im_center = cv2.transform(self.image_center[None, None, :] + offset, rm)[0, 0, :]230 new_center = np.array([self.bound_w / 2, self.bound_h / 2]) + offset - rot_im_center231 # shift the rotation center to the new coordinates232 rm[:, 2] += new_center233 return rm234 235 def inverse(self):236 """237 The inverse is to rotate it back with expand, and crop to get the original shape.238 """239 if not self.expand: # Not possible to inverse if a part of the image is lost240 raise NotImplementedError()241 rotation = RotationTransform(242 self.bound_h, self.bound_w, -self.angle, True, None, self.interp243 )244 crop = CropTransform(245 (rotation.bound_w - self.w) // 2, (rotation.bound_h - self.h) // 2, self.w, self.h246 )247 return TransformList([rotation, crop])248 249 250class ColorTransform(Transform):251 """252 Generic wrapper for any photometric transforms.253 These transformations should only affect the color space and254 not the coordinate space of the image (e.g. annotation255 coordinates such as bounding boxes should not be changed)256 """257 258 def __init__(self, op):259 """260 Args:261 op (Callable): operation to be applied to the image,262 which takes in an ndarray and returns an ndarray.263 """264 if not callable(op):265 raise ValueError("op parameter should be callable")266 super().__init__()267 self._set_attributes(locals())268 269 def apply_image(self, img):270 return self.op(img)271 272 def apply_coords(self, coords):273 return coords274 275 def inverse(self):276 return NoOpTransform()277 278 def apply_segmentation(self, segmentation):279 return segmentation280 281 282class PILColorTransform(ColorTransform):283 """284 Generic wrapper for PIL Photometric image transforms,285 which affect the color space and not the coordinate286 space of the image287 """288 289 def __init__(self, op):290 """291 Args:292 op (Callable): operation to be applied to the image,293 which takes in a PIL Image and returns a transformed294 PIL Image.295 For reference on possible operations see:296 - https://pillow.readthedocs.io/en/stable/297 """298 if not callable(op):299 raise ValueError("op parameter should be callable")300 super().__init__(op)301 302 def apply_image(self, img):303 img = Image.fromarray(img)304 return np.asarray(super().apply_image(img))305 306 307def HFlip_rotated_box(transform, rotated_boxes):308 """309 Apply the horizontal flip transform on rotated boxes.310 311 Args:312 rotated_boxes (ndarray): Nx5 floating point array of313 (x_center, y_center, width, height, angle_degrees) format314 in absolute coordinates.315 """316 # Transform x_center317 rotated_boxes[:, 0] = transform.width - rotated_boxes[:, 0]318 # Transform angle319 rotated_boxes[:, 4] = -rotated_boxes[:, 4]320 return rotated_boxes321 322 323def Resize_rotated_box(transform, rotated_boxes):324 """325 Apply the resizing transform on rotated boxes. For details of how these (approximation)326 formulas are derived, please refer to :meth:`RotatedBoxes.scale`.327 328 Args:329 rotated_boxes (ndarray): Nx5 floating point array of330 (x_center, y_center, width, height, angle_degrees) format331 in absolute coordinates.332 """333 scale_factor_x = transform.new_w * 1.0 / transform.w334 scale_factor_y = transform.new_h * 1.0 / transform.h335 rotated_boxes[:, 0] *= scale_factor_x336 rotated_boxes[:, 1] *= scale_factor_y337 theta = rotated_boxes[:, 4] * np.pi / 180.0338 c = np.cos(theta)339 s = np.sin(theta)340 rotated_boxes[:, 2] *= np.sqrt(np.square(scale_factor_x * c) + np.square(scale_factor_y * s))341 rotated_boxes[:, 3] *= np.sqrt(np.square(scale_factor_x * s) + np.square(scale_factor_y * c))342 rotated_boxes[:, 4] = np.arctan2(scale_factor_x * s, scale_factor_y * c) * 180 / np.pi343 344 return rotated_boxes345 346 347HFlipTransform.register_type("rotated_box", HFlip_rotated_box)348ResizeTransform.register_type("rotated_box", Resize_rotated_box)349 350# not necessary any more with latest fvcore351NoOpTransform.register_type("rotated_box", lambda t, x: x)352