from fastai.vision.all import *
import gradio as gr
from pathlib import Path
import zipfile

# Extraer dataset.zip al inicio
with zipfile.ZipFile("dataset.zip", 'r') as zip_ref:
    zip_ref.extractall(".")  # crea la carpeta dataset/ directamente

# Ruta relativa al dataset extraído
path = Path("dataset")

# Crear DataLoaders desde las subcarpetas de clases
dls = ImageDataLoaders.from_folder(
    path,
    valid_pct=0.2,
    seed=42,
    item_tfms=Resize(224)
)

# Cargar el modelo .pth
learn = vision_learner(dls, resnet34)
learn.load("model_lab")  # si lo pusiste en models/, usar "models/model_lab"

# Tomar automáticamente las clases desde los DataLoaders
labels = learn.dls.vocab

# Función de predicción
def predict(img):
    img = PILImage.create(img)
    _, _, probs = learn.predict(img)
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

# Interfaz Gradio
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="filepath"),
    outputs=gr.Label(num_top_classes=3),
    title="Lab Utensils Classifier"
)

demo.launch()