painter.py

"""Mask and click overlays for the first-frame preview."""

import cv2
import numpy as np

MASK_RGB = np.array([37, 99, 235], np.float32)   # blue
POS_RGB = (16, 185, 129)                         # green
NEG_RGB = (239, 68, 68)                          # red


def paint_mask(frame, mask, alpha=0.5):
    out = frame.astype(np.float32).copy()
    if mask is None:
        return frame.copy()
    m = mask > 127
    out[m] = out[m] * (1 - alpha) + MASK_RGB * alpha
    out = out.astype(np.uint8)
    contours, _ = cv2.findContours(m.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(out, contours, -1, (255, 255, 255), max(2, frame.shape[1] // 600))
    return out


def paint_points(frame, points, labels):
    out = frame.copy()
    r = max(6, frame.shape[1] // 120)
    for (x, y), l in zip(points, labels):
        cv2.circle(out, (int(x), int(y)), r, POS_RGB if l == 1 else NEG_RGB, -1, cv2.LINE_AA)
        cv2.circle(out, (int(x), int(y)), r, (255, 255, 255), max(2, r // 4), cv2.LINE_AA)
    return out