Upload 2 files

My bad scripts I used to make the model/dataset

prepare_data.py

@@ -0,0 +1,299 @@
+import os
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader, Sampler
+from tqdm import tqdm
+import librosa
+import logging
+import argparse
+import json
+import time
+import torchaudio
+from torchvision import transforms
+import pickle
+import random
+
+def configure_logging():
+    logging.basicConfig(level=logging.DEBUG, 
+                        format='%(asctime)s - %(levelname)s - %(message)s', 
+                        handlers=[
+                            logging.StreamHandler()
+                        ])
+    logging.info("Logging is set up.")
+    print("Logging is set up.")
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Spectrogram Dataset Preparation')
+    parser.add_argument('--config', type=str, required=True, help='Path to the config file')
+    return parser.parse_args()
+
+def load_config(config_path):
+    logging.info(f"Loading configuration from {config_path}")
+    print(f"Loading configuration from {config_path}")
+    try:
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        logging.info("Configuration loaded successfully")
+        print("Configuration loaded successfully")
+        return config
+    except Exception as e:
+        logging.error(f"Failed to load config file: {e}", exc_info=True)
+        print(f"Failed to load config file: {e}")
+        raise
+
+def validate_audio(y, sr, target_sr=44100, min_duration=0.1):
+    logging.debug(f"Validating audio with sr={sr}, target_sr={target_sr}, min_duration={min_duration}")
+    print(f"Validating audio with sr={sr}, target_sr={target_sr}, min_duration={min_duration}")
+    if sr != target_sr:
+        logging.warning(f"Resampling from {sr} to {target_sr}")
+        print(f"Resampling from {sr} to {target_sr}")
+        y = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
+    if len(y) < min_duration * target_sr:
+        pad_length = int(min_duration * target_sr - len(y))
+        y = np.pad(y, (0, pad_length), mode='constant')
+        logging.info(f"Audio file padded with {pad_length} samples")
+        print(f"Audio file padded with {pad_length} samples")
+    return y, target_sr
+
+def strip_silence(y, sr, top_db=20, pad_duration=0.1):
+    logging.debug(f"Stripping silence with sr={sr}, top_db={top_db}, pad_duration={pad_duration}")
+    print(f"Stripping silence with sr={sr}, top_db={top_db}, pad_duration={pad_duration}")
+    y_trimmed, _ = librosa.effects.trim(y, top_db=top_db)
+    pad_length = int(pad_duration * sr)
+    y_padded = np.pad(y_trimmed, pad_length, mode='constant')
+    return y_padded
+
+def audio_to_spectrogram(file_path, n_fft=2048, hop_length=256, n_mels=128, target_sr=44100, min_duration=0.1):
+    try:
+        logging.info(f"Loading file: {file_path}")
+        print(f"Loading file: {file_path}")
+        y, sr = librosa.load(file_path, sr=None)
+        logging.debug(f"Loaded file: {file_path} with sr={sr}")
+        print(f"Loaded file: {file_path} with sr={sr}")
+        y, sr = validate_audio(y, sr, target_sr, min_duration)
+        y = strip_silence(y, sr)
+    except Exception as e:
+        logging.error(f"Error reading {file_path}: {e}", exc_info=True)
+        print(f"Error reading {file_path}: {e}")
+        return None
+
+    y = librosa.util.normalize(y)
+    S = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=n_fft, hop_length=hop_length, n_mels=n_mels)
+    S_dB = librosa.power_to_db(S, ref=np.max)
+    logging.debug(f"Generated spectrogram for file: {file_path}")
+    print(f"Generated spectrogram for file: {file_path}")
+    
+    return S_dB
+
+def validate_spectrogram(spectrogram, n_mels=128):
+    logging.debug(f"Validating spectrogram with n_mels={n_mels}")
+    print(f"Validating spectrogram with n_mels={n_mels}")
+    if spectrogram.shape[0] != n_mels:
+        raise ValueError(f"Spectrogram has incorrect number of mel bands: {spectrogram.shape[0]}")
+    if spectrogram.shape[1] == 0:
+        raise ValueError("Spectrogram has zero frames")
+    return True
+
+def save_spectrogram(spectrogram, save_path):
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    np.save(save_path, spectrogram)
+    logging.debug(f"Spectrogram saved at {save_path}")
+    print(f"Spectrogram saved at {save_path}")
+
+class AddNoise(torch.nn.Module):
+    def __init__(self, noise_type='white', snr=10):
+        super(AddNoise, self).__init__()
+        self.noise_type = noise_type
+        self.snr = snr
+
+    def forward(self, waveform):
+        noise = torch.randn_like(waveform)
+        signal_power = waveform.norm(p=2)
+        noise_power = noise.norm(p=2)
+        noise = noise * (signal_power / noise_power) / (10 ** (self.snr / 20))
+        return waveform + noise
+
+class SpectrogramDataset(Dataset):
+    def __init__(self, config, directory, process_new=True):
+        logging.info("Initializing SpectrogramDataset...")
+        print("Initializing SpectrogramDataset...")
+        self.directory = directory
+        self.output_directory = config['output_directory']
+        self.spectrograms = []
+        self.labels = []
+        self.label_to_index = {}
+        self.process_new = process_new
+        self.config = config
+
+        # Paths for saving and loading cache
+        self.cache_path = os.path.join(self.output_directory, 'cache_data.npy')
+        self.dataset_path = os.path.join(self.output_directory, 'spectrogram_dataset.pkl')
+        
+        if os.path.exists(self.dataset_path):
+            self.load_dataset()
+        else:
+            if os.path.exists(self.cache_path):
+                os.remove(self.cache_path)
+                logging.info(f"Cache cleared at {self.cache_path}")
+                print(f"Cache cleared at {self.cache_path}")
+
+            self.load_data()
+            self.save_dataset()
+
+        self.transforms = transforms.Compose([
+            torchaudio.transforms.FrequencyMasking(freq_mask_param=30),
+            torchaudio.transforms.TimeMasking(time_mask_param=30)
+        ]) if self.config['augment'] else None
+
+        self.audio_transforms = torch.nn.Sequential(
+            AddNoise(snr=self.config['noise_snr']),
+            torchaudio.transforms.PitchShift(self.config['sample_rate'], n_steps=self.config['pitch_steps'])
+        ) if self.config['augment'] else None
+        logging.info("SpectrogramDataset initialized successfully")
+        print("SpectrogramDataset initialized successfully")
+
+    def save_dataset(self):
+        with open(self.dataset_path, 'wb') as f:
+            pickle.dump(self, f)
+        logging.info(f"Dataset object saved at {self.dataset_path}")
+        print(f"Dataset object saved at {self.dataset_path}")
+
+    def load_dataset(self):
+        with open(self.dataset_path, 'rb') as f:
+            obj = pickle.load(f)
+        self.__dict__.update(obj.__dict__)
+        logging.info(f"Dataset object loaded from {self.dataset_path}")
+        print(f"Dataset object loaded from {self.dataset_path}")
+
+    def process_file(self, file_path):
+        logging.debug(f"Processing file: {file_path}")
+        print(f"Processing file: {file_path}")
+        try:
+            label = os.path.basename(os.path.dirname(file_path))
+            if label not in self.label_to_index:
+                self.label_to_index[label] = len(self.label_to_index)
+            relative_path = os.path.relpath(file_path, self.directory)
+            spectrogram_path = os.path.join(self.output_directory, os.path.splitext(relative_path)[0] + '_spectrogram.npy')
+            if not os.path.exists(spectrogram_path) and self.process_new:
+                spectrogram = audio_to_spectrogram(file_path, n_fft=self.config['n_fft'], hop_length=self.config['hop_length'], n_mels=self.config['n_mels'], target_sr=self.config['sample_rate'], min_duration=self.config['min_duration'])
+                if spectrogram is not None:
+                    if spectrogram.shape[1] > self.config['max_frames']:
+                        spectrogram = spectrogram[:, :self.config['max_frames']]
+                    try:
+                        validate_spectrogram(spectrogram, n_mels=self.config['n_mels'])
+                        save_spectrogram(spectrogram, spectrogram_path)
+                        logging.debug(f"Spectrogram saved: {spectrogram_path}")
+                        print(f"Spectrogram saved: {spectrogram_path}")
+                    except Exception as e:
+                        logging.error(f"Error validating/saving spectrogram: {e}", exc_info=True)
+                        print(f"Error validating/saving spectrogram: {e}")
+            if os.path.exists(spectrogram_path):
+                try:
+                    spectrogram = np.load(spectrogram_path)
+                    validate_spectrogram(spectrogram, n_mels=self.config['n_mels'])
+                    spectrogram_tensor = torch.tensor(spectrogram, dtype=torch.float32)
+                    self.spectrograms.append(spectrogram_tensor)
+                    self.labels.append(self.label_to_index[label])
+                    logging.debug(f"Spectrogram loaded and appended for file: {file_path}")
+                    print(f"Spectrogram loaded and appended for file: {file_path}")
+                except Exception as e:
+                    logging.error(f"Error loading spectrogram {spectrogram_path}: {e}", exc_info=True)
+                    print(f"Error loading spectrogram {spectrogram_path}: {e}")
+        except Exception as e:
+            logging.error(f"Exception in process_file: {e}", exc_info=True)
+            print(f"Exception in process_file: {e}")
+
+    def load_data(self):
+        start_time = time.time()
+        logging.info("Starting to load and process files...")
+        print("Starting to load and process files...")
+        files_to_process = [os.path.join(root, file) for root, _, files in os.walk(self.directory) for file in files if file.lower().endswith('.wav')]
+        total_files = len(files_to_process)
+        logging.info(f"Total files to process: {total_files}")
+        print(f"Total files to process: {total_files}")
+
+        for file_path in tqdm(files_to_process, desc="Processing files"):
+            self.process_file(file_path)
+
+        end_time = time.time()
+        logging.info(f"Data loading and processing took {end_time - start_time:.2f} seconds")
+        print(f"Data loading and processing took {end_time - start_time:.2f} seconds")
+
+        self.save_cached_data(self.cache_path)
+
+    def save_cached_data(self, cache_path):
+        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
+        np.save(cache_path, {'spectrograms': self.spectrograms, 'labels': self.labels})
+        logging.debug(f"Cached data saved at {cache_path}")
+        print(f"Cached data saved at {cache_path}")
+
+    def __len__(self):
+        return len(self.spectrograms)
+
+    def __getitem__(self, idx):
+        spectrogram, label = self.spectrograms[idx], self.labels[idx]
+        if self.config['augment']:
+            if spectrogram.shape[1] >= 256:  # Ensure sufficient width for PitchShift
+                spectrogram = self.audio_transforms(spectrogram.unsqueeze(0)).squeeze(0)
+            spectrogram = self.transforms(spectrogram.unsqueeze(0)).squeeze(0)
+        return spectrogram, label
+
+def collate_fn(batch):
+    spectrograms, labels = zip(*batch)
+    labels = torch.tensor(labels, dtype=torch.long)
+    max_length = max(s.size(1) for s in spectrograms)
+    max_freq = max(s.size(0) for s in spectrograms)
+    spectrograms_padded = torch.zeros(len(spectrograms), max_freq, max_length)
+    for i, s in enumerate(spectrograms):
+        if s.dim() == 3 and s.size(2) == 1:
+            s = s.squeeze(2)
+        spectrograms_padded[i, :s.size(0), :s.size(1)] = s
+    return spectrograms_padded, labels
+
+class SmartBatchingSampler(Sampler):
+    def __init__(self, data_source, batch_size):
+        self.data_source = data_source
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        sorted_indices = sorted(range(len(self.data_source)), key=lambda i: self.data_source[i][0].shape[1], reverse=True)
+        pooled_indices = [sorted_indices[i:i + self.batch_size] for i in range(0, len(sorted_indices), self.batch_size)]
+        random.shuffle(pooled_indices)
+        for p in pooled_indices:
+            yield from p
+        if len(sorted_indices) % self.batch_size != 0:
+            yield from sorted_indices[-(len(sorted_indices) % self.batch_size):]
+
+    def __len__(self):
+        return len(self.data_source) // self.batch_size
+
+if __name__ == '__main__':
+    print("Starting script")
+    try:
+        args = parse_args()
+        print(f"Arguments parsed: {args}")
+        config = load_config(args.config)
+        print(f"Config loaded: {config}")
+
+        configure_logging()
+        print("Logging configured")
+
+        logging.info("Script started.")
+        dataset = SpectrogramDataset(config, config['directory'], process_new=True)
+        dataloader = DataLoader(dataset, batch_size=config['batch_size'], collate_fn=collate_fn, sampler=SmartBatchingSampler(dataset, config['batch_size']))
+        for batch in dataloader:
+            spectrograms, labels = batch
+            logging.info(f"Spectrograms batch shape: {spectrograms.shape}")
+            logging.info(f"Labels batch shape: {labels.shape}")
+            print(f"Spectrograms batch shape: {spectrograms.shape}")
+            print(f"Labels batch shape: {labels.shape}")
+            break
+
+        logging.info(f"Total files processed: {len(dataset)}")
+        print(f"Total files processed: {len(dataset)}")
+    except Exception as e:
+        logging.error(f"Exception occurred: {e}", exc_info=True)
+        print(f"Exception occurred: {e}")
+    finally:
+        logging.info("Script ended.")
+        print("Script ended")

train_model.py

@@ -0,0 +1,315 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader, WeightedRandomSampler, random_split, RandomSampler, SequentialSampler
+import logging
+import argparse
+import json
+from datetime import datetime
+import optuna
+from prepare_data import SpectrogramDataset, collate_fn
+from sklearn.metrics import classification_report, confusion_matrix
+import matplotlib.pyplot as plt
+import seaborn as sns
+import os
+import numpy as np
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f'Using device: {device}')
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, stride=1):
+        super(ResidualBlock, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        if stride != 1 or in_channels != out_channels:
+            self.shortcut = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+        else:
+            self.shortcut = nn.Identity()
+
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.bn2(self.conv2(out))
+        out += self.shortcut(x)
+        out = F.relu(out)
+        return out
+
+class AudioResNet(nn.Module):
+    def __init__(self, num_classes=6, dropout_rate=0.5):
+        super(AudioResNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(64, 64, num_blocks=2, stride=1)
+        self.layer2 = self._make_layer(64, 128, num_blocks=2, stride=2)
+        self.layer3 = self._make_layer(128, 256, num_blocks=2, stride=2)
+        self.layer4 = self._make_layer(256, 512, num_blocks=2, stride=2)
+
+        self.dropout = nn.Dropout(dropout_rate)
+        self.gap = nn.AdaptiveAvgPool2d((1, 1))  # Global Average Pooling
+        self.fc1 = nn.Linear(512, 1024)
+        self.fc2 = nn.Linear(1024, num_classes)
+
+    def _make_layer(self, in_channels, out_channels, num_blocks, stride):
+        layers = []
+        for i in range(num_blocks):
+            layers.append(ResidualBlock(in_channels if i == 0 else out_channels, out_channels, stride if i == 0 else 1))
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = self.maxpool(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.gap(x)  # Apply Global Average Pooling
+        x = x.view(x.size(0), -1)
+
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+# Example device configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f'Using device: {device}')
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger()
+fh = logging.FileHandler('training.log')
+fh.setLevel(logging.INFO)
+ch = logging.StreamHandler()
+ch.setLevel(logging.INFO)
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+fh.setFormatter(formatter)
+ch.setFormatter(formatter)
+logger.addHandler(fh)
+logger.addHandler(ch)
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Train Sample Classifier Model')
+    parser.add_argument('--config', type=str, required=True, help='Path to the config file')
+    return parser.parse_args()
+
+def load_config(config_path):
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+    with open(config_path, 'r') as f:
+        config = json.load(f)
+    return config
+
+def train_one_epoch(model, train_loader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    total_correct = 0
+
+    for batch_idx, (inputs, labels) in enumerate(train_loader):
+        inputs, labels = inputs.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(inputs.unsqueeze(1))
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item() * inputs.size(0)
+        _, predicted = torch.max(outputs, 1)
+        total_correct += (predicted == labels).sum().item()
+
+    train_loss = running_loss / len(train_loader.dataset)
+    train_accuracy = total_correct / len(train_loader.dataset)
+    return train_loss, train_accuracy
+
+def validate_one_epoch(model, val_loader, criterion, device):
+    model.eval()
+    val_loss = 0.0
+    val_correct = 0
+    with torch.no_grad():
+        for batch_idx, (inputs, labels) in enumerate(val_loader):
+            inputs, labels = inputs.to(device), labels.to(device)
+            outputs = model(inputs.unsqueeze(1))
+            loss = criterion(outputs, labels)
+            val_loss += loss.item() * inputs.size(0)
+            _, predicted = torch.max(outputs, 1)
+            val_correct += (predicted == labels).sum().item()
+
+    val_loss /= len(val_loader.dataset)
+    val_accuracy = val_correct / len(val_loader.dataset)
+    return val_loss, val_accuracy
+
+def train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, device, patience=10, max_epochs=50):
+    best_loss = float('inf')
+    patience_counter = 0
+
+    for epoch in range(max_epochs):
+        train_loss, train_accuracy = train_one_epoch(model, train_loader, criterion, optimizer, device)
+        val_loss, val_accuracy = validate_one_epoch(model, val_loader, criterion, device)
+
+        log_message = (f'Epoch {epoch + 1}:\n'
+                       f'Training Loss: {train_loss:.4f}, Training Accuracy: {train_accuracy:.4f}, '
+                       f'Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.4f}\n')
+        logging.info(log_message)
+
+        scheduler.step(val_loss)
+        current_lr = optimizer.param_groups[0]['lr']
+        logging.info(f'Current learning rate: {current_lr}')
+
+        if val_loss < best_loss:
+            best_loss = val_loss
+            patience_counter = 0
+            torch.save(model.state_dict(), 'best_model.pth')
+        else:
+            patience_counter += 1
+
+        if patience_counter >= patience:
+            logging.info('Early stopping triggered')
+            break
+
+        if (epoch + 1) % 10 == 0:
+            checkpoint_path = f'checkpoint_epoch_{epoch + 1}.pth'
+            torch.save(model.state_dict(), checkpoint_path)
+            logging.info(f'Model saved to {checkpoint_path}')
+
+def evaluate_model(model, test_loader, device, class_names):
+    model.eval()
+    all_labels = []
+    all_preds = []
+    with torch.no_grad():
+        for inputs, labels in test_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            outputs = model(inputs.unsqueeze(1))
+            _, preds = torch.max(outputs, 1)
+            all_labels.extend(labels.cpu().numpy())
+            all_preds.extend(preds.cpu().numpy())
+    logging.info(classification_report(all_labels, all_preds, target_names=class_names))
+    plot_confusion_matrix(all_labels, all_preds, class_names)
+
+def plot_confusion_matrix(labels, preds, class_names, save_path=None):
+    cm = confusion_matrix(labels, preds)
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=class_names, yticklabels=class_names)
+    plt.ylabel('Actual')
+    plt.xlabel('Predicted')
+    plt.title('Confusion Matrix')
+    if save_path:
+        plt.savefig(save_path)
+    plt.show()
+
+def objective(trial, train_loader, val_loader, num_classes):
+    learning_rate = trial.suggest_float('learning_rate', 1e-5, 1e-3, log=True)
+    weight_decay = trial.suggest_float('weight_decay', 1e-5, 1e-3, log=True)
+    dropout_rate = trial.suggest_float('dropout_rate', 0.2, 0.5)
+
+    model = AudioResNet(num_classes=num_classes, dropout_rate=dropout_rate).to(device)
+    criterion = nn.NLLLoss()
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3)
+
+    best_loss = float('inf')
+    patience_counter = 0
+
+    for epoch in range(10):
+        train_loss, _ = train_one_epoch(model, train_loader, criterion, optimizer, device)
+        val_loss, _ = validate_one_epoch(model, val_loader, criterion, device)
+        scheduler.step(val_loss)
+
+        if val_loss < best_loss:
+            best_loss = val_loss
+            patience_counter = 0
+        else:
+            patience_counter += 1
+
+        if patience_counter >= 3:
+            break
+
+    return val_loss
+
+def verify_dataset_and_loader(dataset, train_loader, val_loader, test_loader):
+    try:
+        logger.info(f"Dataset length: {len(dataset)}")
+        logger.info(f"Train dataset length: {len(train_loader.dataset)}")
+        logger.info(f"Validation dataset length: {len(val_loader.dataset)}")
+        logger.info(f"Test dataset length: {len(test_loader.dataset)}")
+
+        for idx in range(len(train_loader.dataset)):
+            _ = train_loader.dataset[idx]
+        logger.info("Train dataset verification passed")
+
+        for idx in range(len(val_loader.dataset)):
+            _ = val_loader.dataset[idx]
+        logger.info("Validation dataset verification passed")
+
+        for idx in range(len(test_loader.dataset)):
+            _ = test_loader.dataset[idx]
+        logger.info("Test dataset verification passed")
+    except IndexError as e:
+        logger.error(f"Dataset index error: {e}")
+
+def verify_sampler_indices(loader, name):
+    indices = list(loader.sampler)
+    logger.info(f"{name} sampler indices: {indices[:10]}... (total: {len(indices)})")
+    max_index = max(indices)
+    if max_index >= len(loader.dataset):
+        logger.error(f"{name} sampler index out of range: {max_index} >= {len(loader.dataset)}")
+    else:
+        logger.info(f"{name} sampler indices within range.")
+
+def main():
+    try:
+        args = parse_args()
+        config = load_config(args.config)
+
+        dataset = SpectrogramDataset(config, config['directory'], process_new=True)
+        if len(dataset) == 0:
+            raise ValueError("The dataset is empty. Please check the data loading process.")
+        num_classes = len(dataset.label_to_index)
+        class_names = list(dataset.label_to_index.keys())
+
+        train_size = int(0.7 * len(dataset))
+        val_size = int(0.15 * len(dataset))
+        test_size = len(dataset) - train_size - val_size
+        train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])
+
+        train_labels = [dataset.labels[i] for i in train_dataset.indices]
+        class_counts = np.bincount(train_labels)
+        class_weights = 1. / class_counts
+        sample_weights = class_weights[train_labels]
+        sampler = WeightedRandomSampler(sample_weights, len(sample_weights))
+
+        train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], collate_fn=collate_fn, sampler=sampler)
+        val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], collate_fn=collate_fn, sampler=RandomSampler(val_dataset))
+        test_loader = DataLoader(test_dataset, batch_size=config['batch_size'], collate_fn=collate_fn, sampler=SequentialSampler(test_dataset))
+
+        verify_dataset_and_loader(dataset, train_loader, val_loader, test_loader)
+        verify_sampler_indices(train_loader, "Train")
+        verify_sampler_indices(val_loader, "Validation")
+        verify_sampler_indices(test_loader, "Test")
+
+        study = optuna.create_study(direction='minimize')
+        study.optimize(lambda trial: objective(trial, train_loader, val_loader, num_classes), n_trials=50)
+
+        print('Best hyperparameters: ', study.best_params)
+
+        best_params = study.best_params
+        model = AudioResNet(num_classes=num_classes, dropout_rate=best_params['dropout_rate']).to(device)
+        criterion = nn.NLLLoss()
+        optimizer = optim.Adam(model.parameters(), lr=best_params['learning_rate'], weight_decay=best_params['weight_decay'])
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=3)
+
+        train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, device, patience=config['patience'])
+
+        model.load_state_dict(torch.load('best_model.pth'))
+        evaluate_model(model, test_loader, device, class_names)
+    except Exception as e:
+        logging.error(f"An error occurred: {e}")
+
+if __name__ == '__main__':
+    main()