diff --git a/.gitignore b/.gitignore index ae1cc6e..bfdfd70 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,5 @@ venv/* config.cfg generated_images/* script.log -**/*.pyc \ No newline at end of file +**/*.pyc +*.rtf \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json index bc61696..7146044 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -13,7 +13,7 @@ "justMyCode": false, "args": [ "--num_inference_steps", - "6", + "10", "--rtf_file", "./NewGen.rtf" ] diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index 5fb2b3f..96c9d92 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -126,21 +126,12 @@ def main(): # Parse the RTF file try: - rtf_file = rtf.parse_rtf(args.rtf_file)[:10] + rtf_file = rtf.parse_rtf(args.rtf_file)[:100] logging.info(f"Parsed RTF file successfully. Found {len(rtf_file)} players.") except FileNotFoundError: logging.error(f"RTF file not found: {args.rtf_file}") sys.exit(1) - - # Extract unique values - try: - hair_length = list(set(item[5] for item in rtf_file)) - hair_colour = list(set(item[6] for item in rtf_file)) - skin_tone = list(set(item[7] for item in rtf_file)) - except IndexError as e: - logging.error(f"Error extracting player attributes: {e}") - sys.exit(1) - + # Load configurations try: with open("config.json", "r") as f: diff --git a/config.json b/config.json index dde05b3..5459879 100644 --- a/config.json +++ b/config.json @@ -123,73 +123,45 @@ "ponytail" ], "skin_tone_map": { - "0": "Lightest", - "1": "Lightest", - "2": "Very Light", - "3": "Light", - "4": "Light with Beige Undertones", - "5": "Beige", - "6": "Light Tan", - "7": "Tan", + "0": "Fair", + "1": "Light Beige", + "2": "Medium Brown", + "3": "Dark Brown", + "4": "Deep Brown", + "5": "Light Tan", + "6": "Dark Tan", + "7": "Pale", "8": "Medium Tan", - "9": "Olive", - "10": "Light Brown", - "11": "Brown", - "12": "Medium Brown", - "13": "Dark Brown", - "14": "Very Dark Brown", - "15": "Deep Brown", - "16": "Light Ebony", - "17": "Ebony", - "18": "Dark Ebony", - "19": "Deep Ebony", - "20": "Darkest" + "9": "Black", + "10": "East Asian" }, "hair_color": { - "0": "Jet Black", - "1": "Jet Black", - "2": "Dark Black", - "3": "Very Dark Brown", - "4": "Dark Brown", - "5": "Medium Brown", - "6": "Light Brown", - "7": "Ash Brown", - "8": "Brown with Red Tint", - "9": "Auburn", - "10": "Chestnut", - "11": "Dark Blonde", - "12": "Medium Blonde", - "13": "Ash Blonde", - "14": "Light Blonde", - "15": "Platinum Blonde", - "16": "Strawberry Blonde", - "17": "Light Ginger", - "18": "Ginger", - "19": "Bright Red", - "20": "Copper Red" - }, + "1": "Blonde", + "2": "Strawberry Blonde", + "3": "Light Brown", + "4": "Bright Orange", + "5": "Black and Grey", + "6": "White", + "7": "Yellow", + "8": "Dark Blonde", + "9": "Light Brown", + "10": "Brown", + "11": "Light Brown", + "12": "Black", + "13": "Orange", + "14": "Ginger", + "15": "Black", + "16": "Black", + "17": "Platinum", + "18": "Grey" + } + , "hair_length": { "0": "Bald", - "1": "Bald", - "2": "Very Short (buzz cut)", - "3": "Short (crew cut)", - "4": "Short with fade", - "5": "Classic Short", - "6": "Medium-Short", - "7": "Medium (textured)", - "8": "Medium (layered)", - "9": "Medium with fringe", - "10": "Medium-Long (shaggy)", - "11": "Long (shoulder-length)", - "12": "Long with parting", - "13": "Long (tied back)", - "14": "Very Long (untied)", - "15": "Curly, Short", - "16": "Curly, Medium", - "17": "Curly, Long", - "18": "Afro-Short", - "19": "Afro-Medium", - "20": "Afro-Long" + "1": "Short", + "2": "Medium", + "3": "Long", + "4": "Bald" }, - "prompt": "Ultra-realistic close-up headshot of a {skin_tone} male soccer player with a white background looking at the camera. The player is {age} years old, from {country}, with {facial_characteristics} and {hair} hair. He is facing the camera with a confident expression, wearing a soccer jersey. The lighting is natural and soft, emphasizing facial features and skin texture" + "prompt": "Ultra-realistic close-up headshot of a {skin_tone} skinned male soccer player with a white background looking at the camera. The player is {age} years old, from {country}, with {facial_characteristics} and {hair} hair. He is facing the camera with a confident expression, wearing a soccer jersey. The lighting is natural and soft, emphasizing facial features and skin texture" } \ No newline at end of file diff --git a/lib/remove_bg.py b/lib/remove_bg.py index b48db1d..d8bd4f0 100644 --- a/lib/remove_bg.py +++ b/lib/remove_bg.py @@ -2,40 +2,60 @@ import os from rembg import remove from PIL import Image from tqdm import tqdm +from concurrent.futures import ThreadPoolExecutor -def remove_bg_from_files_in_dir(directory): +def process_image(input_path, output_path): """ - Process all JPG and JPEG images in the given directory and its subfolders. + Process a single image: remove its background and save the result. + + Args: + input_path (str): Path to the input image. + output_path (str): Path to save the processed image. + + Returns: + bool: True if the image was successfully processed, False otherwise. + """ + try: + with Image.open(input_path) as img: + output = remove(img) + output.save(output_path) + return True + except Exception as e: + print(f"Error processing {input_path}: {str(e)}") + return False + +def remove_bg_from_files_in_dir(directory, max_workers=2): + """ + Process all JPG, JPEG, and PNG images in the given directory and its subfolders using parallel processing. Args: directory (str): Path to the directory containing images. - + max_workers (int): Maximum number of threads to use for parallel processing. + Returns: int: The number of images successfully processed. """ processed_count = 0 - - # Get the total number of files to process - total_files = sum(len(files) for _, _, files in os.walk(directory)) - - # Create a progress bar - with tqdm(total=total_files, desc="Processing images", unit="image") as pbar: - for subdir, dirs, files in os.walk(directory): - for file in files: - if file.lower().endswith(('.jpg', '.jpeg', 'png')): - input_path = os.path.join(subdir, file) - output_filename = os.path.splitext(file)[0] + '.png' - output_path = os.path.join(subdir, output_filename) - - try: - with Image.open(input_path) as img: - output = remove(img) - output.save(output_path) - processed_count += 1 - except Exception as e: - print(f"Error processing {input_path}: {str(e)}") - - # Update the progress bar + files_to_process = [] + + # Gather all the image files to process + for subdir, dirs, files in os.walk(directory): + for file in files: + if file.lower().endswith(('.jpg', '.jpeg', 'png')): + input_path = os.path.join(subdir, file) + output_filename = os.path.splitext(file)[0] + '.png' + output_path = os.path.join(subdir, output_filename) + files_to_process.append((input_path, output_path)) + + # Use ThreadPoolExecutor to process images in parallel + with ThreadPoolExecutor(max_workers=max_workers) as executor: + # Use tqdm to show progress + with tqdm(total=len(files_to_process), desc="Processing images", unit="image") as pbar: + futures = {executor.submit(process_image, input_path, output_path): (input_path, output_path) for input_path, output_path in files_to_process} + + for future in futures: + if future.result(): + processed_count += 1 pbar.update(1) - - return processed_count \ No newline at end of file + + return processed_count diff --git a/mappings.txt b/mappings.txt new file mode 100644 index 0000000..3544c6a --- /dev/null +++ b/mappings.txt @@ -0,0 +1,38 @@ +Hair Lengths with examples (sorted): +1: [{'UID': '2002099259', 'Name': 'Samba Thiandoum'}, {'UID': '2002086544', 'Name': 'Youssef Ramzy'}, {'UID': '2002088012', 'Name': 'William Mukendi'}] SHORT +2: [{'UID': '2002069662', 'Name': 'Ibrahima Touré'}, {'UID': '2002105370', 'Name': 'Marcelo Fernando'}, {'UID': '2000210457', 'Name': 'Goran Begic'}] MEDIUM +3: [{'UID': '2002067856', 'Name': 'Jurdany Straver'}, {'UID': '2002069623', 'Name': 'Aziz Hegab'}, {'UID': '2000038705', 'Name': 'Óscar Palacios'}] LONG +4: [{'UID': '2002099193', 'Name': 'Andrea Marcato'}, {'UID': '2002099411', 'Name': 'Massimo Todaro'}, {'UID': '2000234613', 'Name': 'Marcello Saia'}] BALD + +Hair Colours with examples (sorted): +1: [{'UID': '2000187083', 'Name': 'Marvin Wattiau'}, {'UID': '2000055582', 'Name': 'Håkon Soltvedt'}, {'UID': '19361843', 'Name': 'Petar Andersson'}] BLONDE +2: [{'UID': '2002067750', 'Name': 'Silvio Kovacevic'}, {'UID': '2002069505', 'Name': 'Steve Hubbard'}, {'UID': '2002069418', 'Name': 'Matej Kovač'}] strawberry blonde +3: [{'UID': '2002068177', 'Name': 'Lorenzo Granelli'}, {'UID': '2002096503', 'Name': 'Christian Williams'}, {'UID': '2002096251', 'Name': 'Paulo Cáceres'}] light brown +4: [{'UID': '2002069357', 'Name': 'Dennis Dailly'}, {'UID': '2000199177', 'Name': 'Ian Warner'}, {'UID': '2002069790', 'Name': 'Killian Lenormand'}] bright orange +5: [{'UID': '2002099259', 'Name': 'Samba Thiandoum'}, {'UID': '2002069662', 'Name': 'Ibrahima Touré'}, {'UID': '2002105370', 'Name': 'Marcelo Fernando'}] black and grey +6: [{'UID': '2002121630', 'Name': 'Thiago Tourinho'}, {'UID': '2002085736', 'Name': 'Vinicius'}, {'UID': '2002098155', 'Name': 'Roberto Carlos'}] White +7: [{'UID': '2002069774', 'Name': 'Tómas Sigurðarson'}, {'UID': '2002068148', 'Name': 'Gjergji Jusufi'}, {'UID': '2000148350', 'Name': 'John Ball'}] yellow +8: [{'UID': '2002068153', 'Name': 'Jody Combe'}, {'UID': '2002100076', 'Name': 'Leutrim Ismaili'}, {'UID': '2002068068', 'Name': 'Bernd Heller'}] dark blonde +9: [{'UID': '2000104880', 'Name': 'Oliver Rhodes'}, {'UID': '2002068156', 'Name': 'Jeremy Sutton'}, {'UID': '2002068447', 'Name': 'Oliver Brown'}] light brown +10: [{'UID': '2000249120', 'Name': 'Daniel Williams'}, {'UID': '2002098384', 'Name': 'Kiril Vasilev'}, {'UID': '2000269855', 'Name': 'Aldin Cetin'}] BROWN +11: [{'UID': '2002069421', 'Name': 'Peter Pasfield'}, {'UID': '2002069674', 'Name': 'Fernando Cantera'}, {'UID': '2002098332', 'Name': 'Edson Garza'}] Light Brown +12: [{'UID': '2000210457', 'Name': 'Goran Begic'}, {'UID': '29226266', 'Name': 'Jens Waidner'}, {'UID': '2002068381', 'Name': 'Jonathan Neiland'}] Black +13: [{'UID': '2002098622', 'Name': 'Mikkel Rasmussen'}, {'UID': '2002069251', 'Name': 'Walter Pérez'}, {'UID': '2002069890', 'Name': 'Phil Foster'}] Orange +14: [{'UID': '2002100040', 'Name': 'Luka Jovanović'}, {'UID': '2002068194', 'Name': 'Alex Waldron'}, {'UID': '14250947', 'Name': 'Kylian Langlois'}] Ginger +15: [{'UID': '2002086544', 'Name': 'Youssef Ramzy'}, {'UID': '2002088012', 'Name': 'William Mukendi'}, {'UID': '2002069623', 'Name': 'Aziz Hegab'}] black +16: [{'UID': '2002067856', 'Name': 'Jurdany Straver'}, {'UID': '2002068557', 'Name': 'Patrick Coley'}, {'UID': '2002096137', 'Name': 'Francieudo'}] black +17: [{'UID': '2002086195', 'Name': 'Ramon'}, {'UID': '2002109810', 'Name': 'Jaime Cousillas'}] Platinum +18: [{'UID': '2002096495', 'Name': 'Aridai'}] Grey + +Skin Tones with examples (sorted): +0: [{'UID': '2002067856', 'Name': 'Jurdany Straver'}, {'UID': '2000187083', 'Name': 'Marvin Wattiau'}, {'UID': '2000210457', 'Name': 'Goran Begic'}] WHITE +1: [{'UID': '2002105370', 'Name': 'Marcelo Fernando'}, {'UID': '2002096137', 'Name': 'Francieudo'}, {'UID': '2000038705', 'Name': 'Óscar Palacios'}] Beige +2: [{'UID': '2002086544', 'Name': 'Youssef Ramzy'}, {'UID': '2002069623', 'Name': 'Aziz Hegab'}, {'UID': '2002069687', 'Name': 'Rafik El Hilali'}] Brown +3: [{'UID': '2002099259', 'Name': 'Samba Thiandoum'}, {'UID': '2002069662', 'Name': 'Ibrahima Touré'}, {'UID': '2002088012', 'Name': 'William Mukendi'}] black +4: [{'UID': '2002099204', 'Name': 'Agnelo Ghosh'}, {'UID': '2002068758', 'Name': 'Dipu Singh'}, {'UID': '2002068376', 'Name': 'Jason Benjamin'}] light black +5: [{'UID': '2002069522', 'Name': 'Valgeir Haraldsson'}, {'UID': '2002068519', 'Name': 'Supradit Jantanarm'}, {'UID': '19387456', 'Name': 'Thanik Pathumchai'}] light brown +6: [{'UID': '2002068324', 'Name': 'Timoci Pillay'}, {'UID': '2002105665', 'Name': 'Rusiate Nair'}, {'UID': '2000144679', 'Name': 'Brian Peterson'}] dark brown +7: [{'UID': '2002096094', 'Name': 'Banderley Peña'}, {'UID': '2002105644', 'Name': 'Johan García'}, {'UID': '2002069019', 'Name': 'Dante Acevedo'}] light white +8: [{'UID': '2002105661', 'Name': 'José Matamalai'}, {'UID': '2002121132', 'Name': 'Gaspar Santos'}, {'UID': '2002134581', 'Name': 'Jeremy Goldrick'}] brown +9: [{'UID': '2002068557', 'Name': 'Patrick Coley'}, {'UID': '2002088009', 'Name': 'Abel Díez'}, {'UID': '2002096074', 'Name': 'Josimar'}] black +10: [{'UID': '2002085528', 'Name': 'Noh Chan-Soo'}, {'UID': '19347507', 'Name': 'Kim In-Ho'}, {'UID': '2002069549', 'Name': 'Zhang Ding'}] asian