import gradio as gr
import rembg
from rembg import remove, new_session
from PIL import Image
import numpy as np
import logging
import time
from hashlib import md5
import json
from pathlib import Path
import tempfile
import os
# Set up minimal logging to track usage
logging.basicConfig(level=logging.INFO, filename="user_activity.log", filemode="a")
logging.info(f"rembg version: {rembg.__version__}")
# Initialize user counter file
counter_file = Path("user_count.json")
def init_counter():
if not counter_file.exists():
with open(counter_file, "w") as f:
json.dump({"total_requests": 0, "unique_sessions": []}, f)
def log_user_interaction():
init_counter()
# Generate a unique user ID (hashed timestamp as proxy for session)
user_id = md5(str(time.time()).encode()).hexdigest()
try:
with open(counter_file, "r") as f:
content = f.read().strip()
if not content: # File is empty
data = {"total_requests": 0, "unique_sessions": []}
else:
data = json.loads(content)
except (json.JSONDecodeError, FileNotFoundError):
# Handle corrupted or missing file
data = {"total_requests": 0, "unique_sessions": []}
data["total_requests"] += 1
if user_id not in data["unique_sessions"]:
data["unique_sessions"].append(user_id)
try:
with open(counter_file, "w") as f:
json.dump(data, f)
except Exception as e:
logging.error(f"Failed to write counter file: {e}")
# Log only the request number and timestamp
logging.info(f"Request #{data['total_requests']} at {time.strftime('%Y%m%d-%H%M%S')}")
return data["total_requests"], len(data["unique_sessions"])
# Define model options
MODEL_OPTIONS = {
"": "Select a model",
"u2net": "A pre-trained model for general use cases (default)",
"isnet-general-use": "A new pre-trained model for general use cases",
"isnet-anime": "High-accuracy segmentation for anime characters",
"silueta": "A reduced-size version of u2net (43MB)",
"unet": "Lightweight version of u2net model",
"u2netp": "A lightweight version of u2net model",
"u2net_human_seg": "A pre-trained model for human segmentation",
"u2net_cloth_seg": "A pre-trained model for cloth parsing in human portraits",
}
def hex_to_rgba(hex_color):
if not hex_color:
return None
hex_color = hex_color.lstrip('#')
if len(hex_color) == 6:
hex_color += 'FF' # Add full opacity if no alpha is provided
return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4, 6))
def remove_background(input_path, bg_color=None, transparent_bg=True, model_choice="", alpha_matting=False, post_process_mask=False, only_mask=False):
print(f"DEBUG: Function called with input_path: {input_path}")
print(f"DEBUG: Parameters - bg_color: {bg_color}, transparent_bg: {transparent_bg}, model_choice: {model_choice}")
try:
# Check if input path is valid
if not input_path or not os.path.exists(input_path):
print(f"ERROR: Invalid input path: {input_path}")
return None
print("DEBUG: Starting background removal...")
# Log user interaction (minimal) - with error handling
try:
log_user_interaction()
except Exception as e:
print(f"WARNING: Failed to log user interaction: {e}")
# Continue with background removal even if logging fails
# Open the input image
input_image = Image.open(input_path)
print(f"DEBUG: Opened image with size: {input_image.size}, mode: {input_image.mode}")
# Extract the original filename without extension
original_filename = os.path.splitext(os.path.basename(input_path))[0]
# Create the desired output filename
output_filename = f"{original_filename}_removebg.png"
print(f"DEBUG: Output filename will be: {output_filename}")
# Extract the model name from the choice
model_name = model_choice.split(' | ')[0] if model_choice and ' | ' in model_choice else model_choice
print(f"DEBUG: Using model: {model_name if model_name else 'default'}")
# Set up the session with the chosen model, or None if no model is selected
session = new_session(model_name) if model_name else None
# Use transparent background if selected, otherwise use color
bg_color_rgba = None if transparent_bg else hex_to_rgba(bg_color)
print(f"DEBUG: Background color RGBA: {bg_color_rgba}")
# Prepare additional options
remove_kwargs = {}
# Only add session if we have one
if session is not None:
remove_kwargs["session"] = session
# Only add bgcolor if we have one
if bg_color_rgba is not None:
remove_kwargs["bgcolor"] = bg_color_rgba
# Add other parameters
if alpha_matting:
remove_kwargs.update({
"alpha_matting": True,
"alpha_matting_foreground_threshold": 270,
"alpha_matting_background_threshold": 20,
"alpha_matting_erode_size": 11
})
if post_process_mask:
remove_kwargs["post_process_mask"] = True
if only_mask:
remove_kwargs["only_mask"] = True
print(f"DEBUG: Remove kwargs: {remove_kwargs}")
# Convert PIL Image to numpy array
input_array = np.array(input_image)
print(f"DEBUG: Input array shape: {input_array.shape}")
# Use the remove function
print("DEBUG: Calling rembg.remove()...")
output_array = remove(input_array, **remove_kwargs)
print(f"DEBUG: Output array shape: {output_array.shape}")
# Convert numpy array back to PIL Image
output_image = Image.fromarray(output_array)
print(f"DEBUG: Output image size: {output_image.size}, mode: {output_image.mode}")
# Preserve transparency for transparent background or only_mask
if transparent_bg or only_mask:
if output_image.mode != 'RGBA':
output_image = output_image.convert('RGBA')
print("DEBUG: Converted to RGBA mode")
elif output_image.mode != 'RGB':
output_image = output_image.convert('RGB')
print("DEBUG: Converted to RGB mode")
# Create a temporary directory and save the image with the desired filename
temp_dir = tempfile.mkdtemp()
output_path = os.path.join(temp_dir, output_filename)
output_image.save(output_path, format="PNG")
print(f"DEBUG: Saved output to: {output_path}")
# Verify the file was created
if os.path.exists(output_path):
file_size = os.path.getsize(output_path)
print(f"DEBUG: Output file created successfully, size: {file_size} bytes")
return output_path
else:
print("ERROR: Output file was not created")
return None
except Exception as e:
print(f"ERROR: Exception occurred: {str(e)}")
import traceback
traceback.print_exc()
logging.error(f"An error occurred: {e}")
return None
# Fixed examples with proper values for all inputs
examples = [
[
'scifi_man1.jpg', # input_path
"#FFFFFF", # bg_color (white)
True, # transparent_bg
"", # model_choice (empty string, which is in the choices)
False, # alpha_matting
False, # post_process_mask
False # only_mask
]
]
# Gradio interface
iface = gr.Interface(
fn=remove_background,
inputs=[
gr.Image(type="filepath", label="Input Image"),
gr.ColorPicker(label="Background Color (ignored if transparent is selected)", value="#FFFFFF"),
gr.Checkbox(label="Transparent Background", value=False),
gr.Dropdown(
choices=[""] + [f"{k} | {v}" for k, v in MODEL_OPTIONS.items() if k != ""],
label="Model Selection",
value="", # Changed from empty to match choices
allow_custom_value=False
),
gr.Checkbox(label="Enable Alpha Matting", value=False),
gr.Checkbox(label="Post-Process Mask", value=False),
gr.Checkbox(label="Only Return Mask", value=False)
],
outputs=[
gr.Image(type="filepath", label="Output Image (PNG)")
],
examples=examples,
title="Background Remover v2.9",
description = """
Upload an image to remove the background. Choose a solid color or transparent background, select a model, and customize with alpha matting and other options.
Option to save as PNG or WEBP. Other apps required payment to download in HD here it's free. For the story behind this project checkout this blog post
If you would like to support this project:
""",
allow_flagging="never",
)
if __name__ == "__main__":
print("Starting Background Remover app...")
iface.launch()