mirror of
https://github.com/opendatalab/MinerU.git
synced 2026-09-21 12:42:22 +08:00
@@ -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)
|
||||
|
||||
@@ -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 开始)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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 一致
|
||||
# 一种有缩进的特殊有序 list,start 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']:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
@@ -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.
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
"""初始化 DBHead;v6 模式使用 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):
|
||||
"""初始化多头识别 Head,v6 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
|
||||
Ǵ
|
||||
ǵ
|
||||
Ƕ
|
||||
Ƿ
|
||||
Ǹ
|
||||
ǹ
|
||||
Ǻ
|
||||
ǻ
|
||||
Ǽ
|
||||
ǽ
|
||||
Ǿ
|
||||
ǿ
|
||||
Ȁ
|
||||
ȁ
|
||||
Ȃ
|
||||
ȃ
|
||||
Ȅ
|
||||
ȅ
|
||||
Ȇ
|
||||
ȇ
|
||||
Ȉ
|
||||
ȉ
|
||||
Ȋ
|
||||
ȋ
|
||||
Ȍ
|
||||
ȍ
|
||||
Ȏ
|
||||
ȏ
|
||||
Ȑ
|
||||
ȑ
|
||||
Ȓ
|
||||
ȓ
|
||||
Ȕ
|
||||
ȕ
|
||||
Ȗ
|
||||
ȗ
|
||||
Ș
|
||||
ș
|
||||
Ț
|
||||
ț
|
||||
Ȝ
|
||||
ȝ
|
||||
Ȟ
|
||||
ȟ
|
||||
Ƞ
|
||||
ȡ
|
||||
Ȣ
|
||||
ȣ
|
||||
Ȥ
|
||||
ȥ
|
||||
Ȧ
|
||||
ȧ
|
||||
Ȩ
|
||||
ȩ
|
||||
Ȫ
|
||||
ȫ
|
||||
Ȭ
|
||||
ȭ
|
||||
Ȯ
|
||||
ȯ
|
||||
Ȱ
|
||||
ȱ
|
||||
Ȳ
|
||||
ȳ
|
||||
ȴ
|
||||
ȵ
|
||||
ȶ
|
||||
ȷ
|
||||
ȸ
|
||||
ȹ
|
||||
Ⱥ
|
||||
Ȼ
|
||||
ȼ
|
||||
Ƚ
|
||||
Ⱦ
|
||||
ȿ
|
||||
ɀ
|
||||
Ɂ
|
||||
ɂ
|
||||
Ƀ
|
||||
Ʉ
|
||||
Ʌ
|
||||
Ɇ
|
||||
ɇ
|
||||
Ɉ
|
||||
ɉ
|
||||
Ɋ
|
||||
ɋ
|
||||
Ɍ
|
||||
ɍ
|
||||
Ɏ
|
||||
ɏ
|
||||
!
|
||||
"
|
||||
#
|
||||
$
|
||||
%
|
||||
&
|
||||
'
|
||||
(
|
||||
)
|
||||
*
|
||||
+
|
||||
,
|
||||
-
|
||||
.
|
||||
/
|
||||
:
|
||||
;
|
||||
<
|
||||
=
|
||||
>
|
||||
?
|
||||
@
|
||||
[
|
||||
\
|
||||
]
|
||||
_
|
||||
`
|
||||
{
|
||||
|
|
||||
}
|
||||
^
|
||||
~
|
||||
©
|
||||
®
|
||||
℉
|
||||
№
|
||||
Ω
|
||||
℮
|
||||
™
|
||||
∆
|
||||
✓
|
||||
✔
|
||||
✗
|
||||
✘
|
||||
✕
|
||||
☑
|
||||
☒
|
||||
●
|
||||
▪
|
||||
▫
|
||||
◼
|
||||
▶
|
||||
◀
|
||||
⬆
|
||||
¤
|
||||
¦
|
||||
§
|
||||
¨
|
||||
ª
|
||||
«
|
||||
¬
|
||||
¯
|
||||
°
|
||||
²
|
||||
³
|
||||
´
|
||||
µ
|
||||
¶
|
||||
¸
|
||||
¹
|
||||
º
|
||||
»
|
||||
¼
|
||||
½
|
||||
¾
|
||||
¿
|
||||
×
|
||||
‐
|
||||
‑
|
||||
‒
|
||||
—
|
||||
―
|
||||
‖
|
||||
‗
|
||||
‘
|
||||
’
|
||||
‚
|
||||
‛
|
||||
“
|
||||
”
|
||||
„
|
||||
‟
|
||||
†
|
||||
‡
|
||||
‣
|
||||
․
|
||||
…
|
||||
‧
|
||||
‰
|
||||
‴
|
||||
‵
|
||||
‶
|
||||
‷
|
||||
‸
|
||||
‹
|
||||
›
|
||||
※
|
||||
‼
|
||||
‽
|
||||
‾
|
||||
−
|
||||
₤
|
||||
₡
|
||||
₹
|
||||
₽
|
||||
₴
|
||||
₿
|
||||
¢
|
||||
€
|
||||
£
|
||||
¥
|
||||
Ⅰ
|
||||
Ⅱ
|
||||
Ⅲ
|
||||
Ⅳ
|
||||
Ⅴ
|
||||
Ⅵ
|
||||
Ⅶ
|
||||
Ⅷ
|
||||
Ⅸ
|
||||
Ⅹ
|
||||
Ⅺ
|
||||
Ⅻ
|
||||
ⅰ
|
||||
ⅱ
|
||||
ⅲ
|
||||
ⅳ
|
||||
ⅴ
|
||||
ⅵ
|
||||
ⅶ
|
||||
ⅷ
|
||||
ⅸ
|
||||
ⅹ
|
||||
ⅺ
|
||||
ⅻ
|
||||
➀
|
||||
➁
|
||||
➂
|
||||
➃
|
||||
➄
|
||||
➅
|
||||
➆
|
||||
➇
|
||||
➈
|
||||
➉
|
||||
➊
|
||||
➋
|
||||
➌
|
||||
➍
|
||||
➎
|
||||
➏
|
||||
➐
|
||||
➑
|
||||
➒
|
||||
➓
|
||||
❶
|
||||
❷
|
||||
❸
|
||||
❹
|
||||
❺
|
||||
❻
|
||||
❼
|
||||
❽
|
||||
❾
|
||||
❿
|
||||
①
|
||||
②
|
||||
③
|
||||
④
|
||||
⑤
|
||||
⑥
|
||||
⑦
|
||||
⑧
|
||||
⑨
|
||||
⑩
|
||||
↑
|
||||
→
|
||||
↓
|
||||
↕
|
||||
←
|
||||
↔
|
||||
⇒
|
||||
⇐
|
||||
⇔
|
||||
∀
|
||||
∃
|
||||
∄
|
||||
∴
|
||||
∵
|
||||
∝
|
||||
∞
|
||||
∩
|
||||
∪
|
||||
∂
|
||||
∫
|
||||
∬
|
||||
∭
|
||||
∮
|
||||
∯
|
||||
∰
|
||||
∑
|
||||
∏
|
||||
√
|
||||
∛
|
||||
∜
|
||||
∱
|
||||
∲
|
||||
∳
|
||||
∶
|
||||
∷
|
||||
∼
|
||||
∖
|
||||
∗
|
||||
≈
|
||||
≠
|
||||
≡
|
||||
≤
|
||||
≥
|
||||
⊂
|
||||
⊃
|
||||
⊥
|
||||
⊾
|
||||
⊿
|
||||
□
|
||||
∥
|
||||
∋
|
||||
ƒ
|
||||
′
|
||||
″
|
||||
À
|
||||
Á
|
||||
Â
|
||||
Ã
|
||||
Ä
|
||||
Å
|
||||
Æ
|
||||
Ç
|
||||
È
|
||||
É
|
||||
Ê
|
||||
Ë
|
||||
Ì
|
||||
Í
|
||||
Î
|
||||
Ï
|
||||
Ð
|
||||
Ñ
|
||||
Ò
|
||||
Ó
|
||||
Ô
|
||||
Õ
|
||||
Ö
|
||||
Ø
|
||||
Ù
|
||||
Ú
|
||||
Û
|
||||
Ü
|
||||
Ý
|
||||
Þ
|
||||
à
|
||||
á
|
||||
â
|
||||
ã
|
||||
ä
|
||||
å
|
||||
æ
|
||||
ç
|
||||
è
|
||||
é
|
||||
ê
|
||||
ë
|
||||
ì
|
||||
í
|
||||
î
|
||||
ï
|
||||
ð
|
||||
ñ
|
||||
ò
|
||||
ó
|
||||
ô
|
||||
õ
|
||||
ö
|
||||
ø
|
||||
ù
|
||||
ú
|
||||
û
|
||||
ü
|
||||
ý
|
||||
þ
|
||||
ÿ
|
||||
Α
|
||||
Β
|
||||
Γ
|
||||
Δ
|
||||
Ε
|
||||
Ζ
|
||||
Η
|
||||
Θ
|
||||
Ι
|
||||
Κ
|
||||
Λ
|
||||
Μ
|
||||
Ν
|
||||
Ξ
|
||||
Ο
|
||||
Π
|
||||
Ρ
|
||||
Σ
|
||||
Τ
|
||||
Υ
|
||||
Φ
|
||||
Χ
|
||||
Ψ
|
||||
Ω
|
||||
α
|
||||
β
|
||||
γ
|
||||
δ
|
||||
ε
|
||||
ζ
|
||||
η
|
||||
θ
|
||||
ι
|
||||
κ
|
||||
λ
|
||||
μ
|
||||
ν
|
||||
ξ
|
||||
ο
|
||||
π
|
||||
ρ
|
||||
σ
|
||||
ς
|
||||
τ
|
||||
υ
|
||||
φ
|
||||
χ
|
||||
ψ
|
||||
ω
|
||||
Å
|
||||
ℏ
|
||||
⌀
|
||||
⍺
|
||||
⍵
|
||||
𝑢
|
||||
𝜓
|
||||
०
|
||||
‥
|
||||
︽
|
||||
﹥
|
||||
•
|
||||
÷
|
||||
∕
|
||||
∙
|
||||
⋅
|
||||
·
|
||||
±
|
||||
∓
|
||||
∟
|
||||
∠
|
||||
∡
|
||||
∢
|
||||
℧
|
||||
☺
|
||||
+17967
-17642
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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 = [
|
||||
|
||||
Reference in New Issue
Block a user