Update app.py

app.py

@@ -63,6 +63,29 @@ def cce_dice_loss(y_true, y_pred):
     dice = dice_loss(y_true, y_pred)
     return tf.cast(cce, dtype=tf.float32) + dice
 
+def convertColorToLabel(img):
+    color_to_label = {
+        (115, 178, 115): 0,  # non_mining_land (green)
+        (255, 0, 0): 1,      # illegal_mining_land (red)
+        (0, 0, 0): 2,        # beach (black)
+    }
+
+    # Create empty label array
+    label_img = np.zeros((img.shape[0], img.shape[1]), dtype=np.uint8)
+
+    # Map each RGB color to its corresponding label
+    for color, label in color_to_label.items():
+        mask = np.all(img == color, axis=2)
+        label_img[mask] = label
+
+    # One-hot encode the label image
+    num_classes = len(color_to_label)
+    one_hot = np.zeros((img.shape[0], img.shape[1], num_classes), dtype=np.uint8)
+    for c in range(num_classes):
+        one_hot[:, :, c] = (label_img == c).astype(np.uint8)
+
+    return one_hot
+
 def readImages(data, typeData, width, height):
     images = []
     for img in data:
@@ -84,7 +107,7 @@ def readImages(data, typeData, width, height):
             img = cv2.imread(str(img), cv2.IMREAD_COLOR)
             img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
             img = cv2.resize(img, (width, height), interpolation=cv2.INTER_NEAREST)
-            images.append(self.convertColorToLabel(img))
+            images.append(convertColorToLabel(img))
 
         elif typeData == 'o':  # Optic image
             img = cv2.imread(str(img), cv2.IMREAD_COLOR)