| import argparse |
| import glob |
| import os |
|
|
| import cv2 as cv |
| import numpy as np |
|
|
| here = os.path.dirname(os.path.abspath(__file__)) |
| sz = 512 |
|
|
|
|
| def build_anchors(): |
| scales = [2.0 ** (i / 3.0) for i in range(3)] |
| aspects = [(1.0, 1.0), (1.4, 0.7), (0.7, 1.4)] |
| base = [] |
| for s in scales: |
| for aw, ah in aspects: |
| base.append((32.0 * s * aw, 32.0 * s * ah)) |
| anchors = [] |
| for lvl in range(5): |
| f = sz // (8 * 2 ** lvl) |
| step = 8 * 2 ** lvl |
| m = 2 ** lvl |
| for y in range(f): |
| for x in range(f): |
| cx = (x + 0.5) * step |
| cy = (y + 0.5) * step |
| for bw, bh in base: |
| anchors.append((cx, cy, bw * m, bh * m)) |
| return np.array(anchors, np.float32) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="EfficientDet-D0 (ONNX) object detection demo") |
| parser.add_argument("--model", default=None) |
| 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() |
|
|
| model = args.model |
| if model is None: |
| found = glob.glob(os.path.join(here, "*.onnx")) |
| if not found: |
| raise SystemExit("no onnx, run convert_to_onnx.py") |
| model = found[0] |
|
|
| img = cv.imread(args.image) |
| if img is None: |
| raise SystemExit("could not read image: %s" % args.image) |
|
|
| anchors = build_anchors() |
| acx, acy, aw, ah = anchors[:, 0], anchors[:, 1], anchors[:, 2], anchors[:, 3] |
|
|
| net = cv.dnn.readNetFromONNX(model) |
| inp = cv.resize(cv.cvtColor(img, cv.COLOR_BGR2RGB), (sz, sz)) |
| net.setInput(inp[None].astype(np.uint8)) |
| res = net.forward(net.getUnconnectedOutLayersNames()) |
| box = next(a for a in res if a.shape[-1] == 4).reshape(-1, 4) |
| cls = next(a for a in res if a.shape[-1] != 4).reshape(box.shape[0], -1) |
|
|
| ycenter = box[:, 0] * ah + acy |
| xcenter = box[:, 1] * aw + acx |
| bh = np.exp(box[:, 2]) * ah |
| bw = np.exp(box[:, 3]) * aw |
| boxes = np.stack([xcenter - bw / 2, ycenter - bh / 2, xcenter + bw / 2, ycenter + bh / 2], 1) / sz |
|
|
| prob = 1.0 / (1.0 + np.exp(-cls)) |
| cid = prob.argmax(1) |
| scores = prob.max(1) |
|
|
| keep = scores > args.conf |
| boxes = boxes[keep] |
| scores = scores[keep] |
| cid = cid[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]) |
| iw = np.maximum(0, xx2 - xx1) |
| ih = np.maximum(0, yy2 - yy1) |
| inter = iw * ih |
| 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.6] |
|
|
| print("efficientdet-d0", len(pick), "detections") |
| h, w = img.shape[:2] |
| for i in pick: |
| x1, y1, x2, y2 = boxes[i] |
| print(int(cid[i]), round(float(scores[i]), 3), round(float(x1), 3), round(float(y1), 3), round(float(x2), 3), round(float(y2), 3)) |
| cv.rectangle(img, (int(x1 * w), int(y1 * h)), (int(x2 * w), int(y2 * h)), (0, 255, 0), 2) |
| cv.putText(img, "%d:%.2f" % (int(cid[i]), scores[i]), (int(x1 * w), int(y1 * h) - 5), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) |
| cv.imwrite(args.output, img) |
| print("wrote", args.output) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
