400 lines
14 KiB
Python
400 lines
14 KiB
Python
|
#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)
|