Question

我在PyTorch中为视频中的动作识别建立了CNN模型。我正在使用手电筒数据加载器模块加载数据以进行训练。

train_loader = torch.utils.data.DataLoader(
            training_data,
            batch_size=8,
            shuffle=True,
            num_workers=4,
            pin_memory=True)

然后通过train_loader来训练模型。

train_epoch(i, train_loader, action_detect_model, criterion, optimizer, opt,
                        train_logger, train_batch_logger)

现在，我想添加一条附加路径，该路径将采用视频帧的相应光流。为了计算光通量，我使用cv2.calcOpticalFlowFarneback。但是问题是我不确定如何获得与火车数据加载器张量中的数据相对应的图像，因为它们将被拖曳。我不想预先计算光流，因为存储需求非常大（每个帧需要600 kBs）。

Answer 1

您必须使用自己的数据加载器类来动态计算光流。想法是，此类获得文件名元组的列表（当前图像，下一个图像），该列表包含视频序列的当前和下一帧的文件名，而不是简单的文件名列表。在给文件名列表添加后，可以获取正确的图像对。以下代码为您提供了一个非常简单的示例实现：

from torch.utils.data import Dataset
import cv2
import numpy as np

class FlowDataLoader(Dataset):
def __init__(self,
             filename_tuples):

    random.shuffle(filename_tuples)
    self.lines = filename_tuples

def __getitem__(self, index):
    img_filenames = self.lines[index]
    curr_img = cv2.cvtColor(cv2.imread(img_filenames[0]), cv2.BGR2GRAY)
    next_img = cv2.cvtColor(cv2.imread(img_filenames[1]), cv2.BGR2GRAY)
    flow = cv2.calcOpticalFlowFarneback(curr_img, next_img, ... [parameter])

    # code for loading the class label
    # label = ...
    #
    # this is a very simple data normalization
    curr_img= curr_img.astype(np.float32) / 255
    next_img = next_img .astype(np.float32) / 255
    # you can return the image and flow seperatly 
    return curr_img, flow, label
    # or stacked as follows
    # return np.dstack((curr_img,flow)), label

# at this place you need a function that create a list of training sample filenames
# that look like this
training_filelist = [(img000.png, img001.png), 
                     (img001.png, img002.png),
                     (img002.png, img003.png)] 

training_data = FlowDataLoader(training_filelist)
train_loader = torch.utils.data.DataLoader(
        training_data,
        batch_size=8,
        shuffle=True,
        num_workers=4,
        pin_memory=True)

这只是FlowDataLoader的一个简单示例。从概念上讲，应该对此进行扩展，以便curr_image输出包含归一化的rgb值，并且光流也将归一化和修剪。

计算与torch.utils.data.DataLoader中的数据相对应的光流

1 个答案: