Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import numpy as np
 import tifffile as tiff
 from aicsimageio import AICSImage  # To handle .czi files
+import matplotlib.pyplot as plt
 
 # Placeholder for your 3D model
 def process_3d_image(image):
@@ -36,18 +37,53 @@ def process_file(file):
     output_path = "output_mask.tif"
     tiff.imwrite(output_path, binary_mask)
 
-    return output_path
+    return image, binary_mask, output_path
 
-# Gradio Interface
-def interface(file):
-    return process_file(file)
+# Function to generate the slice visualization
+def visualize_slice(image, mask, slice_index):
+    """
+    Visualizes a 2D slice of the image and the corresponding mask at the given index.
+    """
+    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
+
+    # Extract the 2D slices
+    image_slice = image[slice_index, :, :]
+    mask_slice = mask[slice_index, :, :]
+
+    # Plot image slice
+    axes[0].imshow(image_slice, cmap="gray")
+    axes[0].set_title("Image Slice")
+    axes[0].axis("off")
 
+    # Plot mask slice
+    axes[1].imshow(mask_slice, cmap="gray")
+    axes[1].set_title("Mask Slice")
+    axes[1].axis("off")
+
+    # Return the plot as a Gradio-compatible output
+    plt.tight_layout()
+    plt.close(fig)
+    return fig
+
+# Gradio Interface function
+def interface(file, slice_index):
+    image, mask, output_path = process_file(file)
+    fig = visualize_slice(image, mask, slice_index)
+    return fig, output_path
+
+# Gradio Interface
 iface = gr.Interface(
     fn=interface,
-    inputs=gr.File(label="Upload 3D Image (.tif or .czi)"),
-    outputs=gr.File(label="Download Binary Mask (.tif)"),
+    inputs=[
+        gr.File(label="Upload 3D Image (.tif or .czi)"),
+        gr.Slider(minimum=0, maximum=100, step=1, label="Slice Index")  # Adjust max value based on image size
+    ],
+    outputs=[
+        gr.Plot(label="2D Slice Visualization"),
+        gr.File(label="Download Binary Mask (.tif)")
+    ],
     title="3D Image Processing with Binary Mask Output",
-    description="Upload a 3D image in .tif or .czi format. The model will process the image and output a 3D binary mask."
+    description="Upload a 3D image in .tif or .czi format. The model will process the image and output a 3D binary mask. Use the slider to navigate through the 2D slices."
 )
 
 if __name__ == "__main__":