added dd06 test support, organized 2d viewer and 3d viewer code organization.
This commit is contained in:
ykcheungab
+284
-64
@@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
@@ -8,7 +10,8 @@ import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
WINDOW_NAME = "ROI Polygon Editor"
|
||||
WINDOW_NAME = "ROI Editor"
|
||||
POLYGON_WINDOW_NAME = "ROI Polygon Editor"
|
||||
|
||||
|
||||
def to_uint8_image(image):
|
||||
@@ -49,33 +52,6 @@ def polygon_to_mask(points, image_size):
|
||||
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])
|
||||
@@ -97,20 +73,96 @@ def normalize_mask(mask, shape):
|
||||
return np.where(mask > 0, 255, 0).astype(np.uint8)
|
||||
|
||||
|
||||
def build_roi_payload(image_size, points, mask_path=None, metadata=None):
|
||||
"""Build the shared ``roi`` object written into sensor / FFC JSON configs."""
|
||||
points = np.asarray(points, dtype=np.int32).reshape(-1, 2)
|
||||
if len(points) < 3:
|
||||
raise ValueError("polygon ROI requires at least 3 points")
|
||||
payload = {
|
||||
"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:
|
||||
payload["mask_path"] = str(mask_path)
|
||||
if metadata:
|
||||
payload["metadata"] = metadata
|
||||
return payload
|
||||
|
||||
|
||||
def extract_roi_section(data):
|
||||
"""Return the ROI dict from either a nested ``roi`` field or a legacy top-level ROI file."""
|
||||
if not isinstance(data, dict):
|
||||
raise TypeError("ROI config must be a JSON object")
|
||||
if isinstance(data.get("roi"), dict):
|
||||
return data["roi"]
|
||||
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 save_polygon_roi_config(path, points, image_size, mask_path=None, metadata=None):
|
||||
"""Backward-compatible standalone polygon ROI JSON writer."""
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
data = build_roi_payload(
|
||||
image_size,
|
||||
points,
|
||||
mask_path=mask_path,
|
||||
metadata=metadata,
|
||||
)
|
||||
data["version"] = "1.0"
|
||||
with path.open("w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
return data
|
||||
|
||||
|
||||
def upsert_roi_in_config(path, roi_payload, *, create_if_missing=True):
|
||||
"""Read-modify-write ``roi`` into an existing sensor config, preserving other fields."""
|
||||
path = Path(path)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if path.exists():
|
||||
with path.open("r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError(f"Config must be a JSON object: {path}")
|
||||
elif create_if_missing:
|
||||
data = {"version": "1.0"}
|
||||
else:
|
||||
raise FileNotFoundError(f"Config not found: {path}")
|
||||
|
||||
data["roi"] = roi_payload
|
||||
with path.open("w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
f.write("\n")
|
||||
return data
|
||||
|
||||
|
||||
def load_roi_mask(config_path, frame_shape):
|
||||
config_path = Path(config_path)
|
||||
data = load_roi_config(config_path)
|
||||
roi = extract_roi_section(data)
|
||||
|
||||
height, width = frame_shape[:2]
|
||||
points = data.get("points") or data.get("polygon") or data.get("roi_polygon")
|
||||
target_size = (width, height)
|
||||
source_size = roi.get("image_size") or data.get("image_size") or [width, height]
|
||||
points = roi.get("points")
|
||||
points = roi.get("points") or roi.get("polygon") or roi.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)
|
||||
polygon = scale_points(points, source_size, target_size).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")
|
||||
mask_path = roi.get("mask_path") or data.get("mask_path") or data.get("roi_mask")
|
||||
if mask_path:
|
||||
mask_path = Path(mask_path)
|
||||
if not mask_path.is_absolute():
|
||||
@@ -123,6 +175,54 @@ def load_roi_mask(config_path, frame_shape):
|
||||
raise ValueError(f"Unsupported ROI config: {config_path}")
|
||||
|
||||
|
||||
def resolve_video_source(video):
|
||||
"""Accept camera index (int/str digit) or filesystem path."""
|
||||
if isinstance(video, int):
|
||||
return video
|
||||
text = str(video).strip()
|
||||
if text.isdigit():
|
||||
return int(text)
|
||||
return text
|
||||
|
||||
|
||||
def open_video_capture(video):
|
||||
source = resolve_video_source(video)
|
||||
# Avoid MSMF HW transform stalls on Windows camera capture.
|
||||
os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0")
|
||||
if isinstance(source, int):
|
||||
cap = cv2.VideoCapture(source, cv2.CAP_DSHOW)
|
||||
if not cap.isOpened():
|
||||
cap.release()
|
||||
cap = cv2.VideoCapture(source)
|
||||
else:
|
||||
path = Path(source)
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"Video file not found: {path}")
|
||||
cap = cv2.VideoCapture(str(path))
|
||||
if not cap.isOpened():
|
||||
raise RuntimeError(f"Failed to open video source: {video}")
|
||||
return cap
|
||||
|
||||
|
||||
def capture_frame(cap, *, max_attempts=60):
|
||||
"""Grab the next readable frame from an opened capture."""
|
||||
last_error = None
|
||||
for _ in range(max_attempts):
|
||||
ok, frame = cap.read()
|
||||
if ok and frame is not None and getattr(frame, "size", 0) > 0:
|
||||
return normalize_editor_frame(frame)
|
||||
last_error = "empty frame" if ok else "read failed"
|
||||
raise RuntimeError(f"Failed to capture frame from video source ({last_error})")
|
||||
|
||||
|
||||
def capture_frame_from_source(video):
|
||||
cap = open_video_capture(video)
|
||||
try:
|
||||
return capture_frame(cap)
|
||||
finally:
|
||||
cap.release()
|
||||
|
||||
|
||||
class RoiPolygonEditor:
|
||||
def __init__(self, frame):
|
||||
self.frame = normalize_editor_frame(frame)
|
||||
@@ -187,25 +287,85 @@ class RoiPolygonEditor:
|
||||
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],
|
||||
def build_roi_result(
|
||||
points,
|
||||
image_size,
|
||||
mask,
|
||||
config=None,
|
||||
config_path=None,
|
||||
mask_path=None,
|
||||
):
|
||||
result = {
|
||||
"type": "polygon",
|
||||
"points": [[int(x), int(y)] for x, y in points] if points is not None else None,
|
||||
"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,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def _save_roi_outputs(
|
||||
*,
|
||||
image_size,
|
||||
output,
|
||||
points,
|
||||
metadata=None,
|
||||
save_mask=True,
|
||||
nest_in_config=True,
|
||||
):
|
||||
mask = polygon_to_mask(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
|
||||
|
||||
roi_payload = build_roi_payload(
|
||||
image_size,
|
||||
points,
|
||||
mask_path=saved_mask,
|
||||
metadata=metadata,
|
||||
)
|
||||
if nest_in_config:
|
||||
config = upsert_roi_in_config(config_path, roi_payload)
|
||||
else:
|
||||
# Standalone ROI file (legacy FFC polygon style).
|
||||
standalone = dict(roi_payload)
|
||||
standalone["version"] = "1.0"
|
||||
with config_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(standalone, f, indent=2, ensure_ascii=False)
|
||||
f.write("\n")
|
||||
config = standalone
|
||||
|
||||
return build_roi_result(
|
||||
points,
|
||||
image_size,
|
||||
mask,
|
||||
config=config,
|
||||
config_path=config_path,
|
||||
mask_path=mask_path,
|
||||
)
|
||||
|
||||
|
||||
def select_polygon_roi(
|
||||
frame,
|
||||
output=None,
|
||||
window_name=WINDOW_NAME,
|
||||
window_name=POLYGON_WINDOW_NAME,
|
||||
metadata=None,
|
||||
save_mask=True,
|
||||
destroy_window=True,
|
||||
recapture_callback=None,
|
||||
nest_in_config=False,
|
||||
):
|
||||
"""Open an OpenCV window, draw a polygon ROI, and return points/mask/config info."""
|
||||
editor = RoiPolygonEditor(frame)
|
||||
@@ -230,35 +390,95 @@ def select_polygon_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,
|
||||
return _save_roi_outputs(
|
||||
image_size=image_size,
|
||||
output=output,
|
||||
points=editor.points,
|
||||
metadata=metadata,
|
||||
save_mask=save_mask,
|
||||
nest_in_config=nest_in_config,
|
||||
)
|
||||
finally:
|
||||
if destroy_window:
|
||||
cv2.destroyWindow(window_name)
|
||||
|
||||
|
||||
select_roi = select_polygon_roi
|
||||
|
||||
|
||||
def parse_args(argv=None):
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Capture a frame from a camera/video and save a polygon ROI to JSON."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video",
|
||||
"-v",
|
||||
required=True,
|
||||
help="Camera index (e.g. 0) or path to a video file",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--config",
|
||||
"-c",
|
||||
required=True,
|
||||
help="Output JSON config path (existing sensor configs are updated in place)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--frame-output",
|
||||
default=None,
|
||||
help="Optional path to save the captured reference frame (PNG/JPG)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-mask",
|
||||
action="store_true",
|
||||
help="Do not write a companion *_mask.png next to the config",
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv=None):
|
||||
args = parse_args(argv)
|
||||
video = resolve_video_source(args.video)
|
||||
cap = open_video_capture(video)
|
||||
|
||||
def recapture():
|
||||
# For file sources, rewind so R always gets a usable frame.
|
||||
if not isinstance(video, int):
|
||||
cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
||||
return capture_frame(cap)
|
||||
|
||||
try:
|
||||
frame = capture_frame(cap)
|
||||
if args.frame_output:
|
||||
frame_path = Path(args.frame_output)
|
||||
frame_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
cv2.imwrite(str(frame_path), frame)
|
||||
print(f"Saved reference frame: {frame_path}")
|
||||
|
||||
result = select_polygon_roi(
|
||||
frame,
|
||||
output=args.config,
|
||||
metadata={"source": str(args.video)},
|
||||
save_mask=not args.no_mask,
|
||||
recapture_callback=recapture,
|
||||
nest_in_config=True,
|
||||
)
|
||||
finally:
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
if result is None:
|
||||
print("ROI capture cancelled.")
|
||||
return 1
|
||||
|
||||
print(f"Saved ROI type={result['type']} image_size={result['image_size']}")
|
||||
if result.get("points") is not None:
|
||||
print(f"points={result['points']}")
|
||||
if result.get("config_path"):
|
||||
print(f"config={result['config_path']}")
|
||||
if result.get("mask_path"):
|
||||
print(f"mask={result['mask_path']}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
Reference in New Issue
Block a user