File size: 1,605 Bytes
39bb2df 3e3f00a aa5505d 39bb2df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 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
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) |