Views
No views yet
1import torch
2import torch.nn as nn
3import torch.functional as F
4
5# CNN Block
6class ConvBlock(nn.Module):
7 def __init__(self, c0, c1, k, s, p) -> None:
8 super(ConvBlock, self).__init__()
9 self.net = nn.Sequential(
10 nn.Conv2d(c0, c1, k, s, p),
11 nn.BatchNorm2d(c1),
12 nn.GELU(),
13 )
14 def forward(self, x):
15 return self.net(x)
16
17 # The filter weight of each layer is a Gaussian distribution with zero mean and standard deviation initialized by random extraction 0.001 (deviation is 0).
18 def _initialize_weights(model):
19 """
20 Initializes weights of all layers in a PyTorch model.
21
22 Args:
23 model (nn.Module): The model to initialize weights for.
24 """
25 for m in model.modules():
26 if isinstance(m, nn.Conv2d):
27 nn.init.xavier_normal_(m.weight)
28 elif isinstance(m, nn.Linear):
29 nn.init.xavier_normal_(m.weight)
30 elif isinstance(m, nn.BatchNorm2d):
31 nn.init.constant_(m.weight, 1)
32 nn.init.constant_(m.bias, 0)
33
34# FCN Model
35class FullyConv(nn.Module):
36 def __init__(self, num_classes) -> None:
37 super(FullyConv, self).__init__()
38 # Input 28x28
39 self.net = nn.Sequential(
40 ConvBlock(3, 16, 3, 2, 1), # 14x14
41 ConvBlock(16, 64, 3, 2, 1), # 7x7
42 ConvBlock(64, 128, 3, 2, 1), # 4x4
43 ConvBlock(128, 256, 3, 2, 1), # 2x2
44 nn.Dropout(p=0.5, inplace=True),
45 nn.Conv2d(256, num_classes, 3, 2, 1), # 1x1
46 nn.Flatten()
47 )
48
49 # Initialize model weights.
50 self._initialize_weights()
51
52 def forward(self, x: torch.Tensor) -> torch.Tensor:
53 return self.net(x)
54
55
56 # The filter weight of each layer is a Gaussian distribution with zero mean and standard deviation initialized by random extraction 0.001 (deviation is 0).
57 def _initialize_weights(model):
58 """
59 Initializes weights of all layers in a PyTorch model.
60
61 Args:
62 model (nn.Module): The model to initialize weights for.
63 """
64 for m in model.modules():
65 if isinstance(m, nn.Conv2d):
66 nn.init.xavier_normal_(m.weight)
67 elif isinstance(m, nn.Linear):
68 nn.init.xavier_normal_(m.weight)
69 elif isinstance(m, nn.BatchNorm2d):
70 nn.init.constant_(m.weight, 1)
71 nn.init.constant_(m.bias, 0)
72
73# Init variable
74IMG_SIZE = 28
75transform = transforms.Compose([
76 transforms.Resize((IMG_SIZE, IMG_SIZE)),
77 transforms.Grayscale(3),
78 # transforms.ToDtype(torch.float32, scale=True),
79 transforms.ToTensor(),
80 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
81 ])
82
83class_names = ['あ', 'い', 'う', 'え', 'お', # A
84 'か', 'き', 'く', 'け', 'こ', # Ka
85 'さ', 'し', 'す', 'せ', 'そ', # Sa
86 'た', 'ち', 'つ', 'て', 'と', # Ta
87 'な', 'に', 'ぬ', 'ね', 'の', # Na
88 'は', 'ひ', 'ふ', 'へ', 'ほ', # Ha
89 'ま', 'み', 'む', 'め', 'も', # Ma
90 'や', 'ゆ', 'よ', # Ya
91 'ら', 'り', 'る', 'れ', 'ろ', # Ra
92 'わ', 'ゐ', 'ゑ', # Wa, ?, ?
93 'を', 'ん', 'ゝ'] # Wo, N, ?
94
95# Create model and load weight
96model = FullyConv(len(class_names))
97model = model.from_pretrained("Hendrico/kmnist49-classifier")
98
99# Predict function
100def predict(model, img, transform, class_names):
101 if type(img) == str:
102 img = Image.open(img).convert('RGB')
103 # img = PIL.ImageOps.invert(img)
104 inputs = transform(img).unsqueeze(0)
105 out = model(inputs)
106 act_out = F.softmax(out)
107 prob, pred = act_out.max(axis=1)
108 plt.title(f"{class_names[pred.item()]} ({prob.item()*100:.2f}%)")
109 plt.imshow(np.transpose(vutils.make_grid(inputs, padding=2, normalize=True).cpu(),(1,2,0)))
110 plt.show()
111
112# Called predict function based on image file or PIL image
113predict(model, image_file, transform, class_names)http://www.apache.org/licenses/LICENSE-2.0