Update app.py

app.py

@@ -7,6 +7,8 @@ from matplotlib import pyplot as plt
 import random
 from keras.utils import get_custom_objects
 import os
+import zipfile
+from io import BytesIO
 
 os.environ['SM_FRAMEWORK'] = 'tf.keras'
 
@@ -49,7 +51,6 @@ def label_to_rgb(mask):
     return rgb_mask
 
 
-
 def jaccard_coef(y_true, y_pred):
   y_true_flatten = K.flatten(y_true)
   y_pred_flatten = K.flatten(y_pred)
@@ -64,16 +65,42 @@ total_loss = dice_loss + (1 * focal_loss)
 
 satellite_model = load_model('model/satellite_segmentation_full.h5', custom_objects=({'dice_loss_plus_1focal_loss' : total_loss, 'jaccard_coef': jaccard_coef}))
 
-def process_input_image(image_source):
+def process_input_image(image_source, download_masks=False):
     image = np.expand_dims(image_source, 0)
     prediction = satellite_model.predict(image)
     predicted_image = np.argmax(prediction, axis=3)
     predicted_image = predicted_image[0, :, :]
     rgb_image = label_to_rgb(predicted_image)
-    return "Predicted Masked Image", rgb_image
+    
+    if download_masks:
+        mask_zip, output_name = create_mask_zip(predicted_image)
+        return "Predicted Masked Image", rgb_image, mask_zip, output_name
+    else:
+        return "Predicted Masked Image", rgb_image
+
 
 my_app = gr.Blocks()
 
+
+def create_mask_zip(mask, output_name='masks.zip'):
+    buffer = BytesIO()
+
+    with zipfile.ZipFile(buffer, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
+        for idx, class_name in enumerate(['building', 'land', 'road', 'vegetation', 'water', 'unlabeled']):
+            mask_class = (mask == idx) * 255
+            mask_class = mask_class.astype(np.uint8)
+            mask_image = Image.fromarray(mask_class)
+            
+            mask_buffer = BytesIO()
+            mask_image.save(mask_buffer, format='PNG')
+            mask_buffer.seek(0)
+            
+            zf.writestr(f'{class_name}_mask.png', mask_buffer.getvalue())
+
+    buffer.seek(0)
+    return buffer, output_name
+
+
 # Define the custom legend HTML
 legend_html = '''
 <div style="font-size: 14px; font-weight: bold; display: flex; flex-wrap: wrap; margin-top: 10px;">
@@ -99,25 +126,30 @@ legend_html = '''
 '''
 
 with my_app:
-  gr.Markdown("Image Processing Application UI with Gradio")
-  with gr.Tabs():
-    with gr.TabItem("Select your image"):
-        with gr.Row():
-          with gr.Column():
-              img_source = gr.Image(label="Please select source Image", shape=(256,256))
-              source_image_loader = gr.Button("Load above Image")
-          with gr.Column():
-              output_label = gr.Label(label="Image Info")
-              img_output = gr.Image(label="Image Output")
-              legend = gr.HTML(legend_html) # Add the custom legend component here
-        source_image_loader.click(
-            process_input_image,
-            [
-                img_source
-            ],
-            [
-                output_label,
-                img_output
-            ]
-        )
+    gr.Markdown("Image Processing Application UI with Gradio")
+    with gr.Tabs():
+        with gr.TabItem("Select your image"):
+            with gr.Row():
+                with gr.Column():
+                    img_source = gr.Image(label="Please select source Image", shape=(256, 256))
+                    source_image_loader = gr.Button("Load above Image")
+                    download_masks_checkbox = gr.Checkbox(label="Download masks as separate files")
+                with gr.Column():
+                    output_label = gr.Label(label="Image Info")
+                    img_output = gr.Image(label="Image Output")
+                    masks_zip = gr.File(label="Download Masks", filename="masks.zip")
+                    legend = gr.HTML(legend_html)
+
+            source_image_loader.click(
+                process_input_image,
+                [
+                    img_source,
+                    download_masks_checkbox
+                ],
+                [
+                    output_label,
+                    img_output,
+                    masks_zip,
+                ]
+            )
 my_app.launch(debug=True)