File size: 5,336 Bytes
7d0bd10
 
 
 
39f341b
7d0bd10
 
 
 
 
39f341b
7d0bd10
39f341b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7d0bd10
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39f341b
 
7d0bd10
39f341b
 
 
 
7d0bd10
39f341b
 
 
 
 
7d0bd10
 
 
 
 
 
 
 
 
 
 
 
 
39f341b
7d0bd10
39f341b
 
7d0bd10
39f341b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7d0bd10
39f341b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7d0bd10
39f341b
 
 
7d0bd10
 
 
39f341b
 
 
 
 
 
7d0bd10
 
 
 
 
 
 
 
39f341b
7d0bd10
39f341b
 
 
 
 
 
 
 
 
 
 
7d0bd10
39f341b
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
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 = ["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"):
    """
    Downloads the structure.json file from the Hugging Face Hub and loads it into memory.
    """
    global STRUCTURE_DATA
    if STRUCTURE_DATA is not None:
        return STRUCTURE_DATA
        
    try:
        local_path = hf_hub_download(
            repo_id=repo_id,
            filename=filename,
            repo_type="dataset"
        )
        with open(local_path, 'r') as f:
            STRUCTURE_DATA = json.load(f)
        
        print(f"Successfully loaded {filename}.")
        return STRUCTURE_DATA
        
    except Exception as e:
        print(f"CRITICAL ERROR: Could not load {filename} via hf_hub_download: {e}")
        return {}

def get_hypotheses_for_selection(model, roi):
    """
    Normal ROIs: load from structure.json
    Across ROI: load from best_across.json
    """
    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)):
    """
    Adds padding to the sides of a PIL image.
    """
    old_width, old_height = pil_img.size
    new_width = old_width + 2 * padding_x
    new_height = old_height + 2 * padding_y
    new_img = Image.new(pil_img.mode, (new_width, new_height), color)
    new_img.paste(pil_img, (padding_x, padding_y))
    return new_img

def load_images_and_views(model, roi, hyp):
    """
    Normal ROI: same as before (load FID & brain from folder)
    Across ROI: load single FID and brain as per best_across.json
    """
    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

        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 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

    # 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,
            filename=file_name,
            repo_type="dataset"
        )
        return Image.open(local_path)
    except Exception as 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)


load_structure(HF_REPO)
load_best_across()