from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import os
import sys
import time
import tensorflow as tf

slim = tf.contrib.slim
sys.path.insert(0, './slim/')
from nets import inception, resnet_v2
from preprocessing import inception_preprocessing

data_dir = './data'
# dataset needs to be one of ['ILSVRC2012', 'inat2017', 'cub_200', 'flower_102',
# 'stanford_cars', 'stanford_dogs', 'aircraft', 'nabirds', 'food_101']
dataset = 'cub_200'
# base_network needs to be one of ['InceptionV3', 'InceptionV3SE',
# 'InceptionV4', 'InceptionResnetV2', 'InceptionResnetV2SE', 'ResNet50',
# 'ResNet101', 'ResNet152']
base_network = 'InceptionV3'
checkpoints_path = './checkpoints/inception/inception_v3_iNat_299.ckpt'
# base_network = 'ResNet101'
# checkpoints_path = './checkpoints/resnet/resnet_101_ILSVRC_iNat_299.ckpt'

image_size = 299
moving_average_decay = 0.9999
fea_dim = 2048

# Read train and val list.
train_list = []
val_list = []
for line in open(os.path.join(data_dir, dataset, 'train.txt'), 'r'):
    train_list.append(
        (os.path.join(data_dir, dataset, line.strip().split(': ')[0]),
        int(line.strip().split(': ')[1])))
for line in open(os.path.join(data_dir, dataset, 'val.txt'), 'r'):
    val_list.append(
        (os.path.join(data_dir, dataset, line.strip().split(': ')[0]),
        int(line.strip().split(': ')[1])))
print('Length of train: %d' %len(train_list))
print('Length of val: %d' %len(val_list))

# Base network architecture
if base_network == 'InceptionV3':
    endpoint = 'Mixed_7c'
    arg_scope = inception.inception_v3_arg_scope()
elif base_network == 'InceptionV3SE':
    endpoint = 'Mixed_7c'
    arg_scope = inception.inception_v3_se_arg_scope()
elif base_network == 'InceptionV4':
    endpoint = 'Mixed_7d'
    arg_scope = inception.inception_v4_arg_scope()
elif base_network == 'InceptionResnetV2':
    endpoint = 'Conv2d_7b_1x1'
    arg_scope = inception.inception_resnet_v2_arg_scope()
elif base_network == 'InceptionResnetV2SE':
    endpoint = 'Conv2d_7b_1x1'
    arg_scope = inception.inception_resnet_v2_se_arg_scope()
elif base_network[:6] == 'ResNet':
    layers = base_network.split('ResNet')[1]
    base_network = 'ResNet'

# Feature extraction.
fea_train = np.zeros((len(train_list), fea_dim), dtype=np.float32)
label_train = np.zeros((len(train_list), ), dtype=np.int32)
fea_val = np.zeros((len(val_list), fea_dim), dtype=np.float32)
label_val = np.zeros((len(val_list), ), dtype=np.int32)

with tf.Graph().as_default():
    tf_global_step = tf.train.get_or_create_global_step()
    image_path = tf.placeholder(tf.string)
    image = tf.image.decode_jpeg(tf.read_file(image_path), channels=3)
    image = tf.image.convert_image_dtype(image, tf.float32)
    image = inception_preprocessing.preprocess_image(image,
                                                     image_size,
                                                     image_size,
                                                     is_training=False)
    images  = tf.expand_dims(image, 0)

    if base_network == 'ResNet':
        with slim.arg_scope(resnet_v2.resnet_arg_scope(use_batch_norm=True)):
            if layers == '50':
                net, _ = resnet_v2.resnet_v2_50(images, is_training=False)
            elif layers == '101':
                net, _ = resnet_v2.resnet_v2_101(images, is_training=False)
            elif layers == '152':
                net, _ = resnet_v2.resnet_v2_152(images, is_training=False)
    else:
        with slim.arg_scope(arg_scope):
            slim_args = [slim.batch_norm, slim.dropout]
            with slim.arg_scope(slim_args, is_training=False):
                with tf.variable_scope(base_network, reuse=None) as scope:
                    if base_network == 'InceptionV3':
                        net, _ = inception.inception_v3_base(
                            images, final_endpoint=endpoint, scope=scope)
                    elif base_network == 'InceptionV3SE':
                        net, _ = inception.inception_v3_se_base(
                            images, final_endpoint=endpoint, scope=scope)
                    elif base_network == 'InceptionV4':
                        net, _ = inception.inception_v4_base(
                            images, final_endpoint=endpoint, scope=scope)
                    elif base_network == 'InceptionResnetV2':
                        net, _ = inception.inception_resnet_v2_base(
                            images, final_endpoint=endpoint, scope=scope)
                    elif base_network == 'InceptionResnetV2SE':
                        net, _ = inception.inception_resnet_v2_se_base(
                            images, final_endpoint=endpoint, scope=scope)
    net = tf.reduce_mean(net, [0,1,2])

    variable_averages = tf.train.ExponentialMovingAverage(
        moving_average_decay, tf_global_step)
    variables_to_restore = variable_averages.variables_to_restore()
    init_fn = slim.assign_from_checkpoint_fn(
        checkpoints_path, variables_to_restore)

    config_sess = tf.ConfigProto(allow_soft_placement=True)
    config_sess.gpu_options.allow_growth = True
    with tf.Session(config=config_sess) as sess:
        init_fn(sess)
        start = time.time()
        print('Feature extraction on training set...')
        for i in range(len(fea_train)):
            if i%1000 == 0:
                print('%d/%d %.2fs'%(i, len(fea_train), time.time() - start))
            fea = sess.run(net, feed_dict={image_path:train_list[i][0]})
            fea_train[i, :] = fea
            label_train[i] = train_list[i][1]
        print('Feature extraction on validation set...')
        for i in range(len(fea_val)):
            if i%1000 == 0:
                print('%d/%d %.2fs'%(i, len(fea_val), time.time() - start))
            fea = sess.run(net, feed_dict={image_path:val_list[i][0]})
            fea_val[i, :] = fea
            label_val[i] = val_list[i][1]

model_name = checkpoints_path.split('/')[-1].split('.ckpt')[0]
if not os.path.exists(os.path.join('./feature', model_name)):
    os.makedirs(os.path.join('./feature', model_name))

save_dir = os.path.join('./feature', model_name, dataset)
np.save(os.path.join(save_dir + '_feature_train.npy'), fea_train)
np.save(os.path.join(save_dir + '_label_train.npy'), label_train)
np.save(os.path.join(save_dir + '_feature_val.npy'), fea_val)
np.save(os.path.join(save_dir + '_label_val.npy'), label_val)


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

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

# 加载模型
model = load_model('carplate_CNN_model.h5')

# 加载一张图片
img = cv2.imread('./13-4.jpg')    # 一张 20*20 的图片
# 转换模型所需要的输入格式 单通道灰度  float64 型
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) / 255.0

# 模型需要输入一个四维的数据，所以要把二维的 20*20 的图片  转换为 1*1*20*20 的格式输入模型以预测
temp: np.ndarray = np.ndarray(shape=(1, 1, 20, 20), dtype='float64')
temp[0][0] = img
ret: np.ndarray = model.predict(temp)
# ret 是一个 shape 为(1, 65) 的 ndarray, 不为 0 的那个位置即为类型序号

# 标签位置顺序
# ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
# 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z',
# '云', '京', '冀', '吉', '宁', '川', '新', '晋', '桂', '沪', '津', '浙', '渝', '湘', '琼', '甘',
# '皖','粤', '苏', '蒙', '藏', '豫', '贵', '赣', '辽', '鄂', '闽', '陕', '青', '鲁', '黑']

# 在同个位置找到标签输出即可
(_, index) = np.where(ret == ret.max())
import label_define   # 训练时的标签定义

result = label_define.car_label[index[0]]    # result: str
print(result)


from typing import Any, Dict, List, Optional, Type

import numpy as np
import torch
import torch.nn.functional as F

from tianshou.data import Batch, ReplayBuffer, to_numpy, to_torch
from tianshou.policy import PPOPolicy

class GAILPolicy(PPOPolicy):
    r"""Implementation of Generative Adversarial Imitation Learning. arXiv:1606.03476.

    :param torch.nn.Module actor: the actor network following the rules in
        :class:`~tianshou.policy.BasePolicy`. (s -> logits)
    :param torch.nn.Module critic: the critic network. (s -> V(s))
    :param torch.optim.Optimizer optim: the optimizer for actor and critic network.
    :param dist_fn: distribution class for computing the action.
    :type dist_fn: Type[torch.distributions.Distribution]
    :param ReplayBuffer expert_buffer: the replay buffer contains expert experience.
    :param torch.nn.Module disc_net: the discriminator network with input dim equals
        state dim plus action dim and output dim equals 1.
    :param torch.optim.Optimizer disc_optim: the optimizer for the discriminator
        network.
    :param int disc_update_num: the number of discriminator grad steps per model grad
        step. Default to 4.
    :param float discount_factor: in [0, 1]. Default to 0.99.
    :param float eps_clip: :math:`\epsilon` in :math:`L_{CLIP}` in the original
        paper. Default to 0.2.
    :param float dual_clip: a parameter c mentioned in arXiv:1912.09729 Equ. 5,
        where c > 1 is a constant indicating the lower bound.
        Default to 5.0 (set None if you do not want to use it).
    :param bool value_clip: a parameter mentioned in arXiv:1811.02553 Sec. 4.1.
        Default to True.
    :param bool advantage_normalization: whether to do per mini-batch advantage
        normalization. Default to True.
    :param bool recompute_advantage: whether to recompute advantage every update
        repeat according to https://arxiv.org/pdf/2006.05990.pdf Sec. 3.5.
        Default to False.
    :param float vf_coef: weight for value loss. Default to 0.5.
    :param float ent_coef: weight for entropy loss. Default to 0.01.
    :param float max_grad_norm: clipping gradients in back propagation. Default to
        None.
    :param float gae_lambda: in [0, 1], param for Generalized Advantage Estimation.
        Default to 0.95.
    :param bool reward_normalization: normalize estimated values to have std close
        to 1, also normalize the advantage to Normal(0, 1). Default to False.
    :param int max_batchsize: the maximum size of the batch when computing GAE,
        depends on the size of available memory and the memory cost of the model;
        should be as large as possible within the memory constraint. Default to 256.
    :param bool action_scaling: whether to map actions from range [-1, 1] to range
        [action_spaces.low, action_spaces.high]. Default to True.
    :param str action_bound_method: method to bound action to range [-1, 1], can be
        either "clip" (for simply clipping the action), "tanh" (for applying tanh
        squashing) for now, or empty string for no bounding. Default to "clip".
    :param Optional[gym.Space] action_space: env's action space, mandatory if you want
        to use option "action_scaling" or "action_bound_method". Default to None.
    :param lr_scheduler: a learning rate scheduler that adjusts the learning rate in
        optimizer in each policy.update(). Default to None (no lr_scheduler).
    :param bool deterministic_eval: whether to use deterministic action instead of
        stochastic action sampled by the policy. Default to False.
    :param lr_scheduler: a learning rate scheduler that adjusts the learning rate in
        optimizer in each policy.update(). Default to None (no lr_scheduler).

    .. seealso::

        Please refer to :class:`~tianshou.policy.PPOPolicy` for more detailed
        explanation.
    """

    def __init__(
        self,
        actor: torch.nn.Module,
        critic: torch.nn.Module,
        optim: torch.optim.Optimizer,
        dist_fn: Type[torch.distributions.Distribution],
        expert_buffer: ReplayBuffer,
        disc_net: torch.nn.Module,
        disc_optim: torch.optim.Optimizer,
        disc_update_num: int = 4,
        eps_clip: float = 0.2,
        dual_clip: Optional[float] = None,
        value_clip: bool = False,
        advantage_normalization: bool = True,
        recompute_advantage: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            actor, critic, optim, dist_fn, eps_clip, dual_clip, value_clip,
            advantage_normalization, recompute_advantage, **kwargs
        )
        self.disc_net = disc_net
        self.disc_optim = disc_optim
        self.disc_update_num = disc_update_num
        self.expert_buffer = expert_buffer
        self.action_dim = actor.output_dim

    def process_fn(
        self, batch: Batch, buffer: ReplayBuffer, indices: np.ndarray
    ) -> Batch:
        """Pre-process the data from the provided replay buffer.

        Used in :meth:`update`. Check out :ref:`process_fn` for more information.
        """
        # update reward
        with torch.no_grad():
            batch.rew = to_numpy(-F.logsigmoid(-self.disc(batch)).flatten())
        return super().process_fn(batch, buffer, indices)

    def disc(self, batch: Batch) -> torch.Tensor:
        obs = to_torch(batch.obs, device=self.disc_net.device)  # type: ignore
        act = to_torch(batch.act, device=self.disc_net.device)  # type: ignore
        return self.disc_net(torch.cat([obs, act], dim=1))  # type: ignore

    def learn(  # type: ignore
        self, batch: Batch, batch_size: int, repeat: int, **kwargs: Any
    ) -> Dict[str, List[float]]:
        # update discriminator
        losses = []
        acc_pis = []
        acc_exps = []
        bsz = len(batch) // self.disc_update_num
        for b in batch.split(bsz, merge_last=True):
            logits_pi = self.disc(b)
            exp_b = self.expert_buffer.sample(bsz)[0]
            logits_exp = self.disc(exp_b)
            loss_pi = -F.logsigmoid(-logits_pi).mean()
            loss_exp = -F.logsigmoid(logits_exp).mean()
            loss_disc = loss_pi + loss_exp
            self.disc_optim.zero_grad()
            loss_disc.backward()
            self.disc_optim.step()
            losses.append(loss_disc.item())
            acc_pis.append((logits_pi < 0).float().mean().item())
            acc_exps.append((logits_exp > 0).float().mean().item())
        # update policy
        res = super().learn(batch, batch_size, repeat, **kwargs)
        res["loss/disc"] = losses
        res["stats/acc_pi"] = acc_pis
        res["stats/acc_exp"] = acc_exps
        return res


import random
import shutil
import tempfile
from contextlib import contextmanager
from pathlib import Path

import cv2
import numpy as np
import pyrender
import torch
import trimesh
from FLAME_PyTorch.FLAME import FLAME
from misc.shared import BASE_DIR, CONFIG, flame_config
from tqdm import tqdm

class RenderContext(object):
    def __init__(self, z=0, x=0, y=0, width=1024, light_color=None, f=None):
        if light_color is None:
            light_color = np.array([1.0, 1.0, 1.0])

        if f is None:
            f = np.array([4754.97941935 / 2, 4754.97941935 / 2])

        self.mesh = None
        frustum = {"near": 0.01, "far": 100.0, "height": 1024, "width": width}
        camera_params = {
            "c": np.array([x, y]),
            "k": np.array([-0.19816071, 0.92822711, 0, 0, 0]),
            "f": f,
        }
        intensity = 1.5
        self.rgb_per_v = None

        self.scene = pyrender.Scene(
            ambient_light=[0.2, 0.2, 0.2], bg_color=[255, 255, 255]
        )
        camera = pyrender.IntrinsicsCamera(
            fx=camera_params["f"][0],
            fy=camera_params["f"][1],
            cx=camera_params["c"][0],
            cy=camera_params["c"][1],
            znear=frustum["near"],
            zfar=frustum["far"],
        )

        camera_pose = np.eye(4)
        camera_pose[:3, 3] = np.array([0, 0, 1.0 - z])
        self.scene.add(camera, pose=camera_pose)

        angle = np.pi / 6.0
        pos = [0, 0, 1]

        light = pyrender.PointLight(color=light_color, intensity=intensity)

        light_pose = np.eye(4)
        light_pose[:3, 3] = pos
        self.scene.add(light, pose=light_pose.copy())

        light_pose[:3, 3] = cv2.Rodrigues(np.array([angle, 0, 0]))[0].dot(pos)
        self.scene.add(light, pose=light_pose.copy())

        light_pose[:3, 3] = cv2.Rodrigues(np.array([-angle, 0, 0]))[0].dot(pos)
        self.scene.add(light, pose=light_pose.copy())

        light_pose[:3, 3] = cv2.Rodrigues(np.array([0, -angle, 0]))[0].dot(pos)
        self.scene.add(light, pose=light_pose.copy())

        light_pose[:3, 3] = cv2.Rodrigues(np.array([0, angle, 0]))[0].dot(pos)
        self.scene.add(light, pose=light_pose.copy())

        self.r = pyrender.OffscreenRenderer(
            viewport_width=frustum["width"], viewport_height=frustum["height"]
        )

    def render_meshes(self, meshes):
        scene_meshes = []
        for mesh in meshes:
            render_mesh = pyrender.Mesh.from_trimesh(mesh)
            scene_meshes.append(self.scene.add(render_mesh))  # , pose=np.eye(4)

        flags = pyrender.RenderFlags.SKIP_CULL_FACES
        # flags = None
        img, _ = self.r.render(self.scene)  # , flags=flags
        for scene_mesh in scene_meshes:
            self.scene.remove_node(scene_mesh)
        return img[..., ::-1]

@contextmanager
def create_temp_obj(texture_dir, skin_color):
    with tempfile.TemporaryDirectory() as tmpd:
        obj_file = Path(tmpd) / "file.obj"

        shutil.copy(
            texture_dir / "base_model.mtl", Path(tmpd) / "file.mtl",
        )
        shutil.copy(
            texture_dir / f"texture_{skin_color}.png", Path(tmpd) / "texture.png",
        )
        extra_info = (texture_dir / "base_model.partial_obj").read_text()

        def get_obj(vertices):
            with open(obj_file, "w") as tmpf:

                tmpf.write("mtllib file.mtl\n")
                for vertex in vertices:
                    tmpf.write("v " + " ".join(map(str, vertex.tolist())) + "\n")

                tmpf.write(extra_info)

            return obj_file

        yield get_obj

def render_double_face_video(
    file_name,
    vertices,
    vertices2,
    fps=50,
    skin_color_v1=None,
    skin_color_v2=None,
    width=2048,
):
    rc = RenderContext(
        x=width // 2,
        y=400,
        width=width,
        z=-1,
        f=np.array([4754.97941935, 4754.97941935]),
    )

    writer = cv2.VideoWriter(
        file_name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, 1024)
    )

    random_skin = lambda: random.choice(["black", "white"])

    cto1 = create_temp_obj(
        BASE_DIR / "code/visualize/texture", skin_color_v1 or random_skin()
    )
    cto2 = create_temp_obj(
        BASE_DIR / "code/visualize/texture", skin_color_v2 or random_skin()
    )

    with cto1 as obj_file1, cto2 as obj_file2:

        for vertice, vertice2 in tqdm(list(zip(vertices, vertices2)), leave=False):
            face_width = 0.1

            vertice[:, 0] -= face_width * 2
            vertice2[:, 0] += face_width * 2

            tri_mesh1 = trimesh.load_mesh(
                file_obj=open(obj_file1(vertice)), file_type="obj"
            )
            tri_mesh2 = trimesh.load_mesh(
                file_obj=open(obj_file2(vertice2)), file_type="obj"
            )

            color = rc.render_meshes([tri_mesh1, tri_mesh2])
            writer.write(color)

        writer.release()

def random_shape(seq_len):
    shape = torch.zeros((1, 300))
    shape[:, :100] = torch.rand(100)
    return shape.repeat(seq_len, 1)

def get_vertices(expression, pose, rotation, eyes=None, shape=None, gender="generic"):
    seq_len = expression.shape[0]

    if shape is None:
        shape = random_shape(seq_len).type_as(expression)

    if eyes is None:
        eyes = torch.zeros((seq_len, 6)).type_as(expression)

    flame = FLAME(
        flame_config(
            BASE_DIR / CONFIG["flame"][f"model_path_{gender}"],
            BASE_DIR / CONFIG["flame"]["static_landmark_embedding_path"],
            BASE_DIR / CONFIG["flame"]["dynamic_landmark_embedding_path"],
            seq_len,
            True,
            shape.shape[1],
            expression.shape[1],
            True,
        )
    )

    zero_padding = torch.zeros((seq_len, 3)).type_as(expression)
    pose_params = torch.cat([zero_padding, pose[:, 3:6]], dim=1)

    neck = pose[:, :3] + rotation

    vertices, _ = flame(
        shape_params=shape,
        expression_params=expression,
        pose_params=pose_params,
        neck_pose=neck,
        eye_pose=eyes,
    )
    return vertices.cpu().numpy()


import numpy as np
from deep_utils.utils.box_utils.boxes import Box
import torch
from models.experimental import attempt_load
from utils.general import xywh2xyxy
from utils.datasets import letterbox
import cv2
import time
import torchvision
import torch.nn as nn
from utils.metrics import box_iou

class YOLOV5TorchObjectDetector(nn.Module):
    def __init__(self,
                 model_weight,
                 device,
                 img_size,
                 names=None,
                 mode='eval',
                 confidence=0.4,
                 iou_thresh=0.45,
                 agnostic_nms=False):
        super(YOLOV5TorchObjectDetector, self).__init__()
        self.device = device
        self.model = None
        self.img_size = img_size
        self.mode = mode
        self.confidence = confidence
        self.iou_thresh = iou_thresh
        self.agnostic = agnostic_nms
        self.model = attempt_load(model_weight, map_location=device)
        print("[INFO] Model is loaded")
        self.model.requires_grad_(True)
        self.model.to(device)
        if self.mode == 'train':
            self.model.train()
        else:
            self.model.eval()
        # fetch the names
        if names is None:
            print('[INFO] fetching names from coco file')
            self.names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat',
                          'traffic light',
                          'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
                          'cow',
                          'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase',
                          'frisbee',
                          'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
                          'surfboard',
                          'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana',
                          'apple',
                          'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
                          'couch',
                          'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
                          'keyboard', 'cell phone',
                          'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
                          'teddy bear',
                          'hair drier', 'toothbrush']
        else:
            self.names = names

        # preventing cold start
        img = torch.zeros((1, 3, *self.img_size), device=device)
        self.model(img)

    @staticmethod
    def non_max_suppression(prediction, logits, conf_thres=0.6, iou_thres=0.45, classes=None, agnostic=False,
                            multi_label=False, labels=(), max_det=300):
        """Runs Non-Maximum Suppression (NMS) on inference and logits results

        Returns:
             list of detections, on (n,6) tensor per image [xyxy, conf, cls] and pruned input logits (n, number-classes)
        """

        nc = prediction.shape[2] - 5  # number of classes
        xc = prediction[..., 4] > conf_thres  # candidates

        # Checks
        assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
        assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'

        # Settings
        min_wh, max_wh = 2, 4096  # (pixels) minimum and maximum box width and height
        max_nms = 30000  # maximum number of boxes into torchvision.ops.nms()
        time_limit = 10.0  # seconds to quit after
        redundant = True  # require redundant detections
        multi_label &= nc > 1  # multiple labels per box (adds 0.5ms/img)
        merge = False  # use merge-NMS

        t = time.time()
        output = [torch.zeros((0, 6), device=prediction.device)] * prediction.shape[0]
        logits_output = [torch.zeros((0, 80), device=logits.device)] * logits.shape[0]
        for xi, (x, log_) in enumerate(zip(prediction, logits)):  # image index, image inference
            # Apply constraints
            # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
            x = x[xc[xi]]  # confidence
            log_ = log_[xc[xi]]
            # Cat apriori labels if autolabelling
            if labels and len(labels[xi]):
                l = labels[xi]
                v = torch.zeros((len(l), nc + 5), device=x.device)
                v[:, :4] = l[:, 1:5]  # box
                v[:, 4] = 1.0  # conf
                v[range(len(l)), l[:, 0].long() + 5] = 1.0  # cls
                x = torch.cat((x, v), 0)

            # If none remain process next image
            if not x.shape[0]:
                continue

            # Compute conf
            x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf
            # log_ *= x[:, 4:5]
            # Box (center x, center y, width, height) to (x1, y1, x2, y2)
            box = xywh2xyxy(x[:, :4])

            # Detections matrix nx6 (xyxy, conf, cls)
            if multi_label:
                i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
                x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
            else:  # best class only
                conf, j = x[:, 5:].max(1, keepdim=True)
                # log_ = x[:, 5:]
                x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
                log_ = log_[conf.view(-1) > conf_thres]
            # Filter by class
            if classes is not None:
                x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]

            # Check shape
            n = x.shape[0]  # number of boxes
            if not n:  # no boxes
                continue
            elif n > max_nms:  # excess boxes
                x = x[x[:, 4].argsort(descending=True)[:max_nms]]  # sort by confidence

            # Batched NMS
            c = x[:, 5:6] * (0 if agnostic else max_wh)  # classes
            boxes, scores = x[:, :4] + c, x[:, 4]  # boxes (offset by class), scores
            i = torchvision.ops.nms(boxes, scores, iou_thres)  # NMS
            if i.shape[0] > max_det:  # limit detections
                i = i[:max_det]
            if merge and (1 < n < 3E3):  # Merge NMS (boxes merged using weighted mean)
                # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
                iou = box_iou(boxes[i], boxes) > iou_thres  # iou matrix
                weights = iou * scores[None]  # box weights
                x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  # merged boxes
                if redundant:
                    i = i[iou.sum(1) > 1]  # require redundancy

            output[xi] = x[i]
            logits_output[xi] = log_[i]
            assert log_[i].shape[0] == x[i].shape[0]
            if (time.time() - t) > time_limit:
                print(f'WARNING: NMS time limit {time_limit}s exceeded')
                break  # time limit exceeded

        return output, logits_output

    @staticmethod
    def yolo_resize(img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):

        return letterbox(img, new_shape=new_shape, color=color, auto=auto, scaleFill=scaleFill, scaleup=scaleup)

    def forward(self, img):
        prediction, logits, _ = self.model(img, augment=False)
        prediction, logits = self.non_max_suppression(prediction, logits, self.confidence, self.iou_thresh,
                                                      classes=None,
                                                      agnostic=self.agnostic)
        self.boxes, self.class_names, self.classes, self.confidences = [[[] for _ in range(img.shape[0])] for _ in
                                                                        range(4)]
        for i, det in enumerate(prediction):  # detections per image
            if len(det):
                for *xyxy, conf, cls in det:
                    bbox = Box.box2box(xyxy,
                                       in_source=Box.BoxSource.Torch,
                                       to_source=Box.BoxSource.Numpy,
                                       return_int=True)
                    self.boxes[i].append(bbox)
                    self.confidences[i].append(round(conf.item(), 2))
                    cls = int(cls.item())
                    self.classes[i].append(cls)
                    if self.names is not None:
                        self.class_names[i].append(self.names[cls])
                    else:
                        self.class_names[i].append(cls)
        return [self.boxes, self.classes, self.class_names, self.confidences], logits

    def preprocessing(self, img):
        if len(img.shape) != 4:
            img = np.expand_dims(img, axis=0)
        im0 = img.astype(np.uint8)
        img = np.array([self.yolo_resize(im, new_shape=self.img_size)[0] for im in im0])
        img = img.transpose((0, 3, 1, 2))
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img / 255.0
        return img

if __name__ == '__main__':
    model_path = 'runs/train/cart-detection/weights/best.pt'
    img_path = './16_4322071600_101_0_4160379257.jpg'
    model = YOLOV5TorchObjectDetector(model_path, 'cpu', img_size=(640, 640)).to('cpu')
    img = np.expand_dims(cv2.imread(img_path)[..., ::-1], axis=0)
    img = model.preprocessing(img)
    a = model(img)
    print(model._modules)


import cv2
import numpy as np
import matplotlib.path as mplPath
import time
from datetime import datetime
import os
import threading

import requests
# --------------------- GLOBAL --------------- #
poly_ne = np.array([[793,351],[920,466],[1102,420],[961,329]])
tl = (1106,100)
br = (1153,203)
cap = cv2.VideoCapture("test_video.mp4")
img_folder = "./image/"
# Plate Recognizer
regions = ['vn']
IS_FREE = True
ImageNode = []
#----------------------------------------------#

#---------------------- UTILS FUNCTION ---------------#
def pega_centro(x, y, w, h):
    x1 = int(w / 2)
    y1 = int(h / 2)
    cx = x + x1
    cy = y + y1
    return cx,cy
def inside_poly(pnts, poly):
    bbPath = mplPath.Path(poly)
    return bbPath.contains_point(pnts)

def detect_plate(path):
    with open(path, 'rb') as fp:
        response = requests.post(
            'https://api.platerecognizer.com/v1/plate-reader/',
            data=dict(regions=regions),  # Optional
            files=dict(upload=fp),
            headers={'Authorization': 'Token 11e12906b07074d87046cb0e3b530f55dd85a2c1'})
        result = response.json()['results']
        IS_FREE = True
        print(result)
#------------------------------------------------------#

ret, frame1 = cap.read()
ret, frame2 = cap.read()
while cap.isOpened():
    #Detect Red-light
    if frame1 is None or frame2 is None:
        break
    hsv_Frame = cv2.cvtColor(frame1[tl[1]:br[1], tl[0]:br[0]], cv2.COLOR_BGR2HSV)
    low_red = np.array([161, 155, 84], np.uint8)
    high_red = np.array([179, 255, 255], np.uint8)
    mask_red = cv2.inRange(hsv_Frame, low_red, high_red)
    coord=cv2.findNonZero(mask_red)
    #Detect moving
    try:
    	diff = cv2.absdiff(frame1, frame2)
    except:
    	print("error")
    gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (5,5), 0)
    _, thresh = cv2.threshold(blur, 20, 255, cv2.THRESH_BINARY)
    dilated = cv2.dilate(thresh, None, iterations=3)
    contours, _ = cv2.findContours(dilated, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    if not(coord is None):
        poly = cv2.polylines(frame1, [poly_ne] ,True, (0, 0, 255),3)
        cv2.putText(frame1, "Status: {}".format('Den_do'), (10, 20), cv2.FONT_HERSHEY_SIMPLEX,
                    1, (0, 0, 255), 3)
        for contour in contours:
            (x, y, w, h) = cv2.boundingRect(contour)
            if cv2.contourArea(contour) < 10000:
                continue
            centro = pega_centro(x, y, w, h)
            x1, y1, w1, h1 = x, y, w, h
            circle_o = cv2.circle(frame1, centro, 4, (0, 0, 255), -1)
            rectangel = cv2.rectangle(frame1, (x, y), (x + w, y + h), (0, 255, 0), 2)
            if inside_poly(centro, poly_ne):

                cv2.imshow("crop", frame2[y1: y1 + h1, x1: x1 + w1])

                path = "%s%s_%s\\" % (img_folder, datetime.today().strftime("%Y_%m_%d"),int(time.time()))

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

                #cv2.imwrite(path + "overview.png", frame1)

                cv2.imwrite(path + "detailed.png", frame2[y1: y1 + h1, x1: x1 + w1])
                #detect_plate(path + "detailed.png")
                ImageNode.append(path+"detailed.png")
    else:
    	cv2.putText(frame1, "Status: {}".format('Den_xanh'), (10, 20), cv2.FONT_HERSHEY_SIMPLEX,
                    1, (0, 255, 0), 3)
    print(ImageNode)
    if IS_FREE and len(ImageNode) > 0:
        IS_FREE = False
        thread = threading.Thread(target=detect_plate, args=(ImageNode.pop(0),))
        thread.start()
    cv2.imshow("feed", frame1)
    #cv2.imwrite("image.png", frame1)
    frame1 = frame2
    _, frame2 = cap.read()

    if cv2.waitKey(33) & 0xFF == ord("q"):
        break

cv2.destroyAllWindows()
cap.release()


import tensorflow as tf
from datetime import datetime
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import gym
import argparse
from onail.ONAIL import ONAIL
from argparse import Namespace
from configs.offline_il_configs import OFFLINE_IL_CONFIGS
from sys import argv
import numpy as np
from common.utlis import boolean_string
import random
from gym.envs.registration import register
register(id='AntNoEarlyTerm-v3', entry_point='gym.envs.mujoco.ant_v3:AntEnv',
         kwargs={'terminate_when_unhealthy': False}, max_episode_steps=1000)
register(id='HopperNoEarlyTerm-v3', entry_point='gym.envs.mujoco.hopper_v3:HopperEnv',
         kwargs={'terminate_when_unhealthy': False}, max_episode_steps=1000)
register(id='WalkerNoEarlyTerm-v3', entry_point='gym.envs.mujoco.walker2d_v3:Walker2dEnv',
         kwargs={'terminate_when_unhealthy': False}, max_episode_steps=1000)

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    tf.random.set_seed(seed)

def argparser():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--mle_params', action="append",
                        help="path to npz file storing initial policy params. Parameter can be repeated when learning"
                             "in multiple envinronments.", type=str)
    parser.add_argument('--path', help="path for logging", type=str)
    parser.add_argument('--seed', help="random seed", type=int, default=0)
    parser.add_argument('--max_iter', help="number of learning iterations", type=int, default=int(1e7))
    parser.add_argument('--bc_only', help='dont do any learning after behavioral cloning',
                        type=boolean_string, default=False)
    parser.add_argument('--config_name',
                        help='If not none, unspecified arguments will use the specified preset. The name should correspond'
                             ' to a key in configs/rl_configs.py or (if --il_algo is set) configs/offline_il_configs.py',
                        type=str, default="HALFCHEETAH_ONAIL")
    parser.add_argument('--env_name', action='append', type=str,
                        help='name of the environment. This parameter can be repeated to learn in several environments '
                             'with a shared reward function. The parameter --il_demo_path must be provided as many times'
                             'as --env-name to provide demonstrations for each environment')
    parser.add_argument('--policy_layers', help='list of widths for the policy network', type=int, nargs='+',
                        required=False)
    parser.add_argument('--policy_squash_actions', help='if true the actions are squashed using tanh',
                        type=boolean_string, default=True)
    parser.add_argument('--policy_state_dependent_std', help='should the standard deviation of the Gaussian depend '
                                                              'on the state?', type=boolean_string, default=True)

    parser.add_argument('--il_all_states_are_init_states', help="treat every expert state as an initial state",
                        type=boolean_string, default=True)
    parser.add_argument('--il_vali_steps', help='number of samples used to compute vali/test loss',
                        type=int, default=0)
    parser.add_argument('--il_initialize_bc', help='initilize policy by maximizing likelihood of demonstrations',
                        type=boolean_string, default=True)
    parser.add_argument('--il_bc_vali_split', help='ratio for validation data for behavioral cloning',
                        type=float, default=0.1)
    parser.add_argument('--il_bc_iter', help='maximum number of iterations for behavioral cloning',
                        type=int, default=int(1e7))
    parser.add_argument('--il_bc_learning_rate', help='learning rate for behavioral cloning',
                        type=float, default=1e-4)
    parser.add_argument('--il_bc_policy_reg_coef', help='policy regularization for behavioral cloning',
                        type=float, default=1e-5)

    parser.add_argument('--il_kl_factor', help='weight of the KL wenn optimizing policy w.r.t. Q function',
                        type=float, default=0.0)
    parser.add_argument('--il_gamma', help='discount factor', type=float, default=0.99)
    parser.add_argument('--il_batch_size', help='batch_size for expert and replay samples', type=int, default=256)
    parser.add_argument('--Q_layers', help='list of widths for the nu network', type=int, nargs='+', required=False)
    parser.add_argument('--Q_l2_reg', help='l2 regularization for nu network', type=float, default=0.)
    parser.add_argument('--il_Q_learning_rate', help='learning rate for nu network', type=float, default=1e-3)
    parser.add_argument('--il_nu_gradient_penalty', help='gradient penalty for nu network', type=float, default=10.)
    parser.add_argument('--il_policy_learning_rate', help='learning rate for policy network', type=float, default=1e-5)
    parser.add_argument('--il_max_demos', help='maximum number of expert trajectories to use', type=int, default=int(1e8))
    parser.add_argument('--il_demo_path', action='append',
                        help='path to npz storing expert trajs. This parameter can be repeated when learning in several '
                             'environments. The parameter --il_demo_path must be provided as many times'
                             'as --env-name to provide demonstrations for each environment',
                        type=str)
    parser.add_argument('--il_ent_coef', help='coefficient for policy entropy regularization', type=float, default=0.)
    parser.add_argument('--il_standardize_states', help='should we standardize the states?', type=boolean_string, default=True)
    parser.add_argument('--il_num_nu_steps',
                        help='number of critic gradient steps before updating policy', type=int, default=1)
    parser.add_argument('--il_num_policy_steps',
                        help='number of policy gradient steps before updating critic', type=int, default=1)

    parser.add_argument('--il_use_donsker_varadhan_for_pi_update',
                        help="If False, the policy is updated to maximize the Q-Function",
                        type=boolean_string, default=False)
    parser.add_argument('--il_use_fgan_loss', help="Use F-GAN loss instead of Donsker-Varadhan",
                        type=boolean_string, default=False)
    parser.add_argument('--il_use_lower_bound_reward_for_nu',
                        help="Learn the Q-Function for the lower bound reward (not for the density ratio)",
                        type=boolean_string, default=False)
    return parser

def overwrite_dict_with_kwargs(prefix, current_dict, args, ignore_keys=[""]):
    '''
        extracts all commandline parameters that start with a given prefix (e.g. "Q_", "il_", etc.)
        and stores the corresponding values in the given dictionary under the key that is obtained when removing
        the prefix. Hence args = {"trpo_entcoeff" : 1e-2} would add (or overwrite) {"entcoeff" : 1e-2} to current_dict
    '''
    arg_dict = args.__dict__
    keys = list(arg_dict.keys())
    params = [key.startswith(prefix) for key in keys]
    param_keys = np.array(keys)[np.where(params)]
    for param_key in param_keys:
        if param_key not in ignore_keys:
            current_dict[param_key[len(prefix):]] = arg_dict[param_key]

def build_kwargs_from_command_line_args(args):
    offpolicy_il_kwargs = dict()
    policy_kwargs = dict()
    Q_kwargs = dict()
    overwrite_dict_with_kwargs("il_", offpolicy_il_kwargs, args)
    overwrite_dict_with_kwargs("policy_", policy_kwargs, args)
    overwrite_dict_with_kwargs("Q_", Q_kwargs, args)
    return offpolicy_il_kwargs, policy_kwargs, Q_kwargs,

def get_final_experiment_config():
    def_args = argparser().parse_args()
    if def_args.config_name is not None:
        conf_from_file = OFFLINE_IL_CONFIGS[def_args.config_name]
        final_conf = def_args.__dict__
        for key in conf_from_file.keys():
            # We only want to override those argparse parameters that have not been explicitly specified (defaults)
            # Explicitly specified command line parameters have higher priority than those from the config file
            if "--" + key not in argv:
                final_conf[key] = conf_from_file[key]
        return Namespace(**final_conf)
    else:
        return def_args

def run_experiment(args):
    start_time = datetime.now().strftime('%Y-%m-%d%_H:%M:%S.%f')
    if args.path is None:
        path = os.path.join("/tmp", start_time, "")
    else:
        path = os.path.join(args.path, start_time, "")

    offpolicy_il_kwargs, policy_kwargs, Q_kwargs,= build_kwargs_from_command_line_args(args)
    set_seed(args.seed)

    this_demo_path = offpolicy_il_kwargs['demo_path'][0]
    this_env = gym.make(args.env_name[0])
    this_eval_env = gym.make(args.env_name[0])
    this_path = os.path.join(path, "_"+args.env_name[0], "")
    this_offpolicy_il_kwargs = dict(**offpolicy_il_kwargs)
    this_offpolicy_il_kwargs.update(
            dict(name="il_{}".format(0), policy_kwargs=policy_kwargs, Q_kwargs=Q_kwargs,
                 log_path=this_path, demo_path=this_demo_path, env=this_env, eval_env=this_eval_env))

    il_algorithm = ONAIL(**this_offpolicy_il_kwargs)

    for iters in range(args.max_iter):
      il_algorithm.learn_one_step()

if __name__ == "__main__":
    args = get_final_experiment_config()
    run_experiment(args)


## This code is from https://github.com/pytorch/examples/blob/master/mnist/main.py
## with modification to do training using onnxruntime as backend on cuda device.
## A private PyTorch build from https://aiinfra.visualstudio.com/Lotus/_git/pytorch (ORTTraining branch) is needed to run the demo.
## To run the demo with ORT backend:
##      python mnist_training.py --use-ort
## or
##      python mnist_training.py --use-ort --use-ort-trainer

## When "--use-ort" is not given, it will run training with PyTorch as backend.
## Model testing is not complete.

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import os

# TODO: remove after ready for CV
# import sys
# sys.path.insert(0, '/bert_ort/liqun/onnxruntime/build/Linux/Debug/')
# import onnxruntime as ort

# sys.path.insert(0, '/bert_ort/liqun/onnxruntime/onnxruntime/python/')
# from ort_trainer import IODescription, ModelDescription, ORTTrainer, ORTModel

from onnxruntime.capi.ort_trainer import IODescription, ModelDescription, ORTTrainer, ORTModel
from mpi4py import MPI
from onnxruntime.capi._pybind_state import set_cuda_device_id

class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

def my_loss(x, target):
    return F.nll_loss(F.log_softmax(x, dim=1), target)

def train_with_model(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        data = data.reshape(data.shape[0], -1)
        loss, pred = model.run(data, target)
        if batch_idx % args.log_interval == 0:
            optimizer.step()
            optimizer.zero_grad()
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def test_with_model(args, model, device, test_loader, optimizer, epoch):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            data = data.reshape(data.shape[0], -1)

            pred = model.run(data, target, )

            output = F.log_softmax(model.eval(data), dim=1)
            test_loss += F.nll_loss(output, target, reduction='sum').item()     # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)   # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

def train_with_trainer(args, trainer, device, train_loader, epoch):
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        data = data.reshape(data.shape[0], -1)

        learning_rate = torch.tensor([args.lr])
        loss = trainer.train_step(data, target, learning_rate)

        # Since the output corresponds to [loss_desc, probability_desc], the first value is taken as loss.
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss[0]))

# TODO: comple this once ORT training can do evaluation.
def test_with_trainer(args, trainer, device, test_loader):
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            data = data.reshape(data.shape[0], -1)
            output = F.log_softmax(trainer.eval_step(data, fetches=['probability']), dim=1)
            test_loss += F.nll_loss(output, target, reduction='sum').item()     # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)   # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)

    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

def mnist_model_description():
    input_desc = IODescription('input1', ['batch', 784], torch.float32)
    label_desc = IODescription('label', ['batch', ], torch.int64, num_classes=10)
    loss_desc = IODescription('loss', [], torch.float32)
    probability_desc = IODescription('probability', ['batch', 10], torch.float32)
    return ModelDescription([input_desc, label_desc], [loss_desc, probability_desc])

def main():
#Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
                        help='input batch size for testing (default: 1000)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--no-cuda', action='store_true', default=False,
                        help='disables CUDA training')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')

    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')

    parser.add_argument('--use-ort', action='store_true', default=False,
                        help='to use onnxruntime as training backend')

    parser.add_argument('--use-ort-trainer', action='store_true', default=False,
                        help='to use onnxruntime as training backend')

    args = parser.parse_args()
    use_cuda = not args.no_cuda and torch.cuda.is_available()

    torch.manual_seed(args.seed)

    kwargs = {'num_workers': 0, 'pin_memory': True}
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=True, download=True,
                       transform=transforms.Compose([
                           transforms.ToTensor(),
                           transforms.Normalize((0.1307,), (0.3081,))
                       ])),
        batch_size=args.batch_size, shuffle=True, **kwargs)
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST('../data', train=False, transform=transforms.Compose([
            transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])),
        batch_size=args.test_batch_size, shuffle=True, **kwargs)

    comm = MPI.COMM_WORLD
    args.local_rank = int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK']) if ('OMPI_COMM_WORLD_LOCAL_RANK' in os.environ) else 0
    args.world_rank = int(os.environ['OMPI_COMM_WORLD_RANK']) if ('OMPI_COMM_WORLD_RANK' in os.environ) else 0
    args.world_size=comm.Get_size()
    torch.cuda.set_device(args.local_rank)
    if use_cuda:
        device = torch.device("cuda", args.local_rank)
    else:
        device = torch.device("cpu")
    args.n_gpu = 1
    set_cuda_device_id(args.local_rank)

    input_size = 784
    hidden_size = 500
    num_classes = 10
    model = NeuralNet(input_size, hidden_size, num_classes)

    model_desc = mnist_model_description()
    if args.use_ort_trainer:
        # use log_interval as gradient accumulate steps
        trainer = ORTTrainer(model, my_loss, model_desc, "LambOptimizer", None, IODescription('Learning_Rate', [1,], torch.float32), device, 1, None,
        args.world_rank, args.world_size, use_mixed_precision=False, allreduce_post_accumulation = True)
        print('\nBuild ort model done.')

        for epoch in range(1, args.epochs + 1):
            train_with_trainer(args, trainer, device, train_loader, epoch)
            import pdb
            test_with_trainer(args, trainer, device, test_loader)
    else:
        model = ORTModel(model, my_loss, model_desc, device, None, args.world_rank, args.world_size)
        print('\nBuild ort model done.')

        optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

        for epoch in range(1, args.epochs + 1):
            train_with_model(args, model, device, train_loader, optimizer, epoch)
            # test(args, model, device, test_loader)

if __name__ == '__main__':
    main()


# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from torch.autograd import Variable
from copy import deepcopy
from collections import Counter
import numpy as np

from rlpytorch import Model, ActorCritic
from actor_critic_changed import ActorCriticChanged
from trunk import MiniRTSNet
from trunk import ProcessSet, DataProcess,TargetAttention,TowerAttention
#from td_sampler import TD_Sampler

def flattern(x):
    return x.view(x.size(0), -1)

class Model_ActorCritic(Model):
    def __init__(self, args, sampler = None):
        super(Model_ActorCritic, self).__init__(args)
        self._init(args)
        self.sampler = sampler

    def _init(self, args):
        params = args.params
        assert isinstance(params["num_action"], int), "num_action has to be a number. action = " + str(params["num_action"])
        self.params = params
        #self.net = MiniRTSNet(args, output1d=False)

        self.na = params["num_action"]
        self.num_unit = params["num_unit_type"]
        self.num_planes = params["num_planes"]
        self.num_cmd_type = params["num_cmd_type"]
        self.mapx = params["map_x"]
        self.mapy = params["map_y"]

        out_dim = self.num_planes * self.mapx * self.mapy

        # After trunk, we can predict the unit action and value.
        # Four dimensions. unit loc, target loc,
        #self.unit_locs = nn.Conv2d(self.num_planes, 1, 3, padding=1)
        #self.target_locs = nn.Conv2d(self.num_planes, 1, 3, padding=1)

        # test
        #self.unit_select = nn.Linear(self.mapx * self.mapy,20)
        #self.target_select = nn.Linear(self.mapx * self.mapy,20)

        #self.cmd_types = nn.Linear(out_dim, self.num_cmd_type)
        # self.cmd_types = nn.Linear(out_dim, 3)

        #self.build_types = nn.Linear(out_dim, self.num_unit)
        # self.value = nn.Linear(out_dim, 1)
        self.value_test = nn.Linear(256,1)

        # self.relu = nn.LeakyReLU(0.1)
        self.softmax = nn.Softmax()

        # test
        # self.radarProcess = ProcessSet(args,2, 3, 12)
        #self.towerProcess = ProcessSet(args,16, 8, 32)
        # self.enemyProcess = ProcessSet(args,20, 7, 28)
        # self.baseProcess =  ProcessSet(args,1,3,12)
        # self.dataProcess  = DataProcess(args,baseProcess,radarProcess,towerProcess,enemyProcess)
        self.lstm = nn.LSTM(
            input_size=64,  # 输入数据
            hidden_size=256,  # rnn hidden unit
            num_layers=1,  # 有几层 RNN layers
            batch_first=True,  # input & output 会是以 batch size 为第一维度的特征集 e.g. (batch, time_step, input_size)
        )
        self.dataProcess  = DataProcess(args,ProcessSet(args,1,3,30),ProcessSet(args,2, 9, 100),ProcessSet(args,20, 8, 100),ProcessSet(args,10, 10, 100))
        self.round_select = nn.Linear(256,3)
        self.enemy_select = TargetAttention(args,256,100)
        self.tower_select = TowerAttention(args,256,100,100)
        #self.td_sampler = TD_Sampler(0.1) # epsilon

    def get_define_args():
        return MiniRTSNet.get_define_args()

    def forward(self, x):
        #s, res = x["s"], x["res"]
        # output = self.net(self._var(x["s"]), self._var(x["res"]))

        # s = x["s"]
        # output = self.net(self._var(x["s"])) # 8 10 70 70

        #unit_locs = self.softmax(flattern(self.unit_locs(output)))   # 18
        # target_locs = self.softmax(flattern(self.target_locs(output))) # 8
        #unit_locs = self.softmax(self.unit_select(flattern(self.unit_locs(output))) )   # 20
        #target_locs = self.softmax( self.target_select(flattern(self.target_locs(output))) ) # 20

        # flat_output = flattern(output)
        #cmd_types = self.softmax(self.cmd_types(flat_output))
        #build_types = self.softmax(self.build_types(flat_output))
        #value = self.value(flat_output)

        # import pdb
        # pdb.set_trace()
        self.dataProcess.set_volatile(self.volatile)
        base,radar,tower,tower_embedding,enemy,enemy_embedding,global_stat = self.dataProcess(x)
        self.dataProcess.set_volatile(False)
        r_out, (h_n,h_c) = self.lstm(global_stat,None)  # None表示 h_0 和 c_0 置 0
        global_stat = r_out[:, -1, :]

        value_test = self.value_test(global_stat)
        # import pdb
        # pdb.set_trace()
        round_prob = self.softmax(self.round_select(global_stat))
        round_choose =  self.sampler.sample_with_eps(round_prob)  #  攻击方式
        # import pdb
        # pdb.set_trace()
        enemy_attention = self.enemy_select(global_stat,enemy_embedding,x['s_enemy'],True)# 单位注意力
        enemy_prob = self.softmax(enemy_attention)
        enemy_choose =  self.sampler.sample_with_eps(enemy_prob)  #  攻击目标

        # 收集选择到的目标的特征
        batch_size = enemy_embedding.data.shape[0]
        enemy_embedding_list = []
        enemy_mask_list = []

        for i in range(batch_size):
            enemy_embedding_list.append(enemy_embedding[i][enemy_choose[i].data])
            enemy_mask_list.append(x['s_enemy_mask'][i][enemy_choose[i].data])

        enemy_embedding_select = Variable(torch.cat(enemy_embedding_list).data,volatile = self.volatile)
        enemy_embedding_select = enemy_embedding_select.view(batch_size,-1)
        # import pdb
        # pdb.set_trace()
        enemy_mask_select = torch.cat(enemy_mask_list)
        enemy_mask_select = enemy_mask_select.view(batch_size,-1)

        # 根据目标特征获取防御塔的注意力分布和mask
        tower_attention = self.tower_select(global_stat,enemy_embedding_select,enemy_mask_select,tower_embedding,x['s_tower'],True)
        tower_prob = self.softmax(tower_attention)
        # tower_attention = Variable(tower_attention.data)
        tower_choose =  self.sampler.sample_with_eps(tower_prob)  # 防御塔
        # import pdb
        # pdb.set_trace()
        return dict(V=value_test, uloc_prob=tower_prob, tloc_prob=enemy_prob, ct_prob = round_prob,action_type=1),dict(uloc = tower_choose.data,uloc_prob=tower_prob.data,tloc=enemy_choose.data,tloc_prob=enemy_prob.data,ct=round_choose.data,ct_prob = round_prob.data)
        # return dict(V=value_test, uloc_prob=tower_attention, tloc_prob=enemy_attention, ct_prob = round_prob,action_type=1)
        #return dict(V=value, uloc_prob=unit_locs, tloc_prob=target_locs, ct_prob = cmd_types,action_type=1)

# Format: key, [model, method]
# if method is None, fall back to default mapping from key to method
Models = {
    "actor_critic": [Model_ActorCritic, ActorCritic],
}

Defaults = {
    #"sample_nodes": "uloc_prob,uloc;tloc_prob,tloc;ct_prob,ct;bt_prob,bt",
    #"policy_action_nodes": "uloc_prob,uloc;tloc_prob,tloc;ct_prob,ct;bt_prob,bt",
    "sample_nodes": "uloc_prob,uloc;tloc_prob,tloc;ct_prob,ct",
    "policy_action_nodes": "uloc_prob,uloc;tloc_prob,tloc;ct_prob,ct",
    "arch" : "cccc;-,64,64,64,-"
}


import collections
import numpy as np
import time
from pose_engine import PoseEngine
from ..utils.fps_calculator import convert_infr_time_to_fps

class Detector:
    """
    Perform pose estimation with Coral's PoseNet. The bounding boxes will be extracted by key-points.
    for more information about PoseNet model go to: https://github.com/google-coral/project-posenet
    :param config: Is a ConfigEngine instance which provides necessary parameters.
    """

    def __init__(self, config):
        self.config = config
        # Get the model name from the config
        self.model_name = self.config.get_section_dict('Detector')['Name']
        # Frames Per Second
        self.fps = None
        self.w, self.h, _ = [int(i) for i in self.config.get_section_dict('Detector')['ImageSize'].split(',')]
        self.engine = PoseEngine(
            f"/project-posenet/models/mobilenet/posenet_mobilenet_v1_075_{self.h}_{self.w}_quant_decoder_edgetpu.tflite"
        )

        # Get class id from config
        self.class_id = int(self.config.get_section_dict('Detector')['ClassID'])
        self.score_threshold = float(self.config.get_section_dict('Detector')['MinScore'])
        self.keypoints = (
            'nose',
            'left eye',
            'right eye',
            'left ear',
            'right ear',
            'left shoulder',
            'right shoulder',
            'left elbow',
            'right elbow',
            'left wrist',
            'right wrist',
            'left hip',
            'right hip',
            'left knee',
            'right knee',
            'left ankle',
            'right ankle'
        )

    def inference(self, resized_rgb_image):
        """
        This method will perform inference and return the detected bounding boxes
        Args:
            resized_rgb_image: uint8 numpy array with shape (img_height, img_width, channels)

        Returns:
            result: a dictionary contains of [{"id": 0, "bbox": [x1, y1, x2, y2], "score":s%}, {...}, {...}, ...]
        """
        assert resized_rgb_image.shape == (self.h, self.w, 3)
        t_begin = time.perf_counter()
        poses, _ = self.engine.DetectPosesInImage(resized_rgb_image)
        inference_time = time.perf_counter() - t_begin  # Second
        self.fps = convert_infr_time_to_fps(inference_time)
        result = []
        for i, pose in enumerate(poses):  # number of boxes
            if pose.score > self.score_threshold:
                pose_dict = collections.OrderedDict(sorted(pose.keypoints.items(), key=lambda x: self.keypoints.index(x[0])))
                keypoints = np.array(
                    [[keypoint.yx[1], keypoint.yx[0], keypoint.score] for _, keypoint in pose_dict.items()])
                pred = keypoints[keypoints[:, 2] > .2]
                xs = pred[:, 0]
                ys = pred[:, 1]
                x_min = int(xs.min())
                x_max = int(xs.max())
                y_min = int(ys.min())
                y_max = int(ys.max())
                w = x_max - x_min
                h = y_max - y_min
                xmin = int(max(x_min - .15 * w, 0))
                xmax = int(min(x_max + .15 * w, self.w))
                ymin = int(max(y_min - .1 * h, 0))
                ymax = int(min(y_max + .1 * h, self.h))
                bbox_dict = {
                    "id": f"{self.class_id}-" + str(i), "bbox": [ymin / self.h, xmin / self.w, ymax / self.h, xmax / self.w],
                    "score": pose.score, "face": None
                }
                # extracting face bounding box
                if np.all(keypoints[[0, 1, 2, 5, 6], -1] > 0.15):
                    x_min_face = int(keypoints[6, 0])
                    x_max_face = int(keypoints[5, 0])
                    y_max_face = int((keypoints[5, 1] + keypoints[6, 1]) / 2)
                    y_eyes = int((keypoints[1, 1] + keypoints[2, 1]) / 2)
                    y_min_face = 2 * y_eyes - y_max_face
                    if (y_max_face - y_min_face > 0) and (x_max_face - x_min_face > 0):
                        h_crop = y_max_face - y_min_face
                        x_min_face = int(max(0, x_min_face - 0.1 * h_crop))
                        y_min_face = int(max(0, y_min_face - 0.1 * h_crop))
                        x_max_face = int(min(self.w, x_min_face + 1.1 * h_crop))
                        y_max_face = int(min(self.h, y_min_face + 1.1 * h_crop))
                        bbox_dict["face"] = [y_min_face / self.h, x_min_face / self.w, y_max_face / self.h, x_max_face / self.w]

                result.append(bbox_dict)

        return result


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

# Copyright (C) 2017 Anmol Gulati <anmol01gulati@gmail.com>
# Copyright (C) 2017 Radim Rehurek <radimrehurek@seznam.cz>

"""Python wrapper around `Varembed model <https://github.com/rguthrie3/MorphologicalPriorsForWordEmbeddings>`_.
Original paper:`"Morphological Priors for Probabilistic Neural Word Embeddings" <http://arxiv.org/pdf/1608.01056.pdf>`_.

Notes
-----
* This module allows ability to obtain word vectors for out-of-vocabulary words, for the Varembed model.
* The wrapped model can not be updated with new documents for online training.

"""

import logging
import numpy as np

from gensim import utils
from gensim.models.keyedvectors import KeyedVectors
from gensim.models.word2vec import Vocab

logger = logging.getLogger(__name__)

class VarEmbed(KeyedVectors):
    """Python wrapper using `Varembed <https://github.com/rguthrie3/MorphologicalPriorsForWordEmbeddings>`_.

    Warnings
    --------
    This is **only** python wrapper for `Varembed <https://github.com/rguthrie3/MorphologicalPriorsForWordEmbeddings>`_,
    this allows to load pre-trained models only.

    """
    def __init__(self):
        self.vector_size = 0
        self.vocab_size = 0

    @classmethod
    def load_varembed_format(cls, vectors, morfessor_model=None):
        """Load the word vectors into matrix from the varembed output vector files.

        Parameters
        ----------
        vectors : dict
            Pickle file containing the word vectors.
        morfessor_model : str, optional
            Path to the trained morfessor model.

        Returns
        -------
        :class:`~gensim.models.wrappers.varembed.VarEmbed`
            Ready to use instance.

        """
        result = cls()
        if vectors is None:
            raise Exception("Please provide vectors binary to load varembed model")
        d = utils.unpickle(vectors)
        word_to_ix = d['word_to_ix']
        morpho_to_ix = d['morpho_to_ix']
        word_embeddings = d['word_embeddings']
        morpho_embeddings = d['morpheme_embeddings']
        result.load_word_embeddings(word_embeddings, word_to_ix)
        if morfessor_model:
            try:
                import morfessor
                morfessor_model = morfessor.MorfessorIO().read_binary_model_file(morfessor_model)
                result.add_morphemes_to_embeddings(morfessor_model, morpho_embeddings, morpho_to_ix)
            except ImportError:
                # Morfessor Package not found.
                logger.error('Could not import morfessor. Not using morpheme embeddings')
                raise ImportError('Could not import morfessor.')

        logger.info('Loaded varembed model vectors from %s', vectors)
        return result

    def load_word_embeddings(self, word_embeddings, word_to_ix):
        """Loads the word embeddings.

        Parameters
        ----------
        word_embeddings : numpy.ndarray
            Matrix with word-embeddings.
        word_to_ix : dict of (str, int)
            Mapping word to index.

        """
        logger.info("Loading the vocabulary")
        self.vocab = {}
        self.index2word = []
        counts = {}
        for word in word_to_ix:
            counts[word] = counts.get(word, 0) + 1
        self.vocab_size = len(counts)
        self.vector_size = word_embeddings.shape[1]
        self.syn0 = np.zeros((self.vocab_size, self.vector_size))
        self.index2word = [None] * self.vocab_size
        logger.info("Corpus has %i words", len(self.vocab))
        for word_id, word in enumerate(counts):
            self.vocab[word] = Vocab(index=word_id, count=counts[word])
            self.syn0[word_id] = word_embeddings[word_to_ix[word]]
            self.index2word[word_id] = word
        assert((len(self.vocab), self.vector_size) == self.syn0.shape)
        logger.info("Loaded matrix of %d size and %d dimensions", self.vocab_size, self.vector_size)

    def add_morphemes_to_embeddings(self, morfessor_model, morpho_embeddings, morpho_to_ix):
        """Include morpheme embeddings into vectors.

        Parameters
        ----------
        morfessor_model : :class:`morfessor.baseline.BaselineModel`
            Morfessor model.
        morpho_embeddings : dict
            Pickle file containing morpheme embeddings.
        morpho_to_ix : dict
            Mapping morpheme to index.

        """
        for word in self.vocab:
            morpheme_embedding = np.array(
                [
                    morpho_embeddings[morpho_to_ix.get(m, -1)]
                    for m in morfessor_model.viterbi_segment(word)[0]
                ]
            ).sum(axis=0)
            self.syn0[self.vocab[word].index] += morpheme_embedding
        logger.info("Added morphemes to word vectors")


# I do not claim to own any of the facial recognition algorithm. Though it was modified, the essence of the code is still credited to Will Koehrsen.
# His original code that IsaacGAN was built upon can be found here: https://medium.com/@williamkoehrsen/facial-recognition-using-googles-convolutional-neural-network-5aa752b4240e

"""
Does not work at the moment! This is currently a work in progress, and it still contains many bugs and issues. We are currently working on
making it better, and hopefully we can get a fully working model sometime in the near future. But for now, here's our progress that we've
made so far.
"""

# Import useful Python libraries
import tensorflow as tf
import numpy as np
import os
import sys
import urllib
import tarfile
from distutils.dir_util import copy_tree
import matplotlib.pyplot as plt
import math
from collections import Counter
from scipy.misc import imresize
from nets import inception
from tensorflow.contrib import slim
from datetime import datetime
import time
import scipy

# URL for TensorFlow models
TF_MODELS_URL = "http://download.tensorflow.org/models/"

# Inception V3 CNN, built by Google
INCEPTION_V3_URL = TF_MODELS_URL + "inception_v3_2016_08_28.tar.gz"

# Directory to save model checkpoints
MODELS_DIR = "models/cnn"
INCEPTION_V3_CKPT_PATH = MODELS_DIR + "/inception_v3.ckpt"

# define function since datasets don't work on my computer
def download_and_uncompress_tarball(tarball_url, dataset_dir):
  filename = tarball_url.split('/')[-1]
  filepath = os.path.join(dataset_dir, filename)

  def _progress(count, block_size, total_size):
    sys.stdout.write('\r>> Downloading %s %.1f%%' % (
        filename, float(count * block_size) / float(total_size) * 100.0))
    sys.stdout.flush()
  filepath, _ = urllib.request.urlretrieve(tarball_url, filepath, _progress)
  print()
  statinfo = os.stat(filepath)
  print('Successfully downloaded', filename, statinfo.st_size, 'bytes.')
  tarfile.open(filepath, 'r:gz').extractall(dataset_dir)

# Make the model directory if it does not exist
if not tf.gfile.Exists(MODELS_DIR):
 tf.gfile.MakeDirs(MODELS_DIR)

# Download the appropriate model of Inception
if not os.path.exists(INCEPTION_V3_CKPT_PATH):
    download_and_uncompress_tarball(INCEPTION_V3_URL, MODELS_DIR)

# Full deep-funneled images dataset
FACES_URL = "http://vis-www.cs.umass.edu/lfw/lfw-deepfunneled.tgz"
IMAGES_DOWNLOAD_DIRECTORY = "tmp/faces"
IMAGES_DIRECTORY = "images/faces"

if not os.path.exists(IMAGES_DOWNLOAD_DIRECTORY):
 os.makedirs(IMAGES_DOWNLOAD_DIRECTORY)

# If the file has not already been downloaded, retrieve and extract it
if not os.path.exists(IMAGES_DOWNLOAD_DIRECTORY + "/lfw-deepfunneled.tgz"):
 download_and_uncompress_tarball(FACES_URL, IMAGES_DOWNLOAD_DIRECTORY)

# Count number of photos of each individual
people_number = []

people = IMAGES_DOWNLOAD_DIRECTORY + "/lfw-deepfunneled"

for person in people:
 folder_path = os.path.join(people, person)
 num_images = len([f for f in os.listdir(folder_path) if os.path.isfile(f)])
 people_number.append((person, num_images))

# Sort the list of tuples by the number of images
people_number = sorted(people_number, key=lambda x: x[1], reverse=True)

# Determine number of people with one image
people_with_one_photo = [(person) for person, num_images in people_number if num_images==1]
print("Individuals with one photo: {}".format(len(people_with_one_photo)))

# Number of individuals to retain
num_classes = 10

# Make the images directory if it does not exist
if not os.path.exists(IMAGES_DIRECTORY):
 os.mkdir(IMAGES_DIRECTORY)

 # Take the ten folders with the most images and move to new directory
 # Rename the folders with the number of images and name of individual
 for person in people_number[:num_classes]:
 name = person[0]
 # Original download directory path
 folder_path = IMAGES_DOWNLOAD_DIRECTORY + '/lfw-deepfunneled/' + name
 formatted_num_images = str(person[1]).zfill(3)
 new_folder_name = "{} {}".format(formatted_num_images, name)
 image_new_name = IMAGES_DIRECTORY + "/" + new_folder_name

 # Make a new folder for each individual in the images directory
 os.mkdir(IMAGES_DIRECTORY + '/' + name)
 # Copy the folder from the download location to the new folder
 copy_tree(folder_path, IMAGES_DIRECTORY + '/' + name)
 # Rename the folder with images and individual
 os.rename(IMAGES_DIRECTORY + '/' + name, image_new_name)

# Map each class to an integer label
class_mapping = {}
class_images = {}
# Create dictionary to map integer labels to individuals
# Class_images will record number of images for each class
for index, directory in enumerate(os.listdir("images/faces")):
 class_mapping[index] = directory.split(" ")[1]
 class_images[index] = int(directory.split(' ')[0])
print(class_mapping)

total_num_images = np.sum(list(class_images.values()))
print("Individual \t Composition of Dataset\n")
for label, num_images in class_images.items():
 print("{:20} {:.2f}%".format(
 class_mapping[label], (num_images / total_num_images) * 100))

image_arrays = []
image_labels = []
root_image_directory = "images/faces/"
for label, person in class_mapping.items():
 for directory in os.listdir(root_image_directory):
 if directory.split(" ")[1] == person:
 image_directory = root_image_directory + directory
 break

 for image in os.listdir(image_directory):
 image = plt.imread(os.path.join(image_directory, image))
 image_arrays.append(image)
 image_labels.append(label)
image_arrays = np.array(image_arrays)
image_labels = np.array(image_labels)
print(image_arrays.shape, image_labels.shape)

# Fractions for each dataset
train_frac = 0.70
valid_frac = 0.05
test_frac = 0.25

# This function takes in np arrays of images and labels along with split fractions
# and returns the six data arrays corresponding to each dataset as the appropriate type
def create_data_splits(X, y, train_frac=train_frac, test_frac=test_frac, valid_frac=valid_frac):
 X = np.array(X)
 y = np.array(y)

 # Make sure that the fractions sum to 1.0
 assert (test_frac + valid_frac + train_frac == 1.0), "Test + Valid + Train Fractions must sum to 1.0"

 X_raw_test = []
 X_raw_valid = []
 X_raw_train = []

 y_raw_test = []
 y_raw_valid = []
 y_raw_train = []

 # Randomly order the data and labels
 random_indices = np.random.permutation(len(X))
 X = X[random_indices]
 y = y[random_indices]

 for image, label in zip(X, y):

 # Number of images that correspond to desired fraction
 test_length = math.floor(test_frac * class_images[label])
 valid_length = math.floor(valid_frac * class_images[label])

 # Check to see if the datasets have the right number of labels (and images)
 if Counter(y_raw_test)[label] < test_length:
 X_raw_test.append(image)
 y_raw_test.append(label)
 elif Counter(y_raw_valid)[label] < valid_length:
 X_raw_valid.append(image)
 y_raw_valid.append(label)
 else:
 X_raw_train.append(image)
 y_raw_train.append(label)

 return np.array(X_raw_train, dtype=np.float32), np.array(X_raw_valid, dtype=np.float32), np.array(X_raw_test, dtype=np.float32), np.array(y_raw_train, dtype=np.int32), np.array(y_raw_valid, dtype=np.int32), np.array(y_raw_test, dtype=np.int32)
# Create all the testing splits using the create_splits function
X_train, X_valid, X_test, y_train, y_valid, y_test = create_data_splits(image_arrays, image_labels)

# Check the number of images in each dataset split
print(X_train.shape, X_valid.shape, X_test.shape)
print(y_train.shape, y_valid.shape, y_test.shape)

# Import matplotlib and use magic command to plot in notebook
%matplotlib inline
# Function to plot an array of RGB values
def plot_color_image(image):
 plt.figure(figsize=(4,4))
 plt.imshow(image.astype(np.uint8), interpolation='nearest')
 plt.axis('off')
import random
# PLot 2 examples from each class
num_examples = 2
# Iterate through the classes and plot 2 images from each
for class_number, person in class_mapping.items():
 print('{} Number of Images: {}'.format(person, class_images[class_number]))
 example_images = []
 while len(example_images) < num_examples:
 random_index = np.random.randint(len(X_train))
 if y_train[random_index] == class_number:
 example_images.append(X_train[random_index])

 for i, image in enumerate(example_images):
 plt.subplot(100 + num_examples*10 + i + 1)
 plt.imshow(image.astype(np.uint8), interpolation='nearest')
 plt.axis('off')
 plt.show()

# Function takes in an image array and returns the resized and normalized array
def prepare_image(image, target_height=299, target_width=299):
 image = imresize(image, (target_width, target_height))
 return image.astype(np.float32) / 255

tf.reset_default_graph()
X = tf.placeholder(tf.float32, [None, 299, 299, 3], name='X')
is_training = tf.placeholder_with_default(False, [])
# Run inception function to determine endpoints
with slim.arg_scope(inception.inception_v3_arg_scope()):
 logits, end_points = inception.inception_v3(X, num_classes=1001, is_training=is_training)
# Create saver of network before alterations
inception_saver = tf.train.Saver()
print(end_points)

# Isolate the trainable layer
prelogits = tf.squeeze(end_points['PreLogits'], axis=[1,2])
# Define the training layer and the new output layer
n_outputs = len(class_mapping)
with tf.name_scope("new_output_layer"):
 people_logits = tf.layers.dense(prelogits, n_outputs, name="people_logits")
 probability = tf.nn.softmax(people_logits, name='probability')
# Placeholder for labels
y = tf.placeholder(tf.int32, None)
# Loss function and training operation
# The training operation is passed the variables to train which includes only the single layer
with tf.name_scope("train"):
 xentropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=people_logits, labels=y)
 loss = tf.reduce_mean(xentropy)
 optimizer = tf.train.AdamOptimizer(learning_rate=0.01)
 # Single layer to be trained
 train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope="people_logits")
 # The variables to train are passed to the training operation
 training_op = optimizer.minimize(loss, var_list=train_vars)

# Accuracy for network evaluation
with tf.name_scope("eval"):
 correct = tf.nn.in_top_k(predictions=people_logits, targets=y, k=1)
 accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))

# Intialization function and saver
with tf.name_scope("init_and_saver"):
 init = tf.global_variables_initializer()
 saver = tf.train.Saver()

# Function takes in an array of images and labels and processes the images to create
# a batch of a given size
def create_batch(X, y, start_index=0, batch_size=4):

 stop_index = start_index + batch_size
 prepared_images = []
 labels = []

 for index in range(start_index, stop_index):
 prepared_images.append(prepare_image(X[index]))
 labels.append(y[index])

 # Combine the images into a single array by joining along the 0th axis
 X_batch = np.stack(prepared_images)
 # Combine the labels into a single array
 y_batch = np.array(labels, dtype=np.int32)

 return X_batch, y_batch

X_valid, y_valid = create_batch(X_valid, y_valid, 0, len(X_valid))
print(X_valid.shape, y_valid.shape)

with tf.name_scope("tensorboard_writing"):
 # Track validation accuracy and loss and training accuracy
 valid_acc_summary = tf.summary.scalar(name='valid_acc', tensor=accuracy)
 valid_loss_summary = tf.summary.scalar(name='valid_loss', tensor=loss)
 train_acc_summary = tf.summary.scalar(name='train_acc', tensor=accuracy)
# Merge the validation stats
 valid_merged_summary = tf.summary.merge(inputs=[valid_acc_summary, valid_loss_summary])
# Specify the directory for the FileWriter
now = datetime.now().strftime("%Y%m%d_%H%M%S")
model_dir = "{}_unaugmented".format(now)
logdir = "tensorboard/faces/" + model_dir
file_writer = tf.summary.FileWriter(logdir=logdir, graph=tf.get_default_graph())

# start training model
n_epochs = 100
batch_size = 32
# Early stopping parameters
max_checks_without_progress = 20
checks_without_progress = 0
best_loss = np.float("inf")
# Show progress every show_progress epochs
show_progress = 1
# Want to iterate through the entire training set every epoch
n_iterations_per_epoch = len(X_train) // batch_size
# Specify the directory for the FileWriter
now = datetime.now().strftime("%Y%m%d_%H%M%S")
model_dir = "{}_unaugmented".format(now)
logdir = "tensorboard/faces/" + model_dir
file_writer = tf.summary.FileWriter(logdir=logdir, graph=tf.get_default_graph())
# This is the pre-trained model checkpoint training path
inception_v3_checkpoint_path = "models/cnn/inception_v3.ckpt"
# This is the checkpoint path for our trained model with no dataaugmentation
unaugmented_training_path = "models/cnn/inception_v3_faces_unaugmented.ckpt"
with tf.Session() as sess:
 init.run()
 # Restore all the weights from the original CNN
 inception_saver.restore(sess, inception_v3_checkpoint_path)

 t0 = time.time()
 for epoch in range(n_epochs):
 start_index = 0
 # Each epoch, iterate through all the training instances
 for iteration in range(n_iterations_per_epoch):
 X_batch, y_batch = create_batch(X_train, y_train, start_index, batch_size)
 # Train the trainable layer
 sess.run(training_op, {X: X_batch, y: y_batch})
 start_index += batch_size

 # Display the progress of training and write to the TensorBoard directory
 # for later visualization of the training
 if epoch % show_progress == 0:
 train_summary = sess.run(train_acc_summary, {X: X_batch, y: y_batch})
 file_writer.add_summary(train_summary, (epoch+1))
 #Size for validation limited by GPU memory (68 images will work)
 valid_loss, valid_acc, valid_summary = sess.run([loss, accuracy, valid_merged_summary], {X: X_valid, y: y_valid})
 file_writer.add_summary(valid_summary, (epoch+1))
 print('Epoch: {:4} Validation Loss: {:.4f} Accuracy: {:4f}'.format(
 epoch+1, valid_loss, valid_acc))

 # Check to see if network is still improving, if improved during epoch
 # a snapshot of the model will be saved to retain the best model
 if valid_loss < best_loss:
 best_loss = valid_loss
 checks_without_progess = 0
 save_path = saver.save(sess, unaugmented_training_path)

 # If network is not improving for a specified number of epochs, stop training
 else:
 checks_without_progress += 1
 if checks_without_progress > max_checks_without_progress:
 print('Stopping Early! Loss has not improved in {} epochs'.format(
 max_checks_without_progress))
 break

 t1 = time.time()

print('Total Training Time: {:.2f} minutes'.format( (t1-t0) / 60))

# Evaluate training model
eval_batch_size = 32
n_iterations = len(X_test) // eval_batch_size
with tf.Session() as sess:
 # Restore the new trained model
 saver.restore(sess, unaugmented_training_path)

 start_index = 0
 # Create a dictionary to store all the accuracies
 test_acc = {}

 t0 = time.time()
 # Iterate through entire testing set one batch at a time
 for iteration in range(n_iterations):
 X_test_batch, y_test_batch = create_batch(X_test, y_test, start_index, batch_size=eval_batch_size)
 test_acc[iteration] = accuracy.eval({X: X_test_batch, y:y_test_batch})
 start_index += eval_batch_size
print('Iteration: {} Batch Testing Accuracy: {:.2f}%'.format(
 iteration+1, test_acc[iteration] * 100))

 t1 = time.time()

 # Final accuracy is mean of each batch accuracy
 print('\nFinal Testing Accuracy: {:.4f}% on {} instances.'.format(
 np.mean(list(test_acc.values())) * 100, len(X_test)))
 print('Total evaluation time: {:.4f} seconds'.format((t1-t0)))

# Take in an image as an array and return image with a [dx, dy] shift
def shift_image(image_array, shift):
 return scipy.ndimage.interpolation.shift(image_array, shift, cval=0)
# Four shifts of 30 pixels
shifts = [[30,0], [-30,0], [0, 30], [0,-30]]
shifted_images = []
shifted_labels = []
# Iterate through all training images
for image, label in zip(X_train, y_train):

 # Swap the color channel and height axis
 layers = np.swapaxes(image, 0, 2)

 # Apply four shifts to each original image
 for shift in shifts:
 transposed_image_layers = []

 # Apply the shift to the image one layer at a time
 # Each layer is an RGB color channel
 for layer in layers:
 transposed_image_layers.append(shift_image(layer, shift))

 # Stack the RGB layers to get one image and reswap the axes
 transposed_image = np.stack(transposed_image_layers)
 transposed_image = np.swapaxes(transposed_image, 0, 2)

 # Add the shifted images and the labels to a list
 shifted_images.append(transposed_image)
 shifted_labels.append(label)
# Convert the images and labels to numpy arrays
shifted_images = np.array(shifted_images)
shifted_labels = np.array(shifted_labels)
print(shifted_images.shape,shifted_labels.shape)

ex_index = 5
# Plot original image
plot_color_image(X_train[ex_index])
plt.title("Original Image of {}".format(class_mapping[y_train[ex_index]]))
plt.show()
ex_shifted_images = shifted_images[ex_index*4:(ex_index*4)+ 4]
# Plot four shifted images
for i, image in enumerate(ex_shifted_images):
 shift = shifts[i]
 plt.subplot(2,2,i+1)
 plt.imshow(image.astype(np.uint8), interpolation='nearest')
 plt.title('Shift: {}'.format(shift))
 plt.axis('off')
plt.show()

# Create a new training set with the original and shifted images
X_train_exp = np.concatenate((shifted_images, X_train))
y_train_exp = np.concatenate((shifted_labels, y_train))
print(X_train_exp.shape, y_train_exp.shape)

restart_augmented_training_path = "models/cnn/inception_v3_faces_restart_augmented.ckpt"
with tf.Session() as sess:
 init.run()
 inception_saver.restore(sess, unaugmented_training_path)

images_flipped = []
labels_flipped = []
# Flip every image in the training set
for image, label in zip(X_train, y_train):
 images_flipped.append(np.fliplr(image))
 labels_flipped.append(label)
# Convert the flipped images and labels to arrays
images_flipped = np.array(images_flipped)
labels_flipped = np.array(labels_flipped)

ex_index = 652
plot_color_image(X_train[ex_index])
plt.title('Original Image of {}'.format(class_mapping[y_train[ex_index]]))
plt.show()
plot_color_image(images_flipped[ex_index])
plt.title('Left to Right Flipped Image')
plt.show()

X_train_with_flip = np.concatenate((X_train, images_flipped))
y_train_with_flip = np.concatenate((y_train, labels_flipped))
print(X_train_with_flip.shape, y_train_with_flip.shape)

# Visualize performance of CNN
def classify_image(index, images=X_test, labels=y_test):
 image_array = images[index]
 label = class_mapping[labels[index]]

 prepared_image = prepare_image(image_array)
 prepared_image = np.reshape(prepared_image, newshape=(-1, 299, 299, 3))

 with tf.Session() as sess:
 saver.restore(sess, restart_augmented_training_path)
 predictions = sess.run(probability, {X: prepared_image})

 predictions = [(i, prediction) for i, prediction in enumerate(predictions[0])]
 predictions = sorted(predictions, key=lambda x: x[1], reverse=True)
 print('\nCorrect Answer: {}'.format(label))
 print('\nPredictions:')
 for prediction in predictions:
 class_label = prediction[0]
 probability_value = prediction[1]
 label = class_mapping[class_label]
 print("{:26}: {:.2f}%".format(label, probability_value * 100))

 plot_color_image(image_array)
 return predictions

# this is to test specific images
predictions = classify_image(replace with number)

eval_batch_size = 32
n_iterations = len(X_test) // eval_batch_size
# Evaluation with saved predictions
with tf.Session() as sess:
 saver.restore(sess, restart_augmented_training_path)

 test_predictions = []
 start_index = 0

 # Add each set of predictions to a list
 for iteration in range(n_iterations):
 X_test_batch, y_test_batch = create_batch(X_test, y_test, start_index, batch_size=eval_batch_size)
 test_predictions.append(probability.eval({X: X_test_batch, y: y_test_batch}))
 start_index += eval_batch_size
# Convert list of predictions to np array
test_predictions = np.array(test_predictions)
test_predictions.shape

# Reshape predictions to correct shape and generate label array
test_predictions = np.reshape(test_predictions, (-1, 10))
test_predictions_label = np.argmax(test_predictions, axis=1)
test_predictions_label.shape

# A few of the testing examples are left off by the batching process
y_test_eval = y_test[:352]
# Make sure that accuracy agrees with earlier evaluation
test_accuracy = np.mean(np.equal(test_predictions_label, y_test_eval))
print("Test Accuracy: {:.2f}%".format(test_accuracy*100))

print('Individual: \t Composition of Test Set \t Composition of Predictions')
total_test_images = len(y_test_eval)
for label, individual in class_mapping.items():
 n_test_images = Counter(y_test_eval)[label]
 n_predictions = Counter(test_predictions_label)[label]
 print("{:26} {:5.2f}% {:26.2f}%".format(
 individual, n_test_images / total_test_images * 100, n_predictions / total_test_images * 100))

# Visualize distribution of images in test set
bins = np.arange(11)
fig, ax = plt.subplots(figsize=(14, 8))
plt.subplots_adjust(wspace=0.25)
plt.subplot(1,2,1)
plt.hist(y_test_eval, bins=bins, lw=1.2, edgecolor='black')
plt.title('Frequency of Classes in Test Set'); plt.xlabel('Class'); plt.ylabel('Count');
names = list(class_mapping.values())
xlabels = [name.split("_")[1] for name in names]
plt.xticks(bins+0.5, xlabels, rotation='vertical');
plt.subplot(1,2,2)
# Visualize distribution of images in predictions
bins = np.arange(11)
plt.hist(test_predictions_label, bins=bins, lw=1.2, edgecolor='black')
plt.title('Frequency of Classes in Predictions'); plt.xlabel('Class'); plt.ylabel('Count');
names = list(class_mapping.values())
xlabels = [name.split("_")[1] for name in names]
plt.xticks(bins+0.5, xlabels, rotation='vertical');

relative_differences = []
for label, individual in class_mapping.items():
 n_test_images = Counter(y_test_eval)[label]
 n_predictions = Counter(test_predictions_label)[label]
 n_test_pct = n_test_images / total_test_images
 n_predictions_pct = n_predictions / total_test_images

 relative_differences.append((n_predictions_pct - n_test_pct) / n_test_pct * 100)
colors = ['orange' if difference < 0.0 else 'blue' for difference in relative_differences]
# Visualize relative difference in
bins = np.arange(10)
fig, ax = plt.subplots()
plt.bar(bins, relative_differences, lw=1.2, color=colors, edgecolor='black')
plt.title('Relative Differences between Test and Prediction Composition'); plt.xlabel('Class'); plt.ylabel('Relative Difference (%)');
names = list(class_mapping.values())
xlabels = [name.split("_")[1] for name in names]
plt.xticks(bins, xlabels, rotation='vertical');