145 lines
5.3 KiB
Python
145 lines
5.3 KiB
Python
"""Минимальный конвейер mp4 → face-service /embed → tracks → apps/api.
|
|
|
|
Замысел: smoke-тест для цепочки «видео → эмбеддинг → трек → БД».
|
|
Это НЕ боевой video-ingest из ТЗ §6 (RTSP, ByteTrack, GPU) — это
|
|
скаффолд Фазы 1, чтобы Фаза 0 не блокировала остальную работу.
|
|
|
|
Алгоритм:
|
|
1. Открыть mp4 через OpenCV.
|
|
2. Для каждого N-го кадра (SAMPLE_EVERY): JPEG → face-service /embed.
|
|
3. Если эмбеддинг есть — single-camera трекинг по cos-дистанции:
|
|
если дистанция к last_embedding < TRACK_DISTANCE_THRESH и
|
|
прошло < TRACK_WINDOW_SEC секунд → тот же трек.
|
|
Иначе — новый трек (создаём через apps/api /ingest/tracks).
|
|
4. Сохраняем эмбеддинг через face-service /track-embeddings.
|
|
5. На первом кадре трека постим событие `arrived` через /ingest/track-events.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .api import FaceServiceClient, ReceptionApiClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class IngestConfig:
|
|
source: Path
|
|
camera_name: str
|
|
zone_code: str | None = None
|
|
sample_every: int = int(os.getenv("SAMPLE_EVERY", "10"))
|
|
track_distance_thresh: float = float(os.getenv("TRACK_DISTANCE_THRESH", "0.3"))
|
|
track_window_sec: float = float(os.getenv("TRACK_WINDOW_SEC", "2.0"))
|
|
jpeg_quality: int = 85
|
|
real_time_start: datetime | None = None
|
|
|
|
|
|
def cosine_distance(a: np.ndarray, b: np.ndarray) -> float:
|
|
return float(1.0 - np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
|
|
|
|
|
|
def run(cfg: IngestConfig) -> dict:
|
|
cap = cv2.VideoCapture(str(cfg.source))
|
|
if not cap.isOpened():
|
|
raise RuntimeError(f"Не могу открыть видео: {cfg.source}")
|
|
|
|
fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
|
|
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
logger.info(f"Открыто видео {cfg.source} ({total} кадров, {fps:.1f} fps)")
|
|
|
|
face = FaceServiceClient()
|
|
api = ReceptionApiClient()
|
|
base_time = cfg.real_time_start or datetime.utcnow()
|
|
|
|
current_track_id: str | None = None
|
|
current_track_camera_id: str | None = None
|
|
last_embedding: np.ndarray | None = None
|
|
last_capture_time: datetime | None = None
|
|
tracks_created = 0
|
|
embeddings_saved = 0
|
|
frame_idx = 0
|
|
|
|
while True:
|
|
ok, frame = cap.read()
|
|
if not ok:
|
|
break
|
|
frame_idx += 1
|
|
if frame_idx % cfg.sample_every != 0:
|
|
continue
|
|
|
|
captured_at = base_time + timedelta(seconds=frame_idx / fps)
|
|
|
|
ok_jpg, buf = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), cfg.jpeg_quality])
|
|
if not ok_jpg:
|
|
logger.warning(f"frame {frame_idx}: jpeg encode failed")
|
|
continue
|
|
jpeg_bytes = buf.tobytes()
|
|
|
|
try:
|
|
embed_result = face.embed(jpeg_bytes)
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning(f"frame {frame_idx}: face-service /embed failed: {e}")
|
|
continue
|
|
|
|
if embed_result is None:
|
|
continue
|
|
|
|
embedding_list, _quality = embed_result
|
|
embedding = np.array(embedding_list, dtype=np.float32)
|
|
|
|
new_track = True
|
|
if last_embedding is not None and last_capture_time is not None:
|
|
dist = cosine_distance(embedding, last_embedding)
|
|
age_sec = (captured_at - last_capture_time).total_seconds()
|
|
if dist < cfg.track_distance_thresh and age_sec < cfg.track_window_sec:
|
|
new_track = False
|
|
|
|
if new_track:
|
|
track_info = api.create_track(camera_name=cfg.camera_name, first_seen_at=captured_at)
|
|
current_track_id = track_info["trackId"]
|
|
current_track_camera_id = track_info["cameraId"]
|
|
api.add_event(
|
|
track_id=current_track_id,
|
|
event_type="arrived",
|
|
camera_name=cfg.camera_name,
|
|
occurred_at=captured_at,
|
|
zone_code=cfg.zone_code,
|
|
)
|
|
tracks_created += 1
|
|
logger.info(f"frame {frame_idx}: new track {current_track_id}")
|
|
|
|
assert current_track_id and current_track_camera_id
|
|
|
|
try:
|
|
face.save_track_embedding(
|
|
jpeg_bytes=jpeg_bytes,
|
|
track_id=current_track_id,
|
|
camera_id=current_track_camera_id,
|
|
captured_at=captured_at,
|
|
)
|
|
embeddings_saved += 1
|
|
except Exception as e: # noqa: BLE001
|
|
logger.warning(f"frame {frame_idx}: save_track_embedding failed: {e}")
|
|
|
|
last_embedding = embedding
|
|
last_capture_time = captured_at
|
|
|
|
cap.release()
|
|
logger.info(
|
|
f"Готово: создано {tracks_created} треков, сохранено {embeddings_saved} эмбеддингов "
|
|
f"(всего просмотрено {frame_idx} кадров)"
|
|
)
|
|
return {
|
|
"tracks_created": tracks_created,
|
|
"embeddings_saved": embeddings_saved,
|
|
"frames_processed": frame_idx,
|
|
}
|