基于crnn图像序列预测-pytorch实现——加载自己的数据集

时间：2024-03-23 11:22:53

任务介绍：在这里基于卷积循环神经网络，做一个图像序列的预测。图像序列中每一张图像对应三个状态标签的一个。利用卷积网络对每连续的九张图像进行特征提取，然后将提取的特征序列输入到循环神经网络（LSTM）中，预测下一张图像状态，再与原本图像序列中的第十张图像计算loss。

数据集介绍：数据来自两段工业视频（一段训练，一段测试），将视频按帧提取图像序列（在时间上连续）。对图像进行标记，分为三个状态标签（0,1,2），每张图像对应一个状态。

数据预处理：训练集包括15000张图片，尺寸为3×256×256，命名从00001.jpg-15000.jpg。将图片的路径和标签信息保存到txt文件中，中间以空格分隔。如图：
基于crnn图像序列预测-pytorch实现——加载自己的数据集
数据加载：这里改写了Dataset类中的__getitem__（）函数，使得每次迭代返回连续的九张图像和第十张图像的标签。详细代码如下：

def default_loader(path):
    return Image.open(path).convert('RGB')


class MyDataset(Dataset):
    def __init__(self, txt, transform=None, target_transform=None, loader=default_loader):
        fh = open(txt, 'r')
        imgs = []
        for line in fh:
            line = line.strip('\n')
            line = line.rstrip()
            words = line.split()
            imgs.append((words[0], int(words[1])))
        imgs.sort(key=lambda x: x[0], reverse=False)
        self.num_samples = len(imgs)
        self.num_samples_per_iteration = 9
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader

    def __getitem__(self, index):
        current_index = np.random.choice(range(self.num_samples_per_iteration, self.num_samples))
        current_imgs = []
        current_label = self.imgs[current_index][1]
        for i in range(current_index - self.num_samples_per_iteration, current_index):
            fn, label = self.imgs[i]
            img = self.loader(fn)
            if self.transform is not None:
                img = self.transform(img)
            current_imgs.append(img)
        batch_cur_imgs = np.stack(current_imgs, axis=0)  # [9, 3, 256, 256]
        return batch_cur_imgs, current_label

    def __len__(self):
        return len(self.imgs)


train_data = MyDataset(txt='trainset256.txt', transform=transforms.ToTensor())
train_loader = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=False)

test_data = MyDataset(txt='testset256.txt', transform=transforms.ToTensor())
test_loader = DataLoader(test_data, batch_size=BATCH_SIZE, shuffle=False)
print('num_of_trainData:', len(train_data))
print('num_of_testData:', len(test_data))

相关文章