import numpy as np
import pandas as pd
import importlib
import sys
import random
from tqdm import tqdm
import gc
import argparse
import torch
from torch import optim
from torch.cuda.amp import GradScaler, autocast
from collections import defaultdict
import cv2
from copy import copy
import os
from transformers import get_cosine_schedule_with_warmup
from torch.utils.data import SequentialSampler, DataLoader

cv2.setNumThreads(0)

sys.path.append("configs")
sys.path.append("models")
sys.path.append("data")
sys.path.append("losses")
sys.path.append("utils")

def worker_init_fn(worker_id):
    np.random.seed(np.random.get_state()[1][0] + worker_id)

def get_train_dataloader(train_ds, cfg):
    train_dataloader = DataLoader(
        train_ds,
        sampler=None,
        shuffle=True,
        batch_size=cfg.batch_size,
        num_workers=cfg.num_workers,
        pin_memory=False,
        collate_fn=cfg.tr_collate_fn,
        drop_last=cfg.drop_last,
        worker_init_fn=worker_init_fn,
    )
    print(f"train: dataset {len(train_ds)}, dataloader {len(train_dataloader)}")
    return train_dataloader

def get_val_dataloader(val_ds, cfg):
    sampler = SequentialSampler(val_ds)
    if cfg.batch_size_val is not None:
        batch_size = cfg.batch_size_val
    else:
        batch_size = cfg.batch_size
    val_dataloader = DataLoader(
        val_ds,
        sampler=sampler,
        batch_size=batch_size,
        num_workers=cfg.num_workers,
        pin_memory=False,
        collate_fn=cfg.val_collate_fn,
        worker_init_fn=worker_init_fn,
    )
    print(f"valid: dataset {len(val_ds)}, dataloader {len(val_dataloader)}")
    return val_dataloader

def get_scheduler(cfg, optimizer, total_steps):
    scheduler = get_cosine_schedule_with_warmup(
        optimizer,
        num_warmup_steps=cfg.warmup * (total_steps // cfg.batch_size),
        num_training_steps=cfg.epochs * (total_steps // cfg.batch_size),
    )
    return scheduler

def set_seed(seed=1234):
    random.seed(seed)
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = False
    torch.backends.cudnn.benchmark = True

def get_model(cfg, train_dataset):
    Net = importlib.import_module(cfg.model).Net
    return Net(cfg)

def create_checkpoint(model, optimizer, epoch, scheduler=None, scaler=None):
    checkpoint = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": epoch,
    }

    if scheduler is not None:
        checkpoint["scheduler"] = scheduler.state_dict()

    if scaler is not None:
        checkpoint["scaler"] = scaler.state_dict()
    return checkpoint

parser = argparse.ArgumentParser(description="")

parser.add_argument("-C", "--config", help="config filename")
parser.add_argument("-s", "--seed", type=int, default=-1, help="seed")
parser_args, _ = parser.parse_known_args(sys.argv)

cfg = copy(importlib.import_module(parser_args.config).cfg)

if parser_args.seed > -1:
    cfg.seed = parser_args.seed

os.makedirs(str(cfg.output_dir + "/"), exist_ok=True)

cfg.CustomDataset = importlib.import_module(cfg.dataset).CustomDataset
cfg.tr_collate_fn = importlib.import_module(cfg.dataset).tr_collate_fn
cfg.val_collate_fn = importlib.import_module(cfg.dataset).val_collate_fn
batch_to_device = importlib.import_module(cfg.dataset).batch_to_device

def run_eval(model, val_dataloader, cfg, pre="val"):

    model.eval()
    torch.set_grad_enabled(False)

    val_data = defaultdict(list)

    for data in tqdm(val_dataloader):

        batch = batch_to_device(data, device)

        if cfg.mixed_precision:
            with autocast():
                output = model(batch)
        else:
            output = model(batch)

        for key, val in output.items():
            val_data[key] += [output[key]]

    for key, val in output.items():
        value = val_data[key]
        if len(value[0].shape) == 0:
            val_data[key] = torch.stack(value)
        else:
            val_data[key] = torch.cat(value, dim=0)

    if cfg.save_val_data:
        torch.save(val_data, f"{cfg.output_dir}/{pre}_data_seed{cfg.seed}.pth")

    if "loss" in val_data:
        val_losses = val_data["loss"].cpu().numpy()
        val_loss = np.mean(val_losses)
        print(f"Mean {pre}_loss", np.mean(val_losses))

    else:
        val_loss = 0.0

    print("EVAL FINISHED")

    return val_loss

if __name__ == "__main__":

    if cfg.seed < 0:
        cfg.seed = np.random.randint(1_000_000)
    print("seed", cfg.seed)

    device = "cuda:%d" % cfg.gpu
    cfg.device = device

    set_seed(cfg.seed)

    train_df = pd.read_csv(cfg.train_df)
    val_df = pd.read_csv(cfg.val_df)

    train_dataset = cfg.CustomDataset(train_df, cfg, aug=cfg.train_aug, mode="train")
    val_dataset = cfg.CustomDataset(val_df, cfg, aug=cfg.val_aug, mode="val")

    train_dataloader = get_train_dataloader(train_dataset, cfg)
    val_dataloader = get_val_dataloader(val_dataset, cfg)

    model = get_model(cfg, train_dataset)
    model.to(device)

    total_steps = len(train_dataset)

    params = model.parameters()
    optimizer = optim.Adam(params, lr=cfg.lr, weight_decay=0)

    scheduler = get_scheduler(cfg, optimizer, total_steps)

    if cfg.mixed_precision:
        scaler = GradScaler()
    else:
        scaler = None

    cfg.curr_step = 0
    i = 0
    best_val_loss = np.inf
    optimizer.zero_grad()
    for epoch in range(cfg.epochs):

        set_seed(cfg.seed + epoch)

        cfg.curr_epoch = epoch

        print("EPOCH:", epoch)

        progress_bar = tqdm(range(len(train_dataloader)))
        tr_it = iter(train_dataloader)

        losses = []

        gc.collect()

        if cfg.train:
            # ==== TRAIN LOOP
            for itr in progress_bar:
                i += 1

                cfg.curr_step += cfg.batch_size

                data = next(tr_it)

                model.train()
                torch.set_grad_enabled(True)

                batch = batch_to_device(data, device)

                if cfg.mixed_precision:
                    with autocast():
                        output_dict = model(batch)
                else:
                    output_dict = model(batch)

                loss = output_dict["loss"]

                losses.append(loss.item())

                if cfg.mixed_precision:
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                    optimizer.zero_grad()
                else:
                    loss.backward()
                    optimizer.step()
                    optimizer.zero_grad()

                if scheduler is not None:
                    scheduler.step()

                if cfg.curr_step % cfg.batch_size == 0:
                    progress_bar.set_description(f"loss: {np.mean(losses[-10:]):.4f}")

        if cfg.val:
            if (epoch + 1) % cfg.eval_epochs == 0 or (epoch + 1) == cfg.epochs:
                val_loss = run_eval(model, val_dataloader, cfg)
            else:
                val_score = 0

        if cfg.epochs > 0:
            checkpoint = create_checkpoint(
                model, optimizer, epoch, scheduler=scheduler, scaler=scaler
            )

            torch.save(checkpoint, f"{cfg.output_dir}/checkpoint_last_seed{cfg.seed}.pth")

    if cfg.epochs > 0:
        checkpoint = create_checkpoint(model, optimizer, epoch, scheduler=scheduler, scaler=scaler)

        torch.save(checkpoint, f"{cfg.output_dir}/checkpoint_last_seed{cfg.seed}.pth")


from envs.utils import dict_get, max_n_objects, fill_infos_cache, sample_trajdata_vehicle, load_ngsim_trajdatas
from src.Record.frame import Frame
from src.Record.record import SceneRecord, get_scene
from src.Basic.Vehicle import Vehicle
from feature_extractor.utils import build_feature_extractor
from envs.action import AccelTurnrate, propagate
from src.Vec.VecE2 import norm
import numpy as np
import copy

class AutoEnv:
    def __init__(self, params: dict, trajdatas: list = None, trajinfos: list = None, roadways: list = None,
                 reclength: int = 5, delta_t: float = .1, primesteps: int = 50, H: int = 50,
                 terminate_on_collision: bool = True, terminate_on_off_road: bool = True,
                 render_params: dict = {"zoom": 5., "viz_dir": "/tmp"}):
        '''

        trajdatas::Vector{ListRecord}
        trajinfos::Vector{Dict}
        roadways::Vector{Roadway}
        roadway::Union{Nothing, Roadway} # current roadway
        scene::Scene
        rec::SceneRecord
        ext::MultiFeatureExtractor
        egoid::Int # current id of relevant ego vehicle
        ego_veh::Union{Nothing, Vehicle} # the ego vehicle
        traj_idx::Int # current index into trajdatas
        t::Int # current timestep in the trajdata
        h::Int # current maximum horizon for egoid
        H::Int # maximum horizon
        primesteps::Int # timesteps to prime the scene
        Δt::Float64

        # settings
        terminate_on_collision::Bool
        terminate_on_off_road::Bool

        # metadata
        epid::Int # episode id
        render_params::Dict # rendering options
        infos_cache::Dict # cache for infos intermediate results

        '''
        param_keys = params.keys()
        assert "trajectory_filepaths" in param_keys

        # optionally overwrite defaults
        reclength = dict_get(params, "reclength", reclength)
        primesteps = dict_get(params, "primesteps", primesteps)
        H = dict_get(params, "H", H)

        for (k, v) in dict_get(params, "render_params", render_params).items():
            render_params[k] = v

        terminate_on_collision = dict_get(params, "terminate_on_collision", terminate_on_collision)
        terminate_on_off_road = dict_get(params, "terminate_on_off_road", terminate_on_off_road)

        if trajdatas is None or trajinfos is None or roadways is None:
            trajdatas, trajinfos, roadways = load_ngsim_trajdatas(
                params["trajectory_filepaths"],
                minlength=primesteps + H
            )
        print("traj_data_len: ", len(trajdatas))
        print("trajinfos: ", trajinfos)
        # build components
        scene_length = max_n_objects(trajdatas)
        scene = Frame()
        scene.init(scene_length)
        rec = SceneRecord()
        rec.init(reclength, delta_t, scene_length)
        ext = build_feature_extractor(params)
        infos_cache = fill_infos_cache(ext)

        self.trajdatas = trajdatas
        self.trajinfos = trajinfos
        self.roadways = roadways
        self.roadway = None
        self.scene = scene
        self.rec = rec
        self.ext = ext
        self.egoid = 0
        self.ego_veh = None
        self.traj_idx = 0
        self.t = 0
        self.h = 0
        self.H = H
        self.primesteps = primesteps
        self.delta_t = delta_t
        self.terminate_on_collision = terminate_on_collision
        self.terminate_on_off_road = terminate_on_off_road
        self.epid = 0
        self.render_params = render_params
        self.infos_cache = infos_cache

    def reset(self, offset: int = None, egoid: int = None, start: int = None, traj_idx: int = 0):
        if offset is None:
            offset = self.H + self.primesteps
        self.traj_idx, self.egoid, self.t, self.h = sample_trajdata_vehicle(
            self.trajinfos,
            offset,
            traj_idx,
            egoid,
            start
        )

        self.epid += 1

        self.rec.empty()
        self.scene.empty()

        # prime
        for t in range(self.t, (self.t + self.primesteps + 1)):
            self.scene = get_scene(self.scene, self.trajdatas[self.traj_idx], t)
            # print("step {}:".format(t))
            # for i in range(self.scene.n):
            #     print(self.scene[i].id)
            self.rec.update(self.scene)

        # set the ego vehicle
        print("egoid: ", egoid)
        self.ego_veh = self.scene[self.scene.findfirst(self.egoid)]
        # set the roadway
        self.roadway = self.roadways[self.traj_idx]
        # self.t is the next timestep to load
        self.t += self.primesteps + 1
        # enforce a maximum horizon
        self.h = min(self.h, self.t + self.H)
        return self.get_features()

    def _step(self, action: list):
        # convert action into form
        # action[0] is `a::Float64` longitudinal acceleration [m/s^2]
        # action[1] is `ω::Float64` desired turn rate [rad/sec]
        action = action[0]  # input should be a (n_agents,n_actions) size array
        ego_action = AccelTurnrate(action[0], action[1])

        # propagate the ego vehicle
        ego_state = propagate(
            self.ego_veh,
            ego_action,
            self.roadway,
            self.delta_t
        )
        # update the ego_veh
        self.ego_veh = Vehicle(ego_state, self.ego_veh.definition, self.ego_veh.id)

        # load the actual scene, and insert the vehicle into it
        self.scene = get_scene(self.scene, self.trajdatas[self.traj_idx], self.t)
        vehidx = self.scene.findfirst(self.egoid)
        orig_veh = self.scene[vehidx]  # for infos purposes
        self.scene[vehidx] = self.ego_veh

        # update rec with current scene
        self.rec.update(self.scene)

        # compute info about the step
        step_infos = dict()
        step_infos["rmse_pos"] = norm((orig_veh.state.posG - self.ego_veh.state.posG))
        step_infos["rmse_vel"] = abs((orig_veh.state.v - self.ego_veh.state.v))
        step_infos["rmse_t"] = abs((orig_veh.state.posF.t - self.ego_veh.state.posF.t))
        step_infos["x"] = self.ego_veh.state.posG.x
        step_infos["y"] = self.ego_veh.state.posG.y
        step_infos["orig_x"] = orig_veh.state.posG.x
        step_infos["orig_y"] = orig_veh.state.posG.y
        step_infos["s"] = self.ego_veh.state.posF.s
        step_infos["phi"] = self.ego_veh.state.posF.phi
        return step_infos

    def _extract_rewards(self, infos: dict):
        # rewards design
        r = 0
        if infos["is_colliding"] == 1:
            r -= 1
        if infos["is_offroad"] == 1:
            r -= 1
        return np.array([r]).reshape(1, -1)

    def _compute_feature_infos(self, features: list):
        features = features[0]  # features size is (n_agent, n_features), so here is (1, 66)
        is_colliding = features[self.infos_cache["is_colliding_idx"]]
        markerdist_left = features[self.infos_cache["markerdist_left_idx"]]
        markerdist_right = features[self.infos_cache["markerdist_right_idx"]]
        is_offroad = features[self.infos_cache["out_of_lane_idx"]]
        return {
            "is_colliding": is_colliding,
            "markerdist_left": markerdist_left,
            "markerdist_right": markerdist_right,
            "is_offroad": is_offroad
        }

    def step(self, action: list):

        # action[0] is `a::Float64` longitudinal acceleration [m/s^2]
        # action[1] is `ω::Float64` desired turn rate [rad/sec]
        step_infos = self._step(action)

        # compute features and feature_infos
        features = self.get_features()
        feature_infos = self._compute_feature_infos(features)
        # combine infos
        infos = dict(**step_infos, **feature_infos)

        # update env timestep to be the next scene to load
        self.t += 1

        # compute terminal
        if self.t >= self.h:
            terminal = True
        elif self.terminate_on_collision and infos["is_colliding"] == 1:
            terminal = True
        elif self.terminate_on_off_road and (abs(infos["markerdist_left"]) > 3 and abs(infos["markerdist_right"]) > 3):
            terminal = True
        else:
            terminal = False
        terminal = [terminal]

        reward = self._extract_rewards(infos)
        return features, reward, terminal, infos

    def get_features(self):
        veh_idx = self.scene.findfirst(self.egoid)
        self.ext.pull_features(
            self.rec,
            self.roadway,
            veh_idx
        )
        return np.array(copy.deepcopy(self.ext.features)).reshape(1, -1)

    def observation_space_spec(self):
        low = [0 for i in range(len(self.ext))]
        high = [0 for i in range(len(self.ext))]
        feature_infos = self.ext.feature_info()
        for (i, fn) in enumerate(self.ext.feature_names()):
            low[i] = feature_infos[fn]["low"]
            high[i] = feature_infos[fn]["high"]
        infos = {"high": high, "low": low}
        return (len(self.ext),), "Box", infos

    def action_space_spec(self):
        return (2,), "Box", {"high": [4., .15], "low": [-4., -.15]}

    def obs_names(self):
        return self.ext.feature_names()

    @property
    def action_space(self):
        _, _, action_space = self.action_space_spec()
        return action_space


#!/usr/bin/python3
#coding=utf-8

import os
import cv2
import json
import torch
import numpy as np
import transform
from torch.utils.data import Dataset

class Config(object):
    def __init__(self, **kwargs):
        self.kwargs       = kwargs
        self.mean         = np.array([[[115.75, 120.83, 93.49]]])
        self.std          = np.array([[[ 50.46,  50.08, 49.67]]])
        self.cls_beg      = 0
        self.cls_end      = 1010

        ## different mode
        if self.mode == 'train':
            self.train()
        elif self.mode == 'val':
            self.val()
        elif self.mode == 'test':
            self.test()
        else:
            raise ValueError

        ## each class rate
        with open('../data/info.txt', 'r') as lines:
            self.rate  = torch.zeros(self.cls_end-self.cls_beg).cuda().float()
            for line in lines:
                cls, cnt = line.strip().split(',')
                cls, cnt = int(cls), int(cnt)
                if cls >= self.cls_beg and cls < self.cls_end:
                    self.rate[cls-self.cls_beg] = cnt

        ## print hyper parameters
        print('\nParameters...')
        for k, v in self.kwargs.items():
            print('%-10s: %s'%(k, v))
        print('%-10s: %s'%("sample", len(self.samples)))

    def train(self):
        with open('../data/train.txt', 'r') as lines:
            self.samples = []
            for line in lines:
                name, imgID, cls = line.strip().split(',')
                path = '../data/train/'+cls+'/'+name
                cls  = int(cls)
                if cls >= self.cls_beg and cls < self.cls_end:
                    self.samples.append([path, cls-self.cls_beg])

    def val(self):
        with open('../data/val.txt', 'r') as lines:
            self.samples = []
            for line in lines:
                name, imgID, cls = line.strip().split(',')
                path = '../data/train/'+cls+'/'+name
                cls  = int(cls)
                if cls >= self.cls_beg and cls < self.cls_end:
                    self.samples.append([path, cls-self.cls_beg])

    def test(self):
        with open('../data/test.txt', 'r') as lines:
            self.samples = []
            for line in lines:
                name, imgID = line.strip().split(',')
                path = '../data/test/'+name
                self.samples.append([path, imgID])

    def __getattr__(self, name):
        if name in self.kwargs:
            return self.kwargs[name]
        else:
            return None

class Data(Dataset):
    def __init__(self, cfg):
        self.cfg = cfg
        if self.cfg.mode == 'train':
            self.transform = transform.Compose( transform.Normalize(mean=cfg.mean, std=cfg.std),
                                                transform.Resize(size=448),
                                                transform.RandomHorizontalFlip(),
                                                transform.ToTensor())
        elif self.cfg.mode == 'test' or self.cfg.mode == 'val':
            self.transform = transform.Compose( transform.Normalize(mean=cfg.mean, std=cfg.std),
                                                transform.ToTensor())
        else:
            raise ValueError

    def crop(self, image, mask):
        H, W  = mask.shape
        S     = int(np.sqrt(mask.sum()*2))
        X, Y  = np.where(mask==1)
        if len(X)<10 or len(Y)<10 or S<10:
            return image
        randx = np.random.randint(len(X)//3, len(X)//3*2)
        randy = np.random.randint(len(Y)//3, len(Y)//3*2)
        xmin  = max(int(X[randx]-S/2), 0)
        ymin  = max(int(Y[randy]-S/2), 0)
        xmax  = min(int(X[randx]+S/2), H)
        ymax  = min(int(Y[randy]+S/2), W)
        image = image[xmin:xmax,ymin:ymax,:]
        return image

    def __getitem__(self, idx):
        if self.cfg.mode == 'train':
            path, label   = self.cfg.samples[idx]
            image         = cv2.imread(path).astype(np.float32)[:,:,::-1]
            path          = path.replace('train', 'train-mask').replace('.jpg', '.png')
            mask          = cv2.imread(path, -1).astype(np.float32)/255.0
            image         = self.crop(image, mask)
            image         = self.transform(image)
            return image, label
        elif self.cfg.mode == 'val':
            path, label   = self.cfg.samples[idx]
            image         = cv2.imread(path).astype(np.float32)[:,:,::-1]
            path          = path.replace('train', 'val-mask').replace('.jpg', '.png')
            mask          = cv2.imread(path, -1).astype(np.float32)/255.0
            image         = self.transform(image)
            return image, label, mask
        elif self.cfg.mode == 'test':
            path, imgID = self.cfg.samples[idx]
            image       = cv2.imread(path).astype(np.float32)[:,:,::-1]
            path        = path.replace('test', 'test-mask').replace('.jpg', '.png')
            mask        = cv2.imread(path, -1).astype(np.float32)/255.0
            image       = self.transform(image)
            return image, imgID, mask

    def __len__(self):
        return len(self.cfg.samples)

if __name__=='__main__':
    import matplotlib.pyplot as plt

    cfg    = Config(mode='train')
    data   = Data(cfg)
    for i in range(1000):
        image, label, mask = data[i]
        print(label)
        image = image.permute(1,2,0).numpy()*cfg.std+cfg.mean
        plt.subplot(121)
        plt.imshow(np.uint8(image))
        plt.subplot(122)
        plt.imshow(mask)
        plt.show()


from random import sample
from typing import List, Tuple, BinaryIO

import numpy as np

from pose_format.pose_header import PoseHeader
from pose_format.utils.reader import BufferReader, ConstStructs

POINTS_DIMS = (2, 1, 0, 3)

class PoseBody:
    tensor_reader = 'ABSTRACT-DO-NOT-USE'

    def __init__(self, fps: float, data, confidence):
        self.fps = fps
        self.data = data  # Shape (Frames, People, Points, Dims) - eg (93, 1, 137, 2)
        self.confidence = confidence  # Shape (Frames, People, Points) - eg (93, 1, 137)

    @classmethod
    def read(cls, header: PoseHeader, reader: BufferReader):
        if header.version == 0:
            return cls.read_v0_0(header, reader)
        elif round(header.version, 3) == 0.1:
            return cls.read_v0_1(header, reader)

        raise NotImplementedError("Unknown version - %f" % header.version)

    @classmethod
    def read_v0_0(cls, header: PoseHeader, reader: BufferReader):
        raise NotImplementedError("'read_v0_0' not implemented on '%s'" % cls.__class__)

    @classmethod
    def read_v0_1(cls, header: PoseHeader, reader: BufferReader):
        fps, _frames = reader.unpack(ConstStructs.double_ushort)

        _people = reader.unpack(ConstStructs.ushort)
        _points = sum([len(c.points) for c in header.components])
        _dims = max([len(c.format) for c in header.components]) - 1

        # _frames is defined as short, which sometimes is not enough! TODO change to int
        _frames = int(reader.bytes_left() / (_people * _points * (_dims + 1) * 4))

        tensor_reader = reader.__getattribute__(cls.tensor_reader)
        data = tensor_reader(ConstStructs.float, shape=(_frames, _people, _points, _dims))
        confidence = tensor_reader(ConstStructs.float, shape=(_frames, _people, _points))

        return cls(fps, data, confidence)

    def write(self, version: float, buffer: BinaryIO):
        """
        Write the data to a file based on the version of the spec
        :param version: float
        :param buffer: BinaryIO
        """
        raise NotImplementedError("'write' not implemented on '%s'" % self.__class__)

    def numpy(self):
        """
        Convert the current PoseBody representation to NumpyPoseBody
        :return: NumpyPoseBody
        """
        raise NotImplementedError("'numpy' not implemented on '%s'" % self.__class__)

    def torch(self):
        """
        Convert the current PoseBody representation to TorchPoseBody
        :return: TorchPoseBody
        """
        raise NotImplementedError("'torch' not implemented on '%s'" % self.__class__)

    def tensorflow(self):
        """
        Convert the current PoseBody representation to TensorflowPoseBody
        :return: TensorflowPoseBody
        """
        raise NotImplementedError("'tensorflow' not implemented on '%s'" % self.__class__)

    def flatten(self):
        """
        Convert the data from the (Frames, People, Points, Dims) masked representation to array of points.
        Every item in the result array contains the following dimensions:
        0. Time in milliseconds
        1. Person ID
        2. Point ID
        3. X dimension
        4. Y dimension
        5. Z dimension if exists
        6. Pose estimation confidence
        :return:
        """
        raise NotImplementedError("'flatten' not implemented on '%s'" % self.__class__)

    def slice_step(self, by: int) -> "PoseBody":
        """
        Slice the data by skipping rows.
        This slicing affects the FPS.
        :param by: take one row every "by" rows
        :return: PoseBody
        """
        new_data = self.data[::by]
        new_confidence = self.confidence[::by]
        new_fps = self.fps / by

        return self.__class__(fps=new_fps, data=new_data, confidence=new_confidence)

    def augment2d(self, rotation_std=0.2, shear_std=0.2, scale_std=0.2):
        """
        :param rotation_std: Rotation in radians
        :param shear_std: Shear X in percent
        :param scale_std: Scale X in percent
        :return:
        """
        matrix = np.eye(2)

        # Based on https://en.wikipedia.org/wiki/Shear_matrix
        if shear_std > 0:
            shear_matrix = np.eye(2)
            shear_matrix[0][1] = np.random.normal(loc=0, scale=shear_std, size=1)[0]
            matrix = np.dot(matrix, shear_matrix)

        # Based on https://en.wikipedia.org/wiki/Rotation_matrix
        if rotation_std > 0:
            rotation_angle = np.random.normal(loc=0, scale=rotation_std, size=1)[0]
            rotation_cos = np.cos(rotation_angle)
            rotation_sin = np.sin(rotation_angle)
            rotation_matrix = np.array([[rotation_cos, -rotation_sin], [rotation_sin, rotation_cos]])
            matrix = np.dot(matrix, rotation_matrix)

        # Based on https://en.wikipedia.org/wiki/Scaling_(geometry)
        if scale_std > 0:
            scale_matrix = np.eye(2)
            scale_matrix[1][1] += np.random.normal(loc=0, scale=scale_std, size=1)[0]
            matrix = np.dot(matrix, scale_matrix)

        # Cast to matrix the correct size
        dim_matrix = np.eye(self.data.shape[-1])
        dim_matrix[0:2, 0:2] = matrix

        return self.matmul(dim_matrix.astype(dtype=np.float32))

    def zero_filled(self) -> __qualname__:
        raise NotImplementedError("'zero_filled' not implemented on '%s'" % self.__class__)

    def matmul(self, matrix: np.ndarray) -> __qualname__:
        raise NotImplementedError("'matmul' not implemented on '%s'" % self.__class__)

    def get_points(self, indexes: List[int]) -> __qualname__:
        raise NotImplementedError("'get_points' not implemented on '%s'" % self.__class__)

    def bbox(self, header: PoseHeader) -> __qualname__:
        raise NotImplementedError("'bbox' not implemented on '%s'" % self.__class__)

    def points_perspective(self):
        raise NotImplementedError("'points_perspective' not implemented on '%s'" % self.__class__)

    def select_frames(self, frame_indexes: List[int]):
        data = self.data[frame_indexes]
        confidence = self.confidence[frame_indexes]
        return self.__class__(fps=self.fps, data=data, confidence=confidence)

    def frame_dropout_given_percent(self, dropout_percent: float) -> Tuple["PoseBody", List[int]]:
        data_len = len(self.data)
        dropout_number = min(int(data_len * dropout_percent), int(data_len * 0.99))
        dropout_indexes = set(sample(range(0, data_len), dropout_number))
        select_indexes = [i for i in range(0, data_len) if i not in dropout_indexes]

        return self.select_frames(select_indexes), select_indexes

    def frame_dropout_uniform(self,
                              dropout_min: float = 0.2,
                              dropout_max: float = 1.0) -> Tuple["PoseBody", List[int]]:
        dropout_percent = np.random.uniform(low=dropout_min, high=dropout_max, size=1)[0]

        return self.frame_dropout_given_percent(dropout_percent)

    def frame_dropout_normal(self,
                             dropout_mean: float = 0.5,
                             dropout_std: float = 0.1) -> Tuple["PoseBody", List[int]]:
        dropout_percent = np.abs(np.random.normal(loc=dropout_mean, scale=dropout_std, size=1))[0]

        return self.frame_dropout_given_percent(dropout_percent)


import copy
import time
import collections

import numpy as np

from .. import util
from .. import caching
from .. import exceptions
from .. import transformations

try:
    import networkx as nx
    _ForestParent = nx.DiGraph
except BaseException as E:
    # create a dummy module which will raise the ImportError
    # or other exception only when someone tries to use networkx
    nx = exceptions.ExceptionModule(E)
    _ForestParent = object

class TransformForest(object):
    def __init__(self, base_frame='world'):
        # a graph structure, subclass of networkx DiGraph
        self.transforms = EnforcedForest()
        # hashable, the base or root frame
        self.base_frame = base_frame

        # save paths, keyed with tuple (from, to)
        self._paths = {}
        # cache transformation matrices keyed with tuples
        self._updated = str(np.random.random())
        self._cache = caching.Cache(self.md5)

    def update(self, frame_to, frame_from=None, **kwargs):
        """
        Update a transform in the tree.

        Parameters
        ------------
        frame_from : hashable object
          Usually a string (eg 'world').
          If left as None it will be set to self.base_frame
        frame_to :  hashable object
          Usually a string (eg 'mesh_0')
        matrix : (4,4) float
          Homogeneous transformation matrix
        quaternion :  (4,) float
          Quaternion ordered [w, x, y, z]
        axis : (3,) float
          Axis of rotation
        angle :  float
          Angle of rotation, in radians
        translation : (3,) float
          Distance to translate
        geometry : hashable
          Geometry object name, e.g. 'mesh_0'
        """

        self._updated = str(np.random.random())
        self._cache.clear()

        # if no frame specified, use base frame
        if frame_from is None:
            frame_from = self.base_frame
        # convert various kwargs to a single matrix
        matrix = kwargs_to_matrix(**kwargs)

        # create the edge attributes
        attr = {'matrix': matrix, 'time': time.time()}
        # pass through geometry to edge attribute
        if 'geometry' in kwargs:
            attr['geometry'] = kwargs['geometry']

        # add the edges
        changed = self.transforms.add_edge(frame_from,
                                           frame_to,
                                           **attr)
        # set the node attribute with the geometry information
        if 'geometry' in kwargs:
            nx.set_node_attributes(
                self.transforms,
                name='geometry',
                values={frame_to: kwargs['geometry']})
        # if the edge update changed our structure
        # dump our cache of shortest paths
        if changed:
            self._paths = {}

    def md5(self):
        return self._updated

    def copy(self):
        """
        Return a copy of the current TransformForest

        Returns
        ------------
        copied: TransformForest
        """
        copied = TransformForest()
        copied.base_frame = copy.deepcopy(self.base_frame)
        copied.transforms = copy.deepcopy(self.transforms)

        return copied

    def to_flattened(self, base_frame=None):
        """
        Export the current transform graph as a flattened
        """
        if base_frame is None:
            base_frame = self.base_frame

        flat = {}
        for node in self.nodes:
            if node == base_frame:
                continue
            transform, geometry = self.get(
                frame_to=node, frame_from=base_frame)
            flat[node] = {
                'transform': transform.tolist(),
                'geometry': geometry
            }
        return flat

    def to_gltf(self, scene):
        """
        Export a transforms as the 'nodes' section of a GLTF dict.
        Flattens tree.

        Returns
        --------
        gltf : dict
          with keys:
                  'nodes': list of dicts
        """
        # geometry is an OrderedDict
        # {geometry key : index}
        mesh_index = {name: i for i, name
                      in enumerate(scene.geometry.keys())}
        # save the output
        gltf = collections.deque([])
        for node in self.nodes:
            # don't include edge for base frame
            if node == self.base_frame:
                continue
            # get the transform and geometry from the graph
            transform, geometry = self.get(
                frame_to=node, frame_from=self.base_frame)

            gltf.append({
                'matrix': transform.T.reshape(-1).tolist(),
                'name': node})
            # assign geometry if it exists
            if geometry is not None:
                gltf[-1]['mesh'] = mesh_index[geometry]

            if node == scene.camera.name:
                gltf[-1]['camera'] = 0

        # we have flattened tree, so all nodes will be child of world
        gltf.appendleft({
            'name': self.base_frame,
            'children': list(range(1, 1 + len(gltf)))
        })
        result = {'nodes': list(gltf)}

        return result

    def to_edgelist(self):
        """
        Export the current transforms as a list of edge tuples, with
        each tuple having the format:
        (node_a, node_b, {metadata})

        Returns
        ---------
        edgelist: (n,) list of tuples
        """
        # save cleaned edges
        export = []
        # loop through (node, node, edge attributes)
        for edge in nx.to_edgelist(self.transforms):
            a, b, c = edge
            # geometry is a node property but save it to the
            # edge so we don't need two dictionaries
            if 'geometry' in self.transforms.node[b]:
                c['geometry'] = self.transforms.node[b]['geometry']
            # save the matrix as a float list
            c['matrix'] = np.asanyarray(c['matrix'], dtype=np.float64).tolist()
            export.append((a, b, c))
        return export

    def from_edgelist(self, edges, strict=True):
        """
        Load transform data from an edge list into the current
        scene graph.

        Parameters
        -------------
        edgelist : (n,) tuples
            (node_a, node_b, {key: value})
        strict : bool
            If true, raise a ValueError when a
            malformed edge is passed in a tuple.
        """
        # loop through each edge
        for edge in edges:
            # edge contains attributes
            if len(edge) == 3:
                self.update(edge[1], edge[0], **edge[2])
            # edge just contains nodes
            elif len(edge) == 2:
                self.update(edge[1], edge[0])
            # edge is broken
            elif strict:
                raise ValueError('edge incorrect shape: {}'.format(str(edge)))

    def load(self, edgelist):
        """
        Load transform data from an edge list into the current
        scene graph.

        Parameters
        -------------
        edgelist : (n,) tuples
            (node_a, node_b, {key: value})
        """
        self.from_edgelist(edgelist, strict=True)

    @caching.cache_decorator
    def nodes(self):
        """
        A list of every node in the graph.

        Returns
        -------------
        nodes: (n,) array, of node names
        """
        nodes = np.array(list(self.transforms.nodes()))
        return nodes

    @caching.cache_decorator
    def nodes_geometry(self):
        """
        The nodes in the scene graph with geometry attached.

        Returns
        ------------
        nodes_geometry: (m,) array, of node names
        """

        nodes = np.array([
            n for n in self.transforms.nodes()
            if 'geometry' in self.transforms.node[n]])

        return nodes

    def get(self, frame_to, frame_from=None):
        """
        Get the transform from one frame to another, assuming they are connected
        in the transform tree.

        If the frames are not connected a NetworkXNoPath error will be raised.

        Parameters
        ------------
        frame_to : hashable
          Node name, usually a string (eg 'mesh_0')
        frame_from : hashable
          Node name, usually a string (eg 'world').
          If None it will be set to self.base_frame

        Returns
        ----------
        transform : (4, 4) float
          Homogeneous transformation matrix
        """

        if frame_from is None:
            frame_from = self.base_frame

        # look up transform to see if we have it already
        cache_key = (frame_from, frame_to)
        cached = self._cache[cache_key]
        if cached is not None:
            return cached

        # get the path in the graph
        path = self._get_path(frame_from, frame_to)

        # collect transforms along the path
        transforms = []

        for i in range(len(path) - 1):
            # get the matrix and edge direction
            data, direction = self.transforms.get_edge_data_direction(
                path[i], path[i + 1])
            matrix = data['matrix']
            if direction < 0:
                matrix = np.linalg.inv(matrix)
            transforms.append(matrix)
        # do all dot products at the end
        if len(transforms) == 0:
            transform = np.eye(4)
        elif len(transforms) == 1:
            transform = np.asanyarray(transforms[0], dtype=np.float64)
        else:
            transform = util.multi_dot(transforms)

        geometry = None
        if 'geometry' in self.transforms.node[frame_to]:
            geometry = self.transforms.node[frame_to]['geometry']

        self._cache[cache_key] = (transform, geometry)

        return transform, geometry

    def show(self):
        """
        Plot the graph layout of the scene.
        """
        import matplotlib.pyplot as plt
        nx.draw(self.transforms, with_labels=True)
        plt.show()

    def to_svg(self):
        """
        """
        from ..graph import graph_to_svg
        return graph_to_svg(self.transforms)

    def __contains__(self, key):
        return key in self.transforms.node

    def __getitem__(self, key):
        return self.get(key)

    def __setitem__(self, key, value):
        value = np.asanyarray(value)
        if value.shape != (4, 4):
            raise ValueError('Matrix must be specified!')
        return self.update(key, matrix=value)

    def clear(self):
        self.transforms = EnforcedForest()
        self._paths = {}
        self._updated = str(np.random.random())
        self._cache.clear()

    def _get_path(self, frame_from, frame_to):
        """
        Find a path between two frames, either from cached paths or
        from the transform graph.

        Parameters
        ------------
        frame_from: a frame key, usually a string
                    eg, 'world'
        frame_to:   a frame key, usually a string
                    eg, 'mesh_0'

        Returns
        ----------
        path: (n) list of frame keys
              eg, ['mesh_finger', 'mesh_hand', 'world']
        """
        # store paths keyed as a tuple
        key = (frame_from, frame_to)
        if key not in self._paths:
            # get the actual shortest paths
            path = self.transforms.shortest_path_undirected(
                frame_from, frame_to)
            # store path to avoid recomputing
            self._paths[key] = path
            return path
        return self._paths[key]

class EnforcedForest(_ForestParent):
    """
    A subclass of networkx.DiGraph that will raise an error if an
    edge is added which would make the DiGraph not a forest or tree.
    """

    def __init__(self, *args, **kwargs):
        self.flags = {'strict': False, 'assert_forest': False}

        for k, v in self.flags.items():
            if k in kwargs:
                self.flags[k] = bool(kwargs[k])
                kwargs.pop(k, None)

        super(self.__class__, self).__init__(*args, **kwargs)
        # keep a second parallel but undirected copy of the graph
        # all of the networkx methods for turning a directed graph
        # into an undirected graph are quite slow so we do minor bookkeeping
        self._undirected = nx.Graph()

    def add_edge(self, u, v, *args, **kwargs):
        changed = False
        if u == v:
            if self.flags['strict']:
                raise ValueError('Edge must be between two unique nodes!')
            return changed
        if self._undirected.has_edge(u, v):
            self.remove_edges_from([[u, v], [v, u]])
        elif len(self.nodes()) > 0:
            try:
                path = nx.shortest_path(self._undirected, u, v)
                if self.flags['strict']:
                    raise ValueError(
                        'Multiple edge path exists between nodes!')
                self.disconnect_path(path)
                changed = True
            except (nx.NetworkXError, nx.NetworkXNoPath, nx.NetworkXException):
                pass
        self._undirected.add_edge(u, v)
        super(self.__class__, self).add_edge(u, v, *args, **kwargs)

        if self.flags['assert_forest']:
            # this is quite slow but makes very sure structure is correct
            # so is mainly used for testing
            assert nx.is_forest(nx.Graph(self))

        return changed

    def add_edges_from(self, *args, **kwargs):
        raise ValueError('EnforcedTree requires add_edge method to be used!')

    def add_path(self, *args, **kwargs):
        raise ValueError('EnforcedTree requires add_edge method to be used!')

    def remove_edge(self, *args, **kwargs):
        super(self.__class__, self).remove_edge(*args, **kwargs)
        self._undirected.remove_edge(*args, **kwargs)

    def remove_edges_from(self, *args, **kwargs):
        super(self.__class__, self).remove_edges_from(*args, **kwargs)
        self._undirected.remove_edges_from(*args, **kwargs)

    def disconnect_path(self, path):
        ebunch = np.array([[path[0], path[1]]])
        ebunch = np.vstack((ebunch, np.fliplr(ebunch)))
        self.remove_edges_from(ebunch)

    def shortest_path_undirected(self, u, v):
        try:
            path = nx.shortest_path(self._undirected, u, v)
        except BaseException as E:
            print(u, v)
            raise E
        return path

    def get_edge_data_direction(self, u, v):
        if self.has_edge(u, v):
            direction = 1
        elif self.has_edge(v, u):
            direction = -1
        else:
            raise ValueError('Edge does not exist!')
        data = self.get_edge_data(*[u, v][::direction])
        return data, direction

def kwargs_to_matrix(**kwargs):
    """
    Turn a set of keyword arguments into a transformation matrix.
    """
    if 'matrix' in kwargs:
        # a matrix takes precedence over other options
        matrix = np.asanyarray(kwargs['matrix'], dtype=np.float64)
    elif 'quaternion' in kwargs:
        matrix = transformations.quaternion_matrix(kwargs['quaternion'])
    elif ('axis' in kwargs) and ('angle' in kwargs):
        matrix = transformations.rotation_matrix(kwargs['angle'],
                                                 kwargs['axis'])
    else:
        raise ValueError('Couldn\'t update transform!')

    if 'translation' in kwargs:
        # translation can be used in conjunction with any of the methods of
        # specifying transforms. In the case a matrix and translation are passed,
        # we add the translations together rather than picking one.
        matrix[0:3, 3] += kwargs['translation']
    return matrix


# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
from ppdet.core.workspace import register
from ppdet.modeling.bbox_utils import nonempty_bbox, rbox2poly
from ppdet.modeling.layers import TTFBox
from .transformers import bbox_cxcywh_to_xyxy
try:
    from collections.abc import Sequence
except Exception:
    from collections import Sequence

__all__ = [
    'BBoxPostProcess', 'MaskPostProcess', 'FCOSPostProcess',
    'S2ANetBBoxPostProcess', 'JDEBBoxPostProcess', 'CenterNetPostProcess',
    'DETRBBoxPostProcess', 'SparsePostProcess'
]

@register
class BBoxPostProcess(nn.Layer):
    __shared__ = ['num_classes']
    __inject__ = ['decode', 'nms']

    def __init__(self, num_classes=80, decode=None, nms=None):
        super(BBoxPostProcess, self).__init__()
        self.num_classes = num_classes
        self.decode = decode
        self.nms = nms
        self.fake_bboxes = paddle.to_tensor(
            np.array(
                [[-1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype='float32'))
        self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32'))

    def forward(self, head_out, rois, im_shape, scale_factor):
        """
        Decode the bbox and do NMS if needed.

        Args:
            head_out (tuple): bbox_pred and cls_prob of bbox_head output.
            rois (tuple): roi and rois_num of rpn_head output.
            im_shape (Tensor): The shape of the input image.
            scale_factor (Tensor): The scale factor of the input image.
        Returns:
            bbox_pred (Tensor): The output prediction with shape [N, 6], including
                labels, scores and bboxes. The size of bboxes are corresponding
                to the input image, the bboxes may be used in other branch.
            bbox_num (Tensor): The number of prediction boxes of each batch with
                shape [1], and is N.
        """
        if self.nms is not None:
            bboxes, score = self.decode(head_out, rois, im_shape, scale_factor)
            bbox_pred, bbox_num, _ = self.nms(bboxes, score, self.num_classes)
        else:
            bbox_pred, bbox_num = self.decode(head_out, rois, im_shape,
                                              scale_factor)
        return bbox_pred, bbox_num

    def get_pred(self, bboxes, bbox_num, im_shape, scale_factor):
        """
        Rescale, clip and filter the bbox from the output of NMS to
        get final prediction.

        Notes:
        Currently only support bs = 1.

        Args:
            bboxes (Tensor): The output bboxes with shape [N, 6] after decode
                and NMS, including labels, scores and bboxes.
            bbox_num (Tensor): The number of prediction boxes of each batch with
                shape [1], and is N.
            im_shape (Tensor): The shape of the input image.
            scale_factor (Tensor): The scale factor of the input image.
        Returns:
            pred_result (Tensor): The final prediction results with shape [N, 6]
                including labels, scores and bboxes.
        """

        if bboxes.shape[0] == 0:
            bboxes = self.fake_bboxes
            bbox_num = self.fake_bbox_num

        origin_shape = paddle.floor(im_shape / scale_factor + 0.5)

        origin_shape_list = []
        scale_factor_list = []
        # scale_factor: scale_y, scale_x
        for i in range(bbox_num.shape[0]):
            expand_shape = paddle.expand(origin_shape[i:i + 1, :],
                                         [bbox_num[i], 2])
            scale_y, scale_x = scale_factor[i][0], scale_factor[i][1]
            scale = paddle.concat([scale_x, scale_y, scale_x, scale_y])
            expand_scale = paddle.expand(scale, [bbox_num[i], 4])
            origin_shape_list.append(expand_shape)
            scale_factor_list.append(expand_scale)

        self.origin_shape_list = paddle.concat(origin_shape_list)
        scale_factor_list = paddle.concat(scale_factor_list)

        # bboxes: [N, 6], label, score, bbox
        pred_label = bboxes[:, 0:1]
        pred_score = bboxes[:, 1:2]
        pred_bbox = bboxes[:, 2:]
        # rescale bbox to original image
        scaled_bbox = pred_bbox / scale_factor_list
        origin_h = self.origin_shape_list[:, 0]
        origin_w = self.origin_shape_list[:, 1]
        zeros = paddle.zeros_like(origin_h)
        # clip bbox to [0, original_size]
        x1 = paddle.maximum(paddle.minimum(scaled_bbox[:, 0], origin_w), zeros)
        y1 = paddle.maximum(paddle.minimum(scaled_bbox[:, 1], origin_h), zeros)
        x2 = paddle.maximum(paddle.minimum(scaled_bbox[:, 2], origin_w), zeros)
        y2 = paddle.maximum(paddle.minimum(scaled_bbox[:, 3], origin_h), zeros)
        pred_bbox = paddle.stack([x1, y1, x2, y2], axis=-1)
        # filter empty bbox
        keep_mask = nonempty_bbox(pred_bbox, return_mask=True)
        keep_mask = paddle.unsqueeze(keep_mask, [1])
        pred_label = paddle.where(keep_mask, pred_label,
                                  paddle.ones_like(pred_label) * -1)
        pred_result = paddle.concat([pred_label, pred_score, pred_bbox], axis=1)
        return pred_result

    def get_origin_shape(self, ):
        return self.origin_shape_list

@register
class MaskPostProcess(object):
    def __init__(self, binary_thresh=0.5):
        super(MaskPostProcess, self).__init__()
        self.binary_thresh = binary_thresh

    def paste_mask(self, masks, boxes, im_h, im_w):
        """
        Paste the mask prediction to the original image.
        """
        x0, y0, x1, y1 = paddle.split(boxes, 4, axis=1)
        masks = paddle.unsqueeze(masks, [0, 1])
        img_y = paddle.arange(0, im_h, dtype='float32') + 0.5
        img_x = paddle.arange(0, im_w, dtype='float32') + 0.5
        img_y = (img_y - y0) / (y1 - y0) * 2 - 1
        img_x = (img_x - x0) / (x1 - x0) * 2 - 1
        img_x = paddle.unsqueeze(img_x, [1])
        img_y = paddle.unsqueeze(img_y, [2])
        N = boxes.shape[0]

        gx = paddle.expand(img_x, [N, img_y.shape[1], img_x.shape[2]])
        gy = paddle.expand(img_y, [N, img_y.shape[1], img_x.shape[2]])
        grid = paddle.stack([gx, gy], axis=3)
        img_masks = F.grid_sample(masks, grid, align_corners=False)
        return img_masks[:, 0]

    def __call__(self, mask_out, bboxes, bbox_num, origin_shape):
        """
        Decode the mask_out and paste the mask to the origin image.

        Args:
            mask_out (Tensor): mask_head output with shape [N, 28, 28].
            bbox_pred (Tensor): The output bboxes with shape [N, 6] after decode
                and NMS, including labels, scores and bboxes.
            bbox_num (Tensor): The number of prediction boxes of each batch with
                shape [1], and is N.
            origin_shape (Tensor): The origin shape of the input image, the tensor
                shape is [N, 2], and each row is [h, w].
        Returns:
            pred_result (Tensor): The final prediction mask results with shape
                [N, h, w] in binary mask style.
        """
        num_mask = mask_out.shape[0]
        origin_shape = paddle.cast(origin_shape, 'int32')
        # TODO: support bs > 1 and mask output dtype is bool
        pred_result = paddle.zeros(
            [num_mask, origin_shape[0][0], origin_shape[0][1]], dtype='int32')
        if bbox_num == 1 and bboxes[0][0] == -1:
            return pred_result

        # TODO: optimize chunk paste
        pred_result = []
        for i in range(bboxes.shape[0]):
            im_h, im_w = origin_shape[i][0], origin_shape[i][1]
            pred_mask = self.paste_mask(mask_out[i], bboxes[i:i + 1, 2:], im_h,
                                        im_w)
            pred_mask = pred_mask >= self.binary_thresh
            pred_mask = paddle.cast(pred_mask, 'int32')
            pred_result.append(pred_mask)
        pred_result = paddle.concat(pred_result)
        return pred_result

@register
class FCOSPostProcess(object):
    __inject__ = ['decode', 'nms']

    def __init__(self, decode=None, nms=None):
        super(FCOSPostProcess, self).__init__()
        self.decode = decode
        self.nms = nms

    def __call__(self, fcos_head_outs, scale_factor):
        """
        Decode the bbox and do NMS in FCOS.
        """
        locations, cls_logits, bboxes_reg, centerness = fcos_head_outs
        bboxes, score = self.decode(locations, cls_logits, bboxes_reg,
                                    centerness, scale_factor)
        bbox_pred, bbox_num, _ = self.nms(bboxes, score)
        return bbox_pred, bbox_num

@register
class S2ANetBBoxPostProcess(nn.Layer):
    __shared__ = ['num_classes']
    __inject__ = ['nms']

    def __init__(self, num_classes=15, nms_pre=2000, min_bbox_size=0, nms=None):
        super(S2ANetBBoxPostProcess, self).__init__()
        self.num_classes = num_classes
        self.nms_pre = paddle.to_tensor(nms_pre)
        self.min_bbox_size = min_bbox_size
        self.nms = nms
        self.origin_shape_list = []
        self.fake_pred_cls_score_bbox = paddle.to_tensor(
            np.array(
                [[-1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],
                dtype='float32'))
        self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32'))

    def forward(self, pred_scores, pred_bboxes):
        """
        pred_scores : [N, M]  score
        pred_bboxes : [N, 5]  xc, yc, w, h, a
        im_shape : [N, 2]  im_shape
        scale_factor : [N, 2]  scale_factor
        """
        pred_ploys0 = rbox2poly(pred_bboxes)
        pred_ploys = paddle.unsqueeze(pred_ploys0, axis=0)

        # pred_scores [NA, 16] --> [16, NA]
        pred_scores0 = paddle.transpose(pred_scores, [1, 0])
        pred_scores = paddle.unsqueeze(pred_scores0, axis=0)

        pred_cls_score_bbox, bbox_num, _ = self.nms(pred_ploys, pred_scores,
                                                    self.num_classes)
        # Prevent empty bbox_pred from decode or NMS.
        # Bboxes and score before NMS may be empty due to the score threshold.
        if pred_cls_score_bbox.shape[0] <= 0 or pred_cls_score_bbox.shape[
                1] <= 1:
            pred_cls_score_bbox = self.fake_pred_cls_score_bbox
            bbox_num = self.fake_bbox_num

        pred_cls_score_bbox = paddle.reshape(pred_cls_score_bbox, [-1, 10])
        return pred_cls_score_bbox, bbox_num

    def get_pred(self, bboxes, bbox_num, im_shape, scale_factor):
        """
        Rescale, clip and filter the bbox from the output of NMS to
        get final prediction.
        Args:
            bboxes(Tensor): bboxes [N, 10]
            bbox_num(Tensor): bbox_num
            im_shape(Tensor): [1 2]
            scale_factor(Tensor): [1 2]
        Returns:
            bbox_pred(Tensor): The output is the prediction with shape [N, 8]
                               including labels, scores and bboxes. The size of
                               bboxes are corresponding to the original image.
        """
        origin_shape = paddle.floor(im_shape / scale_factor + 0.5)

        origin_shape_list = []
        scale_factor_list = []
        # scale_factor: scale_y, scale_x
        for i in range(bbox_num.shape[0]):
            expand_shape = paddle.expand(origin_shape[i:i + 1, :],
                                         [bbox_num[i], 2])
            scale_y, scale_x = scale_factor[i][0], scale_factor[i][1]
            scale = paddle.concat([
                scale_x, scale_y, scale_x, scale_y, scale_x, scale_y, scale_x,
                scale_y
            ])
            expand_scale = paddle.expand(scale, [bbox_num[i], 8])
            origin_shape_list.append(expand_shape)
            scale_factor_list.append(expand_scale)

        origin_shape_list = paddle.concat(origin_shape_list)
        scale_factor_list = paddle.concat(scale_factor_list)

        # bboxes: [N, 10], label, score, bbox
        pred_label_score = bboxes[:, 0:2]
        pred_bbox = bboxes[:, 2:]

        # rescale bbox to original image
        pred_bbox = pred_bbox.reshape([-1, 8])
        scaled_bbox = pred_bbox / scale_factor_list
        origin_h = origin_shape_list[:, 0]
        origin_w = origin_shape_list[:, 1]

        bboxes = scaled_bbox
        zeros = paddle.zeros_like(origin_h)
        x1 = paddle.maximum(paddle.minimum(bboxes[:, 0], origin_w - 1), zeros)
        y1 = paddle.maximum(paddle.minimum(bboxes[:, 1], origin_h - 1), zeros)
        x2 = paddle.maximum(paddle.minimum(bboxes[:, 2], origin_w - 1), zeros)
        y2 = paddle.maximum(paddle.minimum(bboxes[:, 3], origin_h - 1), zeros)
        x3 = paddle.maximum(paddle.minimum(bboxes[:, 4], origin_w - 1), zeros)
        y3 = paddle.maximum(paddle.minimum(bboxes[:, 5], origin_h - 1), zeros)
        x4 = paddle.maximum(paddle.minimum(bboxes[:, 6], origin_w - 1), zeros)
        y4 = paddle.maximum(paddle.minimum(bboxes[:, 7], origin_h - 1), zeros)
        pred_bbox = paddle.stack([x1, y1, x2, y2, x3, y3, x4, y4], axis=-1)
        pred_result = paddle.concat([pred_label_score, pred_bbox], axis=1)
        return pred_result

@register
class JDEBBoxPostProcess(nn.Layer):
    __shared__ = ['num_classes']
    __inject__ = ['decode', 'nms']

    def __init__(self, num_classes=1, decode=None, nms=None, return_idx=True):
        super(JDEBBoxPostProcess, self).__init__()
        self.num_classes = num_classes
        self.decode = decode
        self.nms = nms
        self.return_idx = return_idx

        self.fake_bbox_pred = paddle.to_tensor(
            np.array(
                [[-1, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype='float32'))
        self.fake_bbox_num = paddle.to_tensor(np.array([1], dtype='int32'))
        self.fake_nms_keep_idx = paddle.to_tensor(
            np.array(
                [[0]], dtype='int32'))

        self.fake_yolo_boxes_out = paddle.to_tensor(
            np.array(
                [[[0.0, 0.0, 0.0, 0.0]]], dtype='float32'))
        self.fake_yolo_scores_out = paddle.to_tensor(
            np.array(
                [[[0.0]]], dtype='float32'))
        self.fake_boxes_idx = paddle.to_tensor(np.array([[0]], dtype='int64'))

    def forward(self, head_out, anchors):
        """
        Decode the bbox and do NMS for JDE model.

        Args:
            head_out (list): Bbox_pred and cls_prob of bbox_head output.
            anchors (list): Anchors of JDE model.

        Returns:
            boxes_idx (Tensor): The index of kept bboxes after decode 'JDEBox'.
            bbox_pred (Tensor): The output is the prediction with shape [N, 6]
                including labels, scores and bboxes.
            bbox_num (Tensor): The number of prediction of each batch with shape [N].
            nms_keep_idx (Tensor): The index of kept bboxes after NMS.
        """
        boxes_idx, yolo_boxes_scores = self.decode(head_out, anchors)

        if len(boxes_idx) == 0:
            boxes_idx = self.fake_boxes_idx
            yolo_boxes_out = self.fake_yolo_boxes_out
            yolo_scores_out = self.fake_yolo_scores_out
        else:
            yolo_boxes = paddle.gather_nd(yolo_boxes_scores, boxes_idx)
            # TODO: only support bs=1 now
            yolo_boxes_out = paddle.reshape(
                yolo_boxes[:, :4], shape=[1, len(boxes_idx), 4])
            yolo_scores_out = paddle.reshape(
                yolo_boxes[:, 4:5], shape=[1, 1, len(boxes_idx)])
            boxes_idx = boxes_idx[:, 1:]

        if self.return_idx:
            bbox_pred, bbox_num, nms_keep_idx = self.nms(
                yolo_boxes_out, yolo_scores_out, self.num_classes)
            if bbox_pred.shape[0] == 0:
                bbox_pred = self.fake_bbox_pred
                bbox_num = self.fake_bbox_num
                nms_keep_idx = self.fake_nms_keep_idx
            return boxes_idx, bbox_pred, bbox_num, nms_keep_idx
        else:
            bbox_pred, bbox_num, _ = self.nms(yolo_boxes_out, yolo_scores_out,
                                              self.num_classes)
            if bbox_pred.shape[0] == 0:
                bbox_pred = self.fake_bbox_pred
                bbox_num = self.fake_bbox_num
            return _, bbox_pred, bbox_num, _

@register
class CenterNetPostProcess(TTFBox):
    """
    Postprocess the model outputs to get final prediction:
        1. Do NMS for heatmap to get top `max_per_img` bboxes.
        2. Decode bboxes using center offset and box size.
        3. Rescale decoded bboxes reference to the origin image shape.

    Args:
        max_per_img(int): the maximum number of predicted objects in a image,
            500 by default.
        down_ratio(int): the down ratio from images to heatmap, 4 by default.
        regress_ltrb (bool): whether to regress left/top/right/bottom or
            width/height for a box, true by default.
        for_mot (bool): whether return other features used in tracking model.

    """

    __shared__ = ['down_ratio']

    def __init__(self,
                 max_per_img=500,
                 down_ratio=4,
                 regress_ltrb=True,
                 for_mot=False):
        super(TTFBox, self).__init__()
        self.max_per_img = max_per_img
        self.down_ratio = down_ratio
        self.regress_ltrb = regress_ltrb
        self.for_mot = for_mot

    def __call__(self, hm, wh, reg, im_shape, scale_factor):
        heat = self._simple_nms(hm)
        scores, inds, clses, ys, xs = self._topk(heat)
        scores = paddle.tensor.unsqueeze(scores, [1])
        clses = paddle.tensor.unsqueeze(clses, [1])

        reg_t = paddle.transpose(reg, [0, 2, 3, 1])
        # Like TTFBox, batch size is 1.
        # TODO: support batch size > 1
        reg = paddle.reshape(reg_t, [-1, paddle.shape(reg_t)[-1]])
        reg = paddle.gather(reg, inds)
        xs = paddle.cast(xs, 'float32')
        ys = paddle.cast(ys, 'float32')
        xs = xs + reg[:, 0:1]
        ys = ys + reg[:, 1:2]

        wh_t = paddle.transpose(wh, [0, 2, 3, 1])
        wh = paddle.reshape(wh_t, [-1, paddle.shape(wh_t)[-1]])
        wh = paddle.gather(wh, inds)

        if self.regress_ltrb:
            x1 = xs - wh[:, 0:1]
            y1 = ys - wh[:, 1:2]
            x2 = xs + wh[:, 2:3]
            y2 = ys + wh[:, 3:4]
        else:
            x1 = xs - wh[:, 0:1] / 2
            y1 = ys - wh[:, 1:2] / 2
            x2 = xs + wh[:, 0:1] / 2
            y2 = ys + wh[:, 1:2] / 2

        n, c, feat_h, feat_w = hm.shape[:]
        padw = (feat_w * self.down_ratio - im_shape[0, 1]) / 2
        padh = (feat_h * self.down_ratio - im_shape[0, 0]) / 2
        x1 = x1 * self.down_ratio
        y1 = y1 * self.down_ratio
        x2 = x2 * self.down_ratio
        y2 = y2 * self.down_ratio

        x1 = x1 - padw
        y1 = y1 - padh
        x2 = x2 - padw
        y2 = y2 - padh

        bboxes = paddle.concat([x1, y1, x2, y2], axis=1)
        scale_y = scale_factor[:, 0:1]
        scale_x = scale_factor[:, 1:2]
        scale_expand = paddle.concat(
            [scale_x, scale_y, scale_x, scale_y], axis=1)
        boxes_shape = paddle.shape(bboxes)
        boxes_shape.stop_gradient = True
        scale_expand = paddle.expand(scale_expand, shape=boxes_shape)
        bboxes = paddle.divide(bboxes, scale_expand)
        if self.for_mot:
            results = paddle.concat([bboxes, scores, clses], axis=1)
            return results, inds
        else:
            results = paddle.concat([clses, scores, bboxes], axis=1)
            return results, paddle.shape(results)[0:1]

@register
class DETRBBoxPostProcess(object):
    __shared__ = ['num_classes', 'use_focal_loss']
    __inject__ = []

    def __init__(self,
                 num_classes=80,
                 num_top_queries=100,
                 use_focal_loss=False):
        super(DETRBBoxPostProcess, self).__init__()
        self.num_classes = num_classes
        self.num_top_queries = num_top_queries
        self.use_focal_loss = use_focal_loss

    def __call__(self, head_out, im_shape, scale_factor):
        """
        Decode the bbox.

        Args:
            head_out (tuple): bbox_pred, cls_logit and masks of bbox_head output.
            im_shape (Tensor): The shape of the input image.
            scale_factor (Tensor): The scale factor of the input image.
        Returns:
            bbox_pred (Tensor): The output prediction with shape [N, 6], including
                labels, scores and bboxes. The size of bboxes are corresponding
                to the input image, the bboxes may be used in other branch.
            bbox_num (Tensor): The number of prediction boxes of each batch with
                shape [bs], and is N.
        """
        bboxes, logits, masks = head_out

        bbox_pred = bbox_cxcywh_to_xyxy(bboxes)
        origin_shape = paddle.floor(im_shape / scale_factor + 0.5)
        img_h, img_w = origin_shape.unbind(1)
        origin_shape = paddle.stack(
            [img_w, img_h, img_w, img_h], axis=-1).unsqueeze(0)
        bbox_pred *= origin_shape

        scores = F.sigmoid(logits) if self.use_focal_loss else F.softmax(
            logits)[:, :, :-1]

        if not self.use_focal_loss:
            scores, labels = scores.max(-1), scores.argmax(-1)
            if scores.shape[1] > self.num_top_queries:
                scores, index = paddle.topk(
                    scores, self.num_top_queries, axis=-1)
                labels = paddle.stack(
                    [paddle.gather(l, i) for l, i in zip(labels, index)])
                bbox_pred = paddle.stack(
                    [paddle.gather(b, i) for b, i in zip(bbox_pred, index)])
        else:
            scores, index = paddle.topk(
                scores.reshape([logits.shape[0], -1]),
                self.num_top_queries,
                axis=-1)
            labels = index % logits.shape[2]
            index = index // logits.shape[2]
            bbox_pred = paddle.stack(
                [paddle.gather(b, i) for b, i in zip(bbox_pred, index)])

        bbox_pred = paddle.concat(
            [
                labels.unsqueeze(-1).astype('float32'), scores.unsqueeze(-1),
                bbox_pred
            ],
            axis=-1)
        bbox_num = paddle.to_tensor(
            bbox_pred.shape[1], dtype='int32').tile([bbox_pred.shape[0]])
        bbox_pred = bbox_pred.reshape([-1, 6])
        return bbox_pred, bbox_num

@register
class SparsePostProcess(object):
    __shared__ = ['num_classes']

    def __init__(self, num_proposals, num_classes=80):
        super(SparsePostProcess, self).__init__()
        self.num_classes = num_classes
        self.num_proposals = num_proposals

    def __call__(self, box_cls, box_pred, scale_factor_wh, img_whwh):
        """
        Arguments:
            box_cls (Tensor): tensor of shape (batch_size, num_proposals, K).
                The tensor predicts the classification probability for each proposal.
            box_pred (Tensor): tensors of shape (batch_size, num_proposals, 4).
                The tensor predicts 4-vector (x,y,w,h) box
                regression values for every proposal
            scale_factor_wh (Tensor): tensors of shape [batch_size, 2] the scalor of  per img
            img_whwh (Tensor): tensors of shape [batch_size, 4]
        Returns:
            bbox_pred (Tensor): tensors of shape [num_boxes, 6] Each row has 6 values:
            [label, confidence, xmin, ymin, xmax, ymax]
            bbox_num (Tensor): tensors of shape [batch_size] the number of RoIs in each image.
        """
        assert len(box_cls) == len(scale_factor_wh) == len(img_whwh)

        img_wh = img_whwh[:, :2]

        scores = F.sigmoid(box_cls)
        labels = paddle.arange(0, self.num_classes). \
            unsqueeze(0).tile([self.num_proposals, 1]).flatten(start_axis=0, stop_axis=1)

        classes_all = []
        scores_all = []
        boxes_all = []
        for i, (scores_per_image,
                box_pred_per_image) in enumerate(zip(scores, box_pred)):

            scores_per_image, topk_indices = scores_per_image.flatten(
                0, 1).topk(
                    self.num_proposals, sorted=False)
            labels_per_image = paddle.gather(labels, topk_indices, axis=0)

            box_pred_per_image = box_pred_per_image.reshape([-1, 1, 4]).tile(
                [1, self.num_classes, 1]).reshape([-1, 4])
            box_pred_per_image = paddle.gather(
                box_pred_per_image, topk_indices, axis=0)

            classes_all.append(labels_per_image)
            scores_all.append(scores_per_image)
            boxes_all.append(box_pred_per_image)

        bbox_num = paddle.zeros([len(scale_factor_wh)], dtype="int32")
        boxes_final = []

        for i in range(len(scale_factor_wh)):
            classes = classes_all[i]
            boxes = boxes_all[i]
            scores = scores_all[i]

            boxes[:, 0::2] = paddle.clip(
                boxes[:, 0::2], min=0, max=img_wh[i][0]) / scale_factor_wh[i][0]
            boxes[:, 1::2] = paddle.clip(
                boxes[:, 1::2], min=0, max=img_wh[i][1]) / scale_factor_wh[i][1]
            boxes_w, boxes_h = (boxes[:, 2] - boxes[:, 0]).numpy(), (
                boxes[:, 3] - boxes[:, 1]).numpy()

            keep = (boxes_w > 1.) & (boxes_h > 1.)

            if (keep.sum() == 0):
                bboxes = paddle.zeros([1, 6]).astype("float32")
            else:
                boxes = paddle.to_tensor(boxes.numpy()[keep]).astype("float32")
                classes = paddle.to_tensor(classes.numpy()[keep]).astype(
                    "float32").unsqueeze(-1)
                scores = paddle.to_tensor(scores.numpy()[keep]).astype(
                    "float32").unsqueeze(-1)

                bboxes = paddle.concat([classes, scores, boxes], axis=-1)

            boxes_final.append(bboxes)
            bbox_num[i] = bboxes.shape[0]

        bbox_pred = paddle.concat(boxes_final)
        return bbox_pred, bbox_num


from data import DatasetFromObj
from torch.utils.data import DataLoader, TensorDataset
from model import Zi2ZiModel
import os
import argparse
import torch
import random
import time
import math
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import torchvision.transforms as transforms
from torchvision.utils import save_image, make_grid
import time
from model.model import chk_mkdir

writer_dict = {
        '智永': 0, ' 隸書-趙之謙': 1, '張即之': 2, '張猛龍碑': 3, '柳公權': 4, '標楷體-手寫': 5, '歐陽詢-九成宮': 6,
        '歐陽詢-皇甫誕': 7, '沈尹默': 8, '美工-崩雲體': 9, '美工-瘦顏體': 10, '虞世南': 11, '行書-傅山': 12, '行書-王壯為': 13,
        '行書-王鐸': 14, '行書-米芾': 15, '行書-趙孟頫': 16, '行書-鄭板橋': 17, '行書-集字聖教序': 18, '褚遂良': 19, '趙之謙': 20,
        '趙孟頫三門記體': 21, '隸書-伊秉綬': 22, '隸書-何紹基': 23, '隸書-鄧石如': 24, '隸書-金農': 25,  '顏真卿-顏勤禮碑': 26,
        '顏真卿多寶塔體': 27, '魏碑': 28
    }

parser = argparse.ArgumentParser(description='Infer')
parser.add_argument('--experiment_dir', required=True,
                    help='experiment directory, data, samples,checkpoints,etc')
parser.add_argument('--start_from', type=int, default=0)
parser.add_argument('--gpu_ids', default=[], nargs='+', help="GPUs")
parser.add_argument('--image_size', type=int, default=256,
                    help="size of your input and output image")
parser.add_argument('--L1_penalty', type=int, default=100, help='weight for L1 loss')
parser.add_argument('--Lconst_penalty', type=int, default=15, help='weight for const loss')
# parser.add_argument('--Ltv_penalty', dest='Ltv_penalty', type=float, default=0.0, help='weight for tv loss')
parser.add_argument('--Lcategory_penalty', type=float, default=1.0,
                    help='weight for category loss')
parser.add_argument('--embedding_num', type=int, default=40,
                    help="number for distinct embeddings")
parser.add_argument('--embedding_dim', type=int, default=128, help="dimension for embedding")
parser.add_argument('--batch_size', type=int, default=16, help='number of examples in batch')
parser.add_argument('--lr', type=float, default=0.001, help='initial learning rate for adam')
parser.add_argument('--resume', type=int, default=None, help='resume from previous training')
parser.add_argument('--obj_path', type=str, default='./experiment/data/val.obj', help='the obj file you infer')
parser.add_argument('--input_nc', type=int, default=1)

parser.add_argument('--from_txt', action='store_true')
parser.add_argument('--src_txt', type=str, default='大威天龍大羅法咒世尊地藏波若諸佛')
parser.add_argument('--canvas_size', type=int, default=256)
parser.add_argument('--char_size', type=int, default=256)
parser.add_argument('--run_all_label', action='store_true')
parser.add_argument('--label', type=int, default=0)
parser.add_argument('--src_font', type=str, default='charset/gbk/方正新楷体_GBK(完整).TTF')

def draw_single_char(ch, font, canvas_size):
    img = Image.new("RGB", (canvas_size, canvas_size), (255, 255, 255))
    draw = ImageDraw.Draw(img)
    draw.text((0, 0), ch, (0, 0, 0), font=font)
    img = img.convert('L')
    return img

def main():
    args = parser.parse_args()
    data_dir = os.path.join(args.experiment_dir, "data")
    checkpoint_dir = os.path.join(args.experiment_dir, "checkpoint")
    sample_dir = os.path.join(args.experiment_dir, "sample")
    infer_dir = os.path.join(args.experiment_dir, "infer")
    chk_mkdir(infer_dir)

    # train_dataset = DatasetFromObj(os.path.join(data_dir, 'train.obj'), augment=True, bold=True, rotate=True, blur=True)
    # val_dataset = DatasetFromObj(os.path.join(data_dir, 'val.obj'))
    # dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)

    t0 = time.time()

    model = Zi2ZiModel(
        input_nc=args.input_nc,
        embedding_num=args.embedding_num,
        embedding_dim=args.embedding_dim,
        Lconst_penalty=args.Lconst_penalty,
        Lcategory_penalty=args.Lcategory_penalty,
        save_dir=checkpoint_dir,
        gpu_ids=args.gpu_ids,
        is_training=False
    )
    model.setup()
    model.print_networks(True)
    model.load_networks(args.resume)

    t1 = time.time()

    if args.from_txt:
        src = args.src_txt
        font = ImageFont.truetype(args.src_font, size=args.char_size)
        img_list = [transforms.Normalize(0.5, 0.5)(
            transforms.ToTensor()(
                draw_single_char(ch, font, args.canvas_size)
            )
        ).unsqueeze(dim=0) for ch in src]
        label_list = [args.label for _ in src]

        img_list = torch.cat(img_list, dim=0)
        label_list = torch.tensor(label_list)

        dataset = TensorDataset(label_list, img_list, img_list)
        dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)

    else:
        val_dataset = DatasetFromObj(os.path.join(data_dir, 'val.obj'),
                                     input_nc=args.input_nc,
                                     start_from=args.start_from)
        dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False)

    global_steps = 0

    for batch in dataloader:
        if args.run_all_label:
            global writer_dict
            writer_dict_inv = {v: k for k, v in writer_dict.items()}
            for label_idx in range(29):
                model.set_input(torch.ones_like(batch[0]) * label_idx, batch[2], batch[1])
                model.forward()
                tensor_to_plot = torch.cat([model.fake_B, model.real_B], 3)
                # img = vutils.make_grid(tensor_to_plot)
                save_image(tensor_to_plot, os.path.join(infer_dir, "infer_{}".format(writer_dict_inv[label_idx]) + "_construct.png"))
        else:
            # model.set_input(batch[0], batch[2], batch[1])
            # model.optimize_parameters()
            model.sample(batch, infer_dir)
            global_steps += 1

    t_finish = time.time()

    print('cold start time: %.2f, hot start time %.2f' % (t_finish - t0, t_finish - t1))

if __name__ == '__main__':
    with torch.no_grad():
        main()