Browse Source

feat: add DocLayoutYOLOModel for document layout prediction, including batch processing and visualization capabilities, along with a new test image for validation

zhch158_admin 17 hours ago
parent
commit
0f448107b5

BIN
zhch/universal_doc_parser/tests/2023年度报告母公司_page_003_270.png


+ 152 - 0
zhch/universal_doc_parser/tests/test_doclayoutyolo.py

@@ -0,0 +1,152 @@
+import os
+from typing import List, Dict, Union
+
+from doclayout_yolo import YOLOv10
+from tqdm import tqdm
+import numpy as np
+from PIL import Image, ImageDraw
+
+from mineru.utils.enum_class import ModelPath
+from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
+
+
+class DocLayoutYOLOModel:
+    def __init__(
+        self,
+        weight: str,
+        device: str = "cuda",
+        imgsz: int = 1280,
+        conf: float = 0.1,
+        iou: float = 0.45,
+    ):
+        self.model = YOLOv10(weight).to(device)
+        self.device = device
+        self.imgsz = imgsz
+        self.conf = conf
+        self.iou = iou
+
+    def _parse_prediction(self, prediction) -> List[Dict]:
+        layout_res = []
+
+        # 容错处理
+        if not hasattr(prediction, "boxes") or prediction.boxes is None:
+            return layout_res
+
+        for xyxy, conf, cls in zip(
+            prediction.boxes.xyxy.cpu(),
+            prediction.boxes.conf.cpu(),
+            prediction.boxes.cls.cpu(),
+        ):
+            coords = list(map(int, xyxy.tolist()))
+            xmin, ymin, xmax, ymax = coords
+            layout_res.append({
+                "category_id": int(cls.item()),
+                "poly": [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
+                "score": round(float(conf.item()), 3),
+            })
+        return layout_res
+
+    def predict(self, image: Union[np.ndarray, Image.Image]) -> List[Dict]:
+        prediction = self.model.predict(
+            image,
+            imgsz=self.imgsz,
+            conf=self.conf,
+            iou=self.iou,
+            verbose=False
+        )[0]
+        return self._parse_prediction(prediction)
+
+    def batch_predict(
+        self,
+        images: List[Union[np.ndarray, Image.Image]],
+        batch_size: int = 4
+    ) -> List[List[Dict]]:
+        results = []
+        with tqdm(total=len(images), desc="Layout Predict") as pbar:
+            for idx in range(0, len(images), batch_size):
+                batch = images[idx: idx + batch_size]
+                if batch_size == 1:
+                    conf = 0.9 * self.conf
+                else:
+                    conf = self.conf
+                predictions = self.model.predict(
+                    batch,
+                    imgsz=self.imgsz,
+                    conf=conf,
+                    iou=self.iou,
+                    verbose=False,
+                )
+                for pred in predictions:
+                    results.append(self._parse_prediction(pred))
+                pbar.update(len(batch))
+        return results
+
+    # DocLayout-YOLO 类别映射
+    CATEGORY_NAMES = {
+        0: "title",
+        1: "text", 
+        2: "abandon",
+        3: "figure",
+        4: "figure_caption",
+        5: "table",
+        6: "table_caption",
+        7: "table_footnote",
+        8: "isolate_formula",
+        9: "formula_caption",
+    }
+    
+    # 不同类别使用不同颜色
+    CATEGORY_COLORS = {
+        0: "red",           # title
+        1: "blue",          # text
+        2: "gray",          # abandon
+        3: "green",         # figure
+        4: "lightgreen",    # figure_caption
+        5: "orange",        # table
+        6: "yellow",        # table_caption
+        7: "pink",          # table_footnote
+        8: "purple",        # isolate_formula
+        9: "cyan",          # formula_caption
+    }
+
+    def visualize(
+            self,
+            image: Union[np.ndarray, Image.Image],
+            results: List
+    ) -> Image.Image:
+
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+        draw = ImageDraw.Draw(image)
+        for res in results:
+            poly = res['poly']
+            xmin, ymin, xmax, ymax = poly[0], poly[1], poly[4], poly[5]
+            category_id = res['category_id']
+            category_name = self.CATEGORY_NAMES.get(category_id, f"unknown_{category_id}")
+            color = self.CATEGORY_COLORS.get(category_id, "red")
+            
+            print(
+                f"Detected box: {xmin}, {ymin}, {xmax}, {ymax}, Category: {category_name}({category_id}), Score: {res['score']}")
+            # 使用PIL在图像上画框
+            draw.rectangle([xmin, ymin, xmax, ymax], outline=color, width=2)
+            # 在框旁边画类别名和置信度
+            label = f"{category_name} {res['score']:.2f}"
+            draw.text((xmin, ymin - 25), label, fill=color, font_size=20)
+        return image
+
+
+if __name__ == '__main__':
+    image_path = "./2023年度报告母公司_page_003_270.png"
+    doclayout_yolo_weights = os.path.join(auto_download_and_get_model_root_path(ModelPath.doclayout_yolo), ModelPath.doclayout_yolo)
+    device = 'cpu'
+    model = DocLayoutYOLOModel(
+        weight=doclayout_yolo_weights,
+        device=device,
+    )
+    image = Image.open(image_path)
+    results = model.predict(image)
+
+    image = model.visualize(image, results)
+
+    image.show()  # 显示图像