如何平衡 numpy 数组中的类？

Question

我有 2 个 numpy 数组，如下所示：

images 包含图像文件的名称（images.shape 为 (N, 3, 128, 128)）： image_1.jpg image_2.jpg image_3.jpg image_4.jpg

labels 包含对应的标签 (0-3)（labels.shape 是 (N,)）： 1 1 3 2

我面临的问题是类不平衡，类 3 >> 1 > 2 > 0。

我想通过以下方式平衡最终数据集：

• 统计每个类中的图像（样本）数量
• 获取图像数量最少的类的计数
• 使用该计数作为其他 3 个类别的最大图像/标签数
• 随机弹出images和labels中其他3类的多余图片/标签

到目前为止，我使用Counter 来确定每个类的图像数量：

from Collections import Counter
import numpy as np

count = Counter(labels)
print(count)

>>>Counter({'1': 2991, '0': 2953, '2': 2510, '3': 2488})

您如何建议我从 images 和 labels 中随机弹出匹配元素，以便它们包含 0、1 和 2 类的 2488 个样本？

Answer 1

您可以使用np.random.choice 创建一个整数值掩码，您可以将其应用于标签和图像以平衡数据集：

n = 2488

mask = np.hstack([np.random.choice(np.where(labels == l)[0], n, replace=False)
                      for l in np.unique(labels)])