File size: 2,491 Bytes
1b6b67c | 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 | """
Trashpanzee - Local / Hub Inference Verification Script
"""
import os
import sys
import numpy as np
from PIL import Image
# Multi-fallback interpreter import
try:
from ai_edge_litert.interpreter import Interpreter
except ImportError:
try:
from tflite_runtime.interpreter import Interpreter
except ImportError:
try:
import tensorflow as tf
Interpreter = tf.lite.Interpreter
except ImportError:
sys.exit(
"Error: No TFLite runtime found. Please install ai-edge-litert or tflite-runtime or tensorflow:\n"
" pip install ai-edge-litert pillow numpy"
)
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DEFAULT_MODEL = os.path.join(SCRIPT_DIR, "model_unquant.tflite")
DEFAULT_LABELS = os.path.join(SCRIPT_DIR, "labels.txt")
def load_labels(labels_file=DEFAULT_LABELS):
labels = []
if not os.path.exists(labels_file):
return ["Trash Can", "Not Trash"]
with open(labels_file, "r", encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
name = line.split(" ", 1)[-1]
labels.append("Not Trash" if name.lower() == "not tash" else name)
return labels
def classify_image(image_path, model_path=DEFAULT_MODEL, labels_path=DEFAULT_LABELS):
labels = load_labels(labels_path)
interpreter = Interpreter(model_path=model_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# Load and resize
image = Image.open(image_path).convert("RGB").resize((224, 224))
input_data = (np.array(image, dtype=np.float32) / 127.5) - 1.0
input_data = np.expand_dims(input_data, axis=0)
# Inference
interpreter.set_tensor(input_details[0]["index"], input_data)
interpreter.invoke()
output_data = interpreter.get_tensor(output_details[0]["index"])[0]
results = {labels[i]: float(output_data[i]) for i in range(len(labels))}
return results
if __name__ == "__main__":
if len(sys.argv) < 2:
print("Usage: python example_inference.py <path_to_image>")
sys.exit(1)
img_path = sys.argv[1]
scores = classify_image(img_path)
print(f"\nClassification results for '{img_path}':")
for lbl, prob in sorted(scores.items(), key=lambda x: x[1], reverse=True):
print(f" {lbl:12s}: {prob * 100:.2f}%")
|