Image Segmentation - Preparation

# basics.py 
import pandas as pd, numpy as np 
import torch, torchvision, PIL.Image, cv2 
import os, sys
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T
import torch.nn.functional as F
from torchsummary import summary
import time
from tqdm import tqdm

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]

# datautils.py 
from basics import * 
class DroneDataset(Dataset):
    def __init__(self, img_path, mask_path, X, test=False):
        self.img_path = img_path
        self.mask_path = mask_path
        self.X = X
        self.test = test
    def __len__(self):
        return len(self.X)

    def __getitem__(self, idx):
        img = cv2.imread(self.img_path + self.X[idx] + '.jpg')
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, dsize=(600, 400), interpolation=cv2.INTER_NEAREST)
        mask = cv2.imread(self.mask_path + self.X[idx] + '.png', cv2.IMREAD_GRAYSCALE)
        mask = cv2.resize(mask, dsize=(600, 400), interpolation=cv2.INTER_NEAREST)

        t = T.Compose([T.ToTensor(), T.Normalize(mean, std)])
        img = t(img)
        mask = torch.from_numpy(mask).long()

        return img, mask

# datautils.py
from sklearn.model_selection import train_test_split
def import_drone_dataset():
    IMAGE_PATH = "../dataset/semantic_drone_dataset/original_images/"
    MASK_PATH = "../dataset/semantic_drone_dataset/label_images_semantic/"
    name = []
    for dirname, _, filenames in os.walk(IMAGE_PATH):
        for filename in filenames:
            name.append(filename.split('.')[0])
    df = pd.DataFrame({'id': name}, index = np.arange(0, len(name)))
    X_train, X_test = train_test_split(df['id'].values, test_size=0.1, random_state=0)
    train_set = DroneDataset(IMAGE_PATH, MASK_PATH, X_train, test=False)
    test_set = DroneDataset(IMAGE_PATH, MASK_PATH, X_test, test=True)
    return train_set, test_set

# metrics.py 
def pixel_accuracy(output, mask):
    with torch.no_grad():
        output = torch.argmax(output, dim=1)
        correct = torch.eq(output, mask).int()
        accuracy = float(correct.sum()) / float(correct.numel())
    return accuracy

# evaluate.py 
from basics import *

class ModelEvaluation():
    def __init__(self, model, test_dataset, metric):
        self.model = model
        self.test_dataset = test_dataset
        self.metric = metric
    
    def evaluate_single(self, image, mask):
        self.model.eval()
        self.model.to(device)
        image = image.to(device)
        mask = mask.to(device)
        with torch.no_grad():
            image = image.unsqueeze(0)
            mask = mask.unsqueeze(0)
            output = self.model(image)
            acc = self.metric(output, mask)
            masked = torch.argmax(output, dim=1)
            masked = masked.cpu().squeeze(0)
        return masked, acc

    def evaluate_all(self):
        accuracy = [] 
        for i in tqdm(range(len(self.test_dataset))):
            img, mask = self.test_dataset[i]
            pred, acc = self.evaluate_single(img, mask)
            accuracy.append(acc)
        print(f"Mean accruacy = {np.mean(accuracy)}")
        return accuracy

    def show_qualitative(self, ind):
        image, mask = self.test_dataset[ind]
        pred_mask, score = self.evaluate_single(image, mask)
        inv_normalize = T.Normalize(
            mean=[-0.485/0.229, -0.456/0.224, -0.406/0.225],
            std=[1/0.229, 1/0.224, 1/0.225]
        )
        image = inv_normalize(image)
        image = image.cpu().numpy()
        image = image.swapaxes(0, 1)
        image = image.swapaxes(1, 2)
        import matplotlib.pyplot as plt
        fig, (ax1, ax2, ax3) = plt.subplots(1,3, figsize=(20,10))
        ax1.imshow(image)
        ax1.set_title('Picture')
        ax2.imshow(mask)
        ax2.set_title('Ground truth')
        ax2.set_axis_off()
        ax3.imshow(pred_mask)
        ax3.set_title(f'Model | score {score:.3f}')
        ax3.set_axis_off()
        plt.show()