| """Local inference helper for waste classification model.""" |
|
|
| import os |
| import numpy as np |
| import tensorflow as tf |
| from tensorflow.keras.preprocessing import image |
| from tensorflow.keras.applications.efficientnet import preprocess_input |
|
|
| IMG_SIZE = 224 |
| MODEL_PATH = os.path.join(os.getcwd(), "waste_classifier.keras") |
|
|
| CLASS_NAMES = ["dry_waste", "other_waste", "wet_waste"] |
| CLASS_DISPLAY = { |
| "dry_waste": "♻️ Dry Waste", |
| "other_waste": "🔶 Other Waste", |
| "wet_waste": "🍃 Wet Waste", |
| } |
| CLASS_ACTIONS = { |
| "dry_waste": "Can be recycled → Paper, plastic, metal recovery", |
| "other_waste": "Needs special handling → Hazardous / e-waste processing", |
| "wet_waste": "Compostable → Organic composting / biogas generation", |
| } |
|
|
| model = tf.keras.models.load_model(MODEL_PATH) |
|
|
|
|
| def predict(image_path: str): |
| """Predict waste class for an image path.""" |
| img = image.load_img(image_path, target_size=(IMG_SIZE, IMG_SIZE)) |
| img_array = image.img_to_array(img) |
| img_array = np.expand_dims(img_array, axis=0) |
| img_array = preprocess_input(img_array) |
|
|
| predictions = model.predict(img_array, verbose=0)[0] |
|
|
| index = int(np.argmax(predictions)) |
| class_name = CLASS_NAMES[index] |
| confidence = float(predictions[index] * 100) |
|
|
| all_predictions = { |
| CLASS_NAMES[i]: round(float(predictions[i] * 100), 1) |
| for i in range(len(CLASS_NAMES)) |
| } |
|
|
| return { |
| "class": class_name, |
| "display_name": CLASS_DISPLAY.get(class_name, class_name), |
| "confidence": round(confidence, 1), |
| "action": CLASS_ACTIONS.get(class_name, ""), |
| "all_predictions": all_predictions, |
| } |
|
|
