eazymlshubham/DocLayoutYOLO

from PIL import Image
from doclayout_yolo import YOLOv10
from huggingface_hub import hf_hub_download

# Replace with the actual repository ID
repo_id = "eazymlshubham/DocLayoutYOLO"
# Replace with the actual .pt filename
filename = "doclayout_yolo_docstructbench_imgsz1024.pt"      

# Download the model file from Hugging Face Hub
# This will download the file to a local cache directory and return the local path
local_file_path = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="model")

# Load the pre-trained model
# The model is loaded from the local file path obtained from hf_hub_download
model = YOLOv10(local_file_path)

# Perform prediction
det_res = model.predict(
    r"images/clinical_trials9.jpg",   # Image to predict
    imgsz=1024,        # Prediction image size
    conf=0.2,          # Confidence threshold
    device="cpu"    # Device to use (e.g., 'cuda:0' or 'cpu')
)

# Annotate and save the result
annotated_frame = det_res[0].plot(pil=True, line_width=5, font_size=20)


img = Image.fromarray(annotated_frame)
img.show()