remove file after download button creation and scale image input

app.py

@@ -4,15 +4,14 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torchvision import models
-from torchvision.transforms import ToTensor
+from torchvision.transforms import ToTensor, Resize
 
 import numpy as np
 from PIL import Image
 import math
 from obj2html import obj2html
 
-from io import BytesIO
-import base64
+import os
 
 # DEPTH IMAGE TO OBJ
 minDepth=10
@@ -153,7 +152,13 @@ model.load_state_dict(torch.hub.load_state_dict_from_url(path, progress=True))
 model.eval()
 
 def predict(inp):
-    torch_image = ToTensor()(inp)
+    width, height = inp.size
+    if width > height:
+        scale_fn = Resize((640, int((640*width)/height)))
+    else:
+        scale_fn = Resize((int((640*height)/width), 640))
+    res_img = scale_fn(inp)
+    torch_image = ToTensor()(res_img)
     images = torch_image.unsqueeze(0)
 
     with torch.no_grad():
@@ -166,7 +171,7 @@ def predict(inp):
     create_obj(depth, 'model.obj')
     html_string = obj2html('model.obj', html_elements_only=True)
 
-    return img, html_string
+    return res_img, img, html_string
 
 
 # STREAMLIT
@@ -175,19 +180,21 @@ uploader = st.file_uploader('Upload your portrait here',type=['jpg','jpeg','png'
 
 if uploader is not None:
     pil_image = Image.open(uploader)
-    pil_depth, html_string = predict(pil_image)
+    pil_scaled, pil_depth, html_string = predict(pil_image)
 
     components.html(html_string)
     #st.markdown(html_string, unsafe_allow_html=True)
 
     col1, col2, col3 = st.columns(3)
     with col1:
-        st.image(pil_image)
+        st.image(pil_scaled)
     with col2:
         st.image(pil_depth)
     with col3:
         with open('model.obj') as f:
             st.download_button('Download model.obj', f, file_name="model.obj")
+        os.remove('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")
+        os.remove('tmp.png')