Alessio Grancini commited on
Commit
3146d91
·
verified ·
1 Parent(s): 908272b

Update monocular_depth_estimator.py

Browse files
Files changed (1) hide show
  1. monocular_depth_estimator.py +3 -3
monocular_depth_estimator.py CHANGED
@@ -16,9 +16,9 @@ MODEL_FILE_URL = {
16
  }
17
 
18
  class MonocularDepthEstimator:
19
- def __init__(self,
20
- model_type="midas_v21_small_256",
21
- model_weights_path="models/",
22
  optimize=False,
23
  side_by_side=False,
24
  height=None,
 
16
  }
17
 
18
  class MonocularDepthEstimator:
19
+ def __init__(self, model_type="midas_v21_small_256", device="cpu"):
20
+ self.device = device
21
+ self.model, self.transform, *_ = load_model(self.device, f"models/{model_type}.pt", model_type)
22
  optimize=False,
23
  side_by_side=False,
24
  height=None,