import numpy as np
import gradio as gr
from huggingface_hub import from_pretrained_keras

def read_nifti_file(filepath):
    """Read and load volume"""
    # Read file
    scan = nib.load(filepath)
    # Get raw data
    scan = scan.get_fdata()
    return scan

def normalize(volume):
    """Normalize the volume"""
    min = -1000
    max = 400
    volume[volume < min] = min
    volume[volume > max] = max
    volume = (volume - min) / (max - min)
    volume = volume.astype("float32")
    return volume

def resize_volume(img):
    """Resize across z-axis"""
    # Set the desired depth
    desired_depth = 64
    desired_width = 128
    desired_height = 128
    # Get current depth
    current_depth = img.shape[-1]
    current_width = img.shape[0]
    current_height = img.shape[1]
    # Compute depth factor
    depth = current_depth / desired_depth
    width = current_width / desired_width
    height = current_height / desired_height
    depth_factor = 1 / depth
    width_factor = 1 / width
    height_factor = 1 / height
    # Rotate
    img = ndimage.rotate(img, 90, reshape=False)
    # Resize across z-axis
    img = ndimage.zoom(img, (width_factor, height_factor, depth_factor), order=1)
    return img

def process_scan(path):
    """Read and resize volume"""
    # Read scan
    volume = read_nifti_file(path)
    # Normalize
    volume = normalize(volume)
    # Resize width, height and depth
    volume = resize_volume(volume)
    return volume

def infer(filename):
    vol = process_scan(filename.name)
    vol = np.expand_dims(vol, axis=0)
    prediction = model.predict(vol)[0]

    scores = [1 - prediction[0], prediction[0]]

    class_names = ["normal", "abnormal"]
    result = []
    for score, name in zip(scores, class_names):
        result = result + [f"This model is {(100 * score):.2f} percent confident that CT scan is {name}"]

    return result

model = from_pretrained_keras('jalFaizy/3D_CNN')

filepath = gr.inputs.File()
text = gr.outputs.Textbox()

iface = gr.Interface(
    infer, 
    filepath,
    text,
    title='3D CNN for CT scans',
    examples=['example_1_normal.nii.gz']
)

iface.launch()