I’m trying to build a UNet Segmentor for an aerial image of vegetation w/ pytorch. The UNet then gets exported to an ONNX/GIS_ONNX to be usable in the Deepness-Plugin of QGIS. The Unet-Basecode and the training script are mostly derived/copied from Mostafa Wael on toward data science. The Dataloader is a Custom-Build upon the torch.utils.dataloader class. Training data is fed in tile-pairs of cut ortho mosaic and corresponding class mask (512×512).

My Problem is, that Deepness only identifies 100% background class (and no other of the actual classes), despite being used on the training data.

import torch 
import os
import numpy as np 
import glob
import matplotlib.pyplot as plt
import cv2
import math
import random

from torchvision.transforms import transforms
from torchvision.transforms import functional
from torch.utils.data import Dataset, DataLoader, Subset


def largest_rotated_rect(w, h, angle):
    """
    Given a rectangle of size wxh that has been rotated by 'angle' (in
    radians), computes the width and height of the largest possible
    axis-aligned rectangle within the rotated rectangle.

    Original JS code by 'Andri' and Magnus Hoff from Stack Overflow

    Converted to Python by Aaron Snoswell
    """

    quadrant = int(math.floor(angle / (math.pi / 2))) & 3
    sign_alpha = angle if ((quadrant & 1) == 0) else math.pi - angle
    alpha = (sign_alpha % math.pi + math.pi) % math.pi

    bb_w = w * math.cos(alpha) + h * math.sin(alpha)
    bb_h = w * math.sin(alpha) + h * math.cos(alpha)

    gamma = math.atan2(bb_w, bb_w) if (w < h) else math.atan2(bb_w, bb_w)

    delta = math.pi - alpha - gamma

    length = h if (w < h) else w

    d = length * math.cos(alpha)
    a = d * math.sin(alpha) / math.sin(delta)

    y = a * math.cos(gamma)
    x = y * math.tan(gamma)

    return (
        bb_w - 2 * x,
        bb_h - 2 * y
    )

      
def pickTransform(transform_choices): # list must be in torch syntax e.g. transforms.Random... // add metadata img size, so largest_rect is dynamic

    check = np.random.choice([True, False], size=len(transform_choices))
      
    selected_transforms = np.array(transform_choices)[check].tolist()  
    
    if any(isinstance(i, transforms.RandomCrop) for i in selected_transforms):
        selected_transforms.insert(1,transforms.Pad(100, padding_mode="reflect"))
        
        print(selected_transforms)
    return selected_transforms
    
geo_trans = [
    transforms.RandomCrop(400),
    transforms.RandomHorizontalFlip(p=1),
    transforms.RandomVerticalFlip(p=1),
    ]

opt_trans = [
      transforms.ColorJitter(0.05, 0, 0, 0),
      transforms.ColorJitter(0, 0.05, 0, 0),
      transforms.ColorJitter(0, 0, 0.05, 0),
      transforms.ColorJitter(0, 0, 0, 0.05) 
    ]

def rotateCrop(Tensor):
    
    if np.random.choice([True, False]): # rotation is seperated from other transforms, so the actual angle (when object is called) gets extracted for largest rect 
        
        w = Tensor.size(1) #aks Chris, if the x/y Dims are always there
        h = Tensor.size(2) 
        
        angle = float(np.random.choice(range(-45,45,1)))
        print(angle)
        
        crop_x, crop_y = largest_rotated_rect(w, h, angle)
        print(crop_x)
        
        functional.rotate(Tensor, angle)
        
        rect_crop =  transforms.Compose([
            (transforms.CenterCrop(crop_x)),
            (transforms.Resize(w, interpolation=transforms.InterpolationMode.NEAREST))
            ])
        
        TensorRotCrop = rect_crop(Tensor)
        
        return TensorRotCrop
    
    else: 
        return Tensor

def onehot(Tensor):
    
    Tensor = Tensor.long()
    Tensor = torch.nn.functional.one_hot(Tensor,9)
    Tensor = Tensor.squeeze(0)
    Tensor = Tensor.permute(2,0,1)
    
    return Tensor

class customDataset(Dataset):
    def __init__(self, dir_data, transformation=None, train_percent=0.5): #via transformation = None, Val_Data can be created
        super(customDataset,self).__init__()
        self.dir_data = dir_data
        self.img_files = glob.glob(os.path.join(dir_data,'Images','*.*'))
        self.label_files = glob.glob(os.path.join(dir_data,'Labels','*.*'))
        self.transformation = transformation
        
        #Prepare Indices for Training and Validation
        self.indices = list(range(len(self.img_files)))
        random.shuffle(self.indices) 
        num_train = int(train_percent*len(self.indices)) #set threshold
        self.train_indices = self.indices[:num_train] #later on the Training and Val-Data can be acessed via self.XX_indices
        self.val_indices = self.indices[num_train:]
        
    def __len__(self):
        return len(self.img_files)
        
    def __getitem__(self, index): 
        img_path = self.img_files[index]
        label_path = self.label_files[index]
        
        image = np.array(cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
        label = np.array(cv2.imread(label_path, cv2.IMREAD_GRAYSCALE))
        
        
        image_tensor = torch.from_numpy(image).permute(2, 0, 1) #schiebt die Kanäle (2) nach vorne (für torch-syntax Kanäle, Höhe, Breite), muss angepasst werden wenn bilder mit mehr als 3 Kanälen verwendet werdern// sollte Input zw. 0 und 1 vorausgesetzt werdern --> .float() / 255.0
        label_tensor = torch.from_numpy(label).unsqueeze(0)     # add an empty dimension so torch.cat can stack along that dimension
        
        if index in self.train_indices: 
            
            if self.transformation == "opt" or self.transformation == "both":
                
                    transform_opt = transforms.Compose(pickTransform(opt_trans))
                
                    image_tensor = transform_opt(image_tensor)
                    
                    # print(image_tensor.size(), label_tensor.size())
                    
                    image_tensor = image_tensor.float()
                    
                    label_tensor = label_tensor.float()
                    label_tensor = onehot(label_tensor)
                    
                    return image_tensor, label_tensor
                    
            if self.transformation == "geo" or self.transformation == "both":
                
                    stack_tensor = torch.cat((image_tensor, label_tensor)) #concatenate img and mask along dim 0, so transformations affect both 
                
                    transform_geo = transforms.Compose(pickTransform(geo_trans))
                    
                    transformed_stack = transform_geo(stack_tensor)
                    
                    transformed_stack = rotateCrop(transformed_stack)
                    
                    image_tensor = transformed_stack[0:-1] 
                    label_tensor = transformed_stack[-1] #re-adds empty dimension, after Slicing removes it
                    
                    # print(image_tensor.size(), label_tensor.size())
                    
                    image_tensor = image_tensor.float()
                    
                    label_tensor = label_tensor.float()
                    label_tensor = onehot(label_tensor)
                    
                    return image_tensor, label_tensor
            
        else: 
            
            image_tensor = image_tensor.float()
            
            label_tensor = label_tensor.float()
            label_tensor = onehot(label_tensor)
            
            return image_tensor, label_tensor
           
        

I ran between 5-15 Epoch on a high performance cluster. I also switched the optimizer and scheduler step around compared to Mostafa Wael as I tought, that it maybe hindered learning and because I got a torch warning.

def train_model(model, optimizer, scheduler, num_epochs=10):
    dataloaders = get_data_loaders()
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    best_model_wts = copy.deepcopy(model.state_dict())
    best_loss = 1e10

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)

        since = time.time()

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            if phase == 'train':
                
                for param_group in optimizer.param_groups:
                    print("LR", param_group['lr'])

                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            metrics = defaultdict(float)
            epoch_samples = 0

            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    loss = calc_loss(outputs, labels, metrics)

                    # backward + optimize only if in training phase
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()
                        scheduler.step()

                # statistics
                epoch_samples += inputs.size(0)

            print_metrics(metrics, epoch_samples, phase)
            epoch_loss = metrics['loss'] / epoch_samples

            # deep copy the model
            if phase == 'val' and epoch_loss < best_loss:
                print("saving best model")
                best_loss = epoch_loss
                best_model_wts = copy.deepcopy(model.state_dict())

        time_elapsed = time.time() - since
        print('{:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

    print('Best val loss: {:4f}'.format(best_loss))

    # load best model weights
    model.load_state_dict(best_model_wts)

does it just need more training or better training data (the initial polygon tracing was pretty sloppy)? thank you in advance.