Spaces:
Sleeping
Sleeping
| 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 | |
| from utils import load_model_without_flash_attn | |
| # 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' | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| def load_florence_model(): | |
| return AutoModelForCausalLM.from_pretrained( | |
| model_id, | |
| trust_remote_code=True, | |
| torch_dtype=torch.float16 if device == "cuda" else torch.float32 | |
| ).eval().to(device) | |
| 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 | |
| # 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() |