OOP complete and functional rewrite

main.py

@@ -2,18 +2,30 @@ import os
 import numpy as np
 import torch
 import json
+import sys
 import matplotlib.pyplot as plt
 import argparse
 from sklearn.model_selection import train_test_split
 
-from src.data.download import download_clean
-from src.data.augment import create_augmented_datasets, create_log_mel, data_treatment_testing
+from src.data.download import ESC50Downloader
+from src.data.augment import AudioAugment
 from src.models.cnn import CNN
-from src.models.traincnn import train_k_fold_cnn, train_cnn
-from src.models.predict import predict_with_overlapping_patches, predict_top_k, predict_file, load_model
-from src.config.config import sample_rate, cnn_input_length, esc50_labels
+from src.models.predict import AudioPredictor
+from src.models.traincnn import CNNTrainer
+from src.config.config import ProcessingConfig, DatasetConfig, DownloadConfig, TrainConfig
 
-def main():
+def _load_or_preprocess(args) -> tuple[np.ndarray, np.ndarray]:
+    X_path = args.X_path or "data/preprocessed/X.npy"
+    y_path = args.y_path or "data/preprocessed/y.npy"
+    if os.path.exists(X_path) and os.path.exists(y_path):
+        print("Loading existing processed data...")
+        return np.load(X_path, allow_pickle=True), np.load(y_path)
+    print("Processing audio data...")
+    augmenter = AudioAugment()
+    augmenter.run(augment=True, preprocess=True)
+    return np.load(X_path, allow_pickle=True), np.load(y_path)
+
+def main() -> None:
     parser = argparse.ArgumentParser(
         description="ESC50 Audio Classification",
         formatter_class=argparse.RawDescriptionHelpFormatter
@@ -64,8 +76,8 @@ def main():
 
     resume_parser = subparsers.add_parser('resume', help='Resume training from checkpoint')
     resume_parser.add_argument('--resume-from', type=str, required=True, help='Path to checkpoint file')
-    resume_parser.add_argument('--X-path', type=str, help='Path to preprocessed X.npy')
-    resume_parser.add_argument('--y-path', type=str, help='Path to preprocessed y.npy')
+    resume_parser.add_argument('--X-path', type=str, default="data/preprocessed/X.npy")
+    resume_parser.add_argument('--y-path', type=str, default="data/preprocessed/y.npy")
     resume_parser.add_argument('--epochs', type=int, default=100, help='Number of epochs (default: 100)')
     resume_parser.add_argument('--batch-size', type=int, default=100, help='Batch size (default: 100)')
     resume_parser.add_argument('--lr', type=float, default=0.01, help='Learning rate (default: 0.01)')
@@ -76,7 +88,7 @@ def main():
 
     predict_parser = subparsers.add_parser('predict', help='Predict audio file class')
     predict_parser.add_argument('audio_file', type=str, help='Path to .wav file to classify')
-    predict_parser.add_argument('--model', type=str, default='best_model.pt', help='Path to model checkpoint (default: best_model.pt)')
+    predict_parser.add_argument('--model', type=str, default='final_model.pt', help='Path to model checkpoint (default: best_model.pt)')
     predict_parser.add_argument('--top-k', type=int, default=5, help='Number of top predictions (default: 5)')
     predict_parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device (default: auto)')
     predict_parser.set_defaults(func=cmd_predict)
@@ -84,194 +96,78 @@ def main():
     args = parser.parse_args()
     args.func(args)
 
-def cmd_download(args):
-    print("Download ESC50 audio data...")
+def cmd_download(args) -> None:
+    print("Downloading ESC50 audio data...")
+
+    downloader = ESC50Downloader()
 
-    download_clean()
+    downloader.download_clean()
 
-    print("Data downloaded and cleaned.")
+    print("Downloaded and cleaned data.")
 
-def cmd_augment(args):
+def cmd_augment(args) -> None:
     print("Augmenting audio data...")
-    create_augmented_datasets(args.input_dir, args.output_dir)
 
-    print(f"Saved augmented data to {args.output_dir}")
+    augmentater = AudioAugment()
 
-def cmd_preprocess(args):
-    print("Processing audio data...")
-       
-    print("Creating log-mel spectrograms...")
-    X, y = create_log_mel(args.input_dir, args.output_dir)
-    
-    print(f"Dataset size: {len(X)} samples, {len(np.unique(y))} classes")
-    print(f"Saved to {args.output_dir}")
+    augmentater.run(augment=True, preprocess=False)
 
-def cmd_train(args): 
-    device = "cuda" if torch.cuda.is_available() else "cpu" 
-    print(f"Using device: {device}") 
+    print(f"Augmented data and saved to {args.output_dir}")
 
-    X_path = args.X_path or "data/preprocessed/X.npy" 
-    y_path = args.y_path or "data/preprocessed/y.npy" 
+def cmd_preprocess(args) -> None:
+    print("Processing audio data...")
     
-    if os.path.exists(X_path) and os.path.exists(y_path): 
-        print("Loading existing processed data...") 
-        X = np.load(X_path, allow_pickle=True) 
-        y = np.load(y_path) 
-    else: 
-        print("Processing audio data...") 
-        audio_training_path = args.audio_dir or "data/audio/0" 
-        directories = os.listdir(audio_training_path) 
+    augmentater = AudioAugment()
 
-        if len(directories) == 1 and args.augment: 
-            print("Creating augmented datasets...") 
-            create_augmented_datasets(audio_training_path, "data/audio") 
+    augmentater.run(augment=False, preprocess=True)
 
-        print("Creating log-mel spectrograms...") 
-        X, y = create_log_mel(args.audio_dir or "data/audio", args.output_dir or "data/preprocessed") 
-            
-    print(f"Dataset size: {len(X)} samples, {len(np.unique(y))} classes") 
-        
-    X_train, X_val, y_train, y_val = train_test_split( X, y, test_size=0.2, random_state=42, stratify=y ) 
-    print(f"Train: {len(X_train)}, Val: {len(X_val)}") 
-    model = CNN(n_classes=len(np.unique(y))) 
-    best_val_acc = train_cnn( 
-        model, 
-        X_train, y_train, 
-        X_val, y_val, 
-        epochs=args.epochs, 
-        batch_size=args.batch_size, 
-        lr=args.lr, 
-        fold_num=0,
-        device=device, 
-        use_all_patches=True, 
-        samples_per_epoch_fraction=args.sample_fraction, 
-        checkpoint_dir=args.checkpoint_dir, 
-        save_every_n_epoch=args.save_every, 
-        resume_from=None ) 
-    
-    print(f"\nTraining complete! Best validation accuracy: {best_val_acc:.4f}") 
-    return best_val_acc
+    print(f"Preprocessed data and saved to {args.output_dir}")
 
-def cmd_train_cv(args):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-    
-    X_path = args.X_path or "data/preprocessed/X.npy"
-    y_path = args.y_path or "data/preprocessed/y.npy"
-    
-    if os.path.exists(X_path) and os.path.exists(y_path):
-        print("Loading existing processed data...")
-        X = np.load(X_path, allow_pickle=True)
-        y = np.load(y_path)
-    else:
-        print("Processing audio data...")
-        audio_training_path = args.audio_dir or "data/audio/0"
-        directories = os.listdir(audio_training_path)
-        
-        if len(directories) == 1 and args.augment:
-            print("Creating augmented datasets...")
-            create_augmented_datasets(audio_training_path, "data/audio")
-        
-        print("Creating log-mel spectrograms...")
-        X, y = create_log_mel(args.audio_dir or "data/audio", args.output_dir or "data/preprocessed")
-    
-    print(f"Dataset size: {len(X)} samples, {len(np.unique(y))} classes")
-    
-    X_train, X_val, y_train, y_val = train_test_split(
-        X, y, test_size=0.2, random_state=42, stratify=y
-    )
-    print(f"Train: {len(X_train)}, Val: {len(X_val)}")
-    
-    model = CNN(n_classes=len(np.unique(y)))
-    
-    fold_accs, mean_acc = train_k_fold_cnn(
-        model_class=lambda: CNN(),
-        X=X,
-        y=y,
+def cmd_train(args) -> None: 
+    X, y = _load_or_preprocess(args)
+    trainer = CNNTrainer(TrainConfig(
         epochs=args.epochs,
         batch_size=args.batch_size,
         lr=args.lr,
-        k_fold=args.k_fold,
-        device=device,
-        use_all_patches=True,
         samples_per_epoch_fraction=args.sample_fraction,
         checkpoint_dir=args.checkpoint_dir,
-        save_every_n_epoch=args.save_every
-    )
-    
-    print(f"\nTraining complete! Mean validation accuracy: {mean_acc:.4f}")
-    return mean_acc
+        save_every_n_epoch=args.save_every,
+    ))
+    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
+    best_val_acc = trainer.train_cnn(CNN(n_classes=len(np.unique(y))), X_train, y_train, X_val, y_val, fold_num=0)
+    print(f"\nTraining complete! Best validation accuracy: {best_val_acc:.4f}")
 
-def cmd_resume(args):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-    
-    print("Loading processed data...")
-    X = np.load(args.X_path or "data/preprocessed/X.npy", allow_pickle=True)
-    y = np.load(args.y_path or "data/preprocessed/y.npy")
-    
-    X_train, X_val, y_train, y_val = train_test_split(
-        X, y, test_size=0.2, random_state=42, stratify=y
-    )
-    print(f"Train: {len(X_train)}, Val: {len(X_val)}")
-    
-    n_classes = len(np.unique(y))
-    model = CNN(n_classes=n_classes)
-    
-    print(f"Resuming from: {args.resume_from}")
-    
-    best_val_acc = train_cnn(
-        model,
-        X_train, y_train,
-        X_val, y_val,
+def cmd_train_cv(args) -> None:
+    X, y = _load_or_preprocess(args)
+    trainer = CNNTrainer(TrainConfig(
         epochs=args.epochs,
         batch_size=args.batch_size,
         lr=args.lr,
-        device=device,
-        use_all_patches=True,
         samples_per_epoch_fraction=args.sample_fraction,
         checkpoint_dir=args.checkpoint_dir,
         save_every_n_epoch=args.save_every,
-        resume_from=args.resume_from
-    )
-    
-    print(f"\nTraining complete! Best validation accuracy: {best_val_acc:.4f}")
-    return best_val_acc
+    ))
+    fold_accs, mean_acc = trainer.train_k_fold_cnn(CNN, X, y)
+    print(f"\nTraining complete! Mean validation accuracy: {mean_acc:.4f}")
 
-def cmd_predict(args):
+def cmd_predict(args) -> None:
     if not os.path.exists(args.audio_file):
-        print(f"Error: Audio file not found: {args.audio_file}")
-        sys.exit(1)
-    
+        print(f"Error: Audio file not found: {args.audio_file}"); sys.exit(1)
     if not os.path.exists(args.model):
-        print(f"Error: Model file not found: {args.model}")
-        sys.exit(1)
-    
-    try:
-        model = load_model(args.model, device=args.device)
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
-    
+        print(f"Error: Model file not found: {args.model}"); sys.exit(1)
     try:
-        predicted_class, top_probs, top_indices = predict_file(
-            model, args.audio_file, device=args.device, top_k=args.top_k
-        )
-        
+        predictor = AudioPredictor(model_path=args.model, device=args.device)
+        predicted_class, top_probs, top_indices = predictor.predict_file(args.audio_file, top_k=args.top_k)
+        labels = DatasetConfig().esc50_labels
         print("\n" + "=" * 60)
         print(f"Top {args.top_k} Predictions:")
         print("=" * 60)
-        
         for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
-            class_name = esc50_labels[idx]
             marker = "★" if idx == predicted_class else " "
-            print(f"{marker} {i+1}. {class_name:20s} - {prob*100:6.2f}%")
-        
+            print(f"{marker} {i+1}. {labels[idx]:20s} - {prob*100:6.2f}%")
     except Exception as e:
-        print(f"\nError during prediction: {e}")
         import traceback
+        print(f"\nError during prediction: {e}")
         traceback.print_exc()
         sys.exit(1)
 

src/config/config.py

@@ -1,9 +1,9 @@
 import os
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
-@dataclass(frozen=True)
+@dataclass
 class ProcessingConfig:
     audio_path: Path = Path("data/audio/0")
     augmented_path: Path = Path("data/audio/")
@@ -23,7 +23,7 @@ class ProcessingConfig:
     pitch_shift_rates = [-3.5, -2.5, -2, -1, 1, 2.5, 3, 3.5]
     drc_types = ["radio", "filmstandard", "musicstandard", "speech"]
 
-@dataclass(frozen=True)
+@dataclass
 class DatasetConfig:
     cnn_input_length: int = 128
     sample_rate: int = 44100
@@ -40,7 +40,7 @@ class DatasetConfig:
         'train', 'church_bells', 'airplane', 'fireworks', 'hand_saw'
     ])
 
-@dataclass(frozen=True)
+@dataclass
 class DownloadConfig:
     repo_url: str = "https://github.com/karolpiczak/ESC-50/archive/refs/heads/master.zip"
     repo_dst_dir: Path = Path("data")
@@ -55,48 +55,14 @@ class DownloadConfig:
     def __post_init__(self):
         object.__setattr__(self, "audio_dst_dir", self.repo_dst_dir / "audio" / "0")
 
-
-
-parameters = {
-    "n_bands"  : 128,
-    "n_mels" : 128,
-    "frame_size" : 1024,
-    "hop_size": 1024,
-    "sample_rate": 44100,
-    "fft_size": 8192,
-}
-
-cnn_input_length = 128
-
-sample_rate = 44100
-
-esc50_labels = [
-    'dog', 'rooster', 'pig', 'cow', 'frog',
-    'cat', 'hen', 'insects', 'sheep', 'crow',
-    'rain', 'sea_waves', 'crackling_fire', 'crickets', 'chirping_birds',
-    'water_drops', 'wind', 'pouring_water', 'toilet_flush', 'thunderstorm',
-    'crying_baby', 'sneezing', 'clapping', 'breathing', 'coughing',
-    'footsteps', 'laughing', 'brushing_teeth', 'snoring', 'drinking_sipping',
-    'door_wood_knock', 'mouse_click', 'keyboard_typing', 'door_wood_creaks', 'can_opening',
-    'washing_machine', 'vacuum_cleaner', 'clock_alarm', 'clock_tick', 'glass_breaking',
-    'helicopter', 'chainsaw', 'siren', 'car_horn', 'engine',
-    'train', 'church_bells', 'airplane', 'fireworks', 'hand_saw'
-]
-
-# download.py
-repo_url = "https://github.com/karolpiczak/ESC-50/archive/refs/heads/master.zip"
-repo_dst_dir = "data"
-audio_dst_dir = os.path.join(repo_dst_dir, "audio", "0")
-
-paths_to_delete = [
-    ".gitignore",
-    "esc50.gif",
-    "LICENSE",
-    "pytest.ini",
-    "README.md",
-    "requirements.txt",
-    "tests",
-    "meta",
-    ".github",
-    ".circleci"
-]
+@dataclass
+class TrainConfig:
+    epochs: int = 50
+    batch_size: int = 100
+    lr: int = 0.001
+    device = "cuda"
+    use_all_patches: bool = True
+    samples_per_epoch_fraction: float = 1/8
+    checkpoint_dir: str = "models/checkpoints"
+    save_every_n_epoch: int = 1
+    resume_from: Optional[str] = None

src/data/augment.py

@@ -3,173 +3,165 @@ import librosa
 import numpy as np
 import os
 import soundfile as sf
+from typing import Optional
 
-from src.config.config import sample_rate, parameters, cnn_input_length
-
-def data_treatment_training(
-    audio_path, 
-    n_bands, n_mels, frame_size, hop_size, sample_rate, fft_size
-    ):
-    labels = []
-    log_mel_spectrograms = []
-    filenames = os.listdir(audio_path)
-    
-    for filename in tqdm.tqdm(filenames, desc="Processing audio files"):
-        filename_splitted = filename.split("-")
-        label = filename_splitted[-1].split(".")[0]
-        label = label.split("_")[0]
-        labels.append(int(label))
-        
-        file_path = os.path.join(audio_path, filename)
-        audio, sr = librosa.load(file_path, sr=sample_rate)
-        
+from src.config.config import ProcessingConfig
+
+config = ProcessingConfig()
+
+class AudioAugment:
+    def __init__(self, config: ProcessingConfig = config) -> None: 
+        self.config = config
+
+    def _mel_spectrogram(self, audio: np.ndarray) -> np.ndarray:
         mel_spec = librosa.feature.melspectrogram(
             y=audio,
-            sr=sr,
-            n_fft=fft_size,
-            hop_length=hop_size,
-            win_length=frame_size,
-            n_mels=n_bands,
+            sr=self.config.sample_rate,
+            n_fft=self.config.fft_size,
+            hop_length=self.config.hop_size,
+            win_length=self.config.frame_size,
+            n_mels=self.config.n_bands,
             fmin=0,
-            fmax=sample_rate / 2,
+            fmax=self.config.sample_rate / 2,
             window='hann'
         )
-        
+            
         mel_spectrogram_db = 10 * np.log10(mel_spec.T + 1e-10)
         max_db = mel_spectrogram_db.max()
         mel_spectrogram_db = mel_spectrogram_db - max_db
+
+        return mel_spectrogram_db
+
+    def _data_treatment_training(self, audio_path: str) -> tuple[list[np.ndarray], np.ndarray]:
+        labels = []
+        log_mel_spectrograms = []
+        filenames = os.listdir(audio_path)
         
-        log_mel_spectrograms.append(mel_spectrogram_db)
-    
-    return log_mel_spectrograms, np.array(labels)
-
-def data_treatment_testing(
-    file_path, 
-    n_bands, n_mels, frame_size, hop_size, sample_rate, fft_size
-    ):
-    audio, sr = librosa.load(file_path, sr=sample_rate)
-    
-    mel_spec = librosa.feature.melspectrogram(
-        y=audio,
-        sr=sr,
-        n_fft=fft_size,
-        hop_length=hop_size,
-        win_length=frame_size,
-        n_mels=n_bands,
-        fmin=0,
-        fmax=sample_rate / 2,
-        window='hann'
-    )
-    
-    mel_spectrogram_db = 10 * np.log10(mel_spec.T + 1e-10)
-    max_db = mel_spectrogram_db.max()
-    mel_spectrogram_db = mel_spectrogram_db - max_db
-    
-    return [mel_spectrogram_db]
-
-def pad(audio, target_seconds, sample_rate):
-    target_len = int(sample_rate * target_seconds)
-    n = len(audio)
-
-    if n < target_len:
-        audio = np.pad(audio, (0, target_len - n), mode="constant")
-    return audio
-
-def time_stretch_augmentation(file_path, sample_rate, rate):
-    audio, _ = librosa.load(file_path, sr=sample_rate)
-    audio_timestretch = librosa.effects.time_stretch(audio.astype(np.float32), rate=rate)
-    return pad(audio_timestretch, 5, sample_rate)
-
-def pitch_shift_augmentation(file_path, sample_rate, semitones):
-    audio, _ = librosa.load(file_path, sr=sample_rate)
-    return librosa.effects.pitch_shift(audio.astype(np.float32), sr=sample_rate, n_steps=semitones)
-
-def drc_augmentation(file_path, sample_rate, compression):
-    if compression == "musicstandard":   threshold_db=-20; ratio=2.0; attack_ms=5;  release_ms=50
-    elif compression == "filmstandard":  threshold_db=-25; ratio=4.0; attack_ms=10; release_ms= 100
-    elif compression == "speech":         threshold_db=-18; ratio=3.0; attack_ms=2;  release_ms= 40
-    elif compression == "radio":          threshold_db=-15; ratio=3.5; attack_ms=1;  release_ms= 200
-
-    audio, _ = librosa.load(file_path, sr=sample_rate)
-    threshold = 10**(threshold_db / 20)
-
-    attack_coeff  = np.exp(-1.0 / (0.001 * attack_ms * sample_rate))
-    release_coeff = np.exp(-1.0 / (0.001 * release_ms * sample_rate))
-    
-    audio_filtered = np.zeros_like(audio)
-    gain = 1.0
-    
-    for n in range(len(audio)):
-        abs_audio = abs(audio[n])
-        if abs_audio > threshold:
-            desired_gain = (threshold / abs_audio) ** (ratio - 1)
-        else:
-            desired_gain = 1.0
+        for filename in tqdm.tqdm(filenames, desc="Processing audio files"):
+            filename_splitted = filename.split("-")
+            label = filename_splitted[-1].split(".")[0]
+            label = label.split("_")[0]
+            labels.append(int(label))
+            
+            file_path = os.path.join(audio_path, filename)
+            audio, sr = librosa.load(file_path, sr=self.config.sample_rate)
+
+            mel_spectrogram_db = self._mel_spectrogram(audio)
+            log_mel_spectrograms.append(mel_spectrogram_db)
         
-        if desired_gain < gain:
-            gain = attack_coeff * (gain - desired_gain) + desired_gain
-        else:
-            gain = release_coeff * (gain - desired_gain) + desired_gain
+        return log_mel_spectrograms, np.array(labels)
+
+    def _data_treatment_testing(self, file_path: str) -> list[np.ndarray]:
+        audio, sr = librosa.load(file_path, sr=self.config.sample_rate)
         
-        audio_filtered[n] = audio[n] * gain
+        mel_spectrogram_db = self._mel_spectrogram(audio)
 
-    return audio_filtered
+        return [mel_spectrogram_db]
 
-def augment_dataset(audio_path, output_path, probability_list):
-    filenames = os.listdir(audio_path)
+    def _pad(self, audio: np.ndarray) -> np.ndarray:
+        target_len = int(self.config.sample_rate * self.config.target_seconds)
+        n = len(audio)
 
-    p1, p2, p3 = probability_list
-    os.makedirs(output_path, exist_ok=True)
+        if n < target_len:
+            audio = np.pad(audio, (0, target_len - n), mode="constant")
+        return audio
+
+    def _time_stretch_augmentation(self, file_path: str, rate: float) -> np.ndarray:
+        audio, _ = librosa.load(file_path, sr=self.config.sample_rate)
+        audio_timestretch = librosa.effects.time_stretch(audio.astype(np.float32), rate=rate)
+        return self._pad(audio_timestretch)
+
+    def _pitch_shift_augmentation(self, file_path: str, semitones: float) -> np.ndarray:
+        audio, _ = librosa.load(file_path, sr=self.config.sample_rate)
+        return librosa.effects.pitch_shift(audio.astype(np.float32), sr=self.config.sample_rate, n_steps=semitones)
+
+    def _drc_augmentation(self, file_path: str, compression: float) -> np.ndarray:
+        if compression == "musicstandard":   threshold_db=-20; ratio=2.0; attack_ms=5;  release_ms=50
+        elif compression == "filmstandard":  threshold_db=-25; ratio=4.0; attack_ms=10; release_ms= 100
+        elif compression == "speech":         threshold_db=-18; ratio=3.0; attack_ms=2;  release_ms= 40
+        elif compression == "radio":          threshold_db=-15; ratio=3.5; attack_ms=1;  release_ms= 200
+
+        audio, _ = librosa.load(file_path, sr=self.config.sample_rate)
+        threshold = 10**(threshold_db / 20)
+
+        attack_coeff  = np.exp(-1.0 / (0.001 * attack_ms * self.config.sample_rate))
+        release_coeff = np.exp(-1.0 / (0.001 * release_ms * self.config.sample_rate))
+        
+        audio_filtered = np.zeros_like(audio)
+        gain = 1.0
+        
+        for n in range(len(audio)):
+            abs_audio = abs(audio[n])
+            if abs_audio > threshold:
+                desired_gain = (threshold / abs_audio) ** (ratio - 1)
+            else:
+                desired_gain = 1.0
+            
+            if desired_gain < gain:
+                gain = attack_coeff * (gain - desired_gain) + desired_gain
+            else:
+                gain = release_coeff * (gain - desired_gain) + desired_gain
+            
+            audio_filtered[n] = audio[n] * gain
+
+        return audio_filtered
+
+    def _augment_dataset(self, audio_path: str, output_path: str, probability_list: list[float]) -> None:
+        filenames = os.listdir(audio_path)
+
+        p1, p2, p3 = probability_list
+        os.makedirs(output_path, exist_ok=True)
+
+        for filename in tqdm.tqdm(filenames, desc="Augmenting audio files"):        
+            
+            audio, _ = librosa.load(os.path.join(audio_path, filename), sr=self.config.sample_rate)
+            # TS
+            if np.random.rand() > p1:
+                stretch_rates = [0.81, 0.93, 1.07, 1.23]
+                stretch_rate = np.random.choice(stretch_rates)
+                audio = self._time_stretch_augmentation(os.path.join(audio_path, filename), stretch_rate)
+            # PS 
+            if np.random.rand() > p2:
+                semitones = [-3.5, -2.5, -2, -1, 1, 2.5, 3, 3.5]
+                semitone = np.random.choice(semitones)
+                audio = self._pitch_shift_augmentation(os.path.join(audio_path, filename), semitone)
+            # DRC
+            if np.random.rand() > p3:
+                compressions = ["radio", "filmstandard", "musicstandard", "speech"]
+                compression = np.random.choice(compressions)
+                audio = self._drc_augmentation(os.path.join(audio_path, filename), compression)
+
+            sf.write(os.path.join(output_path, filename), audio, self.config.sample_rate)
+
+    def _create_augmented_datasets(self, input_path: str, output_path: str) -> None:
+        probability_lists = self.config.augmentation_probability_lists
+        for i, probability_list in enumerate(probability_lists):
+            augmented_path = os.path.join(output_path, f"{i+1}")
+            os.makedirs(augmented_path, exist_ok=True)
+            self._augment_dataset(input_path, augmented_path, probability_list)
+
+    def _create_log_mel(self, input_path: str, output_path: str) -> tuple[list[np.ndarray], np.ndarray]:
+        directories = os.listdir(input_path)
+        X, y = [], []
+
+        for directory in directories:
+            log_mels, labels = self._data_treatment_training(os.path.join(input_path, directory))
+            X.extend(log_mels)
+            y.extend(labels)
+        
+        X_array = np.empty(len(X), dtype=object)
+        for i, spec in enumerate(X):
+            X_array[i] = spec
 
-    for filename in tqdm.tqdm(filenames, desc="Processing audio files"):        
+        y = np.array(y)
+        os.makedirs(output_path, exist_ok=True)
         
-        audio, _ = librosa.load(os.path.join(audio_path, filename), sr=sample_rate)
-        # TS
-        if np.random.rand() > p1:
-            stretch_rates = [0.81, 0.93, 1.07, 1.23]
-            stretch_rate = np.random.choice(stretch_rates)
-            audio = time_stretch_augmentation(os.path.join(audio_path, filename), sample_rate, stretch_rate)
-        # PS 
-        if np.random.rand() > p2:
-            semitones = [-3.5, -2.5, -2, -1, 1, 2.5, 3, 3.5]
-            semitone = np.random.choice(semitones)
-            audio = pitch_shift_augmentation(os.path.join(audio_path, filename), sample_rate, semitone)
-        # DRC
-        if np.random.rand() > p3:
-            compressions = ["radio", "filmstandard", "musicstandard", "speech"]
-            compression = np.random.choice(compressions)
-            audio = drc_augmentation(os.path.join(audio_path, filename), sample_rate, compression)
-
-        sf.write(os.path.join(output_path, filename), audio, 44100)
-
-def create_augmented_datasets(input_path, output_path):
-    probability_lists = [
-        [0.0 , 1.0, 1.0],
-        [1.0 , 1.0, 0.0],
-        [1.0 , 0.0, 1.0],
-        [0.0 , 0.0, 0.0],
-        [0.5 , 0.5, 0.5]]
-    for i, probability_list in enumerate(probability_lists):
-        augmented_path = os.path.join(output_path, f"{i+1}")
-        os.makedirs(augmented_path, exist_ok=True)
-        augment_dataset(input_path, augmented_path, probability_list)
-
-def create_log_mel(input_path, output_path):
-    directories = os.listdir(input_path)
-    X, y = [], []
-
-    for directory in directories:
-        log_mels, labels = data_treatment_training(os.path.join(input_path, directory), **parameters)
-        X.extend(log_mels)
-        y.extend(labels)
-    
-    X_array = np.empty(len(X), dtype=object)
-    for i, spec in enumerate(X):
-        X_array[i] = spec
-
-    y = np.array(y)
-    os.makedirs(output_path, exist_ok=True)
-    
-    np.save(os.path.join(output_path, "X.npy"), X_array, allow_pickle=True)
-    np.save(os.path.join(output_path, 'y.npy'), y)
-    return X, y
+        np.save(os.path.join(output_path, "X.npy"), X_array, allow_pickle=True)
+        np.save(os.path.join(output_path, 'y.npy'), y)
+        return X, y
+
+    def run(self, augment: bool = True, preprocess : bool = True) -> None:
+        if augment:
+            self._create_augmented_datasets(self.config.audio_path, self.config.augmented_path)
+        if preprocess:
+            self._create_log_mel(self.config.augmented_path, self.config.log_mel_path)

src/data/dataset.py

@@ -0,0 +1,68 @@
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+from typing import Sequence 
+
+from src.config.config import DatasetConfig
+
+config = DatasetConfig()
+            
+class FullTFPatchesDataset(Dataset):
+    def __init__(self, spectrograms: Sequence[np.ndarray], labels: Sequence[int], config: DatasetConfig = config) -> None:
+        self.config = config
+        self.patch_indices = []
+
+        for spec_idx, spec in enumerate(spectrograms):
+            n_frames = spec.shape[0]
+            label = labels[spec_idx]
+
+            if n_frames >= self.config.cnn_input_length:
+                for start_frame in range(n_frames - self.config.cnn_input_length + 1):
+                    self.patch_indices.append((spec_idx, start_frame, label))
+            else:
+                self.patch_indices.append((spec_idx, 0, label))
+            
+        self.spectrograms = spectrograms
+    
+    def __len__(self) -> int:
+        return len(self.patch_indices)
+    
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        spec_idx, start_frame, label = self.patch_indices[idx]
+        spec = self.spectrograms[spec_idx]
+        
+        n_frames = spec.shape[0]
+        
+        if n_frames >= self.config.cnn_input_length:
+            patch = spec[start_frame:start_frame + self.config.cnn_input_length]
+        else:
+            pad = self.config.cnn_input_length - n_frames
+            patch = np.pad(spec, ((0, pad), (0, 0)), mode='constant')
+        
+        patch = patch[np.newaxis, :, :]
+        
+        return torch.tensor(patch, dtype=torch.float32), torch.tensor(label, dtype=torch.long)
+
+class RandomPatchDataset(Dataset):
+    def __init__(self, spectrograms: Sequence[np.ndarray], labels: Sequence[int], config: DatasetConfig = config) -> None:
+        self.config = config
+        self.spectrograms = spectrograms
+        self.labels = labels
+    
+    def __len__(self) -> int:
+        return len(self.labels)
+    
+    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
+        spec = self.spectrograms[idx]
+        label = self.labels[idx]
+        n_frames = spec.shape[0]
+        
+        if n_frames >= self.config.cnn_input_length:
+            start = np.random.randint(0, n_frames - self.config.cnn_input_length + 1)
+            patch = spec[start:start + self.config.cnn_input_length]
+        else:
+            pad = self.config.cnn_input_length - n_frames
+            patch = np.pad(spec, ((0, pad), (0, 0)), mode='constant')
+        
+        patch = patch[np.newaxis, :, :]
+        return torch.tensor(patch, dtype=torch.float32), torch.tensor(label, dtype=torch.long)

src/data/download.py

@@ -1,6 +1,5 @@
 import requests
 import zipfile
-import tarfile
 import io
 import os
 import shutil
@@ -13,11 +12,7 @@ from src.config.config import DownloadConfig
 config = DownloadConfig()
 
 class ESC50Downloader:
-    def __init__(
-        self,
-        repo_url: str = config.repo_url,
-        repo_dst_dir: str = config.repo_dst_dir
-    ):
+    def __init__(self, repo_url: str = config.repo_url, repo_dst_dir: str = config.repo_dst_dir) -> None:
         self.repo_url = repo_url
         self.repo_dst_dir = Path(repo_dst_dir)
         self.audio_dst_dir = config.audio_dst_dir
@@ -25,7 +20,7 @@ class ESC50Downloader:
         self.extracted_dir = config.extracted_dir 
         self.audio_src_dir = config.audio_src_dir 
 
-    def download_and_extract(self):
+    def download_and_extract(self) -> None:
         os.makedirs(self.repo_dst_dir, exist_ok=True)
         print(f"Downloading from {self.repo_url}")
 
@@ -46,7 +41,7 @@ class ESC50Downloader:
             z.extractall(self.repo_dst_dir)
         print("Done extracting.")
 
-    def clean_files(self):
+    def clean_files(self) -> None:
         for f in self.paths_to_delete:
             path = os.path.join(self.extracted_dir, f)
             if os.path.isfile(path):
@@ -56,7 +51,7 @@ class ESC50Downloader:
                 shutil.rmtree(path)
                 print(f"Deleted directory: {path}")
 
-    def move_audio_files(self):
+    def move_audio_files(self) -> None:
         os.makedirs(self.audio_dst_dir, exist_ok=True)
         print(f"Moving audio files from {self.audio_src_dir} to {self.audio_dst_dir}")
         
@@ -67,11 +62,7 @@ class ESC50Downloader:
                 shutil.move(src_file, dst_file)
         print(f"Moved all audio files to {self.audio_dst_dir}")
 
-    def download_clean(self):
+    def download_clean(self) -> None:
         self.download_and_extract()
         self.clean_files()
         self.move_audio_files()
-
-if __name__ == "__main__":
-    downloader = ESC50Downloader()
-    downloader.download_clean()

src/models/cnn.py

@@ -1,19 +1,18 @@
 import torch.nn as nn
+import torch 
 
 class CNN(nn.Module):
-    def __init__(self, n_classes=50):
+    def __init__(self, n_classes: int = 50) -> None:
         super().__init__()
         self.features = nn.Sequential(
             nn.Conv2d(1, 24, kernel_size=(5, 5)),
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=(4, 2), stride=(4, 2)),
             
-
             nn.Conv2d(24, 48, kernel_size=(5, 5)),
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=(4, 2), stride=(4, 2)),
             
-
             nn.Conv2d(48, 48, kernel_size=(5, 5)),
             nn.ReLU(),
         )
@@ -25,8 +24,7 @@ class CNN(nn.Module):
             nn.Linear(64, n_classes)
         )
 
-
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.features(x)
         x = x.flatten(1)
         return self.classifier(x)

src/models/predict.py

@@ -1,193 +1,74 @@
 import numpy as np
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 import argparse
-import os
-import sys
 
 from src.models.cnn import CNN
-from src.data.augment import data_treatment_testing
-from src.config.config import sample_rate, parameters, cnn_input_length, esc50_labels
-
-def predict_with_overlapping_patches(model, spectrogram, patch_length=cnn_input_length, hop=1, batch_size=100, device="cuda"):
-    model.eval()
-
-    n_frames, n_mels = spectrogram.shape
-
-    if n_frames < patch_length:
-        pad = patch_length - n_frames
-        spectrogram = np.pad(spectrogram, ((0, pad), (0, 0)), mode='constant')
-        n_frames = patch_length
-
-    patches = []
-    for start in range(0, n_frames - patch_length + 1, hop):
-        patch = spectrogram[start:start + patch_length]
-        patch = patch[np.newaxis, np.newaxis, :, :]
-        patches.append(patch)
-
-    patches = np.concatenate(patches, axis=0)
-    patches = torch.tensor(patches, dtype=torch.float32).to(device)
-
-    all_outputs = []
-    with torch.no_grad():
-        for i in range(0, len(patches), batch_size):
-            batch = patches[i:i + batch_size]
-            outputs = model(batch)
-            all_outputs.append(outputs)
+from src.data.augment import AudioAugment
+from src.config.config import ProcessingConfig, DatasetConfig, TrainConfig
+
+config = ProcessingConfig()
+
+class AudioPredictor:
+    def __init__(
+        self, 
+        model_path: str, 
+        config: ProcessingConfig = config, 
+        device: str = 'cuda'
+    ) -> None:
+        self.config = config
+        self.audio_dataset = AudioAugment()
+        self.dataset_config = DatasetConfig()
+        self.train_config = TrainConfig()
+        self.device = device
+        self.model = self._load_model(model_path)
+
+    def _load_model(self, model_path: str) -> CNN:
+        model = CNN(n_classes=len(self.dataset_config.esc50_labels))
+        checkpoint = torch.load(model_path, map_location=self.device)
+        state_dict = checkpoint.get("model_state_dict", checkpoint) if isinstance(checkpoint, dict) else checkpoint
+        model.load_state_dict(state_dict)
+        if isinstance(checkpoint, dict) and "best_val_acc" in checkpoint:
+            print(f"Model validation accuracy: {checkpoint['best_val_acc']:.4f}")
+        model.to(self.device).eval()
+        print("Model loaded successfully!\n")
+        return model
+
+    def _extract_patches(self, spectrogram: np.ndarray, hop: int) -> torch.Tensor:
+        n_frames, _ = spectrogram.shape
+        if n_frames < self.dataset_config.cnn_input_length:
+            spectrogram = np.pad(spectrogram, ((0, self.dataset_config.cnn_input_length - n_frames), (0, 0)), mode="constant")
+            n_frames = self.dataset_config.cnn_input_length
+
+        patches = np.concatenate([
+            spectrogram[s:s + self.dataset_config.cnn_input_length][np.newaxis, np.newaxis]
+            for s in range(0, n_frames - self.dataset_config.cnn_input_length + 1, hop)
+        ], axis=0)
+        return torch.tensor(patches, dtype=torch.float32).to(self.device)
     
-    all_outputs = torch.cat(all_outputs, dim=0)
-
-    mean_activations = all_outputs.mean(dim=0)
-    predicted_class = mean_activations.argmax().item()
-
-    return predicted_class
-
-def predict_top_k(model, spectrogram, patch_length=cnn_input_length, hop=1, batch_size=100, device="cpu", top_k=5):  
-    model.eval()
-
-    n_frames, n_mels = spectrogram.shape
-
-    if n_frames < patch_length:
-        pad = patch_length - n_frames
-        spectrogram = np.pad(spectrogram, ((0, pad), (0, 0)), mode='constant')
-        n_frames = patch_length
-
-    patches = []
-    for start in range(0, n_frames - patch_length + 1, hop):
-        patch = spectrogram[start:start + patch_length]
-        patch = patch[np.newaxis, np.newaxis, :, :]
-        patches.append(patch)
-
-    patches = np.concatenate(patches, axis=0)
-    patches = torch.tensor(patches, dtype=torch.float32).to(device)
-
-    all_outputs = []
-    with torch.no_grad():
-        for i in range(0, len(patches), batch_size):
-            batch = patches[i:i + batch_size]
-            outputs = model(batch)
-            all_outputs.append(outputs)
-    
-    all_outputs = torch.cat(all_outputs, dim=0)
-    mean_logits = all_outputs.mean(dim=0)
-    probabilities = torch.nn.functional.softmax(mean_logits, dim=0)
-
-    top_probs, top_indices = torch.topk(probabilities, min(top_k, 50))
-    top_probs = top_probs.cpu().numpy()
-    top_indices = top_indices.cpu().numpy()
-
-    return top_probs, top_indices
+    def _run_inference(self, patches: torch.Tensor, batch_size: int) -> torch.Tensor:
+        all_outputs = []
+        with torch.no_grad():
+            for i in range(0, len(patches), batch_size):
+                all_outputs.append(self.model(patches[i:i + batch_size]))
+        return torch.cat(all_outputs, dim=0).mean(dim=0)
+
+    def predict_class(self, spectrogram: np.ndarray, hop: int = 1) -> int:
+        patches = self._extract_patches(spectrogram, hop)
+        mean_activations = self._run_inference(patches, self.train_config.batch_size)
+        return mean_activations.argmax().item()
+
+    def predict_top_k(self, spectrogram: np.ndarray, hop: int = 1, top_k: int = 5):
+        patches = self._extract_patches(spectrogram, hop)
+        mean_logits = self._run_inference(patches, self.train_config.batch_size)
+        probs = F.softmax(mean_logits, dim=0)
+        top_probs, top_indices = torch.topk(probs, min(top_k, len(self.dataset_config.esc50_labels)))       
+        return top_probs.cpu().numpy(), top_indices.cpu().numpy()
+
+    def predict_file(self, audio_file: str, top_k: int = 5):
+        spectrogram = np.array(self.audio_dataset._data_treatment_testing(audio_file)).squeeze()
+        predicted_class = self.predict_class(spectrogram)
+        top_probs, top_indices = self.predict_top_k(spectrogram, top_k=top_k)
+        return predicted_class, top_probs, top_indices
 
-def predict_file(model, audio_file, device="cpu", top_k=5):
-    parameters = {
-        "n_bands"  : 128,
-        "n_mels" : 128,
-        "frame_size" : 1024,
-        "hop_size": 1024,
-        "sample_rate": sample_rate,
-        "fft_size": 8192,
-    }
-    spectrogram = data_treatment_testing(audio_file, **parameters)
-
-    spectrogram = np.array(spectrogram)
-
-    spectrogram = spectrogram.squeeze()
-
-    predicted_class = predict_with_overlapping_patches(
-        model, spectrogram, patch_length=128, hop=1, batch_size=100, device=device
-    )
-    top_probs, top_indices = predict_top_k(
-        model, spectrogram, patch_length=128, hop=1, batch_size=100, device=device, top_k=top_k
-    )
-
-    return predicted_class, top_probs, top_indices
-
-def load_model(model_path, device='cpu'):
-    print(f"Loading model from {model_path}...")
-    
-    model = CNN(n_classes=50)
-    checkpoint = torch.load(model_path, map_location=device)
-    
-    if isinstance(checkpoint, dict):
-        if 'model_state_dict' in checkpoint:
-            model.load_state_dict(checkpoint['model_state_dict'])
-            if 'best_val_acc' in checkpoint:
-                print(f"Model validation accuracy: {checkpoint['best_val_acc']:.4f}")
-        else:
-            model.load_state_dict(checkpoint)
-    else:
-        model.load_state_dict(checkpoint)
-    
-    model.to(device)
-    model.eval()
-    print("Model loaded successfully!\n")
-    return model
-
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description='Predict environmental sound class using trained ESC-50 model'
-    )
-    parser.add_argument(
-        'audio_file',
-        type=str,
-        help='Path to .wav file to classify'
-    )
-    parser.add_argument(
-        '--model',
-        type=str,
-        default='best_model.pt',
-        help='Path to trained model checkpoint (default: best_model.pt)'
-    )
-    parser.add_argument(
-        '--top-k',
-        type=int,
-        default=5,
-        help='Number of top predictions to show (default: 5)'
-    )
-    parser.add_argument(
-        '--device',
-        type=str,
-        default='cuda' if torch.cuda.is_available() else 'cpu',
-        help='Device to use (default: auto-detect)'
-    )
-    
-    args = parser.parse_args()
-    
-    if not os.path.exists(args.audio_file):
-        print(f"Error: Audio file not found: {args.audio_file}")
-        sys.exit(1)
-    
-    if not os.path.exists(args.model):
-        print(f"Error: Model file not found: {args.model}")
-        sys.exit(1)
-    
-    try:
-        model = load_model(args.model, device=args.device)
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
-    
-    try:
-        predicted_class, top_probs, top_indices = predict_file(
-            model, args.audio_file, device=args.device, top_k=args.top_k
-        )
-                
-        print("\n" + "=" * 60)
-        print(f"Top {args.top_k} Predictions:")
-        print("=" * 60)
-        
-        for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
-            class_name = esc50_labels[idx]
-            marker = "★" if idx == predicted_class else " "
-            print(f"{marker} {i+1}. {class_name:20s} - {prob*100:6.2f}%")
-        
-      
-    except Exception as e:
-        print(f"\nError during prediction: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)

src/models/traincnn.py

@@ -4,293 +4,244 @@ import tqdm
 import json
 import numpy as np
 from torch.utils.data import DataLoader
+from typing import Sequence
 
-from src.models.predict import predict_with_overlapping_patches
+from src.models.predict import AudioPredictor
 from src.data.dataset import FullTFPatchesDataset, RandomPatchDataset
+from src.config.config import TrainConfig
+
+config = TrainConfig()
+
+class CNNTrainer:
+    def __init__(self, config: TrainConfig = config) -> None:
+        self.config = config
+
+    def train_cnn(
+        self,
+        model: torch.nn.Module,
+        X_train: Sequence[np.ndarray], 
+        y_train: Sequence[int],
+        X_val: Sequence[np.ndarray], 
+        y_val: Sequence[int],
+        fold_num: int,
+    ) -> float:
+        device = self.config.device
+        os.makedirs(self.config.checkpoint_dir, exist_ok=True)
+
+        model.to(device)
+        
+        if self.config.use_all_patches:
+            train_dataset = FullTFPatchesDataset(X_train, y_train)
+            print(f"\n{'='*60}\nUsing ALL PATCHES method\n{'='*60}")
+        else:
+            train_dataset = RandomPatchDataset(X_train, y_train)
+            print(f"\n{'='*60}\nUsing ALL PATCHES method\n{'='*60}")
+
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=self.config.batch_size,
+            shuffle=True,
+            num_workers=4,
+            pin_memory=True
+        )
+        
+        total_patches = len(train_dataset)
+        patches_per_epoch = int(total_patches * self.config.samples_per_epoch_fraction)
+        batches_per_epoch = patches_per_epoch // self.config.batch_size
 
-def train_cnn(
-    model,
-    X_train, y_train,
-    X_val, y_val,
-    fold_num,
-    epochs=50,
-    batch_size=100,
-    lr=0.001,
-    device="cuda",
-    use_all_patches=True,
-    samples_per_epoch_fraction=1/8,
-    checkpoint_dir="models/checkpoints",
-    save_every_n_epoch=1,
-    resume_from=None
-    ):
-    os.makedirs(checkpoint_dir, exist_ok=True)
-
-    model.to(device)
-    
-    if use_all_patches:
-        train_dataset = FullTFPatchesDataset(X_train, y_train, patch_length=128)
-        print(f"\n{'='*60}")
-        print("Using ALL PATCHES method (as per paper)")
-        print(f"{'='*60}")
-    else:
-        train_dataset = RandomPatchDataset(X_train, y_train, patch_length=128)
-        print(f"\n{'='*60}")
-        print("Using RANDOM PATCHES method (simpler)")
-        print(f"{'='*60}")
-    
-    # unique, counts = np.unique(y_train, return_counts=True)
-    # print(f"\nClass distribution in y_train:")
-    # print(f"Classes: {len(unique)}")
-    # print(f"Min samples: {counts.min()}, Max samples: {counts.max()}, Mean: {counts.mean():.1f}")
-    # print(f"\nPer-class counts:")
-    # for cls, count in zip(unique, counts):
-    #     print(f"Class {cls}: {count}")
-
-    train_loader = DataLoader(
-        train_dataset,
-        batch_size=batch_size,
-        shuffle=True,
-        num_workers=4,
-        pin_memory=True
-    )
-    
-    total_patches = len(train_dataset)
-    patches_per_epoch = int(total_patches * samples_per_epoch_fraction)
-    batches_per_epoch = patches_per_epoch // batch_size
-    
-    print(f"Total available patches: {total_patches:,}")
-    print(f"Patches per epoch ({samples_per_epoch_fraction}): {patches_per_epoch:,}")
-    print(f"Batches per epoch: {batches_per_epoch:,}")
-    print(f"{'='*60}\n")
-    
-    criterion = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.AdamW([
-        {'params': model.features.parameters(), 'weight_decay': 0.0},
-        {'params': model.classifier.parameters(), 'weight_decay': 0.001}
-    ], lr=lr)#, momentum=0.9)
-    
-
-    start_epoch = 0
-    best_val_acc = 0.0
-    training_history = {
-        'train_loss': [],
-        'train_acc': [],
-        'val_acc': [],
-        'epochs': []
-    }
-
-    if resume_from and os.path.exists(resume_from):
-        print(f"Resuming from checkpoint: {resume_from}")
-        checkpoint = torch.load(resume_from, map_location=device)
-
-        model.load_state_dict(checkpoint['model_state_dict'])
-        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-        start_epoch = checkpoint['epoch'] + 1
-        best_val_acc = checkpoint['best_val_acc']
-        training_history = checkpoint['history']
-
-        print(f"Resuming training from epoch: {checkpoint['epoch']}")
-        print(f"Best val acc: {best_val_acc:.4f}\n")
-
-
-
-    for epoch in range(start_epoch, epochs):
-        model.train()
-        train_loss = 0.0
-        correct = 0
-        total = 0
-        batches_processed = 0
+        print(f"Total available patches: {total_patches:,}")
+        print(f"Patches per epoch ({self.config.samples_per_epoch_fraction}): {patches_per_epoch:,}")
+        print(f"Batches per epoch: {batches_per_epoch:,}\n{'='*60}\n")
         
-        for xb, yb in tqdm.tqdm(train_loader, f"Epoch {epoch+1} Train", leave=False):
-            if batches_processed >= batches_per_epoch:
-                break
+        criterion = torch.nn.CrossEntropyLoss()
+        optimizer = torch.optim.AdamW([
+            {'params': model.features.parameters(), 'weight_decay': 0.0},
+            {'params': model.classifier.parameters(), 'weight_decay': 0.001}
+        ], lr=self.config.lr)
+        
+        start_epoch = 0
+        best_val_acc = 0.0
+        training_history: dict = {'train_loss': [], 'train_acc': [], 'val_acc': [], 'epochs': []}
+
+        if self.config.resume_from and os.path.exists(self.config.resume_from):
+            print(f"Resuming from checkpoint: {self.config.resume_from}")
+            checkpoint = torch.load(self.config.resume_from, map_location=device)
+            model.load_state_dict(checkpoint['model_state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+            start_epoch = checkpoint['epoch'] + 1
+            best_val_acc = checkpoint['best_val_acc']
+            training_history = checkpoint['history']
+            print(f"Resuming from epoch {checkpoint['epoch']}, best val acc: {best_val_acc:.4f}\n")
+
+        for epoch in range(start_epoch, self.config.epochs):
+            model.train()
+            train_loss, correct, total, batches_processed = 0.0, 0, 0, 0
             
-            xb = xb.to(device)
-            yb = yb.to(device)
+            for xb, yb in tqdm.tqdm(train_loader, f"Epoch {epoch+1} Train", leave=False):
+                if batches_processed >= batches_per_epoch:
+                    break
+                xb, yb = xb.to(device), yb.to(device)                
+                optimizer.zero_grad()
+                out = model(xb)
+                loss = criterion(out, yb)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+                optimizer.step()
+
+                train_loss += loss.item() * xb.size(0)
+                _, pred = out.max(1)
+                correct += (pred == yb).sum().item()
+                total += yb.size(0)
+                batches_processed += 1
             
-            optimizer.zero_grad()
-            out = model(xb)
-                       
-            loss = criterion(out, yb)
+            train_loss /= total
+            train_acc = correct / total
             
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
-     
-            optimizer.step()
-
-            train_loss += loss.item() * xb.size(0)
-            _, pred = out.max(1)
-            correct += (pred == yb).sum().item()
-            total += yb.size(0)
-            batches_processed += 1
-        
-        train_loss /= total
-        train_acc = correct / total
-        
-        model.eval()
-        val_correct = 0
-        val_total = len(y_val)
-        
-
-        for i in tqdm.tqdm(range(val_total), desc=f"Epoch {epoch+1} Val", leave=False):
-            spec = X_val[i]
-            true_label = y_val[i]
+            model.eval()
+            val_correct = 0
+            val_total = len(y_val)
             
-            pred_label = predict_with_overlapping_patches(model, spec, device=device)
+            for i in tqdm.tqdm(range(val_total), desc=f"Epoch {epoch+1} Val", leave=False):
+                spec = X_val[i]
+                true_label = y_val[i]
+                
+                pred_label = self._predict_val(model, spec, device)                
+                
+                if pred_label == true_label:
+                    val_correct += 1
             
-            if pred_label == true_label:
-                val_correct += 1
-        
-        val_acc = val_correct / val_total
+            val_acc = val_correct / val_total
 
-        training_history['train_loss'].append(train_loss)
-        training_history['train_acc'].append(train_acc)
-        training_history['val_acc'].append(val_acc)
-        training_history['epochs'].append(epoch + 1)
+            training_history['train_loss'].append(train_loss)
+            training_history['train_acc'].append(train_acc)
+            training_history['val_acc'].append(val_acc)
+            training_history['epochs'].append(epoch + 1)
 
-        is_best = val_acc > best_val_acc
+            is_best = val_acc > best_val_acc
 
-        if is_best:
-            best_val_acc = val_acc
-            torch.save(model.state_dict(), "best_model.pt")
-        
-        print(
-            f"Fold {fold_num} | Epoch {epoch+1}/{epochs} | "
-            f"Train loss: {train_loss:.4f}, Train acc: {train_acc:.4f} | "
-            f"Val acc: {val_acc:.4f} (best: {best_val_acc:.4f})"
-        )
+            if is_best:
+                best_val_acc = val_acc
+                torch.save(model.state_dict(), "best_model.pt")
+            
+            print(
+                f"Fold {fold_num} | Epoch {epoch+1}/{self.config.epochs} | "
+                f"Train loss: {train_loss:.4f}, Train acc: {train_acc:.4f} | "
+                f"Val acc: {val_acc:.4f} (best: {best_val_acc:.4f})"
+            )
 
-        if (epoch + 1) % save_every_n_epoch == 0:
-            checkpoint = {
-                'epoch': epoch,
-                'model_state_dict': model.state_dict(),
-                'optimizer_state_dict': optimizer.state_dict(),
-                'train_loss': train_loss,
-                'train_acc': train_acc,
-                'val_acc': val_acc,
-                'best_val_acc': best_val_acc,
-                'history': training_history,
-                'config': {
-                    'batch_size': batch_size,
-                    'lr': lr,
-                    'total_patches': total_patches,
-                    'patches_per_epoch': patches_per_epoch,
+            if (epoch + 1) % self.config.save_every_n_epoch == 0:
+                checkpoint = {
+                    'epoch': epoch,
+                    'model_state_dict': model.state_dict(),
+                    'optimizer_state_dict': optimizer.state_dict(),
+                    'train_loss': train_loss,
+                    'train_acc': train_acc,
+                    'val_acc': val_acc,
+                    'best_val_acc': best_val_acc,
+                    'history': training_history,
+                    'config': {
+                        'batch_size': self.config.batch_size,
+                        'lr': self.config.lr,
+                        'total_patches': total_patches,
+                        'patches_per_epoch': patches_per_epoch,
+                    }
                 }
+                checkpoint_path = os.path.join(
+                    self.config.checkpoint_dir,
+                    f"checkpoint_epoch_{epoch+1}.pt"
+                )
+                torch.save(checkpoint, checkpoint_path)
+                
+                if is_best:
+                    best_path = os.path.join(self.config.checkpoint_dir, "best_model.pt")
+                    torch.save(checkpoint, best_path)
+
+                latest_path = os.path.join(self.config.checkpoint_dir, "latest_checkpoint.pt")
+                torch.save(checkpoint, latest_path)
+                
+                history_path = os.path.join(self.config.checkpoint_dir, "training_history.json")
+                with open(history_path, 'w') as f:
+                    json.dump(training_history, f, indent=2)
+
+        final_model_dir = "models/saved"
+        os.makedirs(final_model_dir, exist_ok=True)
+        final_model_path = os.path.join(final_model_dir, "final_model.pt")
+        torch.save({
+            'model_state_dict': model.state_dict(),
+            'best_val_acc': best_val_acc,
+            'config': {
+                'batch_size': self.config.batch_size,
+                'lr': self.config.lr,
+                'epochs': self.config.epochs,
             }
-            checkpoint_path = os.path.join(
-                checkpoint_dir,
-                f"checkpoint_epoch_{epoch+1}.pt"
+        }, final_model_path)
+        print(f"\nTraining complete! Final model saved to {final_model_path}")
+
+        return best_val_acc
+
+    def train_k_fold_cnn(
+        self,
+        model_class: type,
+        X: Sequence[np.ndarray],
+        y: Sequence[int],
+    ) -> tuple[list[float], float]:
+
+        X_arr = np.array(X)
+        y_arr = np.array(y)
+        n_samples = len(y_arr)
+        indices = np.arange(n_samples)
+        np.random.shuffle(indices)
+
+        fold_sizes = (n_samples // 5) * np.ones(5, dtype=int)
+        fold_sizes[:n_samples % 5] += 1
+        current = 0
+        fold_accuracies: list[float] = []
+
+        for fold_num, fold_size in enumerate(fold_sizes, 1):
+            start, stop = current, current + fold_size
+            val_idx = indices[start:stop]
+            train_idx = np.concatenate([indices[:start], indices[stop:]])
+            current = stop
+
+            X_train, y_train = X_arr[train_idx].tolist(), y_arr[train_idx]
+            X_val, y_val = X_arr[val_idx].tolist(), y_arr[val_idx]
+
+            print(f"\n{'='*80}\nFOLD {fold_num}/5 | Train: {len(X_train)}, Val: {len(X_val)}\n{'='*80}\n")
+
+            model = model_class()
+            best_acc = self.train_cnn(
+                model=model,
+                X_train=X_train, y_train=y_train,
+                X_val=X_val, y_val=y_val,
+                fold_num=fold_num,
             )
-            torch.save(checkpoint, checkpoint_path)
-            
-            if is_best:
-                best_path = os.path.join(checkpoint_dir, "best_model.pt")
-                torch.save(checkpoint, best_path)
-                #print("Saved best model")
-
-            latest_path = os.path.join(checkpoint_dir, "latest_checkpoint.pt")
-            torch.save(checkpoint, latest_path)
-            
-            history_path = os.path.join(checkpoint_dir, "training_history.json")
-            with open(history_path, 'w') as f:
-                json.dump(training_history, f, indent=2)
-
-    final_model_dir = "models/saved"
-    os.makedirs(final_model_dir, exist_ok=True)
-    final_model_path = os.path.join(final_model_dir, "final_model.pt")
-    torch.save({
-        'model_state_dict': model.state_dict(),
-        'best_val_acc': best_val_acc,
-        'config': {
-            'batch_size': batch_size,
-            'lr': lr,
-            'epochs': epochs,
-        }
-    }, final_model_path)
-    print(f"\nTraining complete! Final model saved to {final_model_path}")
-
-    return best_val_acc
-
-
-def train_k_fold_cnn(
-    model_class,
-    X, y,
-    epochs=50,
-    batch_size=100,
-    lr=0.01,
-    k_fold=5,
-    device="cuda",
-    use_all_patches=True,
-    samples_per_epoch_fraction=1/8,
-    checkpoint_dir="models/checkpoints",
-    save_every_n_epoch=1
-    ):
-
-    X = np.array(X)
-    y = np.array(y)
-    n_samples = len(y)
-    indices = np.arange(n_samples)
-    np.random.shuffle(indices)
-
-    fold_sizes = (n_samples // 5) * np.ones(5, dtype=int)
-    fold_sizes[:n_samples % 5] += 1
-    current = 0
-
-    fold_accuracies = []
-
-    for fold_num, fold_size in enumerate(fold_sizes, 1):
-        start, stop = current, current + fold_size
-        val_idx = indices[start:stop]
-        train_idx = np.concatenate([indices[:start], indices[stop:]])
-        current = stop
-
-        X_train, y_train = X[train_idx].tolist(), y[train_idx]
-        X_val, y_val = X[val_idx].tolist(), y[val_idx]
-
-        print(f"\n{'='*80}")
-        print(f"FOLD {fold_num}/5 | Train: {len(X_train)}, Val: {len(X_val)}")
-        print(f"{'='*80}\n")
-
-        model = model_class()
-
-        best_acc = train_cnn(
-            model=model,
-            X_train=X_train,
-            y_train=y_train,
-            X_val=X_val,
-            y_val=y_val,
-            fold_num=fold_num,
-            epochs=epochs,
-            batch_size=batch_size,
-            lr=lr,
-            device=device,
-            use_all_patches=use_all_patches,
-            samples_per_epoch_fraction=samples_per_epoch_fraction,
-            checkpoint_dir=os.path.join(checkpoint_dir, f"fold_{fold_num}"),
-            save_every_n_epoch=save_every_n_epoch
-        )
-        fold_accuracies.append(best_acc)
-
-        print(f"\nFold {fold_num} Best Accuracy: {best_acc:.4f}\n")
-
-    mean_acc = np.mean(fold_accuracies)
-    std_acc = np.std(fold_accuracies)
-
-    print(f"\n{'='*80}")
-    print("FINAL 5-FOLD CROSS-VALIDATION RESULTS")
-    print(f"Fold Accuracies: {fold_accuracies}")
-    print(f"Mean Accuracy: {mean_acc:.4f} ± {std_acc:.4f}")
-    print(f"{'='*80}\n")
-
-    # Save results
-    results_path = os.path.join(checkpoint_dir, "5fold_cv_results.json")
-    os.makedirs(checkpoint_dir, exist_ok=True)
-    with open(results_path, 'w') as f:
-        json.dump({
-            'fold_accuracies': fold_accuracies,
-            'mean_accuracy': mean_acc,
-            'std_accuracy': std_acc
-        }, f, indent=2)
-    print(f"Results saved to {results_path}")
-
-    return fold_accuracies, mean_acc
+            fold_accuracies.append(best_acc)
+            print(f"\nFold {fold_num} Best Accuracy: {best_acc:.4f}\n")
+
+        mean_acc = float(np.mean(fold_accuracies))
+        std_acc = float(np.std(fold_accuracies))
+
+        print(f"\n{'='*80}\nFINAL 5-FOLD CV RESULTS\nFold Accuracies: {fold_accuracies}\nMean: {mean_acc:.4f} ± {std_acc:.4f}\n{'='*80}\n")
+
+        results_path = os.path.join(self.config.checkpoint_dir, "5fold_cv_results.json")
+        os.makedirs(self.config.checkpoint_dir, exist_ok=True)
+        with open(results_path, 'w') as f:
+            json.dump({'fold_accuracies': fold_accuracies, 'mean_accuracy': mean_acc, 'std_accuracy': std_acc}, f, indent=2)
+
+        return fold_accuracies, mean_acc
+
+    def _predict_val(self, model: torch.nn.Module, spec: np.ndarray, device: str) -> int:
+        from src.config.config import DatasetConfig
+        cfg = DatasetConfig()
+        n_frames = spec.shape[0]
+        if n_frames < cfg.cnn_input_length:
+            spec = np.pad(spec, ((0, cfg.cnn_input_length - n_frames), (0, 0)), mode="constant")
+            n_frames = cfg.cnn_input_length
+        patches = np.stack([
+            spec[s:s + cfg.cnn_input_length]
+            for s in range(0, n_frames - cfg.cnn_input_length + 1)
+        ])[:, np.newaxis]
+        patches_t = torch.tensor(patches, dtype=torch.float32).to(device)
+        with torch.no_grad():
+            out = model(patches_t).mean(dim=0)
+        return out.argmax().item()