Update core/process.py

core/process.py

@@ -93,7 +93,7 @@ def process_image(
         if run_seg:
             logger.info(f"Running segmentation with model: {seg_model}")
             load_start = time.time()
-            model = get_model("segmentation", SEGMENTATION_MODEL_MAP[seg_model], device="cpu")
+            model = get_model("segmentation", seg_model, device="cpu")
             logger.info(f"{seg_model} segmentation model loaded in {time.time() - load_start:.2f} seconds.")
             mask = model.predict(image)
             overlay = model.draw(image, mask, alpha=blend)
@@ -107,7 +107,7 @@ def process_image(
         if run_depth:
             logger.info(f"Running depth estimation with model: {depth_model}")
             load_start = time.time()
-            model = get_model("depth", DEPTH_MODEL_MAP[depth_model], device="cpu")
+            model = get_model("depth", depth_model, device="cpu")
             logger.info(f"{depth_model} depth model loaded in {time.time() - load_start:.2f} seconds.")
             dmap = model.predict(image)
             norm_dmap = ((dmap - dmap.min()) / (dmap.ptp()) * 255).astype(np.uint8)