forked from tensorpack/tensorpack
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathdump-dataflow.py
More file actions
executable file
·55 lines (49 loc) · 1.75 KB
/
dump-dataflow.py
File metadata and controls
executable file
·55 lines (49 loc) · 1.75 KB
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
#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# File: dump-dataflow.py
# Author: Yuxin Wu <ppwwyyxx@gmail.com>
import argparse
import cv2
import imp
import tqdm
import os
from tensorpack.utils import logger
from tensorpack.utils.fs import mkdir_p
from tensorpack.dataflow import RepeatedData
parser = argparse.ArgumentParser()
parser.add_argument(dest='config')
parser.add_argument('-o', '--output',
help='output directory to dump dataset image. If not given, will not dump images.')
parser.add_argument('-s', '--scale',
help='scale the image data (maybe by 255)', default=1, type=int)
parser.add_argument('--index',
help='index of the image component in datapoint',
default=0, type=int)
parser.add_argument('-n', '--number', help='number of images to dump',
default=10, type=int)
args = parser.parse_args()
logger.auto_set_dir(action='d')
get_config_func = imp.load_source('config_script', args.config).get_config
config = get_config_func()
config.dataset.reset_state()
if args.output:
mkdir_p(args.output)
cnt = 0
index = args.index # TODO: as an argument?
for dp in config.dataset.get_data():
imgbatch = dp[index]
if cnt > args.number:
break
for bi, img in enumerate(imgbatch):
cnt += 1
fname = os.path.join(args.output, '{:03d}-{}.png'.format(cnt, bi))
cv2.imwrite(fname, img * args.scale)
NR_DP_TEST = args.number
logger.info("Testing dataflow speed:")
ds = RepeatedData(config.dataset, -1)
with tqdm.tqdm(total=NR_DP_TEST, leave=True, unit='data points') as pbar:
for idx, dp in enumerate(ds.get_data()):
del dp
if idx > NR_DP_TEST:
break
pbar.update()