[FIXED] Count number of classes in a semantic segmented image

Issue

I have an image that is the output of a semantic segmentation algorithm, for example this one

I looked online and tried many pieces of code but none worked for me so far.

It is clear to the human eye that there are 5 different colors in this image: blue, black, red, and white.

I am trying to write a script in python to analyze the image and return the number of colors present in the image but so far it is not working. There are many pixels in the image which contain values that are a mixture of the colors above.

The code I am using is the following but I would like to understand if there is an easier way in your opinion to achieve this goal.

I think that I need to implement some sort of thresholding that has the following logic:

• Is there a similar color to this one? if yes, do not increase the count of colors
• Is this color present for more than N pixels? If not, do not increase the count of colors.

from PIL import Image
 
imgPath = "image.jpg"
 
img = Image.open(imgPath)
uniqueColors = set()
 
w, h = img.size
for x in range(w):
    for y in range(h):
        pixel = img.getpixel((x, y))
        uniqueColors.add(pixel)
 
totalUniqueColors = len(uniqueColors)

print(totalUniqueColors)

print(uniqueColors)

Thanks in advance!

Solution

I solved my issue and I am now able to count colors in images coming from a semantic segmentation dataset (the images must be in .png since it is a lossless format).

Below I try to explain what I have found in the process for a solution and the code I used which should be ready to use (you need to just change the path to the images you want to analyze).

I had two main problems.

The first problem of the color counting was the format of the image. I was using (for some of the tests) .jpeg images that compress the image.

Therefore from something like this

If I would zoom in the top left corner of the glass (marked in green) I was seeing something like this

Which obviously is not good since it will introduce many more colors than the ones "visible to the human eye"

Instead, for my annotated images I had something like the following

If I zoom in the saddle of the bike (marked in green) I had something like this

The second problem was that I did not convert my image into an RGB image.

This is taken care in the code from the line:

img = Image.open(filename).convert('RGB')

The code is below. For sure it is not the most efficient but for me it does the job. Any suggestion to improve its performance is appreciated

import numpy as np
from PIL import Image
import argparse
import os

debug = False

def main(data_dir):
    print("This small script allows you to count the number of different colors in an image")
    print("This code has been written to count the number of classes in images from a semantic segmentation dataset")
    print("Therefore, it is highly recommended to run this code on lossless images (such as .png ones)")
    print("Images are being loaded from: {}".format(data_dir))

    directory = os.fsencode(data_dir)
    interesting_image_format = ".png"
    
    # I will put in the variable filenames all the paths to the images to be analyzed
    filenames = []
    for file in os.listdir(directory):
        filename = os.fsdecode(file)
        if filename.endswith(interesting_image_format): 
            if debug:
                print(os.path.join(directory, filename))
                print("Analyzing image: {}".format(filename))
            filenames.append(os.path.join(data_dir, filename))
        else:
            if debug:
                print("I am not doing much here...")
            continue
    # Sort the filenames in an alphabetical order
    filenames.sort()

    # Analyze the images (i.e., count the different number of colors in the images)
    number_of_colors_in_images = []
    for filename in filenames:
        img = Image.open(filename).convert('RGB')
        if debug: 
            print(img.format)
            print(img.size)
            print(img.mode)
        data_img = np.asarray(img)
        if debug: 
            print(data_img.shape)
        uniques = np.unique(data_img.reshape(-1, data_img.shape[-1]), axis=0)
        # uncomment the following line if you want information for each analyzed image  
        print("The number of different colors in image ({}) {} is: {}".format(interesting_image_format, filename, len(uniques)))
        # print("uniques.shape[0] for image {} is: {}".format(filename, uniques.shape[0]))
        
        # Put the number of colors of each image into an array
        number_of_colors_in_images.append(len(uniques))
    
    print(number_of_colors_in_images)
    # Print the maximum number of colors (classes) of all the analyzed images
    print(np.max(number_of_colors_in_images))
    # Print the average number of colors (classes) of all the analyzed images
    print(np.average(number_of_colors_in_images))

def args_preprocess():
    # Command line arguments
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data_dir", default="default_path_to_images", type=str, help='Specify the directory path from where to take the images of which we want to count the classes')
    args = parser.parse_args()
    main(args.data_dir)

if __name__ == '__main__':
    args_preprocess()

Answered By - desmond13