-
Notifications
You must be signed in to change notification settings - Fork 1
/
gulp_data.py
78 lines (70 loc) · 2.43 KB
/
gulp_data.py
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
"""Program for building GulpIO directory of frames for training
See :ref:`cli_tools_gulp_ingestor` for usage details
Original code from https://github.com/epic-kitchens/C1-Action-Recognition-TSN-TRN-TSM"""
import argparse
from multiprocessing import cpu_count
from pathlib import Path
import pandas as pd
from gulpio2 import GulpIngestor
from video_datasets_api.epic_kitchens_100.gulp_adapter import EpicDatasetAdapter
from video_datasets_api.epic_kitchens_100.gulp_adapter import EpicFlowDatasetAdapter
parser = argparse.ArgumentParser(
"Gulp the EPIC dataset allowing for faster read times during training."
)
parser.add_argument(
"in_folder",
type=Path,
help="Directory where subdirectory is a segment name containing frames for that segment.",
)
parser.add_argument(
"out_folder", type=Path, help="Directory to store the gulped files."
)
parser.add_argument(
"labels",
type=Path,
help="Path to the pickle or CSV file which contains the meta information about the dataset.",
)
parser.add_argument("modality", choices=["flow", "rgb"])
parser.add_argument(
"--extension",
type=str,
default="jpg",
help="Which file extension the frames are saved as.",
)
parser.add_argument("--frame-size", type=int, default=-1, help="Size of frames.")
parser.add_argument(
"--segments-per-chunk",
type=int,
default=100,
help="Number of action segments per chunk to save.",
)
parser.add_argument(
"-j",
"--num-workers",
type=int,
default=cpu_count(),
help="Number of workers to run the task.",
)
def main(args):
if args.labels.suffix.lower() == ".pkl":
labels = pd.read_pickle(args.labels)
elif args.labels.suffix.lower() == ".csv":
labels = pd.read_csv(args.labels, index_col="uid")
else:
raise ValueError("Expected .csv or .pkl suffix for annotation file")
if args.modality.lower() == "flow":
epic_adapter = EpicFlowDatasetAdapter(
str(args.in_folder), labels, args.frame_size, args.extension,
)
elif args.modality.lower() == "rgb":
epic_adapter = EpicDatasetAdapter(
str(args.in_folder), labels, args.frame_size, args.extension,
)
else:
raise ValueError("Modality '{}' not supported".format(args.modality))
ingestor = GulpIngestor(
epic_adapter, str(args.out_folder), args.segments_per_chunk, args.num_workers
)
ingestor()
if __name__ == "__main__":
main(parser.parse_args())