import gradio as gr
import torch
import torchvision.transforms as transforms
from PIL import Image
from model import load_model

# Load model
model_path = 'best_model_augmented.pth'
model = load_model(model_path)

# Define preprocessing transformations
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
])

# Define the class names
class_names = ['Normal', 'Monkeypox', 'Chickenpox', 'Measles']

def predict_image(image):
    # Preprocess the image
    image = Image.fromarray(image)
    image = transform(image).unsqueeze(0)
    
    # Perform inference
    model.eval()
    with torch.no_grad():
        outputs = model(image)
        _, predicted = torch.max(outputs, 1)
        predicted_class = class_names[predicted.item()]
    
    return predicted_class

# HTML for custom styling
title = (
    "<h1 style='text-align: center; font-family: Arial, sans-serif;'>PoxNet</h1>"
    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
    "<a href='https://x.com/ronith_sharmila' target='_blank'>@ronithsharmila</a>"
    "</p>"
)

description = (
    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
    "Upload an image of a skin lesion to classify it into one of the following categories:</p>"
    "<p style='text-align: center; font-family: Arial, sans-serif; font-weight: bold;'>"
    "Normal, Monkeypox, Chickenpox, Measles</p>"
    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
    "Please note that this model is not a substitute for professional medical advice, diagnosis, or treatment. "
    "Always seek the advice of your physician or other qualified health provider with any questions you may have regarding a medical condition."
    "</p>"
    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
    "<a href='https://paypal.me/ronithsharmila?country.x=US&locale.x=en_US' target='_blank'>SUPPORT MY PROJECTS HERE !</a>"
    "</p>"
)

# Gradio app interface
app = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="numpy", label="Upload Image"),
    outputs=gr.Textbox(label="Prediction"),
    title=title,
    description=description,
    theme="default"  # Use the default theme or specify a different theme if available
)

# Launch the app
app.launch()