from torchgeo1.models import FarSeg
from torchgeo.models import ChangeStarFarSeg
import json
import os
import shutil
import sys
import time
import warnings

from skimage.morphology import binary_dilation, disk
import numpy as np
import segmentation_models_pytorch as smp
import torch
import torch.nn as nn
import torch.nn.functional as F
from loguru import logger
from PIL import Image
from pytorch_toolbelt import losses as L
# import torchcontrib
from torch.nn import DataParallel, CrossEntropyLoss
from torch.optim import Adam
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from segmentation_models_pytorch.losses import DiceLoss, SoftCrossEntropyLoss
import loss
from datasets.change_detection import ChangeDetection
from loss import annealing_softmax_focalloss
from models.model_zoo import get_model
from models.pointrend import point_sample
from utils.image import add_image, visualize, visualize_edge
from utils.metric import F1_score
from utils.options import Options
from utils.palette import color_map
warnings.filterwarnings("ignore")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

class Trainer:
    def __init__(self, args):
        self.args = args
        # 加载数据
        trainset = ChangeDetection(
            root=args.data_root, mode="train", use_pseudo_label=args.use_pseudo_label)
        valset = ChangeDetection(root=args.data_root, mode="val")
        self.trainloader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True,
                                      pin_memory=False, num_workers=4, drop_last=True)
        self.valloader = DataLoader(valset, batch_size=args.val_batch_size, shuffle=False,
                                    pin_memory=True, num_workers=4, drop_last=False)
        # *******************************模型选择*************************************
        # self.model = get_model(args.model, args.backbone, args.pretrained,
        #                        len(trainset.CLASSES)-1, args.lightweight)
        self.model = FarSeg(classes=2, backbone=args.backbone)
        # *******************************模型选择*************************************
        # 加载预训练模型
        if args.pretrain_from:
            self.model.load_state_dict(torch.load(
                args.pretrain_from), strict=False)

        if args.load_from:
            self.model.load_state_dict(torch.load(args.load_from), strict=True)

        # *****************************损失函数的选择********************************
        # 二分类
        weight1 = torch.FloatTensor([0.54397636, 6.18487259]).cuda()
        weight2 = torch.FloatTensor([0.54902432, 5.59951011]).cuda()
        weight_bin = torch.FloatTensor([0.51030887, 24.75096457]).cuda()

        self.criterion1 = CrossEntropyLoss(weight=weight1)
        self.criterion2 = CrossEntropyLoss(weight=weight2)
        self.criterion_bin = CrossEntropyLoss(weight=weight_bin)
        # self.criterion_e1 = CrossEntropyLoss(
        #     torch.FloatTensor([0.50848837, 29.95205391]).cuda())
        # self.criterion_e2 = CrossEntropyLoss(
        #     torch.FloatTensor([0.5088186, 28.84916225]).cuda())
        # self.criterion_ec = CrossEntropyLoss(
        #     torch.FloatTensor([0.50165382, 151.66552401]).cuda())
        self.criterion = annealing_softmax_focalloss
        # *****************************损失函数的选择********************************
        #
        # param = [name for name, param in self.model.named_parameters()]
        # print(param)
        # sys.exit()
        # 设置优化器
        self.optimizer = Adam([{"params": [param for name, param in self.model.named_parameters()
                                           if "backbone" in name], "lr": args.lr},
                               {"params": [param for name, param in self.model.named_parameters()
                                           if "backbone" not in name], "lr": args.lr * 10.0}],
                              lr=args.lr, weight_decay=args.weight_decay)
        # 调整学习率
        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
            self.optimizer,
            T_0=3,  # T_0就是初始restart的epoch数目
            T_mult=2,  # T_mult就是重启之后因子,即每个restart后，T_0 = T_0 + T_0 * T_mult
            eta_min=1e-6  # 最低学习率
        )

        # self.model = self.model.cuda()  # 单GPU训练
        # print(next(self.model.parameters()).device)
        self.iters = 0
        # 迭代总次数
        self.total_iters = len(self.trainloader) * args.epochs
        self.previous_best1 = 0.0
        self.previous_best2 = 0.0
        self.previous_best = 0.0

    def training(self):
        tbar = tqdm(self.trainloader)
        self.model.train()
        total_loss = 0.0
        total_loss_sem = 0.0
        total_loss_bin = 0.0
        for i, (img1, img2, mask1, mask2, mask_bin, id) in enumerate(tbar):
            img1, img2 = img1.cuda(), img2.cuda()
            mask1, mask2, mask_bin = mask1.cuda(), mask2.cuda(
            ), mask_bin.cuda()
            # 模型输出
            # out1, out2, out_bin, out_e1, out_e2, out_ec = self.model(
            #     img1, img2)
            out1, out2, out_bin = self.model(img1, img2)
            # print(out1.shape) torch.Size([bz, 2, 512, 512])
            loss1_1 = self.criterion(
                out1, mask1.long(), self.iters, self.total_iters)
            loss2_1 = self.criterion(
                out2, mask2.long(), self.iters, self.total_iters)
            loss_bin_1 = self.criterion(
                out_bin, mask_bin.long(), self.iters, self.total_iters)

            loss1_2 = self.criterion1(
                out1, mask1.long())
            loss2_2 = self.criterion2(
                out2, mask2.long())
            loss_bin_2 = self.criterion_bin(
                out_bin, mask_bin.long())
            # *****************添加边界约束****************
            # 二值化
            # out1 = torch.argmax(out1, dim=1).cpu().numpy()
            # out2 = torch.argmax(out2, dim=1).cpu().numpy()
            # out_bin = torch.argmax(out_bin, dim=1).cpu().numpy()
            # 转换边界
            # out_edge1 = torch.tensor(
            #     [binary_dilation(e1.squeeze(), selem=disk(2)) - e1 for e1 in out1])
            # out_edge2 = torch.tensor(
            #     [binary_dilation(e2.squeeze(), selem=disk(2)) - e2 for e2 in out2])
            # out_ec = torch.tensor(
            #     [binary_dilation(ec.squeeze(), selem=disk(2)) - ec for ec in out_bin])
            # loss_e1 = self.criterion_e1(out_e1, edge1.long())
            # loss_e2 = self.criterion_e2(out_e2, edge2.long())
            # loss_ec = self.criterion_ec(out_ec, edge_bin.long())

            loss1 = loss1_1 + loss1_2
            loss2 = loss2_1 + loss2_2
            loss_bin = loss_bin_1 + loss_bin_2

            loss1 = loss1.mean()
            loss2 = loss2.mean()
            loss_bin = loss_bin.mean()
            # 总loss
            # loss = loss_bin * 2 + loss1 + loss2
            loss = loss1 + loss2 + loss_bin

            total_loss_sem += loss1.item() + loss2.item()
            total_loss_bin += loss_bin.item()
            total_loss += loss.item()

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()
            self.iters += 1

            # lr = self.args.lr * (1 - self.iters / self.total_iters) ** 0.9
            # self.optimizer.param_groups[0]["lr"] = lr
            # self.optimizer.param_groups[1]["lr"] = lr * 10.0

            tbar.set_description("Loss: %.3f, Semantic Loss: %.3f, Binary Loss: %.3f" %
                                 (total_loss / (i + 1), total_loss_sem / (i + 1), total_loss_bin / (i + 1)))
        train_writer.add_scalar("loss", total_loss_bin, epoch)

    def validation(self):
        tbar = tqdm(self.valloader)
        self.model.eval()
        metric1 = F1_score(num_classes=2)
        metric2 = F1_score(num_classes=2)
        metric_change = F1_score(num_classes=2)

        if self.args.save_mask:
            cmap = color_map()

        with torch.no_grad():
            for i, (img1, img2, mask1, mask2, mask_bin, id) in enumerate(tbar):
                img1, img2 = img1.cuda(), img2.cuda()
                # out1, out2, out_bin = self.model(
                #     img1, img2, self.args.tta)
                out1, out2, out_bin = self.model(img1, img2, True)
            #     mask1, mask2, mask_bin = mask1.cuda(), mask2.cuda(
            # ), mask_bin.cuda(), edge1.cuda(), edge2.cuda(), edge_bin.cuda()
                # out1, out2, out_bin, out_e1, out_e2, out_ec = self.model(
                #     img1, img2)
                out1 = torch.argmax(out1, dim=1).cpu().numpy()
                out2 = torch.argmax(out2, dim=1).cpu().numpy()
                out_bin = torch.argmax(out_bin, dim=1).cpu().numpy()

                # out_e1 = torch.argmax(out_e1, dim=1).cpu().numpy()
                # out_e2 = torch.argmax(out_e2, dim=1).cpu().numpy()
                # out_ec = torch.argmax(out_ec, dim=1).cpu().numpy()

                metric1.add_batch(out1, mask1.cpu().numpy())
                metric2.add_batch(out2, mask2.cpu().numpy())
                metric_change.add_batch(out_bin, mask_bin.cpu().numpy())
            f1_1 = metric1.evaluate()
            f1_2 = metric2.evaluate()
            f1_change = metric_change.evaluate()
            Score_seg = 0.2*(f1_1+f1_2)
            Score_bin = 0.6*f1_change
            Score = Score_seg + Score_bin
            # tbar.set_description("f1_1:%.6f  f1_2:%.6f  f1_change:%.6f  Score:%.6f" %
            #                      (f1_1*100.0, f1_2*100.0, f1_change*100.0, Score*100.0))
            logger.info(
                "f1_1_Score:{}  f1_2_Score:{}  f1_change_Score:{}".format(f1_1, f1_2, f1_change))
            logger.info(
                "Score_seg:{}  Score_bin:{}  Score:{}".format(Score_seg, Score_bin, Score))
            val_writer.add_scalar('Score', Score, epoch)
        if self.args.load_from:
            exit(0)
        Score_seg *= 100.0
        Score_bin *= 100.0
        Score *= 100.0
        save_model_path = checkpoint_dir
        if not os.path.exists(save_model_path):
            os.makedirs(save_model_path)

        if Score_seg >= self.previous_best1:
            if self.previous_best1 != 0:
                model_path = "%s/%s_%s_seg_%.2f.pth" % \
                             (save_model_path, self.args.model,
                              self.args.backbone, self.previous_best1)
                if os.path.exists(model_path):
                    os.remove(model_path)
            torch.save(self.model.state_dict(), "%s/%s_%s_seg_%.2f.pth" %
                       (save_model_path, self.args.model, self.args.backbone, Score_seg), _use_new_zipfile_serialization=False)
            self.previous_best1 = Score_seg
            visualize([img1[0].cpu().numpy(), img2[0].cpu().numpy()],
                      [mask1[0].cpu().numpy(),
                       mask2[0].cpu().numpy()],
                      [out1[0], out2[0]],
                      [mask_bin[0].cpu().numpy(), out_bin[0]],
                      imgs_dir, str(epoch)+"_"+id[0].split(".")[0]+("_seg_%.2f" % (Score_seg)) + ".png")
            # visualize_edge([img1[0].cpu().numpy(), img2[0].cpu().numpy()],
            #                [edge1[0].cpu().numpy(),
            #                 edge2[0].cpu().numpy()],
            #                [out_e1[0], out_e2[0]],
            #                [edge_bin[0].cpu().numpy(), out_ec[0]],
            #                imgs_dir, str(epoch)+"_"+id[0].split(".")[0]+("_edge_%.2f" % (Score_seg)) + ".png")

        if Score_bin >= self.previous_best2:
            if self.previous_best2 != 0:
                model_path = "%s/%s_%s_bin_%.2f.pth" % \
                             (save_model_path, self.args.model,
                              self.args.backbone, self.previous_best2)
                if os.path.exists(model_path):
                    os.remove(model_path)
            torch.save(self.model.module.state_dict(), "%s/%s_%s_bin_%.2f.pth" %
                       (save_model_path, self.args.model, self.args.backbone, Score_bin), _use_new_zipfile_serialization=False)
            self.previous_best2 = Score_bin
            visualize([img1[0].cpu().numpy(), img2[0].cpu().numpy()],
                      [mask1[0].cpu().numpy(),
                       mask2[0].cpu().numpy()],
                      [out1[0], out2[0]],
                      [mask_bin[0].cpu().numpy(), out_bin[0]],
                      imgs_dir, str(epoch)+"_"+id[0].split(".")[0]+("_bin_%.2f" % (Score_bin)) + ".png")
            # visualize_edge([img1[0].cpu().numpy(), img2[0].cpu().numpy()],
            #                [edge1[0].cpu().numpy(),
            #                 edge2[0].cpu().numpy()],
            #                [out_e1[0], out_e2[0]],
            #                [edge_bin[0].cpu().numpy(), out_ec[0]],
            #                imgs_dir, str(epoch)+"_"+id[0].split(".")[0]+("_seg_%.2f" % (Score_seg)) + ".png")
        if Score >= self.previous_best:
            if self.previous_best != 0:
                model_path = "%s/%s_%s_%.2f.pth" % \
                             (save_model_path, self.args.model,
                              self.args.backbone, self.previous_best)
                if os.path.exists(model_path):
                    os.remove(model_path)
            torch.save(self.model.module.state_dict(), "%s/%s_%s_%.2f.pth" %
                       (save_model_path, self.args.model, self.args.backbone, Score), _use_new_zipfile_serialization=False)
            visualize([img1[0].cpu().numpy(), img2[0].cpu().numpy()],
                      [mask1[0].cpu().numpy(),
                       mask2[0].cpu().numpy()],
                      [out1[0], out2[0]],
                      [mask_bin[0].cpu().numpy(), out_bin[0]],
                      imgs_dir, str(epoch)+"_"+id[0].split(".")[0]+("_%.2f" % (Score)) + ".png")
            # visualize_edge([img1[0].cpu().numpy(), img2[0].cpu().numpy()],
            #                [edge1[0].cpu().numpy(),
            #                 edge2[0].cpu().numpy()],
            #                [out_e1[0], out_e2[0]],
            #                [edge_bin[0].cpu().numpy(), out_ec[0]],
            #                imgs_dir, str(epoch)+"_"+id[0].split(".")[0]+("_edge_%.2f" % (Score_seg)) + ".png")
            # if self.args.save_mask:
            # for i in range(out_bin.shape[0]):
            #     mask = Image.fromarray(
            #         out_bin[i].astype(np.uint8), mode="P")
            #     mask.putpalette(cmap)
            #     save_path = os.path.join(imgs_dir, "val")
            #     if not os.path.exists(save_path):
            #         os.makedirs(save_path)
            #     mask.save(save_path + "/"+id[i])
            self.previous_best = Score
            logger.info("epoch:{}  Score:{}".format(epoch, Score))

if __name__ == "__main__":
    # 加载训练参数
    args = Options().parse()
    # 设置实验保存文件夹
    time_str = time.strftime('%Y-%m-%d-%H:%M')
    args.exp_name = '{}_{}_{}_{}'.format(
        args.model, args.backbone, args.exp_name, time_str)
    logger.info(args.exp_name)
    # -------------------- set directory for saving files -------------------------
    # 模型文件路径
    checkpoint_dir = os.path.join(
        'exp_result', args.exp_name, 'checkpoints')
    # 日志记录路径
    logs_dir = os.path.join('exp_result', args.exp_name, 'logs')
    # 生成图像路径
    imgs_dir = os.path.join('exp_result', args.exp_name, 'images')
    # 创建多级目录
    if not os.path.isdir(checkpoint_dir):
        os.makedirs(checkpoint_dir)
    if not os.path.isdir(logs_dir):
        os.makedirs(logs_dir)
    if not os.path.isdir(imgs_dir):
        os.makedirs(imgs_dir)

    # 保存训练参数
    with open('{}/model_cfg.json'.format(logs_dir), 'w') as jfile:
        json.dump(vars(args), jfile, indent=2)

    # 设置可视化训练过程
    train_writer = SummaryWriter(os.path.join(
        logs_dir, 'runs', 'training'))
    val_writer = SummaryWriter(os.path.join(
        logs_dir, 'runs', 'val'))
    trainer = Trainer(args)

    if args.load_from:
        trainer.validation()
    # 日志记录文件
    logger.add("{}/val_log.log".format(logs_dir))

    for epoch in range(args.epochs):
        logger.info("\n")
        logger.info("\n==> Epoches %i/%i, learning rate = %.6f\t\t\t\t previous best = %.2f" %
                    (epoch, args.epochs, trainer.optimizer.param_groups[0]["lr"], trainer.previous_best))
        # 训练
        trainer.training()
        # 验证
        trainer.validation()
        # 更新学习率
        trainer.scheduler.step()


from keras.models import load_model
import cv2
import numpy as np

model = load_model('model-013.model')

face_clsfr=cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

source=cv2.VideoCapture(0)

labels_dict={0:'MASK',1:'NO MASK'}
color_dict={0:(0,255,0),1:(0,0,255)}

while(True):

    ret,img=source.read()
    gray=cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    faces=face_clsfr.detectMultiScale(gray,1.3,5)

    for x,y,w,h in faces:

        face_img=gray[y:y+w,x:x+w]
        resized=cv2.resize(face_img,(100,100))
        normalized=resized/255.0
        reshaped=np.reshape(normalized,(1,100,100,1))
        result=model.predict(reshaped)

        label=np.argmax(result,axis=1)[0]

        cv2.rectangle(img,(x,y),(x+w,y+h),color_dict[label],2)
        cv2.rectangle(img,(x,y-40),(x+w,y),color_dict[label],-1)
        cv2.putText(img, labels_dict[label], (x, y-10),cv2.FONT_HERSHEY_SIMPLEX,0.8,(255,255,255),2)

    cv2.imshow('LIVE',img)
    key=cv2.waitKey(1)

    if(key==27):
        break

cv2.destroyAllWindows()
source.release()


# -*- coding:utf-8 -*-

import os
import torch
import matplotlib.pyplot as plt
import numpy as np
import io

from torchvision import transforms
from torchsummary.torchsummary import summary
from PIL import Image, ImageDraw
from utilities.utils import visualize_GT
import yolov2
np.set_printoptions(precision=4, suppress=True)

def test(params):

    input_height = params["input_height"]
    input_width = params["input_width"]

    data_path = params["data_path"]
    datalist_path = params["datalist_path"]
    class_path = params["class_path"]
    num_gpus = [i for i in range(1)]
    checkpoint_path = params["checkpoint_path"]
    USE_GTCHECKER = params["use_gtcheck"]

    num_class = params["num_class"]
    num_achors = 5

    with open(class_path) as f:
        class_list = f.read().splitlines()

    anchors = torch.tensor([[1.3221, 1.73145], [3.19275, 4.00944], [5.05587, 8.09892], [9.47112, 4.84053],
               [11.2364, 10.0071]])
    anchor = anchors[:,:] / 13
    anchor = anchor.cpu().data.numpy()
    objness_threshold = 0.1
    class_threshold = 0.1

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

    net = yolov2.YOLOv2()
    # model = torch.nn.DataParallel(net, device_ids=num_gpus).cuda()
    print("device : ", device)
    if device is "cpu":
        model = torch.nn.DataParallel(net)
    else:
        model = torch.nn.DataParallel(net, device_ids=num_gpus).cuda()

    model.load_state_dict(torch.load(checkpoint_path)["state_dict"])
    # model.load_state_dict(torch.load("/home/madhevan/Yolo_v2/ckpt_noHash_ep00100_loss12.2638_lr0.0001.pth.tar"))
    model.eval()

    # https://discuss.pytorch.org/t/loading-weights-from-dataparallel-models/20570/3

    image_path = os.path.join(data_path, "JPEGImages")

    if not (datalist_path =='./'):
        root = next(os.walk(os.path.abspath(data_path)))[0]
        dir = next(os.walk(os.path.abspath(data_path)))[1]
        #files =['000001','000002','000003','000004','000005','000006','000007','000008','000009','000011','000012','000013','000014','000015','000016','000017','000018','000019']
        #files = next(os.walk(os.path.abspath(data_path)))[2]
        files =[]
        with io.open(datalist_path,encoding='utf8') as f:
         for i in f.readlines():
            files.append(i.splitlines()[0])
        for idx in range(len(files)):
            files[idx] += '.jpg'

    else:
        root, dir, files = next(os.walk(os.path.abspath(image_path)))

    for file in files:
        extension = file.split(".")[-1]
        if extension not in ["jpeg", "jpg", "png", "JPEG", "JPG", "PNG"]:
            continue

        img = Image.open(os.path.join(image_path, file)).convert('RGB')

        # PRE-PROCESSING
        input_img = img.resize((input_width, input_height))
        input_img = transforms.ToTensor()(input_img)
        c, w, h = input_img.shape
        input_img = input_img.view(1, c, w, h)
        input_img = input_img.to(device)

        # INVERSE TRANSFORM IMAGE########
        # inverseTimg = transforms.ToPILImage()(input_img)
        W, H = img.size
        draw = ImageDraw.Draw(img)

        dx = W // 13
        dy = H // 13
        ##################################

        # INFERENCE
        outputs = model(input_img)

        # outputs[:,:,:,0:3] = torch.sigmoid(outputs[:,:,:,0:3])
        # outputs[:,:,:,25:28] = torch.sigmoid(outputs[:,:,:,25:28])
        # outputs[:,:,:,50:53] = torch.sigmoid(outputs[:,:,:,50:53])
        # outputs[:,:,:,75:78] = torch.sigmoid(outputs[:,:,:,75:78])
        # outputs[:,:,:,100:103] = torch.sigmoid(outputs[:,:,:,100:103])

        b, w, h, c = outputs.shape
        outputs = outputs.view(w, h, c)
        threshold_map = torch.zeros(13,13,20,dtype=outputs.dtype,device=outputs.device)
        class_prob = torch.zeros(13,13,5,28,dtype=outputs.dtype,device=outputs.device)
        # class_prob = torch.zeros(13,13,25,28,dtype=outputs.dtype,device=outputs.device) ## TODO need to check this for the threshold error...
        pred_bbox = torch.zeros(outputs.shape[0],outputs.shape[1],num_achors,num_class,dtype=outputs.dtype,device=outputs.device)
        final_bbox = nms(outputs,class_prob,class_threshold,num_class,img.size,anchor)

        ###################################################################################
        #mAP 측정
        '''
        test_dataset = VOC(root=data_path, transform=composed, class_path=class_path, datalist_path=datalist_path)

        test_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                               batch_size=1,
                                               shuffle=True,
                                               collate_fn=detection_collate)

        for i, (images, labels, sizes) in enumerate(train_loader):
            # labels.shape = 1 x 13 x 13 x 6
            labels = labels.to(device)

        precision, recall = compute_mAP(final_bbox,gt_bbox)
        '''
        ###################################################################################

        print("IMAGE SIZE")
        print("width : {}, height : {}".format(W, H))
        print("\n\n\n\n")
        #print(outputs[4,7,25:50])
        #print(outputs[6,6,75:100])

        for i in range(13):
            for j in range(13):
                draw.rectangle(((dx * i, dy * j), (dx * i + dx, dy * j + dy)), outline='#00ff88')

        #if  one >= objness_threshold:
        for idx in final_bbox:
            i = idx[26]
            j = idx[25]
            anchor_idx = int(idx[27])
            x_start_point = dx * i
            y_start_point = dy * j
            center_x = sigmoid(idx[1]) * dx + x_start_point
            center_y = sigmoid(idx[2]) * dx + y_start_point
            w_ratio = idx[3]
            h_ratio = idx[4]
            final_w_ratio = anchor[anchor_idx,0] * np.exp(w_ratio)
            final_h_ratio = anchor[anchor_idx,1] * np.exp(h_ratio)
            # final_w_ratio = final_w_ratio * final_w_ratio
            # final_h_ratio = final_h_ratio * final_h_ratio
            width = int(final_w_ratio * W)
            height = int(final_h_ratio * H)
            #width = int(final_w_ratio * W)
            #height = int(final_w_ratio * H)
            xmin = center_x - (width // 2)
            ymin = center_y - (height // 2)
            xmax = xmin + width
            ymax = ymin + height
            draw.rectangle(((xmin + 2, ymin + 2), (xmax - 2, ymax - 2)), outline="blue")
            #print(np.argmax(idx[5:5+num_class]))
            #print(5 + np.argmax(idx[5:5+num_class]))
            draw.text((xmin + 5, ymin + 5), "{}: {:.2f}".format(class_list[np.argmax(idx[5:5+num_class])], idx[5 + np.argmax(idx[5:5+num_class])]))
            draw.ellipse(((center_x - 2, center_y - 2),
                          (center_x + 2, center_y + 2)),
                         fill='blue')

                        # LOG
        plt.figure(figsize=(8, 8))
        plt.imshow(img)
        plt.show()
        plt.close()

def nms(outputs,class_prob,class_threshold,num_class,img_size,anchor):
    import pdb;pdb.set_trace()
    outputs_np = outputs.cpu().data.numpy()
    final_bbox = []
    for j in range(outputs.shape[0]):
        for i in range(outputs.shape[1]):
            output1 = outputs[i,j,:25]
            output2 = outputs[i,j,25:50]
            output3 = outputs[i,j,50:75]
            output4 = outputs[i,j,75:100]
            output5 = outputs[i,j,100:]

            class_prob[i,j,0,:5] = output1[:5]
            class_prob[i,j,1,:5] = output2[:5]
            class_prob[i,j,2,:5] = output3[:5]
            class_prob[i,j,3,:5] = output4[:5]
            class_prob[i,j,4,:5] = output5[:5]
            # obj_prob x class_prob
            class_prob[i,j,0,5:25] = output1[0] * output1[5:]
            class_prob[i,j,1,5:25] = output2[0] * output2[5:]
            class_prob[i,j,2,5:25] = output3[0] * output3[5:]
            class_prob[i,j,3,5:25] = output4[0] * output4[5:]
            class_prob[i,j,4,5:25] = output5[0] * output5[5:]
            class_prob[i,j,:,25] = i
            class_prob[i,j,:,26] = j
            class_prob[i,j,0,27] = 0
            class_prob[i,j,1,27] = 1
            class_prob[i,j,2,27] = 2
            class_prob[i,j,3,27] = 3
            class_prob[i,j,4,27] = 4

    class_prob_view = class_prob.view(-1,28)
    class_prob_view_np = class_prob_view.cpu().data.numpy()

    for i in range(num_class):
        leaderboard = 1
        idx = np.argsort(class_prob_view_np[:,5+ i])[::-1]
        sort_class_prob = class_prob_view_np[idx]
        test_threshold = sort_class_prob[0,i+5]
        if(test_threshold >= class_threshold):
            # import pdb;pdb.set_trace()
            while(sort_class_prob[leaderboard,i+5] >= class_threshold):
                iou = compute_iou(sort_class_prob[0,:],sort_class_prob[leaderboard,:],img_size,anchor)
                if iou < 0.5:
                    final_bbox.append(sort_class_prob[leaderboard,:])
                leaderboard +=1

            # print(i)
            final_bbox.append(sort_class_prob[0,:])

            print("final box count")
            print(len(final_bbox))
            print(final_bbox.shape)
            print(final_bbox)
    return final_bbox
    #output = 13 x 13 x 25 x n

def compute_iou(pred_bbox, true_bbox, img_size, anchor):
    #truexmin, truexmax, trueymin, trueymax , predboxxmin, predboxxmax, predboxymin,
    #predboxymax
    w, h = img_size

    pred_achor_idx = int(pred_bbox[27])
    pred_bbox_w = anchor[pred_achor_idx,0] * np.exp(pred_bbox[3])
    pred_bbox_h = anchor[pred_achor_idx,1] * np.exp(pred_bbox[4])

    true_achor_idx = int(true_bbox[27])
    true_bbox_w = anchor[true_achor_idx,0] * np.exp(true_bbox[3])
    true_bbox_h = anchor[true_achor_idx,1] * np.exp(true_bbox[4])

    pred_bbox_c_x = ( pred_bbox[1] + pred_bbox[26] ) * w
    pred_bbox_c_y = ( pred_bbox[2] + pred_bbox[25] ) * h
    pred_bbox_w = pred_bbox_w #* pred_bbox_w
    pred_bbox_h = pred_bbox_h #* pred_bbox_h

    pred_bbox_xmin = pred_bbox_c_x - pred_bbox_w/2
    pred_bbox_xmax = pred_bbox_c_x + pred_bbox_w/2
    pred_bbox_ymin = pred_bbox_c_y - pred_bbox_h/2
    pred_bbox_ymax = pred_bbox_c_y + pred_bbox_h/2

    true_bbox_c_x = ( true_bbox[1] + true_bbox[26] ) * w
    true_bbox_c_y = ( true_bbox[2] + true_bbox[25] ) * h
    true_bbox_w = true_bbox_w #* true_bbox_w
    true_bbox_h = true_bbox_h #* true_bbox_h

    true_bbox_xmin = true_bbox_c_x - true_bbox_w/2
    true_bbox_xmax = true_bbox_c_x + true_bbox_w/2
    true_bbox_ymin = true_bbox_c_y - true_bbox_h/2
    true_bbox_ymax = true_bbox_c_y + true_bbox_h/2

    pred_bbox_area = (pred_bbox_xmax - pred_bbox_xmin + 1) * (pred_bbox_ymax - pred_bbox_ymin + 1)
    true_bbox_area = (true_bbox_xmax - true_bbox_xmin + 1) * (true_bbox_ymax - true_bbox_ymin + 1)

    inter_x_min = max(true_bbox_xmin, pred_bbox_xmin)
    inter_y_min = max(true_bbox_ymax, pred_bbox_ymax)
    inter_x_max = min(true_bbox_xmax, pred_bbox_xmax)
    inter_y_max = min(true_bbox_ymax, pred_bbox_ymax)

    inter_area = max(0,inter_x_max - inter_x_min + 1) * max(0,inter_y_max - inter_y_min + 1)

    iou = inter_area / float(pred_bbox_area + true_bbox_area - inter_area)

    return iou
def sigmoid(x, derivative=False):
  return x*(1-x) if derivative else 1/(1+np.exp(-x))


import mxnet as mx
import mxnet.ndarray as nd
import numpy as np
from rcnn.config import config
from rcnn.PY_OP import rpn_fpn_ohem3
from rcnn.PY_OP import focal_loss
from rcnn.symbol.losses import giou_loss, focal_loss_mx

def conv_only(from_layer, name, num_filter, kernel=(1,1), pad=(0,0), \
    stride=(1,1), bias_wd_mult=0.0, shared_weight=None, shared_bias = None):
  if shared_weight is None:
    weight = mx.symbol.Variable(name="{}_weight".format(name),
        init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
    bias = mx.symbol.Variable(name="{}_bias".format(name),
        init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(bias_wd_mult)})
  else:
    weight = shared_weight
    bias = shared_bias
    print('reuse shared var in', name)
  conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad, \
      stride=stride, num_filter=num_filter, name="{}".format(name), weight = weight, bias=bias)
  return conv

def conv_deformable(net, num_filter, num_group=1, act_type='relu',name=''):
  if config.USE_DCN == 1:
    f = num_group*18
    conv_offset = mx.symbol.Convolution(name=name+'_conv_offset', data = net,
                        num_filter=f, pad=(1, 1), kernel=(3, 3), stride=(1, 1))
    net = mx.contrib.symbol.DeformableConvolution(name=name+"_conv", data=net, offset=conv_offset,
                        num_filter=num_filter, pad=(1,1), kernel=(3, 3), num_deformable_group=num_group, stride=(1, 1), no_bias=False)
  else:
    print('use dcnv2 at', name)
    lr_mult = 0.1
    weight_var = mx.sym.Variable(name=name+'_conv2_offset_weight', init=mx.init.Zero(), lr_mult=lr_mult)
    bias_var = mx.sym.Variable(name=name+'_conv2_offset_bias', init=mx.init.Zero(), lr_mult=lr_mult)
    conv2_offset = mx.symbol.Convolution(name=name + '_conv2_offset', data=net, num_filter=27,
      pad=(1, 1), kernel=(3, 3), stride=(1,1), weight=weight_var, bias=bias_var, lr_mult=lr_mult)
    conv2_offset_t = mx.sym.slice_axis(conv2_offset, axis=1, begin=0, end=18)
    conv2_mask =  mx.sym.slice_axis(conv2_offset, axis=1, begin=18, end=None)
    conv2_mask = 2 * mx.sym.Activation(conv2_mask, act_type='sigmoid')

    conv2 = mx.contrib.symbol.ModulatedDeformableConvolution(name=name + '_conv2', data=net, offset=conv2_offset_t, mask=conv2_mask,
        num_filter=num_filter, pad=(1, 1), kernel=(3, 3), stride=(1,1),
        num_deformable_group=num_group, no_bias=True)
    net = conv2
  net = mx.sym.BatchNorm(data=net, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
  if len(act_type)>0:
    net = mx.symbol.Activation(data=net, act_type=act_type, name=name+'_act')
  return net

def conv_act_layer_dw(from_layer, name, num_filter, kernel=(1,1), pad=(0,0), \
    stride=(1,1), act_type="relu", bias_wd_mult=0.0):
    assert kernel[0]==3
    weight = mx.symbol.Variable(name="{}_weight".format(name),
        init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
    bias = mx.symbol.Variable(name="{}_bias".format(name),
        init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(bias_wd_mult)})
    conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad, \
        stride=stride, num_filter=num_filter, num_group=num_filter, name="{}".format(name), weight=weight, bias=bias)
    conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    if len(act_type)>0:
      relu = mx.symbol.Activation(data=conv, act_type=act_type, \
          name="{}_{}".format(name, act_type))
    else:
      relu = conv
    return relu

def conv_act_layer(from_layer, name, num_filter, kernel=(1,1), pad=(0,0), \
    stride=(1,1), act_type="relu", bias_wd_mult=0.0, separable=False, filter_in = -1):

    if config.USE_DCN>1 and kernel==(3,3) and pad==(1,1) and stride==(1,1) and not separable:
      return conv_deformable(from_layer, num_filter, num_group=1, act_type = act_type, name=name)

    if separable:
      assert kernel[0]>1
      assert filter_in>0
    if not separable:
      weight = mx.symbol.Variable(name="{}_weight".format(name),
          init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
      bias = mx.symbol.Variable(name="{}_bias".format(name),
          init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(bias_wd_mult)})
      conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad, \
          stride=stride, num_filter=num_filter, name="{}".format(name), weight=weight, bias=bias)
      conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    else:
      if filter_in<0:
        filter_in = num_filter
      conv = mx.symbol.Convolution(data=from_layer, kernel=kernel, pad=pad, \
          stride=stride, num_filter=filter_in, num_group=filter_in, name="{}_sep".format(name))
      conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_sep_bn')
      conv = mx.symbol.Activation(data=conv, act_type='relu', \
          name="{}_sep_bn_relu".format(name))
      conv = mx.symbol.Convolution(data=conv, kernel=(1,1), pad=(0,0), \
          stride=(1,1), num_filter=num_filter, name="{}".format(name))
      conv = mx.sym.BatchNorm(data=conv, fix_gamma=False, eps=2e-5, momentum=0.9, name=name + '_bn')
    if len(act_type)>0:
      relu = mx.symbol.Activation(data=conv, act_type=act_type, \
          name="{}_{}".format(name, act_type))
    else:
      relu = conv
    return relu

def ssh_context_module(body, num_filter, filter_in, name):
  conv_dimred = conv_act_layer(body, name+'_conv1',
      num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=False, filter_in = filter_in)
  conv5x5 = conv_act_layer(conv_dimred, name+'_conv2',
      num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='', separable=False)
  conv7x7_1 = conv_act_layer(conv_dimred, name+'_conv3_1',
      num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=False)
  conv7x7 = conv_act_layer(conv7x7_1, name+'_conv3_2',
      num_filter, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='', separable=False)
  return (conv5x5, conv7x7)

def ssh_detection_module(body, num_filter, filter_in, name):
  assert num_filter%4==0
  conv3x3 = conv_act_layer(body, name+'_conv1',
      num_filter//2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='', separable=False, filter_in=filter_in)
  #_filter = max(num_filter//4, 16)
  _filter = num_filter//4
  conv5x5, conv7x7 = ssh_context_module(body, _filter, filter_in, name+'_context')
  ret = mx.sym.concat(*[conv3x3, conv5x5, conv7x7], dim=1, name = name+'_concat')
  ret = mx.symbol.Activation(data=ret, act_type='relu', name=name+'_concat_relu')
  out_filter = num_filter//2+_filter*2
  if config.USE_DCN>0:
    ret = conv_deformable(ret, num_filter = out_filter, name = name+'_concat_dcn')
  return ret

#def retina_context_module(body, kernel, num_filter, filter_in, name):
#  conv_dimred = conv_act_layer(body, name+'_conv0',
#      num_filter, kernel=(1,1), pad=(0,0), stride=(1, 1), act_type='relu', separable=False, filter_in = filter_in)
#  conv1 = conv_act_layer(conv_dimred, name+'_conv1',
#      num_filter*6, kernel=(1,1), pad=(0,0), stride=(1, 1), act_type='relu', separable=False, filter_in = filter_in)
#  conv2 = conv_act_layer(conv1, name+'_conv2',
#      num_filter*6, kernel=kernel, pad=((kernel[0]-1)//2, (kernel[1]-1)//2), stride=(1, 1), act_type='relu', separable=True, filter_in = num_filter*6)
#  conv3 = conv_act_layer(conv2, name+'_conv3',
#      num_filter, kernel=(1,1), pad=(0,0), stride=(1, 1), act_type='relu', separable=False)
#  conv3 = conv3 + conv_dimred
#  return conv3

def retina_detection_module(body, num_filter, filter_in, name):
  assert num_filter%4==0
  conv1 = conv_act_layer(body, name+'_conv1',
      num_filter//2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=False, filter_in=filter_in)
  conv2 = conv_act_layer(conv1, name+'_conv2',
      num_filter//2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=False, filter_in=num_filter//2)
  conv3 = conv_act_layer(conv2, name+'_conv3',
      num_filter//2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', separable=False, filter_in=num_filter//2)
  conv4 = conv2 + conv3
  body = mx.sym.concat(*[conv1, conv4], dim=1, name = name+'_concat')
  if config.USE_DCN>0:
    body = conv_deformable(body, num_filter = num_filter, name = name+'_concat_dcn')
  return body

def head_module(body, num_filter, filter_in, name):
  if config.HEAD_MODULE=='SSH':
    return ssh_detection_module(body, num_filter, filter_in, name)
  else:
    return retina_detection_module(body, num_filter, filter_in, name)

def upsampling(data, num_filter, name):
    #ret = mx.symbol.Deconvolution(data=data, num_filter=num_filter, kernel=(4,4),  stride=(2, 2), pad=(1,1),
    #    num_group = num_filter, no_bias = True, attr={'__lr_mult__': '0.0', '__wd_mult__': '0.0'},
    #    name=name)
    #ret = mx.symbol.Deconvolution(data=data, num_filter=num_filter, kernel=(2,2),  stride=(2, 2), pad=(0,0),
    #    num_group = num_filter, no_bias = True, attr={'__lr_mult__': '0.0', '__wd_mult__': '0.0'},
    #    name=name)
    ret = mx.symbol.UpSampling(data, scale=2, sample_type='nearest', workspace=512, name=name, num_args=1)
    return ret

def get_sym_conv(data, sym):
    all_layers = sym.get_internals()
    isize = 640
    _, out_shape, _ = all_layers.infer_shape(data = (1,3,isize,isize))
    last_entry = None
    c1 = None
    c2 = None
    c3 = None
    c1_name = None
    c2_name = None
    c3_name = None
    c1_filter = -1
    c2_filter = -1
    c3_filter = -1
    #print(len(all_layers), len(out_shape))
    #print(all_layers.__class__)
    outputs = all_layers.list_outputs()
    #print(outputs.__class__, len(outputs))
    count = len(outputs)
    stride2name = {}
    stride2layer = {}
    stride2shape = {}
    for i in range(count):
      name = outputs[i]
      shape = out_shape[i]
      if not name.endswith('_output'):
        continue
      if len(shape)!=4:
        continue
      assert isize%shape[2]==0
      if shape[1]>config.max_feat_channel:
        break
      stride = isize//shape[2]
      stride2name[stride] = name
      stride2layer[stride] = all_layers[name]
      stride2shape[stride] = shape
      #print(name, shape)
      #if c1 is None and shape[2]==isize//16:
      #  cname = last_entry[0]
      #  #print('c1', last_entry)
      #  c1 = all_layers[cname]
      #  c1_name = cname
      #if c2 is None and shape[2]==isize//32:
      #  cname = last_entry[0]
      #  #print('c2', last_entry)
      #  c2 = all_layers[cname]
      #  c2_name = cname
      #if shape[2]==isize//32:
      #  c3 = all_layers[name]
      #  #print('c3', name, shape)
      #  c3_name = name

      #last_entry = (name, shape)

    F1 = config.HEAD_FILTER_NUM
    F2 = F1
    if config.SHARE_WEIGHT_BBOX or config.SHARE_WEIGHT_LANDMARK:
      F2 = F1
    strides = sorted(stride2name.keys())
    for stride in strides:
      print('stride', stride, stride2name[stride], stride2shape[stride])
    print('F1_F2', F1, F2)
    #print('cnames', c1_name, c2_name, c3_name, F1, F2)
    _bwm = 1.0
    c0 = stride2layer[4]
    c1 = stride2layer[8]
    c2 = stride2layer[16]
    c3 = stride2layer[32]
    c3 = conv_act_layer(c3, 'rf_c3_lateral',
        F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
    #c3_up = mx.symbol.UpSampling(c3, scale=2, sample_type='nearest', workspace=512, name='ssh_c3_up', num_args=1)
    c3_up = upsampling(c3, F2, 'rf_c3_upsampling')
    c2_lateral = conv_act_layer(c2, 'rf_c2_lateral',
        F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
    if config.USE_CROP:
      c3_up = mx.symbol.Crop(*[c3_up, c2_lateral])
    c2 = c2_lateral+c3_up
    c2 = conv_act_layer(c2, 'rf_c2_aggr',
        F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
    c1_lateral = conv_act_layer(c1, 'rf_c1_red_conv',
        F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
    #c2_up = mx.symbol.UpSampling(c2, scale=2, sample_type='nearest', workspace=512, name='ssh_m2_red_up', num_args=1)
    c2_up = upsampling(c2, F2, 'rf_c2_upsampling')
    #conv4_128 = mx.symbol.Crop(*[conv4_128, conv5_128_up])
    if config.USE_CROP:
      c2_up = mx.symbol.Crop(*[c2_up, c1_lateral])
    c1 = c1_lateral+c2_up
    c1 = conv_act_layer(c1, 'rf_c1_aggr',
        F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
    m1 = head_module(c1, F2*config.CONTEXT_FILTER_RATIO, F2, 'rf_c1_det')
    m2 = head_module(c2, F1*config.CONTEXT_FILTER_RATIO, F2, 'rf_c2_det')
    m3 = head_module(c3, F1*config.CONTEXT_FILTER_RATIO, F2, 'rf_c3_det')
    if len(config.RPN_ANCHOR_CFG)==3:
      ret = {8: m1, 16:m2, 32: m3}
    elif len(config.RPN_ANCHOR_CFG)==1:
      ret = {16:m2}
    elif len(config.RPN_ANCHOR_CFG)==2:
      ret = {8: m1, 16:m2}
    elif len(config.RPN_ANCHOR_CFG)==4:
      c0_lateral = conv_act_layer(c0, 'rf_c0_lateral',
          F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
      c1_up = upsampling(c1, F2, 'rf_c1_upsampling')
      if config.USE_CROP:
        c1_up = mx.symbol.Crop(*[c1_up, c0_lateral])
      c0 = c0_lateral+c1_up
      c0 = conv_act_layer(c0, 'rf_c0_aggr',
          F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

      m0 = head_module(c0, F2*config.CONTEXT_FILTER_RATIO, F2, 'rf_c0_det')
      ret = {4: m0, 8: m1, 16:m2, 32: m3}
    elif len(config.RPN_ANCHOR_CFG)==5:
      c0_lateral = conv_act_layer(c0, 'rf_c0_lateral',
          F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
      c1_up = upsampling(c1, F2, 'rf_c1_upsampling')
      if config.USE_CROP:
        c1_up = mx.symbol.Crop(*[c1_up, c0_lateral])
      c0 = c0_lateral+c1_up
      c0 = conv_act_layer(c0, 'rf_c0_aggr',
          F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

      c4 = conv_act_layer(c3, 'rf_c4',
          F2, kernel=(3, 3), pad=(1, 1), stride=(2, 2), act_type='relu', bias_wd_mult=_bwm)
      m0 = head_module(c0, F2*config.CONTEXT_FILTER_RATIO, F2, 'rf_c0_det')
      m4 = head_module(c4, F1*config.CONTEXT_FILTER_RATIO, F2, 'rf_c4_det')
      ret = {4: m0, 8: m1, 16:m2, 32: m3, 64: m4}
    elif len(config.RPN_ANCHOR_CFG)==6:
      c0_lateral = conv_act_layer(c0, 'rf_c0_lateral',
          F2, kernel=(1, 1), pad=(0, 0), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)
      c1_up = upsampling(c1, F2, 'rf_c1_upsampling')
      if config.USE_CROP:
        c1_up = mx.symbol.Crop(*[c1_up, c0_lateral])
      c0 = c0_lateral+c1_up
      c0 = conv_act_layer(c0, 'rf_c0_aggr',
          F2, kernel=(3, 3), pad=(1, 1), stride=(1, 1), act_type='relu', bias_wd_mult=_bwm)

      c4 = conv_act_layer(c3, 'rf_c4',
          F2, kernel=(3, 3), pad=(1, 1), stride=(2, 2), act_type='relu', bias_wd_mult=_bwm)
      c5 = conv_act_layer(c4, 'rf_c5',
          F2, kernel=(3, 3), pad=(1, 1), stride=(2, 2), act_type='relu', bias_wd_mult=_bwm)
      m0 = head_module(c0, F2*config.CONTEXT_FILTER_RATIO, F2, 'rf_c0_det')
      m4 = head_module(c4, F1*config.CONTEXT_FILTER_RATIO, F2, 'rf_c4_det')
      m5 = head_module(c5, F1*config.CONTEXT_FILTER_RATIO, F2, 'rf_c5_det')
      ret = {4: m0, 8: m1, 16:m2, 32: m3, 64: m4, 128: m5}

    #return {8: m1, 16:m2, 32: m3}
    return ret

def get_out(conv_fpn_feat, prefix, stride, landmark=False, lr_mult=1.0, shared_vars = None):
    A = config.NUM_ANCHORS
    bbox_pred_len = 4
    landmark_pred_len = 10
    if config.USE_BLUR:
      bbox_pred_len = 5
    if config.USE_OCCLUSION:
      landmark_pred_len = 15
    ret_group = []
    num_anchors = config.RPN_ANCHOR_CFG[str(stride)]['NUM_ANCHORS']
    label = mx.symbol.Variable(name='%s_label_stride%d'%(prefix,stride))
    bbox_anchor = mx.symbol.Variable(name='%s_bbox_anchor_stride%d'%(prefix,stride))
    bbox_target = mx.symbol.Variable(name='%s_bbox_target_stride%d'%(prefix,stride))
    bbox_weight = mx.symbol.Variable(name='%s_bbox_weight_stride%d'%(prefix,stride))
    if landmark:
      landmark_target = mx.symbol.Variable(name='%s_landmark_target_stride%d'%(prefix,stride))
      landmark_weight = mx.symbol.Variable(name='%s_landmark_weight_stride%d'%(prefix,stride))
    rpn_relu = conv_fpn_feat[stride]
    maxout_stat = 0
    if config.USE_MAXOUT>=1 and stride==config.RPN_FEAT_STRIDE[-1]:
      maxout_stat = 1
    if config.USE_MAXOUT>=2 and stride!=config.RPN_FEAT_STRIDE[-1]:
      maxout_stat = 2

    if maxout_stat==0:
      rpn_cls_score = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d'%(prefix, stride), 2*num_anchors,
          kernel=(1,1), pad=(0,0), stride=(1, 1), shared_weight = shared_vars[0][0], shared_bias = shared_vars[0][1])
    elif maxout_stat==1:
      cls_list = []
      for a in range(num_anchors):
        rpn_cls_score_bg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_bg'%(prefix,stride,a), 3,
            kernel=(1,1), pad=(0,0), stride=(1, 1))
        rpn_cls_score_bg = mx.sym.max(rpn_cls_score_bg, axis=1, keepdims=True)
        cls_list.append(rpn_cls_score_bg)
        rpn_cls_score_fg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_fg'%(prefix,stride,a), 1,
            kernel=(1,1), pad=(0,0), stride=(1, 1))
        cls_list.append(rpn_cls_score_fg)
      rpn_cls_score = mx.sym.concat(*cls_list, dim=1, name='%s_rpn_cls_score_stride%d'%(prefix,stride))
    else:
      cls_list = []
      for a in range(num_anchors):
        rpn_cls_score_bg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_bg'%(prefix,stride,a), 1,
            kernel=(1,1), pad=(0,0), stride=(1, 1))
        cls_list.append(rpn_cls_score_bg)
        rpn_cls_score_fg = conv_only(rpn_relu, '%s_rpn_cls_score_stride%d_anchor%d_fg'%(prefix,stride,a), 3,
            kernel=(1,1), pad=(0,0), stride=(1, 1))
        rpn_cls_score_fg = mx.sym.max(rpn_cls_score_fg, axis=1, keepdims=True)
        cls_list.append(rpn_cls_score_fg)
      rpn_cls_score = mx.sym.concat(*cls_list, dim=1, name='%s_rpn_cls_score_stride%d'%(prefix,stride))

    rpn_bbox_pred = conv_only(rpn_relu, '%s_rpn_bbox_pred_stride%d'%(prefix,stride), bbox_pred_len*num_anchors,
        kernel=(1,1), pad=(0,0), stride=(1, 1), shared_weight = shared_vars[1][0], shared_bias = shared_vars[1][1])

    # prepare rpn data
    rpn_cls_score_reshape = mx.symbol.Reshape(data=rpn_cls_score,
                                              shape=(0, 2, -1),
                                              name="%s_rpn_cls_score_reshape_stride%s" % (prefix,stride))

    rpn_bbox_pred_reshape = mx.symbol.Reshape(data=rpn_bbox_pred,
                                              shape=(0, 0, -1),
                                              name="%s_rpn_bbox_pred_reshape_stride%s" % (prefix,stride))
    if landmark:
      rpn_landmark_pred = conv_only(rpn_relu, '%s_rpn_landmark_pred_stride%d'%(prefix,stride), landmark_pred_len*num_anchors,
          kernel=(1,1), pad=(0,0), stride=(1, 1), shared_weight = shared_vars[2][0], shared_bias = shared_vars[2][1])
      rpn_landmark_pred_reshape = mx.symbol.Reshape(data=rpn_landmark_pred,
                                              shape=(0, 0, -1),
                                              name="%s_rpn_landmark_pred_reshape_stride%s" % (prefix,stride))

    if config.TRAIN.RPN_ENABLE_OHEM>=2:
      label, anchor_weight, valid_count = mx.sym.Custom(op_type='rpn_fpn_ohem3', stride=int(stride), network=config.network, dataset=config.dataset, prefix=prefix, cls_score=rpn_cls_score_reshape, labels = label)

      _bbox_weight = mx.sym.tile(anchor_weight, (1,1,bbox_pred_len))
      _bbox_weight = _bbox_weight.reshape((0, -1, A * bbox_pred_len)).transpose((0,2,1))
      bbox_weight = mx.sym.elemwise_mul(bbox_weight, _bbox_weight, name='%s_bbox_weight_mul_stride%s'%(prefix,stride))
      valid_count = mx.symbol.sum(valid_count)
      valid_count = valid_count + 0.001 #avoid zero

      if landmark:
        _landmark_weight = mx.sym.tile(anchor_weight, (1,1,landmark_pred_len))
        _landmark_weight = _landmark_weight.reshape((0, -1, A * landmark_pred_len)).transpose((0,2,1))
        landmark_weight = mx.sym.elemwise_mul(landmark_weight, _landmark_weight, name='%s_landmark_weight_mul_stride%s'%(prefix,stride))
      #if not config.FACE_LANDMARK:
      #  label, bbox_weight = mx.sym.Custom(op_type='rpn_fpn_ohem', stride=int(stride), cls_score=rpn_cls_score_reshape, bbox_weight = bbox_weight , labels = label)
      #else:
      #  label, bbox_weight, landmark_weight = mx.sym.Custom(op_type='rpn_fpn_ohem2', stride=int(stride), cls_score=rpn_cls_score_reshape, bbox_weight = bbox_weight, landmark_weight=landmark_weight, labels = label)
    #cls loss
    if config.TRAIN.FOCAL_LOSS < 1:
        rpn_cls_prob = mx.symbol.SoftmaxOutput(data=rpn_cls_score_reshape,
                                               label=label,
                                               multi_output=True,
                                               normalization='valid', use_ignore=True, ignore_label=-1,
                                               grad_scale = lr_mult,
                                               name='%s_rpn_cls_prob_stride%d'%(prefix,stride))

    else:
        rpn_cls_prob = mx.sym.Custom(op_type='FocalLoss', name='%s_rpn_cls_prob_stride%s' %(prefix, stride), data=rpn_cls_score_reshape, labels=label, gamma= 2,alpha = 0.25)

        #rpn_cls_prob, rpn_cls_loss_ = focal_loss_mx(cls_score=rpn_cls_score_reshape,
        #                                            label=label,
        #                                            gamma=2,
        #                                            alpha=0.25)
        #rpn_cls_loss = mx.sym.MakeLoss(name='%s_rpn_cls_loss_stride%d'%(prefix,stride),
        #                                data=rpn_cls_loss_,
        #                                grad_scale=1.0*lr_mult / (config.TRAIN.RPN_BATCH_SIZE))
    ret_group.append(rpn_cls_prob)
    ret_group.append(mx.sym.BlockGrad(label))

    if config.TRAIN.GIOU_LOSS < 1:
        #bbox L1_loss
        bbox_diff = rpn_bbox_pred_reshape-bbox_target
        bbox_diff = bbox_diff * bbox_weight
        rpn_bbox_loss_ = mx.symbol.smooth_l1(name='%s_rpn_bbox_loss_stride%d_'%(prefix,stride), scalar=3.0, data=bbox_diff)
    else:
        #bbox GIOU_loss
        rpn_bbox_loss_ = giou_loss(bbox_anchor, rpn_bbox_pred_reshape, bbox_target, bbox_weight, prefix, stride)

    if config.LR_MODE==0:
        if  config.TRAIN.GIOU_LOSS < 1:
            rpn_bbox_loss = mx.sym.MakeLoss(name='%s_rpn_bbox_loss_stride%d'%(prefix,stride), data=rpn_bbox_loss_, grad_scale=1.0*lr_mult / (config.TRAIN.RPN_BATCH_SIZE))
        else:
            rpn_bbox_loss = mx.sym.MakeLoss(name='%s_rpn_bbox_loss_stride%d'%(prefix,stride), data=rpn_bbox_loss_, grad_scale=3.0*lr_mult / (config.TRAIN.RPN_BATCH_SIZE))
    else:
      rpn_bbox_loss_ = mx.symbol.broadcast_div(rpn_bbox_loss_, valid_count)
      rpn_bbox_loss = mx.sym.MakeLoss(name='%s_rpn_bbox_loss_stride%d'%(prefix,stride), data=rpn_bbox_loss_, grad_scale=0.25*lr_mult)
    ret_group.append(rpn_bbox_loss)
    ret_group.append(mx.sym.BlockGrad(bbox_weight))

    #landmark loss
    if landmark:
      landmark_diff = rpn_landmark_pred_reshape-landmark_target
      landmark_diff = landmark_diff * landmark_weight
      rpn_landmark_loss_ = mx.symbol.smooth_l1(name='%s_rpn_landmark_loss_stride%d_'%(prefix,stride), scalar=3.0, data=landmark_diff)
      if config.LR_MODE==0:
        rpn_landmark_loss = mx.sym.MakeLoss(name='%s_rpn_landmark_loss_stride%d'%(prefix,stride), data=rpn_landmark_loss_, grad_scale=0.4*config.LANDMARK_LR_MULT*lr_mult / (config.TRAIN.RPN_BATCH_SIZE))
      else:
        rpn_landmark_loss_ = mx.symbol.broadcast_div(rpn_landmark_loss_, valid_count)
        rpn_landmark_loss = mx.sym.MakeLoss(name='%s_rpn_landmark_loss_stride%d'%(prefix,stride), data=rpn_landmark_loss_, grad_scale=0.1*config.LANDMARK_LR_MULT*lr_mult)
      ret_group.append(rpn_landmark_loss)
      ret_group.append(mx.sym.BlockGrad(landmark_weight))
    if config.USE_3D:
      from rcnn.PY_OP import rpn_3d_mesh
      pass
    return ret_group

def get_sym_train(sym):
    data = mx.symbol.Variable(name="data")

    # shared convolutional layers
    conv_fpn_feat = get_sym_conv(data, sym)
    ret_group = []
    shared_vars = []
    if config.SHARE_WEIGHT_BBOX:
      assert config.USE_MAXOUT==0
      _name = 'face_rpn_cls_score_share'
      shared_weight = mx.symbol.Variable(name="{}_weight".format(_name),
          init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
      shared_bias = mx.symbol.Variable(name="{}_bias".format(_name),
          init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(0.0)})
      shared_vars.append( [shared_weight, shared_bias] )
      _name = 'face_rpn_bbox_pred_share'
      shared_weight = mx.symbol.Variable(name="{}_weight".format(_name),
          init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
      shared_bias = mx.symbol.Variable(name="{}_bias".format(_name),
          init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(0.0)})
      shared_vars.append( [shared_weight, shared_bias] )
    else:
      shared_vars.append( [None, None] )
      shared_vars.append( [None, None] )
    if config.SHARE_WEIGHT_LANDMARK:
      _name = 'face_rpn_landmark_pred_share'
      shared_weight = mx.symbol.Variable(name="{}_weight".format(_name),
          init=mx.init.Normal(0.01), attr={'__lr_mult__': '1.0'})
      shared_bias = mx.symbol.Variable(name="{}_bias".format(_name),
          init=mx.init.Constant(0.0), attr={'__lr_mult__': '2.0', '__wd_mult__': str(0.0)})
      shared_vars.append( [shared_weight, shared_bias] )
    else:
      shared_vars.append( [None, None] )

    for stride in config.RPN_FEAT_STRIDE:
      ret = get_out(conv_fpn_feat, 'face', stride, config.FACE_LANDMARK, lr_mult=1.0, shared_vars = shared_vars)
      ret_group += ret
      if config.HEAD_BOX:
        assert not config.SHARE_WEIGHT_BBOX and not config.SHARE_WEIGHT_LANDMARK
        shared_vars = [ [None, None], [None, None], [None, None] ]
        ret = get_out(conv_fpn_feat, 'head', stride, False, lr_mult=0.5, shared_vars = shared_vars)
        ret_group += ret

    return mx.sym.Group(ret_group)


import numpy as np
import argparse, time, os

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.optim import lr_scheduler
import torch.optim as optim

from datasets import BalancedBatchSampler, DeepFashionDataset
from networks import ResNetbasedNet
from losses import OnlineTripletLoss, AllTripletLoss
from utils import read_data, pdist, RandomNegativeTripletSelector, HardestNegativeTripletSelector
from metrics import AverageNonzeroTripletsMetric
from trainer import fit
# from trainer_trip_clf import fit
# from inference import get_topK_acc, get_embedding_mtx

def main(args):
    model_path = args.model_path
    save_dir = args.save_dir
    vec_dim = 128

    data_type = ['validation'] if args.phase == 'test' else ['train', 'validation']
    img_list, base_path, item_dict = read_data("DeepFashion2", bbox_gt=True, type_list=data_type)

    # model = ResNetbasedNet(vec_dim=vec_dim, max_pool=True, load_path=model_path, clf2_num=2, adv_eta=1e-4)
    model = ResNetbasedNet(vec_dim=vec_dim, max_pool=True, load_path=model_path, clf2_num=2)

    domain_adap = args.domain_adap
    adv_train = args.adv_train
    is_cud = torch.cuda.is_available()
    device = torch.device("cuda" if is_cud else "cpu")
    if is_cud:
        if torch.cuda.device_count() > 1:
            model = nn.DataParallel(model)
        model.to(device)
    kwargs = {'num_workers': 8, 'pin_memory': True} if is_cud else {}

    if args.phase == 'train':
        train_dataset = DeepFashionDataset(img_list['train'], root=base_path, augment=True)
        train_batch_sampler = BalancedBatchSampler(train_dataset.labels, train_dataset.source, n_classes=64, n_samples=4)
        online_train_loader = torch.utils.data.DataLoader(train_dataset, batch_sampler=train_batch_sampler, **kwargs)

        test_dataset = DeepFashionDataset(img_list['validation'], root=base_path)
        test_batch_sampler = BalancedBatchSampler(test_dataset.labels, test_dataset.source, n_classes=64, n_samples=4)
        online_test_loader = torch.utils.data.DataLoader(test_dataset, batch_sampler=test_batch_sampler, **kwargs)

        margin = 0.2
        loss_fn = OnlineTripletLoss(margin, HardestNegativeTripletSelector(margin), domain_adap)
        # loss_fn = AllTripletLoss(margin)
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=1e-5, weight_decay=5e-4)
        # scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", patience=4, threshold=0.001, cooldown=2, min_lr=1e-4 / (10 * 2),)
        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode="max", patience=4, threshold=1, cooldown=2, min_lr=1e-5 / (10 * 2),)
        n_epochs = 300
        log_interval = 200

        fit(online_train_loader, online_test_loader, model, loss_fn, optimizer, scheduler, n_epochs, is_cud,
            log_interval, save_dir, metrics=[AverageNonzeroTripletsMetric()], start_epoch=200, criterion=criterion,
            domain_adap=domain_adap, adv_train=adv_train)
        # fit(online_train_loader, online_test_loader, model, loss_fn, optimizer, scheduler, n_epochs, is_cud, log_interval,
        #     save_dir, metrics=[AverageNonzeroTripletsMetric()], start_epoch=0, criterion=criterion,
        #     adv_train=True, adv_epsilon=0.01, adv_alph=0.007, adv_iter=1)

    else:
        with torch.no_grad():
            model.eval()
            test_dataset = DeepFashionDataset(img_list['validation'], root=base_path)
            test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=4)
            embedding_mtx = torch.zeros((len(test_dataset), vec_dim))
            labels = np.zeros(len(test_dataset))
            top_k = 500
            idx_ = 0
            start_time = time.time()
            cf_mtx = np.zeros(4, dtype=float) # predict_user_real_user / predict_user_real_shop / predict_shop_real_user / predict_shop_real_shop

            for idx, (data, target, _, source) in enumerate(test_loader):
                emb_vecs = model(data.cuda())
                embedding_mtx[idx_: idx_ + len(data)] = emb_vecs[0]
                predict = torch.argmax(emb_vecs[1], dim=1).cpu().numpy()
                real = source.cpu().numpy()
                cf_mtx[0] += np.sum((predict == 0) & (real == 0))
                cf_mtx[1] += np.sum((predict == 0) & (real == 1))
                cf_mtx[2] += np.sum((predict == 1) & (real == 0))
                cf_mtx[3] += np.sum((predict == 1) & (real == 1))
                labels[idx_:idx_ + len(data)] = np.asarray(target)
                idx_ += len(data)
                if idx % 20 == 0:
                    print(
                        'processing {}/{}... elapsed time {}s'.format(idx + 1, len(test_loader), time.time() - start_time))

        print('Total: {}, Domain Classification Acc: {:.5f}'.format(np.sum(cf_mtx), (cf_mtx[0] + cf_mtx[3]) / np.sum(cf_mtx)))
        print('Recall User Photo: {:.5f}'.format(cf_mtx[0] / (cf_mtx[0] + cf_mtx[2])))
        print('Recall Shop Photo: {:.5f}'.format(cf_mtx[3] / (cf_mtx[1] + cf_mtx[3])))

        np.save(os.path.join(save_dir, 'emb_mtx.npy'), embedding_mtx)
        with open(os.path.join(save_dir, 'file_info.txt'), 'w') as f:
            for i in range(len(test_dataset)):
                f.write('{},{},{},{}\n'.format(img_list['validation'][i][0], test_dataset[i][1], test_dataset[i][2], test_dataset[i][3]))
        print('save files!')

        distance_mtx = pdist(embedding_mtx)
        sorted_idx = torch.argsort(distance_mtx, dim=1).cpu().numpy()
        result_arr = np.zeros((sorted_idx.shape[0], top_k))
        for idx in range(sorted_idx.shape[0]):
            result_arr[idx] = sorted_idx[idx][sorted_idx[idx] != idx][:top_k]
            result_arr[idx] = labels[result_arr[idx].astype(np.int)] == labels[idx]
            if idx % 1000 == 0:
                print(idx)

        for k in [1, 5, 10, 20, 100, 200, 500]:
            topk_accuracy = np.sum(np.sum(result_arr[:, :k], axis=1) > 0) / result_arr.shape[0]
            print('Top-{} Accuracy: {:.5f}'.format(k, topk_accuracy))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='image retrieval system')
    parser.add_argument('-d', '--dataset', required=True, help='DeepFashion / DeepFashion2')
    # parser.add_argument('-b', '--bbox', required=True, help='gt / pred')
    parser.add_argument('--phase', required=False, default='train', help='Train or Inference')
    parser.add_argument('--domain', dest='domain_adap', required=False, default=False, help='Train or Inference')
    parser.add_argument('-mp', dest='model_path', required=False, default=None, help='pretrained model path')
    parser.add_argument('-s', dest='save_dir', required=True, help='Directory to save the models and the results')
    parser.add_argument('--adv_train', required=False, default=False, help='Make Adversarial Sample')
    args = parser.parse_args()

    main(args)


import numpy as np
import tensorflow as tf
from tensorflow.contrib import framework as contrib_framework
from tensorflow.contrib import slim as contrib_slim
from core import resnet_v1_beta, preprocess_utils, cell, attentions
slim = contrib_slim
resnet_v1_beta_block = resnet_v1_beta.resnet_v1_beta_block

def guidance_fusion_method(logits, guid_fuse, num_class, out_node, level):
  if self.guid_fuse == "sum":
    guid = tf.reduce_sum(logits, axis=3, keepdims=True)
  elif self.guid_fuse == "mean":
    guid = tf.reduce_mean(logits, axis=3, keepdims=True)
  elif self.guid_fuse == "entropy":
    guid = tf.clip_by_value(logits, 1e-10, 1.0)
    guid = -tf.reduce_sum(guid * tf.log(guid), axis=3, keepdims=True)
  elif self.guid_fuse == "conv":
    if level < len(self.low_level)-1:
      guid = slim.conv2d(logits, out_node, kernel_size=[3,3], activation_fn=None)
    else:
      guid = tf.reduce_sum(logits, axis=3, keepdims=True)
  elif self.guid_fuse == "sum_dilated":
    size = [8,6,4,2,1]
    kernel = tf.ones((size[level], size[level], num_class))
    guid = tf.nn.dilation2d(logits, filter=kernel, strides=(1,1,1,1),
                              rates=(1,1,1,1), padding="SAME")
    guid = guid - tf.ones_like(guid)
    guid = tf.reduce_sum(guid, axis=3, keepdims=True)
  elif self.guid_fuse == "w_sum":
    w = tf.nn.softmax(tf.reduce_sum(logits, axis=[1,2], keepdims=True), axis=3)
    rev_w = tf.ones_like(w) - w
    guid = tf.reduce_sum(tf.multiply(logits, rev_w), axis=3, keepdims=True)
  elif self.guid_fuse == "conv_sum":
    k_size_list = [1,1,1,3,5]
    k_size = 2 * k_size_list[level] + 1
    guid = slim.conv2d(logits, 1, kernel_size=[k_size,k_size], activation_fn=None,
                      weights_initializer=tf.ones_initializer(), trainable=False, normalizer_fn=None)
    guid = guid / (k_size*k_size*num_class*1)
  elif self.guid_fuse == "w_sum_conv":
    # TODO: make it right
    k_size_list = [1,1,1,2,4]
    k_size = 3 * k_size_list[level] + 1
    w = tf.reduce_sum(logits, axis=[1,2], keepdims=True)
    rev_w = (tf.ones_like(w)+1e-5) / (tf.sqrt(w)+1e-5)
    rev_w = tf.tile(rev_w, [1,k_size,k_size,1])
    rev_w = tf.expand_dims(rev_w, axis=4)

    n, h, w, channels_img = preprocess_utils.resolve_shape(logits, rank=4)
    n, fh, fw, channels, out_channels = preprocess_utils.resolve_shape(rev_w, rank=5)
    # F has shape (n, k_size, k_size, channels, out_channels)

    rev_w = tf.transpose(rev_w, [1, 2, 0, 3, 4])
    rev_w = tf.reshape(rev_w, [fh, fw, channels*n, out_channels])

    guid = tf.transpose(logits, [1, 2, 0, 3]) # shape (H, W, MB, channels_img)
    guid = tf.reshape(guid, [1, h, w, n*channels_img])

    out = tf.nn.depthwise_conv2d(
              guid,
              filter=rev_w,
              strides=[1, 1, 1, 1],
              padding="SAME") # here no requirement about padding being 'VALID', use whatever you want.
    # Now out shape is (1, H-fh+1, W-fw+1, MB*channels*out_channels), because we used "VALID"

    out = tf.reshape(out, [h, w, n, channels, out_channels])
    out = tf.transpose(out, [2, 0, 1, 3, 4])
    out = tf.reduce_sum(out, axis=3)

    guid = out
  elif self.guid_fuse == "sum_wo_back":
    flag = tf.concat([tf.zeros([1,1,1,1]), tf.ones([1,1,1,num_class-1])], axis=3)
    guid = tf.multiply(logits, flag)
    guid = tf.reduce_sum(guid, axis=3, keepdims=True)
  elif self.guid_fuse == "mean_wo_back":
    flag = tf.concat([tf.zeros([1,1,1,1]), tf.ones([1,1,1,num_class-1])], axis=3)
    guid = tf.multiply(logits, flag)
    guid = tf.reduce_mean(guid, axis=3, keepdims=True)
  elif self.guid_fuse == "same":
    pass
  else:
    raise ValueError("Unknown guid fuse")

  tf.add_to_collection("guidance", guid)
  return guid

class Refine(object):
  def __init__(self, low_level, fusions, prior_seg=None, prior_pred=None, stage_pred_loss_name=None, guid_conv_nums=2,
               guid_conv_type="conv2d", embed_node=32, predict_without_background=False,
               num_class=14, weight_decay=0.0, scope=None, is_training=None, **kwargs):
    self.low_level = low_level.values()
    self.fusions = fusions
    self.prior_seg = prior_seg
    self.prior_pred = prior_pred
    self.stage_pred_loss_name = stage_pred_loss_name
    self.embed_node = embed_node
    self.guid_conv_type = guid_conv_type
    self.guid_conv_nums = guid_conv_nums
    self.predict_without_background = predict_without_background
    self.num_class = num_class
    self.weight_decay = weight_decay

    self.scope = scope
    # assert len(self.low_level) == len(self.fusions)
    self.num_stage = len(self.low_level)
    self.is_training = is_training
    self.fine_tune_batch_norm = True

    self.apply_sram2 = kwargs.pop("apply_sram2", False)
    self.guid_fuse = kwargs.pop("guid_fuse", "sum")

    self.model_variants = kwargs.pop("model_variants", None)
    self.height = kwargs.pop("height", None)
    self.width = kwargs.pop("width", None)
    self.g = guidance_fusion_method
    # self.attention = utils.

  def get_fusion_method(self, method):
    if method == "concat":
      return concat_convolution
    elif method == "sum":
      return sum_convolution
    elif method in ("guid", "guid_uni"):
      return guid_attention
    elif method == "guid_class":
      return guid_class_attention
    elif method == "context_att":
      return context_attention
    elif method == "self_att":
      return self_attention

  def simple_decoder(self):
    # TODO: without consider feature size
    for i in range(len(self.low_level)):
      net = slim.conv2d(self.low_level[i], self.embed_node, kernel_size=[1, 1], scope="embed%d" %(self.num_stage-i))
      if i > 0:
        fuse_func = self.get_fusion_method(self.fusions[i])
        net = fuse_func(net, fusion)
      fusion = slim.conv2d(net, self.embed_node, kernel_size=[3, 3], scope="transform%d" %(self.num_stage-i))
    return fusion

  # def refine_decoder(self):
  #   # def func(x, fuse, guid=None, apply_second_att=True):
  #   #   embed = self.e(x)
  #   #   if guid is None:
  #   #     guid = self.g(x)
  #   #   net = attention(embed, guid)
  #   #   net = net + fuse
  #   #   if apply_second_att:
  #   #     net = attention(net, guid)
  #   #   return net

  #   fuse_func = get_fusion_method(self.fusions[0])
  #   # f = fuse_func(self.low_level[0], self.e(self.low_level[0])) # feature_fusion_method

  #   guid = self.g(self.e(self.low_level[0]), self.guid_fuse, self.num_class, self.embed_node, self.num_stage)
  #   fuse = self.t(f)

  #   for i in range(1, len(self.low_level)):
  #     fuse_func = get_fusion_method(self.fusions[i])
  #     # f = fuse_func(self.low_level[i], fuse, guid)
  #     if i < len(self.low_level)-1:
  #       guid = self.g(f, self.guid_fuse, self.num_class, self.embed_node, self.num_stage-i)
  #     fuse = self.e(f)

  #   y = resize_bilinear(fuse, [2*h, 2*w])
  #   y = slim.conv2d(y, self.embed_node, scope="decoder_output")
  #   y = slim.conv2d(y, self.num_class, kernel_size=[1, 1], stride=1, activation_fn=None, scope='logits_pred_class%d' %self.num_class)
  #   return y

  def model(self):
    batch_norm = slim.batch_norm
    batch_norm_params = get_batch_norm_params(decay=0.9997,
                                              epsilon=1e-5,
                                              scale=True,
                                              is_training=(self.is_training and self.fine_tune_batch_norm),
                                              # sync_batch_norm_method=model_options.sync_batch_norm_method
                                              )

    with tf.variable_scope(self.scope, 'Decoder'):
      with slim.arg_scope([slim.conv2d],
                          trainable=True,
                          activation_fn=tf.nn.relu,
                          weights_initializer=tf.initializers.he_normal(),
                          weights_regularizer=slim.l2_regularizer(self.weight_decay),
                          kernel_size=[3, 3],
                          padding='SAME',
                          normalizer_fn=slim.batch_norm):
        with slim.arg_scope([batch_norm], **batch_norm_params):
          with tf.variable_scope(self.model_variants, "decoder"):
            if self.model_variants == "unet" or self.model_variants == "fpn":
              feature = self.simple_decoder()
            elif self.model_variants == "refine":
              feature = self.refine_decoder()
            else:
              raise ValueError("Unknown decoder type")
          y = resize_bilinear(feature, [self.height, self.width])
          y = slim.conv2d(y, self.embed_node, scope="decoder_output")
          y = slim.conv2d(
            y, self.num_class, kernel_size=[1, 1], stride=1, activation_fn=None, scope='logits_pred_class%d' %self.num_class)
    return y

def slim_sram(in_node,
              guidance,
              num_conv=1,
              conv_type="conv",
              conv_node=64,
              scope=None):
    """Single Residual Attention Module"""
    with tf.variable_scope(scope, "sram", reuse=tf.AUTO_REUSE):
        # channel = in_node.get_shape().as_list()[3]
        net = in_node
        if conv_type == "conv":
          conv_op = slim.conv2d
        elif conv_type == "separable_conv":
          conv_op = slim.separable_conv2d
        else:
          raise ValueError("Unknown convolution type")

        for i in range(num_conv-1):
          net = conv_op(net, conv_node, kernel_size=[3,3], scope=conv_type+str(i+1))
        net = conv_op(net, conv_node, kernel_size=[3,3], scope=conv_type+"out", activation_fn=None)
        # for i in range(num_conv):
        #   net = conv_op(net, conv_node, kernel_size=[3,3], scope=conv_type+str(i+1))

        # guidance_filters = guidance.get_shape().as_list()[3]
        guidance_filters = preprocess_utils.resolve_shape(guidance, rank=4)[3]
        if guidance_filters == 1:
            guidance_tile = tf.tile(guidance, [1,1,1,conv_node])
        elif guidance_filters == conv_node:
            guidance_tile = guidance
        else:
            raise ValueError("Unknown guidance filters number")

        # tf.add_to_collection("/sram_embed", {"in_node": in_node,
        #                                      "conv2": conv2,
        #                                      "guidance_tile": guidance_tile,
        #                                      "output": output})

        output = in_node + tf.multiply(net, guidance_tile)
        tf.add_to_collection(scope+"_guided_feature", tf.multiply(net, guidance_tile))
        return output

class Refine(object):
  def __init__(self, low_level, fusions, prior_seg=None, prior_pred=None, stage_pred_loss_name=None, guid_conv_nums=2,
               guid_conv_type="conv2d", embed_node=32, predict_without_background=False,
               num_class=14, weight_decay=0.0, scope=None, is_training=None, **kwargs):
    self.low_level = list(low_level.values())
    self.fusions = fusions
    self.prior_seg = prior_seg
    self.prior_pred = prior_pred
    self.stage_pred_loss_name = stage_pred_loss_name
    self.embed_node = embed_node
    self.guid_conv_type = guid_conv_type
    self.guid_conv_nums = guid_conv_nums
    self.predict_without_background = predict_without_background
    self.num_class = num_class
    self.weight_decay = weight_decay
    self.scope = scope
    # assert len(self.low_level) == len(self.fusions)
    self.num_stage = len(self.low_level)
    self.is_training = is_training
    self.fine_tune_batch_norm = True

    self.apply_sram2 = kwargs.pop("apply_sram2", False)
    self.guid_fuse = kwargs.pop("guid_fuse", "sum")

  def embed(self, x, embed_method, out_node, scope):
    if embed_method == "self_att":
      sa_layer = attentions.self_attention(out_node)
      net = sa_layer(x, x, x, scope)
    else:
      net = slim.conv2d(x, out_node, kernel_size=[1,1], scope=scope)
    return net

  def model(self):
    # TODO: reolve_shape
    # TODO: image size
    # TODO: Remove after finish code
    batch_norm = slim.batch_norm
    batch_norm_params = get_batch_norm_params(decay=0.9997,
                                              epsilon=1e-5,
                                              scale=True,
                                              is_training=(self.is_training and self.fine_tune_batch_norm),
                                              # sync_batch_norm_method=model_options.sync_batch_norm_method
                                              )

    with tf.variable_scope(self.scope, 'Refine_Network'):
      with slim.arg_scope([slim.conv2d],
                          trainable=True,
                          activation_fn=tf.nn.relu,
                          weights_initializer=tf.initializers.he_normal(),
                          weights_regularizer=slim.l2_regularizer(self.weight_decay),
                          kernel_size=[3, 3],
                          padding='SAME',
                          normalizer_fn=slim.batch_norm):
        with slim.arg_scope([batch_norm], **batch_norm_params):
          y_tm1 = self.prior_seg

          # TODO: Would default vars value causes error?
          if self.prior_seg is not None or self.prior_pred is not None:
            if "guid" in self.fusions:
              guid = self.prior_seg
            elif "guid_class" in self.fusions:
              guid = self.prior_pred
            elif "guid_uni" in self.fusions or "context_att" in self.fusions or "self_att" in self.fusions:
              guid = tf.reduce_mean(self.prior_pred, axis=3, keepdims=True)
          out_node = self.embed_node
          tf.add_to_collection("guidance", guid)

          for i, v in enumerate(self.low_level):
            module_order = self.num_stage-i
            fuse_method = self.fusions[i]
            embed = self.embed(v, fuse_method, out_node, scope="embed%d" %module_order)
            tf.add_to_collection("embed", embed)

            fuse_func = self.get_fusion_method(fuse_method)
            h, w = preprocess_utils.resolve_shape(embed, rank=4)[1:3]

            if y_tm1 is not None:
              y_tm1 = resize_bilinear(y_tm1, [h, w])
              tf.add_to_collection("feature", y_tm1)
            else:
              # TODO: remove
              tf.add_to_collection("feature", tf.zeros_like(embed))

            if fuse_method in ("concat", "sum"):
              if y_tm1 is not None:
                y = fuse_func(embed, y_tm1, out_node, fuse_method+str(module_order))
              else:
                y = tf.identity(embed, name="identity%d" %module_order)
            elif fuse_method in ("guid", "guid_class", "guid_uni", "context_att", "self_att"):
              # guid = resize_bilinear(guid, [h, w])
              if guid is not None:
                guid = resize_bilinear(guid, [h, w])
              # tf.add_to_collection("guid", guid)

              fuse = fuse_func(embed, y_tm1, guid, out_node, fuse_method+str(module_order),
                            num_classes=self.num_class, apply_sram2=self.apply_sram2)
              """
              fuse = tf.reshape(fuse, [4, 3, h, w, out_node])
              _, fuse = seq_model(fuse, h, w, out_node, self.weight_decay, self.is_training,
                                  scope="gru"+str(i), cell_type='ConvGRU', output_wo_fuse=True)
              fuse = tf.reshape(fuse, [-1, h, w, out_node])
              """
              y = slim.conv2d(fuse, self.embed_node, scope='fuse'+str(i))
              tf.add_to_collection("refining", y)

            if self.stage_pred_loss_name is not None:

              num_class = self.num_class
              if self.predict_without_background:
                num_class -= 1

              stage_pred =  slim.conv2d(fuse, num_class, kernel_size=[1,1], activation_fn=None,
                                          scope="stage_pred%d_pred_class%d" %(module_order,num_class))

              # preds["guidance%d" %module_order] = stage_pred
              tf.add_to_collection("stage_pred", stage_pred)

            if fuse_method in ("guid"):
              guid = y
              y_tm1 = None

            elif fuse_method in ("guid_class", "guid_uni", "context_att", "self_att"):
              if i < len(self.low_level)-1:
                if "softmax" in self.stage_pred_loss_name:
                  # guid = tf.nn.softmax(stage_pred, axis=3)
                  guid = tf.nn.softmax(stage_pred, axis=3)
                elif "sigmoid" in self.stage_pred_loss_name:
                  guid = tf.nn.sigmoid(stage_pred)

                if self.guid_fuse == "sum":
                  guid = tf.reduce_sum(guid, axis=3, keepdims=True)
                elif self.guid_fuse == "mean":
                  guid = tf.reduce_mean(guid, axis=3, keepdims=True)
                elif self.guid_fuse == "entropy":
                  guid = tf.clip_by_value(guid, 1e-10, 1.0)
                  guid = -tf.reduce_sum(guid * tf.log(guid), axis=3, keepdims=True)
                elif self.guid_fuse == "conv":
                  if i < len(self.low_level)-1:
                    guid = slim.conv2d(guid, out_node, kernel_size=[3,3], activation_fn=None)
                  else:
                    guid = tf.reduce_sum(guid, axis=3, keepdims=True)
                elif self.guid_fuse == "sum_dilated":
                  size = [8,6,4,2,1]
                  kernel = tf.ones((size[i], size[i], num_class))
                  guid = tf.nn.dilation2d(guid, filter=kernel, strides=(1,1,1,1),
                                            rates=(1,1,1,1), padding="SAME")
                  guid = guid - tf.ones_like(guid)
                  guid = tf.reduce_sum(guid, axis=3, keepdims=True)
                elif self.guid_fuse == "w_sum":
                  w = tf.nn.softmax(tf.reduce_sum(guid, axis=[1,2], keepdims=True), axis=3)
                  rev_w = tf.ones_like(w) - w
                  guid = tf.reduce_sum(tf.multiply(guid, rev_w), axis=3, keepdims=True)
                elif self.guid_fuse == "conv_sum":
                  k_size_list = [1,1,1,3,5]
                  k_size = 2 * k_size_list[i] + 1
                  guid = slim.conv2d(guid, 1, kernel_size=[k_size,k_size], activation_fn=None,
                                    weights_initializer=tf.ones_initializer(), trainable=False, normalizer_fn=None)
                  guid = guid / (k_size*k_size*num_class*1)
                elif self.guid_fuse == "w_sum_conv":
                  # TODO: make it right
                  k_size_list = [1,1,1,2,4]
                  k_size = 3 * k_size_list[i] + 1
                  w = tf.reduce_sum(guid, axis=[1,2], keepdims=True)
                  rev_w = (tf.ones_like(w)+1e-5) / (tf.sqrt(w)+1e-5)
                  rev_w = tf.tile(rev_w, [1,k_size,k_size,1])
                  rev_w = tf.expand_dims(rev_w, axis=4)

                  n, h, w, channels_img = preprocess_utils.resolve_shape(guid, rank=4)
                  n, fh, fw, channels, out_channels = preprocess_utils.resolve_shape(rev_w, rank=5)
                  # F has shape (n, k_size, k_size, channels, out_channels)

                  rev_w = tf.transpose(rev_w, [1, 2, 0, 3, 4])
                  rev_w = tf.reshape(rev_w, [fh, fw, channels*n, out_channels])

                  guid = tf.transpose(guid, [1, 2, 0, 3]) # shape (H, W, MB, channels_img)
                  guid = tf.reshape(guid, [1, h, w, n*channels_img])

                  out = tf.nn.depthwise_conv2d(
                            guid,
                            filter=rev_w,
                            strides=[1, 1, 1, 1],
                            padding="SAME") # here no requirement about padding being 'VALID', use whatever you want.
                  # Now out shape is (1, H-fh+1, W-fw+1, MB*channels*out_channels), because we used "VALID"

                  out = tf.reshape(out, [h, w, n, channels, out_channels])
                  out = tf.transpose(out, [2, 0, 1, 3, 4])
                  out = tf.reduce_sum(out, axis=3)

                  guid = out
                elif self.guid_fuse == "sum_wo_back":
                  flag = tf.concat([tf.zeros([1,1,1,1]), tf.ones([1,1,1,num_class-1])], axis=3)
                  guid = tf.multiply(guid, flag)
                  guid = tf.reduce_sum(guid, axis=3, keepdims=True)
                elif self.guid_fuse == "mean_wo_back":
                  flag = tf.concat([tf.zeros([1,1,1,1]), tf.ones([1,1,1,num_class-1])], axis=3)
                  guid = tf.multiply(guid, flag)
                  guid = tf.reduce_mean(guid, axis=3, keepdims=True)
                elif self.guid_fuse == "same":
                  pass
                else:
                  raise ValueError("Unknown guid fuse")

                tf.add_to_collection("guidance", guid)

              y_tm1 = y
            elif fuse_method in ("concat", "sum"):
              y_tm1 = y

          # h, w = y.get_shape().as_list()[1:3]
          y = resize_bilinear(y, [2*h, 2*w])
          y = slim.conv2d(y, self.embed_node, scope="decoder_output")
          y = slim.conv2d(y, self.num_class, kernel_size=[1, 1], stride=1, activation_fn=None, scope='logits_pred_class%d' %self.num_class)

    return y

  def get_fusion_method(self, method):
    if method == "concat":
      return self.concat_convolution
    elif method == "sum":
      return self.sum_convolution
    elif method in ("guid", "guid_uni"):
      return self.guid_attention
    elif method == "guid_class":
      return self.guid_class_attention
    elif method == "context_att":
      return self.guid_attention
      # return self.context_attention
    elif method == "self_att":
      return self.guid_attention
      # return self.self_attention

  def concat_convolution(self, cur, last, out_node, scope):
    with tf.variable_scope(scope, "concat_conv"):
      h, w = preprocess_utils.resolve_shape(cur, rank=4)[1:3]
      last = resize_bilinear(last, [h, w])
      net = slim.conv2d(tf.concat([cur, last], axis=3), out_node, scope="conv1")
      return net

  def sum_convolution(self, cur, last, out_node, scope):
    with tf.variable_scope(scope, "sum_cocnv"):
      h, w = preprocess_utils.resolve_shape(cur, rank=4)[1:3]
      last = resize_bilinear(last, [h, w])
      net = slim.conv2d(cur+last, out_node, scope="conv1")
      return net

  def guid_attention(self, x1, x2, guid, out_node, scope, *args, **kwargs):
    apply_sram2 = kwargs.pop("apply_sram2", False)
    guid_node = preprocess_utils.resolve_shape(guid, rank=4)[3]

    if guid_node != out_node and guid_node != 1:
      raise ValueError("Unknown guidance node number %d, should be 1 or out_node" %guid_node)

    with tf.variable_scope(scope, 'guid'):
      net = slim_sram(x1, guid, self.guid_conv_nums, self.guid_conv_type, self.embed_node, "sram1")
      tf.add_to_collection("sram1", net)
      if x2 is not None:
        net = net + x2
        if apply_sram2:
          net = slim_sram(net, guid, self.guid_conv_nums, self.guid_conv_type, self.embed_node, "sram2")

      tf.add_to_collection("sram2", net)
      return net

  def guid_class_attention(self, x1, x2, guid, out_node, scope, *args, **kwargs):
    num_classes = kwargs.pop("num_classes", None)
    guid_node = guid.get_shape().as_list()[3]
    if guid_node != num_classes:
      raise ValueError("Unknown guidance node number %d, should equal class number" %guid_node)
    with tf.variable_scope(scope, "guid_class"):
      total_att = []
      for i in range(1, num_classes):
        net = slim_sram(x1, guid[...,i:i+1], self.guid_conv_nums, self.guid_conv_type, self.embed_node, "sram1")
        if x2 is not None:
          net = net + x2
          net = slim_sram(net, guid[...,i:i+1], self.guid_conv_nums, self.guid_conv_type, self.embed_node, "sram2")
        total_att.append(net)
      fuse = slim.conv2d(tf.concat(total_att, axis=3), out_node, kernel_size=[1,1], scope="fuse")
    return fuse

  def context_attention(self, x1, x2, guid, out_node, scope, *args, **kwargs):
    guid_node = preprocess_utils.resolve_shape(guid, rank=4)[3]

    if guid_node != out_node and guid_node != 1:
      raise ValueError("Unknown guidance node number %d, should be 1 or out_node" %guid_node)

    with tf.variable_scope(scope, 'guid'):
      context = slim_sram(x1, guid, self.guid_conv_nums, self.guid_conv_type, self.embed_node, "sram1")
      tf.add_to_collection("sram1", context)
      if x2 is not None:
        ca_layer = attentions.self_attention(out_node)
        net = ca_layer(x2, context, x2, "context_att1")
        tf.add_to_collection("context_att1", net)
      return net

  def self_attention(self, x1, x2, guid, out_node, scope, *args, **kwargs):
    guid_node = preprocess_utils.resolve_shape(guid, rank=4)[3]

    if guid_node != out_node and guid_node != 1:
      raise ValueError("Unknown guidance node number %d, should be 1 or out_node" %guid_node)

    with tf.variable_scope(scope, 'guid'):
      net = slim_sram(x1, guid, self.guid_conv_nums, self.guid_conv_type, self.embed_node, "sram1")
      tf.add_to_collection("sram1", net)
      if x2 is not None:
        net = net + x2
        sa_layer = attentions.self_attention(out_node)
        net = sa_layer(net, net, net, "self_att1")
        tf.add_to_collection("self_att1", net)
      return net

def mlp(inputs,
        output_dims,
        num_layers=3,
        decreasing_root=8,
        scope=None,):
  """
  Args:
    inputs:
    output_dims
    num_layers:
    decreasing_root:
  Returns:
  Raises:

  """
  # TODO: raise if input dims smaller than output_dims
  with tf.variable_scope(scope, 'multi_perceptron_network'):
    net = inputs
    for i in range(num_layers):
      dims = net.get_shape().as_list()[1]
      if i < num_layers-1:
        net = fc_layer(net, [dims, dims//decreasing_root], _std=1, reuse=tf.AUTO_REUSE, scope='_'.join(['fc', str(i)]))
      else:
        outputs = fc_layer(net, [dims, output_dims], _std=1, reuse=tf.AUTO_REUSE, scope='_'.join(['fc', str(i)]))
  return outputs

def GCN(x, out_channels, ks=7, scope=None, is_training=None):
    with tf.variable_scope(scope, "global_convolution_network"):
        channels = x.get_shape().as_list()[3]
        x_l1 = conv2d(x, [ks,1,channels,out_channels], scope="x_l1", is_training=False)
        x_l2 = conv2d(x_l1, [1,ks,out_channels,out_channels], scope="x_l2", is_training=False)

        x_r1 = conv2d(x, [1,ks,channels,out_channels], scope="x_r1", is_training=False)
        x_r2 = conv2d(x_r1, [ks,1,out_channels,out_channels], scope="x_r2", is_training=False)

        x = x_l2 + x_r2
        x = batch_norm(x, is_training=is_training, scope='batch_norm')
        x = tf.nn.relu(x)
    return x

# TODO: Organize
def get_batch_norm_fn(sync_batch_norm_method):
  """Gets batch norm function.
  Currently we only support the following methods:
    - `None` (no sync batch norm). We use slim.batch_norm in this case.
  Args:
    sync_batch_norm_method: String, method used to sync batch norm.
  Returns:
    Batchnorm function.
  Raises:
    ValueError: If sync_batch_norm_method is not supported.
  """
  if sync_batch_norm_method == 'None':
    return slim.batch_norm
  else:
    raise ValueError('Unsupported sync_batch_norm_method.')

def get_batch_norm_params(decay=0.9997,
                          epsilon=1e-5,
                          center=True,
                          scale=True,
                          is_training=True,
                          sync_batch_norm_method='None',
                          initialize_gamma_as_zeros=False):
  """Gets batch norm parameters.
  Args:
    decay: Float, decay for the moving average.
    epsilon: Float, value added to variance to avoid dividing by zero.
    center: Boolean. If True, add offset of `beta` to normalized tensor. If
      False,`beta` is ignored.
    scale: Boolean. If True, multiply by `gamma`. If False, `gamma` is not used.
    is_training: Boolean, whether or not the layer is in training mode.
    sync_batch_norm_method: String, method used to sync batch norm.
    initialize_gamma_as_zeros: Boolean, initializing `gamma` as zeros or not.
  Returns:
    A dictionary for batchnorm parameters.
  Raises:
    ValueError: If sync_batch_norm_method is not supported.
  """
  batch_norm_params = {
      'is_training': is_training,
      'decay': decay,
      'epsilon': epsilon,
      'scale': scale,
      'center': center,
  }
  if initialize_gamma_as_zeros:
    if sync_batch_norm_method == 'None':
      # Slim-type gamma_initialier.
      batch_norm_params['param_initializers'] = {
          'gamma': tf.zeros_initializer(),
      }
    else:
      raise ValueError('Unsupported sync_batch_norm_method.')
  return batch_norm_params

def fc_layer(x,
             layer_size,
             _std,
             reuse=None,
             scope=None):
    """Splits a separable conv2d into depthwise and pointwise conv2d.
  This operation differs from `tf.layers.separable_conv2d` as this operation
  applies activation function between depthwise and pointwise conv2d.
  Args:
    inputs: Input tensor with shape [batch, height, width, channels].
    filters: Number of filters in the 1x1 pointwise convolution.
    kernel_size: A list of length 2: [kernel_height, kernel_width] of
      of the filters. Can be an int if both values are the same.
    rate: Atrous convolution rate for the depthwise convolution.
    weight_decay: The weight decay to use for regularizing the model.
    depthwise_weights_initializer_stddev: The standard deviation of the
      truncated normal weight initializer for depthwise convolution.
    pointwise_weights_initializer_stddev: The standard deviation of the
      truncated normal weight initializer for pointwise convolution.
    scope: Optional scope for the operation.
  Returns:
    Computed features after split separable conv2d.
  """
    # TODO: tf.layers.dense()?  SEnet-Tensorflow github
    input_size, output_size = layer_size
    with tf.variable_scope(scope, "fc_layer"):
        w = tf.get_variable(
                "W",
                initializer=tf.truncated_normal(shape=[input_size, output_size], stddev=_std))
        w = w / tf.sqrt(input_size/2)
        b = tf.get_variable(
                "b",
                initializer=tf.constant(0.1, shape=[output_size]))
        output = tf.add(tf.matmul(x, w, name='fc'), b)
    return output

def conv2d(inputs,
          filter_shape,
          strides=[1,1,1,1],
          padding='SAME',
          dilations=[1,1,1,1],
          activate=None,
          is_training=None,
          reuse=False,
          scope=None):
    """Splits a separable conv2d into depthwise and pointwise conv2d.
  This operation differs from `tf.layers.separable_conv2d` as this operation
  applies activation function between depthwise and pointwise conv2d.
  Args:
    inputs: Input tensor with shape [batch, height, width, channels].
    filters: Number of filters in the 1x1 pointwise convolution.
    kernel_size: A list of length 2: [kernel_height, kernel_width] of
      of the filters. Can be an int if both values are the same.
    rate: Atrous convolution rate for the depthwise convolution.
    weight_decay: The weight decay to use for regularizing the model.
    depthwise_weights_initializer_stddev: The standard deviation of the
      truncated normal weight initializer for depthwise convolution.
    pointwise_weights_initializer_stddev: The standard deviation of the
      truncated normal weight initializer for pointwise convolution.
    scope: Optional scope for the operation.
  Returns:
    Computed features after split separable conv2d.
  """
    # TODO: remove numpy dependency
    assert inputs.get_shape().ndims == 4
    kernel_h, kernel_w, in_channels, out_channels = filter_shape
    _std = np.sqrt(2 / (kernel_h * kernel_w * out_channels))

    with tf.variable_scope(scope, reuse=reuse) as scope:
        w = tf.get_variable(
                "W",
                initializer=tf.truncated_normal(shape=filter_shape, stddev=_std))
        conv = tf.nn.conv2d(inputs,
                            w,
                            strides=strides,
                            padding=padding,
#                            dilations=dilations,
                            name='conv',
                            )

        if is_training is not None:
            output = batch_norm(conv, is_training=is_training, scope='batch_norm')
        else:
            b = tf.get_variable(
                "b",
                initializer=tf.constant(0.1, shape=[out_channels]))
            output = tf.nn.bias_add(conv, b)

        if activate is not None:
            output = activate(output)

    return output

def batch_norm(inputs, scope=None, is_training=None):
    """BN for the first input"""
    assert inputs.get_shape().ndims == 4
    with tf.variable_scope(scope):
        output = tf.layers.batch_normalization(inputs, training=is_training)
    return output

def resize_bilinear(images, size, output_dtype=tf.float32):
  """Returns resized images as output_type.
  Args:
    images: A tensor of size [batch, height_in, width_in, channels].
    size: A 1-D int32 Tensor of 2 elements: new_height, new_width. The new size
      for the images.
    output_dtype: The destination type.
  Returns:
    A tensor of size [batch, height_out, width_out, channels] as a dtype of
      output_dtype.
  """
  images = tf.image.resize_bilinear(images, size, align_corners=False)
  return tf.cast(images, dtype=output_dtype)

def scale_dimension(dim, scale):
  """Scales the input dimension.
  Args:
    dim: Input dimension (a scalar or a scalar Tensor).
    scale: The amount of scaling applied to the input.
  Returns:
    Scaled dimension.
  """
  if isinstance(dim, tf.Tensor):
    return tf.cast((tf.to_float(dim) - 1.0) * scale + 1.0, dtype=tf.int32)
  else:
    return int((float(dim) - 1.0) * scale + 1.0)

# def se_block(inputs, node=32, scope=None):
#   with tf.variable_scope(scope, "se_block", reuse=tf.AUTO_REUSE):
#       channel = inputs.get_shape().as_list()[3]
#       net = inputs
#       net = tf.reduce_mean(net, [1,2], keep_dims=False)
#       net = fc_layer(net, [channel, 32], _std=1, scope="fc1")
#       net = tf.nn.relu(net)
#       net = fc_layer(net, [32, channel], _std=1, scope="fc2")
#       net = tf.nn.sigmoid(net)
#   return net

def seq_model(inputs, ny, nx, n_class, weight_decay, is_training, scope, cell_type='ConvGRU', output_wo_fuse=False):
  with tf.variable_scope(scope, "forward_cell"):
    with slim.arg_scope([slim.batch_norm],
                          is_training=is_training):
      with slim.arg_scope([slim.conv2d],
                        weights_initializer=tf.initializers.he_normal(),
                        weights_regularizer=slim.l2_regularizer(weight_decay),
                        normalizer_fn=slim.batch_norm):
      # in_shape = inputs.get_shape().as_list()
        in_shape = preprocess_utils.resolve_shape(inputs, rank=5)
        batch_size = in_shape[0]
        # seq_length = in_shape[1]
        nx = in_shape[2]
        ny = in_shape[3]

        if cell_type =='ConvGRU':
          with tf.variable_scope("forward_cell") as scope:
              cell_forward = cell.ConvGRUCell(shape=[ny, nx], filters=n_class, kernel=[3, 3])
              outputs_forward, state_forward = tf.nn.dynamic_rnn(
                cell=cell_forward, dtype=tf.float32, inputs=inputs,
                initial_state=cell_forward.zero_state(batch_size, dtype=tf.float32))
          feats = state_forward
          output = [outputs_forward]
        elif cell_type =='BiConvGRU':
          with tf.variable_scope("forward_cell") as scope:
              cell_forward = cell.ConvGRUCell(shape=[ny, nx], filters=n_class, kernel=[3, 3])
              outputs_forward, state_forward = tf.nn.dynamic_rnn(
                cell=cell_forward, dtype=tf.float32, inputs=inputs,
                initial_state=cell_forward.zero_state(batch_size, dtype=tf.float32))

          inputs_b = inputs[:,::-1]
          with tf.variable_scope("backward_cell") as scope:
              cell_backward = cell.ConvGRUCell(shape=[ny, nx], filters=n_class, kernel=[3, 3])
              outputs_backward, state_backward = tf.nn.dynamic_rnn(
                cell=cell_backward, dtype=tf.float32, inputs=inputs_b,
                initial_state=cell_backward.zero_state(batch_size, dtype=tf.float32))

          feats = tf.concat([state_forward, state_backward], axis=3)
          output = [outputs_forward, outputs_backward]
        if output_wo_fuse:
          y = None
        else:
          y = slim.conv2d(feats, n_class, kernel_size=[1, 1], stride=1, activation_fn=None, scope='fuse')
    return y, output