Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import math3from typing import List, Tuple, Union4import torch5from fvcore.nn import giou_loss, smooth_l1_loss6from torch.nn import functional as F7 8from detectron2.layers import cat, ciou_loss, diou_loss9from detectron2.structures import Boxes10 11# Value for clamping large dw and dh predictions. The heuristic is that we clamp12# such that dw and dh are no larger than what would transform a 16px box into a13# 1000px box (based on a small anchor, 16px, and a typical image size, 1000px).14_DEFAULT_SCALE_CLAMP = math.log(1000.0 / 16)15 16 17__all__ = ["Box2BoxTransform", "Box2BoxTransformRotated", "Box2BoxTransformLinear"]18 19 20@torch.jit.script21class Box2BoxTransform(object):22    """23    The box-to-box transform defined in R-CNN. The transformation is parameterized24    by 4 deltas: (dx, dy, dw, dh). The transformation scales the box's width and height25    by exp(dw), exp(dh) and shifts a box's center by the offset (dx * width, dy * height).26    """27 28    def __init__(29        self, weights: Tuple[float, float, float, float], scale_clamp: float = _DEFAULT_SCALE_CLAMP30    ):31        """32        Args:33            weights (4-element tuple): Scaling factors that are applied to the34                (dx, dy, dw, dh) deltas. In Fast R-CNN, these were originally set35                such that the deltas have unit variance; now they are treated as36                hyperparameters of the system.37            scale_clamp (float): When predicting deltas, the predicted box scaling38                factors (dw and dh) are clamped such that they are <= scale_clamp.39        """40        self.weights = weights41        self.scale_clamp = scale_clamp42 43    def get_deltas(self, src_boxes, target_boxes):44        """45        Get box regression transformation deltas (dx, dy, dw, dh) that can be used46        to transform the `src_boxes` into the `target_boxes`. That is, the relation47        ``target_boxes == self.apply_deltas(deltas, src_boxes)`` is true (unless48        any delta is too large and is clamped).49 50        Args:51            src_boxes (Tensor): source boxes, e.g., object proposals52            target_boxes (Tensor): target of the transformation, e.g., ground-truth53                boxes.54        """55        assert isinstance(src_boxes, torch.Tensor), type(src_boxes)56        assert isinstance(target_boxes, torch.Tensor), type(target_boxes)57 58        src_widths = src_boxes[:, 2] - src_boxes[:, 0]59        src_heights = src_boxes[:, 3] - src_boxes[:, 1]60        src_ctr_x = src_boxes[:, 0] + 0.5 * src_widths61        src_ctr_y = src_boxes[:, 1] + 0.5 * src_heights62 63        target_widths = target_boxes[:, 2] - target_boxes[:, 0]64        target_heights = target_boxes[:, 3] - target_boxes[:, 1]65        target_ctr_x = target_boxes[:, 0] + 0.5 * target_widths66        target_ctr_y = target_boxes[:, 1] + 0.5 * target_heights67 68        wx, wy, ww, wh = self.weights69        dx = wx * (target_ctr_x - src_ctr_x) / src_widths70        dy = wy * (target_ctr_y - src_ctr_y) / src_heights71        dw = ww * torch.log(target_widths / src_widths)72        dh = wh * torch.log(target_heights / src_heights)73 74        deltas = torch.stack((dx, dy, dw, dh), dim=1)75        assert (src_widths > 0).all().item(), "Input boxes to Box2BoxTransform are not valid!"76        return deltas77 78    def apply_deltas(self, deltas, boxes):79        """80        Apply transformation `deltas` (dx, dy, dw, dh) to `boxes`.81 82        Args:83            deltas (Tensor): transformation deltas of shape (N, k*4), where k >= 1.84                deltas[i] represents k potentially different class-specific85                box transformations for the single box boxes[i].86            boxes (Tensor): boxes to transform, of shape (N, 4)87        """88        deltas = deltas.float()  # ensure fp32 for decoding precision89        boxes = boxes.to(deltas.dtype)90 91        widths = boxes[:, 2] - boxes[:, 0]92        heights = boxes[:, 3] - boxes[:, 1]93        ctr_x = boxes[:, 0] + 0.5 * widths94        ctr_y = boxes[:, 1] + 0.5 * heights95 96        wx, wy, ww, wh = self.weights97        dx = deltas[:, 0::4] / wx98        dy = deltas[:, 1::4] / wy99        dw = deltas[:, 2::4] / ww100        dh = deltas[:, 3::4] / wh101 102        # Prevent sending too large values into torch.exp()103        dw = torch.clamp(dw, max=self.scale_clamp)104        dh = torch.clamp(dh, max=self.scale_clamp)105 106        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]107        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]108        pred_w = torch.exp(dw) * widths[:, None]109        pred_h = torch.exp(dh) * heights[:, None]110 111        x1 = pred_ctr_x - 0.5 * pred_w112        y1 = pred_ctr_y - 0.5 * pred_h113        x2 = pred_ctr_x + 0.5 * pred_w114        y2 = pred_ctr_y + 0.5 * pred_h115        pred_boxes = torch.stack((x1, y1, x2, y2), dim=-1)116        return pred_boxes.reshape(deltas.shape)117 118 119@torch.jit.script120class Box2BoxTransformRotated(object):121    """122    The box-to-box transform defined in Rotated R-CNN. The transformation is parameterized123    by 5 deltas: (dx, dy, dw, dh, da). The transformation scales the box's width and height124    by exp(dw), exp(dh), shifts a box's center by the offset (dx * width, dy * height),125    and rotate a box's angle by da (radians).126    Note: angles of deltas are in radians while angles of boxes are in degrees.127    """128 129    def __init__(130        self,131        weights: Tuple[float, float, float, float, float],132        scale_clamp: float = _DEFAULT_SCALE_CLAMP,133    ):134        """135        Args:136            weights (5-element tuple): Scaling factors that are applied to the137                (dx, dy, dw, dh, da) deltas. These are treated as138                hyperparameters of the system.139            scale_clamp (float): When predicting deltas, the predicted box scaling140                factors (dw and dh) are clamped such that they are <= scale_clamp.141        """142        self.weights = weights143        self.scale_clamp = scale_clamp144 145    def get_deltas(self, src_boxes, target_boxes):146        """147        Get box regression transformation deltas (dx, dy, dw, dh, da) that can be used148        to transform the `src_boxes` into the `target_boxes`. That is, the relation149        ``target_boxes == self.apply_deltas(deltas, src_boxes)`` is true (unless150        any delta is too large and is clamped).151 152        Args:153            src_boxes (Tensor): Nx5 source boxes, e.g., object proposals154            target_boxes (Tensor): Nx5 target of the transformation, e.g., ground-truth155                boxes.156        """157        assert isinstance(src_boxes, torch.Tensor), type(src_boxes)158        assert isinstance(target_boxes, torch.Tensor), type(target_boxes)159 160        src_ctr_x, src_ctr_y, src_widths, src_heights, src_angles = torch.unbind(src_boxes, dim=1)161 162        target_ctr_x, target_ctr_y, target_widths, target_heights, target_angles = torch.unbind(163            target_boxes, dim=1164        )165 166        wx, wy, ww, wh, wa = self.weights167        dx = wx * (target_ctr_x - src_ctr_x) / src_widths168        dy = wy * (target_ctr_y - src_ctr_y) / src_heights169        dw = ww * torch.log(target_widths / src_widths)170        dh = wh * torch.log(target_heights / src_heights)171        # Angles of deltas are in radians while angles of boxes are in degrees.172        # the conversion to radians serve as a way to normalize the values173        da = target_angles - src_angles174        da = (da + 180.0) % 360.0 - 180.0  # make it in [-180, 180)175        da *= wa * math.pi / 180.0176 177        deltas = torch.stack((dx, dy, dw, dh, da), dim=1)178        assert (179            (src_widths > 0).all().item()180        ), "Input boxes to Box2BoxTransformRotated are not valid!"181        return deltas182 183    def apply_deltas(self, deltas, boxes):184        """185        Apply transformation `deltas` (dx, dy, dw, dh, da) to `boxes`.186 187        Args:188            deltas (Tensor): transformation deltas of shape (N, k*5).189                deltas[i] represents box transformation for the single box boxes[i].190            boxes (Tensor): boxes to transform, of shape (N, 5)191        """192        assert deltas.shape[1] % 5 == 0 and boxes.shape[1] == 5193 194        boxes = boxes.to(deltas.dtype).unsqueeze(2)195 196        ctr_x = boxes[:, 0]197        ctr_y = boxes[:, 1]198        widths = boxes[:, 2]199        heights = boxes[:, 3]200        angles = boxes[:, 4]201 202        wx, wy, ww, wh, wa = self.weights203 204        dx = deltas[:, 0::5] / wx205        dy = deltas[:, 1::5] / wy206        dw = deltas[:, 2::5] / ww207        dh = deltas[:, 3::5] / wh208        da = deltas[:, 4::5] / wa209 210        # Prevent sending too large values into torch.exp()211        dw = torch.clamp(dw, max=self.scale_clamp)212        dh = torch.clamp(dh, max=self.scale_clamp)213 214        pred_boxes = torch.zeros_like(deltas)215        pred_boxes[:, 0::5] = dx * widths + ctr_x  # x_ctr216        pred_boxes[:, 1::5] = dy * heights + ctr_y  # y_ctr217        pred_boxes[:, 2::5] = torch.exp(dw) * widths  # width218        pred_boxes[:, 3::5] = torch.exp(dh) * heights  # height219 220        # Following original RRPN implementation,221        # angles of deltas are in radians while angles of boxes are in degrees.222        pred_angle = da * 180.0 / math.pi + angles223        pred_angle = (pred_angle + 180.0) % 360.0 - 180.0  # make it in [-180, 180)224 225        pred_boxes[:, 4::5] = pred_angle226 227        return pred_boxes228 229 230class Box2BoxTransformLinear(object):231    """232    The linear box-to-box transform defined in FCOS. The transformation is parameterized233    by the distance from the center of (square) src box to 4 edges of the target box.234    """235 236    def __init__(self, normalize_by_size=True):237        """238        Args:239            normalize_by_size: normalize deltas by the size of src (anchor) boxes.240        """241        self.normalize_by_size = normalize_by_size242 243    def get_deltas(self, src_boxes, target_boxes):244        """245        Get box regression transformation deltas (dx1, dy1, dx2, dy2) that can be used246        to transform the `src_boxes` into the `target_boxes`. That is, the relation247        ``target_boxes == self.apply_deltas(deltas, src_boxes)`` is true.248        The center of src must be inside target boxes.249 250        Args:251            src_boxes (Tensor): square source boxes, e.g., anchors252            target_boxes (Tensor): target of the transformation, e.g., ground-truth253                boxes.254        """255        assert isinstance(src_boxes, torch.Tensor), type(src_boxes)256        assert isinstance(target_boxes, torch.Tensor), type(target_boxes)257 258        src_ctr_x = 0.5 * (src_boxes[:, 0] + src_boxes[:, 2])259        src_ctr_y = 0.5 * (src_boxes[:, 1] + src_boxes[:, 3])260 261        target_l = src_ctr_x - target_boxes[:, 0]262        target_t = src_ctr_y - target_boxes[:, 1]263        target_r = target_boxes[:, 2] - src_ctr_x264        target_b = target_boxes[:, 3] - src_ctr_y265 266        deltas = torch.stack((target_l, target_t, target_r, target_b), dim=1)267        if self.normalize_by_size:268            stride_w = src_boxes[:, 2] - src_boxes[:, 0]269            stride_h = src_boxes[:, 3] - src_boxes[:, 1]270            strides = torch.stack([stride_w, stride_h, stride_w, stride_h], axis=1)271            deltas = deltas / strides272 273        return deltas274 275    def apply_deltas(self, deltas, boxes):276        """277        Apply transformation `deltas` (dx1, dy1, dx2, dy2) to `boxes`.278 279        Args:280            deltas (Tensor): transformation deltas of shape (N, k*4), where k >= 1.281                deltas[i] represents k potentially different class-specific282                box transformations for the single box boxes[i].283            boxes (Tensor): boxes to transform, of shape (N, 4)284        """285        # Ensure the output is a valid box. See Sec 2.1 of https://arxiv.org/abs/2006.09214286        deltas = F.relu(deltas)287        boxes = boxes.to(deltas.dtype)288 289        ctr_x = 0.5 * (boxes[:, 0] + boxes[:, 2])290        ctr_y = 0.5 * (boxes[:, 1] + boxes[:, 3])291        if self.normalize_by_size:292            stride_w = boxes[:, 2] - boxes[:, 0]293            stride_h = boxes[:, 3] - boxes[:, 1]294            strides = torch.stack([stride_w, stride_h, stride_w, stride_h], axis=1)295            deltas = deltas * strides296 297        l = deltas[:, 0::4]298        t = deltas[:, 1::4]299        r = deltas[:, 2::4]300        b = deltas[:, 3::4]301 302        pred_boxes = torch.zeros_like(deltas)303        pred_boxes[:, 0::4] = ctr_x[:, None] - l  # x1304        pred_boxes[:, 1::4] = ctr_y[:, None] - t  # y1305        pred_boxes[:, 2::4] = ctr_x[:, None] + r  # x2306        pred_boxes[:, 3::4] = ctr_y[:, None] + b  # y2307        return pred_boxes308 309 310def _dense_box_regression_loss(311    anchors: List[Union[Boxes, torch.Tensor]],312    box2box_transform: Box2BoxTransform,313    pred_anchor_deltas: List[torch.Tensor],314    gt_boxes: List[torch.Tensor],315    fg_mask: torch.Tensor,316    box_reg_loss_type="smooth_l1",317    smooth_l1_beta=0.0,318):319    """320    Compute loss for dense multi-level box regression.321    Loss is accumulated over ``fg_mask``.322 323    Args:324        anchors: #lvl anchor boxes, each is (HixWixA, 4)325        pred_anchor_deltas: #lvl predictions, each is (N, HixWixA, 4)326        gt_boxes: N ground truth boxes, each has shape (R, 4) (R = sum(Hi * Wi * A))327        fg_mask: the foreground boolean mask of shape (N, R) to compute loss on328        box_reg_loss_type (str): Loss type to use. Supported losses: "smooth_l1", "giou",329            "diou", "ciou".330        smooth_l1_beta (float): beta parameter for the smooth L1 regression loss. Default to331            use L1 loss. Only used when `box_reg_loss_type` is "smooth_l1"332    """333    if isinstance(anchors[0], Boxes):334        anchors = type(anchors[0]).cat(anchors).tensor  # (R, 4)335    else:336        anchors = cat(anchors)337    if box_reg_loss_type == "smooth_l1":338        gt_anchor_deltas = [box2box_transform.get_deltas(anchors, k) for k in gt_boxes]339        gt_anchor_deltas = torch.stack(gt_anchor_deltas)  # (N, R, 4)340        loss_box_reg = smooth_l1_loss(341            cat(pred_anchor_deltas, dim=1)[fg_mask],342            gt_anchor_deltas[fg_mask],343            beta=smooth_l1_beta,344            reduction="sum",345        )346    elif box_reg_loss_type == "giou":347        pred_boxes = [348            box2box_transform.apply_deltas(k, anchors) for k in cat(pred_anchor_deltas, dim=1)349        ]350        loss_box_reg = giou_loss(351            torch.stack(pred_boxes)[fg_mask], torch.stack(gt_boxes)[fg_mask], reduction="sum"352        )353    elif box_reg_loss_type == "diou":354        pred_boxes = [355            box2box_transform.apply_deltas(k, anchors) for k in cat(pred_anchor_deltas, dim=1)356        ]357        loss_box_reg = diou_loss(358            torch.stack(pred_boxes)[fg_mask], torch.stack(gt_boxes)[fg_mask], reduction="sum"359        )360    elif box_reg_loss_type == "ciou":361        pred_boxes = [362            box2box_transform.apply_deltas(k, anchors) for k in cat(pred_anchor_deltas, dim=1)363        ]364        loss_box_reg = ciou_loss(365            torch.stack(pred_boxes)[fg_mask], torch.stack(gt_boxes)[fg_mask], reduction="sum"366        )367    else:368        raise ValueError(f"Invalid dense box regression loss type '{box_reg_loss_type}'")369    return loss_box_reg370