"""
Simple Gradio app on http://127.0.0.1:7860
Hugging Face: https://huggingface.co/spaces/andreasmlanger/cat_dog
"""

import gradio as gr
import torch
from torch import nn
from torchvision import models, transforms
import os


MODEL_PATH = r'cat_dog_RNN.pth'  # path to trained model, needs to be in same directory
EXAMPLES = [[os.path.join('examples', f)] for f in os.listdir('examples')]  # example images
CLASS_NAMES = ['cat', 'dog']

# Load model architecture and saved weights
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
for param in model.parameters():
    param.requires_grad = False  # freeze pre-trained weights, so they don't get changed during training
model.fc = nn.Sequential(
    nn.Linear(in_features=model.fc.in_features, out_features=512),
    nn.ReLU(),
    nn.Dropout(0.2),
    nn.Linear(512, len(CLASS_NAMES), bias=True)
)
model.load_state_dict(torch.load(MODEL_PATH))


def predict(img):
    transform = transforms.Compose([transforms.Resize(size=(224, 224)), transforms.ToTensor()])
    transformed_image = transform(img).unsqueeze(0)
    model.eval()
    with torch.inference_mode():
        pred_probs = torch.softmax(model(transformed_image), dim=1)
    return {CLASS_NAMES[i]: float(pred_probs[0][i]) for i in range(len(CLASS_NAMES))}


# Create and launch the Gradio demo
demo = gr.Interface(fn=predict,
                    inputs=gr.Image(type='pil', label='Image'),
                    outputs=[gr.Label(num_top_classes=2, label='Predictions')],
                    examples=EXAMPLES,
                    title='Cat vs Dog Image Classifier',
                    description='A residual neuronal network designed to distinguish between cat and dog images.')

demo.launch(debug=False)