init

app-ocr.py

@@ -0,0 +1,82 @@
+import base64
+import os
+
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+model_name = "ucaslcl/GOT-OCR2_0"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map=device)
+model = model.eval().to(device)
+
+
+@spaces.GPU()
+def ocr_process(image, got_mode, ocr_color="", ocr_box="", progress=gr.Progress()):
+    if image is None:
+        return "错误:未提供图片"
+
+    try:
+        image_path = image
+        result_path = f"{os.path.splitext(image_path)[0]}_result.html"
+
+        progress(0, desc="开始处理...")
+
+        if "plain" in got_mode:
+            progress(0.3, desc="执行OCR识别...")
+            if "multi-crop" in got_mode:
+                res = model.chat_crop(tokenizer, image_path, ocr_type="ocr")
+            else:
+                res = model.chat(tokenizer, image_path, ocr_type="ocr", ocr_box=ocr_box, ocr_color=ocr_color)
+            progress(1, desc="处理完成")
+            return res
+        elif "format" in got_mode:
+            progress(0.3, desc="执行OCR识别...")
+            if "multi-crop" in got_mode:
+                res = model.chat_crop(tokenizer, image_path, ocr_type="format", render=True, save_render_file=result_path)
+            else:
+                res = model.chat(tokenizer, image_path, ocr_type="format", ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file=result_path)
+
+            progress(0.7, desc="生成结果...")
+            if os.path.exists(result_path):
+                with open(result_path, "r", encoding="utf-8") as f:
+                    html_content = f.read()
+                encoded_html = base64.b64encode(html_content.encode("utf-8")).decode("utf-8")
+                data_uri = f"data:text/html;charset=utf-8;base64,{encoded_html}"
+                preview = f'<iframe src="{data_uri}" width="100%" height="600px"></iframe>'
+                download_link = f'<a href="{data_uri}" download="result.html">下载完整结果</a>'
+                progress(1, desc="处理完成")
+                return f"{download_link}\n\n{preview}"
+
+        return "错误: 未知的OCR模式"
+    except Exception as e:
+        return f"错误: {str(e)}"
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# OCR 图像识别")
+
+    with gr.Row():
+        image_input = gr.Image(type="filepath", label="上传图片")
+
+    got_mode = gr.Dropdown(
+        choices=["plain texts OCR", "format texts OCR", "plain multi-crop OCR", "format multi-crop OCR", "plain fine-grained OCR", "format fine-grained OCR"],
+        label="OCR模式",
+        value="plain texts OCR",
+    )
+
+    with gr.Row():
+        ocr_color = gr.Textbox(label="OCR颜色 (仅用于fine-grained模式)")
+        ocr_box = gr.Textbox(label="OCR边界框 (仅用于fine-grained模式)")
+
+    submit_button = gr.Button("开始OCR识别")
+
+    output = gr.HTML(label="识别结果")
+
+    submit_button.click(ocr_process, inputs=[image_input, got_mode, ocr_color, ocr_box], outputs=output)
+
+if __name__ == "__main__":
+    demo.launch()

app.py

@@ -1,8 +1,64 @@
-from fastapi import FastAPI
+import base64
+import os
+
+import torch
+from fastapi import FastAPI, File, Form, UploadFile
+from transformers import AutoModel, AutoTokenizer
 
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+# 初始化模型
+model_name = "ucaslcl/GOT-OCR2_0"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map=device)
+model = model.eval().to(device)
+
+
+# OCR处理函数
+async def ocr_process(image_path, got_mode, ocr_color="", ocr_box=""):
+    try:
+        if "plain" in got_mode:
+            if "multi-crop" in got_mode:
+                res = model.chat_crop(tokenizer, image_path, ocr_type="ocr")
+            else:
+                res = model.chat(tokenizer, image_path, ocr_type="ocr", ocr_box=ocr_box, ocr_color=ocr_color)
+            return res
+        elif "format" in got_mode:
+            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
+            if "multi-crop" in got_mode:
+                res = model.chat_crop(tokenizer, image_path, ocr_type="format", render=True, save_render_file=result_path)
+            else:
+                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 os.path.exists(result_path):
+                with open(result_path, "r", encoding="utf-8") as f:
+                    html_content = f.read()
+                encoded_html = base64.b64encode(html_content.encode("utf-8")).decode("utf-8")
+                return {"html_content": encoded_html}
+
+        return {"error": "未知的OCR模式"}
+    except Exception as e:
+        return {"error": str(e)}
 
+
+@app.post("/ocr")
+async def ocr_api(image: UploadFile = File(...), got_mode: str = Form(...), ocr_color: str = Form(""), ocr_box: str = Form("")):
+    # 保存上传的图片
+    image_path = f"temp_{image.filename}"
+    with open(image_path, "wb") as buffer:
+        buffer.write(await image.read())
+
+    # 处理OCR
+    result = await ocr_process(image_path, got_mode, ocr_color, ocr_box)
+
+    # 删除临时文件
+    os.remove(image_path)
+
+    return result
+
+
+@app.get("/")
+def read_root():
+    return {"message": "欢迎使用OCR API"}

requirements.txt

@@ -1,3 +1,14 @@
 fastapi
 uvicorn[standard]
-
+PyMuPDF
+verovio
+gradio
+numpy
+modelscope
+Pillow
+tiktoken
+transformers
+torch
+torchvision
+accelerate
+spaces