Spaces:
Sleeping
Sleeping
File size: 2,465 Bytes
b87d940 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | """
Drawing helpers for pipeline visual outputs.
All functions operate on RGB numpy arrays (uint8) and return RGB numpy arrays.
They are stateless - no model or pipeline state is needed.
"""
import cv2
import numpy as np
# Detection overlay
def make_detection_overlay(
image_rgb : np.ndarray,
boxes : list[tuple],
label : str | None = None,
) -> np.ndarray:
"""
Draw YOLO bounding boxes on the image.
The box colour reflects the final classification label:
- MALIGNANT -> red (#DC3232)
- BENIGN -> green (#32B450)
- Unknown -> grey (#B4B4B4)
Parameters
image_rgb : np.ndarray (H, W, 3) uint8, RGB
boxes : list of (x1, y1, x2, y2, conf) tuples from NoduleDetector
label : "MALIGNANT", "BENIGN", or None (classification not yet done)
Returns
-------
np.ndarray (H, W, 3) uint8, RGB - image with bounding boxes drawn
"""
overlay = image_rgb.copy()
if label == "MALIGNANT":
color = (220, 50, 50)
elif label == "BENIGN":
color = (50, 180, 80)
else:
color = (180, 180, 180)
for (x1, y1, x2, y2, conf) in boxes:
cv2.rectangle(overlay, (x1, y1), (x2, y2), color, 2)
text = f"Nodule {conf:.0%}"
org = (x1, max(y1 - 8, 12))
cv2.putText(
overlay, text, org,
cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2,
cv2.LINE_AA,
)
return overlay
# Segmentation overlay
def make_segmentation_overlay(
image_rgb : np.ndarray,
mask : np.ndarray,
alpha : float = 0.40,
) -> np.ndarray:
"""
Blend the UNet++ binary mask as a semi-transparent orange overlay.
A solid contour is drawn on top for clearer boundary visualisation.
Parameters
image_rgb : np.ndarray (H, W, 3) uint8, RGB
mask : np.ndarray (H, W) uint8, 0/255
alpha : float — opacity of the colour overlay (default 0.40)
Returns
np.ndarray (H, W, 3) uint8, RGB - blended image
"""
color_mask = np.zeros_like(image_rgb)
color_mask[mask > 0] = [255, 140, 0] # orange
blended = cv2.addWeighted(image_rgb, 1.0,
color_mask, alpha, 0)
# Draw contour outline for precise boundary visualisation
contours, _ = cv2.findContours(
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
)
cv2.drawContours(blended, contours, -1, (255, 140, 0), 2)
return blended
|