2026-05-18-19-56-47 重构双视频同步与单帧对比

backend/main.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import subprocess
 import shutil
 import uuid
 from pathlib import Path
@@ -101,6 +102,9 @@ def _save_frame_outputs(
     method: str,
     frame,
     output,
+    source_time: float | None = None,
+    result_time: float | None = None,
+    result_index: int | None = None,
 ) -> dict[str, Any]:
     method_path = job_path / method
     method_path.mkdir(exist_ok=True)
@@ -110,7 +114,7 @@ def _save_frame_outputs(
     cv2.imwrite(str(original_path), frame)
     cv2.imwrite(str(mask_path), output.mask)
     cv2.imwrite(str(overlay_path), output.overlay)
-    return {
+    payload = {
         "frame_index": frame_index,
         "method": method,
         "original_url": _public(original_path),
@@ -118,6 +122,44 @@ def _save_frame_outputs(
         "overlay_url": _public(overlay_path),
         "metrics": output.metrics,
     }
+    if source_time is not None:
+        payload["source_time"] = round(float(source_time), 4)
+    if result_time is not None:
+        payload["result_time"] = round(float(result_time), 4)
+    if result_index is not None:
+        payload["result_index"] = int(result_index)
+    return payload
+
+
+def _browser_video(raw_path: Path, final_path: Path) -> Path:
+    ffmpeg = shutil.which("ffmpeg")
+    if ffmpeg:
+        subprocess.run(
+            [
+                ffmpeg,
+                "-y",
+                "-i",
+                str(raw_path),
+                "-c:v",
+                "libx264",
+                "-pix_fmt",
+                "yuv420p",
+                "-movflags",
+                "+faststart",
+                "-preset",
+                "veryfast",
+                "-crf",
+                "20",
+                str(final_path),
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        raw_path.unlink(missing_ok=True)
+    else:
+        raw_path.replace(final_path)
+    return final_path
 
 
 def _process_image(job_path: Path, source: Path, method: str, sensitivity: float) -> dict[str, Any]:
@@ -126,13 +168,13 @@ def _process_image(job_path: Path, source: Path, method: str, sensitivity: float
         raise HTTPException(status_code=400, detail="无法读取图片")
     if method == "compare":
         results = [
-            _save_frame_outputs(job_path, 0, item.method, frame, item)
+            _save_frame_outputs(job_path, 0, item.method, frame, item, 0.0, 0.0, 0)
             for item in compare_frame(frame, None, sensitivity)
         ]
     else:
         output = segment_frame(frame, method, None, sensitivity)
-        results = [_save_frame_outputs(job_path, 0, method, frame, output)]
-    return {"kind": "image", "frames": results, "video_url": None}
+        results = [_save_frame_outputs(job_path, 0, method, frame, output, 0.0, 0.0, 0)]
+    return {"kind": "image", "frames": results, "video_url": None, "source_fps": 1.0, "result_fps": 1.0, "duration": 0.0}
 
 
 def _selected_frame(index: int, stride: int, selected_count: int, max_frames: int) -> bool:
@@ -156,7 +198,14 @@ def _process_video(
     frame_index = 0
     selected_count = 0
     writer = None
+    result_fps = 8.0
+    raw_video_path = job_path / f"{method}_overlay.raw.mp4"
     video_path = job_path / f"{method}_overlay.mp4"
+    source_fps = float(capture.get(cv2.CAP_PROP_FPS) or 0.0)
+    if source_fps <= 0:
+        source_fps = 12.0
+    frame_count = int(capture.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+    duration = round(frame_count / source_fps, 4) if frame_count else 0.0
 
     try:
         while True:
@@ -171,16 +220,38 @@ def _process_video(
             if method == "compare":
                 outputs = compare_frame(frame, previous, sensitivity)
                 for output in outputs:
-                    frames.append(_save_frame_outputs(job_path, frame_index, output.method, frame, output))
+                    frames.append(
+                        _save_frame_outputs(
+                            job_path,
+                            frame_index,
+                            output.method,
+                            frame,
+                            output,
+                            frame_index / source_fps,
+                            selected_count / result_fps,
+                            selected_count,
+                        )
+                    )
                 video_output = next(item for item in outputs if item.method == "fusion")
             else:
                 video_output = segment_frame(frame, method, previous, sensitivity)
-                frames.append(_save_frame_outputs(job_path, frame_index, method, frame, video_output))
+                frames.append(
+                    _save_frame_outputs(
+                        job_path,
+                        frame_index,
+                        method,
+                        frame,
+                        video_output,
+                        frame_index / source_fps,
+                        selected_count / result_fps,
+                        selected_count,
+                    )
+                )
 
             if writer is None:
                 height, width = video_output.overlay.shape[:2]
                 fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-                writer = cv2.VideoWriter(str(video_path), fourcc, 8.0, (width, height))
+                writer = cv2.VideoWriter(str(raw_video_path), fourcc, result_fps, (width, height))
             writer.write(video_output.overlay)
 
             selected_count += 1
@@ -193,7 +264,39 @@ def _process_video(
 
     if not frames:
         raise HTTPException(status_code=400, detail="视频没有可处理帧")
-    return {"kind": "video", "frames": frames, "video_url": _public(video_path) if video_path.exists() else None}
+    if raw_video_path.exists():
+        _browser_video(raw_video_path, video_path)
+    return {
+        "kind": "video",
+        "frames": frames,
+        "video_url": _public(video_path) if video_path.exists() else None,
+        "source_fps": round(source_fps, 4),
+        "result_fps": result_fps,
+        "duration": duration,
+        "result_duration": round(len({frame["result_index"] for frame in frames}) / result_fps, 4) if frames else 0.0,
+    }
+
+
+def _read_video_frame(source: Path, frame_index: int) -> tuple[Any, Any | None, float]:
+    capture = cv2.VideoCapture(str(source))
+    if not capture.isOpened():
+        raise HTTPException(status_code=400, detail="无法读取视频")
+    source_fps = float(capture.get(cv2.CAP_PROP_FPS) or 0.0)
+    if source_fps <= 0:
+        source_fps = 12.0
+    frame_index = max(0, int(frame_index))
+    previous = None
+    if frame_index > 0:
+        capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index - 1)
+        ok_prev, previous = capture.read()
+        if not ok_prev:
+            previous = None
+    capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
+    ok, frame = capture.read()
+    capture.release()
+    if not ok:
+        raise HTTPException(status_code=400, detail="无法读取指定帧")
+    return frame, previous, source_fps
 
 
 @app.post("/api/segment")
@@ -231,6 +334,44 @@ def segment(
     }
 
 
+@app.post("/api/compare-frame")
+def compare_single_frame(
+    file: UploadFile = File(...),
+    frame_index: int = Form(0),
+    sensitivity: float = Form(0.56),
+) -> dict[str, Any]:
+    ensure_dirs()
+    sensitivity = max(0.05, min(float(sensitivity), 0.95))
+    job_id = uuid.uuid4().hex[:12]
+    job_path = JOB_DIR / job_id
+    job_path.mkdir(parents=True, exist_ok=True)
+    source = _save_upload(file, job_path)
+    suffix = source.suffix.lower()
+    if suffix in IMAGE_SUFFIXES:
+        frame = cv2.imread(str(source), cv2.IMREAD_COLOR)
+        previous = None
+        source_time = 0.0
+    elif suffix in VIDEO_SUFFIXES:
+        frame, previous, source_fps = _read_video_frame(source, frame_index)
+        source_time = int(frame_index) / source_fps
+    else:
+        raise HTTPException(status_code=400, detail="不支持的文件类型")
+    if frame is None:
+        raise HTTPException(status_code=400, detail="无法读取文件")
+    outputs = compare_frame(frame, previous, sensitivity)
+    frames = [
+        _save_frame_outputs(job_path, int(frame_index), output.method, frame, output, source_time, 0.0, 0)
+        for output in outputs
+    ]
+    return {
+        "job_id": job_id,
+        "kind": "compare",
+        "frame_index": int(frame_index),
+        "source_time": round(float(source_time), 4),
+        "frames": frames,
+    }
+
+
 @app.get("/")
 def index() -> FileResponse:
     return FileResponse(FRONTEND_DIR / "index.html")