File size: 3,717 Bytes
e9eb33d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import argparse
import glob
import os
import cv2 as cv
import numpy as np
here = os.path.dirname(os.path.abspath(__file__))
sz = 300
layers = [
(30, 60, [2], 8, 38),
(60, 111, [2, 3], 16, 19),
(111, 162, [2, 3], 32, 10),
(162, 213, [2, 3], 64, 5),
(213, 264, [2], 100, 5),
(264, 315, [2], 300, 5),
]
var = [0.1, 0.1, 0.2, 0.2]
def build_priors():
p = []
for mn, mx, ars, step, fm in layers:
ratios = [1.0]
for a in ars:
ratios += [a, 1.0 / a]
for y in range(fm):
for x in range(fm):
cx = (x + 0.5) * step
cy = (y + 0.5) * step
boxes = [(mn, mn), ((mn * mx) ** 0.5, (mn * mx) ** 0.5)]
for a in ratios[1:]:
boxes.append((mn * a ** 0.5, mn / a ** 0.5))
for bw, bh in boxes:
p.append([cx, cy, bw, bh])
return np.array(p, np.float32)
def default_model():
files = [f for f in glob.glob(os.path.join(here, "*.onnx")) if "known_good" not in os.path.basename(f)]
return files[0] if files else os.path.join(here, "opencv_face_detector_uint8.onnx")
def main():
parser = argparse.ArgumentParser(description="OpenCV SSD face detector (ONNX) demo")
parser.add_argument("--model", default=default_model())
parser.add_argument("--image", default=os.path.join(here, "example_outputs", "input_image.png"))
parser.add_argument("--output", default=os.path.join(here, "example_outputs", "output_image.png"))
parser.add_argument("--conf", type=float, default=0.4)
args = parser.parse_args()
img = cv.imread(args.image)
if img is None:
raise SystemExit("could not read image: %s" % args.image)
inp = cv.resize(img, (sz, sz)).astype(np.float32) - np.array([104.0, 177.0, 123.0], np.float32)
net = cv.dnn.readNetFromONNX(args.model)
onames = net.getUnconnectedOutLayersNames()
net.setInput(inp[None])
res = net.forward(onames)
loc = res[[i for i, n in enumerate(onames) if "mbox_loc" in n][0]].reshape(-1, 4)
conf = res[[i for i, n in enumerate(onames) if "mbox_conf" in n][0]].reshape(-1, 2)
priors = build_priors()
pcx = priors[:, 0] / sz
pcy = priors[:, 1] / sz
pw = priors[:, 2] / sz
ph = priors[:, 3] / sz
e = np.exp(conf - conf.max(1, keepdims=True))
sm = e / e.sum(1, keepdims=True)
scores = sm[:, 1]
cx = pcx + loc[:, 0] * var[0] * pw
cy = pcy + loc[:, 1] * var[1] * ph
bw = pw * np.exp(loc[:, 2] * var[2])
bh = ph * np.exp(loc[:, 3] * var[3])
boxes = np.stack([cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2], 1)
keep = scores > args.conf
boxes = boxes[keep]
scores = scores[keep]
order = scores.argsort()[::-1]
pick = []
while order.size:
i = order[0]
pick.append(i)
xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
ai = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
aj = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1])
iou = inter / (ai + aj - inter + 1e-9)
order = order[1:][iou <= 0.3]
h, w = img.shape[:2]
for i in pick:
x1, y1, x2, y2 = boxes[i]
cv.rectangle(img, (int(x1 * w), int(y1 * h)), (int(x2 * w), int(y2 * h)), (0, 255, 0), 2)
cv.imwrite(args.output, img)
print("opencv_face_detector_uint8", len(pick), "faces")
if __name__ == "__main__":
main()
|