MapYourCity_ExampleTest.py

#!/usr/bin/env python
# coding: utf-8

# Example Test Model for the MapYourCity dataset          

"""
ABOUT SCRIPT: 
This file creates an example Test Model for the validation and test sets using PyTorch  
This code is generated by Nikolaos Dionelis @ESA 
LAST EDITED: 05/02/2024 
"""

# Python library imports 
import imageio.v2 as io
import matplotlib.pyplot as plt 
import torchvision             
import random    
import torchvision.transforms as T                         
import os          
import shutil
from PIL import Image
from torchvision import transforms, datasets
import numpy as np      
from torch.utils.data import Dataset
from sklearn.preprocessing import LabelEncoder 
import torch
from torchsummary import summary
from tqdm import tqdm
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import pandas as pd
import rasterio 
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
import seaborn as sn

# The main lines in this script are: 546 and 585               
# To test: torch.save(model.state_dict(), './modelB.pt')    

# We define the Folder Dataset Path   
FOLDER = '/Data/ndionelis/StreetDataset/'       
FOLDER = '/Data/ndionelis/StreetDataset/'              
MAINFOLDER = '/Data/ndionelis/'   
#NUMWORKERS = 6                    
NUMWORKERS = 0 
#BATCH_SIZE = 256        
BATCH_SIZE = 32
#epochs = 200       
epochs = 100
SEED = random.randint(1, 10000)
print('The random seed is: ' + str(SEED) + '.')  
torch.cuda.empty_cache()
random.seed(SEED)
np.random.seed(SEED) 
torch.manual_seed(SEED)   
torch.cuda.manual_seed_all(SEED)

def get_random_pos(img, window_shape): 
    w, h = window_shape
    W, H = img.shape[-2:]
    x1 = random.randint(0, W - w - 1)
    x2 = x1 + w
    y1 = random.randint(0, H - h - 1)
    y2 = y1 + h 
    return x1, x2, y1, y2
WINDOW_SIZE = (600, 600)        
# image1 = torchvision.io.read_image(FOLDER+'0/0_00_311059203.jpg')    
# x1, x2, y1, y2 = get_random_pos(image1, WINDOW_SIZE)
# image1 = image1[:, x1:x2,y1:y2]
# plt.imshow(image1.permute(1, 2, 0))   

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize((256, 256)),
        transforms.ToTensor(),
    ]),
    'val': transforms.Compose([
        transforms.Resize((256, 256)), 
        transforms.ToTensor(), 
    ]),
    'test': transforms.Compose([ 
        transforms.Resize((256, 256)), 
        transforms.ToTensor(),  
    ]),
} 
data_dir = MAINFOLDER + 'TheNewDataset'   
image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
                                        data_transforms[x])
                for x in ['train', 'val', 'test']}    
train_dataset = image_datasets['train']             
valid_dataset = image_datasets['test']  

dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=BATCH_SIZE,
                                            shuffle=True, num_workers=NUMWORKERS) 
            for x in ['train', 'val', 'test']}
train_dataloader = dataloaders['train']      
valid_dataloader = dataloaders['test'] 

labels = np.array(dataloaders['train'].dataset.targets) 
lb = LabelEncoder() 
labels = lb.fit_transform(labels)
print(f"Total Number of Classes: {len(lb.classes_)}") 
device = torch.device("cuda:0")
print('Device: ' + str(device))   

import torchvision.transforms.functional as fn 
augmentation_train_in = {
    transforms.Compose([
        transforms.ToTensor(),
    ]),
}   
dataloader_kwargs = {'num_workers': 0}  

# We use train_loader and test_loader      
# Define paths to data   
input_path = "/Data/ndionelis/building-age-dataset/" 
train_path = input_path + "train/data/"
test_path = input_path + "test/data/"
test_df = pd.read_csv(input_path + "test/test-set.csv")
train_df = pd.read_csv(input_path + "train/train-set.csv")
train_df.head()
test_df.head() 
names_data = os.listdir(train_path) # to not load all data in a single tensor, load only the names                    
length_names = len(names_data)
perm = torch.randperm(length_names)
#idx = perm[:round(0.8*length_names)] # draw round(0.8*length_names) samples     
#torch.save(idx, 'indexForTrainVal.pt')        
idx = torch.load('indexForTrainVal.pt')  
names_data = np.array(names_data)
idx = idx.numpy() 
training_data = names_data[idx]
mask = np.ones(names_data.size, dtype=bool) 
mask[idx] = False
test_data = names_data[mask]

#from utils import get_activation, get_normalization, SE_Block 
class SE_Block(nn.Module): 
    def __init__(self, channels, reduction=16, activation="relu"):
        super().__init__()
        self.reduction = reduction
        self.squeeze = nn.AdaptiveAvgPool2d(1)
        self.excitation = nn.Sequential(
            nn.Linear(channels, channels // self.reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channels // self.reduction, channels, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        bs, c, _, _ = x.shape
        y = self.squeeze(x).view(bs, c)
        y = self.excitation(y).view(bs, c, 1, 1)
        return x * y.expand_as(x)

def get_activation(activation_name):
    if activation_name == "relu":
        return nn.ReLU6(inplace=True)
    elif isinstance(activation_name, torch.nn.modules.activation.ReLU6):
        return activation_name
    elif activation_name == "gelu":
        return nn.GELU()
    elif isinstance(activation_name, torch.nn.modules.activation.GELU):
        return activation_name
    elif activation_name == "leaky_relu":
        return nn.LeakyReLU(inplace=True)
    elif isinstance(activation_name, torch.nn.modules.activation.LeakyReLU):
        return activation_name
    elif activation_name == "prelu":
        return nn.PReLU()
    elif isinstance(activation_name, torch.nn.modules.activation.PReLU):
        return activation_name
    elif activation_name == "selu":
        return nn.SELU(inplace=True)
    elif isinstance(activation_name, torch.nn.modules.activation.SELU):
        return activation_name
    elif activation_name == "sigmoid":
        return nn.Sigmoid()
    elif isinstance(activation_name, torch.nn.modules.activation.Sigmoid):
        return activation_name
    elif activation_name == "tanh":
        return nn.Tanh()
    elif isinstance(activation_name, torch.nn.modules.activation.Tanh):
        return activation_name
    elif activation_name == "mish":
        return nn.Mish()
    elif isinstance(activation_name, torch.nn.modules.activation.Mish):
        return activation_name
    else:
        raise ValueError(f"activation must be one of leaky_relu, prelu, selu, gelu, sigmoid, tanh, relu. Got: {activation_name}")

def get_normalization(normalization_name, num_channels, num_groups=32, dims=2):
    if normalization_name == "batch":
        if dims == 1:
            return nn.BatchNorm1d(num_channels)
        elif dims == 2:
            return nn.BatchNorm2d(num_channels)
        elif dims == 3:
            return nn.BatchNorm3d(num_channels)
    elif normalization_name == "instance":
        if dims == 1:
            return nn.InstanceNorm1d(num_channels)
        elif dims == 2:
            return nn.InstanceNorm2d(num_channels)
        elif dims == 3:
            return nn.InstanceNorm3d(num_channels)
    elif normalization_name == "layer":
        return nn.LayerNorm(num_channels)
    elif normalization_name == "group":
        return nn.GroupNorm(num_groups=num_groups, num_channels=num_channels)
    elif normalization_name == "bcn":
        if dims == 1:
            return nn.Sequential(
                nn.BatchNorm1d(num_channels),
                nn.GroupNorm(1, num_channels)
            )
        elif dims == 2:
            return nn.Sequential(
                nn.BatchNorm2d(num_channels),
                nn.GroupNorm(1, num_channels)
            )
        elif dims == 3:
            return nn.Sequential(
                nn.BatchNorm3d(num_channels),
                nn.GroupNorm(1, num_channels)
            )    
    elif normalization_name == "none":
        return nn.Identity()
    else:
        raise ValueError(f"normalization must be one of batch, instance, layer, group, none. Got: {normalization_name}") 

class CoreCNNBlock(nn.Module):
    def __init__(self, in_channels, out_channels, *, norm="batch", activation="relu", padding="same", residual=True):
        super(CoreCNNBlock, self).__init__()
        self.activation = get_activation(activation)
        self.residual = residual
        self.padding = padding
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.squeeze = SE_Block(self.out_channels)
        self.match_channels = nn.Identity()
        if in_channels != out_channels:
            self.match_channels = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0, bias=False),
                get_normalization(norm, out_channels),
            )
        self.conv1 = nn.Conv2d(self.in_channels, self.out_channels, 1, padding=0)
        self.norm1 = get_normalization(norm, self.out_channels)
        self.conv2 = nn.Conv2d(self.out_channels, self.out_channels, 3, padding=self.padding, groups=self.out_channels)
        self.norm2 = get_normalization(norm, self.out_channels)
        self.conv3 = nn.Conv2d(self.out_channels, self.out_channels, 3, padding=self.padding, groups=1)
        self.norm3 = get_normalization(norm, self.out_channels)

    def forward(self, x):
        identity = x
        x = self.activation(self.norm1(self.conv1(x)))
        x = self.activation(self.norm2(self.conv2(x)))
        x = self.norm3(self.conv3(x))
        x = x * self.squeeze(x)
        if self.residual:
            x = x + self.match_channels(identity)
        x = self.activation(x) 
        return x

class CoreEncoderBlock(nn.Module): 
    def __init__(self, depth, in_channels, out_channels, norm="batch", activation="relu", padding="same"):
        super(CoreEncoderBlock, self).__init__() 
        self.depth = depth
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.activation = activation
        self.norm = norm
        self.padding = padding
        self.blocks = []
        for i in range(self.depth): 
            _in_channels = self.in_channels if i == 0 else self.out_channels
            block = CoreCNNBlock(_in_channels, self.out_channels, norm=self.norm, activation=self.activation, padding=self.padding)
            self.blocks.append(block)
        self.blocks = nn.Sequential(*self.blocks)
        self.downsample = nn.MaxPool2d(kernel_size=2, stride=2)
    
    def forward(self, x):
        for i in range(self.depth):
            x = self.blocks[i](x)
        before_downsample = x
        x = self.downsample(x)
        return x, before_downsample

class CoreAttentionBlock(nn.Module):
    def __init__(self,
        lower_channels,
        higher_channels, *,
        norm="batch",
        activation="relu",
        padding="same",
    ):
        super(CoreAttentionBlock, self).__init__()
        self.lower_channels = lower_channels
        self.higher_channels = higher_channels
        self.activation = get_activation(activation)
        self.norm = norm
        self.padding = padding
        self.expansion = 4
        self.reduction = 4
        if self.lower_channels != self.higher_channels:
            self.match = nn.Sequential(
                nn.Conv2d(self.higher_channels, self.lower_channels, kernel_size=1, padding=0, bias=False),
                get_normalization(self.norm, self.lower_channels),
            )
        self.compress = nn.Conv2d(self.lower_channels, 1, kernel_size=1, padding=0)
        self.sigmoid = nn.Sigmoid()
        self.attn_c_pool = nn.AdaptiveAvgPool2d(self.reduction)
        self.attn_c_reduction = nn.Linear(self.lower_channels * (self.reduction ** 2), self.lower_channels * self.expansion)
        self.attn_c_extention = nn.Linear(self.lower_channels * self.expansion, self.lower_channels)

    def forward(self, x, skip):
        if x.size(1) != skip.size(1):
            x = self.match(x)
        x = x + skip
        x = self.activation(x)
        attn_spatial = self.compress(x)
        attn_spatial = self.sigmoid(attn_spatial)
        attn_channel = self.attn_c_pool(x)
        attn_channel = attn_channel.reshape(attn_channel.size(0), -1)
        attn_channel = self.attn_c_reduction(attn_channel)
        attn_channel = self.activation(attn_channel)
        attn_channel = self.attn_c_extention(attn_channel)
        attn_channel = attn_channel.reshape(x.size(0), x.size(1), 1, 1)
        attn_channel = self.sigmoid(attn_channel)
        return attn_spatial, attn_channel

class CoreDecoderBlock(nn.Module):
    def __init__(self, depth, in_channels, out_channels, *, norm="batch", activation="relu", padding="same"):
        super(CoreDecoderBlock, self).__init__()
        self.depth = depth
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.activation_blocks = activation
        self.activation = get_activation(activation)
        self.norm = norm
        self.padding = padding
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
        self.match_channels = CoreCNNBlock(self.in_channels * 2, self.out_channels, norm=self.norm, activation=self.activation_blocks, padding=self.padding)
        self.attention = CoreAttentionBlock(self.in_channels, self.in_channels, norm=self.norm, activation=self.activation_blocks, padding=self.padding)
        self.blocks = []
        for _ in range(self.depth):
            block = CoreCNNBlock(self.out_channels, self.out_channels, norm=self.norm, activation=self.activation_blocks, padding=self.padding)
            self.blocks.append(block)
        self.blocks = nn.Sequential(*self.blocks)
    
    def forward(self, x, skip):
        x = self.upsample(x)
        attn_s, attn_c = self.attention(x, skip)
        x = torch.cat([x, (skip * attn_s) + (skip + attn_c)], dim=1)
        x = self.match_channels(x)
        for i in range(self.depth):
            x = self.blocks[i](x)
        return x

class CoreUnet(nn.Module):  
    def __init__(self, *,
        input_dim=10,
        output_dim=1,
        depths=None,
        dims=None,
        activation="relu",
        norm="batch",
        padding="same",
    ): 
        super(CoreUnet, self).__init__() 
        self.depths = [3, 3, 9, 3] if depths is None else depths 
        self.dims = [96, 192, 384, 768] if dims is None else dims
        self.output_dim = output_dim
        self.input_dim = input_dim
        self.activation = activation
        self.norm = norm
        self.padding = padding
        self.dims = [v // 2 for v in self.dims] 
        assert len(self.depths) == len(self.dims), "depths and dims must have the same length. "   
        self.stem = nn.Sequential(
            CoreCNNBlock(self.input_dim, self.dims[0], norm=self.norm, activation=self.activation, padding=self.padding),
        )  
        self.encoder_blocks = []  
        for i in range(len(self.depths)):
            encoder_block = CoreEncoderBlock(
                self.depths[i],
                self.dims[i - 1] if i > 0 else self.dims[0],
                self.dims[i],
                norm=self.norm,
                activation=self.activation,
                padding=self.padding,
            )
            self.encoder_blocks.append(encoder_block)
        self.encoder_blocks = nn.ModuleList(self.encoder_blocks)
        self.decoder_blocks = [] 
        for i in reversed(range(len(self.encoder_blocks))):
            decoder_block = CoreDecoderBlock(
                self.depths[i],
                self.dims[i],
                self.dims[i - 1] if i > 0 else self.dims[0],
                norm=self.norm,
                activation=self.activation,
                padding=self.padding,
            )
            self.decoder_blocks.append(decoder_block)
        self.decoder_blocks = nn.ModuleList(self.decoder_blocks)
        self.bridge = nn.Sequential(
            CoreCNNBlock(self.dims[-1], self.dims[-1], norm=self.norm, activation=self.activation, padding=self.padding),
        )
        self.head = nn.Sequential(
            CoreCNNBlock(self.dims[0], self.dims[0], norm=self.norm, activation=self.activation, padding=self.padding),
            nn.Conv2d(self.dims[0], self.output_dim, kernel_size=1, padding=0),
        )

    def forward(self, x):
        skip_connections = []    
        x = self.stem(x)
        for block in self.encoder_blocks:
            x, skip = block(x)
            skip_connections.append(skip)
        x = self.bridge(x)
        return x

class CoreEncoder(nn.Module):
    def __init__(self, *,
        input_dim=10,
        output_dim=1,
        depths=None,
        dims=None,
        activation="relu",
        norm="batch",
        padding="same",
    ):
        super(CoreEncoder, self).__init__()
        self.depths = [3, 3, 9, 3] if depths is None else depths
        self.dims = [96, 192, 384, 768] if dims is None else dims
        self.output_dim = output_dim
        self.input_dim = input_dim
        self.activation = activation
        self.norm = norm
        self.padding = padding
        assert len(self.depths) == len(self.dims), "depths and dims must have the same length."
        self.stem = CoreCNNBlock(self.input_dim, self.dims[0], norm=self.norm, activation=self.activation, padding=self.padding)
        self.encoder_blocks = []  
        for i in range(len(self.depths)): 
            encoder_block = CoreEncoderBlock(
                self.depths[i],
                self.dims[i - 1] if i > 0 else self.dims[0],
                self.dims[i],
                norm=self.norm,
                activation=self.activation,
                padding=self.padding,
            )
            self.encoder_blocks.append(encoder_block)
        self.encoder_blocks = nn.ModuleList(self.encoder_blocks)
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(),
            nn.Linear(self.dims[-1], self.output_dim),
        )

    def forward(self, x):
        x = self.stem(x)
        for block in self.encoder_blocks:
            x, _ = block(x)
        x = self.head(x)
        return x

class ResNet152(nn.Module):
    def __init__(self, pretrained):
        super(ResNet152, self).__init__() 
        class MyResNet18(nn.Module):
            def __init__(self, resnet, resnet2):
                super().__init__()
                self.features = nn.Sequential(
                    resnet.conv1,
                    resnet.bn1,
                    resnet.relu,
                    resnet.maxpool,
                    resnet.layer1,
                    resnet.layer2,
                    resnet.layer3,
                    resnet.layer4
                ) 
                self.avgpool = resnet.avgpool
                self.fc = resnet.fc
                self.features2 = nn.Sequential(
                    resnet2.conv1,
                    resnet2.bn1,
                    resnet2.relu,
                    resnet2.maxpool,
                    resnet2.layer1,
                    resnet2.layer2,
                    resnet2.layer3,
                    resnet2.layer4
                )
                self.avgpool2 = resnet2.avgpool
                self.fc2 = resnet2.fc

            def _forward_impl(self, x: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
                x = self.features(x)
                x = self.avgpool(x)
                x = torch.flatten(x, 1)
                x = self.fc(x)
                x2 = self.features2(x2)
                x2 = self.avgpool2(x2)
                x2 = torch.flatten(x2, 1)
                x2 = self.fc2(x2) 
                return x, x2

            def forward(self, x: torch.Tensor, x2: torch.Tensor) -> torch.Tensor:
                return self._forward_impl(x, x2) 

        model = torchvision.models.resnet152(pretrained=True)
        model2 = torchvision.models.resnet152(pretrained=True)
        self.model = MyResNet18(model, model2)
        self.l0 = nn.Linear(4480, len(lb.classes_))

    def forward(self, x, x2, x3):
        batch, _, _, _ = x.shape
        x = self.model.features(x)
        x2 = self.model.features2(x2)
        BATCH_SIZE = 32 
        CHANNELS = 12
        HEIGHT = 64
        WIDTH = 64
        model = CoreUnet(
            input_dim=CHANNELS,
            output_dim=1,
        ).to(device)   
        model.train()
        x3 = model(x3) 
        x = F.adaptive_avg_pool2d(x, 1).reshape(batch, -1)    
        x2 = F.adaptive_avg_pool2d(x2, 1).reshape(batch, -1) 
        x3 = F.adaptive_avg_pool2d(x3, 1).reshape(batch, -1) 
        x = torch.cat((x, x2, x3), 1)   
        l0 = self.l0(x)
        return l0 

model = ResNet152(pretrained=True).to(device) 
#model.train() 
# we also use: http://github.com/ESA-PhiLab/AI4EO-Challenge-Building-Sustainability 
criterion = nn.CrossEntropyLoss() 
optimizer = optim.Adam(model.parameters(), lr=1e-5, weight_decay=0.5e-3)     

# to test: torch.save(model.state_dict(), './modelB.pt')          
model.load_state_dict(torch.load('./modelB.pt')) 
model.eval()   

def test(model):      
  print('Now testing')                      
  running_loss = 0.0  
  running_correct = 0
  with torch.no_grad():
      for idx, batch in enumerate(tqdm(test_loader2)): 
          pixel_values, pixel_values2, pixel_values3 = batch[0].to(device, dtype=torch.float32), batch[1].to(device, dtype=torch.float32), batch[2].to(device, dtype=torch.float32)
          pixel_values = pixel_values.permute(0, 3, 1, 2)         
          pixel_values2 = pixel_values2.permute(0, 3, 1, 2)   
          pixel_values3 = pixel_values3.permute(0, 3, 1, 2)
          outputs = model(pixel_values, pixel_values2, pixel_values3)
          _, preds = torch.max(outputs, 1)
          # We also use: http://github.com/ESA-PhiLab/AI4EO-Challenge-Building-Sustainability  
      return

def validate(model):       
  print('Now validating')                        
  y_pred, y_true = [], []  
  running_loss, running_correct = 0.0, 0  
  with torch.no_grad():
      for idx, batch in enumerate(tqdm(valid_dataloader)):
          pixel_values, pixel_values2, pixel_values3, labels = batch[0].to(device, dtype=torch.float32), batch[1].to(device, dtype=torch.float32), batch[2].to(device, dtype=torch.float32), batch[3].to(device)
          pixel_values = pixel_values.permute(0, 3, 1, 2)         
          pixel_values2 = pixel_values2.permute(0, 3, 1, 2)   
          pixel_values3 = pixel_values3.permute(0, 3, 1, 2)
          outputs = model(pixel_values, pixel_values2, pixel_values3)
          _, preds = torch.max(outputs, 1)
          running_correct += (preds == labels).sum().item() 
          y_true.append(labels.cpu()) 
          y_pred.append(preds.cpu())  
      accuracy = 100. * running_correct / len(valid_dataloader.dataset)
      print(f'Val Acc: {accuracy:.2f}') 
      y_true = np.concatenate(y_true) 
      y_pred = np.concatenate(y_pred)
      print(f'Pr: {100.*precision_score(y_true, y_pred, average="weighted"):.2f}, Re: {100.*recall_score(y_true, y_pred, average="weighted"):.2f}, F1: {100.*f1_score(y_true, y_pred, average="weighted"):.2f}')
      classes = ('1920', '1940', '1950', '1970', '1980', '2000', '2010')          
      cf_matrix = confusion_matrix(y_true, y_pred)   
      df_cm = pd.DataFrame(cf_matrix / np.sum(cf_matrix, axis=1)[:, None], index = [i for i in classes], columns = [i for i in classes])
      plt.figure(figsize = (12,7))
      sn.heatmap(df_cm, annot=True) 
      plt.savefig('ConfusionMatrix1.png')          
      sumdiagonals = 0.
      for i in range(len(cf_matrix)):
        sumdiagonals += (cf_matrix / np.sum(cf_matrix, axis=1)[:, None])[i,i]
      sumdiagonals /= len(cf_matrix) 
      print(f'Mean of diagonal items of Confusion Matrix: {100.*sumdiagonals:.2f}') 
      return accuracy     

validate(model)          
#test(model)