本地运行DeepSeek-OCR-2 识别图片文字
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")
更多推荐



所有评论(0)