fffiloni commited on
Commit
95134d2
·
verified ·
1 Parent(s): fd0d56e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -204,17 +204,21 @@ def load_model(checkpoint):
204
  if checkpoint == "tiny":
205
  sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt"
206
  model_cfg = "sam2_hiera_t.yaml"
 
207
  elif checkpoint == "samll":
208
  sam2_checkpoint = "./checkpoints/sam2_hiera_small.pt"
209
  model_cfg = "sam2_hiera_s.yaml"
 
210
  elif checkpoint == "base-plus":
211
  sam2_checkpoint = "./checkpoints/sam2_hiera_base_plus.pt"
212
  model_cfg = "sam2_hiera_b+.yaml"
 
213
  elif checkpoint == "large":
214
  sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt"
215
  model_cfg = "sam2_hiera_l.yaml"
 
216
 
217
- return sam2_checkpoint, model_cfg
218
 
219
  def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir):
220
  # 1. We need to preprocess the video and store frames in the right directory
 
204
  if checkpoint == "tiny":
205
  sam2_checkpoint = "./checkpoints/sam2_hiera_tiny.pt"
206
  model_cfg = "sam2_hiera_t.yaml"
207
+ return sam2_checkpoint, model_cfg
208
  elif checkpoint == "samll":
209
  sam2_checkpoint = "./checkpoints/sam2_hiera_small.pt"
210
  model_cfg = "sam2_hiera_s.yaml"
211
+ return sam2_checkpoint, model_cfg
212
  elif checkpoint == "base-plus":
213
  sam2_checkpoint = "./checkpoints/sam2_hiera_base_plus.pt"
214
  model_cfg = "sam2_hiera_b+.yaml"
215
+ return sam2_checkpoint, model_cfg
216
  elif checkpoint == "large":
217
  sam2_checkpoint = "./checkpoints/sam2_hiera_large.pt"
218
  model_cfg = "sam2_hiera_l.yaml"
219
+ return sam2_checkpoint, model_cfg
220
 
221
+
222
 
223
  def sam_process(input_first_frame_image, checkpoint, tracking_points, trackings_input_label, video_frames_dir):
224
  # 1. We need to preprocess the video and store frames in the right directory