144 lines
5.8 KiB
Python
144 lines
5.8 KiB
Python
import numpy as np
|
|
import torch
|
|
import os
|
|
|
|
from torch import nn
|
|
|
|
with open('data.npy', 'rb') as f:
|
|
data = np.load(f, allow_pickle=True).item()
|
|
X = data['data']
|
|
y = data['label']
|
|
|
|
from torch import nn
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
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):
|
|
super(CNN3D, self).__init__()
|
|
self.conv1 = nn.Conv3d(1, hidden_size, kernel_size=3, stride=1, padding=1)
|
|
self.batchnorm = nn.BatchNorm3d(hidden_size)
|
|
self.conv2 = nn.Conv3d(hidden_size, hidden_size*2, kernel_size=3, stride=1, padding=1)
|
|
self.relu = nn.ReLU()
|
|
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)
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.relu(x)
|
|
x = self.maxpool(x)
|
|
x = self.batchnorm(x)
|
|
x = self.conv2(x)
|
|
x = self.relu(x)
|
|
x = self.maxpool(x)
|
|
# x = self.dropout(x)
|
|
|
|
x = x.view(x.size(0), -1) # Flatten features for fully connected layers
|
|
x = self.fc1(x)
|
|
x = self.relu(x)
|
|
x = self.fc2(x)
|
|
return x
|
|
|
|
def train(model, criterion, optimizer, loader, epochs=5):
|
|
for epoch in range(epochs):
|
|
for idx, (inputs, labels) in enumerate(loader):
|
|
optimizer.zero_grad()
|
|
outputs = model(inputs)
|
|
loss = criterion(outputs, labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
print(f'Epoch {epoch}, Loss: {loss.item()}')
|
|
return model
|
|
|
|
class Model():
|
|
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)
|
|
train_dataset = torch.utils.data.TensorDataset(X, y)
|
|
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=self.batch_size, shuffle=True)
|
|
train(self.model, self.criterion, self.optimizer, train_loader, self.epochs)
|
|
|
|
def predict(self, X):
|
|
self.model.eval()
|
|
with torch.no_grad():
|
|
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)
|
|
|
|
# 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])])
|
|
# 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):
|
|
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)
|
|
|
|
# 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])])
|
|
# Undersample the data for each of the 6 classes. Select max of 300 samples for each class
|
|
# Very much generated with the assitance of chatGPT with some modifications
|
|
# Get the indices of each class
|
|
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], 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]
|
|
|
|
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 tensor_videos, torch.Tensor(y).long()
|
|
|
|
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1)
|
|
|
|
not_nan_indices = np.argwhere(~np.isnan(np.array(y_test))).squeeze()
|
|
y_test = [y_test[i] for i in not_nan_indices]
|
|
X_test = [X_test[i] for i in not_nan_indices]
|
|
|
|
print("init model")
|
|
model = Model()
|
|
model.fit(X_train, y_train)
|
|
|
|
from sklearn.metrics import f1_score
|
|
|
|
y_pred = model.predict(X_test)
|
|
print("F1 Score (macro): {0:.2f}".format(f1_score(y_test, y_pred, average='macro'))) # You may encounter errors, you are expected to figure out what's the issue.
|