| import cv2 |
| import numpy as np |
| from core.config import NUM_CLASSES, COLORS |
|
|
| def render_overlay(img, info): |
| canvas = img.copy() |
| overlay = np.zeros_like(img) |
| |
| for cid in range(NUM_CLASSES): |
| mask = info[cid].get('mask_corrected', info[cid]['mask_original']) |
| overlay[mask] = COLORS[cid] |
| |
| canvas = cv2.addWeighted(canvas, 1.0, overlay, 0.4, 0) |
| |
| |
| for cid in [0, 1]: |
| skel = info[cid].get('skeleton') |
| if skel is not None and skel.sum() > 0: |
| thick = cv2.dilate(skel.astype(np.uint8), np.ones((3, 3), np.uint8)) |
| canvas[thick > 0] = (255, 255, 255) if cid == 0 else (255, 200, 100) |
| |
| return canvas |
