app.py

import gradio as gr
import os
import random
import csv
from pathlib import Path
from datetime import datetime

DATA_DIR = Path("data")
RESULTS_DIR = Path("results")
RESULTS_FILE = RESULTS_DIR / "preferences.csv"
IMAGE_EXTENSIONS = [".png", ".jpg", ".jpeg", ".webp"]

# --- Data Loading ---

def find_image(folder_path: Path, base_name: str) -> Path | None:
    """Finds an image file starting with base_name in a folder."""
    for ext in IMAGE_EXTENSIONS:
        file_path = folder_path / f"{base_name}{ext}"
        if file_path.exists():
            return file_path
    return None

def get_sample_ids() -> list[str]:
    """Scans the data directory for valid sample IDs."""
    sample_ids = []
    if DATA_DIR.is_dir():
        for item in DATA_DIR.iterdir():
            if item.is_dir():
                # Check if required files exist
                prompt_file = item / "prompt.txt"
                input_bg = find_image(item, "input_bg")
                input_fg = find_image(item, "input_fg")
                output_baseline = find_image(item, "baseline")
                output_tficon = find_image(item, "tf-icon")
                if prompt_file.exists() and input_bg and input_fg and output_baseline and output_tficon:
                    sample_ids.append(item.name)
    return sample_ids

def load_sample_data(sample_id: str) -> dict | None:
    """Loads data for a specific sample ID."""
    sample_path = DATA_DIR / sample_id
    if not sample_path.is_dir():
        return None

    prompt_file = sample_path / "prompt.txt"
    input_bg_path = find_image(sample_path, "input_bg")
    input_fg_path = find_image(sample_path, "input_fg")
    output_baseline_path = find_image(sample_path, "baseline")
    output_tficon_path = find_image(sample_path, "tf-icon")

    if not all([prompt_file.exists(), input_bg_path, input_fg_path, output_baseline_path, output_tficon_path]):
        print(f"Warning: Missing files in sample {sample_id}")
        return None

    try:
        prompt = prompt_file.read_text().strip()
    except Exception as e:
        print(f"Error reading prompt for {sample_id}: {e}")
        return None

    return {
        "id": sample_id,
        "prompt": prompt,
        "input_bg": str(input_bg_path),
        "input_fg": str(input_fg_path),
        "output_baseline": str(output_baseline_path),
        "output_tficon": str(output_tficon_path),
    }

# --- State and UI Logic ---

INITIAL_SAMPLE_IDS = get_sample_ids()

def get_next_sample(available_ids: list[str]) -> tuple[dict | None, list[str]]:
    """Selects a random sample ID from the available list."""
    if not available_ids:
        return None, []
    chosen_id = random.choice(available_ids)
    remaining_ids = [id for id in available_ids if id != chosen_id]
    sample_data = load_sample_data(chosen_id)
    return sample_data, remaining_ids


def display_new_sample(state: dict, available_ids: list[str]):
    """Loads and prepares a new sample for display."""
    sample_data, remaining_ids = get_next_sample(available_ids)

    if not sample_data:
        return {
            prompt_display: gr.update(value="No more samples available. Thank you!"),
            input_bg_display: gr.update(value=None, visible=False),
            input_fg_display: gr.update(value=None, visible=False),
            output_a_display: gr.update(value=None, visible=False),
            output_b_display: gr.update(value=None, visible=False),
            choice_button_a: gr.update(visible=False),
            choice_button_b: gr.update(visible=False),
            next_button: gr.update(visible=False),
            status_display: gr.update(value="Completed!"),
            app_state: state,
            available_samples_state: remaining_ids
        }

    outputs = [
        {"model_name": "baseline", "path": sample_data["output_baseline"]},
        {"model_name": "tf-icon", "path": sample_data["output_tficon"]},
    ]
    random.shuffle(outputs)
    output_a = outputs[0]
    output_b = outputs[1]

    state = {
        "current_sample_id": sample_data["id"],
        "output_a_model_name": output_a["model_name"],
        "output_b_model_name": output_b["model_name"],
    }

    return {
        prompt_display: gr.update(value=f"Prompt: {sample_data['prompt']}"),
        input_bg_display: gr.update(value=sample_data["input_bg"], visible=True),
        input_fg_display: gr.update(value=sample_data["input_fg"], visible=True),
        output_a_display: gr.update(value=output_a["path"], visible=True),
        output_b_display: gr.update(value=output_b["path"], visible=True),
        choice_button_a: gr.update(visible=True, interactive=True),
        choice_button_b: gr.update(visible=True, interactive=True),
        next_button: gr.update(visible=False),
        status_display: gr.update(value="Please choose the image you prefer."),
        app_state: state,
        available_samples_state: remaining_ids
    }

def record_preference(choice: str, state: dict, request: gr.Request):
    """Records the user's preference and prepares for the next sample."""
    if not request: # Add a check if request is None
        print("Error: Request object is None. Cannot get session ID.")
        session_id = "unknown_session" # Fallback session ID
    else:
        try:
            session_id = request.client.host # Use IP address as a basic session identifier
        except AttributeError:
             print("Error: request.client is None or has no 'host' attribute.")
             session_id = "unknown_client" # Fallback if client object is weird

    if not state or "current_sample_id" not in state:
        print("Warning: State missing, cannot record preference.")
        return {
            choice_button_a: gr.update(interactive=False),
            choice_button_b: gr.update(interactive=False),
            next_button: gr.update(visible=True, interactive=True),
            status_display: gr.update(value="Error: Session state lost. Click Next Sample."),
            app_state: state # Return unchanged state
        }

    chosen_model_name = state["output_a_model_name"] if choice == "A" else state["output_b_model_name"]

    # Ensure results directory exists
    RESULTS_DIR.mkdir(parents=True, exist_ok=True)

    # Append result to CSV
    file_exists = RESULTS_FILE.exists()
    try:
        with open(RESULTS_FILE, 'a', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            if not file_exists:
                writer.writerow([
                    "timestamp", "session_id", "sample_id",
                    "baseline_displayed_as", "tficon_displayed_as",
                    "chosen_display", "chosen_model_name"
                ]) # Header

            baseline_display = "A" if state["output_a_model_name"] == "baseline" else "B"
            tficon_display = "B" if state["output_a_model_name"] == "baseline" else "A"

            writer.writerow([
                datetime.now().isoformat(),
                session_id,
                state["current_sample_id"],
                baseline_display,
                tficon_display,
                choice, # A or B
                chosen_model_name # baseline or tf-icon
            ])
    except Exception as e:
        print(f"Error writing results: {e}")
        return {
            choice_button_a: gr.update(interactive=False),
            choice_button_b: gr.update(interactive=False),
            next_button: gr.update(visible=True, interactive=True), # Allow user to continue
            status_display: gr.update(value=f"Error saving preference: {e}. Click Next Sample."),
            app_state: state
        }


    # Update UI: disable choice buttons, show next button
    return {
        choice_button_a: gr.update(interactive=False),
        choice_button_b: gr.update(interactive=False),
        next_button: gr.update(visible=True, interactive=True),
        status_display: gr.update(value=f"Preference recorded (Chose {choice}). Click Next Sample."),
        app_state: state # Return unchanged state
    }

# --- New Handler Functions ---
def handle_choice_a(state: dict, request: gr.Request):
    return record_preference("A", state, request)

def handle_choice_b(state: dict, request: gr.Request):
    return record_preference("B", state, request)

# --- Gradio Interface ---

with gr.Blocks(title="Image Composition User Study") as demo:
    gr.Markdown("# Image Composition User Study")
    gr.Markdown(
        "Please look at the input images and the prompt below. "
        "Then, compare the two output images (Output A and Output B) and click the button below the one you prefer."
    )

    # State variables
    app_state = gr.State({}) # Stores current sample info (id, output mapping)
    # Keep track of samples available *for this session*
    available_samples_state = gr.State(INITIAL_SAMPLE_IDS)

    # Displays
    prompt_display = gr.Textbox(label="Prompt", interactive=False)
    status_display = gr.Textbox(label="Status", value="Loading first sample...", interactive=False)

    with gr.Row():
        input_bg_display = gr.Image(label="Input Background", type="filepath", height=300, width=300, interactive=False)
        input_fg_display = gr.Image(label="Input Foreground", type="filepath", height=300, width=300, interactive=False)

    gr.Markdown("---")
    gr.Markdown("## Choose your preferred output:")

    with gr.Row():
        with gr.Column():
            output_a_display = gr.Image(label="Output A", type="filepath", height=400, width=400, interactive=False)
            choice_button_a = gr.Button("Choose Output A", variant="primary")
        with gr.Column():
            output_b_display = gr.Image(label="Output B", type="filepath", height=400, width=400, interactive=False)
            choice_button_b = gr.Button("Choose Output B", variant="primary")

    next_button = gr.Button("Next Sample", visible=False)

    # --- Event Handlers ---

    # Load first sample on page load
    demo.load(
        fn=display_new_sample,
        inputs=[app_state, available_samples_state],
        outputs=[
            prompt_display, input_bg_display, input_fg_display,
            output_a_display, output_b_display,
            choice_button_a, choice_button_b, next_button, status_display,
            app_state, available_samples_state
        ]
    )

    # Handle choice A click - Use the new handler function
    choice_button_a.click(
        fn=handle_choice_a, # Use the dedicated handler
        inputs=[app_state], # Input is still just the state component
        outputs=[choice_button_a, choice_button_b, next_button, status_display, app_state],
        api_name=False,
    )

    # Handle choice B click - Use the new handler function
    choice_button_b.click(
        fn=handle_choice_b, # Use the dedicated handler
        inputs=[app_state], # Input is still just the state component
        outputs=[choice_button_a, choice_button_b, next_button, status_display, app_state],
        api_name=False,
    )

    # Handle next sample click
    next_button.click(
        fn=display_new_sample,
        inputs=[app_state, available_samples_state],
        outputs=[
            prompt_display, input_bg_display, input_fg_display,
            output_a_display, output_b_display,
            choice_button_a, choice_button_b, next_button, status_display,
            app_state, available_samples_state
        ],
        api_name=False,
        # queue=True
    )

if __name__ == "__main__":
    if not INITIAL_SAMPLE_IDS:
        print("Error: No valid samples found in the 'data' directory.")
        print("Please ensure the 'data' directory exists and contains subdirectories")
        print("named like 'sample_id', each with 'prompt.txt', 'input_bg.*',")
        print("'input_fg.*', 'baseline.*', and 'tf-icon.*' files.")
    else:
        print(f"Found {len(INITIAL_SAMPLE_IDS)} samples.")
        print("Starting Gradio app...")
        demo.launch(server_name="0.0.0.0")