File size: 1,918 Bytes
e7ffcd9 | 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 | #!/usr/bin/env python3
"""Inference script for the DCSkyCam helicopter binary classifier.
Usage:
python inference.py <image_path>
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_heli.tflite"
LABELS_PATH = "custom_heli-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):
"""Run inference on an image and return prediction."""
# Load model
interpreter = tflite.Interpreter(model_path=model_path)
interpreter.allocate_tensors()
# Load labels
labels = load_labels(labels_path)
# Prepare image: 224x224, RGB, normalized to [0, 1]
img = Image.open(image_path).convert("RGB")
img = img.resize((224, 224))
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]
pred_idx = int(np.argmax(output))
confidence = float(np.max(output))
return labels[pred_idx], confidence, output.tolist()
if __name__ == "__main__":
if len(sys.argv) < 2:
print(f"Usage: {sys.argv[0]} <image_path>")
sys.exit(1)
image_path = sys.argv[1]
label, confidence, scores = predict(image_path)
print(f"Image: {image_path}")
print(f"Prediction: {label}")
print(f"Confidence: {confidence:.4f}")
for i, (lbl, sc) in enumerate(zip(load_labels(LABELS_PATH), scores)):
print(f" {lbl}: {sc:.4f}")
|