#!/usr/bin/env python3
"""Test RetinaFace RKNN model directly"""

import numpy as np
from PIL import Image
from rknn.api import RKNN

# Create dummy test image (simulating a face region)
test_img = np.random.randint(0, 255, (320, 320, 3), dtype=np.uint8)
# Add a rectangle to simulate face
test_img[100:220, 100:220] = 200  # lighter region

# Save test image
Image.fromarray(test_img).save('test_input.jpg')
print("Created test image")

# Load and test RKNN model
rknn = RKNN(verbose=False)
print("Loading RKNN model...")
ret = rknn.load_rknn('face_det_retinaface_mobile320_rk3588.rknn')
if ret != 0:
    print("Failed to load model")
    exit(1)

print("Initializing runtime...")
ret = rknn.init_runtime(target='rk3588')
if ret != 0:
    print("Failed to init runtime")
    exit(1)

# Prepare input (NHWC -> NCHW)
input_data = np.expand_dims(test_img, 0)  # Add batch dimension
print(f"Input shape: {input_data.shape}, dtype: {input_data.dtype}")

# Inference
print("Running inference...")
outputs = rknn.inference(inputs=[input_data])
print(f"Output count: {len(outputs)}")
for i, out in enumerate(outputs):
    print(f"Output {i}: shape={out.shape}, min={out.min():.4f}, max={out.max():.4f}")

rknn.release()
print("Done!")