import gradio as gr
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import joblib
import torchvision.transforms as transforms
from PIL import Image

class STL10Net(nn.Module):
    def __init__(self):
        super(STL10Net, self).__init__()        
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.bn1 = nn.BatchNorm2d(32)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.bn2 = nn.BatchNorm2d(64)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.conv4 = nn.Conv2d(128, 256, 3, padding=1)
        self.bn4 = nn.BatchNorm2d(256)
        self.fc1 = nn.Linear(256 * 6 * 6, 512) 
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))
        x = self.pool(F.relu(self.bn2(self.conv2(x))))
        x = self.pool(F.relu(self.bn3(self.conv3(x))))
        x = self.pool(F.relu(self.bn4(self.conv4(x))))
        x = x.view(-1, 256 * 6 * 6)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x


original_torch_load = torch.load

def cpu_load(*args, **kwargs):
    kwargs['map_location'] = torch.device('cpu')
    return original_torch_load(*args, **kwargs)


torch.load = cpu_load 
try:
    model = joblib.load('stl10_cnn_model.pkl')
finally:
    torch.load = original_torch_load

model.to('cpu')
model.eval()

classes = joblib.load('stl10_target_names.pkl')

transform = transforms.Compose([
    transforms.Resize((96, 96)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

def predict_image(image):
    if image is None: return None
    image_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(image_tensor)
        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
    confidences = {classes[i]: float(probabilities[i]) for i in range(10)}
    return confidences


image_paths = [
    "Screenshot 2026-01-31 150516.png",
    "Screenshot 2026-01-31 150636.png",
    "Screenshot 2026-01-31 150840.png",
    "Screenshot 2026-01-31 151118.png",
    "Screenshot 2026-01-31 151244.png"
]

example_images = [[path] for path in image_paths]

def load_example_image(evt: gr.SelectData):
    """Load the selected example image into the input component."""
    selected_path = image_paths[evt.index]
    return Image.open(selected_path)

class_colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7", 
                "#DDA0DD", "#98D8C8", "#F7DC6F", "#BB8FCE", "#85C1E9"]
class_badges = " ".join([
    f'<span style="background-color: {class_colors[i]}; color: #000; padding: 4px 12px; border-radius: 15px; margin: 2px; display: inline-block; font-weight: 500;">{cls}</span>'
    for i, cls in enumerate(classes)
])

with gr.Blocks(title="STL-10 Image Classifier") as demo:
    gr.Markdown("# STL-10 Image Classifier")
    gr.Markdown("""
    Upload an image to classify it into one of the 10 STL-10 categories.
    
    The model is based on a CNN architecture trained on the STL-10 dataset.
    """)
    gr.Markdown("**Supported Classes:**")
    gr.HTML(f'<div style="margin: 10px 0; line-height: 2.2;">{class_badges}</div>')
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="Input Image")
            submit_btn = gr.Button("Classify", variant="primary")
            
            gr.Markdown("### Click an example image to select it:")
            example_gallery = gr.Gallery(
                value=image_paths,
                label="Examples",
                columns=5,
                rows=1,
                object_fit="contain",
                height="auto",
                allow_preview=False
            )
        
        with gr.Column():
            label_output = gr.Label(num_top_classes=3, label="Predictions")

    example_gallery.select(fn=load_example_image, inputs=None, outputs=image_input)
    submit_btn.click(fn=predict_image, inputs=image_input, outputs=label_output)

demo.launch()