File size: 3,657 Bytes
3bc6b95 | 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 | from pathlib import Path
import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
class SimpleCNN(nn.Module):
# This must match the network trained in the original notebook.
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 32, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(),
nn.MaxPool2d(2),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128 * 4 * 4, 256),
nn.ReLU(),
nn.Dropout(0.30),
nn.Linear(256, 10),
)
def forward(self, x):
return self.classifier(self.features(x))
MODEL_PATH = Path(__file__).with_name("cifar10_cnn.pt")
DEVICE = torch.device("cpu")
TRANSFORM = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2470, 0.2435, 0.2616),
),
])
def load_model():
if not MODEL_PATH.exists():
raise FileNotFoundError(
"Файл cifar10_cnn.pt не найден. Загрузите его в корень Space рядом с app.py."
)
try:
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=True)
except TypeError:
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
model = SimpleCNN().to(DEVICE)
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()
classes = checkpoint.get(
"classes",
["airplane", "automobile", "bird", "cat", "deer",
"dog", "frog", "horse", "ship", "truck"],
)
return model, classes
MODEL, CLASSES = load_model()
def predict(image):
if image is None:
return {}, "Загрузите изображение."
x = TRANSFORM(image.convert("RGB")).unsqueeze(0).to(DEVICE)
with torch.no_grad():
probabilities = torch.softmax(MODEL(x), dim=1)[0].cpu()
values, indices = torch.topk(probabilities, k=3)
results = {
CLASSES[index.item()]: float(value)
for value, index in zip(values, indices)
}
label = CLASSES[indices[0].item()]
confidence = float(values[0]) * 100
return results, f"Модель считает, что это: **{label}** ({confidence:.1f}%)."
with gr.Blocks() as demo:
gr.Markdown(
"""# 🖼️ Распознавание изображений — CIFAR-10
Загрузите фото. Модель выберет наиболее похожий класс: **самолёт, автомобиль, птица, кот, олень, собака, лягушка, лошадь, корабль или грузовик**.
> Это учебная модель, обученная на маленьких изображениях CIFAR-10. На обычных фото она может ошибаться."""
)
with gr.Row():
image_input = gr.Image(
label="Загрузите изображение",
type="pil",
sources=["upload", "webcam"],
)
result = gr.Label(label="Три наиболее вероятных класса", num_top_classes=3)
explanation = gr.Markdown()
button = gr.Button("Распознать", variant="primary")
button.click(predict, inputs=image_input, outputs=[result, explanation])
if __name__ == "__main__":
demo.launch()
|