--- license: apache-2.0 --- THIS IS WORK IN PROGRESS # Docling Layout Model egret-medium `docling-layout-egret-medium` is a Document Layout Model based on [DFINE-m](https://github.com/Peterande/D-FINE). The model has been trained from scratch on a mix of document datasets. It is part of the [Docling project](https://github.com/docling-project/docling). # Inference code example Prerequisites: ```bash pip install transformers Pillow torch requests ``` Prediction: ```python import requests from transformers import ( DFineForObjectDetection, RTDetrImageProcessor, ) import torch from PIL import Image classes_map = { 0: "Caption", 1: "Footnote", 2: "Formula", 3: "List-item", 4: "Page-footer", 5: "Page-header", 6: "Picture", 7: "Section-header", 8: "Table", 9: "Text", 10: "Title", 11: "Document Index", 12: "Code", 13: "Checkbox-Selected", 14: "Checkbox-Unselected", 15: "Form", 16: "Key-Value Region", } image_url = "https://huggingface.co/spaces/ds4sd/SmolDocling-256M-Demo/resolve/main/example_images/annual_rep_14.png" model_name = "ds4sd/docling-layout-egret-medium" threshold = 0.6 # Download the image image = Image.open(requests.get(image_url, stream=True).raw) image = image.convert("RGB") # Initialize the model image_processor = RTDetrImageProcessor.from_pretrained(model_name) model = DFineForObjectDetection.from_pretrained(model_name) # Run the prediction pipeline inputs = image_processor(images=[image], return_tensors="pt") with torch.no_grad(): outputs = model(**inputs) results = image_processor.post_process_object_detection( outputs, target_sizes=torch.tensor([image.size[::-1]]), threshold=threshold, ) # Get the results for result in results: for score, label_id, box in zip( result["scores"], result["labels"], result["boxes"] ): score = round(score.item(), 2) label = classes_map[label_id.item()] box = [round(i, 2) for i in box.tolist()] print(f"{label}:{score} {box}") ``` # References ``` @techreport{Docling, author = {Deep Search Team}, month = {8}, title = {Docling Technical Report}, url = {https://arxiv.org/abs/2408.09869v4}, eprint = {2408.09869}, doi = {10.48550/arXiv.2408.09869}, version = {1.0.0}, year = {2024} } @misc{peng2024dfine, title={D-FINE: Redefine Regression Task in DETRs as Fine-grained Distribution Refinement}, author={Yansong Peng and Hebei Li and Peixi Wu and Yueyi Zhang and Xiaoyan Sun and Feng Wu}, year={2024}, eprint={2410.13842}, archivePrefix={arXiv}, primaryClass={cs.CV} } ```