import gradio as gr
import os
import torch

#from model import create_densenet_model
import model
from timeit import default_timer as timer


class_names = ["Covid", "Normal", "Viral Pneumonia"]

model_1, transforms = model.create_densenet_model(num_classes=3)


state_dict = torch.load(
    f="efficientv2_43loss_covid_3_classes_model.pth",
    weights_only=False,
    map_location="cpu"
)

model_1.load_state_dict(state_dict())


def predict_img(img):

  start_time = timer()

  img = transforms(img).unsqueeze(0)

  model_1.eval()

  with torch.inference_mode():

    pred_probs = torch.softmax(model_1(img), dim=1)

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

   #pred_time = round(timer() - start_time(),5)
    pred_time = round(timer() - start_time, 5)

    return pred_labels_and_probs, pred_time

title = "Covid Lung Classifier: AI-Driven Pulmonary Assessment"
description = "Upload a chest X-ray for an automated assessment. This system uses EfficientNetV2 deep learning to identify Normal, Viral Pneumonia, and COVID-19 cases with 92% accuracy.\nDisclaimer: These AI tools are for informational and research purposes. Medical diagnoses must be made by qualified healthcare professionals."
article = "Created at Mauaque Resettlement Gonzales Compound 2026"

example_list = [["examples/" + example] for example in os.listdir("examples")]

demo = gr.Interface(
    fn=predict_img,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Label(num_top_classes=3,label="Predictions"),
        gr.Number(label="Prediction Time")
    ],
    examples = example_list,
    title = title,
    description = description,
    article = article
)


demo.launch(debug=True)