cl_estimate.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.ndimage import label, find_objects\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "from tqdm.auto import tqdm\n",
    "import SimpleITK as sitk\n",
    "import pandas as pd\n",
    "import radiomics\n",
    "from radiomics import featureextractor\n",
    "import os\n",
    "from mirp import extract_features\n",
    "\n",
    "\n",
    "def count_blob_pixels(array):\n",
    "    # Label the blobs in the array\n",
    "    labeled_array, num_features = label(array)\n",
    "\n",
    "    # Calculate the number of pixels in each blob\n",
    "    blob_sizes = [(labeled_array == i).sum() for i in range(1, num_features + 1)]\n",
    "\n",
    "    return blob_sizes\n",
    "\n",
    "\n",
    "def remove_small_blobs(array, min_size=100):\n",
    "    # Label the blobs in the array\n",
    "    labeled_array, num_features = label(array)\n",
    "\n",
    "    # Find the sizes of each blob\n",
    "    blob_sizes = [(labeled_array == i + 1).sum() for i in range(num_features)]\n",
    "\n",
    "    # Create a mask to keep blobs that meet the size requirement or if it's the only blob\n",
    "    if num_features == 1:\n",
    "        mask = labeled_array != 0  # Keep the single blob regardless of size\n",
    "    else:\n",
    "        mask = np.isin(\n",
    "            labeled_array,\n",
    "            [i + 1 for i, size in enumerate(blob_sizes) if size >= min_size],\n",
    "        )\n",
    "\n",
    "    # Apply the mask to the original array\n",
    "    filtered_array = array * mask\n",
    "\n",
    "    return filtered_array\n",
    "\n",
    "\n",
    "def extract_radiomics_features(pet_img, mask_array, output_dir, sub_name, label_type):\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "    mask_array = remove_small_blobs(mask_array)\n",
    "    # Label the blobs in the array\n",
    "    labeled_array, num_features = label(mask_array)\n",
    "\n",
    "    # Initialize the PyRadiomics feature extractor\n",
    "    extractor = featureextractor.RadiomicsFeatureExtractor()\n",
    "\n",
    "    for i in range(1, num_features + 1):\n",
    "        # Create the filename\n",
    "\n",
    "        filename = f\"{sub_name}_blob_{i}_label_{label_type}.csv\"\n",
    "\n",
    "        if Path(os.path.join(output_dir, filename)).exists():\n",
    "            continue\n",
    "        blob_mask = (labeled_array == i).astype(np.uint8)\n",
    "        if blob_mask.sum()<=50:\n",
    "            continue\n",
    "        # Ensure the blob mask has the same shape as the original array\n",
    "        assert blob_mask.shape == mask_array.shape\n",
    "        mask_img = sitk.GetImageFromArray(blob_mask)\n",
    "        mask_img.CopyInformation(pet_img)\n",
    "        # Extract features\n",
    "        feature_vector = extract_features(image= sitk.GetArrayFromImage(pet_img), mask = sitk.GetArrayFromImage(mask_img))\n",
    "        print(len(feature_vector))\n",
    "        # Convert the feature vector to a pandas DataFrame\n",
    "        feature_df = pd.DataFrame([feature_vector])\n",
    "\n",
    "        # Remove non-numeric columns\n",
    "        feature_df = feature_df.select_dtypes(include=[np.number])\n",
    "\n",
    "        # Save the DataFrame as a CSV file\n",
    "        feature_df.to_csv(os.path.join(output_dir, filename), index=False)\n",
    "\n",
    "\n",
    "def mean_pred_prob_per_blob(pred_array, gt_array):\n",
    "    gt_labels, num_gt = label(gt_array)\n",
    "    # print(num_gt)\n",
    "    # Get the unique labels in the gt_array (excluding the background, assumed to be label 0)\n",
    "    unique_labels = np.unique(gt_labels)\n",
    "    unique_labels = unique_labels[unique_labels != 0]  # Exclude background\n",
    "\n",
    "    mean_probs = []\n",
    "\n",
    "    # Calculate the mean predicted probability for each blob\n",
    "    for labels in unique_labels:\n",
    "        # Create a mask for the current label (blob)\n",
    "        mask = np.array((gt_labels == labels).astype(int))\n",
    "        # print(np.unique(mask))        # Check for NaNs\n",
    "        # pred_has_nan = np.any(np.isnan(pred_array))\n",
    "        # gt_has_nan = np.any(np.isnan(gt_array))\n",
    "\n",
    "        # print(pred_has_nan)  # True if NaNs are present in pred_array\n",
    "        # print(gt_has_nan)    # True if NaNs are present in gt_array\n",
    "        # Calculate the mean predicted probability for the current blob\n",
    "        mean_prob = np.mean(pred_array[mask==1])\n",
    "\n",
    "        # Store the mean probability with the corresponding blob label\n",
    "        mean_probs.append(mean_prob)\n",
    "\n",
    "    return mean_probs\n",
    "\n",
    "\n",
    "def calculate_tp_fp_fn(gt, pred):\n",
    "    # Label the blobs\n",
    "    gt_labels, num_gt = label(gt)\n",
    "    pred_labels, num_pred = label(pred)\n",
    "\n",
    "    # Find bounding boxes of blobs\n",
    "    gt_slices = find_objects(gt_labels)\n",
    "    pred_slices = find_objects(pred_labels)\n",
    "\n",
    "    # Create empty masks for TP, FP, FN\n",
    "    tp_mask = np.zeros_like(gt)\n",
    "    fp_mask = np.zeros_like(gt)\n",
    "    fn_mask = np.zeros_like(gt)\n",
    "\n",
    "    # Calculate TP, FP, FN\n",
    "    tp, fp, fn = 0, 0, 0\n",
    "\n",
    "    # Helper function to determine if two blobs overlap\n",
    "    def overlap(slice1, slice2):\n",
    "        return not (\n",
    "            slice1[0].stop < slice2[0].start\n",
    "            or slice1[0].start > slice2[0].stop\n",
    "            or slice1[1].stop < slice2[1].start\n",
    "            or slice1[1].start > slice2[1].stop\n",
    "            or slice1[2].stop < slice2[2].start\n",
    "            or slice1[2].start > slice2[2].stop\n",
    "        )\n",
    "\n",
    "    matched_gt = set()\n",
    "    matched_pred = set()\n",
    "\n",
    "    for i, gt_slice in enumerate(gt_slices):\n",
    "        gt_blob = gt_labels == (i + 1)\n",
    "        match_found = False\n",
    "        for j, pred_slice in enumerate(pred_slices):\n",
    "            if overlap(gt_slice, pred_slice):\n",
    "                pred_blob = pred_labels == (j + 1)\n",
    "                tp_mask[gt_blob & pred_blob] = 1  # True positive overlap region\n",
    "                tp += 1\n",
    "                matched_gt.add(i)\n",
    "                matched_pred.add(j)\n",
    "                match_found = True\n",
    "                break\n",
    "        if not match_found:\n",
    "            fn_mask[gt_blob] = 1\n",
    "\n",
    "    for j, pred_slice in enumerate(pred_slices):\n",
    "        if j not in matched_pred:\n",
    "            pred_blob = pred_labels == (j + 1)\n",
    "            fp_mask[pred_blob] = 1\n",
    "\n",
    "    fn = num_gt - tp\n",
    "    fp = num_pred - tp\n",
    "\n",
    "    return tp_mask, fp_mask, fn_mask, tp, fp, fn\n",
    "\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "def etimate_model_confidence(\n",
    "    src_dir,\n",
    "    pred_dir,\n",
    "    csv_path_metrics,\n",
    "    csv_path_mp,\n",
    "    gt_lesion_label=20,\n",
    "    pred_lesion_label=20,\n",
    "):\n",
    "    \"\"\"\n",
    "    Processes image predictions and ground truth labels to calculate metrics\n",
    "    and save results to CSV files.\n",
    "\n",
    "    Parameters:\n",
    "    src_dir (str or Path): Source directory containing images and labels.\n",
    "    pred_dir (str or Path): Directory containing prediction images.\n",
    "    csv_path_metrics (str or Path): Path to save the metrics CSV file.\n",
    "    csv_path_mp (str or Path): Path to save the mean prediction probabilities CSV file.\n",
    "    \"\"\"\n",
    "    src_dir = Path(src_dir)\n",
    "    pred_dir = Path(pred_dir)\n",
    "\n",
    "    metrics = []\n",
    "    mp = []\n",
    "\n",
    "    for pred in tqdm((src_dir / pred_dir).glob(\"*.nii.gz\")):\n",
    "        fname = (pred.name).split(\".nii\")[0]\n",
    "        gt_path = str(pred).replace(\"imagesTs_resXL_pred_proba\", \"labelsTs_combined\")\n",
    "        pred_array = sitk.GetArrayFromImage(sitk.ReadImage(pred))\n",
    "        pred_proba_array = np.load(src_dir / pred_dir / f\"{fname}.npz\")\n",
    "        pred_proba_array = np.array(pred_proba_array[\"probabilities\"])[-1, :, :, :]\n",
    "        gt_array = sitk.GetArrayFromImage(sitk.ReadImage(gt_path))\n",
    "        gt_array = (gt_array == gt_lesion_label).astype(int)\n",
    "        pred_array = (pred_array == pred_lesion_label).astype(float)\n",
    "\n",
    "        mp.extend(mean_pred_prob_per_blob(pred_proba_array, gt_array))\n",
    "\n",
    "        if gt_array.max() == 1:\n",
    "            metrics.append(\n",
    "                {\n",
    "                    \"Subject\": pred.name,\n",
    "                    \"Mean\": pred_proba_array[gt_array != 0].mean(),\n",
    "                    \"Max\": pred_proba_array[gt_array != 0].max(),\n",
    "                    \"Median\": np.median(pred_proba_array[gt_array != 0]),\n",
    "                    \"Label\": gt_array.max(),\n",
    "                }\n",
    "            )\n",
    "        else:\n",
    "            metrics.append(\n",
    "                {\n",
    "                    \"Subject\": pred.name,\n",
    "                    \"Mean\": pred_proba_array.mean(),\n",
    "                    \"Max\": pred_proba_array.max(),\n",
    "                    \"Median\": np.median(pred_proba_array),\n",
    "                    \"Label\": gt_array.max(),\n",
    "                }\n",
    "            )\n",
    "\n",
    "    df_metrics = pd.DataFrame(metrics)\n",
    "    df_metrics.to_csv(csv_path_metrics, index=False)\n",
    "\n",
    "    mp_df = pd.DataFrame({\"mean\": mp})\n",
    "    mp_df.to_csv(csv_path_mp, index=False)\n",
    "\n",
    "\n",
    "def overlapping_labels(mean_positive_predicitons, gt_leasion_label=20):\n",
    "    mp = pd.read_csv(mean_positive_predicitons)\n",
    "    mean_value = mp[\"mean\"].mean()\n",
    "    for idx, pred in enumerate(tqdm((src_dir / pred_dir).glob(\"*.nii.gz\"))):\n",
    "        fname = (pred.name).split(\".nii\")[0]\n",
    "        gt = str(pred).replace(\"imagesTs_resXL_pred_proba\", \"labelsTs_combined\")\n",
    "        pred_array = sitk.GetArrayFromImage(sitk.ReadImage(pred))\n",
    "        pet_img = sitk.ReadImage(src_dir / \"imagesTs\" / str(fname + \"_0001.nii.gz\"))\n",
    "        pred_proba_array = np.load(src_dir / pred_dir / str(fname + \".npz\"))\n",
    "        pred_proba_array = np.array(pred_proba_array[\"probabilities\"])[-1, :, :, :]\n",
    "        gt_array = sitk.GetArrayFromImage(sitk.ReadImage(gt))\n",
    "        gt_array = (gt_array == gt_leasion_label).astype(int)\n",
    "        gt_array[pred_proba_array > 0.50] = 2\n",
    "        gt_img = sitk.GetImageFromArray(gt_array)\n",
    "        gt_img.CopyInformation(pet_img)\n",
    "        output_dir = src_dir / \"confidence_estimates\"\n",
    "        output_dir.mkdir(exist_ok=True)\n",
    "        sitk.WriteImage(gt_img, output_dir / str(fname + \".nii.gz\"))\n",
    "\n",
    "\n",
    "def remove_non_confident_blobs(array, pred_lesion_label=20):\n",
    "    # Label the blobs in the array\n",
    "    labeled_array, num_features = label(array)\n",
    "\n",
    "    # Initialize a mask to keep only the blobs with max value == 2\n",
    "    mask = np.zeros_like(array, dtype=bool)\n",
    "\n",
    "    # Iterate through each blob\n",
    "    for i in range(1, num_features + 1):\n",
    "        # Extract the current blob\n",
    "        blob = array[labeled_array == i]\n",
    "\n",
    "        # Check if the maximum value in the blob is equal to 2\n",
    "        if blob.max() == 2:\n",
    "            mask[labeled_array == i] = True  # Keep this blob\n",
    "\n",
    "    # Apply the mask to the original array\n",
    "    filtered_array = array * mask\n",
    "    filtered_array[filtered_array != 0] = 20\n",
    "    return filtered_array\n",
    "\n",
    "\n",
    "def zero_boundary_slices(arr):\n",
    "    \"\"\"\n",
    "    Set values to zero in the first and last five slices\n",
    "    along all three axes of a 3D NumPy array.\n",
    "\n",
    "    Parameters:\n",
    "    arr (np.ndarray): Input 3D array\n",
    "\n",
    "    Returns:\n",
    "    np.ndarray: The modified 3D array with boundary slices set to zero\n",
    "    \"\"\"\n",
    "    # Validate input array\n",
    "    if arr.ndim != 3:\n",
    "        raise ValueError(\"Input array must be 3-dimensional.\")\n",
    "\n",
    "    # Zero out first and last 5 slices on the first axis\n",
    "    arr[:5, :, :] = 0\n",
    "    arr[-5:, :, :] = 0\n",
    "\n",
    "    # Zero out first and last 5 slices on the second axis\n",
    "    arr[:, :5, :] = 0\n",
    "    arr[:, -5:, :] = 0\n",
    "\n",
    "    # Zero out first and last 5 slices on the third axis\n",
    "    arr[:, :, :5] = 0\n",
    "    arr[:, :, -5:] = 0\n",
    "\n",
    "    return arr\n",
    "\n",
    "\n",
    "def create_confident_labels(\n",
    "    mean_positive_predicitons, gt_leasion_label=20, pred_lesion_label=20\n",
    "):\n",
    "    mp = pd.read_csv(mean_positive_predicitons)\n",
    "    mean_value = mp[\"mean\"].mean()\n",
    "    for idx, pred in enumerate(tqdm((src_dir / pred_dir).glob(\"*.nii.gz\"))):\n",
    "        fname = (pred.name).split(\".nii\")[0]\n",
    "        gt = str(pred).replace(\"imagesTs_resXL_pred_proba\", \"labelsTs_combined\")\n",
    "        pred_array = sitk.GetArrayFromImage(sitk.ReadImage(pred))\n",
    "        pet_img = sitk.ReadImage(src_dir / \"imagesTs\" / str(fname + \"_0001.nii.gz\"))\n",
    "        pred_proba_all = np.load(src_dir / pred_dir / str(fname + \".npz\"))\n",
    "        pred_proba_array = np.array(pred_proba_all[\"probabilities\"])[-1, :, :, :]\n",
    "        pred_proba_all = np.argmax(pred_proba_all[\"probabilities\"], axis=0)\n",
    "        pred_proba_lesion = (pred_proba_all == pred_lesion_label).astype(int)\n",
    "        print(np.unique(pred_proba_all))\n",
    "        gt_array = sitk.GetArrayFromImage(sitk.ReadImage(gt))\n",
    "        gt_array = (gt_array == gt_leasion_label).astype(int)\n",
    "        pred_proba_lesion[pred_proba_array > 0.5] = 2\n",
    "        print(f\"before: {np.unique(pred_proba_lesion)}\")\n",
    "        pred_proba_filterd = remove_non_confident_blobs(pred_proba_lesion)\n",
    "        pred_proba_filterd = zero_boundary_slices(pred_proba_filterd)\n",
    "        print(f\"after: {np.unique(pred_proba_filterd)}\")\n",
    "        pred_img = sitk.GetImageFromArray(pred_proba_filterd)\n",
    "        pred_img.CopyInformation(pet_img)\n",
    "        output_dir = src_dir / \"confidence_estimates_pred\"\n",
    "        output_dir.mkdir(exist_ok=True)\n",
    "        sitk.WriteImage(pred_img, output_dir / str(fname + \".nii.gz\"))\n",
    "    create_confident_labels(csv_path_mp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.ndimage import label, find_objects\n",
    "import matplotlib.pyplot as plt\n",
    "from pathlib import Path\n",
    "from tqdm.auto import tqdm\n",
    "import SimpleITK as sitk\n",
    "import pandas as pd\n",
    "from radiomics import featureextractor\n",
    "import os\n",
    "from mirp import extract_features\n",
    "\n",
    "\n",
    "class BAMF_PET_Processor:\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        src_dir,\n",
    "        pet_dir,\n",
    "        gt_dir,\n",
    "        pred_dir,\n",
    "        gt_lesion_label=20,\n",
    "        pred_lesion_label=20,\n",
    "    ):\n",
    "        self.src_dir = Path(src_dir)\n",
    "        self.gt_dir = Path(gt_dir)\n",
    "        self.pet_dir = Path(pet_dir)\n",
    "        self.pred_dir = Path(pred_dir)\n",
    "        self.extractor = featureextractor.RadiomicsFeatureExtractor()\n",
    "        self.gt_lesion_label = gt_lesion_label\n",
    "        self.pred_lesion_label = pred_lesion_label\n",
    "        self.cle_estimates = f\"{self.pred_dir}_cle_estimates.csv\"\n",
    "        self.blobwise_estimate = f\"{self.pred_dir}_blobwise_estimates.csv\"\n",
    "\n",
    "    def count_blob_pixels(self, array):\n",
    "        labeled_array, num_features = label(array)\n",
    "        blob_sizes = [(labeled_array == i).sum() for i in range(1, num_features + 1)]\n",
    "        return blob_sizes\n",
    "\n",
    "    def remove_small_blobs(self, array, min_size=100):\n",
    "        labeled_array, num_features = label(array)\n",
    "        blob_sizes = [(labeled_array == i + 1).sum() for i in range(num_features)]\n",
    "        mask = (\n",
    "            labeled_array != 0\n",
    "            if num_features == 1\n",
    "            else np.isin(\n",
    "                labeled_array,\n",
    "                [i + 1 for i, size in enumerate(blob_sizes) if size >= min_size],\n",
    "            )\n",
    "        )\n",
    "        filtered_array = array * mask\n",
    "        return filtered_array\n",
    "\n",
    "    def extract_radiomics_features(self, pet_img, mask_array, sub_name, label_type):\n",
    "        if not self.output_dir.exists():\n",
    "            os.makedirs(self.output_dir)\n",
    "        mask_array = self.remove_small_blobs(mask_array)\n",
    "        labeled_array, num_features = label(mask_array)\n",
    "\n",
    "        for i in range(1, num_features + 1):\n",
    "            filename = f\"{sub_name}_blob_{i}_label_{label_type}.csv\"\n",
    "            if (self.output_dir / filename).exists():\n",
    "                continue\n",
    "            blob_mask = (labeled_array == i).astype(np.uint8)\n",
    "            if blob_mask.sum() <= 50:\n",
    "                continue\n",
    "            mask_img = sitk.GetImageFromArray(blob_mask)\n",
    "            mask_img.CopyInformation(pet_img)\n",
    "            feature_vector = extract_features(\n",
    "                image=sitk.GetArrayFromImage(pet_img),\n",
    "                mask=sitk.GetArrayFromImage(mask_img),\n",
    "            )\n",
    "            feature_df = pd.DataFrame([feature_vector]).select_dtypes(\n",
    "                include=[np.number]\n",
    "            )\n",
    "            feature_df.to_csv(self.output_dir / filename, index=False)\n",
    "\n",
    "    def mean_pred_prob_per_blob(self, pred_array, gt_array):\n",
    "        gt_labels, num_gt = label(gt_array)\n",
    "        unique_labels = np.unique(gt_labels)[np.unique(gt_labels) != 0]\n",
    "        mean_probs = [\n",
    "            np.mean(pred_array[(gt_labels == label).astype(int) == 1])\n",
    "            for label in unique_labels\n",
    "        ]\n",
    "        return mean_probs\n",
    "\n",
    "    def calculate_tp_fp_fn(self, gt, pred):\n",
    "        gt_labels, num_gt = label(gt)\n",
    "        pred_labels, num_pred = label(pred)\n",
    "        gt_slices = find_objects(gt_labels)\n",
    "        pred_slices = find_objects(pred_labels)\n",
    "\n",
    "        tp_mask, fp_mask, fn_mask = (\n",
    "            np.zeros_like(gt),\n",
    "            np.zeros_like(gt),\n",
    "            np.zeros_like(gt),\n",
    "        )\n",
    "        tp, fp, fn = 0, 0, 0\n",
    "\n",
    "        def overlap(slice1, slice2):\n",
    "            return not (\n",
    "                slice1[0].stop < slice2[0].start\n",
    "                or slice1[0].start > slice2[0].stop\n",
    "                or slice1[1].stop < slice2[1].start\n",
    "                or slice1[1].start > slice2[1].stop\n",
    "                or slice1[2].stop < slice2[2].start\n",
    "                or slice1[2].start > slice2[2].stop\n",
    "            )\n",
    "\n",
    "        matched_gt, matched_pred = set(), set()\n",
    "\n",
    "        for i, gt_slice in enumerate(gt_slices):\n",
    "            gt_blob = gt_labels == (i + 1)\n",
    "            match_found = False\n",
    "            for j, pred_slice in enumerate(pred_slices):\n",
    "                if overlap(gt_slice, pred_slice):\n",
    "                    pred_blob = pred_labels == (j + 1)\n",
    "                    tp_mask[gt_blob & pred_blob] = 1\n",
    "                    tp += 1\n",
    "                    matched_gt.add(i)\n",
    "                    matched_pred.add(j)\n",
    "                    match_found = True\n",
    "                    break\n",
    "            if not match_found:\n",
    "                fn_mask[gt_blob] = 1\n",
    "\n",
    "        for j, pred_slice in enumerate(pred_slices):\n",
    "            if j not in matched_pred:\n",
    "                pred_blob = pred_labels == (j + 1)\n",
    "                fp_mask[pred_blob] = 1\n",
    "\n",
    "        fn = num_gt - tp\n",
    "        fp = num_pred - tp\n",
    "\n",
    "        return tp_mask, fp_mask, fn_mask, tp, fp, fn\n",
    "\n",
    "    def estimate_model_confidence(\n",
    "        self, \n",
    "    ):\n",
    "        metrics, mp = [], []\n",
    "        for pred in tqdm((self.src_dir / self.pred_dir).glob(\"*.nii.gz\")):\n",
    "            fname = (pred.name).split(\".nii\")[0]\n",
    "            gt_path = str(pred).replace(\n",
    "                str(self.pred_dir), str(self.gt_dir)\n",
    "            )\n",
    "            pred_array = sitk.GetArrayFromImage(sitk.ReadImage(pred))\n",
    "            pred_proba_array = np.load(self.src_dir / self.pred_dir / f\"{fname}.npz\")\n",
    "            pred_proba_array = np.array(pred_proba_array[\"probabilities\"])[-1, :, :, :]\n",
    "            gt_array = sitk.GetArrayFromImage(sitk.ReadImage(gt_path))\n",
    "            gt_array = (gt_array == self.gt_lesion_label).astype(int)\n",
    "            pred_array = (pred_array == self.pred_lesion_label).astype(float)\n",
    "\n",
    "            mp.extend(self.mean_pred_prob_per_blob(pred_proba_array, gt_array))\n",
    "            tp_mask, fp_mask, fn_mask, tp, fp, fn = self.calculate_tp_fp_fn(gt_array, pred_array)\n",
    "            if gt_array.max() == 1:\n",
    "                metrics.append(\n",
    "                    {\n",
    "                        \"Subject\": pred.name,\n",
    "                        \"Mean\": pred_proba_array[gt_array != 0].mean(),\n",
    "                        \"Max\": pred_proba_array[gt_array != 0].max(),\n",
    "                        \"Median\": np.median(pred_proba_array[gt_array != 0]),\n",
    "                        \"tp\": tp,\n",
    "                        \"fp\": fp,\n",
    "                        \"fn\": fn,\n",
    "                        \"Label\": gt_array.max(),\n",
    "                    }\n",
    "                )\n",
    "            else:\n",
    "                metrics.append(\n",
    "                    {\n",
    "                        \"Subject\": pred.name,\n",
    "                        \"Mean\": pred_proba_array.mean(),\n",
    "                        \"Max\": pred_proba_array.max(),\n",
    "                        \"Median\": np.median(pred_proba_array),\n",
    "                        \"Label\": gt_array.max(),\n",
    "                    }\n",
    "                )\n",
    "\n",
    "        pd.DataFrame(metrics).to_csv(self.cle_estimates, index=False)\n",
    "        pd.DataFrame({\"mean\": mp}).to_csv(self.blobwise_estimate, index=False)\n",
    "\n",
    "    def overlapping_labels(self):\n",
    "        if not Path(self.blobwise_estimate).exists():\n",
    "            self.estimate_model_confidence()\n",
    "        mp = pd.read_csv(self.blobwise_estimates)\n",
    "        mean_value = mp[\"mean\"].mean()\n",
    "        for idx, pred in enumerate(\n",
    "            tqdm((self.src_dir / self.pred_dir).glob(\"*.nii.gz\"))\n",
    "        ):\n",
    "            output_dir = self.src_dir / f\"{self.pred_dir}_confidence_estimates\"\n",
    "            output_dir.mkdir(exist_ok=True)\n",
    "            if not (output_dir / f\"{fname}.nii.gz\").exists():\n",
    "                fname = (pred.name).split(\".nii\")[0]\n",
    "                gt = str(pred).replace(str(self.pred_dir), str(self.gt_dir))\n",
    "                pred_array = sitk.GetArrayFromImage(sitk.ReadImage(pred))\n",
    "                pet_img = sitk.ReadImage(self.src_dir / self.pet_dir / f\"{fname}_0001.nii.gz\")\n",
    "                pred_proba_array = np.load(self.src_dir / self.pred_dir / f\"{fname}.npz\")\n",
    "                pred_proba_array = np.array(pred_proba_array[\"probabilities\"])[-1, :, :, :]\n",
    "                gt_array = sitk.GetArrayFromImage(sitk.ReadImage(gt))\n",
    "                gt_array = (gt_array == self.gt_lesion_label).astype(int)\n",
    "                gt_array[pred_proba_array > 0.50] = 2\n",
    "                gt_img = sitk.GetImageFromArray(gt_array)\n",
    "                gt_img.CopyInformation(pet_img)\n",
    "                sitk.WriteImage(gt_img, output_dir / f\"{fname}.nii.gz\")\n",
    "\n",
    "    def remove_non_confident_blobs(self, array):\n",
    "        labeled_array, num_features = label(array)\n",
    "        mask = np.zeros_like(array, dtype=bool)\n",
    "        for i in range(1, num_features + 1):\n",
    "            blob = array[labeled_array == i]\n",
    "            if blob.max() == 2:\n",
    "                mask[labeled_array == i] = True\n",
    "        filtered_array = array * mask\n",
    "        filtered_array[filtered_array != 0] = 20\n",
    "        return filtered_array\n",
    "\n",
    "    def zero_boundary_slices(self, arr):\n",
    "        if arr.ndim != 3:\n",
    "            raise ValueError(\"Input array must be 3-dimensional.\")\n",
    "        arr[:5, :, :] = 0\n",
    "        arr[-5:, :, :] = 0\n",
    "        arr[:, :5, :] = 0\n",
    "        arr[:, -5:, :] = 0\n",
    "        arr[:, :, :5] = 0\n",
    "        arr[:, :, -5:] = 0\n",
    "        return arr\n",
    "\n",
    "    def create_confident_labels(\n",
    "        self,\n",
    "    ):\n",
    "        if not Path(self.blobwise_estimate).exists():\n",
    "            self.estimate_model_confidence()\n",
    "        mp = pd.read_csv(self.blobwise_estimate)\n",
    "        mean_value = mp[\"mean\"].mean()\n",
    "        print(f\"Model CLE: {mean_value}\")\n",
    "        for idx, pred in enumerate(\n",
    "            tqdm((self.src_dir / self.pred_dir).glob(\"*.nii.gz\"))\n",
    "        ):\n",
    "            fname = (pred.name).split(\".nii\")[0]\n",
    "            output_dir = self.src_dir / f\"{self.pred_dir}_confidence_labels\"\n",
    "            output_dir.mkdir(exist_ok=True)\n",
    "\n",
    "            if not (output_dir / f\"{fname}.nii.gz\").exists():\n",
    "                gt = str(pred).replace(str(self.pred_dir), str(self.gt_dir))\n",
    "                pred_array = sitk.GetArrayFromImage(sitk.ReadImage(pred))\n",
    "                pet_img = sitk.ReadImage(self.src_dir / self.pet_dir / f\"{fname}_0001.nii.gz\")\n",
    "                pred_proba_array = np.load(self.src_dir / self.pred_dir / f\"{fname}.npz\")\n",
    "                pred_proba_array = np.array(pred_proba_array[\"probabilities\"])[-1, :, :, :]\n",
    "                gt_array = sitk.GetArrayFromImage(sitk.ReadImage(gt))\n",
    "                gt_array = (gt_array == self.gt_lesion_label).astype(int)\n",
    "                gt_array[pred_proba_array > 0.50] = 2\n",
    "                gt_array = self.remove_non_confident_blobs(gt_array)\n",
    "                gt_array = self.zero_boundary_slices(gt_array)\n",
    "                gt_img = sitk.GetImageFromArray(gt_array)\n",
    "                gt_img.CopyInformation(pet_img)\n",
    "                sitk.WriteImage(gt_img, output_dir / f\"{fname}.nii.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a51c1ba5b9a4f37be1cd7675f5ca23f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model CLE: 0.5817180254671563\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d2e2ccd9e16411bba2ae2c2e3992953",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "src_dir = Path(\n",
    "    \"/mnt/nfs/slow_ai_team/organ_segmentation/nnunet_liverv0.0/nnUNet_raw_database/nnUNet_raw/nnUNet_raw_data/Dataset019_AutoPET2024/\"\n",
    ")\n",
    "pred_dir = \"imagesTr_resXL_pred_proba\"\n",
    "pet_dir = \"imagesTr\"\n",
    "gt_dir = \"labelsTr\"\n",
    "\n",
    "processor = BAMF_PET_Processor(\n",
    "        src_dir,\n",
    "        pet_dir,\n",
    "        gt_dir,\n",
    "        pred_dir,\n",
    "        gt_lesion_label=20,\n",
    "        pred_lesion_label=20,\n",
    "    )\n",
    "\n",
    "processor.create_confident_labels()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Step 1 Estimate model confidence\n",
    "# csv_path_metrics = \"CL_estimates.csv\"\n",
    "# csv_path_mp = \"Mean_positive_prediction.csv\"\n",
    "# src_dir = Path(\n",
    "#     \"/mnt/nfs/slow_ai_team/organ_segmentation/nnunet_liverv0.0/nnUNet_raw_database/nnUNet_raw/nnUNet_raw_data/Dataset019_AutoPET2024/\"\n",
    "# )\n",
    "# pred_dir = \"imagesTs_resXL_pred_proba\"\n",
    "# gt_dir = \"imagesTs\"\n",
    "# etimate_model_confidence(src_dir, pred_dir, csv_path_metrics, csv_path_mp)\n",
    "\n",
    "# # Step 2 Create Overlapping of confident and predicted labels (optional)\n",
    "\n",
    "# overlapping_labels(csv_path_mp)\n",
    "\n",
    "# # Step 3: Create Confident labels as final predicitions\n",
    "\n",
    "# create_confident_labels(csv_path_mp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "aimiv2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}