| 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() |
|
|
