Upload 2 files

utils/image_classification.py

@@ -0,0 +1,307 @@
+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
+
+
+