xdecoder/Instruct-X-Decoder

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 27    Returns:28        Instances: the resized output from the model, based on the output resolution29    """30    if isinstance(output_width, torch.Tensor):31        # This shape might (but not necessarily) be tensors during tracing.32        # Converts integer tensors to float temporaries to ensure true33        # division is performed when computing scale_x and scale_y.34        output_width_tmp = output_width.float()35        output_height_tmp = output_height.float()36        new_size = torch.stack([output_height, output_width])37    else:38        new_size = (output_height, output_width)39        output_width_tmp = output_width40        output_height_tmp = output_height41 42    scale_x, scale_y = (43        output_width_tmp / results.image_size[1],44        output_height_tmp / results.image_size[0],45    )46    results = Instances(new_size, **results.get_fields())47 48    if results.has("pred_boxes"):49        output_boxes = results.pred_boxes50    elif results.has("proposal_boxes"):51        output_boxes = results.proposal_boxes52    else:53        output_boxes = None54    assert output_boxes is not None, "Predictions must contain boxes!"55 56    output_boxes.scale(scale_x, scale_y)57    output_boxes.clip(results.image_size)58 59    results = results[output_boxes.nonempty()]60 61    if results.has("pred_masks"):62        if isinstance(results.pred_masks, ROIMasks):63            roi_masks = results.pred_masks64        else:65            # pred_masks is a tensor of shape (N, 1, M, M)66            roi_masks = ROIMasks(results.pred_masks[:, 0, :, :])67        results.pred_masks = roi_masks.to_bitmasks(68            results.pred_boxes, output_height, output_width, mask_threshold69        ).tensor  # TODO return ROIMasks/BitMask object in the future70 71    if results.has("pred_keypoints"):72        results.pred_keypoints[:, :, 0] *= scale_x73        results.pred_keypoints[:, :, 1] *= scale_y74 75    return results76 77def bbox_postprocess(result, input_size, img_size, output_height, output_width):78    """79    result: [xc,yc,w,h] range [0,1] to [x1,y1,x2,y2] range [0,w], [0,h]80    """81    if result is None:82        return None83    84    scale = torch.tensor([input_size[1], input_size[0], input_size[1], input_size[0]])[None,:].to(result.device)85    result = result.sigmoid() * scale86    x1,y1,x2,y2 = result[:,0] - result[:,2]/2, result[:,1] - result[:,3]/2, result[:,0] + result[:,2]/2, result[:,1] + result[:,3]/287    h,w = img_size88 89    x1 = x1.clamp(min=0, max=w)90    y1 = y1.clamp(min=0, max=h)91    x2 = x2.clamp(min=0, max=w)92    y2 = y2.clamp(min=0, max=h)93 94    box = torch.stack([x1,y1,x2,y2]).permute(1,0)95    scale = torch.tensor([output_width/w, output_height/h, output_width/w, output_height/h])[None,:].to(result.device)96    box = box*scale97    return box98 99def sem_seg_postprocess(result, img_size, output_height, output_width):100    """101    Return semantic segmentation predictions in the original resolution.102 103    The input images are often resized when entering semantic segmentor. Moreover, in same104    cases, they also padded inside segmentor to be divisible by maximum network stride.105    As a result, we often need the predictions of the segmentor in a different106    resolution from its inputs.107 108    Args:109        result (Tensor): semantic segmentation prediction logits. A tensor of shape (C, H, W),110            where C is the number of classes, and H, W are the height and width of the prediction.111        img_size (tuple): image size that segmentor is taking as input.112        output_height, output_width: the desired output resolution.113 114    Returns:115        semantic segmentation prediction (Tensor): A tensor of the shape116            (C, output_height, output_width) that contains per-pixel soft predictions.117    """118    result = result[:, : img_size[0], : img_size[1]].expand(1, -1, -1, -1)119    result = F.interpolate(120        result, size=(output_height, output_width), mode="bilinear", align_corners=False121    )[0]122    return result123