人工智能学习07--pytorch02--数据集取图片+以图片名为文件名的txt文档(label)

该文介绍如何在Python中使用torch.utils.data.Dataset类来构建自定义数据集,通过PIL库读取图片,并结合标签目录获取图片及对应的类别。同时,展示了如何自动化为图片生成标签txt文件,以图片名字命名并写入分类信息。
摘要由CSDN通过智能技术生成

1、取图片

from torch.utils.data import Dataset
# import cv2
from PIL import Image
import os

class MyData(Dataset):
# help() 中说,想要使用 Dataset,需要重写。__getitem__ 和 __len__

    def __init__(self,root_dir,label_dir): # 初始化函数
        # 可以把self指定的变量给后面的函数使用(成为全局变量)
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir,self.label_dir) # 把两个路径拼接在一起(\\)
        self.img_path = os.listdir(self.path) # 获得所有地址

    def __getitem__(self, idx):
        # 要获取图片,首先创建图片地址列表
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir,self.label_dir,img_name)
        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_image'
bees_label_dir = 'bees_image'
ants_dataset = MyData(root_dir,ants_label_dir)
bees_dataset = MyData(root_dir,bees_label_dir)
# 当类中有__getitem__方法,实例化对象A后调用A[0]会自动调用__getitem__的

#img,label=bees_dataset[1]
#img.show()

train_dataset = ants_dataset + bees_dataset  # 把两个数据集进行拼接
#len(train_dataset)
#245
#len(ants_dataset)
#124
#len(bees_dataset)
#121
img,label = train_dataset[123] # ant
img.show()
img,label = train_dataset[124] # bee
img.show()

假设这个train文件夹里面只有ants_image和bees_image
在这里插入图片描述
运行结果
在这里插入图片描述
在这里插入图片描述
可以看到这段代码成功地从数据集里面取出了图片

2、加label

在这里插入图片描述
用图片名字命名txt文件,内容是这个图片的分类。
因为图片太多了,所以用代码自动添加
只以ant为例,bee的就是把这个改改

import os
# 把图片的label 生成以图片名为文件名的txt文档
root_dir = 'dataset/train'
target_dir = 'ants_image'
img_path = os.listdir(os.path.join(root_dir, target_dir))
label = target_dir.split('_')[0]
out_dir = '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:
        f.write(label)

可以看到下面这些都是自动生成的:
在这里插入图片描述

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值