debug fix

src/tenon/detect.py

@@ -29,15 +29,17 @@ def set_mask(radius, check_pixel):
     return mask
 
 
-def cut_image(origin_array, color_type, radius=None, check_pixel=None):
+def cut_image(origin_array, radius=None, check_pixel=None):
     cut_pixel_list = []  # 上, 左, 下, 右
     height, width = origin_array.shape[:2]
     if not radius:
         for rotate_count in range(4):
             cut_pixel = 0
             rotate_array = np.rot90(origin_array, rotate_count)
             for line in rotate_array:
-                pixel_set = set(list(line)) if color_type else set(map(tuple, line))
+                pixel_set = (
+                    set(list(line)) if len(line.shape) == 1 else set(map(tuple, line))
+                )
                 if not pixel_set.issubset({0, 255, (0, 0, 0), (255, 255, 255)}):
                     break
                 cut_pixel += 1
@@ -54,7 +56,11 @@ def cut_image(origin_array, color_type, radius=None, check_pixel=None):
                 for _ in range(p - radius):
                     p_x, p_y = (pos, y) if len(cut_pixel_list) % 2 else (x, pos)
                     pixel_point = origin_array[p_x][p_y]
-                    pixel_set = {pixel_point} if color_type else set(pixel_point)
+                    pixel_set = (
+                        {pixel_point}
+                        if isinstance(pixel_point, np.uint8)
+                        else set(pixel_point)
+                    )
                     if not pixel_set.issubset({0, 255}):
                         break
                     pos += i
@@ -96,59 +102,58 @@ def rotate_image(inner_image, outer_image, anticlockwise):
     return max(rotate_info_list)
 
 
-def image_to_cv2(base_image: str, image_type: int, color_type: bool, proxies=None):
-    if image_type not in [0, 1, 2]:
-        raise Exception("image_type error! 图片类型错误！")
-    image_color_type = cv2.COLOR_RGB2GRAY if color_type else cv2.IMREAD_COLOR
+def image_to_cv2(base_image: str, image_type: int, grayscale: bool, proxies=None):
+    assert image_type in [0, 1, 2]
     if image_type == 0:
         search_base64 = re.search("base64,(.*?)$", base_image)
         base64_image = search_base64.group(1) if search_base64 else base_image
         image_array = np.asarray(
             bytearray(base64.b64decode(base64_image)), dtype="uint8"
         )
-        image = cv2.imdecode(image_array, image_color_type)
+        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
     elif image_type == 1:
         image_content = request_image_content(base_image, proxies)
         if not image_content:
             raise Exception("请求图片链接失败！")
         image_array = np.array(bytearray(image_content), dtype=np.uint8)
-        image = cv2.imdecode(image_array, image_color_type)
+        image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
     else:
-        image = cv2.cvtColor(cv2.imread(base_image), image_color_type)
+        image = cv2.imread(base_image)
+
+    if grayscale:
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+
     return image
 
 
 def rotate_identify(
     small_circle: str,
     big_circle: str,
     image_type: int = 0,
-    color_type: bool = True,
     check_pixel: int = 10,
-    anticlockwise: bool = False,
     speed_ratio: float = 1,
+    grayscale: bool = False,
+    anticlockwise: bool = False,
     proxies: Optional[dict] = None,
 ):
     """
     双图旋转类型滑块验证码识别
     :param small_circle: 小圈图片
     :param big_circle: 大圈图片
     :param image_type: 图片类型: 0: 图片base64; 1: 图片url; 2: 图片文件地址
-    :param color_type: 是否需要灰度化处理: True: 是; False: 否
+    :param grayscale: 是否需要灰度化处理: True: 是; False: 否
     :param check_pixel: 进行图片验证的像素宽度
     :param anticlockwise: 图片旋转的类型: True: 小圈逆时针; False: 小圈顺时针
     :param speed_ratio: 小圈与大圈的转动速率比: 小圈转动360度时/大圈转动的角度
     :param proxies: 代理
     :return: namedtuple -> MatchData
     """
-    inner_image = image_to_cv2(small_circle, image_type, color_type, proxies)
-    outer_image = image_to_cv2(big_circle, image_type, color_type, proxies)
+    inner_image = image_to_cv2(small_circle, image_type, grayscale, proxies)
+    outer_image = image_to_cv2(big_circle, image_type, grayscale, proxies)
 
-    if isinstance(inner_image, bool) or isinstance(outer_image, bool):
-        raise Exception("image_to_cv2 error! 图片解析错误!")
-
-    cut_inner_image = cut_image(inner_image, color_type)
+    cut_inner_image = cut_image(inner_image)
     cut_inner_radius = cut_inner_image.shape[0] // 2
-    cut_outer_image = cut_image(outer_image, color_type, cut_inner_radius, check_pixel)
+    cut_outer_image = cut_image(outer_image, cut_inner_radius, check_pixel)
 
     inner_annulus = mask_image(cut_inner_image, check_pixel)
     outer_annulus = mask_image(cut_outer_image, check_pixel)
@@ -176,8 +181,6 @@ def notch_identify(
     """
     slider_img = image_to_cv2(slider, image_type, color_type, proxies)
     background_img = image_to_cv2(background, image_type, color_type, proxies)
-    if isinstance(slider_img, bool) or isinstance(background_img, bool):
-        raise Exception("notch_detect error! 图片解析错误!")
     background_edge = cv2.Canny(background_img, 100, 200)
     slider_edge = cv2.Canny(slider_img, 100, 200)
     background_pic = cv2.cvtColor(background_edge, cv2.COLOR_GRAY2RGB)
@@ -191,32 +194,30 @@ def rotate_identify_and_show_image(
     small_circle: str,
     big_circle: str,
     image_type: int = 0,
-    color_type: bool = True,
     check_pixel: int = 10,
-    anticlockwise: bool = False,
     image_show_time: int = 0,
+    grayscale: bool = False,
+    anticlockwise: bool = False,
     proxies: Optional[dict] = None,
 ):
     """
     双图旋转类型滑块验证码识别
     :param small_circle: 小圈图片
     :param big_circle: 大圈图片
     :param image_type: 图片类型: 0: 图片base64; 1: 图片url; 2: 图片文件地址
-    :param color_type: 是否需要灰度化处理: True: 是; False: 否
+    :param grayscale: 是否需要灰度化处理: True: 是; False: 否
     :param check_pixel: 进行图片验证的像素宽度
     :param anticlockwise: 图片旋转的类型: True: 小圈逆时针; False: 小圈顺时针
     :param image_show_time: 调试下显式图片时间, 默认常态显式
     :param proxies: 代理
     :return: namedtuple -> MatchData
     """
-    inner_image = image_to_cv2(small_circle, image_type, color_type, proxies)
-    outer_image = image_to_cv2(big_circle, image_type, color_type, proxies)
-    if isinstance(inner_image, bool) or isinstance(outer_image, bool):
-        raise Exception("image_to_cv2 error! 图片解析错误!")
+    inner_image = image_to_cv2(small_circle, image_type, grayscale, proxies)
+    outer_image = image_to_cv2(big_circle, image_type, grayscale, proxies)
 
-    cut_inner_image = cut_image(inner_image, color_type)
+    cut_inner_image = cut_image(inner_image)
     cut_inner_radius = cut_inner_image.shape[0] // 2
-    cut_outer_image = cut_image(outer_image, color_type, cut_inner_radius, check_pixel)
+    cut_outer_image = cut_image(outer_image, cut_inner_radius, check_pixel)
 
     cv2.imshow("cut_inner_image", cut_inner_image)
     cv2.imshow("cut_outer_image", cut_outer_image)