生成点选验证码图片

参考博客：https://blog.csdn.net/sinat_39629323/article/details/121989609

from tqdm import tqdm
from PIL import Image, ImageDraw, ImageFont, ImageOps
import shutil,os
import numpy as np
import cv2
import math
import random
file_path = "/data/lh123/lh/verification_code/generate_data/ppocr_keys_v1.txt"#字典

#返回字外面矩形的4个框
def rotate_rectangle(top_left, bottom_right, angle_degrees):
    # 转换角度为弧度
    angle_rad = math.radians(angle_degrees)
    
    # 矩形的四个角的坐标
    top_right = (bottom_right[0], top_left[1])
    bottom_left = (top_left[0], bottom_right[1])

    # 找到矩形的中心
    center = ((top_left[0]+bottom_right[0])/2, (top_left[1]+bottom_right[1])/2)

    # 定义一个函数来旋转一个点
    def rotate_point(point):
        # 移动到以中心为原点的坐标系
        x = point[0] - center[0]
        y = center[1] - point[1]  # 注意我们在这里翻转y轴，因为图像的原点在左上角

        # 在新的坐标系中进行旋转
        new_x = x * math.cos(angle_rad) - y * math.sin(angle_rad)
        new_y = x * math.sin(angle_rad) + y * math.cos(angle_rad)

        # 再次翻转y轴并加上旋转中心的坐标
        return new_x + center[0], center[1] - new_y

    # 这里我们直接返回没有旋转的矩形的四个角的坐标
    points = [top_left, top_right, bottom_right, bottom_left]

    # 展开点列表并返回
    return [coord for point in points for coord in point]
#---=---img为PIL 对象，将这个图片转为数组
def img_to_array(img,x1, y1, x2, y2, x3, y3, x4, y4):
    width, height = img.size
    pixel_data = list(img.getdata())
    return [pixel_data[n:n+width] for n in range(0, width*height, width)]


# 计算区域内的平均颜色
def calculate_average_color(img_array,x1, y1, x2, y2, x3, y3, x4, y4):
    x1, y1, x2, y2, x3, y3, x4, y4=x1, y1, x2, y2, x3, y3, x4, y4
    total_color = [0, 0, 0, 0]
    count = 0
    for y in range(min(y1, y2, y3, y4), max(y1, y2, y3, y4)):
        for x in range(min(x1, x2, x3, x4), max(x1, x2, x3, x4)):
            total_color = [total_color[i] + img_array[y][x][i] for i in range(3)]
            count += 1
    return [total // count for total in total_color]

# 生成与给定颜色相差较大的颜色
def generate_distinct_colors(avg_color, num_colors,x1, y1, x2, y2, x3, y3, x4, y4):
    colors = []
    for i in range(num_colors):
        random_shift = random.randint(100, 200) + i * 15  # 这里可以调整以获取不同的颜色
        # 对RGB进行更改，保持alpha不变
        rgb = tuple((avg_color[j] + random_shift) % 256 for j in range(3))
        # 将原始的alpha添加到rgb中
        color = rgb 
        colors.append(color)
    return colors

#图像不同连通域填充不同颜色（根据读取字体原来的颜色，遍历每一个像素，不同连通域填充不同的颜色）
def color_regions(img_array, color, colors,x1, y1, x2, y2, x3, y3, x4, y4):
    x1, y1, x2, y2, x3, y3, x4, y4=x1, y1, x2, y2, x3, y3, x4, y4
    directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]
    marked = set()
    region_count = 0
    threshold = 0  # 设置一个阈值

    def color_distance(c1, c2):
        return ((c1[0] - c2[0]) ** 2 + (c1[1] - c2[1]) ** 2 + (c1[2] - c2[2]) ** 2) ** 0.5

    def dfs(x, y, new_color):
        stack = [(x, y)]
        while stack:
            x, y = stack.pop()
            # 判断像素位置是否在指定区域内
            if x < min(x1, x2, x3, x4) or x > max(x1, x2, x3, x4) or y < min(y1, y2, y3, y4) or y > max(y1, y2, y3, y4):
                continue
            if (x, y) in marked or color_distance(img_array[y][x], color) > threshold:
                continue
            marked.add((x, y))
            img_array[y][x] = new_color
            for dx, dy in directions:
                nx, ny = x + dx, y + dy
                if nx >= 0 and ny >= 0 and ny < len(img_array) and nx < len(img_array[0]):
                    stack.append((nx, ny))

    for y in range(len(img_array)):
        for x in range(len(img_array[0])):
            if color_distance(img_array[y][x], color) <= threshold and (x, y) not in marked:
                dfs(x, y, colors[region_count % len(colors)])
                region_count += 1

    return img_array

#---------因为上面函数不能解决字的轮框是一个渐变的颜色，看起来有框，所以这里在起外围的框也填充为字体颜色
def replace_color(img_array, target_color, replace_color, x1, y1, x2, y2, x3, y3, x4, y4):
    img_array = np.array(img_array)
    x1, y1, x2, y2, x3, y3, x4, y4=x1, y1, x2, y2, x3, y3, x4, y4
    a = np.zeros_like(img_array) 
   
    for y in range(y1, y3):
        for x in range(x1, x2):
            
            if np.array_equal(img_array[y][x], target_color):
                img_array[y][x] = replace_color
                
                if (a[y][x] == 1).all():
                    continue
                if y > y1 and not np.array_equal(img_array[y-1][x], target_color):
                 
                    img_array[y-1][x] = replace_color
                    a[y-1][x] = 1
                if y < y3 and y+1<688 and not np.array_equal(img_array[y+1][x], target_color) :
                    img_array[y+1][x] = replace_color
                    a[y+1][x] = 1
                if x > x1 and not np.array_equal(img_array[y][x-1], target_color):
                    img_array[y][x-1] = replace_color
                    a[y][x-1] = 1
                if x < x2 and x<1103 and not np.array_equal(img_array[y][x+1], target_color):
                    img_array[y][x+1] = replace_color
                    a[y][x+1] = 1
                    
    return img_array


class CreateData:
    def __init__(self,create_num):
        self.jay_img_paths=['/data/lh123/lh/verification_code/generate_data/背景图_最终20/' + i for i in os.listdir('/data/lh123/lh/verification_code/generate_data/背景图_最终20/')] # 背景图片路径
        self.img_save_path='/data/lh123/lh/verification_code/generate_data/train_2/' # 生成训练集图片的路径

        # 读取文件内容
        with open(file_path, "r", encoding="utf-8") as file:
            content = file.read()

        # 将单个字保存在列表中
        self.songs = list(content)
    
        self.song2label={song:i for i,song in enumerate(self.songs)}#得到标签，第一个
        self.label2song={i:song for i,song in enumerate(self.songs)}#得到内容
        self.create_num=create_num
        self.image_w=1104  #图片大小
        self.image_h=688#图片大小
        self.max_iou=0.01  # 每首歌名的boxes的iou不能超过0.1，#控制重叠
    #文件夹不存在就造一个，并且文件夹存在会初始化文件夹
    def create_folder(self):
        while True:
            try:
                for path in [self.img_save_path,self.label_save_path,self.test_path]:
                    shutil.rmtree(path,ignore_errors=True)
                    os.makedirs(path,exist_ok=True)
                break
            except:
                pass
    #计算重叠比例
    def bbox_iou(self,box2):
        '''两两计算iou'''
        for box1 in self.tmp_boxes_boxs1:
            inter_x1=max([box1[0],box2[0]])
            inter_y1=max([box1[1],box2[1]])
            inter_x2=min([box1[2],box2[2]])
            inter_y2=min([box1[3],box2[3]])
            inter_area=(inter_x2-inter_x1+1) * (inter_y2-inter_y1+1)
            box1_area=(box1[2]-box1[0]+1) * (box1[3]-box1[1]+1)
            box2_area=(box2[2]-box2[0]+1) * (box2[3]-box2[1]+1)
            iou=inter_area / (box1_area + box2_area - inter_area + 1e-16)
            if iou > self.max_iou:
                # 只要有一个与之的iou大于阈值则重新来过
                return iou
        else:
            return 0
#---------原来的
    def draw_text(self, image, image_draw, song,font_path):
        self.font_path=font_path
        iou = np.inf
        num=0
        while iou > self.max_iou:#判断是否重叠
            if num >= 3000:#最多循环3000防止无线循环，超多就会使得重叠
                break
            random_font_size = np.random.randint(110, 240)#字大小
            random_rotate = np.random.randint(-60, 60)#旋转角度
            random_x = np.random.randint(1, 1104, 1)#写字地址随机数
            random_y = np.random.randint(1, 688, 1)

            font = ImageFont.truetype(self.font_path, random_font_size)
            label = self.song2label[song]#第几个字
            size_wh = font.getsize(song)  #得到字的大小

            img = Image.new('L', size_wh)
            img_draw = ImageDraw.Draw(img)
            img_draw.text((0, 0), song, font=font, fill=255)#写字
            img_rotate = img.rotate(random_rotate, resample=2, expand=True)

            background_color = image.getpixel((int(random_x), int(random_y)))#背景颜色，这里是背景颜色像素随便找一个
            font_color = tuple((np.array(background_color) + np.array([128, 128, 128])) % 256)#得到与背景相差较大的颜色
            img_color = ImageOps.colorize(img_rotate, (0, 0, 0), font_color)
            
            w, h = img_color.size
            xmin = int(random_x)#左上角坐标
            ymin = int(random_y)

            if random_x + w > self.image_w:#判断越界
                xmin = self.image_w - w - 2
            if random_y + h > self.image_h:
                ymin = self.image_h - h - 2
            xmax = xmin + w#右下角坐标
            ymax = ymin + h
            a=rotate_rectangle((xmin, ymin), (xmax, ymax), random_rotate)#得到坐标
            boxes = (a[0], a[1], a[2], a[3],a[4],a[5],a[6],a[7])#整个框的坐标，左上角、右上角、右下角，左下角
            boxes1 = (xmin, ymin,xmax,ymax)
            #----判断重叠是否大
            iou = self.bbox_iou(boxes1)
            #-----判断字体和字是否匹配
            # 对于每个文件，初始化相应的字体对象
            
            fnt = ImageFont.truetype(self.font_path, 15)
           
        
            #检查字是否与字体匹配，如果不匹配就重新那字体
            if not fnt.getmask(song):
               
                #随机字体
                font_directory = '/data/lh123/lh/verification_code/generate_data/fonts'
                font_files = [f for f in os.listdir(font_directory) if f.endswith('.ttf') or f.endswith('.otf') or f.endswith('.ttc') or f.endswith('.TTF') or f.endswith('.OTF') or f.endswith('.TTC')]
                # font_files = [f for f in os.listdir(font_directory) if f.endswith('.otf') ]
                
                random_font_file = random.choice(font_files)
                self.font_path = os.path.join(font_directory, random_font_file)
                iou=1#这为1意思还得执行一下
            num += 1
           
        image.paste(img_color, box=(xmin, ymin), mask=img_rotate)#写字
        return image, boxes, label,boxes1,font_color,song


    def process(self,boxes):   #归一化处理----没用这个---这个是弄成比例
        '''
        将xmin,ymin,xmax,ymax转为x,y,w,h
        以及归一化坐标，生成label
        '''

        x1,y1,x2,y2=boxes
        x=((x1+x2)/2)/self.image_w
        y=((y1+y2)/2)/self.image_h
        w=(x2-x1)/self.image_w
        h=(y2-y1)/self.image_h
        return [x,y,w,h]
    def main(self):
        '''主函数'''
        self.create_folder() # 重置所需文件夹，如果多线程的话，需要删除这个代码
        txt_file='/data/lh123/lh/verification_code/generate_data/train_2.txt'
        with open(txt_file,'w') as f:#在外面速度更快，但是如果多线程会存在数据写串行的问题，但这种数据不多，可以进行数据清理
            num=1
            for i in tqdm(range(self.create_num)):
                self.font_color_list=[]
                random_song_num=np.random.randint(4,6) # 一张图片中4-5个字
                random_jay_img_path=np.random.choice(self.jay_img_paths) # 随机背景
                image=Image.open(random_jay_img_path).convert('RGB').resize((self.image_w,self.image_h))
                image_draw=ImageDraw.Draw(image)
                boxes_list=[]
                label_list=[]
                self.tmp_boxes=[] # 用于计算两两boxes的iou    
                self.tmp_boxes_boxs1=[] #
                self.song_list=[]
                for j in range(random_song_num):
                    song=np.random.choice(self.songs)
                    # song=self.songs[5998]
                    #随机字体
                    font_directory = '/data/lh123/lh/verification_code/generate_data/fonts'
                    font_files = [f for f in os.listdir(font_directory) if f.endswith('.ttf') or f.endswith('.otf') or f.endswith('.ttc') or f.endswith('.TTF') or f.endswith('.OTF') or f.endswith('.TTC')]
                    # font_files = [f for f in os.listdir(font_directory) if f.endswith('.otf') ]
                    random_font_file = random.choice(font_files)
                    self.font_path = os.path.join(font_directory, random_font_file)
    
                    image,boxes,label,boxes1,font_color,self.song=self.draw_text(image,image_draw,song,self.font_path)#图片，框，字体
                    self.font_color_list.append(font_color)
                  
                    self.tmp_boxes.append(boxes)
                    self.tmp_boxes_boxs1.append(boxes1)
                    self.song_list.append(song)
                    # boxes_list.append(self.process(boxes))#存储框对应的字
                    boxes_list.append(boxes)#存储框对应的字

                    label_list.append(label)
                    
                #图片名
                image_filename=self.img_save_path+f'image{num}.jpg' if i < self.create_num else self.test_path+f'test{i}.png'#保存文件
                label_filename=self.label_save_path+f'image{num}.txt' if i < self.create_num else self.test_path+f'test{i}.txt'
                num=num+1
                #输入坐标，图像位置，图像颜色，随机数（0-2)、image
                # 要求不同于图像的每一个颜色，并且存在部分差异，最后返回image 
                
                random_num = random.randint(0, 2) #一张图片中可能有0个、1、2个彩色字
                if(random_num!=0):
                    
                    for i in range(random_num):#主要实现连通域不同字不同
                        x1, y1, x2, y2, x3, y3, x4, y4 = boxes_list[i][0],boxes_list[i][1],boxes_list[i][2],boxes_list[i][3],boxes_list[i][4],boxes_list[i][5],boxes_list[i][6],boxes_list[i][7]
                        img_array = img_to_array(image,x1, y1, x2, y2, x3, y3, x4, y4)
                        # 计算平均颜色
                        avg_color = calculate_average_color(img_array,x1, y1, x2, y2, x3, y3, x4, y4)  
                        # 生成与平均颜色相差较大的颜色
                        colors = generate_distinct_colors(avg_color, 6,x1, y1, x2, y2, x3, y3, x4, y4)
                        

                        # 字的颜色
                        f_color=self.font_color_list[i]+(255,)
                        f_color_list = list(self.font_color_list[i])
                        f_color_nup = np.array(f_color_list)
                        #颜色预处理膨胀
                        img_array=replace_color(img_array, f_color_nup, f_color_nup,x1, y1, x2, y2, x3, y3, x4, y4)
                        img_array=replace_color(img_array, f_color_nup, f_color_nup,x1, y1, x2, y2, x3, y3, x4, y4)


                        # 查找和标记所有的连通区域，并改变每个区域的颜色
                        new_img_array = color_regions(img_array, f_color, colors,x1, y1, x2, y2, x3, y3, x4, y4)
                        image = Image.fromarray(np.uint8(new_img_array))
                    


                image.save(image_filename,format='JPEG')
                #写入便签文件txt文件内容
                f.write(f'labels/{image_filename}\t[')
                number=0
                for k in range(len(label_list)):
                    # label x y w h
                    # f.write(f'{self.song_list[k]} {boxes_list[k][0]} {boxes_list[k][1]} {boxes_list[k][2]} {boxes_list[k][3]} {boxes_list[k][4]} {boxes_list[k][5]} {boxes_list[k][6]} {boxes_list[k][7]}\n')
                    f.write(f'{{"transcription":"{self.song_list[k]}","points":[[{int(boxes_list[k][0])},{int(boxes_list[k][1])}],[{int(boxes_list[k][2])},{int(boxes_list[k][3])}],[{int(boxes_list[k][4])},{int(boxes_list[k][5])}],[{int(boxes_list[k][6])},{int(boxes_list[k][7])}]]}}')
                    if(number!=(len(label_list)-1)):
                        f.write(f',')
                    number=number+1
                f.write(f']\n')
                            
if __name__ == '__main__':
    creator=CreateData(500000)#这里是生成50万张的意思
    creator.main()