generator.py

from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw 

import json
import os
import random
import csv


MAXLINECOUNT=3
FOLDERS = ("en_val","en_train_filtered")
USED_NAME = set()

def create_folders():
    os.makedirs('./all_data',exist_ok=True)
    os.makedirs('./all_data/en_train_filtered',exist_ok=True)
    os.makedirs('./all_data/en_val',exist_ok=True)


def create_csv():
     for foldername in FOLDERS:
        with open(f"all_data/{foldername}/labels.csv",'a') as csvfile:
            spamwriter = csv.writer(csvfile, delimiter=',', lineterminator="\r")
            spamwriter.writerow(["filename","words"])

def write_2_base(filename,line,foldername):
    with open(f"all_data/{foldername}/labels.csv",'a') as csvfile:
         spamwriter = csv.writer(csvfile, delimiter=',', lineterminator="\r")
         spamwriter.writerow([filename,f"{line}"])

with open("config.json") as f:
    config = json.load(f)

def random_font_color():
    return tuple(random.choice(config["color"][0]["font"]))

def random_background_color():
    return tuple(random.choice(config["color"][0]["background"]))

def random_font():
    return random.choice(config["font"])

def random_foldername():
    return random.choice([FOLDERS[0],FOLDERS[1],FOLDERS[1],FOLDERS[1],FOLDERS[1]])

def create_image(line,filename):
    global USED_NAME
    foldername = random_foldername()
    background_color = random_background_color()
    font_raw = random_font()

    font = ImageFont.truetype(font_raw[0],16)
    new_img = Image.new("RGB",(int(font.getlength(line)),21),background_color)
    pencil = ImageDraw.Draw(new_img)

    #pencil.text(font_raw[1],line,  font=font, fill=font_color)
    shift=0
    for word in line.split(' '):
        font_color = random_font_color()
        pencil.text((shift,font_raw[1][1]),word+' ',  font=font, fill=font_color)
        shift+=font.getlength(word+' ')

    #new_img.show()
    while (r:=random.randint(100000,999999)) in USED_NAME:
        pass
    else:
        USED_NAME.add(r)
        filename = filename[:-4]+'.'+str(r) + ".jpg"
        new_img.save(f"all_data/{foldername}/{filename}")
        return filename, foldername

def main():
    linecount=0
    create_folders()
    create_csv()
    for file in os.listdir("code"):
        
        filename = os.fsdecode(file)
        print(filename)

        with open(f"code/{filename}", encoding='utf-8') as f:
            print(linecount)
                
            for line in f.readlines() :
                line = line[:-1].strip().lstrip()
                if line not in ['','\n'] and line.count(' ')!=len(line):
                    for i in range(3):
                        if  linecount<MAXLINECOUNT:
                            image_name,foldername = create_image(line,filename)
                            write_2_base(image_name,line,foldername)
                            linecount+=1
                        else:
                            exit()
main()