Arulkumar03/Fox_Sheep_Detector_Computer_Vision_model
0
1# Copyright (c) Facebook, Inc. and its affiliates.2import collections3import math4from typing import List5import torch6from torch import nn7 8from detectron2.config import configurable9from detectron2.layers import ShapeSpec, move_device_like10from detectron2.structures import Boxes, RotatedBoxes11from detectron2.utils.registry import Registry12 13ANCHOR_GENERATOR_REGISTRY = Registry("ANCHOR_GENERATOR")14ANCHOR_GENERATOR_REGISTRY.__doc__ = """15Registry for modules that creates object detection anchors for feature maps.16 17The registered object will be called with `obj(cfg, input_shape)`.18"""19 20 21class BufferList(nn.Module):22 """23 Similar to nn.ParameterList, but for buffers24 """25 26 def __init__(self, buffers):27 super().__init__()28 for i, buffer in enumerate(buffers):29 # Use non-persistent buffer so the values are not saved in checkpoint30 self.register_buffer(str(i), buffer, persistent=False)31 32 def __len__(self):33 return len(self._buffers)34 35 def __iter__(self):36 return iter(self._buffers.values())37 38 39def _create_grid_offsets(40 size: List[int], stride: int, offset: float, target_device_tensor: torch.Tensor41):42 grid_height, grid_width = size43 shifts_x = move_device_like(44 torch.arange(offset * stride, grid_width * stride, step=stride, dtype=torch.float32),45 target_device_tensor,46 )47 shifts_y = move_device_like(48 torch.arange(offset * stride, grid_height * stride, step=stride, dtype=torch.float32),49 target_device_tensor,50 )51 52 shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)53 shift_x = shift_x.reshape(-1)54 shift_y = shift_y.reshape(-1)55 return shift_x, shift_y56 57 58def _broadcast_params(params, num_features, name):59 """60 If one size (or aspect ratio) is specified and there are multiple feature61 maps, we "broadcast" anchors of that single size (or aspect ratio)62 over all feature maps.63 64 If params is list[float], or list[list[float]] with len(params) == 1, repeat65 it num_features time.66 67 Returns:68 list[list[float]]: param for each feature69 """70 assert isinstance(71 params, collections.abc.Sequence72 ), f"{name} in anchor generator has to be a list! Got {params}."73 assert len(params), f"{name} in anchor generator cannot be empty!"74 if not isinstance(params[0], collections.abc.Sequence): # params is list[float]75 return [params] * num_features76 if len(params) == 1:77 return list(params) * num_features78 assert len(params) == num_features, (79 f"Got {name} of length {len(params)} in anchor generator, "80 f"but the number of input features is {num_features}!"81 )82 return params83 84 85@ANCHOR_GENERATOR_REGISTRY.register()86class DefaultAnchorGenerator(nn.Module):87 """88 Compute anchors in the standard ways described in89 "Faster R-CNN: Towards Real-Time Object Detection with Region Proposal Networks".90 """91 92 box_dim: torch.jit.Final[int] = 493 """94 the dimension of each anchor box.95 """96 97 @configurable98 def __init__(self, *, sizes, aspect_ratios, strides, offset=0.5):99 """100 This interface is experimental.101 102 Args:103 sizes (list[list[float]] or list[float]):104 If ``sizes`` is list[list[float]], ``sizes[i]`` is the list of anchor sizes105 (i.e. sqrt of anchor area) to use for the i-th feature map.106 If ``sizes`` is list[float], ``sizes`` is used for all feature maps.107 Anchor sizes are given in absolute lengths in units of108 the input image; they do not dynamically scale if the input image size changes.109 aspect_ratios (list[list[float]] or list[float]): list of aspect ratios110 (i.e. height / width) to use for anchors. Same "broadcast" rule for `sizes` applies.111 strides (list[int]): stride of each input feature.112 offset (float): Relative offset between the center of the first anchor and the top-left113 corner of the image. Value has to be in [0, 1).114 Recommend to use 0.5, which means half stride.115 """116 super().__init__()117 118 self.strides = strides119 self.num_features = len(self.strides)120 sizes = _broadcast_params(sizes, self.num_features, "sizes")121 aspect_ratios = _broadcast_params(aspect_ratios, self.num_features, "aspect_ratios")122 self.cell_anchors = self._calculate_anchors(sizes, aspect_ratios)123 124 self.offset = offset125 assert 0.0 <= self.offset < 1.0, self.offset126 127 @classmethod128 def from_config(cls, cfg, input_shape: List[ShapeSpec]):129 return {130 "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,131 "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,132 "strides": [x.stride for x in input_shape],133 "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,134 }135 136 def _calculate_anchors(self, sizes, aspect_ratios):137 cell_anchors = [138 self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)139 ]140 return BufferList(cell_anchors)141 142 @property143 @torch.jit.unused144 def num_cell_anchors(self):145 """146 Alias of `num_anchors`.147 """148 return self.num_anchors149 150 @property151 @torch.jit.unused152 def num_anchors(self):153 """154 Returns:155 list[int]: Each int is the number of anchors at every pixel156 location, on that feature map.157 For example, if at every pixel we use anchors of 3 aspect158 ratios and 5 sizes, the number of anchors is 15.159 (See also ANCHOR_GENERATOR.SIZES and ANCHOR_GENERATOR.ASPECT_RATIOS in config)160 161 In standard RPN models, `num_anchors` on every feature map is the same.162 """163 return [len(cell_anchors) for cell_anchors in self.cell_anchors]164 165 def _grid_anchors(self, grid_sizes: List[List[int]]):166 """167 Returns:168 list[Tensor]: #featuremap tensors, each is (#locations x #cell_anchors) x 4169 """170 anchors = []171 # buffers() not supported by torchscript. use named_buffers() instead172 buffers: List[torch.Tensor] = [x[1] for x in self.cell_anchors.named_buffers()]173 for size, stride, base_anchors in zip(grid_sizes, self.strides, buffers):174 shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors)175 shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)176 177 anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))178 179 return anchors180 181 def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):182 """183 Generate a tensor storing canonical anchor boxes, which are all anchor184 boxes of different sizes and aspect_ratios centered at (0, 0).185 We can later build the set of anchors for a full feature map by186 shifting and tiling these tensors (see `meth:_grid_anchors`).187 188 Args:189 sizes (tuple[float]):190 aspect_ratios (tuple[float]]):191 192 Returns:193 Tensor of shape (len(sizes) * len(aspect_ratios), 4) storing anchor boxes194 in XYXY format.195 """196 197 # This is different from the anchor generator defined in the original Faster R-CNN198 # code or Detectron. They yield the same AP, however the old version defines cell199 # anchors in a less natural way with a shift relative to the feature grid and200 # quantization that results in slightly different sizes for different aspect ratios.201 # See also https://github.com/facebookresearch/Detectron/issues/227202 203 anchors = []204 for size in sizes:205 area = size**2.0206 for aspect_ratio in aspect_ratios:207 # s * s = w * h208 # a = h / w209 # ... some algebra ...210 # w = sqrt(s * s / a)211 # h = a * w212 w = math.sqrt(area / aspect_ratio)213 h = aspect_ratio * w214 x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0215 anchors.append([x0, y0, x1, y1])216 return torch.tensor(anchors)217 218 def forward(self, features: List[torch.Tensor]):219 """220 Args:221 features (list[Tensor]): list of backbone feature maps on which to generate anchors.222 223 Returns:224 list[Boxes]: a list of Boxes containing all the anchors for each feature map225 (i.e. the cell anchors repeated over all locations in the feature map).226 The number of anchors of each feature map is Hi x Wi x num_cell_anchors,227 where Hi, Wi are resolution of the feature map divided by anchor stride.228 """229 grid_sizes = [feature_map.shape[-2:] for feature_map in features]230 anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)231 return [Boxes(x) for x in anchors_over_all_feature_maps]232 233 234@ANCHOR_GENERATOR_REGISTRY.register()235class RotatedAnchorGenerator(nn.Module):236 """237 Compute rotated anchors used by Rotated RPN (RRPN), described in238 "Arbitrary-Oriented Scene Text Detection via Rotation Proposals".239 """240 241 box_dim: int = 5242 """243 the dimension of each anchor box.244 """245 246 @configurable247 def __init__(self, *, sizes, aspect_ratios, strides, angles, offset=0.5):248 """249 This interface is experimental.250 251 Args:252 sizes (list[list[float]] or list[float]):253 If sizes is list[list[float]], sizes[i] is the list of anchor sizes254 (i.e. sqrt of anchor area) to use for the i-th feature map.255 If sizes is list[float], the sizes are used for all feature maps.256 Anchor sizes are given in absolute lengths in units of257 the input image; they do not dynamically scale if the input image size changes.258 aspect_ratios (list[list[float]] or list[float]): list of aspect ratios259 (i.e. height / width) to use for anchors. Same "broadcast" rule for `sizes` applies.260 strides (list[int]): stride of each input feature.261 angles (list[list[float]] or list[float]): list of angles (in degrees CCW)262 to use for anchors. Same "broadcast" rule for `sizes` applies.263 offset (float): Relative offset between the center of the first anchor and the top-left264 corner of the image. Value has to be in [0, 1).265 Recommend to use 0.5, which means half stride.266 """267 super().__init__()268 269 self.strides = strides270 self.num_features = len(self.strides)271 sizes = _broadcast_params(sizes, self.num_features, "sizes")272 aspect_ratios = _broadcast_params(aspect_ratios, self.num_features, "aspect_ratios")273 angles = _broadcast_params(angles, self.num_features, "angles")274 self.cell_anchors = self._calculate_anchors(sizes, aspect_ratios, angles)275 276 self.offset = offset277 assert 0.0 <= self.offset < 1.0, self.offset278 279 @classmethod280 def from_config(cls, cfg, input_shape: List[ShapeSpec]):281 return {282 "sizes": cfg.MODEL.ANCHOR_GENERATOR.SIZES,283 "aspect_ratios": cfg.MODEL.ANCHOR_GENERATOR.ASPECT_RATIOS,284 "strides": [x.stride for x in input_shape],285 "offset": cfg.MODEL.ANCHOR_GENERATOR.OFFSET,286 "angles": cfg.MODEL.ANCHOR_GENERATOR.ANGLES,287 }288 289 def _calculate_anchors(self, sizes, aspect_ratios, angles):290 cell_anchors = [291 self.generate_cell_anchors(size, aspect_ratio, angle).float()292 for size, aspect_ratio, angle in zip(sizes, aspect_ratios, angles)293 ]294 return BufferList(cell_anchors)295 296 @property297 def num_cell_anchors(self):298 """299 Alias of `num_anchors`.300 """301 return self.num_anchors302 303 @property304 def num_anchors(self):305 """306 Returns:307 list[int]: Each int is the number of anchors at every pixel308 location, on that feature map.309 For example, if at every pixel we use anchors of 3 aspect310 ratios, 2 sizes and 5 angles, the number of anchors is 30.311 (See also ANCHOR_GENERATOR.SIZES, ANCHOR_GENERATOR.ASPECT_RATIOS312 and ANCHOR_GENERATOR.ANGLES in config)313 314 In standard RRPN models, `num_anchors` on every feature map is the same.315 """316 return [len(cell_anchors) for cell_anchors in self.cell_anchors]317 318 def _grid_anchors(self, grid_sizes):319 anchors = []320 for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):321 shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors)322 zeros = torch.zeros_like(shift_x)323 shifts = torch.stack((shift_x, shift_y, zeros, zeros, zeros), dim=1)324 325 anchors.append((shifts.view(-1, 1, 5) + base_anchors.view(1, -1, 5)).reshape(-1, 5))326 327 return anchors328 329 def generate_cell_anchors(330 self,331 sizes=(32, 64, 128, 256, 512),332 aspect_ratios=(0.5, 1, 2),333 angles=(-90, -60, -30, 0, 30, 60, 90),334 ):335 """336 Generate a tensor storing canonical anchor boxes, which are all anchor337 boxes of different sizes, aspect_ratios, angles centered at (0, 0).338 We can later build the set of anchors for a full feature map by339 shifting and tiling these tensors (see `meth:_grid_anchors`).340 341 Args:342 sizes (tuple[float]):343 aspect_ratios (tuple[float]]):344 angles (tuple[float]]):345 346 Returns:347 Tensor of shape (len(sizes) * len(aspect_ratios) * len(angles), 5)348 storing anchor boxes in (x_ctr, y_ctr, w, h, angle) format.349 """350 anchors = []351 for size in sizes:352 area = size**2.0353 for aspect_ratio in aspect_ratios:354 # s * s = w * h355 # a = h / w356 # ... some algebra ...357 # w = sqrt(s * s / a)358 # h = a * w359 w = math.sqrt(area / aspect_ratio)360 h = aspect_ratio * w361 anchors.extend([0, 0, w, h, a] for a in angles)362 363 return torch.tensor(anchors)364 365 def forward(self, features):366 """367 Args:368 features (list[Tensor]): list of backbone feature maps on which to generate anchors.369 370 Returns:371 list[RotatedBoxes]: a list of Boxes containing all the anchors for each feature map372 (i.e. the cell anchors repeated over all locations in the feature map).373 The number of anchors of each feature map is Hi x Wi x num_cell_anchors,374 where Hi, Wi are resolution of the feature map divided by anchor stride.375 """376 grid_sizes = [feature_map.shape[-2:] for feature_map in features]377 anchors_over_all_feature_maps = self._grid_anchors(grid_sizes)378 return [RotatedBoxes(x) for x in anchors_over_all_feature_maps]379 380 381def build_anchor_generator(cfg, input_shape):382 """383 Built an anchor generator from `cfg.MODEL.ANCHOR_GENERATOR.NAME`.384 """385 anchor_generator = cfg.MODEL.ANCHOR_GENERATOR.NAME386 return ANCHOR_GENERATOR_REGISTRY.get(anchor_generator)(cfg, input_shape)387