Update app.py

app.py

@@ -38,6 +38,17 @@ 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)))
 
+def label_to_rgb(mask):
+    rgb_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
+    rgb_mask[mask == 0] = class_building
+    rgb_mask[mask == 1] = class_land
+    rgb_mask[mask == 2] = class_road
+    rgb_mask[mask == 3] = class_vegetation
+    rgb_mask[mask == 4] = class_water
+    rgb_mask[mask == 5] = class_unlabeled
+    return rgb_mask
+
+
 
 def jaccard_coef(y_true, y_pred):
   y_true_flatten = K.flatten(y_true)
@@ -54,12 +65,12 @@ 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):
-  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 = predicted_image * 50
-  return "Predicted Masked Image", predicted_image
+    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
 
 my_app = gr.Blocks()