图片类型增加为jpeg/jpg/png

cuifengcn · Jan 3, 2023 · 67fcc2c · 67fcc2c
1 parent f1a2a4f
commit 67fcc2c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/methods/image_ocr.py b/methods/image_ocr.py
@@ -11,14 +11,14 @@
 
 class ImageOcrClient:
     def __init__(self):
-        self.model = CnOcr(
-            det_model_name="ch_PP-OCRv3_det",
-        )
+        self.model = CnOcr()
 
     def start(self, keyword, logger=logging.getLogger()):
         main_path = WORK_PATH.joinpath(keyword)
         main_path.mkdir(exist_ok=True, parents=True)
-        files = glob.glob(str(main_path.joinpath("*.jpeg")))
+        files = (glob.glob(str(main_path.joinpath("*.jpeg"))) +
+                 glob.glob(str(main_path.joinpath("*.jpg"))) +
+                 glob.glob(str(main_path.joinpath("*.png"))))
         target_path = main_path.joinpath("ocr")
         target_path.mkdir(exist_ok=True)
         for f in files:

diff --git a/ui/ui.py b/ui/ui.py
@@ -227,7 +227,9 @@ def show_images(self):
         self.images.clear()
         self.ocrs.clear()
         main_path = WORK_PATH.joinpath(self.keyword.value)
-        self.images.extend(glob.glob(str(main_path.joinpath("*.jpeg"))))
+        self.images.extend(glob.glob(str(main_path.joinpath("*.jpeg"))) +
+                           glob.glob(str(main_path.joinpath("*.jpg"))) +
+                           glob.glob(str(main_path.joinpath("*.png"))))
         self.images.sort(key=lambda e: int(os.path.split(e)[-1].split(".")[0]))
         self.index = 0
         self._set_image()