rework country logic

This commit is contained in:
Karl Hudgell 2024-12-13 17:20:49 +00:00
parent 3f19c46c54
commit 00f4fc7b9b
2 changed files with 41 additions and 113 deletions

View File

@ -20,7 +20,7 @@ from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
logging.config.dictConfig(LOGGING_CONFIG) logging.config.dictConfig(LOGGING_CONFIG)
cut = 1000 cut = 500
# Load user configurations # Load user configurations
user_config = configparser.ConfigParser() user_config = configparser.ConfigParser()
@ -63,14 +63,16 @@ def generate_image(uid, comfy_prompt, model, steps):
logging.error(f"Failed to generate image for UID: {uid}. Error: {e}") logging.error(f"Failed to generate image for UID: {uid}. Error: {e}")
def get_country_name_from_code(code): def get_country_name(app_config, country_code):
"""Get country name from 3-letter ISO code.""" # First check if it's a custom mapping
try: if country_code in app_config["facial_characteristics"]:
country = pycountry.countries.get(alpha_3=code.upper()) return app_config["facial_characteristics"][country_code]
return country.name if country else "Unknown country code"
except KeyError: # Use pycountry for standard country codes
logging.warning(f"Invalid country code provided: {code}") country = pycountry.countries.get(alpha_3=country_code)
return "Invalid country code" if country:
return country.name
return "Unknown Country"
def generate_prompts_for_players(players, app_config): def generate_prompts_for_players(players, app_config):
@ -81,7 +83,7 @@ def generate_prompts_for_players(players, app_config):
logging.debug(f"Generating prompt for {player[0]} - {player[8]}") logging.debug(f"Generating prompt for {player[0]} - {player[8]}")
os.makedirs(output_folder, exist_ok=True) os.makedirs(output_folder, exist_ok=True)
country = get_country_name_from_code(player[1]) country = get_country_name(app_config, player[1])
facial_characteristics = random.choice(app_config["facial_characteristics"]) facial_characteristics = random.choice(app_config["facial_characteristics"])
hair_length = app_config["hair_length"][player[5]] hair_length = app_config["hair_length"][player[5]]
hair_colour = app_config["hair_color"][player[6]] hair_colour = app_config["hair_color"][player[6]]
@ -95,7 +97,7 @@ def generate_prompts_for_players(players, app_config):
age=player_age, age=player_age,
country=country, country=country,
facial_characteristics=facial_characteristics or "no facial hair", facial_characteristics=facial_characteristics or "no facial hair",
hair=f"{hair_length} {hair_colour} {hair}", hair=f"{hair_length} {hair_colour}",
) )
logging.debug(f"Generated prompt: {prompt}") logging.debug(f"Generated prompt: {prompt}")
prompt = f"{player[0]}:{prompt}" prompt = f"{player[0]}:{prompt}"
@ -152,8 +154,9 @@ def main():
) )
try: try:
remove_bg_from_files_in_dir(output_folder)
resize_images(output_folder) resize_images(output_folder)
remove_bg_from_files_in_dir(output_folder)
create_config_xml(output_folder) create_config_xml(output_folder)
logging.debug("Post-processing complete.") logging.debug("Post-processing complete.")
except Exception as e: except Exception as e:

View File

@ -1,103 +1,29 @@
{ {
"countries": { "custom_country_codes": {
"African": [ "WAL": "Wales",
"Nigeria", "DEN": "Denmark",
"Kenya", "POR": "Portugal",
"South Africa", "NED": "Netherlands",
"Egypt", "SCO": "Scotland",
"Ghana" "GER": "Germany",
], "CRO": "Croatia",
"Asian": [ "RSA": "South Africa",
"China", "NIR": "Northern Ireland",
"Japan", "GRN": "Grenada",
"India", "URU": "Uruguay",
"South Korea", "SUD": "Sudan",
"Vietnam" "PAR": "Paraguay",
], "TAN": "Tanzania",
"Caucasian": [ "KOS": "Kosovo",
"United States", "GRE": "Greece",
"Canada", "CHI": "Chile",
"Australia", "HAI": "Haiti",
"New Zealand", "GAM": "Gambia",
"United Kingdom" "SKN": "Saint Kitts and Nevis",
], "ANG": "Angola",
"Central European": [ "KSA": "Saudi Arabia",
"Germany", "SUI": "Switzerland",
"Poland", "BUL": "Bulgaria"
"Czech Republic",
"Austria",
"Hungary"
],
"EECA": [
"Ukraine",
"Kazakhstan",
"Belarus",
"Georgia",
"Azerbaijan"
],
"Italmed": [
"Italy",
"Greece",
"Croatia",
"Malta",
"Cyprus"
],
"MENA": [
"Saudi Arabia",
"UAE",
"Iran",
"Israel",
"Morocco"
],
"MESA": [
"Pakistan",
"Afghanistan",
"Bangladesh",
"Nepal",
"Sri Lanka"
],
"SAMed": [
"Tunisia",
"Algeria",
"Libya",
"Lebanon",
"Jordan"
],
"Scandinavian": [
"Sweden",
"Norway",
"Denmark",
"Finland",
"Iceland"
],
"Seasian": [
"Indonesia",
"Malaysia",
"Thailand",
"Philippines",
"Singapore"
],
"South American": [
"Brazil",
"Argentina",
"Colombia",
"Peru",
"Chile"
],
"SpanMed": [
"Spain",
"Portugal",
"Andorra",
"Gibraltar",
"Monaco"
],
"YugoGreek": [
"Serbia",
"Greece",
"North Macedonia",
"Montenegro",
"Albania"
]
}, },
"facial_characteristics": [ "facial_characteristics": [
"beard", "beard",
@ -154,8 +80,7 @@
"16": "Black", "16": "Black",
"17": "Platinum", "17": "Platinum",
"18": "Grey" "18": "Grey"
} },
,
"hair_length": { "hair_length": {
"0": "Bald", "0": "Bald",
"1": "Short", "1": "Short",