detectron2 / app.py
stalyn314's picture
Update app.py
9b6d5b1 verified
raw
history blame
6.21 kB
try:
import detectron2
except:
import os
os.system('pip install git+https://github.com/facebookresearch/detectron2.git')
# Manipulación de Datos y Bibliotecas Científicas
import csv # Leer y escribir archivos CSV
import numpy as np # Cálculos numéricos con soporte para matrices y arrays
import pandas as pd # Manipulación y análisis de datos, especialmente con DataFrames
import random # Generar números aleatorios y realizar operaciones aleatorias
import datetime # Trabajar con fechas y horas
from datetime import datetime
# Visualización
import matplotlib.pyplot as plt # Crear gráficos estáticos, interactivos y animados
from matplotlib.path import Path # Trabajar con caminos geométricos en visualizaciones
from matplotlib.pyplot import axis
# Procesamiento de Imágenes y Visión por Computadora (OpenCV y Scikit-Image)
import cv2 # Para procesamiento de imágenes y tareas de visión por computadora
from skimage import measure, io, color, draw # Funciones de procesamiento de imágenes en skimage
from skimage.measure import regionprops # Para medir propiedades de regiones en imágenes
# Manipulación de Imágenes
from PIL import Image # Para abrir, manipular y guardar imágenes en diferentes formatos
# Serialización de Datos
import json # Para analizar y serializar datos JSON
import yaml # Para leer y escribir archivos YAML (serialización de datos en formato humano)
# Barra de Progreso
from tqdm import tqdm # Para mostrar una barra de progreso en tareas largas
# Diagramas Ternarios
import ternary # Para crear diagramas ternarios (ej. gráficos de 3 variables)
# Generación de Reportes PDF
from fpdf import FPDF # Para crear documentos PDF de forma sencilla
# Librerías de Hugging Face para interactuar con modelos
from transformers import pipeline, AutoModelForImageClassification, AutoTokenizer
import gradio as gr
import requests
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/stalyn314/10xmineralmodel/resolve/main/xplx10x_d2.pth",
"classes": ["minerales", "Afs", "Amp", "Bt", "Ms", "Ol", "Pl", "Px", "Qz"],
"cfg": None,
"metadata": None
},
{
"name": "Version 2 (4-class)",
"model_path": "https://huggingface.co/stalyn314/10xmineralmodel/resolve/main/10xmodel_d2.pth",
"classes": ["minerales", "Afs", "Amp", "Bt", "Ms", "Ol", "Pl", "Px", "Qz"],
"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/mask_rcnn_X_101_32x8d_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_BGR2RGB)
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"))
# Convertir la imagen de salida de BGR a RGB
result_image = out.get_image() # Esto sigue estando en BGR
result_image_rgb = result_image[:, :, ::-1] # Convertir BGR a RGB
return result_image_rgb
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()