import numpy as np import cv2 from rknn.api import RKNN import math import PIL.Image as Image import PIL.ImageDraw as ImageDraw import PIL.ImageFont as ImageFont import re np.set_printoptions(threshold=np.inf) CLASSES = ('__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', '???', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', '???', 'backpack', 'umbrella', '???', '???', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', '???', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', '???', 'dining table', '???', '???', 'toilet', '???', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', '???', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush') NUM_CLS = 91 CONF_THRESH = 0.5 NMS_THRESH = 0.45 TOP_BOXES = 100 max_boxes_to_draw = 100 Y_SCALE = 10.0 X_SCALE = 10.0 H_SCALE = 5.0 W_SCALE = 5.0 prior_file = './box_priors.txt' box_priors_ = [] fp = open(prior_file, 'r') ls = fp.readlines() for s in ls: aList = re.findall('([-+]?\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?', s) for ss in aList: aNum = float((ss[0] + ss[2])) box_priors_.append(aNum) fp.close() def softmax(x): return np.exp(x) / np.sum(np.exp(x), axis=0) 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 + 0.000001) return ovr def ssd_post_process(conf_data, loc_data, imgpath, output_dir='.'): prior_num = int(len(loc_data) / 4) # num prior boxes prior_bboxes = np.array(box_priors_) prior_bboxes = prior_bboxes.reshape(4, prior_num) conf_data = conf_data.reshape(-1, NUM_CLS) for i in range(prior_num): conf_data[i] = softmax(conf_data[i]) idx_class_conf = [] bboxes = [] # conf for prior_idx in range(0, prior_num): conf_data[prior_idx][0] = 0 max_val = np.max(conf_data[prior_idx]) max_idx = np.argmax(conf_data[prior_idx]) if max_val > CONF_THRESH: idx_class_conf.append([prior_idx, max_idx, max_val]) idx_class_conf_sorted = sorted(idx_class_conf, key=lambda x: x[2], reverse=True) idx_class_conf = idx_class_conf_sorted[:min(TOP_BOXES, len(idx_class_conf_sorted))] # boxes for i in range(0, prior_num): bbox_center_x = loc_data[4 * i + 1] / X_SCALE * prior_bboxes[3][i] + prior_bboxes[1][i] bbox_center_y = loc_data[4 * i + 0] / Y_SCALE * prior_bboxes[2][i] + prior_bboxes[0][i] bbox_w = math.exp(loc_data[4 * i + 3] / W_SCALE) * prior_bboxes[3][i] bbox_h = math.exp(loc_data[4 * i + 2] / H_SCALE) * prior_bboxes[2][i] 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) # 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: keep = False break k += 1 else: break if keep: idx_class_conf_.append(idx_class_conf[i]) cur_class_num += 1 idx_class_conf_ = idx_class_conf_[:min(len(idx_class_conf_), max_boxes_to_draw)] 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(imgpath) img_pil = Image.fromarray(img) draw = ImageDraw.Draw(img_pil) font = ImageFont.load_default() name = imgpath.split("/")[-1][:-4] for i in range(0, len(box_class_score)): x1 = box_class_score[i][0] * img.shape[1] y1 = box_class_score[i][1] * img.shape[0] x2 = box_class_score[i][2] * img.shape[1] y2 = box_class_score[i][3] * img.shape[0] # draw rect 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('%.2f' % 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) print('write output image: {}{}_quant.jpg'.format(output_dir, name)) np.copyto(img, np.array(img_pil)) cv2.imwrite("{}{}_quant.jpg".format(output_dir, name), img) print('write output image finished.')