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