Preprocessing RGB image Masks to Segmentation Masks

Preprocessing RGB image Masks to Segmentation Masks

Hi, with this kernel, I will transfer your RGB color image masks to pre-process and make it ready for the semantic segmentation model.

The first step in training our segmentation model is to prepare the dataset. We would need the input RGB images and the corresponding segmentation images. If you want to make your own dataset, a tool like labelme or GIMP can be used to manually generate the ground truth segmentation masks. Assign each class a unique ID. In the segmentation images, the pixel value should denote the class ID of the corresponding pixel. You can examine for Deep Learning based Semantic Segmentation example

If you have RGB color image masks (like Aerial Semantic Segmentation Drone Dataset RGB_color_masks Folder ) , you can follow the steps below. First of all, our goal is to obtain a pixel-based class image from RGB color image masks as follows.RGB color image maskspixel-based class image masks

The dataset has RGB colored segmentation masks, and there are 24 classes in the dataset. The dataset has a file named class_dict.csv and it contains the name of the classes and their respective RGB values. The contents of the file are as follows:

namergb
unlabeled000
paved-area12864128
dirt130760
grass01020
gravel11210387
water2842168
rocks484130
pool05089
vegetation10714235
roof707070
wall102102156
window25422812
door25414812
fence190153153
fence-pole153153153
person2552296
dog102510
car9143150
bicycle1191132
tree51510
bald-tree190250190
ar-marker112150146
obstacle2135115
conflicting25500
import os,re
from PIL import Image
import numpy as np
import pandas as pd
import cv2
import matplotlib.pyplot as plt
%matplotlib inline

def atoi(text) : 
    return int(text) if text.isdigit() else text

def natural_keys(text) :
    return [atoi(c) for c in re.split('(\d+)', text)]


read_csv = pd.read_csv('/kaggle/input/semantic-drone-dataset/class_dict_seg.csv',index_col=False,skipinitialspace=True)
read_csv.head()

def segmantation_images(path,new_path, debug_test_num):
    filenames = []

    for root, dirnames, filenames in os.walk(path):
        filenames.sort(key = natural_keys)
        rootpath = root

    #print(filenames)
    count = 0
    for item in filenames:

        if debug_test_num !=0:
            if debug_test_num <= count:
                break

        count = count + 1

        if os.path.isfile(path+item):
            f, e = os.path.splitext(item)
            image_rgb = Image.open(path+item)
            image_rgb = np.asarray(image_rgb)
            new_image = np.zeros((image_rgb.shape[0],image_rgb.shape[1],3)).astype('int')

            for index, row  in read_csv.iterrows():
                new_image[(image_rgb[:,:,0]==row.r)&
                          (image_rgb[:,:,1]==row.g)&
                          (image_rgb[:,:,2]==row.b)]=np.array([index+1,index+1,index+1]).reshape(1,3)

            new_image = new_image[:,:,0]
            output_filename = new_path+f+'.png'
            cv2.imwrite(output_filename,new_image)
            print('writing file: ',output_filename)

        else:
            print('no file')

    print("number of files written: ",count)
debug_test_num = 5 # 5 samples are enough just to see it working. If you set 0, you can do all the datasets.
segmantation_images("/kaggle/input/semantic-drone-dataset/RGB_color_image_masks/RGB_color_image_masks/", "", debug_test_num=debug_test_num)
from PIL import Image
input_image = "/kaggle/input/semantic-drone-dataset/RGB_color_image_masks/RGB_color_image_masks/"+"002.png"
output_image = "002.png"

fig, axs = plt.subplots(1, 2, figsize=(16, 8), constrained_layout=True)

img_orig = Image.open(input_image)
axs[0].imshow(img_orig)
axs[0].grid(False)

output_image = Image.open(output_image)
out_array = np.asarray(output_image)
axs[1].imshow(out_array)
axs[1].grid(False)