Spaces:
Sleeping
Sleeping
File size: 1,359 Bytes
492a9fd 73cd058 492a9fd 73cd058 492a9fd |
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 |
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import os
# Ensure the model file is in the correct location
model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
if not os.path.exists(model_path):
# Download the model file if it doesn't exist
import requests
model_url = "https://huggingface.co/DILHTWD/documentlayoutsegmentation_YOLOv8_ondoclaynet/resolve/main/yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
response = requests.get(model_url)
with open(model_path, "wb") as f:
f.write(response.content)
# Load the document segmentation model
docseg_model = YOLO(model_path)
def process_image(image):
# Convert image to the format YOLO model expects
image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
results = docseg_model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
# Extract annotated image from results
annotated_img = results[0].plot()
return annotated_img, results[0].boxes
# Define the Gradio interface
interface = gr.Interface(
fn=process_image,
inputs=gr.inputs.Image(type="pil"),
outputs=[gr.outputs.Image(type="pil", label="Annotated Image"),
gr.outputs.Textbox(label="Detected Areas and Labels")]
)
if __name__ == "__main__":
interface.launch()
|