import numpy as np import os import torch from torchvision import transforms, datasets from torch.utils.data import Dataset, DataLoader from PIL import Image import random import torch.nn as nn import torch.nn.functional as F import torch import torch.optim as optim from torch.utils.data import random_split, DataLoader class CustomConcatDataset(Dataset): def __init__(self, audio_dataset_path, image_dataset_path, transform=None): self.audio_dataset_path = audio_dataset_path self.image_dataset_path = image_dataset_path self.transform = transform self.classes = [c for c in sorted(os.listdir(audio_dataset_path)) if c != '.ipynb_checkpoints'] self.class_to_audio_files = {} self.class_to_image_files = {} self.data_pairs = [] for c in self.classes: self.class_to_audio_files[c] = [f for f in sorted(os.listdir(os.path.join(audio_dataset_path, c))) if os.path.isfile(os.path.join(audio_dataset_path, c, f))] self.class_to_image_files[c] = [f for f in sorted(os.listdir(os.path.join(image_dataset_path, c))) if os.path.isfile(os.path.join(image_dataset_path, c, f))] # Create all possible pairs for the class all_pairs = [(c, audio_file, image_file) for audio_file in self.class_to_audio_files[c] for image_file in self.class_to_image_files[c]] # Shuffle and take up to 1000 pairs random.shuffle(all_pairs) selected_pairs = all_pairs[:100] self.data_pairs.extend(selected_pairs) # Print the number of selected pairs for the class print(f"Number of pairs for class {c}: {len(selected_pairs)}") def __len__(self): return len(self.data_pairs) def __getitem__(self, idx): c, audio_file, image_file = self.data_pairs[idx] audio_image_path = os.path.join(self.audio_dataset_path, c, audio_file) image_path = os.path.join(self.image_dataset_path, c, image_file) audio_image = Image.open(audio_image_path) img = Image.open(image_path) if self.transform: audio_image = self.transform(audio_image) img = self.transform(img) concat_image = torch.cat((audio_image, img), dim=0) label = self.classes.index(c) return concat_image, label transform = transforms.Compose([ transforms.Resize((256,256)), transforms.ToTensor() ]) dataset = CustomConcatDataset('./audio_dataset', './image_dataset', transform) dataloader = DataLoader(dataset, batch_size=32, shuffle=True) class CustomCNN(nn.Module): def __init__(self, num_classes): super(CustomCNN, self).__init__() # Convolutional layers self.conv1 = nn.Conv2d(7, 32, kernel_size=3, stride=1, padding=1) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1) self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1) self.conv4 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1) # Max pooling layer self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0) # Fully connected layers self.fc1 = nn.Linear(256 * 16 * 16, 512) self.fc2 = nn.Linear(512, num_classes) # Dropout for regularization self.dropout = nn.Dropout(0.5) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = self.pool(F.relu(self.conv3(x))) x = self.pool(F.relu(self.conv4(x))) # Flatten the tensor x = x.view(-1, 256 * 16 * 16) x = F.relu(self.fc1(x)) x = self.dropout(x) x = self.fc2(x) return x # Number of classes for classification num_classes = len(dataset.classes) # Instantiate the model model = CustomCNN(num_classes=num_classes) # Hyperparameters batch_size = 32 learning_rate = 0.001 num_epochs = 10 val_split = 0.2 # Splitting the dataset into training and validation sets dataset_size = len(dataset) val_size = int(val_split * dataset_size) train_size = dataset_size - val_size train_dataset, val_dataset = random_split(dataset, [train_size, val_size]) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False) # Initialize the CNN model, loss function, and optimizer model = CustomCNN(num_classes=len(dataset.classes)) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) # Move the model to GPU if available device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) for epoch in range(num_epochs): model.train() train_loss = 0.0 correct_train = 0 print(f"\nEpoch {epoch+1}/{num_epochs}") print('-' * 50) for batch_idx, (inputs, labels) in enumerate(train_loader): inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() _, preds = torch.max(outputs, 1) correct_train += (preds == labels).sum().item() # Print batch-level training log # if (batch_idx + 1) % 10 == 0: # batch_accuracy = 100 * (preds == labels).sum().item() / inputs.size(0) # print(f"Batch {batch_idx+1}/{len(train_loader)} - Loss: {loss.item():.4f}, Accuracy: {batch_accuracy:.2f}%") train_accuracy = 100 * correct_train / train_size # Evaluate the model on the validation set model.eval() val_loss = 0.0 correct_val = 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) val_loss += loss.item() _, preds = torch.max(outputs, 1) correct_val += (preds == labels).sum().item() val_accuracy = 100 * correct_val / val_size print(f"Epoch Results - Training Loss: {train_loss/train_size:.4f}, Training Accuracy: {train_accuracy:.2f}%") print(f"Validation Loss: {val_loss/val_size:.4f}, Validation Accuracy: {val_accuracy:.2f}%") print('-' * 50)