Update app.py

app.py

@@ -1,285 +1,247 @@
+import streamlit as st
+import pandas as pd
 import os
+from pathlib import Path
+from datetime import datetime
 import random
-import datetime
-import pandas as pd
-import gradio as gr
-
-ROENTGEN_DIR = "roentgen"
-PEDISYNTH_DIR = "pedisynth"
-RESULTS_DIR = "results"
-
-
-def load_pairs():
-    """Load matching image pairs from both folders, keyed by filename."""
-    ro_files = {f for f in os.listdir(ROENTGEN_DIR)
-                if os.path.isfile(os.path.join(ROENTGEN_DIR, f))}
-    pe_files = {f for f in os.listdir(PEDISYNTH_DIR)
-                if os.path.isfile(os.path.join(PEDISYNTH_DIR, f))}
-
-    common = sorted(ro_files & pe_files)
+from PIL import Image
+import io
+
+# Configure page
+st.set_page_config(
+    page_title="X-Ray Quality Evaluation",
+    page_icon="🏥",
+    layout="wide"
+)
+
+# Initialize session state
+if "session_active" not in st.session_state:
+    st.session_state.session_active = False
+if "image_pairs" not in st.session_state:
+    st.session_state.image_pairs = []
+if "current_index" not in st.session_state:
+    st.session_state.current_index = 0
+if "evaluations" not in st.session_state:
+    st.session_state.evaluations = []
+if "doctor_name" not in st.session_state:
+    st.session_state.doctor_name = ""
+
+# Helper functions
+def get_image_files(folder_path):
+    """Get all image files from a folder."""
+    valid_extensions = {'.jpg', '.jpeg', '.png', '.gif'}
+    files = {}
+    if os.path.exists(folder_path):
+        for file in os.listdir(folder_path):
+            if Path(file).suffix.lower() in valid_extensions:
+                # Extract condition from filename (everything except the last number)
+                parts = file.rsplit(' ', 1)
+                condition = parts[0] if len(parts) > 1 else file
+                if condition not in files:
+                    files[condition] = []
+                files[condition].append(file)
+    return files
+
+def create_image_pairs():
+    """Create paired images from roentgen and pedisynth folders."""
+    roentgen_files = get_image_files("roentgen")
+    pedisynth_files = get_image_files("pedisynth")
+    
     pairs = []
-
-    for fname in common:
-        base, _ext = os.path.splitext(fname)
-        parts = base.split()
-        # "{condition} {number}" -> condition is everything except last token
-        condition = " ".join(parts[:-1]) if len(parts) > 1 else base
-
-        pairs.append(
-            {
-                "filename": fname,
-                "condition": condition,
-                "roentgen_path": os.path.join(ROENTGEN_DIR, fname),
-                "pedisynth_path": os.path.join(PEDISYNTH_DIR, fname),
-            }
-        )
-    return pairs
-
-
-ALL_PAIRS = load_pairs()
-os.makedirs(RESULTS_DIR, exist_ok=True)
-
-
-def start_session(doctor_id):
-    """Initialize a new randomized session for this doctor."""
-    if not ALL_PAIRS:
-        status = "No matching image pairs found. Check your roentgen/ and pedisynth/ folders."
-        return (
-            None,  # img_a
-            None,  # img_b
-            "",    # condition text
-            "",    # progress text
-            None,  # state
-            None,  # results file
-            status,
-            gr.update(value=None),  # clear choice
-        )
-
-    if not doctor_id:
-        doctor_id = "anonymous"
-
-    # Randomize order of pairs
-    pairs = ALL_PAIRS.copy()
+    
+    # Match files by condition across folders
+    for condition in roentgen_files:
+        if condition in pedisynth_files:
+            for roent_file in roentgen_files[condition]:
+                for pedi_file in pedisynth_files[condition]:
+                    # Check if they have the same base name
+                    roent_base = Path(roent_file).stem
+                    pedi_base = Path(pedi_file).stem
+                    if roent_base == pedi_base:
+                        # Randomly decide which image goes left and which goes right
+                        is_roentgen_left = random.choice([True, False])
+                        pairs.append({
+                            'roentgen': roent_file,
+                            'pedisynth': pedi_file,
+                            'condition': condition,
+                            'base_name': roent_base,
+                            'roentgen_left': is_roentgen_left
+                        })
+    
+    # Randomize order
     random.shuffle(pairs)
+    return pairs
 
-    # Randomize which model appears on left (A) vs right (B)
-    session_pairs = []
-    for p in pairs:
-        if random.random() < 0.5:
-            left_model = "roentgen"
-            right_model = "pedisynth"
-            left_path = p["roentgen_path"]
-            right_path = p["pedisynth_path"]
+def load_image(folder, filename):
+    """Load image from folder."""
+    filepath = os.path.join(folder, filename)
+    try:
+        return Image.open(filepath)
+    except Exception as e:
+        st.error(f"Error loading image {filename}: {e}")
+        return None
+
+def save_evaluations(evaluations, doctor_name):
+    """Save evaluations to CSV."""
+    df = pd.DataFrame(evaluations)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"evaluations_{doctor_name}_{timestamp}.csv"
+    
+    # Save to outputs folder if it exists, otherwise current directory
+    output_dir = "outputs"
+    os.makedirs(output_dir, exist_ok=True)
+    filepath = os.path.join(output_dir, filename)
+    
+    df.to_csv(filepath, index=False)
+    return filepath, filename
+
+# Main UI
+st.title("🏥 X-Ray Quality Evaluation System")
+st.markdown("**Pairwise Preference Test for Synthetic X-Ray Images**")
+
+# Sidebar - Session management
+with st.sidebar:
+    st.header("Session Control")
+    
+    if not st.session_state.session_active:
+        doctor_name = st.text_input("Enter your name:", key="doc_input")
+        
+        if st.button("Start Session", type="primary"):
+            if doctor_name.strip():
+                st.session_state.doctor_name = doctor_name
+                st.session_state.image_pairs = create_image_pairs()
+                st.session_state.session_active = True
+                st.session_state.current_index = 0
+                st.session_state.evaluations = []
+                st.rerun()
+            else:
+                st.error("Please enter your name to start a session.")
+    else:
+        st.success(f"✓ Session active for: {st.session_state.doctor_name}")
+        progress = st.session_state.current_index / len(st.session_state.image_pairs)
+        st.progress(progress)
+        st.metric("Progress", f"{st.session_state.current_index}/{len(st.session_state.image_pairs)}")
+        
+        if st.button("End Session & Download Results", type="primary"):
+            if st.session_state.evaluations:
+                filepath, filename = save_evaluations(
+                    st.session_state.evaluations,
+                    st.session_state.doctor_name
+                )
+                st.success(f"✓ Results saved: {filename}")
+                
+                # Provide download link
+                with open(filepath, 'r') as f:
+                    csv_data = f.read()
+                st.download_button(
+                    label="Download CSV",
+                    data=csv_data,
+                    file_name=filename,
+                    mime="text/csv"
+                )
+                st.session_state.session_active = False
+                st.rerun()
+            else:
+                st.warning("No evaluations recorded yet.")
+
+# Main content
+if not st.session_state.session_active:
+    st.info("👈 Enter your name in the sidebar and click 'Start Session' to begin.")
+else:
+    if st.session_state.current_index >= len(st.session_state.image_pairs):
+        st.success("✅ Evaluation complete! All images have been reviewed.")
+        if st.button("End Session & Download Results"):
+            if st.session_state.evaluations:
+                filepath, filename = save_evaluations(
+                    st.session_state.evaluations,
+                    st.session_state.doctor_name
+                )
+                st.success(f"Results saved: {filename}")
+                with open(filepath, 'r') as f:
+                    csv_data = f.read()
+                st.download_button(
+                    label="Download CSV",
+                    data=csv_data,
+                    file_name=filename,
+                    mime="text/csv"
+                )
+                st.session_state.session_active = False
+                st.rerun()
+    else:
+        current_pair = st.session_state.image_pairs[st.session_state.current_index]
+        
+        st.subheading(f"Image {st.session_state.current_index + 1} of {len(st.session_state.image_pairs)}")
+        st.write(f"**Condition:** {current_pair['condition']}")
+        
+        # Display images side by side (blinded - no labels about source)
+        col1, col2 = st.columns(2)
+        
+        # Determine which image goes where
+        if current_pair['roentgen_left']:
+            left_image = load_image("roentgen", current_pair['roentgen'])
+            right_image = load_image("pedisynth", current_pair['pedisynth'])
+            left_source = "roentgen"
+            right_source = "pedisynth"
         else:
-            left_model = "pedisynth"
-            right_model = "roentgen"
-            left_path = p["pedisynth_path"]
-            right_path = p["roentgen_path"]
-
-        session_pairs.append(
-            {
-                **p,
-                "left_model": left_model,
-                "right_model": right_model,
-                "left_path": left_path,
-                "right_path": right_path,
-            }
-        )
-
-    state = {
-        "doctor_id": doctor_id,
-        "pairs": session_pairs,
-        "idx": 0,
-        "records": [],
-    }
-
-    first = session_pairs[0]
-    condition_text = f"**Condition:** {first['condition']}"
-    progress_text = f"Pair 1 of {len(session_pairs)}"
-    status = f"Session started for doctor: {doctor_id}"
-
-    return (
-        first["left_path"],
-        first["right_path"],
-        condition_text,
-        progress_text,
-        state,
-        None,  # no CSV yet
-        status,
-        gr.update(value=None),  # clear choice
-    )
-
-
-def next_pair(choice, state):
-    """Record the doctor's choice and move to the next pair."""
-    # If no active session
-    if state is None or not state.get("pairs"):
-        status = "No active session. Please start a session first."
-        return (
-            None,
-            None,
-            "",
-            "",
-            state,
-            None,
-            status,
-            gr.update(value=None),
-        )
-
-    idx = state["idx"]
-    pairs = state["pairs"]
-
-    # Ensure choice is made
-    if not choice:
-        status = "Please select A or B before proceeding."
-        cur_pair = pairs[idx]
-        condition_text = f"**Condition:** {cur_pair['condition']}"
-        progress_text = f"Pair {idx + 1} of {len(pairs)}"
-        return (
-            cur_pair["left_path"],
-            cur_pair["right_path"],
-            condition_text,
-            progress_text,
-            state,
-            None,
-            status,
-            gr.update(),  # keep current choice
-        )
-
-    cur_pair = pairs[idx]
-
-    chosen_side = choice
-    chosen_model = cur_pair["left_model"] if choice == "A" else cur_pair["right_model"]
-
-    record = {
-        "doctor_id": state["doctor_id"],
-        "timestamp_utc": datetime.datetime.utcnow().isoformat(),
-        "pair_index": idx,
-        "filename": cur_pair["filename"],
-        "condition": cur_pair["condition"],
-        "left_model": cur_pair["left_model"],
-        "right_model": cur_pair["right_model"],
-        "left_image_path": cur_pair["left_path"],
-        "right_image_path": cur_pair["right_path"],
-        "chosen_side": chosen_side,
-        "chosen_model": chosen_model,
-    }
-    state["records"].append(record)
-
-    # Advance index
-    state["idx"] += 1
-    idx = state["idx"]
-
-    # If we've reached the end, save CSV
-    if idx >= len(pairs):
-        df = pd.DataFrame(state["records"])
-        ts = datetime.datetime.utcnow().strftime("%Y%m%d_%H%M%S")
-        safe_doctor_id = "".join(
-            c for c in state["doctor_id"] if c.isalnum() or c in ("_", "-")
-        )
-        csv_name = f"session_{safe_doctor_id}_{ts}.csv"
-        csv_path = os.path.join(RESULTS_DIR, csv_name)
-        df.to_csv(csv_path, index=False)
-
-        status = (
-            f"Session complete. Saved {len(state['records'])} comparisons "
-            f"to `{csv_name}`."
-        )
-
-        # Optionally, reset session index or leave state as-is for audit
-        return (
-            None,             # img_a
-            None,             # img_b
-            "Session finished.",  # condition text
-            "",               # progress
-            state,
-            csv_path,         # file for download
-            status,
-            gr.update(value=None),
-        )
-
-    # Otherwise, show next pair
-    next_p = pairs[idx]
-    condition_text = f"**Condition:** {next_p['condition']}"
-    progress_text = f"Pair {idx + 1} of {len(pairs)}"
-    status = f"Recorded choice for pair {idx}."
-
-    return (
-        next_p["left_path"],
-        next_p["right_path"],
-        condition_text,
-        progress_text,
-        state,
-        None,
-        status,
-        gr.update(value=None),  # clear choice
-    )
-
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Synthetic Chest X-ray A/B Evaluation")
-
-    with gr.Row():
-        doctor_id = gr.Textbox(
-            label="Doctor ID (or initials)",
-            placeholder="e.g., dr_smith",
-        )
-        start_button = gr.Button("Start New Session", variant="primary")
-
-    state = gr.State()
-
-    with gr.Row():
-        condition_text = gr.Markdown("")
-        progress_text = gr.Markdown("")
-
-    with gr.Row():
-        img_a = gr.Image(label="Image A", interactive=False)
-        img_b = gr.Image(label="Image B", interactive=False)
-
-    choice = gr.Radio(
-        choices=["A", "B"],
-        label="Which image better matches the stated condition?",
-        interactive=True,
-    )
-
-    next_button = gr.Button("Next")
-
-    results_file = gr.File(
-        label="Download session CSV",
-        interactive=False,
-    )
-    status_text = gr.Markdown("")
-
-    # Wire callbacks
-    start_button.click(
-        fn=start_session,
-        inputs=doctor_id,
-        outputs=[
-            img_a,
-            img_b,
-            condition_text,
-            progress_text,
-            state,
-            results_file,
-            status_text,
-            choice,
-        ],
-    )
-
-    next_button.click(
-        fn=next_pair,
-        inputs=[choice, state],
-        outputs=[
-            img_a,
-            img_b,
-            condition_text,
-            progress_text,
-            state,
-            results_file,
-            status_text,
-            choice,
-        ],
-    )
-
-if __name__ == "__main__":
-    demo.launch()
+            left_image = load_image("pedisynth", current_pair['pedisynth'])
+            right_image = load_image("roentgen", current_pair['roentgen'])
+            left_source = "pedisynth"
+            right_source = "roentgen"
+        
+        with col1:
+            st.write("### Option A")
+            if left_image:
+                st.image(left_image, use_column_width=True)
+            else:
+                st.error("Could not load image A")
+        
+        with col2:
+            st.write("### Option B")
+            if right_image:
+                st.image(right_image, use_column_width=True)
+            else:
+                st.error("Could not load image B")
+        
+        # Preference selection
+        st.markdown("---")
+        st.write("**Which image is of better quality?**")
+        
+        col_a, col_b, col_equal = st.columns(3)
+        
+        with col_a:
+            if st.button("👈 Prefer Option A", use_container_width=True, key="btn_a"):
+                preferred_source = left_source
+                st.session_state.evaluations.append({
+                    'timestamp': datetime.now().isoformat(),
+                    'condition': current_pair['condition'],
+                    'base_name': current_pair['base_name'],
+                    'preference': preferred_source.capitalize(),
+                    'doctor': st.session_state.doctor_name
+                })
+                st.session_state.current_index += 1
+                st.rerun()
+        
+        with col_b:
+            if st.button("👉 Prefer Option B", use_container_width=True, key="btn_b"):
+                preferred_source = right_source
+                st.session_state.evaluations.append({
+                    'timestamp': datetime.now().isoformat(),
+                    'condition': current_pair['condition'],
+                    'base_name': current_pair['base_name'],
+                    'preference': preferred_source.capitalize(),
+                    'doctor': st.session_state.doctor_name
+                })
+                st.session_state.current_index += 1
+                st.rerun()
+        
+        with col_equal:
+            if st.button("↔️ Equal Quality", use_container_width=True, key="btn_equal"):
+                st.session_state.evaluations.append({
+                    'timestamp': datetime.now().isoformat(),
+                    'condition': current_pair['condition'],
+                    'base_name': current_pair['base_name'],
+                    'preference': 'Equal',
+                    'doctor': st.session_state.doctor_name
+                })
+                st.session_state.current_index += 1
+                st.rerun()