Spaces:
Running
Running
File size: 4,278 Bytes
eb24fe7 8f6b83a eb24fe7 70d5e84 8f6b83a eb24fe7 70d5e84 eb24fe7 8f6b83a eb24fe7 70d5e84 eb24fe7 8f6b83a eb24fe7 8f6b83a 70d5e84 |
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 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
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": ["Illumination", "Illustration"],
"cfg": None,
"metadata": None
},
{
"name": "Version 2 (4-class)",
"model_path": "https://huggingface.co/dbmdz/detectron2-v2-model/resolve/main/model_final.pth",
"classes": ["ILLUSTRATION", "OTHER", "STAMP", "INITIAL"],
"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.
Currently, two models are supported that were trained on manually annotated segments from digitized books:
* [Version 1 (2-class)](https://huggingface.co/dbmdz/detectron2-model): This model can detect *Illustration* or *Illumination* segments on a given page.
* [Version 2 (4-class)](https://huggingface.co/dbmdz/detectron2-v2-model): This model is more powerful and can detect *Illustration*, *Stamp*, *Initial* or *Other* segments on a given page.
"""
footer = "Made in Munich with ❤️ and 🥨."
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()
#gr.Interface(
# inference,
# [gr.inputs.Textbox(label="Image URL", placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg"),
# gr.inputs.Image(type="numpy", label="Input Image"),
# gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score"),
# gr.Radio(choices=[model["name"] for model in models], value=models[0]["name"], label="Select Detectron2 model"),
# ],
# gr.outputs.Image(type="pil", label="Output"),
# title=title,
# description=description,
# article=article,
# examples=[]).launch()
|