Horus: added a pre-trained model - ResNet and checkpoint support.

requirements.txt

@@ -1,7 +1,11 @@
 librosa==0.10.0.post2
 matplotlib==3.5.3
 noisereduce==2.0.1
-numpy==1.23.3
 pydub==0.25.1
+pytorch_lightning==2.0.2
 requests==2.28.1
+scikit_learn==1.2.2
+scipy==1.10.1
+seaborn==0.12.2
 soundfile==0.12.1
+torch==2.0.0

src/models/AudioProcessor.py

@@ -39,7 +39,7 @@ def shift_pitch(self, audio, sr, pitch_factor=3):
 
     def change_speed(self, audio, sr, speed_factor=1.2):
         return librosa.effects.time_stretch(audio, rate=speed_factor)
-    
+
     def trim_audio(self, audio, sr):
         audio_length = len(audio)
         fixed_length_samples = self.fixed_length_seconds * sr
@@ -50,26 +50,30 @@ def trim_audio(self, audio, sr):
             audio = audio[:fixed_length_samples]
 
         return audio
-    
+
     def load_and_trim_audio(self, file_path):
-        audio, sr = librosa.load(file_path, sr=self.sample_rate,  duration=self.fixed_length_seconds)
+        audio, sr = librosa.load(
+            file_path, sr=self.sample_rate, duration=self.fixed_length_seconds
+        )
 
         if len(audio) == 0:
-            return None
+            return None, None
 
         audio = self.trim_audio(audio, sr)
 
         return audio, sr
 
     def preprocess_audio(self, file_path):
-        audio, sr = librosa.load(file_path, sr=self.sample_rate, duration=self.fixed_length_seconds)
+        audio, sr = librosa.load(
+            file_path, sr=self.sample_rate, duration=self.fixed_length_seconds
+        )
 
         if len(audio) == 0:
-            return None
-        
+            return None, None
+
         audio = self.trim_audio(audio, sr)
         if audio is None:
-            return None
+            return None, None
         audio = nr.reduce_noise(y=audio, sr=sr)
         if self.normalize_volume:
             audio = librosa.util.normalize(audio)
@@ -89,7 +93,9 @@ def save_preprocessed_audio(self, audio, sr, file_path):
 
     def process_and_save_audio(self, file_path):
         preprocessed_audio, sr = self.preprocess_audio(file_path)
-        target_file_path = self.save_preprocessed_audio(preprocessed_audio, sr, file_path)
+        target_file_path = self.save_preprocessed_audio(
+            preprocessed_audio, sr, file_path
+        )
         return target_file_path
 
     def create_spectrogram(self, file_path, n_fft=2048, hop_length=512):
@@ -109,7 +115,9 @@ def create_mel_spectrogram_from_audio(self, audio, sr, n_mels=128, hop_length=51
 
     def create_mel_spectrogram(self, file_path, n_mels=128, hop_length=512):
         y, sr = librosa.load(file_path, sr=self.sample_rate)
-        return self.create_mel_spectrogram_from_audio(y, sr, n_mels=n_mels, hop_length=hop_length)
+        return self.create_mel_spectrogram_from_audio(
+            y, sr, n_mels=n_mels, hop_length=hop_length
+        )
 
     def save_mel_spectrogram(self, mel_spectrogram_db, target_path, sr, hop_length):
         plt.figure(figsize=(10, 5))
@@ -132,7 +140,8 @@ def save_mel_spectrogram(self, mel_spectrogram_db, target_path, sr, hop_length):
         mel_spec_path = (
             self.numpy_dir
             / target_path.parent.name
-            / target_path.with_suffix(".npy").name )
+            / target_path.with_suffix(".npy").name
+        )
         plot_path.parent.mkdir(parents=True, exist_ok=True)
         plt.savefig(plot_path)
         plt.close()
@@ -217,7 +226,7 @@ def process_folder(self, folder_path: Path):
 
                 except Exception as e:
                     print(f"Error processing file {file_path}: {e}")
-                    
+
 
 if __name__ == "__main__":
     audio_preprocessor = AudioPreprocessor()

src/models/BirdClassifierResNet.py

@@ -0,0 +1,97 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+import pytorch_lightning as pl
+from torchvision.models.resnet import ResNet, ResNet50_Weights
+from torchvision.models import resnet50
+from torchmetrics import Precision, Accuracy, F1Score,  Recall, ConfusionMatrix, AUROC
+
+class BirdClassifierResNet(pl.LightningModule):
+    def __init__(self, num_classes, learning_rate=0.001):
+        super(BirdClassifierResNet, self).__init__()
+
+        # Load pre-trained ResNet50
+        self.feature_extractor = resnet50(weights=ResNet50_Weights.DEFAULT)
+        # Replace the last fully connected layer
+        # Parameters of newly constructed modules have requires_grad=True by default
+        num_ftrs = self.feature_extractor.fc.in_features
+        self.feature_extractor.fc = nn.Linear(num_ftrs, num_classes)
+
+        self.learning_rate = learning_rate
+
+        # Define metrics
+        self.train_acc = Accuracy()
+        self.val_acc = Accuracy()
+        self.precision_metric = Precision(num_classes=num_classes)
+        self.recall = Recall(num_classes=num_classes)
+        self.f1 = F1Score(num_classes=num_classes)
+        self.auroc = AUROC(num_classes=num_classes)
+
+    def forward(self, x):
+        return self.feature_extractor(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = F.cross_entropy(logits, y)
+        preds = torch.argmax(logits, dim=1)
+
+        # Update metrics
+        self.train_acc(preds, y)
+        self.precision_metric(preds, y)
+        self.recall(preds, y)
+        self.f1(preds, y)
+
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        self.log('train_acc', self.train_acc, on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = F.cross_entropy(logits, y)
+        preds = torch.argmax(logits, dim=1)
+
+        # Update metrics
+        self.val_acc(preds, y)
+        self.precision_metric(preds, y)
+        self.recall(preds, y)
+        self.f1(preds, y)
+        # Convert logits to probabilities
+        probs = torch.softmax(logits, dim=1)
+        self.auroc(probs, y)
+
+        self.log('val_loss', loss, prog_bar=True)
+        self.log('val_acc', self.val_acc, prog_bar=True)
+        self.log('val_precision', self.precision, prog_bar=True)
+        self.log('val_recall', self.recall, prog_bar=True)
+        self.log('val_f1', self.f1, prog_bar=True)
+        self.log('val_auroc', self.auroc, prog_bar=True)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        logits = self(x)
+        loss = F.cross_entropy(logits, y)
+        preds = torch.argmax(logits, dim=1)
+
+        # Update metrics
+        self.val_acc(preds, y)
+        self.precision_metric(preds, y)
+        self.recall(preds, y)
+        self.f1(preds, y)
+        # Convert logits to probabilities
+        probs = torch.softmax(logits, dim=1)
+        self.auroc(probs, y)
+
+        self.log('test_loss', loss, prog_bar=True)
+        self.log('test_acc', self.val_acc, prog_bar=True)
+        self.log('test_precision', self.precision_metric, prog_bar=True)
+        self.log('test_recall', self.recall, prog_bar=True)
+        self.log('test_f1', self.f1, prog_bar=True)
+        self.log('test_auroc', self.auroc, prog_bar=True)
+        return loss
+
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=self.learning_rate)

src/models/bird_spectrogram_classifier.py

@@ -3,13 +3,16 @@
 from torch.nn import functional as F
 import pytorch_lightning as pl
 
+
 class BirdClassifier(pl.LightningModule):
     def __init__(self, num_classes):
         super(BirdClassifier, self).__init__()
         self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
         self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
         self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
-        self.fc1 = nn.Linear(239616, 512) # wartości działające ale z dupy do poprawy TODO
+        self.fc1 = nn.Linear(
+            239616, 512
+        )  # wartości działające ale z dupy do poprawy TODO
         self.fc2 = nn.Linear(512, num_classes)
 
     def forward(self, x):
@@ -26,7 +29,9 @@ def training_step(self, batch, batch_idx):
         x, y = batch
         logits = self(x)
         loss = F.cross_entropy(logits, y)
-        self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
+        self.log(
+            "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
+        )
         return loss
 
     def validation_step(self, batch, batch_idx):

src/models/spectrogram_dataset.py

@@ -7,6 +7,7 @@
 import numpy as np
 from AudioProcessor import AudioPreprocessor
 
+
 class BirdSpectrogramDataset(Dataset):
     def __init__(self, root_dir, split="train"):
         self.root_dir = Path(root_dir)
@@ -19,10 +20,14 @@ def __init__(self, root_dir, split="train"):
             if class_dir.name not in self.label_map:
                 self.label_map[class_dir.name] = len(self.label_map)
             for file in class_dir.glob("*.mp3"):
+                if file is None:
+                    continue
                 trimed_audio, sr = self.audio_preprocessor.load_and_trim_audio(file)
                 if trimed_audio is None:
                     continue
-                mel_spec = self.audio_preprocessor.create_mel_spectrogram_from_audio(trimed_audio, sr)
+                mel_spec = self.audio_preprocessor.create_mel_spectrogram_from_audio(
+                    trimed_audio, sr
+                )
                 label = self.label_map[class_dir.name]
                 self.data.append((mel_spec, label))
 
@@ -32,19 +37,24 @@ def __len__(self):
     def __getitem__(self, idx):
         mel_spec, label = self.data[idx]
         mel_spec = torch.from_numpy(mel_spec).unsqueeze(0)  # Add channel dimension
+        mel_spec = mel_spec.expand(3, -1, -1)  # Expand single channel to three channels
         return mel_spec, label
 
+
 class BirdSpectrogramDataModule(pl.LightningDataModule):
-    def __init__(self, root_dir, batch_size=16):
+    def __init__(self, root_dir, batch_size=16, num_workers=4):
         super().__init__()
         self.root_dir = root_dir
         self.batch_size = batch_size
+        self.num_workers = num_workers
 
     def setup(self, stage=None):
         if stage == "fit" or stage is None:
-            self.train_dataset, self.val_dataset = self._create_datasets("train", "validation")
+            self.train_dataset, self.val_dataset = self._create_datasets(
+                "train", "validation"
+            )
         if stage == "test" or stage is None:
-            self.test_dataset, _ = self._create_datasets("test")
+            self.test_dataset = self._create_datasets("test")[0]
 
     def _create_datasets(self, *splits):
         datasets = []
@@ -54,10 +64,25 @@ def _create_datasets(self, *splits):
         return datasets
 
     def train_dataloader(self):
-        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+        )
 
     def val_dataloader(self):
-        return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False)
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+        )
 
     def test_dataloader(self):
-        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False)
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+        )

src/models/test_models.py

@@ -0,0 +1,52 @@
+from pathlib import Path
+import pytorch_lightning as pl
+from BirdClassifierResNet import BirdClassifierResNet
+from bird_spectrogram_classifier import BirdClassifier
+from spectrogram_dataset import BirdSpectrogramDataModule
+from pytorch_lightning.loggers import WandbLogger
+import torch
+import wandb
+
+root_dirs = {
+    "split1": "data/split_raw",
+    "split2": "data/split_aug",
+    "split3": "data/split_wrong",
+}
+num_classes = 7 
+batch_size = 32  # single batch size
+learning_rate = 0.001  # single learning rate
+
+models = [ 
+    Path(f"split{num}_bird_classifier_rest.ckpt") for num in range(1,4)
+]
+
+
+def main():
+    for split_name, root_dir in root_dirs.items():
+        data_module = BirdSpectrogramDataModule(root_dir, batch_size=batch_size)
+        data_module.setup()
+        data_loders = [("train_subset", data_module.train_dataloader()), ("val_subset", data_module.val_dataloader()), ("test_subset", data_module.test_dataloader())]
+        for model_path in models:
+            # Initialize the model
+            model = BirdClassifierResNet.load_from_checkpoint(model_path, num_classes=num_classes, learning_rate=learning_rate)
+
+            for subset_name, data_loder in data_loders:
+                # Set up the Weights & Biases logger
+                wandb_logger = WandbLogger(
+                    project="bird-classification",
+                    name=f"{split_name}_{model_path.name}_{subset_name}",
+                )
+
+                # Test the model
+                trainer = pl.Trainer(logger=wandb_logger)
+                trainer.test(model, dataloaders=data_loder)
+
+                wandb_logger.experiment.config.update({
+                        "split_name": split_name,
+                        "model_name": model_path.name,
+                        "subset_name": subset_name
+                    })
+                wandb_logger.experiment.finish()
+
+if __name__ == "__main__":
+    main()