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