examples/research_projects/visual_bert/modeling_frcnn.py/ROIOutputs/inference
class ROIOutputs:    def inference(
        self,
        obj_logits,
        attr_logits,
        box_deltas,
        pred_boxes,
        features,
        sizes,
        scales=None,
    ):
        # only the pred boxes is the
        preds_per_image = [p.size(0) for p in pred_boxes]
        boxes_all = self._predict_boxes(pred_boxes, box_deltas, preds_per_image)
        obj_scores_all = self._predict_objs(obj_logits, preds_per_image)  # list of length N
        attr_probs_all, attrs_all = self._predict_attrs(attr_logits, preds_per_image)
        features = features.split(preds_per_image, dim=0)

        # fun for each image too, also I can experiment and do multiple images
        final_results = []
        zipped = zip(boxes_all, obj_scores_all, attr_probs_all, attrs_all, sizes)
        for i, (boxes, obj_scores, attr_probs, attrs, size) in enumerate(zipped):
            for nms_t in self.nms_thresh:
                outputs = do_nms(
                    boxes,
                    obj_scores,
                    size,
                    self.score_thresh,
                    nms_t,
                    self.min_detections,
                    self.max_detections,
                )
                if outputs is not None:
                    max_boxes, max_scores, classes, ids = outputs
                    break

            if scales is not None:
                scale_yx = scales[i]
                max_boxes[:, 0::2] *= scale_yx[1]
                max_boxes[:, 1::2] *= scale_yx[0]

            final_results.append(
                (
                    max_boxes,
                    classes,
                    max_scores,
                    attrs[ids],
                    attr_probs[ids],
                    features[i][ids],
                )
            )
        boxes, classes, class_probs, attrs, attr_probs, roi_features = map(list, zip(*final_results))
        return boxes, classes, class_probs, attrs, attr_probs, roi_features