downscaling input video

app.py

@@ -104,26 +104,47 @@ def remove_directory_contents(directory):
 @spaces.GPU
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
-def process_video(video_path, prompt):
+def process_video(video_path, prompt, target_fps=30, max_dimension=640):
     try:
         # Get video info
         probe = ffmpeg.probe(video_path)
         video_info = next(s for s in probe['streams'] if s['codec_type'] == 'video')
         width = int(video_info['width'])
         height = int(video_info['height'])
-        num_frames = int(video_info['nb_frames'])
-        fps = eval(video_info['r_frame_rate'])
+        original_fps = eval(video_info['r_frame_rate'])
 
-        print(f"Video info: {width}x{height}, {num_frames} frames, {fps} fps")
+        # Calculate new dimensions while maintaining aspect ratio
+        if width > height:
+            if width > max_dimension:
+                new_width = max_dimension
+                new_height = int(height * (max_dimension / width))
+            else:
+                new_width = width
+                new_height = height
+        else:
+            if height > max_dimension:
+                new_height = max_dimension
+                new_width = int(width * (max_dimension / height))
+            else:
+                new_width = width
+                new_height = height
+
+        # Determine target fps
+        fps = min(original_fps, target_fps)
+
+        print(f"Original video: {width}x{height}, {original_fps} fps")
+        print(f"Processing at: {new_width}x{new_height}, {fps} fps")
 
-        # Read frames
+        # Read and resize frames
         out, _ = (
             ffmpeg
             .input(video_path)
+            .filter('fps', fps=fps)
+            .filter('scale', width=new_width, height=new_height)
             .output('pipe:', format='rawvideo', pix_fmt='rgb24')
             .run(capture_stdout=True)
         )
-        frames = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
+        frames = np.frombuffer(out, np.uint8).reshape([-1, new_height, new_width, 3])
 
         print(f"Read {len(frames)} frames")
 
@@ -196,7 +217,7 @@ def process_video(video_path, prompt):
         output_path = "segmented_video.mp4"
         process = (
             ffmpeg
-            .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}', r=fps)
+            .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{new_width}x{new_height}', r=fps)
             .output(output_path, pix_fmt='yuv420p')
             .overwrite_output()
             .run_async(pipe_stdin=True)
@@ -219,7 +240,7 @@ def process_video(video_path, prompt):
         print(traceback.format_exc())  # This will print the full stack trace
         return None
     
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=600)
 def segment_video(video_file, prompt):
     if video_file is None:
         return None