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": ["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() | |