updated the app for DPT model

DPT.py

@@ -14,6 +14,7 @@ def load_model(model_type = 'DPT_Large'):
     midas = torch.hub.load("intel-isl/MiDaS", model_type)
 
     device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    print(f'---DPT will use device: {device}')
     midas.to(device)
     midas.eval()
 
@@ -27,13 +28,13 @@ def load_model(model_type = 'DPT_Large'):
         'midas': midas, 'device': device, 'transform': transform
     }
 
-def inference(img_array_rgb, model_def):
+def inference(img_array_rgb, model_obj):
     '''run DPT model and returns a PIL image'''
     # img = cv2.imread(img.name)
     # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    midas = model_def['midas']
-    transform = model_def['transform']
-    device = model_def['device']
+    midas = model_obj['midas']
+    transform = model_obj['transform']
+    device = model_obj['device']
     input_batch = transform(img_array_rgb).to(device)
 
     with torch.no_grad():

app.py

@@ -1,11 +1,11 @@
 import streamlit as st
 
-import os, sys, io
+import os, sys, io, time
 import urllib.request as urllib
 import numpy as np
 from PIL import Image
 
-import DPT
+import DPT, BTS_infer
 
 ### Some Utils Functions ###
 def get_image(st_asset = st.sidebar, as_np_arr = False, extension_list = ['jpg', 'jpeg', 'png']):
@@ -70,12 +70,31 @@ def im_draw_bbox(pil_im, x0, y0, x1, y1, color = 'black', width = 3, caption = N
 	return Image.fromarray(im_array)
 
 ### Streamlit App ###
+@st.cache(allow_output_mutation = True)
+def get_model_zoo():
+	model_zoo = {
+		'DPT': {'infer_func': DPT.inference,'model': DPT.load_model()},
+		# 'BTS': {'infer_func': BTS_infer.inference,'model': BTS_infer.get_model()}
+	}
+	return model_zoo
 
-def mod_DPT(pil_im, model_def):
-	depth_im = DPT.inference(img_array_rgb = np.array(pil_im), model_def = model_def)
+@st.cache(suppress_st_warning=True)
+def mono_depth(pil_im, model_name):
+	s_time = time.time()
+	model_zoo = get_model_zoo()
+	infer_func = model_zoo[model_name]['infer_func']
+	model_obj = model_zoo[model_name]['model']
+	depth_im = infer_func(img_array_rgb = np.array(pil_im),
+					model_obj = model_obj)
+	st.info(f'''
+		model name: {model_name}\n
+		inference time: `{round(time.time()-s_time,2)}` seconds\n
+		depth image shape: {np.array(depth_im).shape}\n
+		depth image type: {type(depth_im)}
+		''')
 	return depth_im
 
-def Main(model_dict):
+def Main():
 	st.set_page_config(layout = 'wide')
 	l_col, r_col = st.columns(2)
 	show_miro_logo(st_asset = l_col)
@@ -84,20 +103,18 @@ def Main(model_dict):
 		Comparsion of two models: [BTS (CNN)](https://github.com/ErenBalatkan/Bts-PyTorch)
 		and [DPT (Transformer)](https://huggingface.co/Intel/dpt-large)
 		''')
-
+	model_zoo = get_model_zoo()
 	im = get_image(st_asset = r_col.expander('Input Image', expanded = True), extension_list = ['jpg','jpeg'])
-	model_name = r_col.selectbox('Pick Model', options = ['DPT','BTS'])
+	model_name = l_col.selectbox('Pick Model', options = list(model_zoo.keys()))
 
 	if im:
-		model_def = DPT.load_model()
-		d_im = mod_DPT(pil_im = im, model_def=model_def)
+		d_im = mono_depth(pil_im = im, model_name=model_name)
 
 		l_col, r_col = st.columns(2)
 		l_col.image(im, caption = 'Input Image')
-		r_col.image(saliency_im, caption = 'Depth Map')
+		r_col.image(d_im, caption = 'Depth Map')
 	else:
 		st.warning(f'please provide an image :point_up:')
 
 if __name__ == '__main__':
-	model_dict = load_model()
-	Main(model_dict = model_dict)
+	Main()