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

weights = torchvision.models.DenseNet169_Weights.DEFAULT
dense_tranform = weights.transforms()

transfermodel = torchvision.models.densenet169(weights = weights)
transfermodel.classifier = nn.Sequential(nn.Linear(1664, 800), nn.ReLU(),
                                          nn.Linear(800, 400), nn.ReLU(),
                                 nn.Linear(400, 2))
transfermodel.load_state_dict(torch.load('transfermodel.pth'))

class_names=['NORMAL', 'PNEUMONIA']

def predict(img):
    img = dense_tranform(img).unsqueeze(0)

    transfermodel.eval()
    transfermodel.to("cpu")
    with torch.inference_mode():
        pred_probs = torch.softmax(transfermodel(img), dim=1)

    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
    return pred_labels_and_probs


title = "Zatürre Bulucu"
description = "Gönderilen fotoğrafa göre Sağlıklı mı yoksa Zatürre mi olduğunu tahmin eder."

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Label(num_top_classes=2, label="Predictions")],
    title=title,
    description=description
)

demo.launch(debug=False, share=True)