Remove PDF upload and keep image parsing only

README.md

@@ -17,9 +17,7 @@ models:
 
 # MiniCPM-OCR on ZeroGPU
 
-Upload one page image, or a short PDF. Optional upright correction uses
+Upload one page image. Optional upright correction uses
 `PaddlePaddle/PP-LCNet_x1_0_doc_ori_safetensors`, layout detection uses
 `PaddlePaddle/PP-DocLayoutV3_safetensors`, and each region is read by
 `openbmb/MiniCPM-OCR-1.0`. Seal and chart regions can be skipped.
-
-A PDF is limited to the first two pages so one visit fits the ZeroGPU quota.

app.py

@@ -39,7 +39,6 @@ from transformers import (
 OCR_ID = os.environ.get("OCR_MODEL_ID", "openbmb/MiniCPM-OCR-1.0")
 LAYOUT_ID = os.environ.get("LAYOUT_MODEL_ID", "PaddlePaddle/PP-DocLayoutV3_safetensors")
 ORI_ID = os.environ.get("ORI_MODEL_ID", "PaddlePaddle/PP-LCNet_x1_0_doc_ori_safetensors")
-MAX_PDF_PAGES = 2
 
 PROMPTS = {
     "text": ("Text Recognition:", 2048),
@@ -79,9 +78,9 @@ ocr_processor = AutoProcessor.from_pretrained(OCR_ID)
 ocr_model = MiniCPMOcrForConditionalGeneration.from_pretrained(OCR_ID, dtype=torch.bfloat16).to("cuda").eval()
 
 
-def _duration(image, pdf, upright, recognize_seal, recognize_chart):
+def _duration(image, upright, recognize_seal, recognize_chart):
     del image, upright, recognize_seal, recognize_chart
-    return 120 if pdf else 90
+    return 90
 
 
 def _numbers(value):
@@ -367,43 +366,18 @@ def _parse_page(image: Image.Image, upright: bool, recognize_seal: bool, recogni
     return "\n\n".join(pieces) or "这一页没有识别出内容。"
 
 
-def _pdf_pages(pdf) -> list[Image.Image]:
-    import fitz
-
-    path = pdf if isinstance(pdf, str) else getattr(pdf, "name", "")
-    document = fitz.open(path)
-    pages = []
-    try:
-        for index in range(min(len(document), MAX_PDF_PAGES)):
-            pixmap = document[index].get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False)
-            pages.append(Image.frombytes("RGB", (pixmap.width, pixmap.height), pixmap.samples))
-    finally:
-        document.close()
-    return pages
+@spaces.GPU(duration=_duration)
+def parse(image, upright, recognize_seal, recognize_chart):
+    if image is None:
+        raise gr.Error("请上传一张图片。")
+    return _parse_page(image, bool(upright), bool(recognize_seal), bool(recognize_chart))
 
 
-@spaces.GPU(duration=_duration)
-def parse(image, pdf, upright, recognize_seal, recognize_chart):
-    def one(page):
-        return _parse_page(page, bool(upright), bool(recognize_seal), bool(recognize_chart))
-
-    if pdf is not None:
-        pages = _pdf_pages(pdf)
-        if len(pages) > 1:
-            parts = [f"## 第 {index + 1} 页\n\n{one(page)}" for index, page in enumerate(pages)]
-            return "\n\n".join(parts)
-        if pages:
-            return one(pages[0])
-    if image is not None:
-        return one(image)
-    raise gr.Error("请上传一张图片，或一份 PDF。")
-
-
-def run(image, pdf, upright, recognize_seal, recognize_chart):
-    if image is None and pdf is None:
-        raise gr.Error("请上传一张图片，或一份 PDF。")
+def run(image, upright, recognize_seal, recognize_chart):
+    if image is None:
+        raise gr.Error("请上传一张图片。")
     yield "正在解析…"
-    yield parse(image, pdf, upright, recognize_seal, recognize_chart)
+    yield parse(image, upright, recognize_seal, recognize_chart)
 
 
 if os.environ.get("OCR_SKIP_LAUNCH") != "1":
@@ -413,13 +387,12 @@ if os.environ.get("OCR_SKIP_LAUNCH") != "1":
         gr.Markdown(
             """
 # 文档解析
-上传一页图片，或一份不超过两页的 PDF。开启自适应旋转时，先做小角度摆正和 0°/90°/180°/270° 转向，再做版面检测和识别。
+上传一张图片。开启自适应旋转时，先做小角度摆正和 0°/90°/180°/270° 转向，再做版面检测和识别。
 """
         )
         with gr.Row():
             with gr.Column():
                 image = gr.Image(type="pil", label="图片")
-                pdf = gr.File(label="PDF", file_types=[".pdf"])
                 upright = gr.Checkbox(label="自适应旋转", value=True)
                 recognize_seal = gr.Checkbox(label="印章识别", value=True)
                 recognize_chart = gr.Checkbox(label="图表识别", value=True)
@@ -435,7 +408,7 @@ if os.environ.get("OCR_SKIP_LAUNCH") != "1":
             )
         button.click(
             run,
-            inputs=[image, pdf, upright, recognize_seal, recognize_chart],
+            inputs=[image, upright, recognize_seal, recognize_chart],
             outputs=markdown,
             show_progress="full",
         )

requirements.txt

@@ -1,6 +1,5 @@
 gradio==6.31.0
 transformers==5.21.0
 torchvision==0.28.0
-pymupdf
 opencv-python-headless
 pillow