generate script

generate.py

import os 
import random
from PIL import Image
import json
import csv

# Get the current working directory
current_directory = os.path.dirname(os.path.abspath(__file__))

# Set the paths based on the current working directory
OUT_DIR = os.path.join(current_directory, "out")
TRAITS_DIRECTORY = os.path.join(current_directory, "traits")

IMG_SIZE = (2048, 2048)

# Full path to the 'rarity.json' file
rarity_file_path = os.path.join(current_directory, 'rarity.json')

# Open the 'rarity.json' file
with open(rarity_file_path) as rarity_file:
    TRAITS = json.load(rarity_file)

# Define the TRAIT_NAMES
TRAIT_NAMES = [
    "Background",
    "Skin",
    "Clothes",
    "Blop",
    "Eyes",
    "Mouth",
]

trait_counts = {trait_name: {layer: count for layer, count in TRAITS[trait_name].items()} for trait_name in TRAIT_NAMES}

def generate_report():
    headers = ["NAME"]
    headers = headers + TRAIT_NAMES
    files = os.listdir(OUT_DIR)
    csv_data = []
    for file in files:
        if file.endswith(".json"):
            name = file.replace(".json", "")
            file = os.path.join(OUT_DIR, file)
            file = open(file)
            data = json.load(file)
            csv_data_line = [name]
            for trait_name in TRAIT_NAMES:
                for trait in data:
                    if trait_name.lower() == trait.get('trait').lower():
                        csv_data_line.append(trait.get('value'))
                        break
            csv_data.append(csv_data_line)
    with open('report.csv', 'w', encoding='UTF8', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(headers)
        writer.writerows(csv_data)

def pick_trait(trait_name, counts):
    trait = TRAITS.get(trait_name)
    available_layers = [layer for layer in trait.keys() if counts[trait_name][layer] > 0]
    if not available_layers:
        return None
    layer = random.choices(available_layers, weights=[trait[layer] for layer in available_layers])[0]
    counts[trait_name][layer] -= 1
    rarity = trait.get(layer)
    return {
        "layer" : layer,
        "rarity" : rarity,
        "directory" : trait_name,
        "attribute" : {
            "trait" : trait_name.lower(),
            "value" : layer.replace(".png", "").lower(),
            "rarity" : rarity
        }
    }

def generate_avatar2(seq_num):
    layers = []
    rarities = []
    directories = []
    attributes = []
    
    # Determine the final size for the avatar
    final_size = IMG_SIZE

    for trait_name in TRAIT_NAMES:
        trait = pick_trait(trait_name, trait_counts)
        layers.append(trait.get('layer'))
        rarities.append(trait.get('rarity'))
        directories.append(trait.get('directory'))
        attributes.append(trait.get('attribute'))

    image = Image.new("RGBA", final_size)
    for i in range(len(layers)):
        img_filename = os.path.join(TRAITS_DIRECTORY, directories[i], layers[i])
        layer_image = Image.open(img_filename)
        
        # Resize the layer image to the final size
        layer_image = layer_image.resize(final_size)
        
        # Check if the layer image has an alpha channel
        if layer_image.mode != 'RGBA':
            layer_image = layer_image.convert('RGBA')
        
        image = Image.alpha_composite(image, layer_image)

    nft_name = f"blops!_{str(seq_num+1).zfill(4)}"
    out_file = f"{nft_name}.png"
    attr_file = f"{nft_name}.json"

    image.save(os.path.join(OUT_DIR, out_file))

    with open(os.path.join(OUT_DIR, attr_file), "w") as outfile:
        json.dump(attributes, outfile, indent=4)

    print(f"Created {nft_name} with attributes {attributes}")

for i in range(1888):
    generate_avatar2(i)

generate_report()