cheng
265 lines
9.0 KiB
Python
265 lines
9.0 KiB
Python
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)
|