torch autocast

.gitignore

@@ -0,0 +1,3 @@
+temp_frames
+temp_frames_30
+segmented_video.mp4

app.py

@@ -41,6 +41,8 @@ florence_model = load_model_without_flash_attn(load_florence_model)
 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
@@ -61,25 +63,26 @@ def apply_color_mask(frame, mask, obj_id):
  colored_mask = mask * color
  return frame * (1 - mask) + colored_mask * 255
 
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def run_florence(image, text_input):
- 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)
- 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)
- )
+ 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):
@@ -89,7 +92,8 @@ def remove_directory_contents(directory):
  for name in dirs:
  os.rmdir(os.path.join(root, name))
 
-
+@torch.inference_mode()
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process_video(video_path, prompt):
  try:
  # Get video info
@@ -123,14 +127,13 @@ def process_video(video_path, prompt):
  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,
- )
+ 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)
@@ -145,21 +148,20 @@ def process_video(video_path, prompt):
  
  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
- )
- 
- 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)
- }
+ 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')
  print(f"Number of segmented frames: {len(video_segments)}")
@@ -216,12 +218,19 @@ def segment_video(video_file, prompt):
 demo = gr.Interface(
  fn=segment_video,
  inputs=[
- gr.Video(label="Upload Video"),
- gr.Textbox(label="Enter prompt (e.g., 'a gymnast')")
+ gr.Video(label="Upload Video (Keep it under 10 seconds for this demo)"),
+ gr.Textbox(label="Enter text prompt for object detection")
  ],
  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."
+ title="Text-Prompted Video Object Segmentation",
+ description="""
+ This demo uses [Florence-2](https://huggingface.co/microsoft/Florence-2-large), a vision-language model, to enable text-prompted object detection for [SAM2](https://github.com/facebookresearch/segment-anything).
+ Florence-2 interprets your text prompt, allowing SAM2 to segment the described object in the video.
+ 
+ 1. Upload a short video (< 10 sec)
+ 2. Describe the object to segment
+ 3. Get your segmented video!
+ """
 )
 
 demo.launch()