Spaces:
Sleeping
Sleeping
File size: 4,284 Bytes
697180c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | 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() |