Update app.py

app.py

@@ -7,8 +7,6 @@ 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'
 
@@ -51,6 +49,7 @@ 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)
@@ -65,41 +64,16 @@ 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, download_masks=False):
+def process_input_image(image_source):
     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)
-    
-    mask_zip, output_name = create_mask_zip(predicted_image)
-    if download_masks:
-        return {"img_output": rgb_image, "masks_zip": mask_zip.getvalue()}
-    else:
-        return {"img_output": rgb_image}
-
+    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;">
@@ -125,29 +99,25 @@ 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")
-                    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
-                ],
-                [
-                    img_output,
-                    masks_zip,
-                ]
-            )
-my_app.launch(debug=True)
+  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
+            ]
+        )
+my_app.launch(debug=True)