You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
211 lines
7.2 KiB
211 lines
7.2 KiB
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__', |
|
'aeroplane', 'bicycle', 'bird', 'boat', |
|
'bottle', 'bus', 'car', 'cat', 'chair', |
|
'cow', 'diningtable', 'dog', 'horse', |
|
'motorbike', 'person', 'pottedplant', |
|
'sheep', 'sofa', 'train', 'tvmonitor') |
|
|
|
NUM_CLS = 21 |
|
|
|
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('mbox_priorbox_97.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,21) |
|
|
|
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(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('./road_300x300.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__': |
|
|
|
if not os.path.exists('./VGG_VOC0712_SSD_300x300_iter_120000.caffemodel'): |
|
print('!!! Missing VGG_VOC0712_SSD_300x300_iter_120000.caffemodel !!!\n' \ |
|
'1. Download models_VGGNet_VOC0712_SSD_300x300.tar.gz from https://drive.google.com/file/d/0BzKzrI_SkD1_WVVTSmQxU0dVRzA/view\n' \ |
|
'2. Extract the VGG_VOC0712_SSD_300x300_iter_120000.caffemodel from models_VGGNet_VOC0712_SSD_300x300.tar.gz\n' \ |
|
'3. Or you can also download caffemodel from https://eyun.baidu.com/s/3jJhPRzo , password is rknn\n') |
|
exit(-1) |
|
|
|
# Create RKNN object |
|
rknn = RKNN() |
|
|
|
# pre-process config |
|
print('--> config model') |
|
rknn.config(mean_values=[103.94, 116.78, 123.68], std_values=[1, 1, 1], quant_img_RGB2BGR=True) |
|
print('done') |
|
|
|
# Load tensorflow model |
|
print('--> Loading model') |
|
ret = rknn.load_caffe(model='./deploy_rm_detection_output.prototxt', |
|
proto='caffe', |
|
blobs='./VGG_VOC0712_SSD_300x300_iter_120000.caffemodel') |
|
if ret != 0: |
|
print('Load model failed! Ret = {}'.format(ret)) |
|
exit(ret) |
|
print('done') |
|
|
|
# Build model |
|
print('--> Building model') |
|
ret = rknn.build(do_quantization=True, 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('./deploy_rm_detection_output.rknn') |
|
if ret != 0: |
|
print('Export rknn failed!') |
|
exit(ret) |
|
print('done') |
|
|
|
# Set inputs |
|
img = cv2.imread('./road_300x300.jpg') |
|
|
|
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('done') |
|
|
|
outputs[0] = outputs[0].reshape((-1, 1)) |
|
outputs[1] = outputs[1].reshape((-1, 1)) |
|
ssd_post_process(outputs[1], outputs[0]) |
|
|
|
rknn.release()
|
|
|