inference.py

"""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,
    }