【Python实战】随机图像数据集自动生成（原创）

一 、数据集介绍

如上图所示，数据集由一系列随机生成的图像组成，每幅图像包含了40~50个大小相等的“T"图案与一个”L"图案，它们在图像中的位置、方向都是无规则的，因此可以作为目标检测、语义分割等任务的数据集。

二、编程思路

1. 工具：使用turtle库进行绘图（可任意控制画笔的位置、朝向），使用random库生成随机数（调整字母图案的位置和角度）
2. 思路：字母图案逐行生成，行内各字母横向间隔随机，每行字母初始纵向间隔随机，各字母在生成时纵坐标都有随机微调。
3. 结构：1）生成一个“T”的函数和生成一个“L”的函数；2）主绘图函数：嵌套循环，内层循完成一次生成一行图案，外层循环完成则生成全部图案；3）转换图像格式、保存图像的函数

三、完整代码

import os
from turtle import *
from random import *
from PIL import Image
from PIL import EpsImagePlugin
from math import *

def draw_T(x, y, angle,color):# draw a 'T'
    pencolor(color)
    pensize(10)
    pu()# pen up: only move, no drawing
    goto(x, y)
    pd()# pen down: draw while moving
    setheading(angle)# change beginning direction
    forward(40)
    backward(20)
    left(90)# turn left
    forward(35)

def draw_L(x, y, angle,color):# draw an 'L' and record its info
    # draw an 'L'
    pencolor(color)
    pensize(10)
    long_side = 35 # first stroke
    short_side = 20 # second stroke
    pu()
    goto(x, y)
    pd()
    setheading(angle)
    forward(long_side)
    left(90)
    forward(short_side)
    
    # record its info
    sin_angle = sin(angle*pi/180)
    cos_angle = cos(angle*pi/180)
    long_sin = long_side*sin_angle
    long_cos = long_side*cos_angle
    short_sin = short_side*sin_angle
    short_cos = short_side*cos_angle
    
    x2 = long_cos + x # L's second point
    y2 = long_sin + y
    x3 = long_cos - short_sin +x # L's third point
    y3 = long_sin + short_cos + y
    x_right = max(x, x2, x3) + 9 # point of bounding box
    x_left = min(x, x2, x3) - 9
    y_top = max(y, y2, y3) + 9
    y_bottom = min(y, y2, y3) - 9
    Tran = 670/224 # The original image is 670*670 and we need 224*224.
    x_r = 112 + x_right // Tran # central coordinate system --> upper left coordinate system
    x_l = 112 + x_left // Tran
    y_t = 112 - y_top // Tran
    y_b = 112 - y_bottom // Tran
    
    show_bb(x_left,x_right,y_top,y_bottom) # show bounding box on the image
    
    with open('dataset/L_info.txt','a') as f: # store info of bounding box
        f.write('%d,%d,%d,%d\n' % (x_l,y_t,x_r,y_b))

def show_bb(x_left,x_right,y_top,y_bottom): # show bounding box on the image
    if choice == 'Y':
        pencolor('red')
        pensize(2)
        pu()
        goto(x_left,y_top)
        pd()
        goto(x_right,y_top)
        goto(x_right,y_bottom)
        goto(x_left,y_bottom)
        goto(x_left,y_top)

def paint_frame(): # draw background and frame(640*640)
    setup(670,670,0,0)# canvas size(670,670) and location(0,0)
    ht()# hide turtle(symbol of pen)
    tracer(0, 0)# speed up generation
    pencolor('black')
    pensize(10)
    pu()
    goto(-320,-320)
    pd()
    goto(-320,320)
    goto(320,320)
    goto(320,-320)
    goto(-320,-320)

def paint_dataset():# Draw line by line; set random small fluctuations
    paint_frame()
    # draw many T and one L
    num = 0# loop count
    for i in range(7):# 7 columns
        x = randint(-370,-350)# central coordinate system
        y = randint(-300,-280) + i * 90
        for j in range(6):# 6 raws
            x += randint(70,120)
            y += randint(-10,10)
            angle = randint(0,360)
            if num != DRAW_L:# DRAW_L is a random number decides when to draw 'L'
                if -280<x<300 and -280<y<290:# Remove 'T' which may be out of bounds.
                    draw_T(x, y, angle, 'black')# T's color is blue(0,0,255)
            else:
                draw_L(x, y, angle, 'black')# L's color is red(255,0,0)
            num += 1
    
def mkdir(path):# create folder
	folder = os.path.exists(path)
	if not folder:
		os.makedirs(path)           

def save_img(flodername1,flodername2,filename):# eps-->png(502*502)-->png(224*224)
    ts = getscreen().getcanvas()
    ts.postscript(file="%s/%s.eps" % (flodername1,filename))# save image as eps
    img1 = Image.open("%s/%s.eps" % (flodername1,filename))
    img1.save('%s/%s.png' % (flodername1,filename))# eps --> png
    img2 = img1.resize((224,224),Image.ANTIALIAS)
    img2.save('%s/%s.png' % (flodername2,filename))# 502*502 --> 224*224

def create(filename):# once creates one image and its ground_truth
    global DRAW_L # random number decides when to draw 'L'
    DRAW_L = randint(7,35)
    paint_dataset()
    save_img('temporary','dataset',filename)
    clear() # clean the canvas
    print('Finished: %d'%(int(filename)+1))

def generate():
    global filename # e.g.'1','2','3'...
    global choice
    if(os.path.exists('dataset/L_info.txt')):
        os.remove('dataset/L_info.txt') # delete old info
    for folder_name in ['temporary','dataset']:
        mkdir(folder_name)
    print('Tip: The window appears later can be minimized.')
    NUM = input('Please input the quantity of the dataset:\n')
    choice = input('Show the bounding box on images --- input "Y"\nOr generate an usable dataset --- input "N"\n').upper()
    for filename in range(int(NUM)):
        create(str(filename))
        os.remove('temporary/%s.eps'%filename) # Files in 'temporary' are useless.
        os.remove('temporary/%s.png'%filename)

if __name__ == '__main__':
    generate()

四、总结

1. 收获
   1）turtle库、random库的使用
   2）IO编程基本运用
   3）eps文件转png格式方法