1import torch
2import torch.nn as nn
3import torchvision.transforms as transforms
4from PIL import Image
5
6# --------------------------------------------------
7# ResNet20 for CIFAR-10 (exact same architecture)
8# --------------------------------------------------
9def conv3x3(in_planes, out_planes, stride=1):
10 return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)
11
12class BasicBlock(nn.Module):
13 expansion = 1
14 def __init__(self, in_planes, planes, stride=1):
15 super(BasicBlock, self).__init__()
16 self.conv1 = conv3x3(in_planes, planes, stride)
17 self.bn1 = nn.BatchNorm2d(planes)
18 self.relu = nn.ReLU(inplace=True)
19 self.conv2 = conv3x3(planes, planes)
20 self.bn2 = nn.BatchNorm2d(planes)
21 self.shortcut = nn.Sequential()
22 if stride != 1 or in_planes != self.expansion * planes:
23 self.shortcut = nn.Sequential(
24 nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
25 nn.BatchNorm2d(self.expansion * planes)
26 )
27
28 def forward(self, x):
29 out = self.relu(self.bn1(self.conv1(x)))
30 out = self.bn2(self.conv2(out))
31 out += self.shortcut(x)
32 out = self.relu(out)
33 return out
34
35class ResNet(nn.Module):
36 def __init__(self, block, num_blocks, num_classes=10):
37 super(ResNet, self).__init__()
38 self.in_planes = 16
39 self.conv1 = conv3x3(3, 16)
40 self.bn1 = nn.BatchNorm2d(16)
41 self.relu = nn.ReLU(inplace=True)
42 self.layer1 = self._make_layer(block, 16, num_blocks[0], stride=1)
43 self.layer2 = self._make_layer(block, 32, num_blocks[1], stride=2)
44 self.layer3 = self._make_layer(block, 64, num_blocks[2], stride=2)
45 self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
46 self.fc = nn.Linear(64 * block.expansion, num_classes)
47
48 def _make_layer(self, block, planes, num_blocks, stride):
49 strides = [stride] + [1] * (num_blocks - 1)
50 layers = []
51 for stride in strides:
52 layers.append(block(self.in_planes, planes, stride))
53 self.in_planes = planes * block.expansion
54 return nn.Sequential(*layers)
55
56 def forward(self, x):
57 out = self.relu(self.bn1(self.conv1(x)))
58 out = self.layer1(out)
59 out = self.layer2(out)
60 out = self.layer3(out)
61 out = self.avg_pool(out)
62 out = out.view(out.size(0), -1)
63 out = self.fc(out)
64 return out
65
66def ResNet20():
67 return ResNet(BasicBlock, [3, 3, 3])
68
69# --------------------------------------------------
70# Load the model weights
71# --------------------------------------------------
72device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
73model = ResNet20().to(device)
74model.load_state_dict(torch.load("best_resnet20_cifar10.pth", map_location=device))
75model.eval()
76
77# --------------------------------------------------
78# Preprocess an image (must be 32x32 RGB)
79# --------------------------------------------------
80transform = transforms.Compose([
81 transforms.Resize((32, 32)),
82 transforms.ToTensor(),
83 transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
84])
85
86def predict(image_path):
87 image = Image.open(image_path).convert('RGB')
88 input_tensor = transform(image).unsqueeze(0).to(device)
89 with torch.no_grad():
90 output = model(input_tensor)
91 _, predicted = output.max(1)
92 classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
93 return classes[predicted.item()]
94
95# Example usage:
96# print(predict("my_cat.jpg"))