Spaces:
Running
on
Zero
Running
on
Zero
Applying for GPU
Browse files
app.py
DELETED
@@ -1,638 +0,0 @@
|
|
1 |
-
import subprocess
|
2 |
-
import re
|
3 |
-
from typing import List, Tuple, Optional
|
4 |
-
|
5 |
-
import gradio as gr
|
6 |
-
from datetime import datetime
|
7 |
-
import os
|
8 |
-
os.environ["TORCH_CUDNN_SDPA_ENABLED"] = "0,1,2,3,4,5,6,7"
|
9 |
-
import torch
|
10 |
-
import numpy as np
|
11 |
-
import cv2
|
12 |
-
import matplotlib.pyplot as plt
|
13 |
-
from PIL import Image, ImageFilter
|
14 |
-
from sam2.build_sam import build_sam2_video_predictor
|
15 |
-
|
16 |
-
from moviepy.editor import ImageSequenceClip
|
17 |
-
|
18 |
-
# Description
|
19 |
-
title = "<center><strong><font size='8'>Efficient Track Anything (EfficientTAM)<font></strong></center>"
|
20 |
-
|
21 |
-
description_e = """This is a demo of [Efficient Track Anything (EfficientTAM) Model](https://github.com/yformer/EfficientTAM).
|
22 |
-
"""
|
23 |
-
|
24 |
-
description_p = """# Interactive Video Segmentation
|
25 |
-
- Built our demo based on [SAM2-Video-Predictor](https://huggingface.co/spaces/fffiloni/SAM2-Video-Predictor). Thanks to Sylvain Filoni.
|
26 |
-
- Instruction
|
27 |
-
<ol>
|
28 |
-
<li> Upload one video or click one example video</li>
|
29 |
-
<li> Click 'include' point type, select the object to segment and track</li>
|
30 |
-
<li> Click 'exclude' point type (optional), select the area you want to avoid segmenting and tracking</li>
|
31 |
-
<li> Click the 'Segment' button, obtain the mask of the first frame </li>
|
32 |
-
<li> Click the 'coarse' level and the 'Track' button, segment and track the object every 15 frames </li>
|
33 |
-
<li> Click the corresponding frame to add points on the object for mask refining (optional) </li>
|
34 |
-
<li> Click the 'fine' level and the 'Track' button, obtain masklet and masked video </li>
|
35 |
-
<li> Click the 'Reset' button to restart </li>
|
36 |
-
</ol>
|
37 |
-
- Github [link](https://github.com/yformer/EfficientTAM)
|
38 |
-
"""
|
39 |
-
|
40 |
-
# examples
|
41 |
-
examples = [
|
42 |
-
["examples/videos/cat.mp4"],
|
43 |
-
["examples/videos/coffee.mp4"],
|
44 |
-
["examples/videos/car.mp4"],
|
45 |
-
["examples/videos/chick.mp4"],
|
46 |
-
["examples/videos/cups.mp4"],
|
47 |
-
["examples/videos/dog.mp4"],
|
48 |
-
["examples/videos/goat.mp4"],
|
49 |
-
["examples/videos/juggle.mp4"],
|
50 |
-
["examples/videos/street.mp4"],
|
51 |
-
["examples/videos/yacht.mp4"],
|
52 |
-
]
|
53 |
-
|
54 |
-
default_example = examples[0]
|
55 |
-
|
56 |
-
def get_video_fps(video_path):
|
57 |
-
# Open the video file
|
58 |
-
cap = cv2.VideoCapture(video_path)
|
59 |
-
|
60 |
-
if not cap.isOpened():
|
61 |
-
print("Error: Could not open video.")
|
62 |
-
return None
|
63 |
-
|
64 |
-
# Get the FPS of the video
|
65 |
-
fps = cap.get(cv2.CAP_PROP_FPS)
|
66 |
-
|
67 |
-
return fps
|
68 |
-
|
69 |
-
def clear_points(image):
|
70 |
-
# we clean all
|
71 |
-
return [
|
72 |
-
image, # first_frame_path
|
73 |
-
gr.State([]), # tracking_points
|
74 |
-
gr.State([]), # trackings_input_label
|
75 |
-
image, # points_map
|
76 |
-
#gr.State() # stored_inference_state
|
77 |
-
]
|
78 |
-
|
79 |
-
def preprocess_video_in(video_path):
|
80 |
-
if video_path is None:
|
81 |
-
return None, gr.State([]), gr.State([]), None, None, None, None, None, None, gr.update(open=True)
|
82 |
-
|
83 |
-
# Generate a unique ID based on the current date and time
|
84 |
-
unique_id = datetime.now().strftime('%Y%m%d%H%M%S')
|
85 |
-
|
86 |
-
# Set directory with this ID to store video frames
|
87 |
-
extracted_frames_output_dir = f'frames_{unique_id}'
|
88 |
-
|
89 |
-
# Create the output directory
|
90 |
-
os.makedirs(extracted_frames_output_dir, exist_ok=True)
|
91 |
-
|
92 |
-
### Process video frames ###
|
93 |
-
# Open the video file
|
94 |
-
cap = cv2.VideoCapture(video_path)
|
95 |
-
|
96 |
-
if not cap.isOpened():
|
97 |
-
print("Error: Could not open video.")
|
98 |
-
return None
|
99 |
-
|
100 |
-
# Get the frames per second (FPS) of the video
|
101 |
-
fps = cap.get(cv2.CAP_PROP_FPS)
|
102 |
-
|
103 |
-
# Calculate the number of frames to process (10 seconds of video)
|
104 |
-
max_frames = int(fps * 10)
|
105 |
-
|
106 |
-
frame_number = 0
|
107 |
-
first_frame = None
|
108 |
-
|
109 |
-
while True:
|
110 |
-
ret, frame = cap.read()
|
111 |
-
if not ret or frame_number >= max_frames:
|
112 |
-
break
|
113 |
-
|
114 |
-
# Format the frame filename as '00000.jpg'
|
115 |
-
frame_filename = os.path.join(extracted_frames_output_dir, f'{frame_number:05d}.jpg')
|
116 |
-
|
117 |
-
# Save the frame as a JPEG file
|
118 |
-
cv2.imwrite(frame_filename, frame)
|
119 |
-
|
120 |
-
# Store the first frame
|
121 |
-
if frame_number == 0:
|
122 |
-
first_frame = frame_filename
|
123 |
-
|
124 |
-
frame_number += 1
|
125 |
-
|
126 |
-
# Release the video capture object
|
127 |
-
cap.release()
|
128 |
-
|
129 |
-
# scan all the JPEG frame names in this directory
|
130 |
-
scanned_frames = [
|
131 |
-
p for p in os.listdir(extracted_frames_output_dir)
|
132 |
-
if os.path.splitext(p)[-1] in [".jpg", ".jpeg", ".JPG", ".JPEG"]
|
133 |
-
]
|
134 |
-
scanned_frames.sort(key=lambda p: int(os.path.splitext(p)[0]))
|
135 |
-
# print(f"SCANNED_FRAMES: {scanned_frames}")
|
136 |
-
|
137 |
-
return [
|
138 |
-
first_frame, # first_frame_path
|
139 |
-
gr.State([]), # tracking_points
|
140 |
-
gr.State([]), # trackings_input_label
|
141 |
-
first_frame, # input_first_frame_image
|
142 |
-
first_frame, # points_map
|
143 |
-
extracted_frames_output_dir, # video_frames_dir
|
144 |
-
scanned_frames, # scanned_frames
|
145 |
-
None, # stored_inference_state
|
146 |
-
None, # stored_frame_names
|
147 |
-
gr.update(open=False) # video_in_drawer
|
148 |
-
]
|
149 |
-
|
150 |
-
|
151 |
-
def get_point(point_type, tracking_points, trackings_input_label, input_first_frame_image, evt: gr.SelectData):
|
152 |
-
if input_first_frame_image is None:
|
153 |
-
return gr.State([]), gr.State([]), None
|
154 |
-
print(f"You selected {evt.value} at {evt.index} from {evt.target}")
|
155 |
-
|
156 |
-
tracking_points.value.append(evt.index)
|
157 |
-
print(f"TRACKING POINT: {tracking_points.value}")
|
158 |
-
|
159 |
-
if point_type == "include":
|
160 |
-
trackings_input_label.value.append(1)
|
161 |
-
elif point_type == "exclude":
|
162 |
-
trackings_input_label.value.append(0)
|
163 |
-
print(f"TRACKING INPUT LABEL: {trackings_input_label.value}")
|
164 |
-
|
165 |
-
# Open the image and get its dimensions
|
166 |
-
transparent_background = Image.open(input_first_frame_image).convert('RGBA')
|
167 |
-
w, h = transparent_background.size
|
168 |
-
|
169 |
-
# Define the circle radius as a fraction of the smaller dimension
|
170 |
-
fraction = 0.02 # You can adjust this value as needed
|
171 |
-
radius = int(fraction * min(w, h))
|
172 |
-
|
173 |
-
# Create a transparent layer to draw on
|
174 |
-
transparent_layer = np.zeros((h, w, 4), dtype=np.uint8)
|
175 |
-
|
176 |
-
for index, track in enumerate(tracking_points.value):
|
177 |
-
if trackings_input_label.value[index] == 1:
|
178 |
-
cv2.circle(transparent_layer, track, radius, (0, 255, 0, 255), -1)
|
179 |
-
else:
|
180 |
-
cv2.circle(transparent_layer, track, radius, (255, 0, 0, 255), -1)
|
181 |
-
|
182 |
-
# Convert the transparent layer back to an image
|
183 |
-
transparent_layer = Image.fromarray(transparent_layer, 'RGBA')
|
184 |
-
selected_point_map = Image.alpha_composite(transparent_background, transparent_layer)
|
185 |
-
|
186 |
-
return tracking_points, trackings_input_label, selected_point_map
|
187 |
-
|
188 |
-
DEVICE = 'cuda'
|
189 |
-
# use bfloat16 for the entire notebook
|
190 |
-
torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
|
191 |
-
if torch.cuda.get_device_properties(0).major >= 8:
|
192 |
-
# turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
|
193 |
-
torch.backends.cuda.matmul.allow_tf32 = True
|
194 |
-
torch.backends.cudnn.allow_tf32 = True
|
195 |
-
|
196 |
-
def show_mask(mask, ax, obj_id=None, random_color=False):
|
197 |
-
if random_color:
|
198 |
-
color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
|
199 |
-
else:
|
200 |
-
cmap = plt.get_cmap("tab10")
|
201 |
-
cmap_idx = 0 if obj_id is None else obj_id
|
202 |
-
color = np.array([*cmap(cmap_idx)[:3], 0.6])
|
203 |
-
h, w = mask.shape[-2:]
|
204 |
-
mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
|
205 |
-
ax.axis('off')
|
206 |
-
ax.imshow(mask_image)
|
207 |
-
|
208 |
-
|
209 |
-
def show_points(coords, labels, ax, marker_size=200):
|
210 |
-
pos_points = coords[labels==1]
|
211 |
-
neg_points = coords[labels==0]
|
212 |
-
ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
|
213 |
-
ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
|
214 |
-
|
215 |
-
def show_box(box, ax):
|
216 |
-
x0, y0 = box[0], box[1]
|
217 |
-
w, h = box[2] - box[0], box[3] - box[1]
|
218 |
-
ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))
|
219 |
-
|
220 |
-
|
221 |
-
def load_model(checkpoint):
|
222 |
-
# Load model accordingly to user's choice
|
223 |
-
if checkpoint == "efficienttam-s":
|
224 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam_s.pt"
|
225 |
-
model_cfg = "efficienttam-s.yaml"
|
226 |
-
return [efficienttam_checkpoint, model_cfg]
|
227 |
-
elif checkpoint == "efficienttam-ti":
|
228 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-ti.pt"
|
229 |
-
model_cfg = "efficienttam-ti.yaml"
|
230 |
-
return [efficienttam_checkpoint, model_cfg]
|
231 |
-
elif checkpoint == "efficienttam-s-512x512":
|
232 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-s-512x512.pt"
|
233 |
-
model_cfg = "efficienttam_s_512x512.yaml"
|
234 |
-
return [efficienttam_checkpoint, model_cfg]
|
235 |
-
elif checkpoint == "efficienttam-ti-512x512":
|
236 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-ti-512x512.pt"
|
237 |
-
model_cfg = "efficienttam_ti_512x512.yaml"
|
238 |
-
return [efficienttam_checkpoint, model_cfg]
|
239 |
-
elif checkpoint == "efficienttam-s-1":
|
240 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-s-1.pt"
|
241 |
-
model_cfg = "efficienttam-s-1.yaml"
|
242 |
-
return [efficienttam_checkpoint, model_cfg]
|
243 |
-
elif checkpoint == "efficienttam-s-2":
|
244 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-s-2.pt"
|
245 |
-
model_cfg = "efficienttam-s-2.yaml"
|
246 |
-
return [efficienttam_checkpoint, model_cfg]
|
247 |
-
elif checkpoint == "efficienttam-ti-1":
|
248 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-ti-1.pt"
|
249 |
-
model_cfg = "efficienttam-ti-1.yaml"
|
250 |
-
return [efficienttam_checkpoint, model_cfg]
|
251 |
-
elif checkpoint == "efficienttam-ti-2":
|
252 |
-
efficienttam_checkpoint = "./checkpoints/efficienttam-ti-2.pt"
|
253 |
-
model_cfg = "efficienttam-ti-2.yaml"
|
254 |
-
return [efficienttam_checkpoint, model_cfg]
|
255 |
-
else:
|
256 |
-
efficienttam_checkpoint = "./checkpoints/demo/efficienttam_s.pt"
|
257 |
-
model_cfg = "efficienttam-s.yaml"
|
258 |
-
return [efficienttam_checkpoint, model_cfg]
|
259 |
-
|
260 |
-
def get_mask_sam_process(
|
261 |
-
stored_inference_state,
|
262 |
-
input_first_frame_image,
|
263 |
-
checkpoint,
|
264 |
-
tracking_points,
|
265 |
-
trackings_input_label,
|
266 |
-
video_frames_dir, # extracted_frames_output_dir defined in 'preprocess_video_in' function
|
267 |
-
scanned_frames,
|
268 |
-
working_frame: str = None, # current frame being added points
|
269 |
-
available_frames_to_check: List[str] = [],
|
270 |
-
):
|
271 |
-
|
272 |
-
if len(tracking_points.value) == 0:
|
273 |
-
return gr.update(visible=False), None, gr.State(), None, stored_inference_state, working_frame
|
274 |
-
# get model and model config paths
|
275 |
-
print(f"USER CHOSEN CHECKPOINT: {checkpoint}")
|
276 |
-
sam2_checkpoint, model_cfg = load_model(checkpoint)
|
277 |
-
print("MODEL LOADED")
|
278 |
-
|
279 |
-
# set predictor
|
280 |
-
predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
|
281 |
-
print("PREDICTOR READY")
|
282 |
-
|
283 |
-
# `video_dir` a directory of JPEG frames with filenames like `<frame_index>.jpg`
|
284 |
-
# print(f"STATE FRAME OUTPUT DIRECTORY: {video_frames_dir}")
|
285 |
-
video_dir = video_frames_dir
|
286 |
-
|
287 |
-
# scan all the JPEG frame names in this directory
|
288 |
-
frame_names = scanned_frames
|
289 |
-
|
290 |
-
# print(f"STORED INFERENCE STEP: {stored_inference_state}")
|
291 |
-
if stored_inference_state is None:
|
292 |
-
# Init SAM2 inference_state
|
293 |
-
inference_state = predictor.init_state(video_path=video_dir, device="cuda")
|
294 |
-
print("NEW INFERENCE_STATE INITIATED")
|
295 |
-
else:
|
296 |
-
inference_state = stored_inference_state
|
297 |
-
|
298 |
-
# segment and track one object
|
299 |
-
# predictor.reset_state(inference_state) # if any previous tracking, reset
|
300 |
-
|
301 |
-
### HANDLING WORKING FRAME
|
302 |
-
# new_working_frame = None
|
303 |
-
# Add new point
|
304 |
-
if working_frame is None:
|
305 |
-
ann_frame_idx = 0 # the frame index we interact with, 0 if it is the first frame
|
306 |
-
working_frame = "frame_0.jpg"
|
307 |
-
else:
|
308 |
-
# Use a regular expression to find the integer
|
309 |
-
match = re.search(r'frame_(\d+)', working_frame)
|
310 |
-
if match:
|
311 |
-
# Extract the integer from the match
|
312 |
-
frame_number = int(match.group(1))
|
313 |
-
ann_frame_idx = frame_number
|
314 |
-
|
315 |
-
print(f"NEW_WORKING_FRAME PATH: {working_frame}")
|
316 |
-
|
317 |
-
ann_obj_id = 1 # give a unique id to each object we interact with (it can be any integers)
|
318 |
-
|
319 |
-
# Let's add a positive click at (x, y) = (210, 350) to get started
|
320 |
-
points = np.array(tracking_points.value, dtype=np.float32)
|
321 |
-
# for labels, `1` means positive click and `0` means negative click
|
322 |
-
labels = np.array(trackings_input_label.value, np.int32)
|
323 |
-
_, out_obj_ids, out_mask_logits = predictor.add_new_points(
|
324 |
-
inference_state=inference_state,
|
325 |
-
frame_idx=ann_frame_idx,
|
326 |
-
obj_id=ann_obj_id,
|
327 |
-
points=points,
|
328 |
-
labels=labels,
|
329 |
-
)
|
330 |
-
|
331 |
-
# Create the plot
|
332 |
-
plt.figure(figsize=(12, 8))
|
333 |
-
plt.title(f"frame {ann_frame_idx}")
|
334 |
-
plt.imshow(Image.open(os.path.join(video_dir, frame_names[ann_frame_idx])))
|
335 |
-
show_points(points, labels, plt.gca())
|
336 |
-
show_mask((out_mask_logits[0] > 0.0).cpu().numpy(), plt.gca(), obj_id=out_obj_ids[0])
|
337 |
-
|
338 |
-
# Save the plot as a JPG file
|
339 |
-
first_frame_output_filename = "output_first_frame.jpg"
|
340 |
-
plt.savefig(first_frame_output_filename, format='jpg')
|
341 |
-
plt.close()
|
342 |
-
torch.cuda.empty_cache()
|
343 |
-
|
344 |
-
# Assuming available_frames_to_check.value is a list
|
345 |
-
if working_frame not in available_frames_to_check:
|
346 |
-
available_frames_to_check.append(working_frame)
|
347 |
-
print(available_frames_to_check)
|
348 |
-
|
349 |
-
return gr.update(visible=True), "output_first_frame.jpg", frame_names, predictor, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=True)
|
350 |
-
|
351 |
-
def propagate_to_all(tracking_points, video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame):
|
352 |
-
if tracking_points is None or video_in is None or checkpoint is None or stored_inference_state is None:
|
353 |
-
return gr.update(value=None), gr.update(value=None), gr.update(value=None), available_frames_to_check, gr.update(visible=False)
|
354 |
-
#### PROPAGATION ####
|
355 |
-
sam2_checkpoint, model_cfg = load_model(checkpoint)
|
356 |
-
predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device="cuda")
|
357 |
-
|
358 |
-
inference_state = stored_inference_state
|
359 |
-
frame_names = stored_frame_names
|
360 |
-
video_dir = video_frames_dir
|
361 |
-
|
362 |
-
# Define a directory to save the JPEG images
|
363 |
-
frames_output_dir = "frames_output_images"
|
364 |
-
os.makedirs(frames_output_dir, exist_ok=True)
|
365 |
-
|
366 |
-
# Initialize a list to store file paths of saved images
|
367 |
-
jpeg_images = []
|
368 |
-
|
369 |
-
# run propagation throughout the video and collect the results in a dict
|
370 |
-
video_segments = {} # video_segments contains the per-frame segmentation results
|
371 |
-
print("starting propagate_in_video")
|
372 |
-
for out_frame_idx, out_obj_ids, out_mask_logits in predictor.propagate_in_video(inference_state):
|
373 |
-
video_segments[out_frame_idx] = {
|
374 |
-
out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy()
|
375 |
-
for i, out_obj_id in enumerate(out_obj_ids)
|
376 |
-
}
|
377 |
-
|
378 |
-
# obtain the segmentation results every few frames
|
379 |
-
if vis_frame_type == "coarse":
|
380 |
-
vis_frame_stride = 15
|
381 |
-
elif vis_frame_type == "fine":
|
382 |
-
vis_frame_stride = 1
|
383 |
-
|
384 |
-
plt.close("all")
|
385 |
-
for out_frame_idx in range(0, len(frame_names), vis_frame_stride):
|
386 |
-
plt.figure(figsize=(6, 4))
|
387 |
-
plt.title(f"frame {out_frame_idx}")
|
388 |
-
plt.imshow(Image.open(os.path.join(video_dir, frame_names[out_frame_idx])))
|
389 |
-
for out_obj_id, out_mask in video_segments[out_frame_idx].items():
|
390 |
-
show_mask(out_mask, plt.gca(), obj_id=out_obj_id)
|
391 |
-
|
392 |
-
# Define the output filename and save the figure as a JPEG file
|
393 |
-
output_filename = os.path.join(frames_output_dir, f"frame_{out_frame_idx}.jpg")
|
394 |
-
plt.savefig(output_filename, format='jpg')
|
395 |
-
|
396 |
-
# Close the plot
|
397 |
-
plt.close()
|
398 |
-
|
399 |
-
# Append the file path to the list
|
400 |
-
jpeg_images.append(output_filename)
|
401 |
-
|
402 |
-
if f"frame_{out_frame_idx}.jpg" not in available_frames_to_check:
|
403 |
-
available_frames_to_check.append(f"frame_{out_frame_idx}.jpg")
|
404 |
-
|
405 |
-
torch.cuda.empty_cache()
|
406 |
-
print(f"JPEG_IMAGES: {jpeg_images}")
|
407 |
-
|
408 |
-
if vis_frame_type == "coarse":
|
409 |
-
return gr.update(value=jpeg_images), gr.update(value=None), gr.update(choices=available_frames_to_check, value=working_frame, visible=True), available_frames_to_check, gr.update(visible=True)
|
410 |
-
elif vis_frame_type == "fine":
|
411 |
-
# Create a video clip from the image sequence
|
412 |
-
original_fps = get_video_fps(video_in)
|
413 |
-
fps = original_fps # Frames per second
|
414 |
-
total_frames = len(jpeg_images)
|
415 |
-
clip = ImageSequenceClip(jpeg_images, fps=fps)
|
416 |
-
# Write the result to a file
|
417 |
-
final_vid_output_path = "output_video.mp4"
|
418 |
-
|
419 |
-
# Write the result to a file
|
420 |
-
clip.write_videofile(
|
421 |
-
final_vid_output_path,
|
422 |
-
codec='libx264'
|
423 |
-
)
|
424 |
-
|
425 |
-
return gr.update(value=None), gr.update(value=final_vid_output_path), working_frame, available_frames_to_check, gr.update(visible=True)
|
426 |
-
|
427 |
-
def update_ui(vis_frame_type):
|
428 |
-
if vis_frame_type == "coarse":
|
429 |
-
return gr.update(visible=True), gr.update(visible=False)
|
430 |
-
elif vis_frame_type == "fine":
|
431 |
-
return gr.update(visible=False), gr.update(visible=True)
|
432 |
-
|
433 |
-
def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
|
434 |
-
new_working_frame = None
|
435 |
-
if working_frame == None:
|
436 |
-
new_working_frame = os.path.join(video_frames_dir, scanned_frames[0])
|
437 |
-
|
438 |
-
else:
|
439 |
-
# Use a regular expression to find the integer
|
440 |
-
match = re.search(r'frame_(\d+)', working_frame)
|
441 |
-
if match:
|
442 |
-
# Extract the integer from the match
|
443 |
-
frame_number = int(match.group(1))
|
444 |
-
ann_frame_idx = frame_number
|
445 |
-
new_working_frame = os.path.join(video_frames_dir, scanned_frames[ann_frame_idx])
|
446 |
-
return gr.State([]), gr.State([]), new_working_frame, new_working_frame
|
447 |
-
|
448 |
-
def reset_propagation(first_frame_path, predictor, stored_inference_state):
|
449 |
-
predictor.reset_state(stored_inference_state)
|
450 |
-
# print(f"RESET State: {stored_inference_state} ")
|
451 |
-
return first_frame_path, gr.State([]), gr.State([]), gr.update(value=None, visible=False), stored_inference_state, None, ["frame_0.jpg"], first_frame_path, "frame_0.jpg", gr.update(visible=False)
|
452 |
-
|
453 |
-
with gr.Blocks() as demo:
|
454 |
-
first_frame_path = gr.State()
|
455 |
-
tracking_points = gr.State([])
|
456 |
-
trackings_input_label = gr.State([])
|
457 |
-
video_frames_dir = gr.State()
|
458 |
-
scanned_frames = gr.State()
|
459 |
-
loaded_predictor = gr.State()
|
460 |
-
stored_inference_state = gr.State()
|
461 |
-
stored_frame_names = gr.State()
|
462 |
-
available_frames_to_check = gr.State([])
|
463 |
-
with gr.Column():
|
464 |
-
# Title
|
465 |
-
gr.Markdown(title)
|
466 |
-
with gr.Row():
|
467 |
-
|
468 |
-
with gr.Column():
|
469 |
-
# Instructions
|
470 |
-
gr.Markdown(description_p)
|
471 |
-
|
472 |
-
# video_exp = gr.Video(label="Input Example", format="mp4", visible=False)
|
473 |
-
with gr.Accordion("Input Video", open=True) as video_in_drawer:
|
474 |
-
video_in = gr.Video(label="Input Video", format="mp4")
|
475 |
-
|
476 |
-
with gr.Row():
|
477 |
-
point_type = gr.Radio(label="point type", choices=["include", "exclude"], value="include", scale=2)
|
478 |
-
clear_points_btn = gr.Button("Clear Points", scale=1)
|
479 |
-
|
480 |
-
input_first_frame_image = gr.Image(label="input image", interactive=False, type="filepath", visible=False)
|
481 |
-
|
482 |
-
points_map = gr.Image(
|
483 |
-
label="Frame with Point Prompt",
|
484 |
-
type="filepath",
|
485 |
-
interactive=False
|
486 |
-
)
|
487 |
-
|
488 |
-
with gr.Row():
|
489 |
-
checkpoint = gr.Dropdown(label="Checkpoint", choices=["efficienttam-s", "efficienttam-ti", "efficienttam-s-512x512", "efficienttam-ti-512x512", "efficienttam-s-1", "efficienttam-s-2", "efficienttam-ti-1", "efficienttam-ti-2"], value="efficienttam-s")
|
490 |
-
submit_btn = gr.Button("Segment", size="lg")
|
491 |
-
|
492 |
-
|
493 |
-
with gr.Column():
|
494 |
-
gr.Markdown("# Try some of the examples below ⬇️")
|
495 |
-
gr.Examples(
|
496 |
-
examples=examples,
|
497 |
-
inputs=[video_in,],
|
498 |
-
)
|
499 |
-
gr.Markdown('\n\n\n\n\n\n\n\n\n\n\n')
|
500 |
-
gr.Markdown('\n\n\n\n\n\n\n\n\n\n\n')
|
501 |
-
gr.Markdown('\n\n\n\n\n\n\n\n\n\n\n')
|
502 |
-
with gr.Row():
|
503 |
-
working_frame = gr.Dropdown(label="Frame ID", choices=[""], value=None, visible=False, allow_custom_value=False, interactive=True)
|
504 |
-
change_current = gr.Button("change current", visible=False)
|
505 |
-
output_result = gr.Image(label="Reference Mask")
|
506 |
-
with gr.Row():
|
507 |
-
vis_frame_type = gr.Radio(label="Track level", choices=["coarse", "fine"], value="coarse", scale=2)
|
508 |
-
propagate_btn = gr.Button("Track", scale=1)
|
509 |
-
reset_prpgt_brn = gr.Button("Reset", visible=False)
|
510 |
-
output_propagated = gr.Gallery(label="Masklets", columns=4, visible=False)
|
511 |
-
output_video = gr.Video(visible=False)
|
512 |
-
|
513 |
-
|
514 |
-
|
515 |
-
# When new video is uploaded
|
516 |
-
video_in.upload(
|
517 |
-
fn = preprocess_video_in,
|
518 |
-
inputs = [video_in],
|
519 |
-
outputs = [
|
520 |
-
first_frame_path,
|
521 |
-
tracking_points, # update Tracking Points in the gr.State([]) object
|
522 |
-
trackings_input_label, # update Tracking Labels in the gr.State([]) object
|
523 |
-
input_first_frame_image, # hidden component used as ref when clearing points
|
524 |
-
points_map, # Image component where we add new tracking points
|
525 |
-
video_frames_dir, # Array where frames from video_in are deep stored
|
526 |
-
scanned_frames, # Scanned frames by EfficientTAM
|
527 |
-
stored_inference_state, # EfficientTAM inference state
|
528 |
-
stored_frame_names, #
|
529 |
-
video_in_drawer, # Accordion to hide uploaded video player
|
530 |
-
],
|
531 |
-
queue = False
|
532 |
-
)
|
533 |
-
|
534 |
-
video_in.change(
|
535 |
-
fn = preprocess_video_in,
|
536 |
-
inputs = [video_in],
|
537 |
-
outputs = [
|
538 |
-
first_frame_path,
|
539 |
-
tracking_points, # update Tracking Points in the gr.State([]) object
|
540 |
-
trackings_input_label, # update Tracking Labels in the gr.State([]) object
|
541 |
-
input_first_frame_image, # hidden component used as ref when clearing points
|
542 |
-
points_map, # Image component where we add new tracking points
|
543 |
-
video_frames_dir, # Array where frames from video_in are deep stored
|
544 |
-
scanned_frames, # Scanned frames by EfficientTAM
|
545 |
-
stored_inference_state, # EfficientTAM inference state
|
546 |
-
stored_frame_names, #
|
547 |
-
video_in_drawer, # Accordion to hide uploaded video player
|
548 |
-
],
|
549 |
-
queue = False
|
550 |
-
)
|
551 |
-
|
552 |
-
|
553 |
-
# triggered when we click on image to add new points
|
554 |
-
points_map.select(
|
555 |
-
fn = get_point,
|
556 |
-
inputs = [
|
557 |
-
point_type, # "include" or "exclude"
|
558 |
-
tracking_points, # get tracking_points values
|
559 |
-
trackings_input_label, # get tracking label values
|
560 |
-
input_first_frame_image, # gr.State() first frame path
|
561 |
-
],
|
562 |
-
outputs = [
|
563 |
-
tracking_points, # updated with new points
|
564 |
-
trackings_input_label, # updated with corresponding labels
|
565 |
-
points_map, # updated image with points
|
566 |
-
],
|
567 |
-
queue = False
|
568 |
-
)
|
569 |
-
|
570 |
-
# Clear every points clicked and added to the map
|
571 |
-
clear_points_btn.click(
|
572 |
-
fn = clear_points,
|
573 |
-
inputs = input_first_frame_image, # we get the untouched hidden image
|
574 |
-
outputs = [
|
575 |
-
first_frame_path,
|
576 |
-
tracking_points,
|
577 |
-
trackings_input_label,
|
578 |
-
points_map,
|
579 |
-
],
|
580 |
-
queue=False
|
581 |
-
)
|
582 |
-
|
583 |
-
|
584 |
-
change_current.click(
|
585 |
-
fn = switch_working_frame,
|
586 |
-
inputs = [working_frame, scanned_frames, video_frames_dir],
|
587 |
-
outputs = [tracking_points, trackings_input_label, input_first_frame_image, points_map],
|
588 |
-
queue=False
|
589 |
-
)
|
590 |
-
|
591 |
-
|
592 |
-
submit_btn.click(
|
593 |
-
fn = get_mask_sam_process,
|
594 |
-
inputs = [
|
595 |
-
stored_inference_state,
|
596 |
-
input_first_frame_image,
|
597 |
-
checkpoint,
|
598 |
-
tracking_points,
|
599 |
-
trackings_input_label,
|
600 |
-
video_frames_dir,
|
601 |
-
scanned_frames,
|
602 |
-
working_frame,
|
603 |
-
available_frames_to_check,
|
604 |
-
],
|
605 |
-
outputs = [
|
606 |
-
change_current,
|
607 |
-
output_result,
|
608 |
-
stored_frame_names,
|
609 |
-
loaded_predictor,
|
610 |
-
stored_inference_state,
|
611 |
-
working_frame,
|
612 |
-
],
|
613 |
-
concurrency_limit=10,
|
614 |
-
queue=False
|
615 |
-
)
|
616 |
-
|
617 |
-
reset_prpgt_brn.click(
|
618 |
-
fn = reset_propagation,
|
619 |
-
inputs = [first_frame_path, loaded_predictor, stored_inference_state],
|
620 |
-
outputs = [points_map, tracking_points, trackings_input_label, output_propagated, stored_inference_state, output_result, available_frames_to_check, input_first_frame_image, working_frame, reset_prpgt_brn],
|
621 |
-
queue=False
|
622 |
-
)
|
623 |
-
|
624 |
-
propagate_btn.click(
|
625 |
-
fn = update_ui,
|
626 |
-
inputs = [vis_frame_type],
|
627 |
-
outputs = [output_propagated, output_video],
|
628 |
-
queue=False
|
629 |
-
).then(
|
630 |
-
fn = propagate_to_all,
|
631 |
-
inputs = [tracking_points, video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame],
|
632 |
-
outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn],
|
633 |
-
concurrency_limit=10,
|
634 |
-
queue=False
|
635 |
-
)
|
636 |
-
|
637 |
-
demo.queue()
|
638 |
-
demo.launch()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|