Update app.py

app.py

@@ -55,6 +55,29 @@ def process_scan(path):
     volume = resize_volume(volume)
     return volume
 
+def plot_slices(num_rows, num_columns, width, height, data):
+    """Plot a montage of 20 CT slices"""
+    data = np.rot90(np.array(data))
+    data = np.transpose(data)
+    data = np.reshape(data, (num_rows, num_columns, width, height))
+    rows_data, columns_data = data.shape[0], data.shape[1]
+    heights = [slc[0].shape[0] for slc in data]
+    widths = [slc.shape[1] for slc in data[0]]
+    fig_width = 12.0
+    fig_height = fig_width * sum(heights) / sum(widths)
+    f, axarr = plt.subplots(
+        rows_data,
+        columns_data,
+        figsize=(fig_width, fig_height),
+        gridspec_kw={"height_ratios": heights},
+    )
+    for i in range(rows_data):
+        for j in range(columns_data):
+            axarr[i, j].imshow(data[i][j], cmap="gray")
+            axarr[i, j].axis("off")
+
+    return f
+
 def infer(filename):
     vol = process_scan(filename.name)
     vol = np.expand_dims(vol, axis=0)
@@ -67,17 +90,17 @@ def infer(filename):
     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
+    return result, plot_slices(2, 10, 128, 128, vol[0, :, :, :20])
 
 model = from_pretrained_keras('jalFaizy/3D_CNN')
 
-filepath = gr.inputs.File()
-text = gr.outputs.Textbox()
+inputs = gr.inputs.File()
+outputs = [gr.outputs.Textbox(), 'plot']
 
 iface = gr.Interface(
     infer, 
-    filepath,
-    text,
+    inputs,
+    outputs,
     title='3D CNN for CT scans',
     examples=['example_1_normal.nii.gz']
 )