train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2ba33d3d-c17f-4019-94df-347bc4bb8c05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/opt/conda/bin/python\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "print(sys.executable)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7b83b91d-f36a-4603-836d-dc6bc1df4037",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import skimage.io as io\n",
    "import random\n",
    "import os\n",
    "import albumentations as A\n",
    "import cv2\n",
    "import pickle\n",
    "import torch\n",
    "from pycocotools.coco import COCO\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from albumentations.pytorch import ToTensorV2\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2b1478c9-3bbf-4826-b9a1-abc97cbe271d",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "bb7ae67f-8168-4427-8ed4-fb86cc770430",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "def param2theta(param, w, h):\n",
    "    param = np.linalg.inv(param)\n",
    "    theta = np.zeros([2,3])\n",
    "    theta[0,0] = param[0,0]\n",
    "    theta[0,1] = param[0,1]*h/w\n",
    "    theta[0,2] = param[0,2]*2/w + theta[0,0] + theta[0,1] - 1\n",
    "    theta[1,0] = param[1,0]*w/h\n",
    "    theta[1,1] = param[1,1]\n",
    "    theta[1,2] = param[1,2]*2/h + theta[1,0] + theta[1,1] - 1\n",
    "    return theta\n",
    "\n",
    "class MSCOCODataset(Dataset):\n",
    "    def __init__(self, args=None, aug_trans=None, basic_trans=None):\n",
    "        self.args = args\n",
    "        self.images = sorted(glob.glob(args.data_dir + '/*'))\n",
    "        self.ann_file = '{}/annotations/instances_{}.json'.format(args.ann_dir, args.mode)\n",
    "        self.coco = COCO(self.ann_file)\n",
    "        self.aug_trans = aug_trans\n",
    "        self.basic_trans = basic_trans\n",
    "        self.img_size = args.input_size\n",
    "        self.center = (self.img_size/2, self.img_size/2)\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.images)//4\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        annId = int(self.images[idx*4].split('_')[0].split('/')[-1])\n",
    "        # print('annId:', annId)\n",
    "        # file_name, i, img_type = self.imxages[idx*4].split('_')\n",
    "        # print(self.images[idx*4], self.images[idx*4+1], self.images[idx*4+2], self.images[idx*4+3])\n",
    "    \n",
    "        bg = self.getImage(self.images[idx*4])\n",
    "        fg = self.getImage(self.images[idx*4+1])\n",
    "        gt = self.getImage(self.images[idx*4+2])\n",
    "        mask = self.getMask(self.images[idx*4+3])\n",
    "      \n",
    "        ##########\n",
    "\n",
    "        angle_factor = random.randint(-10, 10)\n",
    "        scale_factor = round(random.uniform(0.95, 1.05), 2)\n",
    "        translation_factor_x = round(random.uniform(0, self.img_size*0.05), 2)\n",
    "        translation_factor_y = round(random.uniform(0, self.img_size*0.05), 2)\n",
    "        rotate_matrix = cv2.getRotationMatrix2D(center=self.center, angle=angle_factor, scale=0.95)\n",
    "\n",
    "        new_row = np.array([[0, 0, 1]], dtype=np.float32)\n",
    "        rotate_matrix = np.concatenate((rotate_matrix, new_row), axis=0)\n",
    "        translation_matrix = np.array([[1, 0, translation_factor_x],\n",
    "                                       [0, 1, translation_factor_y],\n",
    "                                       [0, 0, 1],], dtype=np.float32)\n",
    "\n",
    "        matrix = np.matmul(rotate_matrix, translation_matrix)[:2][:3]\n",
    "        matrix = np.concatenate((matrix, new_row), axis=0)\n",
    "        matrix = param2theta(matrix, self.img_size, self.img_size)\n",
    "\n",
    "        # matrix = np.array([[1, 0, 1],\n",
    "        #                 [0, 1, 0],], dtype=np.float32)\n",
    "        # new_row = np.array([[0, 0, 1]], dtype=np.float32)\n",
    "        # matrix = np.concatenate((matrix, new_row), axis=0)\n",
    "        # matrix = param2theta(matrix, self.img_size, self.img_size)\n",
    "\n",
    "        tf_fg = cv2.warpAffine(src=fg, M=matrix, dsize=(self.img_size, self.img_size))\n",
    "        tf_mask = cv2.warpAffine(src=mask, M=matrix, dsize=(self.img_size, self.img_size))\n",
    "        \n",
    "        ###########\n",
    "        \n",
    "        bg = self.basic_trans(image=bg)['image']\n",
    "        gt = self.basic_trans(image=gt)['image']\n",
    "        \n",
    "        tf_item = self.aug_trans(image=tf_fg, mask=tf_mask) #only color aug applied\n",
    "        tf_fg = tf_item['image']\n",
    "        tf_mask = tf_item['mask'].unsqueeze(0)\n",
    "    \n",
    "        tf_item = self.basic_trans(image=fg, mask=mask)\n",
    "        fg = tf_item['image']\n",
    "        mask = tf_item['mask'].unsqueeze(0)\n",
    "        \n",
    "        return gt, fg, bg, mask>0.5, tf_fg, tf_mask>0.5, matrix#, cat\n",
    "    \n",
    "    def getClassName(self, classID, cats):\n",
    "        for i in range(len(cats)):\n",
    "            if cats[i]['id'] == classID:\n",
    "                return cats[i]['name']\n",
    "        return \"None\"\n",
    "    \n",
    "    def getImage(self, file_name):\n",
    "        img = cv2.imread(file_name, cv2.IMREAD_COLOR)\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "        return img\n",
    "    \n",
    "    def getMask(self, file_name):\n",
    "        mask = cv2.imread(file_name, cv2.IMREAD_GRAYSCALE)\n",
    "        return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "73a72737-74a7-4729-a502-1f4896f465ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "\n",
    "class GenCompModel(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super(GenCompModel, self).__init__()\n",
    "        self.args = args\n",
    "        self.img_size = args.input_size\n",
    "        self.batch_size = args.batch_size\n",
    "        self.stn = STN(7) # fg+bg+mask channels\n",
    "        self.colornet = LinearWithChannel(self.batch_size, self.img_size, self.img_size, 3) \n",
    "        self.refinenet = TransformNetwork(3, 3)\n",
    "    \n",
    "    def forward(self, fg, bg, mask):\n",
    "        HI, AI, trans_mat = self.stn(fg, bg, mask)\n",
    "        FI = self.colornet(HI) # changed to inputs as I\n",
    "        AI = AI > 0.5\n",
    "        R_in = torch.multiply(FI, AI) + torch.multiply(bg, ~AI)\n",
    "        R_out = self.refinenet(R_in)\n",
    "        return AI, HI, FI, trans_mat, R_in, R_out\n",
    "    \n",
    "class Discriminator(nn.Module):    \n",
    "    def __init__(self):        \n",
    "        super(Discriminator, self).__init__()\n",
    "        self.imgdisc = ImageDiscriminator(3, norm='spec')\n",
    "        self.segnet = TransformNetwork(3, 1)\n",
    "        \n",
    "    def forward(self, mask, R_out):\n",
    "        img_out = self.imgdisc(R_out)\n",
    "        fg_seg_out = self.segnet(torch.multiply(mask, R_out), last='sigmoid')\n",
    "        bg_seg_out = self.segnet(torch.multiply(~mask, R_out), last='sigmoid')\n",
    "\n",
    "        return img_out, fg_seg_out, bg_seg_out\n",
    "    \n",
    "class TransformNetwork(nn.Module):\n",
    "    def __init__(self, in_ch, out_ch):        \n",
    "        super(TransformNetwork, self).__init__()        \n",
    "        \n",
    "        self.layers = nn.Sequential(            \n",
    "            ConvLayer(in_ch, 32, 9, 1),\n",
    "            ConvLayer(32, 64, 3, 2),\n",
    "            ConvLayer(64, 128, 3, 2),\n",
    "            \n",
    "            ResidualLayer(128, 128, 3, 1),\n",
    "            ResidualLayer(128, 128, 3, 1),\n",
    "            ResidualLayer(128, 128, 3, 1),\n",
    "            ResidualLayer(128, 128, 3, 1),\n",
    "            ResidualLayer(128, 128, 3, 1),\n",
    "            \n",
    "            DeconvLayer(128, 64, 3, 1),\n",
    "            DeconvLayer(64, 32, 3, 1),\n",
    "            ConvLayer(32, out_ch, 9, 1, activation='linear'))\n",
    "        \n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "        \n",
    "    def forward(self, x, last=None):\n",
    "        x = self.layers(x)\n",
    "        if last:\n",
    "            x = self.sigmoid(x)\n",
    "        return x\n",
    "    \n",
    "class LinearWithChannel(nn.Module):\n",
    "    def __init__(self, batch_size, input_size, output_size, channel_size):\n",
    "        super(LinearWithChannel, self).__init__()\n",
    "        self.w = torch.nn.Parameter(torch.empty(channel_size))\n",
    "        self.b = torch.nn.Parameter(torch.zeros(channel_size))\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return x * self.w.view(1, 3, 1, 1) + self.b.view(1, 3, 1, 1)\n",
    "    \n",
    "class STN(nn.Module):\n",
    "    def __init__(self, in_ch):\n",
    "        super(STN, self).__init__()\n",
    "        self.in_ch = in_ch\n",
    "        \n",
    "        # localization-network for STN\n",
    "        self.localization = nn.Sequential(\n",
    "            ConvLayer(self.in_ch, 16, 3, 1), # 256\n",
    "            ConvLayer(16, 16, 3, 1),\n",
    "            ConvLayer(16, 16*2, 3, 2), # 128\n",
    "            ConvLayer(16*2, 16*2, 3, 1),\n",
    "            ConvLayer(16*2, 16*4, 3, 2), # 64\n",
    "            ConvLayer(16*4, 16*4, 3, 1),\n",
    "            ConvLayer(16*4, 16*8, 3, 2), # 32\n",
    "            ConvLayer(16*8, 16*8, 3, 1), \n",
    "            ConvLayer(16*8, 16*16, 3, 2), # 16\n",
    "            ConvLayer(16*16, 16*16, 3, 1), \n",
    "            ConvLayer(16*16, 16*32, 3, 2), # 8\n",
    "            ConvLayer(16*32, 16*32, 3, 1), \n",
    "        )\n",
    "\n",
    "        # [3 * 2] 크기의 아핀(affine) 행렬에 대해 예측\n",
    "        self.fc_loc = nn.Sequential(\n",
    "            nn.Linear(16*32*8*8, 32*8*8),\n",
    "            nn.ReLU(True),\n",
    "            nn.Linear(32*8*8, 8*8),\n",
    "            nn.ReLU(True),\n",
    "            nn.Linear(8*8, 16),\n",
    "            nn.ReLU(True),\n",
    "            nn.Linear(16, 2*3),\n",
    "        )\n",
    "\n",
    "        # 항등 변환(identity transformation)으로 가중치/바이어스 초기화\n",
    "        self.fc_loc[6].weight.data.zero_()\n",
    "        self.fc_loc[6].bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float))\n",
    "\n",
    "    # STN의 forward 함수\n",
    "    def stn(self, fg, bg, mask):\n",
    "        # x = x.type(torch.cuda.FloatTensor)\n",
    "        mask = mask.float()\n",
    "        inputs = torch.cat([fg, bg, mask], dim=1) #[B, (3+3+1), 256, 256]\n",
    "        xs = self.localization(inputs)\n",
    "        # print('xs shape:', xs.shape) #torch.Size([4, 10, 60, 60]) # [1, 128, 32, 32])\n",
    "        xs = xs.view(-1, 16*32*8*8) #xs = xs.view(-1, 10 * 3 * 3)\n",
    "        theta = self.fc_loc(xs)\n",
    "        theta = theta.view(-1, 2, 3)\n",
    "\n",
    "        grid = F.affine_grid(theta, fg.size(), align_corners=False)\n",
    "        fg = F.grid_sample(fg, grid, align_corners=False)\n",
    "\n",
    "        grid = F.affine_grid(theta, mask.size(), align_corners=False)\n",
    "        mask = F.grid_sample(mask, grid, align_corners=False)\n",
    "        \n",
    "        # return fg, mask>0.5, theta\n",
    "        return fg, mask, theta\n",
    "\n",
    "    def forward(self, fg, bg, mask):\n",
    "        # 입력을 변환\n",
    "        fg, mask, trans_mat = self.stn(fg, bg, mask)\n",
    "        return fg, mask, trans_mat\n",
    "    \n",
    "class ConvLayer(nn.Module):    \n",
    "    def __init__(self, in_ch, out_ch, kernel_size, stride, pad='reflect', activation='leaky', normalization='batch'):        \n",
    "        super(ConvLayer, self).__init__()\n",
    "        \n",
    "        # padding\n",
    "        if pad == 'reflect':            \n",
    "            self.pad = nn.ReflectionPad2d(kernel_size//2)\n",
    "        elif pad == 'zero':\n",
    "            self.pad = nn.ZeroPad2d(kernel_size//2)\n",
    "        else:\n",
    "            raise NotImplementedError(\"Not expected pad flag !!!\")\n",
    "    \n",
    "            \n",
    "        # convolution\n",
    "        self.conv_layer = nn.Conv2d(in_ch, out_ch, \n",
    "                                    kernel_size=kernel_size,\n",
    "                                    stride=stride)\n",
    "        if normalization == 'spec':\n",
    "            self.conv_layer = nn.utils.spectral_norm(self.conv_layer)\n",
    "           \n",
    "        \n",
    "        # activation\n",
    "        if activation == 'relu':\n",
    "            self.activation = nn.ReLU()     \n",
    "        elif activation == 'leaky':\n",
    "            self.activation = nn.LeakyReLU(0.2)\n",
    "        elif activation == 'linear':\n",
    "            self.activation = lambda x : x\n",
    "\n",
    "        else:\n",
    "            raise NotImplementedError(\"Not expected activation flag !!!\")\n",
    "\n",
    "        # normalization \n",
    "        if normalization == 'instance':            \n",
    "            self.normalization = nn.InstanceNorm2d(out_ch, affine=True)\n",
    "        elif normalization == 'batch':\n",
    "            self.normalization = nn.BatchNorm2d(out_ch, affine=True)\n",
    "        elif normalization == 'spec':\n",
    "            self.normalization = None\n",
    "        else:\n",
    "            raise NotImplementedError(\"Not expected normalization flag !!!\")\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.pad(x)\n",
    "        x = self.conv_layer(x)\n",
    "        if self.normalization:\n",
    "            x = self.normalization(x)\n",
    "        x = self.activation(x)        \n",
    "        return x\n",
    "    \n",
    "class ResidualLayer(nn.Module):    \n",
    "    def __init__(self, in_ch, out_ch, kernel_size, stride, pad='reflect', normalization='batch'):        \n",
    "        super(ResidualLayer, self).__init__()\n",
    "        \n",
    "        self.conv1 = ConvLayer(in_ch, out_ch, kernel_size, stride, pad, \n",
    "                               activation='relu', \n",
    "                               normalization=normalization)\n",
    "        \n",
    "        self.conv2 = ConvLayer(out_ch, out_ch, kernel_size, stride, pad, \n",
    "                               activation='linear', \n",
    "                               normalization=normalization)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        y = self.conv1(x)\n",
    "        return self.conv2(y) + x\n",
    "    \n",
    "class DeconvLayer(nn.Module):    \n",
    "    def __init__(self, in_ch, out_ch, kernel_size, stride, pad='reflect', activation='leaky', normalization='batch', upsample='nearest'):        \n",
    "        super(DeconvLayer, self).__init__()\n",
    "        \n",
    "        # upsample\n",
    "        self.upsample = upsample\n",
    "        \n",
    "        # pad\n",
    "        if pad == 'reflect':            \n",
    "            self.pad = nn.ReflectionPad2d(kernel_size//2)\n",
    "        elif pad == 'zero':\n",
    "            self.pad = nn.ZeroPad2d(kernel_size//2)\n",
    "        else:\n",
    "            raise NotImplementedError(\"Not expected pad flag !!!\")        \n",
    "        \n",
    "        # conv\n",
    "        self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, stride)\n",
    "        \n",
    "        # activation\n",
    "        if activation == 'relu':\n",
    "            self.activation = nn.ReLU()\n",
    "        elif activation == 'leaky':\n",
    "            self.activation = nn.LeakyReLU(0.2)\n",
    "        elif activation == 'linear':\n",
    "            self.activation = lambda x : x\n",
    "        else:\n",
    "            raise NotImplementedError(\"Not expected activation flag !!!\")\n",
    "        \n",
    "        # normalization\n",
    "        if normalization == 'instance':\n",
    "            self.normalization = nn.InstanceNorm2d(out_ch, affine=True)\n",
    "        elif normalization == 'batch':\n",
    "            self.normalization = nn.BatchNorm2d(out_ch, affine=True)\n",
    "        else:\n",
    "            raise NotImplementedError(\"Not expected normalization flag !!!\")\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = nn.functional.interpolate(x, scale_factor=2, mode=self.upsample)        \n",
    "        x = self.pad(x)\n",
    "        x = self.conv(x)\n",
    "        x = self.normalization(x)        \n",
    "        x = self.activation(x)        \n",
    "        return x\n",
    "    \n",
    "class ImageDiscriminator(nn.Module):\n",
    "    def __init__(self, in_ch, norm='spec'):\n",
    "        super(ImageDiscriminator, self).__init__()\n",
    "        self.layers = nn.Sequential(\n",
    "            ConvLayer(in_ch, 64, 10, 4, normalization=norm, activation='leaky'), #256->64\n",
    "            ConvLayer(64, 128, 10, 4, normalization=norm, activation='leaky'), #64->16\n",
    "            ConvLayer(128, 256, 10, 4, normalization=norm, activation='leaky'), #16->4\n",
    "            nn.Conv2d(256, 1, 5, 1), #4->1\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "    def forward(self, input):\n",
    "        return self.layers(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "98721815-825f-4c38-ab82-cf7bfad03c7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "\n",
    "args = argparse.Namespace(\n",
    "    ann_dir = '/shared/data/COCOdataset2017',\n",
    "    data_dir = '/shared/GCCdataset/alltypes',\n",
    "    save_model_dir = '/shared/GCC-GAN-server/models/',\n",
    "    mode = 'train', # or 'test'\n",
    "    batch_size = 8,\n",
    "    input_size = 256,\n",
    "    epochs = 8,\n",
    "    lr = 2e-5,\n",
    "    lambda_g = 1,\n",
    "    lambda_c = 1,\n",
    "    lambda_a = 0.01,\n",
    "    lambda_s = 0.01,\n",
    "    lambda_s2 = 1,\n",
    "    beta = 0.5,\n",
    "    test_interval = 50,\n",
    "    device_id = 1,#[0, 1, 2, 3]\n",
    "    vis_id = 'HCRS3'# feature_dim = ,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4832f60e-6d2a-463e-9827-d263ebcd4bab",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device: cuda:1\n",
      "# Workers: 8\n",
      "loading annotations into memory...\n",
      "Done (t=20.92s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "import torch.distributed as dist\n",
    "from apex.parallel import DistributedDataParallel as DDP\n",
    "\n",
    "device = 'cuda:'+str(args.device_id) if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "print('Device:', device)\n",
    "\n",
    "aug_transform = A.Compose([\n",
    "        A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),\n",
    "        # A.OneOf([\n",
    "        #     A.ShiftScaleRotate(rotate_limit=10, p=0.5, border_mode=cv2.BORDER_CONSTANT),\n",
    "        #     A.geometric.transforms.Affine(translate_percent=0.05, shear=5, scale = 0.95),\n",
    "        # ], p=1.0),\n",
    "        A.Normalize(\n",
    "            mean=[0.485, 0.456, 0.406],\n",
    "            std=[0.229, 0.224, 0.225],),\n",
    "        ToTensorV2(),\n",
    "])\n",
    "\n",
    "basic_transform = A.Compose([\n",
    "        A.Normalize(\n",
    "            mean=[0.485, 0.456, 0.406],\n",
    "            std=[0.229, 0.224, 0.225],),\n",
    "        ToTensorV2(),\n",
    "])\n",
    "\n",
    "\n",
    "num_workers = 8 # 4 * len(args.device_ids)\n",
    "print('# Workers:', num_workers)\n",
    "\n",
    "dataset = MSCOCODataset(args, aug_trans=aug_transform, basic_trans=basic_transform)\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True, num_workers=num_workers, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9cb52ac8-47b4-4db3-83a1-e03be2693ff5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting up a new session...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[7 epoch 84750/84835]: c:0.000010, g:0.005206, dsc_loss:1.053176, a:0.645480, s:1.064361, total:0.311907\r"
     ]
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "from torch.optim import Adam\n",
    "import visdom\n",
    "import os\n",
    "# from multiprocessing import Pool\n",
    "\n",
    "def denorm(tensor):\n",
    "    mean = torch.Tensor([0.485, 0.456, 0.406]).reshape(-1, 1, 1)\n",
    "    std = torch.Tensor([0.229, 0.224, 0.225]).reshape(-1, 1, 1)\n",
    "    res = torch.clamp(tensor * std + mean, 0, 1)\n",
    "    return res\n",
    "\n",
    "def tensor2im(image_tensor, imtype=np.uint8):\n",
    "    image_numpy = image_tensor.numpy() #image_tensor[0].permute(1,2,0).detach().cpu().float().numpy()\n",
    "    # image_numpy = image_numpy * 0.5 - 0.5\n",
    "    image_numpy = image_numpy * 255\n",
    "    # image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0\n",
    "    # image_numpy = (image_numpy + 1) / 2.0 * 255.0\n",
    "    image_numpy = np.clip(image_numpy, 0, 255)\n",
    "\n",
    "    return image_numpy.astype(imtype)\n",
    "\n",
    "def weights_init(m):\n",
    "    if isinstance(m, nn.Conv2d):\n",
    "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif isinstance(m, nn.BatchNorm2d):\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0)\n",
    "        \n",
    "def save_checkpoint(epoch, model, optimizer, filename):\n",
    "    state = {\n",
    "        'Epoch': epoch,\n",
    "        'State_dict': model.state_dict(),\n",
    "        'optimizer': optimizer.state_dict()\n",
    "    }\n",
    "    torch.save(state, filename)\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    # Visdom display initialization\n",
    "    vis = visdom.Visdom(env=args.vis_id)\n",
    "    vis.close(env=args.vis_id)\n",
    "    plot = vis.line(Y=torch.Tensor(1).zero_(), opts=dict(title='l2norm'), env=args.vis_id) #final loss\n",
    "    plot2 = vis.line(Y=torch.Tensor(1).zero_(), opts=dict(title='Losses', legend=[\"Loss_a\", \"Loss_s\", \"Loss_c\", \"Loss_g\"], showlegend=True), env=args.vis_id)\n",
    "    plot3 = vis.line(Y=torch.Tensor(1).zero_(), opts=dict(title='GAN losses', legend=[\"Discriminator_real\", \"Discriminator_fake\", \"Generator\"], showlegend=True), env=args.vis_id)\n",
    "\n",
    "    vis_inputs = vis.images(np.random.rand(1, 1,256,256), opts=dict(title=\"Inputs\"), env=args.vis_id)\n",
    "    vis_FI = vis.images(np.random.rand(1,3,256,256), opts=dict(title=\"F(I)\"), env=args.vis_id)\n",
    "    vis_HI = vis.images(np.random.rand(1,3,256,256), opts=dict(title=\"H(I)\"), env=args.vis_id)\n",
    "    vis_AI = vis.images(np.random.rand(1,1,256,256), opts=dict(title=\"A(I)\"), env=args.vis_id)\n",
    "    vis_R_in = vis.images(np.random.rand(1,3,256,256), opts=dict(title=\"R_in\"), env=args.vis_id)\n",
    "    vis_R_out = vis.images(np.random.rand(1,3,256,256), opts=dict(title=\"Composite\"), env=args.vis_id)\n",
    "    vis_seg = vis.images(np.random.rand(1,1,256,256), opts=dict(title=\"Segmentation\"), env=args.vis_id)\n",
    "\n",
    "    l1_loss = nn.L1Loss().to(device)\n",
    "    l2_loss = nn.MSELoss().to(device)\n",
    "    criterion = nn.BCEWithLogitsLoss().to(device)\n",
    "    \n",
    "    net = GenCompModel(args) # generator\n",
    "    net.apply(weights_init)\n",
    "    \n",
    "    disc = Discriminator()\n",
    "    disc.apply(weights_init)\n",
    "    \n",
    "    net.to(device)\n",
    "    disc.to(device)\n",
    "    \n",
    "    net.train()\n",
    "    disc.train()\n",
    "\n",
    "    real_label = 1.\n",
    "    fake_label = 0.\n",
    "\n",
    "    optimizerD = Adam(disc.parameters(), lr=args.lr, betas=(args.beta, 0.999))\n",
    "    optimizer = Adam(net.parameters(), lr=args.lr, betas=(args.beta, 0.999))\n",
    "\n",
    "    SMOOTH = 1e-6\n",
    "    \n",
    "    for epoch in range(args.epochs):\n",
    "                \n",
    "        for i, items in enumerate(dataloader):\n",
    "            ############################\n",
    "            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
    "            ###########################\n",
    "            ########## Train with all-real batch\n",
    "            gt = items[0].to(device); fg = items[1].to(device);\n",
    "            bg = items[2].to(device); mask = items[3].to(device);\n",
    "            tf_fg = items[4].to(device); tf_mask = items[5].to(device);\n",
    "            trans_mat = items[6].to(device) # cat = items[6].to(device);\n",
    "            B, C, H, W = gt.shape\n",
    "\n",
    "            real_value = torch.full((B,), real_label, dtype=torch.float, device=device)\n",
    "            fake_value = torch.full((B,), fake_label, dtype=torch.float, device=device)\n",
    "            real_map = torch.full((B, 1, H, W), real_label, dtype=torch.float, device=device)\n",
    "            fake_map = torch.full((B, 1, H, W), fake_label, dtype=torch.float, device=device)\n",
    "            \n",
    "            gt_img_out, gt_fg_seg, gt_bg_seg = disc(mask, gt)#.view(-1) #[b,1,1,1,] -> [b]\n",
    "\n",
    "            d_loss_real = criterion(gt_img_out.view(-1), real_value.clone()) #gt, real\n",
    "            ds_loss_real = criterion(gt_bg_seg, torch.multiply(~mask, real_map.clone())) + criterion(gt_fg_seg, torch.multiply(mask, real_map.clone()))#criterion(gt_seg_out, real_map.clone()) #gt, real\n",
    "\n",
    "            ########## Train with all-fake batch\n",
    "            # Generate batch of latent vectors\n",
    "            \n",
    "            AI, HI, FI, pred_mat, R_in, R_out = net(tf_fg, bg, tf_mask)\n",
    "            R_img_out, R_fg_seg, R_bg_seg = disc(AI, R_out) #torch.Size([8, 1, 1, 1]), torch.Size([8, 1, 256, 256])\n",
    "            \n",
    "            d_loss_fake = criterion(R_img_out.view(-1), fake_value.clone()) #gen, fake\n",
    "            ds_loss_fake = criterion(R_bg_seg, torch.multiply(~AI, real_map.clone())) + criterion(R_fg_seg, torch.multiply(AI, fake_map.clone()))\n",
    "            \n",
    "            dsc_loss = d_loss_real + d_loss_fake + 0.01 * (ds_loss_real + ds_loss_fake) # +d_loss_reg\n",
    "            optimizer.zero_grad()\n",
    "            optimizerD.zero_grad()\n",
    "            dsc_loss.backward()\n",
    "            optimizerD.step()\n",
    "  \n",
    "            ############################\n",
    "            # (2) Update G network: maximize log(D(G(z)))\n",
    "            ###########################\n",
    "            \n",
    "            AI, HI, FI, pred_mat, R_in, R_out = net(tf_fg, bg, tf_mask)\n",
    "            \n",
    "            R_img_out, R_fg_seg, R_bg_seg = disc(AI, R_out)\n",
    "            g_loss_fake = criterion(R_img_out.view(-1), real_value.clone())# + criterion(torch.multiply(AI, R_img_out), real_map.clone()) #criterion(R_img_out.view(-1), real_value.clone()) # err_G = loss_a\n",
    "            \n",
    "            gs_loss_fake = criterion(R_bg_seg, torch.multiply(~AI, real_map.clone())) + criterion(R_fg_seg, torch.multiply(AI, real_map.clone())) #  criterion(R_seg_out, real_map.clone())# + criterion(R_seg_out, fake_map.clone())\n",
    "\n",
    "            lambda_mask = 1\n",
    "            \n",
    "            loss_g = l2_loss(trans_mat.float(), pred_mat.float()) # lambda_mask * torch.exp(-torch.sum(AI, (1,2,3), dtype=torch.float).mean()/args.input_size**2) + torch.linalg.matrix_norm(pred_mat.float(), ord=2).mean()+ + torch.linalg.matrix_norm(pred_mat.float(), ord=2).mean()\n",
    "            loss_c = torch.linalg.matrix_norm(torch.multiply(HI-FI, AI), ord=1).mean() / torch.linalg.matrix_norm(AI.float(), ord=1).mean() + SMOOTH # / (torch.sum(AI, (1,2,3), dtype=torch.float).mean() + SMOOTH)# + torch.linalg.matrix_norm(AI.float(), ord=1)\n",
    "            loss_s2 = l1_loss(gt, R_out)\n",
    "            \n",
    "            loss = args.lambda_g * loss_g + args.lambda_c * loss_c + args.lambda_a * g_loss_fake + args.lambda_s * gs_loss_fake + args.lambda_s2 * loss_s2\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            optimizerD.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            if i % args.test_interval == 0:\n",
    "                print('[{:d} epoch {:d}/{:d}]: c:{:f}, g:{:f}, dsc_loss:{:f}, a:{:f}, s:{:f}, total:{:f}'.format(epoch, i, len(dataloader)-1, loss_c.item(), loss_g.item(), dsc_loss.item(), g_loss_fake.item(), gs_loss_fake.item(), loss.item()), end='\\r')\n",
    "\n",
    "                vis.line(Y=torch.Tensor([torch.linalg.matrix_norm(pred_mat.float(), ord=2).mean().item()]), X=torch.Tensor([epoch*len(dataloader)+i]), win=plot, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([args.lambda_a*g_loss_fake.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Loss_a', win=plot2, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([args.lambda_s*gs_loss_fake.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Loss_s', win=plot2, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([args.lambda_c*loss_c.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Loss_c', win=plot2, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([args.lambda_g*loss_g.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Loss_g', win=plot2, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([args.lambda_s2*loss_s2.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Loss_s2', win=plot2, update='append', env=args.vis_id)\n",
    "\n",
    "                vis.line(Y=torch.Tensor([d_loss_real.item()+ds_loss_real.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Discriminator_real', win=plot3, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([d_loss_fake.item()+ds_loss_fake.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Discriminator_fake', win=plot3, update='append', env=args.vis_id)\n",
    "                vis.line(Y=torch.Tensor([loss.item()]), X=torch.Tensor([epoch*len(dataloader)+i]), name='Generator', win=plot3, update='append', env=args.vis_id)\n",
    "\n",
    "                inputs = tensor2im(denorm((torch.multiply(tf_fg, tf_mask) + torch.multiply(bg, ~tf_mask)).detach().cpu()).float())\n",
    "                FI_out = tensor2im(denorm(FI.detach().cpu()).float())\n",
    "                HI_out = tensor2im(denorm(HI.detach().cpu()).float())\n",
    "                AI_out = tensor2im(AI.detach().cpu())\n",
    "                R_in = tensor2im(denorm(R_in.detach().cpu()).float())\n",
    "                seg_out = tensor2im((R_fg_seg+R_bg_seg).detach().cpu())\n",
    "                R_out = tensor2im(denorm(R_out.detach().cpu()).float())\n",
    "\n",
    "                vis.images(inputs, win=vis_inputs, opts=dict(title=\"Inputs\"), env=args.vis_id)\n",
    "                vis.images(HI_out, win=vis_HI, opts=dict(title=\"H(I)\"), env=args.vis_id)\n",
    "                vis.images(FI_out, win=vis_FI, opts=dict(title=\"F(I)\"), env=args.vis_id)\n",
    "                vis.images(AI_out, win=vis_AI, opts=dict(title=\"A(I)\"), env=args.vis_id)\n",
    "                vis.images(R_in, win=vis_R_in, opts=dict(title=\"R_in\"), env=args.vis_id)\n",
    "                vis.images(R_out, win=vis_R_out, opts=dict(title=\"R_out\"), env=args.vis_id)\n",
    "                vis.images(seg_out, win=vis_seg, opts=dict(title=\"Segmap\"), env=args.vis_id)\n",
    "            \n",
    "            if i % 10000 == 0:\n",
    "                save_checkpoint(epoch, net, optimizer, args.save_model_dir + 'netHCRS-free3.pth')\n",
    "                save_checkpoint(epoch, disc, optimizerD, args.save_model_dir + 'discHCRS-free3.pth')\n",
    "        "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}