from __future__ import division
from __future__ import print_function

from absl import app, flags, logging
from absl.flags import FLAGS
import cv2
import glob
from gluoncv.model_zoo import get_model
from gluoncv.data.batchify import Tuple, Stack, Pad
import mxnet as mx
from mxnet import gluon
import logging
import numpy as np
import os
import random
from tqdm import tqdm

from datasets.pascalvoc import VOCDetection
from datasets.mscoco import COCODetection
from datasets.imgnetdet import ImageNetDetection
from datasets.imgnetvid import ImageNetVidDetection
from datasets.detectset import DetectSet
from datasets.combined import CombinedDetection

from metrics.pascalvoc import VOCMApMetric
from metrics.mscoco import COCODetectionMetric
from metrics.imgnetvid import VIDDetectionMetric

from models.definitions.yolo.transforms import YOLO3VideoInferenceTransform
from models.definitions.yolo.wrappers import yolo3_darknet53, yolo3_3ddarknet

from utils.general import as_numpy
from utils.image import cv_plot_bbox
from utils.video import video_to_frames

# disable autotune
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'

logging.basicConfig(level=logging.INFO)

flags.DEFINE_string('model_path', 'yolo3_darknet53_voc_best.params',
                    'Path to the detection model to use')
flags.DEFINE_string('network', 'darknet53',
                    'Base network name: darknet53')
flags.DEFINE_list('dataset', ['voc'],
                  'Dataset or .jpg image or .mp4 video or .txt image/video list.')
flags.DEFINE_string('trained_on', '',
                    'Dataset the model was trained on.')
flags.DEFINE_string('save_prefix', '0001',
                    'Model save prefix.')
flags.DEFINE_string('save_dir', 'results',
                    'Save directory to save images.')
flags.DEFINE_list('metrics', ['voc', 'coco'],
                  'List of metrics separated by , eg. voc,coco')
flags.DEFINE_integer('batch_size', 1,
                     'Batch size for detection: higher faster, but more memory intensive.')
flags.DEFINE_integer('data_shape', 416,
                     'Input data shape.')
flags.DEFINE_float('detection_threshold', 0.5,
                   'The threshold on detections to them being displayed.')
flags.DEFINE_integer('max_do', -1,
                     'Maximum samples to detect on. -1 is all.')
flags.DEFINE_float('every', 25,
                   'do every this many frames')
flags.DEFINE_list('window', '1, 1',
                  'Temporal window size of frames and the frame gap/stride of the windows samples')
flags.DEFINE_string('k_join_type', None,
                    'way to fuse k type, either max, mean, cat.')
flags.DEFINE_string('k_join_pos', None,
                    'position of k fuse, either early or late.')
flags.DEFINE_string('block_conv_type', '2',
                    "convolution type for the YOLO blocks: '2'2D, '3':3D or '21':2+1D, must be used with 'late' joining")
flags.DEFINE_string('rnn_pos', None,
                    "position of RNN, currently only supports 'late' or 'out")
flags.DEFINE_string('corr_pos', None,
                    "position of correlation features calculation, currently only supports 'early' or 'late")
flags.DEFINE_integer('corr_d', 4,
                     'The d value for the correlation filter.')
flags.DEFINE_string('motion_stream', None,
                    'Add a motion stream? can be flownet or r21d.')
flags.DEFINE_string('stream_gating', None,
                    'Use gating on the appearence stream using the motion stream. can be add or mul.')
flags.DEFINE_list('conv_types', [2, 2, 2, 2, 2, 2],
                  'Darknet Conv types for layers, either 2, 21, or 3 D')
flags.DEFINE_string('h_join_type', None,
                    'Type to join hierarchical darknet. can be max or conv.')
flags.DEFINE_list('hier', [1, 1, 1, 1, 1],
                  'the hierarchical factors, the input must be temporally equal to all these multiplied together')
flags.DEFINE_boolean('mult_out', False,
                     'Have one or multiple outs for timeseries data')
flags.DEFINE_boolean('temp', False,
                     'Use new temporal model')

flags.DEFINE_boolean('visualise', False,
                     'Do you want to display the detections?')
flags.DEFINE_boolean('per_frame_metric', False,
                     'Do you want to save out a per frame metric to the prediction files?')
flags.DEFINE_string('worst_video_path', None,
                    'Path to save video of worst case detections. If not None will require visualise and '
                    'per_frame_metric to have been done previously')
flags.DEFINE_boolean('display_gt', True,
                     'Do you want to display the ground truth boxes on the images?')
flags.DEFINE_boolean('model_agnostic', False,
                     'make the model class agnostic?')
flags.DEFINE_boolean('metric_agnostic', False,
                     'make the metric class agnostic?')

flags.DEFINE_list('gpus', [0],
                  'GPU IDs to use. Use comma or space for multiple eg. 0,1 or 0 1.')
flags.DEFINE_integer('num_workers', 8,
                     'The number of workers should be picked so that it’s equal to number of cores on your machine'
                     ' for max parallelization.')
flags.DEFINE_boolean('new_model', False,
                     'Use features Yolo (new) or stages Yolo (old)?')
flags.DEFINE_integer('offset', 0,
                     'If mult_out specified this selects the offset to test. Can be -2, -1, 0, 1, 2')
flags.DEFINE_integer('hier_level', 10,
                     'What is the hierarchical level cutoff for dets and eval 0,1,2,3,4,5,6?')

def get_dataset(dataset_name):
    datasets = list()
    if 'voc' in dataset_name:
        datasets.append(VOCDetection(splits=[(2007, 'test')], inference=True))

    if 'coco' in dataset_name:
        datasets.append(COCODetection(splits=['instances_val2017'], allow_empty=True, inference=True))

    if 'det' in dataset_name:
        datasets.append(ImageNetDetection(splits=['val'], allow_empty=False, inference=True))

    if 'vid' in dataset_name:
        datasets.append(ImageNetVidDetection(splits=[(2017, 'val')], allow_empty=True, every=FLAGS.every,
                                             window=FLAGS.window, inference=True, mult_out=FLAGS.mult_out))

    if len(datasets) == 0:
        assert len(dataset_name) > 0
        if dataset_name[0][-4:] == '.txt':  # list of images or list of videos
            with open(dataset_name[0], 'r') as f:
                files = [l.rstrip() for l in f.readlines()]
            if files[0][-4:] == '.mp4':  # list of videos
                img_list = list()
                for file in files:  # make frames in tmp folder
                    frames_dir = video_to_frames(file, os.path.join('data', 'tmp'),
                                                 os.path.join('data', 'tmp', 'stats'), overwrite=False)

                    img_list += glob.glob(frames_dir + '/**/*.jpg', recursive=True)

            elif files[0][-4:] == '.jpg':  # list of images
                img_list = files
            dataset = DetectSet(img_list)

        elif dataset_name[0][-4:] == '.jpg':  # single image
            dataset = DetectSet([dataset_name])

        elif dataset_name[0][-4:] == '.mp4':
            # make frames in tmp folder
            frames_dir = video_to_frames(dataset_name[0], os.path.join('data', 'tmp'),
                                         os.path.join('data', 'tmp', 'stats'), overwrite=False)
            img_list = glob.glob(frames_dir + '/**/*.jpg', recursive=True)
            dataset = DetectSet(img_list)
        else:
            raise NotImplementedError('Dataset: {} not implemented.'.format(dataset_name))
    elif len(datasets) == 1:
        dataset = datasets[0]
    else:
        dataset = CombinedDetection(datasets, class_tree=True, inference=True, hier_level=FLAGS.hier_level)

    return dataset

def get_dataloader(dataset, batch_size):
    width, height = FLAGS.data_shape, FLAGS.data_shape
    batchify_fn = Tuple(Stack(), Pad(pad_val=-1), Stack())
    loader = gluon.data.DataLoader(dataset.transform(YOLO3VideoInferenceTransform(width, height)),
                                   batch_size, False, last_batch='keep', num_workers=FLAGS.num_workers,
                                   batchify_fn=batchify_fn)
    return loader

def get_metric(dataset, metric_name, data_shape, save_dir, class_map=None):
    if metric_name.lower() == 'voc':
        metric = VOCMApMetric(iou_thresh=0.5, class_names=dataset.classes, class_map=class_map)

    elif metric_name.lower() == 'coco':
        metric = COCODetectionMetric(dataset, save_dir, cleanup=True, data_shape=None)

    elif metric_name.lower() == 'vid':
        metric = VIDDetectionMetric(dataset, iou_thresh=0.5, class_map=class_map,
                                    agnostic=FLAGS.metric_agnostic, offset=FLAGS.offset)

    else:
        raise NotImplementedError('Mertic: {} not implemented.'.format(metric_name))

    return metric

def detect(net, dataset, loader, ctx, max_do=-1):
    net.collect_params().reset_ctx(ctx)
    net.set_nms(nms_thresh=0.45, nms_topk=400)
    # net.hybridize()
    boxes = dict()
    if FLAGS.mult_out:
        boxes = [dict(), dict(), dict(), dict(), dict()]
    if max_do < 0:
        max_do = len(dataset)
    c = 0
    with tqdm(total=min(max_do, len(dataset)), desc="Detecting") as pbar:
        for ib, batch in enumerate(loader):

            data = gluon.utils.split_and_load(batch[0], ctx_list=ctx, batch_axis=0, even_split=False)
            label = gluon.utils.split_and_load(batch[1], ctx_list=ctx, batch_axis=0, even_split=False)
            idxs = gluon.utils.split_and_load(batch[2], ctx_list=ctx, batch_axis=0, even_split=False)
            det_bboxes = []
            det_ids = []
            det_scores = []
            # gt_bboxes = []
            # gt_ids = []
            # gt_difficults = []
            sidxs = []
            for x, y, sidx in zip(data, label, idxs):
                ids, scores, bboxes = net(x)
                det_ids.append(ids)
                det_scores.append(scores)
                # clip to image size
                det_bboxes.append(bboxes.clip(0, batch[0].shape[-1]))
                # split ground truths
                # gt_ids.append(y.slice_axis(axis=-1, begin=4, end=5))
                # gt_bboxes.append(y.slice_axis(axis=-1, begin=0, end=4))
                # gt_difficults.append(y.slice_axis(axis=-1, begin=5, end=6) if y.shape[-1] > 5 else None)
                sidxs.append(sidx)

            for id, score, box, sidx in zip(*[as_numpy(x) for x in [det_ids, det_scores, det_bboxes, sidxs]]):

                if FLAGS.mult_out:
                    files = dataset.window_paths(int(sidx))

                    for offset, file in enumerate(files):
                        if offset != 2 and file == files[2]:
                            continue  # we skip the offset frames if they are the same as the central frame, prevents repeating the boundary frames

                        valid_pred = np.where(id[offset].flat >= 0)[0]  # get the boxes that have a class assigned
                        box_o = box[offset, valid_pred, :] / batch[0].shape[-1]  # normalise boxes
                        id_o = id[offset].flat[valid_pred].astype(int)
                        score_o = score[offset].flat[valid_pred]

                        for id_, box_, score_ in zip(id_o, box_o, score_o):
                            if file in boxes[offset]:
                                boxes[offset][file].append([id_, score_] + list(box_))
                            else:
                                boxes[offset][file] = [[id_, score_] + list(box_)]

                else:
                    file = dataset.sample_path(int(sidx))

                    valid_pred = np.where(id.flat >= 0)[0]  # get the boxes that have a class assigned
                    box = box[valid_pred, :] / batch[0].shape[-1]  # normalise boxes
                    id = id.flat[valid_pred].astype(int)
                    score = score.flat[valid_pred]

                    for id_, box_, score_ in zip(id, box, score):
                        if file in boxes:
                            boxes[file].append([id_, score_]+list(box_))
                        else:
                            boxes[file] = [[id_, score_]+list(box_)]

            pbar.update(batch[0].shape[0])
            c += batch[0].shape[0]
            if c > max_do:
                break

    return boxes

def save_predictions(save_dir, dataset, boxes, overwrite=True, max_do=-1, agnostic=False):
    if agnostic:
        save_dir = os.path.join(save_dir, 'pred_ag')
    else:
        save_dir = os.path.join(save_dir, 'pred')

    if not overwrite and os.path.exists(save_dir):
        logging.info("Ground truth and prediction files already exist")

    os.makedirs(save_dir, exist_ok=True)

    if max_do < 0:
        max_do = len(dataset)

    for idx in tqdm(range(min(len(dataset), max_do)), desc="Saving out prediction .txts"):

        if FLAGS.mult_out:
            img_paths = dataset.window_paths(idx)

            for offset, img_path in enumerate(img_paths):

                file_id = img_path.split('/')[-1][:-4]
                if FLAGS.dataset == 'vid':
                    file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])
                    os.makedirs(os.path.join(save_dir, img_path.split('/')[-2]), exist_ok=True)

                with open(os.path.join(save_dir, file_id + '_' + str(offset-2) + '.txt'), 'w') as f:
                    if img_path in boxes[offset]:
                        for box in boxes[offset][img_path]:  # sid, class, score, box
                            f.write(
                                "{},{},{},{},{},{},{}\n".format(img_path, box[0], box[1], box[2], box[3], box[4], box[5]))
        else:

            img_path = dataset.sample_path(idx)

            if dataset.name == 'comb':
                dataset_idx, dataset_sample_idx = dataset._samples[dataset.sample_ids[idx]]
                inner_dataset = dataset._datasets[dataset_idx]
                if inner_dataset.name == 'vid':
                    file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])
                    os.makedirs(os.path.join(save_dir, img_path.split('/')[-2]), exist_ok=True)
                else:
                    dir_path, file_id = os.path.split(img_path)
                    file_id = file_id.split('.')[0]

            elif dataset.name == 'vid':
                file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])
                os.makedirs(os.path.join(save_dir, img_path.split('/')[-2]), exist_ok=True)
            else:
                dir_path, file_id = os.path.split(img_path)
                file_id = file_id.split('.')[0]

            with open(os.path.join(save_dir, file_id + '.txt'), 'w') as f:
                if img_path in boxes:
                    for box in boxes[img_path]:  # sid, class, score, box
                        f.write("{},{},{},{},{},{},{}\n".format(img_path, box[0], box[1], box[2], box[3], box[4], box[5]))

def load_predictions(save_dir, dataset, max_do=-1, metric=None, agnostic=False):
    if agnostic:
        save_dir = os.path.join(save_dir, 'pred_ag')
    else:
        save_dir = os.path.join(save_dir, 'pred')

    if metric is None:
        if not os.path.exists(save_dir):
            logging.error("Predictions directory does not exist {}".format(save_dir))
            return None
    else:
        if not os.path.exists(os.path.join(save_dir, 'metric')):
            logging.error("Predictions directory does not exist {}".format(os.path.join(save_dir, 'metric')))
            return None

    if FLAGS.mult_out:
        boxes = [dict(), dict(), dict(), dict(), dict()]
        for idx in tqdm(range(min(len(dataset), max_do)), desc="Loading in prediction .txts"):

            img_paths = dataset.window_paths(idx)

            for offset, img_path in enumerate(img_paths):

                file_id = img_path.split('/')[-1][:-4]
                if FLAGS.dataset == 'vid':
                    file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])

                add_metrics = False
                if metric is None:
                    file_path = os.path.join(save_dir, file_id + '_' + str(offset-2) + '.txt')
                    if not os.path.exists(file_path):
                        logging.error("Prediction file does not exist {}".format(file_path))
                        return None
                else:  # todo allow specific metrics
                    file_path = os.path.join(save_dir, 'metric', file_id + '_' + str(offset-2) + '.txt')
                    if not os.path.exists(file_path):
                        if not os.path.exists(os.path.join(save_dir, 'pred', file_id + '.txt')):
                            logging.error("Prediction file does not exist {}".format(file_path))
                            return None
                        else:
                            add_metrics = True

                with open(file_path, 'r') as f:
                    bb = [line.rstrip().split(',') for line in f.readlines()]

                if metric is not None and not add_metrics:
                    for box in bb:
                        if box[0] in boxes[offset]:
                            boxes[offset][box[0]].append([int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6]), float(box[7])])
                        else:
                            boxes[offset][box[0]] = [[int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6]), float(box[7])]]
                else:
                    for box in bb:
                        if box[0] in boxes[offset]:
                            boxes[offset][box[0]].append([int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6])])
                        else:
                            boxes[offset][box[0]] = [[int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6])]]

    else:
        boxes = dict()

        for idx in tqdm(range(min(len(dataset), max_do)), desc="Loading in prediction .txts"):
            img_path = dataset.sample_path(idx)

            if dataset.name == 'comb':
                dataset_idx, dataset_sample_idx = dataset._samples[dataset.sample_ids[idx]]
                inner_dataset = dataset._datasets[dataset_idx]
                if inner_dataset.name == 'vid':
                    file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])
                    os.makedirs(os.path.join(save_dir, img_path.split('/')[-2]), exist_ok=True)
                else:
                    dir_path, file_id = os.path.split(img_path)
                    file_id = file_id.split('.')[0]

            elif dataset.name == 'vid':
                file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])
                os.makedirs(os.path.join(save_dir, img_path.split('/')[-2]), exist_ok=True)
            else:
                dir_path, file_id = os.path.split(img_path)
                file_id = file_id.split('.')[0]

            add_metrics = False
            if metric is None:
                file_path = os.path.join(save_dir, file_id + '.txt')
                if not os.path.exists(file_path):
                    logging.error("Prediction file does not exist {}".format(file_path))
                    return None
            else:  # todo allow specific metrics
                file_path = os.path.join(save_dir, 'metric', file_id + '.txt')
                if not os.path.exists(file_path):
                    if not os.path.exists(os.path.join(save_dir, 'pred', file_id + '.txt')):
                        logging.error("Prediction file does not exist {}".format(file_path))
                        return None
                    else:
                        add_metrics = True

            with open(file_path, 'r') as f:
                bb = [line.rstrip().split(',') for line in f.readlines()]

            if metric is not None and not add_metrics:
                for box in bb:
                    if box[0] in boxes:
                        boxes[box[0]].append([int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6]), float(box[7])])
                    else:
                        boxes[box[0]] = [[int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6]), float(box[7])]]
            else:
                for box in bb:
                    if box[0] in boxes:
                        boxes[box[0]].append([int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6])])
                    else:
                        boxes[box[0]] = [[int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6])]]

    if add_metrics:
        boxes = add_metrics_to_predictions(save_dir, dataset, metric)

    return boxes

def add_metrics_to_predictions(load_dir, dataset, metric):

    if not os.path.exists(load_dir):
        logging.error("Predictions directory does not exist {}".format(load_dir))
        return None

    summary = dict()
    boxes = dict()
    for idx in tqdm(range(len(dataset)), desc="Adding metrics to predictions .txt"):
        img_path = dataset.sample_path(idx)

        file_id = img_path.split('/')[-1][:-4]
        if FLAGS.dataset == 'vid':
            file_id = os.path.join(img_path.split('/')[-2], img_path.split('/')[-1][:-5])

        if not os.path.exists(os.path.join(load_dir, file_id + '.txt')):
            logging.error("Prediction file does not exist {}".format(os.path.join(load_dir, file_id + '.txt')))
            return None

        # Load the predictions
        with open(os.path.join(load_dir, file_id + '.txt'), 'r') as f:
            bb = [line.rstrip().split(',') for line in f.readlines()]
        for box in bb:
            if box[0] in boxes:
                boxes[box[0]].append([int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6])])
            else:
                boxes[box[0]] = [[int(box[1]), float(box[2]), float(box[3]), float(box[4]), float(box[5]), float(box[6])]]

        # Run the metrics
        # get the gt boxes : [n_gpu, batch_size, samples, dim] : [1, 1, ?, 4 or 1]
        img, y, _ = dataset[idx]
        gt_bboxes = [np.expand_dims(y[:, :4], axis=0)]
        gt_ids = [np.expand_dims(y[:, 4], axis=0)]
        gt_difficults = [np.expand_dims(y[:, 5], axis=0) if y.shape[-1] > 5 else None]

        # get the predictions : [n_gpu, batch_size, samples, dim] : [1, 1, ?, 4 or 1]
        if img_path in boxes:
            det_bboxes = [[[[b[2] * img.shape[-2],  # change pred box dims to match image (unnormalise them)
                             b[3] * img.shape[-3],
                             b[4] * img.shape[-2],
                             b[5] * img.shape[-3]] for b in boxes[img_path]]]]
            det_ids = [[[[b[0]] for b in boxes[img_path]]]]
            det_scores = [[[[b[1]] for b in boxes[img_path]]]]

        metric.reset()
        metric.update(det_bboxes, det_ids, det_scores, gt_bboxes, gt_ids, gt_difficults)
        _, scores = metric.get()
        score = scores[-1]  # take the mean - which is the last score
        if FLAGS.dataset == 'vid':
            vid_id = img_path.split('/')[-2]
            if vid_id in summary:
                summary[vid_id].append(score)
            else:
                summary[vid_id] = [score]
        else:
            summary[img_path] = score

        # Save out the new detection file
        os.makedirs(os.path.join(load_dir, 'metric'), exist_ok=True)

        if FLAGS.dataset == 'vid':
            os.makedirs(os.path.join(load_dir, 'metric', img_path.split('/')[-2]), exist_ok=True)

        with open(os.path.join(load_dir, 'metric', file_id + '.txt'), 'w') as f:
            if img_path in boxes:
                for box in boxes[img_path]:  # sid, class, score, box
                    f.write("{},{},{},{},{},{},{},{}\n".format(img_path, box[0], box[1], box[2], box[3], box[4], box[5], score))
                    box.append(score)

    # generate a summary file listing ranking the worst clips
    if isinstance(summary[list(summary.keys())[0]], list):
        # need to sort on map first then number of frames, more frames ranked higher -> more wrong
        for k in summary.keys():
            summary[k] = [sum(summary[k])/len(summary[k]), len(summary[k])]
        summary_sorted = sorted(summary.items(), key=lambda kv: (kv[1][0], -kv[1][1]))
        for i in range(len(summary_sorted)):
            summary_sorted[i] = (summary_sorted[i][0], summary_sorted[i][1][0])
    else:
        summary_sorted = sorted(summary.items(), key=lambda kv: kv[1])

    with open(os.path.join(load_dir, 'metric', 'summary.txt'), 'w') as f:
        for ss in summary_sorted:
            f.write("{}\t{}\n".format(ss[0], ss[1]))
    return boxes

def visualise_predictions(save_dir, dataset, trained_on_dataset, boxes,
                          max_do=-1, display_gt=False, detection_threshold=0.5):
    colors = dict()
    for i in range(200):
        colors[i] = (int(256 * random.random()), int(256 * random.random()), int(256 * random.random()))
    colors_gt = dict()
    for i in range(200):
        colors_gt[i] = (0, 255, 0)

    if max_do < 0:
        max_do = len(dataset)

    for idx in tqdm(range(min(len(dataset), max_do)), desc="Saving out images"):

        img_path = dataset.sample_path(idx)
        img = cv2.imread(img_path)

        imgb, y, _ = dataset[idx]

        if display_gt and len(y) > 0:
            img = cv_plot_bbox(img=img,
                               bboxes=[list(g) for g in y[:, :4]],
                               scores=None,#[1]*len(y),
                               labels=[g for g in y[:, 4]],
                               thresh=detection_threshold,
                               colors=colors_gt,
                               class_names=dataset.classes,
                               absolute_coordinates=True)

        if img_path in boxes:
            img = cv_plot_bbox(img=img,
                               bboxes=[b[2:] for b in boxes[img_path]],
                               scores=[b[1] for b in boxes[img_path]],
                               labels=[b[0] for b in boxes[img_path]],
                               thresh=detection_threshold,
                               colors=colors,
                               class_names=trained_on_dataset.classes,
                               absolute_coordinates=False)

        os.makedirs(os.path.join(save_dir, 'vis'), exist_ok=True)
        if FLAGS.dataset == 'vid':
            os.makedirs(os.path.join(save_dir, 'vis', img_path.split('/')[-2]), exist_ok=True)
            cv2.imwrite(os.path.join(save_dir, 'vis', '/'.join(img_path.split('/')[-2:])), img)
        else:
            cv2.imwrite(os.path.join(save_dir, 'vis', img_path.split('/')[-1]), img)

def video_of_worst(video_path, frames_dir, summary_file=None, fps=4):
    assert fps < 25
    # add the .mp4 extension if it isn't already there
    if video_path[-4:] != ".mp4":
        video_path += ".mp4"

    files = list()
    summaries_dict = dict()
    # get the frame file paths
    if summary_file is None:
        for ext in [".jpg", ".png", ".jpeg", ".JPG", ".PNG", ".JPEG"]:
            files = glob.glob(frames_dir + "/**/*" + ext, recursive=True)
            if len(files) > 0:
                break
    else:
        with open(summary_file, 'r') as f:
            lines = f.readlines()
        summaries = [l.rstrip().split() for l in lines]

        for vid, score in summaries:
            for ext in [".jpg", ".png", ".jpeg", ".JPG", ".PNG", ".JPEG"]:
                frame_files = glob.glob(os.path.join(frames_dir, vid) + "/**/*" + ext, recursive=True)
                if len(frame_files) > 0:
                    break
            # sort the files alphabetically assuming this will do them in the correct order
            frame_files.sort()
            files += frame_files
            summaries_dict[vid] = score

    # couldn't find any images
    if not len(files) > 0:
        print("Couldn't find any files in {}".format(frames_dir))
        return None

    # make specific frame size and fit all videos in this frame, with rescale and centering
    height = 1080
    width = 1920

    # create the videowriter - will create an .mp4
    video = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), 25, (width, height))

    # load and write the frames to the video
    for filename in tqdm(files, desc="Generating Video {}".format(video_path)):
        blank_image = np.zeros((height, width, 3), np.uint8)
        if os.path.exists(filename):
            image = cv2.imread(filename)  # load the frame
        else:
            continue

        h, w, _ = image.shape
        ratio = min(height/h, width/w)
        hs = int(h*ratio)
        ws = int(w*ratio)
        hm = int(hs/2)
        wm = int(ws/2)
        height_m = int(height/2)
        width_m = int(width/2)

        image = cv2.resize(image, (ws, hs), interpolation=cv2.INTER_AREA)  # resize
        blank_image[height_m-hm:height_m+hm, width_m-wm:width_m+wm, :] = image[:2*hm, :2*wm, :]  # place in centre

        vid_id = filename.split('/')[-2]
        score = 'Clip AP: {:.2f}'.format(float(summaries_dict[vid_id]))

        cv2.putText(blank_image, '{:s}'.format(score), (1650, 1060), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2)
        cv2.putText(blank_image, '{:s}'.format(filename), (10, 1060), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 2)

        for _ in range(int(25/fps)):
            video.write(blank_image)  # write the frame to the video

    video.release()  # release the video

    return video_path

def evaluate(metrics, dataset, predictions):
    for idx in tqdm(range(len(dataset)), desc="Updating metrics with predictions"):

        img_path = dataset.sample_path(idx)
        sid = dataset.get_sample_ids()[idx]
        if FLAGS.mult_out:
            sid = dataset.get_sample_ids()[idx][FLAGS.offset+2]
            img_path = img_path[FLAGS.offset + 2]

        if img_path in predictions:
            gt_bboxes = None
            gt_ids = None
            gt_difficults = None
            if 'voc' in FLAGS.metrics:
                # get the gt boxes : [n_gpu, batch_size, samples, dim] : [1, 1, ?, 4 or 1]
                img, y, _ = dataset[idx]
                if FLAGS.mult_out:
                    img = img[FLAGS.offset+2]
                    y = y[FLAGS.offset+2]
                gt_bboxes = [np.expand_dims(y[:, :4], axis=0)]
                gt_ids = [np.expand_dims(y[:, 4], axis=0)]
                gt_difficults = [np.expand_dims(y[:, 5], axis=0) if y.shape[-1] > 5 else None]
            else:
                img = mx.image.imread(img_path)

            # get the predictions : [n_gpu, batch_size, samples, dim] : [1, 1, ?, 4 or 1]
            det_bboxes = [[[[b[2]*img.shape[-2],  # change pred box dims to match image (unnormalise them)
                             b[3]*img.shape[-3],
                             b[4]*img.shape[-2],
                             b[5]*img.shape[-3]] for b in predictions[img_path]]]]
            det_ids = [[[[b[0]] for b in predictions[img_path]]]]
            det_scores = [[[[b[1]] for b in predictions[img_path]]]]

            for metric in metrics:
                metric.update(det_bboxes, det_ids, det_scores, gt_bboxes, gt_ids, gt_difficults, sid=sid)

    return [metric.get() for metric in metrics]

def get_class_map(trained_on, eval_on):
    toc = trained_on.wn_classes
    eoc = eval_on.wn_classes

    class_map = []
    for c in eoc:
        if c in toc:
            class_map.append(toc.index(c))
        else:
            class_map.append(-1)

    return class_map

def iou(bb, bbgt):
    """
    Single box IoU calculation function for ious

    Args:
        bb: bounding box 1
        bbgt: bounding box 2

    Returns:
        float: the IoU between the two boxes
    """
    ov = 0
    iw = min(bb[2], bbgt[2]) - max(bb[0], bbgt[0]) + 1
    ih = min(bb[3], bbgt[3]) - max(bb[1], bbgt[1]) + 1
    if iw > 0 and ih > 0:
        # compute overlap as area of intersection / area of union
        intersect = iw * ih
        ua = (bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) + \
             (bbgt[2] - bbgt[0] + 1.) * \
             (bbgt[3] - bbgt[1] + 1.) - intersect
        ov = intersect / ua
    return ov

def hierarchical_nms(predictions, dataset, ov_thresh=0.5, conf_thresh=0.0, level_thresh=10):
    """combines boxes along same sub-branch and maxs the confidences as it goes from leaf to root"""
    levels = dataset.get_levels()
    parents = dataset.parents
    cls_map = dataset.wn_classes
    branch_matrix = list()
    for i in range(len(cls_map)):
        branch_list = list()
        for j in range(len(cls_map)):
            branch_list.append(dataset.on_branch(i, j))
        branch_matrix.append(branch_list)
    new_predictions = dict()

    level_thresh = max(0, level_thresh)  # ensure it at least 0, otherwise we could get stuck in infinite while loop

    # for each image/sample
    for img_path, boxes in tqdm(predictions.items(), desc='Performing hierarchical nms'):
        new_predictions[img_path] = list()

        # for each box in this sample
        for box in sorted(boxes, key=lambda x: x[0], reverse=True):  # sort so highest (most leafy) cls first
            cls = box[0]
            conf = box[1]
            coords = box[2:]

            if conf < conf_thresh:
                continue

            # assign the class as the level we want, priority is lowest
            while levels[cls] > level_thresh:
                cls = cls_map.index(parents[cls_map[cls]])

            # check for overlap with other box
            max_ov = 0
            max_idx = -1
            for idx, boxb in enumerate(new_predictions[img_path]):
                overlap = iou(coords, boxb[2:])
                if overlap > ov_thresh and overlap > max_ov:
                    max_ov = overlap
                    max_idx = idx

            if max_idx == -1:  # no overlapping boxes, add it as new box
                new_predictions[img_path].append([cls, conf] + coords)
            else:  # overlapping box
                boxb = new_predictions[img_path][max_idx]  # thanks to initial sorting this will never be parent of box
                # so this basically asks is boxb a child/grangchild/greatgrandchild/etc of box?
                if not branch_matrix[cls][boxb[0]]:  # not on branch, will add box as is a separate class
                    new_predictions[img_path].append([cls, conf] + coords)
                else:  # yes it is
                    if cls == boxb[0]:  # if same cls (eg. boxb's cls has been up'ed to box's cls), max the confs
                        new_predictions[img_path][max_idx][1] = max(new_predictions[img_path][max_idx][1], conf)
                    # otherwise ignore as we already have a child that meets the level and conf threshold reqs

    return new_predictions

def main(_argv):

    FLAGS.window = [int(s) for s in FLAGS.window]
    FLAGS.conv_types = [int(s) for s in FLAGS.conv_types]
    FLAGS.hier = [int(s) for s in FLAGS.hier]
    if FLAGS.model_agnostic:
        FLAGS.metric_agnostic = True

    if FLAGS.motion_stream == 'flownet':
        FLAGS.data_shape = 384  # cause 416 is a nasty shape

    FLAGS.dataset = FLAGS.dataset[0]

    if FLAGS.window[0] > 1:
        assert FLAGS.dataset == 'vid', 'If using window size >1 you can only use the vid dataset'

    # if we aren't given a full path, assume the file is in 'models/save_prefix' directory
    if len(os.path.split(FLAGS.model_path)[0]) > 0:
        model_path = FLAGS.model_path
    else:
        model_path = os.path.join('models', FLAGS.save_prefix, FLAGS.model_path)

    # check model exists
    if not os.path.exists(model_path):
        logging.error("Model doesn't appear where it's expected: {}".format(model_path))

    # get dataset
    dataset = get_dataset(FLAGS.dataset)

    # for use when model predictions are different to evaluation set classes
    if FLAGS.trained_on:
        trained_on_dataset = get_dataset(FLAGS.trained_on)
    else:
        trained_on_dataset = dataset

    # fix for tiny datasets of 1 or few elements
    batch_size = FLAGS.batch_size
    if len(dataset) < batch_size:
        batch_size = len(dataset)

    # handle gpu usage
    gpus = FLAGS.gpus
    if batch_size < len(gpus):
        gpus = [int(gpus[0])]

    # contexts
    ctx = [mx.gpu(int(i)) for i in gpus]
    ctx = ctx if ctx else [mx.cpu()]

    max_do = FLAGS.max_do
    if max_do < 0:
        max_do = len(dataset)

    # organise the save directories for the results
    if len(FLAGS.dataset) > 0:
        str_dataset = '_'.join(FLAGS.dataset)
        save_dir = os.path.join('models', 'experiments', FLAGS.save_prefix, FLAGS.save_dir, str_dataset)
    else:
        save_dir = os.path.join('models', 'experiments', FLAGS.save_prefix, FLAGS.save_dir)
    os.makedirs(save_dir, exist_ok=True)

    # attempt to load predictions

    per_sample_metric = None
    if FLAGS.worst_video_path is not None:
        FLAGS.per_frame_metric = True
    if FLAGS.per_frame_metric:
        per_sample_metric = get_metric(dataset, 'voc', FLAGS.data_shape, save_dir,
                                       class_map=get_class_map(trained_on_dataset, dataset))
    predictions = load_predictions(save_dir, dataset, max_do=max_do, metric=per_sample_metric,
                                   agnostic=FLAGS.model_agnostic)

    if predictions is None:  # id not exist detect and make
        # dataloader
        loader = get_dataloader(dataset, batch_size)

        # setup network
        # net_name = '_'.join(('yolo3', FLAGS.network, 'custom'))
        # net = get_model(net_name, root='models', pretrained_base=True, classes=trained_on_dataset.classes)
        if FLAGS.network == 'darknet53':
            if FLAGS.conv_types[0] is 2:
                net = yolo3_darknet53(trained_on_dataset.classes,
                                      k=FLAGS.window[0], k_join_type=FLAGS.k_join_type, k_join_pos=FLAGS.k_join_pos,
                                      block_conv_type=FLAGS.block_conv_type, rnn_pos=FLAGS.rnn_pos,
                                      corr_pos=FLAGS.corr_pos, corr_d=FLAGS.corr_d, motion_stream=FLAGS.motion_stream,
                                      agnostic=FLAGS.model_agnostic, add_type=FLAGS.stream_gating,
                                      new_model=FLAGS.new_model,
                                      hierarchical=FLAGS.hier, h_join_type=FLAGS.h_join_type, temporal=FLAGS.temp,
                                      t_out=FLAGS.mult_out)
            else:
                net = yolo3_3ddarknet(trained_on_dataset.classes, conv_types=FLAGS.conv_types)
        else:
            raise NotImplementedError('Backbone CNN model {} not implemented.'.format(FLAGS.network))
        net.initialize()
        if FLAGS.window[0] > 1:
            net.summary(mx.nd.random_normal(shape=(1, FLAGS.window[0], 3, FLAGS.data_shape, FLAGS.data_shape)))
        else:
            net.summary(mx.nd.random_normal(shape=(1, 3, FLAGS.data_shape, FLAGS.data_shape)))
        net.load_parameters(model_path)

        predictions = detect(net, dataset, loader, ctx, max_do=max_do)  # todo fix det thresh
        save_predictions(save_dir, dataset, predictions, agnostic=FLAGS.model_agnostic)

    if FLAGS.mult_out:
        predictions = predictions[FLAGS.offset+2]

    if isinstance(FLAGS.dataset, list) and len(FLAGS.dataset) > 1:
        predictions = hierarchical_nms(predictions, dataset, level_thresh=FLAGS.hier_level)

    if FLAGS.visualise:
        visualise_predictions(save_dir, dataset, trained_on_dataset, predictions,
                              max_do, display_gt=FLAGS.display_gt, detection_threshold=FLAGS.detection_threshold)

    if FLAGS.worst_video_path is not None:
        video_of_worst(FLAGS.worst_video_path, os.path.join(save_dir, "vis"),
                       summary_file=os.path.join(save_dir, 'metric', 'summary.txt'), fps=4)

    metrics = list()
    if FLAGS.metrics:
        for metric_name in FLAGS.metrics:
            if FLAGS.trained_on:  # for use when model preds are diff to eval set classes
                metrics.append(get_metric(dataset, metric_name, FLAGS.data_shape, save_dir,
                                          class_map=get_class_map(trained_on_dataset, dataset)))
            else:
                metrics.append(get_metric(dataset, metric_name, FLAGS.data_shape, save_dir))

        results = evaluate(metrics, dataset, predictions)

        for m, metric_name in enumerate(FLAGS.metrics):
            names, values = results[m]
            if FLAGS.metric_agnostic:
                metric_name += '_ag'
                if not FLAGS.model_agnostic:
                    metric_name += '_met'
            if FLAGS.offset != 0:
                metric_name += '_' + str(FLAGS.offset)
            with open(os.path.join(save_dir, metric_name+'.txt'), 'w') as f:
                for k, v in zip(names, values):
                    print(k, v)
                    f.write('{} {}\n'.format(k, v))

if __name__ == '__main__':

    try:
        app.run(main)
    except SystemExit:
        pass


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: train.py

import argparse
import itertools
import numpy as np
import os
import cv2
import six
import shutil

assert six.PY3, "FasterRCNN requires Python 3!"
import tensorflow as tf
import tqdm

import tensorpack.utils.viz as tpviz
from tensorpack import *
from tensorpack.tfutils import optimizer
from tensorpack.tfutils.common import get_tf_version_tuple, get_tensors_by_names
from tensorpack.tfutils.summary import add_moving_summary
from tensorpack.tfutils.varreplace import freeze_variables

import model_frcnn
import model_mrcnn
from basemodel import image_preprocess, resnet_c4_backbone, resnet_conv5, resnet_fpn_backbone, backbone_scope
from dataset import DetectionDataset
from config import finalize_configs, config as cfg
from data import get_all_anchors, get_all_anchors_fpn, get_train_dataflow
from eval_utils import EvalCallback
from model_box import RPNAnchors, clip_boxes, crop_and_resize, roi_align
from model_cascade import CascadeRCNNHead, CascadeRCNNHeadWithHardExamples
from model_fpn import fpn_model, generate_fpn_proposals, multilevel_roi_align, multilevel_rpn_losses
from model_frcnn import BoxProposals, FastRCNNHead, fastrcnn_outputs, fastrcnn_predictions, sample_fast_rcnn_targets
from model_mrcnn import maskrcnn_loss, maskrcnn_upXconv_head
from model_rpn import generate_rpn_proposals, rpn_head, rpn_losses

try:
    import horovod.tensorflow as hvd
except ImportError:
    pass

class DetectionModel(ModelDesc):
    def preprocess(self, image):
        image = tf.expand_dims(image, 0)
        image = image_preprocess(image, bgr=True)
        return tf.transpose(image, [0, 3, 1, 2])

    @property
    def training(self):
        return get_current_tower_context().is_training

    def optimizer(self):
        lr = tf.get_variable('learning_rate', initializer=0.003, trainable=False)
        tf.summary.scalar('learning_rate-summary', lr)

        # The learning rate in the config is set for 8 GPUs, and we use trainers with average=False.
        lr = lr / 8.
        opt = tf.train.MomentumOptimizer(lr, 0.9)
        if cfg.TRAIN.NUM_GPUS < 8:
            opt = optimizer.AccumGradOptimizer(opt, 8 // cfg.TRAIN.NUM_GPUS)
        return opt

    def get_inference_tensor_names(self):
        """
        Returns two lists of tensor names to be used to create an inference callable.

        Returns:
            [str]: input names
            [str]: output names
        """
        if cfg.MODE_THIRD_STAGE:
            out = ['output/boxes', 'output/scores', 'third_stage_features_out', 'ff_gt_tracklet_scores',
                   'sparse_tracklet_scores', 'tracklet_score_indices']
        else:
            out = ['output/boxes', 'output/scores', 'output/labels']
            if cfg.MODE_MASK:
                out.append('output/masks')
        if cfg.EXTRACT_GT_FEATURES:
            return ['image', 'roi_boxes'], ['boxes_for_extraction', 'features_for_extraction']
        else:
            return ['image'], out

    def build_graph(self, *inputs):
        inputs = dict(zip(self.input_names, inputs))

        image = self.preprocess(inputs['image'])  # 1CHW

        features = self.backbone(image)
        anchor_inputs = {k: v for k, v in inputs.items() if k.startswith('anchor_')}
        if cfg.EXTRACT_GT_FEATURES:
            anchor_inputs["roi_boxes"] = inputs["roi_boxes"]
        proposals, rpn_losses = self.rpn(image, features, anchor_inputs)  # inputs?

        targets = [inputs[k] for k in ['gt_boxes', 'gt_labels', 'gt_masks'] if k in inputs]
        head_losses = self.roi_heads(image, features, proposals, targets)

        if self.training:
            wd_cost = regularize_cost(
                '.*/W', l2_regularizer(cfg.TRAIN.WEIGHT_DECAY), name='wd_cost')
            total_cost = tf.add_n(
                rpn_losses + head_losses + [wd_cost], 'total_cost')
            add_moving_summary(total_cost, wd_cost)
            return total_cost

class ResNetC4Model(DetectionModel):
    def inputs(self):
        ret = [
            tf.placeholder(tf.float32, (None, None, 3), 'image'),
            tf.placeholder(tf.int32, (None, None, cfg.RPN.NUM_ANCHOR), 'anchor_labels'),
            tf.placeholder(tf.float32, (None, None, cfg.RPN.NUM_ANCHOR, 4), 'anchor_boxes'),
            tf.placeholder(tf.float32, (None, 4), 'gt_boxes'),
            tf.placeholder(tf.int64, (None,), 'gt_labels')]  # all > 0
        if cfg.MODE_MASK:
            ret.append(
                tf.placeholder(tf.uint8, (None, None, None), 'gt_masks')
            )  # NR_GT x height x width
        return ret

    def backbone(self, image):
        return [resnet_c4_backbone(image, cfg.BACKBONE.RESNET_NUM_BLOCKS[:3])]

    def rpn(self, image, features, inputs):
        featuremap = features[0]
        rpn_label_logits, rpn_box_logits = rpn_head('rpn', featuremap, cfg.RPN.HEAD_DIM, cfg.RPN.NUM_ANCHOR)
        anchors = RPNAnchors(get_all_anchors(), inputs['anchor_labels'], inputs['anchor_boxes'])
        anchors = anchors.narrow_to(featuremap)

        image_shape2d = tf.shape(image)[2:]  # h,w
        pred_boxes_decoded = anchors.decode_logits(rpn_box_logits)  # fHxfWxNAx4, floatbox
        proposal_boxes, proposal_scores = generate_rpn_proposals(
            tf.reshape(pred_boxes_decoded, [-1, 4]),
            tf.reshape(rpn_label_logits, [-1]),
            image_shape2d,
            cfg.RPN.TRAIN_PRE_NMS_TOPK if self.training else cfg.RPN.TEST_PRE_NMS_TOPK,
            cfg.RPN.TRAIN_POST_NMS_TOPK if self.training else cfg.RPN.TEST_POST_NMS_TOPK)

        if self.training:
            losses = rpn_losses(
                anchors.gt_labels, anchors.encoded_gt_boxes(), rpn_label_logits, rpn_box_logits)
        else:
            losses = []

        return BoxProposals(proposal_boxes), losses

    def roi_heads(self, image, features, proposals, targets):
        image_shape2d = tf.shape(image)[2:]  # h,w
        featuremap = features[0]

        gt_boxes, gt_labels, *_ = targets

        if self.training:
            # sample proposal boxes in training
            proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels)
        # The boxes to be used to crop RoIs.
        # Use all proposal boxes in inference

        boxes_on_featuremap = proposals.boxes * (1.0 / cfg.RPN.ANCHOR_STRIDE)
        roi_resized = roi_align(featuremap, boxes_on_featuremap, 14)

        feature_fastrcnn = resnet_conv5(roi_resized, cfg.BACKBONE.RESNET_NUM_BLOCKS[-1])  # nxcx7x7
        # Keep C5 feature to be shared with mask branch
        feature_gap = GlobalAvgPooling('gap', feature_fastrcnn, data_format='channels_first')
        fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs('fastrcnn', feature_gap, cfg.DATA.NUM_CLASS)

        fastrcnn_head = FastRCNNHead(proposals, fastrcnn_box_logits, fastrcnn_label_logits, gt_boxes,
                                     tf.constant(cfg.FRCNN.BBOX_REG_WEIGHTS, dtype=tf.float32))

        if self.training:
            all_losses = fastrcnn_head.losses()

            if cfg.MODE_MASK:
                gt_masks = targets[2]
                # maskrcnn loss
                # In training, mask branch shares the same C5 feature.
                fg_feature = tf.gather(feature_fastrcnn, proposals.fg_inds())
                mask_logits = maskrcnn_upXconv_head(
                    'maskrcnn', fg_feature, cfg.DATA.NUM_CATEGORY, num_convs=0)  # #fg x #cat x 14x14

                target_masks_for_fg = crop_and_resize(
                    tf.expand_dims(gt_masks, 1),
                    proposals.fg_boxes(),
                    proposals.fg_inds_wrt_gt, 14,
                    pad_border=False)  # nfg x 1x14x14
                target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1, 'sampled_fg_mask_targets')
                all_losses.append(maskrcnn_loss(mask_logits, proposals.fg_labels(), target_masks_for_fg))
            return all_losses
        else:
            decoded_boxes = fastrcnn_head.decoded_output_boxes()
            decoded_boxes = clip_boxes(decoded_boxes, image_shape2d, name='fastrcnn_all_boxes')
            label_scores = fastrcnn_head.output_scores(name='fastrcnn_all_scores')
            final_boxes, final_scores, final_labels = fastrcnn_predictions(
                decoded_boxes, label_scores, name_scope='output')

            if cfg.MODE_MASK:
                roi_resized = roi_align(featuremap, final_boxes * (1.0 / cfg.RPN.ANCHOR_STRIDE), 14)
                feature_maskrcnn = resnet_conv5(roi_resized, cfg.BACKBONE.RESNET_NUM_BLOCKS[-1])
                mask_logits = maskrcnn_upXconv_head(
                    'maskrcnn', feature_maskrcnn, cfg.DATA.NUM_CATEGORY, 0)  # #result x #cat x 14x14
                indices = tf.stack([tf.range(tf.size(final_labels)), tf.cast(final_labels, tf.int32) - 1], axis=1)
                final_mask_logits = tf.gather_nd(mask_logits, indices)  # #resultx14x14
                tf.sigmoid(final_mask_logits, name='output/masks')
            return []

class ResNetFPNModel(DetectionModel):

    def inputs(self):
        ret = [
            tf.placeholder(tf.float32, (None, None, 3), 'image')]
        num_anchors = len(cfg.RPN.ANCHOR_RATIOS)
        for k in range(len(cfg.FPN.ANCHOR_STRIDES)):
            ret.extend([
                tf.placeholder(tf.int32, (None, None, num_anchors),
                               'anchor_labels_lvl{}'.format(k + 2)),
                tf.placeholder(tf.float32, (None, None, num_anchors, 4),
                               'anchor_boxes_lvl{}'.format(k + 2))])
        ret.extend([
            tf.placeholder(tf.float32, (None, 4), 'gt_boxes'),
            tf.placeholder(tf.int64, (None,), 'gt_labels')])  # all > 0
        if cfg.MODE_MASK:
            ret.append(
                tf.placeholder(tf.uint8, (None, None, None), 'gt_masks')
            )  # NR_GT x height x width
        if cfg.EXTRACT_GT_FEATURES:
            ret.append(tf.placeholder(tf.float32, (None, 4,), 'roi_boxes'))
        return ret

    def slice_feature_and_anchors(self, p23456, anchors):
        for i, stride in enumerate(cfg.FPN.ANCHOR_STRIDES):
            with tf.name_scope('FPN_slice_lvl{}'.format(i)):
                anchors[i] = anchors[i].narrow_to(p23456[i])

    def backbone(self, image):
        c2345 = resnet_fpn_backbone(image, cfg.BACKBONE.RESNET_NUM_BLOCKS)
        p23456 = fpn_model('fpn', c2345)
        return p23456

    def rpn(self, image, features, inputs):
        if cfg.EXTRACT_GT_FEATURES:
            boxes = inputs['roi_boxes']
            return BoxProposals(boxes), tf.constant(0, dtype=tf.float32)

        assert len(cfg.RPN.ANCHOR_SIZES) == len(cfg.FPN.ANCHOR_STRIDES)

        image_shape2d = tf.shape(image)[2:]  # h,w
        all_anchors_fpn = get_all_anchors_fpn()
        multilevel_anchors = [RPNAnchors(
            all_anchors_fpn[i],
            inputs['anchor_labels_lvl{}'.format(i + 2)],
            inputs['anchor_boxes_lvl{}'.format(i + 2)]) for i in range(len(all_anchors_fpn))]
        self.slice_feature_and_anchors(features, multilevel_anchors)

        # Multi-Level RPN Proposals
        rpn_outputs = [rpn_head('rpn', pi, cfg.FPN.NUM_CHANNEL, len(cfg.RPN.ANCHOR_RATIOS))
                       for pi in features]
        multilevel_label_logits = [k[0] for k in rpn_outputs]
        multilevel_box_logits = [k[1] for k in rpn_outputs]
        multilevel_pred_boxes = [anchor.decode_logits(logits)
                                 for anchor, logits in zip(multilevel_anchors, multilevel_box_logits)]

        proposal_boxes, proposal_scores = generate_fpn_proposals(
            multilevel_pred_boxes, multilevel_label_logits, image_shape2d)

        if self.training:
            losses = multilevel_rpn_losses(
                multilevel_anchors, multilevel_label_logits, multilevel_box_logits)
        else:
            losses = []

        return BoxProposals(proposal_boxes), losses

    def roi_heads(self, image, features, proposals, targets):
        image_shape2d = tf.shape(image)[2:]  # h,w
        assert len(features) == 5, "Features have to be P23456!"
        gt_boxes, gt_labels, *_ = targets

        if self.training:
            proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels)

        fastrcnn_head_func = getattr(model_frcnn, cfg.FPN.FRCNN_HEAD_FUNC)
        if not cfg.FPN.CASCADE:
            roi_feature_fastrcnn = multilevel_roi_align(features[:4], proposals.boxes, 7)

            head_feature = fastrcnn_head_func('fastrcnn', roi_feature_fastrcnn)
            fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(
                'fastrcnn/outputs', head_feature, cfg.DATA.NUM_CLASS)
            fastrcnn_head = FastRCNNHead(proposals, fastrcnn_box_logits, fastrcnn_label_logits,
                                         gt_boxes, tf.constant(cfg.FRCNN.BBOX_REG_WEIGHTS, dtype=tf.float32))
        else:
            def roi_func(boxes):
                return multilevel_roi_align(features[:4], boxes, 7)

            fastrcnn_head = CascadeRCNNHead(
                proposals, roi_func, fastrcnn_head_func,
                (gt_boxes, gt_labels), image_shape2d, cfg.DATA.NUM_CLASS)

        if cfg.EXTRACT_GT_FEATURES:
            roi_feature_fastrcnn = multilevel_roi_align(features[:4], proposals.boxes, 7)
            tf.identity(roi_feature_fastrcnn, "rpn/feature")

        if self.training:
            all_losses = fastrcnn_head.losses()

            if cfg.MODE_MASK:
                gt_masks = targets[2]
                # maskrcnn loss
                roi_feature_maskrcnn = multilevel_roi_align(
                    features[:4], proposals.fg_boxes(), 14,
                    name_scope='multilevel_roi_align_mask')
                maskrcnn_head_func = getattr(model_mrcnn, cfg.FPN.MRCNN_HEAD_FUNC)
                mask_logits = maskrcnn_head_func(
                    'maskrcnn', roi_feature_maskrcnn, cfg.DATA.NUM_CATEGORY)  # #fg x #cat x 28 x 28

                target_masks_for_fg = crop_and_resize(
                    tf.expand_dims(gt_masks, 1),
                    proposals.fg_boxes(),
                    proposals.fg_inds_wrt_gt, 28,
                    pad_border=False)  # fg x 1x28x28
                target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1, 'sampled_fg_mask_targets')
                all_losses.append(maskrcnn_loss(mask_logits, proposals.fg_labels(), target_masks_for_fg))
            return all_losses
        else:
            decoded_boxes = fastrcnn_head.decoded_output_boxes()
            decoded_boxes = clip_boxes(decoded_boxes, image_shape2d, name='fastrcnn_all_boxes')
            label_scores = fastrcnn_head.output_scores(name='fastrcnn_all_scores')
            final_boxes, final_scores, final_labels = fastrcnn_predictions(
                decoded_boxes, label_scores, name_scope='output')
            if cfg.MODE_MASK:
                # Cascade inference needs roi transform with refined boxes.
                roi_feature_maskrcnn = multilevel_roi_align(features[:4], final_boxes, 14)
                maskrcnn_head_func = getattr(model_mrcnn, cfg.FPN.MRCNN_HEAD_FUNC)
                mask_logits = maskrcnn_head_func(
                    'maskrcnn', roi_feature_maskrcnn, cfg.DATA.NUM_CATEGORY)  # #fg x #cat x 28 x 28
                indices = tf.stack([tf.range(tf.size(final_labels)), tf.cast(final_labels, tf.int32) - 1], axis=1)
                final_mask_logits = tf.gather_nd(mask_logits, indices)  # #resultx28x28
                tf.sigmoid(final_mask_logits, name='output/masks')
            return []

class ResNetFPNTrackModel(ResNetFPNModel):
    def inputs(self):
        ret = super().inputs()
        if cfg.USE_PRECOMPUTED_REF_FEATURES:
            ret.append(tf.placeholder(tf.float32, (256, 7, 7), 'ref_features'))
        else:
            ret.append(tf.placeholder(tf.float32, (None, None, 3), 'ref_image'))
            ret.append(tf.placeholder(tf.float32, (4,), 'ref_box'))
        if cfg.MODE_THIRD_STAGE:
            ret.append(tf.placeholder(tf.float32, (256, 7, 7), 'ff_gt_tracklet_feat'))
            ret.append(tf.placeholder(tf.float32, (None, 256, 7, 7), 'active_tracklets_feats'))
            ret.append(tf.placeholder(tf.float32, (None, 4), 'active_tracklets_boxes'))
            ret.append(tf.placeholder(tf.float32, (), 'tracklet_distance_threshold'))
        if cfg.MODE_HARD_MINING:
            ret.append(tf.placeholder(tf.float32, (None, 3, 256, 7, 7), 'hard_negative_features'))
            if cfg.MODE_IF_HARD_MINING_THEN_ALSO_POSITIVES:
                ret.append(tf.placeholder(tf.float32, (None, 3, 256, 7, 7), 'hard_positive_features'))
                ret.append(tf.placeholder(tf.float32, (None, 3), 'hard_positive_ious'))
                ret.append(tf.placeholder(tf.float32, (None, 4), 'hard_positive_gt_boxes'))
                ret.append(tf.placeholder(tf.float32, (None, 3, 4), 'hard_positive_jitter_boxes'))
        if cfg.EXTRACT_GT_FEATURES:
            ret.append(tf.placeholder(tf.float32, (None, 4,), 'roi_boxes'))
        return ret

    def backbone(self, image):
        c2345 = resnet_fpn_backbone(image, cfg.BACKBONE.RESNET_NUM_BLOCKS)
        with backbone_scope(freeze=cfg.BACKBONE.FREEZE_AT > 3):
            p23456 = fpn_model('fpn', c2345)
        return p23456, c2345

    def rpn(self, image, features, inputs):
        if cfg.EXTRACT_GT_FEATURES:
            boxes = inputs['roi_boxes']
            return BoxProposals(boxes), tf.constant(0, dtype=tf.float32)

        if cfg.BACKBONE.FREEZE_AT > 3:
            with freeze_variables(stop_gradient=False, skip_collection=True):
                return super().rpn(image, features, inputs)
        else:
            return super().rpn(image, features, inputs)

    def roi_heads(self, image, ref_features, ref_box, features, proposals, targets, hard_negative_features=None,
                  hard_positive_features=None, hard_positive_ious=None, hard_positive_gt_boxes=None,
                  hard_positive_jitter_boxes=None, precomputed_ref_features=None):
        image_shape2d = tf.shape(image)[2:]  # h,w
        assert len(features) == 5, "Features have to be P23456!"
        gt_boxes, gt_labels, *_ = targets

        if self.training:
            proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels)

        fastrcnn_head_func = getattr(model_frcnn, cfg.FPN.FRCNN_HEAD_FUNC)
        if precomputed_ref_features is None:
            roi_aligned_ref_features = multilevel_roi_align(ref_features[:4], ref_box[tf.newaxis], 7)
        else:
            roi_aligned_ref_features = precomputed_ref_features[tf.newaxis]

        if cfg.MODE_SHARED_CONV_REDUCE:
            scope = tf.get_variable_scope()
        else:
            scope = ""

        assert cfg.FPN.CASCADE

        def roi_func(boxes, already_aligned_features=None):
            if already_aligned_features is None:
                aligned_features = multilevel_roi_align(features[:4], boxes, 7)
            else:
                # for hard example mining
                aligned_features = already_aligned_features
            tiled = tf.tile(roi_aligned_ref_features, [tf.shape(aligned_features)[0], 1, 1, 1])
            concat_features = tf.concat((tiled, aligned_features), axis=1)

            with argscope(Conv2D, data_format='channels_first',
                          kernel_initializer=tf.variance_scaling_initializer(
                              scale=2.0, mode='fan_out',
                              distribution='untruncated_normal' if get_tf_version_tuple() >= (1, 12) else 'normal')):
                with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
                    reduced_features = Conv2D('conv_reduce', concat_features, 256, 1, activation=None)
            return reduced_features

        if cfg.MODE_HARD_MINING and self.training:
            fastrcnn_head = CascadeRCNNHeadWithHardExamples(
                proposals, roi_func, fastrcnn_head_func,
                (gt_boxes, gt_labels), image_shape2d, cfg.DATA.NUM_CLASS, hard_negative_features,
                hard_positive_features, cfg.HARD_NEGATIVE_LOSS_SCALING_FACTOR,
                cfg.HARD_POSITIVE_LOSS_SCALING_FACTOR, hard_positive_ious, hard_positive_gt_boxes,
                hard_positive_jitter_boxes)
        else:
            fastrcnn_head = CascadeRCNNHead(
                proposals, roi_func, fastrcnn_head_func,
                (gt_boxes, gt_labels), image_shape2d, cfg.DATA.NUM_CLASS)

        if cfg.EXTRACT_GT_FEATURES:
            # get boxes and features for each of the three cascade stages!
            b0 = proposals.boxes
            b1, b2, _ = fastrcnn_head._cascade_boxes
            f0 = multilevel_roi_align(features[:4], b0, 7)
            f1 = multilevel_roi_align(features[:4], b1, 7)
            f2 = multilevel_roi_align(features[:4], b2, 7)
            tf.concat([b0, b1, b2], axis=0, name="boxes_for_extraction")
            tf.concat([f0, f1, f2], axis=0, name="features_for_extraction")

        if self.training:
            all_losses = fastrcnn_head.losses()

            if cfg.MODE_MASK:
                gt_masks = targets[2]
                # maskrcnn loss
                roi_feature_maskrcnn = multilevel_roi_align(
                    features[:4], proposals.fg_boxes(), 14,
                    name_scope='multilevel_roi_align_mask')
                maskrcnn_head_func = getattr(model_mrcnn, cfg.FPN.MRCNN_HEAD_FUNC)
                mask_logits = maskrcnn_head_func(
                    'maskrcnn', roi_feature_maskrcnn, cfg.DATA.NUM_CATEGORY)  # #fg x #cat x 28 x 28

                target_masks_for_fg = crop_and_resize(
                    tf.expand_dims(gt_masks, 1),
                    proposals.fg_boxes(),
                    proposals.fg_inds_wrt_gt, 28,
                    pad_border=False)  # fg x 1x28x28
                target_masks_for_fg = tf.squeeze(target_masks_for_fg, 1, 'sampled_fg_mask_targets')
                all_losses.append(maskrcnn_loss(mask_logits, proposals.fg_labels(), target_masks_for_fg))

            if cfg.MEASURE_IOU_DURING_TRAINING:
                decoded_boxes = fastrcnn_head.decoded_output_boxes()
                decoded_boxes = clip_boxes(decoded_boxes, image_shape2d, name='fastrcnn_all_boxes')
                label_scores = fastrcnn_head.output_scores(name='fastrcnn_all_scores')
                final_boxes, final_scores, final_labels = fastrcnn_predictions(
                    decoded_boxes, label_scores, name_scope='output_train')
                # if predictions are empty, this might break...
                # to prevent, stack dummy box
                boxes_for_iou = tf.concat([final_boxes[:1], tf.constant([[0.0, 0.0, 1.0, 1.0]],
                                                                        dtype=tf.float32)], axis=0)
                from examples.FasterRCNN.utils.box_ops import pairwise_iou
                iou_at_1 = tf.identity(pairwise_iou(gt_boxes[:1], boxes_for_iou)[0, 0], name="train_iou_at_1")
                add_moving_summary(iou_at_1)

            return all_losses
        else:
            decoded_boxes = fastrcnn_head.decoded_output_boxes()
            decoded_boxes = clip_boxes(decoded_boxes, image_shape2d, name='fastrcnn_all_boxes')
            label_scores = fastrcnn_head.output_scores(name='fastrcnn_all_scores')
            final_boxes, final_scores, final_labels = fastrcnn_predictions(
                decoded_boxes, label_scores, name_scope='output')
            if cfg.MODE_MASK:
                # Cascade inference needs roi transform with refined boxes.
                roi_feature_maskrcnn = multilevel_roi_align(features[:4], final_boxes, 14)
                maskrcnn_head_func = getattr(model_mrcnn, cfg.FPN.MRCNN_HEAD_FUNC)
                mask_logits = maskrcnn_head_func(
                    'maskrcnn', roi_feature_maskrcnn, cfg.DATA.NUM_CATEGORY)  # #fg x #cat x 28 x 28
                indices = tf.stack([tf.range(tf.size(final_labels)), tf.cast(final_labels, tf.int32) - 1], axis=1)
                final_mask_logits = tf.gather_nd(mask_logits, indices)  # #resultx28x28
                tf.sigmoid(final_mask_logits, name='output/masks')
            return []

    def build_graph(self, *inputs):
        inputs = dict(zip(self.input_names, inputs))
        image = self.preprocess(inputs['image'])  # 1CHW

        fpn_features, backbone_features = self.backbone(image)

        if cfg.USE_PRECOMPUTED_REF_FEATURES:
            ref_features = None
            ref_box = None
        else:
            ref_image = self.preprocess(inputs['ref_image'])  # 1CHW
            ref_box = inputs['ref_box']
            with tf.variable_scope(tf.get_variable_scope(), reuse=True):
                ref_features, _ = self.backbone(ref_image)

        anchor_inputs = {k: v for k, v in inputs.items() if k.startswith('anchor_')}
        if cfg.EXTRACT_GT_FEATURES:
            anchor_inputs["roi_boxes"] = inputs["roi_boxes"]
        proposals, rpn_losses = self.rpn(image, fpn_features, anchor_inputs)  # inputs?

        second_stage_features = fpn_features
        targets = [inputs[k] for k in ['gt_boxes', 'gt_labels', 'gt_masks'] if k in inputs]

        hard_negative_features = None
        hard_positive_features = None
        hard_positive_ious = None
        hard_positive_gt_boxes = None
        hard_positive_jitter_boxes = None
        if cfg.MODE_HARD_MINING:
            hard_negative_features = inputs['hard_negative_features']
            if cfg.MODE_IF_HARD_MINING_THEN_ALSO_POSITIVES:
                hard_positive_features = inputs['hard_positive_features']
                hard_positive_ious = inputs['hard_positive_ious']
                hard_positive_gt_boxes = inputs['hard_positive_gt_boxes']
                hard_positive_jitter_boxes = inputs['hard_positive_jitter_boxes']

        precomputed_ref_features = None
        if cfg.USE_PRECOMPUTED_REF_FEATURES:
            precomputed_ref_features = inputs['ref_features']

        # Extend proposals by previous frame detections
        if not self.training and cfg.MODE_THIRD_STAGE and cfg.EXTEND_PROPOSALS_BY_ACTIVE_TRACKLETS:
            proposal_boxes = proposals.boxes
            tracklet_boxes = inputs['active_tracklets_boxes']
            concat_boxes = tf.concat([proposal_boxes, tracklet_boxes], axis=0)
            proposals = BoxProposals(concat_boxes)

        head_losses = self.roi_heads(image, ref_features, ref_box, second_stage_features, proposals, targets,
                                     hard_negative_features, hard_positive_features, hard_positive_ious,
                                     hard_positive_gt_boxes, hard_positive_jitter_boxes,
                                     precomputed_ref_features=precomputed_ref_features)

        if cfg.MODE_THIRD_STAGE:
            self._run_third_stage(inputs, second_stage_features, tf.shape(image)[2:4])

        if self.training:
            wd_cost = regularize_cost(
                '.*/W', l2_regularizer(cfg.TRAIN.WEIGHT_DECAY), name='wd_cost')
            total_cost = tf.add_n(
                rpn_losses + head_losses + [wd_cost], 'total_cost')
            add_moving_summary(total_cost, wd_cost)
            return total_cost

    def _run_third_stage(self, inputs, second_stage_features, image_hw):
        boxes, scores = get_tensors_by_names(['output/boxes', 'output/scores'])
        # let's fix (as in finalize) the boxes, so we can roi align only one time
        aligned_features_curr = multilevel_roi_align(second_stage_features[:4], boxes, 7)
        # these also need to be extracted!
        aligned_features_curr = tf.identity(aligned_features_curr, name='third_stage_features_out')

        ff_gt_tracklet_scores, _ = self._score_for_third_stage(ref_feats=inputs['ff_gt_tracklet_feat'][tf.newaxis],
                                                               det_feats=aligned_features_curr)
        tf.identity(ff_gt_tracklet_scores, name='ff_gt_tracklet_scores')
        sparse_tracklet_scores, tracklet_score_indices = self._score_for_third_stage(
            ref_feats=inputs['active_tracklets_feats'], det_feats=aligned_features_curr,
            dense=False, ref_boxes=inputs['active_tracklets_boxes'], det_boxes=boxes, image_hw=image_hw,
            tracklet_distance_threshold=inputs['tracklet_distance_threshold'])
        tf.identity(sparse_tracklet_scores, name='sparse_tracklet_scores')
        tf.identity(tracklet_score_indices, name='tracklet_score_indices')

    def _score_for_third_stage(self, ref_feats, det_feats, dense=True, ref_boxes=None, det_boxes=None, image_hw=None,
                               tracklet_distance_threshold=0.08):
        # build all pairs
        n_refs = tf.shape(ref_feats)[0]
        n_dets = tf.shape(det_feats)[0]
        active_tracklets_tiled = tf.tile(ref_feats[:, tf.newaxis], multiples=[1, n_dets, 1, 1, 1])
        dets_tiled = tf.tile(det_feats[tf.newaxis], multiples=[n_refs, 1, 1, 1, 1])
        concated = tf.concat([active_tracklets_tiled, dets_tiled], axis=2)

        if not dense:
            # use boxes to prune the connectivity
            assert ref_boxes is not None
            assert det_boxes is not None
            assert image_hw is not None

            def xyxy_to_cxcywh(boxes_xyxy):
                wh = boxes_xyxy[:, 2:] - boxes_xyxy[:, :2]
                c = boxes_xyxy[:, :2] + wh / 2
                boxes_cwh = tf.concat((c, wh), axis=1)
                return boxes_cwh

            active_tracklets_boxes_cxcywh = xyxy_to_cxcywh(ref_boxes)
            boxes_cxcywh = xyxy_to_cxcywh(det_boxes)
            # normalize by image size
            h = image_hw[0]
            w = image_hw[1]
            norm = tf.cast(tf.stack([w, h, w, h], axis=0), tf.float32)
            diffs = tf.abs(active_tracklets_boxes_cxcywh[:, tf.newaxis] - boxes_cxcywh[tf.newaxis]) / norm[
                tf.newaxis, tf.newaxis]

            # use distances of boxes, first frame scores ("scores") to prune
            thresholds = tf.stack([tracklet_distance_threshold] * 4, axis=0)
            keep_mask = tf.reduce_all(diffs < thresholds, axis=2)

            indices = tf.where(keep_mask)
            flattened = tf.boolean_mask(concated, keep_mask)
        else:
            indices = None
            flattened = tf.reshape(
                concated, [tf.shape(concated)[0] * tf.shape(concated)[1]] + [int(x) for x in concated.shape[2:]])

        fastrcnn_head_func = getattr(model_frcnn, cfg.FPN.FRCNN_HEAD_FUNC)
        if cfg.MODE_SHARED_CONV_REDUCE:
            scope = tf.get_variable_scope()
        else:
            scope = ""
        all_posteriors = []
        # do this for every cascade stage
        for idx in range(3):
            with tf.variable_scope('cascade_rcnn_stage{}'.format(idx + 1), reuse=True):
                with argscope(Conv2D, data_format='channels_first'):
                    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
                        reduced_features = Conv2D('conv_reduce', flattened, 256, 1, activation=None)
                    head_feats = fastrcnn_head_func('head', reduced_features)
                    with tf.variable_scope('outputs_new', reuse=True):
                        classification = FullyConnected('class', head_feats, 2)
                        posteriors = tf.nn.softmax(classification)
                        all_posteriors.append(posteriors)
        posteriors = (all_posteriors[0] + all_posteriors[1] + all_posteriors[2]) / tf.constant(3.0, dtype=tf.float32)
        scores = posteriors[:, 1]
        return scores, indices

    def get_inference_tensor_names(self):
        inp, out = super().get_inference_tensor_names()
        if cfg.USE_PRECOMPUTED_REF_FEATURES:
            inp.append('ref_features')
        else:
            inp.append('ref_image')
            inp.append('ref_box')
        if cfg.MODE_THIRD_STAGE:
            inp.append('ff_gt_tracklet_feat')
            inp.append('active_tracklets_feats')
            inp.append('active_tracklets_boxes')
            inp.append('tracklet_distance_threshold')
        return inp, out

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--load', help='load a model for evaluation or training. Can overwrite BACKBONE.WEIGHTS')
    parser.add_argument('--logdir', help='log directory', default='train_log/siamrcnn')
    parser.add_argument('--config', help="A list of KEY=VALUE to overwrite those defined in config.py",
                        nargs='+')

    if get_tf_version_tuple() < (1, 6):
        # https://github.com/tensorflow/tensorflow/issues/14657
        logger.warn("TF<1.6 has a bug which may lead to crash in FasterRCNN if you're unlucky.")

    args = parser.parse_args()
    if args.config:
        cfg.update_args(args.config)

    MODEL = ResNetFPNTrackModel()
    DetectionDataset()  # initialize the config with information from our dataset

    is_horovod = cfg.TRAINER == 'horovod'
    if is_horovod:
        hvd.init()
        logger.info("Horovod Rank={}, Size={}".format(hvd.rank(), hvd.size()))

    if not is_horovod or hvd.rank() == 0:
        # keep the old log folder if already existing! (before it would just delete it)
        logger.set_logger_dir(args.logdir, 'k')
        # logger.set_logger_dir(args.logdir, 'd')

    finalize_configs(is_training=True)
    stepnum = cfg.TRAIN.STEPS_PER_EPOCH

    # warmup is step based, lr is epoch based
    init_lr = cfg.TRAIN.WARMUP_INIT_LR * min(8. / cfg.TRAIN.NUM_GPUS, 1.)
    warmup_schedule = [(0, init_lr), (cfg.TRAIN.WARMUP, cfg.TRAIN.BASE_LR)]
    warmup_end_epoch = cfg.TRAIN.WARMUP * 1. / stepnum
    lr_schedule = [(int(warmup_end_epoch + 0.5), cfg.TRAIN.BASE_LR)]

    factor = 8. / cfg.TRAIN.NUM_GPUS
    for idx, steps in enumerate(cfg.TRAIN.LR_SCHEDULE[:-1]):
        mult = 0.1 ** (idx + 1)
        lr_schedule.append(
            (steps * factor // stepnum, cfg.TRAIN.BASE_LR * mult))
    logger.info("Warm Up Schedule (steps, value): " + str(warmup_schedule))
    logger.info("LR Schedule (epochs, value): " + str(lr_schedule))
    train_dataflow = get_train_dataflow()
    # This is what's commonly referred to as "epochs"
    total_passes = cfg.TRAIN.LR_SCHEDULE[-1] * 8 / train_dataflow.size()
    logger.info("Total passes of the training set is: {:.5g}".format(total_passes))

    callbacks = [
                    PeriodicCallback(
                        ModelSaver(max_to_keep=10, keep_checkpoint_every_n_hours=1),
                        # every_k_epochs=1),
                        every_k_epochs=20),
                    # linear warmup
                    ScheduledHyperParamSetter(
                        'learning_rate', warmup_schedule, interp='linear', step_based=True),
                    ScheduledHyperParamSetter('learning_rate', lr_schedule),
                    PeakMemoryTracker(),
                    EstimatedTimeLeft(median=True),
                    SessionRunTimeout(60000).set_chief_only(True),  # 1 minute timeout
                ] + [
                    EvalCallback(dataset, *MODEL.get_inference_tensor_names(), args.logdir)
                    for dataset in cfg.DATA.VAL
                ]
    if not is_horovod:
        callbacks.append(GPUUtilizationTracker())

    start_epoch = cfg.TRAIN.STARTING_EPOCH
    if is_horovod and hvd.rank() > 0:
        session_init = None
    else:
        # first try to find existing model
        checkpoint_path = os.path.join(args.logdir, "checkpoint")
        if os.path.exists(checkpoint_path):
            session_init = get_model_loader(checkpoint_path)
            start_step = int(session_init.path.split("-")[-1])
            start_epoch = start_step // stepnum
            logger.info(
                "initializing from existing model, " + session_init.path + ", starting from epoch " + str(start_epoch))
        else:
            if args.load:
                session_init = get_model_loader(args.load)
            else:
                session_init = get_model_loader(cfg.BACKBONE.WEIGHTS) if cfg.BACKBONE.WEIGHTS else None

    max_epoch = min(cfg.TRAIN.LR_SCHEDULE[-1] * factor // stepnum, cfg.TRAIN.MAX_NUM_EPOCHS)

    traincfg = TrainConfig(
        model=MODEL,
        data=QueueInput(train_dataflow),
        callbacks=callbacks,
        steps_per_epoch=stepnum,
        # max_epoch=cfg.TRAIN.LR_SCHEDULE[-1] * factor // stepnum,
        max_epoch=max_epoch,
        session_init=session_init,
        starting_epoch=start_epoch
    )
    if is_horovod:
        trainer = HorovodTrainer(average=False)
    else:
        # nccl mode appears faster than cpu mode
        trainer = SyncMultiGPUTrainerReplicated(cfg.TRAIN.NUM_GPUS, average=False, mode='nccl')
    launch_train_with_config(traincfg, trainer)


import queue
import time

import numpy as np
import voluptuous as vol

from ledfx.color import COLORS, GRADIENTS
from ledfx.effects.audio import AudioReactiveEffect
from ledfx.effects.gradient import GradientEffect

class Strobe(AudioReactiveEffect, GradientEffect):

    NAME = "Real Strobe"
    CONFIG_SCHEMA = vol.Schema(
        {
            vol.Optional(
                "gradient_name",
                description="Color scheme for bass strobe to cycle through",
                default="Dancefloor",
            ): vol.In(list(GRADIENTS.keys())),
            vol.Optional(
                "color_step",
                description="Amount of color change per bass strobe",
                default=0.0625,
            ): vol.All(vol.Coerce(float), vol.Range(min=0, max=0.25)),
            vol.Optional(
                "bass_threshold",
                description="Cutoff for quiet sounds. Higher -> only loud sounds are detected",
                default=0.4,
            ): vol.All(vol.Coerce(float), vol.Range(min=0, max=1)),
            vol.Optional(
                "bass_strobe_decay_rate",
                description="Bass strobe decay rate. Higher -> decays faster.",
                default=0.5,
            ): vol.All(vol.Coerce(float), vol.Range(min=0, max=1)),
            vol.Optional(
                "strobe_color",
                description="Colour for note strobes",
                default="white",
            ): vol.In(list(COLORS.keys())),
            vol.Optional(
                "strobe_width",
                description="Note strobe width, in pixels",
                default=10,
            ): vol.All(vol.Coerce(int), vol.Range(min=0, max=1000)),
            vol.Optional(
                "strobe_decay_rate",
                description="Note strobe decay rate. Higher -> decays faster.",
                default=0.5,
            ): vol.All(vol.Coerce(float), vol.Range(min=0, max=1)),
        }
    )

    def activate(self, pixel_count):
        super().activate(pixel_count)
        self.strobe_overlay = np.zeros(np.shape(self.pixels))
        self.bass_strobe_overlay = np.zeros(np.shape(self.pixels))
        self.onsets_queue = queue.Queue()

    def config_updated(self, config):
        self.bass_threshold = self._config["bass_threshold"]
        self.color_shift_step = self._config["color_step"]

        self.strobe_color = np.array(
            COLORS[self._config["strobe_color"]], dtype=float
        )
        self.last_color_shift_time = 0
        self.strobe_width = self._config["strobe_width"]
        self.color_shift_delay_in_seconds = 1
        self.color_idx = 0

        self.last_strobe_time = 0
        self.strobe_wait_time = 0
        self.strobe_decay_rate = 1 - self._config["strobe_decay_rate"]

        self.last_bass_strobe_time = 0
        self.bass_strobe_wait_time = 0
        self.bass_strobe_decay_rate = (
            1 - self._config["bass_strobe_decay_rate"]
        )

    def get_pixels(self):
        pixels = np.copy(self.bass_strobe_overlay)

        if not self.onsets_queue.empty():
            self.onsets_queue.get()
            strobe_width = min(self.strobe_width, self.pixel_count)
            length_diff = self.pixel_count - strobe_width
            position = (
                0
                if length_diff == 0
                else np.random.randint(self.pixel_count - strobe_width)
            )
            self.strobe_overlay[
                position : position + strobe_width
            ] = self.strobe_color

        pixels += self.strobe_overlay

        self.strobe_overlay *= self.strobe_decay_rate
        self.bass_strobe_overlay *= self.bass_strobe_decay_rate
        self.pixels = pixels
        return self.pixels

    def audio_data_updated(self, data):
        self._dirty = True

        currentTime = time.time()

        if (
            currentTime - self.last_color_shift_time
            > self.color_shift_delay_in_seconds
        ):
            self.color_idx += self.color_shift_step
            self.color_idx = self.color_idx % 1
            self.bass_strobe_color = self.get_gradient_color(self.color_idx)
            self.last_color_shift_time = currentTime

        lows_intensity = np.mean(data.melbank_lows())
        if (
            lows_intensity > self.bass_threshold
            and currentTime - self.last_bass_strobe_time
            > self.bass_strobe_wait_time
        ):
            self.bass_strobe_overlay = np.tile(
                self.bass_strobe_color, (self.pixel_count, 1)
            )
            self.last_bass_strobe_time = currentTime

        onsets = data.onset()
        if (
            onsets["high"]
            and currentTime - self.last_strobe_time > self.strobe_wait_time
        ):
            self.onsets_queue.put(True)
            self.last_strobe_time = currentTime


import numpy as np
import torch.utils.data as data
import torch
import os
from six.moves import xrange  # pylint: disable=redefined-builtin
import PIL.Image as Image
import random
import numpy as np
import cv2
import time
class DatasetProcessing(data.Dataset):
    def __init__(self, data_path):
        #self.feature_path = data_path
        ffp = open(data_path,'r')
        fp = list(ffp)
        self.fullfeature_path = [x.strip().split(' ')[0].replace('HMDB51','full_features')+'.npy'  for x in fp]

        self.partialfeature_path = [x.strip().split(' ')[0].replace('HMDB51','partial_features')+'.npy' for x in fp]
        self.labels = [int(x.strip().split(' ')[1]) for x in fp]
        ffp.close()

    def __getitem__(self, index):
        ff = np.load(self.fullfeature_path[index])
        pf = np.load(self.partialfeature_path[index])
        label = torch.LongTensor([self.labels[index]])
        return ff, pf, label, index
    def __len__(self):
        return len(self.fullfeature_path)
#DatasetProcessing('test.list')

class DatasetProcessing2(data.Dataset):
    def __init__(self, data_path):
        #self.feature_path = data_path
        ffp = open(data_path,'r')
        fp = list(ffp)
        self.fullfeature_path = [x.strip().split(' ')[0]  for x in fp]

        self.partialfeature_path = [x.strip().split(' ')[1]for x in fp]
        self.labels = [int(x.strip().split(' ')[2]) for x in fp]
        ffp.close()

    def __getitem__(self, index):
        ff = np.load(self.fullfeature_path[index])
        pf = np.load(self.partialfeature_path[index])
        label = torch.LongTensor([self.labels[index]])
        return ff, pf, label, index
    def __len__(self):
        return len(self.fullfeature_path)

class DatasetReader(data.Dataset):
    def __init__(self, data_path):
        liness = open(data_path,'r')
        lines = list(liness)

        self.video_path = [x.strip().split(' ')[0] for x in lines]

        self.labels = [int(x.strip().split(' ')[1]) for x in lines]
        liness.close()

        self.num_frames_per_clip = 16
        self.crop_size = 112
        self.np_mean = np.load('crop_mean.npy').reshape([self.num_frames_per_clip, self.crop_size, self.crop_size, 3])

    def __getitem__(self,index):
        filename = self.video_path[index]
        tmp_data_p = []
        tmp_data_f = []
        s_index = 0
        s_index_f  = 0
        k = 0
        #print(filename)
        for parent, dirnames, filenames in os.walk(filename):
            if(len(filenames)<self.num_frames_per_clip):
                return [], s_index
            filenames = sorted(filenames)

            while(k < 10):
                tmp_data = []
                s_index = random.randint(0, len(filenames) - self.num_frames_per_clip)
                for i in range(s_index, s_index + self.num_frames_per_clip):
                    image_name = str(filename) + '/' + str(filenames[i])
                    #print(image_name)
                    img = Image.open(image_name)
                    img_data = np.array(img)
                    tmp_data.append(img_data)

                tmp_data_p.append(tmp_data)
                k+=1
            if self.num_frames_per_clip == 1:
                while (len(ret_arr) < 16):
                    tmp_data.append(img_data)
            while(s_index_f+self.num_frames_per_clip<=len(filenames)):
                arr = []
                for i in range(s_index, s_index + self.num_frames_per_clip):

                    image_name = str(filename) + '/' + str(filenames[i])
                    img = Image.open(image_name)
                    img_data = np.array(img)
                    arr.append(img_data)
                tmp_data_f.append(arr)
                s_index_f += 8

        #print(type(filenames))
        #tmp_data, _ = self.get_frames_data(filenames)
        #print(len(tmp_data_p))
        img_datas = [];
        data_p = []
        #print(len(tmp_data))
        if(len(tmp_data_p)!=0):
           for i in xrange(len(tmp_data_p)):
                imgs_datas_p = []
                for j in xrange(len(tmp_data_p[i])):
                    img = Image.fromarray(tmp_data_p[i][j].astype(np.uint8))

                    if(img.width>img.height):
                        scale = float(self.crop_size)/float(img.height)
                        img = np.array(cv2.resize(np.array(img),(int(img.width * scale + 1), self.crop_size))).astype(np.float32)

                    else:
                        scale = float(self.crop_size)/float(img.width)
                        img = np.array(cv2.resize(np.array(img),(self.crop_size, int(img.height * scale + 1)))).astype(np.float32)
                    crop_x = int((img.shape[0] - self.crop_size)/2)
                    crop_y = int((img.shape[1] - self.crop_size)/2)
                    img = img[crop_x:crop_x+self.crop_size, crop_y:crop_y+self.crop_size,:] - self.np_mean[j]
                    imgs_datas_p.append(img)
                data_p.append(imgs_datas_p)

        data_f = []

        if(len(tmp_data_f)!=0):
            for i in xrange(len(tmp_data_f)):
                imgs_datas_f = []
                for j in xrange(len(tmp_data_f[i])):
                    img = Image.fromarray(tmp_data_f[i][j].astype(np.uint8))

                    if(img.width>img.height):
                        scale = float(self.crop_size)/float(img.height)
                        img = np.array(cv2.resize(np.array(img),(int(img.width * scale + 1), self.crop_size))).astype(np.float32)

                    else:
                        scale = float(self.crop_size)/float(img.width)
                        img = np.array(cv2.resize(np.array(img),(self.crop_size, int(img.height * scale + 1)))).astype(np.float32)
                    crop_x = int((img.shape[0] - self.crop_size)/2)
                    crop_y = int((img.shape[1] - self.crop_size)/2)
                    img = img[crop_x:crop_x+self.crop_size, crop_y:crop_y+self.crop_size,:] - self.np_mean[j]
                    imgs_datas_f.append(img)
                data_f.append(imgs_datas_f)

        label = torch.LongTensor([self.labels[index]])

        data_p = np.array(data_p).astype(np.float32)

        data_f = np.array(data_f).astype(np.float32)

        data_p = data_p.squeeze()
        data_p = data_p.transpose(0,4,1,2,3)

        data_f = data_f.transpose(0,4,1,2,3)

        clip_frames = torch.from_numpy(data_p)
        all_frame = torch.from_numpy(data_f)

        return clip_frames, label, index, all_frame, filename

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

class DatasetReader2(data.Dataset):
    def __init__(self, data_path):
        liness = open(data_path,'r')
        lines = list(liness)
        self.video_path = [x.strip().split(' ')[0] for x in lines]
        self.labels = [int(x.strip().split(' ')[1]) for x in lines]
        liness.close()

        self.num_frames_per_clip = 16
        self.crop_size = 112
        self.np_mean = np.load('crop_mean.npy').reshape([self.num_frames_per_clip, self.crop_size, self.crop_size, 3])

    def __getitem__(self,index):
        filename = self.video_path[index]
        tmp_data = []
        tmp_data_f = []
        s_index = 0
        s_index_f  = 0
        for parent, dirnames, filenames in os.walk(filename):
            if(len(filenames)<self.num_frames_per_clip):
                return [], s_index
            filenames = sorted(filenames)
            s_index = random.randint(0, len(filenames) - self.num_frames_per_clip)
            for i in range(s_index, s_index + self.num_frames_per_clip):
                image_name = str(filename) + '/' + str(filenames[i])
                #print(image_name)
                img = Image.open(image_name)
                img_data = np.array(img)
                tmp_data.append(img_data)
            if self.num_frames_per_clip == 1:
                while (len(ret_arr) < 16):
                    tmp_data.append(img_data)

     #print(type(filenames))
        #tmp_data, _ = self.get_frames_data(filenames)
        img_datas = [];

        #print(len(tmp_data))
        if(len(tmp_data)!=0):
            for j in xrange(len(tmp_data)):
                img = Image.fromarray(tmp_data[j].astype(np.uint8))

                if(img.width>img.height):
                    scale = float(self.crop_size)/float(img.height)
                    img = np.array(cv2.resize(np.array(img),(int(img.width * scale + 1), self.crop_size))).astype(np.float32)

                else:
                    scale = float(self.crop_size)/float(img.width)
                    img = np.array(cv2.resize(np.array(img),(self.crop_size, int(img.height * scale + 1)))).astype(np.float32)
                crop_x = int((img.shape[0] - self.crop_size)/2)
                crop_y = int((img.shape[1] - self.crop_size)/2)
                img = img[crop_x:crop_x+self.crop_size, crop_y:crop_y+self.crop_size,:] - self.np_mean[j]
                img_datas.append(img)

        label = torch.LongTensor([self.labels[index]])

        data = np.array(img_datas).astype(np.float32)

        #print(data_f.shape)
        #print(data.shape,filename)
        #print(data.shape,filename)
        data = data.transpose(3, 0, 1, 2)

        clip_frames = torch.from_numpy(data)

        return clip_frames, label, index

    def __len__(self):
        return len(self.video_path)
DatasetReader('./TrainSplit1.txt')