#!/usr/bin/env python
#Importar paquetes necesarios
import rospy
import actionlib
#Importar elementos para uso de servidor y cliente
from BRT_ACTION.msg import BRTAction, BRTGoal, BRTResult
#Importar tipo de dato usado para mensajes
from std_msgs.msg import Float64
from std_msgs.msg import Float64MultiArray
#Importar algoritmos usados para analisis cinematico
from BraccioDEV import *
import numpy as np
#Importar marcadores AR TAG
from ar_track_alvar_msgs.msg import AlvarMarkers
import time

#Definir variables
posx=[]
posy=[]
posz=[]
posx1=[]
posy1=[]
posz1=[]
#Variar dependiendo la posicion de la camara respecto al robot y Medida de elemento a manipular
camx=-0.02
camy=0.525
camz=0.3
caja=0.05
#Iniciacion de variables
k=0
N=[]
NL=[]
xd=[0,0,0]

#reconocimiento de marcadores
def callback(markers):
        rate= rospy.Rate(1)
	#Declaracion de variables globales
        global posx,camx
        global posy,camy,k
        global posz,camz,caja,xd,NL
        #Reconocimiento y guardado de posiciones en las que se encuentra AR TAG respecto a la camara
        for m in markers.markers:
                    N.append(m.id)
             	    #Identificar TAG
		    NS=set(N)
		    NL=list(NS)
                    marker_pose = m.pose.pose
                    pos = marker_pose.position
                    ori = marker_pose.orientation
		    #Identificar y guardar posicion en una lista
                    posx1.append(pos.x)
	            posx2=set(posx1)
		    posx=list(posx2)
                    posy1.append(pos.y)
                    posy2=set(posy1)
		    posy=list(posy2)
                    posz1.append(pos.z)
	            posz2=set(posz1)
		    posz=list(posz2)
#Main principal
if __name__=="__main__":
    #Inicializar cliente
    rospy.init_node('BRT_action_client')
    #Suscripcion a topico para reconocer AR TAG
    sub = rospy.Subscriber("/ar_pose_marker", AlvarMarkers, callback)
    time.sleep(3)
    print("Cliente activado")
    #La cantidad de posiciones guardadas indican a cantidad de TAG detectados
    tag=len(NL)
    print("TAGs Detectados:")
    print(tag)
    i=0
    xr=[0,0,0]
    #Repeticion de solicitud de lugar detectado dependiendo de cantidad de tag
    for i in range(tag):
	sub = rospy.Subscriber("/ar_pose_marker", AlvarMarkers, callback)
    	time.sleep(3)
	#Inicializacion de pedido y conexion del cliente con server
        client = actionlib.SimpleActionClient('BRT_ACTION', BRTAction)
        client.wait_for_server()
	print("ID Marcador")
	print NL[i]
	#Calculo dimensional para adquirir posiciones de los TAG desde el robot.
        fx=camx-posx[i]
	fz=camz-posy[i]+caja
	fy=-(camy-posz[i])

	fx=np.round(fx,3)
	fy=np.round(fy,3)
	fz=np.round(fz,3)

	#Limitacion de altura
	if fz<0.10:
	       fz=0.1025
	if fz>0.13
	       fz=0.1025

        xd=[fy,fx,fz]
        xm=[posx,posy,posz]
        xcam=[-camy,camx,camz+caja]

        if np.array_equal(xd,xcam) == False:
	       print("AR Detectado")
        else:
               print("Esperando lectura")
        #Variable de solicitud de datos
        goal = BRTGoal()
        #Si el resultado de posicion deseada es distinto que la posicion 0
        if np.array_equal(xd,xr) == False:
		    time.sleep(5)
		    #Se solicita posicion xd
                    goal.xd=xd
                    print("POSICION SOLICITADA:")
                    print(goal.xd)
		    #Envio de solicitud y espera de resultado
                    client.send_goal(goal)
                    client.wait_for_result()
                    #resultados despues de analisis y movimiento
                    MOV=client.get_result().MOVE
                    Q=client.get_result().q
                    MOV=np.round(MOV,3)
                    Q=np.round(Q,3)
                    print("POSICION FINAL:")
                    print(MOV)
                    print("ANGULOS DESPLAZADOS:")
                    print(Q)


import ctypes
import os
import shutil
import random
import sys
import threading
import time
import cv2
import numpy as np
import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt
import torch
import torchvision
import tqdm
import logging

CONF_THRESH = 0.7
IOU_THRESHOLD = 0.4

PLUGIN_LIBRARY = "build/libmyplugins.so"

ctypes.CDLL(PLUGIN_LIBRARY)

class YoLov5TRT(object):
    """
    description: A YOLOv5 class that warps TensorRT ops, preprocess and postprocess ops.
    """

    def __init__(self, checkpoint_path: str, device_num: int):
        self.ctx = cuda.Device(device_num).make_context()
        stream = cuda.Stream()
        TRT_LOGGER = trt.Logger(trt.Logger.INFO)
        runtime = trt.Runtime(TRT_LOGGER)
        self.ENGINE_PATH = checkpoint_path

        with open(self.ENGINE_PATH, "rb") as f:
            engine = runtime.deserialize_cuda_engine(f.read())
        context = engine.create_execution_context()

        host_inputs = []
        cuda_inputs = []
        host_outputs = []
        cuda_outputs = []
        bindings = []

        for binding in engine:
            size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
            dtype = trt.nptype(engine.get_binding_dtype(binding))

            host_mem = cuda.pagelocked_empty(size, dtype)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)

            bindings.append(int(cuda_mem))

            if engine.binding_is_input(binding):
                self.input_w = engine.get_binding_shape(binding)[-1]
                self.input_h = engine.get_binding_shape(binding)[-2]
                host_inputs.append(host_mem)
                cuda_inputs.append(cuda_mem)
            else:
                host_outputs.append(host_mem)
                cuda_outputs.append(cuda_mem)

        self.stream = stream
        self.context = context
        self.engine = engine
        self.host_inputs = host_inputs
        self.cuda_inputs = cuda_inputs
        self.host_outputs = host_outputs
        self.cuda_outputs = cuda_outputs
        self.bindings = bindings
        self.batch_size = engine.max_batch_size
        self.warmup()

    def warmup(self):
        """
        description: Warmup function for TensorRT.
        """
        logging.info('Starting to warm up...')
        for _ in tqdm.tqdm(range(20)):
            standard = np.random.randn(800, 600, 3)
            self.infer((standard + 255).astype(np.uint8))
        logging.info('Warmup has been finished')

    def infer(self, raw_image: np.ndarray):
        """
        description: Entire inference function. Takes raw BGR image (np.ndarray), returns detections.

        param:
            raw_image: np.ndarray (Frame/Image)
        return:
            detection_results: list[list] -> [[x0, y0, x1, y1, confidence], ...]
        """

        self.ctx.push()

        stream = self.stream
        context = self.context
        engine = self.engine
        host_inputs = self.host_inputs
        cuda_inputs = self.cuda_inputs
        host_outputs = self.host_outputs
        cuda_outputs = self.cuda_outputs
        bindings = self.bindings

        batch_image_raw = []
        batch_origin_h = []
        batch_origin_w = []
        batch_input_image = np.empty(shape=[self.batch_size, 3, self.input_h, self.input_w])
        input_image, image_raw, origin_h, origin_w = self.preprocess_image(raw_image)
        batch_image_raw.append(image_raw)
        batch_origin_h.append(origin_h)
        batch_origin_w.append(origin_w)
        np.copyto(batch_input_image[0], input_image)
        batch_input_image = np.ascontiguousarray(batch_input_image)

        np.copyto(host_inputs[0], batch_input_image.ravel())
        cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
        context.execute_async(batch_size=self.batch_size, bindings=bindings, stream_handle=stream.handle)
        cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
        stream.synchronize()
        self.ctx.pop()
        output = host_outputs[0]
        for i in range(1):
            result_boxes, result_scores, result_classid = self.postprocess(
                output[i * 6001: (i + 1) * 6001], batch_origin_h[i], batch_origin_w[i]
            )
        detection_results = []

        for x in range(len(result_boxes)):
            classid = result_classid[x]
            box = result_boxes[x].tolist()
            score = result_scores[x]
            box.append(score)
            detection_results.append(box)

        return detection_results

    def destroy(self):
        self.ctx.pop()

    def preprocess_image(self, raw_bgr_image):
        """
        Image Preprocessor: Letterbox, BGR2RGB, NCHW
        """
        image_raw = raw_bgr_image
        h, w, c = image_raw.shape
        image = cv2.cvtColor(image_raw, cv2.COLOR_BGR2RGB)

        r_w = self.input_w / w
        r_h = self.input_h / h
        if r_h > r_w:
            tw = self.input_w
            th = int(r_w * h)
            tx1 = tx2 = 0
            ty1 = int((self.input_h - th) / 2)
            ty2 = self.input_h - th - ty1
        else:
            tw = int(r_h * w)
            th = self.input_h
            tx1 = int((self.input_w - tw) / 2)
            tx2 = self.input_w - tw - tx1
            ty1 = ty2 = 0

        image = cv2.resize(image, (tw, th))
        image = cv2.copyMakeBorder(
            image, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, (128, 128, 128)
        )
        image = image.astype(np.float32)
        image /= 255.0
        image = np.transpose(image, [2, 0, 1])
        image = np.expand_dims(image, axis=0)
        image = np.ascontiguousarray(image)
        return image, image_raw, h, w

    def xywh2xyxy(self, origin_h, origin_w, x):
        y = torch.zeros_like(x) if isinstance(x, torch.Tensor) else np.zeros_like(x)
        r_w = self.input_w / origin_w
        r_h = self.input_h / origin_h
        if r_h > r_w:
            y[:, 0] = x[:, 0] - x[:, 2] / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2 - (self.input_h - r_w * origin_h) / 2
            y /= r_w
        else:
            y[:, 0] = x[:, 0] - x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 2] = x[:, 0] + x[:, 2] / 2 - (self.input_w - r_h * origin_w) / 2
            y[:, 1] = x[:, 1] - x[:, 3] / 2
            y[:, 3] = x[:, 1] + x[:, 3] / 2
            y /= r_h

        return y

    def postprocess(self, output, origin_h, origin_w):
        num = int(output[0])
        pred = np.reshape(output[1:], (-1, 6))[:num, :]
        pred = torch.Tensor(pred).cuda()
        boxes = pred[:, :4]
        scores = pred[:, 4]
        classid = pred[:, 5]

        chunk = scores > CONF_THRESH
        boxes = boxes[chunk, :]
        scores = scores[chunk]
        classid = classid[chunk]

        boxes = self.xywh2xyxy(origin_h, origin_w, boxes)

        indices = torchvision.ops.nms(boxes, scores, iou_threshold=IOU_THRESHOLD).cpu()
        result_boxes = boxes[indices, :].cpu().numpy().astype(np.int)
        result_scores = scores[indices].cpu().numpy()
        result_classid = classid[indices].cpu().numpy()
        return result_boxes, result_scores, result_classid


"""
Retrain the YOLO model for your own dataset.
"""

import numpy as np
import keras.backend as K
import tensorflow as tf
from keras.layers import Input, Lambda
from keras.models import Model
from keras.optimizers import Adam
from keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping

from yolo3.model import preprocess_true_boxes, yolo_body, tiny_yolo_body, yolo_loss
from yolo3.utils import get_random_data

# utils.py中的get_random_data被修改成读取单通道图片
def _main():
    gpu_options = tf.GPUOptions(allow_growth=True)
    sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
    #读取数据
    annotation_path = '2007_train.txt'
    log_dir = 'logs/033/'
    classes_path = 'model_data/voc_classes.txt'
    anchors_path = 'model_data/yolo_anchors.txt'
    class_names = get_classes(classes_path)
    num_classes = len(class_names)
    anchors = get_anchors(anchors_path)
    #输入图像的高和宽
    input_shape = (416,416) # multiple of 32, hw
    #判断是否是is_tiny_version
    is_tiny_version = len(anchors)==6 # default setting
    if is_tiny_version:
        model = create_tiny_model(input_shape, anchors, num_classes,
            freeze_body=2)
    else:
        model = create_model(input_shape, anchors, num_classes,
            freeze_body=2) # make sure you know what you freeze
    # 参数可视化工具，用于后期画图用的
    logging = TensorBoard(log_dir=log_dir)
    # filename：字符串，保存模型的路径
    # monitor：需要监视的值
    # verbose：信息展示模式，0或1
    # save_best_only：当设置为True时，将只保存在验证集上性能最好的模型
    # mode：‘auto’，‘min’，‘max’之一，在save_best_only=True时决定性能最佳模型的评判准则，
    # 例如，当监测值为val_acc时，模式应为max，
    # 当检测值为val_loss时，模式应为min。在auto模式下，评价准则由被监测值的名字自动推断。
    # save_weights_only：若设置为True，则只保存模型权重，否则将保存整个模型（包括模型结构，配置信息等）
    # period：CheckPoint之间的间隔的epoch数
    checkpoint = ModelCheckpoint(log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5',
        monitor='val_loss', save_weights_only=True, save_best_only=True, period=3)

    # optimer指的是网络的优化器
    # mode (str) ，可选择‘min’或者‘max’，min表示当监控量停止下降的时候，学习率将减小，max表示当监控量停止上升的时候，学习率将减小。默认值为‘min’
    # factor 学习率每次降低多少，new_lr = old_lr * factor
    # patience=10，容忍网路的性能不提升的次数，高于这个次数就降低学习率
    # verbose（bool） - 如果为True，则为每次更新向stdout输出一条消息。 默认值：False
    # threshold（float） - 测量新最佳值的阈值，仅关注重大变化。 默认值：1e-4
    # cooldown： 减少lr后恢复正常操作之前要等待的时期数。 默认值：0。
    # min_lr,学习率的下限
    # eps ，适用于lr的最小衰减。 如果新旧lr之间的差异小于eps，则忽略更新。 默认值：1e-8。
    reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, verbose=1)

    # monitor: 监控的数据接口，有’acc’,’val_acc’,’loss’,’val_loss’等等。正常情况下如果有验证集，就用’val_acc’或者’val_loss’。但是因为笔者用的是5折交叉验证，没有单设验证集，所以只能用’acc’了。
    # min_delta：增大或减小的阈值，只有大于这个部分才算作improvement。这个值的大小取决于monitor，也反映了你的容忍程度。例如笔者的monitor是’acc’，同时其变化范围在70%-90%之间，所以对于小于0.01%的变化不关心。加上观察到训练过程中存在抖动的情况（即先下降后上升），所以适当增大容忍程度，最终设为0.003%。
    # patience：能够容忍多少个epoch内都没有improvement。这个设置其实是在抖动和真正的准确率下降之间做tradeoff。如果patience设的大，那么最终得到的准确率要略低于模型可以达到的最高准确率。如果patience设的小，那么模型很可能在前期抖动，还在全图搜索的阶段就停止了，准确率一般很差。patience的大小和learning rate直接相关。在learning rate设定的情况下，前期先训练几次观察抖动的epoch number，比其稍大些设置patience。在learning rate变化的情况下，建议要略小于最大的抖动epoch number。笔者在引入EarlyStopping之前就已经得到可以接受的结果了，EarlyStopping算是锦上添花，所以patience设的比较高，设为抖动epoch number的最大值。
    # verbose: 信息展示模式
    # mode: 就’auto’, ‘min’, ‘,max’三个可能。如果知道是要上升还是下降，建议设置一下。笔者的monitor是’acc’，所以mode=’max’。
    early_stopping = EarlyStopping(monitor='val_loss', min_delta=0, patience=50, verbose=1)
    # 训练和测试的数据按比例分开
    val_split = 0.1
    with open(annotation_path) as f:
        lines = f.readlines()
    #seed(10101)中的数值相同每次产生的随机数相同，每一次产生随机数之前都要seed(10101)一次
    #seed()中空每次随机数不同
    np.random.seed(10101)
    #打乱行的顺序
    np.random.shuffle(lines)
    np.random.seed(None)
    #测试数量
    num_val = int(len(lines)*val_split)
    #训练数量
    num_train = len(lines) - num_val

    # Train with frozen layers first, to get a stable loss.
    # Adjust num epochs to your dataset. This step is enough to obtain a not bad model.
    if True:
        model.compile(optimizer=Adam(lr=1e-3), loss={
            # use custom yolo_loss Lambda layer.
            'yolo_loss': lambda y_true, y_pred: y_pred})

        batch_size = 1
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))
        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),
                steps_per_epoch=max(1, num_train//batch_size),
                validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),
                validation_steps=max(1, num_val//batch_size),
                epochs=10,
                initial_epoch=0,
                callbacks=[logging, checkpoint])
        model.save(log_dir + 'trained_weights_stage_1.h5')
    print("trained_weights_stage_1 is over")

    # Unfreeze and continue training, to fine-tune.
    # Train longer if the result is not good.
    if True:
        for i in range(len(model.layers)):
            model.layers[i].trainable = True

        # 用于配置训练模型
        # optimizer: 字符串（优化器名）或者优化器实例。 详见 optimizers。
        # loss: 字符串（目标函数名）或目标函数。 详见 losses。 如果模型具有多个输出，则可以通过传递损失函数的字典或列表，在每个输出上使用不同的损失。 模型将最小化的损失值将是所有单个损失的总和。
        # metrics: 在训练和测试期间的模型评估标准。 通常你会使用 metrics = ['accuracy']。 要为多输出模型的不同输出指定不同的评估标准， 还可以传递一个字典，如 metrics = {'output_a'：'accuracy'}。
        # loss_weights: 可选的指定标量系数（Python 浮点数）的列表或字典， 用以衡量损失函数对不同的模型输出的贡献。 模型将最小化的误差值是由 loss_weights 系数加权的加权总和误差。 如果是列表，那么它应该是与模型输出相对应的 1:1 映射。 如果是张量，那么应该把输出的名称（字符串）映到标量系数。
        # sample_weight_mode: 如果你需要执行按时间步采样权重（2D 权重），请将其设置为 temporal。 默认为 None，为采样权重（1D）。 如果模型有多个输出，则可以通过传递 mode 的字典或列表，以在每个输出上使用不同的 sample_weight_mode。
        # weighted_metrics: 在训练和测试期间，由 sample_weight 或 class_weight 评估和加权的度量标准列表。
        # target_tensors: 默认情况下，Keras 将为模型的目标创建一个占位符，在训练过程中将使用目标数据。 相反，如果你想使用自己的目标张量（反过来说，Keras 在训练期间不会载入这些目标张量的外部 Numpy 数据）， 您可以通过 target_tensors 参数指定它们。 它可以是单个张量（单输出模型），张量列表，或一个映射输出名称到目标张量的字典。
        # **kwargs: 当使用 Theano/CNTK 后端时，这些参数被传入 K.function。 当使用 TensorFlow 后端时，这些参数被传递到 tf.Session.run。
        model.compile(optimizer=Adam(lr=1e-4), loss={'yolo_loss': lambda y_true, y_pred: y_pred}) # recompile to apply the change
        print('Unfreeze all of the layers.')

        batch_size = 1 # note that more GPU memory is required after unfreezing the body
        print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size))

        # 使用 Python 生成器（或 Sequence 实例）逐批生成的数据，按批次训练模型。
        # 生成器与模型并行运行，以提高效率。 例如，这可以让你在 CPU 上对图像进行实时数据增强，以在 GPU 上训练模型。
        # keras.utils.Sequence 的使用可以保证数据的顺序， 以及当 use_multiprocessing=True 时 ，保证每个输入在每个 epoch 只使用一次。
        # generator: 一个生成器，或者一个 Sequence (keras.utils.Sequence) 对象的实例， 以在使用多进程时避免数据的重复。 生成器的输出应该为以下之一：
        # 一个 (inputs, targets) 元组
        # 一个 (inputs, targets, sample_weights) 元组。
        # 这个元组（生成器的单个输出）组成了单个的 batch。 因此，这个元组中的所有数组长度必须相同（与这一个 batch 的大小相等）。 不同的 batch 可能大小不同。 例如，一个 epoch 的最后一个 batch 往往比其他 batch 要小， 如果数据集的尺寸不能被 batch size 整除。 生成器将无限地在数据集上循环。当运行到第 steps_per_epoch 时，记一个 epoch 结束。
        # steps_per_epoch: 在声明一个 epoch 完成并开始下一个 epoch 之前从 generator 产生的总步数（批次样本）。 它通常应该等于你的数据集的样本数量除以批量大小。 对于 Sequence，它是可选的：如果未指定，将使用len(generator) 作为步数。
        # epochs: 整数。训练模型的迭代总轮数。一个 epoch 是对所提供的整个数据的一轮迭代，如 steps_per_epoch 所定义。注意，与 initial_epoch 一起使用，epoch 应被理解为「最后一轮」。模型没有经历由 epochs 给出的多次迭代的训练，而仅仅是直到达到索引 epoch 的轮次。
        # verbose: 0, 1 或 2。日志显示模式。 0 = 安静模式, 1 = 进度条, 2 = 每轮一行。
        # callbacks: keras.callbacks.Callback 实例的列表。在训练时调用的一系列回调函数。
        # validation_data: 它可以是以下之一：
        # 验证数据的生成器或 Sequence 实例
        # 一个 (inputs, targets) 元组
        # 一个 (inputs, targets, sample_weights) 元组。
        # 在每个 epoch 结束时评估损失和任何模型指标。该模型不会对此数据进行训练。
        # validation_steps: 仅当 validation_data 是一个生成器时才可用。 在停止前 generator 生成的总步数（样本批数）。 对于 Sequence，它是可选的：如果未指定，将使用 len(generator) 作为步数。
        # class_weight: 可选的将类索引（整数）映射到权重（浮点）值的字典，用于加权损失函数（仅在训练期间）。 这可以用来告诉模型「更多地关注」来自代表性不足的类的样本。
        # max_queue_size: 整数。生成器队列的最大尺寸。 如未指定，max_queue_size 将默认为 10。
        # workers: 整数。使用的最大进程数量，如果使用基于进程的多线程。 如未指定，workers 将默认为 1。如果为 0，将在主线程上执行生成器。
        # use_multiprocessing: 布尔值。如果 True，则使用基于进程的多线程。 如未指定， use_multiprocessing 将默认为 False。 请注意，由于此实现依赖于多进程，所以不应将不可传递的参数传递给生成器，因为它们不能被轻易地传递给子进程。
        # shuffle: 是否在每轮迭代之前打乱 batch 的顺序。 只能与 Sequence (keras.utils.Sequence) 实例同用。
        # initial_epoch: 开始训练的轮次（有助于恢复之前的训练）。
        model.fit_generator(data_generator_wrapper(lines[:num_train], batch_size, input_shape, anchors, num_classes),#训练数据
            steps_per_epoch=max(1, num_train//batch_size),#
            validation_data=data_generator_wrapper(lines[num_train:], batch_size, input_shape, anchors, num_classes),#验证数据
            validation_steps=max(1, num_val//batch_size),
            epochs=20,
            initial_epoch=10,#从哪个epoch开始
            callbacks=[logging, checkpoint, reduce_lr, early_stopping])#回调函数
        model.save(log_dir + 'trained_weights_final.h5')
    print("trained_weights_final is over")
    # Further training if needed.

def get_classes(classes_path):
    '''loads the classes'''
    with open(classes_path) as f:
        class_names = f.readlines()
    #strip() 移除字符串头尾指定的字符序列。
    class_names = [c.strip() for c in class_names]
    return class_names

def get_anchors(anchors_path):
    '''loads the anchors from a file'''
    with open(anchors_path) as f:
        anchors = f.readline()
    anchors = [float(x) for x in anchors.split(',')]
    return np.array(anchors).reshape(-1, 2)
#变成下面这种形式 9X2
# [[ 10.  13.]
#  [ 16.  30.]
#  [ 33.  23.]
#  [ 30.  61.]
#  [ 62.  45.]
#  [ 59. 119.]
#  [116.  90.]
#  [156. 198.]
#  [373. 326.]]

def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
            weights_path='model_data/yolo_weights.h5'):
    '''create the training model'''
    K.clear_session() # get a new session
    image_input = Input(shape=(None, None, 1))
    h, w = input_shape
    num_anchors = len(anchors)

    y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \
        num_anchors//3, num_classes+5)) for l in range(3)]

    # 经过网络后输出y1、y2、y3 并且是以Model的形式进行返回的
    model_body = yolo_body(image_input, num_anchors//3, num_classes)
    print("model_body =",model_body.output[2][...])
    print("y_true =",y_true[2])
    print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    if load_pretrained:
        # 使用keras自带的模型加载
        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
        print('Load weights {}.'.format(weights_path))
        if freeze_body in [1, 2]:
            # Freeze darknet53 body or freeze all but 3 output layers.
            num = (185, len(model_body.layers)-3)[freeze_body-1]
            for i in range(num): model_body.layers[i].trainable = False
            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))

    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
        [*model_body.output, *y_true])
    model = Model([model_body.input, *y_true], model_loss)

    return model

def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=False, freeze_body=2,
            weights_path='model_data/tiny_yolo_weights.h5'):
    '''create the training model, for Tiny YOLOv3'''
    K.clear_session() # get a new session
    image_input = Input(shape=(None, None, 1))
    h, w = input_shape # (416,416)
    num_anchors = len(anchors)

    # Input():用来实例化一个keras张量
    # keras张量是来自底层后端（Theano或Tensorflow）的张量对象，我们增加了某些属性，使我们通过知道模型的输入和输出来构建keras模型。
    # 添加的keras属性有：1)._keras_shape:整型的形状元组通过keras-side 形状推理传播  2)._keras_history: 最后一层应用于张量，整个图层的图可以从那个层，递归地检索出来。
    # #参数：
    # shape: 形状元组（整型），不包括batch size。for instance, shape=(32,) 表示了预期的输入将是一批32维的向量。
    # batch_shape: 形状元组（整型），包括了batch size。for instance, batch_shape=(10,32)表示了预期的输入将是10个32维向量的批次。
    # name: 对于该层是可选的名字字符串。在一个模型中是独一无二的（同一个名字不能复用2次）。如果name没有被特指将会自动生成。
    # dtype: 预期的输入数据类型
    # sparse: 特定的布尔值，占位符是否为sparse
    # tensor: 可选的存在的向量包装到Input层，如果设置了，该层将不会创建一个占位张量。
    # #返回
    # 一个张量
    y_true = [Input(shape=(h//{0:32, 1:16}[l], w//{0:32, 1:16}[l], \
        num_anchors//2, num_classes+5)) for l in range(2)]

    model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)
    print("model_body =",model_body.output[1][...])
    print("y_true =",y_true[1])
    print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))

    if load_pretrained:
        model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
        print('Load weights {}.'.format(weights_path))
        if freeze_body in [1, 2]:
            # Freeze the darknet body or freeze all but 2 output layers.
            num = (20, len(model_body.layers)-2)[freeze_body-1]
            for i in range(num): model_body.layers[i].trainable = False
            print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
    #将任意表达式封装为 Layer 对象。
    # function: 需要封装的函数。 将输入张量作为第一个参数。
    # output_shape: 预期的函数输出尺寸。 只在使用 Theano 时有意义。 可以是元组或者函数。 如果是元组，它只指定第一个维度； 样本维度假设与输入相同： output_shape = (input_shape[0], ) + output_shape 或者，输入是 None 且样本维度也是 None： output_shape = (None, ) + output_shape 如果是函数，它指定整个尺寸为输入尺寸的一个函数： output_shape = f(input_shape)
    # arguments: 可选的需要传递给函数的关键字参数。
    # 输入尺寸
    # 任意。当使用此层作为模型中的第一层时， 使用参数 input_shape （整数元组，不包括样本数的轴）。
    # 输出尺寸
    # 由 output_shape 参数指定 (或者在使用 TensorFlow 时，自动推理得到)
    model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
        arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.7})(
        [*model_body.output, *y_true])
    model = Model([model_body.input, *y_true], model_loss)

    return model

def data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes):
    '''data generator for fit_generator'''
    n = len(annotation_lines)
    i = 0
    while True:
        image_data = []
        box_data = []
        for b in range(batch_size):
            if i==0:
                np.random.shuffle(annotation_lines)
            image, box = get_random_data(annotation_lines[i], input_shape, random=False)
            image_data.append(image)
            box_data.append(box)
            i = (i+1) % n
        image_data = np.array(image_data)
        box_data = np.array(box_data)
        y_true = preprocess_true_boxes(box_data, input_shape, anchors, num_classes)
        yield [image_data, *y_true], np.zeros(batch_size)

def data_generator_wrapper(annotation_lines, batch_size, input_shape, anchors, num_classes):
    n = len(annotation_lines)
    if n==0 or batch_size<=0: return None
    return data_generator(annotation_lines, batch_size, input_shape, anchors, num_classes)

if __name__ == '__main__':
    _main()


import gym
import ray
from ray.rllib.agents import ppo
import random
import numpy as np
from models.util import write_to
from gym.spaces import Discrete, Box, Tuple
from models.station import Station
from models.util import path, sampler, calcReward

'''
## To-DO

**Visualizer to be made - vizualization

**RL env to parallelize

##acceleration is has limits not speed... change the model accordingly
action should be acceleration and speed should be kept track off

'''

PATHS = [[(1,20),(2,20),(3,40), (4,10), (0,10)] for i in range(5)] # paths[i] -> path_of_train[i]
# paths[i][j=current station] = next station, distance

STATIONS = [Station(i) for i in range(5)]

class Controller(gym.Env):
    """
    Action space -> variables that can be controlled and here it is only speed
    Observation space -> (station_to, train_id, num_people_at station, occupancy, distance)
    """

    def __init__(self, train_id=0, num_stations=4,
                    max_acceleration=100, max_speed=100,
                    max_people=500, max_distance=100, initial_state= [1,50,0,0]):
        self.MAX_SPEED = max_speed
        self.MAX_ACC = max_acceleration
        self.MAX_PEOPLE = max_people
        self.MAX_DISTANCE = max_distance
        self.MAX_OCCUPANCY = 1
        self.TRAIN_ID = train_id
        self.DELTA_TIME = 0.1 #10ms
        self.INITIAL_STATE = initial_state
        self.action_space = Box(0,self.MAX_SPEED, shape=(1,))
        self.observation_space = Box(low=np.array([0, 0, 0, 0]),
                                    high=np.array([5, self.MAX_PEOPLE, self.MAX_OCCUPANCY, self.MAX_DISTANCE]),
                                    dtype=np.float64)

        self.curr_state = np.array(self.INITIAL_STATE, dtype=np.float64)
        self.curr_station = 1
        self.prev_action = 0

    def reset(self):
        self.curr_station = 1
        self.curr_state = np.array(self.INITIAL_STATE, dtype=np.float64) #stationid, num_people_at_station,train_occupancy,distance_from_next_station
        return self.curr_state

    def step(self, action):
        if abs(action - self.prev_action) > self.MAX_ACC * self.DELTA_TIME:
            if self.prev_action > action:
                action = self.prev_action - self.MAX_ACC * self.DELTA_TIME
            else:
                action = self.prev_action + self.MAX_ACC * self.DELTA_TIME

        self.prev_action = action
        state = list(self.curr_state)
        #print(f"{action=}")
        done = 0
        # action is current speed
        if action == 0:
            # distance from current station = state[-1]
            if state[-1] == 0:
                # currently at a station
                (next_station, distance), done = path(int(state[0]), self.TRAIN_ID)
                occupancy, next_max_people = sampler(int(state[0]), self.TRAIN_ID, state[1], state[2])
                next_state = [next_station, next_max_people, occupancy, distance]
            else:
                # Between two stations but no movement
                next_state = state
        else:
            if state[-1] == 0:
                #currently leaving the station
                (next_station, distance), done = path(int(state[0]), self.TRAIN_ID)
                occupancy, next_max_people = sampler(int(state[0]), self.TRAIN_ID, state[1], state[2])
                next_state = [next_station, next_max_people, occupancy, distance]
            else:
                #Between two stations and moving
                updated_distance = state[3]- self.DELTA_TIME*action

                if updated_distance > 0:
                    next_state = [state[0], state[1], state[2], updated_distance]
                else:
                    next_state = [state[0], state[1], state[2], 0]

        trains_approaching = STATIONS[next_state[0]].get_distance(self.TRAIN_ID, action, next_state[3])

        rew = calcReward(int(next_state[0]), self.TRAIN_ID, next_state[2], action,next_state[1])
        self.curr_state = tuple(next_state)
        # print(f"{rew=},{next_state=}")
        write_to(self.TRAIN_ID, rew, next_state)
        return np.array(self.curr_state, dtype=np.float64), rew, done, {}

if __name__ == "__main__":
    env = [Controller(train_id=i) for i in range(5)]

    for i in range(4320):
        print(i)
        for j in range(5):
            # print(f"----------------- train ID={j}")
            # env[j].reset()
            env[j].step(env[j].action_space.sample()) # take a random action
            # env[j].close()


import torch
import torch.nn as nn
import copy
import sys
from collections import OrderedDict
import numpy as np
import numba
from ...ops.dcn import DeformConv
from ...ops.iou3d_nms import iou3d_nms_cuda
from ...ops.center_ops import center_ops_cuda

from ...utils import loss_utils
from .target_assigner.mm_assigner import MMAssigner

class MMHead(nn.Module):
    def __init__(self, model_cfg, input_channels, num_class, class_names, grid_size, voxel_size, point_cloud_range,
                 predict_boxes_when_training=True):
        super().__init__()
        self.model_cfg = model_cfg
        self.post_cfg = model_cfg.TEST_CONFIG
        self.in_channels = input_channels
        self.grid_size = grid_size
        self.point_cloud_range = point_cloud_range
        self.predict_boxes_when_training = predict_boxes_when_training

        self.num_classes = [len(t["class_names"]) for t in model_cfg.TASKS]
        self.class_names = [t["class_names"] for t in model_cfg.TASKS]

        self.code_weights = model_cfg.LOSS_CONFIG.code_weights
        self.weight = model_cfg.LOSS_CONFIG.weight # weight between hm loss and loc loss

        self.dataset = model_cfg.DATASET
        self.box_n_dim = 9 if self.dataset == 'nuscenes' else 7

        self.encode_background_as_zeros = True
        self.use_sigmoid_score = True
        self.no_log = False
        self.use_direction_classifier = False
        self.bev_only = True if model_cfg.MODE == "bev" else False

        # a shared convolution
        share_conv_channel = model_cfg.PARAMETERS.share_conv_channel
        self.shared_conv = nn.Sequential(
            nn.Conv2d(self.in_channels, share_conv_channel,
                      kernel_size=3, padding=1, bias=True),
            nn.BatchNorm2d(share_conv_channel),
            nn.ReLU(inplace=True)
        )

        self.common_heads = model_cfg.PARAMETERS.common_heads
        self.init_bias = model_cfg.PARAMETERS.init_bias
        self.tasks = nn.ModuleList()

        self.use_dcn = model_cfg.USE_DCN
        for num_cls in self.num_classes:
            heads = copy.deepcopy(self.common_heads)
            if not self.use_dcn:
                heads.update(dict(hm=(num_cls, 2)))
                self.tasks.append(
                    SepHead(share_conv_channel, heads, bn=True, init_bias=self.init_bias, final_kernel=3, directional_classifier=False)
                )
            else:
                self.tasks.append(
                    DCNSepHead(share_conv_channel, num_cls, heads, bn=True, init_bias=self.init_bias, final_kernel=3,
                                directional_classifier=False)
                )
        self.target_assigner = MMAssigner(
            model_cfg.TARGET_ASSIGNER_CONFIG,
            num_classes = sum(self.num_classes),
            no_log = self.no_log,
            grid_size = grid_size,
            pc_range = point_cloud_range,
            voxel_size = voxel_size
        )

        self.forward_ret_dict = {}
        self.build_losses()

    def assign_targets(self, gt_boxes):
        """
        Args:
            gt_boxes: (B, M, 8)
        Returns:

        """
        targets_dict = self.target_assigner.assign_targets(
            gt_boxes
        )

        return targets_dict

    def forward(self, data_dict):
        multi_head_features = []
        spatial_features_2d = data_dict['spatial_features_2d']
        spatial_features_2d = self.shared_conv(spatial_features_2d)
        for task in self.tasks:
            multi_head_features.append(task(spatial_features_2d))

        self.forward_ret_dict['multi_head_features'] = multi_head_features

        if self.training:
            targets_dict = self.assign_targets(
                gt_boxes=data_dict['gt_boxes']
            )
            self.forward_ret_dict.update(targets_dict)

        if not self.training or self.predict_boxes_when_training:
            data_dict = self.generate_predicted_boxes(data_dict)

        return data_dict

    def build_losses(self):
        self.add_module(
            'crit',
            loss_utils.CenterNetFocalLoss()
        )
        self.add_module(
            'crit_reg',
            loss_utils.CenterNetRegLoss()
        )
        return

    def _sigmoid(self, x):
        y = torch.clamp(x.sigmoid(), min=1e-4, max=1 - 1e-4)
        return y

    def get_loss(self):
        tb_dict = {}
        pred_dicts = self.forward_ret_dict['multi_head_features']
        center_loss = []
        self.forward_ret_dict['pred_box_encoding'] = {}
        for task_id, pred_dict in enumerate(pred_dicts):
            pred_dict['hm'] = self._sigmoid(pred_dict['hm'])
            hm_loss = self.crit(pred_dict['hm'], self.forward_ret_dict['heatmap'][task_id])

            target_box_encoding = self.forward_ret_dict['box_encoding'][task_id]
            # nuscense encoding format [x, y, z, w, l, h, sinr, cosr, vx, vy]
            pred_box_encoding = torch.cat([
                pred_dict['reg'],
                pred_dict['height'],
                pred_dict['dim'],
                pred_dict['rot'],
                pred_dict['vel']
            ], dim = 1).contiguous() # (B, 10, H, W)

            self.forward_ret_dict['pred_box_encoding'][task_id] = pred_box_encoding

            box_loss = self.crit_reg(
                pred_box_encoding,
                self.forward_ret_dict['mask'][task_id],
                self.forward_ret_dict['ind'][task_id],
                target_box_encoding
            )

            loc_loss = (box_loss * box_loss.new_tensor(self.code_weights)).sum()
            loss = hm_loss + self.weight * loc_loss

            tb_key = 'task_' + str(task_id) + '/'
            tb_dict.update({
                tb_key + 'loss': loss.item(), tb_key + 'hm_loss': hm_loss.item(), tb_key + 'loc_loss': loc_loss.item(),
                tb_key + 'x_loss': box_loss[0].item(), tb_key + 'y_loss': box_loss[1].item(), tb_key + 'z_loss': box_loss[2].item(),
                tb_key + 'w_loss': box_loss[3].item(), tb_key + 'l_loss': box_loss[4].item(), tb_key + 'h_loss': box_loss[5].item(),
                tb_key + 'sin_r_loss': box_loss[6].item(), tb_key + 'cos_r_loss': box_loss[7].item(),
                tb_key + 'vx_loss': box_loss[8].item(), tb_key + 'vy_loss': box_loss[9].item(),
                tb_key + 'num_positive': self.forward_ret_dict['mask'][task_id].float().sum(),
            })
            center_loss.append(loss)

        return sum(center_loss), tb_dict

    def _double_flip_process(self, pred_dict, batch_size):
        for k in pred_dict.keys():
            # transform the prediction map back to their original coordinate befor flipping
            # the flipped predictions are ordered in a group of 4. The first one is the original pointcloud
            # the second one is X flip pointcloud(y=-y), the third one is Y flip pointcloud(x=-x), and the last one is
            # X and Y flip pointcloud(x=-x, y=-y).
            # Also please note that pytorch's flip function is defined on higher dimensional space, so dims=[2] means that
            # it is flipping along the axis with H length(which is normaly the Y axis), however in our traditional word, it is flipping along
            # the X axis. The below flip follows pytorch's definition yflip(y=-y) xflip(x=-x)
            _, C, H, W = pred_dict[k].shape
            pred_dict[k] = pred_dict[k].reshape(int(batch_size), 4, C, H, W)
            pred_dict[k][:, 1] = torch.flip(pred_dict[k][:, 1], dims=[2])
            pred_dict[k][:, 2] = torch.flip(pred_dict[k][:, 2], dims=[3])
            pred_dict[k][:, 3] = torch.flip(pred_dict[k][:, 3], dims=[2, 3])

        # batch_hm = pred_dict['hm'].sigmoid_() inplace may cause errors
        batch_hm = pred_dict['hm'].sigmoid()

        batch_reg = pred_dict['reg']
        batch_hei = pred_dict['height']

        batch_dim = pred_dict['dim']

        batch_hm = batch_hm.mean(dim=1)
        batch_hei = batch_hei.mean(dim=1)
        batch_dim = batch_dim.mean(dim=1)

        # y = -y reg_y = 1-reg_y
        batch_reg[:, 1, 1] = 1 - batch_reg[:, 1, 1]
        batch_reg[:, 2, 0] = 1 - batch_reg[:, 2, 0]
        batch_reg[:, 3, 0] = 1 - batch_reg[:, 3, 0]
        batch_reg[:, 3, 1] = 1 - batch_reg[:, 3, 1]
        batch_reg = batch_reg.mean(dim=1)

        batch_rots = pred_dict['rot'][:, :, 0:1]
        batch_rotc = pred_dict['rot'][:, :, 1:2]

        # first yflip
        # y = -y theta = pi -theta ???
        # sin(pi-theta) = sin(theta) cos(pi-theta) = -cos(theta) ???
        # batch_rotc[:, 1] = -batch_rotc[:, 1] ???

        # y = -y; theta = -theta
        # sin(-theta) = -sin(theta); cos(-theta) = cos(theta)
        batch_rots[:, 1] = -batch_rots[:, 1]

        # then xflip x = -x theta = 2pi - theta ???
        # sin(2pi - theta) = -sin(theta) cos(2pi - theta) = cos(theta) ???
        # batch_rots[:, 2] = -batch_rots[:, 2] ???

        # x = -x; theta = pi - theta
        # sin(pi - theta) = sin(theta); cos(pi - theta) = -cos(theta)
        batch_rotc[:, 2] = -batch_rotc[:, 2]

        # double flip
        batch_rots[:, 3] = -batch_rots[:, 3]
        batch_rotc[:, 3] = -batch_rotc[:, 3]

        batch_rotc = batch_rotc.mean(dim=1)
        batch_rots = batch_rots.mean(dim=1)

        batch_vel = pred_dict['vel']
        # flip vy
        batch_vel[:, 1, 1] = - batch_vel[:, 1, 1]
        # flip vx
        batch_vel[:, 2, 0] = - batch_vel[:, 2, 0]
        batch_vel[:, 3] = - batch_vel[:, 3]
        batch_vel = batch_vel.mean(dim=1)

        batch_dir_preds = [None] * batch_size
        return batch_hm, batch_rots, batch_rotc, batch_hei, batch_dim, batch_dir_preds, batch_vel, batch_reg

    def _transpose_and_gather_feat(self, feat, ind):
        feat = feat.permute(0, 2, 3, 1).contiguous()
        feat = feat.view(feat.size(0), -1, feat.size(3))
        feat = self._gather_feat(feat, ind)
        return feat

    def _gather_feat(self, feat, ind, mask=None):
        dim = feat.size(2)
        ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim)
        feat = feat.gather(1, ind)
        if mask is not None:
            mask = mask.unsqueeze(2).expand_as(feat)
            feat = feat[mask]
            feat = feat.view(-1, dim)
        return feat

    def _nms(self, heat, kernel=3):
        pad = (kernel - 1) // 2

        hmax = nn.functional.max_pool2d(
            heat, (kernel, kernel), stride=1, padding=pad)
        keep = (hmax == heat).float()
        return heat * keep

    def _topk(self, scores, K=40):
        batch, cat, height, width = scores.size()

        topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K)

        topk_inds = topk_inds % (height * width)
        topk_ys = (topk_inds / width).int().float()
        topk_xs = (topk_inds % width).int().float()

        topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
        topk_clses = (topk_ind / K).int()
        topk_inds = self._gather_feat(
            topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
        topk_ys = self._gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
        topk_xs = self._gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

        return topk_score, topk_inds, topk_clses, topk_ys, topk_xs

    @numba.jit(nopython=True)
    def circle_nms(self, dets, thresh):
        x1 = dets[:, 0]
        y1 = dets[:, 1]
        scores = dets[:, 2]
        order = scores.argsort()[::-1].astype(np.int32)  # highest->lowest
        ndets = dets.shape[0]
        suppressed = np.zeros((ndets), dtype=np.int32)
        keep = []
        for _i in range(ndets):
            i = order[_i]  # start with highest score box
            if suppressed[i] == 1:  # if any box have enough iou with this, remove it
                continue
            keep.append(i)
            for _j in range(_i + 1, ndets):
                j = order[_j]
                if suppressed[j] == 1:
                    continue
                # calculate center distance between i and j box
                dist = (x1[i] - x1[j]) ** 2 + (y1[i] - y1[j]) ** 2

                # ovr = inter / areas[j]
                if dist <= thresh:
                    suppressed[j] = 1
        return keep

    def _circle_nms(self, boxes, min_radius, post_max_size=83):
        """
        NMS according to center distance
        """
        keep = np.array(self.circle_nms(boxes.cpu().numpy(), thresh=min_radius))[:post_max_size]
        keep = torch.from_numpy(keep).long().to(boxes.device)
        return keep

    def _rotate_nms(self, boxes, scores, thresh, pre_maxsize=None, post_max_size=None):
        """
        :param boxes: (N, 5) [x1, y1, x2, y2, ry]
        :param scores: (N)
        :param thresh:
        :return:
        """
        # areas = (x2 - x1) * (y2 - y1)
        order = scores.sort(0, descending=True)[1]
        if pre_maxsize is not None:
            order = order[:pre_maxsize]

        boxes = boxes[order].contiguous()

        keep = torch.LongTensor(boxes.size(0))
        num_out = center_ops_cuda.center_rotate_nms_gpu(boxes, keep, thresh)
        selected = order[keep[:num_out].cuda()].contiguous()

        if post_max_size is not None:
            selected = selected[:post_max_size]

        return selected

    def _nms_gpu_3d(self, boxes, scores, thresh, pre_maxsize=None, post_max_size = None):
        """
        :param boxes: (N, 7) [x, y, z, dx, dy, dz, heading]
        :param scores: (N)
        :param thresh:
        :return:
        """
        assert boxes.shape[1] == 7
        order = scores.sort(0, descending=True)[1]
        if pre_maxsize is not None:
            order = order[:pre_maxsize]

        boxes = boxes[order].contiguous()
        keep = torch.LongTensor(boxes.size(0))
        num_out = iou3d_nms_cuda.nms_gpu(boxes, keep, thresh)
        selected =  order[keep[:num_out].cuda()].contiguous()

        if post_max_size is not None:
            selected = selected[:post_max_size]

        return selected

    def _boxes3d_to_bevboxes_lidar_torch(self, boxes3d):
        """
        :param boxes3d: (N, 7) [x, y, z, w, l, h, ry] in LiDAR coords
        :return:
            boxes_bev: (N, 5) [x1, y1, x2, y2, ry]
        """
        boxes_bev = boxes3d.new(torch.Size((boxes3d.shape[0], 5)))

        cu, cv = boxes3d[:, 0], boxes3d[:, 1]

        half_w, half_l = boxes3d[:, 3] / 2, boxes3d[:, 4] / 2
        boxes_bev[:, 0], boxes_bev[:, 1] = cu - half_w, cv - half_l
        boxes_bev[:, 2], boxes_bev[:, 3] = cu + half_w, cv + half_l
        boxes_bev[:, 4] = boxes3d[:, -1]
        return boxes_bev

    @torch.no_grad()
    def proposal_layer(self, heat, rots, rotc, hei, dim, vel, reg = None,
                       post_center_range=None, score_threshold=None, cfg=None, raw_rot=False, task_id=-1):

        assert self.encode_background_as_zeros is True
        assert self.use_sigmoid_score is True
        batch, cat, _, _ = heat.size()
        nms_cfg = cfg.nms.train if self.training else cfg.nms.test
        K = nms_cfg.nms_pre_max_size # topK selected
        maxpool = nms_cfg.get('max_pool_nms', False) or (nms_cfg.get('circle_nms', False) and (nms_cfg.min_radius[task_id] == -1))
        use_circle_nms = nms_cfg.get('circle_nms', False) and (nms_cfg.min_radius[task_id] != -1)
        if maxpool:
            heat = self._nms(heat)
        scores, inds, clses, ys, xs = self._topk(heat, K=K)
        # class label
        clses = clses.view(batch, K).float()
        cls_selected = clses.reshape(-1).long() # (batch * K)
        scores = scores.view(batch, K)
        assert reg is not None
        reg = self._transpose_and_gather_feat(reg, inds)
        reg = reg.view(batch, K, 2)
        xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
        ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
        assert raw_rot is False
        rots = self._transpose_and_gather_feat(rots, inds)
        rots = rots.view(batch, K, 1)
        rotc = self._transpose_and_gather_feat(rotc, inds)
        rotc = rotc.view(batch, K, 1)
        rot = torch.atan2(rots, rotc)

        # height in the bev
        hei = self._transpose_and_gather_feat(hei, inds)
        hei = hei.view(batch * K, 1)
        # decode height
        z = self.model_cfg.TASKS[task_id].z
        z = torch.tensor(z).to(hei.device)
        z = z[cls_selected] # (batch * K, 1)
        hei = hei + z
        hei = hei.reshape(batch, K, 1)

        # dim of the box
        dim = self._transpose_and_gather_feat(dim, inds)
        dim = dim.view(batch * K, 3)
        # decode dim
        anchors = self.model_cfg.TASKS[task_id].anchors
        anchors = torch.tensor(anchors).to(dim.device)
        anchors = anchors[cls_selected] # (batch * K, 3)
        dim = torch.exp(dim) * anchors
        dim = dim.reshape(batch, K, 3)

        # center location
        pc_range = cfg.pc_range
        xs = xs.view(batch, K, 1) * cfg.out_size_factor * cfg.voxel_size[0] + pc_range[0]
        ys = ys.view(batch, K, 1) * cfg.out_size_factor * cfg.voxel_size[1] + pc_range[1]
        vel = self._transpose_and_gather_feat(vel, inds)
        vel = vel.view(batch, K, 2)
        # vel after rot
        final_box_preds = torch.cat(
            [xs, ys, hei, dim, rot, vel], dim=2
        )
        final_scores = scores
        final_preds = clses

        # restrict center range
        assert post_center_range is not None
        post_center_range = torch.tensor(post_center_range).to(final_box_preds.device)
        mask = (final_box_preds[..., :3] >= post_center_range[:3]).all(2)
        mask &= (final_box_preds[..., :3] <= post_center_range[3:]).all(2)
        # use score threshold
        assert score_threshold is not None
        thresh_mask = final_scores > score_threshold
        mask &= thresh_mask

        predictions_dicts = []
        for i in range(batch):
            cmask = mask[i, :]
            boxes3d = final_box_preds[i, cmask]
            scores = final_scores[i, cmask]
            labels = final_preds[i, cmask]

            # circle nms
            if use_circle_nms:
                centers = boxes3d[:, [0, 1]]
                boxes = torch.cat([centers, scores.view(-1, 1)], dim=1)
                keep = self._circle_nms(boxes, min_radius=nms_cfg.min_radius[task_id], post_max_size=nms_cfg.post_max_size)

                boxes3d = boxes3d[keep]
                scores = scores[keep]
                labels = labels[keep]

            # rotate nms
            elif nms_cfg.get('use_rotate_nms', False):
                assert not maxpool
                top_scores = scores
                if top_scores.shape[0] != 0:
                    boxes_for_nms =self._boxes3d_to_bevboxes_lidar_torch(boxes3d)
                    selected = self._rotate_nms(boxes_for_nms, top_scores,
                                                thresh=nms_cfg.nms_iou_threshold,
                                                pre_maxsize=nms_cfg.nms_pre_max_size,
                                                post_max_size=nms_cfg.nms_post_max_size
                                               )
                else:
                    selected = []
                boxes3d = boxes3d[selected]
                labels = labels[selected]
                scores = scores[selected]

            # iou 3d nms
            elif nms_cfg.get('use_iou_3d_nms', False):
                assert not maxpool
                top_scores = scores
                if top_scores.shape[0] != 0:
                    selected = self._nms_gpu_3d(boxes3d[:, :7], top_scores,
                                                thresh=nms_cfg.nms_iou_threshold,
                                                pre_maxsize=nms_cfg.nms_pre_max_size,
                                                post_max_size=nms_cfg.nms_post_max_size
                                               )
                else:
                    selected = []
                boxes3d = boxes3d[selected]
                labels = labels[selected]
                scores = scores[selected]

            predictions_dict = {
                "boxes": boxes3d,
                "scores": scores,
                "labels": labels.long()
            }
            predictions_dicts.append(predictions_dict)
        return predictions_dicts

    @torch.no_grad()
    def generate_predicted_boxes(self, data_dict):
        """
        Generate box predictions with decode, topk and circular_nms
        For single-stage-detector, another post-processing (nms) is needed
        For two-stage-detector, no need for proposal layer in roi_head
        Returns:
        """
        double_flip = not self.training and self.post_cfg.get('double_flip', False)
        post_center_range = self.post_cfg.post_center_limit_range
        pred_dicts = self.forward_ret_dict['multi_head_features']

        task_box_preds = {}
        task_score_preds = {}
        task_label_preds = {}
        for task_id, pred_dict in enumerate(pred_dicts):
            batch_size = pred_dict['hm'].shape[0]
            if double_flip:
                assert batch_size % 4 == 0, print(batch_size)
                batch_size = int(batch_size / 4)
                batch_hm, batch_rots, batch_rotc, batch_hei, batch_dim, batch_dir_preds, batch_vel, batch_reg = \
                    self._double_flip_process(pred_dict, batch_size)
                # convert data_dict format
                data_dict['batch_size'] = batch_size
            else:
                batch_hm = pred_dict['hm'].sigmoid()
                batch_reg = pred_dict['reg']
                batch_hei = pred_dict['height']
                batch_dim = pred_dict['dim']
                batch_rots = pred_dict['rot'][:, 0].unsqueeze(1)
                batch_rotc = pred_dict['rot'][:, 1].unsqueeze(1)
                batch_vel = pred_dict['vel']

            #decode
            boxes = self.proposal_layer(
                batch_hm,
                batch_rots,
                batch_rotc,
                batch_hei,
                batch_dim,
                batch_vel,
                reg=batch_reg,
                post_center_range=post_center_range,
                score_threshold=self.post_cfg.score_threshold,
                cfg=self.post_cfg,
                task_id=task_id
            )
            task_box_preds[task_id] = [box['boxes'] for box in boxes]
            task_score_preds[task_id] = [box['scores'] for box in boxes]
            task_label_preds[task_id] = [box['labels'] for box in boxes] #labels are local here

        pred_dicts = []
        batch_size = len(task_box_preds[0])
        rois, roi_scores, roi_labels = [], [], []
        nms_cfg = self.post_cfg.nms.train if self.training else self.post_cfg.nms.test
        num_rois = nms_cfg.nms_post_max_size * len(self.class_names)
        for batch_idx in range(batch_size):
            offset = 1 # class label start from 1
            final_boxes, final_scores, final_labels = [], [], []
            for task_id, class_name in enumerate(self.class_names):
                final_boxes.append(task_box_preds[task_id][batch_idx])
                final_scores.append(task_score_preds[task_id][batch_idx])
                # convert to global labels
                final_global_label = task_label_preds[task_id][batch_idx] + offset
                offset += len(class_name)
                final_labels.append(final_global_label)

            final_boxes = torch.cat(final_boxes)
            final_scores = torch.cat(final_scores)
            final_labels = torch.cat(final_labels)

            roi = final_boxes.new_zeros(num_rois, final_boxes.shape[-1])
            roi_score = final_scores.new_zeros(num_rois)
            roi_label = final_labels.new_zeros(num_rois)
            num_boxes = final_boxes.shape[0]
            roi[:num_boxes] = final_boxes
            roi_score[:num_boxes] = final_scores
            roi_label[:num_boxes] = final_labels
            rois.append(roi)
            roi_scores.append(roi_score)
            roi_labels.append(roi_label)

            record_dict = {
                "pred_boxes": final_boxes,
                "pred_scores": final_scores,
                "pred_labels": final_labels
            }
            pred_dicts.append(record_dict)

        data_dict['pred_dicts'] = pred_dicts
        data_dict['rois'] = torch.stack(rois, dim = 0)
        data_dict['roi_scores'] = torch.stack(roi_scores, dim = 0)
        data_dict['roi_labels'] = torch.stack(roi_labels, dim = 0)
        data_dict['has_class_labels'] = True  # Force to be true
        data_dict.pop('batch_index', None)
        return data_dict

"""
BASIC BUILDING BLOCKS
"""
class Sequential(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super(Sequential, self).__init__()
        if len(args) == 1 and isinstance(args[0], OrderedDict):
            for key, module in args[0].items():
                self.add_module(key, module)
        else:
            for idx, module in enumerate(args):
                self.add_module(str(idx), module)
        for name, module in kwargs.items():
            if sys.version_info < (3, 6):
                raise ValueError("kwargs only supported in py36+")
            if name in self._modules:
                raise ValueError("name exists.")
            self.add_module(name, module)

    def __getitem__(self, idx):
        if not (-len(self) <= idx < len(self)):
            raise IndexError("index {} is out of range".format(idx))
        if idx < 0:
            idx += len(self)
        it = iter(self._modules.values())
        for i in range(idx):
            next(it)
        return next(it)

    def __len__(self):
        return len(self._modules)

    def add(self, module, name=None):
        if name is None:
            name = str(len(self._modules))
            if name in self._modules:
                raise KeyError("name exists")
        self.add_module(name, module)

    def forward(self, input):
        for module in self._modules.values():
            input = module(input)
        return input

def kaiming_init(
    module, a=0, mode="fan_out", nonlinearity="relu", bias=0, distribution="normal"
):
    assert distribution in ["uniform", "normal"]
    if distribution == "uniform":
        nn.init.kaiming_uniform_(
            module.weight, a=a, mode=mode, nonlinearity=nonlinearity
        )
    else:
        nn.init.kaiming_normal_(
            module.weight, a=a, mode=mode, nonlinearity=nonlinearity
        )
    if hasattr(module, "bias") and module.bias is not None:
        nn.init.constant_(module.bias, bias)

class SepHead(nn.Module):
    def __init__(
            self,
            in_channels,
            heads,
            head_conv=64,
            name="",
            final_kernel=1,
            bn=False,
            init_bias=-2.19,
            directional_classifier=False,
            **kwargs,
    ):
        super(SepHead, self).__init__(**kwargs)

        self.heads = heads
        for head in self.heads:
            classes, num_conv = self.heads[head]

            fc = Sequential()
            for i in range(num_conv - 1):
                fc.add(nn.Conv2d(in_channels, head_conv,
                                 kernel_size=final_kernel, stride=1,
                                 padding=final_kernel // 2, bias=True))
                if bn:
                    fc.add(nn.BatchNorm2d(head_conv))
                fc.add(nn.ReLU())

            fc.add(nn.Conv2d(head_conv, classes,
                             kernel_size=final_kernel, stride=1,
                             padding=final_kernel // 2, bias=True))

            if 'hm' in head:
                fc[-1].bias.data.fill_(init_bias)
            else:
                for m in fc.modules():
                    if isinstance(m, nn.Conv2d):
                        kaiming_init(m)

            self.__setattr__(head, fc)

        assert directional_classifier is False, "Doesn't work well with nuScenes in my experiments, please open a pull request if you are able to get it work. We really appreciate contribution for this."

    def forward(self, x):
        ret_dict = dict()
        for head in self.heads:
            ret_dict[head] = self.__getattr__(head)(x)

        return ret_dict

class FeatureAdaption(nn.Module):
    """Feature Adaption Module.

    Feature Adaption Module is implemented based on DCN v1.
    It uses anchor shape prediction rather than feature map to
    predict offsets of deformable conv layer.

    Args:
        in_channels (int): Number of channels in the input feature map.
        out_channels (int): Number of channels in the output feature map.
        kernel_size (int): Deformable conv kernel size.
        deformable_groups (int): Deformable conv group size.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size=3,
                 deformable_groups=4):
        super(FeatureAdaption, self).__init__()
        offset_channels = kernel_size * kernel_size * 2
        self.conv_offset = nn.Conv2d(
            in_channels, deformable_groups * offset_channels, 1, bias=True)
        self.conv_adaption = DeformConv(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            padding=(kernel_size - 1) // 2,
            deformable_groups=deformable_groups)
        self.relu = nn.ReLU(inplace=True)
        self.init_offset()

    def init_offset(self):
        self.conv_offset.weight.data.zero_()
        self.conv_offset.bias.data.zero_()

    def init_weights(self):
        pass

    def forward(self, x,):
        offset = self.conv_offset(x)
        x = self.relu(self.conv_adaption(x, offset))
        return x

class DCNSepHead(nn.Module):
    def __init__(
            self,
            in_channels,
            num_cls,
            heads,
            head_conv=64,
            name="",
            final_kernel=1,
            bn=False,
            init_bias=-2.19,
            directional_classifier=False,
            **kwargs,
    ):
        super(DCNSepHead, self).__init__(**kwargs)

        # feature adaptation with dcn
        # use separate features for classification / regression
        self.feature_adapt_cls = FeatureAdaption(
            in_channels,
            in_channels,
            kernel_size=3,
            deformable_groups=4)

        self.feature_adapt_reg = FeatureAdaption(
            in_channels,
            in_channels,
            kernel_size=3,
            deformable_groups=4)

        # heatmap prediction head
        self.cls_head = Sequential(
            nn.Conv2d(in_channels, head_conv,
                      kernel_size=3, padding=1, bias=True),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(head_conv, num_cls,
                      kernel_size=3, stride=1,
                      padding=1, bias=True)
        )
        self.cls_head[-1].bias.data.fill_(init_bias)

        # other regression target
        self.task_head = SepHead(in_channels, heads, head_conv=head_conv, bn=bn, final_kernel=final_kernel)

    def forward(self, x):
        center_feat = self.feature_adapt_cls(x)
        reg_feat = self.feature_adapt_reg(x)

        cls_score = self.cls_head(center_feat)
        ret = self.task_head(reg_feat)
        ret['hm'] = cls_score

        return ret


import os
import torch
import torch.utils.data as data
import cv2
import numpy as np
from config import DATA_ROOT

class PedestrainDataset(data.Dataset):

    def __init__(self, root, transform=None, dataset_name='train'):
        self.root = root
        self.transform = transform
        self.name = dataset_name
        self.train_anno_list = open(os.path.join(DATA_ROOT, 'train_annotations.txt')).read().split('\n')
        # self.train_ignore_list = open(os.path.join(DATA_ROOT, 'pedestrian_ignore_part_train')).read().split('\n')
        self.train_img_list = open(os.path.join(DATA_ROOT, 'train.txt')).read().split('\n')
        self.train_img_path = os.path.join(DATA_ROOT, 'Images')
        # self.imgs = os.listdir(self.img_path)
        index = 0
        for idx in range(len(self.train_img_list)):
            box_list = self.train_anno_list[index].split(' ')
            if len(box_list) == 1:
                self.train_img_list.remove(self.train_img_list[index])
                self.train_anno_list.remove(self.train_anno_list[index])
                index = index - 1
            index = index + 1
        self.length = len(self.train_img_list)

    def __getitem__(self, index):
        im, gt, h, w = self.pull_item(index)

        return im, gt

    def __len__(self):
        return self.length

    def pull_item(self, index):
        img_path = os.path.join(self.train_img_path, self.train_img_list[index])
        img = cv2.imread(img_path)
        height, width, channels = img.shape
        # print(self.img_list[index])

        anno = self.train_anno_list[index].split(" ")
        target = []
        object_num = int((len(anno) - 1) / 5)
        # print(object_num)
        res = []
        for p in range(object_num):
            if p >= 0:
                bndbox = []
                idx = 1 + 5 * p
                bndbox.append((int(anno[idx + 1]) - 1) / width)
                bndbox.append((int(anno[idx + 2]) - 1) / height)
                bndbox.append((int(anno[idx + 1]) + int(anno[idx + 3]) - 1) / width)
                bndbox.append((int(anno[idx + 2]) + int(anno[idx + 4]) - 1) / height)
                label = int(anno[idx]) - 1
                bndbox.append(label)
                res += [bndbox]
        if self.transform is not None:
            target = np.array(res)
            # print(target.shape)
            img, boxes, labels = self.transform(img, target[:, :4], target[:, 4])
            # to rgb
            img = img[:, :, (2, 1, 0)]
            # img = img.transpose(2, 0, 1)
            target = np.hstack((boxes, np.expand_dims(labels, axis=1)))
        return torch.from_numpy(img).permute(2, 0, 1), target, height, width
        # return torch.from_numpy(img), target, height, width

    def pull_image(self, index):
        '''Returns the original image object at index in PIL form

        Note: not using self.__getitem__(), as any transformations passed in
        could mess up this functionality.

        Argument:
            index (int): index of img to show
        Return:
            PIL img
        '''
        img_path = os.path.join(self.train_img_path, self.train_img_list[index])
        return cv2.imread(img_path, cv2.IMREAD_COLOR)

    def pull_anno(self, index):
        '''Returns the original annotation of image at index

        Note: not using self.__getitem__(), as any transformations passed in
        could mess up this functionality.

        Argument:
            index (int): index of img to get annotation of
        Return:
            list:  [img_id, [(label, bbox coords),...]]
                eg: ('001718', [('dog', (96, 13, 438, 332))])
        '''
        anno = self.train_anno_list[index].split(" ")
        object_num = int((len(anno) - 1) / 5)
        res = []
        for p in range(object_num):
            bndbox = []
            idx = 1 + 5 * p
            bndbox.append(int(anno[idx + 1]) - 1)
            bndbox.append(int(anno[idx + 2]) - 1)
            bndbox.append(int(anno[idx + 1]) + int(anno[idx + 3]) - 1)
            bndbox.append(int(anno[idx + 2]) + int(anno[idx + 4]) - 1)
            label = int(anno[idx]) - 1
            bndbox.append(label)
            res += [bndbox]
        id = index
        return id, res

    def pull_tensor(self, index):
        '''Returns the original image at an index in tensor form

        Note: not using self.__getitem__(), as any transformations passed in
        could mess up this functionality.

        Argument:
            index (int): index of img to show
        Return:
            tensorized version of img, squeezed
        '''
        return torch.Tensor(self.pull_image(index)).unsqueeze_(0)


from __future__ import division, print_function

import argparse
import os
import random

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
# import torch.utils.data
from torch.autograd import Variable

import dataset
import keys, data_dict
import utils
from modeling import crnn

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--root', default='/home/yjwang/Workspaces/data/recognition/images/', help='path to dataset')
    parser.add_argument('--train_label_path', default='/home/yjwang/Workspaces/data/recognition/train.txt', help='path to dataset')
    parser.add_argument('--val_label_path', default='/home/yjwang/Workspaces/data/recognition/test.txt', help='path to dataset')
    parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
    parser.add_argument('--batchSize', type=int, default=32, help='input batch size')
    parser.add_argument('--imgH', type=int, default=32, help='the height of the input image to network')
    parser.add_argument('--imgW', type=int, default=280, help='the width of the input image to network')
    parser.add_argument('--nepoch', type=int, default=25, help='number of epochs to train for')
    # TODO(meijieru): epoch -> iter
    parser.add_argument('--cuda', default='0', action='store_true', help='enables cuda')
    parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
    parser.add_argument('--pretrained', default='expr/denseNet_4_0.9664375.pth', help="path to pretrained model (to continue training)")
    parser.add_argument('--expr_dir', default='expr', help='Where to store samples and models')
    parser.add_argument('--displayInterval', type=int, default=100, help='Interval to be displayed')
    parser.add_argument('--n_test_disp', type=int, default=10, help='Number of samples to display when test')
    parser.add_argument('--valInterval', type=int, default=500, help='Interval to be displayed')
    parser.add_argument('--saveAccuracy', type=int, default=0.5, help='Interval to be displayed')
    parser.add_argument('--lr', type=float, default=1e-5, help='learning rate for Critic, not used by adadealta')
    parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
    parser.add_argument('--adam', action='store_true', help='Whether to use adam (default is rmsprop)')
    parser.add_argument('--adadelta', action='store_true', help='Whether to use adadelta (default is rmsprop)')
    parser.add_argument('--manualSeed', type=int, default=1234, help='reproduce experiemnt')
    parser.add_argument('--random_sample', action='store_true', help='whether to sample the dataset with random sampler')
    opt = parser.parse_args()
    print(opt)
    return opt

# custom weights initialization called on crnn
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('norm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

def val(crnn, valid_loader, criterion, max_iter=1000):
    print('Start val')
    for p in crnn.parameters():
        p.requires_grad = False
    crnn.eval()

    val_iter = iter(valid_loader)

    i = 0
    n_correct = 0
    loss_avg = utils.averager()

    max_iter = min(max_iter, len(valid_loader))
    for i in range(max_iter):
        names, images, texts = val_iter.next()
        batch_size = images.size(0)
        t, l = converter.encode(texts)
        images = images.cuda()
        preds = crnn(images)
        preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
        cost = criterion(preds, t, preds_size, l) / batch_size
        loss_avg.add(cost)

        _, preds = preds.max(2)
        preds = preds.transpose(1, 0).contiguous().view(-1)
        sim_preds = converter.decode(preds.data, preds_size.data, raw=False)
        for pred, target in zip(sim_preds, texts):
            if pred == target:
                n_correct += 1

    raw_preds = converter.decode(preds.data, preds_size.data, raw=True)[:opt.n_test_disp]
    for name, raw_pred, pred, gt in zip(names, raw_preds, sim_preds, texts):
        print('%-20s:%-20s => %-20s, gt: %-20s' % (name, raw_pred, pred, gt))

    accuracy = n_correct / float(max_iter * opt.batchSize)
    print('Test loss: %f, accuray: %f' % (loss_avg.val(), accuracy))
    return accuracy

def train(crnn, train_loader, criterion, optimizer, valid_loader):
    for p in crnn.parameters():
        p.requires_grad = True
    crnn.train()
    train_iter = iter(train_loader)
    # loss averager
    loss_avg = utils.averager()
    for i in range(len(train_loader)):
        data = train_iter.next()
        _, images, texts = data
        batch_size = images.size(0)
        t, l = converter.encode(texts)
        images = images.cuda()
        preds = crnn(images)
        preds_size = torch.IntTensor([preds.size(0)] * batch_size)
        cost = criterion(preds, t, preds_size, l) / batch_size
        crnn.zero_grad()
        cost.backward()
        optimizer.step()
        loss_avg.add(cost)
        if (i+1) % opt.displayInterval == 0:
            print('[%d/%d][%d/%d] Loss: %f' %
                (epoch, opt.nepoch, i, len(train_loader), loss_avg.val()))
            loss_avg.reset()
        # if (i+1) % opt.valInterval:
        #     _ = val(crnn, valid_loader, criterion, max_iter=100)
        #     for p in crnn.parameters():
        #         p.requires_grad = True
        #     crnn.train()

def backward_hook(self, grad_input, grad_output):
    for g in grad_input:
        g[g != g] = 0

if __name__ == "__main__":

    opt = get_args()

    if not os.path.exists(opt.expr_dir):
        os.makedirs(opt.expr_dir)

    random.seed(opt.manualSeed)
    np.random.seed(opt.manualSeed)
    torch.manual_seed(opt.manualSeed)

    cudnn.benchmark = True

    if torch.cuda.is_available() and not opt.cuda:
        print("WARNING: You have a CUDA device, so you should probably run with --cuda")

    train_dataset = dataset.folderDataset(opt.root, opt.train_label_path, transform=dataset.resizeNormalize((opt.imgW, opt.imgH)))
    assert train_dataset

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=opt.batchSize,
        shuffle=True, num_workers=int(opt.workers))

    valid_dataset = dataset.folderDataset(
        opt.root, opt.val_label_path, transform=dataset.resizeNormalize((opt.imgW, opt.imgH)))

    valid_loader = torch.utils.data.DataLoader(
        valid_dataset, shuffle=True, batch_size=opt.batchSize, num_workers=int(opt.workers))

    alphabet = data_dict.alphabet
    nclass = len(alphabet) + 1
    nc = 1

    criterion = torch.nn.MSELoss(reduction='mean')
    ## ResNet18
    # crnn = crnn.CRNN(CNN="ResNet18", RNN='DoubleBiLSTM',  nIn=512, n_class=nclass, nHidden=256, nLayer=1, dropout=0)
    # Densenet18
    CNN = "denseNet"
    RNN = "DoubleBiLSTM"
    crnn = crnn.CRNN(CNN=CNN, RNN=RNN,  nIn=128, n_class=nclass, nHidden=256, nLayer=1, dropout=0)
    # CNN = "DenseNet18_256"
    # RNN = "DoubleBiLSTM"
    # crnn = crnn.CRNN(CNN=CNN, RNN=RNN,  nIn=256, n_class=nclass, nHidden=256, nLayer=1, dropout=0)

    crnn.apply(weights_init)
    if opt.pretrained != '':
        print('loading pretrained model from %s' % opt.pretrained)
        crnn.load_state_dict(torch.load(opt.pretrained))

    from torchsummary import summary
    summary(crnn, (1, 32, 280))

    if opt.cuda:
        crnn = crnn.cuda()
        # crnn = torch.nn.DataParallel(crnn, device_ids=range(opt.ngpu))
        criterion = criterion.cuda()

    # setup optimizer
    if opt.adam:
        optimizer = optim.Adam(crnn.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
    elif opt.adadelta:
        optimizer = optim.Adadelta(crnn.parameters())
    else:
        optimizer = optim.RMSprop(crnn.parameters(), lr=opt.lr)

    crnn.register_backward_hook(backward_hook)

    for epoch in range(opt.nepoch):
        train(crnn, train_loader, criterion, optimizer, valid_loader)
        acc = val(crnn, valid_loader, criterion)
        if acc > opt.saveAccuracy:
            torch.save(
                crnn.state_dict(), '{0}/{1}_{2}_{3}.pth'.format(opt.expr_dir,CNN, epoch, acc))


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import util_misc

import os
import numpy as np
import tensorflow as tf

import util_io

class UtilMiscTest(tf.test.TestCase):
  def test_visualize_landmark_predictions(self):
    num_landmarks = 5
    image_size = 4
    landmark_predictions_np = np.zeros((2, 4, 4, num_landmarks, 3), dtype=np.float32)
    for i in range(num_landmarks):
      landmark_predictions_np[0, 0, 0, i, 2] = 100.0
      landmark_predictions_np[1, 1, 1, i, 2] = 10.0
    landmark_predictions = tf.constant(landmark_predictions_np)

    grid_size = landmark_predictions.shape[1]
    assert grid_size == landmark_predictions.shape[2]
    # Row = [[0, 1, 2, 3...], [0, 1, 2, 3...]...]. Column = [[0, 0, ..], [1, 1,...]...]
    row, column = tf.meshgrid(tf.range(grid_size), tf.range(grid_size))
    dtype = landmark_predictions.dtype
    row = tf.cast(row, dtype) / tf.cast(grid_size, dtype)
    column = tf.cast(column, dtype) / tf.cast(grid_size, dtype)
    rc = tf.stack((row, column), axis=-1)
    rc_reshaped = tf.expand_dims(tf.expand_dims(rc, 0), 3)
    rc_reshaped = util_misc.tf_repeat(rc_reshaped, [landmark_predictions.shape[0], 1, 1, landmark_predictions.shape[3], 1])
    if True:
      landmark_confidence_reshaped = tf.transpose(landmark_predictions[..., 2], (0, 3, 1, 2))
      landmark_confidence_reshaped = tf.reshape(landmark_confidence_reshaped,
                                                tf.TensorShape([landmark_confidence_reshaped.shape[0],
                                                                landmark_confidence_reshaped.shape[1],
                                                                landmark_confidence_reshaped.shape[2] *
                                                                landmark_confidence_reshaped.shape[3], ]))
      c = util_misc.layers.softmax(landmark_confidence_reshaped)
    else:
      c = tf.sigmoid(landmark_predictions[..., 2])
    c = tf.transpose(c, (0, 2, 1))
    # c = tf.reshape(c, tf.TensorShape(
    #   [c.shape[0], grid_size, grid_size, c.shape[2]]))  # [batch size, height, width, num_landmarks]
    xy = rc_reshaped + landmark_predictions[..., 0:2] / int(
      grid_size)  # [batch size, height, width, num_landmarks, 2]
    xy_reshaped = tf.reshape(xy, tf.TensorShape([xy.shape[0], xy.shape[1] * xy.shape[2] * xy.shape[3] * xy.shape[4]]))
    xy_one_hot = util_misc.xy_to_one_hot(xy_reshaped, image_size)
    xy_one_hot_weighed = xy_one_hot * tf.expand_dims(tf.expand_dims(util_misc.layers.flatten(c), 1), 2)
    xy_one_hot_weighed = tf.reduce_sum(xy_one_hot_weighed, axis=-1, keepdims=True)

    landmark_visualization = xy_one_hot_weighed
    with self.test_session():
      actual_output = landmark_visualization.eval()
      print(actual_output)

  def test_is_sketch(self):
    image_dir = u'/mnt/f032b8a5-c186-4fae-b911-bcfdee99a2e9/pixiv_collected_sketches/PixivUtil2/test_samples_tiny'
    image_paths_and_expected = [
      (u'11739035_p1 - レミリア.jpg', False),
      (u'12925016_p2 - 東方系ラクガキ詰め合わせ.jpg', False),
      (u'14362806_p0 - 無題.jpg', False),  # Background not white enough.
      (u'14485122_p0 - 宮子.jpg', False),  # Background not white enough.
      (u'15444948_p4 - ぐ～てん☆もるげんっ！.jpg', False),
      (u'15469173_p0 - 東方魔理沙　塗ってみた.jpg', False),
      (u'17834866_p0 - 以蔵の闘い.jpg', False),
      (u'24774862_p2 - からくりばーすと、メイキング.jpg', False),
      (u'28474957_p0 - 白澤さん。.jpg', True),  # Wierd edge case unrecognizable to human eyes.
      (u'28592442_p0 - 本田菊.jpg', False),
      (u'29908672_p0 - ソードアート・おっぱい.jpg', False),
      (u'30454124_p0 - 水野亜美(セーラーマーキュリー).jpg', False),
      (u'5833646_p0 - Q命病棟でQP化(線画).png', True),
      (u'5943678_p0 - バニー達は塗ってほしそうにこちらを見ている・・.jpg', True),
      (u'5952556_p0 - 死神線画.png', True),
      (u'5981054_p0 - リン・レン線画.jpg', False),
      (u'5987994_p0 - お願いします。.png', True),
      (u'7210533_p0 - カウントダウン2日前.jpg', False),
      (u'7242346_p0 - ミリー【線画】.jpg', True),
      (u'7252006_p0 - ナタネ＆ロズレイド線画.jpg', True),
      (u'7304820_p0 - 少女の宮サマ。.jpg', True),
      (u'7431716_p0 - 今吉さん線画.jpg', True),
      (u'8113468_p0 - あんこくのじょおう【再投稿】.png', True),
      (u'8152845_p0 - ヘルメスさん.jpg', False),
      (u'8425348_p0 - 早苗さん（線画）.png', True),  # PNG bug, minor recall loss.
      (u'8441291_p0 - 【線画】ベルベル.jpg', False),
      (u'8502938_p0 - ベアト線画！.jpg', True),
    ]
    for image_name, expected in image_paths_and_expected:
      image = util_io.imread(os.path.join(image_dir, image_name))
      actual_output = util_misc.is_sketch(image)
      if not actual_output == expected:
        print('unexpected: ', image_name)
      # self.assertAllEqual(actual_output, expected)

  def test_pad_and_break_up_image(self):
    height, width, channel = 16, 16, 1
    image = np.array([[[h * w * (c + 1) for c in range(channel)] for w in range(width)] for h in range(height)], dtype=np.uint8)
    dest_height, dest_width = 3, 4
    broken_image, num_h, num_w, pad_h, pad_w = util_misc.pad_and_break_up_image(image, dest_height, dest_width)
    # Not the best way to test it, but works for now..
    regrouped_image = util_misc.regroup_broken_image(broken_image, num_h, num_w, pad_h, pad_w)
    self.assertAllEqual(regrouped_image, image)

if __name__ == '__main__':
  tf.test.main()