| |
| |
| |
|
|
| 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() |
|
|
| |
| 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']) |
|
|
| |
| 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() |
|
|
