97 lines
3.4 KiB
Python
97 lines
3.4 KiB
Python
import random
|
|
import re
|
|
|
|
def clean_tag(tag):
|
|
return tag.replace("_", " ").strip()
|
|
|
|
def generate_associative_name(tags):
|
|
"""
|
|
Generate a 'real-alike' associative name based on WD tagger tags.
|
|
tags: list of dicts {'tag': str, 'score': float, 'cat': int}
|
|
"""
|
|
if not tags:
|
|
return None
|
|
|
|
# Filter by score
|
|
high_score_tags = [t for t in tags if t['score'] > 0.4]
|
|
if not high_score_tags:
|
|
high_score_tags = tags[:5]
|
|
|
|
# Categories: 0=general, 4=character
|
|
characters = [clean_tag(t['tag']) for t in high_score_tags if t['cat'] == 4]
|
|
general = [clean_tag(t['tag']) for t in high_score_tags if t['cat'] == 0]
|
|
|
|
# Key attributes
|
|
subject = None
|
|
if characters:
|
|
subject = characters[0].title()
|
|
elif "1girl" in [t['tag'] for t in high_score_tags]:
|
|
subject = "Maiden"
|
|
elif "1boy" in [t['tag'] for t in high_score_tags]:
|
|
subject = "Youth"
|
|
else:
|
|
subject = "Subject"
|
|
|
|
# Actions/Poses
|
|
actions = ["standing", "sitting", "lying", "running", "walking", "dancing", "sleeping"]
|
|
found_action = next((clean_tag(t['tag']) for t in high_score_tags if t['tag'] in actions), None)
|
|
|
|
# Setting/Background
|
|
settings = ["forest", "beach", "city", "space", "room", "garden", "ocean", "mountain", "sky", "underwater", "street"]
|
|
found_setting = next((clean_tag(t['tag']) for t in high_score_tags if t['tag'] in settings), None)
|
|
|
|
# Appearance
|
|
colors = ["red", "blue", "green", "white", "black", "gold", "silver", "pink", "purple", "yellow"]
|
|
found_color = next((clean_tag(t['tag']) for t in high_score_tags if clean_tag(t['tag']).split()[0] in colors), None)
|
|
if not found_color:
|
|
found_color = next((clean_tag(t['tag']) for t in high_score_tags if any(c in t['tag'] for c in colors)), None)
|
|
|
|
# Styles/Atmosphere
|
|
styles = ["cyberpunk", "fantasy", "realistic", "ethereal", "dark", "bright", "sketch", "oil painting"]
|
|
found_style = next((clean_tag(t['tag']) for t in high_score_tags if t['tag'] in styles), None)
|
|
|
|
# Build the name
|
|
templates = []
|
|
|
|
if found_style and subject:
|
|
templates.append(f"{found_style.title()} {subject}")
|
|
|
|
if found_color and subject:
|
|
templates.append(f"The {found_color.title()} {subject}")
|
|
|
|
if subject and found_action:
|
|
if found_action.endswith("ing"):
|
|
templates.append(f"{subject} {found_action.title()}")
|
|
else:
|
|
# Basic attempt at present participle
|
|
action_ing = found_action
|
|
if action_ing.endswith("e"):
|
|
action_ing = action_ing[:-1] + "ing"
|
|
else:
|
|
action_ing += "ing"
|
|
templates.append(f"{subject} {action_ing.title()}")
|
|
|
|
if subject and found_setting:
|
|
templates.append(f"{subject} in the {found_setting.title()}")
|
|
|
|
if found_style and found_setting:
|
|
templates.append(f"{found_style.title()} {found_setting.title()}")
|
|
|
|
if not templates:
|
|
# Fallback: combine two random general tags
|
|
if len(general) >= 2:
|
|
return f"{general[0].title()} {general[1].title()}"
|
|
elif general:
|
|
return general[0].title()
|
|
else:
|
|
return "Untitled Artwork"
|
|
|
|
# Return a random template result
|
|
return random.choice(templates)
|
|
|
|
|
|
def get_base_name(name: str) -> str:
|
|
"""Remove timestamp prefixes from filename."""
|
|
# Matches YYYYMMDD_HHMMSS_
|
|
return re.sub(r'^(\d{8}_\d{6}_)+', '', name)
|