update background removal

This commit is contained in:
Karl 2025-09-23 15:15:29 +01:00
parent 497a36bd58
commit 243963b252
4 changed files with 103 additions and 19 deletions

31
gui.py
View File

@ -860,10 +860,35 @@ class FMFaceGeneratorGUI:
# Remove background from images if available # Remove background from images if available
if BG_REMOVAL_AVAILABLE: if BG_REMOVAL_AVAILABLE:
try: try:
remove_bg_from_file_list(self.output_dir_var.get(), [item[0] for item in players_to_process], use_gpu=False) output_dir = self.output_dir_var.get() if self.output_dir_var.get() else "./NewGens/"
self.log_message("Background removed from images.") player_ids = [item[0] for item in players_to_process]
self.log_message(f"Starting background removal for {len(player_ids)} images...")
# Check if images exist before processing
missing_images = []
for uid in player_ids:
image_path = os.path.join(output_dir, f"{uid}.png")
if not os.path.exists(image_path):
missing_images.append(uid)
if missing_images:
self.log_message(f"Warning: {len(missing_images)} images not found: {missing_images[:5]}...")
# Only process existing images
existing_ids = [uid for uid in player_ids if os.path.exists(os.path.join(output_dir, f"{uid}.png"))]
if existing_ids:
self.log_message(f"Processing {len(existing_ids)} images for background removal...")
processed_count = remove_bg_from_file_list(output_dir, existing_ids, use_gpu=False)
self.log_message(f"Background removal completed: {processed_count}/{len(existing_ids)} images processed successfully.")
else:
self.log_message("No images found for background removal.")
except Exception as e: except Exception as e:
self.log_message(f"Background removal failed: {e}") self.log_message(f"Background removal failed: {str(e)}")
import traceback
self.log_message(f"Error details: {traceback.format_exc()}")
else: else:
self.log_message("Background removal not available (rembg not installed). Images will have original backgrounds.") self.log_message("Background removal not available (rembg not installed). Images will have original backgrounds.")

View File

@ -29,6 +29,7 @@ def process_images_in_batch(batch, directory, use_gpu):
int: Number of images successfully processed in this batch. int: Number of images successfully processed in this batch.
""" """
if not REMBG_AVAILABLE: if not REMBG_AVAILABLE:
print("Background removal not available (rembg not installed)")
return 0 return 0
success_count = 0 success_count = 0
@ -37,22 +38,54 @@ def process_images_in_batch(batch, directory, use_gpu):
output_path = os.path.join(directory, f"{filename}.png") output_path = os.path.join(directory, f"{filename}.png")
try: try:
# Check if input file exists
if not os.path.exists(input_path):
print(f"Warning: Input file not found: {input_path}")
continue
# Check if input file is valid
if os.path.getsize(input_path) == 0:
print(f"Warning: Empty file: {input_path}")
continue
with Image.open(input_path) as img: with Image.open(input_path) as img:
# Validate image
if img.size[0] == 0 or img.size[1] == 0:
print(f"Warning: Invalid image dimensions: {input_path}")
continue
# Initialize ONNX session options with GPU support if required # Initialize ONNX session options with GPU support if required
if ONNX_AVAILABLE: if ONNX_AVAILABLE:
try:
session_options = ort.SessionOptions() session_options = ort.SessionOptions()
providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"] providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
ort.set_default_logger_severity(3) # Suppress non-critical logging ort.set_default_logger_severity(3) # Suppress non-critical logging
# Initialize the rembg remove function with appropriate providers # Initialize the rembg remove function with appropriate providers
output = remove(img, session_options=session_options, providers=providers) output = remove(img, session_options=session_options, providers=providers)
else: except Exception as gpu_error:
print(f"GPU processing failed for {input_path}, falling back to CPU: {str(gpu_error)}")
# Fallback to CPU-only processing # Fallback to CPU-only processing
output = remove(img) output = remove(img)
else:
# CPU-only processing
output = remove(img)
output.save(output_path) # Validate output
if output is None:
print(f"Warning: Background removal returned None for {input_path}")
continue
# Save result
output.save(output_path, 'PNG')
success_count += 1 success_count += 1
except Image.UnidentifiedImageError:
print(f"Error: Invalid image format: {input_path}")
except PermissionError:
print(f"Error: Permission denied accessing: {input_path}")
except OSError as e:
print(f"Error: OS error processing {input_path}: {str(e)}")
except Exception as e: except Exception as e:
print(f"Error processing {input_path}: {str(e)}") print(f"Error processing {input_path}: {str(e)}")
@ -76,10 +109,17 @@ def remove_bg_from_file_list(directory, filenames, max_workers=2, batch_size=2,
print("Background removal not available (rembg not installed). Skipping background removal.") print("Background removal not available (rembg not installed). Skipping background removal.")
return 0 return 0
print(f"Starting background removal for {len(filenames)} images in {directory}")
print(f"Configuration: max_workers={max_workers}, batch_size={batch_size}, use_gpu={use_gpu}")
processed_count = 0 processed_count = 0
failed_count = 0
# Divide filenames into batches # Divide filenames into batches
batches = [filenames[i:i + batch_size] for i in range(0, len(filenames), batch_size)] batches = [filenames[i:i + batch_size] for i in range(0, len(filenames), batch_size)]
total_batches = len(batches)
print(f"Processing in {total_batches} batches...")
with ThreadPoolExecutor(max_workers=max_workers) as executor: with ThreadPoolExecutor(max_workers=max_workers) as executor:
with tqdm(total=len(filenames), desc="Removing Backgrounds", unit="image") as pbar: with tqdm(total=len(filenames), desc="Removing Backgrounds", unit="image") as pbar:
@ -88,8 +128,32 @@ def remove_bg_from_file_list(directory, filenames, max_workers=2, batch_size=2,
for batch in batches for batch in batches
} }
for future in futures: for i, future in enumerate(futures, 1):
processed_count += future.result() try:
batch_result = future.result()
processed_count += batch_result
pbar.update(len(futures[future])) pbar.update(len(futures[future]))
if batch_result > 0:
print(f"Batch {i}/{total_batches}: {batch_result} images processed successfully")
else:
print(f"Batch {i}/{total_batches}: No images processed (possible errors)")
except Exception as e:
failed_count += 1
print(f"Batch {i}/{total_batches}: Failed with error: {str(e)}")
print(f"Background removal completed:")
print(f" - Total images: {len(filenames)}")
print(f" - Successfully processed: {processed_count}")
print(f" - Failed batches: {failed_count}")
if processed_count == 0:
print("WARNING: No images were processed successfully!")
print("This might indicate:")
print(" - All images have already been processed")
print(" - Image files are corrupted or missing")
print(" - Background removal is failing due to image format issues")
print(" - GPU/CPU compatibility issues with rembg")
return processed_count return processed_count

View File

@ -236,7 +236,6 @@ python remove_backgrounds.py ./NewGens/ --method simple
- Try installing packages individually: `pip install torch torchvision diffusers transformers` - Try installing packages individually: `pip install torch torchvision diffusers transformers`
- For Python 3.11+, use: `pip install rembg==2.0.59` (instead of 2.0.60) - For Python 3.11+, use: `pip install rembg==2.0.59` (instead of 2.0.60)
- If `striprtf` fails, try: `pip install striprtf==0.0.28` - If `striprtf` fails, try: `pip install striprtf==0.0.28`
- Use alternative requirements: `pip install -r requirements-gui.txt` for GUI-only
**Numba/rembg compatibility issues:** **Numba/rembg compatibility issues:**
- The application now handles missing `rembg` gracefully - The application now handles missing `rembg` gracefully

View File

@ -1,4 +0,0 @@
# Alternative requirements for GUI-only usage (without AI generation)
# Use this if you only want to run the GUI interface
tkinterdnd2==0.3.0
Pillow==10.4.0