Mar2Ding commited on
Commit
11338b3
1 Parent(s): 41971c7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -5
app.py CHANGED
@@ -13,6 +13,11 @@ div#component-18, div#component-25, div#component-35, div#component-41{
13
  }
14
  """
15
 
 
 
 
 
 
16
  @spaces.GPU(duration=120)
17
  def run_install(command):
18
  result = subprocess.run(command, capture_output=True, text=True)
@@ -245,6 +250,7 @@ def get_mask_sam_process(
245
  print("MODEL LOADED")
246
 
247
  # set predictor
 
248
  predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
249
  print("PREDICTOR READY")
250
 
@@ -319,12 +325,13 @@ def get_mask_sam_process(
319
  print(available_frames_to_check)
320
 
321
  # 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)
322
- return "output_first_frame.jpg", frame_names, predictor, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=False)
323
 
324
  @spaces.GPU(duration=120)
325
  def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame, progress=gr.Progress(track_tqdm=True)):
326
  #### PROPAGATION ####
327
  sam2_checkpoint, model_cfg = load_model(checkpoint)
 
328
  predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
329
 
330
  inference_state = stored_inference_state
@@ -425,7 +432,7 @@ def switch_working_frame(working_frame, scanned_frames, video_frames_dir):
425
 
426
 
427
  @spaces.GPU(duration=120)
428
- def reset_propagation(first_frame_path, predictor, stored_inference_state):
429
 
430
  predictor.reset_state(stored_inference_state)
431
  # print(f"RESET State: {stored_inference_state} ")
@@ -438,7 +445,7 @@ with gr.Blocks(css=css) as demo:
438
  trackings_input_label = gr.State([])
439
  video_frames_dir = gr.State()
440
  scanned_frames = gr.State()
441
- loaded_predictor = gr.State()
442
  stored_inference_state = gr.State()
443
  stored_frame_names = gr.State()
444
  available_frames_to_check = gr.State([])
@@ -594,7 +601,7 @@ with gr.Blocks(css=css) as demo:
594
  outputs = [
595
  output_result,
596
  stored_frame_names,
597
- loaded_predictor,
598
  stored_inference_state,
599
  working_frame,
600
  ],
@@ -603,7 +610,7 @@ with gr.Blocks(css=css) as demo:
603
 
604
  reset_prpgt_brn.click(
605
  fn = reset_propagation,
606
- inputs = [first_frame_path, loaded_predictor, stored_inference_state],
607
  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],
608
  queue=False
609
  )
@@ -619,4 +626,6 @@ with gr.Blocks(css=css) as demo:
619
  outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn]
620
  )
621
 
 
 
622
  demo.launch()
 
13
  }
14
  """
15
 
16
+ predictor = None
17
+
18
+
19
+
20
+
21
  @spaces.GPU(duration=120)
22
  def run_install(command):
23
  result = subprocess.run(command, capture_output=True, text=True)
 
250
  print("MODEL LOADED")
251
 
252
  # set predictor
253
+ global predictor
254
  predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
255
  print("PREDICTOR READY")
256
 
 
325
  print(available_frames_to_check)
326
 
327
  # 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)
328
+ return "output_first_frame.jpg", frame_names, inference_state, gr.update(choices=available_frames_to_check, value=working_frame, visible=False)
329
 
330
  @spaces.GPU(duration=120)
331
  def propagate_to_all(video_in, checkpoint, stored_inference_state, stored_frame_names, video_frames_dir, vis_frame_type, available_frames_to_check, working_frame, progress=gr.Progress(track_tqdm=True)):
332
  #### PROPAGATION ####
333
  sam2_checkpoint, model_cfg = load_model(checkpoint)
334
+ global predictor
335
  predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint)
336
 
337
  inference_state = stored_inference_state
 
432
 
433
 
434
  @spaces.GPU(duration=120)
435
+ def reset_propagation(first_frame_path, stored_inference_state):
436
 
437
  predictor.reset_state(stored_inference_state)
438
  # print(f"RESET State: {stored_inference_state} ")
 
445
  trackings_input_label = gr.State([])
446
  video_frames_dir = gr.State()
447
  scanned_frames = gr.State()
448
+ # loaded_predictor = gr.State()
449
  stored_inference_state = gr.State()
450
  stored_frame_names = gr.State()
451
  available_frames_to_check = gr.State([])
 
601
  outputs = [
602
  output_result,
603
  stored_frame_names,
604
+ # loaded_predictor,
605
  stored_inference_state,
606
  working_frame,
607
  ],
 
610
 
611
  reset_prpgt_brn.click(
612
  fn = reset_propagation,
613
+ inputs = [first_frame_path, stored_inference_state],
614
  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],
615
  queue=False
616
  )
 
626
  outputs = [output_propagated, output_video, working_frame, available_frames_to_check, reset_prpgt_brn]
627
  )
628
 
629
+
630
+
631
  demo.launch()