JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
demo.py
Go to the documentation of this file.
1 # This file is part of OpenCV Zoo project.
2 # It is subject to the license terms in the LICENSE file found in the same directory.
3 #
4 # Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
5 # Third party copyrights are property of their respective owners.
6 
7 import os
8 import argparse
9 
10 import numpy as np
11 import cv2 as cv
12 
13 from youtureid import YoutuReID
14 
15 def str2bool(v):
16  if v.lower() in ['on', 'yes', 'true', 'y', 't']:
17  return True
18  elif v.lower() in ['off', 'no', 'false', 'n', 'f']:
19  return False
20  else:
21  raise NotImplementedError
22 
23 backends = [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_CUDA]
24 targets = [cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16]
25 help_msg_backends = "Choose one of the computation backends: {:d}: OpenCV implementation (default); {:d}: CUDA"
26 help_msg_targets = "Chose one of the target computation devices: {:d}: CPU (default); {:d}: CUDA; {:d}: CUDA fp16"
27 try:
28  backends += [cv.dnn.DNN_BACKEND_TIMVX]
29  targets += [cv.dnn.DNN_TARGET_NPU]
30  help_msg_backends += "; {:d}: TIMVX"
31  help_msg_targets += "; {:d}: NPU"
32 except:
33  print('This version of OpenCV does not support TIM-VX and NPU. Visit https://gist.github.com/fengyuentau/5a7a5ba36328f2b763aea026c43fa45f for more information.')
34 
35 parser = argparse.ArgumentParser(
36  description="ReID baseline models from Tencent Youtu Lab")
37 parser.add_argument('--query_dir', '-q', type=str, help='Query directory.')
38 parser.add_argument('--gallery_dir', '-g', type=str, help='Gallery directory.')
39 parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
40 parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
41 parser.add_argument('--topk', type=int, default=10, help='Top-K closest from gallery for each query.')
42 parser.add_argument('--model', '-m', type=str, default='person_reid_youtu_2021nov.onnx', help='Path to the model.')
43 parser.add_argument('--save', '-s', type=str2bool, default=False, help='Set true to save results. This flag is invalid when using camera.')
44 parser.add_argument('--vis', '-v', type=str2bool, default=True, help='Set true to open a window for result visualization. This flag is invalid when using camera.')
45 args = parser.parse_args()
46 
47 def readImageFromDirectory(img_dir, w=128, h=256):
48  img_list = []
49  file_list = os.listdir(img_dir)
50  for f in file_list:
51  img = cv.imread(os.path.join(img_dir, f))
52  img = cv.resize(img, (w, h))
53  img_list.append(img)
54  return img_list, file_list
55 
56 def visualize(results, query_dir, gallery_dir, output_size=(128, 384)):
57  def addBorder(img, color, borderSize=5):
58  border = cv.copyMakeBorder(img, top=borderSize, bottom=borderSize, left=borderSize, right=borderSize, borderType=cv.BORDER_CONSTANT, value=color)
59  return border
60 
61  results_vis = dict.fromkeys(results.keys(), None)
62  for f, topk_f in results.items():
63  query_img = cv.imread(os.path.join(query_dir, f))
64  query_img = cv.resize(query_img, output_size)
65  query_img = addBorder(query_img, [0, 0, 0])
66  cv.putText(query_img, 'Query', (10, 30), cv.FONT_HERSHEY_COMPLEX, 1., (0, 255, 0), 2)
67 
68  gallery_img_list = []
69  for idx, gallery_f in enumerate(topk_f):
70  gallery_img = cv.imread(os.path.join(gallery_dir, gallery_f))
71  gallery_img = cv.resize(gallery_img, output_size)
72  gallery_img = addBorder(gallery_img, [255, 255, 255])
73  cv.putText(gallery_img, 'G{:02d}'.format(idx), (10, 30), cv.FONT_HERSHEY_COMPLEX, 1., (0, 255, 0), 2)
74  gallery_img_list.append(gallery_img)
75 
76  results_vis[f] = np.concatenate([query_img] + gallery_img_list, axis=1)
77 
78  return results_vis
79 
80 if __name__ == '__main__':
81  # Instantiate YoutuReID for person ReID
82  net = YoutuReID(modelPath=args.model, backendId=args.backend, targetId=args.target)
83 
84  # Read images from dir
85  query_img_list, query_file_list = readImageFromDirectory(args.query_dir)
86  gallery_img_list, gallery_file_list = readImageFromDirectory(args.gallery_dir)
87 
88  # Query
89  topk_indices = net.query(query_img_list, gallery_img_list, args.topk)
90 
91  # Index to filename
92  results = dict.fromkeys(query_file_list, None)
93  for f, indices in zip(query_file_list, topk_indices):
94  topk_matches = []
95  for idx in indices:
96  topk_matches.append(gallery_file_list[idx])
97  results[f] = topk_matches
98  # Print
99  print('Query: {}'.format(f))
100  print('\tTop-{} from gallery: {}'.format(args.topk, str(topk_matches)))
101 
102  # Visualize
103  results_vis = visualize(results, args.query_dir, args.gallery_dir)
104 
105  if args.save:
106  for f, img in results_vis.items():
107  cv.imwrite('result-{}'.format(f), img)
108 
109  if args.vis:
110  for f, img in results_vis.items():
111  cv.namedWindow('result-{}'.format(f), cv.WINDOW_AUTOSIZE)
112  cv.imshow('result-{}'.format(f), img)
113  cv.waitKey(0)
114  cv.destroyAllWindows()
115 
demo.str
str
Definition: demo.py:35
demo.str2bool
str2bool
Definition: demo.py:43
youtureid.YoutuReID
Definition: youtureid.py:10
demo.visualize
def visualize(image, results, box_color=(0, 255, 0), text_color=(0, 0, 255), fps=None)
Definition: demo.py:46
demo.readImageFromDirectory
def readImageFromDirectory(img_dir, w=128, h=256)
Definition: demo.py:47