File size: 3,372 Bytes
9bd4818 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | #!/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()
|