divamgupta / image-segmentation-keras

Implementation of Segnet, FCN, UNet , PSPNet and other models in Keras.
https://divamgupta.com/image-segmentation/2019/06/06/deep-learning-semantic-segmentation-keras.html
MIT License
2.91k stars 1.16k forks source link

Mapping VOC label to this repo label style #257

Open bobvo23 opened 3 years ago

bobvo23 commented 3 years ago

Hi, this repository extracts the Blue pixel value to indicate the corresponding class of the pixel. However, the VOC dataset uses a combination of RGB to indicate the correspond class. I'm trying to convert the VOC Segmentation label to this style using the following code. Do we have any other option?

from PIL import Image
#import numpy as np
#from mxnet import gluon, np, image, npx
import cv2 as image
import numpy
import glob
import sys
import numpy as np
np.set_printoptions(threshold=sys.maxsize)

#color mapping
VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]

#@save
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
               'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
               'diningtable', 'dog', 'horse', 'motorbike', 'person',
               'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']

#@save
def build_colormap2label():
    """Build an RGB color to label mapping for segmentation."""
    colormap2label = np.zeros(256 ** 3)
    for i, colormap in enumerate(VOC_COLORMAP):
        colormap2label[(colormap[0]*256 + colormap[1])*256 + colormap[2]] = i
    return colormap2label

#@save
def voc_label_indices(folder_dir, colormap2label):
    explore_path = folder_dir + "*.png"
    path_list = glob.glob(explore_path)
    print(explore_path)
    print(len(path_list))
    for file_dir in path_list:
      #Convert every file in the pathlist to new format
      """Map an RGB color to a label."""
      #colormap is the path to the file
      colormap = image.imread(file_dir)
      colormap = colormap.astype(np.int32)
      idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
            + colormap[:, :, 2])
      output_array = colormap2label[idx]
      #converting the 2d output array to 3d image
      output_image = np.dstack((output_array,output_array,output_array))
      output_image = output_image.astype(np.uint8)
      im = Image.fromarray(output_image)
      os.remove(file_dir)
      im.save(file_dir)
      #print("Done")
boguss1225 commented 3 years ago

For those who the upper code doesn't work. `from PIL import Image

import numpy as np

from mxnet import gluon, np, image, npx

import cv2 import numpy import glob import sys import numpy as np import os np.set_printoptions(threshold=sys.maxsize)

path = '/Users/bogus/Downloads/SegmentationClass/' # Source Folder dstpath = '/Users/bogus/Downloads/SegmentationClass2/' # Destination Folder

color mapping

VOC_COLORMAP = [[176, 112, 32], [48, 240, 32], [48, 112, 32], [176, 240, 160], [176, 240, 32], [48, 240, 160], [176, 112, 160]]

@save

VOC_CLASSES = ['background', 'hyperthalamus', 'cerebral', 'thalamus', 'hippo', 'soma', 'audi']

@convert & save

def convert_to_iii(file_dir): """Build an RGB color to label mapping for segmentation.""" im=cv2.imread(file_dir) print("im",im[1,1]) for i, colormap in enumerate(VOC_COLORMAP):

print(i,".",VOC_CLASSES[i],":",colormap[0],"-",colormap[1],"-",colormap[2])

    im[np.where((im == [colormap[2],colormap[1],colormap[0]]).all(axis = 2))] = [i,i,i]
filename = os.path.split(file_dir)[-1]
cv2.imwrite((os.path.join(dstpath,filename)),im)

@iterate folder

def voc_label_indices(folder_dir): explore_path = folder_dir + "/*.png" path_list = glob.glob(explore_path) print(explore_path) print("len",len(path_list)) for file_dir in path_list: print(file_dir) convert_to_iii(file_dir)

copy_bool = input("Copy target directory before excute. Excute? (Y/N):")

if copy_bool=='y' or copy_bool=='Y': print("Start conversion") voc_label_indices(path) print("Done") else : print("See you") `