%load_ext autoreload
%autoreload 2

# Install brainslicer from GitHub
%%capture
!pip install git+https://github.com/sohaamir/brainslicer.git
!pip install nibabel numpy Pillow matplotlib

# pass the help argument for instructions
!brainslicer --help

# Mount your Google Drive and change to your brainslicer directory (you will have to give permissions)
from google.colab import drive, files

drive.mount('/gdrive')
%cd /gdrive/MyDrive/projects/brainslicer

# drive.mount('/content/drive')
#%cd /content/drive/MyDrive/projects/brainslicer/

# It may be useful to print out the dimensions of the NIFTI first (sagittal, coronal, axial)
import nibabel as nib

# Load the NIFTI file
nifti_file_path = 'niftis/aamir_T1.nii'
nifti_img = nib.load(nifti_file_path)

# Get dimensions of the NIFTI image
nifti_dimensions = nifti_img.header.get_data_shape()
print(nifti_dimensions)

# Run brainslicer using the CLI
!brainslicer niftis/aamir_T1.nii 100 --plane sagittal --brightness 2.0 --contrast 1.5 --colourmap gray

from google.colab import files

uploaded = files.upload()

import os
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

output_dir = '/gdrive/MyDrive/projects/brainslicer/brainslices'  # Your output directory

# Find the latest file in the output directory
latest_file = None
latest_timestamp = 0

for filename in os.listdir(output_dir):
    if filename.endswith('.png'):
        filepath = os.path.join(output_dir, filename)
        file_timestamp = os.path.getmtime(filepath)  # Get file modification timestamp
        if file_timestamp > latest_timestamp:
            latest_timestamp = file_timestamp
            latest_file = filepath

# Check if an image was found
if latest_file:
    img = mpimg.imread(latest_file)

    # Display the image
    plt.figure(figsize=(8,6))
    plt.imshow(img, cmap='gray')
    plt.axis('off')
    plt.title('Latest Brainslicer Output')
    plt.show()
else:
    print("No image found in the 'brainslices' directory.")

# Clone Real-ESRGAN and install requirements
%%capture
!pip install git+https://github.com/sberbank-ai/Real-ESRGAN.git
!git clone https://github.com/boomb0om/Real-ESRGAN-colab
%cd Real-ESRGAN-colab
!pip install -r requirements.txt

!nvidia-smi

import torch
from PIL import Image
import numpy as np
from RealESRGAN import RealESRGAN

# Set up the model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = RealESRGAN(device, scale=4)
model.load_weights('/content/drive/MyDrive/projects/brainslicer/weights/RealESRGAN_x4.pth')

def upscale_image(image_path, output_dir):
    """Upscales a single image and saves the result."""
    image = Image.open(image_path).convert('RGB')
    sr_image = model.predict(image)

    # Create output directory if needed
    os.makedirs(output_dir, exist_ok=True)

    # Save images
    base_filename = os.path.splitext(os.path.basename(image_path))[0]
    image.save(os.path.join(output_dir, f'{base_filename}_original.png'))
    sr_image.save(os.path.join(output_dir, f'{base_filename}_upscaled.png'))

    # Display images side-by-side
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
    ax1.imshow(image)
    ax1.set_title('Original Image')
    ax1.axis('off')
    ax2.imshow(sr_image)
    ax2.set_title('Upscaled Image')
    ax2.axis('off')
    plt.show()

# Find the most recently modified file with a .png extension
most_recent_file = None
most_recent_timestamp = 0

pictures_dir = '/gdrive/MyDrive/projects/brainslicer/brainslices'

for filename in os.listdir(pictures_dir):
    if filename.endswith('.png'):
        filepath = os.path.join(pictures_dir, filename)
        file_timestamp = os.path.getmtime(filepath)
        if file_timestamp > most_recent_timestamp:
            most_recent_timestamp = file_timestamp
            most_recent_file = filepath

if most_recent_file:
    upscale_image(most_recent_file, output_dir)
else:
    print("No PNG files found in the directory.")

uploaded = files.upload()
for filename in uploaded.keys():
  output_dir = 'brainslicer_upscaled'
  upscale_image(filename, output_dir)