From ad3bd3a342b403ecbafa96d155b8020a19e08aaa Mon Sep 17 00:00:00 2001 From: lubin10 Date: Wed, 7 Jul 2021 06:22:21 +0000 Subject: [PATCH] fixbug bs=1 in build_gallery --- aa.txt | 1 - deploy/python/build_gallery.py | 16 ++++++++++++++-- 2 files changed, 14 insertions(+), 3 deletions(-) delete mode 100644 aa.txt diff --git a/aa.txt b/aa.txt deleted file mode 100644 index 68eb95ea6..000000000 --- a/aa.txt +++ /dev/null @@ -1 +0,0 @@ -i have already fix the bug diff --git a/deploy/python/build_gallery.py b/deploy/python/build_gallery.py index 142e3cf23..a7297366d 100644 --- a/deploy/python/build_gallery.py +++ b/deploy/python/build_gallery.py @@ -71,14 +71,26 @@ class GalleryBuilder(object): gallery_features = np.zeros( [len(gallery_images), config['embedding_size']], dtype=np.float32) + #construct batch imgs and do inference + batch_size = config.get("batch_size", 32) + batch_img = [] for i, image_file in enumerate(tqdm(gallery_images)): img = cv2.imread(image_file) if img is None: logger.error("img empty, please check {}".format(image_file)) exit() img = img[:, :, ::-1] - rec_feat = self.rec_predictor.predict(img) - gallery_features[i, :] = rec_feat + batch_img.append(img) + + if (i + 1) % batch_size == 0: + rec_feat = self.rec_predictor.predict(batch_img) + gallery_features[i - batch_size + 1:i + 1, :] = rec_feat + batch_img = [] + + if len(batch_img) > 0: + rec_feat = self.rec_predictor.predict(batch_img) + gallery_features[-len(batch_img):, :] = rec_feat + batch_img = [] # train index self.Searcher = Graph_Index(dist_type=config['dist_type'])