init + inference.py патч

data/MBD/MBD.py

@@ -0,0 +1,110 @@
+import cv2
+import numpy as np
+import MBD_utils
+import torch
+import torch.nn.functional as F
+
+
+def mask_base_dewarper(image,mask):
+    '''
+    input:
+        image -> ndarray HxWx3 uint8
+        mask -> ndarray HxW uint8
+    return
+        dewarped -> ndarray HxWx3 uint8
+        grid (optional) -> ndarray HxWx2 -1~1
+    '''
+
+    ## get contours
+    # _, contours, hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)  ## cv2.__version__ == 3.x
+    contours,hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,method=cv2.CHAIN_APPROX_SIMPLE)  ## cv2.__version__ == 4.x
+
+    ## get biggest contours and four corners based on Douglas-Peucker algorithm
+    four_corners, maxArea, contour= MBD_utils.DP_algorithm(contours)
+    four_corners = MBD_utils.reorder(four_corners)
+
+    ## reserve biggest contours and remove other noisy contours
+    new_mask = np.zeros_like(mask)
+    new_mask = cv2.drawContours(new_mask,[contour],-1,255,cv2.FILLED)
+
+    ## obtain middle points
+    # ratios = [0.25,0.5,0.75]  # ratios = [0.125,0.25,0.375,0.5,0.625,0.75,0.875]
+    ratios = [0.25,0.5,0.75]
+    # ratios = [0.0625,0.125,0.1875,0.25,0.3125,0.375,0.4475,0.5,0.5625,0.625,0.06875,0.75,0.8125,0.875,0.9375]
+    middle = MBD_utils.findMiddle(corners=four_corners,mask=new_mask,points=ratios)
+    
+    ## all points
+    source_points = np.concatenate((four_corners,middle),axis=0) ## all_point = four_corners(topleft,topright,bottom)+top+bottom+left+right
+
+    ## target points
+    h,w = image.shape[:2]
+    padding = 0
+    target_points = [[padding, padding],[w-padding, padding], [padding, h-padding],[w-padding, h-padding]]
+    for ratio in ratios:
+        target_points.append([int((w-2*padding)*ratio)+padding,padding])
+    for ratio in ratios:
+        target_points.append([int((w-2*padding)*ratio)+padding,h-padding])
+    for ratio in ratios:
+        target_points.append([padding,int((h-2*padding)*ratio)+padding])
+    for ratio in ratios:
+        target_points.append([w-padding,int((h-2*padding)*ratio)+padding])
+
+    ## dewarp base on cv2
+    # pts1 = np.float32(source_points)
+    # pts2 = np.float32(target_points)
+    # tps = cv2.createThinPlateSplineShapeTransformer()
+    # matches = []
+    # N = pts1.shape[0]
+    # for i in range(0,N):
+    #     matches.append(cv2.DMatch(i,i,0))
+    # pts1 = pts1.reshape(1,-1,2)
+    # pts2 = pts2.reshape(1,-1,2)
+    # tps.estimateTransformation(pts2,pts1,matches)
+    # dewarped = tps.warpImage(image)
+
+    ## dewarp base on generated grid
+    source_points = source_points.reshape(-1,2)/np.array([image.shape[:2][::-1]]).reshape(1,2)
+    source_points = torch.from_numpy(source_points).float().cuda()
+    source_points = source_points.unsqueeze(0)
+    source_points = (source_points-0.5)*2
+    target_points = np.asarray(target_points).reshape(-1,2)/np.array([image.shape[:2][::-1]]).reshape(1,2)
+    target_points = torch.from_numpy(target_points).float()
+    target_points = (target_points-0.5)*2
+
+    model = MBD_utils.TPSGridGen(target_height=256,target_width=256,target_control_points=target_points)
+    model = model.cuda()
+    grid = model(source_points).view(-1,256,256,2).permute(0,3,1,2)
+    grid = F.interpolate(grid,(h,w),mode='bilinear').permute(0,2,3,1)
+    dewarped = MBD_utils.torch2cvimg(F.grid_sample(MBD_utils.cvimg2torch(image).cuda(),grid))[0]
+    return dewarped,grid[0].cpu().numpy()
+
+def mask_base_cropper(image,mask):
+    '''
+    input:
+        image -> ndarray HxWx3 uint8
+        mask -> ndarray HxW uint8
+    return
+        dewarped -> ndarray HxWx3 uint8
+        grid (optional) -> ndarray HxWx2 -1~1
+    '''
+
+    ## get contours
+    _, contours, hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)  ## cv2.__version__ == 3.x
+    # contours,hierarchy = cv2.findContours(mask,cv2.RETR_EXTERNAL,method=cv2.CHAIN_APPROX_SIMPLE)  ## cv2.__version__ == 4.x
+
+    ## get biggest contours and four corners based on Douglas-Peucker algorithm
+    four_corners, maxArea, contour= MBD_utils.DP_algorithm(contours)
+    four_corners = MBD_utils.reorder(four_corners)
+
+    ## reserve biggest contours and remove other noisy contours
+    new_mask = np.zeros_like(mask)
+    new_mask = cv2.drawContours(new_mask,[contour],-1,255,cv2.FILLED)
+
+    ## 最小外接矩形
+    rect = cv2.minAreaRect(contour) # 得到最小外接矩形的（中心(x,y), (宽,高), 旋转角度）
+    box = cv2.boxPoints(rect) # cv2.boxPoints(rect) for OpenCV 3.x 获取最小外接矩形的4个顶点坐标
+    box = np.int0(box)
+    box = box.reshape((4,1,2))    
+
+
+

data/MBD/MBD_utils.py

@@ -0,0 +1,291 @@
+import cv2
+import numpy as np
+import copy
+import torch
+import torch
+import itertools
+import torch.nn as nn
+from torch.autograd import Function, Variable
+
+def reorder(myPoints):
+    myPoints = myPoints.reshape((4, 2))
+    myPointsNew = np.zeros((4, 1, 2), dtype=np.int32)
+    add = myPoints.sum(1)
+    myPointsNew[0] = myPoints[np.argmin(add)]
+    myPointsNew[3] =myPoints[np.argmax(add)]
+    diff = np.diff(myPoints, axis=1)
+    myPointsNew[1] =myPoints[np.argmin(diff)]
+    myPointsNew[2] = myPoints[np.argmax(diff)]
+    return myPointsNew
+
+
+def findMiddle(corners,mask,points=[0.25,0.5,0.75]):
+    num_middle_points = len(points)
+    top = [np.array([])]*num_middle_points
+    bottom = [np.array([])]*num_middle_points
+    left = [np.array([])]*num_middle_points
+    right = [np.array([])]*num_middle_points
+
+    center_top = []
+    center_bottom = []
+    center_left = []
+    center_right = []
+
+    center = (int((corners[0][0][1]+corners[3][0][1])/2),int((corners[0][0][0]+corners[3][0][0])/2))
+    for ratio in points:
+
+        center_top.append( (center[0],int(corners[0][0][0]*(1-ratio)+corners[1][0][0]*ratio)) )
+
+        center_bottom.append( (center[0],int(corners[2][0][0]*(1-ratio)+corners[3][0][0]*ratio)) )
+
+        center_left.append( (int(corners[0][0][1]*(1-ratio)+corners[2][0][1]*ratio),center[1]) )
+
+        center_right.append( (int(corners[1][0][1]*(1-ratio)+corners[3][0][1]*ratio),center[1]) )
+
+    for i in range(0,center[0],1):
+        for j in range(num_middle_points):
+            if top[j].size==0:
+                if mask[i,center_top[j][1]]==255:
+                    top[j] = np.asarray([center_top[j][1],i])
+                    top[j] = top[j].reshape(1,2)
+
+    for i in range(mask.shape[0]-1,center[0],-1):
+        for j in range(num_middle_points):
+            if bottom[j].size==0:
+                if mask[i,center_bottom[j][1]]==255:
+                    bottom[j] = np.asarray([center_bottom[j][1],i])
+                    bottom[j] = bottom[j].reshape(1,2)
+
+    for i in range(mask.shape[1]-1,center[1],-1):
+        for j in range(num_middle_points):
+            if right[j].size==0:
+                if mask[center_right[j][0],i]==255:
+                    right[j] = np.asarray([i,center_right[j][0]])
+                    right[j] = right[j].reshape(1,2)
+
+    for i in range(0,center[1]):
+        for j in range(num_middle_points):
+            if left[j].size==0:
+                if mask[center_left[j][0],i]==255:
+                    left[j] = np.asarray([i,center_left[j][0]])
+                    left[j] = left[j].reshape(1,2) 
+
+    return np.asarray(top+bottom+left+right)
+
+def DP_algorithmv1(contours):
+    biggest = np.array([])
+    max_area = 0
+    step = 0.001
+    count = 0
+    # while biggest.size==0:
+    while True:
+        for i in contours:
+            # print(i.shape)
+            area = cv2.contourArea(i)
+            # print(area,cv2.arcLength(i, True))
+            if area > cv2.arcLength(i, True)*10:
+                peri = cv2.arcLength(i, True)
+                approx = cv2.approxPolyDP(i, (0.01+step*count) * peri, True)
+                if area > max_area and len(approx) == 4:
+                    max_area = area
+                    biggest_contours = i
+                    biggest = approx
+                    break
+                    if abs(max_area - cv2.contourArea(biggest))/max_area > 0.3:
+                        biggest = np.array([])
+        count += 1
+        if count > 200:
+            break
+    temp = biggest[0]
+    return biggest,max_area, biggest_contours
+
+def DP_algorithm(contours):
+    biggest = np.array([])
+    max_area = 0
+    step = 0.001
+    count = 0
+
+    ### largest contours
+    for i in contours:
+        area = cv2.contourArea(i)
+        if area > max_area:
+            max_area = area
+            biggest_contours = i
+    peri = cv2.arcLength(biggest_contours, True)
+
+    ### find four corners
+    while True:
+        approx = cv2.approxPolyDP(biggest_contours, (0.01+step*count) * peri, True)
+        if len(approx) == 4:
+            biggest = approx
+            break
+            # if abs(max_area - cv2.contourArea(biggest))/max_area > 0.2:
+            # if abs(max_area - cv2.contourArea(biggest))/max_area > 0.4:
+                # biggest = np.array([])
+        count += 1
+        if count > 200:
+            break
+    return biggest,max_area, biggest_contours
+
+def drawRectangle(img,biggest,color,thickness):
+    cv2.line(img, (biggest[0][0][0], biggest[0][0][1]), (biggest[1][0][0], biggest[1][0][1]), color, thickness)
+    cv2.line(img, (biggest[0][0][0], biggest[0][0][1]), (biggest[2][0][0], biggest[2][0][1]), color, thickness)
+    cv2.line(img, (biggest[3][0][0], biggest[3][0][1]), (biggest[2][0][0], biggest[2][0][1]), color, thickness)
+    cv2.line(img, (biggest[3][0][0], biggest[3][0][1]), (biggest[1][0][0], biggest[1][0][1]), color, thickness)
+    return img
+
+def minAreaRect(contours,img):
+    # biggest = np.array([])
+    max_area = 0
+    for i in contours:
+        area = cv2.contourArea(i)
+        if area > max_area:
+            peri = cv2.arcLength(i, True)
+            rect = cv2.minAreaRect(i)
+            points = cv2.boxPoints(rect)
+            max_area = area
+    return points
+
+def cropRectangle(img,biggest):
+    # print(biggest)
+    w = np.abs(biggest[0][0][0] - biggest[1][0][0])
+    h = np.abs(biggest[0][0][1] - biggest[2][0][1])
+    new_img = np.zeros((w,h,img.shape[-1]),dtype=np.uint8)
+    new_img = img[biggest[0][0][1]:biggest[0][0][1]+h,biggest[0][0][0]:biggest[0][0][0]+w]
+    return new_img
+
+def cvimg2torch(img,min=0,max=1):
+    '''
+    input:
+        im -> ndarray uint8 HxWxC 
+    return
+        tensor -> torch.tensor BxCxHxW 
+    '''
+    if len(img.shape)==2:
+        img = np.expand_dims(img,axis=-1)
+    img = img.astype(float) / 255.0
+    img = img.transpose(2, 0, 1) # NHWC -> NCHW
+    img = np.expand_dims(img, 0)
+    img = torch.from_numpy(img).float()
+    return img
+
+def torch2cvimg(tensor,min=0,max=1):
+    '''
+    input:
+        tensor -> torch.tensor BxCxHxW C can be 1,3
+    return
+        im -> ndarray uint8 HxWxC 
+    '''
+    im_list = []
+    for i in range(tensor.shape[0]):
+        im = tensor.detach().cpu().data.numpy()[i]
+        im = im.transpose(1,2,0)
+        im = np.clip(im,min,max)
+        im = ((im-min)/(max-min)*255).astype(np.uint8)
+        im_list.append(im)
+    return im_list
+
+
+
+class TPSGridGen(nn.Module):
+    def __init__(self, target_height, target_width, target_control_points):
+        '''
+        target_control_points -> torch.tensor  num_pointx2 -1~1
+        source_control_points -> torch.tensor  batch_size x num_point x 2 -1~1
+        return:
+            grid -> batch_size x hw x 2 -1~1
+        '''
+        super(TPSGridGen, self).__init__()
+        assert target_control_points.ndimension() == 2
+        assert target_control_points.size(1) == 2
+        N = target_control_points.size(0)
+        self.num_points = N
+        target_control_points = target_control_points.float()
+
+        # create padded kernel matrix
+        forward_kernel = torch.zeros(N + 3, N + 3)
+        target_control_partial_repr = self.compute_partial_repr(target_control_points, target_control_points)
+        forward_kernel[:N, :N].copy_(target_control_partial_repr)
+        forward_kernel[:N, -3].fill_(1)
+        forward_kernel[-3, :N].fill_(1)
+        forward_kernel[:N, -2:].copy_(target_control_points)
+        forward_kernel[-2:, :N].copy_(target_control_points.transpose(0, 1))
+        # compute inverse matrix
+        inverse_kernel = torch.inverse(forward_kernel)
+
+        # create target cordinate matrix
+        HW = target_height * target_width
+        target_coordinate = list(itertools.product(range(target_height), range(target_width)))
+        target_coordinate = torch.Tensor(target_coordinate) # HW x 2
+        Y, X = target_coordinate.split(1, dim = 1)
+        Y = Y * 2 / (target_height - 1) - 1
+        X = X * 2 / (target_width - 1) - 1
+        target_coordinate = torch.cat([X, Y], dim = 1) # convert from (y, x) to (x, y)
+        target_coordinate_partial_repr = self.compute_partial_repr(target_coordinate.to(target_control_points.device), target_control_points)
+        target_coordinate_repr = torch.cat([
+            target_coordinate_partial_repr, torch.ones(HW, 1), target_coordinate
+        ], dim = 1)
+
+        # register precomputed matrices
+        self.register_buffer('inverse_kernel', inverse_kernel)
+        self.register_buffer('padding_matrix', torch.zeros(3, 2))
+        self.register_buffer('target_coordinate_repr', target_coordinate_repr)
+
+    def forward(self, source_control_points):
+        assert source_control_points.ndimension() == 3
+        assert source_control_points.size(1) == self.num_points
+        assert source_control_points.size(2) == 2
+        batch_size = source_control_points.size(0)
+
+        Y = torch.cat([source_control_points, Variable(self.padding_matrix.expand(batch_size, 3, 2))], 1)
+        mapping_matrix = torch.matmul(Variable(self.inverse_kernel), Y)
+        source_coordinate = torch.matmul(Variable(self.target_coordinate_repr), mapping_matrix)
+        return source_coordinate
+    # phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
+    def compute_partial_repr(self, input_points, control_points):
+        N = input_points.size(0)
+        M = control_points.size(0)
+        pairwise_diff = input_points.view(N, 1, 2) - control_points.view(1, M, 2)
+        # original implementation, very slow
+        # pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
+        pairwise_diff_square = pairwise_diff * pairwise_diff
+        pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :, 1]
+        repr_matrix = 0.5 * pairwise_dist * torch.log(pairwise_dist)
+        # fix numerical error for 0 * log(0), substitute all nan with 0
+        mask = repr_matrix != repr_matrix
+        repr_matrix.masked_fill_(mask, 0)
+        return repr_matrix
+
+
+
+
+ 
+    ### deside wheather further process
+    # point_area = cv2.contourArea(np.concatenate((biggest_angle[0].reshape(1,1,2),middle[0:3],biggest_angle[1].reshape(1,1,2),middle[9:12],biggest_angle[3].reshape(1,1,2),middle[3:6][::-1],biggest_angle[2].reshape(1,1,2),middle[6:9][::-1]),axis=0))
+    #### 最小外接矩形
+    # rect = cv2.minAreaRect(contour) # 得到最小外接矩形的（中心(x,y), (宽,高), 旋转角度）
+    # box = cv2.boxPoints(rect) # cv2.boxPoints(rect) for OpenCV 3.x 获取最小外接矩形的4个顶点坐标
+    # box = np.int0(box)
+    # box = box.reshape((4,1,2))
+    # minrect_area = cv2.contourArea(box)
+    # print(abs(minrect_area-point_area)/point_area)
+    #### 四个角点 IOU
+    # biggest_box = np.concatenate((biggest_angle[0,:,:].reshape(1,1,2),biggest_angle[2,:,:].reshape(1,1,2),biggest_angle[3,:,:].reshape(1,1,2),biggest_angle[1,:,:].reshape(1,1,2)),axis=0)
+    # biggest_mask = np.zeros_like(mask)
+    # # corner_area = cv2.contourArea(biggest_box)
+    # cv2.drawContours(biggest_mask,[biggest_box], -1, color=255, thickness=-1)
+    
+    # smooth = 1e-5
+    # biggest_mask_ = biggest_mask > 50
+    # mask_ = mask > 50
+    # intersection = (biggest_mask_ & mask_).sum()
+    # union = (biggest_mask_ | mask_).sum()
+    # iou = (intersection + smooth) / (union + smooth)
+    # if iou > 0.975:
+    #     skip = True
+    # else:
+    #     skip = False
+    # print(iou)
+    # cv2.imshow('mask',cv2.resize(mask,(512,512)))
+    # cv2.imshow('biggest_mask',cv2.resize(biggest_mask,(512,512)))
+    # cv2.waitKey(0)    

data/MBD/infer.py

@@ -0,0 +1,151 @@
+import torch
+import argparse
+import numpy as np
+import torch.nn.functional as F
+import glob
+import cv2
+from tqdm import tqdm
+
+import time
+import os 
+from model.deep_lab_model.deeplab import *
+from MBD import mask_base_dewarper
+import time
+
+from utils import cvimg2torch,torch2cvimg
+
+
+
+def net1_net2_infer(model,img_paths,args):
+
+    ### validate on the real datasets
+    seg_model=model
+    seg_model.eval()
+    for img_path in tqdm(img_paths):
+        if os.path.exists(img_path.replace('_origin','_capture')):
+            continue
+        t1 = time.time()
+        ### segmentation mask predict
+        img_org = cv2.imread(img_path)
+        h_org,w_org = img_org.shape[:2]
+        img = cv2.resize(img_org,(448, 448))       
+        img = cv2.GaussianBlur(img,(15,15),0,0)
+        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
+        img = cvimg2torch(img)
+
+        with torch.no_grad():
+            pred = seg_model(img.cuda())
+            mask_pred = pred[:,0,:,:].unsqueeze(1)
+            mask_pred = F.interpolate(mask_pred,(h_org,w_org))
+            mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy()
+            mask_pred = (mask_pred*255).astype(np.uint8)
+            kernel = np.ones((3,3))
+            mask_pred = cv2.dilate(mask_pred,kernel,iterations=3)
+            mask_pred = cv2.erode(mask_pred,kernel,iterations=3)
+            mask_pred[mask_pred>100] = 255
+            mask_pred[mask_pred<100] = 0
+            ### tps transform base on the mask
+            # dewarp, grid = mask_base_dewarper(img_org,mask_pred)
+        try:
+            dewarp, grid = mask_base_dewarper(img_org,mask_pred)
+        except:
+            print('fail')
+            grid = np.meshgrid(np.arange(w_org),np.arange(h_org))/np.array([w_org,h_org]).reshape(2,1,1)
+            grid = torch.from_numpy((grid-0.5)*2).float().unsqueeze(0).permute(0,2,3,1)
+            dewarp = torch2cvimg(F.grid_sample(cvimg2torch(img_org),grid))[0]
+            grid = grid[0].numpy()
+        # cv2.imshow('in',cv2.resize(img_org,(512,512)))
+        # cv2.imshow('out',cv2.resize(dewarp,(512,512)))
+        # cv2.waitKey(0)
+        cv2.imwrite(img_path.replace('_origin','_capture'),dewarp)
+        cv2.imwrite(img_path.replace('_origin','_mask_new'),mask_pred)
+
+        grid0 = cv2.resize(grid[:,:,0],(128,128))
+        grid1 = cv2.resize(grid[:,:,1],(128,128))
+        grid = np.stack((grid0,grid1),axis=-1)
+        np.save(img_path.replace('_origin','_grid1'),grid)
+
+
+def net1_net2_infer_single_im(img,model_path):
+    seg_model = DeepLab(num_classes=1,
+                    backbone='resnet',
+                    output_stride=16,
+                    sync_bn=None,
+                    freeze_bn=False)
+    seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
+    seg_model.cuda()
+    checkpoint = torch.load(model_path)
+    seg_model.load_state_dict(checkpoint['model_state'])
+    ### validate on the real datasets
+    seg_model.eval()
+    ### segmentation mask predict
+    img_org = img
+    h_org,w_org = img_org.shape[:2]
+    img = cv2.resize(img_org,(448, 448))       
+    img = cv2.GaussianBlur(img,(15,15),0,0)
+    img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)
+    img = cvimg2torch(img)
+
+    with torch.no_grad():
+        # from torchtoolbox.tools import summary
+        # print(summary(seg_model,torch.rand((1, 3, 448, 448)).cuda())) 59.4M 135.6G
+
+        pred = seg_model(img.cuda())
+        mask_pred = pred[:,0,:,:].unsqueeze(1)
+        mask_pred = F.interpolate(mask_pred,(h_org,w_org))
+        mask_pred = mask_pred.squeeze(0).squeeze(0).cpu().numpy()
+        mask_pred = (mask_pred*255).astype(np.uint8)
+        kernel = np.ones((3,3))
+        mask_pred = cv2.dilate(mask_pred,kernel,iterations=3)
+        mask_pred = cv2.erode(mask_pred,kernel,iterations=3)
+        mask_pred[mask_pred>100] = 255
+        mask_pred[mask_pred<100] = 0
+        ### tps transform base on the mask
+        # dewarp, grid = mask_base_dewarper(img_org,mask_pred)
+    # try:
+    #     dewarp, grid = mask_base_dewarper(img_org,mask_pred)
+    # except:
+    #     print('fail')
+    #     grid = np.meshgrid(np.arange(w_org),np.arange(h_org))/np.array([w_org,h_org]).reshape(2,1,1)
+    #     grid = torch.from_numpy((grid-0.5)*2).float().unsqueeze(0).permute(0,2,3,1)
+    #     dewarp = torch2cvimg(F.grid_sample(cvimg2torch(img_org),grid))[0]
+    #     grid = grid[0].numpy()
+    # cv2.imshow('in',cv2.resize(img_org,(512,512)))
+    # cv2.imshow('out',cv2.resize(dewarp,(512,512)))
+    # cv2.waitKey(0)
+    # cv2.imwrite(img_path.replace('_origin','_capture'),dewarp)
+    # cv2.imwrite(img_path.replace('_origin','_mask_new'),mask_pred)
+
+    # grid0 = cv2.resize(grid[:,:,0],(128,128))
+    # grid1 = cv2.resize(grid[:,:,1],(128,128))
+    # grid = np.stack((grid0,grid1),axis=-1)
+    # np.save(img_path.replace('_origin','_grid1'),grid)
+    return mask_pred
+
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Hyperparams')
+    parser.add_argument('--img_folder', nargs='?', type=str, default='./all_data',help='Data path to load data')
+    parser.add_argument('--img_rows', nargs='?', type=int, default=448, 
+                        help='Height of the input image')
+    parser.add_argument('--img_cols', nargs='?', type=int, default=448, 
+                        help='Width of the input image')
+    parser.add_argument('--seg_model_path', nargs='?', type=str, default='checkpoints/mbd.pkl',
+                        help='Path to previous saved model to restart from')
+    args = parser.parse_args()
+
+    seg_model = DeepLab(num_classes=1,
+                    backbone='resnet',
+                    output_stride=16,
+                    sync_bn=None,
+                    freeze_bn=False)
+    seg_model = torch.nn.DataParallel(seg_model, device_ids=range(torch.cuda.device_count()))
+    seg_model.cuda()
+    checkpoint = torch.load(args.seg_model_path)
+    seg_model.load_state_dict(checkpoint['model_state'])
+
+    im_paths = glob.glob(os.path.join(args.img_folder,'*_origin.*'))
+
+    net1_net2_infer(seg_model,im_paths,args)
+

data/MBD/model/__init__.py

@@ -0,0 +1,50 @@
+import torchvision.models as models
+from model.densenetccnl import *
+from model.unetnc import *
+from model.gienet import *
+
+
+def get_model(name, n_classes=1, filters=64,version=None,in_channels=3, is_batchnorm=True, norm='batch', model_path=None, use_sigmoid=True, layers=3,img_size=512):
+    model = _get_model_instance(name)
+
+
+    if name == 'dnetccnl':
+        model = model(img_size=128, in_channels=in_channels, out_channels=n_classes, filters=32)
+    elif name == 'dnetccnl512':
+        model = model(img_size=img_size, in_channels=in_channels, out_channels=n_classes, filters=32)
+    elif name == 'unetnc':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'gie':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giecbam':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'gie2head':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giemask':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giemask2':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'giedilated':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'bmp':
+        model = model(input_nc=in_channels, output_nc=n_classes, num_downs=7)
+    elif name == 'displacement':
+        model = model(n_classes=2, num_filter=32, BatchNorm='GN', in_channels=5)
+    return model
+
+def _get_model_instance(name):
+    try:
+        return {
+            'dnetccnl': dnetccnl,
+            'dnetccnl512': dnetccnl512,
+            'unetnc': UnetGenerator,
+            'gie':GieGenerator,
+            'giecbam':GiecbamGenerator,
+            'giedilated':DilatedSingleUnet,
+            'gie2head':Gie2headGenerator,
+            'giemask':GiemaskGenerator,
+            'giemask2':Giemask2Generator,
+            'bmp':BmpGenerator,
+        }[name]
+    except:
+        print('Model {} not available'.format(name))

data/MBD/model/cbam.py

@@ -0,0 +1,95 @@
+import torch
+import math
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BasicConv(nn.Module):
+    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True, bn=True, bias=False):
+        super(BasicConv, self).__init__()
+        self.out_channels = out_planes
+        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
+        self.bn = nn.BatchNorm2d(out_planes,eps=1e-5, momentum=0.01, affine=True) if bn else None
+        self.relu = nn.ReLU() if relu else None
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.bn is not None:
+            x = self.bn(x)
+        if self.relu is not None:
+            x = self.relu(x)
+        return x
+
+class Flatten(nn.Module):
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+class ChannelGate(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
+        super(ChannelGate, self).__init__()
+        self.gate_channels = gate_channels
+        self.mlp = nn.Sequential(
+            Flatten(),
+            nn.Linear(gate_channels, gate_channels // reduction_ratio),
+            nn.ReLU(),
+            nn.Linear(gate_channels // reduction_ratio, gate_channels)
+            )
+        self.pool_types = pool_types
+    def forward(self, x):
+        channel_att_sum = None
+        for pool_type in self.pool_types:
+            if pool_type=='avg':
+                avg_pool = F.avg_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( avg_pool )
+            elif pool_type=='max':
+                max_pool = F.max_pool2d( x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( max_pool )
+            elif pool_type=='lp':
+                lp_pool = F.lp_pool2d( x, 2, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
+                channel_att_raw = self.mlp( lp_pool )
+            elif pool_type=='lse':
+                # LSE pool only
+                lse_pool = logsumexp_2d(x)
+                channel_att_raw = self.mlp( lse_pool )
+
+            if channel_att_sum is None:
+                channel_att_sum = channel_att_raw
+            else:
+                channel_att_sum = channel_att_sum + channel_att_raw
+
+        scale = F.sigmoid( channel_att_sum ).unsqueeze(2).unsqueeze(3).expand_as(x)
+        return x * scale
+
+def logsumexp_2d(tensor):
+    tensor_flatten = tensor.view(tensor.size(0), tensor.size(1), -1)
+    s, _ = torch.max(tensor_flatten, dim=2, keepdim=True)
+    outputs = s + (tensor_flatten - s).exp().sum(dim=2, keepdim=True).log()
+    return outputs
+
+class ChannelPool(nn.Module):
+    def forward(self, x):
+        return torch.cat( (torch.max(x,1)[0].unsqueeze(1), torch.mean(x,1).unsqueeze(1)), dim=1 )
+
+class SpatialGate(nn.Module):
+    def __init__(self):
+        super(SpatialGate, self).__init__()
+        kernel_size = 7
+        self.compress = ChannelPool()
+        self.spatial = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size-1) // 2, relu=False)
+    def forward(self, x):
+        x_compress = self.compress(x)
+        x_out = self.spatial(x_compress)
+        scale = F.sigmoid(x_out) # broadcasting
+        return x * scale
+
+class CBAM(nn.Module):
+    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max'], no_spatial=False):
+        super(CBAM, self).__init__()
+        self.ChannelGate = ChannelGate(gate_channels, reduction_ratio, pool_types)
+        self.no_spatial=no_spatial
+        if not no_spatial:
+            self.SpatialGate = SpatialGate()
+    def forward(self, x):
+        x_out = self.ChannelGate(x)
+        if not self.no_spatial:
+            x_out = self.SpatialGate(x_out)
+        return x_out

data/MBD/model/deep_lab_model/aspp.py

@@ -0,0 +1,95 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+class _ASPPModule(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size, padding, dilation, BatchNorm):
+        super(_ASPPModule, self).__init__()
+        self.atrous_conv = nn.Conv2d(inplanes, planes, kernel_size=kernel_size,
+                                            stride=1, padding=padding, dilation=dilation, bias=False)
+        self.bn = BatchNorm(planes)
+        self.relu = nn.ReLU()
+
+        self._init_weight()
+
+    def forward(self, x):
+        x = self.atrous_conv(x)
+        x = self.bn(x)
+
+        return self.relu(x)
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+class ASPP(nn.Module):
+    def __init__(self, backbone, output_stride, BatchNorm):
+        super(ASPP, self).__init__()
+        if backbone == 'drn':
+            inplanes = 512
+        elif backbone == 'mobilenet':
+            inplanes = 320
+        else:
+            inplanes = 2048
+        if output_stride == 16:
+            dilations = [1, 6, 12, 18]
+        elif output_stride == 8:
+            dilations = [1, 12, 24, 36]
+        else:
+            raise NotImplementedError
+
+        self.aspp1 = _ASPPModule(inplanes, 256, 1, padding=0, dilation=dilations[0], BatchNorm=BatchNorm)
+        self.aspp2 = _ASPPModule(inplanes, 256, 3, padding=dilations[1], dilation=dilations[1], BatchNorm=BatchNorm)
+        self.aspp3 = _ASPPModule(inplanes, 256, 3, padding=dilations[2], dilation=dilations[2], BatchNorm=BatchNorm)
+        self.aspp4 = _ASPPModule(inplanes, 256, 3, padding=dilations[3], dilation=dilations[3], BatchNorm=BatchNorm)
+
+        self.global_avg_pool = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
+                                             nn.Conv2d(inplanes, 256, 1, stride=1, bias=False),
+                                             BatchNorm(256),
+                                             nn.ReLU())
+        self.conv1 = nn.Conv2d(1280, 256, 1, bias=False)
+        self.bn1 = BatchNorm(256)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(0.5)
+        self._init_weight()
+
+    def forward(self, x):
+        x1 = self.aspp1(x)
+        x2 = self.aspp2(x)
+        x3 = self.aspp3(x)
+        x4 = self.aspp4(x)
+        x5 = self.global_avg_pool(x)
+        x5 = F.interpolate(x5, size=x4.size()[2:], mode='bilinear', align_corners=True)
+        x = torch.cat((x1, x2, x3, x4, x5), dim=1)
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        return self.dropout(x)
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                # m.weight.data.normal_(0, math.sqrt(2. / n))
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+def build_aspp(backbone, output_stride, BatchNorm):
+    return ASPP(backbone, output_stride, BatchNorm)

data/MBD/model/deep_lab_model/backbone/__init__.py

@@ -0,0 +1,13 @@
+from model.deep_lab_model.backbone import resnet, xception, drn, mobilenet
+
+def build_backbone(backbone, output_stride, BatchNorm):
+    if backbone == 'resnet':
+        return resnet.ResNet101(output_stride, BatchNorm)
+    elif backbone == 'xception':
+        return xception.AlignedXception(output_stride, BatchNorm)
+    elif backbone == 'drn':
+        return drn.drn_d_54(BatchNorm)
+    elif backbone == 'mobilenet':
+        return mobilenet.MobileNetV2(output_stride, BatchNorm)
+    else:
+        raise NotImplementedError

data/MBD/model/deep_lab_model/backbone/drn.py

@@ -0,0 +1,402 @@
+import torch.nn as nn
+import math
+import torch.utils.model_zoo as model_zoo
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+webroot = 'http://dl.yf.io/drn/'
+
+model_urls = {
+    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
+    'drn-c-26': webroot + 'drn_c_26-ddedf421.pth',
+    'drn-c-42': webroot + 'drn_c_42-9d336e8c.pth',
+    'drn-c-58': webroot + 'drn_c_58-0a53a92c.pth',
+    'drn-d-22': webroot + 'drn_d_22-4bd2f8ea.pth',
+    'drn-d-38': webroot + 'drn_d_38-eebb45f0.pth',
+    'drn-d-54': webroot + 'drn_d_54-0e0534ff.pth',
+    'drn-d-105': webroot + 'drn_d_105-12b40979.pth'
+}
+
+
+def conv3x3(in_planes, out_planes, stride=1, padding=1, dilation=1):
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=padding, bias=False, dilation=dilation)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 dilation=(1, 1), residual=True, BatchNorm=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride,
+                             padding=dilation[0], dilation=dilation[0])
+        self.bn1 = BatchNorm(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes,
+                             padding=dilation[1], dilation=dilation[1])
+        self.bn2 = BatchNorm(planes)
+        self.downsample = downsample
+        self.stride = stride
+        self.residual = residual
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+        if self.residual:
+            out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None,
+                 dilation=(1, 1), residual=True, BatchNorm=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               padding=dilation[1], bias=False,
+                               dilation=dilation[1])
+        self.bn2 = BatchNorm(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class DRN(nn.Module):
+
+    def __init__(self, block, layers, arch='D',
+                 channels=(16, 32, 64, 128, 256, 512, 512, 512),
+                 BatchNorm=None):
+        super(DRN, self).__init__()
+        self.inplanes = channels[0]
+        self.out_dim = channels[-1]
+        self.arch = arch
+
+        if arch == 'C':
+            self.conv1 = nn.Conv2d(3, channels[0], kernel_size=7, stride=1,
+                                   padding=3, bias=False)
+            self.bn1 = BatchNorm(channels[0])
+            self.relu = nn.ReLU(inplace=True)
+
+            self.layer1 = self._make_layer(
+                BasicBlock, channels[0], layers[0], stride=1, BatchNorm=BatchNorm)
+            self.layer2 = self._make_layer(
+                BasicBlock, channels[1], layers[1], stride=2, BatchNorm=BatchNorm)
+
+        elif arch == 'D':
+            self.layer0 = nn.Sequential(
+                nn.Conv2d(3, channels[0], kernel_size=7, stride=1, padding=3,
+                          bias=False),
+                BatchNorm(channels[0]),
+                nn.ReLU(inplace=True)
+            )
+
+            self.layer1 = self._make_conv_layers(
+                channels[0], layers[0], stride=1, BatchNorm=BatchNorm)
+            self.layer2 = self._make_conv_layers(
+                channels[1], layers[1], stride=2, BatchNorm=BatchNorm)
+
+        self.layer3 = self._make_layer(block, channels[2], layers[2], stride=2, BatchNorm=BatchNorm)
+        self.layer4 = self._make_layer(block, channels[3], layers[3], stride=2, BatchNorm=BatchNorm)
+        self.layer5 = self._make_layer(block, channels[4], layers[4],
+                                       dilation=2, new_level=False, BatchNorm=BatchNorm)
+        self.layer6 = None if layers[5] == 0 else \
+            self._make_layer(block, channels[5], layers[5], dilation=4,
+                             new_level=False, BatchNorm=BatchNorm)
+
+        if arch == 'C':
+            self.layer7 = None if layers[6] == 0 else \
+                self._make_layer(BasicBlock, channels[6], layers[6], dilation=2,
+                                 new_level=False, residual=False, BatchNorm=BatchNorm)
+            self.layer8 = None if layers[7] == 0 else \
+                self._make_layer(BasicBlock, channels[7], layers[7], dilation=1,
+                                 new_level=False, residual=False, BatchNorm=BatchNorm)
+        elif arch == 'D':
+            self.layer7 = None if layers[6] == 0 else \
+                self._make_conv_layers(channels[6], layers[6], dilation=2, BatchNorm=BatchNorm)
+            self.layer8 = None if layers[7] == 0 else \
+                self._make_conv_layers(channels[7], layers[7], dilation=1, BatchNorm=BatchNorm)
+
+        self._init_weight()
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1,
+                    new_level=True, residual=True, BatchNorm=None):
+        assert dilation == 1 or dilation % 2 == 0
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = list()
+        layers.append(block(
+            self.inplanes, planes, stride, downsample,
+            dilation=(1, 1) if dilation == 1 else (
+                dilation // 2 if new_level else dilation, dilation),
+            residual=residual, BatchNorm=BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, residual=residual,
+                                dilation=(dilation, dilation), BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def _make_conv_layers(self, channels, convs, stride=1, dilation=1, BatchNorm=None):
+        modules = []
+        for i in range(convs):
+            modules.extend([
+                nn.Conv2d(self.inplanes, channels, kernel_size=3,
+                          stride=stride if i == 0 else 1,
+                          padding=dilation, bias=False, dilation=dilation),
+                BatchNorm(channels),
+                nn.ReLU(inplace=True)])
+            self.inplanes = channels
+        return nn.Sequential(*modules)
+
+    def forward(self, x):
+        if self.arch == 'C':
+            x = self.conv1(x)
+            x = self.bn1(x)
+            x = self.relu(x)
+        elif self.arch == 'D':
+            x = self.layer0(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+
+        x = self.layer3(x)
+        low_level_feat = x
+
+        x = self.layer4(x)
+        x = self.layer5(x)
+
+        if self.layer6 is not None:
+            x = self.layer6(x)
+
+        if self.layer7 is not None:
+            x = self.layer7(x)
+
+        if self.layer8 is not None:
+            x = self.layer8(x)
+
+        return x, low_level_feat
+
+
+class DRN_A(nn.Module):
+
+    def __init__(self, block, layers, BatchNorm=None):
+        self.inplanes = 64
+        super(DRN_A, self).__init__()
+        self.out_dim = 512 * block.expansion
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = BatchNorm(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0], BatchNorm=BatchNorm)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, BatchNorm=BatchNorm)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=1,
+                                       dilation=2, BatchNorm=BatchNorm)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=1,
+                                       dilation=4, BatchNorm=BatchNorm)
+
+        self._init_weight()
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, BatchNorm=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, BatchNorm=BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes,
+                                dilation=(dilation, dilation, ), BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        return x
+
+def drn_a_50(BatchNorm, pretrained=True):
+    model = DRN_A(Bottleneck, [3, 4, 6, 3], BatchNorm=BatchNorm)
+    if pretrained:
+        model.load_state_dict(model_zoo.load_url(model_urls['resnet50']))
+    return model
+
+
+def drn_c_26(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1], arch='C', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-c-26'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_c_42(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 1, 1], arch='C', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-c-42'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_c_58(BatchNorm, pretrained=True):
+    model = DRN(Bottleneck, [1, 1, 3, 4, 6, 3, 1, 1], arch='C', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-c-58'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_22(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-22'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_24(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 2, 2, 2, 2, 2, 2], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-24'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_38(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-38'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_40(BatchNorm, pretrained=True):
+    model = DRN(BasicBlock, [1, 1, 3, 4, 6, 3, 2, 2], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-40'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_54(BatchNorm, pretrained=True):
+    model = DRN(Bottleneck, [1, 1, 3, 4, 6, 3, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-54'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+
+def drn_d_105(BatchNorm, pretrained=True):
+    model = DRN(Bottleneck, [1, 1, 3, 4, 23, 3, 1, 1], arch='D', BatchNorm=BatchNorm)
+    if pretrained:
+        pretrained = model_zoo.load_url(model_urls['drn-d-105'])
+        del pretrained['fc.weight']
+        del pretrained['fc.bias']
+        model.load_state_dict(pretrained)
+    return model
+
+if __name__ == "__main__":
+    import torch
+    model = drn_a_50(BatchNorm=nn.BatchNorm2d, pretrained=True)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/backbone/mobilenet.py

@@ -0,0 +1,151 @@
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+import math
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+import torch.utils.model_zoo as model_zoo
+
+def conv_bn(inp, oup, stride, BatchNorm):
+    return nn.Sequential(
+        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
+        BatchNorm(oup),
+        nn.ReLU6(inplace=True)
+    )
+
+
+def fixed_padding(inputs, kernel_size, dilation):
+    kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
+    pad_total = kernel_size_effective - 1
+    pad_beg = pad_total // 2
+    pad_end = pad_total - pad_beg
+    padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
+    return padded_inputs
+
+
+class InvertedResidual(nn.Module):
+    def __init__(self, inp, oup, stride, dilation, expand_ratio, BatchNorm):
+        super(InvertedResidual, self).__init__()
+        self.stride = stride
+        assert stride in [1, 2]
+
+        hidden_dim = round(inp * expand_ratio)
+        self.use_res_connect = self.stride == 1 and inp == oup
+        self.kernel_size = 3
+        self.dilation = dilation
+
+        if expand_ratio == 1:
+            self.conv = nn.Sequential(
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 0, dilation, groups=hidden_dim, bias=False),
+                BatchNorm(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, 1, 1, bias=False),
+                BatchNorm(oup),
+            )
+        else:
+            self.conv = nn.Sequential(
+                # pw
+                nn.Conv2d(inp, hidden_dim, 1, 1, 0, 1, bias=False),
+                BatchNorm(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # dw
+                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 0, dilation, groups=hidden_dim, bias=False),
+                BatchNorm(hidden_dim),
+                nn.ReLU6(inplace=True),
+                # pw-linear
+                nn.Conv2d(hidden_dim, oup, 1, 1, 0, 1, bias=False),
+                BatchNorm(oup),
+            )
+
+    def forward(self, x):
+        x_pad = fixed_padding(x, self.kernel_size, dilation=self.dilation)
+        if self.use_res_connect:
+            x = x + self.conv(x_pad)
+        else:
+            x = self.conv(x_pad)
+        return x
+
+
+class MobileNetV2(nn.Module):
+    def __init__(self, output_stride=8, BatchNorm=None, width_mult=1., pretrained=True):
+        super(MobileNetV2, self).__init__()
+        block = InvertedResidual
+        input_channel = 32
+        current_stride = 1
+        rate = 1
+        interverted_residual_setting = [
+            # t, c, n, s
+            [1, 16, 1, 1],
+            [6, 24, 2, 2],
+            [6, 32, 3, 2],
+            [6, 64, 4, 2],
+            [6, 96, 3, 1],
+            [6, 160, 3, 2],
+            [6, 320, 1, 1],
+        ]
+
+        # building first layer
+        input_channel = int(input_channel * width_mult)
+        self.features = [conv_bn(3, input_channel, 2, BatchNorm)]
+        current_stride *= 2
+        # building inverted residual blocks
+        for t, c, n, s in interverted_residual_setting:
+            if current_stride == output_stride:
+                stride = 1
+                dilation = rate
+                rate *= s
+            else:
+                stride = s
+                dilation = 1
+                current_stride *= s
+            output_channel = int(c * width_mult)
+            for i in range(n):
+                if i == 0:
+                    self.features.append(block(input_channel, output_channel, stride, dilation, t, BatchNorm))
+                else:
+                    self.features.append(block(input_channel, output_channel, 1, dilation, t, BatchNorm))
+                input_channel = output_channel
+        self.features = nn.Sequential(*self.features)
+        self._initialize_weights()
+
+        if pretrained:
+            self._load_pretrained_model()
+
+        self.low_level_features = self.features[0:4]
+        self.high_level_features = self.features[4:]
+
+    def forward(self, x):
+        low_level_feat = self.low_level_features(x)
+        x = self.high_level_features(low_level_feat)
+        return x, low_level_feat
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('http://jeff95.me/models/mobilenet_v2-6a65762b.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                # n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                # m.weight.data.normal_(0, math.sqrt(2. / n))
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+if __name__ == "__main__":
+    input = torch.rand(1, 3, 512, 512)
+    model = MobileNetV2(output_stride=16, BatchNorm=nn.BatchNorm2d)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/backbone/resnet.py

@@ -0,0 +1,170 @@
+import math
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, BatchNorm=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
+        self.bn1 = BatchNorm(planes)
+        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
+                               dilation=dilation, padding=dilation, bias=False)
+        self.bn2 = BatchNorm(planes)
+        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
+        self.bn3 = BatchNorm(planes * 4)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self.dilation = dilation
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, output_stride, BatchNorm, pretrained=True):
+        self.inplanes = 64
+        super(ResNet, self).__init__()
+        blocks = [1, 2, 4]
+        if output_stride == 16:
+            strides = [1, 2, 2, 1]
+            dilations = [1, 1, 1, 2]
+        elif output_stride == 8:
+            strides = [1, 2, 1, 1]
+            dilations = [1, 1, 2, 4]
+        else:
+            raise NotImplementedError
+
+        # Modules
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                                bias=False)
+        self.bn1 = BatchNorm(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(block, 64, layers[0], stride=strides[0], dilation=dilations[0], BatchNorm=BatchNorm)
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=strides[1], dilation=dilations[1], BatchNorm=BatchNorm)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=strides[2], dilation=dilations[2], BatchNorm=BatchNorm)
+        self.layer4 = self._make_MG_unit(block, 512, blocks=blocks, stride=strides[3], dilation=dilations[3], BatchNorm=BatchNorm)
+        # self.layer4 = self._make_layer(block, 512, layers[3], stride=strides[3], dilation=dilations[3], BatchNorm=BatchNorm)
+        self._init_weight()
+
+        # if pretrained:
+        #     self._load_pretrained_model()
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, BatchNorm=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, dilation, downsample, BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, dilation=dilation, BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def _make_MG_unit(self, block, planes, blocks, stride=1, dilation=1, BatchNorm=None):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                BatchNorm(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, dilation=blocks[0]*dilation,
+                            downsample=downsample, BatchNorm=BatchNorm))
+        self.inplanes = planes * block.expansion
+        for i in range(1, len(blocks)):
+            layers.append(block(self.inplanes, planes, stride=1,
+                                dilation=blocks[i]*dilation, BatchNorm=BatchNorm))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, input):
+        x = self.conv1(input)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        low_level_feat = x
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x, low_level_feat
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+    def _load_pretrained_model(self):
+
+        import urllib.request
+        import ssl
+        ssl._create_default_https_context = ssl._create_unverified_context
+        response = urllib.request.urlopen('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
+
+        pretrain_dict = model_zoo.load_url('https://download.pytorch.org/models/resnet101-5d3b4d8f.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                # if 'conv1' in k:
+                #     continue 
+                model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+def ResNet101(output_stride, BatchNorm, pretrained=True):
+    """Constructs a ResNet-101 model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    model = ResNet(Bottleneck, [3, 4, 23, 3], output_stride, BatchNorm, pretrained=pretrained)
+    return model
+
+if __name__ == "__main__":
+    import torch
+    model = ResNet101(BatchNorm=nn.BatchNorm2d, pretrained=True, output_stride=8)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/backbone/xception.py

@@ -0,0 +1,288 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.model_zoo as model_zoo
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+def fixed_padding(inputs, kernel_size, dilation):
+    kernel_size_effective = kernel_size + (kernel_size - 1) * (dilation - 1)
+    pad_total = kernel_size_effective - 1
+    pad_beg = pad_total // 2
+    pad_end = pad_total - pad_beg
+    padded_inputs = F.pad(inputs, (pad_beg, pad_end, pad_beg, pad_end))
+    return padded_inputs
+
+
+class SeparableConv2d(nn.Module):
+    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, BatchNorm=None):
+        super(SeparableConv2d, self).__init__()
+
+        self.conv1 = nn.Conv2d(inplanes, inplanes, kernel_size, stride, 0, dilation,
+                               groups=inplanes, bias=bias)
+        self.bn = BatchNorm(inplanes)
+        self.pointwise = nn.Conv2d(inplanes, planes, 1, 1, 0, 1, 1, bias=bias)
+
+    def forward(self, x):
+        x = fixed_padding(x, self.conv1.kernel_size[0], dilation=self.conv1.dilation[0])
+        x = self.conv1(x)
+        x = self.bn(x)
+        x = self.pointwise(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, inplanes, planes, reps, stride=1, dilation=1, BatchNorm=None,
+                 start_with_relu=True, grow_first=True, is_last=False):
+        super(Block, self).__init__()
+
+        if planes != inplanes or stride != 1:
+            self.skip = nn.Conv2d(inplanes, planes, 1, stride=stride, bias=False)
+            self.skipbn = BatchNorm(planes)
+        else:
+            self.skip = None
+
+        self.relu = nn.ReLU(inplace=True)
+        rep = []
+
+        filters = inplanes
+        if grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+            filters = planes
+
+        for i in range(reps - 1):
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(filters, filters, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(filters))
+
+        if not grow_first:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(inplanes, planes, 3, 1, dilation, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if stride != 1:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(planes, planes, 3, 2, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if stride == 1 and is_last:
+            rep.append(self.relu)
+            rep.append(SeparableConv2d(planes, planes, 3, 1, BatchNorm=BatchNorm))
+            rep.append(BatchNorm(planes))
+
+        if not start_with_relu:
+            rep = rep[1:]
+
+        self.rep = nn.Sequential(*rep)
+
+    def forward(self, inp):
+        x = self.rep(inp)
+
+        if self.skip is not None:
+            skip = self.skip(inp)
+            skip = self.skipbn(skip)
+        else:
+            skip = inp
+
+        x = x + skip
+
+        return x
+
+
+class AlignedXception(nn.Module):
+    """
+    Modified Alighed Xception
+    """
+    def __init__(self, output_stride, BatchNorm,
+                 pretrained=True):
+        super(AlignedXception, self).__init__()
+
+        if output_stride == 16:
+            entry_block3_stride = 2
+            middle_block_dilation = 1
+            exit_block_dilations = (1, 2)
+        elif output_stride == 8:
+            entry_block3_stride = 1
+            middle_block_dilation = 2
+            exit_block_dilations = (2, 4)
+        else:
+            raise NotImplementedError
+
+
+        # Entry flow
+        self.conv1 = nn.Conv2d(3, 32, 3, stride=2, padding=1, bias=False)
+        self.bn1 = BatchNorm(32)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(32, 64, 3, stride=1, padding=1, bias=False)
+        self.bn2 = BatchNorm(64)
+
+        self.block1 = Block(64, 128, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False)
+        self.block2 = Block(128, 256, reps=2, stride=2, BatchNorm=BatchNorm, start_with_relu=False,
+                            grow_first=True)
+        self.block3 = Block(256, 728, reps=2, stride=entry_block3_stride, BatchNorm=BatchNorm,
+                            start_with_relu=True, grow_first=True, is_last=True)
+
+        # Middle flow
+        self.block4  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block5  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block6  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block7  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block8  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block9  = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block10 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block11 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block12 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block13 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block14 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block15 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block16 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block17 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block18 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+        self.block19 = Block(728, 728, reps=3, stride=1, dilation=middle_block_dilation,
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=True)
+
+        # Exit flow
+        self.block20 = Block(728, 1024, reps=2, stride=1, dilation=exit_block_dilations[0],
+                             BatchNorm=BatchNorm, start_with_relu=True, grow_first=False, is_last=True)
+
+        self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn3 = BatchNorm(1536)
+
+        self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn4 = BatchNorm(1536)
+
+        self.conv5 = SeparableConv2d(1536, 2048, 3, stride=1, dilation=exit_block_dilations[1], BatchNorm=BatchNorm)
+        self.bn5 = BatchNorm(2048)
+
+        # Init weights
+        self._init_weight()
+
+        # Load pretrained model
+        if pretrained:
+            self._load_pretrained_model()
+
+    def forward(self, x):
+        # Entry flow
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu(x)
+
+        x = self.block1(x)
+        # add relu here
+        x = self.relu(x)
+        low_level_feat = x
+        x = self.block2(x)
+        x = self.block3(x)
+
+        # Middle flow
+        x = self.block4(x)
+        x = self.block5(x)
+        x = self.block6(x)
+        x = self.block7(x)
+        x = self.block8(x)
+        x = self.block9(x)
+        x = self.block10(x)
+        x = self.block11(x)
+        x = self.block12(x)
+        x = self.block13(x)
+        x = self.block14(x)
+        x = self.block15(x)
+        x = self.block16(x)
+        x = self.block17(x)
+        x = self.block18(x)
+        x = self.block19(x)
+
+        # Exit flow
+        x = self.block20(x)
+        x = self.relu(x)
+        x = self.conv3(x)
+        x = self.bn3(x)
+        x = self.relu(x)
+
+        x = self.conv4(x)
+        x = self.bn4(x)
+        x = self.relu(x)
+
+        x = self.conv5(x)
+        x = self.bn5(x)
+        x = self.relu(x)
+
+        return x, low_level_feat
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+                m.weight.data.normal_(0, math.sqrt(2. / n))
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+
+    def _load_pretrained_model(self):
+        pretrain_dict = model_zoo.load_url('http://data.lip6.fr/cadene/pretrainedmodels/xception-b5690688.pth')
+        model_dict = {}
+        state_dict = self.state_dict()
+
+        for k, v in pretrain_dict.items():
+            if k in state_dict:
+                if 'pointwise' in k:
+                    v = v.unsqueeze(-1).unsqueeze(-1)
+                if k.startswith('block11'):
+                    model_dict[k] = v
+                    model_dict[k.replace('block11', 'block12')] = v
+                    model_dict[k.replace('block11', 'block13')] = v
+                    model_dict[k.replace('block11', 'block14')] = v
+                    model_dict[k.replace('block11', 'block15')] = v
+                    model_dict[k.replace('block11', 'block16')] = v
+                    model_dict[k.replace('block11', 'block17')] = v
+                    model_dict[k.replace('block11', 'block18')] = v
+                    model_dict[k.replace('block11', 'block19')] = v
+                elif k.startswith('block12'):
+                    model_dict[k.replace('block12', 'block20')] = v
+                elif k.startswith('bn3'):
+                    model_dict[k] = v
+                    model_dict[k.replace('bn3', 'bn4')] = v
+                elif k.startswith('conv4'):
+                    model_dict[k.replace('conv4', 'conv5')] = v
+                elif k.startswith('bn4'):
+                    model_dict[k.replace('bn4', 'bn5')] = v
+                else:
+                    model_dict[k] = v
+        state_dict.update(model_dict)
+        self.load_state_dict(state_dict)
+
+
+
+if __name__ == "__main__":
+    import torch
+    model = AlignedXception(BatchNorm=nn.BatchNorm2d, pretrained=True, output_stride=16)
+    input = torch.rand(1, 3, 512, 512)
+    output, low_level_feat = model(input)
+    print(output.size())
+    print(low_level_feat.size())

data/MBD/model/deep_lab_model/decoder.py

@@ -0,0 +1,59 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+
+class Decoder(nn.Module):
+    def __init__(self, num_classes, backbone, BatchNorm):
+        super(Decoder, self).__init__()
+        if backbone == 'resnet' or backbone == 'drn':
+            low_level_inplanes = 256
+        elif backbone == 'xception':
+            low_level_inplanes = 128
+        elif backbone == 'mobilenet':
+            low_level_inplanes = 24
+        else:
+            raise NotImplementedError
+
+        self.conv1 = nn.Conv2d(low_level_inplanes, 48, 1, bias=False)
+        self.bn1 = BatchNorm(48)
+        self.relu = nn.ReLU()
+        self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False),
+                                       BatchNorm(256),
+                                       nn.ReLU(),
+                                       nn.Dropout(0.5),
+                                       nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False),
+                                       BatchNorm(256),
+                                       nn.ReLU(),
+                                       nn.Dropout(0.1),
+                                       nn.Conv2d(256, num_classes, kernel_size=1, stride=1),
+                                       nn.Sigmoid()
+                                       )
+        self._init_weight()
+
+
+    def forward(self, x, low_level_feat):
+        low_level_feat = self.conv1(low_level_feat)
+        low_level_feat = self.bn1(low_level_feat)
+        low_level_feat = self.relu(low_level_feat)
+
+        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
+        x = torch.cat((x, low_level_feat), dim=1)
+        x = self.last_conv(x)
+
+        return x
+
+    def _init_weight(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                torch.nn.init.kaiming_normal_(m.weight)
+            elif isinstance(m, SynchronizedBatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+
+def build_decoder(num_classes, backbone, BatchNorm):
+    return Decoder(num_classes, backbone, BatchNorm)

data/MBD/model/deep_lab_model/deeplab.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model.deep_lab_model.sync_batchnorm.batchnorm import SynchronizedBatchNorm2d
+from model.deep_lab_model.aspp import build_aspp
+from model.deep_lab_model.decoder import build_decoder
+from model.deep_lab_model.backbone import build_backbone
+
+class DeepLab(nn.Module):
+    def __init__(self, backbone='resnet', output_stride=16, num_classes=21,
+                 sync_bn=True, freeze_bn=False):
+        super(DeepLab, self).__init__()
+        if backbone == 'drn':
+            output_stride = 8
+
+        if sync_bn == True:
+            BatchNorm = SynchronizedBatchNorm2d
+        else:
+            BatchNorm = nn.BatchNorm2d
+
+        self.backbone = build_backbone(backbone, output_stride, BatchNorm)
+        self.aspp = build_aspp(backbone, output_stride, BatchNorm)
+        self.decoder = build_decoder(num_classes, backbone, BatchNorm)
+
+        self.freeze_bn = freeze_bn
+
+    def forward(self, input):
+        x, low_level_feat = self.backbone(input)
+        x = self.aspp(x)
+        x = self.decoder(x, low_level_feat)
+        x = F.interpolate(x, size=input.size()[2:], mode='bilinear', align_corners=True)
+
+        return x
+
+    def freeze_bn(self):
+        for m in self.modules():
+            if isinstance(m, SynchronizedBatchNorm2d):
+                m.eval()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.eval()
+
+    def get_1x_lr_params(self):
+        modules = [self.backbone]
+        for i in range(len(modules)):
+            for m in modules[i].named_modules():
+                if self.freeze_bn:
+                    if isinstance(m[1], nn.Conv2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+                else:
+                    if isinstance(m[1], nn.Conv2d) or isinstance(m[1], SynchronizedBatchNorm2d) \
+                            or isinstance(m[1], nn.BatchNorm2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+
+    def get_10x_lr_params(self):
+        modules = [self.aspp, self.decoder]
+        for i in range(len(modules)):
+            for m in modules[i].named_modules():
+                if self.freeze_bn:
+                    if isinstance(m[1], nn.Conv2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+                else:
+                    if isinstance(m[1], nn.Conv2d) or isinstance(m[1], SynchronizedBatchNorm2d) \
+                            or isinstance(m[1], nn.BatchNorm2d):
+                        for p in m[1].parameters():
+                            if p.requires_grad:
+                                yield p
+
+if __name__ == "__main__":
+    model = DeepLab(backbone='mobilenet', output_stride=16)
+    model.eval()
+    input = torch.rand(1, 3, 513, 513)
+    output = model(input)
+    print(output.size())
+
+

data/MBD/model/deep_lab_model/sync_batchnorm/__init__.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+# File   : __init__.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d
+from .replicate import DataParallelWithCallback, patch_replication_callback

data/MBD/model/deep_lab_model/sync_batchnorm/batchnorm.py

@@ -0,0 +1,282 @@
+# -*- coding: utf-8 -*-
+# File   : batchnorm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import collections
+
+import torch
+import torch.nn.functional as F
+
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
+
+from .comm import SyncMaster
+
+__all__ = ['SynchronizedBatchNorm1d', 'SynchronizedBatchNorm2d', 'SynchronizedBatchNorm3d']
+
+
+def _sum_ft(tensor):
+    """sum over the first and last dimention"""
+    return tensor.sum(dim=0).sum(dim=-1)
+
+
+def _unsqueeze_ft(tensor):
+    """add new dementions at the front and the tail"""
+    return tensor.unsqueeze(0).unsqueeze(-1)
+
+
+_ChildMessage = collections.namedtuple('_ChildMessage', ['sum', 'ssum', 'sum_size'])
+_MasterMessage = collections.namedtuple('_MasterMessage', ['sum', 'inv_std'])
+
+
+class _SynchronizedBatchNorm(_BatchNorm):
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
+        super(_SynchronizedBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine)
+
+        self._sync_master = SyncMaster(self._data_parallel_master)
+
+        self._is_parallel = False
+        self._parallel_id = None
+        self._slave_pipe = None
+
+    def forward(self, input):
+        # If it is not parallel computation or is in evaluation mode, use PyTorch's implementation.
+        if not (self._is_parallel and self.training):
+            return F.batch_norm(
+                input, self.running_mean, self.running_var, self.weight, self.bias,
+                self.training, self.momentum, self.eps)
+
+        # Resize the input to (B, C, -1).
+        input_shape = input.size()
+        input = input.view(input.size(0), self.num_features, -1)
+
+        # Compute the sum and square-sum.
+        sum_size = input.size(0) * input.size(2)
+        input_sum = _sum_ft(input)
+        input_ssum = _sum_ft(input ** 2)
+
+        # Reduce-and-broadcast the statistics.
+        if self._parallel_id == 0:
+            mean, inv_std = self._sync_master.run_master(_ChildMessage(input_sum, input_ssum, sum_size))
+        else:
+            mean, inv_std = self._slave_pipe.run_slave(_ChildMessage(input_sum, input_ssum, sum_size))
+
+        # Compute the output.
+        if self.affine:
+            # MJY:: Fuse the multiplication for speed.
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std * self.weight) + _unsqueeze_ft(self.bias)
+        else:
+            output = (input - _unsqueeze_ft(mean)) * _unsqueeze_ft(inv_std)
+
+        # Reshape it.
+        return output.view(input_shape)
+
+    def __data_parallel_replicate__(self, ctx, copy_id):
+        self._is_parallel = True
+        self._parallel_id = copy_id
+
+        # parallel_id == 0 means master device.
+        if self._parallel_id == 0:
+            ctx.sync_master = self._sync_master
+        else:
+            self._slave_pipe = ctx.sync_master.register_slave(copy_id)
+
+    def _data_parallel_master(self, intermediates):
+        """Reduce the sum and square-sum, compute the statistics, and broadcast it."""
+
+        # Always using same "device order" makes the ReduceAdd operation faster.
+        # Thanks to:: Tete Xiao (http://tetexiao.com/)
+        intermediates = sorted(intermediates, key=lambda i: i[1].sum.get_device())
+
+        to_reduce = [i[1][:2] for i in intermediates]
+        to_reduce = [j for i in to_reduce for j in i]  # flatten
+        target_gpus = [i[1].sum.get_device() for i in intermediates]
+
+        sum_size = sum([i[1].sum_size for i in intermediates])
+        sum_, ssum = ReduceAddCoalesced.apply(target_gpus[0], 2, *to_reduce)
+        mean, inv_std = self._compute_mean_std(sum_, ssum, sum_size)
+
+        broadcasted = Broadcast.apply(target_gpus, mean, inv_std)
+
+        outputs = []
+        for i, rec in enumerate(intermediates):
+            outputs.append((rec[0], _MasterMessage(*broadcasted[i * 2:i * 2 + 2])))
+
+        return outputs
+
+    def _compute_mean_std(self, sum_, ssum, size):
+        """Compute the mean and standard-deviation with sum and square-sum. This method
+        also maintains the moving average on the master device."""
+        assert size > 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.'
+        mean = sum_ / size
+        sumvar = ssum - sum_ * mean
+        unbias_var = sumvar / (size - 1)
+        bias_var = sumvar / size
+
+        self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data
+        self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data
+
+        return mean, bias_var.clamp(self.eps) ** -0.5
+
+
+class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
+    r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a
+    mini-batch.
+    .. math::
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+    This module differs from the built-in PyTorch BatchNorm1d as the mean and
+    standard-deviation are reduced across all devices during training.
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+    During evaluation, this running mean/variance is used for normalization.
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm
+    Args:
+        num_features: num_features from an expected input of size
+            `batch_size x num_features [x width]`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+    Shape:
+        - Input: :math:`(N, C)` or :math:`(N, C, L)`
+        - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input)
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm1d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 2 and input.dim() != 3:
+            raise ValueError('expected 2D or 3D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm1d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch
+    of 3d inputs
+    .. math::
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+    This module differs from the built-in PyTorch BatchNorm2d as the mean and
+    standard-deviation are reduced across all devices during training.
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+    During evaluation, this running mean/variance is used for normalization.
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+    Shape:
+        - Input: :math:`(N, C, H, W)`
+        - Output: :math:`(N, C, H, W)` (same shape as input)
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm2d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 4:
+            raise ValueError('expected 4D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm2d, self)._check_input_dim(input)
+
+
+class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
+    r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch
+    of 4d inputs
+    .. math::
+        y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta
+    This module differs from the built-in PyTorch BatchNorm3d as the mean and
+    standard-deviation are reduced across all devices during training.
+    For example, when one uses `nn.DataParallel` to wrap the network during
+    training, PyTorch's implementation normalize the tensor on each device using
+    the statistics only on that device, which accelerated the computation and
+    is also easy to implement, but the statistics might be inaccurate.
+    Instead, in this synchronized version, the statistics will be computed
+    over all training samples distributed on multiple devices.
+
+    Note that, for one-GPU or CPU-only case, this module behaves exactly same
+    as the built-in PyTorch implementation.
+    The mean and standard-deviation are calculated per-dimension over
+    the mini-batches and gamma and beta are learnable parameter vectors
+    of size C (where C is the input size).
+    During training, this layer keeps a running estimate of its computed mean
+    and variance. The running sum is kept with a default momentum of 0.1.
+    During evaluation, this running mean/variance is used for normalization.
+    Because the BatchNorm is done over the `C` dimension, computing statistics
+    on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm
+    or Spatio-temporal BatchNorm
+    Args:
+        num_features: num_features from an expected input of
+            size batch_size x num_features x depth x height x width
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Default: 0.1
+        affine: a boolean value that when set to ``True``, gives the layer learnable
+            affine parameters. Default: ``True``
+    Shape:
+        - Input: :math:`(N, C, D, H, W)`
+        - Output: :math:`(N, C, D, H, W)` (same shape as input)
+    Examples:
+        >>> # With Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100)
+        >>> # Without Learnable Parameters
+        >>> m = SynchronizedBatchNorm3d(100, affine=False)
+        >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10))
+        >>> output = m(input)
+    """
+
+    def _check_input_dim(self, input):
+        if input.dim() != 5:
+            raise ValueError('expected 5D input (got {}D input)'
+                             .format(input.dim()))
+        super(SynchronizedBatchNorm3d, self)._check_input_dim(input)

data/MBD/model/deep_lab_model/sync_batchnorm/comm.py

@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+# File   : comm.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import queue
+import collections
+import threading
+
+__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster']
+
+
+class FutureResult(object):
+    """A thread-safe future implementation. Used only as one-to-one pipe."""
+
+    def __init__(self):
+        self._result = None
+        self._lock = threading.Lock()
+        self._cond = threading.Condition(self._lock)
+
+    def put(self, result):
+        with self._lock:
+            assert self._result is None, 'Previous result has\'t been fetched.'
+            self._result = result
+            self._cond.notify()
+
+    def get(self):
+        with self._lock:
+            if self._result is None:
+                self._cond.wait()
+
+            res = self._result
+            self._result = None
+            return res
+
+
+_MasterRegistry = collections.namedtuple('MasterRegistry', ['result'])
+_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result'])
+
+
+class SlavePipe(_SlavePipeBase):
+    """Pipe for master-slave communication."""
+
+    def run_slave(self, msg):
+        self.queue.put((self.identifier, msg))
+        ret = self.result.get()
+        self.queue.put(True)
+        return ret
+
+
+class SyncMaster(object):
+    """An abstract `SyncMaster` object.
+    - During the replication, as the data parallel will trigger an callback of each module, all slave devices should
+    call `register(id)` and obtain an `SlavePipe` to communicate with the master.
+    - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected,
+    and passed to a registered callback.
+    - After receiving the messages, the master device should gather the information and determine to message passed
+    back to each slave devices.
+    """
+
+    def __init__(self, master_callback):
+        """
+        Args:
+            master_callback: a callback to be invoked after having collected messages from slave devices.
+        """
+        self._master_callback = master_callback
+        self._queue = queue.Queue()
+        self._registry = collections.OrderedDict()
+        self._activated = False
+
+    def __getstate__(self):
+        return {'master_callback': self._master_callback}
+
+    def __setstate__(self, state):
+        self.__init__(state['master_callback'])
+
+    def register_slave(self, identifier):
+        """
+        Register an slave device.
+        Args:
+            identifier: an identifier, usually is the device id.
+        Returns: a `SlavePipe` object which can be used to communicate with the master device.
+        """
+        if self._activated:
+            assert self._queue.empty(), 'Queue is not clean before next initialization.'
+            self._activated = False
+            self._registry.clear()
+        future = FutureResult()
+        self._registry[identifier] = _MasterRegistry(future)
+        return SlavePipe(identifier, self._queue, future)
+
+    def run_master(self, master_msg):
+        """
+        Main entry for the master device in each forward pass.
+        The messages were first collected from each devices (including the master device), and then
+        an callback will be invoked to compute the message to be sent back to each devices
+        (including the master device).
+        Args:
+            master_msg: the message that the master want to send to itself. This will be placed as the first
+            message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example.
+        Returns: the message to be sent back to the master device.
+        """
+        self._activated = True
+
+        intermediates = [(0, master_msg)]
+        for i in range(self.nr_slaves):
+            intermediates.append(self._queue.get())
+
+        results = self._master_callback(intermediates)
+        assert results[0][0] == 0, 'The first result should belongs to the master.'
+
+        for i, res in results:
+            if i == 0:
+                continue
+            self._registry[i].result.put(res)
+
+        for i in range(self.nr_slaves):
+            assert self._queue.get() is True
+
+        return results[0][1]
+
+    @property
+    def nr_slaves(self):
+        return len(self._registry)

data/MBD/model/deep_lab_model/sync_batchnorm/replicate.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+# File   : replicate.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import functools
+
+from torch.nn.parallel.data_parallel import DataParallel
+
+__all__ = [
+    'CallbackContext',
+    'execute_replication_callbacks',
+    'DataParallelWithCallback',
+    'patch_replication_callback'
+]
+
+
+class CallbackContext(object):
+    pass
+
+
+def execute_replication_callbacks(modules):
+    """
+    Execute an replication callback `__data_parallel_replicate__` on each module created by original replication.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+    Note that, as all modules are isomorphism, we assign each sub-module with a context
+    (shared among multiple copies of this module on different devices).
+    Through this context, different copies can share some information.
+    We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback
+    of any slave copies.
+    """
+    master_copy = modules[0]
+    nr_modules = len(list(master_copy.modules()))
+    ctxs = [CallbackContext() for _ in range(nr_modules)]
+
+    for i, module in enumerate(modules):
+        for j, m in enumerate(module.modules()):
+            if hasattr(m, '__data_parallel_replicate__'):
+                m.__data_parallel_replicate__(ctxs[j], i)
+
+
+class DataParallelWithCallback(DataParallel):
+    """
+    Data Parallel with a replication callback.
+    An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by
+    original `replicate` function.
+    The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)`
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+        # sync_bn.__data_parallel_replicate__ will be invoked.
+    """
+
+    def replicate(self, module, device_ids):
+        modules = super(DataParallelWithCallback, self).replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+
+def patch_replication_callback(data_parallel):
+    """
+    Monkey-patch an existing `DataParallel` object. Add the replication callback.
+    Useful when you have customized `DataParallel` implementation.
+    Examples:
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallel(sync_bn, device_ids=[0, 1])
+        > patch_replication_callback(sync_bn)
+        # this is equivalent to
+        > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
+        > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
+    """
+
+    assert isinstance(data_parallel, DataParallel)
+
+    old_replicate = data_parallel.replicate
+
+    @functools.wraps(old_replicate)
+    def new_replicate(module, device_ids):
+        modules = old_replicate(module, device_ids)
+        execute_replication_callbacks(modules)
+        return modules
+
+    data_parallel.replicate = new_replicate

data/MBD/model/deep_lab_model/sync_batchnorm/unittest.py

@@ -0,0 +1,29 @@
+# -*- coding: utf-8 -*-
+# File   : unittest.py
+# Author : Jiayuan Mao
+# Email  : maojiayuan@gmail.com
+# Date   : 27/01/2018
+#
+# This file is part of Synchronized-BatchNorm-PyTorch.
+# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch
+# Distributed under MIT License.
+
+import unittest
+
+import numpy as np
+from torch.autograd import Variable
+
+
+def as_numpy(v):
+    if isinstance(v, Variable):
+        v = v.data
+    return v.cpu().numpy()
+
+
+class TorchTestCase(unittest.TestCase):
+    def assertTensorClose(self, a, b, atol=1e-3, rtol=1e-3):
+        npa, npb = as_numpy(a), as_numpy(b)
+        self.assertTrue(
+                np.allclose(npa, npb, atol=atol),
+                'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max())
+        )

data/MBD/model/densenetccnl.py

@@ -0,0 +1,382 @@
+# Densenet decoder encoder with intermediate fully connected layers and dropout
+
+import torch
+import torch.backends.cudnn as cudnn
+import torch.nn as nn
+import torch.nn.functional as F
+import functools
+from torch.autograd import gradcheck
+from torch.autograd import Function
+from torch.autograd import Variable
+from torch.autograd import gradcheck
+from torch.autograd import Function
+import numpy as np
+
+
+def add_coordConv_channels(t):
+    n,c,h,w=t.size()
+    xx_channel=np.ones((h, w))
+    xx_range=np.array(range(h))
+    xx_range=np.expand_dims(xx_range,-1)
+    xx_coord=xx_channel*xx_range
+    yy_coord=xx_coord.transpose()
+
+    xx_coord=xx_coord/(h-1)
+    yy_coord=yy_coord/(h-1)
+    xx_coord=xx_coord*2 - 1
+    yy_coord=yy_coord*2 - 1
+    xx_coord=torch.from_numpy(xx_coord).float()
+    yy_coord=torch.from_numpy(yy_coord).float()
+
+    if t.is_cuda:
+    	xx_coord=xx_coord.cuda()
+    	yy_coord=yy_coord.cuda()
+
+    xx_coord=xx_coord.unsqueeze(0).unsqueeze(0).repeat(n,1,1,1)
+    yy_coord=yy_coord.unsqueeze(0).unsqueeze(0).repeat(n,1,1,1)
+
+    t_cc=torch.cat((t,xx_coord,yy_coord),dim=1)
+
+    return t_cc
+
+
+
+class DenseBlockEncoder(nn.Module):
+    def __init__(self, n_channels, n_convs, activation=nn.ReLU, args=[False]):
+        super(DenseBlockEncoder, self).__init__()
+        assert(n_convs > 0)
+
+        self.n_channels = n_channels
+        self.n_convs    = n_convs
+        self.layers     = nn.ModuleList()
+        for i in range(n_convs):
+            self.layers.append(nn.Sequential(
+                    nn.BatchNorm2d(n_channels),
+                    activation(*args),
+                    nn.Conv2d(n_channels, n_channels, 3, stride=1, padding=1, bias=False),))
+
+    def forward(self, inputs):
+        outputs = []
+
+        for i, layer in enumerate(self.layers):
+            if i > 0:
+                next_output = 0
+                for no in outputs:
+                    next_output = next_output + no
+                outputs.append(next_output)
+            else:
+                outputs.append(layer(inputs))
+        return outputs[-1]
+
+# Dense block in encoder.
+class DenseBlockDecoder(nn.Module):
+    def __init__(self, n_channels, n_convs, activation=nn.ReLU, args=[False]):
+        super(DenseBlockDecoder, self).__init__()
+        assert(n_convs > 0)
+
+        self.n_channels = n_channels
+        self.n_convs    = n_convs
+        self.layers = nn.ModuleList()
+        for i in range(n_convs):
+            self.layers.append(nn.Sequential(
+                    nn.BatchNorm2d(n_channels),
+                    activation(*args),
+                    nn.ConvTranspose2d(n_channels, n_channels, 3, stride=1, padding=1, bias=False),))
+
+    def forward(self, inputs):
+        outputs = []
+
+        for i, layer in enumerate(self.layers):
+            if i > 0:
+                next_output = 0
+                for no in outputs:
+                    next_output = next_output + no
+                outputs.append(next_output)
+            else:
+                outputs.append(layer(inputs))
+        return outputs[-1]
+
+class DenseTransitionBlockEncoder(nn.Module):
+    def __init__(self, n_channels_in, n_channels_out, mp, activation=nn.ReLU, args=[False]):
+        super(DenseTransitionBlockEncoder, self).__init__()
+        self.n_channels_in  = n_channels_in
+        self.n_channels_out = n_channels_out
+        self.mp             = mp
+        self.main           = nn.Sequential(
+                nn.BatchNorm2d(n_channels_in),
+                activation(*args),
+                nn.Conv2d(n_channels_in, n_channels_out, 1, stride=1, padding=0, bias=False),
+                nn.MaxPool2d(mp),
+        )
+    def forward(self, inputs):
+        # print(inputs.shape,'222222222222222',self.main(inputs).shape)
+        return self.main(inputs)
+
+
+class DenseTransitionBlockDecoder(nn.Module):
+    def __init__(self, n_channels_in, n_channels_out, activation=nn.ReLU, args=[False]):
+        super(DenseTransitionBlockDecoder, self).__init__()
+        self.n_channels_in  = n_channels_in
+        self.n_channels_out = n_channels_out
+        self.main           = nn.Sequential(
+                nn.BatchNorm2d(n_channels_in),
+                activation(*args),
+                nn.ConvTranspose2d(n_channels_in, n_channels_out, 4, stride=2, padding=1, bias=False),
+        )
+    def forward(self, inputs):
+        # print(inputs.shape,'333333333333',self.main(inputs).shape)
+        return self.main(inputs)
+
+## Dense encoders and decoders for image of size 128 128
+class waspDenseEncoder128(nn.Module):
+    def __init__(self, nc=1, ndf = 32, ndim = 128, activation=nn.LeakyReLU, args=[0.2, False], f_activation=nn.Tanh, f_args=[]):
+        super(waspDenseEncoder128, self).__init__()
+        self.ndim = ndim
+
+        self.main = nn.Sequential(
+                # input is (nc) x 128 x 128
+                nn.BatchNorm2d(nc),
+                nn.ReLU(True),
+                nn.Conv2d(nc, ndf, 4, stride=2, padding=1),
+
+                # state size. (ndf) x 64 x 64
+                DenseBlockEncoder(ndf, 6),
+                DenseTransitionBlockEncoder(ndf, ndf*2, 2, activation=activation, args=args),
+
+                # state size. (ndf*2) x 32 x 32
+                DenseBlockEncoder(ndf*2, 12),
+                DenseTransitionBlockEncoder(ndf*2, ndf*4, 2, activation=activation, args=args),
+
+                # state size. (ndf*4) x 16 x 16
+                DenseBlockEncoder(ndf*4, 16),
+                DenseTransitionBlockEncoder(ndf*4, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*4) x 8 x 8
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*8) x 4 x 4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndim, 4, activation=activation, args=args),
+                f_activation(*f_args),
+        )
+
+    def forward(self, input):
+        input=add_coordConv_channels(input)
+        output = self.main(input).view(-1,self.ndim)
+        #print(output.size())
+        return output
+
+class waspDenseDecoder128(nn.Module):
+    def __init__(self, nz=128, nc=1, ngf=32, lb=0, ub=1, activation=nn.ReLU, args=[False], f_activation=nn.Hardtanh, f_args=[]):
+        super(waspDenseDecoder128, self).__init__()
+        self.main   = nn.Sequential(
+            # input is Z, going into convolution
+            nn.BatchNorm2d(nz),
+            activation(*args),
+            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
+
+            # state size. (ngf*8) x 4 x 4
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*8),
+
+            # state size. (ngf*4) x 8 x 8
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*4),
+
+            # state size. (ngf*2) x 16 x 16
+            DenseBlockDecoder(ngf*4, 12),
+            DenseTransitionBlockDecoder(ngf*4, ngf*2),
+
+            # state size. (ngf) x 32 x 32
+            DenseBlockDecoder(ngf*2, 6),
+            DenseTransitionBlockDecoder(ngf*2, ngf),
+
+            # state size. (ngf) x 64 x 64
+            DenseBlockDecoder(ngf, 6),
+            DenseTransitionBlockDecoder(ngf, ngf),
+
+            # state size (ngf) x 128 x 128
+            nn.BatchNorm2d(ngf),
+            activation(*args),
+            nn.ConvTranspose2d(ngf, nc, 3, stride=1, padding=1, bias=False),
+            f_activation(*f_args),
+        )
+        # self.smooth=nn.Sequential(
+        #     nn.Conv2d(nc, nc, 1, stride=1, padding=0, bias=False),
+        #     f_activation(*f_args),
+        # )
+    def forward(self, inputs):
+        # return self.smooth(self.main(inputs))
+        return self.main(inputs)
+
+
+
+## Dense encoders and decoders for image of size 512 512
+class waspDenseEncoder512(nn.Module):
+    def __init__(self, nc=1, ndf = 32, ndim = 128, activation=nn.LeakyReLU, args=[0.2, False], f_activation=nn.Tanh, f_args=[]):
+        super(waspDenseEncoder512, self).__init__()
+        self.ndim = ndim
+
+        self.main = nn.Sequential(
+                # input is (nc) x 128 x 128  > *4
+                nn.BatchNorm2d(nc),
+                nn.ReLU(True),
+                nn.Conv2d(nc, ndf, 4, stride=2, padding=1),
+
+                # state size. (ndf) x 64 x 64  > *4
+                DenseBlockEncoder(ndf, 6),
+                DenseTransitionBlockEncoder(ndf, ndf*2, 2, activation=activation, args=args),
+
+                # state size. (ndf*2) x 32 x 32 > *4
+                DenseBlockEncoder(ndf*2, 12),
+                DenseTransitionBlockEncoder(ndf*2, ndf*4, 2, activation=activation, args=args),
+
+                # state size. (ndf*4) x 16 x 16  > *4
+                DenseBlockEncoder(ndf*4, 16),
+                DenseTransitionBlockEncoder(ndf*4, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*8) x 8 x 8   *4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndf*8, 2, activation=activation, args=args),
+
+                # state size. (ndf*8) x 4 x 4 > *4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndf*8, 4, activation=activation, args=args),
+                f_activation(*f_args),
+
+                # state size. (ndf*8) x 2 x 2 > *4
+                DenseBlockEncoder(ndf*8, 16),
+                DenseTransitionBlockEncoder(ndf*8, ndim, 4, activation=activation, args=args),
+                f_activation(*f_args),
+        )
+
+    def forward(self, input):
+        input=add_coordConv_channels(input)
+        output = self.main(input).view(-1,self.ndim)
+        # output = self.main(input).view(8,-1)
+        # print(input.shape,'---------------------')
+        #print(output.size())
+        return output
+
+class waspDenseDecoder512(nn.Module):
+    def __init__(self, nz=128, nc=1, ngf=32, lb=0, ub=1, activation=nn.ReLU, args=[False], f_activation=nn.Tanh, f_args=[]):
+        super(waspDenseDecoder512, self).__init__()
+        self.main   = nn.Sequential(
+            # input is Z, going into convolution
+            nn.BatchNorm2d(nz),
+            activation(*args),
+            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
+
+            # state size. (ngf*8) x 4 x 4
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*8),
+
+            # state size. (ngf*8) x 8 x 8
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*8),
+
+            # state size. (ngf*4) x 16 x 16
+            DenseBlockDecoder(ngf*8, 16),
+            DenseTransitionBlockDecoder(ngf*8, ngf*4),
+
+            # state size. (ngf*2) x 32 x 32
+            DenseBlockDecoder(ngf*4, 12),
+            DenseTransitionBlockDecoder(ngf*4, ngf*2),
+
+            # state size. (ngf) x 64 x 64
+            DenseBlockDecoder(ngf*2, 6),
+            DenseTransitionBlockDecoder(ngf*2, ngf),
+
+            # state size. (ngf) x 128 x 128
+            DenseBlockDecoder(ngf, 6),
+            DenseTransitionBlockDecoder(ngf, ngf),
+
+            # state size. (ngf) x 256 x 256
+            DenseBlockDecoder(ngf, 6),
+            DenseTransitionBlockDecoder(ngf, ngf),
+
+            # state size (ngf) x 512 x 512
+            nn.BatchNorm2d(ngf),
+            activation(*args),
+            nn.ConvTranspose2d(ngf, nc, 3, stride=1, padding=1, bias=False),
+            f_activation(*f_args),
+        )
+        # self.smooth=nn.Sequential(
+        #     nn.Conv2d(nc, nc, 1, stride=1, padding=0, bias=False),
+        #     f_activation(*f_args),
+        # )
+    def forward(self, inputs):
+        # return self.smooth(self.main(inputs))
+        return self.main(inputs)
+
+
+class dnetccnl(nn.Module):
+    #in_channels -> nc      | encoder first layer
+    #filters -> ndf    | encoder first layer
+    #img_size(h,w) -> ndim
+    #out_channels  -> optical flow (x,y)
+
+    def __init__(self, img_size=448, in_channels=3, out_channels=2, filters=32,fc_units=100):
+        super(dnetccnl, self).__init__()
+        self.nc=in_channels
+        self.nf=filters
+        self.ndim=img_size
+        self.oc=out_channels
+        self.fcu=fc_units
+
+        self.encoder=waspDenseEncoder128(nc=self.nc+2,ndf=self.nf,ndim=self.ndim)
+        self.decoder=waspDenseDecoder128(nz=self.ndim,nc=self.oc,ngf=self.nf)
+        # self.fc_layers= nn.Sequential(nn.Linear(self.ndim, self.fcu),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               nn.Linear(self.fcu,self.ndim),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               )
+
+    def forward(self, inputs):
+
+        encoded=self.encoder(inputs)
+        encoded=encoded.unsqueeze(-1).unsqueeze(-1)
+        decoded=self.decoder(encoded)
+        # print torch.max(decoded)
+        # print torch.min(decoded)
+        # print(decoded.shape,'11111111111111111',encoded.shape)
+
+        return decoded
+
+class dnetccnl512(nn.Module):
+    #in_channels -> nc      | encoder first layer
+    #filters -> ndf    | encoder first layer
+    #img_size(h,w) -> ndim
+    #out_channels  -> optical flow (x,y)
+
+    def __init__(self, img_size=448, in_channels=3, out_channels=2, filters=32,fc_units=100):
+        super(dnetccnl512, self).__init__()
+        self.nc=in_channels
+        self.nf=filters
+        self.ndim=img_size
+        self.oc=out_channels
+        self.fcu=fc_units
+
+        self.encoder=waspDenseEncoder512(nc=self.nc+2,ndf=self.nf,ndim=self.ndim)
+        self.decoder=waspDenseDecoder512(nz=self.ndim,nc=self.oc,ngf=self.nf)
+        # self.fc_layers= nn.Sequential(nn.Linear(self.ndim, self.fcu),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               nn.Linear(self.fcu,self.ndim),
+        #                               nn.ReLU(True),
+        #                               nn.Dropout(0.25),
+        #                               )
+
+    def forward(self, inputs):
+
+        encoded=self.encoder(inputs)
+        encoded=encoded.unsqueeze(-1).unsqueeze(-1)
+        decoded=self.decoder(encoded)
+        # print torch.max(decoded)
+        # print torch.min(decoded)
+        # print(decoded.shape,'11111111111111111',encoded.shape)
+
+        return decoded

data/MBD/model/gienet.py

@@ -0,0 +1,742 @@
+from math import log
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+from model.cbam import CBAM
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+class SingleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            # nn.ReflectionPad2d(1),
+            # nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            # nn.BatchNorm2d(out_channels),
+            # nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+class Down_single(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            SingleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+class Up_single(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = SingleConv(in_channels, out_channels)
+        self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
+    def forward(self, x1, x2):
+        x1 = self.deconv(x1)
+        # input is BCHW
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0,stride=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+class Down(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+class Up(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = DoubleConv(in_channels, out_channels)
+        self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
+    def forward(self, x1, x2):
+        x1 = self.deconv(x1)
+        # input is BCHW
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        self.tanh = nn.Tanh()
+        self.hardtanh = nn.Hardtanh()
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x1):
+        x = self.conv(x1)
+        # x = self.sigmoid(x)
+        # x = self.hardtanh(x)
+        # x = (x+1)/2
+        return x
+class GiemaskGenerator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(GiemaskGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(3,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 1)
+        self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        mask = self.outc_1(x_1)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        depth = self.outc(x)
+        return depth,mask
+    """Create a Unet-based generator"""
+class Giemask2Generator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(Giemask2Generator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(3,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 1)
+        self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+        self.outc_2 = OutConv(self.init_channel, 1)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        mask = self.outc_1(x_1)
+        edge = self.outc_2(x_1)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        depth = self.outc(x)
+        return depth,mask,edge
+    """Create a Unet-based generator"""
+class GieGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(GieGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 2)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        logits1 = self.outc(x)
+        return logits1
+
+
+class GiecbamGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(GiecbamGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+        self.cbam = CBAM(gate_channels=self.init_channel*32)
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 2)
+        self.dropout = nn.Dropout(p=0.1)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+        x6 = self.cbam(x6)
+        x = self.up1(x6, x5)
+        x = self.up2(x, x4)
+        x = self.up3(x, x3)
+        x = self.up4(x, x2)
+        x = self.up5(x, x1)
+        x = self.dropout(x)
+        logits1 = self.outc(x)
+        return logits1
+
+
+
+
+class Gie2headGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(Gie2headGenerator, self).__init__()
+        self.init_channel =32
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1_1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+
+        self.up1_2 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2_2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3_2 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4_2 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5_2 = Up(self.init_channel*2, self.init_channel)
+        self.outc_2 = OutConv(self.init_channel, 1)
+
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        logits_1 = self.outc_1(x_1)
+
+        x_2 = self.up1_2(x6, x5)
+        x_2 = self.up2_2(x_2, x4)
+        x_2 = self.up3_2(x_2, x3)
+        x_2 = self.up4_2(x_2, x2)
+        x_2 = self.up5_2(x_2, x1)
+        logits_2 = self.outc_2(x_2)
+
+        logits = torch.cat((logits_1,logits_2),1)
+
+        return logits
+
+
+
+class BmpGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(BmpGenerator, self).__init__()
+        self.init_channel =32
+        self.output_nc = output_nc
+        self.inc = DoubleConv(input_nc,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, self.output_nc)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x = self.up1(x6, x5)
+#        x = self.dropout(x)
+        x = self.up2(x, x4)
+#        x = self.dropout(x)
+        x = self.up3(x, x3)
+#        x = self.dropout(x)
+        x = self.up4(x, x2)
+#        x = self.dropout(x)
+        x = self.up5(x, x1)
+#        x = self.dropout(x)
+        logits1 = self.outc(x)
+        return logits1
+class Bmp2Generator(nn.Module):
+    """Create a Unet-based generator"""
+
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        """Construct a Unet generator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            output_nc (int) -- the number of channels in output images
+            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,
+                                image of size 128x128 will become of size 1x1 # at the bottleneck
+            ngf (int)       -- the number of filters in the last conv layer
+            norm_layer      -- normalization layer
+
+        We construct the U-Net from the innermost layer to the outermost layer.
+        It is a recursive process.
+        """
+        super(Bmp2Generator, self).__init__()
+        #gienet
+        self.init_channel =32
+        self.inc = DoubleConv(3,self.init_channel)
+        self.down1 = Down(self.init_channel, self.init_channel*2)
+        self.down2 = Down(self.init_channel*2, self.init_channel*4)
+        self.down3 = Down(self.init_channel*4, self.init_channel*8)
+        self.down4 = Down(self.init_channel*8, self.init_channel*16)
+        self.down5 = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1 = Up(self.init_channel*32, self.init_channel*16)
+        self.up2 = Up(self.init_channel*16, self.init_channel*8)
+        self.up3 = Up(self.init_channel*8, self.init_channel*4)
+        self.up4 = Up(self.init_channel*4,self.init_channel*2)
+        self.up5 = Up(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, 1)
+        self.up1_1 = Up_single(self.init_channel*32, self.init_channel*16)
+        self.up2_1 = Up_single(self.init_channel*16, self.init_channel*8)
+        self.up3_1 = Up_single(self.init_channel*8, self.init_channel*4)
+        self.up4_1 = Up_single(self.init_channel*4,self.init_channel*2)
+        self.up5_1 = Up_single(self.init_channel*2, self.init_channel)
+        self.outc_1 = OutConv(self.init_channel, 1)
+        self.outc_2 = OutConv(self.init_channel, 1)
+
+        #bpm net
+        self.inc_b = DoubleConv(4,self.init_channel)
+        self.down1_b = Down(self.init_channel, self.init_channel*2)
+        self.down2_b = Down(self.init_channel*2, self.init_channel*4)
+        self.down3_b = Down(self.init_channel*4, self.init_channel*8)
+        self.down4_b = Down(self.init_channel*8, self.init_channel*16)
+        self.down5_b = Down(self.init_channel*16, self.init_channel*32)
+
+        self.up1_b = Up(self.init_channel*32, self.init_channel*16)
+        self.up2_b = Up(self.init_channel*16, self.init_channel*8)
+        self.up3_b = Up(self.init_channel*8, self.init_channel*4)
+        self.up4_b = Up(self.init_channel*4,self.init_channel*2)
+        self.up5_b = Up(self.init_channel*2, self.init_channel)
+        self.outc_b = OutConv(self.init_channel, 2)
+  #      self.dropout = nn.Dropout(p=0.5)
+    def forward(self, input):
+        #gienet
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+
+        x_1 = self.up1_1(x6, x5)
+        x_1 = self.up2_1(x_1, x4)
+        x_1 = self.up3_1(x_1, x3)
+        x_1 = self.up4_1(x_1, x2)
+        x_1 = self.up5_1(x_1, x1)
+        mask = self.outc_1(x_1)
+        edge = self.outc_2(x_1)
+
+        x = self.up1(x6, x5)
+        x = self.up2(x, x4)
+        x = self.up3(x, x3)
+        x = self.up4(x, x2)
+        x = self.up5(x, x1)
+        depth = self.outc(x)
+        
+        #bmpnet
+        mask[mask>0.5]=1.
+        mask[mask<=0.5]=0.
+        image_cat_depth = torch.cat((input*mask,depth*mask),dim=1)
+        x1_b = self.inc_b(image_cat_depth)
+        x2_b = self.down1_b(x1_b)
+        x3_b = self.down2_b(x2_b)
+        x4_b = self.down3_b(x3_b)
+        x5_b = self.down4_b(x4_b)
+        x6_b = self.down5_b(x5_b)
+        x_b = self.up1_b(x6_b, x5_b)
+        x_b = self.up2_b(x_b, x4_b)
+        x_b = self.up3_b(x_b, x3_b)
+        x_b = self.up4_b(x_b, x2_b)
+        x_b = self.up5_b(x_b, x1_b)
+        bm = self.outc_b(x_b)
+        # return depth,mask,edge,bm
+        return bm
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+#class GieGenerator(nn.Module):
+#    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+#                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+#        super(GieGenerator, self).__init__()
+#
+#        # construct unet structure
+#        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+#        for i in range(num_downs - 5):
+#            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+#        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+#        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+#        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+#        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+#
+#        self.model = unet_block
+#
+#    def forward(self, input):
+#        return self.model(input)
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+           # resize = nn.Upsample(scale_factor=2)
+           # conv = nn.Conv2d(inner_nc,outer_nc,kernel_size=4,stride=2,padding=1,bias=use_bias)
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            #up = [uprelu, resize, conv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)
+
+
+
+##===================================================================================================
+class DilatedDoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+class DilatedDown(nn.Module):
+    """Downscaling with maxpool then double conv"""
+
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DilatedDoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+class DilatedUp(nn.Module):
+    """Upscaling then double conv"""
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+        self.up = nn.Upsample(scale_factor=2, mode='nearest')
+        self.conv = DilatedDoubleConv(in_channels, out_channels)
+
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=4,stride=1,dilation=4),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+        # self.deconv = nn.ConvTranspose2d(in_channels, out_channels,kernel_size=4, stride=2,padding=1, bias=True)
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        x1 = self.conv1(x1)
+        # x1 = self.deconv(x1)
+        # input is BCHW
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+class DilatedSingleUnet(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64, biline=True, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(DilatedSingleUnet, self).__init__()
+        self.init_channel = 32
+        self.inc = DilatedDoubleConv(input_nc,self.init_channel)
+        self.down1 = DilatedDown(self.init_channel, self.init_channel*2)
+        self.down2 = DilatedDown(self.init_channel*2, self.init_channel*4)
+        self.down3 = DilatedDown(self.init_channel*4, self.init_channel*8)
+        self.down4 = DilatedDown(self.init_channel*8, self.init_channel*16)
+        self.down5 = DilatedDown(self.init_channel*16, self.init_channel*32)
+        self.cbam = CBAM(gate_channels=self.init_channel*32)
+
+        self.up1 = DilatedUp(self.init_channel*32, self.init_channel*16)
+        self.up2 = DilatedUp(self.init_channel*16, self.init_channel*8)
+        self.up3 = DilatedUp(self.init_channel*8, self.init_channel*4)
+        self.up4 = DilatedUp(self.init_channel*4,self.init_channel*2)
+        self.up5 = DilatedUp(self.init_channel*2, self.init_channel)
+        self.outc = OutConv(self.init_channel, output_nc)
+    def forward(self, input):
+        x1 = self.inc(input)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x6 = self.down5(x5)
+        x6 = self.cbam(x6)
+        x = self.up1(x6, x5)
+        x = self.up2(x, x4)
+        x = self.up3(x, x3)
+        x = self.up4(x, x2)
+        x = self.up5(x, x1)
+        logits1 = self.outc(x)
+        return logits1

data/MBD/model/unetnc.py

@@ -0,0 +1,86 @@
+import torch
+import torch.nn as nn
+from torch.nn import init
+import functools
+
+# Defines the Unet generator.
+# |num_downs|: number of downsamplings in UNet. For example,
+# if |num_downs| == 7, image of size 128x128 will become of size 1x1
+# at the bottleneck
+class UnetGenerator(nn.Module):
+    def __init__(self, input_nc, output_nc, num_downs, ngf=64,
+                 norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetGenerator, self).__init__()
+
+        # construct unet structure
+        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)
+        for i in range(num_downs - 5):
+            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)
+        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)
+
+        self.model = unet_block
+
+    def forward(self, input):
+        return self.model(input)
+
+
+    def forward(self, input):
+        return self.model(input)
+
+# Defines the submodule with skip connection.
+# X -------------------identity---------------------- X
+#   |-- downsampling -- |submodule| -- upsampling --|
+class UnetSkipConnectionBlock(nn.Module):
+    def __init__(self, outer_nc, inner_nc, input_nc=None,
+                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):
+        super(UnetSkipConnectionBlock, self).__init__()
+        self.outermost = outermost
+        if type(norm_layer) == functools.partial:
+            use_bias = norm_layer.func == nn.InstanceNorm2d
+        else:
+            use_bias = norm_layer == nn.InstanceNorm2d
+        if input_nc is None:
+            input_nc = outer_nc
+        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
+                             stride=2, padding=1, bias=use_bias)
+        downrelu = nn.LeakyReLU(0.2, True)
+        downnorm = norm_layer(inner_nc)
+        uprelu = nn.ReLU(True)
+        upnorm = norm_layer(outer_nc)
+
+        if outermost:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1)
+            down = [downconv]
+            up = [uprelu, upconv, nn.Tanh()]
+            model = down + [submodule] + up
+        elif innermost:
+            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv]
+            up = [uprelu, upconv, upnorm]
+            model = down + up
+        else:
+            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
+                                        kernel_size=4, stride=2,
+                                        padding=1, bias=use_bias)
+            down = [downrelu, downconv, downnorm]
+            up = [uprelu, upconv, upnorm]
+
+            if use_dropout:
+                model = down + [submodule] + up + [nn.Dropout(0.5)]
+            else:
+                model = down + [submodule] + up
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        if self.outermost:
+            return self.model(x)
+        else:
+            return torch.cat([x, self.model(x)], 1)

data/MBD/modify_stn_model/stn_head.py

@@ -0,0 +1,123 @@
+from __future__ import absolute_import
+
+import math
+import numpy as np
+import sys
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import init
+
+
+def conv3x3_block(in_planes, out_planes, stride=1):
+  """3x3 convolution with padding"""
+  conv_layer = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1, padding=1)
+
+  block = nn.Sequential(
+    conv_layer,
+    nn.BatchNorm2d(out_planes),
+    nn.ReLU(inplace=True),
+  )
+  return block
+
+
+class STNHead(nn.Module):
+  def __init__(self, in_planes, num_ctrlpoints, activation='none'):
+    super(STNHead, self).__init__()
+
+    self.in_planes = in_planes
+    self.num_ctrlpoints = num_ctrlpoints
+    self.activation = activation
+    self.stn_convnet = nn.Sequential(
+                          conv3x3_block(in_planes, 32), # 32*64
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(32, 64), # 16*32
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(64, 128), # 8*16
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(128, 256), # 4*8
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(256, 256), # 2*4,
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(256, 256)) # 1*2 > 256*8*8
+
+    self.stn_fc1 = nn.Sequential(
+                      # nn.Linear(2*256, 512),
+                      nn.Linear(8*8*256, 512),
+                      nn.BatchNorm1d(512),
+                      nn.ReLU(inplace=True))
+    self.stn_fc2 = nn.Linear(512, num_ctrlpoints*2)
+
+    self.init_weights(self.stn_convnet)
+    self.init_weights(self.stn_fc1)
+    self.init_stn(self.stn_fc2)
+
+  def init_weights(self, module):
+    for m in module.modules():
+      if isinstance(m, nn.Conv2d):
+        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+        m.weight.data.normal_(0, math.sqrt(2. / n))
+        if m.bias is not None:
+          m.bias.data.zero_()
+      elif isinstance(m, nn.BatchNorm2d):
+        m.weight.data.fill_(1)
+        m.bias.data.zero_()
+      elif isinstance(m, nn.Linear):
+        m.weight.data.normal_(0, 0.001)
+        m.bias.data.zero_()
+
+  def init_stn(self, stn_fc2):
+    # margin = 0.01
+    # sampling_num_per_side = int(self.num_ctrlpoints / 2)
+    # ctrl_pts_x = np.linspace(margin, 1.-margin, sampling_num_per_side)
+    # ctrl_pts_y_top = np.ones(sampling_num_per_side) * margin
+    # ctrl_pts_y_bottom = np.ones(sampling_num_per_side) * (1-margin)
+    # ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+    # ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+    # ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0).astype(np.float32)
+  
+    margin_x, margin_y = 0.35,0.35
+    # margin_x, margin_y = 0,0
+    num_ctrl_pts_per_side = (self.num_ctrlpoints-4) // 4 +2
+    ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+    ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+    ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+    ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+    ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+
+    ctrl_pts_x_left = np.ones(num_ctrl_pts_per_side) * margin_x
+    ctrl_pts_x_right = np.ones(num_ctrl_pts_per_side) * (1.0-margin_x)
+    ctrl_pts_left = np.stack([ctrl_pts_x_left[1:-1], ctrl_pts_x[1:-1]], axis=1)
+    ctrl_pts_right = np.stack([ctrl_pts_x_right[1:-1], ctrl_pts_x[1:-1]], axis=1)
+
+    ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom, ctrl_pts_left, ctrl_pts_right],  axis=0).astype(np.float32)
+
+
+    if self.activation is 'none':
+      pass
+    elif self.activation == 'sigmoid':
+      ctrl_points = -np.log(1. / ctrl_points - 1.)
+    stn_fc2.weight.data.zero_()
+    stn_fc2.bias.data = torch.Tensor(ctrl_points).view(-1)
+
+  def forward(self, x):
+    x = self.stn_convnet(x)
+    batch_size, _, h, w = x.size()
+    x = x.view(batch_size, -1)
+    img_feat = self.stn_fc1(x)
+    x = self.stn_fc2(0.1 * img_feat)
+    if self.activation == 'sigmoid':
+      x = F.sigmoid(x)
+    x = x.view(-1, self.num_ctrlpoints, 2)
+    return img_feat, x
+
+
+if __name__ == "__main__":
+  in_planes = 3
+  num_ctrlpoints = 20
+  activation='none' # 'sigmoid'
+  stn_head = STNHead(in_planes, num_ctrlpoints, activation)
+  input = torch.randn(10, 3, 32, 64)
+  control_points = stn_head(input)    
+  print(control_points.size())

data/MBD/modify_stn_model/tps_spatial_transformer.py

@@ -0,0 +1,194 @@
+from __future__ import absolute_import
+
+import numpy as np
+import itertools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def grid_sample(input, grid, canvas = None):
+  output = F.grid_sample(input, grid)
+  if canvas is None:
+    return output
+  else:
+    input_mask = input.data.new(input.size()).fill_(1)
+    output_mask = F.grid_sample(input_mask, grid)
+    padded_output = output * output_mask + canvas * (1 - output_mask)
+    return padded_output
+
+
+# phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
+def compute_partial_repr(input_points, control_points):
+  N = input_points.size(0)
+  M = control_points.size(0)
+  pairwise_diff = input_points.view(N, 1, 2) - control_points.view(1, M, 2)
+  # original implementation, very slow
+  # pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
+  pairwise_diff_square = pairwise_diff * pairwise_diff
+  pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :, 1]
+  repr_matrix = 0.5 * pairwise_dist * torch.log(pairwise_dist)
+  # fix numerical error for 0 * log(0), substitute all nan with 0
+  mask = repr_matrix != repr_matrix
+  repr_matrix.masked_fill_(mask, 0)
+  return repr_matrix
+
+
+# # output_ctrl_pts are specified, according to our task.
+# def build_output_control_points(num_control_points, margins):
+#   margin_x, margin_y = margins
+#   margin_x, margin_y = 0,0
+#   num_ctrl_pts_per_side = num_control_points // 2
+#   ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+#   ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+#   ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+#   ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+#   ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+#   # ctrl_pts_top = ctrl_pts_top[1:-1,:]
+#   # ctrl_pts_bottom = ctrl_pts_bottom[1:-1,:]
+#   output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+#   output_ctrl_pts = torch.Tensor(output_ctrl_pts_arr)
+#   return output_ctrl_pts
+
+# output_ctrl_pts are specified, according to our task.
+# def build_output_control_points(num_control_points, margins):
+#   margin_x, margin_y = margins
+#   # margin_x, margin_y = 0,0
+#   num_ctrl_pts_per_side = (num_control_points-4) // 4 +2
+#   ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+#   ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+#   ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+#   ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+#   ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+
+#   ctrl_pts_x_left = np.ones(num_ctrl_pts_per_side) * margin_x
+#   ctrl_pts_x_right = np.ones(num_ctrl_pts_per_side) * (1.0-margin_x)
+#   ctrl_pts_left = np.stack([ctrl_pts_x_left[1:-1], ctrl_pts_x[1:-1]], axis=1)
+#   ctrl_pts_right = np.stack([ctrl_pts_x_right[1:-1], ctrl_pts_x[1:-1]], axis=1)
+
+#   output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom, ctrl_pts_left, ctrl_pts_right], axis=0)
+#   output_ctrl_pts = torch.Tensor(output_ctrl_pts_arr)
+#   return output_ctrl_pts
+
+def build_output_control_points(num_control_points, margins):
+  points = [0.25,0.5,0.75]
+  pts2 = [[0, 0],[1, 0], [0, 1],[1, 1]]
+  # pts22 = []
+  for ratio in points:
+      pts2.append([1*ratio,0])
+  for ratio in points:
+      pts2.append([1*ratio,1])
+  for ratio in points:
+      pts2.append([0,1*ratio])
+  for ratio in points:
+      pts2.append([1,1*ratio])
+  pts2 = np.float32(pts2)
+  margin_x, margin_y = margins
+  # margin_x, margin_y = 0,0
+  num_ctrl_pts_per_side = (num_control_points-4) // 4 +2
+  ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+  ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+  ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+  ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+  ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+
+  ctrl_pts_x_left = np.ones(num_ctrl_pts_per_side) * margin_x
+  ctrl_pts_x_right = np.ones(num_ctrl_pts_per_side) * (1.0-margin_x)
+  ctrl_pts_left = np.stack([ctrl_pts_x_left[1:-1], ctrl_pts_x[1:-1]], axis=1)
+  ctrl_pts_right = np.stack([ctrl_pts_x_right[1:-1], ctrl_pts_x[1:-1]], axis=1)
+
+  output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom, ctrl_pts_left, ctrl_pts_right], axis=0)
+  # output_ctrl_pts_arr = np.asarray([[0,0],[1,0],[1,1],[0,1],
+  #                                   [],[],[],[],
+  #                                   [],[],[],[],
+  #                                   [],[],[],[]])
+  output_ctrl_pts_arr = pts2
+  # print(output_ctrl_pts_arr.shape,'=================')
+  output_ctrl_pts = torch.FloatTensor(output_ctrl_pts_arr)
+  return output_ctrl_pts
+
+
+
+# demo: ~/test/models/test_tps_transformation.py
+class TPSSpatialTransformer(nn.Module):
+
+  def __init__(self, output_image_size=None, num_control_points=None, margins=None):
+    super(TPSSpatialTransformer, self).__init__()
+    self.output_image_size = output_image_size
+    self.num_control_points = num_control_points
+    self.margins = margins
+
+    self.target_height, self.target_width = output_image_size
+    target_control_points = build_output_control_points(num_control_points, margins)
+    N = num_control_points
+    # N = N - 4
+
+    # create padded kernel matrix
+    forward_kernel = torch.zeros(N + 3, N + 3)
+    target_control_partial_repr = compute_partial_repr(target_control_points, target_control_points)
+    forward_kernel[:N, :N].copy_(target_control_partial_repr)
+    forward_kernel[:N, -3].fill_(1)
+    forward_kernel[-3, :N].fill_(1)
+    forward_kernel[:N, -2:].copy_(target_control_points)
+    forward_kernel[-2:, :N].copy_(target_control_points.transpose(0, 1))
+    # compute inverse matrix
+    # print(forward_kernel.shape)
+    inverse_kernel = torch.inverse(forward_kernel)
+
+    # create target cordinate matrix
+    HW = self.target_height * self.target_width
+    target_coordinate = list(itertools.product(range(self.target_height), range(self.target_width)))
+    target_coordinate = torch.Tensor(target_coordinate) # HW x 2
+    Y, X = target_coordinate.split(1, dim = 1)
+    Y = Y / (self.target_height - 1)
+    X = X / (self.target_width - 1)
+    target_coordinate = torch.cat([X, Y], dim = 1) # convert from (y, x) to (x, y)
+    target_coordinate_partial_repr = compute_partial_repr(target_coordinate, target_control_points)
+    target_coordinate_repr = torch.cat([
+      target_coordinate_partial_repr, torch.ones(HW, 1), target_coordinate
+    ], dim = 1)
+
+    # register precomputed matrices
+    self.register_buffer('inverse_kernel', inverse_kernel)
+    self.register_buffer('padding_matrix', torch.zeros(3, 2))
+    self.register_buffer('target_coordinate_repr', target_coordinate_repr)
+    self.register_buffer('target_control_points', target_control_points)
+
+  def forward(self, input, source_control_points,direction='dewarp'):
+    if direction == 'dewarp':
+      assert source_control_points.ndimension() == 3
+      assert source_control_points.size(1) == self.num_control_points
+      assert source_control_points.size(2) == 2
+      batch_size = source_control_points.size(0)
+
+      Y = torch.cat([source_control_points, self.padding_matrix.expand(batch_size, 3, 2)], 1)
+      mapping_matrix = torch.matmul(self.inverse_kernel, Y)
+      source_coordinate = torch.matmul(self.target_coordinate_repr, mapping_matrix)
+
+      grid = source_coordinate.view(-1, self.target_height, self.target_width, 2)
+      grid = torch.clamp(grid, 0, 1) # the source_control_points may be out of [0, 1].
+      # the input to grid_sample is normalized [-1, 1], but what we get is [0, 1]
+      grid = 2.0 * grid - 1.0
+      output = grid_sample(input, grid, canvas=None)
+      return output, grid
+
+    # elif direction == 'warp':
+    #   target_control_points = source_control_points.clone()
+    #   source_control_points = (build_output_control_points(self.num_control_points, self.margins)).clone()
+    #   source_control_points = source_control_points.unsqueeze(0)
+    #   source_control_points = source_control_points.expand(target_control_points.size(0),self.num_control_points,2)
+    #   assert source_control_points.ndimension() == 3
+    #   assert source_control_points.size(1) == self.num_control_points
+    #   assert source_control_points.size(2) == 2
+    #   batch_size = source_control_points.size(0)
+
+    #   Y = torch.cat([source_control_points.to('cuda'), self.padding_matrix.expand(batch_size, 3, 2)], 1)
+    #   mapping_matrix = torch.matmul(self.inverse_kernel, Y)
+    #   source_coordinate = torch.matmul(self.target_coordinate_repr, mapping_matrix)
+
+    #   grid = source_coordinate.view(-1, self.target_height, self.target_width, 2)
+    #   grid = torch.clamp(grid, 0, 1) # the source_control_points may be out of [0, 1].
+    #   # the input to grid_sample is normalized [-1, 1], but what we get is [0, 1]
+    #   grid = 2.0 * grid - 1.0
+    #   output_maps = grid_sample(input, grid, canvas=None)
+    #   return output_maps, source_coordinate

data/MBD/stn_model/stn_head.py

@@ -0,0 +1,123 @@
+from __future__ import absolute_import
+
+import math
+import numpy as np
+import sys
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.nn import init
+
+
+def conv3x3_block(in_planes, out_planes, stride=1):
+  """3x3 convolution with padding"""
+  conv_layer = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1, padding=1)
+
+  block = nn.Sequential(
+    conv_layer,
+    nn.BatchNorm2d(out_planes),
+    nn.ReLU(inplace=True),
+  )
+  return block
+
+
+class STNHead(nn.Module):
+  def __init__(self, in_planes, num_ctrlpoints, activation='none'):
+    super(STNHead, self).__init__()
+
+    self.in_planes = in_planes
+    self.num_ctrlpoints = num_ctrlpoints
+    self.activation = activation
+    self.stn_convnet = nn.Sequential(
+                          conv3x3_block(in_planes, 32), # 32*64
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(32, 64), # 16*32
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(64, 128), # 8*16
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(128, 256), # 4*8
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(256, 256), # 2*4,
+                          nn.MaxPool2d(kernel_size=2, stride=2),
+                          conv3x3_block(256, 256)) # 1*2 > 256*8*8
+
+    self.stn_fc1 = nn.Sequential(
+                      # nn.Linear(2*256, 512),
+                      nn.Linear(8*8*256, 512),
+                      nn.BatchNorm1d(512),
+                      nn.ReLU(inplace=True))
+    self.stn_fc2 = nn.Linear(512, num_ctrlpoints*2)
+
+    self.init_weights(self.stn_convnet)
+    self.init_weights(self.stn_fc1)
+    self.init_stn(self.stn_fc2)
+
+  def init_weights(self, module):
+    for m in module.modules():
+      if isinstance(m, nn.Conv2d):
+        n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+        m.weight.data.normal_(0, math.sqrt(2. / n))
+        if m.bias is not None:
+          m.bias.data.zero_()
+      elif isinstance(m, nn.BatchNorm2d):
+        m.weight.data.fill_(1)
+        m.bias.data.zero_()
+      elif isinstance(m, nn.Linear):
+        m.weight.data.normal_(0, 0.001)
+        m.bias.data.zero_()
+
+  def init_stn(self, stn_fc2):
+    # margin = 0.01
+    # sampling_num_per_side = int(self.num_ctrlpoints / 2)
+    # ctrl_pts_x = np.linspace(margin, 1.-margin, sampling_num_per_side)
+    # ctrl_pts_y_top = np.ones(sampling_num_per_side) * margin
+    # ctrl_pts_y_bottom = np.ones(sampling_num_per_side) * (1-margin)
+    # ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+    # ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+    # ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0).astype(np.float32)
+  
+    margin_x, margin_y = 0.35,0.35
+    # margin_x, margin_y = 0,0
+    num_ctrl_pts_per_side = (self.num_ctrlpoints-4) // 4 +2
+    ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+    ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+    ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+    ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+    ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+
+    ctrl_pts_x_left = np.ones(num_ctrl_pts_per_side) * margin_x
+    ctrl_pts_x_right = np.ones(num_ctrl_pts_per_side) * (1.0-margin_x)
+    ctrl_pts_left = np.stack([ctrl_pts_x_left[1:-1], ctrl_pts_x[1:-1]], axis=1)
+    ctrl_pts_right = np.stack([ctrl_pts_x_right[1:-1], ctrl_pts_x[1:-1]], axis=1)
+
+    ctrl_points = np.concatenate([ctrl_pts_top, ctrl_pts_bottom, ctrl_pts_left, ctrl_pts_right],  axis=0).astype(np.float32)
+
+
+    if self.activation is 'none':
+      pass
+    elif self.activation == 'sigmoid':
+      ctrl_points = -np.log(1. / ctrl_points - 1.)
+    stn_fc2.weight.data.zero_()
+    stn_fc2.bias.data = torch.Tensor(ctrl_points).view(-1)
+
+  def forward(self, x):
+    x = self.stn_convnet(x)
+    batch_size, _, h, w = x.size()
+    x = x.view(batch_size, -1)
+    img_feat = self.stn_fc1(x)
+    x = self.stn_fc2(0.1 * img_feat)
+    if self.activation == 'sigmoid':
+      x = F.sigmoid(x)
+    x = x.view(-1, self.num_ctrlpoints, 2)
+    return img_feat, x
+
+
+if __name__ == "__main__":
+  in_planes = 3
+  num_ctrlpoints = 20
+  activation='none' # 'sigmoid'
+  stn_head = STNHead(in_planes, num_ctrlpoints, activation)
+  input = torch.randn(10, 3, 32, 64)
+  control_points = stn_head(input)    
+  print(control_points.size())

data/MBD/stn_model/tps_spatial_transformer.py

@@ -0,0 +1,155 @@
+from __future__ import absolute_import
+
+import numpy as np
+import itertools
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def grid_sample(input, grid, canvas = None):
+  output = F.grid_sample(input, grid)
+  if canvas is None:
+    return output
+  else:
+    input_mask = input.data.new(input.size()).fill_(1)
+    output_mask = F.grid_sample(input_mask, grid)
+    padded_output = output * output_mask + canvas * (1 - output_mask)
+    return padded_output
+
+
+# phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
+def compute_partial_repr(input_points, control_points):
+  N = input_points.size(0)
+  M = control_points.size(0)
+  pairwise_diff = input_points.view(N, 1, 2) - control_points.view(1, M, 2)
+  # original implementation, very slow
+  # pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
+  pairwise_diff_square = pairwise_diff * pairwise_diff
+  pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :, 1]
+  repr_matrix = 0.5 * pairwise_dist * torch.log(pairwise_dist)
+  # fix numerical error for 0 * log(0), substitute all nan with 0
+  mask = repr_matrix != repr_matrix
+  repr_matrix.masked_fill_(mask, 0)
+  return repr_matrix
+
+
+# # output_ctrl_pts are specified, according to our task.
+# def build_output_control_points(num_control_points, margins):
+#   margin_x, margin_y = margins
+#   margin_x, margin_y = 0,0
+#   num_ctrl_pts_per_side = num_control_points // 2
+#   ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+#   ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+#   ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+#   ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+#   ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+#   # ctrl_pts_top = ctrl_pts_top[1:-1,:]
+#   # ctrl_pts_bottom = ctrl_pts_bottom[1:-1,:]
+#   output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0)
+#   output_ctrl_pts = torch.Tensor(output_ctrl_pts_arr)
+#   return output_ctrl_pts
+
+# output_ctrl_pts are specified, according to our task.
+def build_output_control_points(num_control_points, margins):
+  margin_x, margin_y = margins
+  # margin_x, margin_y = 0,0
+  num_ctrl_pts_per_side = (num_control_points-4) // 4 +2
+  ctrl_pts_x = np.linspace(margin_x, 1.0 - margin_x, num_ctrl_pts_per_side)
+  ctrl_pts_y_top = np.ones(num_ctrl_pts_per_side) * margin_y
+  ctrl_pts_y_bottom = np.ones(num_ctrl_pts_per_side) * (1.0 - margin_y)
+  ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1)
+  ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1)
+
+  ctrl_pts_x_left = np.ones(num_ctrl_pts_per_side) * margin_x
+  ctrl_pts_x_right = np.ones(num_ctrl_pts_per_side) * (1.0-margin_x)
+  ctrl_pts_left = np.stack([ctrl_pts_x_left[1:-1], ctrl_pts_x[1:-1]], axis=1)
+  ctrl_pts_right = np.stack([ctrl_pts_x_right[1:-1], ctrl_pts_x[1:-1]], axis=1)
+
+  output_ctrl_pts_arr = np.concatenate([ctrl_pts_top, ctrl_pts_bottom, ctrl_pts_left, ctrl_pts_right], axis=0)
+  output_ctrl_pts = torch.Tensor(output_ctrl_pts_arr)
+  return output_ctrl_pts
+
+# demo: ~/test/models/test_tps_transformation.py
+class TPSSpatialTransformer(nn.Module):
+
+  def __init__(self, output_image_size=None, num_control_points=None, margins=None):
+    super(TPSSpatialTransformer, self).__init__()
+    self.output_image_size = output_image_size
+    self.num_control_points = num_control_points
+    self.margins = margins
+
+    self.target_height, self.target_width = output_image_size
+    target_control_points = build_output_control_points(num_control_points, margins)
+    N = num_control_points
+    # N = N - 4
+
+    # create padded kernel matrix
+    forward_kernel = torch.zeros(N + 3, N + 3)
+    target_control_partial_repr = compute_partial_repr(target_control_points, target_control_points)
+    forward_kernel[:N, :N].copy_(target_control_partial_repr)
+    forward_kernel[:N, -3].fill_(1)
+    forward_kernel[-3, :N].fill_(1)
+    forward_kernel[:N, -2:].copy_(target_control_points)
+    forward_kernel[-2:, :N].copy_(target_control_points.transpose(0, 1))
+    # compute inverse matrix
+    # print(forward_kernel.shape)
+    inverse_kernel = torch.inverse(forward_kernel)
+
+    # create target cordinate matrix
+    HW = self.target_height * self.target_width
+    target_coordinate = list(itertools.product(range(self.target_height), range(self.target_width)))
+    target_coordinate = torch.Tensor(target_coordinate) # HW x 2
+    Y, X = target_coordinate.split(1, dim = 1)
+    Y = Y / (self.target_height - 1)
+    X = X / (self.target_width - 1)
+    target_coordinate = torch.cat([X, Y], dim = 1) # convert from (y, x) to (x, y)
+    target_coordinate_partial_repr = compute_partial_repr(target_coordinate, target_control_points)
+    target_coordinate_repr = torch.cat([
+      target_coordinate_partial_repr, torch.ones(HW, 1), target_coordinate
+    ], dim = 1)
+
+    # register precomputed matrices
+    self.register_buffer('inverse_kernel', inverse_kernel)
+    self.register_buffer('padding_matrix', torch.zeros(3, 2))
+    self.register_buffer('target_coordinate_repr', target_coordinate_repr)
+    self.register_buffer('target_control_points', target_control_points)
+
+  def forward(self, input, source_control_points,direction='dewarp'):
+    if direction == 'dewarp':
+      assert source_control_points.ndimension() == 3
+      assert source_control_points.size(1) == self.num_control_points
+      assert source_control_points.size(2) == 2
+      batch_size = source_control_points.size(0)
+
+      Y = torch.cat([source_control_points, self.padding_matrix.expand(batch_size, 3, 2)], 1)
+      mapping_matrix = torch.matmul(self.inverse_kernel, Y)
+      source_coordinate = torch.matmul(self.target_coordinate_repr, mapping_matrix)
+
+      grid = source_coordinate.view(-1, self.target_height, self.target_width, 2)
+      grid = torch.clamp(grid, 0, 1) # the source_control_points may be out of [0, 1].
+      # the input to grid_sample is normalized [-1, 1], but what we get is [0, 1]
+      grid = 2.0 * grid - 1.0
+      output_maps = grid_sample(input, grid, canvas=None)
+      return output_maps, source_coordinate
+
+    # elif direction == 'warp':
+    #   target_control_points = source_control_points.clone()
+    #   source_control_points = (build_output_control_points(self.num_control_points, self.margins)).clone()
+    #   source_control_points = source_control_points.unsqueeze(0)
+    #   source_control_points = source_control_points.expand(target_control_points.size(0),self.num_control_points,2)
+    #   assert source_control_points.ndimension() == 3
+    #   assert source_control_points.size(1) == self.num_control_points
+    #   assert source_control_points.size(2) == 2
+    #   batch_size = source_control_points.size(0)
+
+    #   Y = torch.cat([source_control_points.to('cuda'), self.padding_matrix.expand(batch_size, 3, 2)], 1)
+    #   mapping_matrix = torch.matmul(self.inverse_kernel, Y)
+    #   source_coordinate = torch.matmul(self.target_coordinate_repr, mapping_matrix)
+
+    #   grid = source_coordinate.view(-1, self.target_height, self.target_width, 2)
+    #   grid = torch.clamp(grid, 0, 1) # the source_control_points may be out of [0, 1].
+    #   # the input to grid_sample is normalized [-1, 1], but what we get is [0, 1]
+    #   grid = 2.0 * grid - 1.0
+    #   output_maps = grid_sample(input, grid, canvas=None)
+    #   return output_maps, source_coordinate

data/MBD/tps_grid_gen.py

@@ -0,0 +1,70 @@
+# encoding: utf-8
+
+import torch
+import itertools
+import torch.nn as nn
+from torch.autograd import Function, Variable
+
+class TPSGridGen(nn.Module):
+
+    def __init__(self, target_height, target_width, target_control_points):
+        super(TPSGridGen, self).__init__()
+        assert target_control_points.ndimension() == 2
+        assert target_control_points.size(1) == 2
+        N = target_control_points.size(0)
+        self.num_points = N
+        target_control_points = target_control_points.float()
+
+        # create padded kernel matrix
+        forward_kernel = torch.zeros(N + 3, N + 3)
+        target_control_partial_repr = self.compute_partial_repr(target_control_points, target_control_points)
+        forward_kernel[:N, :N].copy_(target_control_partial_repr)
+        forward_kernel[:N, -3].fill_(1)
+        forward_kernel[-3, :N].fill_(1)
+        forward_kernel[:N, -2:].copy_(target_control_points)
+        forward_kernel[-2:, :N].copy_(target_control_points.transpose(0, 1))
+        # compute inverse matrix
+        inverse_kernel = torch.inverse(forward_kernel)
+
+        # create target cordinate matrix
+        HW = target_height * target_width
+        target_coordinate = list(itertools.product(range(target_height), range(target_width)))
+        target_coordinate = torch.Tensor(target_coordinate) # HW x 2
+        Y, X = target_coordinate.split(1, dim = 1)
+        Y = Y * 2 / (target_height - 1) - 1
+        X = X * 2 / (target_width - 1) - 1
+        target_coordinate = torch.cat([X, Y], dim = 1) # convert from (y, x) to (x, y)
+        target_coordinate_partial_repr = self.compute_partial_repr(target_coordinate, target_control_points)
+        target_coordinate_repr = torch.cat([
+            target_coordinate_partial_repr, torch.ones(HW, 1), target_coordinate
+        ], dim = 1)
+
+        # register precomputed matrices
+        self.register_buffer('inverse_kernel', inverse_kernel)
+        self.register_buffer('padding_matrix', torch.zeros(3, 2))
+        self.register_buffer('target_coordinate_repr', target_coordinate_repr)
+
+    def forward(self, source_control_points):
+        assert source_control_points.ndimension() == 3
+        assert source_control_points.size(1) == self.num_points
+        assert source_control_points.size(2) == 2
+        batch_size = source_control_points.size(0)
+
+        Y = torch.cat([source_control_points, Variable(self.padding_matrix.expand(batch_size, 3, 2))], 1)
+        mapping_matrix = torch.matmul(Variable(self.inverse_kernel), Y)
+        source_coordinate = torch.matmul(Variable(self.target_coordinate_repr), mapping_matrix)
+        return source_coordinate
+    # phi(x1, x2) = r^2 * log(r), where r = ||x1 - x2||_2
+    def compute_partial_repr(self, input_points, control_points):
+        N = input_points.size(0)
+        M = control_points.size(0)
+        pairwise_diff = input_points.view(N, 1, 2) - control_points.view(1, M, 2)
+        # original implementation, very slow
+        # pairwise_dist = torch.sum(pairwise_diff ** 2, dim = 2) # square of distance
+        pairwise_diff_square = pairwise_diff * pairwise_diff
+        pairwise_dist = pairwise_diff_square[:, :, 0] + pairwise_diff_square[:, :, 1]
+        repr_matrix = 0.5 * pairwise_dist * torch.log(pairwise_dist)
+        # fix numerical error for 0 * log(0), substitute all nan with 0
+        mask = repr_matrix != repr_matrix
+        repr_matrix.masked_fill_(mask, 0)
+        return repr_matrix

data/MBD/utils.py

@@ -0,0 +1,234 @@
+'''
+Misc Utility functions
+'''
+from collections import OrderedDict
+import os
+import numpy as np
+import torch
+import random
+import torchvision
+
+def recursive_glob(rootdir='.', suffix=''):
+    """Performs recursive glob with given suffix and rootdir 
+        :param rootdir is the root directory
+        :param suffix is the suffix to be searched
+    """
+    return [os.path.join(looproot, filename)
+        for looproot, _, filenames in os.walk(rootdir)
+        for filename in filenames if filename.endswith(suffix)]
+
+def poly_lr_scheduler(optimizer, init_lr, iter, lr_decay_iter=1, max_iter=30000, power=0.9,):
+    """Polynomial decay of learning rate
+        :param init_lr is base learning rate
+        :param iter is a current iteration
+        :param lr_decay_iter how frequently decay occurs, default is 1
+        :param max_iter is number of maximum iterations
+        :param power is a polymomial power
+
+    """
+    if iter % lr_decay_iter or iter > max_iter:
+        return optimizer
+
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = init_lr*(1 - iter/max_iter)**power
+
+
+def adjust_learning_rate(optimizer, init_lr, epoch):
+    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
+    lr = init_lr * (0.1 ** (epoch // 30))
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def alpha_blend(input_image, segmentation_mask, alpha=0.5):
+    """Alpha Blending utility to overlay RGB masks on RBG images 
+        :param input_image is a np.ndarray with 3 channels
+        :param segmentation_mask is a np.ndarray with 3 channels
+        :param alpha is a float value
+
+    """
+    blended = np.zeros(input_image.size, dtype=np.float32)
+    blended = input_image * alpha + segmentation_mask * (1 - alpha)
+    return blended
+
+def convert_state_dict(state_dict):
+    """Converts a state dict saved from a dataParallel module to normal 
+       module state_dict inplace
+       :param state_dict is the loaded DataParallel model_state
+    
+    """
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        name = k[7:] # remove `module.`
+        new_state_dict[name] = v
+    return new_state_dict
+
+
+class ImagePool():
+    def __init__(self, pool_size):
+        self.pool_size = pool_size
+        if self.pool_size > 0:
+            self.num_imgs = 0
+            self.images = []
+
+    def query(self, images):
+        if self.pool_size == 0:
+            return images
+        return_images = []
+        for image in images:
+            image = torch.unsqueeze(image.data, 0)
+            if self.num_imgs < self.pool_size:
+                self.num_imgs = self.num_imgs + 1
+                self.images.append(image)
+                return_images.append(image)
+            else:
+                p = random.uniform(0, 1)
+                if p > 0.5:
+                    random_id = random.randint(0, self.pool_size - 1)  # randint is inclusive
+                    tmp = self.images[random_id].clone()
+                    self.images[random_id] = image
+                    return_images.append(tmp)
+                else:
+                    return_images.append(image)
+        return_images = torch.cat(return_images, 0)
+        return return_images
+
+
+def set_requires_grad(nets, requires_grad=False):
+    if not isinstance(nets, list):
+        nets = [nets]
+    for net in nets:
+        if net is not None:
+            for param in net.parameters():
+                param.requires_grad = requires_grad
+
+
+
+def get_lr(optimizer):
+    for param_group in optimizer.param_groups:
+        return float(param_group['lr'])
+
+def visualize(epoch,model,layer):    
+    #get conv layers
+    conv_layers=[]
+    for m in model.modules():
+        if isinstance(m,torch.nn.modules.conv.Conv2d):
+            conv_layers.append(m)
+
+    # print conv_layers[layer].weight.data.cpu().numpy().shape
+    tensor=conv_layers[layer].weight.data.cpu()
+    vistensor(tensor, epoch, ch=0, allkernels=False, nrow=8, padding=1)
+
+
+def vistensor(tensor, epoch, ch=0, allkernels=False, nrow=8, padding=1): 
+    '''
+        vistensor: visuzlization tensor
+        @ch: visualization channel 
+        @allkernels: visualization all tensors
+        https://github.com/pedrodiamel/pytorchvision/blob/a14672fe4b07995e99f8af755de875daf8aababb/pytvision/visualization.py#L325
+    ''' 
+    
+    n,c,w,h = tensor.shape
+    if allkernels: tensor = tensor.view(n*c,-1,w,h )
+    elif c != 3: tensor = tensor[:,ch,:,:].unsqueeze(dim=1)
+        
+    rows = np.min( (tensor.shape[0]//nrow + 1, 64 )  )
+    # print rows
+    # print tensor.shape
+    grid = utils.make_grid(tensor, nrow=8, normalize=True, padding=padding)
+    # print grid.shape
+    plt.figure( figsize=(10,10), dpi=200 )
+    plt.imshow(grid.numpy().transpose((1, 2, 0)))
+    plt.savefig('./generated/filters_layer1_dwuv_'+str(epoch)+'.png')
+    plt.close()
+
+
+def show_uloss(uwpred,uworg,inp_img, samples=7):
+    
+    n,c,h,w=inp_img.shape
+    # print(labels.shape)
+    uwpred=uwpred.detach().cpu().numpy()
+    uworg=uworg.detach().cpu().numpy()
+    inp_img=inp_img.detach().cpu().numpy()
+
+    #NCHW->NHWC
+    uwpred=uwpred.transpose((0, 2, 3, 1))
+    uworg=uworg.transpose((0, 2, 3, 1))
+
+    choices=random.sample(range(n), min(n,samples))
+    f, axarr = plt.subplots(samples, 3)
+    for j in range(samples):
+        # print(np.min(labels[j]))
+        # print imgs[j].shape
+        img=inp_img[j].transpose(1,2,0)
+        axarr[j][0].imshow(img[:,:,::-1])
+        axarr[j][1].imshow(uworg[j])
+        axarr[j][2].imshow(uwpred[j])
+    
+    plt.savefig('./generated/unwarp.png')
+    plt.close()
+
+
+def show_uloss_visdom(vis,uwpred,uworg,labels_win,out_win,labelopts,outopts,args):
+    samples=7
+    n,c,h,w=uwpred.shape
+    uwpred=uwpred.detach().cpu().numpy()
+    uworg=uworg.detach().cpu().numpy()
+    out_arr=np.full((samples,3,args.img_rows,args.img_cols),0.0)
+    label_arr=np.full((samples,3,args.img_rows,args.img_cols),0.0)
+    choices=random.sample(range(n), min(n,samples))
+    idx=0
+    for c in choices:
+        out_arr[idx,:,:,:]=uwpred[c]
+        label_arr[idx,:,:,:]=uworg[c]
+        idx+=1
+
+    vis.images(out_arr,
+               win=out_win,
+               opts=outopts)
+    vis.images(label_arr,
+               win=labels_win,
+               opts=labelopts)
+
+def show_unwarp_tnsboard(global_step,writer,uwpred,uworg,grid_samples,gt_tag,pred_tag):
+    idxs=torch.LongTensor(random.sample(range(images.shape[0]), min(grid_samples,images.shape[0])))
+    grid_uworg = torchvision.utils.make_grid(uworg[idxs],normalize=True, scale_each=True)
+    writer.add_image(gt_tag, grid_uworg, global_step)
+    grid_uwpr = torchvision.utils.make_grid(uwpred[idxs],normalize=True, scale_each=True)
+    writer.add_image(pred_tag, grid_uwpr, global_step)
+
+def show_wc_tnsboard(global_step,writer,images,labels, pred, grid_samples,inp_tag, gt_tag, pred_tag):
+    idxs=torch.LongTensor(random.sample(range(images.shape[0]), min(grid_samples,images.shape[0])))
+    grid_inp = torchvision.utils.make_grid(images[idxs],normalize=True, scale_each=True)
+    writer.add_image(inp_tag, grid_inp, global_step)
+    grid_lbl = torchvision.utils.make_grid(labels[idxs],normalize=True, scale_each=True)
+    writer.add_image(gt_tag, grid_lbl, global_step)
+    grid_pred = torchvision.utils.make_grid(pred[idxs],normalize=True, scale_each=True)
+    writer.add_image(pred_tag, grid_pred, global_step)
+def torch2cvimg(tensor,min=0,max=1):
+    '''
+    input:
+        tensor -> torch.tensor BxCxHxW C can be 1,3
+    return
+        im -> ndarray uint8 HxWxC 
+    '''
+    im_list = []
+    for i in range(tensor.shape[0]):
+        im = tensor.detach().cpu().data.numpy()[i]
+        im = im.transpose(1,2,0)
+        im = np.clip(im,min,max)
+        im = ((im-min)/(max-min)*255).astype(np.uint8)
+        im_list.append(im)
+    return im_list
+def cvimg2torch(img,min=0,max=1):
+    '''
+    input:
+        im -> ndarray uint8 HxWxC 
+    return
+        tensor -> torch.tensor BxCxHxW 
+    '''
+    img = img.astype(float) / 255.0
+    img = img.transpose(2, 0, 1) # NHWC -> NCHW
+    img = np.expand_dims(img, 0)
+    img = torch.from_numpy(img).float()
+    return img