# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + id="ExbljdUu0hL6"
import numpy as np
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
from torch.utils.data.sampler import SubsetRandomSampler
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import torchvision.models
from PIL import Image
import torchvision.transforms as transforms
import os

# + colab={"base_uri": "https://localhost:8080/"} id="up8j3Sdx1igI" outputId="e07c5ef5-b469-4f80-ee32-545cb9538bf1"
from google.colab import drive
drive.mount('/content/gdrive')


# + [markdown] id="J2gcf3cd-0PX"
# ### Loading Data

# + id="eEhuq3IW2xuq"
def load_data(batch_size=64):
  # Compose allows us to have multiple transformations to occur
  # and resize all images to 224x224
  transform_it = transforms.Compose([transforms.Resize((224,224)), transforms.ToTensor()])

  # Save the paths of each of the different types of data that are located in my drive 
  train_path = '/content/gdrive/MyDrive/Primary Models/sample_large/train'
  val_path = '/content/gdrive/MyDrive/Primary Models/sample_large/val'
  test_path = '/content/gdrive/MyDrive/Primary Models/sample_large/test'

  # Load all of the data from my google drive
  train_data = torchvision.datasets.ImageFolder(train_path, transform=transform_it)
  val_data = torchvision.datasets.ImageFolder(val_path, transform=transform_it)
  test_data = torchvision.datasets.ImageFolder(test_path, transform=transform_it)

  return train_data, val_data, test_data


# + id="f9zTOTWr15CI"
import torchvision.models
alexnet = torchvision.models.alexnet(pretrained=True)

# + colab={"base_uri": "https://localhost:8080/"} id="YjqultkH3B81" outputId="1c9de438-06c4-4db4-c56f-2681af073b07"
train_data, val_data, test_data = load_data()
# TRAINING EXAMPLES
print("Training Examples:", len(train_data))

#VALIDATION EXMAPLES 
print("Validation Examples:", len(val_data))

#TEST EXAMPLES 
print("Test Examples:", len(test_data))

# + id="pQDKr_gO3b7S"
# Load all of the data such as training, validation and test
train_loader = torch.utils.data.DataLoader(train_data, batch_size=1, num_workers=1, shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data, batch_size=1, num_workers=1, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=1, num_workers=1, shuffle=True)

# Create a function to save all the features into the folders 
def load_features(loader, new_main_dir, folder):
  classes = ['Normal', 'Pneumonia-Bacterial', 'Pneumonia-Viral', 'COVID-19']

  n = 0
  for img, label in loader:
    # Get the features of each image and transform to a tensor 
    feature = alexnet.features(img)
    feature_tensor = torch.from_numpy(feature.detach().numpy())

    # Save those features into folder 
    feature_dir = new_main_dir + '/' + str(folder) + '/' + str(classes[label])
    if not os.path.isdir(feature_dir):
      os.mkdir(feature_dir)
    torch.save(feature_tensor.squeeze(0), feature_dir + '/' + str(n) + '.tensor')

    n+=1

  print("Features done for ", folder)


# + colab={"base_uri": "https://localhost:8080/"} id="aH3gYo8T4Cwf" outputId="66ceb4d1-0b0e-48de-9d67-7ec9b5cd4ddb"
# load_features(train_loader, '/content/gdrive/MyDrive/Primary Models/transfer_learning/AlexNet', 'Train')
# load_features(val_loader,'/content/gdrive/MyDrive/Primary Models/transfer_learning/AlexNet','Validation')
# load_features(test_loader,'/content/gdrive/MyDrive/Primary Models/transfer_learning/AlexNet','Test')

# + [markdown] id="NGzjEhdi-wNC"
# ### Training

# + id="LFmtoWoF-u9m"
class AlexNet_Classifier(nn.Module):
    def __init__(self):
        super(AlexNet_Classifier, self).__init__()
        self.name = "AlexNet"
        self.fc1 = nn.Linear(256 * 6 * 6, 10)
        self.fc2 = nn.Linear(10, 9)

    def forward(self, x):
        x = x.view(-1, 256 * 6 * 6) #flatten feature data
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# + id="iRnn00hCA2Hx"
def get_accuracy(model, data_loader):
    correct = 0
    total = 0
    for imgs, labels in data_loader:
         
        #############################################
        #To Enable GPU Usage
        if use_cuda and torch.cuda.is_available():
          imgs = imgs.cuda()
          labels = labels.cuda()
        #############################################
        
        output = model(imgs)
        
        #select index with maximum prediction score
        pred = output.max(1, keepdim=True)[1]
        correct += pred.eq(labels.view_as(pred)).sum().item()
        total += imgs.shape[0]
    return correct / total


# + id="ENCbtC3E-_GR"
def alex_train(model, batch_size, learning_rate, num_epochs, train_loader, val_loader):

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

    iters, losses, train_acc, val_acc = [], [], [], []

    # training
    n = 0 # the number of iterations
    for epoch in range(num_epochs):
        mini_b=0
        mini_batch_correct = 0
        Mini_batch_total = 0
        for imgs, labels in iter(train_loader):
          
            
            #############################################
            #To Enable GPU Usage
            if use_cuda and torch.cuda.is_available():
              imgs = imgs.cuda()
              labels = labels.cuda()
            #############################################

          #### ALNC is alexNet.features (AlexNet without classifier) ####
          
            out = model(imgs)             # forward pass
            loss = criterion(out, labels) # compute the total loss
            loss.backward()               # backward pass (compute parameter updates)
            optimizer.step()              # make the updates for each parameter
            optimizer.zero_grad()         # a clean up step for PyTorch



            ##### Mini_batch Accuracy ##### We don't compute accuracy on the whole trainig set in every iteration!
            pred = out.max(1, keepdim=True)[1]
            mini_batch_correct = pred.eq(labels.view_as(pred)).sum().item()
            Mini_batch_total = imgs.shape[0]
            train_acc.append((mini_batch_correct / Mini_batch_total))
           ###########################

          # save the current training information
            iters.append(n)
            losses.append(float(loss)/batch_size)             # compute *average* loss
            val_acc.append(get_accuracy(model, val_loader))  # compute validation accuracy
            n += 1
            mini_b += 1

        print ('Epoch: ', epoch)

    # plotting
    plt.title("Training Curve")
    plt.plot(iters, losses, label="Train")
    plt.xlabel("Iterations")
    plt.ylabel("Loss")
    plt.show()

    plt.title("Training Curve")
    plt.plot(iters, train_acc, label="Training")
    plt.plot(iters, val_acc, label="Validation")    
    plt.xlabel("Iterations")
    plt.ylabel("Validation Accuracy")
    plt.legend(loc='best')
    plt.show()

    train_acc.append(get_accuracy(model, train_loader))
    print("Final Training Accuracy: {}".format(train_acc[-1]))
    print("Final Validation Accuracy: {}".format(val_acc[-1]))

# + id="8_D98v02_48d"
# Load features from folder

# Retrieve dataset
features_train = torchvision.datasets.DatasetFolder('/content/gdrive/MyDrive/Primary Models/transfer_learning/AlexNet/Train', loader=torch.load, extensions=('.tensor'))
features_val = torchvision.datasets.DatasetFolder('/content/gdrive/MyDrive/Primary Models/transfer_learning/AlexNet/Validation', loader=torch.load, extensions=('.tensor'))
features_test = torchvision.datasets.DatasetFolder('/content/gdrive/MyDrive/Primary Models/transfer_learning/AlexNet/Test', loader=torch.load, extensions=('.tensor'))

# Prepare loader
features_train_loader = torch.utils.data.DataLoader(features_train, batch_size=64, shuffle=True)
features_val_loader = torch.utils.data.DataLoader(features_val, batch_size=64, shuffle=True)
features_test_loader = torch.utils.data.DataLoader(features_test, batch_size=64, shuffle=True)

# + colab={"base_uri": "https://localhost:8080/", "height": 893} id="erl6eh3a_QTM" outputId="3a6ee895-4ff1-4747-c10b-75ca17253b02"
AlexNet_Model = AlexNet_Classifier()

use_cuda = True
if use_cuda and torch.cuda.is_available():
  AlexNet_Model.cuda()
  print('CUDA is available!  Training on GPU ...')
else:
  print('CUDA is not available.  Training on CPU ...')
  
alex_train(model=AlexNet_Model, batch_size=64, learning_rate=0.001, num_epochs=15, train_loader=features_train_loader, val_loader=features_val_loader)


# + id="lGWxT-9JDDxC"
class Xray_Classifier_MODEL2(nn.Module):
      def __init__(self):
          self.name = "Xray_Classifier_MODEL2"
          super(Xray_Classifier_MODEL2, self).__init__()
          self.conv1 = nn.Conv2d(256, 512, 3) #in_channel=3, out_channel=5, kernel_size=5
          #self.pool = nn.MaxPool2d(2, 2) #kernel_size=2, stride=2 
          #self.conv2 = nn.Conv2d(512, 10, 2) #in_channel=5, out_channel=10, kernel_size=5
          self.fc1 = nn.Linear(512*4*4, 80) #in_features=10*53*53, out_features=80
          self.fc2 = nn.Linear(80, 30) #in_features=80, out_features=30
          self.fc3 = nn.Linear(30, 4) #in_features=30, out_features=4

      def forward(self, x):
          #x = self.pool(F.relu(self.conv1(x))) #apply pooling to 1st convolution layer
          #x = self.pool(F.relu(self.conv2(x))) #apply pooling to 2nd convolution layer
          x = F.relu(self.conv1(x))
          #x = F.relu(self.conv2(x))
          x = x.view(-1, 512*4*4)
          x = F.relu(self.fc1(x))
          x = F.relu(self.fc2(x))
          x = self.fc3(x)
          return x


# + colab={"base_uri": "https://localhost:8080/", "height": 982} id="vRS14STgDqgI" outputId="2f9c9cf7-ef9f-4922-cce7-23586cd09d1f"
AlexNet_CNN_Model = Xray_Classifier_MODEL2()

use_cuda = True
if use_cuda and torch.cuda.is_available():
  AlexNet_CNN_Model.cuda()
  print('CUDA is available!  Training on GPU ...')
else:
  print('CUDA is not available.  Training on CPU ...')
  
alex_train(model=AlexNet_CNN_Model, batch_size=64, learning_rate=0.001, num_epochs=20, train_loader=features_train_loader, val_loader=features_val_loader)

# + colab={"base_uri": "https://localhost:8080/", "height": 982} id="kampFYQbOzRh" outputId="f0c7bb37-d974-4d77-8b8f-4bc018e1d3c9"
# Prepare loader
features_train_loader = torch.utils.data.DataLoader(features_train, batch_size=32, shuffle=True)
features_val_loader = torch.utils.data.DataLoader(features_val, batch_size=32, shuffle=True)
features_test_loader = torch.utils.data.DataLoader(features_test, batch_size=32, shuffle=True)

AlexNet_CNN_Model_2 = Xray_Classifier_MODEL2()

use_cuda = True
if use_cuda and torch.cuda.is_available():
  AlexNet_CNN_Model_2.cuda()
  print('CUDA is available!  Training on GPU ...')
else:
  print('CUDA is not available.  Training on CPU ...')
  
alex_train(model=AlexNet_CNN_Model_2, batch_size=32, learning_rate=0.001, num_epochs=20, train_loader=features_train_loader, val_loader=features_val_loader)

StarcoderdataJupyterScriptsDedupFiltered