downsampler.py

import os

import nibabel as nib
import numpy as np
import torch
import torch.nn.functional as F

from DataSetManager import customDataLoader


def save_safe(obj, filename):
    """Salva l'oggetto su disco in modo sicuro, forzando la scrittura."""
    tmp_filename = filename + ".tmp"
    torch.save(obj, tmp_filename)
    os.replace(tmp_filename, filename)
    with open(filename, 'rb') as f:
        f.flush()
        os.fsync(f.fileno())

def is_torch_file_valid(path):
    """Controlla se un file torch è leggibile e non corrotto."""
    if not os.path.exists(path):
        return False
    try:
        _ = torch.load(path, map_location='cpu')
        return True
    except Exception as e:
        print(f"⚠️ File corrotto o illeggibile: {path} ({e})")
        return False

def loadDownsampledAndSaveTorch(path):
    filename = os.path.basename(path)
    pathroot = os.path.dirname(path) + '/'
    downsampled_path = os.path.join(pathroot, "downsampled_" + filename.replace('.nii.gz', '.pt'))

    # ✅ Controlla se già esiste ed è valido
    if is_torch_file_valid(downsampled_path):
        print(f"⏩ File già valido, salto: {downsampled_path}")
        return

    print(f"🔄 Rigenero: {downsampled_path}")

    try:
        nii = nib.load(path)
        data = nii.get_fdata()

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

        save_safe(data_downsampled, downsampled_path)
        print(f"✅ Salvato: {downsampled_path}")
    except Exception as e:
        print(f"❌ Errore durante il processing di {path}: {e}")

def save_npy_safe(array: np.ndarray, filename: str): #dirs/file.npy
    tmp=filename.replace('.npy', '.tmp') #dirs/file.tmp
    np.save(tmp, array) #dirs/file.tmp.npy su disco
    os.replace(tmp+'.npy', filename)   #dirs/file.tmp.npy --> #dirs/file.npy
    with open(filename, 'rb') as f:
        f.flush()
        os.fsync(f.fileno())

def is_npy_valid(path):
    try:
        return os.path.exists(path) and np.load(path, mmap_mode='r').shape is not None
    except:
        return False

def loadDownsampledAndSaveNPY(path, delete_pt_if_valid=False):
    fname = os.path.basename(path)
    root = os.path.dirname(path)
    pt_path = os.path.join(root, "downsampled_" + fname.replace(".nii.gz", ".pt"))
    npy_path = pt_path.replace(".pt", ".npy")

    if not os.path.exists(pt_path):
        print(f"⚠️ Nessun .pt trovato per {path}")
        return

    if is_npy_valid(npy_path):
        if delete_pt_if_valid: os.remove(pt_path)
        return

    try:
        arr = torch.load(pt_path, map_location='cpu').detach().cpu().numpy()
        save_npy_safe(arr, npy_path)
        if delete_pt_if_valid and is_npy_valid(npy_path):
            os.remove(pt_path)
    except Exception as e:
        print(f"❌ Errore conversione {pt_path}: {e}")

if __name__ == '__main__':
    dlPretrain = customDataLoader(train_ratio=1.0, test_ratio=0, max_same_batch_positives=1)
    dlPretrain.generate_batch_init(0)

    all_entries = []
    while batch := dlPretrain.generate_batch(0, second_chance=0):
        all_entries.extend(entry["image_path"] for entry in batch)

    processed_count = 0
    for path in all_entries:
        loadDownsampledAndSaveTorch(path)
        processed_count += 1
        print(f"✅ Processate {processed_count}/{len(all_entries)} immagini", end='\r')

    print(f"\nTrovate {len(all_entries)} immagini da processare")
    print(f"✅ Tutti i {processed_count} file downsampled salvati.")