import gradio as gr
import torch
import cnn
from torchvision import transforms
from PIL import Image

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor()
])

model = cnn.CNN(2)
model = model.to("cpu")
model.load_state_dict(torch.load("cnn_model.pth", weights_only=True, map_location="cpu"))
model.eval()

label = ["Kucing", "Anjing"]

def inference(image):
    image = transform(image).unsqueeze(0)

    with torch.no_grad():
        output = model(image)
        output = torch.nn.functional.softmax(output, dim=1)

        predicted_class = torch.argmax(output, dim=1).item()
        score = output[0][predicted_class]

    return f'Ini adalah {label[predicted_class]} dengan kecocokan sebesar {score * 100}'

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            inputs = gr.Image(type="pil")
        with gr.Column():
            btn = gr.Button("Cek")
            pred = gr.Text(label="Prediction")

    btn.click(fn=inference, inputs=inputs, outputs=pred)

        

demo.queue().launch()