python pytorch 自制数据集实现图像二分类(包含数据处理过程详细代码)

这是一篇亲试可行的文章，记录了我在完成对自制数据集进行二分类的过程。我很开心自己的能感受到自己的进步～话不多说，直接开始吧。

一、数据处理

我的数据集是两个文件的图片，由于文件的名字十分复杂不太好用，所以我用python对其进行了批量重命名。

file_path = "/Users/***/Tooth-Detection/data/"  # 原文件路径
new_path = "/Users/***/Tooth-Detection/data/"   # 新文件路径
def rename(path):
    i = 0
    filelist = os.listdir(path)
    for files in filelist[:500]:   # 取文件夹中前500个文件
        oldDirpath = os.path.join(path, files)
        # if os.path.isdir(oldDirpath):  # 如果文件夹中还有文件夹，则进行递归操作
        #     rename(oldDirpath)
        filename = os.path.splitext(files)[0]
        filetype = os.path.splitext(files)[1]
        newDirPath = os.path.join(new_path+'/teeth', 'Y'+str(i)+filetype)
        # print(filename, filetype)
        os.rename(oldDirpath, newDirPath)
        i += 1
    print("共有文件数目：", i)

rename(file_path)

效果图：

然后我发现原始数据的后缀不一样，我又只改了名字，没改后缀，所以现在我写个脚本修改后缀名：

file_path = "/Users/***/Tooth-Detection/data/teeth/"
filelist = os.listdir(file_path)
for files in filelist:
    portion = os.path.splitext(files)
    if portion[1] == '.JPG':
        newname = portion[0] + '.jpg'
        print(newname)
        os.rename(files, newname)
print("finished！")

报如下错误：

然后查阅stack overflow 才知道原因，因为listdir returns a list of files without the path，所以我们修改文件名的时候，必须给路径才可以！
将上述倒数第二行代码进行修改：

os.rename(os.path.join(file_path, files), os.path.join(file_path, newname))

修改后效果如下：

万里长征第一步就好啦～

二、数据集准备–打标签

我的目的是做一个二分类，将有牙齿的文件命名为Y+编号，没有牙齿的文件命名为N+编号。
目前没有牙齿共有210+图片，有牙齿的有500张图片。
现在制作训练集、验证集和测试集，由于我的图片是没有标签的，所以我需要自己打标签：
写一个脚本，遍历文件夹下的文件名，并加上其对应的标签，写在.txt文件中：

import os

def generate(dir, label):
    files = os.listdir(dir+'/teeth')
    with open(dir+'val.txt', 'a+') as f:
        for file in files:
            filename = os.path.split(file)[0]
            filetype = os.path.split(file)[1]
            print(filename, filetype)
            if filetype == '.txt':
                continue
            name = '/teeth' + '/' + file + ' ' + str(int(label)) + '\n'
            f.write(name)
    print("finished!")

img_path = '/Users/***/Tooth-Detection/data/val/'

if __name__ == '__main__':
    i = 1
    generate(img_path, i)

效果如下：