| 364 | |
| 365 | |
| 366 | def box_prompt(masks, bbox, target_height, target_width): |
| 367 | h = masks.shape[1] |
| 368 | w = masks.shape[2] |
| 369 | if h != target_height or w != target_width: |
| 370 | bbox = [ |
| 371 | int(bbox[0] * w / target_width), |
| 372 | int(bbox[1] * h / target_height), |
| 373 | int(bbox[2] * w / target_width), |
| 374 | int(bbox[3] * h / target_height), |
| 375 | ] |
| 376 | bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0 |
| 377 | bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0 |
| 378 | bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w |
| 379 | bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h |
| 380 | |
| 381 | # IoUs = torch.zeros(len(masks), dtype=torch.float32) |
| 382 | bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0]) |
| 383 | |
| 384 | masks_area = torch.sum(masks[:, bbox[1] : bbox[3], bbox[0] : bbox[2]], dim=(1, 2)) |
| 385 | orig_masks_area = torch.sum(masks, dim=(1, 2)) |
| 386 | |
| 387 | union = bbox_area + orig_masks_area - masks_area |
| 388 | IoUs = masks_area / union |
| 389 | max_iou_index = torch.argmax(IoUs) |
| 390 | |
| 391 | return masks[max_iou_index].cpu().numpy(), max_iou_index |
| 392 | |
| 393 | |
| 394 | def point_prompt(masks, points, point_label, target_height, target_width): # numpy 处理 |