Team Ai
Apppublic

Arulkumar03/Fox_Sheep_Detector_Computer_Vision_model

sourceHugging Facemitupdated 3y agoView on Hugging Face
0likes
postprocessing.py101 linesDownload Raw Back to modeling
1# Copyright (c) Facebook, Inc. and its affiliates.2import torch3from torch.nn import functional as F4 5from detectron2.structures import Instances, ROIMasks6 7 8# perhaps should rename to "resize_instance"9def detector_postprocess(10    results: Instances, output_height: int, output_width: int, mask_threshold: float = 0.511):12    """13    Resize the output instances.14    The input images are often resized when entering an object detector.15    As a result, we often need the outputs of the detector in a different16    resolution from its inputs.17 18    This function will resize the raw outputs of an R-CNN detector19    to produce outputs according to the desired output resolution.20 21    Args:22        results (Instances): the raw outputs from the detector.23            `results.image_size` contains the input image resolution the detector sees.24            This object might be modified in-place.25        output_height, output_width: the desired output resolution.26    Returns:27        Instances: the resized output from the model, based on the output resolution28    """29    if isinstance(output_width, torch.Tensor):30        # This shape might (but not necessarily) be tensors during tracing.31        # Converts integer tensors to float temporaries to ensure true32        # division is performed when computing scale_x and scale_y.33        output_width_tmp = output_width.float()34        output_height_tmp = output_height.float()35        new_size = torch.stack([output_height, output_width])36    else:37        new_size = (output_height, output_width)38        output_width_tmp = output_width39        output_height_tmp = output_height40 41    scale_x, scale_y = (42        output_width_tmp / results.image_size[1],43        output_height_tmp / results.image_size[0],44    )45    results = Instances(new_size, **results.get_fields())46 47    if results.has("pred_boxes"):48        output_boxes = results.pred_boxes49    elif results.has("proposal_boxes"):50        output_boxes = results.proposal_boxes51    else:52        output_boxes = None53    assert output_boxes is not None, "Predictions must contain boxes!"54 55    output_boxes.scale(scale_x, scale_y)56    output_boxes.clip(results.image_size)57 58    results = results[output_boxes.nonempty()]59 60    if results.has("pred_masks"):61        if isinstance(results.pred_masks, ROIMasks):62            roi_masks = results.pred_masks63        else:64            # pred_masks is a tensor of shape (N, 1, M, M)65            roi_masks = ROIMasks(results.pred_masks[:, 0, :, :])66        results.pred_masks = roi_masks.to_bitmasks(67            results.pred_boxes, output_height, output_width, mask_threshold68        ).tensor  # TODO return ROIMasks/BitMask object in the future69 70    if results.has("pred_keypoints"):71        results.pred_keypoints[:, :, 0] *= scale_x72        results.pred_keypoints[:, :, 1] *= scale_y73 74    return results75 76 77def sem_seg_postprocess(result, img_size, output_height, output_width):78    """79    Return semantic segmentation predictions in the original resolution.80 81    The input images are often resized when entering semantic segmentor. Moreover, in same82    cases, they also padded inside segmentor to be divisible by maximum network stride.83    As a result, we often need the predictions of the segmentor in a different84    resolution from its inputs.85 86    Args:87        result (Tensor): semantic segmentation prediction logits. A tensor of shape (C, H, W),88            where C is the number of classes, and H, W are the height and width of the prediction.89        img_size (tuple): image size that segmentor is taking as input.90        output_height, output_width: the desired output resolution.91 92    Returns:93        semantic segmentation prediction (Tensor): A tensor of the shape94            (C, output_height, output_width) that contains per-pixel soft predictions.95    """96    result = result[:, : img_size[0], : img_size[1]].expand(1, -1, -1, -1)97    result = F.interpolate(98        result, size=(output_height, output_width), mode="bilinear", align_corners=False99    )[0]100    return result101