#!/usr/bin/env python3
# Copyright 2022-2024,2026 NXP
# SPDX-License-Identifier: MIT

import argparse
import time
import cv2
import numpy as np

try:
    import tflite_runtime.interpreter as tflite
except ImportError:
    import tensorflow as tf
    tflite = tf.lite

BOX_COLOR = (255, 128, 0)
THRESHOLD = 0.5
INPUT_WIDTH = 320
INPUT_HEIGHT = 240


def load_image(filename):
    orig_image = cv2.imread(filename, cv2.IMREAD_COLOR)
    if orig_image is None:
        raise FileNotFoundError(f"Could not load image: {filename}")
    image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
    image = cv2.resize(image, (INPUT_WIDTH, INPUT_HEIGHT))
    image = np.expand_dims(image, axis=0)
    image = (image.astype(np.float32) - 127.0) / 128.0
    return orig_image, image


def run_inference(interpreter, image, threshold=THRESHOLD):
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    # Quantize the input if the model expects an integer input
    if input_details[0]['dtype'] in (np.uint8, np.int8):
        input_scale, input_zero_point = input_details[0]['quantization']
        image = image / input_scale + input_zero_point
        image = image.astype(input_details[0]['dtype'])

    interpreter.set_tensor(input_details[0]['index'], image)
    interpreter.invoke()
    out = interpreter.get_tensor(output_details[0]['index'])

    # Dequantize the output if needed
    if output_details[0]['dtype'] in (np.uint8, np.int8):
        scale, zero_point = output_details[0]['quantization']
        out = (out.astype(np.float32) - zero_point) * scale
    out = out.astype(np.float32)

    boxes = out[..., 2:]
    scores = out[..., 1]

    conf_mask = threshold < scores
    boxes, scores = boxes[conf_mask], scores[conf_mask]

    return boxes, scores


def main():
    parser = argparse.ArgumentParser(description="Ultraface slim face detection example")
    parser.add_argument('-m', '--model',
                        default='original_model/ultraface_slim_uint8_float32.tflite',
                        help='Path to the TFLite model file')
    parser.add_argument('-i', '--input', default='example_input.jpg',
                        help='Path to the input image')
    parser.add_argument('-o', '--output', default='example_output.jpg',
                        help='Path to save the annotated output image')
    parser.add_argument('-t', '--threshold', type=float, default=THRESHOLD,
                        help='Confidence threshold for detections')
    args = parser.parse_args()

    interpreter = tflite.Interpreter(model_path=args.model)
    interpreter.allocate_tensors()

    orig_image, processed_image = load_image(args.input)

    start = time.time()
    boxes, scores = run_inference(interpreter, processed_image, args.threshold)
    end = time.time()

    boxes = boxes * np.tile(orig_image.shape[1::-1], 2)
    boxes = boxes.astype(np.int32)

    print(f"Inference time {(end - start) * 1000:.2f} ms")
    print(f"Detected {boxes.shape[0]} faces")
    for i in range(boxes.shape[0]):
        box = boxes[i, :]
        print(f"Box {i}: {box.tolist()} score {scores[i]:.3f}")
        cv2.rectangle(orig_image, (box[0], box[1]), (box[2], box[3]), BOX_COLOR, 4)

    cv2.imwrite(args.output, orig_image)
    print(f"Saved annotated image to {args.output}")


if __name__ == '__main__':
    main()