Update app.py

app.py

@@ -1,31 +1,54 @@
 import gradio as gr
 import numpy as np
-from skimage import io
-import czifile
-
-def load_3d_image(file):
-    if file.name.endswith('.tif'):
-        image = io.imread(file.name)
-    elif file.name.endswith('.czi'):
-        with czifile.CziFile(file.name) as czi:
-            image = czi.asarray()
+import tifffile as tiff
+from aicsimageio import AICSImage  # To handle .czi files
+
+# Placeholder for your 3D model
+def process_3d_image(image):
+    # Dummy model implementation: Replace with your actual model logic
+    binary_mask = (image > image.mean()).astype(np.uint8)
+    return binary_mask
+
+# Function to handle file input and processing
+def process_file(file):
+    """
+    Process the uploaded file and return the binary mask.
+    """
+    if file.name.endswith(".tif"):
+        # Load .tif file as a 3D numpy array
+        image = tiff.imread(file.name)
+    elif file.name.endswith(".czi"):
+        # Load .czi file using AICSImage
+        img = AICSImage(file.name)
+        image = img.get_image_data("CZYX")  # Extracting 3D data
+        image = image[0, 0]  # Assuming single channel, single timepoint
     else:
-        raise ValueError("Unsupported file type. Please upload a .tif or .czi file.")
-    return image
-
-def process_3d_image(input_image):
-    # Your model processing logic here
-    # This is a placeholder function
-    output_mask = np.random.randint(0, 2, input_image.shape).astype(bool)
-    return output_mask
-
-demo = gr.Interface(
-    fn=process_3d_image,
-    inputs=gr.File(label="Input 3D Image (.tif or .czi)", file_types=[".tif", ".czi"]),
-    outputs=gr.Image(label="Output 3D Binary Mask", type="numpy"),
-    title="3D Image to Binary Mask",
-    description="Upload a 3D image (.tif or .czi) and get a 3D binary mask as output."
+        raise ValueError("Unsupported file format. Please upload a .tif or .czi file.")
+
+    # Ensure image is 3D
+    if len(image.shape) != 3:
+        raise ValueError("Input image is not 3D.")
+
+    # Process image through the model
+    binary_mask = process_3d_image(image)
+
+    # Save binary mask to a .tif file to return
+    output_path = "output_mask.tif"
+    tiff.imwrite(output_path, binary_mask)
+
+    return output_path
+
+# Gradio Interface
+def interface(file):
+    return process_file(file)
+
+iface = gr.Interface(
+    fn=interface,
+    inputs=gr.File(label="Upload 3D Image (.tif or .czi)"),
+    outputs=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."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()