vision-agent 0.2.126__tar.gz → 0.2.127__tar.gz

{vision_agent-0.2.126 → vision_agent-0.2.127}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vision-agent
-Version: 0.2.126
+Version: 0.2.127
 Summary: Toolset for Vision Agent
 Author: Landing AI
 Author-email: dev@landing.ai
@@ -10,6 +10,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Requires-Dist: anthropic (>=0.31.0,<0.32.0)
+Requires-Dist: av (>=11.0.0,<12.0.0)
 Requires-Dist: e2b (>=0.17.2a50,<0.18.0)
 Requires-Dist: e2b-code-interpreter (==0.0.11a37)
 Requires-Dist: eva-decord (>=0.6.1,<0.7.0)

{vision_agent-0.2.126 → vision_agent-0.2.127}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "vision-agent"
-version = "0.2.126"
+version = "0.2.127"
 description = "Toolset for Vision Agent"
 authors = ["Landing AI <dev@landing.ai>"]
 readme = "README.md"
@@ -42,6 +42,7 @@ pytube = "15.0.0"
 anthropic = "^0.31.0"
 pydantic = "2.7.4"
 eva-decord = "^0.6.1"
+av = "^11.0.0"
 
 [tool.poetry.group.dev.dependencies]
 autoflake = "1.*"

{vision_agent-0.2.126 → vision_agent-0.2.127}/vision_agent/tools/tools.py

@@ -1518,7 +1518,9 @@ def save_video(
         raise ValueError(f"fps must be greater than 0 got {fps}")
 
     if output_video_path is None:
-        output_video_path = tempfile.NamedTemporaryFile(delete=False).name
+        output_video_path = tempfile.NamedTemporaryFile(
+            delete=False, suffix=".mp4"
+        ).name
 
     output_video_path = video_writer(frames, fps, output_video_path)
     _save_video_to_result(output_video_path)

{vision_agent-0.2.126 → vision_agent-0.2.127}/vision_agent/utils/video.py

@@ -5,6 +5,7 @@ from functools import lru_cache
 from typing import List, Optional, Tuple
 
 import cv2
+import av  # type: ignore
 import numpy as np
 from decord import VideoReader  # type: ignore
 
@@ -43,18 +44,36 @@ def play_video(video_base64: str) -> None:
         cv2.destroyAllWindows()
 
 
+def _resize_frame(frame: np.ndarray) -> np.ndarray:
+    height, width = frame.shape[:2]
+    new_width = width - (width % 2)
+    new_height = height - (height % 2)
+    return cv2.resize(frame, (new_width, new_height))
+
+
 def video_writer(
     frames: List[np.ndarray], fps: float = 1.0, filename: Optional[str] = None
 ) -> str:
     if filename is None:
         filename = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
-
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # type: ignore
+    container = av.open(filename, mode="w")
+    stream = container.add_stream("h264", rate=fps)
     height, width = frames[0].shape[:2]
-    writer = cv2.VideoWriter(filename, fourcc, fps, (width, height))
+    stream.height = height - (height % 2)
+    stream.width = width - (width % 2)
+    stream.pix_fmt = "yuv420p"
     for frame in frames:
-        writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-    writer.release()
+        # Remove the alpha channel (convert RGBA to RGB)
+        frame_rgb = frame[:, :, :3]
+        # Resize the frame to make dimensions divisible by 2
+        frame_rgb = _resize_frame(frame_rgb)
+        av_frame = av.VideoFrame.from_ndarray(frame_rgb, format="rgb24")
+        for packet in stream.encode(av_frame):
+            container.mux(packet)
+
+    for packet in stream.encode():
+        container.mux(packet)
+    container.close()
     return filename