# ---
# 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
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/jerogar/DeepLabV3p-Pytorch-Tutorial_KOR/blob/master/DeepLab_v3_Tutorial.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + [markdown] id="wi0egK65w35f"
# # DeepLab v3 Tutorial
#
# 가짜연구소 Season 2 논문미식회
#
# reference: https://github.com/VainF/DeepLabV3Plus-Pytorch
#
#
#
# ![Main](https://sthalles.github.io/assets/deep_segmentation_network/semantic_segmentation.jpg)
#

# + [markdown] id="1MfKqK5A4v5Z"
# # 1. 환경 설정

# + [markdown] id="wQ1cnSVwXZvj"
# ## (1) Mount Google Drive
# - 구글 드라이브를 연결하고 작업 폴더를 Colab Notebooks/로 설정

# + colab={"base_uri": "https://localhost:8080/"} id="NeGYqxh24oZ-" outputId="8533f30c-8d1e-4edc-e87f-507eb04e5b62"
from google.colab import drive # import drive from google colab

ROOT = "/content/drive"     # default location for the drive
print(ROOT)                 # print content of ROOT (Optional)

drive.mount(ROOT)           # we mount the google drive at /content/drive

# import join used to join ROOT path and MY_GOOGLE_DRIVE_PATH
from os.path import join  

# path to your project on Google Drive
MY_GOOGLE_DRIVE_PATH = '/content/drive/My Drive/Colab Notebooks'

print("MY_GOOGLE_DRIVE_PATH: ", MY_GOOGLE_DRIVE_PATH)
# In case we haven't created the folder already; we will create a folder in the project path  
# %cd "{MY_GOOGLE_DRIVE_PATH}"

# + [markdown] id="qLRclgki42mb"
# ## (2) Set Work directory

# + id="XoRCnewkB77h"
# !git clone DeepLabV3p-Pytorch-Tutorial_KOR

# + id="z6ohDpwJ95R7"
# %cd DeepLabV3p-Pytorch-Tutorial_KOR

# + id="oOi091h8cAbI"
# !pip install torch torchvision numpy pillow scikit-learn tqdm matplotlib opencv-python

# + id="p_Sbl6YTaxtp"
from tqdm import tqdm
import os
import random
import numpy as np

from torch.utils import data
import torch
import torch.nn as nn

import torch.nn.functional as F
from collections import OrderedDict

from torchvision.transforms.functional import normalize
from sklearn.metrics import confusion_matrix


# + [markdown] id="CrHnDiOzJ2CB"
# #2. DeepLab v3 Model
#
# - DeepLab v1 - atrous convolution 및 Conditional Random Field(CRF) 도입
# - DeepLab v2 - A*trous spatial pyramid pooling* (ASPP) 이용하여 다양한 크기의 객체에 대응
# - DeepLab v3 - encoder에 ResNet 구조 도입. ASPP를 보완하고, batch norm. 사용하여 학습이 잘 되도록 함.  CRF 없이 동등 이상 성능 확보
# - DeepLab v3+: Depth-wise seperable Conv.도입하여 런타임 속도 개선. (ASSPP) U-Net 구조를 단순화한 Decoder 구조 사용하여 성능 개선
#
# ![deeplab](https://miro.medium.com/max/1038/0*_Hm_2fqbcnlwLkoz.png)

# + [markdown] id="zacdJ4PLmAM-"
# ## (1) Overall Structure
# ![deeplabv3](https://www.oreilly.com/library/view/hands-on-image-processing/9781789343731/assets/1aa5b349-5a66-456a-8afa-080a7b07a525.png)
#
# - Backbone Network + DeepLab Head 구조
# - Resnet의 layer 4의 feature를 backbone feature로 사용
# - output_stride를 8로 설정하는 경우 
#     - Backborn의 stride=8에 해당하는 layer까지 사용
#     - ResNet을 사용하는 경우 layer2까지 사용하며, stride를 8로 유지하며 layer 3, 4의 atrous  convolution으로 대체.
#     - DeepLab Head의 Atrous Spatial Pyramid Pooling(ASPP)의 dilation rate=[12, 24, 36] 사용
# - output_stride를 16로 설정하는 경우 
#     - Backborn의 stride=16에 해당하는 layer까지 사용
#     - ResNet을 사용하는 경우 layer3까지 사용하며, stride를 16로 유지하며 layer 4를 atrous  convolution으로 대체.
#     - DeepLab Head의 Atrous Spatial Pyramid Pooling(ASPP)의 dilation rate=[6, 12, 18] 사용

# + id="OyzYYks05t9D"
from network.backbone import resnet
from network import utils
from collections import OrderedDict

def ResNetbasedDeepLabV3(num_classes, output_stride, pretrained_backbone):
    if output_stride==8:
        replace_stride_with_dilation=[False, True, True]
        aspp_dilate = [12, 24, 36]
    else:
        replace_stride_with_dilation=[False, False, True]
        aspp_dilate = [6, 12, 18]
        
    backbone = resnet.resnet50(
        pretrained=pretrained_backbone,
        replace_stride_with_dilation=replace_stride_with_dilation)
    
    inplanes = 2048
    low_level_planes = 256

    return_layers = {'layer4': 'out'}
    classifier = DeepLabHead(inplanes , num_classes, aspp_dilate)
    backbone = utils.IntermediateLayerGetter(backbone, return_layers=return_layers)


    # TODO: interpolation parameter 추가
    model = DeepLabV3(backbone, classifier, 'bilinear')
    return model


# + [markdown] id="7o63GzStmvhK"
# ##(3) DeepLab Head
# - ASPP(2048ch→ 256ch) + 3x3 Conv(256ch→256ch) + 1x1 Conv(256ch→21:number of class)
#
# - Backbone + DeepLab v3 Head를 거쳐 도출된 segmentation 결과를, Bilinear upsampling을 이용하여 image-level 해상도로 복원

# + id="E9vmrVu7nv2W"
class DeepLabV3(nn.Module):

    # TODO: interpolation parameter 추가
    #def __init__(self, backbone, classifier):
    def __init__(self, backbone, classifier, upMethod = ''):
        super(DeepLabV3, self).__init__()
        self.backbone = backbone
        self.classifier = classifier
        self.upsampleMethod = upMethod
        
    def forward(self, x):
        input_shape = x.shape[-2:]
        input_ch = x.shape[1] # 256
        features = self.backbone(x)
        x = self.classifier(features)

        if self.upsampleMethod == 'CARAFE':
            # upScale =input_shape[1] / x.shape[-1]
            # print(input_ch)

            carafe =CARAFE(c=21, scale=2)
            if torch.cuda.is_available():
                carafe.cuda()
            x = carafe(x) 
            
            x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)

        else :
            x = F.interpolate(x, size=input_shape, mode='bilinear', align_corners=False)
        
        return x

class DeepLabHead(nn.Module):
    def __init__(self, in_channels, num_classes, aspp_dilate=[12, 24, 36]):
        super(DeepLabHead, self).__init__()

        self.classifier = nn.Sequential(
            ASPP(in_channels, aspp_dilate),
            nn.Conv2d(256, 256, 3, padding=1, bias=False),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, num_classes, 1)
        )
        self._init_weight()

    def forward(self, feature):
        return self.classifier( feature['out'] )

    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)



# + [markdown] id="N7LS-HGQVQou"
# ### ASPP (Atrous Spatial Pyramid Pooling)
# - 초기 Object Detection에 적용된 SPPNet의 Idea 차용 Atrous Conv.로 구성한 SPP 구성
#
#     → ConvLayer의 output을 다양한 Atrous rate를 가지는 kernel을 병렬로 연산하여 multi-scale 특징을 추출하고자 함.
#
# - Atrous rate가 커질수록 유효한 weight의 수가 작아지는 경향을 보임.
#     ⇒  Atorus rate가 커지면 3x3 kernel이 1x1처럼 동작함.
#
#     ⇒ large scale context는 output에 반영 안됨. (receptive field를 크게 만들고 싶은데 안됨..)
#
# - 이런 degenerate 문제 해결을 위해 이미지의 last feature map에 대해서 "Global average pooling"을 적용.
# - 성능 개선을 위해 각 module에 batch normalization 적용
#
# ![ASPP](https://gaussian37.github.io/assets/img/vision/segmentation/aspp/0.png)

# + [markdown] id="E4pFO0OjUTQB"
# ### Atrous Convolution (=Dilated Convolution)
#
# - Kernel의 성분들 사이에 빈 성분(0)을 삽입하여 convolution을 수행
# - Dilation rate는 빈 성분(0)을 몇 개 삽입할지 결정. 기존 일반적인 convolution의 dilation rate=1. dilation rate=2일 경우는 성분들 사이에 빈 성분(0) 1개 추가, 3일 경우는 2개 추가, ...
#
# - 기존 3x3 Kernel에 아래와 같이 dilation rate=2를 사용하면 field-of-view(receptive field)가 5x5의 영역을 커버하게 됨. 
#
# - x: input feature map, w: filter, r: dilation rate, y: output일 때 다음과 같이 표현 가능하다. 
# $$y[i]=\sum_{k}^{K}x[i+r\cdot{k}]w[k]$$
#
# ![https://miro.medium.com/max/395/1*1okwhewf5KCtIPaFib4XaA.gif](https://miro.medium.com/max/395/1*1okwhewf5KCtIPaFib4XaA.gif)
# *convolution*
# |
# ![https://miro.medium.com/max/395/1*SVkgHoFoiMZkjy54zM_SUw.gif](https://miro.medium.com/max/395/1*SVkgHoFoiMZkjy54zM_SUw.gif)*Atrous Conv.*
#
# - **Atrous convolution은 pooling 연산 없이 넓은** field of view를 커버할 수 있음.
# (기존 CNN의 경우 넓은 field of view를 커버하기 위해 conv. pooling을 사용하여 output feature map의 해상도가 감소.)
#

# + id="48poPshQUSPz"
class ASPPConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, dilation):
        modules = [
            nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        ]
        super(ASPPConv, self).__init__(*modules)


# + [markdown] id="WjdyGj9tncj-"
# - Global Average Pooling + 1x1 Conv + Upsampling
#     - Global Average Pooling (nn.AdaptiveAvgPool2d): Channel별로 average를 구하여 pooling (ex. 2048x33x33 → 2048x1x1)
#     - 1x1 Conv: Output channel 수를 256으로  (ex. 2048x1x1 → 256x1x1)
#     - Upsampling: Size를 원래대로 (ex. 256x1x1 → 256x33x33)

# + id="LcFlb6_gnNnM"
class ASPPPooling(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(ASPPPooling, self).__init__(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True))

    def forward(self, x):
        size = x.shape[-2:]
        x = super(ASPPPooling, self).forward(x)
        return F.interpolate(x, size=size, mode='bilinear', align_corners=False)


# + [markdown] id="Gvqtb5ZXnTpE"
# - 병렬화된 1x1 Conv, 3개의 3x3 Dilated Conv, Image pooling(code의 ASPP Pooling)를 concatenate
# - Concatenate된 feature를 1x1 Conv를 통해 output channel을 256으로 생성 (5x256→256)

# + id="pumb6ECinPG_"
class ASPP(nn.Module):
    def __init__(self, in_channels, atrous_rates):
        super(ASPP, self).__init__()
        out_channels = 256
        modules = []
        modules.append(nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)))

        rate1, rate2, rate3 = tuple(atrous_rates)
        modules.append(ASPPConv(in_channels, out_channels, rate1))
        modules.append(ASPPConv(in_channels, out_channels, rate2))
        modules.append(ASPPConv(in_channels, out_channels, rate3))
        modules.append(ASPPPooling(in_channels, out_channels))

        self.convs = nn.ModuleList(modules)

        self.project = nn.Sequential(
            nn.Conv2d(5 * out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1),)

    def forward(self, x):
        res = []
        for conv in self.convs:
            res.append(conv(x))
        res = torch.cat(res, dim=1)
        return self.project(res)


# + [markdown] id="LYfdrVT43EPS" endofcell="--"
# ##(4) CARAFE
# Code from: https://github.com/XiaLiPKU/CARAFE/blob/master/carafe.py#L68
#
# - 
# --

# + id="dnk8vvQN4xVI"
class ConvBNReLU(nn.Module):
    '''Module for the Conv-BN-ReLU tuple.'''
    def __init__(self, c_in, c_out, kernel_size, stride, padding, dilation,
                 use_relu=True):
        super(ConvBNReLU, self).__init__()
        self.conv = nn.Conv2d(
                c_in, c_out, kernel_size=kernel_size, stride=stride, 
                padding=padding, dilation=dilation, bias=False)
        self.bn = nn.BatchNorm2d(c_out)
        if use_relu:
            self.relu = nn.ReLU(inplace=True)
        else:
            self.relu = None

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class CARAFE(nn.Module):
    def __init__(self, c, c_mid=64, scale=2, k_up=5, k_enc=3):
        """ The unofficial implementation of the CARAFE module.
        The details are in "https://arxiv.org/abs/1905.02188".
        Args:
            c: The channel number of the input and the output.
            c_mid: The channel number after compression.
            scale: The expected upsample scale.
            k_up: The size of the reassembly kernel.
            k_enc: The kernel size of the encoder.
        Returns:
            X: The upsampled feature map.
        """
        super(CARAFE, self).__init__()
        self.scale = scale

        self.comp = ConvBNReLU(c, c_mid, kernel_size=1, stride=1, 
                               padding=0, dilation=1)
        self.enc = ConvBNReLU(c_mid, (scale*k_up)**2, kernel_size=k_enc, 
                              stride=1, padding=k_enc//2, dilation=1, 
                              use_relu=False)
        self.pix_shf = nn.PixelShuffle(scale)

        self.upsmp = nn.Upsample(scale_factor=scale, mode='nearest')
        self.unfold = nn.Unfold(kernel_size=k_up, dilation=scale, 
                                padding=k_up//2*scale)

    def forward(self, X):
        b, c, h, w = X.size()
        h_, w_ = h * self.scale, w * self.scale
        
        W = self.comp(X)                                # b * m * h * w
        W = self.enc(W)                                 # b * 100(sigma^2, k_up^2) * h * w
        W = self.pix_shf(W)                             # b * 25(k_up^2) * h_ * w_
        W = F.softmax(W, dim=1)                         # b * 25 * h_ * w_

        X = self.upsmp(X)                               # b * c * h_ * w_
        X = self.unfold(X)                              # b * 25c * h_ * w_
        X = X.view(b, c, -1, h_, w_)                    # b * 25 * c * h_ * w_

        X = torch.einsum('bkhw,bckhw->bchw', [W, X])    # b * c * h_ * w_
        return X


# + [markdown] id="a3vgfeEzqEeC"
# # 3. Dataset

# + [markdown] id="KMUH2XvqcTJ8"
# ###(1) PASCAL VOC 2012 Segmentation Task Dataset
#
# http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
# https://academictorrents.com/details/df0aad374e63b3214ef9e92e178580ce27570e59

# + id="tJFZZQFc2yDz"
# !wget -P datasets/data/ "http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar"
# !tar -xvf datasets/data/VOCtrainval_11-May-2012.tar -C ./datasets/data

# SegmentationClassAug.zip를 다운. (https://www.dropbox.com/s/oeu149j8qtbs1x0/SegmentationClassAug.zip?dl=0)
# !wget -P datasets/data/ "https://www.dropbox.com/s/oeu149j8qtbs1x0/SegmentationClassAug.zip?dl=1"
# !unzip datasets/data/SegmentationClassAug.zip -d datasets/data/VOCdevkit/VOC2012/

# + [markdown] id="htMjRT9LccCX"
# ###(2) Data augmentation

# + [markdown] id="U7MlH7HVl34i"
# - Scaling, Crop, 상하반전으로 Training data augmentation

# + id="Q-Zogp0wa3nX"
import torch
import torch.nn as nn
import torch.nn.functional as F

from datasets import VOCSegmentation, Cityscapes
from utils import ext_transforms as et

def get_dataset(opts):
    """ Dataset And Augmentation
    """
    if opts.dataset == 'voc':
        train_transform = et.ExtCompose([
            #et.ExtResize(size=opts.crop_size),
            et.ExtRandomScale((0.5, 2.0)),
            et.ExtRandomCrop(size=(opts.crop_size, opts.crop_size), pad_if_needed=True),
            et.ExtRandomHorizontalFlip(),
            et.ExtToTensor(),
            et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                            std=[0.229, 0.224, 0.225]),
        ])
        if opts.crop_val:
            val_transform = et.ExtCompose([
                et.ExtResize(opts.crop_size),
                et.ExtCenterCrop(opts.crop_size),
                et.ExtToTensor(),
                et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225]),
            ])
        else:
            val_transform = et.ExtCompose([
                et.ExtToTensor(),
                et.ExtNormalize(mean=[0.485, 0.456, 0.406],
                                std=[0.229, 0.224, 0.225]),
            ])
        train_dst = VOCSegmentation(root=opts.data_root, year=opts.year,
                                    image_set='train', transform=train_transform)
        val_dst = VOCSegmentation(root=opts.data_root, year=opts.year,
                                  image_set='val', transform=val_transform)
    
    return train_dst, val_dst


# + [markdown] id="3paRH68cENLS"
# #4. Train/Test
#

# + [markdown] id="tVkLeySIrfQU"
# ### (1) poly learning rate policy
#
# - 초반엔 선형 감소에 가깝지만, Training step의 마지막에서는 조금 더 가파르게 감소하는 경향성을 나타냄
# $$learning\_rate = \left( 1-\frac{iter}{max\_iter}\right)^{power}$$

# + id="8F-BsVwkMyGS"
from torch.optim.lr_scheduler import _LRScheduler, StepLR

class PolyLR(_LRScheduler):
    def __init__(self, optimizer, max_iters, power=0.9, last_epoch=-1, min_lr=1e-6):
        self.power = power
        self.max_iters = max_iters  # avoid zero lr
        self.min_lr = min_lr
        super(PolyLR, self).__init__(optimizer, last_epoch)
    
    def get_lr(self):
        return [ max( base_lr * ( 1 - self.last_epoch/self.max_iters )**self.power, self.min_lr)
                for base_lr in self.base_lrs]


# + [markdown] id="sxOD7cMcrpVm"
# ###(2) Focal Loss
# - ignore_index: 특정 target value를 무시하며, gradient 계산 할 때 포함시키지 않는다. size_average가 true일 경우 ignore_index되어 있는 target을 제외하고 평균을 계산한다.
#

# + id="qO6caE8yM1fy"
import torch.nn as nn
import torch.nn.functional as F
import torch 

class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=0, size_average=True, ignore_index=255):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.ignore_index = ignore_index
        self.size_average = size_average

    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(
            inputs, targets, reduction='none', ignore_index=self.ignore_index)
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss
        if self.size_average:
            return focal_loss.mean()
        else:
            return focal_loss.sum()


# + id="lY9cLh7vKE48"
import torch.nn as nn
from PIL import Image

def set_bn_momentum(model, momentum=0.1):
    for m in model.modules():
        if isinstance(m, nn.BatchNorm2d):
            m.momentum = momentum

def save_ckpt(path):
    """ save current model
    """
    torch.save({
        "cur_itrs": cur_itrs,
        "model_state": model.module.state_dict(),
        "optimizer_state": optimizer.state_dict(),
        "scheduler_state": scheduler.state_dict(),
        "best_score": best_score,
    }, path)
    print("Model saved as %s" % path)

class Denormalize(object):
    def __init__(self, mean, std):
        mean = np.array(mean)
        std = np.array(std)
        self._mean = -mean/std
        self._std = 1/std

    def __call__(self, tensor):
        if isinstance(tensor, np.ndarray):
            return (tensor - self._mean.reshape(-1,1,1)) / self._std.reshape(-1,1,1)
        return normalize(tensor, self._mean, self._std)


# + [markdown] id="aE4DxXTPizUV"
# ###(3) Validation / Performance Measure

# + id="oRTzZ96Ma6Yl"
from tqdm import tqdm
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image

import cv2
import matplotlib
import matplotlib.pyplot as plt


def validate(opts, model, loader, device, metrics):
    """Do validation and return specified samples"""
    metrics.reset()
    ret_samples = []
    if opts.save_val_results:
        if not os.path.exists('results'):
            os.mkdir('results')
        denorm = Denormalize(mean=[0.485, 0.456, 0.406], 
                              std=[0.229, 0.224, 0.225])
        img_id = 0

    with torch.no_grad():
        for i, (images, labels) in tqdm(enumerate(loader)):
            
            images = images.to(device, dtype=torch.float32)
            labels = labels.to(device, dtype=torch.long)

            outputs = model(images)
            preds = outputs.detach().max(dim=1)[1].cpu().numpy()
            targets = labels.cpu().numpy()

            metrics.update(targets, preds)

            if opts.save_val_results:
                for i in range(len(images)):
                    image = images[i].detach().cpu().numpy()
                    target = targets[i]
                    pred = preds[i]

                    image = (denorm(image) * 255).transpose(1, 2, 0).astype(np.uint8)
                    target = loader.dataset.decode_target(target).astype(np.uint8)
                    pred = loader.dataset.decode_target(pred).astype(np.uint8)

                    Image.fromarray(image).save('results/%d_image.png' % img_id)
                    Image.fromarray(target).save('results/%d_target.png' % img_id)
                    Image.fromarray(pred).save('results/%d_pred.png' % img_id)

                    fig = plt.figure()
                    plt.ion()
                    plt.imshow(image)
                    plt.axis('off')
                    plt.imshow(pred, alpha=0.7)
                    ax = plt.gca()
                    ax.xaxis.set_major_locator(matplotlib.ticker.NullLocator())
                    ax.yaxis.set_major_locator(matplotlib.ticker.NullLocator())
                    plt.savefig('results/%d_overlay.png' % img_id, bbox_inches='tight', pad_inches=0)
                    plt.close()

                    #fig2 = plt.figure()
                    if opts.test_only:
                      imgBGR = cv2.imread('results/%d_image.png' % img_id)

                      # cv2.imread는 BGR로 불러오므로 plt를 이용하려면 RGB로 바꿔줘야 함
                      imgTargetBGR = cv2.imread('results/%d_target.png' % img_id)
                      imgPredBGR = cv2.imread('results/%d_pred.png' % img_id)
                      imgOverBGR = cv2.imread('results/%d_overlay.png' % img_id)
                      
                      fig2 = plt.figure()
                      plt.ion()
                      ax = fig2.add_subplot(1,4,1)
                      plt.axis('off')
                      imgplot = plt.imshow(imgBGR)
                      ax.set_title('input')
                      ax = fig2.add_subplot(1,4,2)
                      plt.axis('off')
                      imgplot = plt.imshow(imgTargetBGR)
                      ax.set_title('GT')
                      ax = fig2.add_subplot(1,4,3)
                      plt.axis('off')
                      imgplot = plt.imshow(imgPredBGR)
                      ax.set_title('Seg.Mask')
                      ax = fig2.add_subplot(1,4,4)
                      plt.axis('off')
                      imgplot = plt.imshow(imgOverBGR)
                      ax.set_title('input with mask')                        
                      plt.show()
                      plt.close()

                    img_id += 1

        score = metrics.get_results()
    return score, ret_samples


# + [markdown] id="rnPKZT2Wr_Th"
# ##(4) Run

# + id="0W4FAzZoeLGB"
class Config():
    # TODO : ADD compound scaling factor
    def __init__(self):
        super(Config, self).__init__()

        self.data_root='./datasets/data'
        self.dataset ='voc'
        self.download =False
        self.year ='2012' # VOC

        self.num_classes = 21
        self.model ='deeplabv3_resnet50_CARAFE'
        self.output_stride =16

        self.pretrained_backbone = True

        # train option
        self.test_only =False
        self.save_val_results =True
        self.total_itrs =30e3
        self.lr =0.01
        self.lr_policy ='poly' # 'step'
        self.step_size =10000
        self.crop_val =False
        self.batch_size =16
        self.val_batch_size =4
        self.crop_size =513
        self.ckpt ='checkpoints/latest_deeplabv3_resnet50_CARAFE_voc_os16.pth'
        self.continue_training =False

        self.loss_type = 'cross_entropy' #'cross_entropy', 'focal_loss'
        self.gpu_id ='0'
        self.weight_decay =1e-4
        self.random_seed =1
        self.print_interval =10
        self.val_interval =100


# + id="6JORB0j-csjm"
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import data

from metrics import StreamSegMetrics

import os
import numpy as np
import random

# Load option
opts =Config()

os.environ['CUDA_VISIBLE_DEVICES'] = opts.gpu_id
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device: %s" % device)

# Setup random seed
torch.manual_seed(opts.random_seed)
np.random.seed(opts.random_seed)
random.seed(opts.random_seed)


# Set dataloader
if opts.dataset.lower() == 'voc':
    opts.num_classes = 21
#elif opts.dataset.lower() == 'cityscapes':
#    opts.num_classes = 19

if opts.dataset=='voc' and not opts.crop_val:
    opts.val_batch_size = 1

train_dst, val_dst = get_dataset(opts)
train_loader = data.DataLoader(
    train_dst, batch_size=opts.batch_size, shuffle=True, num_workers=2)
val_loader = data.DataLoader(
    val_dst, batch_size=opts.val_batch_size, shuffle=True, num_workers=2)
print("Dataset: %s, Train set: %d, Val set: %d" %
      (opts.dataset, len(train_dst), len(val_dst)))

# Set up model
model  = ResNetbasedDeepLabV3(num_classes=opts.num_classes, output_stride=opts.output_stride, pretrained_backbone=opts.pretrained_backbone)
set_bn_momentum(model.backbone, momentum=0.01)

# Set up optimizer
optimizer = torch.optim.SGD(params=[
    {'params': model.backbone.parameters(), 'lr': 0.1*opts.lr},
    {'params': model.classifier.parameters(), 'lr': opts.lr},
], lr=opts.lr, momentum=0.9, weight_decay=opts.weight_decay)

if opts.lr_policy=='poly':
    scheduler = PolyLR(optimizer, opts.total_itrs, power=0.9)
elif opts.lr_policy=='step':
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=opts.step_size, gamma=0.1)

# Set up criterion
if opts.loss_type == 'focal_loss':
    criterion = FocalLoss(ignore_index=255, size_average=True)
elif opts.loss_type == 'cross_entropy':
    criterion = nn.CrossEntropyLoss(ignore_index=255, reduction='mean')

# Set up metrics
metrics = StreamSegMetrics(opts.num_classes)

checkPointPath = 'checkpoints'
if not os.path.exists(checkPointPath):
    os.mkdir(checkPointPath)
# Restore
best_score = 0.0
cur_itrs = 0
cur_epochs = 0

if opts.ckpt is not None and os.path.isfile(opts.ckpt):
    # https://github.com/VainF/DeepLabV3Plus-Pytorch/issues/8#issuecomment-605601402, @PytaichukBohdan
    checkpoint = torch.load(opts.ckpt, map_location=torch.device('cuda'))
    model.load_state_dict(checkpoint["model_state"])
    model = nn.DataParallel(model)
    model.to(device)
    if opts.continue_training:
        optimizer.load_state_dict(checkpoint["optimizer_state"])
        scheduler.load_state_dict(checkpoint["scheduler_state"])
        cur_itrs = checkpoint["cur_itrs"]
        best_score = checkpoint['best_score']
        print("Training state restored from %s" % opts.ckpt)
    print("Model restored from %s" % opts.ckpt)
    del checkpoint  # free memory
else:
    print("[!] Retrain")
    model = nn.DataParallel(model)
    model.to(device)

denorm = Denormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # denormalization for ori images

if opts.test_only:
    model.eval()
    val_score, ret_samples = validate(opts=opts, model=model, loader=val_loader, device=device, metrics=metrics)
    print(metrics.to_str(val_score))

else:
  interval_loss = 0
  while True: #cur_itrs < opts.total_itrs:
      # =====  Train  =====
      model.train()
      cur_epochs += 1
      for (images, labels) in train_loader:
          cur_itrs += 1

          images = images.to(device, dtype=torch.float32)
          labels = labels.to(device, dtype=torch.long)

          optimizer.zero_grad()
          outputs = model(images)
          loss = criterion(outputs, labels)
          loss.backward()
          optimizer.step()

          np_loss = loss.detach().cpu().numpy()
          interval_loss += np_loss

          if (cur_itrs) % 10 == 0:
              interval_loss = interval_loss/10
              print("Epoch %d, Itrs %d/%d, Loss=%f" %
                    (cur_epochs, cur_itrs, opts.total_itrs, interval_loss))
              interval_loss = 0.0

          if (cur_itrs) % opts.val_interval == 0:
              save_ckpt('checkpoints/latest_%s_%s_os%d.pth' %
                        (opts.model, opts.dataset, opts.output_stride))
              
              print("validation...")
              
              model.eval()
              val_score, ret_samples = validate(
                  opts=opts, model=model, loader=val_loader, device=device, metrics=metrics)
              print(metrics.to_str(val_score))
              if val_score['Mean IoU'] > best_score:  # save best model
                  best_score = val_score['Mean IoU']
                  save_ckpt('checkpoints/best_%s_%s_os%d.pth' %
                            (opts.model, opts.dataset,opts.output_stride))

              model.train()
          scheduler.step()  

          #if cur_itrs >=  opts.total_itrs:
              #return


StarcoderdataJupyterScriptsDedupFiltered