#from rknn.api import RKNN from rknnlite.api import RKNNLite import cv2 import numpy as np import cv2 import time import os """ yolov5 预测脚本 for rknn """ SIZE = (640, 640) CLASSES = ("lighting") OBJ_THRESH = 0.1 NMS_THRESH = 0.1 MASKS = [[0, 1, 2], [3, 4, 5], [6, 7, 8]] ANCHORS = [[10, 13], [16, 30], [33, 23], [30, 61], [62, 45], [59, 119], [116, 90], [156, 198], [373, 326]] def sigmoid(x): return 1 / (1 + np.exp(-x)) IMAGE_EXT = [".jpg", "*.JPG", ".jpeg", ".webp", ".bmp", ".png"] def get_image_list(path): image_names = [] for maindir, subdir, file_name_list in os.walk(path): for filename in file_name_list: apath = os.path.join(maindir, filename) ext = os.path.splitext(apath)[1] if ext in IMAGE_EXT: image_names.append(apath) return image_names def filter_boxes(boxes, box_confidences, box_class_probs) -> (np.ndarray, np.ndarray, np.ndarray): box_scores = box_confidences * box_class_probs # 条件概率， 在该cell存在物体的概率的基础上是某个类别的概率 box_classes = np.argmax(box_scores, axis=-1) # 找出概率最大的类别索引 box_class_scores = np.max(box_scores, axis=-1) # 最大类别对应的概率值 pos = np.where(box_class_scores >= OBJ_THRESH) # 找出概率大于阈值的item # pos = box_class_scores >= OBJ_THRESH # 找出概率大于阈值的item boxes = boxes[pos] classes = box_classes[pos] scores = box_class_scores[pos] return boxes, classes, scores def nms_boxes(boxes, scores): x = boxes[:, 0] y = boxes[:, 1] w = boxes[:, 2] h = boxes[:, 3] 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 draw(image, boxes, scores, classes): """Draw the boxes on the image. # Argument: image: original image. boxes: ndarray, boxes of objects. classes: ndarray, classes of objects. scores: ndarray, scores of objects. all_classes: all classes name. """ labels = [] box_ls = [] for box, score, cl in zip(boxes, scores, classes): x, y, w, h = box print('class: {}, score: {}'.format(CLASSES[cl], score)) print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(x, y, x + w, y + h)) x *= image.shape[1] y *= image.shape[0] w *= image.shape[1] h *= image.shape[0] top = max(0, np.floor(x).astype(int)) left = max(0, np.floor(y).astype(int)) right = min(image.shape[1], np.floor(x + w + 0.5).astype(int)) bottom = min(image.shape[0], np.floor(y + h + 0.5).astype(int)) print('class: {}, score: {}'.format(CLASSES[cl], score)) print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom)) labels.append(CLASSES[cl]) box_ls.append((top, left, right, bottom)) 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) return labels, box_ls def load_model0(model_path, npu_id): rknn = RKNNLite() devs = rknn.list_devices() device_id_dict = {} for index, dev_id in enumerate(devs[-1]): if dev_id[:2] != 'TS': device_id_dict[0] = dev_id if dev_id[:2] == 'TS': device_id_dict[1] = dev_id print('-->loading model : ' + model_path) rknn.load_rknn(model_path) print('--> Init runtime environment on: ' + device_id_dict[npu_id]) ret = rknn.init_runtime(device_id=device_id_dict[npu_id]) if ret != 0: print('Init runtime environment failed') exit(ret) print('done') return rknn def load_rknn_model(PATH): # Create RKNN object rknn = RKNNLite() # Load tensorflow model print('--> Loading model') ret = rknn.load_rknn(PATH) if ret != 0: print('load rknn model failed') exit(ret) print('done') #ret = rknn.init_runtime(device_id='TS018083200400178', rknn2precompile=True) ret = rknn.init_runtime() if ret != 0: print('Init runtime environment failed') exit(ret) print('done') return rknn def predict(img_src, rknn): img = cv2.resize(img_src, SIZE) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) t0 = time.time() print("img shape \t:", img.shape) pred_onx = rknn.inference(inputs=[img]) print("time: \t", time.time() - t0) boxes, classes, scores = [], [], [] for t in range(3): input0_data = sigmoid(pred_onx[t][0]) input0_data = np.transpose(input0_data, (1, 2, 0, 3)) grid_h, grid_w, channel_n, predict_n = input0_data.shape anchors = [ANCHORS[i] for i in MASKS[t]] box_confidence = input0_data[..., 4] box_confidence = np.expand_dims(box_confidence, axis=-1) box_class_probs = input0_data[..., 5:] box_xy = input0_data[..., :2] box_wh = input0_data[..., 2:4] col = np.tile(np.arange(0, grid_w), grid_h).reshape(-1, grid_w) row = np.tile(np.arange(0, grid_h).reshape(-1, 1), grid_w) col = col.reshape((grid_h, grid_w, 1, 1)).repeat(3, axis=-2) row = row.reshape((grid_h, grid_w, 1, 1)).repeat(3, axis=-2) grid = np.concatenate((col, row), axis=-1) box_xy = box_xy * 2 - 0.5 + grid box_wh = (box_wh * 2) ** 2 * anchors box_xy /= (grid_w, grid_h) # 计算原尺寸的中心 box_wh /= SIZE # 计算原尺寸的宽高 box_xy -= (box_wh / 2.) # 计算原尺寸的中心 box = np.concatenate((box_xy, box_wh), axis=-1) res = filter_boxes(box, box_confidence, box_class_probs) boxes.append(res[0]) classes.append(res[1]) scores.append(res[2]) boxes, classes, scores = np.concatenate(boxes), np.concatenate(classes), np.concatenate(scores) print("------------------------boxes, classes, scores-----------------------",boxes, classes, scores) 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) keep = [0,1,2] print("--------------keep-------------",keep) nboxes.append(b[keep]) nclasses.append(c[keep]) nscores.append(s[keep]) if len(nboxes) < 1: return [], [], [] boxes = np.concatenate(nboxes) classes = np.concatenate(nclasses) scores = np.concatenate(nscores) return boxes, classes, scores ''' label_list = [] box_list = [] for box, score, cl in zip(boxes, scores, classes): x, y, w, h = box x *= img_src.shape[1] y *= img_src.shape[0] w *= img_src.shape[1] h *= img_src.shape[0] top = max(0, np.floor(x).astype(int)) left = max(0, np.floor(y).astype(int)) right = min(img_src.shape[1], np.floor(x + w + 0.5).astype(int)) bottom = min(img_src.shape[0], np.floor(y + h + 0.5).astype(int)) label_list.append(CLASSES[cl]) box_list.append((top, left, right, bottom)) return label_list, np.array(box_list) ''' def draw(image, boxes, scores, classes): """Draw the boxes on the image. # Argument: image: original image. boxes: ndarray, boxes of objects. classes: ndarray, classes of objects. scores: ndarray, scores of objects. all_classes: all classes name. """ for box, score, cl in zip(boxes, scores, classes): x, y, w, h = box #print('class: {}, score: {}'.format(CLASSES[cl], score)) #print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(x, y, x+w, y+h)) x *= image.shape[1] y *= image.shape[0] w *= image.shape[1] h *= image.shape[0] top = max(0, np.floor(x + 0.5).astype(int)) left = max(0, np.floor(y + 0.5).astype(int)) right = min(image.shape[1], np.floor(x + w + 0.5).astype(int)) bottom = min(image.shape[0], np.floor(y + h + 0.5).astype(int)) # print('class: {}, score: {}'.format(CLASSES[cl], score)) # print('box coordinate left,top,right,down: [{}, {}, {}, {}]'.format(top, left, right, bottom)) 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) return image if __name__ == '__main__': path = "./imgs/" save_folder = "./result/" #RKNN_MODEL_PATH = r"yolov5s-640-640.rknn" #RKNN_MODEL_PATH = r"best_640x640.rknn" RKNN_MODEL_PATH = r"23best_640x640.rknn" rknn = load_rknn_model(RKNN_MODEL_PATH) predict.__defaults__ = (None, rknn) files = get_image_list(path) current_time = time.localtime() for image_name in files: img = cv2.imread(image_name) boxes, classes, scores = predict(img) image = draw(img, boxes, scores, classes) save_file_name = os.path.join(save_folder, os.path.basename(image_name)) cv2.imwrite(save_file_name,image) print("--------------------------res-----------------------",boxes, classes, scores)