mirror of
https://github.com/karl0ss/comfy_fm24_newgens.git
synced 2025-10-03 06:40:06 +01:00
update background removal
This commit is contained in:
parent
497a36bd58
commit
243963b252
31
gui.py
31
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.")
|
||||
|
||||
|
||||
@ -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
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
Loading…
x
Reference in New Issue
Block a user