From 00cb5671dd7cfd6e55ba5b45fe8309f1f251999c Mon Sep 17 00:00:00 2001 From: danleifeng Date: Thu, 13 Feb 2020 07:48:53 +0000 Subject: [PATCH] edit data format in load_bin --- plsc/utils/jpeg_reader.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/plsc/utils/jpeg_reader.py b/plsc/utils/jpeg_reader.py index 554c6e0..dfca8eb 100644 --- a/plsc/utils/jpeg_reader.py +++ b/plsc/utils/jpeg_reader.py @@ -252,7 +252,11 @@ def load_bin(path, image_size, data_format ='NCHW'): bins, issame_list = pickle.load(open(path, 'rb'), encoding='bytes') data_list = [] for flip in [0, 1]: - data = np.empty((len(issame_list) * 2, 3, image_size[0], image_size[1])) + if data_format == 'NCHW': + data = np.empty((len(issame_list) * 2, 3, image_size[0], image_size[1])) + else: + #NHWC + data = np.empty((len(issame_list) * 2, image_size[0], image_size[1], 3)) data_list.append(data) for i in range(len(issame_list) * 2): _bin = bins[i] -- GitLab