experimental

This commit is contained in:
2024-04-28 22:27:25 +08:00
parent 2d95b112c5
commit 215cde2d19
2 changed files with 689 additions and 38 deletions

View File

@@ -17,17 +17,17 @@ import numpy as np
import torch
import os
class CNN3D(nn.Module):
def __init__(self):
def __init__(self, hidden_size=32, dropout=0.0):
super(CNN3D, self).__init__()
self.conv1 = nn.Conv3d(1, 32, kernel_size=3, stride=1, padding=1)
self.conv2 = nn.Conv3d(32, 64, kernel_size=3, stride=1, padding=1)
self.batchnorm = nn.BatchNorm3d(32)
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(1024, 256) # Calculate input size based on output from conv3
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)
@@ -37,6 +37,7 @@ class CNN3D(nn.Module):
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)
@@ -49,7 +50,6 @@ def train(model, criterion, optimizer, loader, epochs=5):
for idx, (inputs, labels) in enumerate(loader):
optimizer.zero_grad()
outputs = model(inputs)
# print(outputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
@@ -57,44 +57,22 @@ def train(model, criterion, optimizer, loader, epochs=5):
return model
def process_data(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
num_samples_to_take = 600
# Get the indices of the samples to take
indices_to_take = [np.random.choice(indices[i], num_samples_to_take, 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)
y = y[indices_to_take]
return torch.Tensor(tensor_videos), torch.Tensor(y).long()
class Model():
def __init__(self):
self.model = CNN3D()
def __init__(self, batch_size=8,lr=0.001,epochs=10, dropout=0.0, hidden_size=32):
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=0.001)
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
def fit(self, X, y):
X, y = process_data(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=128, shuffle=True)
train(self.model, self.criterion, self.optimizer, train_loader, 10)
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()
@@ -103,6 +81,12 @@ 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]):
@@ -111,6 +95,37 @@ class Model():
X = torch.Tensor(tensor_videos.unsqueeze(1))
result = self.model(X)
return torch.max(result, dim=1)[1].numpy()
def process_data(self, X, y, n_samples=600):
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]):
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], 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)
y = y[indices_to_take]
return torch.Tensor(tensor_videos), torch.Tensor(y).long()
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1)
@@ -118,6 +133,7 @@ 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)