def get_data(name, data_dir, meta_dir, gpu_nums):
isTrain = name == 'train'
ds = PascalVOC12(data_dir, meta_dir, name, shuffle=True)
if isTrain:#special augmentation
shape_aug = [RandomResize(xrange=(0.7, 1.5), yrange=(0.7, 1.5),
aspect_ratio_thres=0.15),
RandomCropWithPadding(args.crop_size,IGNORE_LABEL),
Flip(horiz=True),
]
else:
shape_aug = []
ds = AugmentImageComponents(ds, shape_aug, (0, 1), copy=False)
def f(ds):
image, label = ds
m = np.array([104, 116, 122])
const_arr = np.resize(m, (1,1,3)) # NCHW
image = image - const_arr
return image, label
ds = MapData(ds, f)
if isTrain:
ds = BatchData(ds, args.batch_size*gpu_nums)
ds = PrefetchDataZMQ(ds, 1)
else:
ds = BatchData(ds, 1)
return ds
def get_data(name, data_dir, meta_dir, gpu_nums):
isTrain = True if 'train' in name else False
ds = PascalVOC12(data_dir, meta_dir, name, shuffle=True)
if isTrain:
ds = MapData(ds, RandomResize)
if isTrain:
shape_aug = [
RandomCropWithPadding(args.crop_size, IGNORE_LABEL),
Flip(horiz=True),
]
else:
shape_aug = []
ds = AugmentImageComponents(ds, shape_aug, (0, 1), copy=False)
def f(ds):
image, label = ds
m = np.array([104, 116, 122])
const_arr = np.resize(m, (1, 1, 3)) # NCHW
image = image - const_arr
return image, label
ds = MapData(ds, f)
if isTrain:
ds = BatchData(ds, args.batch_size * gpu_nums)
ds = PrefetchData(ds, nr_prefetch=batch_size * 3, nr_proc=2)
else:
ds = BatchData(ds, 1)
return ds
def get_data_tmp(name, base_dir, meta_dir, gpu_nums):
isTrain = True if 'train' in name else False
m = np.array([104, 116, 122])
const_arr = np.resize(m, (1, 1, 3)) # NCHW
const_arr = np.zeros((args.crop_size[0], args.crop_size[1], 3)) + const_arr # broadcast
def imgread(ds):
img, label = ds
img = cv2.imread(img, cv2.IMREAD_COLOR)
label = cv2.imread(label, cv2.IMREAD_GRAYSCALE)
return img, label
if isTrain:
#ds = LMDBData('/data2/dataset/cityscapes/cityscapes_train.lmdb', shuffle=True)
#ds = FakeData([[batch_size, CROP_HEIGHT, CROP_HEIGHT, 3], [batch_size, CROP_HEIGHT, CROP_HEIGHT, 1]], 5000, random=False, dtype='uint8')
ds = PascalVOC12Files(base_dir, meta_dir, name, shuffle=True)
parallel = min(6, multiprocessing.cpu_count())
augmentors = [
RandomCropWithPadding(args.crop_size),
Flip(horiz=True),
]
aug = imgaug.AugmentorList(augmentors)
def mapf(ds):
img, label = ds
img = cv2.imread(img, cv2.IMREAD_COLOR)
label = cv2.imread(label, cv2.IMREAD_GRAYSCALE)
img, params = aug.augment_return_params(img)
label = aug._augment(label, params)
img = img - const_arr # very time-consuming
return img, label
ds = MultiThreadMapData(ds, parallel, mapf, buffer_size=500, strict=True)
ds = BatchData(ds, args.batch_size * gpu_nums)
ds = PrefetchDataZMQ(ds, 1)
else:
ds = PascalVOC12(base_dir, meta_dir, name, shuffle=False)
ds = MapData(ds, imgread)
ds = BatchData(ds, 1)
return ds
def proceed_test():
ds = PascalVOC12(TEST_DATA_DIR, LIST_DIR, "test", shuffle=False)
imagelist = ds.imglist
def f(ds):
image = ds
m = np.array([104, 116, 122])
const_arr = np.resize(m, (1, 1, 3)) # NCHW
image = image - const_arr
return image
ds = MapData(ds, f)
ds = BatchData(ds, 1)
ctx = [mx.gpu(int(i)) for i in args.gpu.split(',')]
sym_instance = resnet101_deeplab_new()
# infer shape
val_provide_data = [[("data", (1, 3, tile_height, tile_width))]]
val_provide_label = [[("softmax_label", (1, 1, tile_height, tile_width))]]
data_shape_dict = {
'data': (1, 3, tile_height, tile_width),
'softmax_label': (1, 1, tile_height, tile_width)
}
eval_sym = sym_instance.get_symbol(NUM_CLASSES,
is_train=False,
use_global_stats=True)
sym_instance.infer_shape(data_shape_dict)
arg_params, aux_params = load_init_param(args.load, process=True)
sym_instance.check_parameter_shapes(arg_params,
aux_params,
data_shape_dict,
is_train=False)
data_names = ['data']
label_names = ['softmax_label']
# create predictor
predictor = Predictor(eval_sym,
data_names,
label_names,
context=ctx,
provide_data=val_provide_data,
provide_label=val_provide_label,
arg_params=arg_params,
aux_params=aux_params)
from mxnetgo.myutils.fs import mkdir_p
vis_dir = "deeplabv2_4gpu_test_result"
check_dir = os.path.join(vis_dir, "check")
import shutil
shutil.rmtree(vis_dir, ignore_errors=True)
mkdir_p(check_dir)
_itr = ds.get_data()
nbatch = 0
for i in tqdm(range(len(imagelist))):
data = next(_itr)
l = imagelist[i]
filename = os.path.basename(l).rsplit(".", 1)[0]
print filename
output_all = predict_scaler(data,
predictor,
scales=[0.5, 0.75, 1.0, 1.25, 1.5],
classes=NUM_CLASSES,
tile_size=(tile_height, tile_width),
is_densecrf=False,
nbatch=nbatch,
val_provide_data=val_provide_data,
val_provide_label=val_provide_label)
output_all = np.argmax(output_all, axis=0).astype(np.uint8)
result = output_all[:, :, None]
cv2.imwrite(os.path.join(vis_dir, "{}.png".format(filename)), result)
cv2.imwrite(
os.path.join(check_dir, "{}.png".format(filename)),
np.concatenate((data[0][0], visualize_label(output_all)), axis=1))
nbatch += 1