Merge pull request #41 from opendatalab/dev

Dev

Author: johnking0099

mineru_vl_utils/__init__.py

@@ -3,7 +3,6 @@
 
 from mineru_vl_utils.version import __version__
 
-
 __lazy_attrs__ = {
     "MinerUClient": (".mineru_client", "MinerUClient"),
     "MinerUSamplingParams": (".mineru_client", "MinerUSamplingParams"),

mineru_vl_utils/mineru_client.py

@@ -154,16 +154,22 @@ def prepare_for_extract(
         self,
         image: Image.Image,
         blocks: list[ContentBlock],
+        not_extract_list: list[str] | None = None,
     ) -> tuple[list[Image.Image | bytes], list[str], list[SamplingParams | None], list[int]]:
         image = get_rgb_image(image)
         width, height = image.size
         block_images: list[Image.Image | bytes] = []
         prompts: list[str] = []
         sampling_params: list[SamplingParams | None] = []
         indices: list[int] = []
+        skip_list = {"image", "list", "equation_block"}
+        if not_extract_list:
+            for not_extract_type in not_extract_list:
+                if not_extract_type in BLOCK_TYPES:
+                    skip_list.add(not_extract_type)
         for idx, block in enumerate(blocks):
-            if block.type in ("image", "list", "equation_block"):
-                continue  # Skip image blocks.
+            if block.type in skip_list:
+                continue  # Skip blocks that should not be extracted.
             x1, y1, x2, y2 = block.bbox
             scaled_bbox = (x1 * width, y1 * height, x2 * width, y2 * height)
             block_image = image.crop(scaled_bbox)
@@ -181,13 +187,17 @@ def prepare_for_extract(
         return block_images, prompts, sampling_params, indices
 
     def post_process(self, blocks: list[ContentBlock]) -> list[ContentBlock]:
-        return post_process(
-            blocks,
-            handle_equation_block=self.handle_equation_block,
-            abandon_list=self.abandon_list,
-            abandon_paratext=self.abandon_paratext,
-            debug=self.debug,
-        )
+        try:
+            return post_process(
+                blocks,
+                handle_equation_block=self.handle_equation_block,
+                abandon_list=self.abandon_list,
+                abandon_paratext=self.abandon_paratext,
+                debug=self.debug,
+            )
+        except Exception as e:
+            print(f"Warning: post-processing failed with error: {e}")
+            return blocks
 
     def batch_prepare_for_layout(
         self,
@@ -212,10 +222,11 @@ def batch_prepare_for_extract(
         executor: Executor | None,
         images: list[Image.Image],
         blocks_list: list[list[ContentBlock]],
+        not_extract_list: list[str] | None = None,
     ) -> list[tuple[list[Image.Image | bytes], list[str], list[SamplingParams | None], list[int]]]:
         if executor is None:
-            return [self.prepare_for_extract(im, bls) for im, bls in zip(images, blocks_list)]
-        return list(executor.map(self.prepare_for_extract, images, blocks_list))
+            return [self.prepare_for_extract(im, bls, not_extract_list) for im, bls in zip(images, blocks_list)]
+        return list(executor.map(self.prepare_for_extract, images, blocks_list, [not_extract_list] * len(images)))
 
     def batch_post_process(
         self,
@@ -247,9 +258,10 @@ async def aio_prepare_for_extract(
         executor: Executor | None,
         image: Image.Image,
         blocks: list[ContentBlock],
+        not_extract_list: list[str] | None = None,
     ) -> tuple[list[Image.Image | bytes], list[str], list[SamplingParams | None], list[int]]:
         loop = asyncio.get_running_loop()
-        return await loop.run_in_executor(executor, self.prepare_for_extract, image, blocks)
+        return await loop.run_in_executor(executor, self.prepare_for_extract, image, blocks, not_extract_list)
 
     async def aio_post_process(
         self,
@@ -607,9 +619,10 @@ def two_step_extract(
         self,
         image: Image.Image,
         priority: int | None = None,
+        not_extract_list: list[str] | None = None,
     ) -> list[ContentBlock]:
         blocks = self.layout_detect(image, priority)
-        block_images, prompts, params, indices = self.helper.prepare_for_extract(image, blocks)
+        block_images, prompts, params, indices = self.helper.prepare_for_extract(image, blocks, not_extract_list)
         outputs = self.client.batch_predict(block_images, prompts, params, priority)
         for idx, output in zip(indices, outputs):
             blocks[idx].content = output
@@ -620,10 +633,16 @@ async def aio_two_step_extract(
         image: Image.Image,
         priority: int | None = None,
         semaphore: asyncio.Semaphore | None = None,
+        not_extract_list: list[str] | None = None,
     ) -> list[ContentBlock]:
         semaphore = semaphore or asyncio.Semaphore(self.max_concurrency)
         blocks = await self.aio_layout_detect(image, priority, semaphore)
-        block_images, prompts, params, indices = await self.helper.aio_prepare_for_extract(self.executor, image, blocks)
+        block_images, prompts, params, indices = await self.helper.aio_prepare_for_extract(
+            self.executor,
+            image,
+            blocks,
+            not_extract_list,
+        )
         outputs = await self.client.aio_batch_predict(block_images, prompts, params, priority, semaphore=semaphore)
         for idx, output in zip(indices, outputs):
             blocks[idx].content = output
@@ -633,13 +652,14 @@ def concurrent_two_step_extract(
         self,
         images: list[Image.Image],
         priority: Sequence[int | None] | int | None = None,
+        not_extract_list: list[str] | None = None,
     ) -> list[list[ContentBlock]]:
         try:
             loop = asyncio.get_running_loop()
         except RuntimeError:
             loop = None
 
-        task = self.aio_concurrent_two_step_extract(images, priority)
+        task = self.aio_concurrent_two_step_extract(images, priority, not_extract_list)
 
         if loop is not None:
             return loop.run_until_complete(task)
@@ -650,6 +670,7 @@ async def aio_concurrent_two_step_extract(
         self,
         images: list[Image.Image],
         priority: Sequence[int | None] | int | None = None,
+        not_extract_list: list[str] | None = None,
         semaphore: asyncio.Semaphore | None = None,
     ) -> list[list[ContentBlock]]:
         if priority is None and self.incremental_priority:
@@ -658,7 +679,7 @@ async def aio_concurrent_two_step_extract(
             priority = [priority] * len(images)
         semaphore = semaphore or asyncio.Semaphore(self.max_concurrency)
         return await gather_tasks(
-            tasks=[self.aio_two_step_extract(*args, semaphore) for args in zip(images, priority)],
+            tasks=[self.aio_two_step_extract(*args, semaphore, not_extract_list) for args in zip(images, priority)],
             use_tqdm=self.use_tqdm,
             tqdm_desc="Two Step Extraction",
         )
@@ -667,6 +688,7 @@ def stepping_two_step_extract(
         self,
         images: list[Image.Image],
         priority: Sequence[int | None] | int | None = None,
+        not_extract_list: list[str] | None = None,
     ) -> list[list[ContentBlock]]:
         if priority is None and self.incremental_priority:
             priority = list(range(len(images)))
@@ -675,7 +697,12 @@ def stepping_two_step_extract(
         all_prompts: list[str] = []
         all_params: list[SamplingParams | None] = []
         all_indices: list[tuple[int, int]] = []
-        prepared_inputs = self.helper.batch_prepare_for_extract(self.executor, images, blocks_list)
+        prepared_inputs = self.helper.batch_prepare_for_extract(
+            self.executor,
+            images,
+            blocks_list,
+            not_extract_list,
+        )
         for img_idx, (block_images, prompts, params, indices) in enumerate(prepared_inputs):
             all_images.extend(block_images)
             all_prompts.extend(prompts)
@@ -690,6 +717,7 @@ async def aio_stepping_two_step_extract(
         self,
         images: list[Image.Image],
         priority: Sequence[int | None] | int | None = None,
+        not_extract_list: list[str] | None = None,
         semaphore: asyncio.Semaphore | None = None,
     ) -> list[list[ContentBlock]]:
         if priority is None and self.incremental_priority:
@@ -701,7 +729,14 @@ async def aio_stepping_two_step_extract(
         all_params: list[SamplingParams | None] = []
         all_indices: list[tuple[int, int]] = []
         prepared_inputs = await gather_tasks(
-            tasks=[self.helper.aio_prepare_for_extract(self.executor, *args) for args in zip(images, blocks_list)],
+            tasks=[
+                self.helper.aio_prepare_for_extract(
+                    self.executor,
+                    *args,
+                    not_extract_list,
+                )
+                for args in zip(images, blocks_list)
+            ],
             use_tqdm=self.use_tqdm,
             tqdm_desc="Extract Preparation",
         )
@@ -731,20 +766,22 @@ def batch_two_step_extract(
         self,
         images: list[Image.Image],
         priority: Sequence[int | None] | int | None = None,
+        not_extract_list: list[str] | None = None,
     ) -> list[list[ContentBlock]]:
         if self.batching_mode == "concurrent":
-            return self.concurrent_two_step_extract(images, priority)
+            return self.concurrent_two_step_extract(images, priority, not_extract_list)
         else:  # self.batching_mode == "stepping"
-            return self.stepping_two_step_extract(images, priority)
+            return self.stepping_two_step_extract(images, priority, not_extract_list)
 
     async def aio_batch_two_step_extract(
         self,
         images: list[Image.Image],
         priority: Sequence[int | None] | int | None = None,
+        not_extract_list: list[str] | None = None,
         semaphore: asyncio.Semaphore | None = None,
     ) -> list[list[ContentBlock]]:
         semaphore = semaphore or asyncio.Semaphore(self.max_concurrency)
         if self.batching_mode == "concurrent":
-            return await self.aio_concurrent_two_step_extract(images, priority, semaphore)
+            return await self.aio_concurrent_two_step_extract(images, priority, not_extract_list, semaphore)
         else:  # self.batching_mode == "stepping"
-            return await self.aio_stepping_two_step_extract(images, priority, semaphore)
+            return await self.aio_stepping_two_step_extract(images, priority, not_extract_list, semaphore)

mineru_vl_utils/post_process/__init__.py

@@ -1,11 +1,11 @@
-from .equation_unbalanced_braces import try_fix_unbalanced_braces
 from ..structs import ContentBlock
+from .equation_big import try_fix_equation_big
 from .equation_block import do_handle_equation_block
 from .equation_double_subscript import try_fix_equation_double_subscript
 from .equation_fix_eqqcolon import try_fix_equation_eqqcolon
-from .equation_big import try_fix_equation_big
-from .equation_leq import try_fix_equation_leq
 from .equation_left_right import try_match_equation_left_right
+from .equation_leq import try_fix_equation_leq
+from .equation_unbalanced_braces import try_fix_unbalanced_braces
 from .otsl2html import convert_otsl_to_html
 
 PARATEXT_TYPES = {
@@ -46,9 +46,17 @@ def post_process(
 ) -> list[ContentBlock]:
     for block in blocks:
         if block.type == "table" and block.content:
-            block.content = convert_otsl_to_html(block.content)
+            try:
+                block.content = convert_otsl_to_html(block.content)
+            except Exception as e:
+                print("Warning: Failed to convert OTSL to HTML: ", e)
+                print("Content: ", block.content)
         if block.type == "equation" and block.content:
-            block.content = _process_equation(block.content, debug=debug)
+            try:
+                block.content = _process_equation(block.content, debug=debug)
+            except Exception as e:
+                print("Warning: Failed to process equation: ", e)
+                print("Content: ", block.content)
 
     if handle_equation_block:
         blocks = do_handle_equation_block(blocks, debug=debug)

mineru_vl_utils/version.py

@@ -1 +1 @@
-__version__ = "0.1.17"
+__version__ = "0.1.17.1"