生成关键点检测的仿真数据 二 ---PAF

import math
from math import fabs, sin, cos, radians
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageDraw
import cv2
import random
import os



'''根据列表点，在图中画折线'''
def draw_Polyline(img, point_list, color=(0, 0, 255)):
    
    pts=np.array(point_list,np.int32)
    pts=pts.reshape((-1,1,2))
  
    cv2.polylines(img,pts=[pts],isClosed=True, color=color, thickness=1)
    
    return img

'''根据列表点集合，在图中画填充的图像'''
def draw_fill_line(img, point_list_sum, color=(255, 255, 255)):
    Point_SUM = []
    for point_list in point_list_sum:
        point_list = np.array(point_list)
        Point_SUM.append(point_list)
    cv2.fillPoly(img, Point_SUM, (255, 255, 255))
    
    return img

'''根据坐标点和半径生成一组列表点，能形成固定形状的5变形（各向异性）'''
def draw_one_small_img(x, y, radius, angle):
    
    angle = angle/180 * math.pi #角度转换成弧度
    
    x1 = x + radius * sin(angle)
    y1 = y - radius * cos(angle)
    
    
    x2 = x + radius * cos(angle)
    y2 = y + radius * sin(angle)
    
    x3 = x - radius * sin(angle)
    y3 = y + radius * cos(angle)

    x4 = x - 0.5*radius * cos(angle)
    y4 = y - 0.5*radius * sin(angle)

    points = []
    points.append([int(x), int(y)])
    points.append([int(x1), int(y1)])
    points.append([int(x2), int(y2)])
    points.append([int(x3), int(y3)])
    points.append([int(x4), int(y4)])
   
    return points

#每次生成一张图
def draw_one_img(num, visual = False):
    cell_count = random.sample([0, 1, 2, 3], num)

    SCR_img = np.ones(shape=[512, 512, 3], dtype=np.uint8)
    SCR_img[:,:,:] = 0
    point_list_sum = []
    
    for count in cell_count:
        
        #生成随机角度
        angle = random.randint(0, 360) 
        
 
        #生成随机半径
        radius = random.randint(50, 120)
        
        #把1200的图像拆成九宫格，根据count的值选择在哪个格子里生成矩形
        X = (count % 2) * 256
        Y = (count // 2) * 256   
        
        x1 = random.randint(10, 246) + X
        y1 = random.randint(10, 246) + Y
        
        if(x1 + radius > X+256):
            x1 = X+256 - radius - 2
        if(x1 - radius < X):
            x1 = X + radius + 2
        if(y1 + radius > Y+256):
            y1 = Y+256 - radius -2
        if(y1 - radius < Y):
            y1 = Y + radius + 2
        
        points = draw_one_small_img(x1, y1, radius, angle)
        
        point_list_sum.append(points)
    
    #画白色填充的五边形图
    result_img = draw_fill_line(SCR_img, point_list_sum) 
    
    #画红色折线图可视化
    if(visual == True):
        for point_list in point_list_sum:
            result_img = draw_Polyline(result_img, point_list)
        
    return result_img, point_list_sum 

if __name__ == '__main__':
    
    path = "./train/"

    for i in range(100):
        img_save_path = path + str(i) + ".png"
        label_save_path = path + str(i) + ".txt"
        
        n = random.sample([1, 2, 3, 4], 1)
    
        img, point_list_sum = draw_one_img(n[0], visual = False)
        
        #保存图片
        cv2.imwrite(img_save_path, img)
        #保存标签
        with open(label_save_path,"w") as f: #设置文件对象
            temp_list = []
            for i in range(len(point_list_sum)):
                if(i == 0):
                    temp_list.append(str(0))
                else:
                    temp_list.append('\n'+str(i))
                for j in range(len(point_list_sum[i])):
                    temp_list.append('\n'+str(point_list_sum[i][j][0]) + ' ' + str(point_list_sum[i][j][1]) + ' ' + str(2))          
            f.writelines(temp_list)
    
#    width = 512
#    height = 512
#    img = np.ones(shape=[width,height,3], dtype=np.uint8)
#    
#    points = draw_one_small_img(265, 265, 50 , -45)
#    
#
#    draw_fill_line(img, [points])
#
#    
#    cv2.imwrite("123.png", img)
#    

#    img = np.zeros((1080, 1920, 3), np.uint8)
#    area1 = np.array([[250, 200], [300, 100], [750, 800], [100, 1000]])
#    area2 = np.array([[1000, 200], [1500, 200], [1500, 400], [1000, 400]])
#     
#    cv2.fillPoly(img, [area1], (255, 255, 255))
#     
#    plt.imshow(img)
#    plt.show()