【点选验证码】生成点选验证码图片--多进程

import os
from concurrent.futures import ThreadPoolExecutor   #定义了一个线程池
from multiprocessing import Pool
#---------------------进程

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/ppocr_keys_v1.txt"
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,file_num):#这个和文件名字有关
        # self.jay_img_paths=['/data/usr/lh123/lh/verification_code/背景图_最终20/')] # 背景图片路径
        self.jay_img_paths=['/data/lh123/lh/verification_code/generate_data/点选文字背景_压缩/' + i for i in os.listdir('/data/lh123/lh/verification_code/generate_data/点选文字背景_压缩/')] # 背景图片路径
        self.file_num=file_num
        # self.font_path='/data/lh123/lh/点选字生成/generate_data/simfang.ttf' # 字体路径
        self.img_save_path='/data/lh123/lh/verification_code/generate_data/trian_多/' # 生成训练集图片的路径
        self.label_save_path='/data/lh123/lh/verification_code/generate_data/labels/' # 生成图片对应label的路径
        self.test_path='/data/lh123/lh/verification_code/generate_data/test/' # 生成测试集图片的路径
        #字体随机选择

        # font_directory = '/data/lh123/lh/点选字生成/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')]
        # random_font_file = random.choice(font_files)
        # self.font_path = os.path.join(font_directory, random_font_file)
        # 100首周杰伦歌曲名称
        # file_path = "/data/usr/lh123/lh/verification_code/ppocr_keys_v1.txt"

        # 读取文件内容
        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=1000
        self.image_w=1104
        self.image_h=688
        self.max_iou=0.01  # 每首歌名的boxes的iou不能超过0.5，#控制重叠
        
    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:
                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
            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() # 重置所需文件夹
        
        # 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) # 随机1~4首
            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)
                
            # save image and label
            image_filename=self.img_save_path+f'image{self.file_num*1000+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) 
            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')
            #写入内容
            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')
                            

def im_process(file_num):#接收一个参数args，它应该是一个包含图像路径和尺寸的元组
    # path, size = args#解析为两个变量
    
    creator=CreateData(file_num)
    creator.main()
    # return im_path  #函数执行完后，返回处理后的图像路径im_path


with open("process.pid", "w") as ij:
    ij.write(str(os.getpid()))


  #将当前进程的id号写入该文件
#最大线程数量max_workers，min最终会有一个线程池
my_list = list(range(0, 501))#这里是500组每组生成1000张图片

txt_file='/data/lh123/lh/verification_code/generate_data/train_多.txt'
with open(txt_file,'w') as f:
    with Pool(processes=os.cpu_count()) as t:#会算一下CPU的核心数，有几个核心就创造多少个线程池
        results = t.map(#在这里的results是返回的列表
            im_process,
            my_list#列表，元组，传入两个参数,要生成50万张图片就要传入50万长度的列表，把参数都传入到里面
        )#使用map应用于多个图像参数，im_process函数会被异步地在多个线程中同时执行，传入的参数是一个包含两个元组的列表，每个元组包含图像的路径和尺寸。