import tensorflow as tf
import json
from PIL import Image
import numpy as np
import os

cfg = None
if os.path.exists("config.json"):
    with open("config.json") as f:
        cfg = json.load(f)
    
class_names = ['battery', 'glass', 'metal', 'organic', 'paper', 'plastic']

base_model = tf.keras.applications.EfficientNetB7(weights=None)
model = models.Sequential([
    base_model,
    layers.GlobalAveragePooling2D(),
    layers.Dropout(0.3),
    layers.Dense(256, activation='relu', kernel_regularizer=regularizers.l2(0.001)),
    layers.Dropout(0.3),
    layers.Dense(128, activation='relu', kernel_regularizer=regularizers.l2(0.001)),
    layers.Dropout(0.3),
    layers.Dense(6, activation='softmax')
])
model.load_weights("model.weights.h5")

def preprocess(image: Image.Image):
    image = image.resize((224, 224))
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)
    return image

def predict(image: Image.Image):
    x = preprocess(image)
    x = model.predict(x)
    class_idx = int(np.argmax(x, axis=1)[0])
    confidence = float(np.max(x))
    return {
        "class": class_names[class_idx],
        "confidence": confidence
    }