openocr-python 0.0.9__py3-none-any.whl → 0.1.0.dev0__py3-none-any.whl

openocr/tools/infer_e2e.py

@@ -75,7 +75,7 @@ def sorted_boxes(dt_boxes):
     return _boxes
 
 
-class OpenOCR(object):
+class OpenOCRE2E(object):
 
     def __init__(self,
                  mode='mobile',
@@ -84,7 +84,7 @@ class OpenOCR(object):
                  onnx_rec_model_path=None,
                  drop_score=0.5,
                  det_box_type='quad',
-                 device='gpu'):
+                 use_gpu='auto'):
         """
         初始化函数，用于初始化OCR引擎的相关配置和组件。
 
@@ -92,11 +92,26 @@ class OpenOCR(object):
             mode (str, optional): 运行模式，可选值为'mobile'或'server'。默认为'mobile'。
             drop_score (float, optional): 检测框的置信度阈值，低于该阈值的检测框将被丢弃。默认为0.5。
             det_box_type (str, optional): 检测框的类型，可选值为'quad' and 'poly'。默认为'quad'。
+            use_gpu (str, optional): GPU使用策略，可选值为'auto'/'true'/'false'。默认为'auto'。
 
         Returns:
             无返回值。
 
         """
+        # Parse use_gpu parameter
+        if use_gpu == 'auto':
+            try:
+                import torch
+                device = 'gpu' if torch.cuda.is_available() else 'cpu'
+            except:
+                device = 'cpu'
+        elif use_gpu == 'true':
+            device = 'gpu'
+        elif use_gpu == 'false':
+            device = 'cpu'
+        else:
+            raise ValueError(f"use_gpu must be 'auto', 'true', or 'false', got '{use_gpu}'")
+
         cfg_det = Config(DEFAULT_CFG_PATH_DET).cfg  # mobile model
         cfg_det['Global']['device'] = device
         if mode == 'server':
@@ -108,9 +123,10 @@ class OpenOCR(object):
 
         self.text_detector = OpenDetector(cfg_det,
                                           backend=backend,
-                                          onnx_model_path=onnx_det_model_path)
+                                          onnx_model_path=onnx_det_model_path,
+                                          use_gpu=use_gpu)
         self.text_recognizer = OpenRecognizer(
-            cfg_rec, backend=backend, onnx_model_path=onnx_rec_model_path)
+            cfg_rec, backend=backend, onnx_model_path=onnx_rec_model_path, use_gpu=use_gpu)
         self.det_box_type = det_box_type
         self.drop_score = drop_score
 
@@ -415,10 +431,11 @@ def main():
                         type=float,
                         default=0.5,
                         help='Score threshold for text recognition.')
-    parser.add_argument('--device',
+    parser.add_argument('--use_gpu',
                         type=str,
-                        default='gpu',
-                        help='Device to use for inference.')
+                        default='auto',
+                        choices=['auto', 'true', 'false'],
+                        help='GPU usage strategy: auto (detect automatically), true (force GPU), false (force CPU)')
     args = parser.parse_args()
 
     img_path = args.img_path
@@ -429,15 +446,15 @@ def main():
     save_dir = args.save_dir
     is_visualize = args.is_vis
     drop_score = args.drop_score
-    device = args.device
+    use_gpu = args.use_gpu
 
-    text_sys = OpenOCR(mode=mode,
+    text_sys = OpenOCRE2E(mode=mode,
                        backend=backend,
                        onnx_det_model_path=onnx_det_model_path,
                        onnx_rec_model_path=onnx_rec_model_path,
                        drop_score=drop_score,
                        det_box_type='quad',
-                       device=device)  # det_box_type: 'quad' or 'poly'
+                       use_gpu=use_gpu)  # det_box_type: 'quad' or 'poly'
     text_sys(img_path=img_path, save_dir=save_dir, is_visualize=is_visualize)
 
 

openocr/tools/infer_rec.py

@@ -127,7 +127,7 @@ def build_rec_process(cfg):
     ratio_resize_flag = True
     for op in cfg['Eval']['dataset']['transforms']:
         op_name = list(op)[0]
-        if 'Resize' in op_name:
+        if 'Resize' in op_name or 'Processor' in op_name:
             ratio_resize_flag = False
         if 'Label' in op_name:
             continue
@@ -149,6 +149,8 @@ def set_device(device, numId=0):
     import torch
     if device == 'gpu' and torch.cuda.is_available():
         device = torch.device(f'cuda:{numId}')
+    elif device == 'mps' and torch.backends.mps.is_available():
+        device = torch.device('mps')
     else:
         logger.info('GPU is not available, using CPU.')
         device = torch.device('cpu')
@@ -162,6 +164,7 @@ class OpenRecognizer:
                  mode='mobile',
                  backend='torch',
                  onnx_model_path=None,
+                 use_gpu='auto',
                  numId=0):
         """
         Args:
@@ -169,12 +172,30 @@ class OpenRecognizer:
             mode (str, optional): 模式，'server' 或 'mobile'。默认为'mobile'。
             backend (str): 'torch' 或 'onnx'
             onnx_model_path (str): ONNX模型路径（仅当backend='onnx'时需要）
+            use_gpu (str, optional): GPU使用策略，可选值为'auto'/'true'/'false'。默认为'auto'。
             numId (int, optional): 设备编号。默认为0。
         """
 
         if config is None:
             config_file = DEFAULT_CFG_PATH_REC_SERVER if mode == 'server' else DEFAULT_CFG_PATH_REC
             config = Config(config_file).cfg
+
+        # Parse use_gpu parameter
+        if use_gpu == 'auto':
+            try:
+                import torch
+                device = 'gpu' if torch.cuda.is_available() else 'cpu'
+            except:
+                device = 'cpu'
+        elif use_gpu == 'true':
+            device = 'gpu'
+        elif use_gpu == 'false':
+            device = 'cpu'
+        else:
+            raise ValueError(f"use_gpu must be 'auto', 'true', or 'false', got '{use_gpu}'")
+
+        config['Global']['device'] = device
+
         self.cfg = config
         # 公共初始化
         self._init_common()
@@ -197,7 +218,7 @@ class OpenRecognizer:
                 else:
                     raise ValueError('ONNX模式需要指定onnx_model_path参数')
             self.onnx_rec_engine = ONNXEngine(
-                onnx_model_path, use_gpu=config['Global']['device'] == 'gpu')
+                onnx_model_path, use_gpu=(device == 'gpu'))
         else:
             raise ValueError("backend参数必须是'torch'或'onnx'")
 
@@ -222,26 +243,43 @@ class OpenRecognizer:
 
     def _init_torch_model(self, numId):
         from tools.utils.ckpt import load_ckpt
-        from tools.infer_det import replace_batchnorm
-        # PyTorch专用初始化
         algorithm_name = self.cfg['Architecture']['algorithm']
-        if algorithm_name in ['SVTRv2_mobile', 'SVTRv2_server']:
-            if not os.path.exists(self.cfg['Global']['pretrained_model']):
-                pretrained_model = check_and_download_model(
-                    MODEL_NAME_REC, DOWNLOAD_URL_REC
-                ) if algorithm_name == 'SVTRv2_mobile' else check_and_download_model(
-                    MODEL_NAME_REC_SERVER, DOWNLOAD_URL_REC_SERVER)
-                self.cfg['Global']['pretrained_model'] = pretrained_model
-
-        from openrec.modeling import build_model as build_rec_model
+        if self.cfg['Global'].get('use_transformers', False):
+            if algorithm_name == 'UniRec':
+                from openrec.modeling.unirec_modeling.modeling_unirec import UniRecForConditionalGenerationNew
+                from openrec.modeling.unirec_modeling.configuration_unirec import UniRecConfig
+                cfg_model = UniRecConfig.from_pretrained(
+                    self.cfg['Global']['vlm_ocr_config'])
+                # cfg_model._attn_implementation = "flash_attention_2"
+                cfg_model._attn_implementation = 'eager'
+                self.model = UniRecForConditionalGenerationNew(
+                    config=cfg_model)
+            elif algorithm_name == 'CMER':
+                from openrec.modeling.cmer_modeling.modeling_cmer import CMER, CMERConfig
+                cfg_model = CMERConfig(
+                    self.cfg['Architecture']['vision_config'],
+                    self.cfg['Architecture']['decoder_config'])
+                self.model = CMER(config=cfg_model)
+        else:
+            # PyTorch专用初始化
+            if algorithm_name in ['SVTRv2_mobile', 'SVTRv2_server']:
+                if not os.path.exists(self.cfg['Global']['pretrained_model']):
+                    pretrained_model = check_and_download_model(
+                        MODEL_NAME_REC, DOWNLOAD_URL_REC
+                    ) if algorithm_name == 'SVTRv2_mobile' else check_and_download_model(
+                        MODEL_NAME_REC_SERVER, DOWNLOAD_URL_REC_SERVER)
+                    self.cfg['Global']['pretrained_model'] = pretrained_model
+
+            from openrec.modeling import build_model as build_rec_model
+            self.model = build_rec_model(self.cfg['Architecture'])
 
-        self.model = build_rec_model(self.cfg['Architecture'])
         load_ckpt(self.model, self.cfg)
 
         self.device = set_device(self.cfg['Global']['device'], numId)
         self.model.to(self.device)
         self.model.eval()
         if algorithm_name == 'SVTRv2_mobile':
+            from tools.infer_det import replace_batchnorm
             replace_batchnorm(self.model.encoder)
 
     def _inference_onnx(self, images):
@@ -329,7 +367,18 @@ class OpenRecognizer:
                 images = self.torch.from_numpy(padded_batch).to(
                     device=self.device)
                 with self.torch.no_grad():
-                    preds = self.model(images, others)  # bs, len, num_classes
+                    if self.cfg['Global'].get('use_transformers', False):
+                        # transformers模型推理
+                        inputs = {
+                            'pixel_values': images,
+                            'input_ids': None,
+                            'attention_mask': None
+                        }
+                        preds = self.model.generate(**inputs)
+                    else:
+                        # PyTorch模型推理
+                        preds = self.model(images,
+                                           others)  # bs, len, num_classes
                 torch_tensor = True
             elif self.backend == 'onnx':
                 # ONNX推理