Update app.py

app.py

@@ -269,6 +269,9 @@ def run_nnunet_predict(nifti_file,hd_bet=False):
             image = extract_middle_slices(input_path, input_slice_path, center=center)
             labeled_mask = extract_middle_slices(new_output_file, output_slice_path, center=center, label_components=True)
 
+            labeled_mask_path = os.path.join(OUTPUT_DIR, f"{base_filename}_LabeledClusters.nii.gz")
+            nib.save(nib.Nifti1Image(labeled_mask.astype(np.int16), img.affine), labeled_mask_path)
+
             # Return paths for the Gradio interface
             return new_output_file, input_slice_path, output_slice_path
         else:
@@ -300,7 +303,8 @@ with gr.Blocks() as demo:
         with gr.Column(scale=2):
             seg_output = gr.File(label="Download the Lesion Segmentation Mask")
             input_img = gr.Image(label="Input: FLAIR image")
-            output_img = gr.Image(label="Output: Lesion Mask")
+            output_img = gr.Image(label="Output: Binary Lesion Mask")
+            clusters_output = gr.Image(label="Output: Labeled Lesion Mask")
 
     gr.Markdown("""
     **If you find this tool useful, please consider citing:**
@@ -319,7 +323,7 @@ with gr.Blocks() as demo:
     submit_button.click(
         fn=run_nnunet_predict,
         inputs=[flair_input, hd_bet],
-        outputs=[seg_output, input_img, output_img]
+        outputs=[seg_output, input_img, output_img, clusters_output]
     )
 
 # Debugging GPU environment