from __future__ import annotations import argparse import sys from collections import deque from datetime import datetime import cv2 import numpy as np from pathlib import Path from typing import Optional, Dict, Tuple, Any, List from ultralytics import YOLO import roi _DIR = Path(__file__).resolve().parent _ROOT = _DIR.parent MODELS_DIR = _ROOT / "models" RESULTS_DIR = _ROOT / "results" DET_MODEL_PATH = MODELS_DIR / "detector" / "yolov8s.pt" DET_CONF = 0.35 DET_CLASSES_DEFAULT = "" # Outside yellow polygon: 1× gauge width on each side (not drawn) _ROI_BUFFER_GAUGE_FRAC = 5.5 OUTPUT_PATH = str(RESULTS_DIR / "videos" / "all_detection.mp4") ORANGE_RED_SHOTS_DIR = RESULTS_DIR / "crossings" DEFAULT_TRACKER_PATH = MODELS_DIR / "trackers" / "bytetrack.yaml" _RU_ALERT_YELLOW = "Сигналь (предупреждение)" _RU_ALERT_ORANGE = "Снизь скорость и сигналь" _RU_ALERT_RED = "Тормози (аварийная остановка)" _RU_ALERT_TRAIN = "Внимание на трамвай/поезд" ALERT_COLORS: Dict[str, Tuple[int, int, int]] = { "red": (0, 0, 255), "orange": (0, 140, 255), "yellow": (0, 220, 255), } _NEUTRAL_COLOR = (180, 180, 180) _ALERT_HOLD_SEC = 1.0 _TRACK_GRID_PX = 80 HISTORY = 5 MIN_MOVE_PX = 1.5 DIST_EPS_PX = 1.0 TRACK_TTL = 30 TRAIL_LEN = 30 _TRAIL_COLOR = (0, 255, 0) _TRAIL_WATCH_COLOR = (0, 0, 255) MIN_REAL_MOVE_PX = 6.0 EGO_MIN_FEATURES = 12 EGO_MAX_CORNERS = 600 _TRAIN_CLASS_NAME = "train" _HAZARD_BASE_NAMES: frozenset[str] = frozenset( { "person", "bicycle", "car", "motorcycle", "bus", "truck", "traffic light", "stop sign", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "bench", } ) _HAZARD_NAME_SUBSTR_EXTRA: Tuple[str, ...] = ( "tree", "trash", "garbage", "waste", "dumpster", ) _TRACK_PERSON_NAMES = {"person"} _TRACK_VEHICLE_NAMES = { "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "skateboard", "skis", "snowboard", "surfboard", } _TRACK_ANIMAL_NAMES = { "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", } _TRACK_BASE_NAMES: frozenset[str] = frozenset( _TRACK_PERSON_NAMES | _TRACK_VEHICLE_NAMES | _TRACK_ANIMAL_NAMES ) _TRACK_NAME_SUBSTR: Tuple[str, ...] = ("bike", "scooter", "skate") def _norm_det_name(model: YOLO, cls_id: int) -> str: return str(model.names.get(cls_id, f"class_{cls_id}")).strip().lower() def _build_track_class_names(model: YOLO) -> frozenset[str]: """Transport (incl. ridables), people and animals — the trackable set.""" extra: set[str] = set() for v in model.names.values(): lv = str(v).strip().lower() if lv and any(s in lv for s in _TRACK_NAME_SUBSTR): extra.add(lv) return frozenset(_TRACK_BASE_NAMES | extra) def _build_hazard_class_names(model: YOLO) -> frozenset[str]: extra: set[str] = set() for v in model.names.values(): lv = str(v).strip().lower() if not lv or lv == _TRAIN_CLASS_NAME: continue if any(s in lv for s in _HAZARD_NAME_SUBSTR_EXTRA): extra.add(lv) return frozenset(_HAZARD_BASE_NAMES | extra) def _det_class_ids(model: YOLO, names_csv: str) -> list[int]: want = {s.strip().lower() for s in names_csv.split(",") if s.strip()} if not want: raise ValueError("Empty class list") found: dict[str, int] = {} for k, v in model.names.items(): key = str(v).strip().lower() if key in want: found[key] = int(k) missing = want - set(found.keys()) if missing: sys.stderr.write( f"Classes not found in model: {sorted(missing)}. " f"Model names: {list(model.names.values())[:30]}\n" ) sys.exit(1) return sorted(found.values()) class _AlertTracker: def __init__(self, fps: int, hold_sec: float) -> None: self._fps = max(1, fps) self._min_frames = max(1, int(fps * hold_sec)) self._timers: Dict[Tuple[int, int, int], int] = {} @staticmethod def _key(cls_id: int, cx: int, cy: int) -> Tuple[int, int, int]: return (cls_id, cx // _TRACK_GRID_PX, cy // _TRACK_GRID_PX) def tick(self, cls_id: int, cx: int, cy: int) -> Tuple[bool, int]: k = self._key(cls_id, cx, cy) self._timers[k] = self._timers.get(k, 0) + 1 nf = self._timers[k] return nf >= self._min_frames, nf def dwell_sec(self, cls_id: int, cx: int, cy: int) -> float: k = self._key(cls_id, cx, cy) return self._timers.get(k, 0) / float(self._fps) def end_frame(self, seen: set) -> None: for k in list(self._timers): if k not in seen: del self._timers[k] class _CrossingTracker: def __init__(self, fps: int) -> None: self._fps = max(1, fps) self._cnt: Dict[Tuple[int, int, int], int] = {} @staticmethod def _key(cls_id: int, cx: int, cy: int) -> Tuple[int, int, int]: return (cls_id, cx // _TRACK_GRID_PX, cy // _TRACK_GRID_PX) def tick(self, cls_id: int, cx: int, cy: int, crossing: bool) -> int: k = self._key(cls_id, cx, cy) if crossing: self._cnt[k] = self._cnt.get(k, 0) + 1 return self._cnt[k] self._cnt.pop(k, None) return 0 def dwell_sec(self, cls_id: int, cx: int, cy: int) -> float: k = self._key(cls_id, cx, cy) return self._cnt.get(k, 0) / float(self._fps) def end_frame(self, seen: set) -> None: for k in list(self._cnt): if k not in seen: del self._cnt[k] def estimate_distance_m( y_px: int, ys: np.ndarray, ls: np.ndarray, rs: np.ndarray, f_px: float, ) -> Optional[float]: if len(ys) < 2: return None widths = np.maximum(rs - ls, 1).astype(float) if y_px >= ys[np.argmax(ys)]: return 0.0 w_at_y = float(np.interp(y_px, ys, widths)) if w_at_y < 1: return None d_obj = f_px * roi.GAUGE_MM / w_at_y d_bot = f_px * roi.GAUGE_MM / float(widths[np.argmax(ys)]) return max(0.0, (d_obj - d_bot) / 1000.0) def classify_zone( x: int, y: int, ys: np.ndarray, ls: np.ndarray, rs: np.ndarray, zones: dict, ) -> Optional[str]: if len(ys) == 0: return None i = int(np.argmin(np.abs(ys - y))) if abs(int(ys[i]) - y) > 15: return None if int(ls[i]) <= x <= int(rs[i]): return "red" for name in ("red", "orange", "yellow"): if name not in zones: continue zl, zr = zones[name] if int(zl[i]) <= x <= int(zr[i]): return name return None def _outer_bounds( zones: dict, ls: np.ndarray, rs: np.ndarray ) -> Tuple[np.ndarray, np.ndarray]: for name in ("yellow", "orange", "red"): if name in zones: return zones[name] return ls, rs def _roi_band_at_y( y: int, ys: np.ndarray, L_arr: np.ndarray, R_arr: np.ndarray ) -> Tuple[float, float]: yy = float(np.clip(y, int(ys.min()), int(ys.max()))) ysf = ys.astype(np.float64) L = float(np.interp(yy, ysf, L_arr.astype(np.float64))) R = float(np.interp(yy, ysf, R_arr.astype(np.float64))) if L > R: L, R = R, L return L, R def _horiz_gap(x1: int, x2: int, L: float, R: float) -> float: if x2 < L: return L - x2 if x1 > R: return x1 - R return 0.0 def _sign(v: float) -> int: if v > 0: return 1 if v < 0: return -1 return 0 def _warp_affine_pt( M: Optional[np.ndarray], x: float, y: float ) -> Tuple[float, float]: if M is None: return float(x), float(y) nx = float(M[0, 0] * x + M[0, 1] * y + M[0, 2]) ny = float(M[1, 0] * x + M[1, 1] * y + M[1, 2]) return nx, ny class _EgoMotion: _QUALITY = 0.01 _MIN_DIST = 8 _BOX_MARGIN = 12 _LK = dict( winSize=(21, 21), maxLevel=3, criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 30, 0.01), ) def __init__(self) -> None: self.prev_gray: Optional[np.ndarray] = None def update( self, gray: np.ndarray, exclude_boxes: List[Tuple[int, int, int, int]], ) -> Optional[np.ndarray]: prev = self.prev_gray self.prev_gray = gray if prev is None: return None h, w = gray.shape[:2] mask = np.full((h, w), 255, dtype=np.uint8) m = self._BOX_MARGIN for (x1, y1, x2, y2) in exclude_boxes: xa, ya = max(0, int(x1) - m), max(0, int(y1) - m) xb, yb = min(w, int(x2) + m), min(h, int(y2) + m) if xb > xa and yb > ya: mask[ya:yb, xa:xb] = 0 p0 = cv2.goodFeaturesToTrack( prev, maxCorners=EGO_MAX_CORNERS, qualityLevel=self._QUALITY, minDistance=self._MIN_DIST, mask=mask, ) if p0 is None or len(p0) < EGO_MIN_FEATURES: return None p1, st, _ = cv2.calcOpticalFlowPyrLK(prev, gray, p0, None, **self._LK) if p1 is None or st is None: return None st = st.reshape(-1).astype(bool) a = p0.reshape(-1, 2)[st] b = p1.reshape(-1, 2)[st] if len(a) < EGO_MIN_FEATURES: return None M, inl = cv2.estimateAffinePartial2D( a, b, method=cv2.RANSAC, ransacReprojThreshold=3.0, maxIters=2000, confidence=0.99, ) if M is None or inl is None or int(inl.sum()) < EGO_MIN_FEATURES // 2: return None return M def _bbox_spans_red_and_orange( x1: int, y1: int, x2: int, y2: int, ys: np.ndarray, ls: np.ndarray, rs: np.ndarray, zones: dict, ) -> bool: if y2 <= y1 or "orange" not in zones or "red" not in zones: return False seen: set[str] = set() cx = (x1 + x2) // 2 for t in (0.0, 0.12, 0.28, 0.45, 0.62, 0.78, 0.92, 1.0): yy = int(y1 + (y2 - y1) * t) z = classify_zone(cx, yy, ys, ls, rs, zones) if z: seen.add(z) if "red" in seen and "orange" in seen: return True y_foot = max(y1, y2 - 3) for xx in (x1 + 2, cx, x2 - 2): xx = int(min(max(xx, x1), x2)) z = classify_zone(xx, y_foot, ys, ls, rs, zones) if z: seen.add(z) if "red" in seen and "orange" in seen: return True return "red" in seen and "orange" in seen def _expanded_interest_band( ys: np.ndarray, ls: np.ndarray, rs: np.ndarray, zones: dict, w: int, buffer_gauge_frac: float = _ROI_BUFFER_GAUGE_FRAC, ) -> Tuple[np.ndarray, np.ndarray]: """Yellow outer edge + buffer_gauge_frac * gauge width OUTSIDE on left/right.""" widths = np.maximum(rs - ls, 1).astype(float) buf_px = (buffer_gauge_frac * widths).astype(np.int32) yl, yr = _outer_bounds(zones, ls, rs) el = np.clip(yl.astype(np.int32) - buf_px, 0, w - 1) er = np.clip(yr.astype(np.int32) + buf_px, 0, w - 1) return el, er def _expanded_band_poly( ys: np.ndarray, exp_l: np.ndarray, exp_r: np.ndarray, ) -> np.ndarray: lp = np.column_stack([exp_l, ys]) rp = np.column_stack([exp_r, ys]) return np.vstack([lp, rp[::-1]]).astype(np.int32) def _mask_frame_for_roi_det( frame: np.ndarray, ys: np.ndarray, exp_l: np.ndarray, exp_r: np.ndarray, ) -> np.ndarray: """Black-out everything outside the expanded band (invisible, not drawn).""" poly = _expanded_band_poly(ys, exp_l, exp_r) mask = np.zeros(frame.shape[:2], dtype=np.uint8) cv2.fillPoly(mask, [poly], 255, lineType=cv2.LINE_8) out = np.zeros_like(frame) out[mask > 0] = frame[mask > 0] return out def _bbox_in_expanded_roi( x1: int, y1: int, x2: int, y2: int, ys: np.ndarray, exp_l: np.ndarray, exp_r: np.ndarray, ) -> bool: if len(ys) < 2: return False ymn, ymx = int(ys.min()), int(ys.max()) y_lo, y_hi = int(min(y1, y2)), int(max(y1, y2)) for yy in (y_lo, (y_lo + y_hi) // 2, y_hi): yy = int(np.clip(yy, ymn, ymx)) fl = float(np.interp(yy, ys.astype(np.float64), exp_l.astype(np.float64))) fr = float(np.interp(yy, ys.astype(np.float64), exp_r.astype(np.float64))) el, er = int(round(fl)), int(round(fr)) if el > er: el, er = er, el if not (x2 < el or x1 > er): return True return False def alert_distance_in_band(speed_kmh: int, distance_m: Optional[float]) -> bool: if distance_m is None: return False d = max(0.0, float(distance_m)) v = max(0, min(75, int(speed_kmh))) if v <= 30: return d <= 15.0 if v < 44: return 2.0 <= d <= 10.0 return d >= 10.0 def compute_alert( zone: Optional[str], obj_name: str, distance_m: Optional[float], speed: int, ) -> Optional[Tuple[str, str]]: if zone is None: return None if not alert_distance_in_band(speed, distance_m): return None d = distance_m if d is None: return None dist_human = f"{d:.1f} м" if zone == "red": ru = _RU_ALERT_RED console = ( f"{ru} | объект={obj_name} | дистанция={dist_human} | скорость={speed} км/ч" ) return (console, ru) if zone == "orange": ru = _RU_ALERT_ORANGE console = ( f"{ru} | объект={obj_name} | дистанция={dist_human} | скорость={speed} км/ч" ) return (console, ru) if zone == "yellow": ru = _RU_ALERT_YELLOW console = ( f"{ru} | объект={obj_name} | дистанция={dist_human} | скорость={speed} км/ч" ) return (console, ru) return None def _write_danger_log_line( danger_f, ts_iso: str, obj_name: str, distance_m: Optional[float], speed_kmh: int, action_ru: str, frame_idx: int, ) -> None: date_s, _sep, time_s = ts_iso.partition(" ") dist_s = f"{distance_m:.2f}" if distance_m is not None else "na" danger_f.write( f"{ts_iso}\tdate={date_s}\ttime={time_s}\tclass={obj_name}\t" f"dist_m={dist_s}\tspeed_kmh={speed_kmh}\taction={action_ru}\tframe={frame_idx}\n" ) danger_f.flush() _LT = cv2.LINE_8 _FONT = cv2.FONT_HERSHEY_PLAIN _FONT_SCALE = 1 def _draw_detection( frame: np.ndarray, x1: int, y1: int, x2: int, y2: int, obj_name: str, zone: Optional[str], distance_m: Optional[float], override_color: Optional[Tuple[int, int, int]] = None, override_thickness: Optional[int] = None, track_id: Optional[int] = None, ) -> None: if override_color is not None: color = override_color thickness = override_thickness if override_thickness is not None else 3 else: color = ALERT_COLORS.get(zone, _NEUTRAL_COLOR) if zone else _NEUTRAL_COLOR thickness = 2 if zone else 1 cv2.rectangle(frame, (x1, y1), (x2, y2), color, thickness) label = obj_name if track_id is not None and track_id >= 0: label += f" #{int(track_id)}" if distance_m is not None: label += f" {distance_m:.1f}m" (tw, th), bl = cv2.getTextSize(label, _FONT, _FONT_SCALE, 1) pad_y = th + bl + 6 cv2.rectangle(frame, (x1, y1 - pad_y), (x1 + tw + 4, y1), color, -1) cv2.putText( frame, label, (x1 + 2, y1 - bl - 3), _FONT, _FONT_SCALE, (255, 255, 255), 1, _LT, ) def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("-s", "--speed", type=int, required=True) ap.add_argument("-v", "--video", type=str, default=roi.VIDEO_PATH) ap.add_argument("-o", "--output", type=str, default=OUTPUT_PATH) ap.add_argument("--det-conf", type=float, default=DET_CONF) ap.add_argument("--det-model", type=str, default=str(DET_MODEL_PATH)) ap.add_argument( "--det-classes", type=str, default=DET_CLASSES_DEFAULT ) ap.add_argument( "--roi", action="store_true", help=( "Detect/track in colored zones + " f"{_ROI_BUFFER_GAUGE_FRAC:.1f}× gauge width outside yellow (masked, not drawn)" ), ) ap.add_argument("--seg-model", type=str, default=None) ap.add_argument( "-m", "--morph", action="store_true", dest="morph", ) ap.add_argument("--hold-sec", type=float, default=_ALERT_HOLD_SEC) ap.add_argument("--fov", type=float, default=roi.DEFAULT_FOV_DEG) ap.add_argument( "--tracker", type=str, default=str(DEFAULT_TRACKER_PATH) if DEFAULT_TRACKER_PATH.exists() else "bytetrack.yaml", ) args = ap.parse_args() speed = max(1, min(75, args.speed)) seg_model = roi.load_model(args.seg_model) det_model_path = Path(args.det_model) if not det_model_path.exists() and args.det_model == str(DET_MODEL_PATH): print( f"[ERROR] Модель детектора не найдена: {det_model_path}\n" "Запустите: python scripts/download_models.py" ) sys.exit(1) det_model = YOLO(str(det_model_path)) hazard_class_names = _build_hazard_class_names(det_model) track_class_names = _build_track_class_names(det_model) _dc = (args.det_classes or "").strip().lower() if not _dc or _dc == "all": det_class_ids = None else: det_class_ids = _det_class_ids(det_model, args.det_classes) cap = cv2.VideoCapture(args.video) if not cap.isOpened(): raise OSError(f"Failed to open video: {args.video}") fps = int(cap.get(cv2.CAP_PROP_FPS)) or 25 W = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) H = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) hold_frames = max(1, int(fps * args.hold_sec)) _log_date = datetime.now().strftime("%d_%m_%Y") logs_dir = RESULTS_DIR / "logs" logs_dir.mkdir(parents=True, exist_ok=True) log_path = logs_dir / f"detection_log_{_log_date}.txt" danger_log_path = logs_dir / f"DANGER_LOG_{_log_date}.txt" writer = None if args.output: Path(args.output).parent.mkdir(parents=True, exist_ok=True) writer = cv2.VideoWriter( args.output, cv2.VideoWriter_fourcc(*"mp4v"), fps, (W, H), ) f_px = roi.get_focal_px(W, args.fov) ORANGE_RED_SHOTS_DIR.mkdir(parents=True, exist_ok=True) temporal = roi._TemporalState() tracker = _AlertTracker(fps, args.hold_sec) cross_tracker = _CrossingTracker(fps) train_warn_last: Dict[Tuple[int, int, int], int] = {} traj_hist: Dict[int, deque] = {} trail_pts: Dict[int, deque] = {} track_last_seen: Dict[int, int] = {} ego = _EgoMotion() prev_obj_boxes: List[Tuple[int, int, int, int]] = [] track_prev_pos: Dict[int, Tuple[int, int]] = {} track_resid: Dict[int, deque] = {} n = 0 shot_n = 0 with open(log_path, "a", encoding="utf-8") as log_f, open( danger_log_path, "a", encoding="utf-8" ) as danger_f: _run_hdr = ( f"\n!!!!!! run {datetime.now().isoformat()} speed_kmh={speed} " f"video={args.video} hold_sec={args.hold_sec} fps={fps} " f"roi_filter={args.roi} roi_buffer_gauge_frac={_ROI_BUFFER_GAUGE_FRAC} !!!!!!!\n" ) log_f.write(_run_hdr) log_f.flush() danger_f.write( f"{_run_hdr}" "ts\tdate\ttime\tclass\tdist_m\tspeed_kmh\taction\tframe\n" ) danger_f.flush() while True: ok, frame = cap.read() if not ok: break h, w = frame.shape[:2] # Движение камеры prev->cur (по чистому кадру, до отрисовки зон). gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) ego_M = ego.update(gray, prev_obj_boxes) seg_res = seg_model(frame, conf=roi.CONF, device=roi.DEVICE, verbose=False)[0] roi._maybe_save_low_conf_frame(frame, roi._max_rail_confidence(seg_res)) full_mask = roi.extract_rail_mask(seg_res, h, w, morph=args.morph) if np.any(full_mask): groups = roi.find_groups(full_mask) gmask = roi._select_rail_group(groups, h, temporal.center()) ys, ls, rs = roi.rail_rows(gmask) else: ys = np.array([], dtype=np.int32) ls = np.array([], dtype=np.int32) rs = np.array([], dtype=np.int32) ys, ls, rs = temporal.update(ys, ls, rs, w) has_roi = len(ys) >= 4 zones: dict = {} exp_l: Optional[np.ndarray] = None exp_r: Optional[np.ndarray] = None outer_l: Optional[np.ndarray] = None outer_r: Optional[np.ndarray] = None if has_roi: zones = roi.zone_bounds(ys, ls, rs, w) grid = roi.grid_positions(ys, ls, rs, f_px) roi.draw_zones(frame, ys, zones, ls, rs, grid) outer_l, outer_r = _outer_bounds(zones, ls, rs) if args.roi: exp_l, exp_r = _expanded_interest_band(ys, ls, rs, zones, w) _dkw: dict = { "conf": args.det_conf, "device": roi.DEVICE, "verbose": False, } if det_class_ids is not None: _dkw["classes"] = det_class_ids det_res = None if not args.roi or has_roi: det_frame = frame if ( args.roi and has_roi and exp_l is not None and exp_r is not None ): det_frame = _mask_frame_for_roi_det(frame, ys, exp_l, exp_r) det_res = det_model.track( det_frame, tracker=args.tracker, persist=True, **_dkw )[0] seen_keys: set = set() all_det_keys: set = set() seen_ids: set = set() going_now: set = set() cur_obj_boxes: List[Tuple[int, int, int, int]] = [] save_orange_red_shot = False ts_log = datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")[:-3] if det_res is not None and det_res.boxes is not None: boxes_xyxy = det_res.boxes.xyxy.cpu().numpy() classes = det_res.boxes.cls.cpu().numpy().astype(int) if det_res.boxes.id is not None: track_ids = det_res.boxes.id.cpu().numpy().astype(int) else: track_ids = np.full(len(classes), -1, dtype=int) for box, cls_id, tid in zip(boxes_xyxy, classes, track_ids): x1, y1, x2, y2 = box.astype(int) obj_name = det_model.names.get(cls_id, f"class_{cls_id}") cx = (x1 + x2) // 2 cy_foot = int(y2) cur_obj_boxes.append((int(x1), int(y1), int(x2), int(y2))) cls_name_l = _norm_det_name(det_model, cls_id) is_trackable = cls_name_l in track_class_names eff_tid = int(tid) if (is_trackable and tid >= 0) else -1 really_moving = False net_dx = 0.0 if eff_tid >= 0: prev_pos = track_prev_pos.get(eff_tid) if prev_pos is not None: pcx, pcy = _warp_affine_pt( ego_M, prev_pos[0], prev_pos[1] ) rdx = cx - pcx rdy = cy_foot - pcy else: rdx = rdy = 0.0 rb = track_resid.setdefault(eff_tid, deque(maxlen=HISTORY)) rb.append((rdx, rdy)) track_prev_pos[eff_tid] = (cx, cy_foot) net_dx = float(sum(d[0] for d in rb)) net_dy = float(sum(d[1] for d in rb)) really_moving = ( (net_dx * net_dx + net_dy * net_dy) ** 0.5 >= MIN_REAL_MOVE_PX ) going = False if really_moving and has_roi and outer_l is not None and eff_tid >= 0: L, R = _roi_band_at_y(cy_foot, ys, outer_l, outer_r) dist = _horiz_gap(x1, x2, L, R) roi_cx = 0.5 * (L + R) hist = traj_hist.setdefault(eff_tid, deque(maxlen=HISTORY)) if len(hist) >= 1: ref_cx, ref_cy, ref_dist, ref_roi_cx = hist[0] h_dir = _sign(ref_roi_cx - ref_cx) if h_dir == 0: moving_toward = True else: moving_toward = ( _sign(net_dx) == h_dir and abs(net_dx) >= MIN_MOVE_PX ) if moving_toward and dist < ref_dist - DIST_EPS_PX: going = True print(f"WATCH THE OBJECT {obj_name}") hist.append((cx, cy_foot, dist, roi_cx)) if eff_tid >= 0: track_last_seen[eff_tid] = n seen_ids.add(eff_tid) if really_moving: trail_pts.setdefault( eff_tid, deque(maxlen=TRAIL_LEN) ).append((cx, cy_foot)) if going: going_now.add(eff_tid) if args.roi: if ( exp_l is None or exp_r is None or not _bbox_in_expanded_roi( x1, y1, x2, y2, ys, exp_l, exp_r ) ): continue if cls_name_l == _TRAIN_CLASS_NAME: all_det_keys.add(tracker._key(cls_id, cx, cy_foot)) zone_t: Optional[str] = None dist_t: Optional[float] = None crossing_t = False if has_roi and zones: zone_t = classify_zone( cx, cy_foot, ys, ls, rs, zones ) if zone_t is not None: dist_t = estimate_distance_m( cy_foot, ys, ls, rs, f_px ) crossing_t = _bbox_spans_red_and_orange( x1, y1, x2, y2, ys, ls, rs, zones ) if crossing_t: save_orange_red_shot = True nf_cross_t = cross_tracker.tick( cls_id, cx, cy_foot, crossing_t ) cross_sec_t = nf_cross_t / float(fps) d_ts = ( f"{dist_t:.2f}" if dist_t is not None else "na" ) z_ts = zone_t if zone_t is not None else "na" log_f.write( f"[TRAIN] {ts_log} frame={n + 1} obj={obj_name} " f"zone={z_ts} dist_m={d_ts} crossing={crossing_t} " f"dwell_in_crossing_s={cross_sec_t:.3f} " f"speed_kmh={speed}\n" ) log_f.flush() wk = tracker._key(cls_id, cx, cy_foot) _tw_gap = max(1, fps // 2) if n - train_warn_last.get(wk, -10**9) >= _tw_gap: if alert_distance_in_band(speed, dist_t): _dist_h = ( f"{dist_t:.1f} м" if dist_t is not None else "н/д" ) _tcons = ( f"{_RU_ALERT_TRAIN} | объект={obj_name} | " f"дистанция={_dist_h} | скорость={speed} км/ч" ) print( f"[DANGER {ts_log}] frame={n + 1} {_tcons}" ) train_warn_last[wk] = n _write_danger_log_line( danger_f, ts_log, str(obj_name), dist_t, speed, _RU_ALERT_TRAIN, n + 1, ) _draw_detection( frame, x1, y1, x2, y2, obj_name, zone_t, dist_t, override_color=(255, 0, 255), track_id=eff_tid, ) continue if cls_name_l not in hazard_class_names: continue all_det_keys.add(tracker._key(cls_id, cx, cy_foot)) zone: Optional[str] = None dist_m: Optional[float] = None crossing = False if has_roi and zones: zone = classify_zone(cx, cy_foot, ys, ls, rs, zones) if zone is not None: dist_m = estimate_distance_m(cy_foot, ys, ls, rs, f_px) crossing = _bbox_spans_red_and_orange( x1, y1, x2, y2, ys, ls, rs, zones ) if crossing: save_orange_red_shot = True nf_cross = cross_tracker.tick(cls_id, cx, cy_foot, crossing) cross_sec = nf_cross / float(fps) if zone is not None: key = tracker._key(cls_id, cx, cy_foot) seen_keys.add(key) _confirmed, nf_zone = tracker.tick(cls_id, cx, cy_foot) dwell_zone_s = nf_zone / float(fps) msg_pair = compute_alert(zone, obj_name, dist_m, speed) d_s = f"{dist_m:.2f}" if dist_m is not None else "na" if crossing: log_f.write( f"[CROSS] {ts_log} frame={n + 1} obj={obj_name} " f"dist_m={d_s} zone={zone} dwell_in_crossing_s={cross_sec:.3f} " f"dwell_in_zone_s={dwell_zone_s:.3f} speed_kmh={speed}\n" ) else: log_f.write( f"[ZONE] {ts_log} frame={n + 1} obj={obj_name} " f"zone={zone} dist_m={d_s} dwell_in_zone_s={dwell_zone_s:.3f} " f"speed_kmh={speed}\n" ) log_f.flush() if ( msg_pair is not None and nf_zone == hold_frames ): console_line, action_ru = msg_pair print( f"[DANGER {ts_log}] frame={n + 1} {console_line} " f"(зона ≥{args.hold_sec:.1f} с)" ) _write_danger_log_line( danger_f, ts_log, str(obj_name), dist_m, speed, action_ru, n + 1, ) _draw_detection( frame, x1, y1, x2, y2, obj_name, zone, dist_m, track_id=eff_tid ) cross_tracker.end_frame(all_det_keys) if save_orange_red_shot: shot_n += 1 ts = datetime.now().strftime("%d_%m_%Y_%H%M%S_%f")[:-3] out_png = ORANGE_RED_SHOTS_DIR / f"frame{n + 1:06d}_{shot_n:04d}_{ts}.png" cv2.imwrite(str(out_png), frame) tracker.end_frame(seen_keys) for tid in seen_ids: pts = trail_pts.get(tid) if not pts or len(pts) < 2: continue arr = np.array(pts, dtype=np.int32).reshape(-1, 1, 2) tc = _TRAIL_WATCH_COLOR if tid in going_now else _TRAIL_COLOR cv2.polylines(frame, [arr], False, tc, 2, cv2.LINE_AA) px, py = pts[-1] cv2.circle(frame, (int(px), int(py)), 3, tc, -1) for tid in list(track_last_seen.keys()): if n - track_last_seen[tid] > TRACK_TTL: track_last_seen.pop(tid, None) traj_hist.pop(tid, None) trail_pts.pop(tid, None) track_prev_pos.pop(tid, None) track_resid.pop(tid, None) prev_obj_boxes = cur_obj_boxes cv2.imshow("All detection + zones", frame) if writer: writer.write(frame) if cv2.waitKey(1) & 0xFF == ord("q"): break n += 1 cap.release() if writer: writer.release() cv2.destroyAllWindows() if __name__ == "__main__": main()