app.py

import gradio as gr
import huggingface_hub
import onnxruntime as rt
import numpy as np
import cv2
import os
import csv
import datetime
import time

# --- Constants ---
LOG_FILE = "processing_log.csv"
LOG_HEADER = [
    "Timestamp", "Repository", "Model Filename", "Model Size (MB)", 
    "Image Resolution (WxH)", "Execution Provider", "Processing Time (s)"
]

# Global variables for model and providers
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
model_repo_default = "skytnt/anime-seg"

# --- Logging Functions ---

def initialize_log_file():
    """Creates the log file and writes the header if it doesn't exist."""
    if not os.path.exists(LOG_FILE):
        try:
            with open(LOG_FILE, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                writer.writerow(LOG_HEADER)
            print(f"Log file initialized: {LOG_FILE}")
        except IOError as e:
            print(f"Error initializing log file {LOG_FILE}: {e}")

def log_processing_event(timestamp, repo, model_filename, model_size_mb, 
                         resolution, provider, processing_time):
    """Appends a processing event to the CSV log file."""
    try:
        with open(LOG_FILE, 'a', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow([
                timestamp, repo, model_filename, f"{model_size_mb:.2f}", 
                resolution, provider, f"{processing_time:.4f}"
            ])
    except IOError as e:
        print(f"Error writing to log file {LOG_FILE}: {e}")
    except Exception as e:
        print(f"An unexpected error occurred during logging: {e}")

def read_log_file():
    """Reads the entire log file content."""
    try:
        if not os.path.exists(LOG_FILE):
            return "Log file not found."
        with open(LOG_FILE, 'r', encoding='utf-8') as f:
            # Read all lines and join them for display
            return "".join(f.readlines())
            # Alternatively, for cleaner display of CSV in a textbox:
            # reader = csv.reader(f)
            # rows = list(reader)
            # # Format header and rows nicely
            # header = rows[0]
            # data_rows = rows[1:]
            # formatted_rows = [", ".join(header)] # Join header elements
            # for row in data_rows:
            #     formatted_rows.append(", ".join(row)) # Join data elements
            # return "\n".join(formatted_rows)
    except IOError as e:
        print(f"Error reading log file {LOG_FILE}: {e}")
        return f"Error reading log file: {e}"
    except Exception as e:
        print(f"An unexpected error occurred reading log file: {e}")
        return f"Error reading log file: {e}"

# --- Helper Functions ---

def get_model_details_from_choice(choice_string: str) -> tuple[str, float | None]:
    """
    Extracts filename and size (MB) from the dropdown choice string.
    Returns (filename, size_mb) or (filename, None) if size is not parseable.
    """
    if not choice_string:
        return "", None
    parts = choice_string.split(" (")
    filename = parts[0]
    size_mb = None
    if len(parts) > 1 and parts[1].endswith(" MB)"):
        try:
            size_str = parts[1].replace(" MB)", "")
            size_mb = float(size_str)
        except ValueError:
            pass # Size couldn't be parsed
    return filename, size_mb
    
# --- Model Loading and UI Functions (Mostly unchanged, modifications marked) ---

def update_onnx_files(repo: str):
    """
    Lists .onnx files in the Hugging Face repository and updates the Dropdown with file sizes.
    """
    onnx_files_with_size = []
    try:
        api = huggingface_hub.HfApi()
        repo_info = api.model_info(repo_id=repo, files_metadata=True)
        
        for file_info in repo_info.siblings:
            if file_info.rfilename.endswith('.onnx'):
                try:
                    # Use file_info.size which is in bytes
                    size_mb = file_info.size / (1024 * 1024) if file_info.size else 0
                    onnx_files_with_size.append(f"{file_info.rfilename} ({size_mb:.2f} MB)")
                except Exception:
                    onnx_files_with_size.append(f"{file_info.rfilename} (Size N/A)")
        
        if onnx_files_with_size:
            onnx_files_with_size.sort()
            return gr.update(choices=onnx_files_with_size, value=onnx_files_with_size[0])
        else:
            return gr.update(choices=[], value="", warning=f"No .onnx files found in repository '{repo}'")
            
    except huggingface_hub.utils.RepositoryNotFoundError:
        return gr.update(choices=[], value="", error=f"Repository '{repo}' not found or access denied.")
    except Exception as e:
        print(f"Error fetching repo files for {repo}: {e}")
        return gr.update(choices=[], value="", error=f"Error fetching files: {str(e)}")

# Get default choices and filename
default_onnx_files_with_size = []
default_model_filename = ""
try:
    initial_update = update_onnx_files(model_repo_default)
    if isinstance(initial_update, gr.update) and initial_update.choices:
        default_onnx_files_with_size = initial_update.choices
        default_model_filename, _ = get_model_details_from_choice(default_onnx_files_with_size[0]) # Use helper
    else:
        default_onnx_files_with_size = ["isnetis.onnx (Size N/A)"]
        default_model_filename = "isnetis.onnx"
        print(f"Warning: Could not fetch initial ONNX files from {model_repo_default}. Using fallback '{default_model_filename}'.")
except Exception as e:
    default_onnx_files_with_size = ["isnetis.onnx (Size N/A)"]
    default_model_filename = "isnetis.onnx"
    print(f"Error during initial model fetch: {e}. Using fallback '{default_model_filename}'.")

# Global variables for current model state
current_model_repo = model_repo_default
current_model_filename = default_model_filename

# Initial download and model load
model_path = None
rmbg_model = None
try:
    print(f"Attempting initial download: {current_model_repo}/{current_model_filename}")
    if current_model_filename: # Only download if we have a filename
        model_path = huggingface_hub.hf_hub_download(current_model_repo, current_model_filename)
        rmbg_model = rt.InferenceSession(model_path, providers=providers)
        print(f"Initial model loaded successfully: {model_path}")
        print(f"Available Execution Providers: {rt.get_available_providers()}")
        print(f"Using Provider(s): {rmbg_model.get_providers()}")
    else:
         print("FATAL: No default model filename determined. Cannot load initial model.")
except Exception as e:
    print(f"FATAL: Could not download or load initial model '{current_model_repo}/{current_model_filename}'. Error: {e}")

# --- Inference Functions (Unchanged get_mask, rmbg_fn) ---
def get_mask(img, s=1024):
    if rmbg_model is None:
        raise gr.Error("Model is not loaded. Please check model selection and update status.")
    img_normalized = (img / 255.0).astype(np.float32)
    h0, w0 = img.shape[:2]
    if h0 >= w0: h, w = (s, int(s * w0 / h0))
    else: h, w = (int(s * h0 / w0), s)
    ph, pw = s - h, s - w
    img_input = np.zeros([s, s, 3], dtype=np.float32)
    resized_img = cv2.resize(img_normalized, (w, h), interpolation=cv2.INTER_AREA)
    img_input[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = resized_img
    img_input = np.transpose(img_input, (2, 0, 1))[np.newaxis, :]
    input_name = rmbg_model.get_inputs()[0].name
    mask_output = rmbg_model.run(None, {input_name: img_input})[0][0]
    mask_processed = np.transpose(mask_output, (1, 2, 0))
    mask_processed = mask_processed[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
    mask_resized = cv2.resize(mask_processed, (w0, h0), interpolation=cv2.INTER_LINEAR)
    if mask_resized.ndim == 2: mask_resized = mask_resized[:, :, np.newaxis]
    mask_final = np.clip(mask_resized, 0, 1)
    return mask_final

def rmbg_fn(img):
    if img is None: raise gr.Error("Please provide an input image.")
    mask = get_mask(img)
    if img.dtype != np.uint8: img = (img * 255).clip(0, 255).astype(np.uint8) if img.max() <= 1.0 else img.clip(0, 255).astype(np.uint8)
    alpha_channel = (mask * 255).astype(np.uint8)
    if img.shape[2] == 3: img_out_rgba = np.concatenate([img, alpha_channel], axis=2)
    else: img_out_rgba = img.copy(); img_out_rgba[:, :, 3] = alpha_channel[:,:,0]
    mask_img_display = (mask * 255).astype(np.uint8).repeat(3, axis=2)
    return mask_img_display, img_out_rgba

# --- Model Update Function ---
def update_model(model_repo, model_filename_with_size):
    global rmbg_model, current_model_repo, current_model_filename
    model_filename, _ = get_model_details_from_choice(model_filename_with_size) # Use helper
    if not model_filename: return "Error: No model filename selected or extracted."
    if model_repo == current_model_repo and model_filename == current_model_filename:
        # Even if it's the same, report the provider being used
        current_provider = rmbg_model.get_providers()[0] if rmbg_model else "N/A"
        return f"Model already loaded: {current_model_repo}/{current_model_filename}\nUsing Provider: {current_provider}"
        
    try:
        print(f"Updating model to: {model_repo}/{model_filename}")
        model_path = huggingface_hub.hf_hub_download(model_repo, model_filename)
        new_rmbg_model = rt.InferenceSession(model_path, providers=providers)
        rmbg_model = new_rmbg_model
        current_model_repo = model_repo
        current_model_filename = model_filename
        active_provider = rmbg_model.get_providers()[0] # Get the provider actually used
        print(f"Model updated successfully: {model_path}")
        print(f"Using Provider: {active_provider}")
        return f"Model updated: {current_model_repo}/{current_model_filename}\nUsing Provider: {active_provider}"
    except huggingface_hub.utils.HfHubHTTPError as e:
         print(f"Error downloading model: {e}")
         return f"Error downloading model: {model_repo}/{model_filename}. ({e.response.status_code})"
    except rt.ONNXRuntimeException as e:
        print(f"Error loading ONNX model: {e}")
        # Attempt to provide more specific feedback if it's a provider issue
        if "CUDAExecutionProvider" in str(e):
             return f"Error loading ONNX model '{model_filename}'. CUDA unavailable or setup issue? Falling back might require restart or different build. Error: {e}"
        return f"Error loading ONNX model '{model_filename}'. Incompatible or corrupted? Error: {e}"
    except Exception as e:
        print(f"Error updating model: {e}")
        return f"Error updating model: {str(e)}"

# --- Main Processing Function (MODIFIED FOR LOGGING) ---
def process_and_update(img, model_repo, model_filename_with_size, history):
    global current_model_repo, current_model_filename, rmbg_model
    
    # --- Pre-checks ---
    if img is None:
        return None, [], history, "generated", "Please upload an image first.", read_log_file() # Return current log
    if rmbg_model is None:
         return None, [], history, "generated", "ERROR: Model not loaded. Update model first.", read_log_file() # Return current log
         
    selected_model_filename, selected_model_size_mb = get_model_details_from_choice(model_filename_with_size) # Use helper
    status_message = ""

    # --- Model Update Check ---
    if model_repo != current_model_repo or selected_model_filename != current_model_filename:
        status_message = update_model(model_repo, model_filename_with_size)
        if "Error" in status_message:
             return None, [], history, "generated", f"Model Update Failed:\n{status_message}", read_log_file() # Return current log
        if rmbg_model is None:
             return None, [], history, "generated", "ERROR: Model failed to load after update.", read_log_file() # Return current log

    # --- Processing & Logging ---
    try:
        start_time = time.time() # Start timer
        mask_img, generated_img_rgba = rmbg_fn(img) # Run inference
        end_time = time.time() # End timer
        
        processing_time = end_time - start_time # Calculate duration
        
        # --- Gather Log Information ---
        timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        h, w = img.shape[:2]
        resolution = f"{w}x{h}"
        # Get the *actually used* provider from the loaded session
        active_provider = rmbg_model.get_providers()[0] 
        
        # Log the event
        log_processing_event(
            timestamp=timestamp,
            repo=current_model_repo, # Use the confirmed current repo
            model_filename=current_model_filename, # Use the confirmed current filename
            model_size_mb=selected_model_size_mb if selected_model_size_mb is not None else 0.0, # Use extracted size
            resolution=resolution,
            provider=active_provider,
            processing_time=processing_time
        )
        
        # --- Prepare Outputs ---
        new_history = history + [generated_img_rgba]
        output_pair = [mask_img, generated_img_rgba]
        current_log_content = read_log_file() # Read updated log
        
        status_message = f"{status_message}\nProcessing complete ({processing_time:.2f}s)".strip()
        
        return generated_img_rgba, output_pair, new_history, "generated", status_message, current_log_content

    except Exception as e:
        print(f"Error during processing: {e}")
        import traceback
        traceback.print_exc()
        # Still return the log content even if processing fails
        return None, [], history, "generated", f"Error during processing: {str(e)}", read_log_file()


# --- UI Interaction Functions (Unchanged toggle_view, clear_all needs slight modification) ---
def toggle_view(view_state, output_pair):
    if not output_pair or len(output_pair) != 2:
        return None, view_state, "View Mask" if view_state == "generated" else "View Generated"
    if view_state == "generated":
        return output_pair[0], "mask", "View Generated"
    else:
        return output_pair[1], "generated", "View Mask"

def clear_all():
    """ Resets inputs, outputs, states, status, but keeps log view """
    # Keeps the log viewer content, as history shouldn't be wiped by clearing inputs
    initial_log_content = read_log_file() # Read log to display upon clearing
    return None, None, [], [], "generated", "Interface cleared.", "View Mask", [], initial_log_content

# --- Gradio UI Definition ---
if __name__ == "__main__":
    initialize_log_file() # Ensure log file exists before launching app

    app = gr.Blocks(css=".gradio-container { max-width: 95% !important; }") # Wider
    with app:
        gr.Markdown("# Image Background Removal (Segmentation) with Logging")
        gr.Markdown("Test ONNX models, view performance logs.")

        with gr.Row():
            # Left Column: Controls and Input
            with gr.Column(scale=2):
                 with gr.Group():
                    gr.Markdown("### Model Selection")
                    model_repo_input = gr.Textbox(value=model_repo_default, label="Hugging Face Repository")
                    model_filename_dropdown = gr.Dropdown(
                        choices=default_onnx_files_with_size,
                        value=default_onnx_files_with_size[0] if default_onnx_files_with_size else "",
                        label="ONNX Model File (.onnx)"
                    )
                    update_btn = gr.Button("🔄 Update/Load Model")
                    model_status_textbox = gr.Textbox(label="Status", value="Initial model loaded." if rmbg_model else "ERROR: Initial model failed to load.", interactive=False, lines=2)

                 gr.Markdown("#### Source Image")
                 input_img = gr.Image(label="Upload Image", type="numpy")

                 with gr.Row():
                     run_btn = gr.Button("▶️ Run Background Removal", variant="primary")
                     clear_btn = gr.Button("🗑️ Clear Inputs/Outputs")

            # Right Column: Output and Logs
            with gr.Column(scale=3):
                 gr.Markdown("#### Output Image")
                 output_img = gr.Image(label="Output", image_mode="RGBA", format="png", type="numpy")
                 toggle_btn = gr.Button("View Mask")

                 gr.Markdown("---")
                 gr.Markdown("### Processing History")
                 history_gallery = gr.Gallery(label="Generated Image History", show_label=False, columns=8, object_fit="contain", height="auto")

                 gr.Markdown("---")
                 gr.Markdown("### Processing Log (`processing_log.csv`)")
                 # Use gr.Code for better viewing of CSV/text data
                 log_display = gr.Code(
                     value=read_log_file(), # Initial content
                     label="Log Viewer",
                     lines=10,
                     interactive=False
                 )
                 # Optional: Add a manual refresh button if auto-update isn't sufficient
                 # refresh_log_btn = gr.Button("🔄 Refresh Log View")

        # Hidden states
        output_pair_state = gr.State([])
        view_state = gr.State("generated")
        history_state = gr.State([])

        # --- Event Listeners ---
        model_repo_input.submit(fn=update_onnx_files, inputs=model_repo_input, outputs=model_filename_dropdown)
        model_repo_input.blur(fn=update_onnx_files, inputs=model_repo_input, outputs=model_filename_dropdown)
        update_btn.click(fn=update_model, inputs=[model_repo_input, model_filename_dropdown], outputs=model_status_textbox)

        # Run includes updating the log display
        run_btn.click(
            fn=process_and_update,
            inputs=[input_img, model_repo_input, model_filename_dropdown, history_state],
            outputs=[output_img, output_pair_state, history_state, view_state, model_status_textbox, log_display] # ADD log_display here
        )

        toggle_btn.click(fn=toggle_view, inputs=[view_state, output_pair_state], outputs=[output_img, view_state, toggle_btn])

        # Clear resets inputs/outputs/status, but re-reads log for display
        clear_btn.click(
            fn=clear_all,
            outputs=[input_img, output_img, output_pair_state, history_state, view_state, model_status_textbox, toggle_btn, history_gallery, log_display] # ADD log_display here
        )

        # Manual log refresh button (optional, as run/clear update it)
        # refresh_log_btn.click(fn=read_log_file, inputs=None, outputs=log_display)

        history_state.change(fn=lambda history: history, inputs=history_state, outputs=history_gallery)

    app.launch(debug=True)