统计图像分割训练集中的类别分布

本系列文章由 @yhl_leo 出品，转载请注明出处。

文章链接： http://blog.csdn.net/yhl_leo/article/details/52225600

对于一个语义分割数据集，可以使用如下方法统计样本集ground truth的类别分布情况：

import cv2, os
import numpy as np

#amount of classer
CLASSES_NUM = 21

#find imagee in folder dir
def findImages(dir,topdown=True):
im_list = []
if not os.path.exists(dir):
print "Path for {} not exist!".format(dir)
raise
else:
for root, dirs, files in os.walk(dir, topdown):
for fl in files:
im_list.append(fl)
return im_list

# amount of images corresponding to each classes
images_count = [0]*CLASSES_NUM
# amount of pixels corresponding to each class
class_pixels_count = [0]*CLASSES_NUM
# amount of pixels corresponding to the images of each class
image_pixels_count = [0]*CLASSES_NUM

image_folder = '../data/gt'
im_list = findImages(image_folder)

for im in im_list:
print im
cv_img = cv2.imread(os.path.join(image_folder, im), cv2.IMREAD_UNCHANGED)
size_img = cv_img.shape
colors = set([])
for i in range(size_img[0]):
for j in range(size_img[1]):
p_value = cv_img.item(i,j)
if not p_value < CLASSES_NUM: # check
print p_value
else:
class_pixels_count[p_value] = class_pixels_count[p_value] + 1
colors.add(p_value)
im_size = size_img[0]*size_img[1]
for n in range(CLASSES_NUM):
if n in colors:
images_count
= images_count
+ 1
image_pixels_count
= image_pixels_count
+ im_size

print images_count
print class_pixels_count
print image_pixels_count

上述代码，主要统计了每一类别所包含的图像数量（

images_count

），每一类别的像素数目（

class_pixels_count

）和每一类别对应的图像的总像素数目（

image_pixels_count

），有了这三组统计结果，就可以进一步计算训练时每一类别的

loss

的

class_weight

。