Upload 5 files

app.py

@@ -0,0 +1,79 @@
+import streamlit as st
+import torch
+from PIL import Image
+import torchvision.transforms as transforms
+from train import UNet
+import numpy as np 
+
+# Load the trained model
+model_path = '/teamspace/studios/this_studio/Aerial-Segmentation/model.pth'
+model = UNet(n_channels=3, n_classes=6)
+model.load_state_dict(torch.load(model_path))
+model.eval()
+
+# Create a Streamlit app
+st.title('Aerial Image Segmentation')
+
+# Add a file uploader to the app
+uploaded_file = st.file_uploader("Choose an image...", type="jpg")
+
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+
+    # Display the original image
+    st.image(image, caption='Uploaded Image.', use_column_width=True)
+
+    # Preprocess the image
+    data_transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor()]
+    )
+    image = data_transform(image)
+    image = image.unsqueeze(0)  # add a batch dimension
+
+    # Pass the image through the model
+    with torch.no_grad():
+        output = model(image)
+
+    # Postprocess the output
+    # Define the color map
+    color_map = {
+        0: np.array([155, 155, 155]),  # Unlabeled
+        1: np.array([60, 16, 152]),  # Building
+        2: np.array([132, 41, 246]),  # Land
+        3: np.array([110, 193, 228]),  # Road
+        4: np.array([254, 221, 58]),  # Vegetation
+        5: np.array([226, 169, 41])  # Water
+    }
+    class_labels = {
+    0: 'Unlabeled',
+    1: 'Building',
+    2: 'Land',
+    3: 'Road',
+    4: 'Vegetation',
+    5: 'Water'
+    }
+
+    # Display the class labels and their colors in a sidebar
+    for k, v in class_labels.items():
+        st.sidebar.markdown(f'<div style="color:rgb{tuple(color_map[k])};">{v}</div>', unsafe_allow_html=True)
+
+    # Pass the image through the model
+    with torch.no_grad():
+        output = model(image)
+
+    # Postprocess the output
+    output = torch.argmax(output.squeeze(), dim=0).detach().cpu().numpy()
+
+    # Squeeze the batch dimension
+    output = np.squeeze(output)
+
+    # Now you can create the RGB image
+    output_rgb = np.zeros((output.shape[0], output.shape[1], 3), dtype=np.uint8)
+    for k, v in color_map.items():
+        output_rgb[output == k] = v
+
+    # Display the segmented image
+    st.image(output_rgb, caption='Segmented Image.', use_column_width=True)
+
+

dataloader.py

@@ -0,0 +1,51 @@
+import os
+from torch.utils.data import Dataset
+from PIL import Image
+import torchvision.transforms as transforms
+import numpy as np
+
+class AerialImageDataset(Dataset):
+    def __init__(self, image_dir, mask_dir, transform=None):
+        self.image_dir = image_dir
+        self.mask_dir = mask_dir
+        self.transform = transform
+        self.images = os.listdir(self.image_dir)
+        self.Hex_Classes = [
+            ('Unlabeled', '#9B9B9B'),
+            ('Building','#3C1098'),
+            ('Land', '#8429F6'),
+            ('Road', '#6EC1E4'),
+            ('Vegetation', '#FEDD3A'),
+            ('Water', '#E2A929'),
+        ]
+
+    def __len__(self):
+        return len(self.images)
+   
+    def __getitem__(self, idx):
+        
+        img_path = os.path.join(self.image_dir, self.images[idx])
+        mask_path = os.path.join(self.mask_dir, self.images[idx].replace('.jpg', '.png'))
+
+        image = Image.open(img_path)
+        mask = Image.open(mask_path)
+
+        mask = np.array(mask)
+        mask = self.encode_segmap(mask)
+        mask = Image.fromarray(mask)
+
+        if self.transform:
+            image = self.transform(image)
+            mask = self.transform(mask)
+
+        return image, mask
+
+    def encode_segmap(self, mask):
+        mask = mask.astype(int)
+        label_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int16) 
+        for i, (name, color) in enumerate(self.Hex_Classes):
+            if mask.ndim == 3:
+                label_mask[(mask[:,:,0] == int(color[1:3], 16)) & (mask[:,:,1] == int(color[3:5], 16)) & (mask[:,:,2] == int(color[5:7], 16))] = i
+            elif mask.ndim == 2:
+                label_mask[(mask == int(color[1:3], 16))] = i
+        return label_mask

dataprep.py

@@ -0,0 +1,71 @@
+import os
+import shutil
+import random
+
+dataset_path = "/teamspace/studios/this_studio/Aerial-Segmentation/Semantic segmentation dataset"
+new_dataset_path = "/teamspace/studios/this_studio/Aerial-Segmentation"
+train_path = os.path.join(new_dataset_path, "train")
+val_path = os.path.join(new_dataset_path, "val")
+
+os.makedirs(train_path, exist_ok=True)
+os.makedirs(val_path, exist_ok=True)
+
+train_image_path = os.path.join(train_path, "images")
+train_mask_path = os.path.join(train_path, "masks")
+
+val_image_path = os.path.join(val_path, "images")
+val_mask_path = os.path.join(val_path, "masks")
+
+os.makedirs(train_image_path, exist_ok=True)
+os.makedirs(val_image_path, exist_ok=True)
+os.makedirs(train_mask_path, exist_ok=True)
+os.makedirs(val_mask_path, exist_ok=True)
+
+
+tile_folders = [folder for folder in os.listdir(dataset_path) if os.path.isdir(os.path.join(dataset_path, folder))]
+
+n_train_images = 8
+n_val_images = 1
+
+
+def copy(train_status):
+    if train_status:
+        images = train_images
+        path_image = train_image_path
+        path_mask = train_mask_path
+    else:
+        images = val_images
+        path_image = val_image_path
+        path_mask = val_mask_path
+
+    
+
+    for image in images:
+        tile_image_name = f'{tile_folder}_{image}'
+        shutil.copy(os.path.join(images_path, image), os.path.join(path_image, tile_image_name))
+
+        mask_name = image.split('.')[0]+'.png'
+        tile_mask_name = f'{tile_folder}_{mask_name}'
+        shutil.copy(os.path.join(masks_path, mask_name), os.path.join(path_mask,  tile_mask_name))
+        
+
+for tile_folder in tile_folders:
+    images_path = os.path.join(dataset_path, tile_folder, 'images')
+    masks_path = os.path.join(dataset_path, tile_folder, 'masks')
+
+    images = os.listdir(images_path)
+    masks = os.listdir(masks_path)
+
+    random.shuffle(images)
+    random.shuffle(masks)
+
+    train_images = images[:n_train_images]
+    val_images = images[n_train_images:]
+
+    copy(train_status=True)
+    copy(train_status=False)
+
+
+shutil.rmtree(dataset_path)
+
+print(f"Data organization and split completed successfully. Total Training Files is {len(os.listdir(train_image_path))} and Validation Files is {len(os.listdir(val_image_path))}")

main.py

@@ -0,0 +1,28 @@
+import os
+import lightning as L
+from dataloader import AerialImageDataset
+from train import UNet
+from torch.utils.data import DataLoader
+from torchvision.transforms import transforms
+import torch
+
+train_path = "/teamspace/studios/this_studio/Aerial-Segmentation/train"
+val_path = "/teamspace/studios/this_studio/Aerial-Segmentation/val"
+
+data_transform = transforms.Compose([
+    transforms.Resize((512, 512)),
+    transforms.ToTensor()]
+    )
+
+train_dataset = AerialImageDataset(os.path.join(train_path, 'images'), os.path.join(train_path, 'masks'), transform=data_transform)
+val_dataset = AerialImageDataset(os.path.join(val_path, 'images'), os.path.join(val_path, 'masks'), transform=data_transform)
+
+train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
+val_loader = DataLoader(val_dataset, batch_size=4, shuffle=False)
+
+model = UNet(n_channels=3, n_classes=6)
+
+trainer = L.Trainer(max_epochs=100)
+trainer.fit(model, train_loader, val_loader)
+
+torch.save(model.state_dict(), "/teamspace/studios/this_studio/Aerial-Segmentation/model.pth")

model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c7bbcc5453e099cdc1cf71edae30b611fed8ce893de98e35631a9a765832bbe
+size 53651114