diff --git a/gui.py b/gui.py index 4cc9adb..26191c9 100644 --- a/gui.py +++ b/gui.py @@ -860,10 +860,35 @@ class FMFaceGeneratorGUI: # Remove background from images if available if BG_REMOVAL_AVAILABLE: try: - remove_bg_from_file_list(self.output_dir_var.get(), [item[0] for item in players_to_process], use_gpu=False) - self.log_message("Background removed from images.") + output_dir = self.output_dir_var.get() if self.output_dir_var.get() else "./NewGens/" + 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: - 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: self.log_message("Background removal not available (rembg not installed). Images will have original backgrounds.") diff --git a/lib/remove_bg.py b/lib/remove_bg.py index 20fbfc8..1b9f6fc 100644 --- a/lib/remove_bg.py +++ b/lib/remove_bg.py @@ -29,6 +29,7 @@ def process_images_in_batch(batch, directory, use_gpu): int: Number of images successfully processed in this batch. """ if not REMBG_AVAILABLE: + print("Background removal not available (rembg not installed)") return 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") 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: + # 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 if ONNX_AVAILABLE: - session_options = ort.SessionOptions() - providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"] - ort.set_default_logger_severity(3) # Suppress non-critical logging + try: + session_options = ort.SessionOptions() + providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"] + ort.set_default_logger_severity(3) # Suppress non-critical logging - # Initialize the rembg remove function with appropriate providers - output = remove(img, session_options=session_options, providers=providers) + # Initialize the rembg remove function with appropriate providers + output = remove(img, session_options=session_options, providers=providers) + 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 + output = remove(img) else: - # Fallback to CPU-only processing + # 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 - success_count += 1 + # Save result + output.save(output_path, 'PNG') + 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: 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.") 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 + failed_count = 0 # Divide filenames into batches 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 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 future in futures: - processed_count += future.result() - pbar.update(len(futures[future])) + for i, future in enumerate(futures, 1): + try: + batch_result = future.result() + processed_count += batch_result + 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 \ No newline at end of file diff --git a/readme.md b/readme.md index e8d7a56..c3fad5c 100644 --- a/readme.md +++ b/readme.md @@ -236,7 +236,6 @@ python remove_backgrounds.py ./NewGens/ --method simple - 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) - 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:** - The application now handles missing `rembg` gracefully diff --git a/requirements-gui.txt b/requirements-gui.txt deleted file mode 100644 index 8ea4d2b..0000000 --- a/requirements-gui.txt +++ /dev/null @@ -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 \ No newline at end of file