Resize the output instances. The input images are often resized when entering an object detector. As a result, we often need the outputs of the detector in a different resolution from its inputs. This function will resize the raw outputs of an R-CNN detector to produce outp
(
results: Instances, output_height: int, output_width: int, mask_threshold: float = 0.5
)
| 9 | |
| 10 | # perhaps should rename to "resize_instance" |
| 11 | def detector_postprocess( |
| 12 | results: Instances, output_height: int, output_width: int, mask_threshold: float = 0.5 |
| 13 | ): |
| 14 | """ |
| 15 | Resize the output instances. |
| 16 | The input images are often resized when entering an object detector. |
| 17 | As a result, we often need the outputs of the detector in a different |
| 18 | resolution from its inputs. |
| 19 | |
| 20 | This function will resize the raw outputs of an R-CNN detector |
| 21 | to produce outputs according to the desired output resolution. |
| 22 | |
| 23 | Args: |
| 24 | results (Instances): the raw outputs from the detector. |
| 25 | `results.image_size` contains the input image resolution the detector sees. |
| 26 | This object might be modified in-place. |
| 27 | output_height, output_width: the desired output resolution. |
| 28 | |
| 29 | Returns: |
| 30 | Instances: the resized output from the model, based on the output resolution |
| 31 | """ |
| 32 | # Change to 'if is_tracing' after PT1.7 |
| 33 | if isinstance(output_height, torch.Tensor): |
| 34 | # Converts integer tensors to float temporaries to ensure true |
| 35 | # division is performed when computing scale_x and scale_y. |
| 36 | output_width_tmp = output_width.float() |
| 37 | output_height_tmp = output_height.float() |
| 38 | new_size = torch.stack([output_height, output_width]) |
| 39 | else: |
| 40 | new_size = (output_height, output_width) |
| 41 | output_width_tmp = output_width |
| 42 | output_height_tmp = output_height |
| 43 | |
| 44 | scale_x, scale_y = ( |
| 45 | output_width_tmp / results.image_size[1], |
| 46 | output_height_tmp / results.image_size[0], |
| 47 | ) |
| 48 | results = Instances(new_size, **results.get_fields()) |
| 49 | |
| 50 | if results.has("pred_boxes"): |
| 51 | output_boxes = results.pred_boxes |
| 52 | elif results.has("proposal_boxes"): |
| 53 | output_boxes = results.proposal_boxes |
| 54 | else: |
| 55 | output_boxes = None |
| 56 | assert output_boxes is not None, "Predictions must contain boxes!" |
| 57 | |
| 58 | output_boxes.scale(scale_x, scale_y) |
| 59 | output_boxes.clip(results.image_size) |
| 60 | |
| 61 | results = results[output_boxes.nonempty()] |
| 62 | |
| 63 | if results.has("pred_masks"): |
| 64 | results.pred_masks = retry_if_cuda_oom(paste_masks_in_image)( |
| 65 | results.pred_masks[:, 0, :, :], # N, 1, M, M |
| 66 | results.pred_boxes, |
| 67 | results.image_size, |
| 68 | threshold=mask_threshold, |