Started separating content in notebook into seperate files.

Author: KennethTM

data_classes.py

@@ -0,0 +1,93 @@
+import torch
+from torch.utils.data import Dataset
+import random
+from helpers import mask_morph_trans, dem_aug, mask_aug, dem_scale
+
+#Classes (train and validation) for loading DEM data for unsupervised training
+
+class DEMTrain(Dataset):
+
+    def __init__(self, array, masks, dem_transform = dem_aug, mask_transform = mask_aug):
+      
+        self.array = array
+        self.masks = masks
+        self.n_masks = masks.shape[0]
+                
+        self.dem_transform = dem_transform
+        
+        self.mask_transform = mask_transform
+
+    def __getitem__(self, idx):
+
+        target = self.array[idx]
+        target_transformed = self.dem_transform(image=target)
+        target_trans = target_transformed["image"]
+        target_tensor = torch.from_numpy(target_trans).unsqueeze(0)
+        
+        mask = self.masks[random.choice(range(self.n_masks))]
+        mask_transformed = self.mask_transform(image=mask)
+        mask_trans = mask_transformed["image"]
+        mask_trans_morph = mask_morph_trans(mask_trans, p=0.25)
+        mask_tensor = torch.from_numpy(mask_trans_morph).unsqueeze(0)
+        
+        input_tensor = target_tensor*(1 - mask_tensor)
+        
+        return input_tensor, target_tensor, mask_tensor
+
+    def __len__(self):
+        return self.array.shape[0]
+
+class DEMValid(Dataset):
+
+    def __init__(self, array, masks):
+      
+        self.array = array
+        self.masks = masks
+
+    def __getitem__(self, idx):
+
+        target = self.array[idx]
+        target_tensor = torch.from_numpy(target).unsqueeze(0)
+        
+        mask = self.masks[idx]
+        mask_tensor = torch.from_numpy(mask).unsqueeze(0)
+        
+        input_tensor = target_tensor*(1 - mask_tensor)
+        
+        return input_tensor, target_tensor, mask_tensor
+
+    def __len__(self):
+        return self.array.shape[0]
+
+#Class for loading lake data from dicts
+
+class Lakes(Dataset):
+
+    def __init__(self, lakes_list, transform = None):
+      
+        self.lakes_list = lakes_list
+        self.transform = transform
+
+    def __getitem__(self, idx):
+        
+        item = self.lakes_list[idx]
+        lake = item["lake"]
+        mask = item["mask"]
+                
+        lake = dem_scale(lake)
+
+        if self.transform is not None:
+            arrays_trans = self.transform(image = lake, mask = mask)
+            lake = arrays_trans["image"]
+            mask = arrays_trans["mask"]
+            mask = mask_morph_trans(mask)
+        
+        target_tensor = torch.from_numpy(lake).unsqueeze(0)
+        mask_tensor = torch.from_numpy(mask).unsqueeze(0)
+        
+        input_tensor = target_tensor * (1-mask_tensor)
+        
+        return input_tensor, target_tensor, mask_tensor
+
+    def __len__(self):
+        return len(self.lakes_list)

dem_autoencoder.ipynb

@@ -0,0 +1,44 @@
+import numpy as np
+import random
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+from pytorch_lightning.callbacks import ModelCheckpoint
+from lightning_module import AutoEncoder
+from helpers import dem_scale
+from data_classes import DEMTrain, DEMValid
+
+pl.seed_everything(9999)
+random.seed(9999)
+
+#Load DEM data
+data = np.load("data/data.npz")
+dem_data = dem_scale(data["dem"])
+
+train = dem_data[:8000]
+valid = dem_data[8000:]
+
+#Load mask data
+random_mask = np.load("data/lakes_random.npz")
+mask_data = random_mask["mask"]
+
+random_mask_idx = random.choices(range(mask_data.shape[0]), k=valid.shape[0])
+valid_mask = mask_data[random_mask_idx]
+
+#Create datasets
+train_dataset = DEMTrain(train, mask_data)
+valid_dataset = DEMValid(valid, valid_mask)
+
+#Create dataloaders
+train_loader = DataLoader(train_dataset, batch_size=32, num_workers=8, shuffle = True)
+val_loader = DataLoader(valid_dataset, batch_size=32, num_workers=8, shuffle = False)
+
+#Initiate model
+lake_model = AutoEncoder(init_features=8)
+
+#Initiate trainer
+checkpoint_callback = ModelCheckpoint(monitor="val_loss", save_last=True)
+
+trainer = pl.Trainer(gpus=1, max_epochs=1000, callbacks=checkpoint_callback, precision=16)
+
+#Train model
+trainer.fit(lake_model, train_loader, val_loader)

dem_model_train.py

@@ -0,0 +1,62 @@
+import random
+import numpy as np
+import cv2
+import albumentations as A
+
+#Helper functions 
+
+#Scale data range to -1 to 1
+def dem_scale(dem, min_val=-25, max_val=175):
+    zero_one = (dem - min_val)/(max_val - min_val)
+    minus_one = (zero_one*2)-1
+    
+    return(minus_one)
+
+#Scale range back to original (from -1 to 1)
+def dem_inv_scale(dem_scale, min_val=-25, max_val=175):
+    zero_one = (dem_scale + 1)/2
+    orig_scale = zero_one*(max_val - min_val) + min_val
+    
+    return(orig_scale)
+
+#Mask augmentation using erotion and dilation
+def mask_morph_trans(mask, p=0.5, min_iters = 1, max_iters = 10):
+    
+    rand = random.uniform(0, 1)
+    
+    if rand > p:
+        return(mask)
+        
+    kernel = np.ones((3,3),np.uint8)
+    morph_op = random.choice([cv2.erode, cv2.dilate])
+    iters = random.randint(min_iters, max_iters)
+    
+    mask_copy = mask.copy()
+    mask_morph = morph_op(mask_copy, kernel, iterations=iters)
+    
+    if mask_morph.sum()/mask.sum() < 0.1:
+        return(mask)
+    
+    return(mask_morph)
+
+#Augmentations
+dem_aug = A.Compose([
+    A.RandomRotate90(p=0.25),
+    A.Flip(p=0.25),
+    A.RandomResizedCrop(p=0.25, height=256, width=256, scale=(0.5, 1), interpolation=cv2.INTER_LINEAR)
+])
+
+mask_aug = A.Compose([
+    A.ShiftScaleRotate(p=0.25, scale_limit=0.2, shift_limit=0.2,
+                        interpolation=cv2.INTER_NEAREST, border_mode=cv2.BORDER_CONSTANT),
+    A.RandomRotate90(p=0.25),
+    A.Flip(p=0.25)
+])
+
+lake_aug = A.Compose([
+    A.ShiftScaleRotate(p=0.25, border_mode=cv2.BORDER_CONSTANT, interpolation=cv2.INTER_LINEAR),
+    A.RandomRotate90(p=0.25),
+    A.Flip(p=0.25),
+    A.GaussNoise(p=0.25, var_limit=(0, 1e-4)),
+    A.GaussianBlur(p=0.25)
+])

helpers.py

@@ -0,0 +1,25 @@
+import cv2
+import os
+import pickle
+import numpy as np 
+
+#Load random lake masks created in "lakes_random_shapes.R" script and save as numpy array
+
+#Load lake masks and convert to array
+mask_random_dir = "data/lakes_random/"
+mask_random_paths = os.listdir(mask_random_dir)
+
+mask_random_list = []
+
+for i in mask_random_paths:
+    lake_path = mask_random_dir+i
+    lake_array = cv2.imread(lake_path, cv2.IMREAD_UNCHANGED)
+    lake_array_float = lake_array.astype("float32") / 255.0
+    mask_random_list.append(lake_array_float)
+
+mask_random_np = np.array(mask_random_list)
+
+#Save file
+file_name = "data/lakes_random"
+
+np.savez(file_name, mask = mask_random_np)

lakes_random_shapes_list.py

@@ -0,0 +1,37 @@
+import pytorch_lightning as pl
+from loss import MAELossHole
+import torch 
+from model import UNet
+
+class AutoEncoder(pl.LightningModule):
+    def __init__(self, init_features = 8, lr = 1e-4):
+        super().__init__()
+        
+        self.save_hyperparameters()
+                
+        self.lr = lr
+                
+        self.unet = UNet(in_channels=1, out_channels=1, mask_channels=1, init_features=init_features)
+        
+        self.loss = MAELossHole()
+
+    def forward(self, x_in, x_mask):
+        x_hat = self.unet(x_in, x_mask)
+        return x_hat
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)
+        return optimizer
+
+    def training_step(self, train_batch, batch_idx):
+        x_in, x_obs, x_mask = train_batch
+        x_hat = self.unet(x_in, x_mask)
+        loss = self.loss(x_hat, x_obs, x_mask)
+        self.log('train_loss', loss, on_epoch=True, prog_bar=True)
+        return {'loss': loss}
+
+    def validation_step(self, val_batch, batch_idx):
+        x_in, x_obs, x_mask = val_batch
+        x_hat = self.unet(x_in, x_mask)
+        loss = self.loss(x_hat, x_obs, x_mask)
+        self.log('val_loss', loss, on_epoch=True, prog_bar=True)

lightning_module.py

@@ -0,0 +1,46 @@
+import torch.nn as nn
+
+#Loss functions 
+
+class MSELossWeighted(nn.Module):
+    def __init__(self, w_hole = 10, w_nonhole = 1):
+        super().__init__()
+        self.l2 = nn.MSELoss()
+        self.w_hole = w_hole
+        self.w_nonhole = w_nonhole
+
+    def forward(self, hat, obs, mask):
+        
+        l2_hole = self.l2(hat[mask == 1], obs[mask == 1])
+        l2_nonhole = self.l2(hat[mask == 0], obs[mask == 0])
+        
+        l2_total = (l2_hole*self.w_hole) + (l2_nonhole*self.w_nonhole)
+
+        return l2_total
+    
+class MAELossWeighted(nn.Module):
+    def __init__(self, w_hole = 10, w_nonhole = 1):
+        super().__init__()
+        self.l1 = nn.L1Loss()
+        self.w_hole = w_hole
+        self.w_nonhole = w_nonhole
+
+    def forward(self, hat, obs, mask):
+        
+        l1_hole = self.l1(hat[mask == 1], obs[mask == 1])
+        l1_nonhole = self.l1(hat[mask == 0], obs[mask == 0])
+        
+        l1_total = (l1_hole*self.w_hole) + (l1_nonhole*self.w_nonhole)
+
+        return l1_total
+    
+class MAELossHole(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.l1 = nn.L1Loss()
+
+    def forward(self, hat, obs, mask):
+        
+        l1_hole = self.l1(hat[mask == 1], obs[mask == 1])
+
+        return l1_hole