image_output_1.png
from PIL import Image

from matplotlib import patches, patheffects

from torch.utils.data import Dataset

from torch.utils.data import Dataset, DataLoader

from torchvision import datasets, transforms

from torchvision.transforms.functional import to_tensor

from tqdm.notebook import trange, tqdm

import matplotlib.pyplot as plt

import numpy as np

import os

import torch as tc

import torch.nn as nn

import torch.optim as optim

import xml.etree.ElementTree as ET

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 calculate_iou(box_true, box_pred):
    x1_tr, y1_tr, w_tr, h_tr = box_true
    x1_pr, y1_pr, w_pr, h_pr = box_pred
    intersection_x1 = max(x1_tr, x1_pr)
    intersection_y1 = max(y1_tr, y1_pr)
    intersection_x2 = min(x1_tr + w_tr, x1_pr + w_pr)
    intersection_y2 = min(y1_tr + h_tr, y1_pr + h_pr)
    intersection_w = max(0, intersection_x2 - intersection_x1)
    intersection_h = max(0, intersection_y2 - intersection_y1)
    intersection_area = intersection_w * intersection_h
    area_true = w_tr * h_tr
    area_pred = w_pr * h_pr
    union_area = area_true + area_pred - intersection_area
    iou = intersection_area / union_area
    return iou
def batch_iou(box_true_batch, box_pred_batch):
    iou_list = []
    for box_true, box_pred in zip(box_true_batch, box_pred_batch):
        iou = calculate_iou(box_true, box_pred)
        iou_list.append(iou.item())
        iou_array = np.array(iou_list).mean()
    return iou_array

class Data2tensor(Dataset):
    def __init__(self, data_dir, transforms=None):
        self.data_dir = data_dir
        self.transforms = transforms
        self.img_paths = []
        self.xml_paths = []
        for filename in os.listdir(data_dir):
            name, ext = os.path.splitext(filename)
            if ext == '.jpg':
                img_path = os.path.join(data_dir, filename)
                xml_path = os.path.join(data_dir, name + '.xml')
                if os.path.isfile(xml_path):
                    self.img_paths.append(img_path)
                    self.xml_paths.append(xml_path)
    def __len__(self):
        return len(self.img_paths)
    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        xml_path = self.xml_paths[idx]
        img = Image.open(img_path).convert('RGB')
        w, h = img.size
        tree = ET.parse(xml_path)
        root = tree.getroot()
        bboxes = []
        for obj in root.findall('object'):
            bbox = obj.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)
            y_bbox = np.array([xmin, ymin, xmax, ymax])
            y_bbox = y_bbox / np.array([w, h, w, h])
            y_bbox = [*y_bbox[:2], *y_bbox[2:] - y_bbox[:2]]
            bboxes.append(y_bbox)
        if self.transforms:
            img = self.transforms(img)
        else:
            img = to_tensor(img)
        bboxes = [tc.tensor(bbox, dtype=tc.float32) for bbox in bboxes]
        return (img, bboxes)

def show_img_and_bbox(x, y, ax=None):
    # plot the image:
    if not ax: _, ax = plt.subplots(1, 1, figsize=(5, 5))
    if len(x.shape) == 3:
        H,W = x.shape[1:]
        x = x.numpy().transpose(1, 2, 0)    
    ax.imshow(x)
    ax.axis('off')
    
    # showing bounding box
    bbox = y
    x, y, w, h = bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]
    draw_rect(ax, [x*W, y*H], x*W+ w*W, y*H+h*H)

x, y = ds_tr0[3]
show_img_and_bbox(x, y[0])