temporary.py

import os
import re
import sys
import ast
import yaml
import json
import torch
import utils
import random
import logging
import numpy as np
import pandas as pd
import torch.nn as nn
import nibabel as nib
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
# from sklearn.model_selection import KFold
from sklearn.model_selection import StratifiedKFold

import monai.transforms as mt
from utils import evaluate_segmentation, numpy_to_list, list_to_numpy
from monai.inferers import SlidingWindowInferer
from unetr_pp.network_architecture.synapse.unetr_pp_synapse import UNETR_PP
from monai.networks.nets import UNet
from dataset_seg import CRCDataset_seg

# SET UP LOGGING -------------------------------------------------------------
logger = logging.getLogger(__name__)
logger_radiomics = logging.getLogger("radiomics")
logging.basicConfig(level=logging.ERROR)
logging.basicConfig(stream=sys.stdout, level=logging.INFO)

# MAKE PARSER AND LOAD PARAMS FROM CONFIG FILE--------------------------------
parser = utils.get_args_parser('config.yml')
# parser.add_argument("--fold", type=int, default=None)
args, unknown = parser.parse_known_args()
with open(args.config_path) as file:
    config = yaml.load(file, Loader=yaml.FullLoader)

# USE CONFIG PARAMS ----------------------------------------------------------
root = config['dir']['root']
patch_size = ast.literal_eval(config['training']['patch_size'])
stride = config['training']['stride']
batch_size = config['training']['batch_size']
num_workers = config['training']['num_workers']
num_classes = config['training']['num_classes']
mode = config['training']['mode']
patch_mode = config['training']['patch_mode']

# SET FIXED SEED FOR REPRODUCIBILITY --------------------------------
seed = config['seed']
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
random.seed(seed)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def enable_last_mc_dropout(model):
    dropout3d_layers = [m for m in model.modules() if isinstance(m, nn.Dropout3d)]
    
    if dropout3d_layers:
        # Set only the last Dropout3d layer to train mode
        dropout3d_layers[-1].train()


def enable_mc_dropout(model):
    dropout3d_layers = [m for m in model.modules() if isinstance(m, nn.Dropout3d)]
    
    for layer in dropout3d_layers:
        layer.train()


@torch.no_grad()
def mc_forward(model, inputs, inferer, T=20, body_mask=None):
    model.eval()
    enable_mc_dropout(model)
    preds_list = []
    metrics_list = []

    for _ in range(T):
        logits = inferer(inputs=inputs, network=model)
        if body_mask is not None:
            logits[body_mask == 0] = -1e10
        preds = torch.sigmoid(logits)
        preds_list.append(preds)
        metrics = evaluate_segmentation(preds, targets, num_classes=num_classes, prob_thresh=probs, logits_input=False)
        metrics_list.append(metrics)

    preds_stack = torch.stack(preds_list, dim=0)  # Shape: [T, B, C, H, W, D]
    mean_pred = preds_stack.mean(dim=0)
    std_pred = preds_stack.std(dim=0)

    metrics_mean = {}
    metrics_std = {}

    for key in metrics_list[0].keys():
        values = [m[key] for m in metrics_list]
        if isinstance(values[0], np.ndarray):  
            arr = np.stack(values, axis=0)
            mean_arr = arr.mean(axis=0)
            metrics_mean[key] = (mean_arr >= 0.5).astype(int)
            metrics_std[key] = arr.std(axis=0)
        else:
            metrics_mean[key] = np.mean(values)
            metrics_std[key] = np.std(values)
    return mean_pred, std_pred, metrics_mean, metrics_std
# %%
train_transforms = mt.Compose([mt.ToTensord(keys=["image", "mask"])])
val_transforms = mt.Compose([
    mt.CenterSpatialCropd(keys=["image", "mask"], roi_size=patch_size) if mode == '2d' else mt.Lambda(lambda x: x),
    mt.ToTensord(keys=["image", "mask"]),
])
transforms = [train_transforms, val_transforms]
# %%
df_path_u = os.path.join(root, config['dir']['processed'], 'unhealthy_df.pkl')
df_path_h = os.path.join(root, config['dir']['processed'], 'healthy_df.pkl')
# %%
df_u = pd.read_pickle(df_path_u)
df_h = pd.read_pickle(df_path_h)
stair_step_artifact_ids = [1, 19, 98]
slice_thickness_ids = [97, 128, 137]  # slice thickness > 5mm
colon_blockage_ids = [64, 77, 173]
rectal_cancer_ids = [76]
explicit_ids_test = [31, 32, 47, 54, 78, 109, 73, 197, 204] # > tx t0
# drop specific ids
df_u = df_u[~df_u['id'].astype(int).isin(stair_step_artifact_ids + slice_thickness_ids + colon_blockage_ids + rectal_cancer_ids)]
df_u_explicit_test = df_u[df_u['id'].astype(int).isin(explicit_ids_test)]
df_u = df_u[~df_u['id'].astype(int).isin(explicit_ids_test)]
# %%
# load clinical data
def extract_T(value):
    if value is None:
        return None
    s = str(value).strip()
    m = re.match(r"^(T\d)", s, flags=re.IGNORECASE)
    if m:
        return m.group(1)
    print(f"Warning: Could not extract T from value '{value}'")
    return None

clinical_data_dir = os.path.join(root, config['dir']['clinical_data'])
default_missing = pd._libs.parsers.STR_NA_VALUES
clinical_data = pd.read_excel(
            clinical_data_dir,
            index_col=False,
            na_filter=True,
            na_values=default_missing)
tnm_data = clinical_data.rename(columns={'TNM wg mnie': 'TNM', 'Nr pacjenta': 'id'})[['id', 'TNM']].dropna(subset=['id'])
tnm_data = tnm_data[tnm_data['id'].astype(int).isin(df_u['id'].astype(int))]
tnm_data["T_extracted"] = tnm_data["TNM"].apply(extract_T)
most_frequent_T = tnm_data["T_extracted"].dropna().mode()[0]
tnm_data["T_clean"] = tnm_data["T_extracted"].fillna(most_frequent_T)
# %%
dataset_u = CRCDataset_seg(root_dir=root,
                         df=df_u,
                         config=config,
                         transforms=transforms,
                         patch_size=patch_size,
                         stride=stride,
                         num_patches=100,
                         mode=mode,
                         patch_mode=patch_mode)
dataset_h = CRCDataset_seg(root_dir=root,
                            df=df_h,
                            config=config,
                            transforms=transforms,
                            patch_size=patch_size,
                            stride=stride,
                            num_patches=100,
                            mode=mode,
                            patch_mode=patch_mode)
# %%
ids_train_val = dataset_u.df['id'].astype(int).unique().tolist()
SPLITS = 10

stratification_labels = tnm_data.set_index('id').reindex(ids_train_val)['T_clean'].values
skf = StratifiedKFold(n_splits=SPLITS, shuffle=True, random_state=seed)
folds = list(skf.split(ids_train_val, stratification_labels))

# %%
paths = [
    "best_model_ae559da8-a149-4a42-b1f0-3d1618f79206.pth",  # 1
    "best_model_2c5cbe95-534d-4d13-b640-d06f9cb4902f.pth",  # 2
    "best_model_f6f767bb-ef12-47a3-bbb4-91e7f86901a0.pth",  # 3
    "best_model_940c745a-b772-41c6-9097-81293ba46fc2.pth",  # 4
    "best_model_63bda0be-35cd-45f4-b8c8-ba3023cb1e94.pth",  # 5
    "best_model_9fd71df9-2974-4ec6-9d9d-dbbd9770b475.pth",  # 6
    "best_model_be989c74-6ab5-4046-98be-b5d67622c2c4.pth",  # 7
    "best_model_814ccb96-72da-43c5-8551-b1b16120e0fb.pth",  # 8
    "best_model_af77cf3e-eb1e-49c9-b667-610e714719df.pth",  # 9
    "best_model_ec84797f-9555-4ca9-ae43-f34083d33df2.pth"  # 10
]


# Aggregated containers across folds
agg_validation_all = {}
agg_test_all = {}

all_patients_metrics = {}


# %%
for i, path in enumerate(paths):
    fold = i+1
    weights_path = path
    for fold_idx, (train_idx, val_idx) in enumerate(folds):
        if fold_idx + 1 != fold:
            continue
        val_ids = [ids_train_val[i] for i in val_idx]
        test_ids = dataset_h.df['id'].astype(int).unique().tolist()
        val_dataset = torch.utils.data.Subset(dataset_u, [i for i in range(len(dataset_u)) if int(dataset_u.get_patient_id(i)) in val_ids])
        test_dataset = torch.utils.data.Subset(dataset_h, range(len(dataset_h)))
        val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=num_workers)
        test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=num_workers)

        weights_path = os.path.join(config['dir']['root'], 'models', weights_path)
        print(fold, val_ids, weights_path)
 
        model = UNETR_PP(
            in_channels=1,
            out_channels=14,
            img_size=tuple(patch_size),
            depths=[3, 3, 3, 3],
            dims=[32, 64, 128, 256],
            do_ds=False
        )
        model.out1.conv.conv = torch.nn.Conv3d(16, 1, kernel_size=(1, 1, 1), stride=(1, 1, 1))
        checkpoint = torch.load(weights_path, weights_only=False, map_location=device)
        model.load_state_dict(checkpoint)
        model = model.to(device)

        # gaussian or not?
        inferer = SlidingWindowInferer(roi_size=tuple(patch_size), sw_batch_size=36, overlap=0.75, mode="constant", device=torch.device('cpu'))
        model.eval()
        probs = 0.5

        
        # Run validation then test sequentially and save per-phase JSON
        for phase_name, dataloader in [('validation', val_dataloader), ('test', test_dataloader)]:
            all_patients_metrics = {}
            dataloader.dataset.dataset.set_mode(train_mode=False)

            num_samples = len(dataloader)

            with torch.no_grad():
                totals = None
                for i, (_, _, image, mask, id) in enumerate(dataloader):

                    inputs = image.to(device, dtype=torch.float32)
                    targets = mask["mask"].to(torch.device('cpu'), dtype=torch.long)
                    body_mask = mask["body_mask"].to(torch.device('cpu'), dtype=torch.long)
                    inputs = inputs.unsqueeze(0)

                    targets = targets.unsqueeze(0)
                    body_mask = body_mask.unsqueeze(0)

                    use_mc = False
                    if use_mc:
                        mean_pred, std_pred, metrics, metrics_std = mc_forward(model, inputs, inferer, T=20, body_mask=body_mask)
                        metrics_str = " | ".join([f"{key}: {float(value):.3f}" for key, value in metrics.items() if 'patient' not in key])
                        print(f"Patient_ID: {str(id[0]):<7} | {metrics_str}")
                        metrics_str = " | ".join([f"{key}: {float(value):.3f}" for key, value in metrics_std.items() if 'patient' not in key])
                        final_output = mean_pred
                        final_output = (final_output.squeeze(0) > probs).cpu().numpy().astype(np.uint8)
                       
                    else:
                        logits = inferer(inputs=inputs, network=model)
                        logits[body_mask == 0] = -1e10
                        metrics = evaluate_segmentation(logits, targets, num_classes=num_classes, prob_thresh=probs)
                        metrics_str = " | ".join([f"{key}: {float(value):.3f}" for key, value in metrics.items() if 'patient' not in key])
                        print(f"Patient_ID: {str(id[0]):<7} | {metrics_str}")
                        final_output_logit = torch.sigmoid(logits)
                        final_output = (final_output_logit.squeeze(0) > probs).cpu().numpy().astype(np.uint8)

                    mask = targets.squeeze(0).cpu().numpy().astype(np.uint8)
                    combined_output = np.zeros_like(mask)
                    combined_output[mask == 1] = 1
                    combined_output[final_output == 1] = 2
                    combined_output[(mask == 1) & (final_output == 1)] = 3

                    if phase_name == 'validation':
                        output_dir = os.path.join("inference_output_last", "validation", "mc" if use_mc else "no_mc", f"patient_{id}")
                    else:
                        output_dir = os.path.join("inference_output_last", "test", f"fold_{fold}",  "mc" if use_mc else "no_mc", f"patient_{id}")
                    os.makedirs(output_dir, exist_ok=True)

                    mask = mask.squeeze(0)
                    combined_output = combined_output.squeeze(0)

                    image_nifti = nib.Nifti1Image(image.squeeze(0).cpu().numpy(), np.eye(4))
                    mask_nifti = nib.Nifti1Image(mask, np.eye(4))
                    combined_output_nifti = nib.Nifti1Image(combined_output, np.eye(4))
                    nib.save(image_nifti, os.path.join(output_dir, f"{id}_image.nii.gz"))
                    nib.save(mask_nifti, os.path.join(output_dir, f"{id}_mask.nii.gz"))
                    nib.save(combined_output_nifti, os.path.join(output_dir, f"{id}_result.nii.gz"))

                    if use_mc is False:
                        # logits may be a MONAI MetaTensor or a torch.Tensor — convert robustly to numpy
                        try:
                            arr = logits.squeeze(0).cpu().numpy()
                        except Exception:
                            # support MONAI MetaTensor which may store underlying tensor in .tensor
                            try:
                                arr = logits.squeeze(0).tensor.cpu().numpy()
                            except Exception:
                                # fallback: try converting via torch.as_tensor
                                try:
                                    import torch as _torch
                                    arr = _torch.as_tensor(logits.squeeze(0)).cpu().numpy()
                                except Exception:
                                    raise
                        # ensure float32 for nibabel
                        arr = arr.astype(np.float32)
                        nib.save(nib.Nifti1Image(arr, np.eye(4)), os.path.join(output_dir, f"{id}_logits.nii.gz"))
                        # save the sigmoid scores robustly (handle MetaTensor)
                        try:
                            scores_arr = final_output_logit.squeeze(0).cpu().numpy()
                        except Exception:
                            try:
                                scores_arr = final_output_logit.squeeze(0).tensor.cpu().numpy()
                            except Exception:
                                try:
                                    import torch as _torch
                                    scores_arr = _torch.as_tensor(final_output_logit.squeeze(0)).cpu().numpy()
                                except Exception:
                                    raise
                        scores_arr = scores_arr.astype(np.float32)
                        nib.save(nib.Nifti1Image(scores_arr, np.eye(4)), os.path.join(output_dir, f"{id}_scores.nii.gz"))

                    if totals is None:
                        totals = {key: 0 for key in metrics.keys()}
                    for key, value in metrics.items():
                        totals[key] += value

                    all_patients_metrics[str(id[0])] = {
                        'fold': fold,
                        'metrics': metrics
                    }
                averages = {key: total / num_samples for key, total in totals.items()}
                avg_metrics_str = ", ".join([f"Average {key}: {avg:.4f}" for key, avg in averages.items() if 'patient' not in key])
                print(f"[{phase_name}] {avg_metrics_str}")

            # save per-phase all_patients_metrics JSON
            if phase_name == 'validation':
                save_path = os.path.join("inference_output_last", "figures", "all_patients_metrics_mc.json" if use_mc else "all_patients_metrics.json")
            else:
                save_path = os.path.join("inference_output_last", "figures", "test", f"fold_{fold}", "all_patients_metrics_mc.json" if use_mc else "all_patients_metrics.json")
            os.makedirs(os.path.dirname(save_path), exist_ok=True)
            with open(save_path, 'w') as f:
                json.dump(numpy_to_list(all_patients_metrics), f, indent=4)
            # append to aggregated across-fold containers
            if phase_name == 'validation':
                agg_validation_all[f'fold_{fold}'] = all_patients_metrics
            else:
                agg_test_all[f'fold_{fold}'] = all_patients_metrics
  
# %%

try:
    val_agg_path = os.path.join("inference_output_last", "figures", "validation", 'all_patients_metrics_aggregated.json')
    os.makedirs(os.path.dirname(val_agg_path), exist_ok=True)
    with open(val_agg_path, 'w') as f:
        json.dump(numpy_to_list(agg_validation_all), f, indent=4)
except Exception as e:
    print(f"Warning: failed to save aggregated validation JSON: {e}")

try:
    test_agg_path = os.path.join("inference_output_last", "figures", "test", 'all_patients_metrics_aggregated.json')
    os.makedirs(os.path.dirname(test_agg_path), exist_ok=True)
    with open(test_agg_path, 'w') as f:
        json.dump(numpy_to_list(agg_test_all), f, indent=4)
except Exception as e:
    print(f"Warning: failed to save aggregated test JSON: {e}")