gui updates

This commit is contained in:
Karl 2025-09-23 14:25:59 +01:00
parent ac4ff3204b
commit c5793c9940
6 changed files with 573 additions and 108 deletions

View File

@ -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")

345
gui.py
View File

@ -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:

View File

@ -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)

View File

@ -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:**

137
remove_backgrounds.py Normal file
View File

@ -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()

Binary file not shown.