From 024e3ba0295777b6a6cab7652c95801dc9554bec Mon Sep 17 00:00:00 2001
From: raiots <raiot.lee@outlook.com>
Date: Fri, 2 Jun 2023 22:37:51 +0800
Subject: [PATCH] feat: beauty the function

---
 utils/ocr.py | 39 ++++++++++++++++++++++++++++++++-------
 1 file changed, 32 insertions(+), 7 deletions(-)

diff --git a/utils/ocr.py b/utils/ocr.py
index 6491aee..1742bca 100644
--- a/utils/ocr.py
+++ b/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):
-    source_img = cv2.imread(img_path)
-    d =  pytesseract.image_to_data(source_img, output_type=Output.DICT, config='--psm 10')
+def identy_char(img_path=None, img_source=None):
+    '''
+    @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():
-            print(text)
+        if len(text) == 1:
+            ident_char = text 
+            break
+    if ident_char:
+        return ident_char
+    else:
+        return None
 
 if __name__ == "__main__":
-    identy_char('./img.png')
\ No newline at end of file
+    # identy_char('./img.png')
+    source_img = cv2.imread('./img.png')
+    print(identy_char(img_source=source_img))
+    # print(identy_char('./coffee-ocr.jpg'))
\ No newline at end of file