DeeplabV3 custom dataset, inference problem black images

Question

Good morning,

I want to train a custom dataset using deeplabV3.

I'm following this tutorial (https://sanjayparajuli27.medium.com/how-to-train-deeplab-on-custom-dataset-a40c41c4c6a3) for this dataset (https://www.kaggle.com/datasets/dansbecker/cityscapes-image-pairs) that I found on Kaggle, based on cityscapes.

There are images of 256x256 pixel in RGB colors, divided in 2975 imgs for training and 500 for validation, and I created the respective mask using this script

import tensorflow as tf
from PIL import Image
from tqdm import tqdm
import numpy as np

import os, shutil

# palette (color map) describes the (R, G, B): Label pair
palette = {(0,   0,   0) : 0 ,
           (128,  0, 0) : 1,
           (0, 128, 0): 2,
           (128, 128, 0): 3,
           (0, 0, 128): 4,
           (0, 128, 128): 5,
           (128, 128, 128): 6,
           (64, 0, 0): 7,
           (192, 0, 0): 8,
           (64, 128, 0): 9,
           (192, 128, 0): 10,
           (64, 0, 128): 11,
           (192, 0, 128): 12,
           (64, 128, 128): 13,
           (0, 64, 0): 14,
           (128, 64, 0): 15,
           (0, 192, 0): 16,
           (128, 192, 0): 17,
           (0, 64, 128): 18,
           (128, 0, 128): 19
         }

def convert_from_color_segmentation(arr_3d):
    arr_2d = np.zeros((arr_3d.shape[0], arr_3d.shape[1]), dtype=np.uint8)

    for c, i in palette.items():
        m = np.all(arr_3d == np.array(c).reshape(1, 1, 3), axis=2)
        arr_2d[m] = i
    return arr_2d


label_dir = "C:/Users/paolo.david/Desktop/Datasets/final/Validation/Mask_real/" #don't forget the '/' at the end
new_label_dir = "C:/Users/paolo.david/Desktop/Datasets/final/Validation/Mask_RAW2/"

if not os.path.isdir(new_label_dir):
    print("creating folder: ",new_label_dir)
    os.mkdir(new_label_dir)
else:
    print("Folder alread exists. Delete the folder and re-run the code!!!")


label_files = os.listdir(label_dir)

for l_f in tqdm(label_files):
    #arr = np.array(Image.open(l_f))
    arr = np.array(Image.open(label_dir + l_f))
    arr = arr[:,:,0:3]
    arr_2d = convert_from_color_segmentation(arr)
    #Image.fromarray(arr_2d).save(label_dir)
    Image.fromarray(arr_2d).save(new_label_dir + l_f)

Each image in the dataset contain its same mask, so before to launch the new notebook I divided the image and the mask to have a situation like in the tutorial.

You can find my code here: https://drive.google.com/drive/folders/105JMDmujY6lknH3D74WM8R8S7jTb51qX?usp=sharing and this is the notebook https://drive.google.com/file/d/1xmUtLB-XPj4mZdqbx9SAQOXKSwxtxCLX/view?usp=sharing

I have a problem with the inference. Every time I launch the notebook with few epochs (less then 10) I receive good results, but trying to increase the number of epoch I have all black images.

These are the parameters that I used for the train:

--model_variant="xception_65" \
--atrous_rates=6 \
--atrous_rates=12 \
--atrous_rates=18 \
--output_stride=16 \
--decoder_output_stride=4 \
--train_crop_size="256,256" \
--train_batch_size=4 \
--training_number_of_steps=30 \
--initialize_last_layer=False \
--last_layers_contain_logits_only=True \
--fine_tune_batch_norm=False \

I edited the data_generator.py file putting

_CUSTOM_INFORMATION = DatasetDescriptor(
    splits_to_sizes={
        'train': 270,  # num of samples in train.txt
        'val': 30,  # num of samples in val.txt
    },
    num_classes=21, # classes+bg+ignore_label
    ignore_label=255,
)

_DATASETS_INFORMATION = {
    'cityscapes': _CITYSCAPES_INFORMATION,
    'pascal_voc_seg': _PASCAL_VOC_SEG_INFORMATION,
    'ade20k': _ADE20K_INFORMATION,
    'custom': _CUSTOM_INFORMATION  # custom dataset
}

I'm using a pretrained model downloaded from here: http://download.tensorflow.org/models/deeplabv3_cityscapes_train_2018_02_06.tar.gz

I keep the batch size at 4, and I don't know if it is correct or not. Can you tell me where could be the possible error?