import gradio as gr
import os
import torch

# BINARY
from timeit import default_timer as timer
import model

class_names = ["Normal","Pneumonia"]

densenet121, transforms = model.create_densenet_model(num_classes=2)


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

densenet121.load_state_dict(state_dict)


def predict(img):
      class_names = ["Normal", "Pneumonia"]
      start_time = timer()
                          
      img = transforms(img).unsqueeze(0)
      densenet121.eval()
                                                                
      with torch.inference_mode():
      # Get the probability for the positive class (Pneumonia)
          prob_pneumonia = torch.sigmoid(densenet121(img)).item()
                                                                                                                                                          # Calculate the probability for the negative class (Normal)
          prob_normal = 1.0 - prob_pneumonia
                                                                                                                                                                                                                                              
                                                                                                                                                                                                                                                        # Create the dictionary Gradio expects
          pred_labels_and_probs = {"Normal": float(prob_normal),"Pneumonia": float(prob_pneumonia)}                                                                                                         
          pred_time = round(timer() - start_time, 5)
          return pred_labels_and_probs, pred_time


example_list = 'examples' # The path to your directory

import gradio as gr

# Create title, description and article strings
title = "AI-Driven Diagnostic Assistant: Breast Cancer & Pneumonia Classification"
description = " Engineered a high-precision computer vision pipeline using DenseNet121 to assist in the automated screening of medical imaging. The model achieves 90% accuracy in identifying pathologies across MRI and X-ray datasets. To ensure accessibility, I deployed the model via a Gradio web interface, allowing for real-time inference and a streamlined 'human-in-the-loop' diagnostic workflow.\nDisclaimer: These AI tools are for informational and research purposes. Medical diagnoses must be made by qualified healthcare professionals."
article = "Created at Mauaque Resettlement Center Gonzales Compound"

# Create the Gradio demo
demo = gr.Interface(fn=predict, # mapping function from input to output
                    inputs=gr.Image(type="pil"), # what are the inputs?
                    outputs=[gr.Label(num_top_classes=2, label="Predictions Result"), # what are the outputs?
                    gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
                    examples=example_list,
                    title=title,
                    description=description,
                    article=article)

                                                                                                                                                                                                                                                                                                 # Launch the demo!
                                                                                                                                                                                                                                                                                                 
demo.launch(debug=True, # print errors locally?
            share=True) # generate a publically shareable URL?