feat: beauty the function

utils/ocr.py

@@ -4,7 +4,7 @@ import numpy as np
 from pytesseract import Output
 
 # img_source = cv2.VideoCapture(0)
-img_source = cv2.imread('./img.png')
+
 
 
 def get_grayscale(image):
@@ -25,6 +25,7 @@ def canny(image):
  
  
 def test_performance():
+    img_source = cv2.imread('./img.png')
     gray = get_grayscale(img_source)
     thresh = thresholding(gray)
     opening = opening(gray)
@@ -53,14 +54,38 @@ def test_performance():
 
 # print(pytesseract.image_to_data(img_source, output_type=Output.DICT)['text'])
 
-def identy_char(img_path):
+def identy_char(img_path=None, img_source=None):
-    source_img = cv2.imread(img_path)
+    '''
-    d =  pytesseract.image_to_data(source_img, output_type=Output.DICT, config='--psm 10')
+    @description: 识别单个字符
+    @param {img_path} 图片路径
+    @param {img_source} 图片源
+    @return {ident_char} 识别到的字符
+    '''
 
+    ident_char = None
+    if img_path != None and img_source.any() == None:
+        source_img = cv2.imread(img_path)
+
+    elif img_path == None and img_source.any() != None:
+        source_img = img_source
+
+    else:
+        raise ValueError('img_path and img_source cannot be None OR Exist at the same time')
+    
+    d =  pytesseract.image_to_data(source_img, output_type=Output.DICT, config='--psm 10')  # --psm 10: single char
+    print(d['text'])
     for text in d['text']:
         # print(text)
-        if len(text) == 1 and text.isalpha():
+        if len(text) == 1:
-            print(text)
+            ident_char = text 
+            break
+    if ident_char:
+        return ident_char
+    else:
+        return None
 
 if __name__ == "__main__":
-    identy_char('./img.png')
+    # identy_char('./img.png')
+    source_img = cv2.imread('./img.png')
+    print(identy_char(img_source=source_img))
+    # print(identy_char('./coffee-ocr.jpg'))