import gradio as gr, numpy as np, json, tensorflow as tf
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
import os

# Load model - try local first, then HF Hub
MODEL_PATH = "corrosion_classifier.keras"
CLASSES_PATH = "classes.json"

M = tf.keras.models.load_model(MODEL_PATH, compile=False)
CLASSES = json.load(open(CLASSES_PATH))
SZ = M.input_shape[1]

def predict(img):
    if img is None:
                return {}, "Unggah foto permukaan logam."
    x = tf.image.resize(tf.cast(img, tf.float32), [SZ, SZ])[None]
    p = M.predict(x, verbose=0)[0]
    conf = {CLASSES[i]: float(p[i]) for i in range(len(CLASSES))}
    top = CLASSES[int(p.argmax())]
    note = "Triase/pra-skrining - BUKAN pengganti NDT."
    return conf, f"Prediksi: {top} - {note}"

demo = gr.Interface(
        fn=predict,
        inputs=gr.Image(type="numpy", label="Foto permukaan"),
        outputs=[
                    gr.Label(num_top_classes=3, label="Klasifikasi korosi"),
                    gr.Textbox(label="Catatan")
        ],
        title="FerroScan — Klasifikasi Korosi",
        description="MobileNetV2 + Transfer Learning. Triase visual, bukan pengganti NDT.",
        examples=[]
)

if __name__ == "__main__":
        demo.launch()