File size: 1,942 Bytes
39bb2df af0a65d 39bb2df 3e3f00a aa5505d af0a65d 39bb2df af0a65d 39bb2df 3f45de4 39bb2df e9beee9 |
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 50 51 |
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) |