RingoDingo commited on
Commit
1f6d2ce
1 Parent(s): d02bda9

Upload 2 files

Browse files
Files changed (2) hide show
  1. maid_classifier_model.py +94 -0
  2. maidel_E22.pth +3 -0
maid_classifier_model.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ from PIL import Image
4
+ from torchvision.transforms import ToTensor
5
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
6
+
7
+ class MCNN(nn.Module):
8
+ def __init__(self):
9
+ super(MCNN, self).__init__()
10
+
11
+ # Convolution layers
12
+ self.conv1 = nn.Conv2d(3, 64, 3, 1, 1)
13
+ self.bn1 = nn.BatchNorm2d(64)
14
+
15
+ self.conv2 = nn.Conv2d(64, 128, 3, 1, 1)
16
+ self.bn2 = nn.BatchNorm2d(128)
17
+
18
+ self.conv3 = nn.Conv2d(128, 256, 3, 1, 1)
19
+ self.bn3 = nn.BatchNorm2d(256)
20
+
21
+ self.conv4 = nn.Conv2d(256, 512, 3, 1, 1) # Added another convolutional layer
22
+ self.bn4 = nn.BatchNorm2d(512)
23
+
24
+ # Pooling layer
25
+ self.pool = nn.MaxPool2d(2, 2)
26
+
27
+ # Fully connected layers
28
+ self.fc1 = nn.Linear(100352, 2048)
29
+ self.fc2 = nn.Linear(2048, 1024)
30
+ self.fc3 = nn.Linear(1024, 512)
31
+ self.fc4 = nn.Linear(512, 256)
32
+ self.fc5 = nn.Linear(256, 2) # Two classes
33
+
34
+ # Activation and dropout
35
+ self.relu = nn.ReLU()
36
+ self.dropout = nn.Dropout(0.2)
37
+
38
+ def forward(self, pixel_values, labels=None):
39
+ x = self.pool(self.relu(self.bn1(self.conv1(pixel_values))))
40
+ x = self.pool(self.relu(self.bn2(self.conv2(x))))
41
+ x = self.pool(self.relu(self.bn3(self.conv3(x))))
42
+ x = self.pool(self.relu(self.bn4(self.conv4(x)))) # Pass through the added conv layer
43
+
44
+ x = x.view(x.size(0), -1) # flatten
45
+ x = self.dropout(self.relu(self.fc1(x)))
46
+ x = self.dropout(self.relu(self.fc2(x)))
47
+ x = self.dropout(self.relu(self.fc3(x)))
48
+ x = self.dropout(self.relu(self.fc4(x)))
49
+ logits = self.fc5(x)
50
+
51
+ loss = None
52
+ if labels is not None:
53
+ loss_fct = nn.CrossEntropyLoss()
54
+ loss = loss_fct(logits.view(-1, 2), labels.view(-1))
55
+
56
+ if loss is not None:
57
+ return logits, loss.item()
58
+ else:
59
+ return logits, None
60
+
61
+ def preprocess_image(img, desired_size=224):
62
+ im = img
63
+
64
+ # Resize and pad the image
65
+ old_size = im.size
66
+ ratio = float(desired_size) / max(old_size)
67
+ new_size = tuple([int(x*ratio) for x in old_size])
68
+ im = im.resize(new_size)
69
+
70
+ # Create a new image and paste the resized on it
71
+ new_im = Image.new("RGB", (desired_size, desired_size), "white")
72
+ new_im.paste(im, ((desired_size-new_size[0])//2,
73
+ (desired_size-new_size[1])//2))
74
+ return new_im
75
+
76
+ def predict_image(image, model):
77
+ # Ensure model is in eval mode
78
+ model.eval()
79
+
80
+ # Convert image to tensor
81
+ transform = ToTensor()
82
+ input_tensor = transform(image)
83
+ input_batch = input_tensor.unsqueeze(0)
84
+
85
+ # Move tensors to the right device
86
+ input_batch = input_batch.to(device)
87
+
88
+ # Forward pass of the image through the model
89
+ output = model(input_batch)
90
+
91
+ # Convert model output to probabilities using softmax
92
+ probabilities = torch.nn.functional.softmax(output[0], dim=1)
93
+
94
+ return probabilities.cpu().detach().numpy()
maidel_E22.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c785bfd1ecccce3fb08e55eb91b11a05d122577a100e20f2d4acb0564cd234b4
3
+ size 839342059