from __future__ import annotations import json from datetime import datetime from pathlib import Path import cv2 import numpy as np WINDOW_NAME = "ROI Polygon Editor" def to_uint8_image(image): array = np.asarray(image) if array.dtype == np.uint8: return array array = np.nan_to_num(array.astype(np.float32), copy=False) if array.size > 0 and float(np.max(array)) <= 1.0: array = array * 255.0 return np.clip(array, 0, 255).astype(np.uint8) def normalize_editor_frame(frame): if not isinstance(frame, np.ndarray): raise TypeError("frame must be a numpy ndarray with shape HxW or HxWxC") frame = to_uint8_image(frame) if frame.ndim == 2: return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR) if frame.ndim != 3: raise ValueError("frame must have shape HxW or HxWxC") channels = frame.shape[2] if channels == 1: return cv2.cvtColor(frame[:, :, 0], cv2.COLOR_GRAY2BGR) if channels == 3: return np.ascontiguousarray(frame) if channels == 4: return cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR) raise ValueError("frame channel count must be 1, 3, or 4") def polygon_to_mask(points, image_size): width, height = int(image_size[0]), int(image_size[1]) mask = np.zeros((height, width), dtype=np.uint8) polygon = np.asarray(points, dtype=np.int32).reshape(-1, 1, 2) cv2.fillPoly(mask, [polygon], 255) return mask def save_polygon_roi_config(path, points, image_size, mask_path=None, metadata=None): path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) data = { "version": "1.0", "type": "polygon", "image_size": [int(image_size[0]), int(image_size[1])], "points": [[int(x), int(y)] for x, y in points], "created_at": datetime.now().isoformat(timespec="seconds"), } if mask_path is not None: data["mask_path"] = str(mask_path) if metadata: data["metadata"] = metadata with path.open("w", encoding="utf-8") as f: json.dump(data, f, indent=2, ensure_ascii=False) return data def load_roi_config(path): path = Path(path) with path.open("r", encoding="utf-8") as f: data = json.load(f) data["_config_path"] = str(path) return data def scale_points(points, source_size, target_size): source_w, source_h = int(source_size[0]), int(source_size[1]) target_w, target_h = int(target_size[0]), int(target_size[1]) pts = np.asarray(points, dtype=np.float32).reshape(-1, 2) if source_w > 0 and source_h > 0 and (source_w, source_h) != (target_w, target_h): pts[:, 0] *= target_w / float(source_w) pts[:, 1] *= target_h / float(source_h) pts[:, 0] = np.clip(pts[:, 0], 0, max(target_w - 1, 0)) pts[:, 1] = np.clip(pts[:, 1], 0, max(target_h - 1, 0)) return np.rint(pts).astype(np.int32) def normalize_mask(mask, shape): if mask is None: return np.ones(shape, dtype=np.uint8) * 255 mask = np.asarray(mask) if mask.shape[:2] != shape: mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST) return np.where(mask > 0, 255, 0).astype(np.uint8) def load_roi_mask(config_path, frame_shape): config_path = Path(config_path) data = load_roi_config(config_path) height, width = frame_shape[:2] points = data.get("points") or data.get("polygon") or data.get("roi_polygon") if points: source_size = data.get("image_size", [width, height]) polygon = scale_points(points, source_size, (width, height)).reshape(-1, 1, 2) output = np.zeros((height, width), dtype=np.uint8) cv2.fillPoly(output, [polygon], 255) return output mask_path = data.get("mask_path") or data.get("roi_mask") if mask_path: mask_path = Path(mask_path) if not mask_path.is_absolute(): mask_path = config_path.parent / mask_path loaded = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) if loaded is None: raise FileNotFoundError(f"Failed to read ROI mask: {mask_path}") return normalize_mask(loaded, (height, width)) raise ValueError(f"Unsupported ROI config: {config_path}") class RoiPolygonEditor: def __init__(self, frame): self.frame = normalize_editor_frame(frame) self.points = [] def set_frame(self, frame): self.frame = normalize_editor_frame(frame) self.points = [] def on_mouse(self, event, x, y, flags, param): if event == cv2.EVENT_LBUTTONDOWN: self.points.append((int(x), int(y))) elif event == cv2.EVENT_RBUTTONDOWN and self.points: self.points.pop() def clear(self): self.points = [] def undo(self): if self.points: self.points.pop() def render(self): view = self.frame.copy() if len(self.points) >= 3: overlay = view.copy() polygon = np.asarray(self.points, dtype=np.int32).reshape(-1, 1, 2) cv2.fillPoly(overlay, [polygon], (0, 128, 255)) view = cv2.addWeighted(overlay, 0.25, view, 0.75, 0.0) cv2.polylines(view, [polygon], True, (0, 255, 255), 2, cv2.LINE_AA) elif len(self.points) >= 2: cv2.polylines( view, [np.asarray(self.points, dtype=np.int32).reshape(-1, 1, 2)], False, (0, 255, 255), 2, cv2.LINE_AA, ) for idx, (x, y) in enumerate(self.points): cv2.circle(view, (x, y), 4, (0, 255, 0), -1, cv2.LINE_AA) cv2.putText( view, str(idx + 1), (x + 6, y - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 255, 0), 1, cv2.LINE_AA, ) lines = [ "LClick add RClick/U undo C clear R recapture S/Enter save Q/Esc quit", f"points={len(self.points)} resolution={view.shape[1]}x{view.shape[0]}", ] for i, text in enumerate(lines): y = 24 + i * 24 cv2.putText(view, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3, cv2.LINE_AA) cv2.putText(view, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA) return view def build_roi_result(points, image_size, mask, config=None, config_path=None, mask_path=None): return { "points": [[int(x), int(y)] for x, y in points], "image_size": [int(image_size[0]), int(image_size[1])], "mask": mask, "config": config, "config_path": str(config_path) if config_path is not None else None, "mask_path": str(mask_path) if mask_path is not None else None, } def select_polygon_roi( frame, output=None, window_name=WINDOW_NAME, metadata=None, save_mask=True, destroy_window=True, recapture_callback=None, ): """Open an OpenCV window, draw a polygon ROI, and return points/mask/config info.""" editor = RoiPolygonEditor(frame) cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE) cv2.setMouseCallback(window_name, editor.on_mouse) try: while True: cv2.imshow(window_name, editor.render()) key = cv2.waitKey(16) & 0xFF if key in (27, ord("q")): return None if key in (8, ord("u")): editor.undo() elif key == ord("c"): editor.clear() elif key == ord("r") and recapture_callback is not None: editor.set_frame(recapture_callback()) elif key in (13, ord("s")): if len(editor.points) < 3: print("Need at least 3 points to save ROI.") continue image_size = (editor.frame.shape[1], editor.frame.shape[0]) mask = polygon_to_mask(editor.points, image_size) config = None config_path = None mask_path = None if output is not None: config_path = Path(output) config_path.parent.mkdir(parents=True, exist_ok=True) saved_mask = None if save_mask: mask_path = config_path.with_name(config_path.stem + "_mask.png") cv2.imwrite(str(mask_path), mask) saved_mask = mask_path.name config = save_polygon_roi_config( config_path, editor.points, image_size, mask_path=saved_mask, metadata=metadata, ) return build_roi_result( editor.points, image_size, mask, config=config, config_path=config_path, mask_path=mask_path, ) finally: if destroy_window: cv2.destroyWindow(window_name)