updated page for uploaded files

app.py

@@ -7,6 +7,7 @@ import gdown
 import numpy as np
 import streamlit as st
 import torch
+from PIL import Image
 
 
 def load_classes(csv_reader):
@@ -105,7 +106,7 @@ def load_model(model_path, prefix: str = 'model/'):
     return model
 
 
-def process_img(model, image, labels, caption: bool = True):
+def process_img(model, image, labels, caption: bool = True, thickness=2):
     """
     Process img given a model.
 
@@ -113,6 +114,7 @@ def process_img(model, image, labels, caption: bool = True):
     :param image: image to process
     :param model: inference model
     :param labels: given labels
+    :param thickness: thickness of bboxes
     :return:
     """
 
@@ -174,7 +176,8 @@ def process_img(model, image, labels, caption: bool = True):
             }
             cap = '{}'.format(label_name) if caption else ''
             draw_caption(image_orig, (x1, y1, x2, y2), cap)
-            cv2.rectangle(image_orig, (x1, y1), (x2, y2), color=colors[label_name], thickness=2)
+            cv2.rectangle(image_orig, (x1, y1), (x2, y2), color=colors[label_name],
+                          thickness=int(1 * (smallest_side / 100)))
 
     return image_orig
 
@@ -194,6 +197,7 @@ ids = {
 # Download all models from drive
 download_models(ids)
 
+# Split page into columns
 left, right = st.columns([5, 3])
 
 # Model selection
@@ -204,15 +208,23 @@ model = load_model(model_path=model_path) if model_path != '' else None
 # Display example selection
 index = left.number_input('', min_value=0, max_value=852, value=495, help='Choose an image. ')
 
-# Get corresponding image and transform it
-image = cv2.imread(f'data/validation/image/maksssksksss{str(index)}.jpg')
-image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+# Uploader
+uploaded = st.file_uploader("Try it out with your own image!", type=['.jpg', '.png', '.jfif'])
+
+if uploaded is not None:
+    # Convert file to image
+    image = Image.open(uploaded)
+    image = np.array(image)
+else:
+    # Get corresponding image and transform it
+    image = cv2.imread(f'data/validation/image/maksssksksss{str(index)}.jpg')
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 # Process img
 with st.spinner('Please wait while the image is being processed... This may take a while. '):
     image = process_img(model, image, labels, caption=False)
 
-left.image(image)
+left.image(cv2.resize(image, (450, 300)))
 
 # Write labels dict and device on right
 right.write({
@@ -222,3 +234,9 @@ right.write({
 })
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 right.write(device)
+
+captions = [image for image in os.listdir('data/examples/')]
+images = [Image.open(f'data/examples/{image}') for image in os.listdir('data/examples/')]
+
+# Display examples
+st.image(images, width=350)