Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import math3import numpy as np4from enum import IntEnum, unique5from typing import List, Tuple, Union6import torch7from torch import device8 9_RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray]10 11 12@unique13class BoxMode(IntEnum):14 """15 Enum of different ways to represent a box.16 """17 18 XYXY_ABS = 019 """20 (x0, y0, x1, y1) in absolute floating points coordinates.21 The coordinates in range [0, width or height].22 """23 XYWH_ABS = 124 """25 (x0, y0, w, h) in absolute floating points coordinates.26 """27 XYXY_REL = 228 """29 Not yet supported!30 (x0, y0, x1, y1) in range [0, 1]. They are relative to the size of the image.31 """32 XYWH_REL = 333 """34 Not yet supported!35 (x0, y0, w, h) in range [0, 1]. They are relative to the size of the image.36 """37 XYWHA_ABS = 438 """39 (xc, yc, w, h, a) in absolute floating points coordinates.40 (xc, yc) is the center of the rotated box, and the angle a is in degrees ccw.41 """42 43 @staticmethod44 def convert(box: _RawBoxType, from_mode: "BoxMode", to_mode: "BoxMode") -> _RawBoxType:45 """46 Args:47 box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 548 from_mode, to_mode (BoxMode)49 50 Returns:51 The converted box of the same type.52 """53 if from_mode == to_mode:54 return box55 56 original_type = type(box)57 is_numpy = isinstance(box, np.ndarray)58 single_box = isinstance(box, (list, tuple))59 if single_box:60 assert len(box) == 4 or len(box) == 5, (61 "BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"62 " where k == 4 or 5"63 )64 arr = torch.tensor(box)[None, :]65 else:66 # avoid modifying the input box67 if is_numpy:68 arr = torch.from_numpy(np.asarray(box)).clone()69 else:70 arr = box.clone()71 72 assert to_mode not in [BoxMode.XYXY_REL, BoxMode.XYWH_REL] and from_mode not in [73 BoxMode.XYXY_REL,74 BoxMode.XYWH_REL,75 ], "Relative mode not yet supported!"76 77 if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:78 assert (79 arr.shape[-1] == 580 ), "The last dimension of input shape must be 5 for XYWHA format"81 original_dtype = arr.dtype82 arr = arr.double()83 84 w = arr[:, 2]85 h = arr[:, 3]86 a = arr[:, 4]87 c = torch.abs(torch.cos(a * math.pi / 180.0))88 s = torch.abs(torch.sin(a * math.pi / 180.0))89 # This basically computes the horizontal bounding rectangle of the rotated box90 new_w = c * w + s * h91 new_h = c * h + s * w92 93 # convert center to top-left corner94 arr[:, 0] -= new_w / 2.095 arr[:, 1] -= new_h / 2.096 # bottom-right corner97 arr[:, 2] = arr[:, 0] + new_w98 arr[:, 3] = arr[:, 1] + new_h99 100 arr = arr[:, :4].to(dtype=original_dtype)101 elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS:102 original_dtype = arr.dtype103 arr = arr.double()104 arr[:, 0] += arr[:, 2] / 2.0105 arr[:, 1] += arr[:, 3] / 2.0106 angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)107 arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)108 else:109 if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:110 arr[:, 2] += arr[:, 0]111 arr[:, 3] += arr[:, 1]112 elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS:113 arr[:, 2] -= arr[:, 0]114 arr[:, 3] -= arr[:, 1]115 else:116 raise NotImplementedError(117 "Conversion from BoxMode {} to {} is not supported yet".format(118 from_mode, to_mode119 )120 )121 122 if single_box:123 return original_type(arr.flatten().tolist())124 if is_numpy:125 return arr.numpy()126 else:127 return arr128 129 130class Boxes:131 """132 This structure stores a list of boxes as a Nx4 torch.Tensor.133 It supports some common methods about boxes134 (`area`, `clip`, `nonempty`, etc),135 and also behaves like a Tensor136 (support indexing, `to(device)`, `.device`, and iteration over all boxes)137 138 Attributes:139 tensor (torch.Tensor): float matrix of Nx4. Each row is (x1, y1, x2, y2).140 """141 142 def __init__(self, tensor: torch.Tensor):143 """144 Args:145 tensor (Tensor[float]): a Nx4 matrix. Each row is (x1, y1, x2, y2).146 """147 if not isinstance(tensor, torch.Tensor):148 tensor = torch.as_tensor(tensor, dtype=torch.float32, device=torch.device("cpu"))149 else:150 tensor = tensor.to(torch.float32)151 if tensor.numel() == 0:152 # Use reshape, so we don't end up creating a new tensor that does not depend on153 # the inputs (and consequently confuses jit)154 tensor = tensor.reshape((-1, 4)).to(dtype=torch.float32)155 assert tensor.dim() == 2 and tensor.size(-1) == 4, tensor.size()156 157 self.tensor = tensor158 159 def clone(self) -> "Boxes":160 """161 Clone the Boxes.162 163 Returns:164 Boxes165 """166 return Boxes(self.tensor.clone())167 168 def to(self, device: torch.device):169 # Boxes are assumed float32 and does not support to(dtype)170 return Boxes(self.tensor.to(device=device))171 172 def area(self) -> torch.Tensor:173 """174 Computes the area of all the boxes.175 176 Returns:177 torch.Tensor: a vector with areas of each box.178 """179 box = self.tensor180 area = (box[:, 2] - box[:, 0]) * (box[:, 3] - box[:, 1])181 return area182 183 def clip(self, box_size: Tuple[int, int]) -> None:184 """185 Clip (in place) the boxes by limiting x coordinates to the range [0, width]186 and y coordinates to the range [0, height].187 188 Args:189 box_size (height, width): The clipping box's size.190 """191 assert torch.isfinite(self.tensor).all(), "Box tensor contains infinite or NaN!"192 h, w = box_size193 x1 = self.tensor[:, 0].clamp(min=0, max=w)194 y1 = self.tensor[:, 1].clamp(min=0, max=h)195 x2 = self.tensor[:, 2].clamp(min=0, max=w)196 y2 = self.tensor[:, 3].clamp(min=0, max=h)197 self.tensor = torch.stack((x1, y1, x2, y2), dim=-1)198 199 def nonempty(self, threshold: float = 0.0) -> torch.Tensor:200 """201 Find boxes that are non-empty.202 A box is considered empty, if either of its side is no larger than threshold.203 204 Returns:205 Tensor:206 a binary vector which represents whether each box is empty207 (False) or non-empty (True).208 """209 box = self.tensor210 widths = box[:, 2] - box[:, 0]211 heights = box[:, 3] - box[:, 1]212 keep = (widths > threshold) & (heights > threshold)213 return keep214 215 def __getitem__(self, item) -> "Boxes":216 """217 Args:218 item: int, slice, or a BoolTensor219 220 Returns:221 Boxes: Create a new :class:`Boxes` by indexing.222 223 The following usage are allowed:224 225 1. `new_boxes = boxes[3]`: return a `Boxes` which contains only one box.226 2. `new_boxes = boxes[2:10]`: return a slice of boxes.227 3. `new_boxes = boxes[vector]`, where vector is a torch.BoolTensor228 with `length = len(boxes)`. Nonzero elements in the vector will be selected.229 230 Note that the returned Boxes might share storage with this Boxes,231 subject to Pytorch's indexing semantics.232 """233 if isinstance(item, int):234 return Boxes(self.tensor[item].view(1, -1))235 b = self.tensor[item]236 assert b.dim() == 2, "Indexing on Boxes with {} failed to return a matrix!".format(item)237 return Boxes(b)238 239 def __len__(self) -> int:240 return self.tensor.shape[0]241 242 def __repr__(self) -> str:243 return "Boxes(" + str(self.tensor) + ")"244 245 def inside_box(self, box_size: Tuple[int, int], boundary_threshold: int = 0) -> torch.Tensor:246 """247 Args:248 box_size (height, width): Size of the reference box.249 boundary_threshold (int): Boxes that extend beyond the reference box250 boundary by more than boundary_threshold are considered "outside".251 252 Returns:253 a binary vector, indicating whether each box is inside the reference box.254 """255 height, width = box_size256 inds_inside = (257 (self.tensor[..., 0] >= -boundary_threshold)258 & (self.tensor[..., 1] >= -boundary_threshold)259 & (self.tensor[..., 2] < width + boundary_threshold)260 & (self.tensor[..., 3] < height + boundary_threshold)261 )262 return inds_inside263 264 def get_centers(self) -> torch.Tensor:265 """266 Returns:267 The box centers in a Nx2 array of (x, y).268 """269 return (self.tensor[:, :2] + self.tensor[:, 2:]) / 2270 271 def scale(self, scale_x: float, scale_y: float) -> None:272 """273 Scale the box with horizontal and vertical scaling factors274 """275 self.tensor[:, 0::2] *= scale_x276 self.tensor[:, 1::2] *= scale_y277 278 @classmethod279 def cat(cls, boxes_list: List["Boxes"]) -> "Boxes":280 """281 Concatenates a list of Boxes into a single Boxes282 283 Arguments:284 boxes_list (list[Boxes])285 286 Returns:287 Boxes: the concatenated Boxes288 """289 assert isinstance(boxes_list, (list, tuple))290 if len(boxes_list) == 0:291 return cls(torch.empty(0))292 assert all([isinstance(box, Boxes) for box in boxes_list])293 294 # use torch.cat (v.s. layers.cat) so the returned boxes never share storage with input295 cat_boxes = cls(torch.cat([b.tensor for b in boxes_list], dim=0))296 return cat_boxes297 298 @property299 def device(self) -> device:300 return self.tensor.device301 302 # type "Iterator[torch.Tensor]", yield, and iter() not supported by torchscript303 # https://github.com/pytorch/pytorch/issues/18627304 @torch.jit.unused305 def __iter__(self):306 """307 Yield a box as a Tensor of shape (4,) at a time.308 """309 yield from self.tensor310 311 312def pairwise_intersection(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:313 """314 Given two lists of boxes of size N and M,315 compute the intersection area between __all__ N x M pairs of boxes.316 The box order must be (xmin, ymin, xmax, ymax)317 318 Args:319 boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.320 321 Returns:322 Tensor: intersection, sized [N,M].323 """324 boxes1, boxes2 = boxes1.tensor, boxes2.tensor325 width_height = torch.min(boxes1[:, None, 2:], boxes2[:, 2:]) - torch.max(326 boxes1[:, None, :2], boxes2[:, :2]327 ) # [N,M,2]328 329 width_height.clamp_(min=0) # [N,M,2]330 intersection = width_height.prod(dim=2) # [N,M]331 return intersection332 333 334# implementation from https://github.com/kuangliu/torchcv/blob/master/torchcv/utils/box.py335# with slight modifications336def pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:337 """338 Given two lists of boxes of size N and M, compute the IoU339 (intersection over union) between **all** N x M pairs of boxes.340 The box order must be (xmin, ymin, xmax, ymax).341 342 Args:343 boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.344 345 Returns:346 Tensor: IoU, sized [N,M].347 """348 area1 = boxes1.area() # [N]349 area2 = boxes2.area() # [M]350 inter = pairwise_intersection(boxes1, boxes2)351 352 # handle empty boxes353 iou = torch.where(354 inter > 0,355 inter / (area1[:, None] + area2 - inter),356 torch.zeros(1, dtype=inter.dtype, device=inter.device),357 )358 return iou359 360 361def pairwise_ioa(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:362 """363 Similar to :func:`pariwise_iou` but compute the IoA (intersection over boxes2 area).364 365 Args:366 boxes1,boxes2 (Boxes): two `Boxes`. Contains N & M boxes, respectively.367 368 Returns:369 Tensor: IoA, sized [N,M].370 """371 area2 = boxes2.area() # [M]372 inter = pairwise_intersection(boxes1, boxes2)373 374 # handle empty boxes375 ioa = torch.where(376 inter > 0, inter / area2, torch.zeros(1, dtype=inter.dtype, device=inter.device)377 )378 return ioa379 380 381def pairwise_point_box_distance(points: torch.Tensor, boxes: Boxes):382 """383 Pairwise distance between N points and M boxes. The distance between a384 point and a box is represented by the distance from the point to 4 edges385 of the box. Distances are all positive when the point is inside the box.386 387 Args:388 points: Nx2 coordinates. Each row is (x, y)389 boxes: M boxes390 391 Returns:392 Tensor: distances of size (N, M, 4). The 4 values are distances from393 the point to the left, top, right, bottom of the box.394 """395 x, y = points.unsqueeze(dim=2).unbind(dim=1) # (N, 1)396 x0, y0, x1, y1 = boxes.tensor.unsqueeze(dim=0).unbind(dim=2) # (1, M)397 return torch.stack([x - x0, y - y0, x1 - x, y1 - y], dim=2)398 399 400def matched_pairwise_iou(boxes1: Boxes, boxes2: Boxes) -> torch.Tensor:401 """402 Compute pairwise intersection over union (IOU) of two sets of matched403 boxes that have the same number of boxes.404 Similar to :func:`pairwise_iou`, but computes only diagonal elements of the matrix.405 406 Args:407 boxes1 (Boxes): bounding boxes, sized [N,4].408 boxes2 (Boxes): same length as boxes1409 Returns:410 Tensor: iou, sized [N].411 """412 assert len(boxes1) == len(413 boxes2414 ), "boxlists should have the same" "number of entries, got {}, {}".format(415 len(boxes1), len(boxes2)416 )417 area1 = boxes1.area() # [N]418 area2 = boxes2.area() # [N]419 box1, box2 = boxes1.tensor, boxes2.tensor420 lt = torch.max(box1[:, :2], box2[:, :2]) # [N,2]421 rb = torch.min(box1[:, 2:], box2[:, 2:]) # [N,2]422 wh = (rb - lt).clamp(min=0) # [N,2]423 inter = wh[:, 0] * wh[:, 1] # [N]424 iou = inter / (area1 + area2 - inter) # [N]425 return iou426