github下载地址: https://github.com/deepseek-ai/DeepSeek-OCR-2    

按照文档操作安装环境即可

conda create -n deepseek-ocr2 python=3.12.9 -y
conda activate deepseek-ocr2

pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu118
pip install vllm-0.8.5+cu118-cp38-abi3-manylinux1_x86_64.whl
pip install -r requirements.txt
pip install flash-attn==2.7.3 --no-build-isolation

有一个问题在windows系统下比较难搞:pip install flash-attn==2.7.3 --no-build-isolation

flash-attn 安装比较费劲,我多次尝试未能安装成功,可以选择别人已经弄好whl文件,一定要搭配好自己安装好的环境!!感谢好心人!!!!

flash-attn 可用的windows版本下载whl地址: https://github.com/kingbri1/flash-attention/releases  

环境弄好后,就可以运行代码了:

cd DeepSeek-OCR2-master/DeepSeek-OCR2-hf
python run_dpsk_ocr2.py

我运行稍微有些报错,代码我稍做了修改,首次运行,需要梯子下载模型!之后便可以访问本地已经下载好的模型即可。

默认下载地址:C:\Users\用户名\.cache\huggingface\hub\models--deepseek-ai--DeepSeek-OCR-2\snapshots\aaa02xxxxxx 目录下

简版代码:

import os
import torch
from transformers import AutoModel, AutoTokenizer

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

model_name = (
    r"C:\Users\用户名\.cache\huggingface\hub\models--deepseek-ai--DeepSeek-OCR-2\snapshots\aaa02f3"
)

tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    trust_remote_code=True
)

model = AutoModel.from_pretrained(
    model_name,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2"
).to("cuda").eval()

prompt = "<image>\n<|grounding|>Convert the document to markdown."

image_file = r"F:\DeepSeek-OCR2-master\DeepSeek-OCR2-hf\jingdong_images\jingdong_40c4bf6b.jpeg"
output_path = r"F:\DeepSeek-OCR-2-main\DeepSeek-OCR2-master\DeepSeek-OCR2-hf\output"

assert os.path.exists(image_file), f"Image not found: {image_file}"

res = model.infer(
    tokenizer,
    prompt=prompt,
    image_file=image_file,
    output_path=output_path,
    base_size=1024,
    image_size=768,
    crop_mode=True,
    save_results=True
)

拓展代码,添加读取文件夹和转换为txt

import os
import shutil
import time
import torch
from transformers import AutoModel, AutoTokenizer


class DeepSeekOCR:
    def __init__(
        self,
        model_path: str,
        device: str = "cuda",
        base_size: int = 1024,
        image_size: int = 768,
    ):
        self.device = device
        self.base_size = base_size
        self.image_size = image_size

        init_start = time.perf_counter()
        print("Initializing DeepSeekOCR...")

        # -------- tokenizer --------
        t0 = time.perf_counter()
        print("Loading tokenizer...")
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path, trust_remote_code=True
        )
        print(f"Tokenizer loaded in {time.perf_counter() - t0:.2f}s")

        # -------- model --------
        t1 = time.perf_counter()
        print("Loading model...")
        self.model = AutoModel.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            _attn_implementation="flash_attention_2",
            use_safetensors=True,
        ).to(device).eval()
        print(f"Model loaded in {time.perf_counter() - t1:.2f}s")

        self.prompt = "<image>\n<|grounding|>Convert the document to markdown."

        total_init_cost = time.perf_counter() - init_start
        print(f"Model initialization done. Total init time: {total_init_cost:.2f}s")


    # -----------------------------
    # 内部:单次推理(含计时)
    # -----------------------------

    def _infer_once(self, image_path: str, output_dir: str) -> float:
        start = time.perf_counter()

        self.model.infer(
            self.tokenizer,
            prompt=self.prompt,
            image_file=image_path,
            output_path=output_dir,
            base_size=self.base_size,
            image_size=self.image_size,
            crop_mode=True,
            save_results=True, # 后期关闭可减少时间3s左右
        )

        return time.perf_counter() - start

    

    # -----------------------------
    # 内部:重命名输出文件
    # -----------------------------
    def _rename_outputs(self, output_dir: str, base_name: str):
        # result.mmd
        src_mmd = os.path.join(output_dir, "result.mmd")
        dst_mmd = os.path.join(output_dir, f"{base_name}.mmd")
        if os.path.exists(src_mmd):
            shutil.move(src_mmd, dst_mmd)

        # result_with_boxes.jpg
        src_box = os.path.join(output_dir, "result_with_boxes.jpg")
        dst_box = os.path.join(output_dir, f"{base_name}_with_boxes.jpg")
        if os.path.exists(src_box):
            shutil.move(src_box, dst_box)

    # -----------------------------
    # 对外:单张 OCR
    # -----------------------------
    def ocr_image(self, image_path: str, output_dir: str):
        assert os.path.exists(image_path), f"Image not found: {image_path}"
        os.makedirs(output_dir, exist_ok=True)

        image_name = os.path.basename(image_path)
        base_name = os.path.splitext(image_name)[0]

        print(f"Start OCR: {image_name}")

        cost = self._infer_once(image_path, output_dir)
        self._rename_outputs(output_dir, base_name)

        mmd_path = os.path.join(output_dir, f"{base_name}.mmd")
        self._mmd_to_txt(mmd_path)

        print(f"OCR done: {image_name} | Time: {cost:.3f}s")

    # -----------------------------
    # 对外:批量 OCR
    # -----------------------------
    def ocr_folder(
        self,
        image_dir: str,
        output_dir: str,
        exts=(".jpg", ".jpeg", ".png", ".bmp", ".tiff"),
    ):
        assert os.path.isdir(image_dir), f"Dir not found: {image_dir}"
        os.makedirs(output_dir, exist_ok=True)

        images = [
            f for f in os.listdir(image_dir)
            if f.lower().endswith(exts)
        ]

        if not images:
            print("No images found.")
            return

        print(f"Found {len(images)} images")
        total_start = time.perf_counter()

        for idx, img in enumerate(images, 1):
            img_path = os.path.join(image_dir, img)
            base_name = os.path.splitext(img)[0]

            print(f"[{idx}/{len(images)}] OCR: {img}")

            try:
                cost = self._infer_once(img_path, output_dir)
                self._rename_outputs(output_dir, base_name)

                mmd_path = os.path.join(output_dir, f"{base_name}.mmd")
                self._mmd_to_txt(mmd_path)
                print(f"Done: {img} | Time: {cost:.3f}s")

            except Exception as e:
                print(f"Failed: {img}, Reason: {e}")

        total_cost = time.perf_counter() - total_start
        print("-" * 50)
        print(f"Total images: {len(images)}")
        print(f"Total time  : {total_cost:.2f}s")
        print(f"Avg per img : {total_cost / len(images):.2f}s")

    # -----------------------------
    # 内部:mmd 转 txt
    # -----------------------------
    def _mmd_to_txt(self, mmd_path: str):
        if not os.path.exists(mmd_path):
            return

        txt_path = os.path.splitext(mmd_path)[0] + ".txt"

        with open(mmd_path, "r", encoding="utf-8") as f:
            content = f.read()

        # 简单 Markdown 清洗(保留纯文本)
        lines = []
        for line in content.splitlines():
            line = line.strip()

            # 跳过图片、表格分隔线
            if line.startswith("![]"):
                continue
            if set(line) <= {"|", "-", " "}:
                continue

            # 去掉 markdown 符号
            for ch in ["#", "*", "`", "_", ">"]:
                line = line.replace(ch, "")

            lines.append(line)

        text = "\n".join(lines).strip()

        with open(txt_path, "w", encoding="utf-8") as f:
            f.write(text)



if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    MODEL_PATH = r"F:\DeepSeek-OCR-2-main\DeepSeek-OCR2-master\DeepSeek-OCR2-hf\model"
    OUTPUT_DIR = r"F:\DeepSeek-OCR-2-main\DeepSeek-OCR2-master\DeepSeek-OCR2-hf\output"

    ocr = DeepSeekOCR(MODEL_PATH)

    # 单张
    SINGLE_IMAGE = r"F:\DeepSeek-OCR-2-main\DeepSeek-OCR2-master\DeepSeek-OCR2-hf\zhongtong_images\zhongtong_5a1d2892.jpeg"
    ocr.ocr_image(SINGLE_IMAGE, OUTPUT_DIR)

    # 批量
    # IMAGE_DIR = r"F:\DeepSeek-OCR-2-main\DeepSeek-OCR2-master\DeepSeek-OCR2-hf\zhongtong_images"
    # ocr.ocr_folder(IMAGE_DIR, OUTPUT_DIR)

    print("All OCR done.")

ocr+llm硅基流动整合json

import os
import time
import torch
import json
import re
import cv2
from transformers import AutoModel, AutoTokenizer
from openai import OpenAI


class DeepSeekOCR:

    def __init__(
        self,
        model_path,
        device="cuda",
        base_size=1024,
        image_size=768,
        api_key=None,
        api_base="https://api.siliconflow.cn/v1",
        llm_model="deepseek-ai/DeepSeek-V3",
        enable_resize=True,
        save_intermediate_files=False
    ):

        self.device = device
        self.base_size = base_size
        self.image_size = image_size
        self.enable_resize = enable_resize
        self.save_intermediate_files = save_intermediate_files
        self.api_key = api_key # 记录状态
        self.api_base = api_base
        self.llm_model = llm_model

        print("\n===== 初始化 DeepSeekOCR =====")
        start = time.time()

        # === 优化逻辑:动态配置加速参数 ===
        is_cpu_mode = (device.lower() == 'cpu')
        
        if is_cpu_mode:
            current_dtype = torch.float32
            attn_implementation = "eager"
            torch.set_num_threads(8) 
            print(f"运行配置 -> 设备: CPU | 计算线程: {torch.get_num_threads()}")
        else:
            current_dtype = torch.bfloat16
            attn_implementation = "flash_attention_2"
            print(f"运行配置 -> 设备: GPU ({device})")

        # tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_path,
            trust_remote_code=True
        )

        # model
        self.model = AutoModel.from_pretrained(
            model_path,
            trust_remote_code=True,
            torch_dtype=current_dtype,
            _attn_implementation=attn_implementation,
            use_safetensors=True
        ).to(device).eval()

        print(f"模型加载完成: {time.time() - start:.2f}s")

        self.prompt = "<image>\n<|grounding|>Convert the document to markdown."

        # 按需初始化 LLM
        if api_key:
            try:
                self.client = OpenAI(
                    api_key=api_key,
                    base_url=self.api_base
                )
                print("LLM 客户端初始化成功 (OCR + LLM 模式)")
            except Exception as e:
                print(f"LLM 客户端初始化失败: {e}")
                self.client = None
        else:
            self.client = None
            print("未提供 API Key (仅 OCR 模式)")

    
    # 图像预处理
    def _preprocess(self, image_path):
        if not self.enable_resize:
            return image_path, False

        img = cv2.imread(image_path)
        if img is None:
            return image_path, False

        h, w = img.shape[:2]
        max_side = 1280

        if max(h, w) <= max_side:
            return image_path, False

        scale = max_side / max(h, w)
        new_w = int(w * scale)
        new_h = int(h * scale)
        img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)

        tmp_path = os.path.join(os.path.dirname(image_path), f"__tmp_resize_{int(time.time())}.jpg")
        cv2.imwrite(tmp_path, img)
        
        return tmp_path, True

    
    # OCR (融合 IO 开关控制)
    def run_ocr(self, image_path, output_dir):
        os.makedirs(output_dir, exist_ok=True)
        
        # 1. 预处理
        input_path, is_temp = self._preprocess(image_path)

        print("开始OCR...")
        start = time.time()

        # 2. 推理
        self.model.infer(
            self.tokenizer,
            prompt=self.prompt,
            image_file=input_path,
            output_path=output_dir,
            base_size=self.base_size,
            image_size=self.image_size,
            crop_mode=True,
            save_results=True,
        )
        
        ocr_time = time.time() - start
        print(f"OCR完成: {ocr_time:.2f}s")

        # 3. 结果处理
        md_path = os.path.join(output_dir, "result.mmd")
        box_path = os.path.join(output_dir, "result_with_boxes.jpg")
        text_content = ""

        if os.path.exists(md_path):
            with open(md_path, "r", encoding="utf-8") as f:
                text_content = f.read()

            # IO 开关控制 
            if not self.save_intermediate_files:
                # 模式B: 删除中间文件
                try:
                    if os.path.exists(md_path): os.remove(md_path)
                    if os.path.exists(box_path): os.remove(box_path)
                except Exception: pass
            else:
                # 模式A: 重命名保存
                base_name = os.path.splitext(os.path.basename(image_path))[0]
                try:
                    if os.path.exists(md_path):
                        os.rename(md_path, os.path.join(output_dir, f"{base_name}.mmd"))
                    if os.path.exists(box_path):
                        os.rename(box_path, os.path.join(output_dir, f"{base_name}_box.jpg"))
                except Exception: pass
        
        # 清理预处理临时图
        if is_temp and os.path.exists(input_path):
            try: os.remove(input_path)
            except: pass

        return text_content.strip(), ocr_time

    
    # JSON安全解析
    def _safe_json(self, content):
        content = re.sub(r"```json|```", "", content)
        start = content.find("{")
        end = content.rfind("}") + 1
        try:
            return json.loads(content[start:end])
        except:
            return {"raw_llm_output": content}

    
    # LLM解析
    def parse_waybill(self, text):
        # 再次检查客户端
        if not self.client:
            return None, 0.0

        text_chunk = re.sub(r"\s+", " ", text)[:4000]

        prompt = f"""
从OCR文本中提取快递面单信息,输出JSON。不解释。
OCR:
{text_chunk}

JSON格式:
{{
"express_company": "",
"tracking_number": "",
"receiver": {{"name": "", "phone": "", "address": ""}},
"sender": {{"name": "", "phone": "", "address": ""}},
"product_name": "",
"amount": ""
}}
"""

        print("开始LLM解析...")
        start = time.time()

        try:
            resp = self.client.chat.completions.create(
                model=self.llm_model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0,
                response_format={"type": "json_object"}
            )
            result = self._safe_json(resp.choices[0].message.content)
        except Exception as e:
            print(f"LLM调用失败: {e}")
            result = {"error": str(e)}

        llm_time = time.time() - start
        print(f"LLM解析完成: {llm_time:.2f}s")
        
        return result, llm_time

    
    # 主流程
    def process(self, image_path, output_dir):
        total_start = time.time()
        base_name = os.path.splitext(os.path.basename(image_path))[0]

        # 1. OCR (必选)
        text, ocr_time = self.run_ocr(image_path, output_dir)

        # 2. LLM (可选)
        llm_time = 0.0
        result = None

        if self.client:
            # 模式A: OCR + LLM 
            result, llm_time = self.parse_waybill(text)
            
            # 保存 JSON
            json_path = os.path.join(output_dir, f"{base_name}.json")
            with open(json_path, "w", encoding="utf-8") as f:
                json.dump(result, f, ensure_ascii=False, indent=2)
            print("JSON保存:", json_path)
        
        else:
            # 模式B: 仅 OCR 
            print("未配置API Key,跳过LLM解析 (仅OCR模式)")
            
            # 如果用户关闭了中间文件,这里强制保存一份清洗后的txt
            # 如果开启了中间文件,用户已经有了 .mmd,这里就不重复保存了
            if not self.save_intermediate_files:
                txt_path = os.path.join(output_dir, f"{base_name}.txt")
                with open(txt_path, "w", encoding="utf-8") as f:
                    f.write(text)
                print("纯文本保存:", txt_path)
            
            # 封装返回结果
            result = {"raw_text": text}

        total_cost = time.time() - total_start

        print("\n============== 耗时统计 ==============")
        print(f"OCR耗时: {ocr_time:.2f}s")
        if self.client:
            print(f"LLM耗时: {llm_time:.2f}s")
        print(f"总耗时: {total_cost:.2f}s")
        print("=====================================")

        return result


if __name__ == "__main__":

    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    MODEL_PATH = r"model/deepseek"
    OUTPUT_DIR = r"outputdeepseek"
    IMAGE_PATH = r"images/jitu_4e844581.jpeg"

    start = time.time()

    # True: OCR + LLM模式 | False: 仅OCR模式
    ENABLE_LLM = False 

    # 调用llm的apikey是付费的
    MY_API_KEY = ""  if ENABLE_LLM else None # 硅基流动的apikey。如果为 None,则仅运行 OCR
    MY_API_BASE = "https://api.siliconflow.cn/v1" # 硅基流动调用api地址
    MY_LLM_MODEL = "deepseek-ai/DeepSeek-V3"      # 硅基流动的模型

    # 仅 OCR 模式 (不传 api_key)
    # OCR + LLM 模式 (传入 api_key)
    ocr_only = DeepSeekOCR(
        MODEL_PATH,
        device="cuda",
        api_key=MY_API_KEY,
        api_base=MY_API_BASE,
        llm_model=MY_LLM_MODEL,
        enable_resize=True,
        save_intermediate_files=False 
    )
    res1 = ocr_only.process(IMAGE_PATH, OUTPUT_DIR)
    print("返回结果:", res1)
    print(f"\n全部完成: {time.time()-start:.2f}s")

Logo

这里是“一人公司”的成长家园。我们提供从产品曝光、技术变现到法律财税的全栈内容,并连接云服务、办公空间等稀缺资源,助你专注创造,无忧运营。

更多推荐