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

Question

我已经在 PyTorch 中构建了一个用于视频动作识别的 CNN 模型。我正在使用 torch 数据加载器模块加载训练数据。

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

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

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中数据对应的光流

Compute Optical Flow corresponding to data in the torch.utils.data.DataLoader

python

pytorch

opencv-python