File size: 8,283 Bytes
9b87d5a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
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()