nft-gen-1.py

class GeneratorController:

    def __init__(self, base_dir, img_dir, meta_dir, weights_map, traits_order,
                 percent_appearance=None, matching_traits=None):
        self.base_dir = base_dir
        self.img_dir = img_dir
        self.meta_dir = meta_dir
        self.weights_map = weights_map
        self.traits_order = traits_order
        self.percent_appearance = percent_appearance or {}
        self.matching_traits = matching_traits or {}

    def generate_all_images(self, total_count, start=0):
        i = start
        unique_names = set()
        while i < start + total_count:
            metadata = {}
            unique_name = None
            print(f'Generating image {i}')
            for trait_type in self.traits_order:
                print(f'Trait {trait_type}')
                if trait_type in self.percent_appearance:
                    if self.should_skip(trait_type):
                        continue
                if trait_type in self.matching_traits:
                    category_id, trait_val = self.pick_matching_trait(trait_type, metadata)
                else:
                    category_id = self.get_random_category()
                    trait_val = self.gen_random_trait_by_category(category_id, trait_type)
                if not unique_name:
                    unique_name = trait_val
                else:
                    unique_name += f'-{trait_val}'
                metadata[trait_type] = {
                    'value': trait_val,
                    'category': category_id
                }
            if unique_name in unique_names:
                continue
            else:
                i += 1
                unique_names.add(unique_name)
            image_name = f'{i}.png'
            result_img = self.create_image(metadata)
            result_img.save(f'{self.img_dir}/{image_name}')
            self.store_metadata(i, metadata)