big update

transfer_learning/extractors.py

@@ -0,0 +1,149 @@
+import torch
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import Dataset
+import torch.nn as nn
+from PIL import Image
+from tqdm import tqdm
+import pandas as pd
+
+
+class FeatureExtractor:
+    supported_extractors = ['resnet18', 'resnet34', 'resnet50',
+                                 'alexnet_fc6', 'alexnet_fc7', 'vgg16',
+                                 'densenet121', 'inception_v3', 'squeezenet']
+
+    def __init__(self, version='resnet18', device='cpu'):
+        assert version.lower() in self.supported_extractors
+        self.device = device
+        self.version = version
+        self.F = self.__choose_feature_extractor(version)
+        for param in self.F.parameters():
+            param.requires_grad = False
+        self.F.eval()
+        self.input_size = (299, 299) if version.lower() == 'inception' else (224, 224)
+        self.transforms = transforms.Compose([
+            transforms.Resize(self.input_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+
+    def to(self, device):
+        self.device = device
+        self.F = self.F.to(self.device)
+        return self
+
+    def __choose_feature_extractor(self, version):
+        if 'resnet' in version.lower():
+            v = int(version[-2:])
+            if v == 18:
+                resnet = torchvision.models.resnet18(pretrained=True)
+            elif v == 34:
+                resnet = torchvision.models.resnet34(pretrained=True)
+            elif v == 50:
+                resnet = torchvision.models.resnet50(pretrained=True)
+            return nn.Sequential(*list(resnet.children())[:-1])
+        elif 'alexnet' in version.lower():
+            v = int(version[-1])
+            alexnet = torchvision.models.alexnet(pretrained=True)
+            if v == 7:
+                f = nn.Sequential(*list(alexnet.classifier.children())[:-2])
+            elif v == 6:
+                f = nn.Sequential(*list(alexnet.classifier.children())[:-5])
+            alexnet.classifier = f
+            return alexnet
+        elif version.lower() == 'vgg16':
+            vgg = torchvision.models.vgg16_bn(pretrained=True)
+            classifier = list(
+                vgg.classifier.children())[:4]
+            vgg.classifier = nn.Sequential(*classifier)
+            return vgg
+        elif version.lower() == 'densenet121':
+            densenet = torchvision.models.densenet121(pretrained=True)
+            avg_pool = nn.AvgPool2d(kernel_size=7)
+            densenet = nn.Sequential(*list(densenet.children())[:-1])
+            densenet.add_module('avg_pool', avg_pool)
+            return densenet
+        elif version.lower() == 'inception_v3':
+            inception = torchvision.models.inception_v3(pretrained=True)
+            f = nn.Sequential(*list(inception.children())[:-1])
+            f._modules.pop('13')
+            f.add_module('global average', nn.AvgPool2d(26))
+            return f
+        elif version.lower() == 'squeezenet':
+            squeezenet = torchvision.models.squeezenet1_1(pretrained=True)
+            f = torch.nn.Sequential(
+                squeezenet.features,
+                torch.nn.AdaptiveAvgPool2d(output_size=(2, 2))
+            )
+            return f
+        else:
+            raise NotImplementedError('The feature extractor you requested is not yet supported')
+
+
+    @property
+    def feature_size(self):
+        x = torch.randn(size=(1, 3, *self.input_size)).to(self.device)
+        return self.F(x).squeeze().shape[0]
+
+    def __call__(self, batch):
+        batch = self.transforms(batch)
+        if len(batch.shape) <= 3:
+            batch = batch.unsqueeze(0)
+        return self.F(batch).view(batch.shape[0], -1).squeeze()
+
+    def from_image_folder(self, folder_path, extension='jpg'):
+        sorted_files = sorted(list(folder_path.glob(f'*.{extension}')))
+        split_names = [x.stem.split('_') for x in sorted_files]
+        names = [x[0] for x in split_names]
+        seq_ids = [x[1] for x in split_names]
+        X = []
+        for i, p_img in enumerate(tqdm(sorted_files)):
+            x = Image.open(p_img)
+            features = self(x)
+            X.append([names[i], seq_ids[i]] + features.tolist())
+        return pd.DataFrame(X, columns=['name', 'seq_id', *(f'feature_{i}' for i in range(self.feature_size))])
+
+
+class AudioTransferLearningImageDataset(Dataset):
+    def __init__(self, root_or_files, extension='jpg', input_size=224):
+        self.root_or_files = root_or_files
+        if type(root_or_files) == list:
+            self.files = root_or_files
+        else:
+            self.files = list(self.root.glob(f'*.{extension}'))
+        self.transforms = transforms.Compose([
+            transforms.Resize(input_size),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+
+    def process_name(self, name):
+        split_name = name.stem.split('_')
+        return split_name[0], split_name[1]  #name, seq_id
+
+    def __getitem__(self, item):
+        p_img = self.files[item]
+        x = Image.open(p_img)
+        x = self.transforms(x)
+        name, seq_id = self.process_name(p_img)
+        return x, name, seq_id
+
+    def __len__(self):
+        return len(self.files)
+
+if __name__ == '__main__':
+    from pathlib import Path
+    version='resnet18'
+    extractor = FeatureExtractor(version=version)
+    models = ['slider', 'pump', 'fan']
+    model_ids = [0, 2, 4, 6]
+    for model in models:
+        for model_id in model_ids:
+            df = extractor.from_image_folder(Path( f'/home/robert/coding/audio_anomaly_detection/data/mimii/-6_dB_{model}/id_0{model_id}/normal/melspec_images/'))
+            df.to_csv(Path(f'/home/robert/coding/audio_anomaly_detection/data/mimii/-6_dB_{model}/id_0{model_id}/normal/{version}_features.csv'), index=False)
+            del df
+            df = extractor.from_image_folder(Path( f'/home/robert/coding/audio_anomaly_detection/data/mimii/-6_dB_{model}/id_0{model_id}/abnormal/melspec_images/'))
+            df.to_csv(Path(f'/home/robert/coding/audio_anomaly_detection/data/mimii/-6_dB_{model}/id_0{model_id}/abnormal/{version}_features.csv'), index=False)