from PIL import Image
from model import yolox
from os import listdir
import os.path
import requests
import gradio as gr
import numpy as np
import streamlit as st

DMINITY_MODEL_URL = "https://github.com/Dolpheyn/dminity/releases/download/v-1.0.0/dminity.onnx"
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):
        print("Downloading dminity model weight from {}...".format(DMINITY_MODEL_URL))
        r = requests.get(DMINITY_MODEL_URL, allow_redirects=True)        
        print("Writing to {}".format(MODEL_PATH))
        open(MODEL_PATH, 'wb').write(r.content)        
        print("Done!")
    
    # 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. Please note that the first detection will take around 35 seconds because the model is being downloaded."
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)