yunyangx commited on
Commit
c05437e
·
verified ·
1 Parent(s): 8932a03

Applying for GPU

Browse files
Files changed (1) hide show
  1. app.py +0 -638
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()