download buttons

app.py

@@ -11,6 +11,10 @@ from PIL import Image
 import math
 from obj2html import obj2html
 
+from io import BytesIO
+import base64
+
+# DEPTH IMAGE TO OBJ
 minDepth=10
 maxDepth=1000
 def my_DepthNorm(x, maxDepth):
@@ -87,7 +91,7 @@ def create_obj(img, objPath='model.obj', mtlPath='model.mtl', matName='colored',
                         vete(v2, v2) + " " + vete(v3, v3) + "\n")
                 f.write("f " + vete(v3, v3) + " " +
                         vete(v2, v2) + " " + vete(v4, v4) + "\n")
-
+# MODEL
 class UpSample(nn.Sequential):
     def __init__(self, skip_input, output_features):
         super(UpSample, self).__init__()        
@@ -164,7 +168,8 @@ def predict(inp):
 
     return img, html_string
 
-st.title("Monocular Depth Estimation")
+
+# STREAMLIT
 
 uploader = st.file_uploader('Upload your portrait here',type=['jpg','jpeg','png'])
 
@@ -172,11 +177,17 @@ if uploader is not None:
     pil_image = Image.open(uploader)
     pil_depth, html_string = predict(pil_image)
 
-    col1, col2 = st.columns(2)
+    components.html(html_string)
+    #st.markdown(html_string, unsafe_allow_html=True)
+
+    col1, col2, col3 = st.columns(3)
     with col1:
         st.image(pil_image)
     with col2:
         st.image(pil_depth)
-
-    components.html(html_string)
-    st.markdown(html_string, unsafe_allow_html=True)
+    with col3:
+        with open('model.obj') as f:
+            st.download_button('Download model.obj', f, file_name="model.obj")
+        pil_depth.save('tmp.png')
+        with open('tmp.png', "rb") as f:
+            st.download_button('Download depth.png', f,file_name="depth.png", mime="image/png")