data.py
2.79 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
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
# --------------------------------------------------------
# Dragon
# Copyright(c) 2017 SeetaTech
# Written by Ting Pan
# --------------------------------------------------------
import numpy as np
from dragon.core.tensor import Tensor
from dragon.operators.utils import Run
def Imagenet(**kwargs):
"""
:param kwargs: a dict of imagenet data param
:param --> mean_value: a list of mean values for channles [B-G-R]
:param --> source: a str of the images root directory
:param --> imageset: a str of text file contains image name / label
:param --> prefetch: a int of the prefetching size
:param --> batch_size: a int of the batch size
:param --> force_gray a bool of whether to use only 1 channel
:param --> shuffle a bool of whether to use shuffle
:param --> scale a float of the coeff to scale
:return: 2 Tensors of data and label
"""
args = locals(); kwargs = args['kwargs']
del args['kwargs']; kwargs = dict(args, **kwargs)
kwargs['module'] = 'dragon.vm.caffe.io.data_layer'
kwargs['op'] = 'DataLayer'
return Run([], param_str=str(kwargs), nout=2, **kwargs)
def LMDBData(**kwargs):
"""
:param kwargs: a dict of imagenet data param
:param --> mean_value: a list of mean values for channles [B-G-R]
:param --> source: a str of the images root directory
:param --> imageset: a str of text file contains image name / label
:param --> prefetch: a int of the prefetching size
:param --> batch_size: a int of the batch size
:param --> force_gray a bool of whether to use only 1 channel
:param --> shuffle a bool of whether to use shuffle
:param --> crop_size a int
:param --> mirror a bool
:param --> color_augmentation a bool
:param --> min_random_scale a float, defualt is 1.0
:param --> max_random_scale a float, default is 1.0
:param --> scale a float of the coeff to scale
:return: 2 Tensors of data and label
"""
args = locals(); kwargs = args['kwargs']
del args['kwargs']; kwargs = dict(args, **kwargs)
kwargs['module'] = 'dragon.vm.caffe.io.data_layer'
kwargs['op'] = 'DataLayer'
return Run([], param_str=str(kwargs), nout=2, **kwargs)
def MemoryData(inputs, dtype=np.float32, **kwargs):
args = locals(); kwargs = args['kwargs']
del args['kwargs']; kwargs = dict(args, **kwargs)
if dtype is np.float32: kwargs['dtype'] = 1
elif dtype is np.float16: kwargs['dtype'] = 12
else: raise TypeError('unsupported data type')
return Tensor.CreateOperator(nout=1, op_type='MemoryData', **kwargs)