DataSetManager.py

import hashlib
import random
import numpy as np
import  os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from dataLoaderLazyADNI import load_adni

OLD_ROOTDATASET=os.environ.get("OLD_ROOTDATASET","/scratch")
old_rootDataset=OLD_ROOTDATASET
NEW_ROOTDATASET=os.environ.get("NEW_ROOTDATASET","/shared/neuroimaging")
NEW_ROOTDATASET_FAST=os.environ.get("NEW_ROOTDATASET_FAST","/shared/neuroimaging")
new_rootDataset=NEW_ROOTDATASET
new_rootDataset_FAST=NEW_ROOTDATASET_FAST

FOMO_PATH=os.environ.get("FOMO_PATHS_JSONL","/scratch/progetto/FomoPaths.jsonl")
OPENNEURO_JSON=os.environ.get("OPENNEURO_PATHS_JSON","/scratch/progetto/qrawopenneuro_new.json")
class DatasetManager:
    def __init__(self, downsample ):

        self.must_downsample = downsample
        self.fomo_path = Path(FOMO_PATH)
        self.openneuro_path = Path(OPENNEURO_JSON)

        self.raw_data = []
        self.subject_view = []
        self.hash_map = {}

        self._load_and_process()

    def _load_jsonl(self, jsonl_path):
        data = []
        with open(jsonl_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    data.append(json.loads(line))
        print(f"[INFO] Caricate {len(data)} voci da JSONL: {jsonl_path}")
        return data

    def _load_json(self, json_path):
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        print(f"[INFO] Caricate {len(data)} voci da JSON: {json_path}")
        return data

    def _extract_paths(self, dataset_json, downsample=True):
        result = []
        for subject_entry in dataset_json:
            subject_id = subject_entry.get("subject")
            for item in subject_entry.get("data", []):

                img_path=item.get("image_path")
                if img_path.startswith(f"{old_rootDataset}/OpenMind"):
                    img_path = img_path.replace(old_rootDataset, new_rootDataset_FAST, 1)
                else:
                    img_path = img_path.replace(old_rootDataset, new_rootDataset)
                img_path=img_path.replace("Fomo","FOMO")
                if downsample:
                    filename = img_path.split('/')[-1]
                    pathroot = img_path.replace(filename, '')
                    img_path = pathroot + "downsampled_" + filename.replace('.nii.gz', '.pt')
                result.append({
                    "subject": subject_id,
                    "image_path": img_path,
                    "mask_path": item.get("mask_path")
                })
        return result

    def _hash_subject(self, subject):
        return hashlib.md5(subject.encode('utf-8')).hexdigest()

    def _build_subject_view(self):
        seen = set()
        view = []
        for entry in self.raw_data:
            subject = entry["subject"]
            if subject not in seen:
                seen.add(subject)
                view.append({
                    "subject": subject,
                    "hash": self._hash_subject(subject)
                })
        return view

    def _build_hash_map(self):
        subject_map = {}
        for entry in self.raw_data:
            subject = entry["subject"]
            h = self._hash_subject(subject)
            # In _build_hash_map:
            if h not in subject_map:
                subject_map[h] = {}
            if subject not in subject_map[h]:
                subject_map[h][subject] = []
            subject_map[h][subject].append({
                "subject": subject,
                "image_path": entry["image_path"],
            })
        return subject_map

    def _load_and_process(self):
        fomo_data = self._load_jsonl(self.fomo_path)
        openneuro_data = self._load_json(self.openneuro_path)

        print("[INFO] Estrazione dei path...")
        self.raw_data = self._extract_paths(fomo_data,downsample=self.must_downsample) + self._extract_paths(openneuro_data,downsample=self.must_downsample)
        print(f"[INFO] Totale elementi combinati: {len(self.raw_data)}")
        # print(self.raw_data)
        self.subject_view = self._build_subject_view()
        self.hash_map = self._build_hash_map()
        print(f"[INFO] Totale soggetti unici: {len(self.subject_view)}")

    def get_subject_view(self):
        return self.subject_view

    def get_entries_by_hash(self, subject_hash):
        return self.hash_map.get(subject_hash)

    def get_all_data(self):
        return self.raw_data




import copy


from dataLoaderLazyOpenBHB import *

class customDataLoader(DatasetManager):

    def __init__(self, train_ratio=0.8, test_ratio=0.2, batch_size=32, seed=None,
                 max_same_batch_positives=4, max_real_same_batch_positives=2,
                 onSubjects=False, second_chance=0, ddp_enabled=True,downsample=True):
        super().__init__(downsample=downsample)

        self.second_chance = second_chance
        self.max_real_same_batch_positives = max_real_same_batch_positives
        self.max_same_batch_positives = max_same_batch_positives
        self.test_batches = None
        self.train_batches = None
        self.count_warn = 0
        if train_ratio + test_ratio > 1:
            raise ValueError("Train and test ratios must sum to less than or equal to 1.0")

        self.train_ratio = train_ratio
        self.test_ratio = test_ratio
        self.batch_size = batch_size
        self.train_set = []
        self.test_set = []
        self.train_hashes = []
        self.test_hashes = []
        self.hash_map_copy = []
        self.subject_view_copy = []
        self.ddp_enabled = ddp_enabled

        if self.ddp_enabled and dist.is_initialized():
            self.rank = dist.get_rank()
            self.world_size = dist.get_world_size()
        else:
            self.rank = 0
            self.world_size = 1

        if self.ddp_enabled and dist.is_initialized():

            if self.rank == 0:
                self._split_data(seed, onSubjects)

            # Broadcast the split dataset
            payload = (self.hash_map, self.subject_view, self.train_set) if self.rank == 0 else None
            obj_list = [payload]
            dist.broadcast_object_list(obj_list, src=0)
            self.hash_map, self.subject_view, self.train_set = obj_list[0]


        self.hash_map_original = copy.deepcopy(self.hash_map)
        self.subject_view_original = copy.deepcopy(self.subject_view)


        if self.ddp_enabled and dist.is_initialized():
            self._apply_ddp_split()
        else:

            self._split_data(seed, onSubjects)
            self.hash_map_original = copy.deepcopy(self.hash_map)
            self.subject_view_original = copy.deepcopy(self.subject_view)




    def _apply_ddp_split(self):

        self._apply_ddp_split_temp(self.rank, self.world_size, self.hash_map_original)

    def _apply_ddp_split_temp(self, rank, world_size, hash_map_to_split):
        all_entries = [(h, subj, e) for h, subj_dict in hash_map_to_split.items()
                       for subj, entries in subj_dict.items() for e in entries]

        all_entries.sort(key=lambda x: x[1])
        total_images = len(all_entries)
        images_per_rank = (total_images + world_size - 1) // world_size
        start = rank * images_per_rank
        end = min(start + images_per_rank, total_images)
        shard = all_entries[start:end]

        new_hash_map = {}
        new_subject_view = []
        seen_subjects = set()

        for h, subj, entry in shard:
            if h not in new_hash_map:
                new_hash_map[h] = {}
            if subj not in new_hash_map[h]:
                new_hash_map[h][subj] = []
                if subj not in seen_subjects:
                    new_subject_view.append({"subject": subj, "hash": h})
                    seen_subjects.add(subj)
            new_hash_map[h][subj].append(entry)

        self.hash_map = new_hash_map
        self.train_set = new_subject_view
        self.subject_view = new_subject_view

        print(f"[Rank {self.rank}] Got {len(shard)} images out of {total_images} total (world_size={world_size})")

    def random_add_n_pop_from_old_hashmap(self,rand_subject,h):
        entries = self.hash_map.get(h, {}).get(rand_subject, [])

        random_entry=random.choice(entries)

        entries.remove(random_entry)



        return random_entry, len(entries)==0


    def _split_data(self, seed, on_subjects):
        if seed is not None:
            random.seed(seed)
        else:
            raise Exception("Seed is None")

        if on_subjects:
            subjects = self.subject_view.copy()
            random.shuffle(subjects)

            total = len(subjects)
            train_cutoff = int(total * self.train_ratio)


            self.train_set = subjects[:train_cutoff]


            self.train_hashes = {s["hash"] for s in self.train_set}
            self.test_hashes = {s["hash"] for s in self.test_set}
            print(f"[INFO] Train set: {len(self.train_set)} subjects")

        else:
            subjects = self.subject_view.copy()
            total_images = len(self.raw_data)
            image_to_generate = int(total_images * self.train_ratio)

            print("Target train images:", image_to_generate)

            if self.train_ratio <= 0.5:
                # --- sample & add  ---
                new_subject_view = []
                new_hash_map = {}
                random_element = None
                rand_subject = None

                print("Sampling", image_to_generate, "images")

                for i in range(0, image_to_generate):
                    if random_element is None or random.random() > self.second_chance:
                        random_element = random.choice(subjects)
                        rand_subject = random_element["subject"]

                    hashcode = self._hash_subject(rand_subject)

                    if random_element not in new_subject_view:
                        new_subject_view.append(random_element)

                    if hashcode not in new_hash_map:
                        new_hash_map[hashcode] = {}
                    if rand_subject not in new_hash_map[hashcode]:
                        new_hash_map[hashcode][rand_subject] = []

                    entry, need_to_remove = self.random_add_n_pop_from_old_hashmap(rand_subject, hashcode)

                    if need_to_remove:
                        subjects.remove(random_element)
                        random_element = None

                    new_hash_map[hashcode][rand_subject].append({
                        "subject": rand_subject,
                        "image_path": entry["image_path"],
                    })

                total_images_added = sum(
                    len(images)
                    for subject_dict in new_hash_map.values()
                    for images in subject_dict.values()
                )
                print(f"[INFO] Train set: {len(self.train_set)} subjects")
                print(f"[INFO] Train set: {total_images_added} unique images")

            else:
                # --- Reverse logic: start full, remove random images ---
                print("Removing", total_images - image_to_generate, "images instead of sampling")
                random_element=None
                new_subject_view = subjects.copy()
                new_hash_map = self.hash_map.copy()

                images_to_remove = total_images - image_to_generate

                for _ in range(images_to_remove):
                    if random_element is None or random.random() > (1-self.second_chance):
                        random_element = random.choice(new_subject_view)
                        rand_subject = random_element["subject"]

                    hashcode = self._hash_subject(rand_subject)

                    # Remove a random image from the hash map
                    if hashcode in new_hash_map and rand_subject in new_hash_map[hashcode]:
                        if new_hash_map[hashcode][rand_subject]:
                            idx = random.randrange(len(new_hash_map[hashcode][rand_subject]))
                            del new_hash_map[hashcode][rand_subject][idx]

                            # If subject has no images left, remove it entirely
                            if not new_hash_map[hashcode][rand_subject]:
                                del new_hash_map[hashcode][rand_subject]
                                new_subject_view.remove(random_element)
                                random_element = None  # Force new subject

                    # clean up empty hash entries
                new_hash_map = {h: s for h, s in new_hash_map.items() if s}


            remaining_images = sum(
                len(images)
                for subject_dict in new_hash_map.values()
                for images in subject_dict.values()
            )


            self.train_set = new_subject_view
            self.subject_view = new_subject_view
            self.hash_map = new_hash_map
            if image_to_generate > 0:
                print(self.subject_view[-1])
                print(self.hash_map[self.subject_view[-1]["hash"]])
            print(f"[INFO] Train set: {len(self.train_set)} subjects")
            print(f"[INFO] Train set: {remaining_images} unique images")




    def get_test_data(self, forModelTraining=False):
        if forModelTraining:
            return load_train()
        return load_test()

    def get_adni_test_data(self, forModelTraining=False):
        return load_adni(forModelTraining)


    def load_nii_as_tensor(self, path, normalize=False,downsample=True, to_float32=True):
        nii = nib.load(path)
        data = nii.get_fdata()


        if downsample:
            D, H, W = data.shape
            new_size = (D // 2, H // 2, W // 2)
            data_t = torch.from_numpy(data).unsqueeze(0).unsqueeze(0).float()
            data_t = F.interpolate(data_t, size=new_size, mode='trilinear', align_corners=False)
            data = data_t.squeeze().numpy()
        else:
            print("no downsample")


        if normalize:
            min_val = np.min(data)
            max_val = np.max(data)
            if max_val > min_val:  # Evita divisione per zero
                data = (data - min_val) / (max_val - min_val)

        tensor = torch.from_numpy(data)

        if to_float32:
            tensor = tensor.float()

        return tensor.unsqueeze(0).unsqueeze(0)

    def load_tensor(self, path, normalize=False, to_float32=True):
        """
        Carica un tensor da un file .pt.
        Se il file non esiste, ricostruisce il tensor dal file NIfTI originale e lo salva downsampled.
        """
        try:
            tensor = torch.load(path)
        except FileNotFoundError:
            print(f"⚠ File {path} non trovato, ricostruisco dal NIfTI originale.")
            # Ricava il path NIfTI originale
            filename = os.path.basename(path)
            pathroot = os.path.dirname(path)
            nii_filename = filename.replace("downsampled_", "").replace(".pt", ".nii.gz")
            nii_path = os.path.join(pathroot, nii_filename)

            if not os.path.exists(nii_path):
                raise FileNotFoundError(f"Nemmeno il file NIfTI originale {nii_path} esiste!")

            # Carica NIfTI
            nii = nib.load(nii_path)
            data = nii.get_fdata()

            # Downsample
            D, H, W = data.shape
            new_size = (D // 2, H // 2, W // 2)
            data_t = torch.from_numpy(data).unsqueeze(0).unsqueeze(0).float()
            data_t = F.interpolate(data_t, size=new_size, mode='trilinear', align_corners=False)
            data_downsampled = data_t.squeeze()

            # Salva il tensor downsampled
            torch.save(data_downsampled, path)
            tensor = data_downsampled

        # Normalizzazione opzionale
        if normalize:
            data = tensor.cpu().numpy()
            min_val = np.min(data)
            max_val = np.max(data)
            if max_val > min_val:
                data = (data - min_val) / (max_val - min_val)
            tensor = torch.from_numpy(data)

        # Conversione a float32
        if to_float32:
            tensor = tensor.float()

        return tensor.unsqueeze(0).unsqueeze(0)

    def getPositiveExamplePath(self, subject, current_image):
        # esempio entry {'subject': 'sub-5537','image_path': '/data03/cartellaGB/quasiRaw/quasiRawFomo60k/quasi-raw/sub-5537/ses-1/anat/sub-5537_ses-1_preproc-quasiraw_T1w.nii.gz'}
        hash_key = self._hash_subject(subject)
        entries = self.hash_map_copy.get(hash_key, {}).get(subject, [])
        for entry in entries:
            if entry != current_image:
                return entry
        return None

    def _add_entry_to_batch(self, batch, subject, subject_hash, entry, available_subjects):
        batch.append(entry)
        entries = self.hash_map_copy.get(subject_hash, {}).get(subject, [])
        if entry in entries:
            entries.remove(entry)
        if not entries:
            if subject_hash in self.hash_map_copy:
                self.hash_map_copy[subject_hash].pop(subject, None)
                if not self.hash_map_copy[subject_hash]:
                    self.hash_map_copy.pop(subject_hash, None)
            available_subjects.discard(subject)

    def set_epoch(self, epoch):


        self.epoch = epoch
        random.seed(epoch + self.rank)
        np.random.seed(epoch + self.rank)


        self.hash_map_copy = copy.deepcopy(self.hash_map_original)
        self.subject_view_copy = copy.deepcopy(self.subject_view_original)



    def reshuffle_ddp_split(self, epoch):


        # Group images by subject
        subjects_dict = {
            subj: entries
            for h, subj_dict in self.hash_map_original.items()
            for subj, entries in subj_dict.items()
        }


        subject_keys = list(subjects_dict.keys())
        rng = random.Random(epoch)  # all ranks agree
        rng.shuffle(subject_keys)


        all_images = []
        for subj in subject_keys:
            h = None
            for hashcode, subj_dict in self.hash_map_original.items():
                if subj in subj_dict:
                    h = hashcode
                    break
            for entry in subjects_dict[subj]:
                all_images.append((h, subj, entry))


        total_images = len(all_images)
        images_per_rank = (total_images + self.world_size - 1) // self.world_size
        start = self.rank * images_per_rank
        end = min(start + images_per_rank, total_images)
        shard = all_images[start:end]


        new_hash_map = {}
        new_subject_view = []
        seen_subjects = set()

        for h, subj, entry in shard:
            if h not in new_hash_map:
                new_hash_map[h] = {}
            if subj not in new_hash_map[h]:
                new_hash_map[h][subj] = []
                if subj not in seen_subjects:
                    new_subject_view.append({"subject": subj, "hash": h})
                    seen_subjects.add(subj)
            new_hash_map[h][subj].append(entry)

        self.hash_map = new_hash_map
        self.train_set = new_subject_view
        self.subject_view = new_subject_view

        # Logging
        total_subjects = len(subject_keys)
        shard_subjects = len(new_subject_view)
        if self.rank == 0:
            print(f"[Reshuffle Epoch {epoch}] Total subjects={total_subjects}, subjects_per_rank~={shard_subjects}")
        print(f"[Rank {self.rank}] Epoch {epoch}: Got {shard_subjects} subjects, {len(shard)} images")

    def generate_batch_init(self, epoch):

        if self.ddp_enabled and dist.is_initialized():
            self.reshuffle_ddp_split(epoch)

        self.hash_map_copy = copy.deepcopy(self.hash_map)
        self.subject_view_copy = copy.deepcopy(self.subject_view)

        self.available_subjects = {s["subject"] for s in self.subject_view_copy}
        self.subject_to_hash = {s["subject"]: s["hash"] for s in self.subject_view_copy}

        subjects_list = list(self.available_subjects)
        random.shuffle(subjects_list)  # shuffle order within this shard
        self.available_subjects = set(subjects_list)

    def generate_batch(self, epoch, second_chance=0,):

        if not hasattr(self, 'available_subjects'):
            self.generate_batch_init(epoch)

        available_subjects = self.available_subjects
        subject_to_hash = self.subject_to_hash
        batch = []
        subject = None

        while len(batch) < self.batch_size and available_subjects:
            if subject is None or random.random() > second_chance:
                subject = random.choice(list(available_subjects))

            subject_hash = subject_to_hash.get(subject)
            subject_entries = self.hash_map_copy.get(subject_hash, {}).get(subject, [])

            if not subject_entries:
                self.hash_map_copy.get(subject_hash, {}).pop(subject, None)
                if not self.hash_map_copy.get(subject_hash):
                    self.hash_map_copy.pop(subject_hash, None)
                available_subjects.discard(subject)
                subject = None
                continue

            first_entry = random.choice(subject_entries)
            self._add_entry_to_batch(batch, subject, subject_hash, first_entry, available_subjects)


            positives_added = 1
            while len(batch) < self.batch_size and positives_added < self.max_same_batch_positives:
                positive_entry = None
                if positives_added < self.max_real_same_batch_positives:
                    positive_entry = self.getPositiveExamplePath(subject, first_entry)
                if positive_entry is None:
                    positive_entry = dict(first_entry)
                    positive_entry["require_aug"] = True

                self._add_entry_to_batch(batch, subject, subject_hash, positive_entry, available_subjects)
                positives_added += 1

        return batch



    def load_npz_as_tensor(self, path, normalize=False, to_float32=True,downsample=False):

        with np.load(path) as npz_file:
            data = npz_file["data"]

        if normalize:

            min_val = np.min(data)
            max_val = np.max(data)
            if max_val > min_val:
                data = (data - min_val) / (max_val - min_val)

        tensor = torch.from_numpy(data)

        if to_float32:
            tensor = tensor.float()

        # Reshape to [1, 1, D, H, W] if necessary
        if tensor.dim() == 3:
            tensor = tensor.unsqueeze(0).unsqueeze(0)
        elif tensor.dim() == 4:
            tensor = tensor.unsqueeze(0)
        elif tensor.dim() != 5:
            raise ValueError(f"Unexpected tensor shape: {tensor.shape}")
        if downsample:
            tensor = F.interpolate(tensor, scale_factor=0.5, mode='trilinear', align_corners=False)
            print("downsample")


        return tensor


    #used for openBHB or any dataset to test on
    class SimpleBatchSampler:
        def __init__(self, dataset):
            self.dataset_copy = None
            self.dataset_original = dataset
            self.reset()

        def reset(self):
            self.dataset_copy = copy.deepcopy(self.dataset_original)

        def sample_batch(self, batch_size):
            if not self.dataset_copy:
                return []

            batch = random.sample(self.dataset_copy, min(batch_size, len(self.dataset_copy)))
            sampled_paths = set(entry["image_path"] for entry in batch)
            self.dataset_copy = [entry for entry in self.dataset_copy if entry["image_path"] not in sampled_paths]


            return batch

        def load_npz_as_tensor(self, path, normalize=False, to_float32=True,downsample=False):

            with np.load(path) as npz_file:
                data = npz_file["data"]

            if normalize:

                min_val = np.min(data)
                max_val = np.max(data)
                if max_val > min_val:
                    data = (data - min_val) / (max_val - min_val)

            tensor = torch.from_numpy(data)

            if to_float32:
                tensor = tensor.float()

            # Reshape to [1, 1, D, H, W] if necessary
            if tensor.dim() == 3:
                tensor = tensor.unsqueeze(0).unsqueeze(0)
            elif tensor.dim() == 4:
                tensor = tensor.unsqueeze(0)
            elif tensor.dim() != 5:
                raise ValueError(f"Unexpected tensor shape: {tensor.shape}")
            if downsample:
                tensor = F.interpolate(tensor, scale_factor=0.5, mode='trilinear', align_corners=False)
                print("downsample")


            return tensor