update background removal

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)
+                            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:
-                            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.")
 

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()
+                    try:
-                    providers = ["CUDAExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
+                        session_options = ort.SessionOptions()
-                    ort.set_default_logger_severity(3)  # Suppress non-critical logging
+                        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
+                        # 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)
+                    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:
+            for i, future in enumerate(futures, 1):
-                processed_count += future.result()
+                try:
-                pbar.update(len(futures[future]))
+                    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

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

requirements-gui.txt

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