#from rknn.api import RKNN from rknnlite.api import RKNNLite import cv2 import numpy as np import cv2 import time import os from PIL import Image """ yolov5 预测脚本 for rknn """ SIZE = (640, 640) Width = 640 Height = 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]] IMAGE_EXT = [".jpg", "*.JPG", ".jpeg", ".webp", ".bmp", ".png"] def sigmoid(x): return 1 / (1 + np.exp(-x)) def letterbox_image(image, size): iw, ih = image.size w, h = size scale = min(w / iw, h / ih) nw = int(iw * scale) nh = int(ih * scale) image = np.array(image) image = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_LINEAR) image = Image.fromarray(image) new_image = Image.new('RGB', size, (128, 128, 128)) new_image.paste(image, ((w - nw) // 2, (h - nh) // 2)) return new_image 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) # Set inputs #image = Image.open(img_src) #img = letterbox_image(img_src, (Width, Height)) #img = np.array(img) 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 print("-------------------input0_data.shape----------------",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.) # 计算原尺寸的xy 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) print("------------------------boxes, classes, scores-----------------------",boxes, classes, scores) 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 def clip_coords(boxes, img_shape): # Clip bounding xyxy bounding boxes to image shape (height, width) #boxes[:, 0].clp(0, img_shape[1]) # x1 #boxes[:, 1].clp(0, img_shape[0]) # y1 #boxes[:, 2].clp(0, img_shape[1]) # x2 #boxes[:, 3].clp(0, img_shape[0]) # y2 np.clip(boxes[:, 0],0,img_shape[1]) np.clip(boxes[:, 1],0,img_shape[0]) np.clip(boxes[:, 2],0,img_shape[1]) np.clip(boxes[:, 3],0,img_shape[0]) return boxes def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None): # Rescale coords (xyxy) from img1_shape to img0_shape if ratio_pad is None: # calculate from img0_shape gain = min(img1_shape[0]/img0_shape[0], img1_shape[1]/img0_shape[1]) # gain = old / new print("------------gain-----------",gain) pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding else: gain = ratio_pad[0][0] pad = ratio_pad[1] print("-----------old-coords-----------",coords) coords[:, [2]] = (coords[:, [0]] + coords[:, [2]]) * img1_shape[1] - pad[0] # x padding coords[:, [3]] = (coords[:, [1]] + coords[:, [3]]) * img1_shape[0] - pad[1] # y padding coords[:, [0]] = coords[:, [0]] * img1_shape[1] - pad[0] # x padding coords[:, [1]] = coords[:, [1]] * img1_shape[0] - pad[1] # y padding print("-----------new-coords-----------",coords) print("------------pad-----------",pad) coords[:, :4] /= gain coords = clip_coords(coords, img0_shape) return coords def display(boxes=None, classes=None, scores=None, image_src=None, input_size=(640, 640), line_thickness=None, text_bg_alpha=0.0): labels = classes boxs = boxes confs = scores h, w, c = image_src.shape if len(boxes) <= 0: return image_src boxs[:, :] = scale_coords(input_size, boxs[:, :], (h, w)).round() tl = line_thickness or round(0.002 * (w + h) / 2) + 1 for i, box in enumerate(boxs): x1, y1, x2, y2 = box ratio = (y2-y1)/(x2-x1) x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) np.random.seed(int(labels[i]) + 2020) color = (np.random.randint(0, 255), 0, np.random.randint(0, 255)) cv2.rectangle(image_src, (x1, y1), (x2, y2), color, max(int((w + h) / 600), 1), cv2.LINE_AA) label = '{0:.3f}'.format(confs[i]) t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=1)[0] c2 = x1 + t_size[0] + 3, y1 - t_size[1] - 5 if text_bg_alpha == 0.0: cv2.rectangle(image_src, (x1 - 1, y1), c2, color, cv2.FILLED, cv2.LINE_AA) else: # 透明文本背景 alphaReserve = text_bg_alpha # 0：不透明 1：透明 BChannel, GChannel, RChannel = color xMin, yMin = int(x1 - 1), int(y1 - t_size[1] - 3) xMax, yMax = int(x1 + t_size[0]), int(y1) image_src[yMin:yMax, xMin:xMax, 0] = image_src[yMin:yMax, xMin:xMax, 0] * alphaReserve + BChannel * (1 - alphaReserve) image_src[yMin:yMax, xMin:xMax, 1] = image_src[yMin:yMax, xMin:xMax, 1] * alphaReserve + GChannel * (1 - alphaReserve) image_src[yMin:yMax, xMin:xMax, 2] = image_src[yMin:yMax, xMin:xMax, 2] * alphaReserve + RChannel * (1 - alphaReserve) cv2.putText(image_src, label, (x1 + 3, y1 - 4), 0, tl / 3, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA) return image_src 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: image_src = cv2.imread(image_name) #image_src = Image.open(image_name) boxes, classes, scores = predict(image_src) ''' 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) ''' image = np.array(image_src) save_image = display(boxes, classes, scores, image) save_image = cv2.cvtColor(save_image, cv2.COLOR_BGR2RGB) save_file_name = os.path.join(save_folder, os.path.basename(image_name)) cv2.imwrite(save_file_name,save_image) print("--------------------------res-----------------------",boxes, classes, scores)