shift to ffmpeg

app.py

@@ -1,4 +1,3 @@
-
 import os
 import torch
 import numpy as np
@@ -10,6 +9,7 @@ from sam2.sam2_image_predictor import SAM2ImagePredictor
 import cv2
 import traceback
 import matplotlib.pyplot as plt
+import ffmpeg
 from utils import load_model_without_flash_attn
 
 
@@ -62,7 +62,7 @@ def apply_color_mask(frame, mask, obj_id):
  return frame * (1 - mask) + colored_mask * 255
 
 def run_florence(image, text_input):
- with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+ with torch.amp.autocast(dtype=torch.bfloat16):
  task_prompt = '<OPEN_VOCABULARY_DETECTION>'
  prompt = task_prompt + text_input
  inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to('cuda', torch.bfloat16)
@@ -89,125 +89,135 @@ def remove_directory_contents(directory):
  for name in dirs:
  os.rmdir(os.path.join(root, name))
 
-def process_video(video_path, prompt, chunk_size=30):
+
+def process_video(video_path, prompt):
  try:
- video = cv2.VideoCapture(video_path)
- if not video.isOpened():
- raise ValueError("Unable to open video file")
+ # 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'])
+
+ print(f"Video info: {width}x{height}, {num_frames} frames, {fps} fps")
+
+ # Read frames
+ out, _ = (
+ ffmpeg
+ .input(video_path)
+ .output('pipe:', format='rawvideo', pix_fmt='rgb24')
+ .run(capture_stdout=True)
+ )
+ frames = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
+
+ print(f"Read {len(frames)} frames")
+
+ # Florence detection on first frame
+ first_frame = Image.fromarray(frames[0])
+ mask_box = run_florence(first_frame, prompt)
+ print("Original mask box:", mask_box)
  
- fps = video.get(cv2.CAP_PROP_FPS)
- frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+ # Convert mask_box to numpy array
+ mask_box = np.array(mask_box)
+ print("Reshaped mask box:", mask_box)
  
- # Process video in chunks
- all_segmented_frames = []
- for chunk_start in range(0, frame_count, chunk_size):
- chunk_end = min(chunk_start + chunk_size, frame_count)
- 
- frames = []
- video.set(cv2.CAP_PROP_POS_FRAMES, chunk_start)
- for _ in range(chunk_end - chunk_start):
- ret, frame = video.read()
- if not ret:
- break
- frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
- 
- if not frames:
- print(f"No frames extracted for chunk starting at {chunk_start}")
- continue
- 
- # Florence detection on first frame of the chunk
- first_frame = Image.fromarray(frames[0])
- mask_box = run_florence(first_frame, prompt)
- print("Original mask box:", mask_box)
- 
- # Convert mask_box to numpy array and ensure it's in the correct format
- mask_box = np.array(mask_box)
- print("Reshaped mask box:", mask_box)
- 
- # SAM2 segmentation on first frame
- with torch.cuda.amp.autocast(dtype=torch.bfloat16):
- image_predictor.set_image(first_frame)
- masks, _, _ = image_predictor.predict(
- point_coords=None,
- point_labels=None,
- box=mask_box[None, :],
- multimask_output=False,
- )
- print("masks.shape",masks.shape)
- 
- mask = masks.squeeze().astype(bool)
- print("Mask shape:", mask.shape)
- print("Frame shape:", frames[0].shape)
- 
- # SAM2 video propagation
- temp_dir = f"temp_frames_{chunk_start}"
- os.makedirs(temp_dir, exist_ok=True)
- for i, frame in enumerate(frames):
- cv2.imwrite(os.path.join(temp_dir, f"{i:04d}.jpg"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
- 
- with torch.cuda.amp.autocast(dtype=torch.bfloat16):
- inference_state = video_predictor.init_state(video_path=temp_dir)
- _, _, _ = video_predictor.add_new_mask(
- inference_state=inference_state,
- frame_idx=0,
- obj_id=1,
- mask=mask
- )
- 
- video_segments = {}
- for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor.propagate_in_video(inference_state):
- video_segments[out_frame_idx] = {
- out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
- for i, out_obj_id in enumerate(out_obj_ids)
- }
-
- print('segmenting for main vid done')
- 
- # Apply segmentation masks to frames
- for i, frame in enumerate(frames):
- if i in video_segments:
- for out_obj_id, mask in video_segments[i].items():
- frame = apply_color_mask(frame, mask, out_obj_id)
- all_segmented_frames.append(frame.astype(np.uint8))
- else:
- all_segmented_frames.append(frame)
+ # SAM2 segmentation on first frame
+ with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+ image_predictor.set_image(first_frame)
+ masks, _, _ = image_predictor.predict(
+ point_coords=None,
+ point_labels=None,
+ box=mask_box[None, :],
+ multimask_output=False,
+ )
+ print("masks.shape", masks.shape)
+ 
+ mask = masks.squeeze().astype(bool)
+ print("Mask shape:", mask.shape)
+ print("Frame shape:", frames[0].shape)
+ 
+ # SAM2 video propagation
+ temp_dir = "temp_frames"
+ os.makedirs(temp_dir, exist_ok=True)
+ for i, frame in enumerate(frames):
+ Image.fromarray(frame).save(os.path.join(temp_dir, f"{i:04d}.jpg"))
+ 
+ print(f"Saved {len(frames)} temporary frames")
+ 
+ with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+ inference_state = video_predictor.init_state(video_path=temp_dir)
+ _, _, _ = video_predictor.add_new_mask(
+ inference_state=inference_state,
+ frame_idx=0,
+ obj_id=1,
+ mask=mask
+ )
  
- # Clean up temporary files
- remove_directory_contents(temp_dir)
- os.rmdir(temp_dir)
+ video_segments = {}
+ for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor.propagate_in_video(inference_state):
+ video_segments[out_frame_idx] = {
+ out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+ for i, out_obj_id in enumerate(out_obj_ids)
+ }
+
+ print('Segmenting for main vid done')
+ print(f"Number of segmented frames: {len(video_segments)}")
  
- video.release()
+ # Apply segmentation masks to frames
+ all_segmented_frames = []
+ for i, frame in enumerate(frames):
+ if i in video_segments:
+ for out_obj_id, mask in video_segments[i].items():
+ frame = apply_color_mask(frame, mask, out_obj_id)
+ all_segmented_frames.append(frame.astype(np.uint8))
+ else:
+ all_segmented_frames.append(frame)
  
- if not all_segmented_frames:
- raise ValueError("No frames were processed successfully")
-
- # Create video from segmented frames
+ print(f"Applied masks to {len(all_segmented_frames)} frames")
+ 
+ # Clean up temporary files
+ remove_directory_contents(temp_dir)
+ os.rmdir(temp_dir)
+ 
+ # Write output video using ffmpeg
  output_path = "segmented_video.mp4"
- out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, 
- (all_segmented_frames[0].shape[1], all_segmented_frames[0].shape[0]))
+ process = (
+ ffmpeg
+ .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}', r=fps)
+ .output(output_path, pix_fmt='yuv420p')
+ .overwrite_output()
+ .run_async(pipe_stdin=True)
+ )
+ 
  for frame in all_segmented_frames:
- out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
- out.release()
+ process.stdin.write(frame.tobytes())
+ 
+ process.stdin.close()
+ process.wait()
  
+ if not os.path.exists(output_path):
+ raise ValueError(f"Output video file was not created: {output_path}")
+ 
+ print(f"Successfully created output video: {output_path}")
  return output_path
 
  except Exception as e:
  print(f"Error in process_video: {str(e)}")
  print(traceback.format_exc()) # This will print the full stack trace
  return None
-
-def segment_video(video_file, prompt, chunk_size):
+ 
+def segment_video(video_file, prompt):
  if video_file is None:
  return None
- output_video = process_video(video_file, prompt, int(chunk_size))
+ output_video = process_video(video_file, prompt)
  return output_video
 
 demo = gr.Interface(
  fn=segment_video,
  inputs=[
  gr.Video(label="Upload Video"),
- gr.Textbox(label="Enter prompt (e.g., 'a gymnast')"),
- gr.Slider(minimum=10, maximum=100, step=10, value=30, label="Chunk Size (frames)")
+ gr.Textbox(label="Enter prompt (e.g., 'a gymnast')")
  ],
  outputs=gr.Video(label="Segmented Video"),
  title="Video Object Segmentation with Florence and SAM2",

myapp2.py

@@ -0,0 +1,204 @@
+import os
+import torch
+import numpy as np
+import gradio as gr
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForCausalLM
+from sam2.build_sam import build_sam2_video_predictor, build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+import cv2
+import traceback
+import matplotlib.pyplot as plt
+
+# CUDA optimizations
+torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
+if torch.cuda.get_device_properties(0).major >= 8:
+ torch.backends.cuda.matmul.allow_tf32 = True
+ torch.backends.cudnn.allow_tf32 = True
+
+# Initialize models
+sam2_checkpoint = "../checkpoints/sam2_hiera_large.pt"
+model_cfg = "sam2_hiera_l.yaml"
+
+video_predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
+sam2_model = build_sam2(model_cfg, sam2_checkpoint, device="cuda")
+image_predictor = SAM2ImagePredictor(sam2_model)
+
+model_id = 'microsoft/Florence-2-large'
+florence_model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.bfloat16).eval().cuda()
+florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+
+def apply_color_mask(frame, mask, obj_id):
+ cmap = plt.get_cmap("tab10")
+ color = np.array(cmap(obj_id % 10)[:3]) # Use modulo 10 to cycle through colors
+ 
+ # Ensure mask has the correct shape
+ if mask.ndim == 4:
+ mask = mask.squeeze() # Remove singleton dimensions
+ if mask.ndim == 3 and mask.shape[0] == 1:
+ mask = mask[0] # Take the first channel if it's a single-channel 3D array
+ 
+ # Reshape mask to match frame dimensions
+ mask = cv2.resize(mask.astype(np.float32), (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR)
+ 
+ # Expand dimensions of mask and color for broadcasting
+ mask = np.expand_dims(mask, axis=2)
+ color = color.reshape(1, 1, 3)
+ 
+ colored_mask = mask * color
+ return frame * (1 - mask) + colored_mask * 255
+
+def run_florence(image, text_input):
+ with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+ task_prompt = '<OPEN_VOCABULARY_DETECTION>'
+ prompt = task_prompt + text_input
+ inputs = florence_processor(text=prompt, images=image, return_tensors="pt").to('cuda', torch.bfloat16)
+ generated_ids = florence_model.generate(
+ input_ids=inputs["input_ids"].cuda(),
+ pixel_values=inputs["pixel_values"].cuda(),
+ max_new_tokens=1024,
+ early_stopping=False,
+ do_sample=False,
+ num_beams=3,
+ )
+ generated_text = florence_processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+ parsed_answer = florence_processor.post_process_generation(
+ generated_text, 
+ task=task_prompt, 
+ image_size=(image.width, image.height)
+ )
+ return parsed_answer[task_prompt]['bboxes'][0]
+
+def remove_directory_contents(directory):
+ for root, dirs, files in os.walk(directory, topdown=False):
+ for name in files:
+ os.remove(os.path.join(root, name))
+ for name in dirs:
+ os.rmdir(os.path.join(root, name))
+
+def process_video(video_path, prompt, chunk_size=30):
+ try:
+ video = cv2.VideoCapture(video_path)
+ if not video.isOpened():
+ raise ValueError("Unable to open video file")
+ 
+ fps = video.get(cv2.CAP_PROP_FPS)
+ frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+ 
+ # Process video in chunks
+ all_segmented_frames = []
+ for chunk_start in range(0, frame_count, chunk_size):
+ chunk_end = min(chunk_start + chunk_size, frame_count)
+ 
+ frames = []
+ video.set(cv2.CAP_PROP_POS_FRAMES, chunk_start)
+ for _ in range(chunk_end - chunk_start):
+ ret, frame = video.read()
+ if not ret:
+ break
+ frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+ 
+ if not frames:
+ print(f"No frames extracted for chunk starting at {chunk_start}")
+ continue
+ 
+ # Florence detection on first frame of the chunk
+ first_frame = Image.fromarray(frames[0])
+ mask_box = run_florence(first_frame, prompt)
+ print("Original mask box:", mask_box)
+ 
+ # Convert mask_box to numpy array and ensure it's in the correct format
+ mask_box = np.array(mask_box)
+ print("Reshaped mask box:", mask_box)
+ 
+ # SAM2 segmentation on first frame
+ with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+ image_predictor.set_image(first_frame)
+ masks, _, _ = image_predictor.predict(
+ point_coords=None,
+ point_labels=None,
+ box=mask_box[None, :],
+ multimask_output=False,
+ )
+ print("masks.shape",masks.shape)
+ 
+ mask = masks.squeeze().astype(bool)
+ print("Mask shape:", mask.shape)
+ print("Frame shape:", frames[0].shape)
+ 
+ # SAM2 video propagation
+ temp_dir = f"temp_frames_{chunk_start}"
+ os.makedirs(temp_dir, exist_ok=True)
+ for i, frame in enumerate(frames):
+ cv2.imwrite(os.path.join(temp_dir, f"{i:04d}.jpg"), cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+ 
+ with torch.cuda.amp.autocast(dtype=torch.bfloat16):
+ inference_state = video_predictor.init_state(video_path=temp_dir)
+ _, _, _ = video_predictor.add_new_mask(
+ inference_state=inference_state,
+ frame_idx=0,
+ obj_id=1,
+ mask=mask
+ )
+ 
+ video_segments = {}
+ for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor.propagate_in_video(inference_state):
+ video_segments[out_frame_idx] = {
+ out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
+ for i, out_obj_id in enumerate(out_obj_ids)
+ }
+
+ print('segmenting for main vid done')
+ 
+ # Apply segmentation masks to frames
+ for i, frame in enumerate(frames):
+ if i in video_segments:
+ for out_obj_id, mask in video_segments[i].items():
+ frame = apply_color_mask(frame, mask, out_obj_id)
+ all_segmented_frames.append(frame.astype(np.uint8))
+ else:
+ all_segmented_frames.append(frame)
+ 
+ # Clean up temporary files
+ remove_directory_contents(temp_dir)
+ os.rmdir(temp_dir)
+ 
+ video.release()
+ 
+ if not all_segmented_frames:
+ raise ValueError("No frames were processed successfully")
+
+ # Create video from segmented frames
+ output_path = "segmented_video.mp4"
+ out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, 
+ (all_segmented_frames[0].shape[1], all_segmented_frames[0].shape[0]))
+ for frame in all_segmented_frames:
+ out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+ out.release()
+ 
+ return output_path
+
+ except Exception as e:
+ print(f"Error in process_video: {str(e)}")
+ print(traceback.format_exc()) # This will print the full stack trace
+ return None
+
+def segment_video(video_file, prompt, chunk_size):
+ if video_file is None:
+ return None
+ output_video = process_video(video_file, prompt, int(chunk_size))
+ return output_video
+
+demo = gr.Interface(
+ fn=segment_video,
+ inputs=[
+ gr.Video(label="Upload Video"),
+ gr.Textbox(label="Enter prompt (e.g., 'a gymnast')"),
+ gr.Slider(minimum=10, maximum=100, step=10, value=30, label="Chunk Size (frames)")
+ ],
+ outputs=gr.Video(label="Segmented Video"),
+ title="Video Object Segmentation with Florence and SAM2",
+ description="Upload a video and provide a text prompt to segment a specific object throughout the video."
+)
+
+demo.launch()

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -8,4 +8,5 @@ opencv-python
 matplotlib
 einops
 timm
-pytest
+pytest
+ffmpeg-python