Python获取标签

掌纹数据集A每类前三张作为训练集，数据集B所有样本作为测试集，Python获取训练集和测试集标签，以数组的形式保存为.npy文件

import os
from pathlib import Path
import numpy as np

#-------------------获取label---------------------
def get_train_label():
    train_lab = []
    class_num = 202    # 类别总数
    per_num = 3    # 训练集数目
    for i in range(1, class_num + 1):
        for j in range(0, per_num):
            train_lab.append(i)
    train_lab = np.array(train_lab).astype('int64')
    np.save('data/train_lab.npy', train_lab)
    return train_lab


def get_test_label(path):    #
    test_lab = []
    path_list = os.listdir(path)
    path_list.sort(key=lambda x: int(x.split('_')[0]))    # 按类名排序
    k = 0
    for img_path in path_list:
        img_path = Path(os.path.join(path, img_path))    # 字符串转化为路径
        # print(img_path)
        _class = img_path.stem.split('_')[0]
        _num = img_path.stem.split('_')[1]
        test_lab.append(_class)

    test_lab = np.array(test_lab).astype('int64')
    np.save('data/test_lab.npy', test_lab)
    return test_lab