Update app.py

app.py

@@ -1,74 +1,111 @@
+class_building = '#3C1098'
+class_building = class_building.lstrip('#')
+class_building = np.array(tuple(int(class_building[i:i+2], 16) for i in (0,2,4)))
+
+class_land = '#8429F6'
+class_land = class_land.lstrip('#')
+class_land = np.array(tuple(int(class_land[i:i+2], 16) for i in (0,2,4)))
+
+class_road = '#6EC1E4'
+class_road = class_road.lstrip('#')
+class_road = np.array(tuple(int(class_road[i:i+2], 16) for i in (0,2,4)))
+
+class_vegetation = '#FEDD3A'
+class_vegetation = class_vegetation.lstrip('#')
+class_vegetation = np.array(tuple(int(class_vegetation[i:i+2], 16) for i in (0,2,4)))
+
+class_water = '#E2A929'
+class_water = class_water.lstrip('#')
+class_water = np.array(tuple(int(class_water[i:i+2], 16) for i in (0,2,4)))
+
+class_unlabeled = '#9B9B9B'
+class_unlabeled = class_unlabeled.lstrip('#')
+class_unlabeled = np.array(tuple(int(class_unlabeled[i:i+2], 16) for i in (0,2,4)))
+
+
 import gradio as gr
 import os
 import cv2
-from PIL import Image 
-import numpy as np 
+from PIL import Image
+import numpy as np
 from matplotlib import pyplot as plt
 import random
 from keras.utils import get_custom_objects
 import os
 
-os.environ['SM_FRAMEWORK'] = 'tf.keras'
-
-
+#os.environ['SM_FRAMEWORK'] = 'tf.keras'
 
-#%env SM_FRAMEWORK=tf.keras
 import segmentation_models as sm
 
 from keras import backend as K
 from keras.models import load_model
 
+
 def jaccard_coef(y_true, y_pred):
-  y_true_flatten = K.flatten(y_true)
-  y_pred_flatten = K.flatten(y_pred)
-  intersection = K.sum(y_true_flatten * y_pred_flatten)
-  final_coef_value = (intersection + 1.0) / (K.sum(y_true_flatten) + K.sum(y_pred_flatten) - intersection + 1.0)
-  return final_coef_value
+    y_true_flatten = K.flatten(y_true)
+    y_pred_flatten = K.flatten(y_pred)
+    intersection = K.sum(y_true_flatten * y_pred_flatten)
+    final_coef_value = (intersection + 1.0) / (
+            K.sum(y_true_flatten) + K.sum(y_pred_flatten) - intersection + 1.0)
+    return final_coef_value
+
 
-    #six class for six weights
+# six class for six weights
 weights = [0.1666, 0.1666, 0.1666, 0.1666, 0.1666, 0.1666]
-dice_loss = sm.losses.DiceLoss(class_weights = weights)
+dice_loss = sm.losses.DiceLoss(class_weights=weights)
 focal_loss = sm.losses.CategoricalFocalLoss()
 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}))
+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):
-  #image = image_source
-  #image = image.resize((256,256))
-  #image = np.array(image)
-  image = np.expand_dims(image_source, 0)
 
-  prediction = satellite_model.predict(image)
+def label_to_rgb(label_segment):
+    rgb_image = np.zeros((label_segment.shape[0], label_segment.shape[1], 3), dtype=np.uint8)
+
+    rgb_image[label_segment == 0] = class_water
+    rgb_image[label_segment == 1] = class_land
+    rgb_image[label_segment == 2] = class_road
+    rgb_image[label_segment == 3] = class_building
+    rgb_image[label_segment == 4] = class_vegetation
+    rgb_image[label_segment == 5] = class_unlabeled
+
+    return rgb_image
+
+
+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, :, :]
 
-  predicted_image = np.argmax(prediction, axis=3)
-  predicted_image = predicted_image[0,:,:]
-  predicted_image = predicted_image * 50
+    # Convert the predicted image labels to RGB
+    colored_predicted_image = label_to_rgb(predicted_image)
 
-  return "Predicted Masked Image", predicted_image
+    return "Predicted Masked Image", colored_predicted_image
 
 
 my_app = gr.Blocks()
 
 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")
-        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")
+                with gr.Column():
+                    output_label = gr.Label(label="Image Info")
+                    img_output = gr.Image(label="Image Output")
+            source_image_loader.click(
+                process_input_image,
+                [
+                    img_source
+                ],
+                [
+                    output_label,
+                    img_output
+                ]
+            )
 my_app.launch(debug=True)