Create app.py

app.py

@@ -0,0 +1,297 @@
+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")