Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import copy3import itertools4import numpy as np5from typing import Any, Iterator, List, Union6import pycocotools.mask as mask_util7import torch8from torch import device9 10from detectron2.layers.roi_align import ROIAlign11from detectron2.utils.memory import retry_if_cuda_oom12 13from .boxes import Boxes14 15 16def polygon_area(x, y):17 # Using the shoelace formula18 # https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates19 return 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))20 21 22def polygons_to_bitmask(polygons: List[np.ndarray], height: int, width: int) -> np.ndarray:23 """24 Args:25 polygons (list[ndarray]): each array has shape (Nx2,)26 height, width (int)27 28 Returns:29 ndarray: a bool mask of shape (height, width)30 """31 if len(polygons) == 0:32 # COCOAPI does not support empty polygons33 return np.zeros((height, width)).astype(bool)34 rles = mask_util.frPyObjects(polygons, height, width)35 rle = mask_util.merge(rles)36 return mask_util.decode(rle).astype(bool)37 38 39def rasterize_polygons_within_box(40 polygons: List[np.ndarray], box: np.ndarray, mask_size: int41) -> torch.Tensor:42 """43 Rasterize the polygons into a mask image and44 crop the mask content in the given box.45 The cropped mask is resized to (mask_size, mask_size).46 47 This function is used when generating training targets for mask head in Mask R-CNN.48 Given original ground-truth masks for an image, new ground-truth mask49 training targets in the size of `mask_size x mask_size`50 must be provided for each predicted box. This function will be called to51 produce such targets.52 53 Args:54 polygons (list[ndarray[float]]): a list of polygons, which represents an instance.55 box: 4-element numpy array56 mask_size (int):57 58 Returns:59 Tensor: BoolTensor of shape (mask_size, mask_size)60 """61 # 1. Shift the polygons w.r.t the boxes62 w, h = box[2] - box[0], box[3] - box[1]63 64 polygons = copy.deepcopy(polygons)65 for p in polygons:66 p[0::2] = p[0::2] - box[0]67 p[1::2] = p[1::2] - box[1]68 69 # 2. Rescale the polygons to the new box size70 # max() to avoid division by small number71 ratio_h = mask_size / max(h, 0.1)72 ratio_w = mask_size / max(w, 0.1)73 74 if ratio_h == ratio_w:75 for p in polygons:76 p *= ratio_h77 else:78 for p in polygons:79 p[0::2] *= ratio_w80 p[1::2] *= ratio_h81 82 # 3. Rasterize the polygons with coco api83 mask = polygons_to_bitmask(polygons, mask_size, mask_size)84 mask = torch.from_numpy(mask)85 return mask86 87 88class BitMasks:89 """90 This class stores the segmentation masks for all objects in one image, in91 the form of bitmaps.92 93 Attributes:94 tensor: bool Tensor of N,H,W, representing N instances in the image.95 """96 97 def __init__(self, tensor: Union[torch.Tensor, np.ndarray]):98 """99 Args:100 tensor: bool Tensor of N,H,W, representing N instances in the image.101 """102 if isinstance(tensor, torch.Tensor):103 tensor = tensor.to(torch.bool)104 else:105 tensor = torch.as_tensor(tensor, dtype=torch.bool, device=torch.device("cpu"))106 assert tensor.dim() == 3, tensor.size()107 self.image_size = tensor.shape[1:]108 self.tensor = tensor109 110 @torch.jit.unused111 def to(self, *args: Any, **kwargs: Any) -> "BitMasks":112 return BitMasks(self.tensor.to(*args, **kwargs))113 114 @property115 def device(self) -> torch.device:116 return self.tensor.device117 118 @torch.jit.unused119 def __getitem__(self, item: Union[int, slice, torch.BoolTensor]) -> "BitMasks":120 """121 Returns:122 BitMasks: Create a new :class:`BitMasks` by indexing.123 124 The following usage are allowed:125 126 1. `new_masks = masks[3]`: return a `BitMasks` which contains only one mask.127 2. `new_masks = masks[2:10]`: return a slice of masks.128 3. `new_masks = masks[vector]`, where vector is a torch.BoolTensor129 with `length = len(masks)`. Nonzero elements in the vector will be selected.130 131 Note that the returned object might share storage with this object,132 subject to Pytorch's indexing semantics.133 """134 if isinstance(item, int):135 return BitMasks(self.tensor[item].unsqueeze(0))136 m = self.tensor[item]137 assert m.dim() == 3, "Indexing on BitMasks with {} returns a tensor with shape {}!".format(138 item, m.shape139 )140 return BitMasks(m)141 142 @torch.jit.unused143 def __iter__(self) -> torch.Tensor:144 yield from self.tensor145 146 @torch.jit.unused147 def __repr__(self) -> str:148 s = self.__class__.__name__ + "("149 s += "num_instances={})".format(len(self.tensor))150 return s151 152 def __len__(self) -> int:153 return self.tensor.shape[0]154 155 def nonempty(self) -> torch.Tensor:156 """157 Find masks that are non-empty.158 159 Returns:160 Tensor: a BoolTensor which represents161 whether each mask is empty (False) or non-empty (True).162 """163 return self.tensor.flatten(1).any(dim=1)164 165 @staticmethod166 def from_polygon_masks(167 polygon_masks: Union["PolygonMasks", List[List[np.ndarray]]], height: int, width: int168 ) -> "BitMasks":169 """170 Args:171 polygon_masks (list[list[ndarray]] or PolygonMasks)172 height, width (int)173 """174 if isinstance(polygon_masks, PolygonMasks):175 polygon_masks = polygon_masks.polygons176 masks = [polygons_to_bitmask(p, height, width) for p in polygon_masks]177 if len(masks):178 return BitMasks(torch.stack([torch.from_numpy(x) for x in masks]))179 else:180 return BitMasks(torch.empty(0, height, width, dtype=torch.bool))181 182 @staticmethod183 def from_roi_masks(roi_masks: "ROIMasks", height: int, width: int) -> "BitMasks":184 """185 Args:186 roi_masks:187 height, width (int):188 """189 return roi_masks.to_bitmasks(height, width)190 191 def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:192 """193 Crop each bitmask by the given box, and resize results to (mask_size, mask_size).194 This can be used to prepare training targets for Mask R-CNN.195 It has less reconstruction error compared to rasterization with polygons.196 However we observe no difference in accuracy,197 but BitMasks requires more memory to store all the masks.198 199 Args:200 boxes (Tensor): Nx4 tensor storing the boxes for each mask201 mask_size (int): the size of the rasterized mask.202 203 Returns:204 Tensor:205 A bool tensor of shape (N, mask_size, mask_size), where206 N is the number of predicted boxes for this image.207 """208 assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))209 device = self.tensor.device210 211 batch_inds = torch.arange(len(boxes), device=device).to(dtype=boxes.dtype)[:, None]212 rois = torch.cat([batch_inds, boxes], dim=1) # Nx5213 214 bit_masks = self.tensor.to(dtype=torch.float32)215 rois = rois.to(device=device)216 output = (217 ROIAlign((mask_size, mask_size), 1.0, 0, aligned=True)218 .forward(bit_masks[:, None, :, :], rois)219 .squeeze(1)220 )221 output = output >= 0.5222 return output223 224 def get_bounding_boxes(self) -> Boxes:225 """226 Returns:227 Boxes: tight bounding boxes around bitmasks.228 If a mask is empty, it's bounding box will be all zero.229 """230 boxes = torch.zeros(self.tensor.shape[0], 4, dtype=torch.float32)231 x_any = torch.any(self.tensor, dim=1)232 y_any = torch.any(self.tensor, dim=2)233 for idx in range(self.tensor.shape[0]):234 x = torch.where(x_any[idx, :])[0]235 y = torch.where(y_any[idx, :])[0]236 if len(x) > 0 and len(y) > 0:237 boxes[idx, :] = torch.as_tensor(238 [x[0], y[0], x[-1] + 1, y[-1] + 1], dtype=torch.float32239 )240 return Boxes(boxes)241 242 @staticmethod243 def cat(bitmasks_list: List["BitMasks"]) -> "BitMasks":244 """245 Concatenates a list of BitMasks into a single BitMasks246 247 Arguments:248 bitmasks_list (list[BitMasks])249 250 Returns:251 BitMasks: the concatenated BitMasks252 """253 assert isinstance(bitmasks_list, (list, tuple))254 assert len(bitmasks_list) > 0255 assert all(isinstance(bitmask, BitMasks) for bitmask in bitmasks_list)256 257 cat_bitmasks = type(bitmasks_list[0])(torch.cat([bm.tensor for bm in bitmasks_list], dim=0))258 return cat_bitmasks259 260 261class PolygonMasks:262 """263 This class stores the segmentation masks for all objects in one image, in the form of polygons.264 265 Attributes:266 polygons: list[list[ndarray]]. Each ndarray is a float64 vector representing a polygon.267 """268 269 def __init__(self, polygons: List[List[Union[torch.Tensor, np.ndarray]]]):270 """271 Arguments:272 polygons (list[list[np.ndarray]]): The first273 level of the list correspond to individual instances,274 the second level to all the polygons that compose the275 instance, and the third level to the polygon coordinates.276 The third level array should have the format of277 [x0, y0, x1, y1, ..., xn, yn] (n >= 3).278 """279 if not isinstance(polygons, list):280 raise ValueError(281 "Cannot create PolygonMasks: Expect a list of list of polygons per image. "282 "Got '{}' instead.".format(type(polygons))283 )284 285 def _make_array(t: Union[torch.Tensor, np.ndarray]) -> np.ndarray:286 # Use float64 for higher precision, because why not?287 # Always put polygons on CPU (self.to is a no-op) since they288 # are supposed to be small tensors.289 # May need to change this assumption if GPU placement becomes useful290 if isinstance(t, torch.Tensor):291 t = t.cpu().numpy()292 return np.asarray(t).astype("float64")293 294 def process_polygons(295 polygons_per_instance: List[Union[torch.Tensor, np.ndarray]]296 ) -> List[np.ndarray]:297 if not isinstance(polygons_per_instance, list):298 raise ValueError(299 "Cannot create polygons: Expect a list of polygons per instance. "300 "Got '{}' instead.".format(type(polygons_per_instance))301 )302 # transform each polygon to a numpy array303 polygons_per_instance = [_make_array(p) for p in polygons_per_instance]304 for polygon in polygons_per_instance:305 if len(polygon) % 2 != 0 or len(polygon) < 6:306 raise ValueError(f"Cannot create a polygon from {len(polygon)} coordinates.")307 return polygons_per_instance308 309 self.polygons: List[List[np.ndarray]] = [310 process_polygons(polygons_per_instance) for polygons_per_instance in polygons311 ]312 313 def to(self, *args: Any, **kwargs: Any) -> "PolygonMasks":314 return self315 316 @property317 def device(self) -> torch.device:318 return torch.device("cpu")319 320 def get_bounding_boxes(self) -> Boxes:321 """322 Returns:323 Boxes: tight bounding boxes around polygon masks.324 """325 boxes = torch.zeros(len(self.polygons), 4, dtype=torch.float32)326 for idx, polygons_per_instance in enumerate(self.polygons):327 minxy = torch.as_tensor([float("inf"), float("inf")], dtype=torch.float32)328 maxxy = torch.zeros(2, dtype=torch.float32)329 for polygon in polygons_per_instance:330 coords = torch.from_numpy(polygon).view(-1, 2).to(dtype=torch.float32)331 minxy = torch.min(minxy, torch.min(coords, dim=0).values)332 maxxy = torch.max(maxxy, torch.max(coords, dim=0).values)333 boxes[idx, :2] = minxy334 boxes[idx, 2:] = maxxy335 return Boxes(boxes)336 337 def nonempty(self) -> torch.Tensor:338 """339 Find masks that are non-empty.340 341 Returns:342 Tensor:343 a BoolTensor which represents whether each mask is empty (False) or not (True).344 """345 keep = [1 if len(polygon) > 0 else 0 for polygon in self.polygons]346 return torch.from_numpy(np.asarray(keep, dtype=bool))347 348 def __getitem__(self, item: Union[int, slice, List[int], torch.BoolTensor]) -> "PolygonMasks":349 """350 Support indexing over the instances and return a `PolygonMasks` object.351 `item` can be:352 353 1. An integer. It will return an object with only one instance.354 2. A slice. It will return an object with the selected instances.355 3. A list[int]. It will return an object with the selected instances,356 correpsonding to the indices in the list.357 4. A vector mask of type BoolTensor, whose length is num_instances.358 It will return an object with the instances whose mask is nonzero.359 """360 if isinstance(item, int):361 selected_polygons = [self.polygons[item]]362 elif isinstance(item, slice):363 selected_polygons = self.polygons[item]364 elif isinstance(item, list):365 selected_polygons = [self.polygons[i] for i in item]366 elif isinstance(item, torch.Tensor):367 # Polygons is a list, so we have to move the indices back to CPU.368 if item.dtype == torch.bool:369 assert item.dim() == 1, item.shape370 item = item.nonzero().squeeze(1).cpu().numpy().tolist()371 elif item.dtype in [torch.int32, torch.int64]:372 item = item.cpu().numpy().tolist()373 else:374 raise ValueError("Unsupported tensor dtype={} for indexing!".format(item.dtype))375 selected_polygons = [self.polygons[i] for i in item]376 return PolygonMasks(selected_polygons)377 378 def __iter__(self) -> Iterator[List[np.ndarray]]:379 """380 Yields:381 list[ndarray]: the polygons for one instance.382 Each Tensor is a float64 vector representing a polygon.383 """384 return iter(self.polygons)385 386 def __repr__(self) -> str:387 s = self.__class__.__name__ + "("388 s += "num_instances={})".format(len(self.polygons))389 return s390 391 def __len__(self) -> int:392 return len(self.polygons)393 394 def crop_and_resize(self, boxes: torch.Tensor, mask_size: int) -> torch.Tensor:395 """396 Crop each mask by the given box, and resize results to (mask_size, mask_size).397 This can be used to prepare training targets for Mask R-CNN.398 399 Args:400 boxes (Tensor): Nx4 tensor storing the boxes for each mask401 mask_size (int): the size of the rasterized mask.402 403 Returns:404 Tensor: A bool tensor of shape (N, mask_size, mask_size), where405 N is the number of predicted boxes for this image.406 """407 assert len(boxes) == len(self), "{} != {}".format(len(boxes), len(self))408 409 device = boxes.device410 # Put boxes on the CPU, as the polygon representation is not efficient GPU-wise411 # (several small tensors for representing a single instance mask)412 boxes = boxes.to(torch.device("cpu"))413 414 results = [415 rasterize_polygons_within_box(poly, box.numpy(), mask_size)416 for poly, box in zip(self.polygons, boxes)417 ]418 """419 poly: list[list[float]], the polygons for one instance420 box: a tensor of shape (4,)421 """422 if len(results) == 0:423 return torch.empty(0, mask_size, mask_size, dtype=torch.bool, device=device)424 return torch.stack(results, dim=0).to(device=device)425 426 def area(self):427 """428 Computes area of the mask.429 Only works with Polygons, using the shoelace formula:430 https://stackoverflow.com/questions/24467972/calculate-area-of-polygon-given-x-y-coordinates431 432 Returns:433 Tensor: a vector, area for each instance434 """435 436 area = []437 for polygons_per_instance in self.polygons:438 area_per_instance = 0439 for p in polygons_per_instance:440 area_per_instance += polygon_area(p[0::2], p[1::2])441 area.append(area_per_instance)442 443 return torch.tensor(area)444 445 @staticmethod446 def cat(polymasks_list: List["PolygonMasks"]) -> "PolygonMasks":447 """448 Concatenates a list of PolygonMasks into a single PolygonMasks449 450 Arguments:451 polymasks_list (list[PolygonMasks])452 453 Returns:454 PolygonMasks: the concatenated PolygonMasks455 """456 assert isinstance(polymasks_list, (list, tuple))457 assert len(polymasks_list) > 0458 assert all(isinstance(polymask, PolygonMasks) for polymask in polymasks_list)459 460 cat_polymasks = type(polymasks_list[0])(461 list(itertools.chain.from_iterable(pm.polygons for pm in polymasks_list))462 )463 return cat_polymasks464 465 466class ROIMasks:467 """468 Represent masks by N smaller masks defined in some ROIs. Once ROI boxes are given,469 full-image bitmask can be obtained by "pasting" the mask on the region defined470 by the corresponding ROI box.471 """472 473 def __init__(self, tensor: torch.Tensor):474 """475 Args:476 tensor: (N, M, M) mask tensor that defines the mask within each ROI.477 """478 if tensor.dim() != 3:479 raise ValueError("ROIMasks must take a masks of 3 dimension.")480 self.tensor = tensor481 482 def to(self, device: torch.device) -> "ROIMasks":483 return ROIMasks(self.tensor.to(device))484 485 @property486 def device(self) -> device:487 return self.tensor.device488 489 def __len__(self):490 return self.tensor.shape[0]491 492 def __getitem__(self, item) -> "ROIMasks":493 """494 Returns:495 ROIMasks: Create a new :class:`ROIMasks` by indexing.496 497 The following usage are allowed:498 499 1. `new_masks = masks[2:10]`: return a slice of masks.500 2. `new_masks = masks[vector]`, where vector is a torch.BoolTensor501 with `length = len(masks)`. Nonzero elements in the vector will be selected.502 503 Note that the returned object might share storage with this object,504 subject to Pytorch's indexing semantics.505 """506 t = self.tensor[item]507 if t.dim() != 3:508 raise ValueError(509 f"Indexing on ROIMasks with {item} returns a tensor with shape {t.shape}!"510 )511 return ROIMasks(t)512 513 @torch.jit.unused514 def __repr__(self) -> str:515 s = self.__class__.__name__ + "("516 s += "num_instances={})".format(len(self.tensor))517 return s518 519 @torch.jit.unused520 def to_bitmasks(self, boxes: torch.Tensor, height, width, threshold=0.5):521 """522 Args: see documentation of :func:`paste_masks_in_image`.523 """524 from detectron2.layers.mask_ops import paste_masks_in_image, _paste_masks_tensor_shape525 526 if torch.jit.is_tracing():527 if isinstance(height, torch.Tensor):528 paste_func = _paste_masks_tensor_shape529 else:530 paste_func = paste_masks_in_image531 else:532 paste_func = retry_if_cuda_oom(paste_masks_in_image)533 bitmasks = paste_func(self.tensor, boxes.tensor, (height, width), threshold=threshold)534 return BitMasks(bitmasks)535