File size: 2,051 Bytes
61196f3 | 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 | #!/usr/bin/env python3
"""Inference script for the DCSkyCam helicopter type multi-class classifier.
Usage:
python inference.py <image_path> [--top-k 3]
Requires: tflite-runtime (on Raspberry Pi) or tensorflow (on desktop).
"""
import sys
import numpy as np
from PIL import Image
try:
import tflite_runtime.interpreter as tflite
except ImportError:
import tensorflow.lite as tflite
MODEL_PATH = "custom_multi_e2m.tflite"
LABELS_PATH = "custom_multi_e2m-labels.txt"
def load_labels(path):
with open(path, "r") as f:
return [line.strip() for line in f.readlines()]
def predict(image_path, model_path=MODEL_PATH, labels_path=LABELS_PATH, top_k=3):
"""Run inference on an image and return predictions."""
# Load model
interpreter = tflite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()
# Load labels
labels = load_labels(labels_path)
# Prepare image: 480x480, RGB, normalized to [0, 1]
img = Image.open(image_path).convert("RGB")
img = img.resize((480, 480))
input_data = np.expand_dims(np.array(img, dtype=np.float32) / 255.0, axis=0)
# Run inference
input_details = interpreter.get_input_details()[0]
interpreter.set_tensor(input_details["index"], input_data)
interpreter.invoke()
# Get results
output = interpreter.get_tensor(interpreter.get_output_details()[0]["index"])[0]
# Get top-k predictions
top_k_idx = np.argsort(output)[::-1][:top_k]
return labels, top_k_idx, output.tolist()
if __name__ == "__main__":
if len(sys.argv) < 2:
print(f"Usage: {sys.argv[0]} <image_path> [--top-k N]")
sys.exit(1)
image_path = sys.argv[1]
top_k = 3
if "--top-k" in sys.argv:
idx = sys.argv.index("--top-k")
if idx + 1 < len(sys.argv):
top_k = int(sys.argv[idx + 1])
labels, top_k_idx, scores = predict(image_path, top_k=top_k)
print(f"Image: {image_path}")
for i, idx in enumerate(top_k_idx):
print(f" #{i+1} {labels[idx]}: {scores[idx]:.4f}")
|