关于裁剪后的mobilenet ssd模型转换问题 - Toybrick

import os
import math
import numpy as np
import cv2
from rknn.api import RKNN
import PIL.Image as Image
import PIL.ImageDraw as ImageDraw
import PIL.ImageFont as ImageFont
np.set_printoptions(threshold=np.inf)
CLASSES = ('background','head')
NUM_CLS = 2
CONF_THRESH = 0.5
NMS_THRESH = 0.45
def IntersectBBox(box1, box2):
if box1[0]>box2[2] or box1[2]<box2[0] or box1[1]>box2[3] or box1[3]<box2[1]:
return 0
else:
area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
xx1 = max(box1[0], box2[0])
yy1 = max(box1[1], box2[1])
xx2 = min(box1[2], box2[2])
yy2 = min(box1[3], box2[3])
w = max(0, xx2-xx1)
h = max(0, yy2-yy1)
ovr = w*h / (area1 + area2 - w*h)
return ovr
def ssd_post_process(conf_data, loc_data):
prior_data = np.loadtxt('priorbox_ylc.txt', dtype = np.float32)
prior_bboxes = prior_data[:len(loc_data)]
prior_variances = prior_data[len(loc_data):]
prior_num = int(len(loc_data) / 4) # 8732
conf_data = conf_data.reshape(-1,2)
idx_class_conf = []
bboxes = []
# conf
for prior_idx in range(0,prior_num):
max_val = np.max(conf_data[prior_idx])
max_idx = np.argmax(conf_data[prior_idx])
if max_val > CONF_THRESH and max_idx != 0:
idx_class_conf.append([prior_idx, max_idx, max_val])
#print(len(idx_class_conf))
# boxes
for i in range(0,prior_num):
prior_w = prior_bboxes[4*i+2] - prior_bboxes[4*i]
prior_h = prior_bboxes[4*i+3] - prior_bboxes[4*i+1]
prior_center_x = (prior_bboxes[4*i+2] + prior_bboxes[4*i]) / 2
prior_center_y = (prior_bboxes[4*i+3] + prior_bboxes[4*i+1]) / 2
bbox_center_x = prior_variances[4*i+0] * loc_data[4*i+0] * prior_w + prior_center_x
bbox_center_y = prior_variances[4*i+1] * loc_data[4*i+1] * prior_h + prior_center_y
bbox_w = math.exp(prior_variances[4*i+2] * loc_data[4*i+2]) * prior_w
bbox_h = math.exp(prior_variances[4*i+3] * loc_data[4*i+3]) * prior_h
tmp = []
tmp.append(max(min(bbox_center_x - bbox_w / 2., 1), 0))
tmp.append(max(min(bbox_center_y - bbox_h / 2., 1), 0))
tmp.append(max(min(bbox_center_x + bbox_w / 2., 1), 0))
tmp.append(max(min(bbox_center_y + bbox_h / 2., 1), 0))
bboxes.append(tmp)
print('idx_class_conf = ',len(idx_class_conf))
#nms
cur_class_num = 0
idx_class_conf_ = []
for i in range(0, len(idx_class_conf)):
keep = True
k = 0
while k < cur_class_num:
if keep:
ovr = IntersectBBox(bboxes[idx_class_conf[i][0]], bboxes[idx_class_conf_[k][0]])
if idx_class_conf_[k][1]==idx_class_conf[i][1] and ovr > NMS_THRESH:
if idx_class_conf_[k][2]<idx_class_conf[i][2]:
idx_class_conf_.pop(k)
idx_class_conf_.append(idx_class_conf[i])
keep = False
break
k += 1
else:
break
if keep:
idx_class_conf_.append(idx_class_conf[i])
cur_class_num += 1
print(idx_class_conf_)
box_class_score = []
for i in range(0, len(idx_class_conf_)):
bboxes[idx_class_conf_[i][0]].append(idx_class_conf_[i][1])
bboxes[idx_class_conf_[i][0]].append(idx_class_conf_[i][2])
box_class_score.append( bboxes[idx_class_conf_[i][0]])
img = cv2.imread('./123.jpg')
img_pil = Image.fromarray(img)
draw = ImageDraw.Draw(img_pil)
font = ImageFont.load_default()
for i in range(0, len(box_class_score)):
x1 = int(box_class_score[i][0]*img.shape[1])
y1 = int(box_class_score[i][1]*img.shape[0])
x2 = int(box_class_score[i][2]*img.shape[1])
y2 = int(box_class_score[i][3]*img.shape[0])
color = (0, int(box_class_score[i][4]/20.0*255), 255)
draw.line([(x1, y1), (x1, y2), (x2 , y2),
(x2 , y1), (x1, y1)], width=2, fill=color)
display_str = CLASSES[box_class_score[i][4]] + ":" + str(box_class_score[i][5])
display_str_height = np.ceil((1 + 2 * 0.05) * font.getsize(display_str)[1])+1
if y1 > display_str_height:
text_bottom = y1
else:
text_bottom = y1 + display_str_height
text_width, text_height = font.getsize(display_str)
margin = np.ceil(0.05 * text_height)
draw.rectangle([(x1, text_bottom-text_height-2*margin), (x1+text_width, text_bottom)],fill=color)
draw.text((x1+margin, text_bottom-text_height-margin), display_str, fill='black', font=font)
np.copyto(img, np.array(img_pil))
cv2.imwrite("result.jpg", img)
if __name__ == '__main__':
# Create RKNN object
rknn = RKNN(verbose=False)
# pre-process config
print('--> config model')
rknn.config(channel_mean_value='103.94 116.78 123.68 1', reorder_channel='2 1 0')
print('done')
# Load tensorflow model
print('--> Loading model')
ret = rknn.load_caffe(model='./mobilenet_v2_ssd_truncated2.prototxt',
proto='caffe',
blobs='./mobilenet_iter_38200.caffemodel')
if ret != 0:
print("---------------------------------------------")
print('Load model failed! Ret = {}'.format(ret))
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False, dataset='./dataset.txt')
if ret != 0:
print('Build model failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export RKNN model')
ret = rknn.export_rknn('./mobilenet_v2_ssd.rknn')
if ret != 0:
print('Export rknn failed!')
exit(ret)
print('done')
# Set inputs
img = cv2.imread('./123.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
print('--> Init runtime environment')
ret = rknn.init_runtime()
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
# Inference
print('--> Running model')
outputs = rknn.inference(inputs=[img])
print('outputs = ',len(outputs),type(outputs),len(outputs[0].reshape((-1, 1))),len(outputs[1].reshape((-1, 1))))
print('done')
outputs[0] = outputs[0].reshape((-1, 1))
outputs[1] = outputs[1].reshape((-1, 1))
ssd_post_process(outputs[1], outputs[0])
rknn.release()