Update app.py

app.py

@@ -24,174 +24,6 @@ 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)
-        )
-
 @tf.keras.saving.register_keras_serializable()
 def dice_score(y_true, y_pred, threshold=0.5, smooth=1.0):
     #determine binary or multiclass segmentation
@@ -230,39 +62,83 @@ def cce_dice_loss(y_true, y_pred):
     dice = dice_loss(y_true, y_pred)
     return tf.cast(cce, dtype=tf.float32) + dice
 
+def readImages(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)
+
+            # Contrast 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.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA)
+            images.append(img)
+
+    print(f"(INFO..) Read {len(images)} '{typeData}' image(s)")
+    return np.array(images)
+
+
+def normalizeImages(images, typeData):
+    normalized_images = []
+    for img in images:
+        img = img.astype(np.uint8)
+        if typeData in ['s', 'o']:
+            img = img / 255.
+        normalized_images.append(img)
+
+    print("(INFO..) Normalization Image Done")
+    return np.array(normalized_images)
 
+        
 # Streamlit App Title
 st.title("Satellite Mining Segmentation: SAR + Optic Image Inference")
 
-# Sidebar inputs
-#dataset_path = st.text_input("Enter dataset path", "path/to/your/dataset")
-#model_path = st.text_input("Enter path to trained model (.h5)", "model.h5")
-dataset_path = "bangka_testing_data"
-model_path = "Residual_UNET_Bilinear.keras"
+
+sar_file = st.file_uploader("Upload SAR Image", type=["tiff"])
+optic_file = st.file_uploader("Upload Optical Image", type=["tiff"])
+mask_file = st.file_uploader("Upload Mask Image", type=["tiff"])
+
 num_samples = st.slider("Number of test samples to visualize", 1, 10, 3)
 
+if sar_file is not None and optic_file is not None and mask_file is not None:
+    st.success("All files uploaded successfully!")
+    st.write(f"Number of samples selected for visualization: {num_samples}")
+else:
+    st.warning("Please upload all three .tiff files to proceed.")
+    
+sarImages = [sar_file]
+opticImages = [optic_file]
+masks = [mask_file]
+model_path = "Residual_UNET_Bilinear.keras"
+
+
 if st.button("Run Inference"):
     with st.spinner("Loading data and model..."):
 
-        # Prepare paths
-        sar_path = os.path.join(dataset_path, 'sar_images')
-        optic_path = os.path.join(dataset_path, 'optic_images')
-        mask_path = os.path.join(dataset_path, 'masks')
-
-        # Read dataset
-        dataset = readDataset(
-            sarPathes=sar_path,
-            opticPathes=optic_path,
-            masksPathes=mask_path
-        )
-        dataset.readPathes()
-
-        sar_images = dataset.readImages(dataset.sarImages, typeData='s', width=WIDTH, height=HEIGHT)
-        optic_images = dataset.readImages(dataset.opticImages, typeData='o', width=WIDTH, height=HEIGHT)
-        masks = dataset.readImages(dataset.masks, typeData='m', width=WIDTH, height=HEIGHT)
-
-        sar_images = dataset.normalizeImages(sar_images, 's')
-        optic_images = dataset.normalizeImages(optic_images, 'i')
+
+        sar_images = readImages(sarImages, typeData='s', width=WIDTH, height=HEIGHT)
+        optic_images = readImages(opticImages, typeData='o', width=WIDTH, height=HEIGHT)
+        masks = readImages(masks, typeData='m', width=WIDTH, height=HEIGHT)
+
+        sar_images = normalizeImages(sar_images, 's')
+        optic_images = normalizeImages(optic_images, 'i')
 
         # Load model
         model = tf.keras.models.load_model(model_path,