Update app.py

app.py

@@ -3,7 +3,10 @@ import os
 import numpy as np
 import matplotlib.pyplot as plt
 from tensorflow.keras.models import load_model
-from your_dataset_module import readDataset  # Replace with your actual dataset module
+import pathlib
+import natsort
+import datetime
+import shutil
 
 # Configuration
 HEIGHT = WIDTH = 256
@@ -18,6 +21,173 @@ CLASS_COLORS = [
     [0, 0, 0]
 ]
 
+class readDataset:
+    def __init__(self, sarPathes, opticPathes, masksPathes):
+        self.sarPathes = sarPathes
+        self.opticPathes = opticPathes
+        self.masksPathes = masksPathes
+        self.sarImages = None
+        self.opticImages = None
+        self.masks = None
+        self.testSarImages = None
+        self.testopticImages = None
+        self.testMasks = None
+
+    def readPathes(self):
+        # Get all file paths
+        all_sar_images = natsort.natsorted(list(pathlib.Path(self.sarPathes).glob('*.*')))
+        all_optic_images = natsort.natsorted(list(pathlib.Path(self.opticPathes).glob('*.*')))
+        all_mask_images = natsort.natsorted(list(pathlib.Path(self.masksPathes).glob('*.*')))
+
+        # Clean up .ipynb_checkpoints
+        for directory in [self.sarPathes, self.opticPathes, self.masksPathes]:
+            try:
+                shutil.rmtree(os.path.join(directory, ".ipynb_checkpoints"))
+                print(f".ipynb_checkpoints directory deleted successfully from {directory}.")
+            except Exception:
+                pass
+
+        # Extract image IDs - just getting the filename without extension
+        def extract_id(filepath):
+            return pathlib.Path(str(filepath)).stem
+
+        # Create dictionaries mapping IDs to paths for efficient lookup
+        sar_dict = {extract_id(f): f for f in all_sar_images}
+        optic_dict = {extract_id(f): f for f in all_optic_images}
+        mask_dict = {extract_id(f): f for f in all_mask_images}
+
+        # Find common IDs across all three datasets
+        common_ids = set(sar_dict.keys()) & set(optic_dict.keys()) & set(mask_dict.keys())
+
+        # Create matched file lists using sorted common IDs
+        sorted_common_ids = natsort.natsorted(list(common_ids))
+        self.sarImages = [sar_dict[id] for id in sorted_common_ids]
+        self.opticImages = [optic_dict[id] for id in sorted_common_ids]
+        self.masks = [mask_dict[id] for id in sorted_common_ids]
+
+        print(f"(INFO..) Found {len(all_sar_images)} SAR, {len(all_optic_images)} optical, {len(all_mask_images)} mask images")
+        print(f"(INFO..) Complete triplets: {len(common_ids)}")
+
+    def convertColorToLabel(self, 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 colors to labels
+        for color, label in color_to_label.items():
+            mask = np.all(img == color, axis=2)
+            label_img[mask] = label
+
+        # One-hot encode labels
+        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(self, data, typeData, width, height):
+        images = []
+        for img in data:
+            if typeData == 's':  # SAR image
+                with rasterio.open(str(img)) as src:
+                    sar_bands = [src.read(i) for i in range(1, src.count + 1)]
+                    sar_image = np.stack(sar_bands, axis=-1)
+                # Stretching
+                p2, p98 = np.percentile(sar_image, (2, 98))
+                sar_image = np.clip(sar_image, p2, p98)
+                sar_image = ((sar_image - p2) / (p98 - p2) * 255).astype(np.uint8)
+                # Resize
+                sar_image = cv2.resize(sar_image, (width, height), interpolation=cv2.INTER_AREA)
+                images.append(np.expand_dims(sar_image, axis=-1))
+            elif typeData == 'm':  # Mask image
+                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))
+            elif typeData == 'o':  # Optic image
+                img = cv2.imread(str(img), cv2.IMREAD_COLOR)
+                img = cv2.resize(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), (width, height), interpolation=cv2.INTER_AREA)
+                images.append(img)
+        print(f"(INFO..) Read {len(images)} {typeData} images")
+        return np.array(images)
+
+    def normalizeImages(self, images, typeData):
+        normalized_images = []
+        for img in images:
+            img = img.astype(np.uint8)
+            if typeData == 's':
+                img = img / 255.
+            if typeData == 'o':
+                img = img / 255.
+            normalized_images.append(img)
+        print("(INFO..) Normalization Image Done")
+        return np.array(normalized_images)
+
+    def dataAugmentation(self, sar_images, optic_images, masks, n_augments, size=WIDTH):
+        # Define augmentation pipeline once
+        augmentation = A.ReplayCompose([
+            A.RandomResizedCrop(size=(size, size), scale=(0.2, 0.9), ratio=(1, 1),
+                                interpolation=cv2.INTER_AREA, mask_interpolation=cv2.INTER_NEAREST, p=0.5),
+            A.HorizontalFlip(p=0.5),
+            A.ShiftScaleRotate(scale_limit=(0.0, 0.15), rotate_limit=(-90, 90),
+                              interpolation=cv2.INTER_AREA, mask_interpolation=cv2.INTER_NEAREST,
+                              border_mode=cv2.BORDER_REFLECT, p=0.5),
+            A.RandomGamma(p=0.5),
+            A.RandomBrightnessContrast(brightness_limit=(-0.25, 0.25), contrast_limit=(-0.25, 0.25), p=0.5)
+        ], additional_targets={'sar': 'image'})
+
+        if not (len(sar_images) == len(optic_images) == len(masks)):
+            raise ValueError("Number of SAR images, optic images, and masks must be the same.")
+
+        # Initialize lists with original data
+        augmented_sar = list(sar_images)
+        augmented_optic = list(optic_images)
+        augmented_masks = list(masks)
+
+        # Perform augmentations
+        for i, (sar, optic, mask) in enumerate(zip(sar_images, optic_images, masks)):
+            for _ in range(n_augments):
+                augmented = augmentation(image=optic.astype(np.uint8),
+                                        mask=mask.astype(np.uint8),
+                                        sar=sar.astype(np.uint8))
+                augmented_sar.append(augmented['sar'])
+                augmented_optic.append(augmented['image'])
+                augmented_masks.append(augmented['mask'])
+
+        # Print statistics
+        total_original = len(optic_images)
+        total_augmented = len(augmented_optic)
+        print(f"(INFO..) Original Train Optic Images: {total_original}")
+        print(f"(INFO..) Total Augmented Train Optic Images: {total_augmented}")
+        print(f"(INFO..) Augmentation Multiplier: {total_augmented / total_original:.2f}x")
+        print("(INFO..) Augmentation Image Done \n")
+
+        return (np.array(augmented_sar), np.array(augmented_optic), np.array(augmented_masks))
+
+    def splitDataset(self, sar_images, optic_images, masks, test_size=0.1, n_augments=10):
+        data = list(zip(sar_images, optic_images, masks))
+        train_data, test_data = train_test_split(data, test_size=test_size, random_state=42)
+
+        # Unpack the training and test data
+        train_sar, train_optic, train_masks = zip(*train_data)
+        test_sar, test_optic, test_masks = zip(*test_data)
+
+        # Augment train data
+        train_sar_aug, train_optic_aug, train_masks_aug = self.dataAugmentation(
+            np.array(train_sar), np.array(train_optic), np.array(train_masks),
+            n_augments=n_augments
+        )
+
+        print("(INFO..) Splitting and Saving Data Done \n")
+        return (
+            np.array(train_sar_aug), np.array(train_optic_aug), np.array(train_masks_aug),
+            np.array(test_sar), np.array(test_optic), np.array(test_masks)
+        )
 # Streamlit App Title
 st.title("Satellite Mining Segmentation: SAR + Optic Image Inference")