Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# -*- coding: utf-8 -*-2# Copyright (c) Facebook, Inc. and its affiliates.3 4import torch5from torchvision.ops import boxes as box_ops6from torchvision.ops import nms  # noqa . for compatibility7 8 9def batched_nms(10    boxes: torch.Tensor, scores: torch.Tensor, idxs: torch.Tensor, iou_threshold: float11):12    """13    Same as torchvision.ops.boxes.batched_nms, but with float().14    """15    assert boxes.shape[-1] == 416    # Note: Torchvision already has a strategy (https://github.com/pytorch/vision/issues/1311)17    # to decide whether to use coordinate trick or for loop to implement batched_nms. So we18    # just call it directly.19    # Fp16 does not have enough range for batched NMS, so adding float().20    return box_ops.batched_nms(boxes.float(), scores, idxs, iou_threshold)21 22 23# Note: this function (nms_rotated) might be moved into24# torchvision/ops/boxes.py in the future25def nms_rotated(boxes: torch.Tensor, scores: torch.Tensor, iou_threshold: float):26    """27    Performs non-maximum suppression (NMS) on the rotated boxes according28    to their intersection-over-union (IoU).29 30    Rotated NMS iteratively removes lower scoring rotated boxes which have an31    IoU greater than iou_threshold with another (higher scoring) rotated box.32 33    Note that RotatedBox (5, 3, 4, 2, -90) covers exactly the same region as34    RotatedBox (5, 3, 4, 2, 90) does, and their IoU will be 1. However, they35    can be representing completely different objects in certain tasks, e.g., OCR.36 37    As for the question of whether rotated-NMS should treat them as faraway boxes38    even though their IOU is 1, it depends on the application and/or ground truth annotation.39 40    As an extreme example, consider a single character v and the square box around it.41 42    If the angle is 0 degree, the object (text) would be read as 'v';43 44    If the angle is 90 degrees, the object (text) would become '>';45 46    If the angle is 180 degrees, the object (text) would become '^';47 48    If the angle is 270/-90 degrees, the object (text) would become '<'49 50    All of these cases have IoU of 1 to each other, and rotated NMS that only51    uses IoU as criterion would only keep one of them with the highest score -52    which, practically, still makes sense in most cases because typically53    only one of theses orientations is the correct one. Also, it does not matter54    as much if the box is only used to classify the object (instead of transcribing55    them with a sequential OCR recognition model) later.56 57    On the other hand, when we use IoU to filter proposals that are close to the58    ground truth during training, we should definitely take the angle into account if59    we know the ground truth is labeled with the strictly correct orientation (as in,60    upside-down words are annotated with -180 degrees even though they can be covered61    with a 0/90/-90 degree box, etc.)62 63    The way the original dataset is annotated also matters. For example, if the dataset64    is a 4-point polygon dataset that does not enforce ordering of vertices/orientation,65    we can estimate a minimum rotated bounding box to this polygon, but there's no way66    we can tell the correct angle with 100% confidence (as shown above, there could be 4 different67    rotated boxes, with angles differed by 90 degrees to each other, covering the exactly68    same region). In that case we have to just use IoU to determine the box69    proximity (as many detection benchmarks (even for text) do) unless there're other70    assumptions we can make (like width is always larger than height, or the object is not71    rotated by more than 90 degrees CCW/CW, etc.)72 73    In summary, not considering angles in rotated NMS seems to be a good option for now,74    but we should be aware of its implications.75 76    Args:77        boxes (Tensor[N, 5]): Rotated boxes to perform NMS on. They are expected to be in78           (x_center, y_center, width, height, angle_degrees) format.79        scores (Tensor[N]): Scores for each one of the rotated boxes80        iou_threshold (float): Discards all overlapping rotated boxes with IoU < iou_threshold81 82    Returns:83        keep (Tensor): int64 tensor with the indices of the elements that have been kept84        by Rotated NMS, sorted in decreasing order of scores85    """86    return torch.ops.detectron2.nms_rotated(boxes, scores, iou_threshold)87 88 89# Note: this function (batched_nms_rotated) might be moved into90# torchvision/ops/boxes.py in the future91 92 93@torch.jit.script_if_tracing94def batched_nms_rotated(95    boxes: torch.Tensor, scores: torch.Tensor, idxs: torch.Tensor, iou_threshold: float96):97    """98    Performs non-maximum suppression in a batched fashion.99 100    Each index value correspond to a category, and NMS101    will not be applied between elements of different categories.102 103    Args:104        boxes (Tensor[N, 5]):105           boxes where NMS will be performed. They106           are expected to be in (x_ctr, y_ctr, width, height, angle_degrees) format107        scores (Tensor[N]):108           scores for each one of the boxes109        idxs (Tensor[N]):110           indices of the categories for each one of the boxes.111        iou_threshold (float):112           discards all overlapping boxes113           with IoU < iou_threshold114 115    Returns:116        Tensor:117            int64 tensor with the indices of the elements that have been kept118            by NMS, sorted in decreasing order of scores119    """120    assert boxes.shape[-1] == 5121 122    if boxes.numel() == 0:123        return torch.empty((0,), dtype=torch.int64, device=boxes.device)124    boxes = boxes.float()  # fp16 does not have enough range for batched NMS125    # Strategy: in order to perform NMS independently per class,126    # we add an offset to all the boxes. The offset is dependent127    # only on the class idx, and is large enough so that boxes128    # from different classes do not overlap129 130    # Note that batched_nms in torchvision/ops/boxes.py only uses max_coordinate,131    # which won't handle negative coordinates correctly.132    # Here by using min_coordinate we can make sure the negative coordinates are133    # correctly handled.134    max_coordinate = (135        torch.max(boxes[:, 0], boxes[:, 1]) + torch.max(boxes[:, 2], boxes[:, 3]) / 2136    ).max()137    min_coordinate = (138        torch.min(boxes[:, 0], boxes[:, 1]) - torch.max(boxes[:, 2], boxes[:, 3]) / 2139    ).min()140    offsets = idxs.to(boxes) * (max_coordinate - min_coordinate + 1)141    boxes_for_nms = boxes.clone()  # avoid modifying the original values in boxes142    boxes_for_nms[:, :2] += offsets[:, None]143    keep = nms_rotated(boxes_for_nms, scores, iou_threshold)144    return keep145