utils/image_classification.py

import torch
import torch.nn as nn
from torch.nn import functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
from PIL import Image
from tempfile import TemporaryDirectory
import streamlit as st

cudnn.benchmark = True
plt.ion()   # interactive mode

class classifier():
    def __init__(self):
        self.data_transforms = None
        self.data_dir = None
        self.image_datasets = None
        self.dataloaders = None
        self.dataset_sizes = None
        self.class_names = None
        self.device = None
        self.num_classes = None
    def data_loader(self,path,batch_size=4):
        # Data augmentation and normalization for training
        # Just normalization for validation
        self.data_transforms = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ]),
            'val': transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ]),
            'test': transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        }

        self.data_dir = path
        self.image_datasets = {x: datasets.ImageFolder(os.path.join(self.data_dir, x),
                                                  self.data_transforms[x])
                          for x in ['train', 'val','test']}
        self.dataloaders = {x: torch.utils.data.DataLoader(self.image_datasets[x], batch_size=batch_size,
                                                     shuffle=True, num_workers=4)
                      for x in ['train', 'val','test']}
        self.dataset_sizes = {x: len(self.image_datasets[x]) for x in ['train', 'val','test']}
        self.class_names = self.image_datasets['train'].classes
        self.num_classes = len(self.class_names)
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        
    def train(self,model, criterion, optimizer, scheduler, num_epochs=25):
        since = time.time()

        # Create a temporary directory to save training checkpoints
        with TemporaryDirectory() as tempdir:
            best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')

            torch.save(model.state_dict(), best_model_params_path)
            best_acc = 0.0

            for epoch in range(num_epochs):
                print(f'Epoch {epoch+1}/{num_epochs}')
                print('-' * 10)
                st.sidebar.subheader(f':blue[Epoch {epoch+1}/{num_epochs}]', divider='blue')
                # st.sidebar.code('-' * 10)
                # Each epoch has a training and validation phase
                for phase in ['train', 'val']:
                    if phase == 'train':
                        model.train()  # Set model to training mode
                    else:
                        model.eval()   # Set model to evaluate mode

                    running_loss = 0.0
                    running_corrects = 0

                    # Iterate over data.
                    for inputs, labels in self.dataloaders[phase]:
                        inputs = inputs.to(self.device)
                        labels = labels.to(self.device)

                        # zero the parameter gradients
                        optimizer.zero_grad()

                        # forward
                        # track history if only in train
                        with torch.set_grad_enabled(phase == 'train'):
                            outputs = model(inputs)
                            _, preds = torch.max(outputs, 1)
                            loss = criterion(outputs, labels)

                            # backward + optimize only if in training phase
                            if phase == 'train':
                                loss.backward()
                                optimizer.step()

                        # statistics
                        running_loss += loss.item() * inputs.size(0)
                        running_corrects += torch.sum(preds == labels.data)
                    if phase == 'train':
                        scheduler.step()

                    epoch_loss = running_loss / self.dataset_sizes[phase]
                    epoch_acc = running_corrects.double() / self.dataset_sizes[phase]

                    print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')
                    st.sidebar.caption(f':blue[{phase[0].upper() + phase[1:]} Loss:] {epoch_loss:.4f} :blue[   Accuracy:] {epoch_acc:.4f}')
                    # deep copy the model
                    if phase == 'val' and epoch_acc > best_acc:
                        best_acc = epoch_acc
                        torch.save(model.state_dict(), best_model_params_path)

                print()

            time_elapsed = time.time() - since
            print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
            print(f'Best val Accuracy: {best_acc:4f}')
            st.sidebar.caption(f':green[Training complete in] {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
            st.sidebar.subheader(f':blue[Best val Accuracy:] {best_acc:4f}')
            # load best model weights
            model.load_state_dict(torch.load(best_model_params_path))
        return model
    
    def train_model(self,model_name,epochs):
        num_classes = self.num_classes
        if model_name == 'EfficientNet_B0':
            model = models.efficientnet_b0(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        #         model.classifier[1].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)

        elif model_name == 'EfficientNet_B1':
            model = models.efficientnet_b1(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        #         model.classifier[1].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)
        elif model_name == 'MnasNet0_5':
            model = models.mnasnet0_5(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        #         model.classifier[1].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)

        elif model_name == 'MnasNet0_75':
            model = models.mnasnet0_75(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        #         model.classifier[1].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'MnasNet1_0':
            model = models.mnasnet1_0(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        #         model.classifier[1].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'MobileNet_v2':
            model = models.mobilenet_v2(pretrained=True)
            model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
        #         model.classifier[1].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'MobileNet_v3_small':
            model = models.mobilenet_v3_small(pretrained=True)
            model.classifier[3] = nn.Linear(model.classifier[3].in_features, num_classes)
        #         model.classifier[3].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[3].parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'MobileNet_v3_large':
            model = models.mobilenet_v3_large(pretrained=True)
            model.classifier[3] = nn.Linear(model.classifier[3].in_features, num_classes)
        #         model.classifier[3].out_features = num_classes
            optimizer = torch.optim.SGD(model.classifier[3].parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'RegNet_y_400mf':
            model = models.regnet_y_400mf(pretrained=True)
            model.fc = nn.Linear(model.fc.in_features, num_classes)
        #         model.fc.out_features = num_classes
            optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'ShuffleNet_v2_x0_5':
            model = models.shufflenet_v2_x0_5(pretrained=True)
            model.fc = nn.Linear(model.fc.in_features, num_classes)
        #         model.fc.out_features = num_classes
            optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'ShuffleNet_v2_x1_0':
            model = models.shufflenet_v2_x1_0(pretrained=True)
            model.fc = nn.Linear(model.fc.in_features, num_classes)
        #         model.fc.out_features = num_classes
            optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'ShuffleNet_v2_x1_5':
            model = models.shufflenet_v2_x1_5(pretrained=True)
            model.fc = nn.Linear(model.fc.in_features, num_classes)
        #         model.fc.out_features = num_classes
            optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'SqueezeNet 1_0':
            model = models.squeezenet1_0(pretrained=True)
            model.classifier[1] = nn.Conv2d(model.classifier[1].in_channels, num_classes,model.classifier[1].kernel_size, model.classifier[1].stride)
        #         model.classifier[1].out_channels = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)


        elif model_name == 'SqueezeNet 1_1':
            model = models.squeezenet1_1(pretrained=True)
            model.classifier[1] = nn.Conv2d(model.classifier[1].in_channels, num_classes,model.classifier[1].kernel_size, model.classifier[1].stride)
        #         model.classifier[1].out_channels = num_classes
            optimizer = torch.optim.SGD(model.classifier[1].parameters(), lr=0.001, momentum=0.9)

        exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
        criterion = nn.CrossEntropyLoss()
        model_ft = self.train(model, criterion, optimizer, exp_lr_scheduler,
                           num_epochs=epochs)
        torch.save(model.state_dict(), 'model.pt')
        return model_ft
    
    def imshow(self,inp, title=None):
        """Display image for Tensor."""
        inp = inp.numpy().transpose((1, 2, 0))
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        inp = std * inp + mean
        inp = np.clip(inp, 0, 1)
        plt.imshow(inp)
        if title is not None:
            plt.title(title)
        plt.pause(0.001)
    
    def visualize_model(self,model, num_images=6):
        was_training = model.training
        model.eval()
        images_so_far = 0
        fig = plt.figure()

        with torch.no_grad():
            for i, (inputs, labels) in enumerate(self.dataloaders['val']):
                inputs = inputs.to(self.device)
                labels = labels.to(self.device)

                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)

                for j in range(inputs.size()[0]):
                    images_so_far += 1
                    ax = plt.subplot(num_images//2, 2, images_so_far)
                    ax.axis('off')
                    ax.set_title(f'predicted: {self.class_names[preds[j]]}')
                    self.imshow(inputs.cpu().data[j])

                    if images_so_far == num_images:
                        model.train(mode=was_training)
                        return
            model.train(mode=was_training)
            
    def pytorch_predict(self,model):
        '''
        Make prediction from a pytorch model 
        '''
        # set model to evaluate model
        
        model.eval()

        y_true = torch.tensor([], dtype=torch.long, device=self.device)
        all_outputs = torch.tensor([], device=self.device)

        # deactivate autograd engine and reduce memory usage and speed up computations
        with torch.no_grad():
            for data in self.dataloaders['test']:
                inputs = [i.to(self.device) for i in data[:-1]]
                labels = data[-1].to(self.device)

                outputs = model(*inputs)
                y_true = torch.cat((y_true, labels), 0)
                all_outputs = torch.cat((all_outputs, outputs), 0)

        y_true = y_true.cpu().numpy()  
        _, y_pred = torch.max(all_outputs, 1)
        y_pred = y_pred.cpu().numpy()
        y_pred_prob = F.softmax(all_outputs, dim=1).cpu().numpy()

        return y_true, y_pred, y_pred_prob