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

MODEL_PATH = 'banana_classifier.pth'
CLASS_NAMES = ['overripe', 'ripe', 'rotten', 'unripe']
IMG_SIZE = (224, 224)
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


model = models.efficientnet_b0(weights='IMAGENET1K_V1')
num_ftrs = model.classifier[1].in_features
model.classifier[1] = nn.Linear(num_ftrs, len(CLASS_NAMES))

model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
model = model.to(DEVICE)
model.eval()
print("Model loaded and ready for prediction.")


transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(IMG_SIZE),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])


def predict_image(image: Image.Image):
    try:
        image = image.convert('RGB')
        image_tensor = transform(image).unsqueeze(0).to(DEVICE)

        with torch.no_grad():
            outputs = model(image_tensor)
            probs = torch.nn.functional.softmax(outputs[0], dim=0)

        confidences = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
        return confidences

    except Exception as e:
        print("Error during prediction:", e)
        return {"error": str(e)}

iface = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="Upload a Banana Image"),
    outputs=gr.Label(num_top_classes=2, label="Prediction Results"),
    title="Banana Ripeness Classifier",
    description=(
        "Upload an image of a banana and the model will predict its ripeness level: "
        "**unripe**, **ripe**, **overripe**, or **rotten**."
    ),
    examples=[
        '2.jpg',
        '1.jpeg'
    ],
    allow_flagging="never"
)

if __name__ == "__main__":
    iface.launch()