NYUDv2数据集40class的Class Weight Balance

统计数据集label图中每一类对象的像素, 在总像素数量中所占比重, 以及计算的class weight

---

• 代码:

import numpy as np

from PIL import Image

import os
import math

classes = ['background', 'wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
           'window', 'bookshelf', 'picture', 'counter', 'blinds', 'desk', 'shelves', 'curtain',
           'dresser', 'pillow', 'mirror', 'floor mat', 'clothes', 'ceiling', 'books', 'refridgerator',
           'television', 'paper', 'towel', 'shower curtain', 'box', 'whiteboard', 'person', 'night stand',
           'toilet', 'sink', 'lamp', 'bathtub', 'bag', 'otherstructure', 'otherfurniture', 'otherprop']

# 每一类的像素数量
class_count = [0 for i in range(41)]
# 每一类的像素所占比例
class_count_percent = [0 for i in range(41)]

root = './dataset/NYUDv2/'
txt_fname = ('./dataset/NYUDv2/train.txt')
with open(txt_fname, 'r') as f:
    images = f.read().split()
label_list = [os.path.join(root, 'nyu_labels', i + '.png') for i in images]

for label_name in label_list:
    label = np.array(Image.open(label_name))
    label = label.flatten()
    for i in label:
        # 遍历label图计算每一类的像素数量
        class_count[int(i)] = class_count[int(i)] + 1
    print(label_name)

# 所有像素的总数
all_class_pixel_count = sum(class_count)

for i in range(41):
    # 计算每一类的像素所占比例
    class_count_percent[i] = class_count[i] / all_class_pixel_count

# 按比例进行排序,找出比例的中位数以及对应类
class_percent_sort = [i for i in class_count_percent]
class_percent_sort.sort()
median = 0
middle_count = 0
middle_class = 0
# for i in range(41):
#     middle_count += class_percent_sort[i]
#     if midd