| from __future__ import annotations |
|
|
| import numpy as np |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| from src.config.endpoints import BBOX_NORMALIZATION_SCALE |
| from src.schemas.detection import LocalizedFinding |
|
|
|
|
| def draw_view_localizations( |
| image: Image.Image, |
| image_index: int, |
| findings: list[LocalizedFinding], |
| ) -> Image.Image: |
| annotated = image.copy() |
| draw = ImageDraw.Draw(annotated) |
| font = ImageFont.load_default() |
|
|
| width, height = annotated.size |
|
|
| for finding in findings: |
| for view_localization in finding.localizations: |
| if view_localization.image_index != image_index: |
| continue |
|
|
| for localization in view_localization.boxes: |
| y0, x0, y1, x1 = localization.box_2d |
|
|
| left = int(x0 / BBOX_NORMALIZATION_SCALE * width) |
| top = int(y0 / BBOX_NORMALIZATION_SCALE * height) |
| right = int(x1 / BBOX_NORMALIZATION_SCALE * width) |
| bottom = int(y1 / BBOX_NORMALIZATION_SCALE * height) |
|
|
| draw.rectangle( |
| [(left, top), (right, bottom)], |
| outline="red", |
| width=max(2, width // 256), |
| ) |
|
|
| label = localization.label |
| text_bbox = draw.textbbox((left, top), label, font=font) |
| text_width = text_bbox[2] - text_bbox[0] |
| text_height = text_bbox[3] - text_bbox[1] |
| label_top = max(0, top - text_height - 6) |
|
|
| draw.rectangle( |
| [ |
| (left, label_top), |
| (left + text_width + 6, label_top + text_height + 6), |
| ], |
| fill="black", |
| ) |
| draw.text( |
| (left + 3, label_top + 3), |
| label, |
| fill="white", |
| font=font, |
| ) |
|
|
| return annotated |
|
|
|
|
| def draw_segmentation_overlay( |
| base_images: list[Image.Image], |
| masks: list, |
| ) -> list[Image.Image]: |
| """Overlay polygon segmentation masks onto images, matching the MedSAM notebook style. |
| |
| Draws a semi-transparent yellow polygon fill + outline for each successful |
| MaskResult. Falls back to a pixel-level mask blend when no polygon is available. |
| Also draws the bounding box that prompted the segmentation. |
| """ |
| result = [img.convert("RGBA") for img in base_images] |
|
|
| for mask_result in masks: |
| if mask_result.status != "success": |
| continue |
| idx = mask_result.image_index |
| if idx >= len(result): |
| continue |
|
|
| img = result[idx] |
| W, H = img.size |
| overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0)) |
| draw = ImageDraw.Draw(overlay) |
|
|
| if mask_result.polygon and len(mask_result.polygon) >= 3: |
| pts = [(x, y) for x, y in mask_result.polygon] |
| |
| draw.polygon(pts, fill=(251, 252, 30, 153), outline=(251, 220, 0, 230)) |
| elif mask_result.mask is not None: |
| |
| base_np = np.array(img, dtype=np.float32) |
| mask = mask_result.mask |
| if mask.shape != (H, W): |
| mask_pil = Image.fromarray((mask * 255).astype(np.uint8)).resize( |
| (W, H), Image.NEAREST |
| ) |
| mask = (np.array(mask_pil) > 127).astype(np.uint8) |
| yellow = np.array([251, 252, 30, 255], dtype=np.float32) |
| alpha = 0.6 |
| mask_bool = mask.astype(bool) |
| base_np[mask_bool, :3] = base_np[mask_bool, :3] * (1 - alpha) + yellow[:3] * alpha |
| img = Image.fromarray(base_np.clip(0, 255).astype(np.uint8), mode="RGBA") |
| overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0)) |
| draw = ImageDraw.Draw(overlay) |
|
|
| |
| y0_n, x0_n, y1_n, x1_n = mask_result.bbox_normalized |
| x0 = int(x0_n / BBOX_NORMALIZATION_SCALE * W) |
| y0 = int(y0_n / BBOX_NORMALIZATION_SCALE * H) |
| x1 = int(x1_n / BBOX_NORMALIZATION_SCALE * W) |
| y1 = int(y1_n / BBOX_NORMALIZATION_SCALE * H) |
| lw = max(2, W // 256) |
| draw.rectangle([(x0, y0), (x1, y1)], outline=(0, 100, 255, 255), width=lw) |
|
|
| result[idx] = Image.alpha_composite(img, overlay) |
|
|
| return [img.convert("RGB") for img in result] |
|
|
