jhj0517 commited on
Commit
b109f8e
1 Parent(s): 23ff121

Fix missing config bug on colab

Browse files
Files changed (1) hide show
  1. modules/sam_inference.py +5 -3
modules/sam_inference.py CHANGED
@@ -62,7 +62,9 @@ class SamInference:
62
  if model_type is None:
63
  model_type = DEFAULT_MODEL_TYPE
64
 
65
- config = MODEL_CONFIGS[model_type]
 
 
66
  filename, url = AVAILABLE_MODELS[model_type]
67
 
68
  model_path = os.path.join(self.model_dir, filename)
@@ -76,7 +78,7 @@ class SamInference:
76
  try:
77
  self.model = None
78
  self.video_predictor = build_sam2_video_predictor(
79
- config_file=config,
80
  ckpt_path=model_path,
81
  device=self.device
82
  )
@@ -86,7 +88,7 @@ class SamInference:
86
 
87
  try:
88
  self.model = build_sam2(
89
- config_file=config,
90
  ckpt_path=model_path,
91
  device=self.device
92
  )
 
62
  if model_type is None:
63
  model_type = DEFAULT_MODEL_TYPE
64
 
65
+ config_path = MODEL_CONFIGS[model_type]
66
+ config_dir, config_name = os.path.split(config_path)
67
+
68
  filename, url = AVAILABLE_MODELS[model_type]
69
 
70
  model_path = os.path.join(self.model_dir, filename)
 
78
  try:
79
  self.model = None
80
  self.video_predictor = build_sam2_video_predictor(
81
+ config_file=config_name,
82
  ckpt_path=model_path,
83
  device=self.device
84
  )
 
88
 
89
  try:
90
  self.model = build_sam2(
91
+ config_file=config_name,
92
  ckpt_path=model_path,
93
  device=self.device
94
  )