275 lines
10 KiB
Python
275 lines
10 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 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 = 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
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
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"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:
|
|||
|
|
print("--------------------------image_name-----------------------", image_name)
|
|||
|
|
image_src = Image.open(image_name)
|
|||
|
|
boxes, classes, scores = predict(image_src)
|
|||
|
|
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)
|