text_plain_output_1.png
from matplotlib import patches, patheffects

from torch.utils.data import ConcatDataset

from torch.utils.data import Dataset, DataLoader

from torch.utils.data import Subset

from torchvision import datasets, transforms

from tqdm.notebook import trange, tqdm

import math

import matplotlib.pyplot as plt

import numpy as np

import torch as tc

import torch.nn as nn

import torch.optim as optim

imagenet_stats = ([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

def one_epoch(net, loss, dl, opt=None, metric=None):
    
    if opt: net.train()  # only affects some layers
    else:   net.eval()
    
    L, M = [], []
    dl_it = iter(dl)
    for xb, yb in tqdm(dl_it, leave=False):
        xb = xb.cuda()
        if not isinstance(yb, list): yb = [yb]  # this is new(!)
        yb = [yb_.cuda() for yb_ in yb]
        
        if opt:
            y_ = net(xb)
            l = loss(y_, yb)
            opt.zero_grad()
            l.backward()
            opt.step()
        else:
            with tc.no_grad():
                y_ = net(xb)
                l = loss(y_, yb)
        L.append(l.detach().cpu().numpy())
        if isinstance(metric, list):
            for m in metric:
                M.append(m(tc.sigmoid(y_), yb[0]))
                
        elif metric:
            M.append(metric(tc.sigmoid(y_), yb[0]))       
    return L, M

def fit(net, tr_dl, val_dl, loss=nn.CrossEntropyLoss(), epochs=3, lr=3e-3, wd=1e-3, plot=True):   
    
    opt = optim.Adam(net.parameters(), lr=lr, weight_decay=wd)
    
    Ltr_hist, Lval_hist = [], []
    for epoch in trange(epochs):
        Ltr,  _    = one_epoch(net, loss, tr_dl,  opt)
        Lval, Aval = one_epoch(net, loss, val_dl, None, batch_iou)
        Ltr_hist.append(np.mean(Ltr))
        Lval_hist.append(np.mean(Lval))
        #print(f'epoch: {epoch}\ttraining loss: {np.mean(Ltr):0.4f}\tvalidation loss: {np.mean(Lval):0.4f} \tvalidation accuracy: {mean(Aval):0.2f}')
        print(f'epoch: {epoch}\ttraining loss: {np.mean(Ltr):0.4f}\tvalidation loss: {np.mean(Lval):0.4f}, overlap accuracy: {np.array(Aval).mean():0.2f}')
        
    # plot the losses     
    if plot:
        _,ax = plt.subplots(1,1,figsize=(16,4))
        ax.plot(1+np.arange(len(Ltr_hist)),Ltr_hist)
        ax.plot(1+np.arange(len(Lval_hist)),Lval_hist)
        ax.grid('on')
        ax.set_xlim(left=1, right=len(Ltr_hist))
        ax.legend(['training loss', 'validation loss']);
        
    return Ltr_hist, Lval_hist

def denorm(x, stats=imagenet_stats):
    return x * tc.Tensor(stats[1])[:,None,None] + tc.Tensor(stats[0])[:,None,None]

transform = transforms.Compose([
    transforms.RandomApply([transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1)], p=1),
    transforms.ToTensor()])

def draw_rect(ax, xy, w, h):
    patch = ax.add_patch(patches.Rectangle(xy, w, h, fill=False, edgecolor='yellow', lw=2))
    patch.set_path_effects([patheffects.Stroke(linewidth=6, foreground='black'), patheffects.Normal()])
    
    
def _freeze(md, fr=True):
    ch = list(md.children())
    for c in ch: _freeze(c, fr)
    if not ch and not isinstance(md, tc.nn.modules.batchnorm.BatchNorm2d):  # not freezing the BatchNorm layers!
        for p in md.parameters(): 
            p.requires_grad = not fr

def freeze_to(md, ix=-1):
    ch_all = list(md.children())
    for ch in ch_all[:ix]: _freeze(ch, True)

def unfreeze_to(md, ix=-1):
    ch_all = list(md.children())
    for ch in ch_all[:ix]: _freeze(ch, False)
        

def ds_train_val(train_dir, val_dir, p, transforms):
    ds_full = Data2tensor(train_dir, transforms)
    num_data = len(ds_full)
    num_train = math.ceil(num_data * p)
    train_indices = list(range(num_train))
    ds_tr = Subset(ds_full, train_indices)
    ds_val = Data2tensor(val_dir, None)
    return (ds_tr, ds_val)

train_dir = '/kaggle/input/spot250/train'
val_dir = '/kaggle/input/spot250/valid'
ds_tr0, ds_val = ds_train_val(train_dir, val_dir, p=1, transforms=transform)

from torch.utils.data import ConcatDataset
transform1 = transforms.Compose([transforms.RandomApply([transforms.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1)], p=1), transforms.ToTensor()])
transform2 = transforms.Compose([transforms.RandomApply([transforms.ColorJitter(brightness=0.4, contrast=0.5, saturation=0.2, hue=0.1)], p=1), transforms.ToTensor()])
ds_tr, _ = ds_train_val(train_dir, val_dir, p=1, transforms=None)
transf_ds1, _ = ds_train_val(train_dir, val_dir, p=0.4, transforms=transform1)
ds_tr1 = ConcatDataset([ds_tr, transf_ds1])
transf_ds2, _ = ds_train_val(train_dir, val_dir, p=0.4, transforms=transform2)
transformed_ds_tr = ConcatDataset([ds_tr1, transf_ds2])

bs = 500
train_and_transformed = DataLoader(transformed_ds_tr, batch_size=bs, shuffle=True, num_workers=0)
train_dl = DataLoader(ds_tr, batch_size=bs, shuffle=True, num_workers=0)
val_dl = DataLoader(ds_val, batch_size=2 * bs, shuffle=True, num_workers=0)
xt, yt = next(iter(train_and_transformed))
xb, yb = next(iter(train_dl))
x_v, y_v = next(iter(val_dl))
print(f'Train and transformed data:  x shape = {xt.shape}, bbox shape = {yt[0].shape} \n')
print(f'Training data:  image shape = {xb.shape}, bbox shape = {yb[0].shape} \n')
print(f'Validation data: image shape = {x_v.shape}, bbox shape = {y_v[0].shape}')