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

# Define device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Define classes
classes = ["Eyelid", "Normal", "Cataract", "Uveitis", "Conjunctivitis"]
num_classes = len(classes)

# Define image transformations (same as training)
data_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

# Load the trained model
def load_model():
    model = mobilenet_v2(weights=None)
    model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
    model.load_state_dict(torch.load("best_mobilenetv2.pth", map_location=device))
    model = model.to(device)
    model.eval()
    return model

# Inference function
def predict(image):
    try:
        # Load model
        model = load_model()
        
        # Convert Gradio image input (numpy array) to PIL Image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image).convert('RGB')
        else:
            image = image.convert('RGB')
        
        # Apply transformations
        image = data_transforms(image).unsqueeze(0).to(device)
        
        # Get model predictions
        with torch.no_grad():
            outputs = model(image)
            probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
            predicted_idx = np.argmax(probs)
            predicted_class = classes[predicted_idx]
        
        # Prepare output
        result = {
            "Predicted Class": predicted_class,
            "Probabilities": {classes[i]: f"{probs[i]:.4f}" for i in range(num_classes)}
        }
        return result
    except Exception as e:
        return f"Error during prediction: {str(e)}"

# Create Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload Eye Image"),
    outputs=gr.JSON(label="Prediction Results"),
    title="Eye Disease Classification",
    description="Upload an eye image to classify it as one of: Eyelid, Normal, Cataract, Uveitis, or Conjunctivitis."
)

# Launch the interface (not needed for Hugging Face Spaces, included for local testing)
if __name__ == "__main__":
    interface.launch()