digital_reception_ra/apps/video-ingest/video_ingest/ingestor.py

"""Минимальный конвейер 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,
    }