166 lines
5.9 KiB
Python
166 lines
5.9 KiB
Python
# coding: utf-8
|
||
# 2019-12-10
|
||
"""
|
||
YOlo相关的预处理api；
|
||
"""
|
||
import cv2
|
||
import time
|
||
import numpy as np
|
||
|
||
|
||
# 加载label names；
|
||
def get_labels(names_file):
|
||
names = list()
|
||
with open(names_file, 'r') as f:
|
||
lines = f.read()
|
||
for name in lines.splitlines():
|
||
names.append(name)
|
||
f.close()
|
||
return names
|
||
|
||
|
||
# 照片预处理
|
||
def process_img(img_path, input_shape):
|
||
ori_img = cv2.imread(img_path)
|
||
img = cv2.resize(ori_img, input_shape)
|
||
image = img[:, :, ::-1].transpose((2, 0, 1))
|
||
image = image[np.newaxis, :, :, :] / 255
|
||
image = np.array(image, dtype=np.float32)
|
||
return ori_img, ori_img.shape, image
|
||
|
||
|
||
# 视频预处理
|
||
def frame_process(frame, input_shape):
|
||
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
image = cv2.resize(image, input_shape)
|
||
# image = cv2.resize(image, (640, 480))
|
||
image_mean = np.array([127, 127, 127])
|
||
image = (image - image_mean) / 128
|
||
image = np.transpose(image, [2, 0, 1])
|
||
image = np.expand_dims(image, axis=0)
|
||
image = image.astype(np.float32)
|
||
return image
|
||
|
||
|
||
# sigmoid函数
|
||
def sigmoid(x):
|
||
s = 1 / (1 + np.exp(-1 * x))
|
||
return s
|
||
|
||
|
||
# 获取预测正确的类别，以及概率和索引;
|
||
def get_result(class_scores):
|
||
class_score = 0
|
||
class_index = 0
|
||
for i in range(len(class_scores)):
|
||
if class_scores[i] > class_score:
|
||
class_index += 1
|
||
class_score = class_scores[i]
|
||
return class_score, class_index
|
||
|
||
|
||
# 通过置信度筛选得到bboxs
|
||
def get_bbox(feat, anchors, image_shape, confidence_threshold=0.25):
|
||
box = list()
|
||
for i in range(len(anchors)):
|
||
for cx in range(feat.shape[0]):
|
||
for cy in range(feat.shape[1]):
|
||
tx = feat[cx][cy][0 + 85 * i]
|
||
ty = feat[cx][cy][1 + 85 * i]
|
||
tw = feat[cx][cy][2 + 85 * i]
|
||
th = feat[cx][cy][3 + 85 * i]
|
||
cf = feat[cx][cy][4 + 85 * i]
|
||
cp = feat[cx][cy][5 + 85 * i:85 + 85 * i]
|
||
|
||
bx = (sigmoid(tx) + cx) / feat.shape[0]
|
||
by = (sigmoid(ty) + cy) / feat.shape[1]
|
||
bw = anchors[i][0] * np.exp(tw) / image_shape[0]
|
||
bh = anchors[i][1] * np.exp(th) / image_shape[1]
|
||
b_confidence = sigmoid(cf)
|
||
b_class_prob = sigmoid(cp)
|
||
b_scores = b_confidence * b_class_prob
|
||
b_class_score, b_class_index = get_result(b_scores)
|
||
|
||
if b_class_score >= confidence_threshold:
|
||
box.append([bx, by, bw, bh, b_class_score, b_class_index])
|
||
return box
|
||
|
||
|
||
# 采用nms算法筛选获取到的bbox
|
||
def nms(boxes, nms_threshold=0.6):
|
||
l = len(boxes)
|
||
if l == 0:
|
||
return []
|
||
else:
|
||
b_x = boxes[:, 0]
|
||
b_y = boxes[:, 1]
|
||
b_w = boxes[:, 2]
|
||
b_h = boxes[:, 3]
|
||
scores = boxes[:, 4]
|
||
areas = (b_w + 1) * (b_h + 1)
|
||
order = scores.argsort()[::-1]
|
||
keep = list()
|
||
while order.size > 0:
|
||
i = order[0]
|
||
keep.append(i)
|
||
xx1 = np.maximum(b_x[i], b_x[order[1:]])
|
||
yy1 = np.maximum(b_y[i], b_y[order[1:]])
|
||
xx2 = np.minimum(b_x[i] + b_w[i], b_x[order[1:]] + b_w[order[1:]])
|
||
yy2 = np.minimum(b_y[i] + b_h[i], b_y[order[1:]] + b_h[order[1:]])
|
||
|
||
# 相交面积,不重叠时面积为0
|
||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||
inter = w * h
|
||
# 相并面积,面积1+面积2-相交面积
|
||
union = areas[i] + areas[order[1:]] - inter
|
||
# 计算IoU：交 /（面积1+面积2-交）
|
||
IoU = inter / union
|
||
# 保留IoU小于阈值的box
|
||
inds = np.where(IoU <= nms_threshold)[0]
|
||
order = order[inds + 1] # 因为IoU数组的长度比order数组少一个,所以这里要将所有下标后移一位
|
||
|
||
final_boxes = [boxes[i] for i in keep]
|
||
return final_boxes
|
||
|
||
|
||
# 绘制预测框
|
||
def draw_box(boxes, img, img_shape):
|
||
label = ["background", "person",
|
||
"bicycle", "car", "motorbike", "aeroplane",
|
||
"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", "sofa", "potted plant", "bed", "dining table",
|
||
"toilet", "TV monitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
|
||
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
|
||
"scissors", "teddy bear", "hair drier", "toothbrush"]
|
||
for box in boxes:
|
||
x1 = int((box[0] - box[2] / 2) * img_shape[1])
|
||
y1 = int((box[1] - box[3] / 2) * img_shape[0])
|
||
x2 = int((box[0] + box[2] / 2) * img_shape[1])
|
||
y2 = int((box[1] + box[3] / 2) * img_shape[0])
|
||
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
cv2.putText(img, label[int(box[5])] + ":" + str(round(box[4], 3)), (x1 + 5, y1 + 10), cv2.FONT_HERSHEY_SIMPLEX,
|
||
0.5, (0, 0, 255), 1)
|
||
print(label[int(box[5])] + ":" + "概率值：%.3f" % box[4])
|
||
cv2.imshow('image', img)
|
||
cv2.waitKey(10)
|
||
cv2.destroyAllWindows()
|
||
|
||
|
||
# 获取预测框
|
||
def get_boxes(prediction, anchors, img_shape, confidence_threshold=0.25, nms_threshold=0.6):
|
||
boxes = []
|
||
for i in range(len(prediction)):
|
||
feature_map = prediction[i][0].transpose((2, 1, 0))
|
||
box = get_bbox(feature_map, anchors[i], img_shape, confidence_threshold)
|
||
boxes.extend(box)
|
||
Boxes = nms(np.array(boxes), nms_threshold)
|
||
return Boxes
|