import argparse
import os

import cv2 as cv
import numpy as np
import onnxruntime as ort

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


def main():
    parser = argparse.ArgumentParser(description="TensorFlow Inception (ONNX) image classification demo")
    parser.add_argument("--model", default=os.path.join(here, "tensorflow_inception_graph_2026jul.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("--labels", help="optional ImageNet label file, one class name per line")
    args = parser.parse_args()

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

    rgb = cv.resize(cv.cvtColor(img, cv.COLOR_BGR2RGB), (224, 224)).astype(np.float32)

    sess = ort.InferenceSession(args.model, providers=["CPUExecutionProvider"])
    scores = sess.run(None, {sess.get_inputs()[0].name: rgb[None]})[0].ravel()

    top = int(np.argmax(scores))
    conf = float(scores[top])
    label = str(top)
    if args.labels:
        names = open(args.labels).read().splitlines()
        if top < len(names):
            label = names[top]

    print("class", top, label, "confidence", round(conf, 4))

    out = img.copy()
    cv.putText(out, "%s (%.2f)" % (label, conf), (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)
    cv.imwrite(args.output, out)
    print("wrote", args.output)


if __name__ == "__main__":
    main()