| """ |
| Trashpanzee - Local / Hub Inference Verification Script |
| """ |
| import os |
| import sys |
| import numpy as np |
| from PIL import Image |
|
|
| |
| 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() |
| |
| |
| 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) |
| |
| |
| 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}%") |
|
|
