Update app.py

app.py

@@ -269,8 +269,11 @@ 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)
 
+            # Load the binary lesion mask to get its affine
+            output_img = nib.load(new_output_file)
+
             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)
+            nib.save(nib.Nifti1Image(labeled_mask.astype(np.int16), output_img.affine), labeled_mask_path)
 
             # Return paths for the Gradio interface
             return new_output_file, input_slice_path, output_slice_path, labeled_mask_path
@@ -302,9 +305,10 @@ with gr.Blocks() as demo:
             submit_button = gr.Button("Submit")
         with gr.Column(scale=2):
             seg_output = gr.File(label="Download the Lesion Segmentation Mask")
+            clusters_output = gr.File(label="Download the Labeled Lesion Segmentation Mask")
             input_img = gr.Image(label="Input: FLAIR image")
             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:**