!pip install opendatasets --upgrade --quiet

import opendatasets as od

od.download('https://www.kaggle.com/meemr5/vadapav')

import os
import torch
import pandas as pd
import numpy as np
from torch.utils.data import Dataset, random_split, DataLoader
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import torchvision.models as models
import matplotlib.pyplot as plt
import torchvision.transforms as T
import torch.nn.functional as F
import torch.nn as nn
from torchvision.utils import make_grid
from torchvision import  datasets
%matplotlib inline

DATA_DIR = './vadapav/VadaPav'

image_size=224

def target_to_oh(target):
    NUM_CLASS = 2  # hard code here, can do partial
    one_hot = torch.eye(NUM_CLASS)[target]
    return one_hot

def decode_target(target, threshold=0.5):
    result = []
    for i, x in enumerate(target):
        if (x >= threshold):
            return i

transform =T.Compose([ T.Resize(image_size),T.CenterCrop(image_size),T.ToTensor()])

dataset=datasets.ImageFolder(DATA_DIR,target_transform = target_to_oh ,transform= transform)

print("Dataset size :",len(dataset) )

print(" Classes :",dataset.classes)

def show_sample(img, target):
    plt.imshow(img.permute(1, 2, 0))
    print('Labels:', dataset.classes[decode_target(target)])

show_sample(*dataset[103])

torch.manual_seed(10)

val_pct , test_pct = 0.19, 0.01
val_size = int(val_pct * len(dataset))
test_size = int(test_pct*len(dataset))
train_size = len(dataset) - val_size - test_size

train_ds, val_ds ,test_ds= random_split(dataset, [train_size, val_size,test_size])
len(train_ds), len(val_ds),len(test_ds)

batch_size = 32

train_dl = DataLoader(train_ds, batch_size, shuffle=True)
val_dl = DataLoader(val_ds, batch_size, shuffle=True)
test_dl = DataLoader(test_ds, batch_size, shuffle=True)

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

show_batch(train_dl)

class BunClassificationBase(nn.Module):
    def training_step(self, batch):
        images, targets = batch
        out = self(images)         
        loss = F.binary_cross_entropy(out, targets.float())      
        return loss
    
    def validation_step(self, batch):
        images, targets = batch 
        out = self(images)                                   # Generate predictions
        loss = F.binary_cross_entropy(out, targets.float())  # Calculate loss
        score = F_score(out, targets)
        return {'val_loss': loss.detach(), 'val_score': score.detach() }
        
    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_scores = [x['val_score'] for x in outputs]
        epoch_score = torch.stack(batch_scores).mean()      # Combine accuracies
        return {'val_loss': epoch_loss.item(), 'val_score': epoch_score.item()}
    
    def epoch_end(self, epoch, result):
        print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}, val_score: {:.4f}".format(
            epoch, result['train_loss'], result['val_loss'], result['val_score']))

def F_score(output, targets, threshold=0.5):    
    pred = output
    pred[output>=threshold]=1
    pred[output<threshold]= 0
    
    res= (pred==targets).float()

    return torch.mean(res)
    

class BunCnnModel(BunClassificationBase):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            
            nn.Conv2d(3, 32, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=32),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),  

            nn.Conv2d(32, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=64),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),    

            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=128),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(num_features=256),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),

            nn.AdaptiveAvgPool2d(1),

            nn.Flatten(), 
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 2),
            nn.Sigmoid()
        )     
    def forward(self, xb):
        return self.network(xb)

model = BunCnnModel()

def get_default_device():
    """Pick GPU if available, else CPU"""
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        return torch.device('cpu')
    
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)

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)

device = get_default_device()
device

train_dl = DeviceDataLoader(train_dl, device)
val_dl = DeviceDataLoader(val_dl, device)
test_dl = DeviceDataLoader(test_dl, device)
to_device(model, device);

def try_batch(dl):
    for images, labels in dl:
        print('images.shape:', images.shape)
        out = model(images)
        print('out.shape:', out.shape)
        print('out[0]:', out[0])
        break

try_batch(train_dl)

from tqdm.notebook import tqdm

@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 fit(epochs, lr, model, train_loader, val_loader, opt_func=torch.optim.SGD):
    torch.cuda.empty_cache()
    history = []
    optimizer = opt_func(model.parameters(), lr)
    for epoch in range(epochs):
        # Training Phase 
        model.train()
        train_losses = []
        for batch in tqdm(train_loader):
            loss = model.training_step(batch)
            train_losses.append(loss)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        # Validation phase
        result = evaluate(model, val_loader)
        result['train_loss'] = torch.stack(train_losses).mean().item()
        model.epoch_end(epoch, result)
        history.append(result)
    return history

model = to_device(BunCnnModel(), device)

evaluate(model, val_dl)

num_epochs = 10
opt_func = torch.optim.Adam
lr = 1e-3

%%time
history = fit(num_epochs, lr, model, train_dl, val_dl, opt_func)

def predict_single(image):
    xb = image.unsqueeze(0)
    xb = to_device(xb, device)
    preds = model(xb)
    prediction = preds[0]
    print("Prediction: ", prediction)
    show_sample(image, prediction)

predict_single(test_ds[0][0])