489 lines
19 KiB
Python
489 lines
19 KiB
Python
"""Per-frame door-mask loading and YOLO/RTMW-only 2-D risk scoring."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import re
|
|
import time
|
|
from collections import Counter
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
RASTER_SUFFIXES = {".png", ".bmp", ".tif", ".tiff"}
|
|
DEFAULT_DOOR_LABELS = ("door", "车门")
|
|
|
|
|
|
@dataclass
|
|
class LoadedDoorMask:
|
|
status: str
|
|
annotation_index: int
|
|
path: Path | None
|
|
mask: np.ndarray | None
|
|
source_format: str | None
|
|
|
|
def metadata(self) -> dict:
|
|
return {
|
|
"status": self.status,
|
|
"annotation_index": self.annotation_index,
|
|
"path": str(self.path.resolve()) if self.path is not None else None,
|
|
"format": self.source_format,
|
|
"area_px": int(np.count_nonzero(self.mask)) if self.mask is not None else 0,
|
|
}
|
|
|
|
|
|
def _frame_index_from_path(path: Path) -> int | None:
|
|
matches = re.findall(r"\d+", path.stem)
|
|
return int(matches[-1]) if matches else None
|
|
|
|
|
|
class DoorMaskSequence:
|
|
"""Index a directory of X-AnyLabeling JSON polygons or raster masks."""
|
|
|
|
def __init__(
|
|
self,
|
|
directory: Path | None = None,
|
|
mask_format: str = "auto",
|
|
labels: tuple[str, ...] = DEFAULT_DOOR_LABELS,
|
|
frame_offset: int = 0,
|
|
require_every_frame: bool = False,
|
|
constant_mask: Path | None = None,
|
|
):
|
|
if (directory is None) == (constant_mask is None):
|
|
raise ValueError("Provide exactly one of directory or constant_mask")
|
|
self.constant_path = (
|
|
Path(constant_mask).expanduser().resolve()
|
|
if constant_mask is not None
|
|
else None
|
|
)
|
|
self.directory = (
|
|
self.constant_path.parent if self.constant_path is not None else Path(directory)
|
|
)
|
|
if self.constant_path is not None and not self.constant_path.is_file():
|
|
raise ValueError(f"Constant door mask does not exist: {self.constant_path}")
|
|
if self.constant_path is None and not self.directory.is_dir():
|
|
raise ValueError(f"Door mask directory does not exist: {self.directory}")
|
|
self.labels = {label.casefold() for label in labels}
|
|
self.frame_offset = int(frame_offset)
|
|
self.require_every_frame = require_every_frame
|
|
|
|
if self.constant_path is not None:
|
|
if mask_format == "auto":
|
|
selected_format = (
|
|
"xanylabeling"
|
|
if self.constant_path.suffix.lower() == ".json"
|
|
else "raster"
|
|
)
|
|
else:
|
|
selected_format = mask_format
|
|
self.source_format = selected_format
|
|
self.paths = {0: self.constant_path}
|
|
return
|
|
|
|
json_paths = sorted(self.directory.rglob("*.json"))
|
|
raster_paths = sorted(
|
|
path
|
|
for path in self.directory.rglob("*")
|
|
if path.is_file() and path.suffix.lower() in RASTER_SUFFIXES
|
|
)
|
|
if mask_format == "auto":
|
|
selected_format = "xanylabeling" if json_paths else "raster"
|
|
else:
|
|
selected_format = mask_format
|
|
candidates = json_paths if selected_format == "xanylabeling" else raster_paths
|
|
if not candidates:
|
|
expected = "JSON" if selected_format == "xanylabeling" else "PNG/BMP/TIFF"
|
|
raise ValueError(f"No {expected} door masks found in {self.directory}")
|
|
|
|
self.source_format = selected_format
|
|
self.paths: dict[int, Path] = {}
|
|
for path in candidates:
|
|
index = _frame_index_from_path(path)
|
|
if index is None and selected_format == "xanylabeling":
|
|
data = json.loads(path.read_text(encoding="utf-8-sig"))
|
|
explicit = data.get("frame_index", data.get("frameIndex"))
|
|
index = int(explicit) if explicit is not None else None
|
|
if index is None:
|
|
continue
|
|
if index in self.paths:
|
|
raise ValueError(
|
|
f"Multiple door masks map to annotation frame {index}: "
|
|
f"{self.paths[index]} and {path}"
|
|
)
|
|
self.paths[index] = path
|
|
if not self.paths:
|
|
raise ValueError(
|
|
"Door mask filenames need a frame number, for example 000000.json"
|
|
)
|
|
|
|
def load(self, video_frame_index: int, frame_shape: tuple[int, int]) -> LoadedDoorMask:
|
|
if self.constant_path is not None:
|
|
annotation_index = video_frame_index
|
|
path = self.constant_path
|
|
else:
|
|
annotation_index = video_frame_index + self.frame_offset
|
|
path = self.paths.get(annotation_index)
|
|
if path is None:
|
|
if self.require_every_frame:
|
|
raise RuntimeError(
|
|
f"No door mask for video frame {video_frame_index} "
|
|
f"(expected annotation index {annotation_index})"
|
|
)
|
|
return LoadedDoorMask("missing", annotation_index, None, None, self.source_format)
|
|
|
|
if self.source_format == "xanylabeling":
|
|
mask = self._load_xanylabeling(path, frame_shape)
|
|
else:
|
|
mask = self._load_raster(path, frame_shape)
|
|
status = "loaded" if np.any(mask) else "empty"
|
|
return LoadedDoorMask(status, annotation_index, path, mask, self.source_format)
|
|
|
|
@staticmethod
|
|
def _load_raster(path: Path, frame_shape: tuple[int, int]) -> np.ndarray:
|
|
image = cv2.imread(str(path), cv2.IMREAD_UNCHANGED)
|
|
if image is None:
|
|
raise RuntimeError(f"Cannot read door mask: {path}")
|
|
if image.ndim == 3:
|
|
image = np.max(image[:, :, :3], axis=2)
|
|
height, width = frame_shape
|
|
if image.shape != (height, width):
|
|
image = cv2.resize(image, (width, height), interpolation=cv2.INTER_NEAREST)
|
|
return image > 0
|
|
|
|
def _load_xanylabeling(
|
|
self, path: Path, frame_shape: tuple[int, int]
|
|
) -> np.ndarray:
|
|
data = json.loads(path.read_text(encoding="utf-8-sig"))
|
|
height, width = frame_shape
|
|
source_width = float(data.get("imageWidth") or width)
|
|
source_height = float(data.get("imageHeight") or height)
|
|
scale = np.asarray(
|
|
[width / max(source_width, 1.0), height / max(source_height, 1.0)],
|
|
dtype=np.float32,
|
|
)
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
for shape in data.get("shapes", []):
|
|
if str(shape.get("label", "")).casefold() not in self.labels:
|
|
continue
|
|
points = np.asarray(shape.get("points", []), dtype=np.float32).reshape(-1, 2)
|
|
if len(points) < 2:
|
|
continue
|
|
points = np.rint(points * scale).astype(np.int32)
|
|
shape_type = str(shape.get("shape_type", "polygon")).lower()
|
|
if shape_type == "rectangle" and len(points) >= 2:
|
|
x1, y1 = points[0]
|
|
x2, y2 = points[1]
|
|
cv2.rectangle(mask, (x1, y1), (x2, y2), 1, -1)
|
|
elif shape_type == "circle" and len(points) >= 2:
|
|
radius = int(round(float(np.linalg.norm(points[1] - points[0]))))
|
|
cv2.circle(mask, tuple(points[0]), radius, 1, -1)
|
|
elif len(points) >= 3:
|
|
cv2.fillPoly(mask, [points], 1)
|
|
return mask > 0
|
|
|
|
|
|
def _fit_mask(mask: np.ndarray, frame_shape: tuple[int, int]) -> np.ndarray:
|
|
height, width = frame_shape
|
|
result = np.asarray(mask) > 0
|
|
if result.shape != (height, width):
|
|
result = cv2.resize(
|
|
result.astype(np.uint8), (width, height), interpolation=cv2.INTER_NEAREST
|
|
) > 0
|
|
return result
|
|
|
|
|
|
def _box_mask(box: list[float] | None, frame_shape: tuple[int, int]) -> np.ndarray | None:
|
|
if box is None or len(box) < 4:
|
|
return None
|
|
height, width = frame_shape
|
|
x1, y1, x2, y2 = [int(round(float(value))) for value in box[:4]]
|
|
x1, x2 = sorted((max(0, min(width - 1, x1)), max(0, min(width - 1, x2))))
|
|
y1, y2 = sorted((max(0, min(height - 1, y1)), max(0, min(height - 1, y2))))
|
|
if x2 <= x1 or y2 <= y1:
|
|
return None
|
|
mask = np.zeros((height, width), dtype=bool)
|
|
mask[y1 : y2 + 1, x1 : x2 + 1] = True
|
|
return mask
|
|
|
|
|
|
def _hand_mask(points: list, frame_shape: tuple[int, int]) -> np.ndarray | None:
|
|
hand = np.asarray(points, dtype=np.float32).reshape(-1, 2)
|
|
if len(hand) < 3 or not np.isfinite(hand).all():
|
|
return None
|
|
hull = cv2.convexHull(np.rint(hand).astype(np.int32))
|
|
mask = np.zeros(frame_shape, dtype=np.uint8)
|
|
cv2.fillConvexPoly(mask, hull, 1)
|
|
return mask > 0 if np.any(mask) else None
|
|
|
|
|
|
def _decision(distance_relative: float, overlap_px: int, warning: float, release: float) -> tuple[str, str]:
|
|
if overlap_px > 0:
|
|
return "danger", "DANGER_INTERSECTION"
|
|
if distance_relative <= warning:
|
|
return "warning", "WARNING_NEAR"
|
|
if distance_relative <= release:
|
|
return "watch", "WATCH"
|
|
return "clear", "CLEAR"
|
|
|
|
|
|
class DoorRiskEvaluator:
|
|
"""Score RTMW hand regions, with one YOLO body-box fallback per person."""
|
|
|
|
def __init__(
|
|
self,
|
|
sequence: DoorMaskSequence,
|
|
warning_relative: float = 0.30,
|
|
release_relative: float = 0.45,
|
|
):
|
|
if warning_relative < 0 or release_relative <= warning_relative:
|
|
raise ValueError("Risk thresholds require 0 <= warning < release")
|
|
self.sequence = sequence
|
|
self.warning_relative = float(warning_relative)
|
|
self.release_relative = float(release_relative)
|
|
self.mask_status_counts: Counter = Counter()
|
|
self.level_counts: Counter = Counter()
|
|
self.source_counts: Counter = Counter()
|
|
self.maximum_score = 0.0
|
|
self.timing_seconds: Counter = Counter()
|
|
self.timing_counts: Counter = Counter()
|
|
self._constant_cache_shape: tuple[int, int] | None = None
|
|
self._constant_loaded: LoadedDoorMask | None = None
|
|
self._constant_distance_map: np.ndarray | None = None
|
|
self.distance_map_cache_builds = 0
|
|
|
|
def _record_timing(self, name: str, started: float) -> None:
|
|
self.timing_seconds[name] += time.perf_counter() - started
|
|
self.timing_counts[name] += 1
|
|
|
|
def _measure_region(
|
|
self,
|
|
region: np.ndarray,
|
|
door_mask: np.ndarray,
|
|
distance_map: np.ndarray,
|
|
source: str,
|
|
track_id: int,
|
|
side: str,
|
|
state: str | None,
|
|
) -> dict:
|
|
area_px = int(np.count_nonzero(region))
|
|
overlap_px = int(np.count_nonzero(region & door_mask))
|
|
distances = distance_map[region]
|
|
distance_px = float(np.min(distances)) if len(distances) else math.inf
|
|
scale_px = max(math.sqrt(area_px), 1.0)
|
|
distance_relative = distance_px / scale_px
|
|
score = float(np.clip(1.0 - distance_relative / self.release_relative, 0.0, 1.0))
|
|
if overlap_px:
|
|
score = 1.0
|
|
level, decision = _decision(
|
|
distance_relative,
|
|
overlap_px,
|
|
self.warning_relative,
|
|
self.release_relative,
|
|
)
|
|
result = {
|
|
"status": "calculated",
|
|
"risk_score": round(score, 5),
|
|
"risk_level": level,
|
|
"decision": decision,
|
|
"source": source,
|
|
"track_id": int(track_id),
|
|
"side": side,
|
|
"state": state,
|
|
"region_area_px": area_px,
|
|
"door_overlap_px": overlap_px,
|
|
"region_overlap_ratio": round(overlap_px / max(area_px, 1), 6),
|
|
"min_distance_px": round(distance_px, 3),
|
|
"region_scale_px": round(scale_px, 3),
|
|
"distance_relative": round(distance_relative, 6),
|
|
}
|
|
self.level_counts[level] += 1
|
|
self.source_counts[source] += 1
|
|
self.maximum_score = max(self.maximum_score, score)
|
|
return result
|
|
|
|
@staticmethod
|
|
def _unavailable(source: str, track_id: int, side: str, status: str) -> dict:
|
|
return {
|
|
"status": status,
|
|
"risk_score": None,
|
|
"risk_level": "unavailable",
|
|
"decision": "EVIDENCE_UNAVAILABLE",
|
|
"source": source,
|
|
"track_id": int(track_id),
|
|
"side": side,
|
|
}
|
|
|
|
def evaluate(
|
|
self,
|
|
frame_index: int,
|
|
record: dict,
|
|
frame_shape: tuple[int, int],
|
|
) -> tuple[dict, np.ndarray | None]:
|
|
total_started = time.perf_counter()
|
|
stage_started = time.perf_counter()
|
|
use_constant_cache = self.sequence.constant_path is not None
|
|
if use_constant_cache and self._constant_cache_shape == frame_shape:
|
|
cached = self._constant_loaded
|
|
assert cached is not None
|
|
loaded = LoadedDoorMask(
|
|
cached.status,
|
|
frame_index,
|
|
cached.path,
|
|
cached.mask,
|
|
cached.source_format,
|
|
)
|
|
else:
|
|
loaded = self.sequence.load(frame_index, frame_shape)
|
|
if use_constant_cache:
|
|
self._constant_cache_shape = frame_shape
|
|
self._constant_loaded = loaded
|
|
self._constant_distance_map = None
|
|
self._record_timing("door_mask_load", stage_started)
|
|
self.mask_status_counts[loaded.status] += 1
|
|
mask_ready = loaded.status == "loaded" and loaded.mask is not None
|
|
door_mask = loaded.mask if mask_ready else None
|
|
distance_map = None
|
|
if door_mask is not None and use_constant_cache and self._constant_distance_map is not None:
|
|
distance_map = self._constant_distance_map
|
|
elif door_mask is not None:
|
|
stage_started = time.perf_counter()
|
|
distance_map = cv2.distanceTransform(
|
|
(~door_mask).astype(np.uint8), cv2.DIST_L2, cv2.DIST_MASK_PRECISE
|
|
)
|
|
self._record_timing("distance_transform", stage_started)
|
|
if use_constant_cache:
|
|
self._constant_distance_map = distance_map
|
|
self.distance_map_cache_builds += 1
|
|
|
|
stage_started = time.perf_counter()
|
|
regions: list[tuple[dict, np.ndarray, str, int, str, str | None]] = []
|
|
track_ids = [int(value) for value in record.get("track_ids", [])]
|
|
person_positions_with_hands: set[int] = set()
|
|
for hand in record.get("global_hands", []):
|
|
position = int(hand.get("person_id", 0)) - 1
|
|
track_id = track_ids[position] if 0 <= position < len(track_ids) else position + 1
|
|
side = str(hand.get("side", "unknown"))
|
|
region = _hand_mask(hand.get("keypoints", []), frame_shape)
|
|
if region is None:
|
|
continue
|
|
person_positions_with_hands.add(position)
|
|
regions.append(
|
|
(
|
|
hand,
|
|
region,
|
|
"RTMW_HAND_KEYPOINTS",
|
|
track_id,
|
|
side,
|
|
str(hand.get("state", "unknown")),
|
|
)
|
|
)
|
|
|
|
boxes = record.get("boxes", [])
|
|
body_fallbacks = []
|
|
for position, box in enumerate(boxes):
|
|
if position in person_positions_with_hands:
|
|
continue
|
|
region = _box_mask(box, frame_shape)
|
|
if region is None:
|
|
continue
|
|
track_id = track_ids[position] if position < len(track_ids) else position + 1
|
|
fallback = {
|
|
"track_id": track_id,
|
|
"side": "body",
|
|
"state": "rtmw_hand_unavailable",
|
|
"body_zone_box": box,
|
|
}
|
|
body_fallbacks.append(fallback)
|
|
regions.append(
|
|
(
|
|
fallback,
|
|
region,
|
|
"YOLO_BODY_BOX_FALLBACK",
|
|
track_id,
|
|
"body",
|
|
"rtmw_hand_unavailable",
|
|
)
|
|
)
|
|
record["body_fallbacks"] = body_fallbacks
|
|
self._record_timing("region_build", stage_started)
|
|
|
|
stage_started = time.perf_counter()
|
|
measurements = []
|
|
unavailable_status = f"door_mask_{loaded.status}"
|
|
for owner, region, source, track_id, side, state in regions:
|
|
if door_mask is None or distance_map is None:
|
|
measurement = self._unavailable(
|
|
source, track_id, side, unavailable_status
|
|
)
|
|
else:
|
|
measurement = self._measure_region(
|
|
region, door_mask, distance_map, source, track_id, side, state
|
|
)
|
|
owner["risk"] = measurement
|
|
measurements.append(measurement)
|
|
self._record_timing("region_measure", stage_started)
|
|
|
|
calculated = [item for item in measurements if item["risk_score"] is not None]
|
|
nearest = max(calculated, key=lambda item: item["risk_score"]) if calculated else None
|
|
if loaded.status != "loaded":
|
|
status = unavailable_status
|
|
elif not calculated:
|
|
status = "person_region_missing"
|
|
else:
|
|
status = "calculated"
|
|
frame_risk = {
|
|
"status": status,
|
|
"risk_score": nearest["risk_score"] if nearest else None,
|
|
"risk_level": nearest["risk_level"] if nearest else "unavailable",
|
|
"decision": nearest["decision"] if nearest else "EVIDENCE_UNAVAILABLE",
|
|
"source": nearest["source"] if nearest else None,
|
|
"track_id": nearest["track_id"] if nearest else None,
|
|
"side": nearest["side"] if nearest else None,
|
|
"door_mask": loaded.metadata(),
|
|
"warning_distance_relative": self.warning_relative,
|
|
"release_distance_relative": self.release_relative,
|
|
"regions": measurements,
|
|
}
|
|
self._record_timing("risk_total", total_started)
|
|
return frame_risk, loaded.mask
|
|
|
|
def summary(self) -> dict:
|
|
return {
|
|
"mask_directory": str(self.sequence.directory.resolve()),
|
|
"constant_mask": (
|
|
str(self.sequence.constant_path)
|
|
if self.sequence.constant_path is not None
|
|
else None
|
|
),
|
|
"constant_for_all_frames": self.sequence.constant_path is not None,
|
|
"mask_format": self.sequence.source_format,
|
|
"door_labels": sorted(self.sequence.labels),
|
|
"frame_offset": self.sequence.frame_offset,
|
|
"warning_distance_relative": self.warning_relative,
|
|
"release_distance_relative": self.release_relative,
|
|
"mask_status_counts": dict(self.mask_status_counts),
|
|
"risk_level_counts": dict(self.level_counts),
|
|
"source_counts": dict(self.source_counts),
|
|
"maximum_risk_score": round(self.maximum_score, 5),
|
|
"constant_distance_map_cached": self.sequence.constant_path is not None,
|
|
"distance_map_cache_builds": self.distance_map_cache_builds,
|
|
"timing_ms": {
|
|
name: {
|
|
"count": int(self.timing_counts[name]),
|
|
"total": round(float(seconds) * 1000.0, 3),
|
|
"mean": round(
|
|
float(seconds) * 1000.0
|
|
/ max(int(self.timing_counts[name]), 1),
|
|
3,
|
|
),
|
|
}
|
|
for name, seconds in self.timing_seconds.items()
|
|
},
|
|
"score_semantics": "2-D image overlap/proximity research score; not millimeters or safety certification",
|
|
"region_policy": "RTMW alarm-eligible hand hulls; YOLO body box only when a person has no usable RTMW hand region",
|
|
}
|