222 lines
6.1 KiB
Python
222 lines
6.1 KiB
Python
import sys
|
|
import os
|
|
import time
|
|
import math
|
|
import numpy as np
|
|
|
|
import itertools
|
|
import struct # get_image_size
|
|
import imghdr # get_image_size
|
|
|
|
|
|
def sigmoid(x):
|
|
return 1.0 / (np.exp(-x) + 1.)
|
|
|
|
|
|
def softmax(x):
|
|
x = np.exp(x - np.expand_dims(np.max(x, axis=1), axis=1))
|
|
x = x / np.expand_dims(x.sum(axis=1), axis=1)
|
|
return x
|
|
|
|
|
|
def bbox_iou(box1, box2, x1y1x2y2=True):
|
|
|
|
# print('iou box1:', box1)
|
|
# print('iou box2:', box2)
|
|
|
|
if x1y1x2y2:
|
|
mx = min(box1[0], box2[0])
|
|
Mx = max(box1[2], box2[2])
|
|
my = min(box1[1], box2[1])
|
|
My = max(box1[3], box2[3])
|
|
w1 = box1[2] - box1[0]
|
|
h1 = box1[3] - box1[1]
|
|
w2 = box2[2] - box2[0]
|
|
h2 = box2[3] - box2[1]
|
|
else:
|
|
w1 = box1[2]
|
|
h1 = box1[3]
|
|
w2 = box2[2]
|
|
h2 = box2[3]
|
|
|
|
mx = min(box1[0], box2[0])
|
|
Mx = max(box1[0] + w1, box2[0] + w2)
|
|
my = min(box1[1], box2[1])
|
|
My = max(box1[1] + h1, box2[1] + h2)
|
|
uw = Mx - mx
|
|
uh = My - my
|
|
cw = w1 + w2 - uw
|
|
ch = h1 + h2 - uh
|
|
carea = 0
|
|
if cw <= 0 or ch <= 0:
|
|
return 0.0
|
|
|
|
area1 = w1 * h1
|
|
area2 = w2 * h2
|
|
carea = cw * ch
|
|
uarea = area1 + area2 - carea
|
|
return carea / uarea
|
|
|
|
|
|
def nms_cpu(boxes, confs, nms_thresh=0.5, min_mode=False):
|
|
# print(boxes.shape)
|
|
x1 = boxes[:, 0]
|
|
y1 = boxes[:, 1]
|
|
x2 = boxes[:, 2]
|
|
y2 = boxes[:, 3]
|
|
|
|
areas = (x2 - x1) * (y2 - y1)
|
|
order = confs.argsort()[::-1]
|
|
|
|
keep = []
|
|
while order.size > 0:
|
|
idx_self = order[0]
|
|
idx_other = order[1:]
|
|
|
|
keep.append(idx_self)
|
|
|
|
xx1 = np.maximum(x1[idx_self], x1[idx_other])
|
|
yy1 = np.maximum(y1[idx_self], y1[idx_other])
|
|
xx2 = np.minimum(x2[idx_self], x2[idx_other])
|
|
yy2 = np.minimum(y2[idx_self], y2[idx_other])
|
|
|
|
w = np.maximum(0.0, xx2 - xx1)
|
|
h = np.maximum(0.0, yy2 - yy1)
|
|
inter = w * h
|
|
|
|
if min_mode:
|
|
over = inter / np.minimum(areas[order[0]], areas[order[1:]])
|
|
else:
|
|
over = inter / (areas[order[0]] + areas[order[1:]] - inter)
|
|
|
|
inds = np.where(over <= nms_thresh)[0]
|
|
order = order[inds + 1]
|
|
|
|
return np.array(keep)
|
|
|
|
|
|
|
|
def plot_boxes_cv2(img, boxes, savename=None, class_names=None, color=None):
|
|
import cv2
|
|
img = np.copy(img)
|
|
colors = np.array([[1, 0, 1], [0, 0, 1], [0, 1, 1], [0, 1, 0], [1, 1, 0], [1, 0, 0]], dtype=np.float32)
|
|
|
|
def get_color(c, x, max_val):
|
|
ratio = float(x) / max_val * 5
|
|
i = int(math.floor(ratio))
|
|
j = int(math.ceil(ratio))
|
|
ratio = ratio - i
|
|
r = (1 - ratio) * colors[i][c] + ratio * colors[j][c]
|
|
return int(r * 255)
|
|
|
|
width = img.shape[1]
|
|
height = img.shape[0]
|
|
for i in range(len(boxes)):
|
|
box = boxes[i]
|
|
x1 = int(box[0] * width)
|
|
y1 = int(box[1] * height)
|
|
x2 = int(box[2] * width)
|
|
y2 = int(box[3] * height)
|
|
|
|
if color:
|
|
rgb = color
|
|
else:
|
|
rgb = (255, 0, 0)
|
|
if len(box) >= 7 and class_names:
|
|
cls_conf = box[5]
|
|
cls_id = box[6]
|
|
print('%s: %f' % (class_names[cls_id], cls_conf))
|
|
classes = len(class_names)
|
|
offset = cls_id * 123457 % classes
|
|
red = get_color(2, offset, classes)
|
|
green = get_color(1, offset, classes)
|
|
blue = get_color(0, offset, classes)
|
|
if color is None:
|
|
rgb = (red, green, blue)
|
|
img = cv2.putText(img, class_names[cls_id], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 1.2, rgb, 1)
|
|
img = cv2.rectangle(img, (x1, y1), (x2, y2), rgb, 1)
|
|
if savename:
|
|
print("save plot results to %s" % savename)
|
|
cv2.imwrite(savename, img)
|
|
return img
|
|
|
|
|
|
def read_truths(lab_path):
|
|
if not os.path.exists(lab_path):
|
|
return np.array([])
|
|
if os.path.getsize(lab_path):
|
|
truths = np.loadtxt(lab_path)
|
|
truths = truths.reshape(truths.size / 5, 5) # to avoid single truth problem
|
|
return truths
|
|
else:
|
|
return np.array([])
|
|
|
|
def post_processing(img, conf_thresh, nms_thresh, output):
|
|
|
|
# anchors = [12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243, 459, 401]
|
|
# num_anchors = 9
|
|
# anchor_masks = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
|
|
# strides = [8, 16, 32]
|
|
# anchor_step = len(anchors) // num_anchors
|
|
|
|
# [batch, num, 1, 4]
|
|
box_array = output[0]
|
|
# [batch, num, num_classes]
|
|
confs = output[1]
|
|
|
|
t1 = time.time()
|
|
|
|
if type(box_array).__name__ != 'ndarray':
|
|
box_array = box_array.cpu().detach().numpy()
|
|
confs = confs.cpu().detach().numpy()
|
|
|
|
num_classes = confs.shape[2]
|
|
|
|
# [batch, num, 4]
|
|
box_array = box_array[:, :, 0]
|
|
|
|
# [batch, num, num_classes] --> [batch, num]
|
|
max_conf = np.max(confs, axis=2)
|
|
max_id = np.argmax(confs, axis=2)
|
|
|
|
t2 = time.time()
|
|
|
|
bboxes_batch = []
|
|
for i in range(box_array.shape[0]):
|
|
|
|
argwhere = max_conf[i] > conf_thresh
|
|
l_box_array = box_array[i, argwhere, :]
|
|
l_max_conf = max_conf[i, argwhere]
|
|
l_max_id = max_id[i, argwhere]
|
|
|
|
bboxes = []
|
|
# nms for each class
|
|
for j in range(num_classes):
|
|
|
|
cls_argwhere = l_max_id == j
|
|
ll_box_array = l_box_array[cls_argwhere, :]
|
|
ll_max_conf = l_max_conf[cls_argwhere]
|
|
ll_max_id = l_max_id[cls_argwhere]
|
|
|
|
keep = nms_cpu(ll_box_array, ll_max_conf, nms_thresh)
|
|
|
|
if (keep.size > 0):
|
|
ll_box_array = ll_box_array[keep, :]
|
|
ll_max_conf = ll_max_conf[keep]
|
|
ll_max_id = ll_max_id[keep]
|
|
|
|
for k in range(ll_box_array.shape[0]):
|
|
bboxes.append([ll_box_array[k, 0], ll_box_array[k, 1], ll_box_array[k, 2], ll_box_array[k, 3], ll_max_conf[k], ll_max_conf[k], ll_max_id[k]])
|
|
|
|
bboxes_batch.append(bboxes)
|
|
|
|
t3 = time.time()
|
|
|
|
print('-----------------------------------')
|
|
print(' max and argmax : %f' % (t2 - t1))
|
|
print(' nms : %f' % (t3 - t2))
|
|
print('Post processing total : %f' % (t3 - t1))
|
|
print('-----------------------------------')
|
|
|
|
return bboxes_batch
|