diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index f8d4d45..552b53d 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -21,12 +21,7 @@ except ImportError: from lib.generate_xml import create_config_xml, append_to_config_xml from lib.resize_images import resize_images from lib.xml_reader import extract_from_values -from lib.general import ( - choose_profile, - create_or_update, - process_player_or_file, - get_player_input, -) +from lib.general import list_profiles from lib.logging import LOGGING_CONFIG # from simple_term_menu import TerminalMenu @@ -45,7 +40,23 @@ process_player = False user_config = configparser.ConfigParser() try: user_config.read("./user_config.cfg") - selected_profile = choose_profile("./user_config.cfg") + + # Check if we're running in GUI mode (look for GUI environment variable or check if GUI is importing us) + running_in_gui = os.environ.get('FM_NEWGEN_GUI_MODE', 'false').lower() == 'true' + + if running_in_gui: + # GUI mode: don't do CLI profile selection, let GUI handle it + # Use default profile or first available profile + profiles = [section.split(':', 1)[1] for section in user_config.sections() if section.startswith('profile:')] + if profiles: + selected_profile = profiles[0] # Use first available profile + else: + selected_profile = "NewGens" # Default fallback + logging.debug(f"GUI mode: using profile '{selected_profile}'") + else: + # CLI mode: do normal profile selection + selected_profile = choose_profile("./user_config.cfg") + selected_profile = f"profile:{selected_profile}" output_folder = user_config[selected_profile]["output_dir"] logging.debug("Configuration loaded successfully.") @@ -267,13 +278,14 @@ def main(): logging.error(f"RTF file not found: {rtf_location}") sys.exit(1) - update = create_or_update() - process_player = process_player_or_file() - if process_player: - player_uuid = get_player_input() - # Check for processed + # Get parameters from environment variables (set by GUI) + update_mode = os.environ.get('FM_NEWGEN_UPDATE_MODE', 'false').lower() == 'true' + process_specific_player = os.environ.get('FM_NEWGEN_PROCESS_PLAYER', 'false').lower() == 'true' + specific_player_uid = os.environ.get('FM_NEWGEN_PLAYER_UID', '') + + # Check for processed players try: - if update: + if update_mode: values_from_config = extract_from_values( f"{user_config[selected_profile]['output_dir']}config.xml" ) @@ -282,22 +294,22 @@ def main(): # Filter list_a to remove inner lists whose first item matches an ID in list_b players_to_process = [item for item in rtf_file if item[0] not in ids_in_b] - if process_player: + if process_specific_player and specific_player_uid: players_to_process = [ inner_list for inner_list in players_to_process - if int(inner_list[0]) == player_uuid + if int(inner_list[0]) == int(specific_player_uid) ] - elif process_player: + elif process_specific_player and specific_player_uid: players_to_process = [ inner_list for inner_list in rtf_file - if int(inner_list[0]) == player_uuid + if int(inner_list[0]) == int(specific_player_uid) ] else: players_to_process = rtf_file except FileNotFoundError: - logging.error("config.json file not found.") + logging.error("config.xml file not found.") sys.exit(1) if len(players_to_process) > 0: print(f"Processing {len(players_to_process)} players") diff --git a/gui.py b/gui.py index 9b361dc..295c47b 100644 --- a/gui.py +++ b/gui.py @@ -46,6 +46,9 @@ class FMFaceGeneratorGUI: self.root.geometry("1000x700") self.root.minsize(800, 600) + # Set environment variable to indicate GUI mode + os.environ['FM_NEWGEN_GUI_MODE'] = 'true' + # Initialize variables self.config = None self.app_config = None @@ -53,14 +56,280 @@ class FMFaceGeneratorGUI: self.generation_thread = None self.stop_generation = False self.image_queue = queue.Queue() + self.current_profile = None # Set up the GUI self.setup_gui() - self.load_config() + + # Show profile selection dialog first + self.show_profile_selection() # Start image preview update loop self.update_image_preview() + def show_profile_selection(self): + """Show profile selection dialog at startup""" + # Load existing profiles + profiles = self.list_profiles() + + if not profiles: + # No profiles exist, create default one + self.create_default_profile() + self.current_profile = "NewGens" + self.load_config() + return + + # Create selection dialog + dialog = tk.Toplevel(self.root) + dialog.title("Select Profile") + dialog.geometry("400x300") + dialog.transient(self.root) + dialog.grab_set() + + # Center the dialog + dialog.update_idletasks() + x = (dialog.winfo_screenwidth() - dialog.winfo_width()) // 2 + y = (dialog.winfo_screenheight() - dialog.winfo_height()) // 2 + dialog.geometry(f"+{x}+{y}") + + # Profile selection + ttk.Label(dialog, text="Select a profile to use:").pack(pady=10) + + profile_var = tk.StringVar() + profile_combo = ttk.Combobox(dialog, textvariable=profile_var, width=30) + profile_combo['values'] = profiles + profile_combo.pack(pady=10) + profile_combo.set(profiles[0]) # Select first profile by default + + def on_select(): + selected = profile_var.get() + if selected: + self.current_profile = selected + dialog.destroy() + self.load_config() + self.log_message(f"Selected profile: {selected}") + + def on_create_new(): + dialog.destroy() + self.show_create_profile_dialog() + + def on_manage(): + dialog.destroy() + self.notebook.select(1) # Switch to profiles tab + + # Buttons + btn_frame = ttk.Frame(dialog) + btn_frame.pack(pady=20) + + ttk.Button(btn_frame, text="Select", command=on_select).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Create New", command=on_create_new).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Manage Profiles", command=on_manage).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Cancel", command=self.root.quit).pack(side=tk.LEFT, padx=5) + + # Wait for dialog to close + self.root.wait_window(dialog) + + def show_create_profile_dialog(self): + """Show dialog to create a new profile""" + dialog = tk.Toplevel(self.root) + dialog.title("Create New Profile") + dialog.geometry("400x250") + dialog.transient(self.root) + dialog.grab_set() + + # Center the dialog + dialog.update_idletasks() + x = (dialog.winfo_screenwidth() - dialog.winfo_width()) // 2 + y = (dialog.winfo_screenheight() - dialog.winfo_height()) // 2 + dialog.geometry(f"+{x}+{y}") + + # Profile name + ttk.Label(dialog, text="Profile Name:").grid(row=0, column=0, sticky=tk.W, pady=10, padx=10) + profile_name_var = tk.StringVar() + profile_name_entry = ttk.Entry(dialog, textvariable=profile_name_var, width=30) + profile_name_entry.grid(row=0, column=1, sticky=tk.W, padx=10, pady=10) + + # FM Version + ttk.Label(dialog, text="FM Version:").grid(row=1, column=0, sticky=tk.W, pady=5, padx=10) + fm_version_var = tk.StringVar(value="2024") + fm_version_entry = ttk.Entry(dialog, textvariable=fm_version_var, width=30) + fm_version_entry.grid(row=1, column=1, sticky=tk.W, padx=10, pady=5) + + # RTF File + ttk.Label(dialog, text="RTF File:").grid(row=2, column=0, sticky=tk.W, pady=5, padx=10) + rtf_file_var = tk.StringVar() + rtf_entry = ttk.Entry(dialog, textvariable=rtf_file_var, width=30) + rtf_entry.grid(row=2, column=1, sticky=tk.W, padx=10, pady=5) + ttk.Button(dialog, text="Browse...", command=lambda: self.browse_file(rtf_file_var, "RTF files", "*.rtf")).grid(row=2, column=2, padx=10, pady=5) + + # Output Directory + ttk.Label(dialog, text="Output Directory:").grid(row=3, column=0, sticky=tk.W, pady=5, padx=10) + output_dir_var = tk.StringVar() + output_entry = ttk.Entry(dialog, textvariable=output_dir_var, width=30) + output_entry.grid(row=3, column=1, sticky=tk.W, padx=10, pady=5) + ttk.Button(dialog, text="Browse...", command=lambda: self.browse_directory(output_dir_var)).grid(row=3, column=2, padx=10, pady=5) + + def on_create(): + name = profile_name_var.get().strip() + if not name: + messagebox.showerror("Error", "Profile name cannot be empty") + return + + if self.profile_exists(name): + messagebox.showerror("Error", f"Profile '{name}' already exists") + return + + # Create the profile + self.create_profile(name, fm_version_var.get(), rtf_file_var.get(), output_dir_var.get()) + self.current_profile = name + dialog.destroy() + self.load_config() + self.log_message(f"Created and selected profile: {name}") + + def on_cancel(): + dialog.destroy() + self.show_profile_selection() # Go back to profile selection + + # Buttons + btn_frame = ttk.Frame(dialog) + btn_frame.grid(row=4, column=0, columnspan=3, pady=20) + + ttk.Button(btn_frame, text="Create", command=on_create).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Cancel", command=on_cancel).pack(side=tk.LEFT, padx=5) + + def list_profiles(self): + """List all available profiles""" + if not os.path.exists('user_config.cfg'): + return [] + + config = configparser.ConfigParser() + config.read('user_config.cfg') + + profiles = [section.split(':', 1)[1] for section in config.sections() if section.startswith('profile:')] + return profiles + + def profile_exists(self, profile_name): + """Check if a profile exists""" + profiles = self.list_profiles() + return profile_name in profiles + + def create_profile(self, name, fm_version, rtf_file, output_dir): + """Create a new profile""" + if not self.config: + self.config = configparser.ConfigParser() + self.config.add_section('models') + + section_name = f'profile:{name}' + self.config.add_section(section_name) + self.config.set(section_name, 'football_manager_version', fm_version) + self.config.set(section_name, 'rtf_file', rtf_file) + self.config.set(section_name, 'output_dir', output_dir) + + # Save to file + with open('user_config.cfg', 'w') as configfile: + self.config.write(configfile) + + def create_default_profile(self): + """Create a default profile if none exists""" + self.config = configparser.ConfigParser() + self.config.add_section('models') + self.config.add_section('profile:NewGens') + + self.config.set('models', 'model_name', 'SG161222/Realistic_Vision_V6.0_B1') + self.config.set('profile:NewGens', 'football_manager_version', '2024') + self.config.set('profile:NewGens', 'rtf_file', '') + self.config.set('profile:NewGens', 'output_dir', './NewGens/') + + with open('user_config.cfg', 'w') as configfile: + self.config.write(configfile) + + def browse_file(self, var, title, filetypes): + """Browse for a file""" + file_path = filedialog.askopenfilename(title=title, filetypes=[(title, filetypes), ("All files", "*.*")]) + if file_path: + var.set(file_path) + + def browse_directory(self, var): + """Browse for a directory""" + directory = filedialog.askdirectory(title="Select Directory") + if directory: + var.set(directory) + + def setup_profiles_tab(self): + """Set up the profiles management tab""" + # Profile list + list_frame = ttk.LabelFrame(self.profiles_frame, text="Available Profiles", padding=10) + list_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5) + + # Profile listbox + self.profiles_listbox = tk.Listbox(list_frame, height=10) + self.profiles_listbox.pack(fill=tk.BOTH, expand=True, pady=5) + + # Buttons + btn_frame = ttk.Frame(list_frame) + btn_frame.pack(fill=tk.X, pady=5) + + ttk.Button(btn_frame, text="Refresh List", command=self.refresh_profiles_list).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Select Profile", command=self.select_profile_from_list).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Create New", command=self.show_create_profile_dialog).pack(side=tk.LEFT, padx=5) + ttk.Button(btn_frame, text="Delete", command=self.delete_profile).pack(side=tk.RIGHT, padx=5) + + # Profile details + details_frame = ttk.LabelFrame(self.profiles_frame, text="Profile Details", padding=10) + details_frame.pack(fill=tk.X, padx=10, pady=5) + + # Current profile display + self.current_profile_var = tk.StringVar() + ttk.Label(details_frame, text="Current Profile:").grid(row=0, column=0, sticky=tk.W, pady=5) + ttk.Label(details_frame, textvariable=self.current_profile_var).grid(row=0, column=1, sticky=tk.W, padx=5, pady=5) + + # Refresh the profiles list + self.refresh_profiles_list() + + def refresh_profiles_list(self): + """Refresh the profiles listbox""" + self.profiles_listbox.delete(0, tk.END) + profiles = self.list_profiles() + for profile in profiles: + self.profiles_listbox.insert(tk.END, profile) + + if self.current_profile: + self.current_profile_var.set(self.current_profile) + + def select_profile_from_list(self): + """Select a profile from the list""" + selection = self.profiles_listbox.curselection() + if selection: + profile_name = self.profiles_listbox.get(selection[0]) + self.current_profile = profile_name + self.load_config() + self.log_message(f"Selected profile: {profile_name}") + self.refresh_profiles_list() + + def delete_profile(self): + """Delete the selected profile""" + selection = self.profiles_listbox.curselection() + if not selection: + messagebox.showerror("Error", "Please select a profile to delete") + return + + profile_name = self.profiles_listbox.get(selection[0]) + + if profile_name == self.current_profile: + messagebox.showerror("Error", "Cannot delete the currently active profile") + return + + if messagebox.askyesno("Confirm Delete", f"Are you sure you want to delete profile '{profile_name}'?"): + # Remove from config + if self.config: + self.config.remove_section(f'profile:{profile_name}') + + with open('user_config.cfg', 'w') as configfile: + self.config.write(configfile) + + self.log_message(f"Deleted profile: {profile_name}") + self.refresh_profiles_list() + def setup_gui(self): """Set up the main GUI components""" # Create main notebook (tabs) @@ -71,6 +340,10 @@ class FMFaceGeneratorGUI: self.config_frame = ttk.Frame(self.notebook) self.notebook.add(self.config_frame, text="Configuration") + # Profiles tab + self.profiles_frame = ttk.Frame(self.notebook) + self.notebook.add(self.profiles_frame, text="Profiles") + # Generation tab self.generation_frame = ttk.Frame(self.notebook) self.notebook.add(self.generation_frame, text="Generation") @@ -81,6 +354,7 @@ class FMFaceGeneratorGUI: # Setup each tab self.setup_config_tab() + self.setup_profiles_tab() self.setup_generation_tab() self.setup_preview_tab() @@ -257,22 +531,28 @@ class FMFaceGeneratorGUI: if not self.config: self.config = configparser.ConfigParser() self.config.add_section('models') - self.config.add_section('profile:NewGens') + + # Use current profile or create NewGens if none selected + profile_name = self.current_profile or "NewGens" + profile_section = f'profile:{profile_name}' + + if not self.config.has_section(profile_section): + self.config.add_section(profile_section) self.config.set('models', 'model_name', self.model_var.get()) if self.model_dir_var.get(): self.config.set('models', 'model_dir', self.model_dir_var.get()) - self.config.set('profile:NewGens', 'football_manager_version', self.fm_version_var.get()) - self.config.set('profile:NewGens', 'rtf_file', self.rtf_file_var.get()) - self.config.set('profile:NewGens', 'output_dir', self.output_dir_var.get()) + self.config.set(profile_section, 'football_manager_version', self.fm_version_var.get()) + self.config.set(profile_section, 'rtf_file', self.rtf_file_var.get()) + self.config.set(profile_section, 'output_dir', self.output_dir_var.get()) # Save to file with open('user_config.cfg', 'w') as configfile: self.config.write(configfile) - messagebox.showinfo("Success", "Configuration saved successfully!") - self.log_message("Configuration saved") + messagebox.showinfo("Success", f"Configuration saved to profile '{profile_name}'!") + self.log_message(f"Configuration saved to profile '{profile_name}'") except Exception as e: messagebox.showerror("Error", f"Failed to save configuration: {str(e)}") @@ -290,13 +570,27 @@ class FMFaceGeneratorGUI: self.model_var.set(self.config.get('models', 'model_name', fallback='SG161222/Realistic_Vision_V6.0_B1')) self.model_dir_var.set(self.config.get('models', 'model_dir', fallback='')) - # Load FM settings - if 'profile:NewGens' in self.config: - self.fm_version_var.set(self.config.get('profile:NewGens', 'football_manager_version', fallback='2024')) - self.rtf_file_var.set(self.config.get('profile:NewGens', 'rtf_file', fallback='')) - self.output_dir_var.set(self.config.get('profile:NewGens', 'output_dir', fallback='')) - - self.log_message("Configuration loaded") + # Load current profile settings + if self.current_profile: + profile_section = f'profile:{self.current_profile}' + if profile_section in self.config: + self.fm_version_var.set(self.config.get(profile_section, 'football_manager_version', fallback='2024')) + self.rtf_file_var.set(self.config.get(profile_section, 'rtf_file', fallback='')) + self.output_dir_var.set(self.config.get(profile_section, 'output_dir', fallback='')) + self.log_message(f"Configuration loaded for profile '{self.current_profile}'") + else: + # Profile doesn't exist, create it with defaults + self.create_profile(self.current_profile, '2024', '', './NewGens/') + self.fm_version_var.set('2024') + self.rtf_file_var.set('') + self.output_dir_var.set('./NewGens/') + self.log_message(f"Created new profile '{self.current_profile}' with defaults") + else: + # No profile selected, use defaults + self.fm_version_var.set('2024') + self.rtf_file_var.set('') + self.output_dir_var.set('./NewGens/') + self.log_message("No profile selected, using defaults") else: self.log_message("No configuration file found, using defaults") @@ -310,7 +604,19 @@ class FMFaceGeneratorGUI: self.model_dir_var.set('') self.fm_version_var.set('2024') self.rtf_file_var.set('') - self.output_dir_var.set('') + self.output_dir_var.set('./NewGens/') + + # Update current profile with defaults + if self.current_profile: + profile_section = f'profile:{self.current_profile}' + if self.config and self.config.has_section(profile_section): + self.config.set(profile_section, 'football_manager_version', '2024') + self.config.set(profile_section, 'rtf_file', '') + self.config.set(profile_section, 'output_dir', './NewGens/') + + with open('user_config.cfg', 'w') as configfile: + self.config.write(configfile) + self.log_message("Configuration reset to defaults") def log_message(self, message): @@ -327,6 +633,12 @@ class FMFaceGeneratorGUI: if not self.validate_config(): return + # Set environment variables for the main script + os.environ['FM_NEWGEN_UPDATE_MODE'] = 'true' if self.update_mode_var.get() == "true" else 'false' + os.environ['FM_NEWGEN_PROCESS_PLAYER'] = 'true' if self.player_mode_var.get() == "specific" else 'false' + if self.player_mode_var.get() == "specific": + os.environ['FM_NEWGEN_PLAYER_UID'] = self.player_uid_var.get() + # Disable start button, enable stop button self.start_button.config(state=tk.DISABLED) self.stop_button.config(state=tk.NORMAL) @@ -369,7 +681,8 @@ class FMFaceGeneratorGUI: elif update: # Filter out already processed players try: - values_from_config = extract_from_values(f"{self.output_dir_var.get()}/config.xml") + output_dir = self.output_dir_var.get() if self.output_dir_var.get() else "./NewGens/" + values_from_config = extract_from_values(f"{output_dir}/config.xml") ids_in_b = [item for item in values_from_config] players_to_process = [item for item in rtf_file if item[0] not in ids_in_b] except FileNotFoundError: diff --git a/lib/general.py b/lib/general.py index 922a6a1..1e7a0c6 100644 --- a/lib/general.py +++ b/lib/general.py @@ -1,80 +1,9 @@ import configparser def list_profiles(cfg_file): + """List all available profiles from config file""" config = configparser.ConfigParser() config.read(cfg_file) - + profiles = [section.split(':', 1)[1] for section in config.sections() if section.startswith('profile:')] - return profiles - - -def choose_profile(cfg_file): - profiles = list_profiles(cfg_file) - if not profiles: - print("No profiles found.") - return None - - print("Available Profiles:") - for i, profile in enumerate(profiles, 1): - print(f"{i}. {profile}") - - while True: - try: - choice = int(input("Enter the number of the profile you want to use: ")) - if 1 <= choice <= len(profiles): - return profiles[choice - 1] - else: - print("Invalid number. Please try again.") - except ValueError: - print("Invalid input. Please enter a number.") - -def create_or_update(): - values = ["Create", "Update"] - print("Create or Update:") - for i, profile in enumerate(values, 1): - print(f"{i}. {profile}") - - while True: - try: - choice = int(input("Enter the number of the choice: ")) - if 1 <= choice <= len(values): - if choice == 1: - return False - else: - return True - else: - print("Invalid number. Please try again.") - except ValueError: - print("Invalid input. Please enter a number.") - - -def process_player_or_file(): - values = ["Whole File", "Specific Player"] - print("Process whole rtf file or a specific player?:") - for i, profile in enumerate(values, 1): - print(f"{i}. {profile}") - - while True: - try: - choice = int(input("Enter the number of the choice: ")) - if 1 <= choice <= len(values): - if choice == 1: - return False - else: - return True - else: - print("Invalid number. Please try again.") - except ValueError: - print("Invalid input. Please enter a number.") - - -def get_player_input(): - # Prompt the user to enter a UUID - player_uuid = input("Please enter the player UUID: ").strip() - - # Validate that the UUID is not empty - while not player_uuid: - print("Player UUID cannot be empty. Please try again.") - player_uuid = input("Please enter the player UUID: ").strip() - - return int(player_uuid) \ No newline at end of file + return profiles \ No newline at end of file diff --git a/readme.md b/readme.md index cf642e3..e8d7a56 100644 --- a/readme.md +++ b/readme.md @@ -78,6 +78,20 @@ python test_gui.py python run_gui.py ``` +### **🔧 Background Removal Only** +If you only need to remove backgrounds from existing images: + +```bash +# Remove backgrounds from generated images +python remove_backgrounds.py ./NewGens/ + +# Use specific method +python remove_backgrounds.py ./NewGens/ --method rembg + +# Save to different directory +python remove_backgrounds.py ./NewGens/ --output-dir ./Processed/ +``` + **GUI Features:** - Easy configuration with file browsers - Real-time progress monitoring @@ -149,6 +163,66 @@ The tool uses Hugging Face models for image generation. You can specify which mo - **Error Handling**: Clear error messages and troubleshooting tips - **Multi-threaded**: Non-blocking interface during generation +### **🔧 Background Removal Features** +- **Automatic Processing**: Backgrounds removed automatically after generation +- **Multiple Methods**: AI-powered (rembg) and simple fallback methods +- **Standalone Tool**: `remove_backgrounds.py` for manual processing +- **Flexible Options**: Works with or without GPU acceleration +- **Batch Processing**: Process entire directories of images + +## Background Removal + +### **Automatic Background Removal (Built-in)** +The application automatically removes backgrounds from generated images during post-processing: + +**When it works:** +- ✅ Images are processed automatically after generation +- ✅ Transparent backgrounds for better FM integration +- ✅ GPU acceleration when available + +**When it doesn't work:** +- ⚠️ Images keep original backgrounds +- ⚠️ May need manual processing + +### **Manual Background Removal** +If automatic removal fails, use the standalone tool: + +```bash +# Remove backgrounds from all images in a directory +python remove_backgrounds.py ./NewGens/ + +# Use specific method +python remove_backgrounds.py ./NewGens/ --method rembg + +# Save to different directory +python remove_backgrounds.py ./NewGens/ --output-dir ./NewGens_Processed/ +``` + +### **Background Removal Methods** + +| Method | Quality | Speed | Requirements | Notes | +|--------|---------|-------|--------------|-------| +| **rembg** | ⭐⭐⭐⭐⭐ | ⭐⭐⭐ | `rembg[gpu]` | Best quality, AI-powered | +| **Simple** | ⭐⭐ | ⭐⭐⭐⭐⭐ | Built-in | Fallback method, basic results | + +### **Installing Background Removal** + +**For best quality (recommended):** +```bash +pip install rembg[gpu]==2.0.59 +``` + +**For CPU-only:** +```bash +pip install rembg==2.0.59 +``` + +**If rembg fails:** +```bash +# The simple method works without additional dependencies +python remove_backgrounds.py ./NewGens/ --method simple +``` + ## GUI Troubleshooting ### **Common Issues:** diff --git a/remove_backgrounds.py b/remove_backgrounds.py new file mode 100644 index 0000000..ed436e3 --- /dev/null +++ b/remove_backgrounds.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python3 +""" +Standalone Background Removal Tool +Remove backgrounds from generated FM face images +""" + +import os +import sys +from pathlib import Path +from PIL import Image +import argparse + +def remove_background_simple(image_path, output_path=None): + """ + Simple background removal using PIL and basic image processing + This is a fallback method when rembg is not available + """ + try: + with Image.open(image_path) as img: + # Convert to RGBA if not already + if img.mode != 'RGBA': + img = img.convert('RGBA') + + # Get image data + data = img.getdata() + + # Simple background removal based on color similarity + # This is a basic implementation - results may vary + new_data = [] + for item in data: + # Consider pixels with high brightness/low saturation as background + r, g, b, a = item + + # Calculate brightness and saturation + brightness = (r + g + b) / 3 + saturation = max(r, g, b) - min(r, g, b) + + # If pixel is very bright and low saturation (likely background) + if brightness > 240 and saturation < 30: + # Make transparent + new_data.append((r, g, b, 0)) + else: + new_data.append(item) + + # Create new image with transparent background + new_img = Image.new('RGBA', img.size) + new_img.putdata(new_data) + + # Save result + output_path = output_path or image_path + new_img.save(output_path, 'PNG') + return True + + except Exception as e: + print(f"Error processing {image_path}: {e}") + return False + +def find_face_images(directory): + """Find all PNG images in directory that might be face images""" + image_extensions = {'.png', '.jpg', '.jpeg'} + face_images = [] + + for file_path in Path(directory).rglob('*'): + if file_path.suffix.lower() in image_extensions: + # Skip if already processed (has _no_bg in name) + if '_no_bg' not in file_path.stem: + face_images.append(str(file_path)) + + return sorted(face_images) + +def main(): + parser = argparse.ArgumentParser(description="Remove backgrounds from FM face images") + parser.add_argument("directory", help="Directory containing face images") + parser.add_argument("--method", choices=['rembg', 'simple', 'auto'], + default='auto', help="Background removal method") + parser.add_argument("--output-dir", help="Output directory for processed images") + parser.add_argument("--suffix", default="_no_bg", + help="Suffix to add to processed images") + + args = parser.parse_args() + + if not os.path.exists(args.directory): + print(f"Error: Directory {args.directory} does not exist") + sys.exit(1) + + # Find images + images = find_face_images(args.directory) + if not images: + print(f"No images found in {args.directory}") + sys.exit(1) + + print(f"Found {len(images)} images to process") + + # Determine method + method = args.method + if method == 'auto': + try: + from rembg import remove + method = 'rembg' + print("Using rembg for background removal") + except ImportError: + method = 'simple' + print("Using simple background removal (rembg not available)") + + # Process images + success_count = 0 + output_dir = args.output_dir or args.directory + + for i, image_path in enumerate(images, 1): + print(f"Processing {i}/{len(images)}: {os.path.basename(image_path)}") + + if method == 'rembg': + try: + with Image.open(image_path) as img: + output = remove(img) + output_path = os.path.join(output_dir, + f"{Path(image_path).stem}{args.suffix}.png") + output.save(output_path, 'PNG') + success_count += 1 + except Exception as e: + print(f" Error: {e}") + else: + # Simple method + output_path = os.path.join(output_dir, + f"{Path(image_path).stem}{args.suffix}.png") + if remove_background_simple(image_path, output_path): + success_count += 1 + + print(f"\nCompleted: {success_count}/{len(images)} images processed successfully") + + if success_count < len(images): + print(f"\nNote: {len(images) - success_count} images failed to process") + print("You may need to install rembg for better results:") + print("pip install rembg[gpu]==2.0.59") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 621668f..e86b6f4 100644 Binary files a/requirements.txt and b/requirements.txt differ