|
|
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/amaurisarchibald/model_damage/resolve/main/model_damage.pth", |
|
|
"classes": ["damage", "scratch"], |
|
|
"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: |
|
|
|
|
|
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() |
|
|
|
|
|
def inferenceScore(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: |
|
|
|
|
|
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 = outputs["instances"].scores[1] |
|
|
return float(out.item()) |
|
|
|
|
|
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") |
|
|
|
|
|
Prediccion = gr.Number( label="Score") |
|
|
|
|
|
|
|
|
|
|
|
inference_button.click(fn=inferenceScore, inputs=[url_input, image_input, min_score, model_name], outputs=Prediccion) |
|
|
|
|
|
|
|
|
gr.Markdown(footer) |
|
|
|
|
|
demo.launch() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
