410 lines
14 KiB
Python
410 lines
14 KiB
Python
|
"""
|
|||
|
|
输入：原图
|
|||
|
|
输出：图片中face框
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
import onnxruntime as ort
|
|||
|
|
|
|||
|
|
|
|||
|
|
class Box:
|
|||
|
|
def __init__(self, x1, y1, x2, y2, score, label=1, label_text = 'face' ,flag=True):
|
|||
|
|
self.x1 = x1
|
|||
|
|
self.y1 = y1
|
|||
|
|
self.x2 = x2
|
|||
|
|
self.y2 = y2
|
|||
|
|
self.score = score
|
|||
|
|
self.label = label
|
|||
|
|
self.label_text = label_text
|
|||
|
|
self.flag = flag
|
|||
|
|
|
|||
|
|
def iou_of(self, other):
|
|||
|
|
inter_x1 = max(self.x1, other.x1)
|
|||
|
|
inter_y1 = max(self.y1, other.y1)
|
|||
|
|
inter_x2 = min(self.x2, other.x2)
|
|||
|
|
inter_y2 = min(self.y2, other.y2)
|
|||
|
|
|
|||
|
|
if inter_x1 < inter_x2 and inter_y1 < inter_y2:
|
|||
|
|
inter_area = (inter_x2 - inter_x1 + 1.0) * (inter_y2 - inter_y1 + 1.0)
|
|||
|
|
self_area = (self.x2 - self.x1 + 1.0) * (self.y2 - self.y1 + 1.0)
|
|||
|
|
other_area = (other.x2 - other.x1 + 1.0) * (other.y2 - other.y1 + 1.0)
|
|||
|
|
union_area = self_area + other_area - inter_area
|
|||
|
|
return inter_area / union_area
|
|||
|
|
else:
|
|||
|
|
return 0
|
|||
|
|
|
|||
|
|
def area(self):
|
|||
|
|
return (self.x2 - self.x1 + 1) * (self.y2 - self.y1 + 1)
|
|||
|
|
|
|||
|
|
def hard_nms(boxes, iou_threshold, topk):
|
|||
|
|
if not boxes:
|
|||
|
|
return []
|
|||
|
|
boxes.sort(key=lambda x: x.score, reverse=True)
|
|||
|
|
|
|||
|
|
merged = [0] * len(boxes)
|
|||
|
|
output = []
|
|||
|
|
|
|||
|
|
count = 0
|
|||
|
|
for i in range(len(boxes)):
|
|||
|
|
if merged[i]:
|
|||
|
|
continue
|
|||
|
|
buf = [boxes[i]]
|
|||
|
|
merged[i] = 1
|
|||
|
|
|
|||
|
|
for j in range(i + 1, len(boxes)):
|
|||
|
|
if merged[j]:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
iou = boxes[i].iou_of(boxes[j])
|
|||
|
|
if iou > iou_threshold:
|
|||
|
|
merged[j] = 1
|
|||
|
|
buf.append(boxes[j])
|
|||
|
|
|
|||
|
|
output.append(buf[0])
|
|||
|
|
|
|||
|
|
count += 1
|
|||
|
|
if count >= topk:
|
|||
|
|
break
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
def blending_nms(boxes, iou_threshold, topk):
|
|||
|
|
if not boxes:
|
|||
|
|
return []
|
|||
|
|
boxes.sort(key=lambda x: x.score, reverse=True)
|
|||
|
|
merged = [0] * len(boxes)
|
|||
|
|
output = []
|
|||
|
|
|
|||
|
|
count = 0
|
|||
|
|
for i in range(len(boxes)):
|
|||
|
|
if merged[i]:
|
|||
|
|
continue
|
|||
|
|
buf = [boxes[i]]
|
|||
|
|
merged[i] = 1
|
|||
|
|
|
|||
|
|
for j in range(i + 1, len(boxes)):
|
|||
|
|
if merged[j]:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
iou = boxes[i].iou_of(boxes[j])
|
|||
|
|
if iou > iou_threshold:
|
|||
|
|
merged[j] = 1
|
|||
|
|
buf.append(boxes[j])
|
|||
|
|
|
|||
|
|
total = sum([np.exp(box.score) for box in buf])
|
|||
|
|
rects = Box(0, 0, 0, 0, 0)
|
|||
|
|
for box in buf:
|
|||
|
|
rate = np.exp(box.score) / total
|
|||
|
|
rects.x1 += box.x1 * rate
|
|||
|
|
rects.y1 += box.y1 * rate
|
|||
|
|
rects.x2 += box.x2 * rate
|
|||
|
|
rects.y2 += box.y2 * rate
|
|||
|
|
rects.score += box.score * rate
|
|||
|
|
rects.flag = True
|
|||
|
|
output.append(rects)
|
|||
|
|
|
|||
|
|
count += 1
|
|||
|
|
if count >= topk:
|
|||
|
|
break
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
def offset_nms(boxes, iou_threshold, topk):
|
|||
|
|
if not boxes:
|
|||
|
|
return []
|
|||
|
|
boxes.sort(key=lambda x: x.score, reverse=True)
|
|||
|
|
merged = [0] * len(boxes)
|
|||
|
|
offset = 4096.0
|
|||
|
|
|
|||
|
|
for box in boxes:
|
|||
|
|
box.x1 += box.label * offset
|
|||
|
|
box.y1 += box.label * offset
|
|||
|
|
box.x2 += box.label * offset
|
|||
|
|
box.y2 += box.label * offset
|
|||
|
|
|
|||
|
|
output = []
|
|||
|
|
count = 0
|
|||
|
|
for i in range(len(boxes)):
|
|||
|
|
if merged[i]:
|
|||
|
|
continue
|
|||
|
|
buf = [boxes[i]]
|
|||
|
|
merged[i] = 1
|
|||
|
|
|
|||
|
|
for j in range(i + 1, len(boxes)):
|
|||
|
|
if merged[j]:
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
iou = boxes[i].iou_of(boxes[j])
|
|||
|
|
if iou > iou_threshold:
|
|||
|
|
merged[j] = 1
|
|||
|
|
buf.append(boxes[j])
|
|||
|
|
|
|||
|
|
output.append(buf[0])
|
|||
|
|
|
|||
|
|
count += 1
|
|||
|
|
if count >= topk:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
for box in output:
|
|||
|
|
box.x1 -= box.label * offset
|
|||
|
|
box.y1 -= box.label * offset
|
|||
|
|
box.x2 -= box.label * offset
|
|||
|
|
box.y2 -= box.label * offset
|
|||
|
|
|
|||
|
|
return output
|
|||
|
|
|
|||
|
|
def draw_rectface(img, box):
|
|||
|
|
x = max(0,int(box.x1))
|
|||
|
|
y = max(0,int(box.y1))
|
|||
|
|
w = min(img.shape[1]-x, int(box.x2-x+1))
|
|||
|
|
h = min(img.shape[0]-y, int(box.y2-y+1))
|
|||
|
|
cv2.rectangle(img,(x,y),(x+w,y+h),(0,0,255),3)
|
|||
|
|
# return img
|
|||
|
|
|
|||
|
|
def cut_rectface(img, box):
|
|||
|
|
x = max(0,int(box.x1))
|
|||
|
|
y = max(0,int(box.y1))
|
|||
|
|
w = min(img.shape[1]-x, int(box.x2-x+1))
|
|||
|
|
h = min(img.shape[0]-y, int(box.y2-y+1))
|
|||
|
|
return img[y:y+h,x:x+w]
|
|||
|
|
|
|||
|
|
def normalize_inplace(mat, mean, scale):
|
|||
|
|
mat = mat.astype(np.float32)
|
|||
|
|
mat -= mean
|
|||
|
|
mat *= scale
|
|||
|
|
return mat
|
|||
|
|
|
|||
|
|
def create_tensor(mat, tensor_dims, memory_info_handler, data_format):
|
|||
|
|
rows, cols, channels = mat.shape
|
|||
|
|
if len(tensor_dims) != 4:
|
|||
|
|
raise RuntimeError("dims mismatch.")
|
|||
|
|
if tensor_dims[0] != 1:
|
|||
|
|
raise RuntimeError("batch != 1")
|
|||
|
|
|
|||
|
|
if data_format == "CHW":
|
|||
|
|
target_height = tensor_dims[2]
|
|||
|
|
target_width = tensor_dims[3]
|
|||
|
|
target_channel = tensor_dims[1]
|
|||
|
|
# target_tensor_size = target_channel * target_height * target_width
|
|||
|
|
if target_channel != channels:
|
|||
|
|
raise RuntimeError("channel mismatch.")
|
|||
|
|
|
|||
|
|
if target_height != rows or target_width != cols:
|
|||
|
|
print("in create_tensor, resize mat...")
|
|||
|
|
mat = cv2.resize(mat, (target_width, target_height))
|
|||
|
|
|
|||
|
|
mat = mat.transpose(2, 0, 1) # HWC -> CHW # 这儿存疑。
|
|||
|
|
mat = np.expand_dims(mat, axis=0)
|
|||
|
|
return ort.OrtValue.ortvalue_from_numpy(mat, 'cpu')
|
|||
|
|
|
|||
|
|
elif data_format == "HWC":
|
|||
|
|
target_height = tensor_dims[1]
|
|||
|
|
target_width = tensor_dims[2]
|
|||
|
|
target_channel = tensor_dims[3]
|
|||
|
|
target_tensor_size = target_channel * target_height * target_width
|
|||
|
|
if target_channel != channels:
|
|||
|
|
raise RuntimeError("channel mismatch.")
|
|||
|
|
|
|||
|
|
if target_height != rows or target_width != cols:
|
|||
|
|
mat = cv2.resize(mat, (target_width, target_height))
|
|||
|
|
|
|||
|
|
return ort.OrtValue.ortvalue_from_numpy(mat, 'cpu')
|
|||
|
|
|
|||
|
|
class BasicOrtHandler:
|
|||
|
|
def __init__(self, onnx_path, num_threads=1):
|
|||
|
|
self.onnx_path = onnx_path
|
|||
|
|
self.num_threads = num_threads
|
|||
|
|
self.initialize_handler()
|
|||
|
|
|
|||
|
|
def initialize_handler(self):
|
|||
|
|
# self.ort_env = ort.Env(ort.logging.ERROR)
|
|||
|
|
session_options = ort.SessionOptions()
|
|||
|
|
session_options.intra_op_num_threads = self.num_threads
|
|||
|
|
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
|
|||
|
|
|
|||
|
|
# self.ort_session = ort.InferenceSession(self.onnx_path, session_options)
|
|||
|
|
# self.memory_info_handler = ort.OrtMemoryInfo("cpu", ort.OrtAllocatorType.ORT_ARENA_ALLOCATOR)
|
|||
|
|
|
|||
|
|
# Initialize session
|
|||
|
|
self.ort_session = ort.InferenceSession(self.onnx_path, session_options)
|
|||
|
|
self.memory_info_handler = ort.OrtMemoryInfo("Cpu", ort.OrtAllocatorType.ORT_ARENA_ALLOCATOR, 0, ort.OrtMemType.DEFAULT)
|
|||
|
|
|
|||
|
|
|
|||
|
|
self.input_node_names = [self.ort_session.get_inputs()[0].name]
|
|||
|
|
self.input_node_dims = self.ort_session.get_inputs()[0].shape # 获取输入张量的shape
|
|||
|
|
self.input_tensor_size = np.prod(self.input_node_dims)
|
|||
|
|
|
|||
|
|
self.output_node_names = [out.name for out in self.ort_session.get_outputs()]
|
|||
|
|
self.output_node_dims = [out.shape for out in self.ort_session.get_outputs()]
|
|||
|
|
self.num_outputs = len(self.output_node_names)
|
|||
|
|
|
|||
|
|
def __del__(self):
|
|||
|
|
del self.ort_session
|
|||
|
|
|
|||
|
|
class FaceBoxesV2(BasicOrtHandler):
|
|||
|
|
def __init__(self, onnx_path, num_threads=1):
|
|||
|
|
super().__init__(onnx_path, num_threads)
|
|||
|
|
self.mean_vals = np.array([104.0, 117.0, 123.0], dtype=np.float32)
|
|||
|
|
self.scale_vals = np.array([1.0, 1.0, 1.0], dtype=np.float32)
|
|||
|
|
self.variance = [0.1, 0.2]
|
|||
|
|
self.steps = [32, 64, 128]
|
|||
|
|
self.min_sizes = [
|
|||
|
|
[32, 64, 128],
|
|||
|
|
[256],
|
|||
|
|
[512]
|
|||
|
|
]
|
|||
|
|
self.max_nms = 30000
|
|||
|
|
|
|||
|
|
def transform(self, mat):
|
|||
|
|
canvas = cv2.resize(mat, (self.input_node_dims[3], self.input_node_dims[2]))
|
|||
|
|
canvas = normalize_inplace(canvas, self.mean_vals, self.scale_vals)
|
|||
|
|
return create_tensor(canvas, self.input_node_dims, self.memory_info_handler, "CHW")
|
|||
|
|
|
|||
|
|
def detect(self, mat, score_threshold=0.35, iou_threshold=0.45, topk=300, nms_type=0):
|
|||
|
|
if mat is None or mat.size == 0:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
img_height = float(mat.shape[0])
|
|||
|
|
img_width = float(mat.shape[1])
|
|||
|
|
|
|||
|
|
# 1. make input tensor
|
|||
|
|
input_tensor = self.transform(mat)
|
|||
|
|
# 2. inference scores & boxes.
|
|||
|
|
output_tensors = self.ort_session.run(self.output_node_names, {self.input_node_names[0]: input_tensor})
|
|||
|
|
|
|||
|
|
# 3. rescale & exclude.
|
|||
|
|
bbox_collection = []
|
|||
|
|
bbox_collection = self.generate_bboxes(output_tensors, score_threshold, img_height, img_width)
|
|||
|
|
|
|||
|
|
# 4. hard|blend|offset nms with topk. return detected_boxes
|
|||
|
|
return self.nms(bbox_collection, iou_threshold, topk, nms_type)
|
|||
|
|
|
|||
|
|
def generate_bboxes(self, output_tensors, score_threshold, img_height, img_width):
|
|||
|
|
bboxes = output_tensors[0] # e.g (1,n,4)
|
|||
|
|
probs = output_tensors[1] # e.g (1,n,2) after softmax
|
|||
|
|
bbox_dims = self.output_node_dims[0] # (1,n,4)
|
|||
|
|
bbox_num = bbox_dims[1] # n = ?
|
|||
|
|
input_height = self.input_node_dims[2] # e.g 640
|
|||
|
|
input_width = self.input_node_dims[3] # e.g 640
|
|||
|
|
|
|||
|
|
anchors = self.generate_anchors(input_height, input_width)
|
|||
|
|
|
|||
|
|
num_anchors = len(anchors)
|
|||
|
|
if num_anchors != bbox_num:
|
|||
|
|
print(f"num_anchors={num_anchors} but detected bbox_num={bbox_num}")
|
|||
|
|
raise RuntimeError("mismatch num_anchors != bbox_num")
|
|||
|
|
|
|||
|
|
bbox_collection = []
|
|||
|
|
count = 0
|
|||
|
|
for i in range(num_anchors):
|
|||
|
|
conf = probs[0, i, 1]
|
|||
|
|
if conf < score_threshold:
|
|||
|
|
continue # filter first.
|
|||
|
|
|
|||
|
|
# prior_cx = anchors[i].cx
|
|||
|
|
# prior_cy = anchors[i].cy
|
|||
|
|
# prior_s_kx = anchors[i].s_kx
|
|||
|
|
# prior_s_ky = anchors[i].s_ky
|
|||
|
|
prior_cx, prior_cy, prior_s_kx, prior_s_ky = anchors[i]
|
|||
|
|
|
|||
|
|
dx = bboxes[0, i, 0]
|
|||
|
|
dy = bboxes[0, i, 1]
|
|||
|
|
dw = bboxes[0, i, 2]
|
|||
|
|
dh = bboxes[0, i, 3]
|
|||
|
|
|
|||
|
|
cx = prior_cx + dx * self.variance[0] * prior_s_kx
|
|||
|
|
cy = prior_cy + dy * self.variance[0] * prior_s_ky
|
|||
|
|
w = prior_s_kx * np.exp(dw * self.variance[1])
|
|||
|
|
h = prior_s_ky * np.exp(dh * self.variance[1]) # norm coor (0.,1.)
|
|||
|
|
|
|||
|
|
box = Box(
|
|||
|
|
x1=(cx - w / 2.0) * img_width,
|
|||
|
|
y1=(cy - h / 2.0) * img_height,
|
|||
|
|
x2=(cx + w / 2.0) * img_width,
|
|||
|
|
y2=(cy + h / 2.0) * img_height,
|
|||
|
|
score=conf,
|
|||
|
|
label=1,
|
|||
|
|
label_text="face",
|
|||
|
|
flag=True
|
|||
|
|
)
|
|||
|
|
bbox_collection.append(box)
|
|||
|
|
|
|||
|
|
count += 1 # limit boxes for nms.
|
|||
|
|
if count > self.max_nms:
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
return bbox_collection
|
|||
|
|
|
|||
|
|
def nms(self, input_boxes, iou_threshold, topk, nms_type):
|
|||
|
|
if nms_type == 1:
|
|||
|
|
output_boxes = blending_nms(input_boxes, iou_threshold, topk)
|
|||
|
|
elif nms_type == 2:
|
|||
|
|
output_boxes = offset_nms(input_boxes, iou_threshold, topk)
|
|||
|
|
elif nms_type == 0:
|
|||
|
|
output_boxes = hard_nms(input_boxes, iou_threshold, topk)
|
|||
|
|
else:
|
|||
|
|
raise NotImplementedError
|
|||
|
|
return output_boxes
|
|||
|
|
|
|||
|
|
def generate_anchors(self, target_height, target_width):
|
|||
|
|
feature_maps = []
|
|||
|
|
for step in self.steps:
|
|||
|
|
feature_maps.append([
|
|||
|
|
int(np.ceil(target_height / step)),
|
|||
|
|
int(np.ceil(target_width / step))
|
|||
|
|
])
|
|||
|
|
|
|||
|
|
anchors = []
|
|||
|
|
for k, f_map in enumerate(feature_maps):
|
|||
|
|
tmp_min_sizes = self.min_sizes[k]
|
|||
|
|
f_h, f_w = f_map
|
|||
|
|
|
|||
|
|
offset_32 = [0.0, 0.25, 0.5, 0.75]
|
|||
|
|
offset_64 = [0.0, 0.5]
|
|||
|
|
|
|||
|
|
for i in range(f_h):
|
|||
|
|
for j in range(f_w):
|
|||
|
|
for min_size in tmp_min_sizes:
|
|||
|
|
s_kx = min_size / target_width
|
|||
|
|
s_ky = min_size / target_height
|
|||
|
|
|
|||
|
|
if min_size == 32:
|
|||
|
|
for offset_y in offset_32:
|
|||
|
|
for offset_x in offset_32:
|
|||
|
|
cx = (j + offset_x) * self.steps[k] / target_width
|
|||
|
|
cy = (i + offset_y) * self.steps[k] / target_height
|
|||
|
|
anchors.append([cx, cy, s_kx, s_ky])
|
|||
|
|
elif min_size == 64:
|
|||
|
|
for offset_y in offset_64:
|
|||
|
|
for offset_x in offset_64:
|
|||
|
|
cx = (j + offset_x) * self.steps[k] / target_width
|
|||
|
|
cy = (i + offset_y) * self.steps[k] / target_height
|
|||
|
|
anchors.append([cx, cy, s_kx, s_ky])
|
|||
|
|
else:
|
|||
|
|
cx = (j + 0.5) * self.steps[k] / target_width
|
|||
|
|
cy = (i + 0.5) * self.steps[k] / target_height
|
|||
|
|
anchors.append([cx, cy, s_kx, s_ky])
|
|||
|
|
|
|||
|
|
return anchors
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
# Usage example
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
import sys
|
|||
|
|
import os
|
|||
|
|
img_path = sys.argv[1]
|
|||
|
|
reta = FaceBoxesV2(r"./checkpoints/faceboxesv2-640x640.onnx",4)
|
|||
|
|
img = cv2.imread(img_path)
|
|||
|
|
detected_boxes = reta.detect(img)
|
|||
|
|
count = 0
|
|||
|
|
for box in detected_boxes:
|
|||
|
|
print(f"({box.x1:.3f},{box.y1:.3f},{box.x2:.3f},{box.y2:.3f})", end=" ")
|
|||
|
|
count += 1
|
|||
|
|
print("total face number:",count)
|
|||
|
|
|
|||
|
|
for box in detected_boxes:
|
|||
|
|
draw_rectface(img, box)
|
|||
|
|
|
|||
|
|
filename = os.path.basename(img_path)
|
|||
|
|
cv2.imwrite("./" + filename, img)
|