How to analyze document layout by YOLO

Why we need document layout analysis

Analyzing document layout is critical because it aids in the proper interpretation and understanding of document content. It becomes more significant with the rise of RAG, which relies heavily on the ability to parse documents accurately.

RAG systems frequently interact with a variety of documents. Scientific papers, for instance, typically have a complex layout that includes figures, tables, references, and structured sections. Proper parsing is important to avoid content disarray. If not done correctly, the LLM could fail due to the 'garbage in, garbage out' principle.

However, due to the complexity of this problem, it's impossible to apply handcrafted rules as a solution. The best approach is to train a machine learning model.

My solution

You can find my solution in yolo-doclaynet. After examining several models and datasets, I've chosen YOLO as the base model and DocLayNet as the training data. Let's delve into more details.

1. YOLO is the most advanced vision detection model. It is maintained by Ultralytics, a leading computer vision team. The model is easy to train, evaluate, and deploy. Plus, its size is compact enough to run in a browser or on a smartphone.
2. DocLayNet is a human-annotated document layout segmentation dataset, containing 80,863 pages from a wide variety of document sources. To the best of my knowledge, it is the highest-quality dataset for document layout analysis. You can download and find more information from this link.

Live Demo

1. Download the pretrained model from huggingface. Your options include yolov8n-doclaynet, yolov8s-doclaynet, and yolov8m-doclaynet.
2. Install Ultralytics by executing pip install ultralytics. If you encounter any issues, please refer to https://docs.ultralytics.com/quickstart/#install-ultralytics.

3. Copy and modify this code snippet to output the detection result:

   from ultralytics import YOLO
   
   img = cv2.imread(your_image_path, cv2.IMREAD_COLOR)
   model = YOLO(your_model_path)
   result = model.predict(img)[0]
   print(result)
   

4. Debugging the model can be challenging when you can only check the plain-text output. Fortunately, visualizing the results is simple:

   from ultralytics.utils.plotting import Annotator, Colors
   
   colors = Colors()
   annotator = Annotator(img, line_width=line_width, font_size=font_size)
   for label, box in zip(result.boxes.cls.tolist(), result.boxes.xyxyn.tolist()):
       label = int(label)
       annotator.box_label(
           [box[0] * width, box[1] * height, box[2] * width, box[3] * height],
           result.names[label],
           color=colors(label, bgr=True),
       )
   annotator.save(
       os.path.join(os.path.dirname(image), "annotated-" + os.path.basename(your_image_path))
   )
   

5. Examine the annotated image. Here's an example:

Benchmark

I also evaluate the mAP50-95 performance of the entire yolov8 series on the DocLayNet test set.

Here is an overview of the mAP50-95 performance with different model sizes.