import cv2 import os import numpy as np from rknnlite.api import RKNNLite import sys from copy import copy OBJ_THRESH = 0.25 NMS_THRESH = 0.45 # The follew two param is for map test # OBJ_THRESH = 0.001 # NMS_THRESH = 0.65 IMG_SIZE = (640, 640) # (width, height), such as (1280, 736) CLASSES = ("person", "bicycle", "car","motorbike ","aeroplane ","bus ","train","truck ","boat","traffic light", "fire hydrant","stop sign ","parking meter","bench","bird","cat","dog ","horse ","sheep","cow","elephant", "bear","zebra ","giraffe","backpack","umbrella","handbag","tie","suitcase","frisbee","skis","snowboard","sports ball","kite", "baseball bat","baseball glove","skateboard","surfboard","tennis racket","bottle","wine glass","cup","fork","knife ", "spoon","bowl","banana","apple","sandwich","orange","broccoli","carrot","hot dog","pizza ","donut","cake","chair","sofa", "pottedplant","bed","diningtable","toilet ","tvmonitor","laptop ","mouse ","remote ","keyboard ","cell phone","microwave ", "oven ","toaster","sink","refrigerator ","book","clock","vase","scissors ","teddy bear ","hair drier", "toothbrush ") coco_id_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] def get_real_box(self, box, in_format='xyxy'): bbox = copy(box) if self.enable_ltter_box == True: # unletter_box result if in_format=='xyxy': bbox[:,0] -= self.letter_box_info_list[-1].dw bbox[:,0] /= self.letter_box_info_list[-1].w_ratio bbox[:,0] = np.clip(bbox[:,0], 0, self.letter_box_info_list[-1].origin_shape[1]) bbox[:,1] -= self.letter_box_info_list[-1].dh bbox[:,1] /= self.letter_box_info_list[-1].h_ratio bbox[:,1] = np.clip(bbox[:,1], 0, self.letter_box_info_list[-1].origin_shape[0]) bbox[:,2] -= self.letter_box_info_list[-1].dw bbox[:,2] /= self.letter_box_info_list[-1].w_ratio bbox[:,2] = np.clip(bbox[:,2], 0, self.letter_box_info_list[-1].origin_shape[1]) bbox[:,3] -= self.letter_box_info_list[-1].dh bbox[:,3] /= self.letter_box_info_list[-1].h_ratio bbox[:,3] = np.clip(bbox[:,3], 0, self.letter_box_info_list[-1].origin_shape[0]) return bbox def filter_boxes(boxes, box_confidences, box_class_probs): """Filter boxes with object threshold. """ box_confidences = box_confidences.reshape(-1) class_max_score = np.max(box_class_probs, axis=-1) classes = np.argmax(box_class_probs, axis=-1) _class_pos = np.where(class_max_score* box_confidences >= OBJ_THRESH) scores = (class_max_score* box_confidences)[_class_pos] boxes = boxes[_class_pos] classes = classes[_class_pos] return boxes, classes, scores def nms_boxes(boxes, scores): """Suppress non-maximal boxes. # Returns keep: ndarray, index of effective boxes. """ x = boxes[:, 0] y = boxes[:, 1] w = boxes[:, 2] - boxes[:, 0] h = boxes[:, 3] - boxes[:, 1] areas = w * h order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(i) xx1 = np.maximum(x[i], x[order[1:]]) yy1 = np.maximum(y[i], y[order[1:]]) xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]]) yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]]) w1 = np.maximum(0.0, xx2 - xx1 + 0.00001) h1 = np.maximum(0.0, yy2 - yy1 + 0.00001) inter = w1 * h1 ovr = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(ovr <= NMS_THRESH)[0] order = order[inds + 1] keep = np.array(keep) return keep def box_process(position, anchors): grid_h, grid_w = position.shape[2:4] col, row = np.meshgrid(np.arange(0, grid_w), np.arange(0, grid_h)) col = col.reshape(1, 1, grid_h, grid_w) row = row.reshape(1, 1, grid_h, grid_w) grid = np.concatenate((col, row), axis=1) stride = np.array([IMG_SIZE[1]//grid_h, IMG_SIZE[0]//grid_w]).reshape(1,2,1,1) col = col.repeat(len(anchors), axis=0) row = row.repeat(len(anchors), axis=0) anchors = np.array(anchors) anchors = anchors.reshape(*anchors.shape, 1, 1) box_xy = position[:,:2,:,:]*2 - 0.5 box_wh = pow(position[:,2:4,:,:]*2, 2) * anchors box_xy += grid box_xy *= stride box = np.concatenate((box_xy, box_wh), axis=1) # Convert [c_x, c_y, w, h] to [x1, y1, x2, y2] xyxy = np.copy(box) xyxy[:, 0, :, :] = box[:, 0, :, :] - box[:, 2, :, :]/ 2 # top left x xyxy[:, 1, :, :] = box[:, 1, :, :] - box[:, 3, :, :]/ 2 # top left y xyxy[:, 2, :, :] = box[:, 0, :, :] + box[:, 2, :, :]/ 2 # bottom right x xyxy[:, 3, :, :] = box[:, 1, :, :] + box[:, 3, :, :]/ 2 # bottom right y return xyxy def post_process(input_data, anchors): boxes, scores, classes_conf = [], [], [] # 1*255*h*w -> 3*85*h*w input_data = [_in.reshape([len(anchors[0]),-1]+list(_in.shape[-2:])) for _in in input_data] for i in range(len(input_data)): boxes.append(box_process(input_data[i][:,:4,:,:], anchors[i])) scores.append(input_data[i][:,4:5,:,:]) classes_conf.append(input_data[i][:,5:,:,:]) def sp_flatten(_in): ch = _in.shape[1] _in = _in.transpose(0,2,3,1) return _in.reshape(-1, ch) boxes = [sp_flatten(_v) for _v in boxes] classes_conf = [sp_flatten(_v) for _v in classes_conf] scores = [sp_flatten(_v) for _v in scores] boxes = np.concatenate(boxes) classes_conf = np.concatenate(classes_conf) scores = np.concatenate(scores) # filter according to threshold boxes, classes, scores = filter_boxes(boxes, scores, classes_conf) # nms nboxes, nclasses, nscores = [], [], [] for c in set(classes): inds = np.where(classes == c) b = boxes[inds] c = classes[inds] s = scores[inds] keep = nms_boxes(b, s) if len(keep) != 0: nboxes.append(b[keep]) nclasses.append(c[keep]) nscores.append(s[keep]) if not nclasses and not nscores: return None, None, None boxes = np.concatenate(nboxes) classes = np.concatenate(nclasses) scores = np.concatenate(nscores) return boxes, classes, scores def draw(image, boxes, scores, classes): for box, score, cl in zip(boxes, scores, classes): top, left, right, bottom = [int(_b) for _b in box] print("%s @ (%d %d %d %d) %.3f" % (CLASSES[cl], top, left, right, bottom, score)) cv2.rectangle(image, (top, left), (right, bottom), (255, 0, 0), 2) cv2.putText(image, '{0} {1:.2f}'.format(CLASSES[cl], score), (top, left - 6), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2) INPUT_SIZE = 640 if __name__ == '__main__': rknn_lite = RKNNLite() model = sys.argv[1] image_name = sys.argv[2] target = None # load RKNN model print('--> Load RKNN model') ret = rknn_lite.load_rknn(model) if ret != 0: print('Load RKNN model failed') exit(ret) print('done') ori_img = cv2.imread(image_name) img = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (IMG_SIZE,IMG_SIZE)) # init runtime environment print('--> Init runtime environment') ret = rknn_lite.init_runtime(target=target) if ret != 0: print('Init runtime environment failed') exit(ret) print('done') with open("anchors_yolov5.txt", 'r') as f: values = [float(_v) for _v in f.readlines()] anchors = np.array(values).reshape(3,-1,2).tolist() print("use anchors from '{}', which is {}".format("anchors_yolov5.txt", anchors)) # Inference print('--> Running model') outputs = rknn_lite.inference(inputs=[img]) boxes, classes, scores = post_process(outputs, anchors) img_p = ori_img.copy() if boxes is not None: draw(img_p, get_real_box(boxes), scores, classes) if not os.path.exists('./result'): os.mkdir('./result') result_path = os.path.join('./result', image_name) cv2.imwrite(result_path, img_p) print('Detection result save to {}'.format(result_path)) rknn_lite.release()