Das deutsche Python-Forum

Code: Alles auswählen

from os import path, mkdir
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import pandas as pd
import numpy as np

output_folder = "ABC_NFT"
if not path.exists(output_folder):
    mkdir(output_folder)

backgrounds = ["background_blackestnight", "background_bluenight", "background_cloudyblue", "background_cloudygrey", "background_largemoon", "background_largesun", "background_lightning", "background_snow_grey", "background_snowblue", "background_sunrise", "background_sunset", "background_meteor"]
characters = ["alle verschiedenen characters"]
mouths = ["mouth_happy1", "mouth_happy2", "mouth_happy3", "mouth_sad", "mouth_scared", "mouth_shout1", "mouth_shout2", "mouth_shout3", "mouth_smile", "mouth_smirk1", "mouth_smirk2", "mouth_smirk3", "mouth_teeth"]
brows = ["brows_angry1", "brows_angry2", "brows_angry3", "brows_neutral1", "brows_neutral2", "brows_neutral3", "brows_oh1", "brows_oh2", "brows_oh3", "brows_sus1", "brows_sus2", "brows_sus3"]
eyes = ["eyes_straight", "eyes_left", "eyes_right", "eyes_down", "eyes_angry", "eyes_angry_right", "eyes_angry_left", "eyes_sad", "eyes_sad_left", "eyes_sad_right", "eyes_sus", "eyes_sus_left", "eyes_sus_right", "eyes_red", "eyes_blue", "eyes_brown", "eyes_cash", "eyes_gold", "eyes_laser_blue"]
file_name = ["test"]

def generate_image(background, character, mouth, brow, eye, file_name):
    """Generate image with given background, given character and given object and save it with the given file name
    Args:
        background (str): background name
        character (str): character name
        mouth (str): mouth name
        brow (str): brow name
        eye (str): eye name
        file_name (str): file name
    """
    background_file = path.join("backgrounds", f"{background}.png")
    background_image = Image.open(background_file)

    #Create character
    character_file = path.join("characters", f"{character}.png")
    character_image = Image.open(character_file)

    background_image.paste(character_image, mask=character_image)

    #Create mouth
    mouth_file = path.join("mouths", f"{mouth}.png")
    mouth_image = Image.open(mouth_file)

    background_image.paste(mouth_image, mask=mouth_image)

    #Create brow
    brow_file = path.join("brows", f"{brow}.png")
    brow_image = Image.open(brow_file)

    background_image.paste(brow_image, mask=brow_image)

     #Create eye
    eye_file = path.join("eyes", f"{eye}.png")
    eye_image = Image.open(eye_file)

    background_image.paste(eye_image, mask=eye_image)
    
    output_file = path.join(output_folder, f"{file_name}.jpg")
    background_image.convert("RGB").save(output_file)

def generate_random_imgs(total_imgs):
    df = pd.DataFrame(columns = ["background", "character", "mouth", "brow", "eye", "file_name"])
    for num in range(total_imgs):
        background = np.random.choice(np.arange(0,len(backgrounds)), p=[0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.05, 0.03, 0.01, 0.01])
        background = backgrounds[background]
        
        character = np.random.choice(np.arange(0,len(characters)), p=[0.2, 0.2, 0.2, 0.2, 0.05, 0.05, 0.025, 0.025, 0.025, 0.025])
        character = characters[character]
        
        mouth = np.random.choice(np.arange(0,len(mouths)), p=[0.2, 0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.04, 0.005, 0.005])
        mouth = mouths[mouth]
        
        brow = np.random.choice(np.arange(0,len(brows)), p=[0.20, 0.10, 0.10, 0.10, 0.10, 0.10, 0.10, 0.05, 0.05, 0.05, 0.025, 0.025])
        brow = brows[brow]
        
        eye = np.random.choice(np.arange(0,len(eyes)), p=[0.10, 0.10, 0.10, 0.10, 0.10, 0.09, 0.05, 0.05, 0.04, 0.05, 0.03, 0.03, 0.03, 0.03, 0.03, 0.03, 0.02, 0.01, 0.01])
        eye = eyes[eye]

        generate_image(background, character, mouth, brow, eye, f"ABC_NFT{num}")
        data = [background, character, mouth, brow, eye, f"ABC_NFT{num}"]
        s = pd.Series(data, index=df.columns)
        df = df.append(s, ignore_index=True)
        
    df.to_csv('data.csv', index=False)

if __name__ == "__main__":
    #generate_all_imgs()
    generate_random_imgs(200)