Spaces:
Build error
Build error
import cv2 | |
from doc_ufcn.main import DocUFCN | |
import gradio as gr | |
from doc_ufcn import models | |
from PIL import Image | |
from PIL import ImageDraw | |
model_path, parameters = models.download_model("generic-page") | |
model = DocUFCN(len(parameters["classes"]), parameters["input_size"], "cpu") | |
model.load(model_path, parameters["mean"], parameters["std"]) | |
def query_image(image): | |
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
detected_polygons, probabilities, mask, overlap = model.predict( | |
image, raw_output=True, mask_output=True, overlap_output=True | |
) | |
max_confidence = 0 | |
max_prediction = None | |
for prediction in detected_polygons[1]: | |
if prediction["confidence"] > max_confidence: | |
max_confidence = prediction["confidence"] | |
max_prediction = prediction | |
image = Image.fromarray(image) | |
img2 = image.copy() | |
xy = max_prediction["polygon"] | |
draw = ImageDraw.Draw(img2) | |
draw.polygon(xy, fill="blue") | |
return Image.blend(image, img2, 0.5) | |
demo = gr.Interface( | |
query_image, | |
inputs=[gr.Image()], | |
outputs="image", | |
title="doc-ufcn Page Detection Demo", | |
description="A demo showing the top prediction from the [Teklia/doc-ufcn-generic-page](" | |
"https://huggingface.co/Teklia/doc-ufcn-generic-page) model. This generic page detection model predicts single pages from document images.", | |
examples=[ | |
"v2_bsb00046516_00100_full_full_0_default.jpg", | |
"512px-Page_from_the_Arthurian_Romances_illuminated_manuscript.jpg", | |
], | |
) | |
demo.launch() | |