Update app.py

app.py

@@ -1,4 +1,49 @@
+from PIL import Image
+from model import yolox
+from os import listdir
+import os.path
+import gdown
+import gradio as gr
+import numpy as np
 import streamlit as st
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+DMINITY_MODEL_URL = "https://drive.google.com/uc?id=1HnBFOXoRsqn6LiX63hZqZFSEOic2-71M"
+MODEL_PATH = "dminity.onnx"
+
+@st.cache(allow_output_mutation=True, show_spinner=True)
+def get_model():
+    # Download model from Google Drive if it does not exist
+    if not os.path.isfile(MODEL_PATH):
+        gdown.download(DMINITY_MODEL_URL, MODEL_PATH)
+    
+    # Load model with OpenCV
+    model = yolox(MODEL_PATH, p6=False, confThreshold=0.3)
+
+    return model
+
+
+def dminity(im, size=640):
+    model = get_model()
+    # Resize image
+    g = (size / max(im.size))
+    im = im.resize((int(x * g) for x in im.size))
+    im = np.array(im)
+
+    # Detect and get back rendered image and amenities list
+    image, amenities = model.detect(im)
+    return image, amenities
+
+inputs = gr.inputs.Image(type='pil', label="Original Image")
+outputs = [gr.outputs.Image(type='pil', label="Output Image"), "text"]
+
+title = "Dminity"
+description = "Dminity demo for amenity object detection. Upload a house interior image with amenities or click an example image to use."
+article = "<p style='text-align: center'>Dminity is a YOLOX object detection model trained to detect home amenities.</p>"
+
+# List of example images
+files = ['images/' + f for f in listdir('images')]
+examples=[[f] for f in files]
+
+gr.Interface(dminity, inputs, outputs, title=title, description=description,
+        article=article, examples=examples,
+        theme="huggingface").launch(debug=True)