update

app.py

@@ -1,82 +1,134 @@
 # app.py
 
 import gradio as gr
-from utils import HF_REPO, ROIS, load_roi_image, load_images, get_hypotheses_for_selection
+from utils import BEST_ACROSS, ROIS, load_roi_image, get_hypotheses_for_selection, load_images_and_views, load_roi_for_selection
 
-# -------- CONFIGURATION --------
 MODELS = ["SAE", "SAE_ICA", "ICA", "PCA", "NMF", "Voxel"]
 
-# Initial values derived from the helper functions after the structure is loaded
 INITIAL_HYPOTHESES = get_hypotheses_for_selection(MODELS[0], ROIS[0])
-INITIAL_ROI_IMAGE = load_roi_image(ROIS[0]) 
+INITIAL_FID_IMAGES, INITIAL_VIEW_IMAGES = load_images_and_views(
+    MODELS[0], ROIS[0], INITIAL_HYPOTHESES[0]
+)
+INITIAL_VIEW_IMAGE = None
+INITIAL_ROI_IMAGE = None
 
-# -------- GRADIO UI AND CALLBACK LOGIC --------
+
+# -------- CALLBACKS --------
 def update_hypotheses(model, roi):
     new_hypotheses = get_hypotheses_for_selection(model, roi)
     return gr.Dropdown(choices=new_hypotheses, value=new_hypotheses[0] if new_hypotheses else None)
 
-with gr.Blocks() as demo:
+
+def update_activating_images(model, roi, hyp):
+    fid_images, brain_images = load_images_and_views(model, roi, hyp)
+
+    # For Across, fetch the ROI image dynamically
+    roi_img = None
+    if roi.lower() == "across":
+        override_roi = BEST_ACROSS.get(model, {}).get(hyp, {}).get("roi")
+        roi_img = load_roi_image(roi_name=override_roi)
+    else:
+        roi_img = load_roi_image(roi_name=roi)
+
+    initial_view = brain_images[0] if brain_images else None
+    return fid_images, brain_images, initial_view, roi_img
+
+
+
+def update_view_on_click(evt: gr.SelectData, brain_images):
+    idx = evt.index
+    if brain_images and idx is not None:
+        return brain_images[idx]
+    return None
+
+
+# -------- UI --------
+with gr.Blocks(fill_width=True) as demo:
+
+    brain_state = gr.State(value=INITIAL_VIEW_IMAGES)
+
     gr.Markdown("## BrainExplore: Visual Concept Explorer")
 
-    # 1. TOP ROW: Dropdowns (Menus)
     with gr.Row():
-        model_dropdown = gr.Dropdown(label="Select Model", choices=MODELS, value=MODELS[0], scale=1)
-        roi_dropdown = gr.Dropdown(label="Select ROI", choices=ROIS, value=ROIS[0], scale=1)
-        
+        model_dropdown = gr.Dropdown(
+            label="Select Model", choices=MODELS, value=MODELS[0], scale=1
+        )
+        roi_dropdown = gr.Dropdown(
+            label="Select ROI", choices=ROIS, value=ROIS[0], scale=1
+        )
         hyp_dropdown = gr.Dropdown(
-            label="Select Hypothesis", 
+            label="Select Hypothesis",
             choices=INITIAL_HYPOTHESES,
-            value=INITIAL_HYPOTHESES[0] if INITIAL_HYPOTHESES else None,
-            scale=1
+            value=INITIAL_HYPOTHESES[0],
+            scale=1,
         )
 
-    # 2. SECOND ROW: Visualizations (ROI Image and Gallery)
     with gr.Row(equal_height=True):
-        
-        roi_image = gr.Image(
-            label="Selected Brain Region (ROI)",
-            value=INITIAL_ROI_IMAGE,
-            interactive=False, 
-            height=600, 
-            scale=1 
-        )
-        
-        image_gallery = gr.Gallery(
-            label="Top Activating Images", 
-            columns=None, rows=1, scale=3, 
-            object_fit="contain", preview=True, height=600
-        )
 
-    # -------- CALLBACKS --------
+        # LEFT — ROI + VIEW image
+        with gr.Column(scale=1):
+            roi_image = gr.Image(
+                label="Selected Brain Region (ROI)",
+                value=None,   # empty initially
+                interactive=False,
+                height=300,
+            )
+
+            view_image = gr.Image(
+                label="Viewing Angle",
+                value=None,   # empty initially
+                interactive=False,
+                height=300,
+            )
+
+        # RIGHT — gallery
+        with gr.Column(scale=3):
+            image_gallery = gr.Gallery(
+                label="Top Activating Images",
+                columns=None,
+                rows=1,
+                object_fit="contain",
+                preview=True,
+                height=600,
+                allow_preview=True,
+                elem_id="gallery"
+            )
+
+
+    # -------- CALLBACK WIRING --------
+
+    # Update hypotheses when model/ROI changes
     model_dropdown.change(
-        fn=update_hypotheses, 
-        inputs=[model_dropdown, roi_dropdown], 
-        outputs=[hyp_dropdown], 
-        queue=False 
-    ).then(
-        fn=load_images, 
-        inputs=[model_dropdown, roi_dropdown, hyp_dropdown], 
-        outputs=[image_gallery]
+        fn=update_hypotheses,
+        inputs=[model_dropdown, roi_dropdown],
+        outputs=[hyp_dropdown],
     )
 
-    # ROI Change: Update Hypotheses, THEN update ROI Image, THEN update Activating Images
     roi_dropdown.change(
-        fn=update_hypotheses, 
-        inputs=[model_dropdown, roi_dropdown], 
-        outputs=[hyp_dropdown],
-        queue=False
+    fn=update_hypotheses,
+    inputs=[model_dropdown, roi_dropdown],
+    outputs=[hyp_dropdown],
     ).then(
-        fn=load_roi_image,
-        inputs=[roi_dropdown],
-        outputs=[roi_image]
-    ).then(
-        fn=load_images, 
-        inputs=[model_dropdown, roi_dropdown, hyp_dropdown], 
-        outputs=[image_gallery]
+        fn=lambda model, roi, hyp: load_roi_for_selection(model, roi, hyp),
+        inputs=[model_dropdown, roi_dropdown, hyp_dropdown],
+        outputs=[roi_image],
+    )
+
+    # Update gallery + brain state + initial view image
+    for cb in [model_dropdown, roi_dropdown, hyp_dropdown]:
+        cb.change(
+            fn=update_activating_images,
+            inputs=[model_dropdown, roi_dropdown, hyp_dropdown],
+            outputs=[image_gallery, brain_state, view_image, roi_image],
+        )
+
+    # Update view image when clicking a different gallery item
+    image_gallery.select(
+        fn=update_view_on_click,
+        inputs=[brain_state],
+        outputs=[view_image],
     )
 
-    # Hypothesis Change: Update Activating Images only
-    hyp_dropdown.change(fn=load_images, inputs=[model_dropdown, roi_dropdown, hyp_dropdown], outputs=[image_gallery])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

utils.py

@@ -2,13 +2,31 @@ import json
 from datasets import load_dataset
 from PIL import Image
 from huggingface_hub import hf_hub_download
+from pathlib import Path
 import os
 
 # -------- CONFIG & STATE --------
 HF_REPO = "mcosarinsky/BrainExplore-Images"
 STRUCTURE_DATA = None
-ROIS = ["EBA","FBA-1","FBA-2","FFA-1","FFA-2","OFA","OPA","OWFA","PPA","RSC","VWFA-1","VWFA-2","hV4"]
+ROIS = ["Across","EBA","FBA-1","FBA-2","FFA-1","FFA-2","OFA","OPA","OWFA","PPA","RSC","VWFA-1","VWFA-2","hV4"]
 
+# Load best_across.json once
+BEST_ACROSS = None
+
+def load_best_across():
+    global BEST_ACROSS
+    if BEST_ACROSS is not None:
+        return BEST_ACROSS
+    try:
+        path = hf_hub_download(repo_id=HF_REPO, filename="best_across.json", repo_type="dataset")
+        with open(path, "r") as f:
+            BEST_ACROSS = json.load(f)
+        return BEST_ACROSS
+    except Exception as e:
+        print(f"Error loading best_across.json: {e}")
+        BEST_ACROSS = {}
+        return BEST_ACROSS
+    
 # -------- 1. JSON STRUCTURE HELPERS --------
 def load_structure(repo_id, filename="structure.json"):
     """
@@ -36,18 +54,19 @@ def load_structure(repo_id, filename="structure.json"):
 
 def get_hypotheses_for_selection(model, roi):
     """
-    Retrieves the list of hypotheses based on the selected model and ROI 
-    from the pre-loaded STRUCTURE_DATA.
+    Normal ROIs: load from structure.json
+    Across ROI: load from best_across.json
     """
-    if not STRUCTURE_DATA:
-        return []
-        
-    try:
-        hypotheses = STRUCTURE_DATA.get(model, {}).get(roi, [])
-        return sorted(hypotheses)
-    except Exception as e:
-        print(f"Error accessing structure data for {model}/{roi}: {e}")
+    model = model  # assume already mapped
+    if roi.lower() == "across":
+        if model in BEST_ACROSS:
+            return sorted(BEST_ACROSS[model].keys())
         return []
+    else:
+        # existing STRUCTURE_DATA logic
+        if not STRUCTURE_DATA:
+            return []
+        return sorted(STRUCTURE_DATA.get(model, {}).get(roi, []))
 
 # -------- 2. IMAGE PROCESSING HELPERS --------
 def pad_image(pil_img, padding_x=10, padding_y=0, color=(255, 255, 255)):
@@ -61,28 +80,63 @@ def pad_image(pil_img, padding_x=10, padding_y=0, color=(255, 255, 255)):
     new_img.paste(pil_img, (padding_x, padding_y))
     return new_img
 
-
-# -------- 3. DATA LOADING HELPERS --------
-def load_images(model, roi, hyp):
+def load_images_and_views(model, roi, hyp):
     """
-    Loads top activating images for the given parameters and applies padding.
+    Normal ROI: same as before (load FID & brain from folder)
+    Across ROI: load single FID and brain as per best_across.json
     """
-    folder_path = f"{model}/{roi}/{hyp}"
-    ds = load_dataset(HF_REPO, data_dir=folder_path, split="train", streaming=True)
-    image_paths = [img["image"] for img in ds.take(2)]
-    padded_images = [pad_image(img, padding_x=10, padding_y=0) for img in image_paths]
-    return padded_images
+    if roi.lower() == "across":
+        model_dict = BEST_ACROSS.get(model, {})
+        hyp_entry = model_dict.get(hyp)
+        if not hyp_entry:
+            return [], []
+
+        roi_name = hyp_entry["roi"]
+        fid_file = hyp_entry["fid_file"]
+        brain_file = hyp_entry.get("brain_file")
+
+        ds = load_dataset(HF_REPO, data_dir=f"{model}/{roi_name}/{hyp}", split="train")
+        fid_image = None
+        brain_image = None
+
+        for item in ds:
+            img = item["image"]
+            img_name = Path(img.filename).name
+            if img_name == fid_file:
+                fid_image = pad_image(img, padding_x=10, padding_y=0)
+            elif brain_file and img_name == brain_file:
+                brain_image = img
 
-def load_roi_image(roi_name):
+        return [fid_image] if fid_image else [], [brain_image] if brain_image else []
+
+    else:
+        # normal loading
+        folder_path = f"{model}/{roi}/{hyp}"
+        ds = load_dataset(HF_REPO, data_dir=folder_path, split="train")
+        fid_images, brain_images = [], []
+        for item in ds:
+            img = item["image"]
+            img_name = Path(img.filename).name
+            if "fid" in img_name.lower():
+                fid_images.append(pad_image(img, padding_x=10, padding_y=0))
+            elif "brain" in img_name.lower():
+                brain_images.append(img)
+        return fid_images, brain_images
+
+def load_roi_image(roi_name, override_roi=None):
     """
-    Downloads the ROI PNG from the Hugging Face Hub, looking inside the 'ROI/' folder,
-    and loads it as a PIL Image.
+    Downloads the ROI PNG from the Hugging Face Hub and loads it as a PIL Image.
+    
+    If roi_name is 'Across' and override_roi is provided, it loads that ROI instead.
     """
     if not roi_name:
         return None
-        
-    file_name = f"ROI/{roi_name}.png"
-    
+
+    # Use override if provided (for Across)
+    actual_roi = override_roi if override_roi else roi_name
+
+    file_name = f"ROI/{actual_roi}.png"
+
     try:
         local_path = hf_hub_download(
             repo_id=HF_REPO,
@@ -91,9 +145,19 @@ def load_roi_image(roi_name):
         )
         return Image.open(local_path)
     except Exception as e:
-        # A 404 error here likely means the file doesn't exist under that path
-        print(f"Error loading ROI image {roi_name} via hf_hub_download. Path used: {file_name}. Details: {e}")
+        print(f"Error loading ROI image {actual_roi} via hf_hub_download. Path used: {file_name}. Details: {e}")
         return None
+    
+def load_roi_for_selection(model, roi, hyp):
+    """
+    Loads ROI image. If 'Across' is selected, use the original ROI
+    from BEST_ACROSS for the selected hypothesis.
+    """
+    override_roi = None
+    if roi.lower() == "across":
+        override_roi = BEST_ACROSS.get(model, {}).get(hyp, {}).get("roi")
+    return load_roi_image(roi_name=roi, override_roi=override_roi)
+
 
-# Perform initial structure load for dynamic dependency resolution in app.py
-load_structure(HF_REPO)
+load_structure(HF_REPO)
+load_best_across()