1# Tencent is pleased to support the open source community by making ncnn available. 2# 3# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved. 4# 5# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except 6# in compliance with the License. You may obtain a copy of the License at 7# 8# https://opensource.org/licenses/BSD-3-Clause 9# 10# Unless required by applicable law or agreed to in writing, software distributed 11# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 12# CONDITIONS OF ANY KIND, either express or implied. See the License for the 13# specific language governing permissions and limitations under the License. 14 15import sys 16import cv2 17import numpy as np 18import ncnn 19from ncnn.model_zoo import get_model 20from ncnn.utils import draw_detection_objects 21 22 23def draw_result(image, class_names, boxes, masks, classes, scores): 24 colors = [ 25 [56, 0, 255], 26 [226, 255, 0], 27 [0, 94, 255], 28 [0, 37, 255], 29 [0, 255, 94], 30 [255, 226, 0], 31 [0, 18, 255], 32 [255, 151, 0], 33 [170, 0, 255], 34 [0, 255, 56], 35 [255, 0, 75], 36 [0, 75, 255], 37 [0, 255, 169], 38 [255, 0, 207], 39 [75, 255, 0], 40 [207, 0, 255], 41 [37, 0, 255], 42 [0, 207, 255], 43 [94, 0, 255], 44 [0, 255, 113], 45 [255, 18, 0], 46 [255, 0, 56], 47 [18, 0, 255], 48 [0, 255, 226], 49 [170, 255, 0], 50 [255, 0, 245], 51 [151, 255, 0], 52 [132, 255, 0], 53 [75, 0, 255], 54 [151, 0, 255], 55 [0, 151, 255], 56 [132, 0, 255], 57 [0, 255, 245], 58 [255, 132, 0], 59 [226, 0, 255], 60 [255, 37, 0], 61 [207, 255, 0], 62 [0, 255, 207], 63 [94, 255, 0], 64 [0, 226, 255], 65 [56, 255, 0], 66 [255, 94, 0], 67 [255, 113, 0], 68 [0, 132, 255], 69 [255, 0, 132], 70 [255, 170, 0], 71 [255, 0, 188], 72 [113, 255, 0], 73 [245, 0, 255], 74 [113, 0, 255], 75 [255, 188, 0], 76 [0, 113, 255], 77 [255, 0, 0], 78 [0, 56, 255], 79 [255, 0, 113], 80 [0, 255, 188], 81 [255, 0, 94], 82 [255, 0, 18], 83 [18, 255, 0], 84 [0, 255, 132], 85 [0, 188, 255], 86 [0, 245, 255], 87 [0, 169, 255], 88 [37, 255, 0], 89 [255, 0, 151], 90 [188, 0, 255], 91 [0, 255, 37], 92 [0, 255, 0], 93 [255, 0, 170], 94 [255, 0, 37], 95 [255, 75, 0], 96 [0, 0, 255], 97 [255, 207, 0], 98 [255, 0, 226], 99 [255, 245, 0], 100 [188, 255, 0], 101 [0, 255, 18], 102 [0, 255, 75], 103 [0, 255, 151], 104 [255, 56, 0], 105 [245, 255, 0], 106 ] 107 108 color_index = 0 109 110 for box, mask, label, score in zip(boxes, masks, classes, scores): 111 if score < 0.15: 112 continue 113 114 print( 115 "%s = %.5f at %.2f %.2f %.2f x %.2f\n" 116 % (label, score, box[0], box[1], box[2], box[3]) 117 ) 118 119 cv2.rectangle( 120 image, 121 (int(box[0]), int(box[1])), 122 (int(box[0] + box[2]), int(int(box[1] + box[3]))), 123 (255, 0, 0), 124 ) 125 126 text = "%s %.1f%%" % (class_names[int(label)], score * 100) 127 128 label_size, baseLine = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) 129 130 x = box[0] 131 y = box[1] - label_size[1] - baseLine 132 if y < 0: 133 y = 0 134 if x + label_size[0] > image.shape[1]: 135 x = image.shape[1] - label_size[0] 136 137 cv2.rectangle( 138 image, 139 (int(x), int(y)), 140 (int(x + label_size[0]), int(y + label_size[1] + baseLine)), 141 (255, 255, 255), 142 -1, 143 ) 144 145 cv2.putText( 146 image, 147 text, 148 (int(x), int(y + label_size[1])), 149 cv2.FONT_HERSHEY_SIMPLEX, 150 0.5, 151 (0, 0, 0), 152 ) 153 154 image[mask] = image[mask] * 0.5 + np.array(colors[color_index]) * 0.5 155 color_index += 1 156 157 cv2.imshow("image", image) 158 cv2.waitKey(0) 159 160 161if __name__ == "__main__": 162 if len(sys.argv) != 2: 163 print("Usage: %s [imagepath]" % (sys.argv[0])) 164 sys.exit(0) 165 166 imagepath = sys.argv[1] 167 m = cv2.imread(imagepath) 168 if m is None: 169 print("cv2.imread %s failed\n" % (imagepath)) 170 sys.exit(0) 171 172 net = get_model( 173 "yolact", 174 target_size=550, 175 confidence_threshold=0.05, 176 nms_threshold=0.5, 177 keep_top_k=200, 178 num_threads=4, 179 use_gpu=True, 180 ) 181 182 boxes, masks, classes, scores = net(m) 183 184 draw_result(m, net.class_names, boxes, masks, classes, scores) 185