update test_stitch.py

2026-03-25 21:32:09 +08:00
parent 966fc74313
commit 04a47cc939
2 changed files with 393 additions and 32 deletions
@@ -1,46 +1,81 @@
+import argparse
 import av
 import cv2
 import time

 DEVICE = "/dev/video4"
 TARGET_FPS = 30
-WIDTH = 1280
-HEIGHT = 720
+WIDTH = 1920
+HEIGHT = 1080

-# Open v4l2 device with same style options as ffmpeg CLI
-container = av.open(
-    DEVICE,
-    format="v4l2",
-    options={
-        "input_format": "mjpeg",
-        "video_size": f"{WIDTH}x{HEIGHT}",
-        "framerate": str(TARGET_FPS),
-    },
-)

-stream = container.streams.video[0]
-frame_count = 0
-t_start = time.perf_counter()
+def main():
+    parser = argparse.ArgumentParser(description="Capture from V4L2 via PyAV, optionally save to MP4.")
+    parser.add_argument(
+        "-o",
+        "--output",
+        type=str,
+        default="",
+        help="Output MP4 path (e.g. capture.mp4). If empty, only preview, no file.",
+    )
+    args = parser.parse_args()

-for frame in container.decode(stream):
-    t0 = time.perf_counter()
-    img = frame.to_ndarray(format="bgr24")  # OpenCV-compatible
-    t1 = time.perf_counter()
+    # Open v4l2 device with same style options as ffmpeg CLI
+    container = av.open(
+        DEVICE,
+        format="v4l2",
+        options={
+            "input_format": "mjpeg",
+            "video_size": f"{WIDTH}x{HEIGHT}",
+            "framerate": str(TARGET_FPS),
+        },
+    )

-    cv2.imshow("pyav_cam", img)
-    key = cv2.waitKey(1) & 0xFF
-    t2 = time.perf_counter()
+    stream = container.streams.video[0]
+    frame_count = 0
+    t_start = time.perf_counter()

-    frame_count += 1
-    if frame_count % 30 == 0:
-        elapsed = t2 - t_start
-        print(
-            f"Frame {frame_count}: convert={1000*(t1-t0):.1f}ms, "
-            f"imshow+waitKey={1000*(t2-t1):.1f}ms, avg_fps={frame_count/elapsed:.1f}"
+    writer = None
+    if args.output:
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        writer = cv2.VideoWriter(
+            args.output, fourcc, float(TARGET_FPS), (WIDTH, HEIGHT)
        )
+        if not writer.isOpened():
+            container.close()
+            raise RuntimeError(f"Could not open VideoWriter for {args.output}")
+        print(f"Recording to {args.output} @ {TARGET_FPS} fps, {WIDTH}x{HEIGHT}")

-    if key == ord("q"):
-        break
+    try:
+        for frame in container.decode(stream):
+            t0 = time.perf_counter()
+            img = frame.to_ndarray(format="bgr24")  # OpenCV-compatible
+            t1 = time.perf_counter()

-container.close()
-cv2.destroyAllWindows()
+            if writer is not None:
+                writer.write(img)
+
+            cv2.imshow("pyav_cam", img)
+            key = cv2.waitKey(1) & 0xFF
+            t2 = time.perf_counter()
+
+            frame_count += 1
+            if frame_count % 30 == 0:
+                elapsed = t2 - t_start
+                print(
+                    f"Frame {frame_count}: convert={1000*(t1-t0):.1f}ms, "
+                    f"imshow+waitKey={1000*(t2-t1):.1f}ms, avg_fps={frame_count/elapsed:.1f}"
+                )
+
+            if key == ord("q"):
+                break
+    finally:
+        container.close()
+        if writer is not None:
+            writer.release()
+            print(f"Saved {frame_count} frames to {args.output}")
+        cv2.destroyAllWindows()
+
+
+if __name__ == "__main__":
+    main()