loader.py

from torch.utils.data import Dataset, DataLoader
import cv2
import os
import torch
import numpy as np


def read_video(filename):
    video = cv2.VideoCapture(filename)
    vid = []
    while True:
        ret, img = video.read()
        if ret == True:
            img = cv2.resize(img,(100,100))
            vid.append(img)
        else:
            break
    return np.array(vid)


class Firedataset(Dataset):
    """Fire dataset."""

    def __init__(self, video_dir, label_dir):
        """
        Args:
            video_dir (string): Directory with all the videos.
            label_dir (string): Path to the label data.

        """
        self.video_dir = video_dir
        self.filename = os.listdir(video_dir)
        self.labels = np.loadtxt(label_dir)

    def __len__(self):
        return len(os.listdir(self.video_dir))

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        video_name = os.path.join(self.video_dir,self.filename[idx])
        video = read_video(video_name)
        label = self.labels[idx].repeat(video.shape[0])

        sample = {'label':label, 'video':video}
        return sample