bugfix for visualize and demo (#468)

Summary minor bugs for visualize and demo caused by revised DataLoader

demo/demo.py

@@ -21,8 +21,10 @@ from fastreid.utils.logger import setup_logger
 from fastreid.utils.file_io import PathManager
 
 from predictor import FeatureExtractionDemo
+
 # import some modules added in project like this below
-# from projects.PartialReID.partialreid import *
+# sys.path.append('../projects/PartialReID')
+# from partialreid import *
 
 cudnn.benchmark = True
 setup_logger(name="fastreid")
@@ -84,4 +86,4 @@ if __name__ == '__main__':
             img = cv2.imread(path)
             feat = demo.run_on_image(img)
             feat = feat.numpy()
-            np.save(os.path.join(args.output, path.replace('.jpg', '.npy').split('/')[-1]), feat)
+            np.save(os.path.join(args.output, os.path.basename(path) + '.npy'), feat)

demo/predictor.py

@@ -6,11 +6,11 @@
 
 import atexit
 import bisect
+from collections import deque
 
 import cv2
 import torch
 import torch.multiprocessing as mp
-from collections import deque
 
 from fastreid.engine import DefaultPredictor
 
@@ -70,16 +70,16 @@ class FeatureExtractionDemo(object):
                 if cnt >= buffer_size:
                     batch = batch_data.popleft()
                     predictions = self.predictor.get()
-                    yield predictions, batch["targets"].numpy(), batch["camids"].numpy()
+                    yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
 
             while len(batch_data):
                 batch = batch_data.popleft()
                 predictions = self.predictor.get()
-                yield predictions, batch["targets"].numpy(), batch["camids"].numpy()
+                yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
         else:
             for batch in data_loader:
                 predictions = self.predictor(batch["images"])
-                yield predictions, batch["targets"].numpy(), batch["camids"].numpy()
+                yield predictions, batch["targets"].cpu().numpy(), batch["camids"].cpu().numpy()
 
 
 class AsyncPredictor:

demo/visualize_result.py

@@ -24,7 +24,8 @@ from fastreid.utils.visualizer import Visualizer
 
 # import some modules added in project
 # for example, add partial reid like this below
-# from projects.PartialReID.partialreid import *
+# sys.path.append("projects/PartialReID")
+# from partialreid import *
 
 cudnn.benchmark = True
 setup_logger(name="fastreid")
@@ -35,6 +36,7 @@ logger = logging.getLogger('fastreid.visualize_result')
 def setup_cfg(args):
     # load config from file and command-line arguments
     cfg = get_cfg()
+    # add_partialreid_config(cfg)
     cfg.merge_from_file(args.config_file)
     cfg.merge_from_list(args.opts)
     cfg.freeze()
@@ -100,7 +102,7 @@ def get_parser():
 if __name__ == '__main__':
     args = get_parser().parse_args()
     cfg = setup_cfg(args)
-    test_loader, num_query = build_reid_test_loader(cfg, args.dataset_name)
+    test_loader, num_query = build_reid_test_loader(cfg, dataset_name=args.dataset_name)
     demo = FeatureExtractionDemo(cfg, parallel=args.parallel)
 
     logger.info("Start extracting image features")