project_name='10-famous-personality-classification'

import torch
import torchvision
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams['figure.facecolor'] = '#ffffff'

%matplotlib inline

from google.colab import drive
drive.mount('/content/gdrive')

import os
os.environ['KAGGLE_CONFIG_DIR'] = "/content/gdrive/My Drive/Kaggle"

%cd /content/gdrive/My Drive/Kaggle

if os.path.exists('Famous Personality'):
  os.system("rm -r 'Famous Personality'")
!kaggle datasets download -p "Famous Personality" -d tanishgupta26/famous-personalities-image-dataset --unzip

root="Famous Personality/"

os.listdir(root)

os.listdir(root+'Dataset/Dataset')

os.listdir(root+'cropped/cropped')

import cv2
face_cascade = cv2.CascadeClassifier(root+'haarcascade_frontalface_default.xml')

def get_face(image_path):
    img = cv2.imread(image_path,0)
    face = None
    if img is not None:
        faces = face_cascade.detectMultiScale(img,1.3,5)
        if len(faces)==1:
          x,y,w,h=faces[0]
          face = img[y:y+h, x:x+w]
    return face

path=root+'Dataset/Dataset/Bill_Gates/bigates_067.jpg'
plt.imshow(cv2.imread(path));

crop=get_face(path)
if crop is not None:plt.imshow(crop)
else:plt.imshow(cv2.imread(path))

%%time
src='Dataset/Dataset/'
des='cropped/'
if os.path.exists(root+des):
  os.system(f"rm -r '{root+des}'")
os.mkdir(root+des)
for folder in os.listdir(root+src):
  os.mkdir(root+des+folder)
  c=0
  for f in os.listdir(root+src+folder):
    src_path=os.path.join(root,src,folder,f)
    crop=get_face(src_path)
    if crop is not None:
      c+=1
      des_path=os.path.join(root,des,folder,f"{c:03d}.png")
      cv2.imwrite(des_path,crop)
  print(f"Folder: {folder}, {c} images cropped")

for folder in os.listdir(root+des):
  print(folder, len(os.listdir(root+des+folder)))

import torchvision.transforms as tt

stats = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
train_tfms = tt.Compose([tt.Resize((48,48)),
                         tt.RandomCrop(48, padding=6, padding_mode='reflect'),
                         tt.RandomHorizontalFlip(p=.5),
                         tt.ToTensor(),
                         tt.Normalize(*stats,inplace=True)])
valid_tfms = tt.Compose([tt.ToPILImage(),
                         tt.ToTensor(),
                         tt.Resize((48,48)),
                         tt.Normalize(*stats)])

src="Dataset/Dataset/"
des="cropped/"

from torchvision.datasets import ImageFolder

dataset = ImageFolder(root+des, train_tfms)

from random import choice

img,label=choice(dataset)
plt.imshow(img.permute(1,2,0).clamp(0,1))
print(dataset.classes[label].replace('_',' '))

from torch.utils.data import random_split

train_len=int(len(dataset)*.9)
test_len=len(dataset)-train_len
train_ds, valid_ds=random_split(dataset,[train_len, test_len])

from torch.utils.data import DataLoader

batch_size = 200

train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=3, pin_memory=True)
valid_dl = DataLoader(valid_ds, batch_size*2, num_workers=3, pin_memory=True)

from torchvision.utils import make_grid

def denormalize(images, means, stds):
    means = torch.tensor(means).reshape(1, 3, 1, 1)
    stds = torch.tensor(stds).reshape(1, 3, 1, 1)
    return images * stds + means

def show_batch(dl):
    for images, labels in dl:
        fig, ax = plt.subplots(figsize=(12, 12))
        ax.set_xticks([]); ax.set_yticks([])
        denorm_images = denormalize(images, *stats)
        ax.imshow(make_grid(denorm_images[:64], nrow=8).permute(1, 2, 0).clamp(0,1))
        break

show_batch(train_dl)

torch.cuda.is_available()

def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')

device = get_default_device()
device

def to_device(data, device):
    """Move tensor(s) to chosen device"""
    if isinstance(data, (list,tuple)):
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

for images, labels in train_dl:
    print(images.shape)
    images = to_device(images, device)
    print(images.device)
    break

class DeviceDataLoader():
    """Wrap a dataloader to move data to a device"""
    def __init__(self, dl, device):
        self.dl = dl
        self.device = device
        
    def __iter__(self):
        """Yield a batch of data after moving it to device"""
        for b in self.dl: 
            yield to_device(b, self.device)

    def __len__(self):
        """Number of batches"""
        return len(self.dl)

train_dl = DeviceDataLoader(train_dl, device)
valid_dl = DeviceDataLoader(valid_dl, device)

import torch.nn as nn
import torch.nn.functional as F

def accuracy(outputs, labels):
    _, preds = torch.max(outputs, dim=1)
    return torch.tensor(torch.sum(preds == labels).item() / len(preds))

class ImageClassificationBase(nn.Module):
    def training_step(self, batch):
        images, labels = batch 
        out = self(images)                  # Generate predictions
        loss = F.cross_entropy(out, labels) # Calculate loss
        return loss
    
    def validation_step(self, batch):
        images, labels = batch 
        out = self(images)                    # Generate predictions
        loss = F.cross_entropy(out, labels)   # Calculate loss
        acc = accuracy(out, labels)           # Calculate accuracy
        return {'val_loss': loss.detach(), 'val_acc': acc}
        
    def validation_epoch_end(self, outputs):
        batch_losses = [x['val_loss'] for x in outputs]
        epoch_loss = torch.stack(batch_losses).mean()   # Combine losses
        batch_accs = [x['val_acc'] for x in outputs]
        epoch_acc = torch.stack(batch_accs).mean()      # Combine accuracies
        return {'val_loss': epoch_loss.item(), 'val_acc': epoch_acc.item()}
    
    def epoch_end(self, epoch, result):
        print("Epoch [{}], last_lr: {:.5f}, train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(
            epoch, result['lrs'][-1], result['train_loss'], result['val_loss'], result['val_acc']))

def conv_block(in_channels, out_channels, pool=False):
    layers = [nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), 
              nn.BatchNorm2d(out_channels), 
              nn.ReLU(inplace=True)]
    if pool: layers.append(nn.MaxPool2d(2))
    return nn.Sequential(*layers)

class ResNet9(ImageClassificationBase):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        # 3 X 48 X 48
        self.conv1 = conv_block(in_channels, 64)  # 64 X 48 X 48
        self.conv2 = conv_block(64, 128, pool=True) # 128 X 24 X 24
        self.res1 = nn.Sequential(conv_block(128, 128),
                                  conv_block(128, 128)) # 128 X 24 X 24
        
        self.conv3 = conv_block(128, 256, pool=True)  # 256 X 12 X 12
        self.conv4 = conv_block(256, 512, pool=True)  # 512 X 6 X 6
        self.res2 = nn.Sequential(conv_block(512, 512),
                                  conv_block(512, 512)) # 512 X 6 X 6
        
        self.classifier = nn.Sequential(nn.MaxPool2d(6), # 512 X 1 X 1
                                        nn.Flatten(),
                                        nn.Dropout(0.2),
                                        nn.Linear(512, num_classes))
        
    def forward(self, xb):
        out = self.conv1(xb)
        out = self.conv2(out)
        out = self.res1(out) + out
        out = self.conv3(out)
        out = self.conv4(out)
        out = self.res2(out) + out
        out = self.classifier(out)
        return out

model = to_device(ResNet9(3, 10), device)
model

@torch.no_grad()
def evaluate(model, val_loader):
    model.eval()
    outputs = [model.validation_step(batch) for batch in val_loader]
    return model.validation_epoch_end(outputs)

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

def fit_one_cycle(epochs, max_lr, model, train_loader, val_loader, 
                  weight_decay=0, grad_clip=None, opt_func=torch.optim.SGD):
    torch.cuda.empty_cache()
    history = []
    
    # Set up custom optimizer with weight decay
    optimizer = opt_func(model.parameters(), max_lr, weight_decay=weight_decay)
    # Set up one-cycle learning rate scheduler
    sched = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr, epochs=epochs, 
                                                steps_per_epoch=len(train_loader))
    
    for epoch in range(epochs):
        # Training Phase 
        model.train()
        train_losses = []
        lrs = []
        for batch in train_loader:
            loss = model.training_step(batch)
            train_losses.append(loss)
            loss.backward()
            
            # Gradient clipping
            if grad_clip: 
                nn.utils.clip_grad_value_(model.parameters(), grad_clip)
            
            optimizer.step()
            optimizer.zero_grad()
            
            # Record & update learning rate
            lrs.append(get_lr(optimizer))
            sched.step()
        
        # Validation phase
        result = evaluate(model, val_loader)
        result['train_loss'] = torch.stack(train_losses).mean().item()
        result['lrs'] = lrs
        model.epoch_end(epoch, result)
        history.append(result)
    return history

history = [evaluate(model, valid_dl)]
history

epochs = 10
max_lr = 1e-3
grad_clip = 0.1
weight_decay = 1e-4
opt_func = torch.optim.Adam

%%time
history += fit_one_cycle(epochs, max_lr, model, train_dl, valid_dl, 
                             grad_clip=grad_clip, 
                             weight_decay=weight_decay, 
                             opt_func=opt_func)

def plot_accuracies(history):
    accuracies = [x['val_acc'] for x in history]
    plt.plot(accuracies, '-x')
    plt.xlabel('epoch')
    plt.ylabel('accuracy')
    plt.title('Accuracy vs. No. of epochs');

plot_accuracies(history)

def plot_losses(history):
    train_losses = [x.get('train_loss') for x in history]
    val_losses = [x['val_loss'] for x in history]
    plt.plot(train_losses, '-bx')
    plt.plot(val_losses, '-rx')
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend(['Training', 'Validation'])
    plt.title('Loss vs. No. of epochs');

plot_losses(history)

def plot_lrs(history):
    lrs = np.concatenate([x.get('lrs', []) for x in history])
    plt.plot(lrs)
    plt.xlabel('Batch no.')
    plt.ylabel('Learning rate')
    plt.title('Learning Rate vs. Batch no.');

plot_lrs(history)

def predict_image(img, model):
    # Convert to a batch of 1
    xb = to_device(img.unsqueeze(0), device)
    # Get predictions from model
    yb = model(xb)
    # Pick index with highest probability
    _, preds  = torch.max(yb, dim=1)
    # Retrieve the class label
    return dataset.classes[preds[0].item()]

img, label=choice(valid_ds)
plt.imshow(img.permute(1,2,0).clamp(0,1))
plt.imshow(denormalize(img,*stats).squeeze(0).permute(1,2,0).clamp(0,1))
predict_image(img, model), dataset.classes[label]

person=choice(os.listdir(root+src)) # random person
img=choice(os.listdir(root+src+person)) # random person's image
path=os.path.join(root,src,person,img)
crop=get_face(path)
if crop is not None:
  plt.imshow(cv2.imread(path))
  print(predict_image(valid_tfms(cv2.cvtColor(crop,cv2.COLOR_GRAY2BGR)), model))
else:
  print('Face not recognized')
  plt.imshow(cv2.imread(path))