Merge pull request #36 from myhloli/add_ppocrv6

Add ppocrv6
This commit is contained in:
Xiaomeng Zhao
2026-06-16 16:16:59 +08:00
committed by GitHub
29 changed files with 19159 additions and 41473 deletions
+1 -1
View File
@@ -15,7 +15,7 @@ Options:
-b, --backend [pipeline|vlm-engine|hybrid-engine|vlm-http-client|hybrid-http-client]
Parsing backend (default: hybrid-engine)
--effort [medium|high] Hybrid parsing effort (default: medium)
-l, --lang [ch|ch_server|ch_lite|en|korean|japan|chinese_cht|ta|te|ka|th|el|latin|arabic|east_slavic|cyrillic|devanagari]
-l, --lang [ch|ch_server|korean|ta|te|ka|th|el|arabic|east_slavic|cyrillic|devanagari]
Specify document language (improves OCR accuracy, pipeline backend only)
-u, --url TEXT OpenAI-compatible backend URL passed through to the server when using http-client
-s, --start INTEGER Starting page number for parsing (0-based)
+1 -1
View File
@@ -15,7 +15,7 @@ Options:
-b, --backend [pipeline|vlm-engine|hybrid-engine|vlm-http-client|hybrid-http-client]
解析后端(默认为 hybrid-engine
--effort [medium|high] Hybrid 解析强度(默认:medium
-l, --lang [ch|ch_server|ch_lite|en|korean|japan|chinese_cht|ta|te|ka|th|el|latin|arabic|east_slavic|cyrillic|devanagari]
-l, --lang [ch|ch_server|korean|ta|te|ka|th|el|arabic|east_slavic|cyrillic|devanagari]
指定文档语言(可提升 OCR 准确率,仅用于 pipeline 后端)
-u, --url TEXT 当使用 http-client 时,传给服务端后端的 OpenAI 兼容地址
-s, --start INTEGER 开始解析的页码(从 0 开始)
+1 -1
View File
@@ -74,7 +74,7 @@ def table_orientation_cls_model_init():
atom_model_name=AtomicModel.OCR,
det_db_box_thresh=0.5,
det_db_unclip_ratio=1.6,
lang="ch_lite",
lang="ch",
enable_merge_det_boxes=False
)
cls_model = MineruTableOrientationClsModel(ocr_engine)
+4 -2
View File
@@ -1,6 +1,5 @@
# Copyright (c) Opendatalab. All rights reserved.
import copy
from loguru import logger
from mineru.utils.enum_class import ContentType, BlockType, SplitFlag
from mineru.utils.language import detect_lang
@@ -239,7 +238,8 @@ def __is_list_or_index_block(block):
):
line[ListLineTag.IS_LIST_END_LINE] = True
line_start_flag = True
# 一种有缩进的特殊有序list,start line 左侧不贴边且以数字开头,end line 以 IS_LIST_END_FLAG 结尾且数量和start line 一致
# 一种有缩进的特殊有序 liststart line 左侧不贴边且以数字开头,
# end line 以 IS_LIST_END_FLAG 结尾且数量和 start line 一致。
elif num_start_count >= 2 and num_start_count == flag_end_count:
for i, line in enumerate(block['lines']):
if len(lines_text_list[i]) > 0:
@@ -331,6 +331,8 @@ def __merge_2_vertical_text_blocks(block1, block2):
and abs(block1_height - block2_height) < min_block_height
and not span_start_with_num
and not span_start_with_big_char
# 下一个纵排块的右边界要进入上一个纵排块左边界右侧
and block1['bbox'][2] > block2['bbox'][0]
):
if block1['page_num'] != block2['page_num']:
for line in block1['lines']:
+4 -3
View File
@@ -9,7 +9,8 @@ from mineru.utils.pdf_image_tools import get_crop_img
OCR_DET_PADDING = 50
def get_ch_lite_ocr_det_model():
def get_ch_ocr_det_model():
"""获取默认中文 OCR 检测模型,当前 ch 已对应轻量 PP-OCRv6 配置。"""
try:
from mineru.backend.pipeline.model_init import AtomModelSingleton
except Exception as e:
@@ -24,7 +25,7 @@ def get_ch_lite_ocr_det_model():
atom_model_name='ocr',
ocr_show_log=False,
det_db_box_thresh=0.3,
lang='ch_lite'
lang='ch'
)
@@ -44,7 +45,7 @@ def detect_ocr_boxes_from_padded_crop(bbox, page_pil_img, scale, ocr_model=None,
crop_img = cv2.cvtColor(crop_np_img, cv2.COLOR_RGB2BGR)
if ocr_model is None:
ocr_model = get_ch_lite_ocr_det_model()
ocr_model = get_ch_ocr_det_model()
ocr_det_res = ocr_model.ocr(crop_img, rec=False)[0]
return ocr_det_res or [], padding
+8 -1
View File
@@ -457,7 +457,14 @@ def _can_auto_merge_vertical_text_blocks(
min_metric_height = min(current_metric_height, previous_metric_height)
if min_metric_height <= 0:
return False
return abs(current_metric_height - previous_metric_height) < min_metric_height
if abs(current_metric_height - previous_metric_height) >= min_metric_height:
return False
return _has_mergeable_vertical_block_bbox_relation(current_block, previous_block)
def _has_mergeable_vertical_block_bbox_relation(current_block, previous_block):
"""复刻横排同构关系:当前纵排块右边界需要进入前块左边界右侧。"""
return current_block["bbox"][2] > previous_block["bbox"][0]
def _cleanup_block_internal_metadata(block):
+16 -20
View File
@@ -13,6 +13,11 @@ from mineru.cli.backend_options import (
validate_effort as validate_public_effort,
)
from mineru.cli.public_http_client_policy import validate_public_http_client_request
from mineru.utils.ocr_language import (
PUBLIC_OCR_LANGUAGE_SCHEMA_EXTRA,
format_public_ocr_lang_description,
validate_public_ocr_lang_list,
)
ALLOWED_PARSE_METHODS = {"auto", "txt", "ocr"}
SWAGGER_UI_FILE_ARRAY_SCHEMA_EXTRA = {
@@ -76,6 +81,14 @@ def validate_parse_effort(effort: str) -> str:
raise HTTPException(status_code=400, detail=str(exc)) from exc
def validate_parse_lang_list(lang_list: list[str]) -> list[str]:
"""校验公开 API 允许的 OCR 语言列表,避免旧语言入口进入解析链路。"""
try:
return validate_public_ocr_lang_list(lang_list)
except ValueError as exc:
raise HTTPException(status_code=400, detail=str(exc)) from exc
async def parse_request_form(
request: Request,
files: Annotated[
@@ -88,25 +101,8 @@ async def parse_request_form(
lang_list: Annotated[
list[str],
Form(
description="""(Adapted for pipeline backend only) Input the languages in the pdf to improve OCR accuracy. Options:
- ch: Chinese, English, Chinese Traditional.
- ch_lite: Chinese, English, Chinese Traditional, Japanese.
- ch_server: Chinese, English, Chinese Traditional, Japanese.
- en: English.
- korean: Korean, English.
- japan: Chinese, English, Chinese Traditional, Japanese.
- chinese_cht: Chinese, English, Chinese Traditional, Japanese.
- ta: Tamil, English.
- te: Telugu, English.
- ka: Kannada.
- th: Thai, English.
- el: Greek, English.
- latin: French, German, Afrikaans, Italian, Spanish, Bosnian, Portuguese, Czech, Welsh, Danish, Estonian, Irish, Croatian, Uzbek, Hungarian, Serbian (Latin), Indonesian, Occitan, Icelandic, Lithuanian, Maori, Malay, Dutch, Norwegian, Polish, Slovak, Slovenian, Albanian, Swedish, Swahili, Tagalog, Turkish, Latin, Azerbaijani, Kurdish, Latvian, Maltese, Pali, Romanian, Vietnamese, Finnish, Basque, Galician, Luxembourgish, Romansh, Catalan, Quechua.
- arabic: Arabic, Persian, Uyghur, Urdu, Pashto, Kurdish, Sindhi, Balochi, English.
- east_slavic: Russian, Belarusian, Ukrainian, English.
- cyrillic: Russian, Belarusian, Ukrainian, Serbian (Cyrillic), Bulgarian, Mongolian, Abkhazian, Adyghe, Kabardian, Avar, Dargin, Ingush, Chechen, Lak, Lezgin, Tabasaran, Kazakh, Kyrgyz, Tajik, Macedonian, Tatar, Chuvash, Bashkir, Malian, Moldovan, Udmurt, Komi, Ossetian, Buryat, Kalmyk, Tuvan, Sakha, Karakalpak, English.
- devanagari: Hindi, Marathi, Nepali, Bihari, Maithili, Angika, Bhojpuri, Magahi, Santali, Newari, Konkani, Sanskrit, Haryanvi, English.
""",
description=format_public_ocr_lang_description(),
json_schema_extra=PUBLIC_OCR_LANGUAGE_SCHEMA_EXTRA,
),
] = ["ch"],
backend: Annotated[
@@ -237,7 +233,7 @@ async def parse_request_form(
effective_return_original_file = return_original_file and response_format_zip
return ParseRequestOptions(
files=files,
lang_list=lang_list,
lang_list=validate_parse_lang_list(lang_list),
backend=backend,
effort=effort,
parse_method=validate_parse_method(parse_method),
+16 -21
View File
@@ -30,6 +30,7 @@ from mineru.utils.config_reader import (
get_max_concurrent_requests as read_max_concurrent_requests,
)
from mineru.utils.guess_suffix_or_lang import guess_suffix_by_path
from mineru.utils.ocr_language import PUBLIC_OCR_LANGUAGES, validate_public_ocr_lang
from mineru.utils.pdf_page_id import get_end_page_id
from mineru.utils.pdfium_guard import (
close_pdfium_document,
@@ -117,6 +118,18 @@ def normalize_effort_option(
raise click.BadParameter(str(exc), ctx=ctx, param=param) from exc
def normalize_ocr_lang_option(
ctx: click.Context,
param: click.Parameter,
value: str,
) -> str:
"""校验 CLI OCR 语言参数,并将兼容别名归一到实际模型语言。"""
try:
return validate_public_ocr_lang(value)
except ValueError as exc:
raise click.BadParameter(str(exc), ctx=ctx, param=param) from exc
@dataclass(frozen=True)
class TaskFailure:
task_index: int
@@ -1096,28 +1109,10 @@ async def run_orchestrated_cli(
"-l",
"--lang",
"lang",
type=click.Choice(
[
"ch",
"ch_server",
"ch_lite",
"en",
"korean",
"japan",
"chinese_cht",
"ta",
"te",
"ka",
"th",
"el",
"latin",
"arabic",
"east_slavic",
"cyrillic",
"devanagari",
]
),
type=str,
default="ch",
callback=normalize_ocr_lang_option,
metavar="[" + "|".join(PUBLIC_OCR_LANGUAGES) + "]",
help="""
Input the languages in the pdf (if known) to improve OCR accuracy.
Without languages specified, 'ch' will be used by default.
+3 -23
View File
@@ -49,6 +49,7 @@ from mineru.cli.client_side_output import regenerate_client_side_outputs
from mineru.cli.output_paths import resolve_parse_dir
from mineru.cli.vlm_preload import resolve_gradio_local_api_cli_args
from mineru.cli.visualization import VisualizationJob, run_visualization_job
from mineru.utils.ocr_language import PUBLIC_OCR_LANGUAGE_CHOICES
_gradio_local_api_server = _api_client.ReusableLocalAPIServer()
@@ -1339,28 +1340,7 @@ def render_header_html(i18n):
)
return rendered_header
other_lang = [
'ch (Chinese, English, Chinese Traditional)',
'ch_lite (Chinese, English, Chinese Traditional, Japanese)',
'ch_server (Chinese, English, Chinese Traditional, Japanese)',
'en (English)',
'korean (Korean, English)',
'japan (Chinese, English, Chinese Traditional, Japanese)',
'chinese_cht (Chinese, English, Chinese Traditional, Japanese)',
'ta (Tamil, English)',
'te (Telugu, English)',
'ka (Kannada)',
'el (Greek, English)',
'th (Thai, English)'
]
add_lang = [
'latin (French, German, Afrikaans, Italian, Spanish, Bosnian, Portuguese, Czech, Welsh, Danish, Estonian, Irish, Croatian, Uzbek, Hungarian, Serbian (Latin), Indonesian, Occitan, Icelandic, Lithuanian, Maori, Malay, Dutch, Norwegian, Polish, Slovak, Slovenian, Albanian, Swedish, Swahili, Tagalog, Turkish, Latin, Azerbaijani, Kurdish, Latvian, Maltese, Pali, Romanian, Vietnamese, Finnish, Basque, Galician, Luxembourgish, Romansh, Catalan, Quechua)',
'arabic (Arabic, Persian, Uyghur, Urdu, Pashto, Kurdish, Sindhi, Balochi, English)',
'east_slavic (Russian, Belarusian, Ukrainian, English)',
'cyrillic (Russian, Belarusian, Ukrainian, Serbian (Cyrillic), Bulgarian, Mongolian, Abkhazian, Adyghe, Kabardian, Avar, Dargin, Ingush, Chechen, Lak, Lezgin, Tabasaran, Kazakh, Kyrgyz, Tajik, Macedonian, Tatar, Chuvash, Bashkir, Malian, Moldovan, Udmurt, Komi, Ossetian, Buryat, Kalmyk, Tuvan, Sakha, Karakalpak, English)',
'devanagari (Hindi, Marathi, Nepali, Bihari, Maithili, Angika, Bhojpuri, Magahi, Santali, Newari, Konkani, Sanskrit, Haryanvi, English)'
]
all_lang = [*other_lang, *add_lang]
all_lang = list(PUBLIC_OCR_LANGUAGE_CHOICES)
def safe_stem(file_path):
@@ -1988,7 +1968,7 @@ def main(ctx,
language = gr.Dropdown(
all_lang,
label=i18n("ocr_language"),
value='ch (Chinese, English, Chinese Traditional)',
value=all_lang[0],
info=i18n("ocr_language_info"),
)
is_ocr = gr.Checkbox(
+54
View File
@@ -1094,6 +1094,56 @@ class PPDocLayoutV2LayoutModel:
box["label"] = label
box["cls_id"] = cls_id
@staticmethod
def _set_header_footer_label(box: Dict, label: str) -> None:
"""同步设置页眉/页脚相关标签及其类别编号。"""
label_cls_ids = {
"footer": 8,
"footer_image": 9,
"header": 12,
"header_image": 13,
}
if label not in label_cls_ids:
raise ValueError(f"Unsupported header/footer label: {label}")
box["label"] = label
box["cls_id"] = label_cls_ids[label]
@classmethod
def _reclassify_header_footer_by_page_half(
cls,
boxes: List[Dict],
image_size: Optional[Tuple[int, int]],
) -> List[Dict]:
"""按页面上下半区重新校正页眉/页脚锚点,避免跨半页误触发边界规则。"""
if image_size is None:
return boxes
page_height = float(image_size[0])
if page_height <= 0:
return boxes
page_middle = page_height * 0.5
upper_half_labels = {
"footer": "header",
"footer_image": "header_image",
}
lower_half_labels = {
"header": "footer",
"header_image": "footer_image",
}
for box in boxes:
bbox = box.get("bbox")
if not bbox or len(bbox) < 4:
continue
label = box.get("label")
y_mid = (float(bbox[1]) + float(bbox[3])) / 2
if y_mid < page_middle and label in upper_half_labels:
cls._set_header_footer_label(box, upper_half_labels[label])
elif y_mid >= page_middle and label in lower_half_labels:
cls._set_header_footer_label(box, lower_half_labels[label])
return boxes
@staticmethod
def _union_bbox(box1: Sequence[float], box2: Sequence[float]) -> List[int]:
x1_min, y1_min, x1_max, y1_max = [float(v) for v in box1]
@@ -1232,6 +1282,10 @@ class PPDocLayoutV2LayoutModel:
footer_labels = {"footer", "footer_image"}
exempt_labels = {"aside_text", "footnote", "number"}
ordered_boxes = sorted(boxes, key=lambda box: box["index"])
ordered_boxes = cls._reclassify_header_footer_by_page_half(
ordered_boxes,
image_size=image_size,
)
boundary_anchor_ids = {
id(box)
for box in ordered_boxes
+23 -133
View File
@@ -14,6 +14,7 @@ from mineru.model.ocr.seal_crop import CropByPolys, SortPolyBoxes
from mineru.utils.config_reader import get_device
from mineru.utils.enum_class import ModelPath
from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
from mineru.utils.ocr_language import normalize_ocr_model_lang
from mineru.utils.ocr_utils import (
check_img,
preprocess_image,
@@ -27,111 +28,6 @@ from mineru.model.utils.tools.infer import pytorchocr_utility as utility
import argparse
latin_lang = [
"af",
"az",
"bs",
"cs",
"cy",
"da",
"de",
"es",
"et",
"fr",
"ga",
"hr",
"hu",
"id",
"is",
"it",
"ku",
"la",
"lt",
"lv",
"mi",
"ms",
"mt",
"nl",
"no",
"oc",
"pi",
"pl",
"pt",
"ro",
"rs_latin",
"sk",
"sl",
"sq",
"sv",
"sw",
"tl",
"tr",
"uz",
"vi",
"french",
"german",
"fi",
"eu",
"gl",
"lb",
"rm",
"ca",
"qu",
]
arabic_lang = ["ar", "fa", "ug", "ur", "ps", "ku", "sd", "bal"]
cyrillic_lang = [
"ru",
"rs_cyrillic",
"be",
"bg",
"uk",
"mn",
"abq",
"ady",
"kbd",
"ava",
"dar",
"inh",
"che",
"lbe",
"lez",
"tab",
"kk",
"ky",
"tg",
"mk",
"tt",
"cv",
"ba",
"mhr",
"mo",
"udm",
"kv",
"os",
"bua",
"xal",
"tyv",
"sah",
"kaa",
]
east_slavic_lang = ["ru", "be", "uk"]
devanagari_lang = [
"hi",
"mr",
"ne",
"bh",
"mai",
"ang",
"bho",
"mah",
"sck",
"new",
"gom",
"sa",
"bgc",
]
def get_model_params(lang, config):
if lang in config['lang']:
params = config['lang'][lang]
@@ -156,44 +52,38 @@ class PytorchPaddleOCR(TextSystem):
parser = utility.init_args()
args = parser.parse_args(args)
self.lang = kwargs.get('lang', 'ch')
self.is_seal = self.lang in ['seal', 'seal_lite']
requested_lang = kwargs.get('lang', 'ch')
self.lang = requested_lang
self.is_seal = requested_lang in ['seal', 'seal_lite']
self.enable_merge_det_boxes = kwargs.get("enable_merge_det_boxes", True)
device = get_device()
if device == 'cpu':
if self.lang in ['ch', 'ch_server', 'japan', 'chinese_cht']:
# logger.warning("The current device in use is CPU. To ensure the speed of parsing, the language is automatically switched to ch_lite.")
self.lang = 'ch_lite'
elif self.lang in ['seal']:
self.lang = 'seal_lite'
if self.lang in latin_lang:
self.lang = 'latin'
elif self.lang in east_slavic_lang:
self.lang = 'east_slavic'
elif self.lang in arabic_lang:
self.lang = 'arabic'
elif self.lang in cyrillic_lang:
self.lang = 'cyrillic'
elif self.lang in devanagari_lang:
self.lang = 'devanagari'
else:
pass
models_config_path = os.path.join(root_dir, 'pytorchocr', 'utils', 'resources', 'models_config.yml')
with open(models_config_path) as file:
models_config_path = os.path.join(
root_dir, 'pytorchocr', 'utils', 'resources', 'models_config.yml'
)
with open(models_config_path, encoding='utf-8') as file:
config = yaml.safe_load(file)
self.lang = normalize_ocr_model_lang(
requested_lang,
device=device,
supported_langs=config['lang'],
)
det, rec, dict_file = get_model_params(self.lang, config)
ocr_models_dir = ModelPath.pytorch_paddle
det_model_path = f"{ocr_models_dir}/{det}"
det_model_path = os.path.join(auto_download_and_get_model_root_path(det_model_path), det_model_path)
det_model_path = os.path.join(
auto_download_and_get_model_root_path(det_model_path), det_model_path
)
rec_model_path = f"{ocr_models_dir}/{rec}"
rec_model_path = os.path.join(auto_download_and_get_model_root_path(rec_model_path), rec_model_path)
rec_model_path = os.path.join(
auto_download_and_get_model_root_path(rec_model_path), rec_model_path
)
kwargs['det_model_path'] = det_model_path
kwargs['rec_model_path'] = rec_model_path
kwargs['rec_char_dict_path'] = os.path.join(root_dir, 'pytorchocr', 'utils', 'resources', 'dict', dict_file)
kwargs['rec_char_dict_path'] = os.path.join(
root_dir, 'pytorchocr', 'utils', 'resources', 'dict', dict_file
)
kwargs['rec_batch_num'] = 6
if self.is_seal:
kwargs['det_limit_side_len'] = 736
@@ -292,7 +182,7 @@ class PytorchPaddleOCR(TextSystem):
tqdm_progress_bar=None,
):
assert isinstance(img, (np.ndarray, list, str, bytes))
if isinstance(img, list) and det == True:
if isinstance(img, list) and det:
logger.error('When input a list of images, det must be false')
exit(0)
img = check_img(img)
+41 -3
View File
@@ -1,6 +1,9 @@
# Copyright (c) Opendatalab. All rights reserved.
import os
from pathlib import Path
import torch
from .modeling.architectures.base_model import BaseModel
class BaseOCRV20:
@@ -13,13 +16,47 @@ class BaseOCRV20:
def build_net(self, **kwargs):
self.net = BaseModel(self.config, **kwargs)
@staticmethod
def _is_safetensors_path(weights_path):
"""判断权重文件是否为 safetensors 格式。"""
return Path(weights_path).suffix == ".safetensors"
@staticmethod
def _load_weight_file(weights_path):
"""根据文件后缀选择 safetensors 或 torch 原生加载方式。"""
if BaseOCRV20._is_safetensors_path(weights_path):
from safetensors.torch import load_file
return load_file(str(weights_path), device="cpu")
try:
return torch.load(weights_path, map_location="cpu", weights_only=True)
except TypeError:
return torch.load(weights_path, map_location="cpu")
@staticmethod
def _normalize_ppocrv6_state_dict(weights, weights_path):
"""归一化 HF OCR safetensors 的外层 `model.` 前缀。"""
if not BaseOCRV20._is_safetensors_path(weights_path):
return weights
if not any(key.startswith("model.") for key in weights.keys()):
return weights
return {
key.removeprefix("model."): value
for key, value in weights.items()
}
def read_pytorch_weights(self, weights_path):
"""读取 PyTorch OCR 权重,并兼容 PP-OCRv6 safetensors。"""
if not os.path.exists(weights_path):
raise FileNotFoundError('{} is not existed.'.format(weights_path))
weights = torch.load(weights_path)
return weights
weights = self._load_weight_file(weights_path)
return self._normalize_ppocrv6_state_dict(weights, weights_path)
def get_out_channels(self, weights):
"""从权重结构推断识别输出通道数。"""
if "head.head.weight" in weights:
# PP-OCRv6 safetensors 的识别分类层固定命名为 head.head。
return weights["head.head.weight"].shape[0]
if list(weights.keys())[-1].endswith('.weight') and len(list(weights.values())[-1].shape) == 2:
out_channels = list(weights.values())[-1].numpy().shape[1]
else:
@@ -31,7 +68,8 @@ class BaseOCRV20:
# print('weights is loaded.')
def load_pytorch_weights(self, weights_path):
self.net.load_state_dict(torch.load(weights_path, weights_only=True))
"""加载 PyTorch OCR 权重,按后缀兼容 safetensors。"""
self.net.load_state_dict(self.read_pytorch_weights(weights_path))
# print('model is loaded: {}'.format(weights_path))
def inference(self, inputs):
@@ -20,39 +20,34 @@ def build_backbone(config, model_type):
from .det_mobilenet_v3 import MobileNetV3
from .rec_hgnet import PPHGNet_small
from .rec_lcnetv3 import PPLCNetV3
from .rec_lcnetv4 import PPLCNetV4
from .rec_pphgnetv2 import PPHGNetV2_B4
support_dict = [
"MobileNetV3",
"ResNet",
"ResNet_vd",
"ResNet_SAST",
"PPLCNetV3",
"PPHGNet_small",
'PPHGNetV2_B4',
]
support_dict = {
"MobileNetV3": MobileNetV3,
"PPLCNetV3": PPLCNetV3,
"PPLCNetV4": PPLCNetV4,
"PPHGNet_small": PPHGNet_small,
"PPHGNetV2_B4": PPHGNetV2_B4,
}
elif model_type == "rec" or model_type == "cls":
from .rec_hgnet import PPHGNet_small
from .rec_lcnetv3 import PPLCNetV3
from .rec_lcnetv4 import PPLCNetV4
from .rec_mobilenet_v3 import MobileNetV3
from .rec_svtrnet import SVTRNet
from .rec_mv1_enhance import MobileNetV1Enhance
from .rec_pphgnetv2 import PPHGNetV2_B4, PPHGNetV2_B6_Formula
support_dict = [
"MobileNetV1Enhance",
"MobileNetV3",
"ResNet",
"ResNetFPN",
"MTB",
"ResNet31",
"SVTRNet",
"ViTSTR",
"DenseNet",
"PPLCNetV3",
"PPHGNet_small",
"PPHGNetV2_B4",
"PPHGNetV2_B6_Formula"
]
support_dict = {
"MobileNetV1Enhance": MobileNetV1Enhance,
"MobileNetV3": MobileNetV3,
"SVTRNet": SVTRNet,
"PPLCNetV3": PPLCNetV3,
"PPLCNetV4": PPLCNetV4,
"PPHGNet_small": PPHGNet_small,
"PPHGNetV2_B4": PPHGNetV2_B4,
"PPHGNetV2_B6_Formula": PPHGNetV2_B6_Formula,
}
else:
raise NotImplementedError
@@ -62,5 +57,5 @@ def build_backbone(config, model_type):
model_type, support_dict
)
)
module_class = eval(module_name)(**config)
module_class = support_dict[module_name](**config)
return module_class
@@ -0,0 +1,311 @@
# Copyright (c) Opendatalab. All rights reserved.
import torch
import torch.nn.functional as F
from torch import nn
NET_CONFIG_DET = {
"small": {
"stem_channels": [3, 24, 48],
"block_configs": [
[[3, 48, 48, 1, True], [3, 48, 48, 1, False]],
[[3, 48, 96, 2, False], [3, 96, 96, 1, True], [3, 96, 96, 1, False]],
[
[3, 96, 192, 2, False],
[3, 192, 192, 1, True],
[3, 192, 192, 1, False],
[3, 192, 192, 1, True],
[3, 192, 192, 1, False],
],
[[3, 192, 384, 2, False], [3, 384, 384, 1, True], [3, 384, 384, 1, False]],
],
},
}
NET_CONFIG_REC = {
"small": {
"stem_channels": [3, 48, 96],
"block_configs": [
[[3, 96, 96, 1, True]],
[[3, 96, 96, 1, False], [3, 96, 96, 1, False]],
[
[3, 96, 192, (2, 1), False],
[3, 192, 192, 1, True],
[3, 192, 192, 1, False],
[3, 192, 192, 1, True],
[3, 192, 192, 1, False],
[3, 192, 192, 1, True],
[3, 192, 192, 1, False],
],
[[3, 192, 384, (2, 1), False], [3, 384, 384, 1, True], [3, 384, 384, 1, False]],
],
},
"medium": {
"stem_channels": [3, 64, 128],
"block_configs": [
[[3, 128, 128, 1, True]],
[[3, 128, 256, 1, False], [3, 256, 256, 1, False], [3, 256, 256, 1, True]],
[
[3, 256, 512, (2, 1), False],
[3, 512, 512, 1, True],
[3, 512, 512, 1, False],
[3, 512, 512, 1, True],
[3, 512, 512, 1, False],
[3, 512, 512, 1, True],
[3, 512, 512, 1, False],
],
[[3, 512, 768, (2, 1), False], [3, 768, 768, 1, True], [3, 768, 768, 1, False]],
],
},
}
def _build_activation(name):
"""按 PP-OCRv6 配置创建无参数激活层,便于和 safetensors 权重命名保持解耦。"""
if name is None:
return nn.Identity()
if name == "relu":
return nn.ReLU()
if name == "gelu":
return nn.GELU()
if name in {"silu", "swish"}:
return nn.SiLU()
if name == "hardsigmoid":
return nn.Hardsigmoid()
raise ValueError(f"Unsupported activation: {name}")
def _to_stride(stride):
"""把 Paddle/Transformers 配置里的 stride 统一成 PyTorch 可接受的格式。"""
if isinstance(stride, list):
return tuple(stride)
return stride
class PPLCNetV4ConvLayer(nn.Module):
"""PP-LCNetV4 的 Conv-BN-Act 基础层,属性名对齐 HF 权重。"""
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
groups=1,
activation="relu",
):
"""初始化卷积、归一化和激活层。"""
super().__init__()
self.convolution = nn.Conv2d(
in_channels,
out_channels,
kernel_size=kernel_size,
stride=_to_stride(stride),
groups=groups,
padding=(kernel_size - 1) // 2,
bias=False,
)
self.normalization = nn.BatchNorm2d(out_channels)
self.activation = _build_activation(activation)
def forward(self, hidden_states):
"""执行 Conv-BN-Act 前向计算。"""
hidden_states = self.convolution(hidden_states)
hidden_states = self.normalization(hidden_states)
hidden_states = self.activation(hidden_states)
return hidden_states
class PPLCNetV4SqueezeExcitationModule(nn.Module):
"""PP-LCNetV4 的 SE 模块,保留 `convolutions.0/2` 权重命名。"""
def __init__(self, channel, reduction=4):
"""初始化全局池化和两层 1x1 卷积。"""
super().__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(1)
self.convolutions = nn.ModuleList(
[
nn.Conv2d(channel, channel // reduction, kernel_size=1, stride=1, padding=0, bias=True),
nn.ReLU(),
nn.Conv2d(channel // reduction, channel, kernel_size=1, stride=1, padding=0, bias=True),
nn.Hardsigmoid(),
]
)
def forward(self, hidden_states):
"""根据通道注意力缩放输入特征。"""
residual = hidden_states
hidden_states = self.avg_pool(hidden_states)
for layer in self.convolutions:
hidden_states = layer(hidden_states)
return residual * hidden_states
class PPLCNetV4LargeStem(nn.Module):
"""PP-LCNetV4 branch stem,属性名对齐 `encoder.convolution.stem*`。"""
def __init__(self, stem_channels):
"""初始化 v6 small/medium 使用的分支 stem。"""
super().__init__()
self.stem1 = PPLCNetV4ConvLayer(stem_channels[0], stem_channels[1], kernel_size=3, stride=2)
self.stem2a = PPLCNetV4ConvLayer(stem_channels[1], stem_channels[1] // 2, kernel_size=2, stride=1)
self.stem2b = PPLCNetV4ConvLayer(stem_channels[1] // 2, stem_channels[1], kernel_size=2, stride=1)
self.stem3 = PPLCNetV4ConvLayer(stem_channels[1] * 2, stem_channels[1], kernel_size=3, stride=2)
self.stem4 = PPLCNetV4ConvLayer(stem_channels[1], stem_channels[2], kernel_size=1, stride=1)
self.pool = nn.MaxPool2d(kernel_size=2, stride=1, ceil_mode=True)
def forward(self, pixel_values):
"""执行分支 stem 的 pad、pool 和 concat 流程。"""
embedding = self.stem1(pixel_values)
embedding = F.pad(embedding, (0, 1, 0, 1))
emb_stem_2a = self.stem2a(embedding)
emb_stem_2a = F.pad(emb_stem_2a, (0, 1, 0, 1))
emb_stem_2a = self.stem2b(emb_stem_2a)
pooled_emb = self.pool(embedding)
embedding = torch.cat([pooled_emb, emb_stem_2a], dim=1)
embedding = self.stem3(embedding)
embedding = self.stem4(embedding)
return embedding
class PPLCNetV4DepthwiseSeparableConvLayer(nn.Module):
"""PP-LCNetV4 block 中的 token mixer 和 channel mixer。"""
def __init__(
self,
in_channels,
out_channels,
stride,
kernel_size,
use_squeeze_excitation,
reduction=4,
):
"""按 v6 配置初始化深度卷积、SE 和两层 point-wise 卷积。"""
super().__init__()
self.has_residual = in_channels == out_channels and stride == 1
self.use_rep_dw = stride == 1 and in_channels == out_channels
if self.use_rep_dw:
self.token_conv = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=kernel_size // 2,
groups=in_channels,
bias=True,
)
else:
self.token_conv = PPLCNetV4ConvLayer(
in_channels=in_channels,
out_channels=in_channels,
kernel_size=kernel_size,
stride=stride,
groups=in_channels,
activation=None,
)
self.token_squeeze_excitation = (
PPLCNetV4SqueezeExcitationModule(in_channels, reduction) if use_squeeze_excitation else nn.Identity()
)
self.channel_conv1 = PPLCNetV4ConvLayer(
in_channels=in_channels,
out_channels=in_channels * 2,
kernel_size=1,
stride=1,
activation=None,
)
self.channel_act_fn = nn.GELU()
self.channel_conv2 = PPLCNetV4ConvLayer(
in_channels=in_channels * 2,
out_channels=out_channels,
kernel_size=1,
stride=1,
activation=None,
)
def forward(self, hidden_states):
"""执行 token mixing、channel mixing 和可选残差连接。"""
hidden_states = self.token_conv(hidden_states)
hidden_states = self.token_squeeze_excitation(hidden_states)
residual = hidden_states
hidden_states = self.channel_conv1(hidden_states)
hidden_states = self.channel_act_fn(hidden_states)
hidden_states = self.channel_conv2(hidden_states)
if self.has_residual:
hidden_states = residual + hidden_states
return hidden_states
class PPLCNetV4Block(nn.Module):
"""PP-LCNetV4 的一个 stage,内部包含多个 depthwise separable block。"""
def __init__(self, block_configs):
"""根据 stage 配置创建 block 列表。"""
super().__init__()
self.blocks = nn.ModuleList()
for kernel_size, in_channels, out_channels, stride, use_squeeze_excitation in block_configs:
self.blocks.append(
PPLCNetV4DepthwiseSeparableConvLayer(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=_to_stride(stride),
use_squeeze_excitation=use_squeeze_excitation,
)
)
def forward(self, hidden_states):
"""顺序执行当前 stage 的所有 block。"""
for block in self.blocks:
hidden_states = block(hidden_states)
return hidden_states
class PPLCNetV4Encoder(nn.Module):
"""PP-LCNetV4 编码器,保留 `encoder.convolution` 和 `encoder.blocks` 命名。"""
def __init__(self, stem_channels, block_configs):
"""初始化 stem 和四个 stage。"""
super().__init__()
self.convolution = PPLCNetV4LargeStem(stem_channels)
self.blocks = nn.ModuleList([PPLCNetV4Block(stage_configs) for stage_configs in block_configs])
def forward(self, pixel_values):
"""返回四个 stage 的输出特征,供 det/rec 上层按需使用。"""
hidden_states = self.convolution(pixel_values)
feature_maps = []
for block in self.blocks:
hidden_states = block(hidden_states)
feature_maps.append(hidden_states)
return feature_maps
class PPLCNetV4(nn.Module):
"""PP-OCRv6 使用的 PPLCNetV4 backbone,支持 det small 和 rec small/medium。"""
def __init__(self, det=False, model_size="small", in_channels=3, **kwargs):
"""按 det/rec 模式选择 v6 的固定网络配置。"""
super().__init__()
self.det = det
if in_channels != 3:
raise ValueError(f"PPLCNetV4 only supports 3 input channels, got {in_channels}.")
config_dict = NET_CONFIG_DET if det else NET_CONFIG_REC
if model_size not in config_dict:
mode = "det" if det else "rec"
raise ValueError(f"PPLCNetV4 {mode} model_size must be one of {list(config_dict)}, got {model_size}.")
config = config_dict[model_size]
self.encoder = PPLCNetV4Encoder(config["stem_channels"], config["block_configs"])
stage_out_channels = [stage[-1][2] for stage in config["block_configs"]]
self.out_channels = stage_out_channels if det else stage_out_channels[-1]
def forward(self, x):
"""det 返回四级特征列表,rec 返回高度池化后的识别特征。"""
feature_maps = self.encoder(x)
if self.det:
return feature_maps
x = feature_maps[-1]
if self.training:
return F.adaptive_avg_pool2d(x, [1, 40])
if x.shape[2] < 3:
raise ValueError(f"Feature height {x.shape[2]} < pool kernel 3.")
return F.avg_pool2d(x, [3, 2])
@@ -49,6 +49,49 @@ class Head(nn.Module):
return x
class PPOCRV6DBConvBatchnormLayer(nn.Module):
"""PP-OCRv6 DBHead 使用的 Conv-BN-Act 基础层,命名对齐 safetensors。"""
def __init__(
self,
in_channels,
out_channels,
kernel_size,
stride=1,
padding=1,
activation="relu",
bias=False,
convolution_transpose=False,
):
"""初始化普通卷积或反卷积、BN 和激活层。"""
super().__init__()
if convolution_transpose:
self.convolution = nn.ConvTranspose2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
)
else:
self.convolution = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias,
)
self.norm = nn.BatchNorm2d(out_channels)
self.act_fn = nn.ReLU() if activation == "relu" else nn.Identity()
def forward(self, hidden_states):
"""执行 DBHead v6 分支的卷积、BN 和激活。"""
hidden_states = self.convolution(hidden_states)
hidden_states = self.norm(hidden_states)
hidden_states = self.act_fn(hidden_states)
return hidden_states
class DBHead(nn.Module):
"""
Differentiable Binarization (DB) for text detection:
@@ -57,24 +100,51 @@ class DBHead(nn.Module):
params(dict): super parameters for build DB network
"""
def __init__(self, in_channels, k=50, **kwargs):
def __init__(self, in_channels, k=50, mode=None, kernel_list=None, fix_nan=False, **kwargs):
"""初始化 DBHeadv6 模式使用 safetensors 对齐的三层上采样 head。"""
super(DBHead, self).__init__()
self.k = k
binarize_name_list = [
'conv2d_56', 'batch_norm_47', 'conv2d_transpose_0', 'batch_norm_48',
'conv2d_transpose_1', 'binarize'
]
thresh_name_list = [
'conv2d_57', 'batch_norm_49', 'conv2d_transpose_2', 'batch_norm_50',
'conv2d_transpose_3', 'thresh'
]
self.binarize = Head(in_channels, **kwargs)# binarize_name_list)
self.thresh = Head(in_channels, **kwargs)#thresh_name_list)
self.mode = mode
self.fix_nan = fix_nan
if mode == "ppocrv6":
kernel_list = kernel_list or [3, 2, 2]
self.conv_down = PPOCRV6DBConvBatchnormLayer(
in_channels=in_channels,
out_channels=in_channels // 4,
kernel_size=kernel_list[0],
padding=int(kernel_list[0] // 2),
)
self.conv_up = PPOCRV6DBConvBatchnormLayer(
in_channels=in_channels // 4,
out_channels=in_channels // 4,
kernel_size=kernel_list[1],
stride=2,
convolution_transpose=True,
)
self.conv_final = nn.ConvTranspose2d(
in_channels=in_channels // 4,
out_channels=1,
kernel_size=kernel_list[2],
stride=2,
)
return
self.binarize = Head(in_channels, **kwargs)
self.thresh = Head(in_channels, **kwargs)
def step_function(self, x, y):
"""计算 DB 二值化近似阶跃函数。"""
return torch.reciprocal(1 + torch.exp(-self.k * (x - y)))
def forward(self, x):
"""推理时返回统一的 `maps` 字段,兼容现有 OCR-det 后处理。"""
if self.mode == "ppocrv6":
shrink_maps = self.conv_down(x)
shrink_maps = self.conv_up(shrink_maps)
shrink_maps = self.conv_final(shrink_maps)
shrink_maps = torch.sigmoid(shrink_maps)
if self.fix_nan:
shrink_maps = torch.nan_to_num(shrink_maps)
return {'maps': shrink_maps}
shrink_maps = self.binarize(x)
return {'maps': shrink_maps}
@@ -107,4 +177,4 @@ class PFHeadLocal(DBHead):
base_maps = shrink_maps
cbn_maps = self.cbn_layer(self.up_conv(f), shrink_maps, None)
cbn_maps = F.sigmoid(cbn_maps)
return {'maps': 0.5 * (base_maps + cbn_maps), 'cbn_maps': cbn_maps}
return {'maps': 0.5 * (base_maps + cbn_maps), 'cbn_maps': cbn_maps}
@@ -1,7 +1,8 @@
# Copyright (c) Opendatalab. All rights reserved.
import torch.nn.functional as F
from torch import nn
from ..necks.rnn import Im2Seq, SequenceEncoder
from ..necks.rnn import EncoderWithLightSVTR, Im2Seq, SequenceEncoder
from .rec_ctc_head import CTCHead
@@ -21,8 +22,10 @@ class FCTranspose(nn.Module):
class MultiHead(nn.Module):
def __init__(self, in_channels, out_channels_list, **kwargs):
"""初始化多头识别 Headv6 LightSVTR 分支使用 HF safetensors 命名。"""
super().__init__()
self.head_list = kwargs.pop("head_list")
self.use_light_svtr_head = False
self.gtc_head = "sar"
assert len(self.head_list) >= 2
@@ -38,22 +41,37 @@ class MultiHead(nn.Module):
self.encoder_reshape = Im2Seq(in_channels)
neck_args = self.head_list[idx][name]["Neck"]
encoder_type = neck_args.pop("name")
self.ctc_encoder = SequenceEncoder(
in_channels=in_channels, encoder_type=encoder_type, **neck_args
)
# ctc head
head_args = self.head_list[idx][name].get("Head", {})
if head_args is None:
head_args = {}
if encoder_type == "lightsvtr":
# v6 safetensors 中 CTC 分支直接命名为 head.encoder/head.head。
self.encoder = EncoderWithLightSVTR(in_channels=in_channels, **neck_args)
self.head = nn.Linear(self.encoder.out_channels, out_channels_list["CTCLabelDecode"], bias=True)
self.out_channels = out_channels_list["CTCLabelDecode"]
self.use_light_svtr_head = True
else:
self.ctc_encoder = SequenceEncoder(
in_channels=in_channels, encoder_type=encoder_type, **neck_args
)
# ctc head
head_args = self.head_list[idx][name].get("Head", {})
if head_args is None:
head_args = {}
self.ctc_head = CTCHead(
in_channels=self.ctc_encoder.out_channels,
out_channels=out_channels_list["CTCLabelDecode"],
**head_args,
)
self.ctc_head = CTCHead(
in_channels=self.ctc_encoder.out_channels,
out_channels=out_channels_list["CTCLabelDecode"],
**head_args,
)
else:
raise NotImplementedError(f"{name} is not supported in MultiHead yet")
def forward(self, x, data=None):
"""根据配置执行 v6 LightSVTR 或历史 CTC 分支。"""
if self.use_light_svtr_head:
ctc_encoder = self.encoder(x)
ctc_encoder = ctc_encoder.squeeze(dim=2).permute(0, 2, 1)
predicts = self.head(ctc_encoder)
if not self.training:
predicts = F.softmax(predicts, dim=2)
return predicts
ctc_encoder = self.ctc_encoder(x)
return self.ctc_head(ctc_encoder)
@@ -16,14 +16,20 @@ __all__ = ["build_neck"]
def build_neck(config):
from .db_fpn import DBFPN, LKPAN, RSEFPN
from .db_fpn import DBFPN, LKPAN, RSEFPN, RepLKFPN
from .rnn import SequenceEncoder
support_dict = ["DBFPN", "SequenceEncoder", "RSEFPN", "LKPAN"]
support_dict = {
"DBFPN": DBFPN,
"SequenceEncoder": SequenceEncoder,
"RSEFPN": RSEFPN,
"LKPAN": LKPAN,
"RepLKFPN": RepLKFPN,
}
module_name = config.pop("name")
assert module_name in support_dict, Exception(
"neck only support {}".format(support_dict)
)
module_class = eval(module_name)(**config)
module_class = support_dict[module_name](**config)
return module_class
@@ -25,7 +25,7 @@ class DSConv(nn.Module):
**kwargs
):
super(DSConv, self).__init__()
if groups == None:
if groups is None:
groups = in_channels
self.if_act = if_act
self.act = act
@@ -285,6 +285,136 @@ class RSEFPN(nn.Module):
return fuse
class RepLKFPNSqueezeExcitationModule(nn.Module):
"""PP-OCRv6 RepLKFPN 使用的轻量 SE 模块,命名对齐 safetensors。"""
def __init__(self, in_channels, reduction, activation="relu"):
"""初始化通道压缩与恢复卷积。"""
super().__init__()
self.avg_pool = nn.AdaptiveAvgPool2d(output_size=1)
self.conv1 = nn.Conv2d(in_channels, in_channels // reduction, kernel_size=1, stride=1, padding=0)
self.conv2 = nn.Conv2d(in_channels // reduction, in_channels, kernel_size=1, stride=1, padding=0)
if activation == "relu":
self.act_fn = nn.ReLU()
else:
raise ValueError(f"Unsupported RepLKFPN SE activation: {activation}")
def forward(self, hidden_states):
"""计算 SE 权重并缩放输入特征。"""
residual = hidden_states
hidden_states = self.avg_pool(hidden_states)
hidden_states = self.conv2(self.act_fn(self.conv1(hidden_states)))
hidden_states = torch.clamp(0.2 * hidden_states + 0.5, min=0.0, max=1.0)
return residual * hidden_states
class RepLKFPNDepthwiseSeparableConvLayer(nn.Module):
"""RepLKFPN 的大核深度卷积 + point-wise 压缩分支。"""
def __init__(self, in_channels, out_channels, kernel_size, reduction):
"""初始化 depthwise、pointwise 和 SE 子层。"""
super().__init__()
self.depthwise_convolution = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=1,
padding=kernel_size // 2,
groups=in_channels,
bias=True,
)
self.squeeze_excitation_module = RepLKFPNSqueezeExcitationModule(out_channels // 4, reduction)
self.pointwise_convolution = nn.Conv2d(
in_channels=out_channels,
out_channels=out_channels // 4,
kernel_size=1,
bias=False,
)
def forward(self, hidden_states):
"""执行大核 DW 卷积、PW 压缩和 SE 残差增强。"""
hidden_states = self.depthwise_convolution(hidden_states)
hidden_states = self.pointwise_convolution(hidden_states)
hidden_states = hidden_states + self.squeeze_excitation_module(hidden_states)
return hidden_states
class RepLKFPNResidualSqueezeExcitationLayer(nn.Module):
"""RepLKFPN 的输入投影层,属性名对齐 `insert_conv.*` 权重。"""
def __init__(self, in_channels, out_channels, kernel_size, reduction, shortcut=True):
"""初始化输入投影和 SE 分支。"""
super().__init__()
self.in_conv = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
padding=int(kernel_size // 2),
bias=False,
)
self.squeeze_excitation_block = RepLKFPNSqueezeExcitationModule(out_channels, reduction)
self.shortcut = shortcut
def forward(self, hidden_states):
"""执行 1x1 投影并按配置叠加 SE 输出。"""
hidden_states = self.in_conv(hidden_states)
if self.shortcut:
return hidden_states + self.squeeze_excitation_block(hidden_states)
return self.squeeze_excitation_block(hidden_states)
class RepLKFPN(nn.Module):
"""PP-OCRv6 small det 使用的 RepLKFPN neck。"""
def __init__(self, in_channels, out_channels, shortcut=True, dilated_kernel_size=7, reduction=4, **kwargs):
"""按四级 backbone 通道创建 v6 FPN 投影和大核融合层。"""
super().__init__()
self.out_channels = out_channels
self.interpolate_mode = kwargs.get("interpolate_mode", "nearest")
self.insert_conv = nn.ModuleList()
self.input_conv = nn.ModuleList()
for channels in in_channels:
self.insert_conv.append(
RepLKFPNResidualSqueezeExcitationLayer(
in_channels=channels,
out_channels=out_channels,
kernel_size=1,
reduction=reduction,
shortcut=shortcut,
)
)
self.input_conv.append(
RepLKFPNDepthwiseSeparableConvLayer(
in_channels=out_channels,
out_channels=out_channels,
kernel_size=dilated_kernel_size,
reduction=reduction,
)
)
def forward(self, feature_maps):
"""融合四级特征并返回 DBHead 需要的单张特征图。"""
fused = []
for conv, feature in zip(self.insert_conv, feature_maps):
fused.append(conv(feature))
for idx in range(2, -1, -1):
fused[idx] = fused[idx] + F.interpolate(
fused[idx + 1],
scale_factor=2,
mode=self.interpolate_mode,
)
features = [conv(feat) for conv, feat in zip(self.input_conv, fused)]
processed = []
for feat, scale in zip(features, [1, 2, 4, 8]):
if scale == 1:
processed.append(feat)
else:
processed.append(F.interpolate(feat, scale_factor=scale, mode=self.interpolate_mode))
return torch.cat(processed[::-1], dim=1)
class LKPAN(nn.Module):
def __init__(self, in_channels, out_channels, mode="large", **kwargs):
super(LKPAN, self).__init__()
@@ -200,6 +200,185 @@ class EncoderWithSVTR(nn.Module):
return z
class LightSVTRConvLayer(nn.Module):
"""PP-OCRv6 LightSVTR 使用的 Conv-BN-SiLU 基础层。"""
def __init__(
self,
in_channels,
out_channels,
kernel_size=(3, 3),
activation="silu",
groups=1,
):
"""初始化与 safetensors key 对齐的卷积、归一化和激活层。"""
super().__init__()
if isinstance(kernel_size, int):
kernel_size = (kernel_size, kernel_size)
self.convolution = nn.Conv2d(
in_channels,
out_channels,
kernel_size=kernel_size,
stride=1,
padding=(kernel_size[0] // 2, kernel_size[1] // 2),
bias=False,
groups=groups,
)
self.normalization = nn.BatchNorm2d(out_channels)
self.activation = nn.SiLU() if activation in {"silu", "swish"} else nn.Identity()
def forward(self, hidden_states):
"""执行卷积、BN 和激活。"""
hidden_states = self.convolution(hidden_states)
hidden_states = self.normalization(hidden_states)
hidden_states = self.activation(hidden_states)
return hidden_states
class LightSVTRAttention(nn.Module):
"""LightSVTR 的多头自注意力,属性名对齐 `self_attn.*` 权重。"""
def __init__(self, hidden_size, num_heads=8, qkv_bias=True, attention_dropout=0.1):
"""初始化 qkv 投影和输出投影。"""
super().__init__()
if hidden_size % num_heads != 0:
raise ValueError(f"hidden_size {hidden_size} must be divisible by num_heads {num_heads}.")
self.hidden_size = hidden_size
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.scale = self.head_dim**-0.5
self.attention_dropout = attention_dropout
self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=qkv_bias)
self.projection = nn.Linear(hidden_size, hidden_size)
def forward(self, hidden_states):
"""计算全局自注意力并返回投影后的序列特征。"""
batch_size, seq_len, embed_dim = hidden_states.shape
mixed_qkv = self.qkv(hidden_states)
mixed_qkv = mixed_qkv.reshape(batch_size, seq_len, 3, self.num_heads, embed_dim // self.num_heads)
mixed_qkv = mixed_qkv.permute(2, 0, 3, 1, 4)
query_states, key_states, value_states = mixed_qkv[0], mixed_qkv[1], mixed_qkv[2]
attn_weights = torch.matmul(query_states, key_states.transpose(-1, -2)) * self.scale
attn_weights = nn.functional.softmax(attn_weights, dim=-1)
attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, embed_dim).contiguous()
return self.projection(attn_output)
class LightSVTRMLP(nn.Module):
"""LightSVTR block 内的前馈网络,属性名对齐 `mlp.fc*` 权重。"""
def __init__(self, hidden_size, mlp_ratio=4.0, drop_rate=0.1):
"""初始化两层线性层、SiLU 激活和 dropout。"""
super().__init__()
self.fc1 = nn.Linear(hidden_size, int(hidden_size * mlp_ratio))
self.activation = nn.SiLU()
self.fc2 = nn.Linear(int(hidden_size * mlp_ratio), hidden_size)
self.drop = nn.Dropout(drop_rate)
def forward(self, hidden_states):
"""执行 MLP 前向计算。"""
hidden_states = self.fc1(hidden_states)
hidden_states = self.activation(hidden_states)
hidden_states = self.drop(hidden_states)
hidden_states = self.fc2(hidden_states)
hidden_states = self.drop(hidden_states)
return hidden_states
class LightSVTRBlock(nn.Module):
"""LightSVTR 的 Transformer block,属性名对齐 v6 safetensors。"""
def __init__(
self,
hidden_size,
num_heads=8,
qkv_bias=True,
mlp_ratio=4.0,
drop_rate=0.1,
attn_drop_rate=0.1,
layer_norm_eps=1e-6,
):
"""初始化注意力、MLP 和两层 LayerNorm。"""
super().__init__()
self.self_attn = LightSVTRAttention(hidden_size, num_heads, qkv_bias, attn_drop_rate)
self.layer_norm1 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
self.mlp = LightSVTRMLP(hidden_size, mlp_ratio, drop_rate)
self.layer_norm2 = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
def forward(self, hidden_states):
"""执行 pre-norm attention 和 pre-norm MLP 残差结构。"""
residual = hidden_states
hidden_states = self.layer_norm1(hidden_states)
hidden_states = residual + self.self_attn(hidden_states)
residual = hidden_states
hidden_states = self.layer_norm2(hidden_states)
hidden_states = residual + self.mlp(hidden_states)
return hidden_states
class EncoderWithLightSVTR(nn.Module):
"""PP-OCRv6 使用的 LightSVTR neck,输出仍保持 4D 特征。"""
def __init__(
self,
in_channels,
dims=64,
depth=1,
num_heads=8,
qkv_bias=True,
mlp_ratio=4.0,
drop_rate=0.1,
attn_drop_rate=0.1,
drop_path=0.0,
qk_scale=None,
local_kernel=7,
use_guide=False,
**kwargs,
):
"""初始化 skip/reduce/local conv、LightSVTR block 和归一化层。"""
super().__init__()
self.use_guide = use_guide
self.conv_block = nn.ModuleList(
[
LightSVTRConvLayer(in_channels, dims, kernel_size=(1, 1), activation="silu"),
LightSVTRConvLayer(in_channels, dims, kernel_size=(1, 1), activation="silu"),
LightSVTRConvLayer(dims, dims, kernel_size=(1, local_kernel), activation="silu", groups=dims),
]
)
self.svtr_block = nn.ModuleList(
[
LightSVTRBlock(
hidden_size=dims,
num_heads=num_heads,
qkv_bias=qkv_bias,
mlp_ratio=mlp_ratio,
drop_rate=drop_rate,
attn_drop_rate=attn_drop_rate,
)
for _ in range(depth)
]
)
self.norm = nn.LayerNorm(dims, eps=1e-6)
self.out_channels = dims
def forward(self, x):
"""执行轻量局部卷积增强、全局注意力和 skip 残差融合。"""
if self.use_guide:
x = x.detach()
residual = self.conv_block[0](x)
hidden_states = self.conv_block[1](x)
hidden_states = hidden_states + self.conv_block[2](hidden_states)
batch_size, channels, height, width = hidden_states.shape
hidden_states = hidden_states.flatten(2).permute(0, 2, 1)
for block in self.svtr_block:
hidden_states = block(hidden_states)
hidden_states = self.norm(hidden_states)
hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
return hidden_states + residual
class SequenceEncoder(nn.Module):
def __init__(self, in_channels, encoder_type, hidden_size=48, **kwargs):
super(SequenceEncoder, self).__init__()
@@ -214,12 +393,13 @@ class SequenceEncoder(nn.Module):
"fc": EncoderWithFC,
"rnn": EncoderWithRNN,
"svtr": EncoderWithSVTR,
"lightsvtr": EncoderWithLightSVTR,
}
assert encoder_type in support_encoder_dict, "{} must in {}".format(
encoder_type, support_encoder_dict.keys()
)
if encoder_type == "svtr":
if encoder_type in ("svtr", "lightsvtr"):
self.encoder = support_encoder_dict[encoder_type](
self.encoder_reshape.out_channels, **kwargs
)
@@ -231,7 +411,7 @@ class SequenceEncoder(nn.Module):
self.only_reshape = False
def forward(self, x):
if self.encoder_type != "svtr":
if self.encoder_type not in ("svtr", "lightsvtr"):
x = self.encoder_reshape(x)
if not self.only_reshape:
x = self.encoder(x)
@@ -92,6 +92,27 @@ ch_PP-OCRv5_det_server_infer:
k: 50
mode: "large"
ch_PP-OCRv6_small_det_infer:
model_type: det
algorithm: DB
Transform: null
Backbone:
name: PPLCNetV4
det: True
model_size: small
Neck:
name: RepLKFPN
out_channels: 96
dilated_kernel_size: 7
shortcut: True
Head:
name: DBHead
k: 50
mode: ppocrv6
fix_nan: True
aux_in_channels: 96
kernel_list: [3, 2, 2]
ch_PP-OCRv4_det_server_infer:
model_type: det
algorithm: DB
@@ -236,6 +257,57 @@ ch_PP-OCRv5_rec_infer:
nrtr_dim: 384
max_text_length: 25
ch_PP-OCRv6_small_rec_infer:
model_type: rec
algorithm: SVTR_LCNet
Transform:
Backbone:
name: PPLCNetV4
model_size: small
Head:
name: MultiHead
out_channels_list:
CTCLabelDecode: 18710
head_list:
- CTCHead:
Neck:
name: lightsvtr
dims: 120
depth: 2
mlp_ratio: 2.0
local_kernel: 7
Head:
fc_decay: 0.00001
- NRTRHead:
nrtr_dim: 384
max_text_length: 25
ch_PP-OCRv6_medium_rec_infer:
model_type: rec
algorithm: SVTR_LCNet
Transform:
Backbone:
name: PPLCNetV4
model_size: medium
Head:
name: MultiHead
out_channels_list:
CTCLabelDecode: 18710
head_list:
- CTCHead:
Neck:
name: lightsvtr
dims: 192
depth: 2
mlp_ratio: 4.0
local_kernel: 7
use_guide: False
Head:
fc_decay: 0.00001
- NRTRHead:
nrtr_dim: 512
max_text_length: 25
ka_PP-OCRv3_rec_infer:
model_type: rec
algorithm: SVTR
@@ -541,4 +613,4 @@ te_PP-OCRv5_rec_infer:
fc_decay: 0.00001
- NRTRHead:
nrtr_dim: 384
max_text_length: 25
max_text_length: 25
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -1,436 +0,0 @@
0
1
2
3
4
5
6
7
8
9
A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z
a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z
!
"
#
$
%
&
'
(
)
*
+
,
-
.
/
:
;
<
=
>
?
@
[
\
]
_
`
{
|
}
^
~
©
®
¤
¦
§
¨
ª
«
¬
¯
°
²
³
´
µ
¸
¹
º
»
¼
½
¾
¿
×
¢
£
¥
ƒ
À
Á
Â
Ã
Ä
Å
Æ
Ç
È
É
Ê
Ë
Ì
Í
Î
Ï
Ð
Ñ
Ò
Ó
Ô
Õ
Ö
Ø
Ù
Ú
Û
Ü
Ý
Þ
à
á
â
ã
ä
å
æ
ç
è
é
ê
ë
ì
í
î
ï
ð
ñ
ò
ó
ô
õ
ö
ø
ù
ú
û
ü
ý
þ
ÿ
Α
Β
Γ
Δ
Ε
Ζ
Η
Θ
Ι
Κ
Λ
Μ
Ν
Ξ
Ο
Π
Ρ
Σ
Τ
Υ
Φ
Χ
Ψ
Ω
α
β
γ
δ
ε
ζ
η
θ
ι
κ
λ
μ
ν
ξ
ο
π
ρ
σ
ς
τ
υ
φ
χ
ψ
ω
𝑢
𝜓
÷
·
±
@@ -1,836 +0,0 @@
0
1
2
3
4
5
6
7
8
9
A
B
C
D
E
F
G
H
I
J
K
L
M
N
O
P
Q
R
S
T
U
V
W
X
Y
Z
a
b
c
d
e
f
g
h
i
j
k
l
m
n
o
p
q
r
s
t
u
v
w
x
y
z
À
Á
Â
Ã
Ä
Å
Æ
Ç
È
É
Ê
Ë
Ì
Í
Î
Ï
Ð
Ñ
Ò
Ó
Ô
Õ
Ö
×
Ø
Ù
Ú
Û
Ü
Ý
Þ
ß
à
á
â
ã
ä
å
æ
ç
è
é
ê
ë
ì
í
î
ï
ð
ñ
ò
ó
ô
õ
ö
÷
ø
ù
ú
û
ü
ý
þ
ÿ
Ā
ā
Ă
ă
Ą
ą
Ć
ć
Ĉ
ĉ
Ċ
ċ
Č
č
Ď
ď
Đ
đ
Ē
ē
Ĕ
ĕ
Ė
ė
Ę
ę
Ě
ě
Ĝ
ĝ
Ğ
ğ
Ġ
ġ
Ģ
ģ
Ĥ
ĥ
Ħ
ħ
Ĩ
ĩ
Ī
ī
Ĭ
ĭ
Į
į
İ
ı
IJ
ij
Ĵ
ĵ
Ķ
ķ
ĸ
Ĺ
ĺ
Ļ
ļ
Ľ
ľ
Ŀ
ŀ
Ł
ł
Ń
ń
Ņ
ņ
Ň
ň
ʼn
Ŋ
ŋ
Ō
ō
Ŏ
ŏ
Ő
ő
Œ
œ
Ŕ
ŕ
Ŗ
ŗ
Ř
ř
Ś
ś
Ŝ
ŝ
Ş
ş
Š
š
Ţ
ţ
Ť
ť
Ŧ
ŧ
Ũ
ũ
Ū
ū
Ŭ
ŭ
Ů
ů
Ű
ű
Ų
ų
Ŵ
ŵ
Ŷ
ŷ
Ÿ
Ź
ź
Ż
ż
Ž
ž
ſ
ƀ
Ɓ
Ƃ
ƃ
Ƅ
ƅ
Ɔ
Ƈ
ƈ
Ɖ
Ɗ
Ƌ
ƌ
ƍ
Ǝ
Ə
Ɛ
Ƒ
ƒ
Ɠ
Ɣ
ƕ
Ɩ
Ɨ
Ƙ
ƙ
ƚ
ƛ
Ɯ
Ɲ
ƞ
Ɵ
Ơ
ơ
Ƣ
ƣ
Ƥ
ƥ
Ʀ
Ƨ
ƨ
Ʃ
ƪ
ƫ
Ƭ
ƭ
Ʈ
Ư
ư
Ʊ
Ʋ
Ƴ
ƴ
Ƶ
ƶ
Ʒ
Ƹ
ƹ
ƺ
ƻ
Ƽ
ƽ
ƾ
ƿ
ǀ
ǁ
ǂ
ǃ
DŽ
Dž
dž
LJ
Lj
lj
NJ
Nj
nj
Ǎ
ǎ
Ǐ
ǐ
Ǒ
ǒ
Ǔ
ǔ
Ǖ
ǖ
Ǘ
ǘ
Ǚ
ǚ
Ǜ
ǜ
ǝ
Ǟ
ǟ
Ǡ
ǡ
Ǣ
ǣ
Ǥ
ǥ
Ǧ
ǧ
Ǩ
ǩ
Ǫ
ǫ
Ǭ
ǭ
Ǯ
ǯ
ǰ
DZ
Dz
dz
Ǵ
ǵ
Ƕ
Ƿ
Ǹ
ǹ
Ǻ
ǻ
Ǽ
ǽ
Ǿ
ǿ
Ȁ
ȁ
Ȃ
ȃ
Ȅ
ȅ
Ȇ
ȇ
Ȉ
ȉ
Ȋ
ȋ
Ȍ
ȍ
Ȏ
ȏ
Ȑ
ȑ
Ȓ
ȓ
Ȕ
ȕ
Ȗ
ȗ
Ș
ș
Ț
ț
Ȝ
ȝ
Ȟ
ȟ
Ƞ
ȡ
Ȣ
ȣ
Ȥ
ȥ
Ȧ
ȧ
Ȩ
ȩ
Ȫ
ȫ
Ȭ
ȭ
Ȯ
ȯ
Ȱ
ȱ
Ȳ
ȳ
ȴ
ȵ
ȶ
ȷ
ȸ
ȹ
Ⱥ
Ȼ
ȼ
Ƚ
Ⱦ
ȿ
ɀ
Ɂ
ɂ
Ƀ
Ʉ
Ʌ
Ɇ
ɇ
Ɉ
ɉ
Ɋ
ɋ
Ɍ
ɍ
Ɏ
ɏ
!
"
#
$
%
&
'
(
)
*
+
,
-
.
/
:
;
<
=
>
?
@
[
\
]
_
`
{
|
}
^
~
©
®
¤
¦
§
¨
ª
«
¬
¯
°
²
³
´
µ
¸
¹
º
»
¼
½
¾
¿
×
¢
£
¥
ƒ
À
Á
Â
Ã
Ä
Å
Æ
Ç
È
É
Ê
Ë
Ì
Í
Î
Ï
Ð
Ñ
Ò
Ó
Ô
Õ
Ö
Ø
Ù
Ú
Û
Ü
Ý
Þ
à
á
â
ã
ä
å
æ
ç
è
é
ê
ë
ì
í
î
ï
ð
ñ
ò
ó
ô
õ
ö
ø
ù
ú
û
ü
ý
þ
ÿ
Α
Β
Γ
Δ
Ε
Ζ
Η
Θ
Ι
Κ
Λ
Μ
Ν
Ξ
Ο
Π
Ρ
Σ
Τ
Υ
Φ
Χ
Ψ
Ω
α
β
γ
δ
ε
ζ
η
θ
ι
κ
λ
μ
ν
ξ
ο
π
ρ
σ
ς
τ
υ
φ
χ
ψ
ω
𝑢
𝜓
÷
·
±
@@ -1,77 +1,57 @@
lang:
ch_lite:
det: ch_PP-OCRv5_det_infer.pth
rec: ch_PP-OCRv5_rec_infer.pth
dict: ppocrv5_dict.txt
ch_server:
det: ch_PP-OCRv5_det_infer.pth
rec: ch_PP-OCRv5_rec_server_infer.pth
dict: ppocrv5_dict.txt
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: ch_PP-OCRv6_medium_rec_infer.safetensors
dict: ppocrv6_dict.txt
ch:
det: ch_PP-OCRv5_det_infer.pth
rec: ch_PP-OCRv4_rec_server_doc_infer.pth
dict: ppocrv4_doc_dict.txt
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: ch_PP-OCRv6_small_rec_infer.safetensors
dict: ppocrv6_dict.txt
korean:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: korean_PP-OCRv5_rec_infer.pth
dict: ppocrv5_korean_dict.txt
japan:
det: ch_PP-OCRv5_det_infer.pth
rec: ch_PP-OCRv5_rec_server_infer.pth
dict: ppocrv5_dict.txt
chinese_cht:
det: ch_PP-OCRv5_det_infer.pth
rec: ch_PP-OCRv5_rec_server_infer.pth
dict: ppocrv5_dict.txt
ta:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: ta_PP-OCRv5_rec_infer.pth
dict: ppocrv5_ta_dict.txt
te:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: te_PP-OCRv5_rec_infer.pth
dict: ppocrv5_te_dict.txt
ka:
det: Multilingual_PP-OCRv3_det_infer.pth
rec: ka_PP-OCRv3_rec_infer.pth
dict: ka_dict.txt
latin:
det: ch_PP-OCRv5_det_infer.pth
rec: latin_PP-OCRv5_rec_infer.pth
dict: ppocrv5_latin_dict.txt
arabic:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: arabic_PP-OCRv5_rec_infer.pth
dict: ppocrv5_arabic_dict.txt
cyrillic:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: cyrillic_PP-OCRv5_rec_infer.pth
dict: ppocrv5_cyrillic_dict.txt
devanagari:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: devanagari_PP-OCRv5_rec_infer.pth
dict: ppocrv5_devanagari_dict.txt
east_slavic:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: eslav_PP-OCRv5_rec_infer.pth
dict: ppocrv5_eslav_dict.txt
el:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: el_PP-OCRv5_rec_infer.pth
dict: ppocrv5_el_dict.txt
th:
det: ch_PP-OCRv5_det_infer.pth
det: ch_PP-OCRv6_small_det_infer.safetensors
rec: th_PP-OCRv5_rec_infer.pth
dict: ppocrv5_th_dict.txt
en:
det: ch_PP-OCRv5_det_infer.pth
rec: en_PP-OCRv5_rec_infer.pth
dict: ppocrv5_en_dict.txt
seal:
det: seal_PP-OCRv4_det_server_infer.pth
rec: ch_PP-OCRv4_rec_server_infer.pth
dict: ppocr_keys_v1.txt
rec: ch_PP-OCRv6_medium_rec_infer.safetensors
dict: ppocrv6_dict.txt
seal_lite:
det: seal_PP-OCRv4_det_infer.pth
rec: ch_PP-OCRv4_rec_infer.pth
dict: ppocr_keys_v1.txt
rec: ch_PP-OCRv6_small_rec_infer.safetensors
dict: ppocrv6_dict.txt
@@ -9,6 +9,9 @@ import argparse
root_dir = Path(__file__).resolve().parent.parent.parent
DEFAULT_CFG_PATH = root_dir / "pytorchocr" / "utils" / "resources" / "arch_config.yaml"
DEFAULT_REC_CHAR_DICT_PATH = (
root_dir / "pytorchocr" / "utils" / "resources" / "dict" / "ppocrv6_dict.txt"
)
def init_args():
@@ -88,8 +91,7 @@ def init_args():
parser.add_argument(
"--rec_char_dict_path",
type=str,
default=os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
'pytorchocr/utils/ppocr_keys_v1.txt'))
default=str(DEFAULT_REC_CHAR_DICT_PATH))
# params for text classifier
parser.add_argument("--use_angle_cls", type=str2bool, default=False)
+157
View File
@@ -0,0 +1,157 @@
# Copyright (c) Opendatalab. All rights reserved.
PUBLIC_OCR_LANGUAGES = (
"ch",
"ch_server",
"korean",
"ta",
"te",
"ka",
"th",
"el",
"arabic",
"east_slavic",
"cyrillic",
"devanagari",
)
_PUBLIC_OCR_LANGUAGE_DESCRIPTIONS = {
"ch": "Chinese, English, Japanese, Chinese Traditional, Latin",
"ch_server": "Chinese, English, Japanese, Chinese Traditional, Latin",
"korean": "Korean, English",
"ta": "Tamil, English",
"te": "Telugu, English",
"ka": "Kannada",
"th": "Thai, English",
"el": "Greek, English",
"arabic": (
"Arabic, Persian, Uyghur, Urdu, Pashto, Kurdish, Sindhi, Balochi, English"
),
"east_slavic": "Russian, Belarusian, Ukrainian, English",
"cyrillic": (
"Russian, Belarusian, Ukrainian, Serbian (Cyrillic), Bulgarian, Mongolian, "
"Abkhazian, Adyghe, Kabardian, Avar, Dargin, Ingush, Chechen, Lak, Lezgin, "
"Tabasaran, Kazakh, Kyrgyz, Tajik, Macedonian, Tatar, Chuvash, Bashkir, "
"Malian, Moldovan, Udmurt, Komi, Ossetian, Buryat, Kalmyk, Tuvan, Sakha, "
"Karakalpak, English"
),
"devanagari": (
"Hindi, Marathi, Nepali, Bihari, Maithili, Angika, Bhojpuri, Magahi, "
"Santali, Newari, Konkani, Sanskrit, Haryanvi, English"
),
}
PUBLIC_OCR_LANGUAGE_CHOICES = tuple(
f"{lang} ({_PUBLIC_OCR_LANGUAGE_DESCRIPTIONS[lang]})"
for lang in PUBLIC_OCR_LANGUAGES
)
PUBLIC_OCR_LANGUAGE_SCHEMA_EXTRA = {"items": {"enum": list(PUBLIC_OCR_LANGUAGES)}}
_ARABIC_LANG_ALIASES = {"ar", "fa", "ug", "ur", "ps", "ku", "sd", "bal"}
_CH_LANG_ALIASES = {"en", "japan", "chinese_cht", "latin"}
_EAST_SLAVIC_LANG_ALIASES = {"ru", "be", "uk"}
_CYRILLIC_LANG_ALIASES = {
"rs_cyrillic",
"bg",
"mn",
"abq",
"ady",
"kbd",
"ava",
"dar",
"inh",
"che",
"lbe",
"lez",
"tab",
"kk",
"ky",
"tg",
"mk",
"tt",
"cv",
"ba",
"mhr",
"mo",
"udm",
"kv",
"os",
"bua",
"xal",
"tyv",
"sah",
"kaa",
}
_DEVANAGARI_LANG_ALIASES = {
"hi",
"mr",
"ne",
"bh",
"mai",
"ang",
"bho",
"mah",
"sck",
"new",
"gom",
"sa",
"bgc",
}
def format_public_ocr_lang_description() -> str:
"""生成公开 API 使用的 OCR 语言说明,避免入口文案各自维护。"""
option_lines = [
f"- {lang}: {_PUBLIC_OCR_LANGUAGE_DESCRIPTIONS[lang]}."
for lang in PUBLIC_OCR_LANGUAGES
]
return (
"(Adapted for pipeline backend only) Input the languages in the pdf "
"to improve OCR accuracy. Options:\n"
+ "\n".join(option_lines)
)
def validate_public_ocr_lang(lang: str) -> str:
"""校验公开入口允许的 OCR 语言,并将兼容入口规范到实际模型 key。"""
if lang in _CH_LANG_ALIASES:
return "ch"
if lang not in PUBLIC_OCR_LANGUAGES:
raise ValueError(
f"Language {lang} not supported. Allowed values: "
+ ", ".join(PUBLIC_OCR_LANGUAGES)
)
return lang
def validate_public_ocr_lang_list(lang_list: list[str]) -> list[str]:
"""校验公开 API 的语言列表,返回可安全传入下游的副本。"""
effective_lang_list = lang_list or ["ch"]
return [validate_public_ocr_lang(lang) for lang in effective_lang_list]
def normalize_ocr_model_lang(
lang: str | None,
*,
device: str | None = None,
supported_langs=None,
) -> str:
"""将 OCR 语言参数归一为模型配置 key,保留内部 seal 与语系短码能力。"""
normalized_lang = lang or "ch"
if device == "cpu" and normalized_lang == "seal":
normalized_lang = "seal_lite"
elif normalized_lang in _CH_LANG_ALIASES:
normalized_lang = "ch"
elif normalized_lang in _EAST_SLAVIC_LANG_ALIASES:
normalized_lang = "east_slavic"
elif normalized_lang in _ARABIC_LANG_ALIASES:
normalized_lang = "arabic"
elif normalized_lang in _CYRILLIC_LANG_ALIASES:
normalized_lang = "cyrillic"
elif normalized_lang in _DEVANAGARI_LANG_ALIASES:
normalized_lang = "devanagari"
if supported_langs is not None and normalized_lang not in supported_langs:
raise ValueError(f"Language {lang} not supported")
return normalized_lang
+1
View File
@@ -98,6 +98,7 @@ pipeline = [
"torch>=2.6.0,<3",
"torchvision",
"transformers>=4.57.3,<5.0.0",
"safetensors>=0.4.0,<1",
"onnxruntime>1.17.0",
]
gradio = [