|
|
|
from ultralytics import ASSETS, YOLO, RTDETR |
|
import gradio as gr |
|
from huggingface_hub import snapshot_download |
|
from PIL import Image |
|
|
|
model_dir = snapshot_download("omarelsayeed/DETR-ARABIC-DOCUMENT-LAYOUT-ANALYSIS") + "/Model.pt" |
|
model = RTDETR(model_dir) |
|
|
|
|
|
def predict_image(img, conf_threshold, iou_threshold): |
|
"""Predicts objects in an image using a YOLO11 model with adjustable confidence and IOU thresholds.""" |
|
results = model.predict( |
|
source=img, |
|
conf=conf_threshold, |
|
iou=iou_threshold, |
|
show_labels=True, |
|
show_conf=True, |
|
imgsz=640, |
|
) |
|
|
|
for r in results: |
|
im_array = r.plot() |
|
im = Image.fromarray(im_array[..., ::-1]) |
|
|
|
return im |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict_image, |
|
inputs=[ |
|
gr.Image(type="pil", label="Upload Image"), |
|
gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"), |
|
gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"), |
|
], |
|
outputs=gr.Image(type="pil", label="Result"), |
|
title="Ultralytics Gradio", |
|
description="Upload images for inference. The Ultralytics YOLO11n model is used by default.", |
|
examples=[ |
|
["kashida.png", 0.2, 0.45], |
|
["image.jpg", 0.2, 0.45], |
|
["Screenshot 2024-11-06 130230.png" , 0.25 , 0.45] |
|
], |
|
theme=gr.themes.Default() |
|
) |
|
|
|
if __name__ == "__main__": |
|
iface.launch() |
