JeVoisBase  1.21
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
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
7import os
8import argparse
9
10import numpy as np
11import cv2 as cv
12
13from youtureid import YoutuReID
14
15def 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
23backends = [cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_CUDA]
24targets = [cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16]
25help_msg_backends = "Choose one of the computation backends: {:d}: OpenCV implementation (default); {:d}: CUDA"
26help_msg_targets = "Chose one of the target computation devices: {:d}: CPU (default); {:d}: CUDA; {:d}: CUDA fp16"
27try:
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"
32except:
33 print('This version of OpenCV does not support TIM-VX and NPU. Visit https://gist.github.com/fengyuentau/5a7a5ba36328f2b763aea026c43fa45f for more information.')
34
35parser = argparse.ArgumentParser(
36 description="ReID baseline models from Tencent Youtu Lab")
37parser.add_argument('--query_dir', '-q', type=str, help='Query directory.')
38parser.add_argument('--gallery_dir', '-g', type=str, help='Gallery directory.')
39parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
40parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
41parser.add_argument('--topk', type=int, default=10, help='Top-K closest from gallery for each query.')
42parser.add_argument('--model', '-m', type=str, default='person_reid_youtu_2021nov.onnx', help='Path to the model.')
43parser.add_argument('--save', '-s', type=str2bool, default=False, help='Set true to save results. This flag is invalid when using camera.')
44parser.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.')
45args = parser.parse_args()
46
47def 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
56def 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
80if __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
readImageFromDirectory(img_dir, w=128, h=256)
Definition demo.py:47
visualize(image, results, box_color=(0, 255, 0), text_color=(0, 0, 255), fps=None)
Definition demo.py:46
str
Definition demo.py:35