feat: idek anymore

cs2109s/labs/final/final.py

@@ -16,6 +16,7 @@ from torch import nn
 import numpy as np
 import torch
 import os
+from torchvision.transforms.functional import equalize
 
 class CNN3D(nn.Module):
     def __init__(self, hidden_size=32, dropout=0.0):
@@ -27,7 +28,7 @@ class CNN3D(nn.Module):
         self.maxpool = nn.MaxPool3d(kernel_size=2, stride=2)
         self.fc1 = nn.Linear(hidden_size*32, 256)  # Calculate input size based on output from conv3
         self.fc2 = nn.Linear(256, 6)
-        self.dropout = nn.Dropout(dropout)
+        # self.dropout = nn.Dropout(dropout)
 
     def forward(self, x):
         x = self.conv1(x)
@@ -37,7 +38,7 @@ class CNN3D(nn.Module):
         x = self.conv2(x)
         x = self.relu(x)
         x = self.maxpool(x)
-        x = self.dropout(x)
+        # x = self.dropout(x)
 
         x = x.view(x.size(0), -1)  # Flatten features for fully connected layers
         x = self.fc1(x)
@@ -56,17 +57,16 @@ def train(model, criterion, optimizer, loader, epochs=5):
         print(f'Epoch {epoch}, Loss: {loss.item()}')
     return model
 
-
-
-
 class Model():
-    def __init__(self, batch_size=8,lr=0.001,epochs=10, dropout=0.0, hidden_size=32):
+    def __init__(self, batch_size=64,lr=0.001,epochs=5, dropout=0.0, hidden_size=32, n_samples=900):
+        print(batch_size, epochs, lr, dropout, hidden_size, n_samples)
         self.batch_size = batch_size
         self.lr = lr
         self.epochs = epochs
         self.model = CNN3D(dropout=dropout, hidden_size=hidden_size)
         self.criterion = nn.CrossEntropyLoss()
         self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
+        self.n_samples = n_samples
 
     def fit(self, X, y):
         X, y = self.process_data(X, y)
@@ -81,31 +81,25 @@ class Model():
             tensor_videos = torch.tensor(X, dtype=torch.float32)
             # Clip values to 0 and 255
             tensor_videos = np.clip(tensor_videos, 0, 255)
-            # TEMP
-            threshold = 180
-            tensor_videos[tensor_videos > threshold] = 255
-            tensor_videos[tensor_videos < threshold] = 0
-            # END TEMP
 
             # Replace NaNs in each frame, with the average of the frame. This was generated with GPT
             for i in range(tensor_videos.shape[0]):
                 for j in range(tensor_videos.shape[1]):
                     tensor_videos[i][j][torch.isnan(tensor_videos[i][j])] = torch.mean(
                         tensor_videos[i][j][~torch.isnan(tensor_videos[i][j])])
-            X = torch.Tensor(tensor_videos.unsqueeze(1))
-            result = self.model(X)
+            # tensor_videos = torch.Tensor(tensor_videos).to(torch.uint8).reshape(-1, 1, 16, 16)
+            # tensor_videos = equalize(tensor_videos).float().reshape(-1, 1, 6, 16, 16)
+            tensor_videos = torch.Tensor(tensor_videos).reshape(-1, 1, 6, 16, 16)
+            # some funky code to make the features more prominent
+
+            result = self.model(tensor_videos)
         return torch.max(result, dim=1)[1].numpy()
-    def process_data(self, X, y, n_samples=600):
+    def process_data(self, X, y):
         y = np.array(y)
         X = np.array([video[:6] for video in X])
         tensor_videos = torch.tensor(X, dtype=torch.float32)
         # Clip values to 0 and 255
         tensor_videos = np.clip(tensor_videos, 0, 255)
-        # TEMP
-        threshold = 180
-        tensor_videos[tensor_videos > threshold] = 255
-        tensor_videos[tensor_videos < threshold] = 0
-        # END TEMP
 
         # Replace NaNs in each frame, with the average of the frame. This was generated with GPT
         for i in range(tensor_videos.shape[0]):
@@ -118,13 +112,19 @@ class Model():
         indices = [np.argwhere(y == i).squeeze(1) for i in range(6)]
         # Get the number of samples to take for each class
         # Get the indices of the samples to take
-        indices_to_take = [np.random.choice(indices[i], n_samples, replace=True) for i in range(6)]
+        indices_to_take = [np.random.choice(indices[i], self.n_samples, replace=True) for i in range(6)]
         # Concatenate the indices
         indices_to_take = np.concatenate(indices_to_take)
         # Select the samples
-        tensor_videos = tensor_videos[indices_to_take].unsqueeze(1)
+        tensor_videos = tensor_videos[indices_to_take]
+
+        tensor_videos = torch.Tensor(tensor_videos).reshape(-1, 1, 6, 16, 16)
+        # Reshape the tensor to int for image processing
+        # tensor_videos = torch.Tensor(tensor_videos).to(torch.uint8).reshape(-1, 1, 16, 16)
+        # tensor_videos = equalize(tensor_videos).float().reshape(-1, 1, 6, 16, 16)
+
         y = y[indices_to_take]
-        return torch.Tensor(tensor_videos), torch.Tensor(y).long()
+        return tensor_videos, torch.Tensor(y).long()
 
 
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1)