import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np
import gradio as gr
from torch import nn
from gradio import components
from PIL import Image

class BrainTumorClassifier(nn.Module):
    def __init__(self, num_classes):
        super(BrainTumorClassifier, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 20, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
            nn.Conv2d(20, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2, 2)
        )
        self.classifier = nn.Sequential(
            nn.Linear(32 * 56 * 56, 128),  # Adjust input size based on image size
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )

    def forward(self, x):
        x = self.features(x)
        x = x.view(-1, 32 * 56 * 56)
        x = self.classifier(x)
        return x
    
def predict(image):
    image = Image.fromarray(np.uint8(image)).convert('RGB')
    ## give the weights trained
    model_path = 'cnn_tumorbrain_classifier_self.pth'
    model_load = BrainTumorClassifier(4)
    model_load.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    ## put the model in evaluation mode
    model_load.eval()

    transform_pipeline = transforms.Compose([
        transforms.Resize((224,224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))
    ])

    ## transform the img like the training image
    input_img = transform_pipeline(image).unsqueeze(0)
    # input_img

    ## define the label by index
    class_to_label = {0: 'glioma', 1: 'meningioma', 2: 'notumor', 3: 'pituitary'}

    ## run the model
    with torch.no_grad():
        output = model_load(input_img)

    ## convert to the softmax for getting percent each label
    probabilities = F.softmax(output, dim=1)

    ## get predicted label with highest value
    _, predicted_label = torch.max(probabilities,1)
    # confidence_percent = probabilities[0].tolist()[predicted_label.item()]
    conf, _ = torch.max(probabilities, 1)

    result = "{}, with confidence level in {}%".format(class_to_label[predicted_label.item()], conf.item()*100)
    return result

iface = gr.Interface(fn=predict,
                        inputs=gr.Image(),
                        outputs="textbox")

iface.launch(share=True)