| 211 | self.nms = nms_cfg |
| 212 | |
| 213 | def __call__( |
| 214 | self, predict, rev_tensor: Optional[Tensor] = None, image_size: Optional[List[int]] = None |
| 215 | ) -> List[Tensor]: |
| 216 | if image_size is not None: |
| 217 | self.converter.update(image_size) |
| 218 | prediction = self.converter(predict["Main"]) |
| 219 | pred_class, _, pred_bbox = prediction[:3] |
| 220 | pred_conf = prediction[3] if len(prediction) == 4 else None |
| 221 | if rev_tensor is not None: |
| 222 | pred_bbox = (pred_bbox - rev_tensor[:, None, 1:]) / rev_tensor[:, 0:1, None] |
| 223 | pred_bbox = bbox_nms(pred_class, pred_bbox, self.nms, pred_conf) |
| 224 | return pred_bbox |
| 225 | |
| 226 | |
| 227 | def collect_prediction(predict_json: List, local_rank: int) -> List: |