| --- |
| library_name: transformers |
| license: apache-2.0 |
| pipeline_tag: image-segmentation |
| datasets: |
| - ds4sd/DocLayNet |
| --- |
| |
| # DIT-base-layout-detection |
|
|
| We present the model cmarkea/dit-base-layout-detection, which allows extracting different layouts (Text, Picture, Caption, Footnote, etc.) from an image of a document. |
| This is a fine-tuning of the model [dit-base](https://huggingface.co/microsoft/dit-base) on the [DocLayNet](https://huggingface.co/datasets/ds4sd/DocLayNet) |
| dataset. It is ideal for processing documentary corpora to be ingested into an |
| ODQA system. |
|
|
| This model allows extracting 11 entities, which are: Caption, Footnote, Formula, List-item, Page-footer, Page-header, Picture, Section-header, Table, Text, and Title. |
|
|
| ## Performance |
|
|
| In this section, we will assess the model's performance by separately considering semantic segmentation and object detection. We did not perform any post-processing |
| for the semantic segmentation. As for object detection, we only applied OpenCV's `findContours` without any further post-processing. |
|
|
| For semantic segmentation, we will use the F1-score to evaluate the classification of each pixel. For object detection, we will assess performance based on the |
| Generalized Intersection over Union (GIoU) and the accuracy of the predicted bounding box class. The evaluation is conducted on 500 pages from the PDF evaluation |
| dataset of DocLayNet. |
|
|
| | Class | f1-score (x100) | GIoU (x100) | accuracy (x100) | |
| |:--------------:|:---------------:|:-----------:|:---------------:| |
| | Background | 94.98 | NA | NA | |
| | Caption | 75.54 | 55.61 | 72.62 | |
| | Footnote | 72.29 | 50.08 | 70.97 | |
| | Formula | 82.29 | 49.91 | 94.48 | |
| | List-item | 67.56 | 35.19 | 69 | |
| | Page-footer | 83.93 | 57.99 | 94.06 | |
| | Page-header | 62.33 | 65.25 | 79.39 | |
| | Picture | 78.32 | 58.22 | 92.71 | |
| | Section-header | 69.55 | 56.64 | 78.29 | |
| | Table | 83.69 | 63.03 | 90.13 | |
| | Text | 90.94 | 51.89 | 88.09 | |
| | Title | 61.19 | 52.64 | 70 | |
|
|
| ## Benchmark |
|
|
| Now, let's compare the performance of this model with other models. |
|
|
| | Model | f1-score (x100) | GIoU (x100) | accuracy (x100) | |
| |:---------------------------------------------------------------------------------------------:|:---------------:|:-----------:|:---------------:| |
| | cmarkea/dit-base-layout-detection | 90.77 | 56.29 | 85.26 | |
| | [cmarkea/detr-layout-detection](https://huggingface.co/cmarkea/detr-layout-detection) | 91.27 | 80.66 | 90.46 | |
|
|
| ### Direct Use |
|
|
| ```python |
| import torch |
| from transformers import AutoImageProcessor, BeitForSemanticSegmentation |
| |
| img_proc = AutoImageProcessor.from_pretrained( |
| "cmarkea/dit-base-layout-detection" |
| ) |
| model = BeitForSemanticSegmentation.from_pretrained( |
| "cmarkea/dit-base-layout-detection" |
| ) |
| |
| img: PIL.Image |
| |
| with torch.inference_mode(): |
| input_ids = img_proc(img, return_tensors='pt') |
| output = model(**input_ids) |
| |
| segmentation = img_proc.post_process_semantic_segmentation( |
| output, |
| target_sizes=[img.size[::-1]] |
| ) |
| ``` |
|
|
| Here is a simple method for detecting bounding boxes from semantic segmentation. This is the method used to calculate the model's performance in object |
| detection, as described in the "Performance" section. The method is provided without any additional post-processing. |
|
|
| ```python |
| import cv2 |
| |
| def detect_bboxes(masks: np.ndarray): |
| r""" |
| A simple bounding box detection function |
| """ |
| detected_blocks = [] |
| contours, _ = cv2.findContours( |
| masks.astype(np.uint8), |
| cv2.RETR_EXTERNAL, |
| cv2.CHAIN_APPROX_SIMPLE |
| ) |
| for contour in list(contours): |
| if len(list(contour)) >= 4: |
| # smallest rectangle containing all points |
| x, y, width, height = cv2.boundingRect(contour) |
| bounding_box = [x, y, x + width, y + height] |
| detected_blocks.append(bounding_box) |
| return detected_blocks |
| |
| bbox_pred = [] |
| for segment in segmentation: |
| boxes, labels = [], [] |
| for ii in range(1, len(model.config.label2id)): |
| mm = segment == ii |
| if mm.sum() > 0: |
| bbx = detect_bboxes(mm.numpy()) |
| boxes.extend(bbx) |
| labels.extend([ii]*len(bbx)) |
| bbox_pred.append(dict(boxes=boxes, labels=labels)) |
| ``` |
|
|
| ### Example |
|
|
|  |
|
|
| ### Citation |
|
|
| ``` |
| @online{DeDitLay, |
| AUTHOR = {Cyrile Delestre}, |
| URL = {https://huggingface.co/cmarkea/dit-base-layout-detection}, |
| YEAR = {2024}, |
| KEYWORDS = {Image Processing ; Transformers ; Layout}, |
| } |
| ``` |
