import numpy as np
from PIL import Image
import cv2


# __all__ = []

def ver_code_rect_split(rect, max_width=50):
 """
 验证码中的矩形框切分
 :param rect: 矩形框(left, upper, width, height)
 :param max_width: 当矩形框超过此值,则会被切分
 :return: list(rect)
 """
 rect_num = rect[2] // max_width
 if rect_num <= 1:
 return [rect]

 rect_lst = []
 rect_width = rect[2] // rect_num
 for i in range(rect_num):
 rect_temp = list(rect)
 rect_temp[0] = rect_temp[0] + i * rect_width
 rect_temp[2] = rect_width
 rect_lst.append(tuple(rect_temp))
 return rect_lst


def ver_code_split(img, out_size=(32, 32), channel_first: bool = True, mode=None):
 if isinstance(img, str):
 img = Image.open(img)
 img_height = 132
 img_width = int(img.size[0] * (img_height / img.size[1]))
 # 先将图像数据放大,再做高斯滤波和中值滤波,最后二值化处理
 img = img.resize((img_width, img_height), resample=Image.BICUBIC).convert('RGB')
 img_pro = cv2.GaussianBlur(np.array(img.convert('L')), (3, 3), 0)
 img_pro = cv2.medianBlur(img_pro, 11)
 ret, img_pro = cv2.threshold(img_pro, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)

 # 先做边缘检测,再做轮廓检测(只需要外轮廓)
 img_pro = cv2.Canny(img_pro, 80, 200, L2gradient=True)
 cnts = cv2.findContours(img_pro.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)[0]
 # 可过滤一些非常小的轮廓,轮廓小于50的过滤,!注意:这个参数应当随图像尺寸大小变化面变化
 cnts = [c for c in cnts if len(c) >= 50]
 rects = [cv2.boundingRect(cf) for cf in cnts] # 找出轮廓的外接矩形

 # 有的字母连接在一起,矩形会很宽,将其分割开来
 filter_rects = []
 for rt in rects:
 if rt[2] * rt[3] <= 1500: # 矩形面积小于1500的过滤掉
 continue
 filter_rects.extend(ver_code_rect_split(rt)) # 将大的矩形框分割开

 # 将所有矩形内的字母提取出来
 letter_lst = []
 for x, y, w, h in sorted(filter_rects, key=lambda px: px[0]):
 img_temp = img.crop((x, y, x + w, y + h))
 if out_size is not None: # 更改尺寸
 img_temp = img_temp.resize(out_size)
 if mode == "L":
 img_array = np.array(img_temp.convert(mode))
 else:
 img_array = np.array(img_temp)
 if channel_first:
 img_array = img_array.transpose((2, 0, 1))
 letter_lst.append(img_array)
 return letter_lst