Spaces:
Runtime error
Runtime error
try: | |
import detectron2 | |
except: | |
import os | |
os.system("pip install git+https://github.com/facebookresearch/detectron2.git") | |
import cv2 | |
from matplotlib.pyplot import axis | |
import gradio as gr | |
import requests | |
import numpy as np | |
from torch import nn | |
import requests | |
import torch | |
from detectron2 import model_zoo | |
from detectron2.engine import DefaultPredictor | |
from detectron2.config import get_cfg | |
from detectron2.utils.visualizer import Visualizer | |
from detectron2.data import MetadataCatalog | |
models = [ | |
{ | |
"name": "Version 1 (2-class)", | |
"model_path": "https://huggingface.co/dbmdz/detectron2-model/resolve/main/model_final.pth", | |
"classes": ["footer"], | |
"cfg": None, | |
"metadata": None, | |
}, | |
] | |
model_name_to_id = {model["name"]: id_ for id_, model in enumerate(models)} | |
for model in models: | |
model["cfg"] = get_cfg() | |
model["cfg"].merge_from_file("./configs/detectron2/faster_rcnn_R_50_FPN_3x.yaml") | |
model["cfg"].MODEL.ROI_HEADS.NUM_CLASSES = len(model["classes"]) | |
model["cfg"].MODEL.WEIGHTS = model["model_path"] | |
model["metadata"] = MetadataCatalog.get(model["name"]) | |
model["metadata"].thing_classes = model["classes"] | |
if not torch.cuda.is_available(): | |
model["cfg"].MODEL.DEVICE = "cpu" | |
def inference(image_url, image, min_score, model_name): | |
if image_url: | |
r = requests.get(image_url) | |
if r: | |
im = np.frombuffer(r.content, dtype="uint8") | |
im = cv2.imdecode(im, cv2.IMREAD_COLOR) | |
else: | |
# Model expect BGR! | |
im = image[:, :, ::-1] | |
model_id = model_name_to_id[model_name] | |
models[model_id]["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score | |
predictor = DefaultPredictor(models[model_id]["cfg"]) | |
outputs = predictor(im) | |
v = Visualizer(im, models[model_id]["metadata"], scale=1.2) | |
out = v.draw_instance_predictions(outputs["instances"].to("cpu")) | |
return out.get_image() | |
title = "# DBMDZ Detectron2 Model Demo" | |
description = """ | |
This demo introduces an interactive playground for our trained Detectron2 model. | |
""" | |
footer = "Made in Egypt with ❤️." | |
with gr.Blocks() as demo: | |
gr.Markdown(title) | |
gr.Markdown(description) | |
with gr.Tab("From URL"): | |
url_input = gr.Textbox( | |
label="Image URL", | |
placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg", | |
) | |
with gr.Tab("From Image"): | |
image_input = gr.Image(type="numpy", label="Input Image") | |
min_score = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score") | |
model_name = gr.Radio( | |
choices=[model["name"] for model in models], | |
value=models[0]["name"], | |
label="Select Detectron2 model", | |
) | |
output_image = gr.Image(type="pil", label="Output") | |
inference_button = gr.Button("Submit") | |
inference_button.click( | |
fn=inference, | |
inputs=[url_input, image_input, min_score, model_name], | |
outputs=output_image, | |
) | |
gr.Markdown(footer) | |
demo.launch() | |