Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import torch3from torch import nn4from torch.autograd import Function5from torch.autograd.function import once_differentiable6from torch.nn.modules.utils import _pair7 8 9class _ROIAlignRotated(Function):10    @staticmethod11    def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio):12        ctx.save_for_backward(roi)13        ctx.output_size = _pair(output_size)14        ctx.spatial_scale = spatial_scale15        ctx.sampling_ratio = sampling_ratio16        ctx.input_shape = input.size()17        output = torch.ops.detectron2.roi_align_rotated_forward(18            input, roi, spatial_scale, output_size[0], output_size[1], sampling_ratio19        )20        return output21 22    @staticmethod23    @once_differentiable24    def backward(ctx, grad_output):25        (rois,) = ctx.saved_tensors26        output_size = ctx.output_size27        spatial_scale = ctx.spatial_scale28        sampling_ratio = ctx.sampling_ratio29        bs, ch, h, w = ctx.input_shape30        grad_input = torch.ops.detectron2.roi_align_rotated_backward(31            grad_output,32            rois,33            spatial_scale,34            output_size[0],35            output_size[1],36            bs,37            ch,38            h,39            w,40            sampling_ratio,41        )42        return grad_input, None, None, None, None, None43 44 45roi_align_rotated = _ROIAlignRotated.apply46 47 48class ROIAlignRotated(nn.Module):49    def __init__(self, output_size, spatial_scale, sampling_ratio):50        """51        Args:52            output_size (tuple): h, w53            spatial_scale (float): scale the input boxes by this number54            sampling_ratio (int): number of inputs samples to take for each output55                sample. 0 to take samples densely.56 57        Note:58            ROIAlignRotated supports continuous coordinate by default:59            Given a continuous coordinate c, its two neighboring pixel indices (in our60            pixel model) are computed by floor(c - 0.5) and ceil(c - 0.5). For example,61            c=1.3 has pixel neighbors with discrete indices [0] and [1] (which are sampled62            from the underlying signal at continuous coordinates 0.5 and 1.5).63        """64        super(ROIAlignRotated, self).__init__()65        self.output_size = output_size66        self.spatial_scale = spatial_scale67        self.sampling_ratio = sampling_ratio68 69    def forward(self, input, rois):70        """71        Args:72            input: NCHW images73            rois: Bx6 boxes. First column is the index into N.74                The other 5 columns are (x_ctr, y_ctr, width, height, angle_degrees).75        """76        assert rois.dim() == 2 and rois.size(1) == 677        orig_dtype = input.dtype78        if orig_dtype == torch.float16:79            input = input.float()80            rois = rois.float()81        output_size = _pair(self.output_size)82 83        # Scripting for Autograd is currently unsupported.84        # This is a quick fix without having to rewrite code on the C++ side85        if torch.jit.is_scripting() or torch.jit.is_tracing():86            return torch.ops.detectron2.roi_align_rotated_forward(87                input, rois, self.spatial_scale, output_size[0], output_size[1], self.sampling_ratio88            ).to(dtype=orig_dtype)89 90        return roi_align_rotated(91            input, rois, self.output_size, self.spatial_scale, self.sampling_ratio92        ).to(dtype=orig_dtype)93 94    def __repr__(self):95        tmpstr = self.__class__.__name__ + "("96        tmpstr += "output_size=" + str(self.output_size)97        tmpstr += ", spatial_scale=" + str(self.spatial_scale)98        tmpstr += ", sampling_ratio=" + str(self.sampling_ratio)99        tmpstr += ")"100        return tmpstr101