通过FastAPI把Nanonets OCR服务作为接口服务使用

1、Ubuntu本地部署好Nanonets环境和服务

(参考之前的文章https://frankz1998.blogspot.com/2025/10/ubuntuhuggingfacespacenanonets-ocr2.html)

2、编辑app.py文件

import os

from fastapi import FastAPI, HTTPException

from pydantic import BaseModel

from PIL import Image

import torch

from transformers import AutoTokenizer, AutoProcessor, AutoModelForImageTextToText

# -------- 初始化模型 -------- #

model_path = "nanonets/Nanonets-OCR2-3B"

print("Loading model...")

model = AutoModelForImageTextToText.from_pretrained(

    model_path,

    torch_dtype="auto",

    device_map="auto"

)

model.eval()

tokenizer = AutoTokenizer.from_pretrained(model_path)

processor = AutoProcessor.from_pretrained(model_path)

print("Model loaded successfully.")

# -------- 定义 FastAPI 应用 -------- #

app = FastAPI(title="OCR Markdown Generator API", description="Convert images to Markdown files using Nanonets OCR")

class OCRRequest(BaseModel):

    folder_path: str

    max_new_tokens: int = 4096

def ocr_image(image_path: str, max_new_tokens: int = 4096) -> str:

    """

    处理单张图片并返回识别结果

    """

    prompt = """Extract the text from the above document as if you were reading it naturally. Return the tables in html format. Return the equations in LaTeX representation. If there is an image in the document and image caption is not present, add a small description of the image inside the <img></img> tag; otherwise, add the image caption inside <img></img>. Watermarks should be wrapped in brackets. Ex: <watermark>OFFICIAL COPY</watermark>. Page numbers should be wrapped in brackets. Ex: <page_number>14</page_number> or <page_number>9/22</page_number>. Prefer using ☐ and ☑ for check boxes."""

    image = Image.open(image_path)

    messages = [

        {"role": "system", "content": "You are a helpful assistant."},

        {"role": "user", "content": [

            {"type": "image", "image": f"file://{image_path}"},

            {"type": "text", "text": prompt},

        ]},

    ]

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = processor(text=[text], images=[image], padding=True, return_tensors="pt").to(model.device)

    output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)

    generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(inputs.input_ids, output_ids)]

    output_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)

    return output_text[0]

@app.post("/ocr_folder")

def process_folder(request: OCRRequest):

    folder_path = request.folder_path

    max_tokens = request.max_new_tokens

    if not os.path.exists(folder_path):

        raise HTTPException(status_code=404, detail=f"Folder not found: {folder_path}")

    # 输出路径：md/文件夹名/

    folder_name = os.path.basename(folder_path.rstrip("/"))

    output_dir = os.path.join("md", folder_name)

    os.makedirs(output_dir, exist_ok=True)

    # 找出所有 PNG 文件

    png_files = sorted([f for f in os.listdir(folder_path) if f.lower().endswith(".png")])

    if not png_files:

        raise HTTPException(status_code=400, detail="No PNG files found in the given folder.")

    results = []

    for filename in png_files:

        image_path = os.path.join(folder_path, filename)

        md_filename = os.path.splitext(filename)[0] + ".md"

        md_output_path = os.path.join(output_dir, md_filename)

        try:

            print(f"Processing {filename} ...")

            text = ocr_image(image_path, max_tokens)

            with open(md_output_path, "w", encoding="utf-8") as f:

                f.write(text)

            results.append({"file": filename, "output": md_output_path})

        except Exception as e:

            results.append({"file": filename, "error": str(e)})

    return {

        "status": "completed",

        "input_folder": folder_path,

        "output_folder": output_dir,

        "results": results

    }

# -------- 启动服务 -------- #

# 运行命令：uvicorn app:app --host 0.0.0.0 --port 8000

3、编写run.py文件

import sys

import requests

import os

# ----------- 可自定义基础路径 -----------

BASE_PATH = "/home/frank1998/env/Nanonets-OCR2-3B/images/"

# ----------- 主程序开始 -----------

def main():

    if len(sys.argv) < 2:

        print("❌ 用法: python run.py <folder_name>")

        print("例如: python run.py 003")

        sys.exit(1)

    folder_name = sys.argv[1]

    folder_path = os.path.join(BASE_PATH, folder_name)

    # 打印确认路径

    print(f"🧭 正在请求 OCR 处理: {folder_path}")

    # 检查文件夹是否存在

    if not os.path.exists(folder_path):

        print(f"❌ 错误: 文件夹不存在: {folder_path}")

        sys.exit(1)

    # 调用 FastAPI 接口

    try:

        resp = requests.post(

            "http://localhost:8000/ocr_folder",

            json={"folder_path": folder_path}

        )

        if resp.status_code == 200:

            print("✅ 处理完成！结果如下：")

            print(resp.json())

        else:

            print(f"❌ 调用失败: {resp.status_code}")

            print(resp.text)

    except Exception as e:

        print(f"⚠️ 调用接口出错: {e}")

# ----------- 程序入口 -----------

if __name__ == "__main__":

    main()

4、运行fastapi接口服务:

进入项目文件夹,激活虚拟环境

source .venv/bin/activate

export HF_ENDPOINT=https://hf-mirror.com

uvicorn app:app --host 0.0.0.0 --port 8000

5、在另一个命令行终端中运行:

进入项目文件夹,激活虚拟环境

source .venv/bin/activate

(对项目目录下images文件夹内的007文件夹内的图片进行OCR识别)

python run.py 007

6、查看识别后的markdown格式文件

进入项目目录下的md文件夹内,对应着images中的图片文件夹的名称就会生成存储了md文件的对应文件夹.

(例如,images中包含001,002,003文件夹,里面都是png格式的图片,则通过python run.py 001接口调用后,在md文件夹中会生成一个名为001的文件夹,里面包含了对应的png图片OCR识别后的md文件)