PVCD-Pytorch/create_dataset.py at master · vanhao-le/PVCD-Pytorch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import torch
import torch.utils.data as data
from PIL import Image
import os
import os.path
import numpy as np
from numpy.random import randint
import pandas as pd

class VideoRecord(object):
    def __init__(self, row):
        self._data = row

    @property
    def path(self):
        return self._data[0]

    @property
    def total_frames(self):
        return int(self._data[1])

    @property
    def label(self):
        return int(self._data[2])


# Creating a Custom Dataset for your files
# https://pytorch.org/tutorials/beginner/basics/data_tutorial.html

class TSNDataSet(data.Dataset):
    """
    Args:
    data_dir: file path of images for training / testing
    annotation_dir: file path of groundtruth e.g., train.csv / test.csv
    file_categories: file path of lables e.g., category.txt
    """
    def __init__(self, data_dir, annotation_dir, file_categories, num_segments=8, transform=None):
        self.data_dir = data_dir
        self.annotation_dir = annotation_dir
        self.num_segments = num_segments
        self.transform = transform
        self.classes = list(np.loadtxt(file_categories, dtype=np.str, delimiter=',')) #load category.txt
        self.video_list = list()
        anno_array = np.loadtxt(annotation_dir, dtype=np.str, delimiter=',')
        for anno in anno_array:
            path = str(anno[0]).split('.')[0]
            label_name = anno[1]
            label = self.classes.index(label_name)
            data_path = os.path.join(data_dir, path)
            total_frames = len(os.listdir(data_path))
            # print(path, total_frames, label)
            self.video_list.append(VideoRecord([path, total_frames, label]))

    def _sample(self, num_total, num_segments):
        # sample = np.random.choice(range(num_total), size=num_segs, replace=None)
        # sample = np.sort(sample, axis=-1, kind='quicksort', order=None)
        sample = np.linspace(0, num_total-1, num_segments, endpoint=True, retstep=True, dtype=int)[0]
        return sample

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

    # the __getitem__ method also allows you to turn your object into an iterable.

    def __getitem__(self, index):
        assert index < len(self.video_list)
        info = self.video_list[index]
        target = info.label
        total_frames = info.total_frames
        data_path = os.path.join(self.data_dir, info.path)
        image_path_list = os.listdir(data_path)
        image_list = list()
        sample = self._sample(total_frames, self.num_segments)
        for i in sample:
            img_path = os.path.join(data_path, image_path_list[i])
            img = Image.open(img_path)
            if self.transform:
                img = self.transform(img)
            image_list.append(img)
        image = torch.stack(image_list)
        return image, target

# Creating a Custom Dataset for your files
# https://pytorch.org/tutorials/beginner/basics/data_tutorial.html

class VCDBDataset(data.Dataset):
    """
    Args:
    data_dir: file path of images for training / validation
    annotation_dir: file path of groundtruth e.g., train.csv / val.csv
    file_categories: file path of lables e.g., category.csv
    """
    def __init__(self, data_dir, annotation_dir, file_categories, num_segments=8, transform=None):
        self.data_dir = data_dir
        self.annotation_dir = annotation_dir
        self.num_segments = num_segments
        self.transform = transform
        self.classes = pd.read_csv(file_categories)
        self.video_list = list()
        df_annotation = pd.read_csv(annotation_dir)
        for item in df_annotation.itertuples():
            path = item.video_name.split('.')[0]
            label_name = item.classIDx
            label = label_name
            data_path = os.path.join(data_dir, path)
            total_frames = len(os.listdir(data_path))
            # print(path, total_frames, label)
            self.video_list.append(VideoRecord([path, total_frames, label]))

    def _sample(self, num_total, num_segments):
        # sample = np.random.choice(range(num_total), size=num_segs, replace=None)
        # sample = np.sort(sample, axis=-1, kind='quicksort', order=None)
        # numpy.linspace(start, stop, num=50, endpoint=True, retstep=False, dtype=None, axis=0)
        # Returns num evenly spaced samples, calculated over the interval [start, stop].
        # Ouput: samples, step -> put dim=0 is to get only samples
        sample = np.linspace(0, num_total-1, num_segments, endpoint=True, retstep=True, dtype=int)[0]
        return sample

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

    # the __getitem__ method also allows you to turn your object into an iterable.

    def __getitem__(self, index):
        assert index < len(self.video_list)
        info = self.video_list[index]
        target = info.label
        total_frames = info.total_frames
        data_path = os.path.join(self.data_dir, info.path)
        image_path_list = os.listdir(data_path)
        image_list = list()
        sample = self._sample(total_frames, self.num_segments)
        for i in sample:
            img_path = os.path.join(data_path, image_path_list[i])
            img = Image.open(img_path)
            if self.transform:
                img = self.transform(img)
            image_list.append(img)
        image = torch.stack(image_list)
        return image, target