readme update

README.md

@@ -71,19 +71,47 @@ This repository contains a **custom convolutional neural network** trained on sa
 class Net(nn.Module):
     def __init__(self):
         super().__init__()
-        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
-        self.conv2 = nn.Conv2d(16, 8, kernel_size=3, padding=1)
-        self.fc1 = nn.Linear(8 * 64 * 64, 32)
-        self.fc2 = nn.Linear(32, 4)
+        
+        # Feature extractor
+        self.conv1 = nn.Conv2d(4, 16, kernel_size=3, padding=1) # (RGB + EDGE: 3 + 1)
+        self.bn1 = nn.BatchNorm2d(16)
+        
+        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(32)
+        
+        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm2d(64)
+        
+        # After 3x maxpool (stride=2), 256 -> 128 -> 64 -> 32
+        self.fc1 = nn.Linear(64 * 32 * 32, 256)
+        self.fc2 = nn.Linear(256, 64)
+        self.fc3 = nn.Linear(64, 4)  # 4 classes
+
+        self.dropout = nn.Dropout(0.5)
 
     def forward(self, x):
-        out = F.max_pool2d(torch.tanh(self.conv1(x)), 2)
-        out = F.max_pool2d(torch.tanh(self.conv2(out)), 2)
-        out = out.view(-1, 8 * 64 * 64)
-        out = torch.tanh(self.fc1(out))
-        out = self.fc2(out)
+        # Conv layers
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = F.max_pool2d(out, 2)  # 256 -> 128
+        
+        out = F.relu(self.bn2(self.conv2(out)))
+        out = F.max_pool2d(out, 2)  # 128 -> 64
+        
+        out = F.relu(self.bn3(self.conv3(out)))
+        out = F.max_pool2d(out, 2)  # 64 -> 32
+        
+        # Flatten
+        out = out.view(out.size(0), -1)
+        
+        # Fully connected layers
+        out = F.relu(self.fc1(out))
+        out = self.dropout(out)
+        
+        out = F.relu(self.fc2(out))
+        out = self.fc3(out)  # logits, apply CrossEntropyLoss
         return out
 
+
 ```
 ## Example Notebook
 
@@ -93,7 +121,7 @@ We provide an example notebook that demonstrates how to train the model:
 
 This notebook is intended as a starting point for experimentation and helps you quickly see how to use the dataset in practice.
 
-Accuracy is 0.9728
+Accuracy is 0.9745
 
 ## IT infrastructure
 

install.sh

@@ -2,3 +2,4 @@
 # Install Git LFS
 sudo apt-get update && sudo apt-get install -y git-lfs
 git lfs install
+sudo apt-get install -y libgl1