Spaces:
Running
Running
| 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() |