两种文件的数据处理

第一种

文件结构：二分类

from torch.utils.data import Dataset
import cv2
from PIL import Image#可以快速的用图片数据
import os#引入文件

class MyData(Dataset):
    #1.创建一个列表
    def __init__(self,root_dir,label_dir):
        self.root_dir = root_dir
        #self的功能是将变量 转变为全局变量,获取'dataset/train'地址
        self.label_dir  =  label_dir
        # 获取'ants'的地址
        self.path = os.path.join(self.root_dir, self.label_dir)
        #获取全的地址'dataset/train\\ants'
        self.img_path = os.listdir(self.path)
        #让地址内的文件成为列表



    #2.从列表中获取其中的某一个图片
    def __getitem__(self, idx):
        img_name= self.img_path[idx]
        #idex =0,图片'0013035.jpg'
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        #全的地址'dataset/train\\ants\\0013035.jpg'
        img = Image.open(img_item_path)
        label = self.label_dir
        return img, label

    #获取数据集长度
    def __len__(self):
        return len(self.img_path)

root_dir = "../dataset/train"
ants_label_dir = "ants"
bees_label_dir = "bees"
ants_dataset = MyData(root_dir,ants_label_dir)
bees_dataset = MyData(root_dir,bees_label_dir)

# 整个训练集是两个数据记得拼接
train_dataset = ants_dataset + bees_dataset
print(len(ants_dataset))
print(len(train_dataset))

第二种

文件结构：两种对象的数据，标签

制作图片标签txt:图片名称.txt

import os
root_dir = "../dataset_2/train"
target_dir = "ants_image"


img_path = os.listdir(os.path.join(root_dir, target_dir))
#拼接完整路径
label = target_dir.split('_')[0]
#label标签为ants
out_dir = "ants_label"
#目标文件为ants_label

# 循环从图片文件中依次修改
for i in img_path:
    file_name = i.split('.jpg')[0]
    # 图片名称为文件名
    with open(os.path.join(root_dir, out_dir, "{}.txt".format(file_name)), 'w') as f:
        # 将文件名传入，使用形式.txt，写到拼接目录中，每个文件内写入 ants
        f.write(label)


import os
root_dir = "../dataset_2/train"
target_dir2 = "bees_image"

image_path = os.listdir(os.path.join(root_dir, target_dir2))
label = target_dir2.split('_')[0]

out_dir = "bees_label"

for i in image_path:
    file_name = i.split('.jpg')[0]
    with open(os.path.join(root_dir,out_dir,"{}.txt".format(file_name)),'w') as f:
        # 打开（目标文件，文件名格式，写入）作为f，每个文件夹内写入“bees
        f.write(label)