feat: support to log_level

Author: SWHL

demo.py

@@ -21,7 +21,7 @@
 # 只识别
 from rapid_videocr import RapidVideOCR, RapidVideOCRInput
 
-input_args = RapidVideOCRInput(is_batch_rec=False)
+input_args = RapidVideOCRInput(is_batch_rec=False, log_level="critical")
 extractor = RapidVideOCR(input_args)
 
 rgb_dir = "tests/test_files/RGBImages"

rapid_videocr/main.py

@@ -8,14 +8,14 @@
 from .export import ExportStrategyFactory, OutputFormat
 from .ocr_processor import OCRProcessor
 from .utils.crop_by_project import CropByProject
-from .utils.logger import Logger
+from .utils.logger import logger
 from .utils.typings import IMAGE_EXTENSIONS, RapidVideOCRInput
 from .utils.utils import mkdir
 
 
 class RapidVideOCR:
     def __init__(self, input_params: RapidVideOCRInput):
-        self.logger = Logger(logger_name=__name__).get_log()
+        logger.setLevel(input_params.log_level.upper())
 
         self.ocr_processor = OCRProcessor(
             input_params.ocr_params, input_params.batch_size
@@ -78,12 +78,11 @@ def export_file(
     ):
         try:
             strategy = ExportStrategyFactory.create_strategy(self.out_format)
-
             mkdir(save_dir)
             strategy.export(save_dir, save_name, srt_result, ass_result, txt_result)
-            self.logger.info("[OCR] Results saved to directory: %s", save_dir)
+            logger.info(f"[OCR] Results saved to directory: {save_dir}")
         except ValueError as e:
-            self.logger.error("Export failed: %s", str(e))
+            logger.error(f"Export failed: {e}")
             raise
 
     def print_console(self, txt_result: List):

rapid_videocr/ocr_processor.py

@@ -9,7 +9,8 @@
 from rapidocr import RapidOCR
 from tqdm import tqdm
 
-from .utils.logger import Logger
+from .utils.logger import logger
+from .utils.typings import LOG_LEVEL_MAP
 from .utils.utils import (
     compute_centroid,
     compute_poly_iou,
@@ -21,11 +22,15 @@
 
 class OCRProcessor:
     def __init__(self, ocr_params: Optional[Dict] = None, batch_size: int = 10):
-        self.logger = Logger(logger_name=__name__).get_log()
         self.ocr_engine = self._init_ocr_engine(ocr_params)
         self.batch_size = batch_size
 
     def _init_ocr_engine(self, ocr_params: Optional[Dict] = None) -> RapidOCR:
+        log_level_dict = {"Global.log_level": LOG_LEVEL_MAP[logger.level]}
+        if ocr_params is None:
+            return RapidOCR(params=log_level_dict)
+
+        ocr_params.update(log_level_dict)
         return RapidOCR(params=ocr_params)
 
     def __call__(
@@ -40,7 +45,7 @@ def __call__(
         return srt_results, ass_results, txt_results
 
     def single_rec(self, img_list: List[Path]) -> List[Tuple[int, str, str, str]]:
-        self.logger.info("[OCR] Running with single recognition.")
+        logger.info("[OCR] Running with single recognition.")
 
         rec_results = []
         for i, img_path in enumerate(tqdm(img_list, desc="OCR")):
@@ -112,7 +117,7 @@ def _preprocess_image(img_path: Path) -> np.ndarray:
     def _generate_srt_results(
         rec_results: List[Tuple[int, str, str, str]],
     ) -> List[str]:
-        return [f"{i+1}\n{time_str}\n{txt}\n" for i, time_str, txt, _ in rec_results]
+        return [f"{i + 1}\n{time_str}\n{txt}\n" for i, time_str, txt, _ in rec_results]
 
     @staticmethod
     def _generate_ass_results(
@@ -128,7 +133,7 @@ def _generate_txt_result(rec_results: List[Tuple[int, str, str, str]]) -> List[s
         return [f"{txt}\n" for _, _, txt, _ in rec_results]
 
     def batch_rec(self, img_list: List[Path]) -> List[Tuple[int, str, str, str]]:
-        self.logger.info("[OCR] Running with concat recognition.")
+        logger.info("[OCR] Running with concat recognition.")
 
         img_nums = len(img_list)
         rec_results = []

rapid_videocr/utils/logger.py

@@ -13,7 +13,7 @@ def __init__(self, log_level=logging.DEBUG, logger_name=None):
         self.logger.propagate = False
 
         formatter = colorlog.ColoredFormatter(
-            "%(log_color)s[%(levelname)s] %(asctime)s [RapidVideOCR] %(filename)s:%(lineno)d: %(message)s",
+            "%(log_color)s[%(levelname)s] %(asctime)s %(filename)s:%(lineno)d: %(message)s",
             log_colors={
                 "DEBUG": "cyan",
                 "INFO": "green",
@@ -35,3 +35,6 @@ def __init__(self, log_level=logging.DEBUG, logger_name=None):
 
     def get_log(self):
         return self.logger
+
+
+logger = Logger(log_level=logging.INFO).get_log()

rapid_videocr/utils/typings.py

@@ -28,3 +28,14 @@ class RapidVideOCRInput:
     batch_size: int = 10
     out_format: str = OutputFormat.ALL.value
     ocr_params: Optional[Dict[str, Any]] = None
+    log_level: str = "info"  # debug / info / warning / error / critical
+
+
+LOG_LEVEL_MAP = {
+    50: "CRITICAL",
+    40: "ERROR",
+    30: "WARNING",
+    20: "INFO",
+    10: "DEBUG",
+    0: "NOTSET",
+}

rapid_videocr/vsf_ocr_cli.py

@@ -5,7 +5,7 @@
 from pathlib import Path
 
 from .main import OutputFormat, RapidVideOCR, RapidVideOCRInput
-from .utils.logger import Logger
+from .utils.logger import logger
 from .utils.typings import VideoFormat
 from .utils.utils import float_range
 from .vsf_cli import VideoSubFinder, VideoSubFinderInput
@@ -17,7 +17,6 @@ def __init__(
         vsf_input_params: VideoSubFinderInput,
         ocr_input_params: RapidVideOCRInput,
     ):
-        self.logger = Logger(logger_name=__name__).get_log()
         self.vsf = VideoSubFinder(vsf_input_params)
         self.video_ocr = RapidVideOCR(ocr_input_params)
         self.video_formats = [VideoFormat[v].value for v in VideoFormat.__members__]
@@ -31,13 +30,13 @@ def __call__(self, video_path: str, output_dir: str = "outputs"):
         else:
             video_list = [video_path]
 
-        self.logger.info(
+        logger.info(
             "Extracting subtitle images with VideoSubFinder (takes quite a long time) ..."
         )
         video_num = len(video_list)
         for i, one_video in enumerate(video_list):
-            self.logger.info(
-                "[%s/%s] Starting to extract %s key frame", i + 1, video_num, one_video
+            logger.info(
+                f"[{i + 1}{video_num}] Starting to extract {one_video} key frame"
             )
 
             save_name = Path(one_video).stem
@@ -47,16 +46,14 @@ def __call__(self, video_path: str, output_dir: str = "outputs"):
             try:
                 self.vsf(str(one_video), str(save_vsf_dir))
             except Exception as e:
-                self.logger.error("Extract %s error, %s, skip", one_video, e)
+                logger.error(f"Extract {one_video} error, {e}, skip")
                 continue
 
-            self.logger.info(
-                "[%s/%s] Starting to run %s ocr", i + 1, video_num, one_video
-            )
+            logger.info(f"{i + 1}/{video_num}] Starting to run {one_video} ocr")
 
             rgb_dir = Path(save_vsf_dir) / "RGBImages"
             if not list(rgb_dir.iterdir()):
-                self.logger.warning("Extracting frames from %s is 0, skip", one_video)
+                logger.warning(f"Extracting frames from {one_video} is 0, skip")
                 continue
             self.video_ocr(rgb_dir, save_dir, save_name=save_name)