app.py

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. Right-click the output image to save as PNG.",
    allow_flagging="never",
)

if __name__ == "__main__":
    print("Starting Background Remover app...")
    iface.launch()