import argparse
import os

import cv2 as cv
import numpy as np

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


def main():
    parser = argparse.ArgumentParser(description="EAST scene-text detection (ONNX) demo")
    parser.add_argument("--model", default=os.path.join(here, "east_text_detection_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("--conf", type=float, default=0.5, help="confidence threshold")
    parser.add_argument("--nms", type=float, default=0.4, help="NMS threshold")
    args = parser.parse_args()

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

    model = cv.dnn.TextDetectionModel_EAST(args.model)
    model.setConfidenceThreshold(args.conf).setNMSThreshold(args.nms)
    model.setInputParams(1.0, (320, 320), (123.68, 116.78, 103.94), True, False)

    boxes, confidences = model.detectTextRectangles(img)
    print("detections", len(boxes))

    out = img.copy()
    for box in boxes:
        pts = cv.boxPoints(box).astype(np.int32)
        cv.polylines(out, [pts], True, (0, 255, 0), 2)
    cv.imwrite(args.output, out)
    print("wrote", args.output)


if __name__ == "__main__":
    main()