Update app.py

app.py

@@ -33,19 +33,34 @@ from mono.utils.transform import gray_to_colormap
 from mono.utils.visualization import vis_surface_normal
 import gradio as gr
 
-cfg = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.large.py')
-
 torch.hub.download_url_to_file('https://images.unsplash.com/photo-1437622368342-7a3d73a34c8f', 'turtle.jpg')
 torch.hub.download_url_to_file('https://images.unsplash.com/photo-1519066629447-267fffa62d4b', 'lions.jpg')
 
-model = get_configured_monodepth_model(cfg, )
-model, _,  _, _ = load_ckpt('./weight/metric_depth_vit_large_800k.pth', model, strict_match=False)
-model.eval()
+cfg_large = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.large.py')
+model_large = get_configured_monodepth_model(cfg_large, )
+model_large, _,  _, _ = load_ckpt('./weight/metric_depth_vit_large_800k.pth', model_large, strict_match=False)
+model_large.eval()
 
-device = "cpu"
-model.to(device)
+cfg_small = Config.fromfile('./mono/configs/HourglassDecoder/vit.raft5.small.py')
+model_small = get_configured_monodepth_model(cfg_small, )
+model_small, _,  _, _ = load_ckpt('./weight/metric_depth_vit_small_800k.pth', model_small, strict_match=False)
+model_small.eval()
 
-def depth_normal(img):
+device = "cpu"
+model_large.to(device)
+model_small.to(device)
+
+def depth_normal(img, model_selection="vit-small"):
+    if model_selection == "vit_small":
+        model = model_small
+        cfg = cfg_small
+    elif model == "vit_large":
+        model = model_large
+        cfg = cfg_large
+
+    else:
+        raise NotImplementedError
+    
     cv_image = np.array(img) 
     img = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
     intrinsic = [1000.0, 1000.0, img.shape[1]/2, img.shape[0]/2]
@@ -86,7 +101,7 @@ def depth_normal(img):
 #normal = gr.outputs.Image(type="pil",label="Output Normal")
 
 title = "Metric3D"
-description = "Gradio demo for Metric3D (vit-large) running on CPU which takes in a single image for computing metric depth and surface normal. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
+description = "Gradio demo for Metric3D (v2, more diverse models) running on CPU which takes in a single image for computing metric depth and surface normal. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
 article = "<p style='text-align: center'><a href='https://arxiv.org/pdf/2307.10984.pdf'>Metric3D: Towards Zero-shot Metric 3D Prediction from A Single Image</a> | <a href='https://github.com/YvanYin/Metric3D'>Github Repo</a></p>"
 
 examples = [
@@ -96,6 +111,6 @@ examples = [
 
 gr.Interface(
     depth_normal, 
-    inputs=[gr.Image(type='pil', label="Original Image")], 
+    inputs=[gr.Image(type='pil', label="Original Image"), gr.Dropdown(["vit-small", "vit-large"], label="Model", info="Will support more models later!"], 
     outputs=[gr.Image(type="pil",label="Output Depth"), gr.Image(type="pil",label="Output Normal")], 
     title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch()