from __future__ import (
    division,
    absolute_import,
    with_statement,
    print_function,
    unicode_literals,
)
import open3d as o3d
import sys
import rospy
import ros_numpy
from sensor_msgs.msg import PointCloud2, Image
from lib.utils import pcl_helper
from geometry_msgs.msg import Pose, PoseArray, Point, Quaternion
from std_msgs.msg import Header
from scipy.spatial.transform import Rotation as R
sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')
import open3d as o3d
import torch
import torch.optim as optim
import torch.optim.lr_scheduler as lr_sched
import torch.nn as nn
import numpy as np
from torch.utils.data import DataLoader
import pprint
import os.path as osp
import os
import argparse
import time
import shutil
import pcl
import tqdm
from lib.utils.etw_pytorch_utils.viz import *
from lib import PVN3D
from datasets.openDR.openDR_dataset import openDR_Dataset
from lib.loss import OFLoss, FocalLoss
from common import Config
from lib.utils.sync_batchnorm import convert_model
from lib.utils.warmup_scheduler import CyclicLR
from lib.utils.pvn3d_eval_utils import TorchEval
import lib.utils.etw_pytorch_utils as pt_utils
import resource
from collections import namedtuple
from lib.utils.basic_utils import Basic_Utils

rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (30000, rlimit[1]))
config = Config(dataset_name='openDR')
bs_utl = Basic_Utils(config)

parser = argparse.ArgumentParser(description="Arg parser")
parser.add_argument(
    "-weight_decay",
    type=float,
    default=0,
    help="L2 regularization coeff [default: 0.0]",
)
parser.add_argument(
    "-lr", type=float, default=1e-2, help="Initial learning rate [default: 1e-2]"
)
parser.add_argument(
    "-lr_decay",
    type=float,
    default=0.5,
    help="Learning rate decay gamma [default: 0.5]",
)
parser.add_argument(
    "-decay_step",
    type=float,
    default=2e5,
    help="Learning rate decay step [default: 20]",
)
parser.add_argument(
    "-bn_momentum",
    type=float,
    default=0.9,
    help="Initial batch norm momentum [default: 0.9]",
)
parser.add_argument(
    "-bn_decay",
    type=float,
    default=0.5,
    help="Batch norm momentum decay gamma [default: 0.5]",
)
parser.add_argument(
    "-checkpoint", type=str, default=None, help="Checkpoint to start from"
)
parser.add_argument(
    "-epochs", type=int, default=1000, help="Number of epochs to train for"
)
parser.add_argument(
    "-run_name",
    type=str,
    default="sem_seg_run_1",
    help="Name for run in tensorboard_logger",
)

parser.add_argument("--test", action="store_true")
parser.add_argument("--cal_metrics", action="store_true")

lr_clip = 1e-5
bnm_clip = 1e-2

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

def checkpoint_state(model=None, optimizer=None, best_prec=None, epoch=None, it=None):
    optim_state = optimizer.state_dict() if optimizer is not None else None
    if model is not None:
        if isinstance(model, torch.nn.DataParallel):
            model_state = model.module.state_dict()
        else:
            model_state = model.state_dict()
    else:
        model_state = None

    return {
        "epoch": epoch,
        "it": it,
        "best_prec": best_prec,
        "model_state": model_state,
        "optimizer_state": optim_state,
    }

def save_checkpoint(
        state, is_best, filename="checkpoint", bestname="model_best",
        bestname_pure='pvn3d_best'
):
    filename = "{}_{}_kps.pth.tar".format(filename, config.n_keypoints)
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, "{}.pth.tar".format(bestname))
        shutil.copyfile(filename, "{}.pth.tar".format(bestname_pure))

def load_checkpoint(model=None, optimizer=None, filename="checkpoint"):
    filename = "{}.pth.tar".format(filename)

    if os.path.isfile(filename):
        print("==> Loading from checkpoint '{}'".format(filename))
        checkpoint = torch.load(filename)
        epoch = checkpoint["epoch"]
        it = checkpoint.get("it", 0.0)
        best_prec = checkpoint["best_prec"]
        if model is not None and checkpoint["model_state"] is not None:
            model.load_state_dict(checkpoint["model_state"])
        if optimizer is not None and checkpoint["optimizer_state"] is not None:
            optimizer.load_state_dict(checkpoint["optimizer_state"])
        print("==> Done")
        return it, epoch, best_prec
    else:
        print("==> Checkpoint '{}' not found".format(filename))
        return None

def model_fn_decorator(
    criterion, criterion_of, test=False
):
    modelreturn = namedtuple("modelreturn", ["preds", "loss", "acc"])
    teval = TorchEval('openDR')

    def model_fn(
        model, data, epoch=0, is_eval=False, is_test=False, finish_test=False
    ):
        if finish_test:
            teval.cal_auc('openDR')
            return None
        if is_eval:
            model.eval()
        with torch.set_grad_enabled(not is_eval):
            cu_dt = [item.to("cuda", non_blocking=True) for item in data]
            rgb, pcld, cld_rgb_nrm, choose, kp_targ_ofst, ctr_targ_ofst, \
                cls_ids, rts, labels, kp_3ds, ctr_3ds = cu_dt

            pred_kp_of, pred_rgbd_seg, pred_ctr_of = model(
                cld_rgb_nrm, rgb, choose
            )

            #print('labels '+ str(np.unique(labels.cpu().numpy())))
            loss_rgbd_seg = criterion(
                pred_rgbd_seg.view(labels.numel(), -1),
                labels.view(-1)
            ).sum()
            loss_kp_of = criterion_of(
                pred_kp_of, kp_targ_ofst, labels,
            ).sum()
            loss_ctr_of = criterion_of(
                pred_ctr_of, ctr_targ_ofst, labels,
            ).sum()
            w = [2.0, 1.0, 1.0]
            loss = loss_rgbd_seg * w[0] + loss_kp_of * w[1] + \
                   loss_ctr_of * w[2]

            _, classes_rgbd = torch.max(pred_rgbd_seg, -1)
            acc_rgbd = (
                classes_rgbd == labels
            ).float().sum() / labels.numel()

            if is_test:
                pred_pose_lst,pred_kps_lst,_,_2 = teval.eval_pose_parallel(
                    pcld, rgb, classes_rgbd,
                    pred_ctr_of, ctr_targ_ofst,
                    labels, epoch,
                    cls_ids, rts, pred_kp_of,
                    min_cnt=1, use_p2d=False, use_ctr_clus_flter=True,
                    use_ctr=True, ds_type="openDR"
                )

        return modelreturn(
            (pred_kp_of, pred_rgbd_seg, pred_ctr_of,  pred_pose_lst, cld_rgb_nrm, rts, rgb, pred_kps_lst), loss,
            {
                "acc_rgbd": acc_rgbd.item(),
                "loss": loss.item(),
                "loss_rgbd_seg": loss_rgbd_seg.item(),
                "loss_kp_of": loss_kp_of.item(),
                "loss_ctr_of": loss_ctr_of.item(),
                "loss_target": loss.item(),
            }
        )

    return model_fn

class Trainer(object):
    r"""
        Reasonably generic trainer for pytorch models

    Parameters
    ----------
    model : pytorch model
        Model to be trained
    model_fn : function (model, inputs, labels) -> preds, loss, accuracy
    optimizer : torch.optim
        Optimizer for model
    checkpoint_name : str
        Name of file to save checkpoints to
    best_name : str
        Name of file to save best model to
    lr_scheduler : torch.optim.lr_scheduler
        Learning rate scheduler.  .step() will be called at the start of every epoch
    bnm_scheduler : BNMomentumScheduler
        Batchnorm momentum scheduler.  .step() will be called at the start of every epoch
    """

    def __init__(
        self,
        model,
        model_fn,
        optimizer,
        checkpoint_name="ckpt",
        best_name="best",
        lr_scheduler=None,
        bnm_scheduler=None,
        viz=None,
    ):
        self.model, self.model_fn, self.optimizer, self.lr_scheduler, self.bnm_scheduler = (
            model,
            model_fn,
            optimizer,
            lr_scheduler,
            bnm_scheduler,
        )

        self.checkpoint_name, self.best_name = checkpoint_name, best_name

        self.training_best, self.eval_best = {}, {}
        self.viz = viz

    def eval_epoch(self, d_loader, is_test=False):
        self.model.eval()

        eval_dict = {}
        total_loss = 0.0
        count = 1.0
        for i, data in tqdm.tqdm(
            enumerate(d_loader), leave=False, desc="val"
        ):

            while not rospy.is_shutdown():
                self.optimizer.zero_grad()

                preds, loss, eval_res = self.model_fn(
                    self.model, data, is_eval=True, is_test=is_test
                )

                pcld = preds[4].cpu().numpy()[0]
                pcld = np.array(pcld, dtype=np.float32)
                rgb_img = preds[6].cpu().numpy()[0].transpose()
                rgb_img = rgb_img.transpose(1, 0, 2)
                pose = np.array(preds[3])
                gt_pose = preds[5].cpu().numpy()

                pred_kps = preds[7]#.cpu().numpy()[0]
                #print(pred_kps)
                _ , classes_rgbd = torch.max(preds[1], -1)
                #classes_rgbd  = preds[1]
                classes_rgbd = classes_rgbd.cpu().numpy()
                rgb_img = np.array(rgb_img, dtype=np.uint8)
                '''
                print('gt_pose '+str(gt_pose.shape))
                print('pose '+str(np.array(pose).shape))
                print('kps '+ str(np.array(pred_kps).shape))
                print('classes rgbd '+str(classes_rgbd.shape))
                print('unique labels '+str(np.unique(classes_rgbd)))
                print('')'''
                print('pose '+str(np.array(pose).shape))
                #print('pose[0] '+str(pose[0,:,:]) )
                rgb_labeled = rgb_img.copy()
                pose_msg = []
                pose_msg2 = []
                mask = np.zeros((1), dtype=np.uint32)
                rgb = np.zeros((pose.shape[0],3), dtype=np.uint32)
                rgb_mask = []
                for cls in range(pose.shape[0]):

                    rot = R.from_dcm(pose[cls, 0:3,0:3])
                    quat = rot.as_quat()

                    rot_gt = R.from_dcm(gt_pose[0,cls,0:3,0:3])
                    quat_gt = rot_gt.as_quat()

                    registered_corners = np.matmul(pose[cls, 0:3,0:3], corners_t[cls,:,:].T)
                    registered_corners += pose[cls,:3,3].reshape(3,1)

                    # Image with all the Bounding boxes - Takes longer per iteration
                    corners_2d = bs_utl.project_p3d(registered_corners.T, 1 , K = config.intrinsic_matrix['custom'] )
                    rgb_labeled_bb = bs_utl.draw_bounding_box(rgb_labeled, corners_2d)
                    '''
                    pose_msg.append(Pose(position = Point(x=pose[cls,0,3],y=pose[cls,1,3],z=pose[cls, 2,3]), orientation = Quaternion(x=quat[0],y=quat[1],z=quat[2],w=quat[3])))
                    pose_msg2.append(Pose(position = Point(x=gt_pose[0,cls,0,3],y=gt_pose[0,cls,1,3],z=gt_pose[0,cls,2,3]), orientation = Quaternion(x=quat_gt[0],y=quat_gt[1],z=quat_gt[2],w=quat_gt[3])))
                    '''

                    this_mask = np.where(classes_rgbd[0,:]==cls+1)[0]
                    rgb_mask.append(this_mask.shape[0])
                    this_rgb = bs_utl.get_label_color( cls_id=cls+1 )
                    rgb[cls,:] = np.array([this_rgb[0], this_rgb[1], this_rgb[2] ])
                    mask = np.concatenate((mask, this_mask), axis = 0  )

                cls = 3
                if pose.shape[0]>cls:
                    rot = R.from_dcm(pose[cls, 0:3,0:3])
                    quat = rot.as_quat()
                    rot_gt = R.from_dcm(gt_pose[0,cls,0:3,0:3])
                    quat_gt = rot_gt.as_quat()

                    pose_msg.append(Pose(position = Point(x=pose[cls,0,3],y=pose[cls,1,3],z=pose[cls, 2,3]), orientation = Quaternion(x=quat[0],y=quat[1],z=quat[2],w=quat[3])))
                    pose_msg2.append(Pose(position = Point(x=gt_pose[0,cls,0,3],y=gt_pose[0,cls,1,3],z=gt_pose[0,cls,2,3]), orientation = Quaternion(x=quat_gt[0],y=quat_gt[1],z=quat_gt[2],w=quat_gt[3])))
                    rgb = np.repeat(rgb, rgb_mask, axis=0)
                    p3d_rgbs = np.concatenate( (pcld[mask[1:],0:3], rgb) , axis= 1 )
                    #p2ds = bs_utl.project_p3d(pcld[mask,0:3], 1 , K = config.intrinsic_matrix['openDR'] )
                    p2ds_rgb = bs_utl.project_p3d(p3d_rgbs, 1 , K = config.intrinsic_matrix['openDR'] )
                    #print(p2ds.shape)
                    r = rgb[0,0]#np.random.randint(0,255)
                    g = rgb[1,0]#np.random.randint(0,255)
                    b = rgb[2,0]#np.random.randint(0,255)
                    #registered_pts = (np.matmul(pose[:3,:3],mesh_pts_t.T)+pose[:,3].reshape(3,1)).T
                    '''
                    # Image with single bounding box of the selected the cls - for faster visualization
                    registered_corners = np.matmul(pose[cls, 0:3,0:3], corners_t[cls,:,:].T)
                    registered_corners+= pose[cls,:3,3].reshape(3,1)
                    corners_2d = bs_utl.project_p3d(registered_corners.T, 1 , K = config.intrinsic_matrix['custom'] )'''
                    #mesh_p2ds = bs_utl.project_p3d(registered_pts, 1 , K = config.intrinsic_matrix['custom'] )
                    kps_2d = bs_utl.project_p3d(pred_kps[cls], 1 , K = config.intrinsic_matrix['custom'] )

                    ## Projects 2D points from all labels together with their respective colors
                    rgb_labeled = bs_utl.draw_p2ds(rgb_labeled, p2ds_rgb, (r, g, b), 1)
                    #rgb_labeled_kps = bs_utl.draw_p2ds(rgb_labeled, kps_2d, (0, 255, 0), 2)
                    #rgb_labeled_bb = bs_utl.draw_bounding_box(rgb_labeled_kps, corners_2d)
                    #rgb_labeled_axis = bs_utl.draw_axis(rgb_labeled, gt_pose[0:3,0:3], gt_pose[:,3], config.intrinsic_matrix['custom'])
                    #rgb_labeled = rgb_img
                    #rgb_labeled_kps = bs_utl.draw_p2ds( rgb_labeled ,kps_2d, (255,0,0),2 )
                    #print(rgb_labeled.max())
                    #print(rgb_labeled.min())
                    #print(pcld.shape)
                    #rgb_labeled.dtype = np.uint16
                    #print(rgb_labeled.shape)

                    cloud = pcl.PointCloud()
                    kps = pcl.PointCloud()
                    kps.from_array(pred_kps[cls])
                    cloud.from_array(pcld[:,0:3])
                    #colors = 200*np.ones((pcld.shape[0],3))
                    #colors[mask[1:],:] = p3d_rgbs[:,3:6]
                    colors = pcld[:,3:6]
                    #colors[mask,:]= [r,g,b]
                    cloud_out = pcl_helper.XYZ_to_XYZRGB(cloud, colors.tolist())
                    cloud_msg = pcl_helper.pcl_to_ros(cloud_out, 'world')
                    kp_cld_msg = pcl_helper.pcl_to_ros(kps, 'world')
                    #cld_registered = pcl.PointCloud()

                    #cld_registered.from_array(np.array(registered_pts, dtype=np.float32) )

                    pub.publish(cloud_msg)
                    pub2.publish(PoseArray(header=Header(stamp=rospy.Time.now(), frame_id='world'), poses = pose_msg))
                    pub3.publish(PoseArray(header=Header(stamp=rospy.Time.now(), frame_id='world'), poses = pose_msg2))
                    pub4.publish(ros_numpy.msgify(Image,rgb_labeled, encoding='rgb8'))
                    pub5.publish(kp_cld_msg)
                    #pub6.publish(pcl_helper.pcl_to_ros(cld_registered, 'world'))
                    '''
                    try:
                        usrIn = input('Enter to continue to next Test Image...')
                    except:
                        usrIn = None'''
                break
            if 'loss_target' in eval_res.keys():
                total_loss += eval_res['loss_target']
            else:
                total_loss += loss.item()

            count += 1
            for k, v in eval_res.items():
                if v is not None:
                    eval_dict[k] = eval_dict.get(k, []) + [v]
        if is_test:
            self.model_fn(
                self.model, data, is_eval=True, is_test=is_test, finish_test=True
            )

        return total_loss / count, eval_dict

    def train(
        self,
        start_it,
        start_epoch,
        n_epochs,
        train_loader,
        test_loader=None,
        best_loss=0.0,
        log_epoch_f = None
    ):
        r"""
           Call to begin training the model

        Parameters
        ----------
        start_epoch : int
            Epoch to start at
        n_epochs : int
            Number of epochs to train for
        test_loader : torch.utils.data.DataLoader
            DataLoader of the test_data
        train_loader : torch.utils.data.DataLoader
            DataLoader of training data
        best_loss : float
            Testing loss of the best model
        """

        def is_to_eval(epoch, it):
            if it < 300 * 100:
                eval_frequency = 870
            elif it < 400 * 100:
                eval_frequency = (20 * 100)
            elif it < 500 * 100:
                eval_frequency = (12 * 100)
            elif it < 600 * 100:
                eval_frequency = (8 * 100)
            elif it < 800 * 100:
                eval_frequency = (4 * 100)
            else:
                eval_frequency = (2 * 100)
            to_eval = (it % eval_frequency) == 0
            return to_eval, eval_frequency

        it = start_it
        _, eval_frequency = is_to_eval(0, it)

        with tqdm.trange(start_epoch, n_epochs + 1, desc="epochs") as tbar, tqdm.tqdm(
            total=eval_frequency, leave=False, desc="train"
        ) as pbar:

            for epoch in tbar:
                # Reset numpy seed.
                # REF: https://github.com/pytorch/pytorch/issues/5059
                np.random.seed()
                if log_epoch_f is not None:
                    print('Epoch Started...')
                    os.system("echo {} > {}".format(epoch, log_epoch_f))
                for ibs, batch in enumerate(train_loader):
                    print('Batch '+ str(ibs))
                    self.model.train()

                    if self.lr_scheduler is not None:
                        self.lr_scheduler.step(it)

                    if self.bnm_scheduler is not None:
                        self.bnm_scheduler.step(it)

                    self.optimizer.zero_grad()
                    _, loss, res = self.model_fn(self.model, batch)
                    print(res)
                    loss.backward()
                    self.optimizer.step()

                    it += 1

                    pbar.update()
                    pbar.set_postfix(dict(total_it=it))
                    tbar.refresh()

                    if self.viz is not None:
                        self.viz.update("train", it, res)

                    eval_flag, eval_frequency = is_to_eval(epoch, it)
                    if eval_flag:
                        pbar.close()

                        if test_loader is not None:
                            val_loss, res = self.eval_epoch(test_loader)

                            if self.viz is not None:
                                self.viz.update("val", it, res)

                            is_best = val_loss < best_loss
                            best_loss = min(best_loss, val_loss)
                            save_checkpoint(
                                checkpoint_state(
                                    self.model, self.optimizer, val_loss, epoch, it
                                ),
                                is_best,
                                filename=self.checkpoint_name,
                                bestname=self.best_name+'_%.4f'%val_loss,
                                bestname_pure=self.best_name,
                            )
                            info_p = self.checkpoint_name.replace(
                                '.pth.tar','_epoch.txt'
                            )
                            os.system(
                                'echo {} {} >> {}'.format(
                                    it, val_loss, info_p
                                )
                            )

                        pbar = tqdm.tqdm(
                            total=eval_frequency, leave=False, desc="train"
                        )
                        pbar.set_postfix(dict(total_it=it))

                    self.viz.flush()

        return best_loss

if __name__ == "__main__":
    args = parser.parse_args()
    rospy.init_node('pvn3d_openDR_pred')
    pub = rospy.Publisher('/eval/pvn3d_cloud', PointCloud2)
    pub2 = rospy.Publisher('/eval/pvn3d_pose', PoseArray)
    pub3 = rospy.Publisher('/eval/pvn3d_gt_pose', PoseArray )
    pub4 = rospy.Publisher('/eval/pvn3d_label_img', Image )
    pub5 = rospy.Publisher('/eval/pvn3d_kps', PointCloud2)
    pub6 = rospy.Publisher('/eval/pvn3d_mesh_registered', PointCloud2)
    #pub6 = rospy.Publisher('/pvn3d_mesh_registered', PointCloud2)
    #mesh = pcl.load('/home/ahmad3/piston_from_cad_poisson.pcd')
    #sor = mesh.make_voxel_grid_filter()
    #sor.set_leaf_size(0.02, 0.02, 0.02)
    #mesh_ds = sor.filter()
    #mesh_pts = np.array(mesh, dtype= np.float32)
    cam2optical = R.from_euler('zyx',[1.57, 0,1.57])
    cam2optical = cam2optical.as_dcm()
    corners_t = np.zeros((10,8,3))
    for i in range(0,10):
        corners = np.loadtxt('./PVN3D/pvn3d/datasets/openDR/openDR_object_kps/'+str(i+1)+'/corners.txt')
        corners_t[i] = np.matmul(cam2optical, corners.transpose()).transpose()
        corners_t[i] = np.concatenate(( corners_t[i,:,2].reshape(8,1) , -corners_t[i,:, 0].reshape(8,1), -corners_t[i,:, 1].reshape(8,1) ), axis=1)
    '''

    mesh_pts_t = np.matmul(cam2optical, mesh_pts.transpose()).transpose()
    corners_t = np.matmul(cam2optical, corners.transpose()).transpose()
    n = mesh_pts_t.shape[0]
    mesh_pts_t = np.concatenate(( mesh_pts_t[:,2].reshape(n,1) , -mesh_pts_t[:, 0].reshape(n,1), -mesh_pts_t[:, 1].reshape(n,1) ), axis=1)
    corners_t = np.concatenate(( corners_t[:,2].reshape(8,1) , -corners_t[:, 0].reshape(8,1), -corners_t[:, 1].reshape(8,1) ), axis=1)
    '''

    test_ds = openDR_Dataset('test')
    test_loader = torch.utils.data.DataLoader(
        test_ds, batch_size=config.test_mini_batch_size, shuffle=False,
        num_workers=20
    )

    model = PVN3D(
        num_classes=config.n_classes, pcld_input_channels=6, pcld_use_xyz=True,
        num_points=config.n_sample_points, num_kps=config.n_keypoints
    ).cuda()
    model = convert_model(model)
    model.cuda()
    optimizer = optim.Adam(
        model.parameters(), lr=args.lr, weight_decay=args.weight_decay
    )

    # default value
    it = -1  # for the initialize value of `LambdaLR` and `BNMomentumScheduler`
    best_loss = 1e10
    start_epoch = 1

    cur_mdl_pth = os.path.join(config.log_model_dir, 'pvn3d.pth.tar')
    if args.checkpoint is None and os.path.exists(cur_mdl_pth):
        args.checkpoint = cur_mdl_pth
    # load status from checkpoint
    if args.checkpoint is not None:

        checkpoint_status = load_checkpoint(
            model, optimizer, filename=args.checkpoint[:-8]
        )
        if checkpoint_status is not None:
            it, start_epoch, best_loss = checkpoint_status
            print('Loading from checkpoint...' + 'It: '+str(it) + 'Epoch: '+str(start_epoch)+'Loss: '+str(best_loss))
    model = nn.DataParallel(
        model
    )

    lr_scheduler = CyclicLR(
        optimizer, base_lr=1e-5, max_lr=1e-3,
        step_size=config.n_total_epoch * config.num_mini_batch_per_epoch // 6,
        mode = 'triangular'
    )

    bnm_lmbd = lambda it: max(
        args.bn_momentum
        * args.bn_decay ** (int(it * config.mini_batch_size / args.decay_step)),
        bnm_clip,
    )
    bnm_scheduler = pt_utils.BNMomentumScheduler(
        model, bn_lambda=bnm_lmbd, last_epoch=it
    )

    it = max(it, 0)  # for the initialize value of `trainer.train`

    model_fn = model_fn_decorator(
        nn.DataParallel(FocalLoss(gamma=2)),
        nn.DataParallel(OFLoss()),
        args.test,
    )

    viz = CmdLineViz()

    viz.text(pprint.pformat(vars(args)))

    checkpoint_fd = config.log_model_dir
    if not os.path.exists(checkpoint_fd):
        os.system('mkdir -p {}'.format(checkpoint_fd))

    trainer = Trainer(
        model,
        model_fn,
        optimizer,
        checkpoint_name = os.path.join(checkpoint_fd, "pvn3d"),
        best_name = os.path.join(checkpoint_fd, "pvn3d_best"),
        lr_scheduler = lr_scheduler,
        bnm_scheduler = bnm_scheduler,
        viz = viz,
    )

    start = time.time()
    val_loss, res = trainer.eval_epoch(
        test_loader, is_test=True
    )


#!/usr/bin/env python
import random
import flask
import gi
import multiprocessing
import io
import detect_tflite
import base64
import binascii
import time
import math
import argparse
gi.require_version('Gst', '1.0')
gi.require_version('GstApp', '1.0')
gi.require_version('GstVideo', '1.0')
from gi.repository import GObject, Gst, GstApp, GstVideo
import sys
import typing
import queue
import numpy as np
import pickle
from PIL import Image
import cv2
import datetime
import logging

# For metrics exposition
from prometheus_client import start_http_server, Summary, Counter, Gauge, make_wsgi_app
from prometheus_client import generate_latest
app = flask.Flask(__name__)

start_time = None
IMG_BUF_LEN = 10
coco_labels = detect_tflite.load_labels('coco_labels.txt')

mgr = multiprocessing.Manager()
img_q = mgr.Queue(maxsize=100)
res_q = mgr.Queue(maxsize=100)
# Shared rolling buffer of images that will be consolidated into the primary
# processed and rendered as a video
shr_tf_img = mgr.list()
# The corresponding list of results computed for the images. Will be provided by
# the worker as a list of results (see detect_objects for structure)
shr_tf_res = mgr.list()
# Null out the post-processing shared buffers
for i in range(0, IMG_BUF_LEN):
    shr_tf_img.append((-1, Image.fromarray(np.zeros((600,800), dtype=np.uint8))))
    shr_tf_res.append({})

shr_args = mgr.Namespace()
gen_fps = mgr.Value('f', 12.0)
# Track two different counters - the number of the frame emitted from the camera, and the seq of the
# image that we put on the queue for processing
cam_seq = mgr.Value('i', 0)
vid_seq = mgr.Value('i', 0)
tf_seq = mgr.Value('i', 0)

# Prometheus metrics
prom_camera_images_captured = Counter('nsbeetle_camera_images_captured', 'Number of camera frames captured')
prom_vid_images_processed = Counter('nsbeetle_vid_images_processed', 'Number of video stills processed')
# TODO Create a histogram or summary that includes allows us to bucket the predictions based on the
# confidence value from TF. For now just use "low", "medium" and "high"
prom_objects_seen = Counter('nsbeetle_objects_seen', 'Number of objects seen', ['class', 'confidence'])

# Trying to adapt the ideas from:
# http://lifestyletransfer.com/how-to-use-gstreamer-appsink-in-python/

def on_buffer(sink: GstApp.AppSink, data: typing.Any) -> Gst.FlowReturn:
    global start_time
    """Callback on 'new-sample' signal"""
    if not start_time:
        start_time = datetime.datetime.now().timestamp()
    # Emit 'pull-sample' signal
    # https://lazka.github.io/pgi-docs/GstApp-1.0/classes/AppSink.html#GstApp.AppSink.signals.pull_sample

    # It seems like, even if we don't intend to do anything with the sample, we still have to emit it
    sample = sink.emit("pull-sample")  # Gst.Sample

    elapsed = datetime.datetime.now().timestamp() - start_time
    gen_fps.value = cam_seq.value / elapsed

    if not isinstance(sample, Gst.Sample):
        logging.error("Sample retrieved was corrupted")
        return Gst.FlowReturn.ERROR

    cam_seq.value += 1
    prom_camera_images_captured.inc()
    # FIXME Adjust based on incoming rate
    if cam_seq.value % math.ceil(12 / shr_args.target_fps) != 0:
        return Gst.FlowReturn.OK

    vid_seq.value += 1
    prom_vid_images_processed.inc()
    img_arr = extract_buffer(sample)
    logging.debug("Pushing image sample size: {}, incr: {}, q_sz: {}, est. fps: {:.1f}".format(
        img_arr.shape, cam_seq.value, img_q.qsize(), gen_fps.value
    ))
    img_q.put((vid_seq.value, img_arr))
    return Gst.FlowReturn.OK

def extract_buffer(sample: Gst.Sample) -> np.ndarray:
    """Extracts Gst.Buffer from Gst.Sample and converts to np.ndarray"""

    buffer = sample.get_buffer()  # Gst.Buffer
    caps = sample.get_caps()
    caps_format = sample.get_caps().get_structure(0)  # Gst.Structure

    # GstVideo.VideoFormat
    video_format = GstVideo.VideoFormat.from_string(caps_format.get_value('format'))
    w, h = caps_format.get_value('width'), caps_format.get_value('height')
    # c = utils.get_num_channels(video_format)
    buffer_size = buffer.get_size()

    # Change back to 800x900 when YUV decoding in place
    shape = (int(h * 1.5), w) # One channel? 800 * 600 = 720000
    # The YUV420 format will be uint8
    return np.ndarray(shape=shape, buffer=buffer.extract_dup(0, buffer_size), dtype=np.uint8)

def make_pipeline():
    """ Adapted from
    https://stackoverflow.com/questions/34688897/python-gstreamer-getting-meta-api-for-appsink-buffer"""

    # Need to experiment more with use of queues - these allow GStreamer to internally spin off threads which can
    # help extra more performance from a device like an Orange Pi which has 4 relatively weak cores, but it may be
    # creating overhead which ultimately doesn't improve performance.
    src =  Gst.ElementFactory.make("v4l2src")
    src.set_property("device", "/dev/video0")
    src.set_property("do-timestamp", 1)
    filt = Gst.ElementFactory.make("capsfilter")
    # filt.set_property("caps", Gst.caps_from_string("video/x-raw,format=NV12,width=640,height=480,framerate=20/1"))
    filt.set_property("caps", Gst.caps_from_string("video/x-raw,format=NV12,width=800,height=600,framerate=12/1"))
    p1 = Gst.ElementFactory.make("cedar_h264enc")
    # p1_q = Gst.ElementFactory.make("queue")
    p2 = Gst.ElementFactory.make("h264parse")
    p3 = Gst.ElementFactory.make("rtph264pay")
    p3.set_property("config-interval", 1)
    p3.set_property("pt", 96)
    p4 = Gst.ElementFactory.make("rtph264depay")
    p5 = Gst.ElementFactory.make("avdec_h264")
    sink = Gst.ElementFactory.make("appsink", "sink")
    pipeline_elements = [src, filt, p1, p2, p3, p4, p5, sink]

    sink.set_property("max-buffers", 10) # prevent the app to consume huge part of memory
    sink.set_property('emit-signals', True) #tell sink to emit signals
    sink.set_property('sync', False) #no sync to make decoding as fast as possible
    sink.connect("new-sample", on_buffer, sink)

    # Create an empty pipeline & add/link elements
    pipeline = Gst.Pipeline.new("test-pipeline")
    for elem in pipeline_elements:
        pipeline.add(elem)
    for i in range(len(pipeline_elements[:-1])):
        if not Gst.Element.link(pipeline_elements[i], pipeline_elements[i+1]):
            raise Exception("Elements {} and {} could not be linked.".format(
                pipeline_elements[i], pipeline_elements[i+1]))
    return pipeline

@app.route("/")
def index():
    ctx = {
        "raw_img_seq": cam_seq.value,
        "tf_seq": tf_seq.value,
        "gen_fps": gen_fps.value

    }
    return flask.render_template('index.html', **ctx)

# New experimental feed based on shared memory buffer
@app.route('/stream')
def stream_shm():
    """Video streaming home page."""
    return flask.render_template('stream.html')

def convert_yuv_to_rgb(img_arr):
    """ These 3 lines of code caused a lot of pain. Could not figure
    out any way of doing the YUV to RGB conversion without using OpenCV,
    which has convenience functions for this """
    rgb = cv2.cvtColor(img_arr, cv2.COLOR_YUV2BGR_I420)
    rgb = cv2.cvtColor(rgb, cv2.COLOR_BGR2RGB)
    return Image.fromarray(rgb)

@app.route("/pic")
def shm_pic():
    s = tf_seq.value - 1
    seq, img = shr_tf_img[s % IMG_BUF_LEN]
    logging.debug("Cur tf_seq: {}, seq {} from final image buffer".format(s, seq))
    img_io = io.BytesIO()
    img.save(img_io, 'JPEG', quality=85)
    img_io.seek(0)
    return flask.send_file(img_io, mimetype='image/jpeg', cache_timeout=0)

def bus_call(bus, message, loop):
    t = message.type
    if t == Gst.MessageType.EOS:
        logging.info("End-of-stream\n")
        loop.quit()
    elif t == Gst.MessageType.ERROR:
        err, debug = message.parse_error()
        logging.error("Trying to restart event loop", err, debug)
        time.sleep(1)
        loop.run()

    return True

################################################
# Streaming functions

def pil_image_to_base64(pil_image):
    buf = io.BytesIO()
    pil_image.save(buf, format="JPEG")
    b = base64.b64encode(buf.getvalue())
    return binascii.a2b_base64(b)

def frame_gen_shm():
    """Video streaming generator function based on frames in shared memory. It seems like
a limitation in Flask prevents this from running in more than one thread; will need to look
at something like Flask-SocketIO to get parallel video streams to work"""

    def calc_sleep():
        return 1 / shr_args.target_fps

    logging.info("Generating frames from shared memory! shr_args: {}".format(shr_args))
    while True:
        seq = (tf_seq.value - 1) % IMG_BUF_LEN
        img_seq, img = shr_tf_img[seq]
        frame = pil_image_to_base64(img)

        s = calc_sleep()
        logging.info("Cur tf_seq {}, calc'd seq {}, got seq from tf {}, sleep time {} ".format(tf_seq.value, seq, img_seq, s))
        time.sleep(s)

        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')

@app.route('/video_feed')
def video_feed():
    """Video streaming route. Put this in the src attribute of an img tag."""
    return flask.Response(frame_gen_shm(), mimetype='multipart/x-mixed-replace; boundary=frame')

@app.route('/metrics')
def mymetrics():
    """ Collect metrics manually rather than try to get DispatcherMiddleware to work """
    _update_metric_counters()
    logging.debug(prom_objects_seen.collect())
    return flask.Response(generate_latest(), mimetype='text/plain')

########

def worker():
    while True:
        seq, img_arr = img_q.get()
        img = convert_yuv_to_rgb(img_arr)
        if not shr_args.skip_od:
            results = detect_tflite.obj_detect_from_pil(img, shr_args.threshold, shr_args.classes)
            detect_tflite.draw_boxes(img, results)
        targ_seq = seq % IMG_BUF_LEN
        logging.debug("Processed img seq {} from queue, putting in pos {} , prev tf_seq {}".format(seq, targ_seq, tf_seq.value))
        tf_seq.value = seq
        shr_tf_img[targ_seq] = (seq, img)
        shr_tf_res[targ_seq] = results
        res_q.put(results)

def _update_metric_counters():
    while True:
        try:
            results = res_q.get_nowait()
            prom_camera_images_captured.inc()
            for r in results:
                if r['score'] > 0.70:
                    conf = 'High'
                elif r['score'] > 0.50:
                    conf = 'Medium'
                else:
                    conf = 'Low'
                prom_objects_seen.labels( coco_labels[r['class_id']], conf).inc()
        except queue.Empty:
            break

def main():
    GObject.threads_init()
    Gst.init(None)

    pipe = make_pipeline()

    # create and event loop and feed gstreamer bus mesages to it
    loop = GObject.MainLoop()

    bus = pipe.get_bus()
    bus.add_signal_watch()
    bus.connect("message", bus_call, loop)

    # start play back and listed to events
    pipe.set_state(Gst.State.PLAYING)
    try:
        loop.run()
    except:
        pass

    # cleanup
    pipe.set_state(Gst.State.NULL)

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Run nsbeetle')
    parser.add_argument('--target-fps', type=float, default=1, help='Target FPS')
    parser.add_argument('--skip-od', action='store_true', help='Whether to run object detection')
    parser.add_argument('--classes', nargs='+', type=int, default=None, help='List of object classes to show')
    parser.add_argument('--threshold', type=float, default=0.5, help='Probability threshold')
    parser.add_argument('--workers', type=int, default=1, help='Number of TF workers to spawn')
    parser.add_argument('--log-level', default='info', help='Logging debug level')

    shr_args = parser.parse_args()
    logging.basicConfig(level=shr_args.log_level.upper(), format='%(asctime)s [%(name)s] %(filename)s:%(lineno)s %(levelname)s: %(message)s')
    logging.debug("Shared args: {}".format(shr_args))

    with app.app_context():
        print("Attempting to run gstreamer in a separate Process...")
        # Running gstreamer in a separate process, then communicating via Queue with
        # worker processes that do the Tensorflow object detection heavy lifting
        # This allows much better use of the 4 CPU cores available on the orange pi in addition
        # to making it easier to terminate the forked process via the daemon flag
        t1 = multiprocessing.Process(target=main, args=(), daemon=True)
        for i in range(0, shr_args.workers):
            w = multiprocessing.Process(target=worker, args=(), daemon=True)
            w.start()
        t1.start()

    app.run(debug=True, use_reloader=False, port=8771, host="0.0.0.0", threaded=False)


import numpy as np
import cv2
import argparse
import Namespace

#class Namespace:
#    def __init__(self, **kwargs):
#        self.__dict__.update(kwargs)
#

class Object_Detection(object):

    def __init__(self):

        self.cnt = 0
        self.classes = None
        self.COLORS = None

    # function to get the output layer names
    # in the architecture
    def get_output_layers(self, net):

        layer_names = net.getLayerNames()

        output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

        return output_layers

    # function to draw bounding box on the detected object with class name

    def draw_prediction(self, img, class_id, confidence, x, y, x_plus_w, y_plus_h):

        label = str(self.classes[class_id])

        color = self.COLORS[class_id]

        cv2.rectangle(img, (x,y), (x_plus_w,y_plus_h), color, 2)

        cv2.putText(img, label, (x-10,y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

    # ## Webcam

    # # generate different colors for different classes
    # COLORS = np.random.uniform(0, 255, size=(len(classes), 3))

    def Detection(self, image):

        #args =    {'classes': 'yolov3.txt', 'config': 'yolov2-tiny.cfg', 'image': 'dog.jpg', 'weights'='yolov2-tiny.weights'}

        args = Namespace.Namespace(classes='yolov3.txt', config='yolov2-tiny.cfg', image='dog.jpg', weights='yolov2-tiny.weights')

#        cap = cv2.VideoCapture(0)
#        cnt = 0
#        while(True):
        # Capture frame-by-frame
#        ret, image = cap.read()
        self.cnt += 1
        # Our operations on the frame come here
        #gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        self.image = image

        Width = self.image.shape[1]
        Height = self.image.shape[0]
        scale = 0.00392

        # read class names from text file
        #classes = None

        with open(args.classes, 'r') as f:
            self.classes = [line.strip() for line in f.readlines()]

        # generate different colors for different classes
        self.COLORS = np.random.uniform(0, 255, size=(len(self.classes), 3))
        # read pre-trained model and config file
        net = cv2.dnn.readNet(args.weights, args.config)
        #create input blob
        blob = cv2.dnn.blobFromImage(self.image, scale, (416,416), (0,0,0), True, crop=False)
        # set input blob for the network
        net.setInput(blob)
        if self.cnt%1 == 0:
            # run inference through the network
            # and gather predictions from output layers
            outs = net.forward(self.get_output_layers(net))
            # initialization
            class_ids = []
            confidences = []
            boxes = []
            conf_threshold = 0.5
            nms_threshold = 0.4
            # for each detetion from each output layer
            # get the confidence, class id, bounding box params
            # and ignore weak detections (confidence < 0.5)

            for out in outs:
                for detection in out:
                    scores = detection[5:]
                    class_id = np.argmax(scores)
                    confidence = scores[class_id]
                    if confidence > 0.5:
                        center_x = int(detection[0] * Width)
                        center_y = int(detection[1] * Height)
                        w = int(detection[2] * Width)
                        h = int(detection[3] * Height)
                        x = center_x - w / 2
                        y = center_y - h / 2
                        class_ids.append(class_id)
                        confidences.append(float(confidence))
                        boxes.append([x, y, w, h])

            # apply non-max suppression
            indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)

            # go through the detections remaining
            # after nms and draw boundng box
            for i in indices:
                i = i[0]
                box = boxes[i]
                x = box[0]
                y = box[1]
                w = box[2]
                h = box[3]
                self.draw_prediction(self.image, class_ids[i], confidences[i], round(x), round(y), round(x+w), round(y+h))

        # Display the resulting frame (note : if change position, it doesn't draw bounding box)
        cv2.imshow('frame',image)

#        if cv2.waitKey(1) & 0xFF == ord('q'):
#            break
#
#    # When everything done, release the capture
#    cap.release()
#    cv2.destroyAllWindows()


"""
Plot file for the polar radar scatter chart
"""

import plotly
import plotly.graph_objs as go
import numpy as np

def plotPolarScatterChart(emotions, search_query, total_tweets):
    start = 1
    title = "Search Term: " + search_query
    divide = 1

    if total_tweets > 125:
        divide = 3

    trace1 = go.Scatter(
        r=np.random.uniform(start,emotions['joy'],size=emotions['joy']/divide),
        t=np.random.uniform(5,40,size=emotions['joy']),
        mode='markers',
        name='Joy',
        marker=dict(
            color='rgb(102,166,30)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace2 = go.Scatter(
        r=np.random.uniform(start,emotions['surprise'],size=emotions['surprise']/divide),
        t=np.random.uniform(50,85,size=emotions['surprise']),
        mode='markers',
        name='Surprise',
        marker=dict(
            color='rgb(27,158,119)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace3 = go.Scatter(
        r=np.random.uniform(start,emotions['anticipation'],size=emotions['anticipation']/divide),
        t=np.random.uniform(95,130,size=emotions['anticipation']),
        mode='markers',
        name='Anticipation',
        marker=dict(
            color='rgb(231,41,138)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace4 = go.Scatter(
        r=np.random.uniform(start,emotions['trust'],size=emotions['trust']/divide),
        t=np.random.uniform(140,175,size=emotions['trust']),
        mode='markers',
        name='Trust',
        marker=dict(
            color='rgb(117,112,179)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace5 = go.Scatter(
        r=np.random.uniform(start,emotions['sadness'],size=emotions['sadness']/divide),
        t=np.random.uniform(185,220,size=emotions['sadness']),
        mode='markers',
        name='Sadness',
        marker=dict(
            color='rgb(230,171,2)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace6 = go.Scatter(
        r=np.random.uniform(start,emotions['fear'],size=emotions['fear']/divide),
        t=np.random.uniform(230,265,size=emotions['fear']),
        mode='markers',
        name='Fear',
        marker=dict(
            color='rgb(81,81,81)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace7 = go.Scatter(
        r=np.random.uniform(start,emotions['anger'],size=emotions['anger']/divide),
        t=np.random.uniform(275,310,size=emotions['anger']),
        mode='markers',
        name='Anger',
        marker=dict(
            color='rgb(176,23,1)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    trace8 = go.Scatter(
        r = np.random.uniform(start,emotions['disgust'],size=emotions['disgust']/divide),
        t = np.random.uniform(320,355,size=emotions['disgust']),
        mode='markers',
        name='Disgust',
        marker=dict(
            color='rgb(217,95,2)',
            size=110,
            line=dict(
                color='white'
            ),
            opacity=0.7
        )
    )
    data = [trace1, trace2, trace3, trace4, trace5, trace6, trace7, trace8]
    layout = go.Layout(
        title=title,
        xaxis=dict(
            showticklabels='FALSE',
            showgrid='FALSE',
            showtickprefix='none',
            rangeselector=dict(
                visible='FALSE')),
        font=dict(
            size=15
        ),
        plot_bgcolor='rgb(223, 223, 223)',
        angularaxis=dict(
            tickcolor='rgb(253,253,253)'
        )
    )
    fig = go.Figure(data=data, layout=layout)
    plotly.offline.plot(fig, filename='static/graphs/polarScatterChart.html')


import cv2
import numpy as np
import time

class Yolo_detector():
    def __init__(self):
        print("Initializing Yolo detector class")

        self.CONFIDENCE = 0.5
        self.SCORE_THRESHOLD = 0.5
        self.IOU_THRESHOLD = 0.5

        # the neural network configuration
        config_path = r"F:\Git\MSc\CV_DL\yolo_custom\next_attempt\darknet\cfg\yolov3-custom.cfg"
        # the YOLO net weights file
        weights_path = r"F:\Git\MSc\CV_DL\yolo_custom\next_attempt\darknet\backup\yolov3-custom_final.weights"

        # loading all the class labels (objects)
        self.labels = open(r"F:\Git\MSc\CV_DL\yolo_custom\next_attempt\darknet\data\obj.names").read(
        ).strip().split("\n")
        # generating colors for each object for later plotting
        # colors = np.random.randint(
        # 0, 255, size=(len(labels), 3), dtype="uint8")

        # load the YOLO network
        self.net = cv2.dnn.readNetFromDarknet(config_path, weights_path)

    """
    # Name: detect_objects(self, image)
    # Description: receives and image and looks for a beetle or a dung ball with Yolo.
    # @args:
    #    all_points: the frame on which beetle or ball are to be searched for
    # @return
        list of objects:
        [{'label': 'Ball', 'box': [342, 174, 417, 234]}, {'label': 'Beetle', 'box': [356, 224, 391, 270]}]
                #x1, y1, x2, y2 = objects[0]["box"]
    """

    def detect_objects(self, image):
        # from https://cloudxlab.com/blog/object-detection-yolo-and-python-pydarknet/
        h, w = image.shape[:2]
        # create 4D blob
        blob = cv2.dnn.blobFromImage(
            image, 1/255.0, (416, 416), swapRB=True, crop=False)

        # sets the blob as the input of the network
        self.net.setInput(blob)

        # get all the layer names
        ln = self.net.getLayerNames()
        ln = [ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
        # feed forward (inference) and get the network output
        # measure how much it took in seconds
        start = time.perf_counter()
        layer_outputs = self.net.forward(ln)
        time_took = time.perf_counter() - start
        print(f"Time took to detect objects: {time_took:.2f}s")

        boxes, confidences, class_ids = [], [], []

        # loop over each of the layer outputs
        for output in layer_outputs:
            # loop over each of the object detections
            for detection in output:
                # extract the class id (label) and confidence (as a probability) of
                # the current object detection
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                # discard weak predictions by ensuring the detected
                # probability is greater than the minimum probability
                if confidence > self.CONFIDENCE:
                    # scale the bounding box coordinates back relative to the
                    # size of the image, keeping in mind that YOLO actually
                    # returns the center (x, y)-coordinates of the bounding
                    # box followed by the boxes' width and height
                    box = detection[:4] * np.array([w, h, w, h])
                    (centerX, centerY, width, height) = box.astype("int")

                    # use the center (x, y)-coordinates to derive the top and
                    # and left corner of the bounding box
                    x = int(centerX - (width / 2))
                    y = int(centerY - (height / 2))

                    # update our list of bounding box coordinates, confidences,
                    # and class IDs
                    boxes.append([x, y, int(width), int(height)])
                    confidences.append(float(confidence))
                    class_ids.append(class_id)

        # perform the non maximum suppression given the scores defined before
        idxs = cv2.dnn.NMSBoxes(
            boxes, confidences, self.SCORE_THRESHOLD, self.IOU_THRESHOLD)

        objects = []

        if len(idxs) > 0:
            for i in idxs.flatten():
                x, y = boxes[i][0], boxes[i][1]
                w, h = boxes[i][2], boxes[i][3]
                bounding_box = [x, y, x + w, y + h]
                objects.append({
                    "label": self.labels[class_ids[i]],
                    "box": bounding_box,
                    "confidence": confidences[i]
                })

        return objects


# Copyright 2022 The TEMPO Collaboration
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file  in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Module on for the original time evolving matrix product operator (TEMPO)
algorithm. This module is based on [Strathearn2017] and [Strathearn2018].

**[Strathearn2017]**
A. Strathearn, B.W. Lovett, and P. Kirton, *Efficient real-time path integrals
for non-Markovian spin-boson models*. New Journal of Physics, 19(9),
p.093009 (2017).

**[Strathearn2018]**
A. Strathearn, P. Kirton, D. Kilda, J. Keeling and
B. W. Lovett,  *Efficient non-Markovian quantum dynamics using
time-evolving matrix product operators*, Nat. Commun. 9, 3322 (2018).
"""

import sys
from typing import Callable, Dict, Optional, Text
import warnings
from copy import copy

import numpy as np
from numpy import ndarray
from scipy.linalg import expm
from oqupy.correlations import BaseCorrelations

from oqupy.bath import Bath
from oqupy.base_api import BaseAPIClass
from oqupy.config import NpDtype, MAX_DKMAX, DEFAULT_TOLERANCE
from oqupy.config import TEMPO_BACKEND_CONFIG
from oqupy.dynamics import Dynamics
from oqupy.system import BaseSystem
from oqupy.tempo.backends.tempo_backend import TempoBackend
from oqupy.operators import commutator, acommutator
from oqupy.util import get_progress

class TempoParameters(BaseAPIClass):
    r"""
    Parameters for the TEMPO computation.

    Parameters
    ----------
    dt: float
        Length of a time step :math:`\delta t`. - It should be small enough
        such that a trotterisation between the system Hamiltonian and the
        environment it valid, and the environment auto-correlation function
        is reasonably well sampled.
    dkmax: int
        Number of time steps :math:`K\in\mathbb{N}` that should be included in
        the non-Markovian memory. - It must be large
        enough such that :math:`\delta t \times K` is larger than the
        necessary memory time :math:`\tau_\mathrm{cut}`.
    epsrel: float
        The maximal relative error in the singular value truncation (done
        in the underlying tensor network algorithm). - It must be small enough
        such that the numerical compression (using tensor network algorithms)
        does not truncate relevant correlations.
    add_correlation_time: float
        Additional correlation time to include in the last influence
        functional as explained in [Strathearn2017].
    name: str (default = None)
        An optional name for the tempo parameters object.
    description: str (default = None)
        An optional description of the tempo parameters object.
    """
    def __init__(
            self,
            dt: float,
            dkmax: int,
            epsrel: float,
            add_correlation_time: Optional[float] = None,
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Create a TempoParameters object."""
        self.dt = dt
        self.dkmax = dkmax
        self.epsrel = epsrel
        self.add_correlation_time = add_correlation_time
        super().__init__(name, description)

    def __str__(self) -> Text:
        ret = []
        ret.append(super().__str__())
        ret.append("  dt                   = {} \n".format(self.dt))
        ret.append("  dkmax                = {} \n".format(self.dkmax))
        ret.append("  epsrel               = {} \n".format(self.epsrel))
        ret.append("  add_correlation_time = {} \n".format(
            self.add_correlation_time))
        return "".join(ret)

    @property
    def dt(self) -> float:
        """Length of a time step."""
        return self._dt

    @dt.setter
    def dt(self, new_dt: float) -> None:
        try:
            tmp_dt = float(new_dt)
        except Exception as e:
            raise AssertionError("Argument 'dt' must be float.") from e
        assert tmp_dt > 0.0, \
            "Argument 'dt' must be bigger than 0."
        self._dt = tmp_dt

    @property
    def dkmax(self) -> float:
        """Number of time steps that should be included in the non-Markovian
        memory. """
        return self._dkmax

    @dkmax.setter
    def dkmax(self, new_dkmax: float) -> None:
        try:
            if new_dkmax is None:
                tmp_dkmax = None
            else:
                tmp_dkmax = int(new_dkmax)
        except Exception as e:
            raise AssertionError("Argument 'dkmax' must be int or None.") \
                from e
        assert tmp_dkmax is None or tmp_dkmax > 0, \
            "Argument 'dkmax' must be bigger than or equal to 0 or None."
        self._dkmax = tmp_dkmax

    @dkmax.deleter
    def dkmax(self) -> None:
        self._dkmax = None

    @property
    def epsrel(self) -> float:
        """The maximal relative error in the singular value truncation."""
        return self._epsrel

    @epsrel.setter
    def epsrel(self, new_epsrel: float) -> None:
        try:
            tmp_epsrel = float(new_epsrel)
        except Exception as e:
            raise AssertionError("Argument 'epsrel' must be float.") from e
        assert tmp_epsrel > 0.0, \
            "Argument 'epsrel' must be bigger than 0."
        self._epsrel = tmp_epsrel

    @property
    def add_correlation_time(self) -> float:
        """
        Additional correlation time to include in the last influence
        functional.
        """
        return self._add_correlation_time

    @add_correlation_time.setter
    def add_correlation_time(self, new_tau: Optional[float] = None) -> None:
        if new_tau is None:
            del self.add_correlation_time
        else:
            # check input: cutoff
            try:
                tmp_new_tau = float(new_tau)
            except Exception as e:
                raise AssertionError( \
                    "Additional correlation time must be a float.") from e
            if tmp_new_tau < 0:
                raise ValueError(
                    "Additional correlation time must be non-negative.")
            self._add_correlation_time = tmp_new_tau

    @add_correlation_time.deleter
    def add_correlation_time(self) -> None:
        self._add_correlation_time = None

class Tempo(BaseAPIClass):
    """
    Class representing the entire TEMPO tensornetwork as introduced in
    [Strathearn2018].

    Parameters
    ----------
    system: BaseSystem
        The system.
    bath: Bath
        The Bath (includes the coupling operator to the system).
    parameters: TempoParameters
        The parameters for the TEMPO computation.
    initial_state: ndarray
        The initial density matrix of the system.
    start_time: float
        The start time.
    backend: str (default = None)
        The name of the backend to use for the computation. If
        `backend` is ``None`` then the default backend is used.
    backend_config: dict (default = None)
        The configuration of the backend. If `backend_config` is
        ``None`` then the default backend configuration is used.
    name: str (default = None)
        An optional name for the tempo object.
    description: str (default = None)
        An optional description of the tempo object.
    """
    def __init__(
            self,
            system: BaseSystem,
            bath: Bath,
            parameters: TempoParameters,
            initial_state: ndarray,
            start_time: float,
            backend_config: Optional[Dict] = None,
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Create a Tempo object. """

        assert isinstance(system, BaseSystem), \
            "Argument 'system' must be an instance of BaseSystem."
        self._system = system

        assert isinstance(bath, Bath), \
            "Argument 'bath' must be an instance of Bath."
        self._bath = bath

        self._correlations = self._bath.correlations

        assert isinstance(parameters, TempoParameters), \
            "Argument 'parameters' must be an instance of TempoParameters."
        self._parameters = parameters

        try:
            tmp_initial_state = np.array(initial_state, dtype=NpDtype)
            tmp_initial_state.setflags(write=False)
        except Exception as e:
            raise AssertionError("Initial state must be numpy array.") from e
        assert len(tmp_initial_state.shape) == 2, \
            "Initial state is not a matrix."
        assert tmp_initial_state.shape[0] == \
            tmp_initial_state.shape[1], \
            "Initial state is not a square matrix."
        self._initial_state = tmp_initial_state
        self._dimension = self._initial_state.shape[0]

        try:
            tmp_start_time = float(start_time)
        except Exception as e:
            raise AssertionError("Start time must be a float.") from e
        self._start_time = tmp_start_time

        if backend_config is None:
            self._backend_config = TEMPO_BACKEND_CONFIG
        else:
            self._backend_config = backend_config

        assert self._bath.dimension == self._dimension and \
            self._system.dimension == self._dimension, \
            "Hilbertspace dimensions are unequal: " \
            + "system ({}), ".format(self._system.dimension) \
            + "initial state ({}), ".format(self._dimension) \
            + "and bath coupling ({}), ".format(self._bath.dimension)

        super().__init__(name, description)

        tmp_coupling_comm = commutator(self._bath._coupling_operator)
        tmp_coupling_acomm = acommutator(self._bath._coupling_operator)
        self._coupling_comm = tmp_coupling_comm.diagonal()
        self._coupling_acomm = tmp_coupling_acomm.diagonal()

        self._dynamics = None
        self._backend_instance = None

        self._init_tempo_backend()

    def _init_tempo_backend(self):
        """Create and initialize the tempo backend. """
        dim = self._dimension
        initial_state = self._initial_state.reshape(dim**2)
        influence = self._influence
        unitary_transform = self._bath.unitary_transform
        propagators = self._propagators
        sum_north = np.array([1.0]*(dim**2))
        sum_west = np.array([1.0]*(dim**2))
        dkmax = self._parameters.dkmax
        epsrel = self._parameters.epsrel
        self._backend_instance = TempoBackend(
                initial_state,
                influence,
                unitary_transform,
                propagators,
                sum_north,
                sum_west,
                dkmax,
                epsrel,
                config=self._backend_config)

    def _init_dynamics(self):
        """Create a Dynamics object with metadata from the Tempo object. """
        name = None
        description = "computed from '{}' tempo".format(self.name)
        self._dynamics = Dynamics(name=name,
                                  description=description)

    def _influence(self, dk: int):
        """Create the influence functional matrix for a time step distance
        of dk. """
        return influence_matrix(
            dk,
            parameters=self._parameters,
            correlations=self._correlations,
            coupling_acomm=self._coupling_acomm,
            coupling_comm=self._coupling_comm)

    def _propagators(self, step: int):
        """Create the system propagators (first and second half) for the time
        step `step`. """
        dt = self._parameters.dt
        t = self._time(step)
        first_step = expm(self._system.liouvillian(t+dt/4.0)*dt/2.0)
        second_step = expm(self._system.liouvillian(t+dt*3.0/4.0)*dt/2.0)
        return first_step, second_step

    def _time(self, step: int):
        """Return the time that corresponds to the time step `step`. """
        return self._start_time + float(step)*self._parameters.dt

    @property
    def dimension(self) -> ndarray:
        """Hilbert space dimension. """
        return copy(self._dimension)

    def compute(
            self,
            end_time: float,
            progress_type: Text = None) -> Dynamics:
        """
        Propagate (or continue to propagate) the TEMPO tensor network to
        time `end_time`.

        Parameters
        ----------
        end_time: float
            The time to which the TEMPO should be computed.
        progress_type: str (default = None)
            The progress report type during the computation. Types are:
            {``'silent'``, ``'simple'``, ``'bar'``}. If `None` then
            the default progress type is used.

        Returns
        -------
        dynamics: Dynamics
            The instance of Dynamics associated with the TEMPO object.
        """
        try:
            tmp_end_time = float(end_time)
        except Exception as e:
            raise AssertionError("End time must be a float.") from e

        dim = self._dimension
        if self._backend_instance.step is None:
            step, state = self._backend_instance.initialize()
            self._init_dynamics()
            self._dynamics.add(self._time(step), state.reshape(dim, dim))

        start_step = self._backend_instance.step
        end_step = int((end_time - self._start_time)/self._parameters.dt)
        num_step = max(0, end_step - start_step)

        progress = get_progress(progress_type)
        with progress(num_step) as prog_bar:
            while self._time(self._backend_instance.step) < tmp_end_time:
                step, state = self._backend_instance.compute_step()
                self._dynamics.add(self._time(step), state.reshape(dim, dim))
                prog_bar.update(self._backend_instance.step - start_step)
            prog_bar.update(self._backend_instance.step - start_step)

        return self._dynamics

    def get_dynamics(self) -> Dynamics:
        """Returns the instance of Dynamics associated with the Tempo object.
        """
        return self._dynamics

def influence_matrix(
        dk: int,
        parameters: TempoParameters,
        correlations: BaseCorrelations,
        coupling_acomm: ndarray,
        coupling_comm: ndarray):
    """Compute the influence functional matrix. """
    dt = parameters.dt
    dkmax = parameters.dkmax

    if dk == 0:
        time_1 = 0.0
        time_2 = None
        shape = "upper-triangle"
    elif dk < 0:
        time_1 = float(dkmax) * dt
        if parameters.add_correlation_time is not None:
            time_2 = float(dkmax) * dt \
                + np.min([float(-dk) * dt,
                            1.0*dt + parameters.add_correlation_time])
        else:
            return None
        shape = "rectangle"
    else:
        time_1 = float(dk) * dt
        time_2 = None
        shape = "square"

    eta_dk = correlations.correlation_2d_integral( \
        delta=dt,
        time_1=time_1,
        time_2=time_2,
        shape=shape,
        epsrel=parameters.epsrel)
    op_p = coupling_acomm
    op_m = coupling_comm

    if dk == 0:
        infl = np.diag(np.exp(-op_m*(eta_dk.real*op_m \
                                        + 1j*eta_dk.imag*op_p)))
    else:
        infl = np.exp(-np.outer(eta_dk.real*op_m \
                                + 1j*eta_dk.imag*op_p, op_m))

    return infl

def _analyse_correlation(
        corr_func: Callable[[np.ndarray],np.ndarray],
        times: np.ndarray,
        corr_vals: np.ndarray):
    """Check correlation function on a finer grid."""
    additional_times = (times[:-1] + times[1:])/2.0
    additional_corr_vals = corr_func(additional_times)
    new_times = list(times)
    new_corr_vals = list(corr_vals)
    for i in range(len(additional_times)):
        new_times.insert(2*i+1,additional_times[i])
        new_corr_vals.insert(2*i+1,additional_corr_vals[i])

    errors = []
    integrals = []
    integral = 0.0

    for i in range(len(times)-1):
        dt = new_times[2*i+2] - new_times[2*i]

        rough_int = 0.5 * dt * (new_corr_vals[2*i] + new_corr_vals[2*i+2])
        fine_int = 0.5 * (rough_int + dt * new_corr_vals[2*i+1])
        error = np.abs(rough_int-fine_int)
        errors.append(error)

        rough_abs_int = 0.5 * dt \
                * (np.abs(new_corr_vals[2*i]) + np.abs(new_corr_vals[2*i+2]))
        fine_abs_int = 0.5 * (rough_abs_int + dt * np.abs(new_corr_vals[2*i+1]))
        integral += fine_abs_int
        integrals.append(integral)

    full_abs_integral = integrals[-1]

    new_times = np.array(new_times)
    new_corr_val = np.array(new_corr_vals)
    errors = np.array(errors) / full_abs_integral
    integrals = np.array(integrals) / full_abs_integral

    return new_times, new_corr_val, errors, integrals

def _estimate_epsrel(
        dkmax: int,
        tolerance: float) -> float:
    """Heuristic estimation of appropriate epsrel for TEMPO."""
    power = np.log(dkmax)/np.log(4)-np.log(tolerance)/np.log(10)
    return np.power(10,-power)

GUESS_WARNING_MSG = "Estimating parameters for TEMPO computation. " \
    + "No guarantee that resulting TEMPO computation converges towards " \
    + "the correct dynamics! " \
    + "Please refer to the TEMPO documentation and check convergence by " \
    + "varying the parameters for TEMPO manually."

MAX_DKMAX_WARNING_MSG = f"Reached maximal recommended `dkmax` ({MAX_DKMAX})! " \
    + "Interrupt TEMPO parameter estimation. "\
    + "Please choose a lower tolerance, or analyse the correlation function " \
    + "to choose TEMPO parameters manually. " \
    + "Could not reach specified tolerance! "

def guess_tempo_parameters(
        bath: Bath,
        start_time: float,
        end_time: float,
        system: Optional[BaseSystem] = None,
        tolerance: Optional[float] = DEFAULT_TOLERANCE) -> TempoParameters:
    """
    Function to roughly estimate appropriate parameters for a TEMPO
    computation.

    .. warning::

        No guarantee that resulting TEMPO calculation converges towards the
        correct dynamics! Please refer to the TEMPO documentation and check
        convergence by varying the parameters for TEMPO manually.

    Parameters
    ----------
    bath: Bath
        The bath.
    start_time: float
        The start time.
    end_time: float
        The time to which the TEMPO should be computed.
    system: BaseSystem
        The system.
    tolerance: float
        Tolerance for the parameter estimation.

    Returns
    -------
    tempo_parameters : TempoParameters
        Estimate of appropriate tempo parameters.
    """
    assert isinstance(bath, Bath), \
        "Argument 'bath' must be a oqupy.Bath object."
    try:
        tmp_start_time = float(start_time)
        tmp_end_time = float(end_time)
    except Exception as e:
        raise AssertionError("Start and end time must be a float.") from e
    if tmp_end_time <= tmp_start_time:
        raise ValueError("End time must be bigger than start time.")
    assert isinstance(system, (type(None), BaseSystem)), \
        "Argument 'system' must be 'None' or a oqupy.BaseSystem object."
    try:
        tmp_tolerance = float(tolerance)
    except Exception as e:
        raise AssertionError("Argument 'tolerance' must be float.") from e
    assert tmp_tolerance > 0.0, \
        "Argument 'tolerance' must be larger then 0."
    warnings.warn(GUESS_WARNING_MSG, UserWarning)
    print("WARNING: "+GUESS_WARNING_MSG, file=sys.stderr, flush=True)

    max_tau = tmp_end_time - tmp_start_time

    corr_func = np.vectorize(bath.correlations.correlation)
    new_times = np.linspace(0, max_tau, 11, endpoint=True)
    new_corr_vals = corr_func(new_times)
    times = new_times
    corr_vals = new_corr_vals

    while True:
        if len(new_times) > MAX_DKMAX:
            warnings.warn(MAX_DKMAX_WARNING_MSG, UserWarning)
            break
        times = new_times
        corr_vals = new_corr_vals
        new_times, new_corr_vals, errors, integrals = \
                _analyse_correlation(corr_func, times, corr_vals)
        cut = np.where(integrals>(1-tolerance))[0][0]
        cut = cut+2 if cut+2<=len(times) else len(times)
        times = times[:cut]
        corr_vals = corr_vals[:cut]
        new_times = new_times[:2*cut-1]
        new_corr_vals = new_corr_vals[:2*cut-1]
        if (errors < tolerance).all():
            break

    dt = np.min(times[1:] - times[:-1])
    dkmax = len(times)
    epsrel = _estimate_epsrel(dkmax, tolerance)
    sys.stderr.flush()

    return TempoParameters(
        dt=dt,
        dkmax=dkmax,
        epsrel=epsrel,
        name="Roughly estimated parameters",
        description="Estimated with 'guess_tempo_parameters()'")

def tempo_compute(
        system: BaseSystem,
        bath: Bath,
        initial_state: ndarray,
        start_time: float,
        end_time: float,
        parameters: Optional[TempoParameters] = None,
        tolerance: Optional[float] = DEFAULT_TOLERANCE,
        backend_config: Optional[Dict] = None,
        progress_type: Optional[Text] = None,
        name: Optional[Text] = None,
        description: Optional[Text] = None) -> Dynamics:
    """
    Shortcut for creating a Tempo object and running the computation.

    Parameters
    ----------
    system: BaseSystem
        The system.
    bath: Bath
        The Bath (includes the coupling operator to the system).
    initial_state: ndarray
        The initial density matrix of the system.
    start_time: float
        The start time.
    end_time: float
        The time to which the TEMPO should be computed.
    parameters: TempoParameters
        The parameters for the TEMPO computation.
    tolerance: float
        Tolerance for the parameter estimation (only applicable if
        `parameters` is None).
    backend_config: dict (default = None)
        The configuration of the backend. If `backend_config` is
        ``None`` then the default backend configuration is used.
    progress_type: str (default = None)
        The progress report type during the computation. Types are:
        {``'silent'``, ``'simple'``, ``'bar'``}.  If `None` then
        the default progress type is used.
    name: str (default = None)
        An optional name for the tempo object.
    description: str (default = None)
        An optional description of the tempo object.
    """
    if parameters is None:
        assert tolerance is not None, \
            "If 'parameters' is 'None' then 'tolerance' must be " \
            + "a positive float."
        parameters = guess_tempo_parameters(bath=bath,
                                            start_time=start_time,
                                            end_time=end_time,
                                            system=system,
                                            tolerance=tolerance)
    tempo = Tempo(system,
                  bath,
                  parameters,
                  initial_state,
                  start_time,
                  backend_config,
                  name,
                  description)
    tempo.compute(end_time, progress_type=progress_type)
    return tempo.get_dynamics()


import sys
import time
import os
import numpy as np
from math import ceil
from PIL import Image, ImageDraw
from models.tiny_yolo import TinyYoloNet
from utils import *
from darknet import Darknet
from torchvision import datasets, transforms
import dataset
from torch.utils.data import Dataset
import cv2

class MonopoleLoader(Dataset):
    """Custom dataset that includes image file paths. Extends
    torchvision.datasets.ImageFolder
    """
    def __init__(self, root, transform=None, batch_size=1, shape=(0, 0)):
        self.root = root
        self.transform = transform

        tmp = []
        sorted_files = sorted(os.listdir(root))
        for f in sorted_files:
            if f.split('.')[-1] != "png":
                continue
            tmp += [f]
        sorted_files = tmp

        #self.files = [[sorted_files[i + j*batch_size] for i in range(batch_size)] for j in range(ceil(len(sorted_files)/batch_size) - 1)]
        #self.files += [[sorted_files[-i-1] for i in range(len(sorted_files)%batch_size)]]
        self.files = [sorted_files[i] for i in range(len(sorted_files))]
        self.batch_size = batch_size

        if not (shape[0] or shape[1]):
            raise ValueError("Shape must have positive width or height")
        self.shape = shape

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

    def __getitem__(self, index):

        data = []
        #for f in self.files[index]:
        img = cv2.imread(os.path.join(self.root, self.files[index]))
        sized = cv2.resize(img, (self.shape[0], self.shape[1]))
        sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB)

        return torch.tensor(np.transpose(sized, (2,0,1))).float().div(255.0), self.files[index]

def dir_setup(imgfile, weight_num):
    if os.getcwd() != "/home/ashwin/pytorch-yolo2":
        print("predict.py must be run from /home/ashwin/pytorch-yolo2\nExiting...")
        sys.exit(0)

    # make a folder to put our predictions in if it does not yet exist
    if not os.path.isdir("./predictions"):
        os.mkdir("./predictions")
    print(imgfile)
    if os.path.isdir(imgfile):
        imgfolder = imgfile.strip("/").split('/')[-1] + "-" + weight_num
        if not os.path.isdir(os.path.join("predictions", imgfolder)):
            os.mkdir(os.path.join("predictions", imgfolder))
    elif imgfile.split("/")[-1].split('.')[1] == 'txt':
        filename = imgfile.split("/")[-1].split('.')[0]
        if not os.path.isdir(os.path.join("predictions", filename)):
            os.mkdir(os.path.join("predictions", filename))

def detect_cv2(cfgfile, weightfile, imgfile):
    m = Darknet(cfgfile)

    weight_num = weightfile.split('/')[-1].split('.')[0]

    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    if m.num_classes == 20:
        namesfile = 'data/voc.names'
    elif m.num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/monopoles.names'

    device = "cuda:0"

    use_cuda = 1
    if use_cuda:
        m.cuda().to(device)

    class_names = load_class_names(namesfile)

    dir_setup(imgfile, weight_num)

    if os.path.isdir(imgfile):

        batch_size = 32
        num_workers = 1
        init_width        = m.width
        init_height       = m.height

        kwargs = {'num_workers': num_workers, 'pin_memory': True} if use_cuda else {}

        data_loader = torch.utils.data.DataLoader(
            MonopoleLoader(imgfile,
                           transform=transforms.Compose([
                               transforms.ToTensor(),
                           ]),
            batch_size=batch_size, shape=(init_width, init_height)),
            batch_size=batch_size, **kwargs)

        imgfolder = imgfile.strip("/").split('/')[-1] + "-" + weight_num
        with open(os.path.join('predictions', imgfolder, '{}_list.txt'.format(imgfolder)), 'w') as out:

            for data, files in data_loader:
                if len(data.shape) == 3:
                    data = data.view(1, *data.shape)
                data = data.to(device)

                for i in range(2):
                    start = time.time()
                    boxes = do_detect(m, data, 0.5, 0.4, use_cuda)
                    finish = time.time()
                    if i == 0:
                        print('%i images predicted in %f seconds.' % (len(files), (finish-start)))
                for bx, f in zip(boxes, files):
                    if bx:
                        out.write(os.path.join(imgfile, f) + '\n')
                        imgname = f.split("/")[-1].split('.')[0]
                        plot_boxes_cv2(cv2.imread(os.path.join(imgfile, f)), bx, savename=os.path.join('predictions', imgfolder, '{}.jpg'.format(imgname)), class_names=class_names)

def detect(cfgfile, weightfile, imgfile):

    dir_setup(imgfile)

    m = Darknet(cfgfile)

    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    if m.num_classes == 20:
        namesfile = 'data/voc.names'
    elif m.num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/monopoles.names'

    use_cuda = 1
    if use_cuda:
        m.cuda()

    class_names = load_class_names(namesfile)

    if os.path.isfile(imgfile):
        filehandle = imgfile.split("/")[-1].split('.')[1]
        if filehandle == 'txt':
            filename = imgfile.split("/")[-1].split('.')[0]
            with open(imgfile, 'r') as infile:
                with open(os.path.join('predictions', filename, '{}_list.txt'.format(filename)), 'w') as f:
                    for files in infile:
                        files = files.strip('\n')
                        imgname = files.split("/")[-1].split('.')[0]
                        img = Image.open(files).convert('RGB')

                        sized = img.resize((m.width, m.height))

                        for i in range(2):
                            start = time.time()
                            boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
                            finish = time.time()
                            if i == 1:
                                print('%s: Predicted in %f seconds.' % (files, (finish-start)))

                        if boxes:
                            f.write(files + '\n')
                        plot_boxes(img, boxes, os.path.join('predictions', filename, '{}.jpg'.format(imgname)), class_names)
        else:
            imgname = imgfile.split("/")[-1].split('.')[0]
            img = Image.open(imgfile).convert('RGB')
            sized = img.resize((m.width, m.height))

            for i in range(2):
                start = time.time()
                boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
                finish = time.time()
                if i == 1:
                    print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))

            plot_boxes(img, boxes, os.path.join('predictions', '{}.jpg'.format(imgname)), class_names)
    elif os.path.isdir(imgfile):
        imgfolder = imgfile.strip("/").split('/')[-1]
        with open(os.path.join('predictions', imgfolder, '{}_list.txt'.format(imgfolder)), 'w') as f:
            for files in os.listdir(imgfile):
                if os.path.isdir(os.path.join(imgfile, files)):
                    continue
                imgname = files.split("/")[-1].split('.')[0]
                img = Image.open(os.path.join(imgfile, files)).convert('RGB')

                print(type(img))
                sys.exit(0)

                sized = img.resize((m.width, m.height))

                for i in range(2):
                    start = time.time()
                    boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
                    finish = time.time()
                    if i == 1:
                        print('%s: Predicted in %f seconds.' % (files, (finish-start)))
                if boxes:
                    f.write(os.path.join(imgfile, files) + '\n')
                plot_boxes(img, boxes, os.path.join('predictions', imgfolder, '{}.jpg'.format(imgname)), class_names)

def detect_skimage(cfgfile, weightfile, imgfile):
    from skimage import io
    from skimage.transform import resize
    m = Darknet(cfgfile)

    m.print_network()
    m.load_weights(weightfile)
    print('Loading weights from %s... Done!' % (weightfile))

    if m.num_classes == 20:
        namesfile = 'data/voc.names'
    elif m.num_classes == 80:
        namesfile = 'data/coco.names'
    else:
        namesfile = 'data/monopoles.names'

    use_cuda = 1
    if use_cuda:
        m.cuda()

    img = io.imread(imgfile)
    sized = resize(img, (m.width, m.height)) * 255

    for i in range(2):
        start = time.time()
        boxes = do_detect(m, sized, 0.5, 0.4, use_cuda)
        finish = time.time()
        if i == 1:
            print('%s: Predicted in %f seconds.' % (imgfile, (finish-start)))

    class_names = load_class_names(namesfile)
    plot_boxes_cv2(img, boxes, savename='predictions.jpg', class_names=class_names)

if __name__ == '__main__':
    if len(sys.argv) == 4:
        cfgfile = sys.argv[1]
        weightfile = sys.argv[2]
        imgfile = sys.argv[3]
        if imgfile[-1] == '/':
            imgfile = imgfile[:-1]

        # detect(cfgfile, weightfile, imgfile)
        detect_cv2(cfgfile, weightfile, imgfile)
        #detect_skimage(cfgfile, weightfile, imgfile)
    else:
        print('Usage: ')
        print('  python3 detect.py cfgfile weightfile imgfile')
        #detect('cfg/tiny-yolo-voc.cfg', 'tiny-yolo-voc.weights', 'data/person.jpg', version=1)