pytorch学习（一）数据加载之前的预处理（UCSD数据集）

最近在做有关视频异常检测方面的实验，需要用到UCSD数据集，pytorch自定义加载自己的数据集时需要将自己的数据的路径以及标签存放到txt文档中，方便后续的数据加载。
最后我会给出生成好的UCSD数据集txt文档下载链接，有需要的可以自行下载。

1、UCSD异常检测数据集

UCSD数据集是通过固定摄像头俯瞰人行道拍摄的，它有两个子文件夹，分别是UCSDped1和UCSDped2，每个文件夹下都有剪辑好的视频片段，每个视频片段大约在200帧左右。所有的训练样本都是正常样本，每一个测试片段都会包含某些异常帧，测试集中的异常帧的帧数区间在.m文件中有相应的标注。

• UCSDped1是行人远离和走进相机的场景的剪辑片段，训练集有34个视频片段，测试集有36个视频片段
• UCSDped2是行人行走路径平行于相机场景的视频剪辑片段，训练集有16个视频片段，测试集有12个视频片段

下面就将UCSD数据集中的图片信息以及标签信息（测试集中的异常样本）读取到4个txt文档中。

2、记录UCSD数据集图片路径以及标签信息

将数据信息记录到4个txt文档中，分别是UCSDped1_train.txt， UCSDped1_test.txt， UCSDped2_train.txt和UCSDped2_test.txt，其中正常样本标签为0，异常样本标签设置为1，设置标签是为了模型进行测试时方便计算出一些模型评估指标。。

import os
import re
"""
purpose:将UCSD数据集中图片路径以及标签写入txt文件，正常样本为0， 异常样本为1，为pytorch的数据加载做准备
author: Cui
time: 2020.5.31
"""


# UCSDped1路径和UCSDped2路径
train_path1 = r'G:\UCSD\UCSD_Anomaly_Dataset\UCSDped1\Train'
test_path1 = r'G:\UCSD\UCSD_Anomaly_Dataset\UCSDped1\Test'
train_path2 = r'G:\UCSD\UCSD_Anomaly_Dataset\UCSDped2\Train'
test_path2 = r'G:\UCSD\UCSD_Anomaly_Dataset\UCSDped2\Test'

# 写入的txt文件
ped1_train = open('./UCSDped1_train.txt', 'w')
ped1_test = open('./UCSDped1_test.txt', 'w')
ped2_train = open('./UCSDped2_train.txt', 'w')
ped2_test = open('./UCSDped2_test.txt', 'w')

paths = [train_path1, test_path1, train_path2, test_path2]
txt_files = [ped1_train, ped1_test, ped2_train, ped2_test]

label1_path = r'G:\UCSD\UCSD_Anomaly_Dataset\UCSDped1\Test\label.txt'
label2_path = r'G:\UCSD\UCSD_Anomaly_Dataset\UCSDped2\Test\label.txt'


# 将测试集中的异常样本帧帧数提取出来
def get_label(label_path):
    file = open(label_path)
    labels = []
    for f in file:
        f = f.strip(';\n')
        f = f.rstrip('[]')
        f_str = re.split(r',|:', f[33:])
        f_int = list(map(int, f_str))
        labels.append(f_int)
    file.close()
    return labels


# 将图片路径和标签写入文件
def txt_write(path, txt, label=None):
    lists = os.listdir(path)
    print(lists)
    if label:  # 测试集
        count = 0
        for i in lists:
            if (('Train' in i) or ('Test' in i)) and ('gt' not in i):
                files = os.listdir(path + '\\' + i)
                c = 1
                for file in files:
                    if 'tif' in file:
                        if len(label[count]) == 2:
                            if c in list(range(label[count][0], label[count][1] + 1)):
                                name = path + '\\' + i + '\\' + file + ' ' + str(int(1)) + '\n'
                                c += 1
                            else:
                                name = path + '\\' + i + '\\' + file + ' ' + str(int(0)) + '\n'
                                c += 1
                            txt.write(name)
                        else:
                            print(label[count])
                            if (c in list(range(label[count][0], label[count][1] + 1))) or (c in list(
                                    range(label[count][2], label[count][3] + 1))):
                                name = path + '\\' + i + '\\' + file + ' ' + str(int(1)) + '\n'
                                c += 1
                            else:
                                name = path + '\\' + i + '\\' + file + ' ' + str(int(0)) + '\n'
                                c += 1
                            txt.write(name)
                count += 1
    else:  # 训练集
        for i in lists:
            if (('Train' in i) or ('Test' in i)) and ('gt' not in i):
                files = os.listdir(path + '\\' + i)
                for file in files:
                    if 'tif' in file:
                        name = path + '\\' + i + '\\' + file + ' ' + str(int(0)) + '\n'
                        txt.write(name)
    txt.close()


label1 = get_label(label1_path)
label2 = get_label(label2_path)

print('starting')
for index, path in enumerate(paths):
    if index == 1:
        txt_write(path, txt_files[index], label=label1)
    elif index == 3:
        txt_write(path, txt_files[index], label=label2)
    else:
        txt_write(path, txt_files[index])
print('finished')

生成好的UCSDped1_train.txt文档：

UCSDped1_test.txt文档：

——————————————————————————
更新：
因为原来的txt文档中保存的是绝对路径，所以不方便大家的使用，我稍微修改了一下，保存了图片的相对路径，这样可以方便大家的调用。
链接：https://pan.baidu.com/s/19ZJQXT-NrWZe0w5ZnHssVg
提取码：itow