Update app.py as per #1

app.py

@@ -2,11 +2,70 @@ import numpy as np
 import gradio as gr
 from huggingface_hub import from_pretrained_keras
 
-def infer(filepath):
-    image_tensor = np.load(filepath)
-    prediction = model.predict(np.expand_dims(image_tensor, axis=0))[0]
-    confidence = 100 * (1 - prediction[0])
-    return confidence
+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')
 
@@ -18,7 +77,7 @@ iface = gr.Interface(
     filepath,
     text,
     title='3D CNN for CT scans',
-    examples=['example_1_normal.npy']
+    examples=['example_1_normal.nii.gz']
 )
 
 iface.launch()