279 lines
9.8 KiB
Python
279 lines
9.8 KiB
Python
#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)
|