Initial commit

Author: licksylick

.gitignore

@@ -0,0 +1,20 @@
+__pycache__/
+saves/
+*.mp4
+*.csv
+*.pt
+
+# virtualenv
+venv/
+ENV/
+env/
+
+# IPython Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+.idea/
+dataset/
+classification_models_pytorch/dataset/

classification_models_pytorch/config.yaml

@@ -0,0 +1,58 @@
+common:
+  exp_name: test_exp
+  seed: 666
+  batch_size: 8
+  num_workers: 1
+  max_epochs: 4
+  use_cross_validation: False
+  num_splits: 4
+
+dataset:
+  path: dataset
+  img_size: 224
+  val_size: 0.3
+
+trainer:
+  target: pytorch_lightning.Trainer
+  params:
+#    gpus: [0]
+    max_epochs: 20
+
+model:
+  target: models.ClassificationModel
+  params:
+    arch: resnet18
+    pretrained: True
+    num_classes: 2
+
+criterions:
+- target: torch.nn.CrossEntropyLoss
+  weight: 1.0
+  name: cross_entropy
+
+optimizers:
+- target: Adam # optimizer name from torch.optim
+  params:
+    lr: 0.001
+  use_lookahead: True
+
+scheduler:
+  - target: ReduceLROnPlateau
+    monitor: val_loss
+
+metrics:
+- target: torch.nn.functional.cross_entropy
+
+callbacks:
+- target: pytorch_lightning.callbacks.ModelCheckpoint
+  params:
+    filename: best-{epoch:02d}-min-{val_loss:2.2f}
+    monitor: val_loss
+    mode: min
+    save_top_k: 1
+    save_last: true
+- target: pytorch_lightning.callbacks.EarlyStopping
+  params:
+    monitor: val_loss
+    patience: 5
+    mode: min

classification_models_pytorch/data.py

@@ -0,0 +1,35 @@
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+from utils import make_weights_for_balanced_classes
+
+class ClassificationDataset(Dataset):
+    def __init__(self, file_paths, labels, transform=None):
+        self.file_paths = file_paths
+        self.labels = labels
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx):
+        image = cv2.imread(self.file_paths[idx])
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        label = torch.tensor(self.labels[idx], dtype=torch.long)
+
+        if self.transform:
+            image = self.transform(image)
+
+        return image, label
+
+
+def create_dataloader(file_paths, labels, batch_size, transform, num_workers):
+    dataset = ClassificationDataset(file_paths, labels, transform)
+
+    weights = make_weights_for_balanced_classes(dataset.labels)
+    weights = torch.DoubleTensor(weights)
+    sampler = torch.utils.data.sampler.WeightedRandomSampler(weights, len(weights))
+
+    dataloader = DataLoader(dataset, batch_size=batch_size, sampler=sampler, num_workers=num_workers)
+    return dataloader

classification_models_pytorch/inference.py

@@ -0,0 +1,43 @@
+import os
+import warnings
+import torch
+import cv2
+from argparse import ArgumentParser
+from models.model import ClassificationModel
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from transforms import Transforms
+
+
+warnings.filterwarnings('ignore')
+
+
+def predict(model, image, config):
+    transforms = Transforms(config['dataset']['img_size'])
+    test_transform = transforms.val_test_transform()
+    img_tensor = test_transform(image).unsqueeze_(0)
+    prediction = torch.argmax(F.softmax(model(img_tensor))).item()
+    return prediction
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument('--config', type=str, required=True, help='Model name (resnet18|resnet34|resnet50|efficientnet')
+    parser.add_argument('--model_path', required=True, type=str, help='Path to model checkpoint')
+    parser.add_argument('--image', type=str, required=True, help='Path to input image')
+    args = parser.parse_args()
+
+    checkpoint = torch.load(args.model_path)
+    config = OmegaConf.load(args.config)
+    model = ClassificationModel(config)
+    model.load_state_dict(checkpoint['state_dict'])
+    model.eval()
+
+    image = cv2.imread(os.path.join(args.image))
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+    print(f'Prediction is: {predict(model, image, config)}')
+
+
+if __name__ == '__main__':
+    main()

classification_models_pytorch/models/model.py

@@ -0,0 +1,116 @@
+import os
+import sys
+sys.path.append('models')
+import torch
+import timm
+import pytorch_lightning as pl
+import torch.nn.functional as F
+from torchmetrics.functional import f1_score, precision, recall
+from sklearn.metrics import confusion_matrix
+import matplotlib.pyplot as plt
+import seaborn as sns
+from optimizer import Lookahead
+
+
+class ClassificationModel(pl.LightningModule):
+
+    def __init__(self, config: dict):
+        super().__init__()
+        self.config = config
+        model_name = config['model']['params']['arch']
+        try:
+            self.backbone = timm.create_model(model_name,
+                                              pretrained=config['model']['params']['pretrained'],
+                                              num_classes=config['model']['params']['num_classes'])
+        except:
+            raise ValueError(f'Undefined value of model name: {model_name}')
+
+        self.num_classes = config['model']['params']['num_classes']
+
+    def forward(self, x):
+        return self.backbone(x)
+
+    def configure_optimizers(self):
+        optimizer_name = self.config['optimizers'][0]['target']
+        optimizer_params = self.config['optimizers'][0]['params']
+
+        optimizer_class = getattr(torch.optim, optimizer_name)
+
+        if self.config['optimizers'][0].get('use_lookahead', True):
+            base_optim = optimizer_class(self.parameters(), **optimizer_params)
+            optimizer = Lookahead(base_optim)
+        else:
+            optimizer = optimizer_class(self.parameters(), **optimizer_params)
+
+        scheduler_name = self.config['scheduler'][0]['target']
+        scheduler_params = getattr(torch.optim.lr_scheduler, scheduler_name)
+        scheduler = scheduler_params(optimizer)
+
+        monitor = self.config['scheduler'][0].get('monitor', '')
+        # return [optimizer], [scheduler]
+        return  {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": monitor}
+
+    def compute_metrics(self, pred, target):
+        metrics = dict()
+        metrics['f1_score'] = f1_score(pred, target, num_classes=self.num_classes, task='multiclass')
+        metrics['precision'] = precision(pred, target, num_classes=self.num_classes, task='multiclass')
+        metrics['recall'] = recall(pred, target, num_classes=self.num_classes, task='multiclass')
+        return metrics
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        output = self(x)
+        loss = F.cross_entropy(output, y)
+        metrics = self.compute_metrics(output, y)
+        self.log('train_f1', metrics['f1_score'], prog_bar=True, on_step=False, on_epoch=True)
+        self.log('train_prec', metrics['precision'], prog_bar=True, on_step=False, on_epoch=True)
+        self.log('train_rec', metrics['recall'], prog_bar=True, on_step=False, on_epoch=True)
+        self.log('train_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x, y = batch
+        output = self(x)
+        loss = F.cross_entropy(output, y)
+        metrics = self.compute_metrics(output, y)
+        self.log('val_f1', metrics['f1_score'], on_step=False, on_epoch=True)
+        self.log('val_prec', metrics['precision'], on_step=False, on_epoch=True)
+        self.log('val_rec', metrics['recall'], on_step=False, on_epoch=True)
+        self.log('val_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        return loss
+
+    def test_step(self, batch, batch_idx):
+        x, y = batch
+        output = self(x)
+        loss = F.cross_entropy(output, y)
+        metrics = self.compute_metrics(output, y)
+        self.log('test_f1', metrics['f1_score'], on_step=False, on_epoch=True)
+        self.log('test_prec', metrics['precision'], on_step=False, on_epoch=True)
+        self.log('test_rec', metrics['recall'], on_step=False, on_epoch=True)
+        self.log('test_loss', loss, prog_bar=True, on_step=False, on_epoch=True)
+        print(f'test_metrics: {metrics}')
+
+        # Convert predictions to class indices
+        _, preds = torch.max(output, 1)
+
+        return {"loss": loss, "preds": preds, "labels": y}
+
+    def test_epoch_end(self, outputs):
+        all_preds = torch.cat([out["preds"] for out in outputs])
+        all_labels = torch.cat([out["labels"] for out in outputs])
+
+        all_preds = all_preds.cpu().numpy()
+        all_labels = all_labels.cpu().numpy()
+
+        conf_matrix = confusion_matrix(all_labels, all_preds)
+
+        plt.figure(figsize=(self.num_classes, self.num_classes))
+        sns.heatmap(conf_matrix, annot=True, fmt="d", cmap="Blues",
+                    xticklabels=range(self.num_classes), yticklabels=range(self.num_classes))
+        plt.xlabel("Predicted")
+        plt.ylabel("True")
+        plt.title("Confusion Matrix")
+
+        # Save the confusion matrix as an image
+        plt.savefig(os.path.join(self.config['common'].get('exp_name', 'exp0') ,"confusion_matrix.png"))
+        plt.close()