pytorvhDataSet获取自己的数据集

最新推荐文章于 2023-05-05 10:39:00 发布

原创最新推荐文章于 2023-05-05 10:39:00 发布 · 190 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch

python 专栏收录该内容

112 篇文章

订阅专栏

本文介绍了一个使用PyTorch和OpenCV实现的自定义数据集加载类，该类能够从指定路径读取图像及其对应的边界框标签，并在图像上绘制这些边界框。通过遍历整个数据集，每张图像上的目标都被正确标注。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

from torch.utils import data
import os
import cv2

class datasest(data.Dataset):
    def __init__(self,path1,path2):
        self.img_path = path1
        self.label_path = path2
    def __len__(self):
        return len(os.listdir(self.img_path))
    def __getitem__(self, index):
        file = open(self.label_path,'r')
        f = file.readlines()
        labels = [fi.replace('\n','').replace('     ',' ').replace('    ',' ').replace('   ',' ').replace('  ',' ').split(' ') for fi in f][2:]
        img = cv2.imread(os.path.join(self.img_path,labels[index][0]))
        return img,labels[index]

data = datasest(r'D:\CelebA\img_celeba',r'D:\CelebA\list_bbox_celeba.txt')
for i in range(data.__len__()):
    x1 = int(data[i][1][1])
    y1 = int(data[i][1][2])
    w = int(data[i][1][3])
    h = int(data[i][1][4])
    x2 = x1 + w
    y2 = y1 + h
    img = data[i][0]
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
    cv2.imshow('i',img)
    cv2.waitKey(1)