diff --git a/src/tenon/detect.py b/src/tenon/detect.py index 487ad32..fabba97 100644 --- a/src/tenon/detect.py +++ b/src/tenon/detect.py @@ -1,8 +1,13 @@ +import math import re import cv2 import base64 import numpy as np from typing import Optional +from collections import namedtuple + +RotateData = namedtuple("RotateData", "similar, angle, start, end, step") +MatchData = namedtuple("MatchData", "similar, inner_rotate_angle, total_rotate_angle") def request_image_content(image_url, proxies: Optional[dict] = None): @@ -16,58 +21,79 @@ def request_image_content(image_url, proxies: Optional[dict] = None): return response.content -def set_mask(shape, r1, r2): - (x, y) = (shape[0] // 2, shape[1] // 2) - mask = np.zeros(shape[:2], dtype=np.uint8) - mask = cv2.circle(mask, (x, y), r1, (255, 255, 255), -1) - mask = cv2.circle(mask, (x, y), r2, (0, 0, 0), -1) +def set_mask(radius, check_pixel): + center_point = (radius, radius) + mask = np.zeros((radius * 2, radius * 2), dtype=np.uint8) + mask = cv2.circle(mask, center_point, radius, (255, 255, 255), -1) + mask = cv2.circle(mask, center_point, radius - check_pixel, (0, 0, 0), -1) return mask -def cut_image(origin_image, crop_pixel): - height, width = origin_image.shape[:2] - h_cut_size = (height - crop_pixel * 2) // 2 - w_cut_size = (width - crop_pixel * 2) // 2 - image = origin_image[ - h_cut_size : height - h_cut_size, w_cut_size : width - w_cut_size - ] - return image - - -def mask_image(origin_image, r1, r2): - deal_image = cut_image(origin_image, r1) - mask = set_mask(deal_image.shape, r1, r2) - image = cv2.add(deal_image, np.zeros(deal_image.shape, dtype=np.uint8), mask=mask) - return image - - -def rotate_image(inner_image, outer_image, rotate_type): - angle = 0 - max_val = 0 - start = 0 - end = 360 +def cut_image(origin_array, color_type, 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)) + if not pixel_set.issubset({0, 255, (0, 0, 0), (255, 255, 255)}): + break + cut_pixel += 1 + cut_pixel_list.append(cut_pixel) + cut_pixel_list[2] = height - cut_pixel_list[2] + cut_pixel_list[3] = width - cut_pixel_list[3] + + elif check_pixel: + y, x = height // 2, width // 2 + resize_check_pixel = math.ceil(radius / (radius - check_pixel) * check_pixel) + for i in -1, 1: + for p in y, x: + pos = p + i * radius + 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) + if not pixel_set.issubset({0, 255}): + break + pos += i + cut_pixel_list.append(pos + i * resize_check_pixel) + + up, left, down, right = cut_pixel_list + cut_array = origin_array[up:down, left:right] + diameter = (radius or min(cut_array.shape[:2]) // 2) * 2 + cut_result = cv2.resize(cut_array, dsize=(diameter, diameter)) + return cut_result + + +def mask_image(origin_array, check_pixel): + radius = origin_array.shape[0] // 2 + mask = set_mask(radius, check_pixel) + src_array = np.zeros(origin_array.shape, dtype=np.uint8) + mask_result = cv2.add(origin_array, src_array, mask=mask) + return mask_result + + +def rotate_image(inner_image, outer_image, anticlockwise): + rotate_info_list = [RotateData(0, 0, 1, 361, 10)] + rtype = int(anticlockwise) or -1 h, w = inner_image.shape[:2] - rotate_type = 1 if rotate_type else -1 - for i in range(0, 2): - rotate_max_similar = 0 - step = 10 ** (1 - i) - rotate = start - while rotate < end: - rotate += step - mat_rotate = cv2.getRotationMatrix2D( - (h * 0.5, w * 0.5), rotate_type * rotate, 1 - ) + for item in rotate_info_list: + min_similar_rotate_info = item + for angle in range(*item[2:]): + mat_rotate = cv2.getRotationMatrix2D((h * 0.5, w * 0.5), rtype * angle, 1) dst = cv2.warpAffine(inner_image, mat_rotate, (h, w)) - result = cv2.matchTemplate(outer_image, dst, cv2.TM_CCOEFF_NORMED) - min_max_loc = cv2.minMaxLoc(result) - if min_max_loc[1] > max_val: - max_val = min_max_loc[1] - if max_val >= rotate_max_similar: # 获取最高匹配值 - rotate_max_similar = max_val - angle = rotate - start = angle - step - end = angle + step - return max_val, angle + ret = cv2.matchTemplate(outer_image, dst, cv2.TM_CCOEFF_NORMED) + similar_value = cv2.minMaxLoc(ret)[1] + if similar_value < min_similar_rotate_info.similar: + continue + rotate_info = RotateData(similar_value, angle, angle - 10, angle + 10, 10) + rotate_info_list.append(rotate_info) + if len(rotate_info_list) > 5: + rotate_info_list.remove(min_similar_rotate_info) + min_similar_rotate_info = min(rotate_info_list) + return max(rotate_info_list) def image_to_cv2(base_image: str, image_type: int, color_type: bool, proxies=None): @@ -98,9 +124,7 @@ def rotate_identify( image_type: int = 0, color_type: bool = True, check_pixel: int = 10, - rotate_type: bool = False, - big_circle_empty_radius: int = 0, - small_circle_crop_pixel: int = 0, + anticlockwise: bool = False, speed_ratio: float = 1, proxies: Optional[dict] = None, ): @@ -111,34 +135,27 @@ def rotate_identify( :param image_type: 图片类型: 0: 图片base64; 1: 图片url; 2: 图片文件地址 :param color_type: 是否需要灰度化处理: True: 是; False: 否 :param check_pixel: 进行图片验证的像素宽度 - :param rotate_type: 图片旋转的类型: True: 小圈逆时针; False: 小圈顺时针 - :param big_circle_empty_radius: 大圈内部留白部分半径（在留白部分大于内圈实际图片时传值） - :param small_circle_crop_pixel: 小圈外部留白的像素:（图片宽度 - 有图部分的直径) / 2 + :param anticlockwise: 图片旋转的类型: True: 小圈逆时针; False: 小圈顺时针 :param speed_ratio: 小圈与大圈的转动速率比: 小圈转动360度时/大圈转动的角度 :param proxies: 代理 - :return: dict(相似度, 总共旋转的角度, 小圈图片旋转的角度) + :return: namedtuple -> MatchData """ - small_circle_image = image_to_cv2(small_circle, image_type, color_type, proxies) - big_circle_image = image_to_cv2(big_circle, image_type, color_type, proxies) - if isinstance(small_circle_image, bool) or isinstance(big_circle_image, bool): + 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! 图片解析错误!") - small_circle_r1 = small_circle_image.shape[:2][1] // 2 - small_circle_crop_pixel - small_circle_r2 = small_circle_r1 - check_pixel + cut_inner_image = cut_image(inner_image, color_type) + cut_inner_radius = cut_inner_image.shape[0] // 2 + cut_outer_image = cut_image(outer_image, color_type, cut_inner_radius, check_pixel) - big_circle_r2 = ( - big_circle_empty_radius if big_circle_empty_radius else small_circle_r1 - ) - big_circle_r1 = big_circle_r2 + check_pixel + inner_annulus = mask_image(cut_inner_image, check_pixel) + outer_annulus = mask_image(cut_outer_image, check_pixel) - outer_image_before_resize = mask_image( - big_circle_image, big_circle_r1, big_circle_r2 - ) - inner_image = mask_image(small_circle_image, small_circle_r1, small_circle_r2) - outer_image = cv2.resize(outer_image_before_resize, inner_image.shape[:2]) - similar, total_angle = rotate_image(inner_image, outer_image, rotate_type) - inner_angle = round(total_angle * speed_ratio / (speed_ratio + 1), 2) - return dict(similar=similar, total_angle=total_angle, inner_angle=inner_angle) + rotate_info = rotate_image(inner_annulus, outer_annulus, anticlockwise) + inner_angle = round(rotate_info.angle * speed_ratio / (speed_ratio + 1), 2) + return MatchData(rotate_info.similar, inner_angle, rotate_info.angle) def notch_identify( @@ -176,41 +193,47 @@ def rotate_identify_and_show_image( image_type: int = 0, color_type: bool = True, check_pixel: int = 10, - rotate_type: bool = False, - big_circle_empty_radius: int = 0, - small_circle_crop_pixel: int = 0, + anticlockwise: bool = False, image_show_time: int = 0, proxies: Optional[dict] = None, ): - small_circle_image = image_to_cv2(small_circle, image_type, color_type, proxies) - big_circle_image = image_to_cv2(big_circle, image_type, color_type, proxies) - if isinstance(small_circle_image, bool) or isinstance(big_circle_image, bool): + """ + 双图旋转类型滑块验证码识别 + :param small_circle: 小圈图片 + :param big_circle: 大圈图片 + :param image_type: 图片类型: 0: 图片base64; 1: 图片url; 2: 图片文件地址 + :param color_type: 是否需要灰度化处理: 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! 图片解析错误!") - small_circle_r1 = small_circle_image.shape[:2][1] // 2 - small_circle_crop_pixel - small_circle_r2 = small_circle_r1 - check_pixel - big_circle_r2 = ( - big_circle_empty_radius if big_circle_empty_radius else small_circle_r1 - ) - big_circle_r1 = big_circle_r2 + check_pixel - - outer_image_before_resize = mask_image( - big_circle_image, big_circle_r1, big_circle_r2 - ) - inner_image = mask_image(small_circle_image, small_circle_r1, small_circle_r2) - outer_image = cv2.resize(outer_image_before_resize, inner_image.shape[:2]) - - similar, total_rotate_angle = rotate_image(inner_image, outer_image, rotate_type) - - cut_small_image = cut_image(small_circle_image, small_circle_r1) - height, width = cut_small_image.shape[:2] - if rotate_type: - rotate_angle = total_rotate_angle - else: - rotate_angle = -total_rotate_angle + cut_inner_image = cut_image(inner_image, color_type) + cut_inner_radius = cut_inner_image.shape[0] // 2 + cut_outer_image = cut_image(outer_image, color_type, cut_inner_radius, check_pixel) + + cv2.imshow("cut_inner_image", cut_inner_image) + cv2.imshow("cut_outer_image", cut_outer_image) + + inner_annulus = mask_image(cut_inner_image, check_pixel) + outer_annulus = mask_image(cut_outer_image, check_pixel) + # outer_annulus = cv2.resize(outer_annulus, inner_annulus.shape[:2]) + + cv2.imshow("inner_annulus", inner_annulus) + cv2.imshow("outer_annulus", outer_annulus) + rotate_info = rotate_image(inner_annulus, outer_annulus, anticlockwise) + height, width = cut_inner_image.shape[:2] + rotate_angle = rotate_info.angle * (int(anticlockwise) or -1) mat_rotate = cv2.getRotationMatrix2D((height * 0.5, width * 0.5), rotate_angle, 1) - rotate_small_image = cv2.warpAffine(cut_small_image, mat_rotate, (height, width)) + rotated_inner_image = cv2.warpAffine(cut_inner_image, mat_rotate, (height, width)) - cv2.imshow("big_circle_image", big_circle_image) - cv2.imshow("rotate_small_image", rotate_small_image) + cv2.imshow("cut_outer_image", cut_outer_image) + cv2.imshow("outer_image", outer_image) + cv2.imshow("rotated_inner_image", rotated_inner_image) cv2.waitKey(image_show_time * 1000) diff --git a/tests/test_detect.py b/tests/test_detect.py index 91a4a6a..5ef0e01 100644 --- a/tests/test_detect.py +++ b/tests/test_detect.py @@ -6,28 +6,24 @@ @pytest.mark.parametrize( - "small_circle, big_circle, image_type, small_circle_crop_pixel, expected", + "small_circle, big_circle, image_type, expected", [ ( current_work_dir + "/image/inner.png", current_work_dir + "/image/outer.png", 2, - 102, - 313, + 311, ), ], ) -def test_rotate_identify( - small_circle, big_circle, image_type, small_circle_crop_pixel, expected -): +def test_rotate_identify(small_circle, big_circle, image_type, expected): result = rotate_identify( small_circle, big_circle, image_type, - small_circle_crop_pixel=small_circle_crop_pixel, ) - assert result["total_angle"] == expected + assert result.total_rotate_angle == expected @pytest.mark.parametrize(