Views
No views yet
1import torch
2from PIL import Image
3from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
4
5# Load the model and processor
6image_processor = RTDetrImageProcessor.from_pretrained("HuggingPanda/docling-layout")
7model = RTDetrForObjectDetection.from_pretrained("HuggingPanda/docling-layout")
8
9# Load an image
10image = Image.open("hocr_output_page-0001.jpg")
11
12# Preprocess the image
13resize = {"height":640, "width":640}
14inputs = image_processor(
15 images=image,
16 return_tensors="pt",
17 size=resize,
18)
19
20# Perform inference
21with torch.no_grad():
22 outputs = model(**inputs)
23
24# Post-process results
25results = image_processor.post_process_object_detection(
26 outputs,
27 target_sizes=torch.tensor([image.size[::-1]]),
28 threshold=0.3
29)
30
31# Print detected objects
32for result in results:
33 for score, label_id, box in zip(result["scores"], result["labels"], result["boxes"]):
34 score, label = score.item(), label_id.item()
35 box = [round(i, 2) for i in box.tolist()]
36 print(f"{model.config.id2label[label+1]}: {score:.2f} {box}")
37@misc{docling2024, title={Docling Models for Document Layout Analysis}, author={DS4SD Team}, year={2024}, howpublished={Hugging Face Repository}, url={https://huggingface.co/ds4sd/docling-models} }