initial commit

app/video_pipeline.py

@@ -0,0 +1,373 @@
+from __future__ import annotations
+
+import shutil
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import cv2
+
+from app.config import settings
+from app.pipeline import BirdCascadePipeline
+from app.utils import (
+    draw_results,
+    now_timestamp_minute,
+    safe_stem,
+    save_json,
+    write_video_xml,
+)
+
+
+def make_unique_result_dir(parent_dir: Path, timestamp: str, original_filename: str) -> tuple[Path, str, str]:
+    """
+    创建单次视频识别结果目录。
+
+    目录名：
+        时间戳_原文件名stem
+
+    例如：
+        runs/video/202605181431_DJI_20260423155914_0006_Z/
+
+    返回：
+        result_dir, result_base_name, original_stem
+    """
+    original_stem = safe_stem(original_filename)
+    base_name = f"{timestamp}_{original_stem}"
+
+    result_dir = parent_dir / base_name
+
+    if not result_dir.exists():
+        result_dir.mkdir(parents=True, exist_ok=True)
+        return result_dir, base_name, original_stem
+
+    idx = 1
+    while True:
+        candidate_base = f"{base_name}_{idx:02d}"
+        candidate_dir = parent_dir / candidate_base
+        if not candidate_dir.exists():
+            candidate_dir.mkdir(parents=True, exist_ok=True)
+            return candidate_dir, candidate_base, original_stem
+        idx += 1
+
+
+def iou_xyxy(a: List[float], b: List[float]) -> float:
+    ax1, ay1, ax2, ay2 = a
+    bx1, by1, bx2, by2 = b
+
+    x1 = max(ax1, bx1)
+    y1 = max(ay1, by1)
+    x2 = min(ax2, bx2)
+    y2 = min(ay2, by2)
+
+    inter = max(0.0, x2 - x1) * max(0.0, y2 - y1)
+    area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
+    area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
+
+    return inter / (area_a + area_b - inter + 1e-7)
+
+
+class VideoResultSmoother:
+    """
+    降低视频框闪烁：
+    1. 非推理帧复用上一推理结果；
+    2. 短暂漏检时保留上一结果若干帧；
+    3. 对匹配目标的坐标做指数平滑；
+    4. 对同一目标累计物种置信度，避免标签和颜色逐帧跳变。
+    """
+
+    def __init__(self, hold_frames: int = 8, alpha: float = 0.65):
+        self.hold_frames = hold_frames
+        self.alpha = alpha
+        self.last_results: List[Dict] = []
+        self.last_update_frame: int = -10**9
+
+    @staticmethod
+    def _public_results(results: List[Dict]) -> List[Dict]:
+        return [
+            {
+                key: value
+                for key, value in item.items()
+                if not key.startswith("_smoother_")
+            }
+            for item in results
+        ]
+
+    def update(self, current_results: List[Dict], frame_index: int) -> List[Dict]:
+        if len(current_results) == 0:
+            if frame_index - self.last_update_frame <= self.hold_frames:
+                return self._public_results(self.last_results)
+            self.last_results = []
+            return []
+
+        if len(self.last_results) == 0:
+            self.last_results = []
+            for current in current_results:
+                item = dict(current)
+                class_key = (
+                    int(item.get("cls_class_id", -1)),
+                    str(item.get("cls_class_name", "") or ""),
+                )
+                confidence = max(float(item.get("cls_confidence", 0.0)), 0.01)
+                item["_smoother_class_votes"] = {class_key: confidence}
+                self.last_results.append(item)
+            self.last_update_frame = frame_index
+            return self._public_results(self.last_results)
+
+        smoothed_results = []
+        used_last = set()
+
+        for cur in current_results:
+            cur_box = cur["bbox_xyxy"]
+
+            best_idx = -1
+            best_iou = 0.0
+
+            for idx, prev in enumerate(self.last_results):
+                if idx in used_last:
+                    continue
+
+                score = iou_xyxy(cur_box, prev["bbox_xyxy"])
+                if score > best_iou:
+                    best_iou = score
+                    best_idx = idx
+
+            new_item = dict(cur)
+
+            if best_idx >= 0 and best_iou >= 0.1:
+                prev = self.last_results[best_idx]
+                prev_box = prev["bbox_xyxy"]
+
+                smooth_box = []
+                for p, c in zip(prev_box, cur_box):
+                    smooth_box.append(round(self.alpha * float(p) + (1 - self.alpha) * float(c), 2))
+
+                new_item["bbox_xyxy"] = smooth_box
+
+                votes = {
+                    key: float(value) * 0.90
+                    for key, value in prev.get(
+                        "_smoother_class_votes",
+                        {},
+                    ).items()
+                }
+                current_key = (
+                    int(cur.get("cls_class_id", -1)),
+                    str(cur.get("cls_class_name", "") or ""),
+                )
+                current_confidence = max(
+                    float(cur.get("cls_confidence", 0.0)),
+                    0.01,
+                )
+                votes[current_key] = votes.get(current_key, 0.0) + current_confidence
+                chosen_key = max(votes, key=votes.get)
+                new_item["_smoother_class_votes"] = votes
+                new_item["cls_class_id"] = chosen_key[0]
+                new_item["cls_class_name"] = chosen_key[1]
+                if chosen_key != current_key:
+                    new_item["show_cls_name"] = bool(
+                        prev.get("show_cls_name", False)
+                    )
+                used_last.add(best_idx)
+            else:
+                class_key = (
+                    int(cur.get("cls_class_id", -1)),
+                    str(cur.get("cls_class_name", "") or ""),
+                )
+                confidence = max(float(cur.get("cls_confidence", 0.0)), 0.01)
+                new_item["_smoother_class_votes"] = {class_key: confidence}
+
+            smoothed_results.append(new_item)
+
+        self.last_results = smoothed_results
+        self.last_update_frame = frame_index
+
+        return self._public_results(smoothed_results)
+
+    def get_for_skipped_frame(self, frame_index: int) -> List[Dict]:
+        if frame_index - self.last_update_frame <= self.hold_frames:
+            return self._public_results(self.last_results)
+        return []
+
+
+def process_video_file(
+    video_path: str,
+    pipeline: BirdCascadePipeline,
+    original_filename: str,
+    sample_interval: Optional[int] = None,
+    save_annotated_video: Optional[bool] = None,
+) -> Dict:
+    """
+    视频结果自动保存到：
+        runs/video/时间戳_原视频名/
+
+    目录内部生成：
+        原视频名.mp4
+        时间戳_原视频名.json
+        时间戳_原视频名.xml
+        原视频名_source.原后缀
+        frames/
+
+    示例：
+        runs/video/202605181431_DJI_20260423155914_0006_Z/
+        ├── DJI_20260423155914_0006_Z.mp4
+        ├── 202605181431_DJI_20260423155914_0006_Z.json
+        ├── 202605181431_DJI_20260423155914_0006_Z.xml
+        ├── DJI_20260423155914_0006_Z_source.MP4
+        └── frames/
+    """
+    sample_interval = settings.video_sample_interval if sample_interval is None else int(sample_interval)
+    save_annotated_video = settings.save_annotated_video if save_annotated_video is None else bool(save_annotated_video)
+
+    video_root = Path(settings.output_dir) / settings.video_output_subdir
+    video_root.mkdir(parents=True, exist_ok=True)
+
+    timestamp = now_timestamp_minute()
+    result_dir, result_base_name, original_stem = make_unique_result_dir(
+        parent_dir=video_root,
+        timestamp=timestamp,
+        original_filename=original_filename,
+    )
+
+    original_suffix = Path(original_filename).suffix or ".mp4"
+
+    source_video_path = result_dir / f"{original_stem}_source{original_suffix}"
+
+    # 输出视频固定为 mp4，文件名不带时间戳
+    annotated_video_path = result_dir / f"{original_stem}.mp4"
+
+    # JSON 和 XML 带时间戳
+    json_path = result_dir / f"{result_base_name}.json"
+    xml_path = result_dir / f"{result_base_name}.xml"
+
+    frames_dir = result_dir / "frames"
+    frames_dir.mkdir(parents=True, exist_ok=True)
+
+    shutil.copy2(video_path, source_video_path)
+
+    cap = cv2.VideoCapture(str(source_video_path))
+    if not cap.isOpened():
+        raise ValueError(f"视频打开失败: {source_video_path}")
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    if fps is None or fps <= 0:
+        fps = 25.0
+
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0)
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0)
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+
+    writer = None
+    if save_annotated_video:
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        writer = cv2.VideoWriter(str(annotated_video_path), fourcc, fps, (width, height))
+
+    smoother = VideoResultSmoother(
+        hold_frames=settings.smooth_hold_frames,
+        alpha=settings.box_smooth_alpha,
+    )
+
+    frame_results = []
+    processed_frames = 0
+    detected_frames = 0
+    frame_index = 0
+
+    start_time = time.time()
+
+    try:
+        while True:
+            ok, frame = cap.read()
+            if not ok:
+                break
+
+            should_process = frame_index % sample_interval == 0
+
+            if should_process:
+                output = pipeline.predict(frame, return_image=False)
+                raw_results = output["results"]
+                processed_frames += 1
+
+                if settings.smooth_video_boxes:
+                    results = smoother.update(raw_results, frame_index)
+                else:
+                    results = raw_results
+
+                if len(results) > 0:
+                    detected_frames += 1
+
+                frame_results.append(
+                    {
+                        "frame_index": frame_index,
+                        "time_sec": round(frame_index / fps, 3),
+                        "num_objects": len(results),
+                        "results": results,
+                    }
+                )
+            else:
+                if settings.smooth_video_boxes:
+                    results = smoother.get_for_skipped_frame(frame_index)
+                else:
+                    results = []
+
+            annotated_frame = draw_results(frame, results)
+
+            if should_process and len(results) > 0:
+                frame_img_path = frames_dir / f"frame_{frame_index:08d}.jpg"
+                cv2.imwrite(str(frame_img_path), annotated_frame)
+
+            if writer is not None:
+                writer.write(annotated_frame)
+
+            frame_index += 1
+
+    finally:
+        cap.release()
+        if writer is not None:
+            writer.release()
+
+    elapsed = round(time.time() - start_time, 3)
+
+    summary = {
+        "success": True,
+        "type": "video",
+        "original_filename": original_filename,
+        "result_dir": str(result_dir),
+        "source_video": str(source_video_path),
+        "annotated_video": str(annotated_video_path) if annotated_video_path.exists() else None,
+        "json": str(json_path),
+        "xml": str(xml_path),
+        "frames_dir": str(frames_dir),
+        "video_info": {
+            "width": width,
+            "height": height,
+            "fps": fps,
+            "total_frames": total_frames,
+        },
+        "sample_interval": sample_interval,
+        "smooth_video_boxes": settings.smooth_video_boxes,
+        "smooth_hold_frames": settings.smooth_hold_frames,
+        "box_smooth_alpha": settings.box_smooth_alpha,
+        "processed_frames": processed_frames,
+        "detected_frames": detected_frames,
+        "elapsed_sec": elapsed,
+    }
+
+    json_data = {
+        "summary": summary,
+        "frames": frame_results,
+    }
+
+    save_json(json_path, json_data)
+
+    write_video_xml(
+        xml_path=xml_path,
+        video_filename=annotated_video_path.name,
+        video_path=str(annotated_video_path),
+        width=width,
+        height=height,
+        fps=fps,
+        total_frames=total_frames,
+        sample_interval=sample_interval,
+        frame_results=frame_results,
+    )
+
+    return summary