图像分类每个标签按比例划分数据

问题

在做图像分类时候，会收集一批相应的数据，这里将其称为总数据集total-data，

按照一般的做法，会将总数据集划分为训练集（train-data）、验证集（valid-data）以及测试集（test-data）。这里为了方便，将总的数据集划分为训练集和验证集。

有时候使用随机分配的算法会导致每个标签下样本的个数分布不是很均匀，有的标签下样本个数很多，有的标签下样本个数很少，这就导致了一种数据不均衡问题，使得训练的模型偏向于数据样本多的标签。那么我们能不能按照相应的比例，也将每一个标签下的数据按照对应的比例进行划分呢？这其实也是比较好实现的。

解决办法

假设项目目录为dataset，下面有total-data文件夹、train-data以及test-data文件夹

total-data目录下存放的是所有的图像数据集，图像命名样式为label_xxxx.jpg

• 1.获取所有的图像样本名称：

  ls total-data/ > total_data.txt
  

• 2.按照比例将total_data.txt划分为train_data.txt以及test_data.txt:

  python split_train_test.py total_data.txt
  

  其中，split_train_test.py样式如下：

  import sys
  import random
  
  
  total_img_file = sys.argv[1]
  res = {}   # stastic label and the numbers
  ratio = 0.8  # the ratio of train and test is 8:2
  
  
  with open(total_img_file) as f:
  	for ln in f:
  		img_name = ln.rstrip('\n')
  		img_label = img_name.split('_')[0]
  		if img_label in res:
  			res[img_label].append(img_name)
  		else:
  			res[img_label] = [img_name]
  
  
  train_data = open('train_data.txt', 'w')
  test_data = open('test_data.txt', 'w')
  
  
  for k, v in res.items():
  	total_nums = len(v)
  	train_nums = int(total_nums * ratio)
  	test_nums = total_nums - train_nums  # 每个标签下都按照比例划分样本数据
  	# random choice train_nums example from img_name list
  	train_imgs = random.sample(v, train_nums)
  	test_imgs = list(set(v) - set(train_imgs))  # the reverse
  	for train_img in train_imgs:
  		print(train_img, file=train_data)
  	for test_img in test_imgs:
  		print(test_img, file=test_data)
  
  
  train_data.close()
  test_data.close()
  

• 3.统计分析下每个标签下样本的数量

  python statistic.py

其中，statistic.py:

  import os
  import sys
  
  res = {}  # use dict to collect the nums of every labels
  
  img_file = sys.argv[1]
  with open(img_file) as f:
  	for ln in f:
  		img_name = ln.rstrip('\n')
  		img_label = img_name.split('_')[0]
  		if img_label in res:
  			res[img_label] += 1
  		else:
  			res[img_label] = 1
  
  print(res)

• 从总数据集中复制对应的图像文件到训练集和测试集

  python cp_file.py test_data.txt

其中，cp_file.py：

  import sys
  import shutil
  
  img_files = sys.argv[1]
  
  with open(img_files) as f:
  	for ln in f:
  		img_file = ln.rstrip('\n')
  		print(img_file)
  		src_path = 'total-data/' + img_file
  		dst_path = 'test-data/' + img_file
  		shutil.copy(src_path, dst_path)

或者使用shell

cat train_data.txt | xargs -i cp total_data/{} train_data/

后续就是准备标签，以及定义模型及训练了，这里不做过多介绍。上述有些功能在一些深度学习框架中有实现，但自己实现一下可以更加清楚一下整个流程。