import argparse
import os

import cv2 as cv
import numpy as np

here = os.path.dirname(os.path.abspath(__file__))

# TF Object Detection API heads, in the order the graph declares them.
OUTPUTS = ["detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0"]


def main():
    parser = argparse.ArgumentParser(description="Macbeth colour chart detection (ONNX) demo")
    parser.add_argument("--model", default=os.path.join(here, "macbeth_chart_detector_2026sep.onnx"))
    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.3, help="confidence threshold")
    args = parser.parse_args()

    img = cv.imread(args.image)
    if img is None:
        raise SystemExit("could not read image: %s" % args.image)

    # The network takes raw uint8 RGB in NHWC layout at the image's own resolution:
    # the graph resizes internally (keep-aspect-ratio, min side 768 / max side 1024),
    # so do not pre-resize and do not use blobFromImage (that would give NCHW).
    blob = cv.cvtColor(img, cv.COLOR_BGR2RGB)[None]

    net = cv.dnn.readNetFromONNX(args.model, cv.dnn.ENGINE_OPENCV)
    net.setInput(blob)
    boxes, scores, _classes, _num = net.forward(OUTPUTS)

    boxes = boxes.reshape(-1, 4)
    scores = scores.reshape(-1)

    # Threshold on the scores rather than on num_detections: that head reports the
    # post-NMS slot count (100), not the number of charts actually found. The single
    # class is the chart itself, so detection_classes carries no information here.
    keep = np.nonzero(scores >= args.conf)[0]

    h, w = img.shape[:2]
    out = img.copy()
    print("macbeth_chart_detector", len(keep), "detections")
    for i in keep:
        ymin, xmin, ymax, xmax = boxes[i]          # normalized, y first
        p1 = (int(round(xmin * w)), int(round(ymin * h)))
        p2 = (int(round(xmax * w)), int(round(ymax * h)))
        cv.rectangle(out, p1, p2, (0, 255, 0), 2)
        cv.putText(out, "chart %.2f" % scores[i], (p1[0], max(0, p1[1] - 5)),
                   cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
        print("  %.3f  x1=%d y1=%d x2=%d y2=%d" % (scores[i], p1[0], p1[1], p2[0], p2[1]))

    cv.imwrite(args.output, out)
    print("wrote", args.output)


if __name__ == "__main__":
    main()