• 重识别算法基准测试
• • 重试别算法基准测试分类
  • 脚本样例
  • • 代码仓库目录结构
    • 配置文件 re_ID.yaml
    • 训练脚本 train.py
    • 测试脚本 test.py

重识别算法基准测试

重试别算法基准测试分类

Name	名称
vehicle reid	车辆重识别
image-based person reid	基于图像的行人重识别
person search	行人搜索
video-based person reid	基于跟踪序列的行人重识别

脚本样例

代码仓库目录结构

1. train.py 和 test.py 需要在主目录下。
2. 主目录下可以包含其他文件或子文件夹。
3. 需要存在一个config目录，并包含re_ID.yaml 文件，该文件内容在下一节说明。
4. 评测脚本中的子目录名称由bk_开头，为保证脚本正确运行，避免以bk_为开头命名文件或子目录。
     |--reID_alg_sample
      |--config
         |--re_ID.yaml
         |--XXX
      |--XXX
         |--XXXX.py
      |--train.py
      |--test.py

配置文件 re_ID.yaml

1. 需要包含重识别子类型。
2. 需要包含子类型对应的数据集名称。数据集均在/dataset 路径下。
3. 其他配置信息可以根据算法需求自定义。

配置文件务必保留以下部分

# scene: type of re-ID
# four types in total:
# 1. vehicle reid
# 2. image_based person reid
# 3. person search
# 4. video-based person reid
scene: 2 # for image_based person re_ID
attr_name: 
attr_value: 

input:
  dataset:
    data_dir: /dataset/benchmark_person_reID # for image_based person re_ID
    # data_dir: /dataset/benchmark_vehicle_reID # for vehicle reid

训练脚本 train.py

1. 需要包含train函数, train函数的输入包括配置文件和logger。
2. 函数中需要通过logger记录模型参数总数，计算复杂度，模型训练总周期数。
3. 训练完成的模型需要存储到指定文件夹。

训练脚本务必保留以下部分

#bk_utils, bk_config 模块由评测脚本提供
from bk_utils.model_complexity import compute_model_complexity
from bk_utils.util import save_network, get_stream_logger
from bk_config.mainconfig import OUTPUT_RESULT_DIR

from torchvision import datasets

def train(config_file_path: str, logger):

    #phrase yaml file
    with open(config_file_path, encoding='utf-8') as f:
        opts = yaml.load(f, Loader=yaml.SafeLoader)

    #记录数据集名称和re_ID子类型
    data_dir = opts['input']['dataset']['data_dir']
    data_name = data_dir.split('/')[-1]
    logger.info("dataset name: %s"%(data_name))

    '''
    解析其他配置文件中的变量
    '''

    '''
    加载模型，transform
    '''

    #加载数据集
    image_datasets = {}

    image_datasets['train'] =datasets.ImageFolder(os.path.join(data_dir, 'bounding_box_train'), data_transforms['train'])
    image_datasets['val'] =datasets.ImageFolder(os.path.join(data_dir, 'val'), data_transforms['val'])

    dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batchsize,
                                                shuffle=True, num_workers=8, pin_memory=True) # 8 workers may work faster
                                                for x in ['train', 'val']}

    dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}


    #计算模型参数总数和计算复杂度
    params, FLOPs = compute_model_complexity(model, (1,3,input_h,input_w), verbose=False, only_conv_linear=True)
    logger.info('number of params (M): %.2f'%(params/1e6))
    logger.info('FLOPs (G): %.2f'%(FLOPs/1e9))

    '''
    训练
    '''

    #存储最优模型
    model.load_state_dict(best_model_wts)
    save_network(model, save_path, name, 'last')

    # 记录训练时间和训练总周期数
    logger.info('total train time: %.2f minutes'%(time_elapsed / 60))
    logger.info('total train epochs: %d epochs'%num_epochs)

测试脚本 test.py

1. 需要包含test函数, test函数的输入包括配置文件和logger。
2. 存储query和gallery之间的特征的距离矩阵、特征样例、gallery车辆和摄像头ID、query车辆和摄像头ID到指定文件夹。
3. 函数中需要通过logger记录特征提取的时间。

测试脚本务必保留以下部分

import yaml
import pickle
from bk_config.mainconfig import OUTPUT_RESULT_DIR
from torchvision import datasets

def test(config_file_path:str, logger):

    #解析配置文件
    with open(config_file_path, encoding='utf-8') as f:
        opts = yaml.load(f, Loader=yaml.SafeLoader)

    data_dir = opts['input']['dataset']['data_dir']
    name = "trained_" + opts['input']['config']['name']
    trained_model_name = name + "_last.pth"
    save_path = OUTPUT_RESULT_DIR

    since = time.time()
    with torch.no_grad():
        '''
        特征提取
        '''
    process_time = time.time() - since
    logger.info('total forward time: %.2f minutes'%(process_time/60))

    return_dict = {}

    # 特征距离矩阵 dist.shape = (len(query), len(gallery))
    return_dict['dist'] = dist.numpy()
    # 特征样例
    return_dict['feature_example'] = query_feature[0].numpy()
    # gallery车辆和摄像头ID
    return_dict['gallery_label'] = gallery_label
    return_dict['gallery_cam'] = gallery_cam
    # query车辆和摄像头ID
    return_dict['query_label'] = query_label
    return_dict['query_cam'] = query_cam

    pickle.dump(return_dict, open(OUTPUT_RESULT_DIR+'test_result.pkl', 'wb'), protocol=4)