(annotations, text, img_path, device, wider=False, threshold=0.9)
| 416 | |
| 417 | |
| 418 | def text_prompt(annotations, text, img_path, device, wider=False, threshold=0.9): |
| 419 | cropped_boxes, cropped_images, not_crop, origin_id, annotations_ = crop_image( |
| 420 | annotations, img_path |
| 421 | ) |
| 422 | clip_model, preprocess = clip.load("ViT-B/32", device=device) |
| 423 | scores = retriev( |
| 424 | clip_model, preprocess, cropped_boxes, text, device=device |
| 425 | ) |
| 426 | max_idx = scores.argsort() |
| 427 | max_idx = max_idx[-1] |
| 428 | max_idx = origin_id[int(max_idx)] |
| 429 | |
| 430 | # find the biggest mask which contains the mask with max score |
| 431 | if wider: |
| 432 | mask0 = annotations_[max_idx]["segmentation"] |
| 433 | area0 = np.sum(mask0) |
| 434 | areas = [(i, np.sum(mask["segmentation"])) for i, mask in enumerate(annotations_) if i in origin_id] |
| 435 | areas = sorted(areas, key=lambda area: area[1], reverse=True) |
| 436 | indices = [area[0] for area in areas] |
| 437 | for index in indices: |
| 438 | if index == max_idx or np.sum(annotations_[index]["segmentation"] & mask0) / area0 > threshold: |
| 439 | max_idx = index |
| 440 | break |
| 441 | |
| 442 | return annotations_[max_idx]["segmentation"], max_idx |
no test coverage detected