pix2pix_stroke_conf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nibabel as nib\n",
    "from totalsegmentator.python_api import totalsegmentator\n",
    "import os\n",
    "import SimpleITK as sitk\n",
    "from tqdm.notebook import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from image_utils import convert_series_to_nifti, quantize_maps, quantize_maps_top_quarter, clear_directory, convert_and_copy_with_labels_and_rename, generate_dataset_json\n",
    "import logging\n",
    "from image_utils import ImageVisualizer, show_progress\n",
    "from skimage.restoration import inpaint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "fixed_images_dir = \"/mnt/d/CTH_archive/CTH_STRIPPED\"  \n",
    "moving_images_dir = \"/mnt/d/CTH_archive/CTP_STRIPPED\" \n",
    "\n",
    "\n",
    "transforms_dir =  \"/mnt/d/CTH_archive/CTH_TMAX_TRANSFORMS_COLOR\"\n",
    "transformed_img_dir = \"/mnt/d/CTH_archive/CTH_STRIPPED_REG_COLOR\"\n",
    "\n",
    "\n",
    "def register_images(fixed_image_path, moving_image_path, transforms_dir):\n",
    "    # Extract patient identifier from the file name, ensuring .nii is not included\n",
    "    patient = os.path.splitext(os.path.basename(moving_image_path))[0]\n",
    "    patient = os.path.splitext(patient)[0]  # Remove .nii if present\n",
    "\n",
    "    # Construct the transform file path\n",
    "    transform_file = os.path.join(transforms_dir, f'{patient}.h5')\n",
    "\n",
    "    # Check if the transform file already exists and skip registration if it does\n",
    "    if os.path.exists(transform_file):\n",
    "        print(f\"Transform file already exists for patient {patient}, skipping registration.\")\n",
    "        return\n",
    "\n",
    "    # Load the fixed and moving images\n",
    "    fixed_image = sitk.ReadImage(fixed_image_path)\n",
    "    moving_image = sitk.ReadImage(moving_image_path)\n",
    "\n",
    "    # Initialize the registration method\n",
    "    registration_method = sitk.ImageRegistrationMethod()\n",
    "    registration_method.SetMetricAsMattesMutualInformation(numberOfHistogramBins=50)\n",
    "    registration_method.SetMetricSamplingPercentage(0.10)\n",
    "    registration_method.SetMetricSamplingStrategy(registration_method.RANDOM)\n",
    "    registration_method.SetOptimizerAsGradientDescentLineSearch(learningRate=0.5, numberOfIterations=500)\n",
    "    registration_method.SetOptimizerScalesFromPhysicalShift()\n",
    "    registration_method.SetShrinkFactorsPerLevel(shrinkFactors=[8, 4, 2])\n",
    "    registration_method.SetSmoothingSigmasPerLevel(smoothingSigmas=[4, 2, 1])\n",
    "    registration_method.SmoothingSigmasAreSpecifiedInPhysicalUnitsOn()\n",
    "\n",
    "    initial_transform = sitk.CenteredTransformInitializer(sitk.Cast(fixed_image, moving_image.GetPixelID()), \n",
    "                                                      moving_image, \n",
    "                                                      sitk.AffineTransform(fixed_image.GetDimension()),\n",
    "                                                      sitk.CenteredTransformInitializerFilter.MOMENTS)\n",
    "    \n",
    "    registration_method.SetInitialTransform(initial_transform, True)\n",
    "\n",
    "    try:\n",
    "        # Execute the registration\n",
    "        final_transform = registration_method.Execute(fixed_image, moving_image)\n",
    "        \n",
    "        # Save the transform\n",
    "        sitk.WriteTransform(final_transform, transform_file)\n",
    "        sitk.WriteImage(sitk.Resample(moving_image, fixed_image, final_transform,  sitk.sitkNearestNeighbor), f\"{transformed_img_dir}/{patient}.nii\")\n",
    "        print(\"Final metric value: {0}\".format(registration_method.GetMetricValue()))\n",
    "        print(\n",
    "            \"Optimizer's stopping condition, {0}\".format(\n",
    "                registration_method.GetOptimizerStopConditionDescription()\n",
    "            )\n",
    "        )\n",
    "        print(f\"Registration successful for patient: {patient}. Transform saved to {transform_file}\")\n",
    "    except RuntimeError as e:\n",
    "        print(f\"Registration failed for patient {patient}: {e}\")\n",
    "\n",
    "\n",
    "for filename in tqdm(os.listdir(fixed_images_dir)):\n",
    "    fixed_image_path = os.path.join(fixed_images_dir, filename)\n",
    "    moving_image_path = os.path.join(moving_images_dir, filename)\n",
    "\n",
    "    if os.path.isfile(fixed_image_path) and os.path.isfile(moving_image_path):\n",
    "        if not os.path.exists(transforms_dir):\n",
    "            os.makedirs(transforms_dir, exist_ok=True)\n",
    "        register_images(fixed_image_path, moving_image_path, transforms_dir)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import pydicom\n",
    "\n",
    "# Define the root directory and output directory\n",
    "root_dir = '/mnt/d/CTH_archive/DICOM/'\n",
    "output_root_dir = '/mnt/d/CTH_archive/CBF_COLORED'\n",
    "\n",
    "# Define the target series descriptions\n",
    "primary_series = \"VIZ CBF COLORED\"\n",
    "secondary_series = \"VIZ CBF\"\n",
    "\n",
    "# Function to copy a DICOM file to the output directory\n",
    "def copy_dicom_file(filepath, output_dir):\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    shutil.copy2(filepath, output_dir)\n",
    "    print(f\"Copied {os.path.basename(filepath)} to {output_dir}\")\n",
    "\n",
    "# Process each patient directory\n",
    "for patient_name in os.listdir(root_dir):\n",
    "    patient_dir = os.path.join(root_dir, patient_name)\n",
    "    if os.path.isdir(patient_dir):\n",
    "        print(f\"Processing {patient_name}...\")\n",
    "        \n",
    "        primary_found = False\n",
    "        \n",
    "        # First, search for and copy primary series\n",
    "        for dirpath, _, filenames in os.walk(patient_dir):\n",
    "            for filename in filenames:\n",
    "                if filename.lower().endswith(\".dcm\"):\n",
    "                    filepath = os.path.join(dirpath, filename)\n",
    "                    try:\n",
    "                        dicom_file = pydicom.dcmread(filepath)\n",
    "                        if 'SeriesDescription' in dicom_file and dicom_file.SeriesDescription == primary_series:\n",
    "                            output_dir = os.path.join(output_root_dir, patient_name)\n",
    "                            copy_dicom_file(filepath, output_dir)\n",
    "                            primary_found = True\n",
    "                    except Exception as e:\n",
    "                        print(f\"Could not process file {filepath}: {e}\")\n",
    "        \n",
    "        # If no primary series was found, search for and copy secondary series\n",
    "        if not primary_found:\n",
    "            print(f\"No primary series found for {patient_name}, searching for secondary series...\")\n",
    "            for dirpath, _, filenames in os.walk(patient_dir):\n",
    "                for filename in filenames:\n",
    "                    if filename.lower().endswith(\".dcm\"):\n",
    "                        filepath = os.path.join(dirpath, filename)\n",
    "                        try:\n",
    "                            dicom_file = pydicom.dcmread(filepath)\n",
    "                            if 'SeriesDescription' in dicom_file and dicom_file.SeriesDescription == secondary_series:\n",
    "                                output_dir = os.path.join(output_root_dir, patient_name)\n",
    "                                copy_dicom_file(filepath, output_dir)\n",
    "                        except Exception as e:\n",
    "                            print(f\"Could not process file {filepath}: {e}\")\n",
    "\n",
    "print(\"Processing complete.\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pydicom\n",
    "import numpy as np\n",
    "import SimpleITK as sitk\n",
    "\n",
    "def convert_dicom_series_to_nii(dicom_dir, output_dir):\n",
    "    # Get patient name and ID from the directory name\n",
    "    patient_name_id = os.path.basename(dicom_dir)\n",
    "\n",
    "    # Load DICOM series\n",
    "    reader = sitk.ImageSeriesReader()\n",
    "    dicom_names = reader.GetGDCMSeriesFileNames(dicom_dir)\n",
    "    reader.SetFileNames(dicom_names)\n",
    "\n",
    "    # Read the image series\n",
    "    image = reader.Execute()\n",
    "\n",
    "    # Ensure the output image has RGB channels\n",
    "    image_array = sitk.GetArrayFromImage(image)  # shape: (slices, height, width)\n",
    "    if image_array.ndim == 3:  # Single channel, need to convert to RGB\n",
    "        rgb_image_array = np.stack((image_array,) * 3, axis=-1)  # shape: (slices, height, width, 3)\n",
    "    else:\n",
    "        rgb_image_array = image_array\n",
    "\n",
    "    # Resize to 256x256 if necessary\n",
    "    original_size = rgb_image_array.shape[1:3]\n",
    "    if original_size != (256, 256):\n",
    "        resized_image_array = np.zeros((rgb_image_array.shape[0], 256, 256, 3), dtype=rgb_image_array.dtype)\n",
    "        for i in range(rgb_image_array.shape[0]):\n",
    "            slice_image = sitk.GetImageFromArray(rgb_image_array[i])\n",
    "            resized_slice = sitk.Resample(slice_image, [256, 256], sitk.Transform(), sitk.sitkLinear)\n",
    "            resized_image_array[i] = sitk.GetArrayFromImage(resized_slice)\n",
    "    else:\n",
    "        resized_image_array = rgb_image_array\n",
    "\n",
    "    # Create SimpleITK image from RGB array\n",
    "    rgb_image = sitk.GetImageFromArray(resized_image_array)\n",
    "    rgb_image.CopyInformation(image)  # Preserve orientation, direction, origin\n",
    "\n",
    "    # Save as .nii file\n",
    "    output_file = os.path.join(output_dir, f\"{patient_name_id}.nii\")\n",
    "    sitk.WriteImage(rgb_image, output_file)\n",
    "\n",
    "    print(f\"Converted {patient_name_id} to {output_file}\")\n",
    "\n",
    "def process_all_dicom_series(base_dir, output_dir):\n",
    "    # Recursively find all patient directories\n",
    "    for root, dirs, files in os.walk(base_dir):\n",
    "        if any(file.endswith('.dcm') for file in files):\n",
    "            convert_dicom_series_to_nii(root, output_dir)\n",
    "\n",
    "# Directory paths\n",
    "base_dir = \"/mnt/d/CTH_archive/CBF_COLORED\"\n",
    "output_dir = \"/mnt/d/CTH_archive/CBF_COLORED_NIFTI\"\n",
    "\n",
    "# Convert all series in the base directory\n",
    "process_all_dicom_series(base_dir, output_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import cv2\n",
    "import os\n",
    "import shutil\n",
    "from tqdm import tqdm\n",
    "import logging\n",
    "import re\n",
    "\n",
    "# Set up logging\n",
    "logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
    "\n",
    "def create_rectangle_mask(image_shape, top_left, bottom_right):\n",
    "    mask = np.zeros(image_shape[:2], dtype=np.uint8)\n",
    "    cv2.rectangle(mask, top_left, bottom_right, 255, -1)\n",
    "    return mask\n",
    "\n",
    "def sanitize_filename(filename):\n",
    "    return re.sub(r'[^a-zA-Z0-9_.-]', '_', filename)\n",
    "\n",
    "def remove_text_and_inpaint(nii_path, output_path):\n",
    "    logging.info(f\"Processing file: {os.path.basename(nii_path)}\")\n",
    "    logging.info(f\"Output path: {output_path}\")\n",
    "    \n",
    "    if not os.path.exists(nii_path):\n",
    "        raise FileNotFoundError(f\"Input file not found: {nii_path}\")\n",
    "\n",
    "    nii_img = nib.load(nii_path)\n",
    "    data = nii_img.get_fdata()\n",
    "    assert data.shape[3:] == (1, 3), f\"Unexpected data shape: {data.shape}\"\n",
    "\n",
    "    top_left = (67, 205)\n",
    "    bottom_right = (43, 267)\n",
    "    mask = create_rectangle_mask(data.shape[:2], top_left, bottom_right)\n",
    "\n",
    "    for i in range(data.shape[2]):\n",
    "        slice_rgb = data[:, :, i, 0, :]\n",
    "        if slice_rgb.max() > 1:\n",
    "            slice_rgb = slice_rgb / 255.0\n",
    "        slice_uint8 = (slice_rgb * 255).astype(np.uint8)\n",
    "        result_uint8 = cv2.inpaint(slice_uint8, mask, 3, cv2.INPAINT_TELEA)\n",
    "        result = result_uint8.astype(float) / 255.0\n",
    "        data[:, :, i, 0, :] = result\n",
    "\n",
    "    data = np.clip(data, 0, 1)\n",
    "    data = data.astype(np.float32)\n",
    "    new_img = nib.Nifti1Image(data, nii_img.affine, header=nii_img.header)\n",
    "    new_img.set_data_dtype(np.float32)\n",
    "\n",
    "    os.makedirs(os.path.dirname(output_path), exist_ok=True)\n",
    "\n",
    "    logging.info(f\"Saving processed file to: {output_path}\")\n",
    "    try:\n",
    "        nib.save(new_img, output_path)\n",
    "        logging.info(f\"File saved successfully: {output_path}\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Error saving file: {str(e)}\")\n",
    "        alternative_output_path = os.path.join(os.path.dirname(output_path), sanitize_filename(os.path.basename(output_path)))\n",
    "        logging.info(f\"Attempting to save with sanitized filename: {alternative_output_path}\")\n",
    "        try:\n",
    "            nib.save(new_img, alternative_output_path)\n",
    "            logging.info(f\"File saved successfully with sanitized filename: {alternative_output_path}\")\n",
    "        except Exception as e2:\n",
    "            logging.error(f\"Error saving file with sanitized filename: {str(e2)}\")\n",
    "            logging.info(\"Attempting to copy the original file\")\n",
    "            try:\n",
    "                shutil.copy2(nii_path, output_path)\n",
    "                logging.info(f\"Copied original file to: {output_path}\")\n",
    "            except Exception as e3:\n",
    "                logging.error(f\"Error copying original file: {str(e3)}\")\n",
    "                raise\n",
    "\n",
    "def process_all_nifti_files(input_dir, output_dir):\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    \n",
    "    nii_files = [f for f in os.listdir(input_dir) if f.endswith('.nii')]\n",
    "    \n",
    "    for nii_file in tqdm(nii_files, desc=\"Processing NIfTI files\"):\n",
    "        input_path = os.path.join(input_dir, nii_file)\n",
    "        output_path = os.path.join(output_dir, nii_file)\n",
    "        try:\n",
    "            remove_text_and_inpaint(input_path, output_path)\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Error processing {nii_file}: {str(e)}\")\n",
    "            if nii_file == \"CLARKSON-FARRELL_EDWARD.nii\":\n",
    "                logging.error(f\"Detailed error for CLARKSON-FARRELL_EDWARD.nii:\")\n",
    "                logging.error(f\"Input path: {input_path}\")\n",
    "                logging.error(f\"Output path: {output_path}\")\n",
    "                logging.error(f\"Input path exists: {os.path.exists(input_path)}\")\n",
    "                logging.error(f\"Output directory exists: {os.path.exists(os.path.dirname(output_path))}\")\n",
    "                logging.error(f\"Can write to output directory: {os.access(os.path.dirname(output_path), os.W_OK)}\")\n",
    "                logging.error(f\"File permissions of input file: {oct(os.stat(input_path).st_mode)[-3:]}\")\n",
    "                logging.error(f\"File permissions of output directory: {oct(os.stat(os.path.dirname(output_path)).st_mode)[-3:]}\")\n",
    "                \n",
    "                # Try to list the contents of the output directory\n",
    "                try:\n",
    "                    logging.info(f\"Contents of output directory: {os.listdir(os.path.dirname(output_path))}\")\n",
    "                except Exception as dir_error:\n",
    "                    logging.error(f\"Error listing output directory contents: {str(dir_error)}\")\n",
    "\n",
    "# Usage\n",
    "input_dir = '/mnt/d/CTH_archive/CBF_COLORED_NIFTI_RESIZED_Processed'\n",
    "output_dir = '/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT'\n",
    "\n",
    "process_all_nifti_files(input_dir, output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "import os\n",
    "\n",
    "def resize_to_512_if_needed(input_path, output_path):\n",
    "    # Read the image\n",
    "    image = sitk.ReadImage(input_path)\n",
    "    \n",
    "    # Get current size\n",
    "    size = image.GetSize()\n",
    "    \n",
    "    # Check if resizing is needed\n",
    "    if size[0] != 512 or size[1] != 512:\n",
    "        # Create a resampler\n",
    "        resampler = sitk.ResampleImageFilter()\n",
    "        resampler.SetInterpolator(sitk.sitkLinear)\n",
    "        resampler.SetOutputDirection(image.GetDirection())\n",
    "        resampler.SetOutputOrigin(image.GetOrigin())\n",
    "        \n",
    "        # Calculate new spacing to maintain physical size\n",
    "        old_spacing = image.GetSpacing()\n",
    "        new_spacing = [\n",
    "            old_spacing[0] * (size[0] / 512),\n",
    "            old_spacing[1] * (size[1] / 512),\n",
    "            old_spacing[2]\n",
    "        ]\n",
    "        resampler.SetOutputSpacing(new_spacing)\n",
    "        \n",
    "        # Set new size\n",
    "        new_size = [512, 512, size[2]]\n",
    "        resampler.SetSize(new_size)\n",
    "        \n",
    "        # Resample image\n",
    "        resampled_image = resampler.Execute(image)\n",
    "        \n",
    "        # Save resampled image\n",
    "        sitk.WriteImage(resampled_image, output_path)\n",
    "        print(f\"Resized and saved: {output_path}\")\n",
    "    else:\n",
    "        # If no resizing needed, just copy the file\n",
    "        sitk.WriteImage(image, output_path)\n",
    "        print(f\"No resizing needed, copied: {output_path}\")\n",
    "\n",
    "def process_directory(input_dir, output_dir):\n",
    "    # Ensure the output directory exists\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "    \n",
    "    # Process each .nii file in the input directory\n",
    "    for filename in os.listdir(input_dir):\n",
    "        if filename.endswith('.nii'):\n",
    "            input_path = os.path.join(input_dir, filename)\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "            resize_to_512_if_needed(input_path, output_path)\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    input_dir = r\"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT\"\n",
    "    output_dir = r\"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT_512\"\n",
    "    \n",
    "    process_directory(input_dir, output_dir)\n",
    "    print(\"Processing complete.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "import os\n",
    "from scipy import ndimage\n",
    "\n",
    "def resize_and_match_properties(color_image_path, grayscale_image_path, output_path):\n",
    "    # Read the color and grayscale images\n",
    "    color_image = sitk.ReadImage(color_image_path)\n",
    "    grayscale_image = sitk.ReadImage(grayscale_image_path)\n",
    "\n",
    "    # Convert to numpy arrays\n",
    "    color_array = sitk.GetArrayFromImage(color_image)\n",
    "    grayscale_array = sitk.GetArrayFromImage(grayscale_image)\n",
    "\n",
    "    # Get the target size\n",
    "    target_size = grayscale_array.shape\n",
    "\n",
    "    # Resize the color image array\n",
    "    resized_color_array = np.zeros(target_size + (3,), dtype=color_array.dtype)\n",
    "    for i in range(3):  # For each color channel\n",
    "        resized_color_array[..., i] = ndimage.zoom(color_array[..., i], \n",
    "                                                   (target_size[0] / color_array.shape[0],\n",
    "                                                    target_size[1] / color_array.shape[1],\n",
    "                                                    target_size[2] / color_array.shape[2]))\n",
    "\n",
    "    # Flip the image along the anterior-posterior axis (assuming it's the second axis)\n",
    "    #resized_color_array = np.flip(resized_color_array, axis=1)\n",
    "\n",
    "    # Create a new SimpleITK image from the resized and flipped array\n",
    "    resized_color_image = sitk.GetImageFromArray(resized_color_array, isVector=True)\n",
    "\n",
    "    # Copy all properties from the grayscale image\n",
    "    resized_color_image.CopyInformation(grayscale_image)\n",
    "\n",
    "    # Save the resized, flipped, and matched color image\n",
    "    sitk.WriteImage(resized_color_image, output_path)\n",
    "\n",
    "def process_directory(color_dir, grayscale_dir, output_dir):\n",
    "    # Ensure the output directory exists\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "    # Process each .nii file in the color directory\n",
    "    for filename in os.listdir(color_dir):\n",
    "        if filename.endswith('.nii'):\n",
    "            color_path = os.path.join(color_dir, filename)\n",
    "            grayscale_path = os.path.join(grayscale_dir, filename)\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "\n",
    "            if os.path.exists(grayscale_path):\n",
    "                print(f\"Processing {filename}...\")\n",
    "                resize_and_match_properties(color_path, grayscale_path, output_path)\n",
    "                print(f\"Saved processed file to {output_path}\")\n",
    "            else:\n",
    "                print(f\"Warning: No matching grayscale file found for {filename}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    color_dir = \"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT_512\"\n",
    "    grayscale_dir = \"/mnt/d/CTH_archive/TMAX_NIFTI_512\"\n",
    "    output_dir = \"/mnt/d/CTH_archive//CBF_COLORED_NIFTI_RESIZED_Processed\"\n",
    "\n",
    "    process_directory(color_dir, grayscale_dir, output_dir)\n",
    "    print(\"Processing complete.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import logging\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "from matplotlib import cm\n",
    "\n",
    "FINAL_SIZE = 512 \n",
    "\n",
    "logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
    "\n",
    "import os\n",
    "import numpy as np\n",
    "from matplotlib import cm\n",
    "import SimpleITK as sitk\n",
    "\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "\n",
    "def process_3d_rgb_image(image):\n",
    "    try:\n",
    "        # Convert SimpleITK image to numpy array\n",
    "        array = sitk.GetArrayFromImage(image)\n",
    "        \n",
    "        # Ensure the array is in the correct shape (depth, height, width, channels)\n",
    "        if array.shape[-1] != 3:\n",
    "            array = np.moveaxis(array, 0, -1)\n",
    "        \n",
    "        # Normalize the array to 0-255 range if it's not already\n",
    "        if array.max() <= 1.0:\n",
    "            array = (array * 255).astype(np.uint8)\n",
    "        \n",
    "        # Convert back to SimpleITK image (RGB)\n",
    "        processed_image = sitk.GetImageFromArray(array, isVector=True)\n",
    "        \n",
    "        # Set metadata for the new image\n",
    "        processed_image.SetSpacing(image.GetSpacing())\n",
    "        processed_image.SetDirection(image.GetDirection())\n",
    "        processed_image.SetOrigin(image.GetOrigin())\n",
    "        \n",
    "        return processed_image\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Error in process_3d_rgb_image: {str(e)}\")\n",
    "        logging.error(f\"Image shape: {image.GetSize()}, Pixel type: {image.GetPixelID()}\")\n",
    "        logging.exception(\"Stack trace:\")\n",
    "        raise\n",
    "    \n",
    "def resample_image(moving_image, ctp_image):\n",
    "    desired_size = [FINAL_SIZE, FINAL_SIZE, ctp_image.GetSize()[2]]  # Use the same number of slices as the CTP image\n",
    "    \n",
    "    resampler = sitk.ResampleImageFilter()\n",
    "    resampler.SetReferenceImage(moving_image)\n",
    "    resampler.SetSize(desired_size)\n",
    "    resampler.SetOutputSpacing([moving_image.GetSpacing()[i] * (moving_image.GetSize()[i] / desired_size[i]) for i in range(3)])\n",
    "    resampler.SetInterpolator(sitk.sitkLinear)\n",
    "    \n",
    "    resized_image = resampler.Execute(moving_image)\n",
    "    resized_image.SetSpacing(ctp_image.GetSpacing())\n",
    "    resized_image.SetOrigin(ctp_image.GetOrigin())\n",
    "    resized_image.SetDirection(ctp_image.GetDirection())\n",
    "    \n",
    "    return resized_image\n",
    "\n",
    "def correct_orientation(image, reference_image):\n",
    "    image_direction = np.array(image.GetDirection()).reshape(3, 3)\n",
    "    reference_direction = np.array(reference_image.GetDirection()).reshape(3, 3)\n",
    "    \n",
    "    if not np.allclose(image_direction, reference_direction):\n",
    "        transform = sitk.AffineTransform(3)\n",
    "        transform.SetMatrix(np.dot(reference_direction, np.linalg.inv(image_direction)).flatten())\n",
    "        \n",
    "        corrected_image = sitk.Resample(image, reference_image, transform, sitk.sitkLinear, 0.0, image.GetPixelID())\n",
    "        \n",
    "        corrected_image.SetOrigin(reference_image.GetOrigin())\n",
    "        corrected_image.SetSpacing(reference_image.GetSpacing())\n",
    "        corrected_image.SetDirection(reference_image.GetDirection())\n",
    "        \n",
    "        return corrected_image\n",
    "    else:\n",
    "        return image\n",
    "\n",
    "def apply_final_transform(resized_moving_image, fixed_image, transform_file_path):\n",
    "    try:\n",
    "        final_transform = sitk.ReadTransform(transform_file_path)\n",
    "        \n",
    "        if resized_moving_image is not None:\n",
    "            resampled_image = sitk.Resample(resized_moving_image, \n",
    "                                            fixed_image, \n",
    "                                            final_transform, \n",
    "                                            sitk.sitkBSpline, \n",
    "                                            0.0, \n",
    "                                            resized_moving_image.GetPixelID())\n",
    "            \n",
    "            resampled_image = correct_orientation(resampled_image, fixed_image)\n",
    "            \n",
    "            return resampled_image\n",
    "        else:\n",
    "            logging.error(\"Resized moving image is invalid. Cannot apply final transform.\")\n",
    "            return None\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Failed to apply final transform: {str(e)}\")\n",
    "        return None\n",
    "\n",
    "def main():\n",
    "    tmax_nifti_dir = r\"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT_512\"\n",
    "    transforms_dir = r\"/mnt/d/CTH_archive/CTH_TMAX_TRANSFORMS_COLOR\"\n",
    "    cth_stripped_dir = r\"/mnt/d/CTH_archive/CTH_STRIPPED\"\n",
    "    ctp_stripped_dir = r\"/mnt/d/CTH_archive/CTP_STRIPPED\"\n",
    "    \n",
    "    output_dir = r\"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT_REG\"\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "\n",
    "    for transform_file in os.listdir(transforms_dir):\n",
    "        transform_path = os.path.join(transforms_dir, transform_file)\n",
    "        base_filename = transform_file.replace('.h5', '')\n",
    "\n",
    "        moving_image_path = os.path.join(tmax_nifti_dir, base_filename + '.nii')\n",
    "        fixed_image_path = os.path.join(cth_stripped_dir, base_filename + '.nii')\n",
    "        ctp_image_path = os.path.join(ctp_stripped_dir, base_filename + '.nii')\n",
    "        logging.info(f\"Processing {base_filename}...\")\n",
    "\n",
    "        if os.path.exists(moving_image_path) and os.path.exists(fixed_image_path) and os.path.exists(ctp_image_path):\n",
    "            try:\n",
    "                moving_image = sitk.ReadImage(moving_image_path)\n",
    "                fixed_image = sitk.ReadImage(fixed_image_path)\n",
    "                ctp_image = sitk.ReadImage(ctp_image_path)\n",
    "\n",
    "                logging.info(f\"Moving image: size={moving_image.GetSize()}, pixel type={moving_image.GetPixelID()}\")\n",
    "                logging.info(f\"Fixed image: size={fixed_image.GetSize()}, pixel type={fixed_image.GetPixelID()}\")\n",
    "                logging.info(f\"CTP image: size={ctp_image.GetSize()}, pixel type={ctp_image.GetPixelID()}\")\n",
    "\n",
    "                resized_moving_image = resample_image(moving_image, ctp_image)\n",
    "                logging.info(f\"Resized moving image: size={resized_moving_image.GetSize()}, pixel type={resized_moving_image.GetPixelID()}\")\n",
    "\n",
    "                resampled_image = apply_final_transform(resized_moving_image, fixed_image, transform_path)\n",
    "\n",
    "                if resampled_image is not None:\n",
    "                    logging.info(f\"Resampled image: size={resampled_image.GetSize()}, pixel type={resampled_image.GetPixelID()}\")\n",
    "                    logging.info(f\"Resampled image origin: {resampled_image.GetOrigin()}\")\n",
    "                    logging.info(f\"Resampled image spacing: {resampled_image.GetSpacing()}\")\n",
    "                    logging.info(f\"Resampled image direction: {resampled_image.GetDirection()}\")\n",
    "                    \n",
    "                    # Process the 3D RGB image\n",
    "                    processed_image = process_3d_rgb_image(resampled_image)\n",
    "                    \n",
    "                    output_image_path = os.path.join(output_dir, base_filename + '.nii')\n",
    "                    sitk.WriteImage(processed_image, output_image_path)\n",
    "                    logging.info(f\"Processed and saved: {output_image_path}\")\n",
    "                else:\n",
    "                    logging.error(f\"Failed to process {base_filename}. Skipping.\")\n",
    "            except Exception as e:\n",
    "                logging.error(f\"Error processing {base_filename}: {str(e)}\")\n",
    "                logging.exception(\"Stack trace:\")\n",
    "        else:\n",
    "            logging.warning(f\"Required files for {base_filename} are not available.\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import SimpleITK as sitk\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Set up logging\n",
    "logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\n",
    "\n",
    "def normalize_array(array):\n",
    "    \"\"\"Normalize the array to 0-255 range.\"\"\"\n",
    "    min_val = np.min(array)\n",
    "    max_val = np.max(array)\n",
    "    if min_val == max_val:\n",
    "        return np.zeros_like(array)\n",
    "    return ((array - min_val) / (max_val - min_val) * 255).astype(np.uint8)\n",
    "\n",
    "def adjust_brightness(array, factor=1.2):\n",
    "    \"\"\"Adjust brightness of the image.\"\"\"\n",
    "    return np.clip(array * factor, 0, 255).astype(np.uint8)\n",
    "\n",
    "def enhance_contrast(array, low_percentile=2, high_percentile=98):\n",
    "    \"\"\"Enhance contrast using percentile-based contrast stretching.\"\"\"\n",
    "    low = np.percentile(array, low_percentile)\n",
    "    high = np.percentile(array, high_percentile)\n",
    "    return np.clip((array - low) * (255.0 / (high - low)), 0, 255).astype(np.uint8)\n",
    "\n",
    "def process_rgb_array(array, brightness_factor=1.2, contrast_low=2, contrast_high=98):\n",
    "    \"\"\"Process the RGB array: normalize, enhance contrast, adjust brightness, and ensure black background.\"\"\"\n",
    "    # Normalize each channel\n",
    "    normalized = np.stack([normalize_array(array[..., i]) for i in range(3)], axis=-1)\n",
    "    \n",
    "    # Enhance contrast\n",
    "    contrasted = np.stack([enhance_contrast(normalized[..., i], contrast_low, contrast_high) for i in range(3)], axis=-1)\n",
    "    \n",
    "    # Adjust brightness\n",
    "    brightened = adjust_brightness(contrasted, brightness_factor)\n",
    "    \n",
    "    # Ensure background remains black\n",
    "    black_mask = np.all(array == 0, axis=-1)\n",
    "    brightened[black_mask] = 0\n",
    "    \n",
    "    return brightened\n",
    "\n",
    "def visualize_slice(original, processed, slice_num):\n",
    "    \"\"\"Visualize original and processed slices side by side.\"\"\"\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))\n",
    "    ax1.imshow(original)\n",
    "    ax1.set_title(f\"Original Slice {slice_num}\")\n",
    "    ax2.imshow(processed)\n",
    "    ax2.set_title(f\"Processed Slice {slice_num}\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "def nifti_to_png_slices(input_dir, output_dir, brightness_factor=1.2, contrast_low=2, contrast_high=98):\n",
    "    # Ensure output directory exists\n",
    "    os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "    # Iterate through all .nii files in the input directory\n",
    "    for filename in os.listdir(input_dir):\n",
    "        if filename.endswith('.nii'):\n",
    "            nifti_path = os.path.join(input_dir, filename)\n",
    "            base_name = os.path.splitext(filename)[0]\n",
    "\n",
    "            try:\n",
    "                # Read the NIFTI file\n",
    "                image = sitk.ReadImage(nifti_path)\n",
    "                \n",
    "                # Log image information\n",
    "                logging.info(f\"Processing {filename}\")\n",
    "                logging.info(f\"Image size: {image.GetSize()}\")\n",
    "                logging.info(f\"Image pixel type: {image.GetPixelIDTypeAsString()}\")\n",
    "                \n",
    "                # Convert to numpy array\n",
    "                array = sitk.GetArrayFromImage(image)\n",
    "                logging.info(f\"Array shape: {array.shape}\")\n",
    "                logging.info(f\"Array min: {array.min()}, max: {array.max()}, mean: {array.mean()}\")\n",
    "\n",
    "                # Iterate through axial slices\n",
    "                for i in range(array.shape[0]):  # Assuming axial slices are along the first dimension\n",
    "                    slice_array = array[i, :, :, :]\n",
    "\n",
    "                    # Process the RGB slice\n",
    "                    slice_array_processed = process_rgb_array(slice_array, brightness_factor, contrast_low, contrast_high)\n",
    "\n",
    "                    # Visualize every 10th slice\n",
    "                    if i % 10 == 0:\n",
    "                        visualize_slice(normalize_array(slice_array), slice_array_processed, i)\n",
    "\n",
    "                    # Create PIL Image\n",
    "                    slice_image = Image.fromarray(slice_array_processed)\n",
    "\n",
    "                    # Resize to 512x512 if not already that size\n",
    "                    if slice_image.size != (512, 512):\n",
    "                        slice_image = slice_image.resize((512, 512), Image.LANCZOS)\n",
    "\n",
    "                    # Save as PNG\n",
    "                    output_filename = f\"{base_name}_{i+1:06d}.png\"\n",
    "                    output_path = os.path.join(output_dir, output_filename)\n",
    "                    slice_image.save(output_path)\n",
    "\n",
    "                logging.info(f\"Processed {filename} - saved {array.shape[0]} slices\")\n",
    "\n",
    "            except Exception as e:\n",
    "                logging.error(f\"Error processing {filename}: {str(e)}\")\n",
    "\n",
    "# Define input and output directories\n",
    "input_dir = r\"/mnt/d/CTH_ARCHIVE/CBF_COLORED_NIFIT_NOTEXT_REG_BLACK/\"\n",
    "output_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/CBF_COLORED_PNG_NOTEXT_REG/\"\n",
    "\n",
    "# Run the conversion with contrast enhancement\n",
    "nifti_to_png_slices(input_dir, output_dir, brightness_factor=1.2, contrast_low=40, contrast_high=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import random\n",
    "\n",
    "def split_datasets(tmax_dir, cth_dir, train_a_dir, test_a_dir, train_b_dir, test_b_dir, split_ratio=0.9):\n",
    "    # Create output directories if they don't exist\n",
    "    for dir in [train_a_dir, test_a_dir, train_b_dir, test_b_dir]:\n",
    "        os.makedirs(dir, exist_ok=True)\n",
    "\n",
    "    # Get all files from TMAX_PNG_REG\n",
    "    tmax_files = [f for f in os.listdir(tmax_dir) if f.endswith('.png')]\n",
    "\n",
    "    # Randomly shuffle the files\n",
    "    random.shuffle(tmax_files)\n",
    "\n",
    "    # Calculate the split index\n",
    "    split_index = int(len(tmax_files) * split_ratio)\n",
    "\n",
    "    # Split the files\n",
    "    train_files = tmax_files[:split_index]\n",
    "    test_files = tmax_files[split_index:]\n",
    "\n",
    "    # Function to copy files\n",
    "    def copy_files(file_list, src_dir_a, dst_dir_a, src_dir_b, dst_dir_b):\n",
    "        for file in file_list:\n",
    "            # Copy TMAX file\n",
    "            shutil.copy(os.path.join(src_dir_a, file), os.path.join(dst_dir_a, file))\n",
    "            \n",
    "            # Copy corresponding CTH file\n",
    "            cth_file = file  # Assuming the filenames are the same\n",
    "            if os.path.exists(os.path.join(src_dir_b, cth_file)):\n",
    "                shutil.copy(os.path.join(src_dir_b, cth_file), os.path.join(dst_dir_b, cth_file))\n",
    "            else:\n",
    "                print(f\"Warning: Corresponding file {cth_file} not found in {src_dir_b}\")\n",
    "\n",
    "    # Copy training files\n",
    "    copy_files(train_files, tmax_dir, train_a_dir, cth_dir, train_b_dir)\n",
    "\n",
    "    # Copy test files\n",
    "    copy_files(test_files, tmax_dir, test_a_dir, cth_dir, test_b_dir)\n",
    "\n",
    "    print(f\"Split complete. {len(train_files)} files in training set, {len(test_files)} files in test set.\")\n",
    "\n",
    "# Define directories\n",
    "tmax_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/CBF_COLORED_PNG_NOTEXT_REG\"\n",
    "cth_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/CTH_PNG\"\n",
    "train_a_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A\"\n",
    "test_a_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/test_A\"\n",
    "train_b_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_B\"\n",
    "test_b_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/test_B\"\n",
    "\n",
    "# Run the splitting process\n",
    "split_datasets(tmax_dir, cth_dir, train_a_dir, test_a_dir, train_b_dir, test_b_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "\n",
    "def create_prompt_json(dir_a, dir_b, output_json):\n",
    "    # Get all PNG files from both directories\n",
    "    files_a = set(f for f in os.listdir(dir_a) if f.endswith('.png'))\n",
    "    files_b = set(f for f in os.listdir(dir_b) if f.endswith('.png'))\n",
    "\n",
    "    # Find common files (in case there are any mismatches)\n",
    "    common_files = sorted(list(files_a.intersection(files_b)))\n",
    "\n",
    "    # Create a dictionary with file names as keys and empty strings as values\n",
    "    prompt_dict = {file: \"\" for file in common_files}\n",
    "\n",
    "    # Write the dictionary to a JSON file\n",
    "    with open(output_json, 'w') as f:\n",
    "        json.dump(prompt_dict, f, indent=2)\n",
    "\n",
    "    print(f\"Created {output_json} with {len(common_files)} entries.\")\n",
    "\n",
    "# Define directories and output files\n",
    "train_a_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/train_A\"\n",
    "train_b_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/train_B\"\n",
    "test_a_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/test_A\"\n",
    "test_b_dir = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/test_B\"\n",
    "\n",
    "train_json = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/train_prompts.json\"\n",
    "test_json = r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/test_prompts.json\"\n",
    "\n",
    "# Create JSON files for train and test sets\n",
    "create_prompt_json(train_a_dir, train_b_dir, train_json)\n",
    "create_prompt_json(test_a_dir, test_b_dir, test_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "def convert_grayscale_to_3channel(directory):\n",
    "    for filename in os.listdir(directory):\n",
    "        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
    "            file_path = os.path.join(directory, filename)\n",
    "            \n",
    "            with Image.open(file_path) as img:\n",
    "                # Check if the image is grayscale\n",
    "                if img.mode == 'L':\n",
    "                    print(f\"Converting {filename} to 3-channel image...\")\n",
    "                    \n",
    "                    # Convert to numpy array\n",
    "                    img_array = np.array(img)\n",
    "                    \n",
    "                    # Stack the single channel three times to create a 3-channel image\n",
    "                    rgb_array = np.stack((img_array,)*3, axis=-1)\n",
    "                    \n",
    "                    # Create a new image from the 3-channel array\n",
    "                    rgb_img = Image.fromarray(rgb_array.astype('uint8'), 'RGB')\n",
    "                    \n",
    "                    # Save the image, overwriting the original file\n",
    "                    rgb_img.save(file_path)\n",
    "                else:\n",
    "                    print(f\"{filename} is already a multi-channel image. Skipping...\")\n",
    "\n",
    "# Directories to process\n",
    "directories = [\n",
    "    r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/test_B\",\n",
    "    r\"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/train_B\"\n",
    "]\n",
    "\n",
    "# Process each directory\n",
    "for directory in directories:\n",
    "    print(f\"Processing directory: {directory}\")\n",
    "    convert_grayscale_to_3channel(directory)\n",
    "    print(f\"Finished processing {directory}\")\n",
    "\n",
    "print(\"All directories processed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RUN CODE MODEL HERE "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity\n",
    "\n",
    "def calculate_metrics(output_dir, ground_truth_dir):\n",
    "    mse_values = []\n",
    "    psnr_values = []\n",
    "    ssim_values = []\n",
    "    \n",
    "    for filename in os.listdir(output_dir):\n",
    "        if filename.lower().endswith('.png'):\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "            ground_truth_path = os.path.join(ground_truth_dir, filename)\n",
    "            \n",
    "            if os.path.exists(ground_truth_path):\n",
    "                # Open images\n",
    "                output_img = np.array(Image.open(output_path).convert('RGB'))\n",
    "                ground_truth_img = np.array(Image.open(ground_truth_path).convert('RGB'))\n",
    "                \n",
    "                # Ensure images have the same shape\n",
    "                if output_img.shape != ground_truth_img.shape:\n",
    "                    print(f\"Shape mismatch for {filename}. Skipping.\")\n",
    "                    continue\n",
    "                \n",
    "                # Calculate MSE\n",
    "                mse = mean_squared_error(ground_truth_img, output_img)\n",
    "                mse_values.append(mse)\n",
    "                \n",
    "                # Calculate PSNR\n",
    "                psnr = peak_signal_noise_ratio(ground_truth_img, output_img, data_range=255)\n",
    "                psnr_values.append(psnr)\n",
    "                \n",
    "                # Calculate SSIM\n",
    "                ssim = structural_similarity(ground_truth_img, output_img, channel_axis=2, data_range=255)\n",
    "                ssim_values.append(ssim)\n",
    "                \n",
    "                print(f\"Metrics for {filename}:\")\n",
    "                print(f\"  MSE: {mse:.4f}\")\n",
    "                print(f\"  PSNR: {psnr:.4f} dB\")\n",
    "                print(f\"  SSIM: {ssim:.4f}\")\n",
    "            else:\n",
    "                print(f\"Ground truth file not found for {filename}. Skipping.\")\n",
    "    \n",
    "    if mse_values and psnr_values and ssim_values:\n",
    "        average_mse = np.mean(mse_values)\n",
    "        average_psnr = np.mean(psnr_values)\n",
    "        average_ssim = np.mean(ssim_values)\n",
    "        print(f\"\\nAverage Metrics:\")\n",
    "        print(f\"  MSE: {average_mse:.4f}\")\n",
    "        print(f\"  PSNR: {average_psnr:.4f} dB\")\n",
    "        print(f\"  SSIM: {average_ssim:.4f}\")\n",
    "    else:\n",
    "        print(\"No valid image pairs found for metric calculation.\")\n",
    "\n",
    "# Set your paths\n",
    "output_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/output/pix2pix_turbo/CTH_CBF_AUG/\"\n",
    "ground_truth_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/test_B\"\n",
    "\n",
    "# Calculate metrics\n",
    "calculate_metrics(output_dir, ground_truth_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity\n",
    "\n",
    "def calculate_metrics(output_dir, ground_truth_dir):\n",
    "    mse_values = []\n",
    "    psnr_values = []\n",
    "    ssim_values = []\n",
    "    \n",
    "    for filename in os.listdir(output_dir):\n",
    "        if filename.lower().endswith('.png'):\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "            ground_truth_path = os.path.join(ground_truth_dir, filename)\n",
    "            \n",
    "            if os.path.exists(ground_truth_path):\n",
    "                try:\n",
    "                    # Open images\n",
    "                    output_img = np.array(Image.open(output_path).convert('RGB'))\n",
    "                    ground_truth_img = np.array(Image.open(ground_truth_path).convert('RGB'))\n",
    "                except Exception as e:\n",
    "                    print(f\"Error loading images for {filename}: {str(e)}\")\n",
    "                    continue\n",
    "                \n",
    "                # Ensure images have the same shape\n",
    "                if output_img.shape != ground_truth_img.shape:\n",
    "                    print(f\"Shape mismatch for {filename}. Skipping.\")\n",
    "                    continue\n",
    "                \n",
    "                # Calculate MSE\n",
    "                mse = mean_squared_error(ground_truth_img, output_img)\n",
    "                if mse == 0:\n",
    "                    print(f\"MSE is zero for {filename}. Skipping.\")\n",
    "                    continue\n",
    "                \n",
    "                mse_values.append(mse)\n",
    "                \n",
    "                # Calculate PSNR\n",
    "                psnr = peak_signal_noise_ratio(ground_truth_img, output_img, data_range=255)\n",
    "                psnr_values.append(psnr)\n",
    "                \n",
    "                # Calculate SSIM\n",
    "                ssim = structural_similarity(ground_truth_img, output_img, channel_axis=2, data_range=255)\n",
    "                ssim_values.append(ssim)\n",
    "                \n",
    "                print(f\"Metrics for {filename}:\")\n",
    "                print(f\"  MSE: {mse:.4f}\")\n",
    "                print(f\"  PSNR: {psnr:.4f} dB\")\n",
    "                print(f\"  SSIM: {ssim:.4f}\")\n",
    "            else:\n",
    "                print(f\"Ground truth file not found for {filename}. Skipping.\")\n",
    "    \n",
    "    if mse_values and psnr_values and ssim_values:\n",
    "        average_mse = np.mean(mse_values)\n",
    "        average_psnr = np.mean(psnr_values)\n",
    "        average_ssim = np.mean(ssim_values)\n",
    "        print(f\"\\nAverage Metrics:\")\n",
    "        print(f\"  MSE: {average_mse:.4f}\")\n",
    "        print(f\"  PSNR: {average_psnr:.4f} dB\")\n",
    "        print(f\"  SSIM: {average_ssim:.4f}\")\n",
    "        \n",
    "        print(\"\\nDetailed Statistics:\")\n",
    "        print(f\"MSE - Min: {np.min(mse_values):.4f}, Max: {np.max(mse_values):.4f}, Median: {np.median(mse_values):.4f}\")\n",
    "        print(f\"PSNR - Min: {np.min(psnr_values):.4f}, Max: {np.max(psnr_values):.4f}, Median: {np.median(psnr_values):.4f}\")\n",
    "        print(f\"SSIM - Min: {np.min(ssim_values):.4f}, Max: {np.max(ssim_values):.4f}, Median: {np.median(ssim_values):.4f}\")\n",
    "    else:\n",
    "        print(\"No valid image pairs found for metric calculation.\")\n",
    "\n",
    "# Set your paths\n",
    "output_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/output/pix2pix_turbo/CTH_CBF_AUG/\"\n",
    "ground_truth_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG/test_B\"\n",
    "\n",
    "# Calculate metrics\n",
    "calculate_metrics(output_dir, ground_truth_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from skimage.metrics import mean_squared_error, peak_signal_noise_ratio, structural_similarity\n",
    "from scipy.linalg import sqrtm\n",
    "from tensorflow.keras.applications.inception_v3 import InceptionV3, preprocess_input\n",
    "\n",
    "def calculate_metrics(output_dir, ground_truth_dir):\n",
    "    mse_values = []\n",
    "    psnr_values = []\n",
    "    ssim_values = []\n",
    "    output_images = []\n",
    "    ground_truth_images = []\n",
    "    \n",
    "    for filename in os.listdir(output_dir):\n",
    "        if filename.lower().endswith('.png'):\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "            ground_truth_path = os.path.join(ground_truth_dir, filename)\n",
    "            \n",
    "            if os.path.exists(ground_truth_path):\n",
    "                # Open images\n",
    "                output_img = np.array(Image.open(output_path).convert('RGB'))\n",
    "                ground_truth_img = np.array(Image.open(ground_truth_path).convert('RGB'))\n",
    "                \n",
    "                # Ensure images have the same shape\n",
    "                if output_img.shape != ground_truth_img.shape:\n",
    "                    print(f\"Shape mismatch for {filename}. Skipping.\")\n",
    "                    continue\n",
    "                \n",
    "                # Calculate MSE\n",
    "                mse = mean_squared_error(ground_truth_img, output_img)\n",
    "                mse_values.append(mse)\n",
    "                \n",
    "                # Calculate PSNR\n",
    "                psnr = peak_signal_noise_ratio(ground_truth_img, output_img, data_range=255)\n",
    "                psnr_values.append(psnr)\n",
    "                \n",
    "                # Calculate SSIM\n",
    "                ssim = structural_similarity(ground_truth_img, output_img, channel_axis=2, data_range=255)\n",
    "                ssim_values.append(ssim)\n",
    "                \n",
    "                # Store images for FID calculation\n",
    "                output_images.append(output_img)\n",
    "                ground_truth_images.append(ground_truth_img)\n",
    "                \n",
    "                print(f\"Metrics for {filename}:\")\n",
    "                print(f\"  MSE: {mse:.4f}\")\n",
    "                print(f\"  PSNR: {psnr:.4f} dB\")\n",
    "                print(f\"  SSIM: {ssim:.4f}\")\n",
    "            else:\n",
    "                print(f\"Ground truth file not found for {filename}. Skipping.\")\n",
    "    \n",
    "    # Calculate FID\n",
    "    if output_images and ground_truth_images:\n",
    "        fid_value = calculate_fid(output_images, ground_truth_images)\n",
    "        print(f\"\\nFID: {fid_value:.4f}\")\n",
    "    \n",
    "    if mse_values and psnr_values and ssim_values:\n",
    "        average_mse = np.mean(mse_values)\n",
    "        average_psnr = np.mean(psnr_values)\n",
    "        average_ssim = np.mean(ssim_values)\n",
    "        print(f\"\\nAverage Metrics:\")\n",
    "        print(f\"  MSE: {average_mse:.4f}\")\n",
    "        print(f\"  PSNR: {average_psnr:.4f} dB\")\n",
    "        print(f\"  SSIM: {average_ssim:.4f}\")\n",
    "    else:\n",
    "        print(\"No valid image pairs found for metric calculation.\")\n",
    "\n",
    "def calculate_fid(output_images, ground_truth_images):\n",
    "    # Load the InceptionV3 model\n",
    "    model = InceptionV3(include_top=False, pooling='avg', input_shape=(299, 299, 3))\n",
    "    \n",
    "    # Prepare the images\n",
    "    output_images_resized = np.array([np.array(Image.fromarray(img).resize((299, 299))) for img in output_images])\n",
    "    ground_truth_images_resized = np.array([np.array(Image.fromarray(img).resize((299, 299))) for img in ground_truth_images])\n",
    "    \n",
    "    # Preprocess images\n",
    "    output_images_resized = preprocess_input(output_images_resized)\n",
    "    ground_truth_images_resized = preprocess_input(ground_truth_images_resized)\n",
    "    \n",
    "    # Get features\n",
    "    output_features = model.predict(output_images_resized)\n",
    "    ground_truth_features = model.predict(ground_truth_images_resized)\n",
    "    \n",
    "    # Calculate mean and covariance for both sets of features\n",
    "    mu1, sigma1 = np.mean(output_features, axis=0), np.cov(output_features, rowvar=False)\n",
    "    mu2, sigma2 = np.mean(ground_truth_features, axis=0), np.cov(ground_truth_features, rowvar=False)\n",
    "    \n",
    "    # Calculate FID\n",
    "    diff = mu1 - mu2\n",
    "    covmean = sqrtm(sigma1.dot(sigma2))\n",
    "    \n",
    "    # Numerical stability check\n",
    "    if np.iscomplexobj(covmean):\n",
    "        covmean = covmean.real\n",
    "    \n",
    "    fid_value = diff.dot(diff) + np.trace(sigma1 + sigma2 - 2 * covmean)\n",
    "    return fid_value\n",
    "\n",
    "# Set your paths\n",
    "output_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/output/pix2pix_turbo/CTH_TMAX_COLOR/\"\n",
    "ground_truth_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/test_B\"\n",
    "\n",
    "# Calculate metrics\n",
    "calculate_metrics(output_dir, ground_truth_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "def is_mostly_black(image_path, threshold=0.995):\n",
    "    with Image.open(image_path) as img:\n",
    "        img_array = np.array(img)\n",
    "        if len(img_array.shape) == 3:  # RGB image\n",
    "            return np.mean(img_array == 0) >= threshold\n",
    "        else:  # Grayscale image\n",
    "            return np.mean(img_array == 0) >= threshold\n",
    "\n",
    "def remove_files_and_update_json(dir_a, dir_b, json_file):\n",
    "    with open(json_file, 'r') as f:\n",
    "        prompts = json.load(f)\n",
    "\n",
    "    files_to_remove = []\n",
    "\n",
    "    for filename in os.listdir(dir_a):\n",
    "        if filename.lower().endswith('.png'):\n",
    "            file_path_a = os.path.join(dir_a, filename)\n",
    "            file_path_b = os.path.join(dir_b, filename)\n",
    "\n",
    "            if is_mostly_black(file_path_a):\n",
    "                print(f\"Removing {filename}\")\n",
    "                os.remove(file_path_a)\n",
    "                if os.path.exists(file_path_b):\n",
    "                    os.remove(file_path_b)\n",
    "                files_to_remove.append(filename)\n",
    "\n",
    "    # Update JSON\n",
    "    for file in files_to_remove:\n",
    "        prompts.pop(file, None)\n",
    "\n",
    "    with open(json_file, 'w') as f:\n",
    "        json.dump(prompts, f, indent=2)\n",
    "\n",
    "    print(f\"Removed {len(files_to_remove)} files and updated {json_file}\")\n",
    "\n",
    "# Paths\n",
    "test_a = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF/test_A\"\n",
    "test_b = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF/test_B\"\n",
    "train_a = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/train_A\"\n",
    "train_b = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/train_B\"\n",
    "test_json = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/test_prompts.json\"\n",
    "train_json = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/train_prompts.json\"\n",
    "\n",
    "# Process test and train sets\n",
    "remove_files_and_update_json(test_a, test_b, test_json)\n",
    "remove_files_and_update_json(train_a, train_b, train_json)\n",
    "\n",
    "print(\"Processing complete.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import os\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "def apply_clahe(image_path, output_path, clip_limit=2.0, tile_grid_size=(8, 8)):\n",
    "    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)\n",
    "    \n",
    "    # Create a mask for the brain tissue (assuming brain tissue is not black)\n",
    "    brain_mask = image > 0\n",
    "\n",
    "    # Apply CLAHE to the brain tissue\n",
    "    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)\n",
    "    clahe_image = clahe.apply(image)\n",
    "    \n",
    "    # Apply the mask to keep the background black\n",
    "    final_image = np.where(brain_mask, clahe_image, 0)\n",
    "\n",
    "    cv2.imwrite(output_path, final_image)\n",
    "\n",
    "def process_directory(input_dir, output_dir, clip_limit=2.0, tile_grid_size=(8, 8)):\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.makedirs(output_dir)\n",
    "\n",
    "    for filename in tqdm(os.listdir(input_dir)):\n",
    "        if filename.endswith('.png'):\n",
    "            input_path = os.path.join(input_dir, filename)\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "            apply_clahe(input_path, output_path, clip_limit, tile_grid_size)\n",
    "\n",
    "# Paths to your datasets\n",
    "input_dir_train = '/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/train_A'\n",
    "output_dir_train = '/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/train_A_enhanced'\n",
    "\n",
    "input_dir_test = '/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/test_A'\n",
    "output_dir_test = '/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_TMAX/test_A_enhanced'\n",
    "\n",
    "# Process the directories\n",
    "process_directory(input_dir_train, output_dir_train)\n",
    "process_directory(input_dir_test, output_dir_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pydicom\n",
    "import numpy as np\n",
    "import SimpleITK as sitk\n",
    "\n",
    "def convert_dicom_series_to_nii(dicom_dir, output_dir):\n",
    "    # Get patient name and ID from the directory name\n",
    "    patient_name_id = os.path.basename(dicom_dir)\n",
    "\n",
    "    # Load DICOM series\n",
    "    reader = sitk.ImageSeriesReader()\n",
    "    dicom_names = reader.GetGDCMSeriesFileNames(dicom_dir)\n",
    "    reader.SetFileNames(dicom_names)\n",
    "\n",
    "    # Read the image series\n",
    "    image = reader.Execute()\n",
    "\n",
    "    # Ensure the output image has RGB channels\n",
    "    image_array = sitk.GetArrayFromImage(image)  # shape: (slices, height, width)\n",
    "    if image_array.ndim == 3:  # Single channel, need to convert to RGB\n",
    "        rgb_image_array = np.stack((image_array,) * 3, axis=-1)  # shape: (slices, height, width, 3)\n",
    "    else:\n",
    "        rgb_image_array = image_array\n",
    "\n",
    "    # Resize to 256x256 if necessary\n",
    "    original_size = rgb_image_array.shape[1:3]\n",
    "    if original_size != (256, 256):\n",
    "        resized_image_array = np.zeros((rgb_image_array.shape[0], 256, 256, 3), dtype=rgb_image_array.dtype)\n",
    "        for i in range(rgb_image_array.shape[0]):\n",
    "            slice_image = sitk.GetImageFromArray(rgb_image_array[i])\n",
    "            resized_slice = sitk.Resample(slice_image, [256, 256], sitk.Transform(), sitk.sitkLinear)\n",
    "            resized_image_array[i] = sitk.GetArrayFromImage(resized_slice)\n",
    "    else:\n",
    "        resized_image_array = rgb_image_array\n",
    "\n",
    "    # Create SimpleITK image from RGB array\n",
    "    rgb_image = sitk.GetImageFromArray(resized_image_array)\n",
    "    rgb_image.CopyInformation(image)  # Preserve orientation, direction, origin\n",
    "\n",
    "    # Save as .nii file\n",
    "    output_file = os.path.join(output_dir, f\"{patient_name_id}.nii\")\n",
    "    sitk.WriteImage(rgb_image, output_file)\n",
    "\n",
    "    print(f\"Converted {patient_name_id} to {output_file}\")\n",
    "\n",
    "def process_all_dicom_series(base_dir, output_dir):\n",
    "    # Recursively find all patient directories\n",
    "    for root, dirs, files in os.walk(base_dir):\n",
    "        if any(file.endswith('.dcm') for file in files):\n",
    "            convert_dicom_series_to_nii(root, output_dir)\n",
    "\n",
    "# Directory paths\n",
    "base_dir = \"/mnt/d/CTH_archive/CBF_COLORED\"\n",
    "output_dir = \"/mnt/d/CTH_archive/CBF_COLORED_NIFTI\"\n",
    "\n",
    "# Convert all series in the base directory\n",
    "process_all_dicom_series(base_dir, output_dir)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nibabel as nib\n",
    "import numpy as np\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "def make_background_black_fast(input_path, output_path, threshold=0.05):\n",
    "    # Load the NIfTI image\n",
    "    nii_img = nib.load(input_path)\n",
    "    data = nii_img.get_fdata()\n",
    "\n",
    "    # Check if the data is 5D with the expected shape\n",
    "    if data.shape[-2:] == (1, 3):\n",
    "        # Create a mask where all channels are below the threshold\n",
    "        mask = np.all(data < threshold, axis=-1)\n",
    "        \n",
    "        # Expand the mask to match the original shape\n",
    "        mask = np.expand_dims(mask, axis=-1)\n",
    "        mask = np.repeat(mask, 3, axis=-1)\n",
    "        \n",
    "        # Apply the mask\n",
    "        data[mask] = 0\n",
    "    else:\n",
    "        raise ValueError(f\"Unexpected data shape: {data.shape}. Expected a 5D image with shape (..., 1, 3).\")\n",
    "\n",
    "    # Create a new NIfTI image\n",
    "    new_img = nib.Nifti1Image(data, nii_img.affine, header=nii_img.header)\n",
    "\n",
    "    # Save the modified image\n",
    "    nib.save(new_img, output_path)\n",
    "    print(f\"Processed image saved to {output_path}\")\n",
    "\n",
    "def process_directory(input_dir, output_dir):\n",
    "    # Create the output directory if it doesn't exist\n",
    "    Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    # Process all .nii files in the input directory\n",
    "    for filename in os.listdir(input_dir):\n",
    "        if filename.endswith('.nii'):\n",
    "            input_path = os.path.join(input_dir, filename)\n",
    "            output_path = os.path.join(output_dir, filename)\n",
    "            \n",
    "            print(f\"Processing {filename}...\")\n",
    "            try:\n",
    "                make_background_black_fast(input_path, output_path)\n",
    "            except Exception as e:\n",
    "                print(f\"Error processing {filename}: {str(e)}\")\n",
    "\n",
    "# Usage\n",
    "input_dir = \"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT_REG\"\n",
    "output_dir = \"/mnt/d/CTH_archive/CBF_COLORED_NIFIT_NOTEXT_REG_BLACK\"\n",
    "process_directory(input_dir, output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import os\n",
    "import random\n",
    "from collections import Counter\n",
    "from scipy import ndimage\n",
    "\n",
    "def replace_most_common_color(image, tolerance=1):\n",
    "    flat_image = image.reshape(-1, image.shape[-1])\n",
    "    color_counts = Counter(map(tuple, flat_image))\n",
    "    most_common_color = np.array(max(color_counts, key=color_counts.get))\n",
    "    mask = np.all(np.abs(image - most_common_color) <= tolerance, axis=-1)\n",
    "    modified_img = image.copy()\n",
    "    modified_img[mask] = [0, 0, 0]\n",
    "    return modified_img\n",
    "\n",
    "def fill_isolated_black_pixels(image, iterations=3):\n",
    "    filled_img = image.copy()\n",
    "    \n",
    "    for _ in range(iterations):\n",
    "        black_pixels = np.all(filled_img == [0, 0, 0], axis=-1)\n",
    "        non_black_mask = ~black_pixels\n",
    "        \n",
    "        for channel in range(3):\n",
    "            channel_data = filled_img[..., channel]\n",
    "            filled_channel = ndimage.generic_filter(\n",
    "                channel_data, \n",
    "                function=lambda values: np.mean(values[values > 0]) if np.any(values > 0) else 0,\n",
    "                size=3,\n",
    "                mode='constant',\n",
    "                cval=0\n",
    "            )\n",
    "            filled_img[..., channel] = np.where(non_black_mask, channel_data, filled_channel)\n",
    "    \n",
    "    return filled_img\n",
    "\n",
    "def hard_edge_crop(image, crop_width=3):\n",
    "    h, w = image.shape[:2]\n",
    "    cropped = image[crop_width:-crop_width, crop_width:-crop_width]\n",
    "    return cropped\n",
    "\n",
    "def process_and_display_image(image_path):\n",
    "    img = Image.open(image_path)\n",
    "    img_array = np.array(img)\n",
    "\n",
    "    # Step 1: Replace most common color with black\n",
    "    modified_img = replace_most_common_color(img_array, tolerance=1)\n",
    "    \n",
    "    # Step 2: Fill isolated black pixels\n",
    "    filled_img = fill_isolated_black_pixels(modified_img, iterations=3)\n",
    "    \n",
    "    # Step 3: Hard edge crop\n",
    "    final_img = hard_edge_crop(filled_img, crop_width=3)\n",
    "\n",
    "    fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(40, 10))\n",
    "\n",
    "    ax1.imshow(img_array)\n",
    "    ax1.set_title(\"Original Image\")\n",
    "    ax1.axis('off')\n",
    "\n",
    "    ax2.imshow(modified_img)\n",
    "    ax2.set_title(\"Most Common Color Replaced\")\n",
    "    ax2.axis('off')\n",
    "\n",
    "    ax3.imshow(filled_img)\n",
    "    ax3.set_title(\"Isolated Black Pixels Filled\")\n",
    "    ax3.axis('off')\n",
    "\n",
    "    ax4.imshow(final_img)\n",
    "    ax4.set_title(\"Hard Edge Cropped\")\n",
    "    ax4.axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Directory containing the images\n",
    "directory = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A/\"\n",
    "\n",
    "# Get all PNG files in the directory\n",
    "all_files = [f for f in os.listdir(directory) if f.endswith('.png')]\n",
    "\n",
    "# Randomly select 15 files\n",
    "selected_files = random.sample(all_files, 15)\n",
    "\n",
    "# Process and display each selected file\n",
    "for file in selected_files:\n",
    "    image_path = os.path.join(directory, file)\n",
    "    print(f\"Processing: {file}\")\n",
    "    process_and_display_image(image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import os\n",
    "import random\n",
    "from collections import Counter\n",
    "from scipy import ndimage\n",
    "\n",
    "def replace_most_common_color(image, tolerance=1):\n",
    "    flat_image = image.reshape(-1, image.shape[-1])\n",
    "    color_counts = Counter(map(tuple, flat_image))\n",
    "    most_common_color = np.array(max(color_counts, key=color_counts.get))\n",
    "    mask = np.all(np.abs(image - most_common_color) <= tolerance, axis=-1)\n",
    "    modified_img = image.copy()\n",
    "    modified_img[mask] = [0, 0, 0]\n",
    "    return modified_img\n",
    "\n",
    "def fill_non_edge_black_pixels(image):\n",
    "    filled_img = image.copy()\n",
    "    \n",
    "    # Create a mask of non-black pixels\n",
    "    non_black_mask = np.any(image != [0, 0, 0], axis=-1)\n",
    "    \n",
    "    # Use a 3x3 structuring element to find black pixels surrounded by non-black pixels\n",
    "    struct = np.ones((3, 3), dtype=bool)\n",
    "    isolated_black = ndimage.binary_hit_or_miss(~non_black_mask, structure1=struct)\n",
    "    \n",
    "    # For each isolated black pixel, calculate the average of surrounding non-black pixels\n",
    "    for y, x in zip(*np.where(isolated_black)):\n",
    "        neighborhood = image[max(0, y-1):y+2, max(0, x-1):x+2]\n",
    "        non_black_neighbors = neighborhood[np.any(neighborhood != [0, 0, 0], axis=-1)]\n",
    "        if len(non_black_neighbors) > 0:\n",
    "            filled_img[y, x] = np.mean(non_black_neighbors, axis=0).astype(np.uint8)\n",
    "    \n",
    "    return filled_img\n",
    "\n",
    "def process_and_display_image(image_path):\n",
    "    img = Image.open(image_path)\n",
    "    img_array = np.array(img)\n",
    "\n",
    "    # Step 1: Replace most common color with black\n",
    "    modified_img = replace_most_common_color(img_array, tolerance=1)\n",
    "    \n",
    "    # Step 2: Fill non-edge black pixels\n",
    "    final_img = fill_non_edge_black_pixels(modified_img)\n",
    "\n",
    "    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(30, 10))\n",
    "\n",
    "    ax1.imshow(img_array)\n",
    "    ax1.set_title(\"Original Image\")\n",
    "    ax1.axis('off')\n",
    "\n",
    "    ax2.imshow(modified_img)\n",
    "    ax2.set_title(\"Most Common Color Replaced\")\n",
    "    ax2.axis('off')\n",
    "\n",
    "    ax3.imshow(final_img)\n",
    "    ax3.set_title(\"Non-Edge Black Pixels Filled\")\n",
    "    ax3.axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Directory containing the images\n",
    "directory = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A/\"\n",
    "\n",
    "# Get all PNG files in the directory\n",
    "all_files = [f for f in os.listdir(directory) if f.endswith('.png')]\n",
    "\n",
    "# Randomly select 15 files\n",
    "selected_files = random.sample(all_files, 15)\n",
    "\n",
    "# Process and display each selected file\n",
    "for file in selected_files:\n",
    "    image_path = os.path.join(directory, file)\n",
    "    print(f\"Processing: {file}\")\n",
    "    process_and_display_image(image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import os\n",
    "import random\n",
    "from collections import Counter\n",
    "\n",
    "def replace_most_common_color(image, tolerance=1):\n",
    "    flat_image = image.reshape(-1, image.shape[-1])\n",
    "    color_counts = Counter(map(tuple, flat_image))\n",
    "    most_common_color = np.array(max(color_counts, key=color_counts.get))\n",
    "    mask = np.all(np.abs(image - most_common_color) <= tolerance, axis=-1)\n",
    "    modified_img = image.copy()\n",
    "    modified_img[mask] = [0, 0, 0]\n",
    "    return modified_img\n",
    "\n",
    "def fill_isolated_black_pixels(image, neighborhood_size=1):\n",
    "    filled_img = image.copy()\n",
    "    height, width = image.shape[:2]\n",
    "    \n",
    "    for y in range(height):\n",
    "        for x in range(width):\n",
    "            if np.all(image[y, x] == [0, 0, 0]):\n",
    "                # Check the neighborhood\n",
    "                y_start = max(0, y - neighborhood_size)\n",
    "                y_end = min(height, y + neighborhood_size + 1)\n",
    "                x_start = max(0, x - neighborhood_size)\n",
    "                x_end = min(width, x + neighborhood_size + 1)\n",
    "                \n",
    "                neighborhood = image[y_start:y_end, x_start:x_end]\n",
    "                non_black_pixels = neighborhood[np.any(neighborhood != [0, 0, 0], axis=-1)]\n",
    "                \n",
    "                # If the neighborhood contains non-black pixels, fill the current pixel\n",
    "                if len(non_black_pixels) > 0:\n",
    "                    filled_img[y, x] = np.mean(non_black_pixels, axis=0).astype(np.uint8)\n",
    "    \n",
    "    return filled_img\n",
    "\n",
    "def process_and_display_image(image_path):\n",
    "    img = Image.open(image_path)\n",
    "    img_array = np.array(img)\n",
    "\n",
    "    # Step 1: Replace most common color with black\n",
    "    modified_img = replace_most_common_color(img_array, tolerance=1)\n",
    "    \n",
    "    # Step 2: Fill isolated black pixels\n",
    "    final_img = fill_isolated_black_pixels(modified_img, neighborhood_size=1)\n",
    "\n",
    "    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(30, 10))\n",
    "\n",
    "    ax1.imshow(img_array)\n",
    "    ax1.set_title(\"Original Image\")\n",
    "    ax1.axis('off')\n",
    "\n",
    "    ax2.imshow(modified_img)\n",
    "    ax2.set_title(\"Most Common Color Replaced\")\n",
    "    ax2.axis('off')\n",
    "\n",
    "    ax3.imshow(final_img)\n",
    "    ax3.set_title(\"Isolated Black Pixels Filled\")\n",
    "    ax3.axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Directory containing the images\n",
    "directory = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A/\"\n",
    "\n",
    "# Get all PNG files in the directory\n",
    "all_files = [f for f in os.listdir(directory) if f.endswith('.png')]\n",
    "\n",
    "# Randomly select 15 files\n",
    "selected_files = random.sample(all_files, 15)\n",
    "\n",
    "# Process and display each selected file\n",
    "for file in selected_files:\n",
    "    image_path = os.path.join(directory, file)\n",
    "    print(f\"Processing: {file}\")\n",
    "    process_and_display_image(image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import os\n",
    "import random\n",
    "from collections import Counter\n",
    "from skimage import exposure\n",
    "\n",
    "def replace_most_common_color(image, tolerance=1):\n",
    "    flat_image = image.reshape(-1, image.shape[-1])\n",
    "    color_counts = Counter(map(tuple, flat_image))\n",
    "    most_common_color = np.array(max(color_counts, key=color_counts.get))\n",
    "    mask = np.all(np.abs(image - most_common_color) <= tolerance, axis=-1)\n",
    "    modified_img = image.copy()\n",
    "    modified_img[mask] = [0, 0, 0]\n",
    "    return modified_img\n",
    "\n",
    "def fill_isolated_black_pixels(image, neighborhood_size=1):\n",
    "    filled_img = image.copy()\n",
    "    height, width = image.shape[:2]\n",
    "    \n",
    "    for y in range(height):\n",
    "        for x in range(width):\n",
    "            if np.all(image[y, x] == [0, 0, 0]):\n",
    "                y_start = max(0, y - neighborhood_size)\n",
    "                y_end = min(height, y + neighborhood_size + 1)\n",
    "                x_start = max(0, x - neighborhood_size)\n",
    "                x_end = min(width, x + neighborhood_size + 1)\n",
    "                \n",
    "                neighborhood = image[y_start:y_end, x_start:x_end]\n",
    "                non_black_pixels = neighborhood[np.any(neighborhood != [0, 0, 0], axis=-1)]\n",
    "                \n",
    "                if len(non_black_pixels) > 0:\n",
    "                    filled_img[y, x] = np.mean(non_black_pixels, axis=0).astype(np.uint8)\n",
    "    \n",
    "    return filled_img\n",
    "\n",
    "def enhance_contrast(image):\n",
    "    # Create a mask for non-black pixels\n",
    "    non_black_mask = np.any(image != [0, 0, 0], axis=-1)\n",
    "    \n",
    "    # Perform contrast limited adaptive histogram equalization (CLAHE)\n",
    "    enhanced = np.zeros_like(image)\n",
    "    for i in range(3):  # For each color channel\n",
    "        channel = image[..., i]\n",
    "        # Apply CLAHE only to non-black areas\n",
    "        channel_eq = exposure.equalize_adapthist(channel, clip_limit=0.03)\n",
    "        enhanced[..., i] = np.where(non_black_mask, (channel_eq * 255).astype(np.uint8), 0)\n",
    "    \n",
    "    return enhanced\n",
    "\n",
    "def process_and_display_image(image_path):\n",
    "    img = Image.open(image_path)\n",
    "    img_array = np.array(img)\n",
    "\n",
    "    # Step 1: Replace most common color with black\n",
    "    modified_img = replace_most_common_color(img_array, tolerance=1)\n",
    "    \n",
    "    # Step 2: Fill isolated black pixels\n",
    "    filled_img = fill_isolated_black_pixels(modified_img, neighborhood_size=1)\n",
    "    \n",
    "    # Step 3: Enhance contrast\n",
    "    enhanced_img = enhance_contrast(filled_img)\n",
    "\n",
    "    fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(40, 10))\n",
    "\n",
    "    ax1.imshow(img_array)\n",
    "    ax1.set_title(\"Original Image\")\n",
    "    ax1.axis('off')\n",
    "\n",
    "    ax2.imshow(modified_img)\n",
    "    ax2.set_title(\"Most Common Color Replaced\")\n",
    "    ax2.axis('off')\n",
    "\n",
    "    ax3.imshow(filled_img)\n",
    "    ax3.set_title(\"Isolated Black Pixels Filled\")\n",
    "    ax3.axis('off')\n",
    "\n",
    "    ax4.imshow(enhanced_img)\n",
    "    ax4.set_title(\"Contrast Enhanced\")\n",
    "    ax4.axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# Directory containing the images\n",
    "directory = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A/\"\n",
    "\n",
    "# Get all PNG files in the directory\n",
    "all_files = [f for f in os.listdir(directory) if f.endswith('.png')]\n",
    "\n",
    "# Randomly select 15 files\n",
    "selected_files = random.sample(all_files, 15)\n",
    "\n",
    "# Process and display each selected file\n",
    "for file in selected_files:\n",
    "    image_path = os.path.join(directory, file)\n",
    "    print(f\"Processing: {file}\")\n",
    "    process_and_display_image(image_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from PIL import Image\n",
    "import os\n",
    "from collections import Counter\n",
    "from skimage import exposure\n",
    "from tqdm import tqdm\n",
    "\n",
    "def replace_most_common_color(image, tolerance=1):\n",
    "    flat_image = image.reshape(-1, image.shape[-1])\n",
    "    color_counts = Counter(map(tuple, flat_image))\n",
    "    most_common_color = np.array(max(color_counts, key=color_counts.get))\n",
    "    mask = np.all(np.abs(image - most_common_color) <= tolerance, axis=-1)\n",
    "    modified_img = image.copy()\n",
    "    modified_img[mask] = [0, 0, 0]\n",
    "    return modified_img\n",
    "\n",
    "def fill_isolated_black_pixels(image, neighborhood_size=1):\n",
    "    filled_img = image.copy()\n",
    "    height, width = image.shape[:2]\n",
    "    \n",
    "    for y in range(height):\n",
    "        for x in range(width):\n",
    "            if np.all(image[y, x] == [0, 0, 0]):\n",
    "                y_start = max(0, y - neighborhood_size)\n",
    "                y_end = min(height, y + neighborhood_size + 1)\n",
    "                x_start = max(0, x - neighborhood_size)\n",
    "                x_end = min(width, x + neighborhood_size + 1)\n",
    "                \n",
    "                neighborhood = image[y_start:y_end, x_start:x_end]\n",
    "                non_black_pixels = neighborhood[np.any(neighborhood != [0, 0, 0], axis=-1)]\n",
    "                \n",
    "                if len(non_black_pixels) > 0:\n",
    "                    filled_img[y, x] = np.mean(non_black_pixels, axis=0).astype(np.uint8)\n",
    "    \n",
    "    return filled_img\n",
    "\n",
    "def enhance_contrast(image):\n",
    "    non_black_mask = np.any(image != [0, 0, 0], axis=-1)\n",
    "    \n",
    "    enhanced = np.zeros_like(image)\n",
    "    for i in range(3):\n",
    "        channel = image[..., i]\n",
    "        channel_eq = exposure.equalize_adapthist(channel, clip_limit=0.03)\n",
    "        enhanced[..., i] = np.where(non_black_mask, (channel_eq * 255).astype(np.uint8), 0)\n",
    "    \n",
    "    return enhanced\n",
    "\n",
    "def process_image(image_path):\n",
    "    img = Image.open(image_path)\n",
    "    img_array = np.array(img)\n",
    "\n",
    "    # Step 1: Replace most common color with black\n",
    "    modified_img = replace_most_common_color(img_array, tolerance=1)\n",
    "    \n",
    "    # Step 2: Fill isolated black pixels\n",
    "    filled_img = fill_isolated_black_pixels(modified_img, neighborhood_size=1)\n",
    "    \n",
    "    # Step 3: Enhance contrast\n",
    "    enhanced_img = enhance_contrast(filled_img)\n",
    "    \n",
    "    # Step 4: Resize to 512x512\n",
    "    resized_img = Image.fromarray(enhanced_img).resize((512, 512), Image.LANCZOS)\n",
    "    \n",
    "    return resized_img\n",
    "\n",
    "def process_directory(directory):\n",
    "    for filename in tqdm(os.listdir(directory), desc=f\"Processing {os.path.basename(directory)}\"):\n",
    "        if filename.endswith('.png'):\n",
    "            image_path = os.path.join(directory, filename)\n",
    "            processed_img = process_image(image_path)\n",
    "            processed_img.save(image_path)  # Overwrite the original file\n",
    "\n",
    "# Directories to process\n",
    "train_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A/\"\n",
    "test_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/test_A/\"\n",
    "\n",
    "# Process both directories\n",
    "process_directory(train_dir)\n",
    "process_directory(test_dir)\n",
    "\n",
    "print(\"All images have been processed, resized, and saved.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "def is_mostly_black(image_path, threshold=0.995):\n",
    "    with Image.open(image_path) as img:\n",
    "        img_array = np.array(img)\n",
    "        if len(img_array.shape) == 3:  # RGB image\n",
    "            return np.mean(img_array == 0) >= threshold\n",
    "        else:  # Grayscale image\n",
    "            return np.mean(img_array == 0) >= threshold\n",
    "\n",
    "def remove_files_and_update_json(dir_a, dir_b, json_file):\n",
    "    with open(json_file, 'r') as f:\n",
    "        prompts = json.load(f)\n",
    "\n",
    "    files_to_remove = []\n",
    "\n",
    "    for filename in os.listdir(dir_a):\n",
    "        if filename.lower().endswith('.png'):\n",
    "            file_path_a = os.path.join(dir_a, filename)\n",
    "            file_path_b = os.path.join(dir_b, filename)\n",
    "\n",
    "            if is_mostly_black(file_path_a):\n",
    "                print(f\"Removing {filename}\")\n",
    "                os.remove(file_path_a)\n",
    "                if os.path.exists(file_path_b):\n",
    "                    os.remove(file_path_b)\n",
    "                files_to_remove.append(filename)\n",
    "\n",
    "    # Update JSON\n",
    "    for file in files_to_remove:\n",
    "        prompts.pop(file, None)\n",
    "\n",
    "    with open(json_file, 'w') as f:\n",
    "        json.dump(prompts, f, indent=2)\n",
    "\n",
    "    print(f\"Removed {len(files_to_remove)} files and updated {json_file}\")\n",
    "\n",
    "# Paths\n",
    "test_a = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A/\"\n",
    "test_b = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/test_B\"\n",
    "train_a = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_A\"\n",
    "train_b = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_B\"\n",
    "test_json = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/test_prompts.json\"\n",
    "train_json = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK/train_prompts.json\"\n",
    "\n",
    "# Process test and train sets\n",
    "remove_files_and_update_json(test_a, test_b, test_json)\n",
    "remove_files_and_update_json(train_a, train_b, train_json)\n",
    "\n",
    "print(\"Processing complete.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import shutil\n",
    "from PIL import Image, ImageEnhance\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "def apply_augmentation(img_a, img_b):\n",
    "    # Ensure images are in PIL format\n",
    "    img_a = Image.fromarray(np.uint8(img_a)) if isinstance(img_a, np.ndarray) else img_a\n",
    "    img_b = Image.fromarray(np.uint8(img_b)) if isinstance(img_b, np.ndarray) else img_b\n",
    "\n",
    "    # Horizontal flip\n",
    "    if random.random() > 0.5:\n",
    "        img_a = img_a.transpose(Image.FLIP_LEFT_RIGHT)\n",
    "        img_b = img_b.transpose(Image.FLIP_LEFT_RIGHT)\n",
    "\n",
    "    # Vertical flip\n",
    "    if random.random() > 0.5:\n",
    "        img_a = img_a.transpose(Image.FLIP_TOP_BOTTOM)\n",
    "        img_b = img_b.transpose(Image.FLIP_TOP_BOTTOM)\n",
    "\n",
    "    # Rotation\n",
    "    angle = random.uniform(-15, 15)\n",
    "    img_a = img_a.rotate(angle, resample=Image.BICUBIC, expand=True)\n",
    "    img_b = img_b.rotate(angle, resample=Image.BICUBIC, expand=True)\n",
    "\n",
    "    # Scaling\n",
    "    scale = random.uniform(0.8, 1.2)\n",
    "    new_size = (int(img_a.width * scale), int(img_a.height * scale))\n",
    "    img_a = img_a.resize(new_size, Image.BICUBIC)\n",
    "    img_b = img_b.resize(new_size, Image.BICUBIC)\n",
    "\n",
    "    # Crop or pad to original size\n",
    "    original_size = (512, 512)  # Assuming original size is 512x512\n",
    "    img_a = crop_or_pad(img_a, original_size)\n",
    "    img_b = crop_or_pad(img_b, original_size)\n",
    "\n",
    "    return img_a, img_b\n",
    "\n",
    "def crop_or_pad(img, target_size):\n",
    "    width, height = img.size\n",
    "    target_width, target_height = target_size\n",
    "\n",
    "    if width > target_width:\n",
    "        left = (width - target_width) // 2\n",
    "        img = img.crop((left, 0, left + target_width, height))\n",
    "    elif width < target_width:\n",
    "        new_img = Image.new('RGB', (target_width, height), (0, 0, 0))\n",
    "        left = (target_width - width) // 2\n",
    "        new_img.paste(img, (left, 0))\n",
    "        img = new_img\n",
    "\n",
    "    if height > target_height:\n",
    "        top = (height - target_height) // 2\n",
    "        img = img.crop((0, top, target_width, top + target_height))\n",
    "    elif height < target_height:\n",
    "        new_img = Image.new('RGB', (target_width, target_height), (0, 0, 0))\n",
    "        top = (target_height - height) // 2\n",
    "        new_img.paste(img, (0, top))\n",
    "        img = new_img\n",
    "\n",
    "    return img\n",
    "\n",
    "def process_directory_pair(dir_a, dir_b, output_dir_a, output_dir_b):\n",
    "    os.makedirs(output_dir_a, exist_ok=True)\n",
    "    os.makedirs(output_dir_b, exist_ok=True)\n",
    "    \n",
    "    files = [f for f in os.listdir(dir_a) if f.endswith('.png')]\n",
    "    \n",
    "    for file in tqdm(files, desc=f\"Processing {os.path.basename(dir_a)} and {os.path.basename(dir_b)}\"):\n",
    "        a_path = os.path.join(dir_a, file)\n",
    "        b_path = os.path.join(dir_b, file)\n",
    "        \n",
    "        if not os.path.exists(b_path):\n",
    "            print(f\"Warning: No matching B image for {file}\")\n",
    "            continue\n",
    "        \n",
    "        # Copy original images\n",
    "        shutil.copy2(a_path, os.path.join(output_dir_a, file))\n",
    "        shutil.copy2(b_path, os.path.join(output_dir_b, file))\n",
    "        \n",
    "        # Create two augmented versions\n",
    "        for i in range(2):\n",
    "            img_a = Image.open(a_path)\n",
    "            img_b = Image.open(b_path)\n",
    "            \n",
    "            aug_a, aug_b = apply_augmentation(img_a, img_b)\n",
    "            \n",
    "            aug_name = f\"aug_{i+1}_{file}\"\n",
    "            aug_a.save(os.path.join(output_dir_a, aug_name))\n",
    "            aug_b.save(os.path.join(output_dir_b, aug_name))\n",
    "\n",
    "# Define paths\n",
    "base_input_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_BLACK\"\n",
    "base_output_dir = \"/home/nicholasjprimiano/ML/img2img-turbo/data/CTH_CBF_AUG\"\n",
    "\n",
    "for subset in ['train']:\n",
    "    input_a = os.path.join(base_input_dir, f\"{subset}_A\")\n",
    "    input_b = os.path.join(base_input_dir, f\"{subset}_B\")\n",
    "    output_a = os.path.join(base_output_dir, f\"{subset}_A\")\n",
    "    output_b = os.path.join(base_output_dir, f\"{subset}_B\")\n",
    "    \n",
    "    process_directory_pair(input_a, input_b, output_a, output_b)\n",
    "\n",
    "print(\"Augmentation complete!\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}