import torch
from torch import nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr
import numpy as np

# Define model (pretrained ResNet50)
model = models.resnet50(pretrained=True)
model.fc = nn.Linear(model.fc.in_features, 3)  # 3 output classes: Normal, Hypothyroidism, Hyperthyroidism

# Define image transformations (resizing and normalization)
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # Resize image for input
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Example classification function
def classify_thyroid_condition(image):
    image = Image.fromarray(image.astype('uint8'), 'RGB')  # Convert numpy array to Pillow Image
    image = transform(image).unsqueeze(0)  # Add batch dimension

    model.eval()  # Set the model to evaluation mode
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output, 1)

    # Map prediction to class labels
    if predicted.item() == 0:
        diagnosis = "Normal"
    elif predicted.item() == 1:
        diagnosis = "Hypothyroidism"
    else:
        diagnosis = "Hyperthyroidism"

    return diagnosis

# Create Gradio interface for image input
gr.Interface(fn=classify_thyroid_condition, inputs="image", outputs="text").launch()