add: pdf2images

.gitignore

@@ -0,0 +1,2 @@
+results/*
+uploads/*

app.py

@@ -1,52 +1,95 @@
 import os
+import shutil
 import uuid
 
 import fitz  # PyMuPDF
 import gradio as gr
-from PIL import Image
+from modelscope import AutoModel, AutoTokenizer
+from PIL import Image, ImageEnhance
+
+from got_ocr import got_ocr
+
+# 初始化模型和分词器
+tokenizer = AutoTokenizer.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True)
+model = AutoModel.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True, low_cpu_mem_usage=True, device_map="cuda", use_safetensors=True)
+model = model.eval().cuda()
 
 UPLOAD_FOLDER = "./uploads"
 RESULTS_FOLDER = "./results"
 
+# 确保必要的文件夹存在
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(RESULTS_FOLDER, exist_ok=True)
+
 
 def pdf_to_images(pdf_path):
     images = []
     pdf_document = fitz.open(pdf_path)
     for page_num in range(len(pdf_document)):
         page = pdf_document.load_page(page_num)
-        pix = page.get_pixmap()
+        # 进一步增加分辨率和缩放比例
+        zoom = 4  # 增加缩放比例到4
+        mat = fitz.Matrix(zoom, zoom)
+        pix = page.get_pixmap(matrix=mat, alpha=False)
         img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+
+        # 增加对比度
+        enhancer = ImageEnhance.Contrast(img)
+        img = enhancer.enhance(1.5)  # 增加50%的对比度
+
         images.append(img)
     pdf_document.close()
     return images
 
 
 def process_pdf(pdf_file):
+    if pdf_file is None:
+        return None
+
     temp_pdf_path = os.path.join(UPLOAD_FOLDER, f"{uuid.uuid4()}.pdf")
-    pdf_file.save(temp_pdf_path)
+
+    # 使用 shutil 复制上传的件到临时位置
+    shutil.copy(pdf_file.name, temp_pdf_path)
+
     images = pdf_to_images(temp_pdf_path)
     os.remove(temp_pdf_path)
-    return images
+
+    # 将图像保存为临时文件并返回文件路径列表
+    image_paths = []
+    for i, img in enumerate(images):
+        img_path = os.path.join(RESULTS_FOLDER, f"page_{i+1}.png")
+        img.save(img_path, "PNG")
+        image_paths.append(img_path)
+
+    return image_paths
 
 
-def display_images(images):
-    image_elements = [gr.Image(value=img, type="pil") for img in images]
-    return gr.Gallery(value=image_elements)
+def on_image_select(evt: gr.SelectData):
+    if evt.index is not None:
+        return evt.index
+    return None
 
 
-def on_image_select(image):
-    return image
+def perform_ocr(selected_index, image_paths):
+    if selected_index is not None and image_paths and 0 <= selected_index < len(image_paths):
+        selected_image = image_paths[selected_index][0]
+        # 这里添加OCR处理逻辑
+        return got_ocr(model, selected_image)
+    return "请先选择一个图片"
 
 
 with gr.Blocks() as demo:
     pdf_input = gr.File(label="上传PDF文件")
     image_gallery = gr.Gallery(label="PDF页面预览", columns=3, height="auto")
-    selected_image = gr.Image(label="选中的图片", type="pil")
+    selected_index = gr.State(None)
+    ocr_button = gr.Button("开始OCR识别")
+    ocr_result = gr.Textbox(label="OCR结果")
 
     pdf_input.upload(fn=process_pdf, inputs=pdf_input, outputs=image_gallery)
-    image_gallery.select(fn=on_image_select, inputs=image_gallery, outputs=selected_image)
+    image_gallery.select(fn=on_image_select, inputs=[], outputs=selected_index)
+    ocr_button.click(fn=perform_ocr, inputs=[selected_index, image_gallery], outputs=ocr_result)
 
-    # 这里可以添加OCR转换功能的相关组件和逻辑
+    # 移除了选中图片的显示部分
 
 if __name__ == "__main__":
     demo.launch()

got_ocr.py

@@ -0,0 +1,53 @@
+import base64
+import os
+
+
+def got_ocr(model, image_path, got_mode="plain texts OCR", fine_grained_mode="", ocr_color="", ocr_box=""):
+    # 执行OCR
+    try:
+        if got_mode == "plain texts OCR":
+            res = model.chat(tokenizer, image_path, ocr_type="ocr")
+            return res, None
+        elif got_mode == "format texts OCR":
+            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
+            res = model.chat(tokenizer, image_path, ocr_type="format", render=True, save_render_file=result_path)
+        elif got_mode == "plain multi-crop OCR":
+            res = model.chat_crop(tokenizer, image_path, ocr_type="ocr")
+            return res, None
+        elif got_mode == "format multi-crop OCR":
+            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
+            res = model.chat_crop(tokenizer, image_path, ocr_type="format", render=True, save_render_file=result_path)
+        elif got_mode == "plain fine-grained OCR":
+            res = model.chat(tokenizer, image_path, ocr_type="ocr", ocr_box=ocr_box, ocr_color=ocr_color)
+            return res, None
+        elif got_mode == "format fine-grained OCR":
+            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
+            res = model.chat(tokenizer, image_path, ocr_type="format", ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file=result_path)
+
+        # 处理格式化结果
+        if "format" in got_mode and os.path.exists(result_path):
+            with open(result_path, "r") as f:
+                html_content = f.read()
+            encoded_html = base64.b64encode(html_content.encode("utf-8")).decode("utf-8")
+            return res, encoded_html
+        else:
+            return res, None
+
+    except Exception as e:
+        return f"错误: {str(e)}", None
+
+
+# 使用示例
+if __name__ == "__main__":
+    from modelscope import AutoModel, AutoTokenizer
+
+    # 初始化模型和分词器
+    tokenizer = AutoTokenizer.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True)
+    model = AutoModel.from_pretrained("stepfun-ai/GOT-OCR2_0", trust_remote_code=True, low_cpu_mem_usage=True, device_map="cuda", use_safetensors=True)
+    model = model.eval().cuda()
+
+    image_path = "path/to/your/image.png"
+    result, html = got_ocr(model, image_path, got_mode="format texts OCR")
+    print("OCR结果:", result)
+    if html:
+        print("HTML结果可用")