import torch
import torch.nn as nn
import torchvision.transforms as transforms
from PIL import Image
import gradio as gr

# Load model
class_names = {0: 'Budian', 1: 'Deri', 2: 'Putra', 3: 'Unknown', 4: 'Uqi', 5: 'Uwa'}
model = torch.load("model_deri.pth", map_location=torch.device('cpu'))
model.eval()

# Image preprocessing
def transform_image(image):
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5], std=[0.5])
    ])
    return transform(image).unsqueeze(0)  # Add batch dimension

# Prediction function
def predict(image):
    image = transform_image(image)
    with torch.no_grad():
        output = model(image)
        predicted_class = torch.argmax(output, dim=1).item()
    return class_names[predicted_class]

# Gradio UI
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Face Recognition Model",
    description="Upload an image and get the predicted class."
)

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