import re, os, subprocess, json, soundfile
import numpy as np

workingdir = os.path.dirname(os.path.abspath(__file__));
os.chdir(workingdir);

global GLOBAL_VARS;
GLOBAL_VARS = {
    "ffmpeg_path" : "ffmpeg"
};

# "convert" will also convert the music file, generating wavfile.wav in local path
def read_osu_file(path, convert=False, wav_name="wavfile.wav", json_name="temp_json_file.json"):
    file_dir = os.path.dirname(os.path.abspath(path));

    # ask node.js to convert the .osu file to .json format
    # this is my own converter (not the one on npm) made a few years ago
    # tweaked a little to get it running in node.
    # ... no one cares about that anyways.
    subprocess.call(["node", "load_map.js", "jq", path, json_name]);

    with open(json_name, encoding="utf-8") as map_json:
        map_dict = json.load(map_json); # not "loads" it is not a string

        if convert:
            mp3_file = os.path.join(file_dir, map_dict["general"]["AudioFilename"]);
            subprocess.call([GLOBAL_VARS["ffmpeg_path"], "-y", "-i", mp3_file, wav_name]);

    # delete the temp json here!!!
    if json_name == "temp_json_file.json":
        os.remove(json_name);

    return map_dict, wav_name;

def get_map_timing_array(map_json, length=-1, divisor=4):
    if length == -1:
        length = map_json["obj"][-1]["time"] + 1000; # it has an extra time interval after the last note
        if map_json["obj"][-1]["type"] & 8: # spinner end
            length = map_json["obj"][-1]["spinnerEndTime"] + 1000;
    uts_a =  map_json["timing"]["uts"];
    out = [];
    for i, uts in enumerate(uts_a):
        begin_time = uts["beginTime"];
        mspb = uts["tickLength"];
        if i < len(uts_a)-1:
            end_time = uts_a[i+1]["beginTime"];
        else:
            end_time = length;
        arr = np.floor(np.arange(begin_time, end_time, mspb / divisor));
        out = out + list(map(lambda f: int(f), arr));
    return out;

def get_tick_len(map_json, tick):
    uts_a = map_json["timing"]["uts"];
    if tick < uts_a[0]["beginTime"]:
        return uts_a[0]["tickLength"];
    for uts in reversed(uts_a):
        if tick >= uts["beginTime"]:
            return uts["tickLength"];

def get_slider_len(map_json, tick):
    ts_a = map_json["timing"]["ts"];
    if tick < ts_a[0]["beginTime"]:
        return ts_a[0]["sliderLength"];
    for ts in reversed(ts_a):
        if tick >= ts["beginTime"]:
            return ts["sliderLength"];

def get_slider_len_ts(ts_a, tick):
    if tick < ts_a[0]["beginTime"]:
        return ts_a[0]["sliderLength"];
    for ts in reversed(ts_a):
        if tick >= ts["beginTime"]:
            return ts["sliderLength"];

def get_end_time(note):
    if note["type"] & 8:
        return note["spinnerEndTime"];
    elif note["type"] & 2:
        return note["sliderData"]["endTime"];
    #elif note["type"] & 128:
    #    return note["holdEndTime"];
    else:
        return note["time"];

# edited from uts to ts wwww
def get_all_ticks_and_lengths_from_ts(uts_array, ts_array, end_time):
    # Returns array of all timestamps, ticklens and sliderlens.
    endtimes = ([uts["beginTime"] for uts in uts_array] + [end_time])[1:];
    ticks = [np.arange(uts["beginTime"], endtimes[i], uts["tickLength"] / 4) for i, uts in enumerate(uts_array)];
    tick_len = [[uts["tickLength"]] * len(np.arange(uts["beginTime"], endtimes[i], uts["tickLength"] / 4)) for i, uts in enumerate(uts_array)];
    # slider_len = [[ts["sliderLength"]] * len(np.arange(ts["beginTime"], endtimes[i], ts["tickLength"] / 4)) for i, ts in enumerate(ts_array)];
    slider_len = [get_slider_len_ts(ts_array, tick) for tick in np.concatenate(ticks)];
    return np.round(np.concatenate(ticks)).astype(int), np.concatenate(tick_len), np.array(slider_len);

def get_end_point(note):
    if note["type"] & 8:
        return np.array([256, 192]);
    elif note["type"] & 2:
        return np.array(note["sliderData"]["endpoint"]);
    else:
        return np.array([note["x"], note["y"]]);

def get_input_vector(note, prev_note):
    if note["type"] & 8:
        return None;
    #elif note["type"] & 2:
    #    return np.array(note["sliderData"]["dIn"]);
    else:
        vec = np.array([note["x"], note["y"]]) - get_end_point(prev_note);
        return vec / max(0.001, np.sqrt(vec.dot(vec)));

def get_output_vector(note, prev_note):
    if note["type"] & 8:
        return None;
    elif note["type"] & 2:
        return np.array(note["sliderData"]["dOut"]);
    else:
        vec = np.array([note["x"], note["y"]]) - get_end_point(prev_note);
        return vec / max(0.001, np.sqrt(vec.dot(vec)));

def get_angular_momentum(note, prev_note, prev2_note):
    v1 = get_input_vector(note, prev_note);
    v0 = get_output_vector(prev_note, prev2_note);
    if v0 is None or v1 is None or v0[0] is None or v1[0] is None:
        return 0;
    return (np.arctan2(v1[1], v1[0]) - np.arctan2(v0[1], v0[0])) / max(10, (note["time"] - get_end_time(prev_note)));

def get_momentum(note, prev_note):
    v1 = np.array([note["x"], note["y"]]);
    v0 = get_end_point(prev_note);
    v = v1 - v0;
    if note["time"] - get_end_time(prev_note) == 0 or note["time"] - prev_note["time"] == 0:
        # it has the same time the previous note ends. either a bugged sliderend or a double note
        return 0;
    end_type_momentum = np.sqrt(v.dot(v)) / (note["time"] - get_end_time(prev_note));

    # Since slider jumps in maps cause this parameters to learned too high
    # we try to deal with slider leniency by using the beginning of slider
    v2 = np.array([prev_note["x"], prev_note["y"]]);
    v3 = v1 - v2;
    start_type_momentum = np.sqrt(v3.dot(v3)) / (note["time"] - prev_note["time"]);
    return np.min([end_type_momentum, start_type_momentum]);

def get_map_notes(map_json, **kwargs):
    length = kwargs.get("length", -1);
    divisor = kwargs.get("divisor", 4);
    tick_times = get_map_timing_array(map_json, length = length, divisor = divisor);

    objs = map_json["obj"];
    obj_times = list(map(lambda obj: obj["time"], objs));

    def get_note_type(obj):
        if not obj:
            return 0;
        if obj["type"] & 2:
            return 2;
        elif obj["type"] & 8:
            return 3;
        return 1;

    po = 0;
    note_max_wait_time = kwargs.get("note_max_wait_time", 1000);
    start_time = obj_times[0] - note_max_wait_time;
    last_obj_time = start_time;
    sliding = 0;
    slider_end_time = 0;
    spinning = 0;
    spinner_end_time = 0;
    data = [];
    flow_data = [];

    tlen_mp = 1/500;
    tlen_s = 1;
    bpm_mp = 1/120;
    bpm_s = 1;
    slen_mp = 1/150;
    slen_s = 1;

    for i, tick in enumerate(tick_times):

        # Attach extra vars at the end of each note datum
        tlen = get_tick_len(map_json, tick);
        bpm = 60000 / tlen;
        slen = get_slider_len(map_json, tick);
        ex1 = tlen * tlen_mp - tlen_s;
        ex2 = bpm * bpm_mp - bpm_s;
        ex3 = slen * slen_mp - slen_s;

        while obj_times[po] < tick - 5 and po < len(obj_times) - 1:
            po += 1;
        if obj_times[po] >= tick - 5 and obj_times[po] <= tick + 5: # found note
            last_obj_time = tick;
            note_type = get_note_type(objs[po]);

            # calculate momentum
            if po >= 1:
                momentum = get_momentum(objs[po], objs[po-1]);
            else:
                momentum = 0;

            # calculate angular momentum
            if po >= 2:
                angular_momentum = get_angular_momentum(objs[po], objs[po-1], objs[po-2]);
            else:
                angular_momentum = 0;

            # flow data
            if po >= 1:
                input_vector = get_input_vector(objs[po], objs[po-1]);
                output_vector = get_output_vector(objs[po], objs[po-1]);
            else:
                input_vector = [0, 0];
                output_vector = [0, 0];
            if input_vector is None or input_vector[0] is None or input_vector[1] is None:
                input_vector = [0, 0];
            if output_vector is None or output_vector[0] is None or output_vector[1] is None:
                output_vector = [0, 0];

            # end point
            endpoint = get_end_point(objs[po]);
            flow_data.append([i, tick, note_type, objs[po]["x"], objs[po]["y"], input_vector[0], input_vector[1], output_vector[0], output_vector[1], endpoint[0], endpoint[1]]);

            # put data
            if note_type == 1:
                spinning = 0;
                sliding = 0;
            elif note_type == 2:
                sliding = 1;
                slider_end_time = objs[po]["sliderData"]["endTime"];
            elif note_type == 3:
                spinning = 1;
                spinner_end_time = objs[po]["spinnerEndTime"];
                # because the spinner sometimes get over 3 secs
                last_obj_time = spinner_end_time;

            # TICK, TIME, NOTE, NOTE_TYPE, SLIDING, SPINNING, MOMENTUM, ANGULAR_MOMENTUM, Ex1, Ex2, Ex3
            data.append([i, tick, 1, note_type, sliding, spinning, momentum, angular_momentum, ex1, ex2, ex3]);
        elif spinning == 1:
            if tick >= spinner_end_time - 5:
                # somehow it is better not to put a note here (?)
                spinning = 0;
                data.append([i, tick, 1, 5, 0, 0, 0, 0, ex1, ex2, ex3]);
            else:
                data.append([i, tick, 0, 0, 0, 1, 0, 0, ex1, ex2, ex3]);
        elif sliding == 1:
            if tick >= slider_end_time - 5:
                # somehow it is better not to put a note here (?)
                sliding = 0;
                data.append([i, tick, 1, 4, 0, 0, 0, 0, ex1, ex2, ex3]);
            else:
                data.append([i, tick, 0, 0, 1, 0, 0, 0, ex1, ex2, ex3]);
        else: # not found
            if tick - last_obj_time < note_max_wait_time and tick >= start_time:
                data.append([i, tick, 0, 0, 0, 0, 0, 0, ex1, ex2, ex3]);
    return data, flow_data;

def get_freqs(sig, fft_size):
    Lf = np.fft.fft(sig, fft_size);
    Lc = Lf[0:fft_size//2];
    La = np.abs(Lc[0:fft_size//2]);
    Lg = np.angle(Lc[0:fft_size//2]);
    return La, Lg;

def slice_wave_at(ms, sig, samplerate, size):
    ind = (ms/1000 * samplerate)//1;
    return sig[max(0, int(ind - size//2)):int(ind + size - size//2)];

def lrmix(sig):
    return (sig[:,0]+sig[:,1])/2;

def get_wav_data_at(ms, sig, samplerate, size=1024, fft_size=2048, freq_low=0, freq_high=-1):
    if freq_high == -1:
        freq_high = samplerate//2;
    waveslice = slice_wave_at(ms, sig, samplerate, size);

    # since osu! maps are usually not mapped to stereo wave, let's mix it to reduce 50% of data
    #waveslice_lr = lrmix(waveslice);
    waveslice_lr = waveslice

    # do a nice FFT
    La, Lg = get_freqs(waveslice_lr, fft_size);

    # cut the frequency bins
    # freq_step * size == samplerate
    # freq_step * k == frequency
    # k/freq == size/smr
    # T O T A L L Y I D E N T I C A L
    La = La[fft_size*freq_low//samplerate:fft_size*freq_high//samplerate];
    Lg = Lg[fft_size*freq_low//samplerate:fft_size*freq_high//samplerate];

    return La, Lg;

def read_wav_data(timestamps, wavfile, snapsize=[1, 2, 4, 8, 16], fft_size = 1024):
    sig, samplerate = soundfile.read(wavfile);
    data = list();

    # normalize sound wave
    # sig = sig / np.sqrt(np.mean(sig**2, axis=0));
    sig = sig / np.max(np.max(np.abs(sig), axis=0));

    # calc a length array
    tmpts = np.array(timestamps);
    timestamp_interval = tmpts[1:] - tmpts[:-1];
    timestamp_interval = np.append(timestamp_interval, timestamp_interval[-1]);

    for sz in snapsize:
        data_r = np.array([get_wav_data_at(coord, sig, samplerate, size=timestamp_interval[i] * sz, fft_size=fft_size, freq_high=samplerate//4) for i, coord in enumerate(timestamps)]);
        data.append(data_r);

    raw_data = np.array(data);
    norm_data = np.tile(np.expand_dims(np.mean(raw_data, axis=1), 1), (1, raw_data.shape[1], 1, 1));
    std_data = np.tile(np.expand_dims(np.std(raw_data, axis=1), 1), (1, raw_data.shape[1], 1, 1));
    return (raw_data - norm_data) / std_data;

#
# Main function
# Generated data shape:
#     - "lst" file, table of [TICK, TIME, NOTE, IS_CIRCLE, IS_SLIDER, IS_SPINNER, IS_SLIDER_END, IS_SPINNER_END, SLIDING, SPINNING, MOMENTUM, ANGULAR_MOMENTUM, EX1, EX2, EX3], length MAPTICKS
#                                0,    1,    2,         3,         4,          5,             6,              7,        8,       9,       10,               11
#     - "wav" file, np.array, shape of [len(snapsize), MAPTICKS, 2, fft_size//4]
#     - "flow" file, table of [TICK, TIME, TYPE, X, Y, IN_DX, IN_DY, OUT_DX, OUT_DY] notes only
#
# MAPTICKS = (Total map time + 3000) / tickLength / (divisor = 4) - EMPTY_TICKS
# EMPTY_TICKS = ticks where no note around in 5 secs
#
def read_and_save_osu_file(path, filename = "saved"):
    osu_dict, wav_file = read_osu_file(path, convert = True);
    data, flow_data = get_map_notes(osu_dict);
    timestamps = [c[1] for c in data];
    wav_data = read_wav_data(timestamps, wav_file, snapsize=[0.25, 0.5, 1, 2, 4, 8, 16], fft_size = 128);
    # in order to match first dimension
    wav_data = np.swapaxes(wav_data, 0, 1);

    # change the representation of note_type
    # a bit of copypaste code because I changed the data structure many times here
    transformed_data = [];
    for d in data:
        if d[3] == 1:
            transformed_data.append([d[0], d[1], d[2], 1, 0, 0, 0, 0, d[4], d[5], d[6], d[7], d[8], d[9], d[10]]);
        elif d[3] == 2:
            transformed_data.append([d[0], d[1], d[2], 0, 1, 0, 0, 0, d[4], d[5], d[6], d[7], d[8], d[9], d[10]]);
        elif d[3] == 3:
            transformed_data.append([d[0], d[1], d[2], 0, 0, 1, 0, 0, d[4], d[5], d[6], d[7], d[8], d[9], d[10]]);
        elif d[3] == 4:
            transformed_data.append([d[0], d[1], d[2], 0, 0, 0, 1, 0, d[4], d[5], d[6], d[7], d[8], d[9], d[10]]);
        elif d[3] == 5:
            transformed_data.append([d[0], d[1], d[2], 0, 0, 0, 0, 1, d[4], d[5], d[6], d[7], d[8], d[9], d[10]]);
        else:
            transformed_data.append([d[0], d[1], d[2], 0, 0, 0, 0, 0, d[4], d[5], d[6], d[7], d[8], d[9], d[10]]);

    # import pandas as pd;
    # pd.DataFrame(flow_data).to_csv("flow.csv", header=["TICK", "TIME", "TYPE", "X", "Y", "IN_DX", "IN_DY", "OUT_DX", "OUT_DY"]);

    np.savez_compressed(filename, lst = transformed_data, wav = wav_data, flow = flow_data);

def read_and_save_osu_tester_file(path, filename = "saved", json_name="mapthis.json"):
    osu_dict, wav_file = read_osu_file(path, convert = True, json_name=json_name);
    sig, samplerate = soundfile.read(wav_file);
    file_len = (sig.shape[0] / samplerate * 1000 - 3000);

    timestamps, tick_lengths, slider_lengths = get_all_ticks_and_lengths_from_ts(osu_dict["timing"]["uts"], osu_dict["timing"]["ts"], file_len);
    ticks = np.array([i for i,k in enumerate(timestamps)]);
    # uts = osu_dict["timing"]["uts"][0];
    # ticks = np.array(list(range(0,int((file_len - uts["beginTime"])/uts["tickLength"]*4))));
    # timestamps = np.floor(ticks * uts["tickLength"]/4 + uts["beginTime"]);
    extra = np.array([60000 / tick_lengths, slider_lengths]);

    wav_data = read_wav_data(timestamps, wav_file, snapsize=[0.25, 0.5, 1, 2, 4, 8, 16], fft_size = 128);
    # in order to match first dimension
    wav_data = np.swapaxes(wav_data, 0, 1);

    np.savez_compressed(filename, ticks = ticks, timestamps = timestamps, wav = wav_data, extra = extra);

def read_and_return_osu_file(path):
    osu_dict, wav_file = read_osu_file(path, convert = True);
    data, flow_data = get_map_notes(osu_dict);
    timestamps = [c[1] for c in data];
    wav_data = read_wav_data(timestamps, wav_file, snapsize=[0.25, 0.5, 1, 2, 4, 8, 16], fft_size = 128);
    return data, wav_data, flow_data;

# import pandas as pd;
# pd.DataFrame(data).to_csv("hmmmm.csv", header=["TICK", "TIME", "NOTE", "NOTE_TYPE", "SLIDING", "SPINNING", "MOMENTUM", "ANGULAR_MOMENTUM"]);
# print(subprocess.call("node load_map.js j pinkheart.osu t.json"));
# print(subprocess.call(["node", "load_map.js", "o", "t.json", "z.osu"]));


import argparse
import os
import sys
import math
import numpy as np
import multiprocessing
from sklearn.metrics import confusion_matrix
import time
import pickle
import multiprocessing as mp
from ops.sequence_funcs import *
from ops.anet_db import ANetDB
from ops.thumos_db import THUMOSDB
from ops.detection_metrics import get_temporal_proposal_recall, name_proposal
from ops.sequence_funcs import temporal_nms
from ops.io import dump_window_list
parser = argparse.ArgumentParser()
parser.add_argument('score_files', type=str, nargs='+')
parser.add_argument("--anet_version", type=str, default='1.2', help='')
parser.add_argument("--dataset", type=str, default='activitynet', choices=['activitynet', 'thumos14'])
parser.add_argument("--cls_scores", type=str, default=None,
                    help='classification scores, if set to None, will use groundtruth labels')
parser.add_argument("--subset", type=str, default='validation', choices=['training', 'validation', 'testing'])
parser.add_argument("--iou_thresh", type=float, nargs='+', default=[0.5, 0.75, 0.95])
parser.add_argument("--score_weights", type=float, nargs='+', default=None, help='')
parser.add_argument("--write_proposals", type=str, default=None, help='')
parser.add_argument("--minimum_len", type=float, default=0, help='minimum length of a proposal, in second')
parser.add_argument("--reg_score_files", type=str, nargs='+', default=None)
parser.add_argument("--frame_path", type=str, default='/mnt/SSD/ActivityNet/anet_v1.2_extracted_340/')

args = parser.parse_args()

if args.dataset == 'activitynet':
    db = ANetDB.get_db(args.anet_version)
    db.try_load_file_path('/mnt/SSD/ActivityNet/anet_v1.2_extracted_340/')
elif args.dataset == 'thumos14':
    db = THUMOSDB.get_db()
    db.try_load_file_path('/mnt/SSD/THUMOS14/')

    # rename subset test
    if args.subset == 'testing':
        args.subset = 'test'
else:
    raise ValueError("unknown dataset {}".format(args.dataset))

video_list = db.get_subset_videos(args.subset)
video_list = [v for v in video_list if v.instances != []]
print("video list size: {}".format(len(video_list)))
# load scores
print('loading scores...')
score_list = []
for fname in args.score_files:
    score_list.append(pickle.load(open(fname, 'rb')))
print('load {} piles of scores'.format(len(score_list)))

# load classification scores if specified
if args.cls_scores:
    cls_scores = cPickle.load(open(args.cls_scores, 'rb'))
else:
    cls_scores = None
print('done')

# load regression scores
if args.reg_score_files is not None:
    print('loading regression scores')
    reg_score_list = []
    for fname in args.reg_score_files:
        reg_score_list.append(cPickle.load(open(fname, 'rb')))
    print('load {} piles of regression scores'.format(len(reg_score_list)))
else:
    reg_score_list = None

# merge scores
print('merging scores')
score_dict = {}
for key in score_list[0].keys():
    out_score = score_list[0][key].mean(axis=1) * (1.0 if args.score_weights is None else args.score_weights[0])
    for i in range(1, len(score_list)):
        add_score = score_list[i][key].mean(axis=1)
        if add_score.shape[0] < out_score.shape[0]:
            out_score = out_score[:add_score.shape[0], :]
        elif add_score.shape[0] > out_score.shape[0]:
            tick = add_score.shape[0] / float(out_score.shape[0])
            indices = [int(x * tick) for x in range(out_score.shape[0])]
            add_score = add_score[indices, :]
        out_score += add_score * (1.0 if args.score_weights is None else args.score_weights[i])
    score_dict[key] = out_score
print('done')

# merge regression scores
if reg_score_list is not None:
    print('merging regression scores')
    reg_score_dict = {}
    for key in reg_score_list[0].keys():
        out_score = reg_score_list[0][key].mean(axis=1)
        for i in range(1, len(reg_score_list)):
            add_score = reg_score_list[i][key].mean(axis=1)
            if add_score.shape[0] < out_score.shape[0]:
                out_score = out_score[:add_score.shape[0], :]
            out_score += add_score
        reg_score_dict[key] = out_score / len(reg_score_list)
    print('done')
else:
    reg_score_dict = None

# bottom-up generate proposals
print('generating proposals')
pr_dict = {}
pr_score_dict = {}
topk = 1

def gen_prop(v):
    if (args.dataset == 'activitynet') or (args.dataset == 'thumos14'):
        vid = v.id
    else:
        vid = v.path.split('/')[-1].split('.')[0]
    scores = score_dict[vid]
    frm_duration = len(scores)
    topk_cls = [0]
    topk_labels = label_frame_by_threshold(scores, topk_cls, bw=3, thresh=[0.01, 0.05, 0.1, .15, 0.25, .4, .5, .6, .7, .8, .9, .95, ], multicrop=False)

    bboxes = []
    tol_lst = [0.05, .1, .2, .3, .4, .5, .6, 0.8, 1.0]

    bboxes.extend(build_box_by_search(topk_labels, np.array(tol_lst)))
    if reg_score_dict:
        reg_scores = reg_score_dict[score_id]
        bboxes = regress_box(bboxes, reg_scores, len(scores))

    # print len(bboxes)
    bboxes = temporal_nms(bboxes, 0.9)

    pr_box = [(x[0] / float(frm_duration) * v.duration, x[1] / float(frm_duration) * v.duration) for x in bboxes]

    # filter out too short proposals
    pr_box = list(filter(lambda b: b[1] - b[0] > args.minimum_len, pr_box))
    return v.id, pr_box, [x[3] for x in bboxes]

def call_back(rst):
    pr_dict[rst[0]] = rst[1]
    pr_score_dict[rst[0]] = rst[2]
    import sys
    print(rst[0], len(pr_dict), len(rst[1]))
    sys.stdout.flush()

pool = mp.Pool(processes = 32)
lst = []
handle = [pool.apply_async(gen_prop, args=(x, ), callback=call_back) for x in video_list]
pool.close()
pool.join()

# evaluate proposal info
proposal_list = [pr_dict[v.id] for v in video_list if v.id in pr_dict]
gt_spans_full = [[(x.num_label, x.time_span) for x in v.instances] for v in video_list if v.id in pr_dict]
gt_spans = [[item[1] for item in x] for x in gt_spans_full]
score_list = [score_dict[v.id] for v in video_list if v.id in pr_dict]
duration_list = [v.duration for v in video_list if v.id in pr_dict]
proposal_score_list = [pr_score_dict[v.id] for v in video_list if v.id in pr_dict]
print('{} groundtruth boxes from'.format(sum(map(len, gt_spans))))

print('average # of proposals: {}'.format(np.mean(list(map(len, proposal_list)))))
IOU_thresh = np.arange(0.5, 1, 0.2)
p_list = []
for th in IOU_thresh:
    pv, pi = get_temporal_proposal_recall(proposal_list, gt_spans, th)
    print('IOU threshold {}. per video recall: {:02f}, per instance recall: {:02f}'.format(th, pv * 100, pi * 100))
    p_list.append((pv, pi))
print('Average Recall: {:.04f} {:.04f}'.format(*(np.mean(p_list, axis=0)*100)))

if args.write_proposals:

    name_pattern = 'img_*.jpg'
    frame_path = args.frame_path

    named_proposal_list = [name_proposal(x, y) for x, y in zip(gt_spans_full, proposal_list)]
    allow_empty = args.dataset == 'activitynet' and args.subset == 'testing'
    dumped_list = [dump_window_list(v, prs, frame_path, name_pattern, score=score, allow_empty=allow_empty) for v, prs, score in
                   zip(filter(lambda x: x.id in pr_dict, video_list), named_proposal_list, score_list)]

    with open(args.write_proposals, 'w') as of:
        for i, e in enumerate(dumped_list):
            of.write('# {}\n'.format(i + 1))
            of.write(e)

    print('list written. got {} videos'.format(len(dumped_list)))


#!/usr/bin/env python

import os
import time
import numpy as np
import socket
import logging
import tempfile
import shutil
from glob import glob
from datetime import datetime
from dlnpyutils import utils as dln, coords
import photred
from photred import io,daomatch
from photred import allframe as alf
import dill as pickle
from astropy.io import fits
from astropy.wcs import WCS
from astropy.table import Table
from astropy.coordinates import SkyCoord
from dustmaps.sfd import SFDQuery
import subprocess

def make_brick_wcs(brickstr):
    """
    This creates a brick WCS given the brick structure

    Each brick covers 0.25 x 0.25 deg at a pixel scale of 0.262"/pix
    with 3600x3600 pixels that gives us an overlap of ~82 pixels on
    each side.  The *UNIQUE* brick area (0.25x0.25 deg) is defined by
    the BRAMIN/MAX and BDECMIN/MAX keywords.
    """

    # Make the tiling file
    #---------------------
    nx = 3600
    ny = 3600
    step = 0.262 / 3600    # 0.262" per pixel, DECam pixel scale
    xref = nx/2
    yref = ny/2

    #  Make the header as well
    tilehead = fits.Header()
    tilehead['NAXIS1'] = nx
    tilehead['CDELT1'] = step
    tilehead['CRPIX1'] = xref+1
    tilehead['CRVAL1'] = brickstr['ra']
    tilehead['CTYPE1'] = 'RA---TAN'
    tilehead['NAXIS2'] = ny
    tilehead['CDELT2'] = step
    tilehead['CRPIX2'] = yref+1
    tilehead['CRVAL2'] = brickstr['dec']
    tilehead['CTYPE2'] = 'DEC--TAN'
    tilehead['BRAMIN'] = brickstr['ra1'],'RA min of unique brick area'
    tilehead['BRAMAX'] = brickstr['ra2'],'RA max of unique brick area'
    tilehead['BDECMIN'] = brickstr['dec1'],'DEC min of unique brick area'
    tilehead['BDECMAX'] = brickstr['dec2'],'DEC max of unique brick area'
    wcs = WCS(tilehead)

    # Create the TILE structure
    tile = {'type':'WCS','naxis':np.array([nx,ny]),'cdelt':np.array([step,step]).astype(float),
            'crpix':np.array([xref+1,yref+1]).astype(float),
            'crval':np.array([brickstr['ra'],brickstr['dec']]).astype(float),'ctype':['RA--TAN','DEC--TAN'],
            'head':tilehead,'wcs':wcs,'xrange':[0,nx],'yrange':[0,ny],'nx':nx,'ny':ny}

    return tile

def forcebrick(brick,scriptsdir=None,irafdir=None,workdir=None,redo=False,
               update=False,logfile=None,delvedir=None,delvereddir=None):
    """
    Process a single DELVE brick and perform ALLFRAME FORCED photometry

    Parameters
    ----------
    brick : str
      The DELVE brick name, e.g. 1234m045

    By D. Nidever  August 2019
    Translated to Python by D. Nidever,  April 2022
    """

    # Limit the number of threads
    #CPU,tpool_nthreads=4

    # This bricks pre-processing script gets DELVE and community MC data ready
    # to run PHOTRED ALLFRAME on it.

    t0 = time.time()
    curdir = os.path.abspath(os.curdir)

    # Defaults
    if delvedir is not None:
        if delvedir.endswith('/')==False:
            delvedir += '/'
    else:
        delvedir = '/net/dl2/dnidever/delve/'
    if os.path.exists(delvedir)==False:
        os.makedirs(delvedir)
    if delvereddir is not None:
        if delvereddir.endswith('/')==False:
            delvereddir += '/'
    else:
        delvereddir = '/'.join(os.path.abspath(__file__).split('/')[:-3])+'/'
    if scriptsdir is not None:
        if scriptsdir.endswith('/')==False:
            scriptsdir += '/'
    else:
        scriptsdir = '/'.join(photred.__file__.split('/')[:-3])+'/scripts/'
    if irafdir is not None:
        if irafdir.endswith('/')==False:
            irafdir += '/'
    else:
        irafdir = os.path.expanduser('~/iraf/')
    tempdir = '/tmp/'
    if workdir is None:
        host = socket.gethostname()
        hostname = host.split('.')[0]
        workdir = '/data0/dnidever/delve/'
        if os.path.exists(workdir)==False:
            os.makedirs(workdir)
    # Exposures directory
    expdir = delvedir+'exposures/'
    # Bricks directory
    brickdir = delvedir+'bricks/'
    if os.path.exists(brickdir)==False:
        os.makedirs(brickdir)
    logsdir = brickdir+'logs/'
    if os.path.exists(logsdir)==False:
        os.makedirs(logsdir)

    # Start the logfile
    #------------------
    # format is delvered_forcebrick.brick.host.DATETIME.log
    host = socket.gethostname()
    hostname = host.split('.')[0]
    logtime = datetime.now().strftime("%Y%m%d%H%M%S")
    # Set up logging to screen and logfile
    logFormatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s]  %(message)s")
    logger = logging.getLogger()
    while logger.hasHandlers(): # some existing loggers, remove them
        logger.removeHandler(logger.handlers[0])
    logger = logging.getLogger()
    logtime = datetime.now().strftime("%Y%m%d%H%M%S")
    if logfile is None:
        logfile = logsdir+'delvered_brick.'+brick+'.'+hostname+'.'+logtime+'.log'
    if os.path.exists(logfile): os.remove(logfile)
    fileHandler = logging.FileHandler(logfile)
    fileHandler.setFormatter(logFormatter)
    logger.addHandler(fileHandler)
    consoleHandler = logging.StreamHandler()
    consoleHandler.setFormatter(logFormatter)
    logger.addHandler(consoleHandler)
    logger.setLevel(logging.NOTSET)

    # Chip name and CCDNUM relations
    decam = Table.read(delvereddir+'data/decam.txt',format='ascii')
    for n in decam.colnames: decam[n].name = n.lower()  # lowercase column names

    # Load the brick information
    brickstr = Table.read(delvereddir+'data/delvemc_bricks_0.25deg.fits.gz')
    for n in brickstr.colnames: brickstr[n].name = n.lower()  # lowercase column names

    # Get the brick information
    bind, = np.where(brickstr['brickname'] == brick)
    if len(bind) == 0:
        logger.info(brick+' not in DELVE-MC brick list')
        return
    brickstr1 = brickstr[bind[0]]

    # Subdirectory is the first 4 digits of the brickname, e.g., 0952 of 0952m462, the RA portion of the name
    subdir = brickdir+brick[0:4]+'/'
    if os.path.exists(subdir)==False:
        os.makedirs(subdir)
    # Create brick directory if it doesn't exist yet
    bdir = subdir+brick+'/'
    if os.path.exists(bdir)==False:
        os.makedirs(bdir)
    logfile = bdir+brick+'.'+logtime+'.log'

    # Check the output file
    photfile = bdir+brick+'.fits'
    if os.path.exists(photfile+'.gz') and redo==False and update==False:
        logger.info(photfile+'.gz EXISTS and redo or update NOT set')
        return

    # DECam imager
    thisimager = {'telescope':'BLANCO','instrument':'DECam','namps':62,'separator':'_'}

    # Print information
    logger.info('--------------------------------------------------------------------')
    logger.info(' Run ALLFRAME forced photometry on DELVE Brick = '+brick)
    logger.info('--------------------------------------------------------------------')
    logger.info('DELVEDIR = '+delvedir)
    logger.info('SCRIPTSDIR = '+scriptsdir)
    logger.info('IRAFDIR = '+irafdir)
    logger.info('EXPOSUREDIR = '+expdir)
    logger.info('WORKINGDIR = '+workdir)
    logger.info('HOST = '+host)
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

    # The photed.setup file
    setup = ['##### REQUIRED #####',
             'scriptsdir  '+scriptsdir,
             'irafdir     '+irafdir,
             'telescope   Blanco',
             'instrument  DECAM',
             'observatory CTIO',
             'nmulti      10',
             'nmulti_wcs       40',
             'nmulti_daophot   30',
             'nmulti_allframe  10',
             'filtref     g,i,r,z,u',
             'trans       delve.trans',
             '##### OPTIONAL #####',
             'sepfielddir  1',
             'sepchipdir   1',
             'keepmef      0',
             'catformat    FITS',
             'workdir      '+workdir,
             'clean        1',
             'skipcheck    1',
             'redo         0',
             'wcsrefname   GAIADR2',
             'searchdist   20',
             '#wcsrmslim   1.0',
             'hyperthread   1',
             'daopsfva      1',
             'daofitradfwhm 1.0',
             'psfcomsrc     0',
             'psfcomglobal  0',
             'psfcomgauss   0',
             '#mchmaxshift  50.0',
             'finditer      1',       # 2->1 on 7/22/20
             'alfdetprog  sextractor',
             '#alfnocmbimscale 0',
             'alftrimcomb   0',
             '#ddo51radoffset  1',
             'cmbforce      1',
             'keepinstr     1',
             'avgmag        1',
             'avgonlymag    0',
             'todered       u,g,r,i,z,g-i',
             '##### STAGES #####',
             '#rename',
             '#split',
             '#wcs',
             '#daophot',
             '#zeropoint',
             '#match',
             ' allframe',
             '#apcor',
             '#astrom',
             '#calib',
             '#combine',
             '#deredden',
             '#save',
             '#html']
    dln.writelines(bdir+'photred.setup',setup)
    setup = io.readsetup(setupdir=bdir)

    # Print out some information about the brick
    logger.info('RA = %.5f' % brickstr1['ra'])
    logger.info('DEC = %.3f' % brickstr1['dec'])
    logger.info('RA range  = [ %.5f,%.5f ]' % (brickstr1['ra1'],brickstr1['ra2']))
    logger.info('DEC range = [ %.5f,%.5f ]' % (brickstr1['dec1'],brickstr1['dec2']))

    # Get the Brick WCS information
    tile = make_brick_wcs(brickstr1)

    # Step 1: Get the list of exposures/chips that overlap this brick
    #----------------------------------------------------------------
    logger.info('Step 1: Get list of chip files that overlap this brick')
    #logger.info('Getting chips that overlap this brick'
    cenra = brickstr1['ra']
    cendec = brickstr1['dec']
    tid,tmpfile = tempfile.mkstemp(prefix="tmp",suffix='.fits',dir=tempdir)
    dln.touch(tmpfile)
    for f in [tmpfile]:
        if os.path.exists(f): os.remove(f)
    # /noshell causes problems on gp09 because it gives python2 instead of python3
    out = subprocess.check_output(delvereddir+'bin/query_delvered_summary_table '+str(cenra)+' '+str(cendec)+' '+tmpfile+' --lim 0.5',shell=True)
    exists = os.path.exists(tmpfile)
    size = os.path.getsize(tmpfile)
    if size == 0:
        logger.info('No overlapping chips found')
        if os.path.exists(tmpfile): os.remove(tmpfile)
        return
    chstr = Table.read(tmpfile)
    if os.path.exists(tmpfile): os.remove(tmpfile)
    nchstr = len(chstr)
    logger.info('Found '+str(nchstr)+' overlapping chips within 0.5 deg of brick center')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

    # Make sure the chips are unique,  some were duplicated on SMASH nights
    chid = np.char.array(chstr['expnum'].astype(str))+'-'+np.char.array(chstr['chip'].astype(str))
    uchstr,ui = np.unique(chid,return_index=True)
    chstr = chstr[ui]
    nchstr = len(chstr)

    # Do more rigorous overlap checking
    #  the brick region with overlap
    print('Performing more rigorous overlap checking')

    bcoo = tile['wcs'].pixel_to_world([0,tile['nx']-1,tile['nx']-1,0],[0,0,tile['ny']-1,tile['ny']-1])
    bvra = bcoo.ra.deg
    bvdec = bcoo.dec.deg
    #HEAD_XYAD,tile.head,[0,tile.nx-1,tile.nx-1,0],[0,0,tile.ny-1,tile.ny-1],bvra,bvdec,/deg
    olap = np.zeros(nchstr,bool)
    vxarr = np.zeros((nchstr,4),float)
    vyarr = np.zeros((nchstr,4),float)
    for i in range(nchstr):
        if (i % 100 == 0) and (i > 0):
            print(i)
        hd1 = io.readfile(chstr['file'][i],header=True)
        nx = hd1['naxis1']
        ny = hd1['naxis2']
        w1 = WCS(hd1)
        vcoo = w1.pixel_to_world([0,nx-1,nx-1,0],[0,0,ny-1,ny-1])
        vra = vcoo.ra.deg
        vdec = vcoo.dec.deg
        olap[i] = coords.doPolygonsOverlap(bvra,bvdec,vra,vdec)
        vx,vy = tile['wcs'].world_to_pixel(vcoo)
        vxarr[i,:] = vx
        vyarr[i,:] = vy
    # Require at least a 2 pixel overlap in X and Y
    g, = np.where((olap==True) & (np.max(vxarr,axis=1) >= 2) & (np.max(vyarr,axis=1) >= 2) &
                  (np.min(vxarr,axis=1) <= tile['nx']-3) & (np.min(vyarr,axis=1) <= tile['ny']-3))
    if len(g) == 0:
        logger.info('No chips overlap this brick')
        return
    logger.info(str(len(g))+' chips overlap this brick')
    chstr = chstr[g]
    nchstr = len(chstr)

    # APPLY cuts on the exposures
    #-----------------------------
    logger.info('Applying quality, zero-point, filter and exptime cuts')

    # Zero-point structure, from NSC
    zpstr = np.zeros(7,dtype=np.dtype([('instrument',np.str,10),('filter',np.str,10),('amcoef',(float,2)),('thresh',float)]))
    #zpstr = replicate({instrument:'',filter:'',amcoef:fltarr(2),thresh:0.5},7)
    zpstr['instrument'] = 'c4d'
    zpstr['filter'] = ['u','g','r','i','z','Y','VR']
    zpstr['amcoef'][0] = [-1.60273, -0.375253]  # c4d-u
    zpstr['amcoef'][1] = [0.277124, -0.198037]  # c4d-g
    zpstr['amcoef'][2] = [0.516382, -0.115443]  # c4d-r
    zpstr['amcoef'][3] = [0.380338, -0.067439]  # c4d-i
    zpstr['amcoef'][4] = [0.074517, -0.067031]  # c4d-z
    zpstr['amcoef'][5] = [-1.07800, -0.060014]  # c4d-Y
    zpstr['amcoef'][6] = [1.111859, -0.083630]  # c4d-VR

    # Convert to additive zero-point as used in NSC
    zpterm = -chstr['calib_zpterm']
    # Fix early DES exposures that used different units/gain
    gdes, = np.where(chstr['gain'] < 2)
    if len(gdes) > 0:
        zpterm[gdes] -= 1.55
    # global zpterm and airmass correction
    for i in range(len(zpstr)):
        ind, = np.where(chstr['filter']==zpstr['filter'][i])
        if len(ind)>0:
            zpterm[ind] -= np.polyval(np.flip(zpstr['amcoef'][i]),chstr['airmass'][ind])

    fwhmthresh = 2.0  # seeing 2.0" threshold
    filt = np.char.array(chstr['filter'].astype(str)).ljust(1)  # first characater
    gdch, = np.where((chstr['fwhm']*chstr['pixscale'] <= fwhmthresh) & (chstr['exptime'] >= 90.) & (zpterm >= -0.5) &
                     np.isfinite(zpterm) & np.isfinite(chstr['apcor']) &
                     ((filt== 'u') | (filt=='g') | (filt=='r') | (filt=='i') | (filt=='z') | (filt=='Y')))
    if len(gdch) == 0:
        logger.info('No chips passed the cuts')
        return
    logger.info(str(len(gdch))+' chips passed the cuts')
    chstr = chstr[gdch]
    nchstr = len(chstr)
    chstr['file'] = np.char.array(chstr['file'].astype(str)).strip()
    chstr['base'] = np.char.array(chstr['base'].astype(str)).strip()

    # Check if we need to update
    if update and redo==False:
        # Load previous meta
        metafile = bdir+brick+'_meta.fits'
        if os.path.exists(metafile):
            meta0 = Table.read(metafile,1)
            meta0['base'] = str(meta0['base'])
            ind1,ind2 = dln.match(meta0['base'],chstr['base'])
            if nchstr == len(meta0) and nchstr == nmatch:
                logger.info('Nothing to UPDATE')
                return
            # Moving previous catalogs to a backup
            oldfiles = glob(bdir+brick+'*fits*')
            oldfiles += glob(bdir+brick+'allframe.opt')
            oldfiles += glob(bdir+brick+'default.*')
            if len(oldfiles)>0:
                bakdir = bdir+'bak'+logtime
                logger.info('Backing up old files to '+bakdir)
                os.makedirs(bakdir)
                for f in oldfiles:
                    shutil.move(f,bakdir)
            else:
                logger.info('No old files to backup')
            logger.info('There are new exposures to include.  UPDATING')

    # Create temporary local directory to perform the work/processing
    #  copy everything to it
    if os.path.exists(workdir)==False:
        os.makedirs(workdir)

    procdir = tempfile.mkdtemp(prefix="brk",dir=workdir)
    procdir += '/'
    os.chmod(procdir,0o755)
    logger.info('Working in temporary directory '+procdir)
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

    # Copy the files
    #----------------
    logger.info('Copying over the necessary files to '+procdir)
    # Chip loop
    for i in range(nchstr):
        logger.info('Copying over files for '+chstr['file'][i])
        # New filename
        #chstr[i].newbase = chstr[i].base
        odir = os.path.dirname(chstr['file'][i])+'/'
        obase = photred.utils.fitsext(chstr['file'][i],basename=True)
        # Need symlinks to .psf, .als
        if os.path.exists(odir+obase+'.psf') == False:
            logger.info(odir+obase+'.psf NOT FOUND.  Skipping this chip')
            continue
        if os.path.exists(odir+obase+'.als') == False:
            logger.info(odir+obase+'.als NOT FOUND.  Skipping this chip')
            continue
        for e in ['.psf','.als','.ap','.opt','.als.opt','.log']:
            dln.remove(procdir+str(chstr['base'][i])+e,allow=True)
        for e in ['.psf','.als','.ap','.opt','.als.opt','.log']:
            dln.remove(procdir+'/'+obase+e,allow=True)
            shutil.copyfile(odir+obase+e,procdir+obase+e)
            os.chmod(procdir+chstr['base'][i]+e,0o755)  # make sure they are writable
        # Copy the fits, fits resource file and header files locally
        if os.path.exists(odir+obase+'.fits') == False:
            logger.info(odir+obase+'.fits NOT FOUND.  Skipping this chip')
            continue
        dln.remove(procdir+obase+'.fits',allow=True)
        shutil.copyfile(odir+obase+'.fits',procdir+obase+'.fits')
        if os.path.exists(odir+'.'+obase+'.fits'):
            dln.remove(procdir+'.'+obase+'.fits',allow=True)
            shutil.copyfile(odir+'.'+obase+'.fits',procdir+'.'+obase+'.fits')
        if os.path.exists(odir+obase+'.fits.head'):
            dln.remove(procdir+obase+'.fits.head',allow=True)
            shutil.copyfile(odir+obase+'.fits.head',procdir+obase+'.fits.head')
        os.chmod(procdir+chstr['base'][i]+'.fits',0o755)

    # Step 2: Run DAOMATCH_TILE.PRO on the files
    #--------------------------------------------
    logger.info('Step 2: Matching up objects with DAOMATCH_TILE')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    os.chdir(procdir)
    group = {'x0':0,'y0':0}
    daomatch.daomatch_tile(np.char.array(chstr['base'])+'.als',tile,group)

    # Step 3: Run ALLFRAME
    #----------------------
    logger.info('Step 3: Run ALLFRAME')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    # DO I NEED TO HAVE IT TRIM THE COMBINED IMAGE???
    mchbase = procdir+chstr['base'][0]
    mchfile = mchbase+'.mch'# allframe needs absolute path
    import pdb; pdb.set_trace()
    alf.allframe(mchfile,tile=tile,setupdir=bdir,scriptsdir=scriptsdir,irafdir=irafdir,
                 logfile=logfile,catformat='FITS',imager=thisimager,geocoef=0)
    magfile = chstr['base'][0]+'.mag'
    if os.path.exists(magfile) == False:
        logger.info(magfile+' NOT FOUND')
        return

    # Step 4: Calculate coordinates
    #-------------------------------
    logger.info('Step 4: Adding coordinates')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    # Load the MCH file
    alsfiles,trans,magoff = io.readmch(mchfile)
    nalsfiles = len(alsfiles)
    # Load the photometry file
    instphot = io.readfile(magfile)
    ninstphot = len(instphot)
    logger.info('Nstars = '+str(ninstphot))
    # Converting to IDL X/Y convention, starting at (0,0)
    # DAOPHOT has X/Y start at (1,1)
    ra,dec = tile['wcs'].pixel_to_world(instphot.x-1.0,instphot.y-1.0)
    #HEAD_XYAD,tile.head,instphot.x-1.0,instphot.y-1.0,ra,dec,/degree

    # Step 5: Calibrating photometry with zero-points
    #-------------------------------------------------
    logger.info('Step 5: Calibrating photometry with zero-points')
    logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
    cmag = np.zeros((ninstphot,nchstr),float)+99.99
    cerr = np.zeros((ninstphot,nchstr),float)+9.99
    # Chip loop
    for i in range(nchstr):
        # id, x, y, ra, dec, unsolved magnitudes, chi, sharp
        imag = instphot['mag'+str(i+1)]
        ierr = instphot['err'+str(i+1)]
        gdmag, = np.where(imag < 50)
        if len(gdmag) > 0:
            # exptime, aperture correction, zero-point
            # aperture correction is SUBTRACTIVE, makes it brighter
            # ZPTERM is a SUBTRACTIVE constant offset
            cmag[gdmag,i] = imag[gdmag] + 2.5*np.log10(chstr['exptime'][i]) - chstr['apcor'][i] - chstr['calib_zpterm'][i]
            # Add zero-point error in quadrature
            cerr[gdmag,i] = np.sqrt(ierr[gdmag]**2+chstr['calib_zptermsig'][i]**2)
    # Calculate average photometry per filter
    ufilt = np.unique(chstr['filter'])
    nufilt = len(ufilt)
    avgmag = np.zeros((ninstphot,nufilt),float)
    avgerr = np.zeros((ninstphot,nufilt),float)
    ndet = np.zeros((ninstphot,nufilt),int)
    for i in range(nufilt):
        gdf, = np.where(chstr['filter'] == ufilt[i])
        # Single exposure
        if len(gdf) == 1:
            avgmag[:,i] = cmag[:,gdf[0]]
            avgerr[:,i] = cerr[:,gdf[0]]
            ndet[:,i] = np.sum(cmag[:,gdf[0]] < 50)
            # Multiple exposures
        else:
            # Loop through all of the exposures and add up the flux, totalwt, etc.
            totalwt = np.zeros(ninstphot,float)
            totalfluxwt = np.zeros(ninstphot,float)
            for k in range(ngdf):
                gdmag, = np.where(cmag[:,gdf[k]] < 50)
                if len(gdmag) > 0:
                    totalwt[gdmag] += 1.0/cerr[gdmag,gdf[k]]**2
                    totalfluxwt[gdmag] += 2.5118864**cmag[gdmag,gdf[k]] * (1.0/cerr[gdmag,gdf[k]]**2)
                    ndet[gdmag,i] += 1
            newflux = totalfluxwt/totalwt
            newmag = 2.5*np.log10(newflux)
            newerr = np.sqrt(1.0/totalwt)
            bdmag, = np.where(~np.isfinite(newmag))
            if len(bdmag) > 0:
                newmag[bdmag] = 99.99
                newerr[bdmag] = 9.99
            avgmag[:,i] = newmag
            avgerr[:,i] = newerr
    # Measure scatter
    scatter = np.zeros((ninstphot,nufilt),float)+99.99
    for i in range(nufilt):
        gdf, = np.where(chstr['filter'] == ufilt[i])
        if len(gdf) > 1:
            totaldiff = np.zeros(ninstphot,float)
            for k in range(ngdf):
                gdmag, = np.where(cmag[:,gdf[k]] < 50)
                if len(gdmag) > 0:
                    totaldiff[gdmag] += (avgmag[gdmag,i]-cmag[gdmag,gdf[k]])**2
            scatter[:,i] = np.sqrt( totaldiff/np.maximum(ndet[:,i],1) )
            bd, = np.where(ndet[:,i] <= 1)
            if len(bd) > 0 :
                scatter[bd,i]=99.99
    # Create final catalog schema
    photdt = [('objid',(np.str,100)),('x',float),('y',float),('ra',float),('dec',float)]

    # Add columns for calibrated single-epoch photometry columns
    cmagnames = np.zeros(nchstr,(np.str,100))
    cerrnames = np.zeros(nchstr,(np.str,100))
    for i in range(nufilt):
        ind, = np.where(chstr['filter'] == ufilt[i])
        cmagnames[ind] = [ufilt[i].upper()+str(j+1)+'MAG' for j in np.arange(nind)]
        cerrnames[ind] = [ufilt[i].upper()+str(j+1)+'ERR' for j in np.arange(nind)]
    for i in range(nchstr):
        photdt += [(cmagnames[i],np.float32),(cerrnames[i],np.float32)]
    # Add columns for average photometry per filter
    for i in range(nufilt):
        photdt += [(ufilt[i]+'MAG',np.float32),(ufilt[i]+'ERR',np.float32),(ufilt[i]+'SCATTER',np.float32),('NDET'+ufilt[i],int)]
    # Extra columns
    photdt += [('chi',np.float32),('sharp',np.float32),('prob',np.float32),('ebv',np.float32)]
    # other SE columns
    photdt += [('mag_auto',np.float32),('magerr_auto',np.float32),('asemi',np.float32),('bsemi',np.float32),
               ('theta',np.float32),('ellipticity',np.float32),('fwhm',np.float32)]
    # in unique brick area
    dt += [('brickuniq',boolean)]
    # Create final catalog
    phot = np.zeros(ninstphot,dtype=np.dtype(photdt))
    for n in instphot:
        phot[n] = instphot[n]
    phtags = phot.colnames
    # object IDs
    phot['objid'] = brick+'.'+np.char.array(instphot['id'])
    # Stuff in the coordinates calculated above
    phot['ra'] = ra
    phot['dec'] = dec
    # Stuff in the calibrated single-epoch photometry columns
    for i in range(nchstr):
        phot[cmagnames[i]] = cmag[:,i]
        phot[cerrnames[i]] = cerr[:,i]
    # Stuff in the average photometry per filter
    for i in range(nufilt):
        phot[ufilt[i].upper()+'MAG'] = avgmag[:,i]
        phot[ufilt[i].upper()+'ERR'] = avgerr[:,i]
        phot[ufilt[i].upper()+'SCATTER'] = scatter[:,i]
        phot['NDET'+ufilt[i].upper()] = ndet[:,i]

    # Calculate SFD E(B-V)
    #GLACTC,phot.ra,phot.dec,2000.0,glon,glat,1,/deg
    coo = SkyCoord(phot['ra'], phot['dec'], unit='deg', frame='icrs')
    coo_gal = coo.transform_to('galactic')
    sfd = SFDQuery()
    phot['ebv'] = sfd(coords_gal)

    # THIS IS NOW BEING DONE IN DELVERED_FINALCAT.PRO THAT COMBINES ALL CATALOGS
    # Only include objects that are INSIDE the UNIQUE brick area
    #   INCLUSIVE at the lower RA and DEC limit
    # Getting objects that are in the UNIQUE brick area
    if brickstr1.dec == -90:
        # the brick right at the pole does not have any RA limits
        ginside, = np.where(phot['dec'] < brickstr1['dec2'])
    else:
        ginside, = np.where((phot['ra'] >= brickstr1['ra1']) & (phot['ra'] < brickstr1['ra2']) &
                            (phot['dec'] >= brickstr1['dec1']) & (phot['dec'] < brickstr1['dec2']))
    if len(ginside)>0:
        phot['brickuniq'][ginside] = True

    # Get some meta-data
    for i in range(nchstr):
        alffile = chstr['base'][i]+'.alf'
        if os.path.exists(alffile):
            chstr['alf_nsources'][i] = dln.numlines(alffile)-3

    # Make the exposure-level forced photometry catalog
    #--------------------------------------------------
    # Load the individual ALF files to get chi, sharp
    # Load TFR file to conver ALF IDs to final object ID
    tfrfile = mchbase+'_comb.tfr'
    alffiles,tfrstr = io.readtfr(tfrfile)
    expdt = [('id',(np.str,100)),('objid',(np.str,100)),('exposure',(np.str,100)),('ccdnum',int),('filter',(np.str,10)),
             ('mjd',float),('x',float),('y',float),('ra',float),('dec',float),('imag',float),('ierr',float),
             ('mag',float),('err',float),('sky',float),('chi',float),('sharp',float)]
    expcat = np.zeros(int(np.sum(cmag < 50))+10000,dtype=np.dtype(expdt))
    cnt = 0
    for i in range(nchstr):
        base1 = chstr['base'][i]
        fitsfile = base1+'.fits'
        if os.path.exists(alffiles[i]) and os.path.exists(fitsfile):
            alf = io.readals(alffiles[i])
            if len(alf)==0:
                continue
            # Sometimes the rows are duplicated in the ALF file
            ui = np.uniq(alf.id,np.argsort(alf.id))
            if len(ui) < nalf:
                alf = alf[ui]
                nalf = len(alf)
            head = io.readfile(fitsfile,header=True)

            # Calibrate the photometry
            # exptime, aperture correction, zero-point
            # aperture correction is SUBTRACTIVE, makes it brighter
            # ZPTERM is a SUBTRACTIVE constant offset
            cmag1 = alf['mag'] + 2.5*np.log10(chstr['exptime'][i]) - chstr['apcor'][i] - chstr['calib_zpterm'][i]
            # Add zero-point error in quadrature
            cerr1 = np.sqrt(alf['err']**2+chstr['calib_zptermsig'][i]**2)

            # Coordinates
            ra1,dec1 = head.pixel_to_world(alf['x']-1,alf['y']-1)
            #HEAD_XYAD,head,alf.x-1,alf.y-1,ra1,dec1,/deg

            # MATCH up ALF IDs to TFR INDEX
            # One row per unique object
            # the INDEX values are 1-based indices into the ALF files
            ind1,ind2 = dln.match(tfrstr.index[i],lindgen(nalf)+1)
            #MATCH,tfrstr.index[i],lindgen(nalf)+1,ind1,ind2,/sort,count=nmatch
            objid = brick+'.'+np.char.array(tfrstr['id'][ind1])

            # Create the new catalog
            newcat = np.zeros(nalf,dtype=np.dtype(dt))
            newcat['objid'] = objid
            newcat['id'] = chstr['expnum'][i]+'_'+chstr['chip'][i]+'.'+np.char.array(alf['id'])
            newcat['exposure'] = chstr['base'][i]
            newcat['ccdnum'] = chstr['chip'][i]
            newcat['filter'] = chstr['filter'][i]
            dateobs = chstr['utdate'][i]+'T'+chstr['uttime'][i]
            t = Time(dateobs)
            newcat['mjd'] = t.mjd
            newcat['x'] = alf['x']
            newcat['y'] = alf['y']
            newcat['ra'] = ra1
            newcat['dec'] = dec1
            newcat['imag'] = alf['mag']
            newcat['ierr'] = alf['err']
            newcat['mag'] = cmag1
            newcat['err'] = cerr1
            newcat['sky'] = alf['sky']
            newcat['chi'] = alf['chi']
            newcat['sharp'] = alf['sharp']

            # Add more elements
            if cnt+nalf > len(expcat):
                expcat = add_elements(expcat,100000>nalf)

            # Add to global catalog
            expcat[cnt:cnt+nalf-1] = newcat
            cnt += nalf

        else:
            logger.info(alffile+' NOT FOUND')
        expcat = expcat[0:cnt]   # this should not be needed

        # Object catalog
        #---------------
        # do not want all of the individual epoch photometry
        #  these are between the DEC and the first unique mean magnitude (e.g. GMAG, and GERR)
        lo = where(phtags=='DEC')[0][0]
        hi = where(np.char.array(phtags).upper() == ufilt[0].upper()+'MAG')[0][0]
        objdt = photdt[0:lo]
        objdt += photdt[hi:]
        obj = np.zeros(ninstphot,dtype=np.dtype(objdt))
        for n in obj.colnames:
            obj[n] = phot[h]

        # Saving final catalog
        logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
        photfile = bdir+brick+'.fits'
        if len(phot.colnames) <= 999:
            logger.info('Writing photometry to '+photfile+'.gz')
            phot.writeto(photfile,overwrite=True)
            out = subprocess.check_output(['gzip','-f',photfile],shell=False)
        else:
            logger.info('Too many columns for FITS.  Saving as pickle file instead. '+bdir+brick+'.pkl')
            with open(bdir+brick+'.pkl','wb') as f:
                pickle.dump(phot,f)

        # Saving object photometry catalog
        objfile = bdir+brick+'_object.fits'
        obj.writeto(objfile,overwrite=True)
        out = subprocess.check_output(['gzip','-f',objfile],shell=False)

        # Saving exposure-level forced photometry catalog
        expfile = bdir+brick+'_expforced.fits'
        expcat.writeto(expfile,overwrite=True)
        out = subprocess.check_output(['gzip','-f',expfile],shell=False)

        # Save metadata
        metafile = bdir+brick+'_meta.fits'
        logger.info('Writing meta-data to '+metafile)
        chstr.writeto(metafile,overwrite=True)

        # Delete files we don't want to keep
        #-----------------------------------
        # Individual fits files
        alsbase = os.path.basename(alsfiles,'.als')
        # if fits files have resources files then replace the fits file by a
        #  dummy file
        for i in range(len(alsbase)):
            exists = os.path.exists(alsbase[i]+'.fits')
            size = os.path.getsize(alsbase[i]+'.fits')
            if exists and size > 1:
                os.remove(alsbase[i]+'.fits')
                if os.path.exists('.'+alsbase[i]+'.fits'):
                    dln.writelines(alsbase[i]+'.fits','')
        # Combined files
        #   _comb  lst, lst1, lst2, lst1.chi, grp, nst, lst2.chi, plst.chi, psfini.ap
        #   nei, als.inp, a.fits, cmn.log, cmn.coo, cmn.ap, cmn.lst,
        #   _sub.fits, _sub.cat, _sub.als, _all.coo, makemag
        base = os.path.basename(mchfile,'.mch')
        ext = ['.lst','.lst1','.lst2','.lst1.chi','.lst2.chi','.grp','.nst','.plst.chi','.nei',
               '.als.inp','.cmn.log','.cmn.coo','.cmn.ap','.cmn.lst','a.fits','a.fits.fz',
               '_sub.fits','_sub.cat','_sub.als','_all.coo','.makemag']
        for e in ext:
            if os.path.exists(base+e): os.remove(base+e)
        if os.path.exists('check.fits'): os.remove('check.fits')

        # fpack _comb.fits and _combs.fits
        for f in [base+'_comb.fits.fz',base+'_combs.fits.fz']:
            if os.path.exists(f): os.remove(f)
        out = subprocess.check_output(['fpack','-D','-Y',base+'_comb.fits'],shell=False)
        out = subprocess.check_output(['fpack','-D','-Y',base+'_combs.fits'],shell=False)
        # gzip _comb.mask.fits and _comb.bpm.fits
        out = subprocess.check_output(['gzip','-f',base+'_comb.mask.fits'],shell=False)
        out = subprocess.check_output(['gzip','-f',base+'_comb.bpm.fits'],shell=False)

        # Copy everything left back to the original directory
        logger.info('Copying files back to '+bdir)
        allfiles = glob(procdir+'*')
        allfiles += glob(procdir+'.*.fits')
        nallfiles = len(allfiles)
        for f in allfiles:
            base = os.path.basename(f)
            if os.path.exists(bdir+'/'+base): os.remove(bdir+'/'+base)
            shutil.move(f,bdir)
        os.rmdir(procdir)  # delete temporary processing directory

        logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))
        logger.info('DELVERED_FORCEBRICK:ne after '+str(time.time()-t0,2)+' sec.')

        os.chdir(curdir)  # back to original directory

        # Create JOINT catalogs
        logger.info('')
        logger.info('CREATE JOINT CATALOGS')
        logger.info('')
        logger.info(datetime.now().strftime("%a %b %d %H:%M:%S %Y"))

        if redo or update:
            jntredo = True
        else:
            jntredo = False
        delvered_jointbrickcats(brick,logfile=logfile,redo=jntredo)


import numpy as np
import time

from urh.dev.native.Device import Device
from urh.dev.native.lib import hackrf
from urh.util.Logger import logger

class HackRF(Device):
    BYTES_PER_SAMPLE = 2  # HackRF device produces 8 bit unsigned IQ data

    def __init__(self, bw, freq, gain, srate, is_ringbuffer=False):
        super().__init__(bw, freq, gain, srate, is_ringbuffer)
        self.success = 0
        self.error_not_open = -4242

        self._max_bandwidth = 28e6
        self._max_frequency = 6e9
        self._max_sample_rate = 20e6
        self._max_gain = 40

        self.error_codes = {
            0: "HACKRF_SUCCESS",
            1: "HACKRF_TRUE",
            1337: "TIMEOUT ERROR",
            -2: "HACKRF_ERROR_INVALID_PARAM",
            -5: "HACKRF_ERROR_NOT_FOUND",
            -6: "HACKRF_ERROR_BUSY",
            -11: "HACKRF_ERROR_NO_MEM",
            -1000: "HACKRF_ERROR_LIBUSB",
            -1001: "HACKRF_ERROR_THREAD",
            -1002: "HACKRF_ERROR_STREAMING_THREAD_ERR",
            -1003: "HACKRF_ERROR_STREAMING_STOPPED",
            -1004: "HACKRF_ERROR_STREAMING_EXIT_CALLED",
            -4242: "HACKRF NOT OPEN",
            -9999: "HACKRF_ERROR_OTHER"
        }

        # self.__lut = np.zeros(0xffff + 1, dtype=np.complex64)
        # self.little_endian = False
        # for i in range(0, 0xffff + 1):
        #     if self.little_endian:
        #         real = (float(np.int8(i & 0xff))) * (1.0 / 128.0)
        #         imag = (float(np.int8(i >> 8))) * (1.0 / 128.0)
        #     else:
        #         real = (float(np.int8(i >> 8))) * (1.0 / 128.0)
        #         imag = (float(np.int8(i & 0xff))) * (1.0 / 128.0)
        #
        #     self.__lut[i] = complex(real, imag)

    def reopen(self):
        if self.is_open:
            hackrf.reopen()

    def open(self, init=True):
        if not self.is_open:
            if init:
                ret = hackrf.setup()
            else:
                ret = hackrf.open()

            self.is_open = ret == self.success
            self.log_retcode(ret, "open")

    def close(self, exit=True):
        if self.is_open:
            logger.info("HackRF: Attempting to close...")
            time.sleep(0.01)
            ret = hackrf.close()
            self.is_open = ret != self.success

            if self.is_open:
                logger.error("Failed to close HackRF")
            else:
                logger.info("Successfully closed HackRF")

            self.log_retcode(ret, "close")
            if exit:
                self.exit()

    def exit(self):
        return hackrf.exit()

    def start_rx_mode(self):
        if self.is_open:
            self.init_recv_buffer()
            self.set_device_parameters()
            ret = hackrf.start_rx_mode(self.callback_recv)
            self.is_receiving = ret == self.success

            if self.is_receiving:
                logger.info("HackRF: Starting receiving thread")
                self._start_readqueue_thread()

            self.log_retcode(ret, "start_rx_mode")
        else:
            self.log_retcode(self.error_not_open, "start_rx_mode")

    def stop_rx_mode(self, msg):
        self.is_receiving = False

        logger.info("HackRF: Stopping RX Mode: "+msg)

        if hasattr(self, "read_queue_thread") and self.read_queue_thread.is_alive():
            try:
                self.read_queue_thread.join(0.001)
                logger.info("HackRF: Joined read_queue_thread")
            except RuntimeError:
                logger.error("HackRF: Could not join read_queue_thread")

        if self.is_open:
            logger.info("stopping HackRF rx mode ({0})".format(msg))
            logger.warning("closing because stop_rx_mode of HackRF is bugged and will not allow re receive without close")
            self.close(exit=False)

    def switch_from_rx2tx(self):
        # https://github.com/mossmann/hackrf/pull/246/commits/4f9665fb3b43462e39a1592fc34f3dfb50de4a07
        self.reopen()

    def start_tx_mode(self, samples_to_send: np.ndarray = None, repeats=None, resume=False):
        if self.is_open:
            self.init_send_parameters(samples_to_send, repeats, resume=resume)
            retcode = hackrf.start_tx_mode(self.callback_send)

            if retcode == self.success:
                self.is_transmitting = True
                self._start_sendbuffer_thread()
            else:
                self.is_transmitting = False
        else:
            retcode = self.error_not_open
        self.log_retcode(retcode, "start_tx_mode")

    def stop_tx_mode(self, msg):
        self.is_transmitting = False
        try:
            self.send_buffer_reader.close()
            self.send_buffer.close()
        except AttributeError:
            logger.warning("HackRF: Could not close send buffer, because it was not open yet")

        if self.is_open:
            logger.info("stopping HackRF tx mode ({0})".format(msg))
            logger.info("closing because stop_tx_mode of HackRF is bugged and never returns")
            self.close(exit=False)

    def set_device_bandwidth(self, bw):
        if self.is_open:
            retcode = hackrf.set_baseband_filter_bandwidth(bw)
        else:
            retcode = self.error_not_open
        self.log_retcode(retcode, "set_bandwidth", bw)

    def set_device_frequency(self, value):
        if self.is_open:
            retcode = hackrf.set_freq(value)
        else:
            retcode = self.error_not_open
        self.log_retcode(retcode, "set_frequency", value)

    def set_device_gain(self, gain):
        if self.is_open:
            hackrf.set_lna_gain(gain)
            hackrf.set_vga_gain(gain)
            hackrf.set_txvga_gain(gain)

    def set_device_sample_rate(self, sample_rate):
        if self.is_open:
            retcode = hackrf.set_sample_rate(sample_rate)
        else:
            retcode = self.error_not_open

        self.log_retcode(retcode, "set_sample_rate", sample_rate)

    def unpack_complex(self, buffer, nvalues: int):
        result = np.empty(nvalues, dtype=np.complex64)
        unpacked = np.frombuffer(buffer, dtype=[('r', np.int8), ('i', np.int8)])
        result.real = unpacked['r'] / 128.0
        result.imag = unpacked['i'] / 128.0
        return result

    def pack_complex(self, complex_samples: np.ndarray):
        assert complex_samples.dtype == np.complex64
        # tostring() is a compatibility (numpy<1.9) alias for tobytes(). Despite its name it returns bytes not strings.
        return (128 * complex_samples.view(np.float32)).astype(np.int8).tostring()


"""Thermostate Module"""
#####Imports#####
import numpy as np
import copy

class ThermodynamicState(object):
    """
    Dummy super to wrap the functionality of various classes that define invariant distributions
    ThermodynamicStates define invariants of target distributions and proposal kernels.
    """
    def __init__(self, **kwargs):
        """
        Dummy init method.
        """
        #update other kwargs that have not been set
        self.__dict__.update(kwargs)

    def set_parameters(self, parameters, **kwargs):
        """
        Dummy wrapper to set parameters to define ThermodynamicState invariants
        args
            parameters : np.array or float
                parameters defining new invariant thermodynamic_state
        """
        pass

class PersesCompoundAlchemicalThermodynamicState(ThermodynamicState):
    """
    Wrapper subclass for openmmtools.states.CompoundThermodynamicState object as defined by Perses
    """
    def __init__(self, thermodynamic_state, **kwargs):
        """
        Wrap an openmmtools thermodynamic_state or CompoundThermodynamicState
        """
        from perses.annihilation.lambda_protocol import LambdaProtocol
        self.thermodynamic_state = copy.deepcopy(thermodynamic_state)
        super(OpenMMToolsThermodynamicState, self).__init__(**kwargs)

    def set_parameters(self, parameters, lambda_protocol = LambdaProtocol(), **kwargs):
        """
        wrapper method to update parameters
        args
            parameters : np.array or float
                parameters defining new invariant thermodynamic_state
            lambda_protocol : perses.annihilation.lambda_protocol.LambdaProtocol
        """
        self.thermodynamic_state.set_alchemical_parameters(parameters, lambda_protocol)


import numpy as np
import cv2, tqdm
import tensorflow.keras.backend as K
K.set_image_data_format('channels_last')

import keras.backend as K
K.set_image_data_format('channels_last')

from PIL import Image
# retinanet
from libs.keras_retinanet.keras_retinanet.bin.train import *
from libs.keras_retinanet.keras_retinanet import models
from libs.keras_retinanet.keras_retinanet.utils.image import preprocess_image, resize_image
from libs.keras_retinanet.keras_retinanet.utils.visualization import draw_box, draw_caption
from libs.keras_retinanet.keras_retinanet.utils.gpu import setup_gpu
from libs.yolo3.train import yolo_det
from libs.keras_retinanet.keras_retinanet.utils.colors import label_color
from libs.EfficientDet.train import efficientdet_train, efficientdet_parse_args
from libs.segmentation import Segmentation
import matplotlib.pyplot as plt

class Detection(object):
    def __init__(self, img_dir, label_csv_path, split_rate=0.2, batch_size=32,
                 resized_shape=(480, 720), base='detection',yolo=False):
        """

        :param img_dir:
        :param label_csv_path:
        :param split_rate:
        :param batch_size:
        :param resized_shape:
        :param base:
        """
        self.img_dir = img_dir
        self.resized_shape = resized_shape
        self.label_csv_path = label_csv_path
        self.base = base
        self.yolo=yolo
        if base.lower() == 'detection':
            if not yolo:
                self.split_rate = split_rate
                self.batch_size = batch_size
                self.train_path, self.val_path, self.cls_path = self.process()
                self.train_steps = len(open(self.train_path, 'r', encoding='utf-8').readlines()) // batch_size
                self.val_steps = len(open(self.val_path, 'r', encoding='utf-8').readlines()) // batch_size
            else:
                self.app=yolo_det()
        elif base.lower() == 'segmentation':
            self.app = Segmentation(img_dir=self.img_dir,
                                    label_csv_path=self.label_csv_path,
                                    batch_size=self.batch_size, split_rate=self.split_rate,
                                    shrink=True, use_aug=False)

    def process(self):
        """

        :return:
        """
        shape = self.resized_shape
        os.makedirs('labels/detection/', exist_ok=True)
        os.makedirs('saved_models/detection', exist_ok=True)
        if os.path.exists(
                'labels/detection/anno_detection_train_{}x{}.csv'.format(shape[0], shape[1])) and os.path.exists(
            'labels/detection/anno_detection_val_{}x{}.csv'.format(shape[0], shape[1])) and os.path.exists(
            'labels/detection/detection_class.csv'):
            return 'labels/detection/anno_detection_train_{}x{}.csv'.format(shape[0], shape[
                1]), 'labels/detection/anno_detection_val_{}x{}.csv'.format(shape[0], shape[
                1]), 'labels/detection/detection_class.csv'
        else:
            f = open(self.label_csv_path, 'r', encoding='utf-8').readlines()
            train_writer = open('labels/detection/anno_detection_train_{}x{}.csv'.format(shape[0], shape[1]), 'w',
                                encoding='utf-8')
            val_writer = open('labels/detection/anno_detection_val_{}x{}.csv'.format(shape[0], shape[1]), 'w',
                              encoding='utf-8')
            class_writer = open('labels/detection/detection_class.csv', 'w', encoding='utf-8')
            all_index = np.array(range(len(f)))
            val_index = np.random.choice(all_index, size=int(self.split_rate * len(f)), replace=False)
            for idx, line in tqdm.tqdm(enumerate(f), total=len(f)):
                line = line.strip('\n').split(',')
                img_path = os.path.join(self.img_dir, line[0])
                img = None
                try:
                    img = Image.open(img_path)
                except Exception as E:
                    print(E)
                if img is None:
                    continue
                original_w, original_h = img.size
                h_rate = shape[0] / original_h
                w_rate = shape[1] / original_w
                line[1], line[3] = int(int(line[1]) * w_rate), int(int(line[3]) * w_rate)
                line[2], line[4] = int(int(line[2]) * h_rate), int(int(line[4]) * h_rate)
                if idx not in val_index:
                    train_writer.write('{},{},{},{},{},{}\n'.format(
                        img_path, line[1], line[2], line[3], line[4], 'face'
                    ))
                else:
                    val_writer.write('{},{},{},{},{},{}\n'.format(
                        img_path, line[1], line[2], line[3], line[4], 'face'))
            class_writer.write('face,0\n')
            class_writer.close()
            val_writer.close()
            train_writer.close()
            return 'labels/detection/anno_detection_train_{}x{}.csv'.format(shape[0], shape[
                1]), 'labels/detection/anno_detection_val_{}x{}.csv'.format(shape[0], shape[
                1]), 'labels/detection/detection_class.csv'

    def train_model(self, gpu=3, directly_train=True, backbone='resnet152', method='retinanet',
                    model_path=None, augmentation=True, gpu_fraction=0.3):
        """

        :param gpu:
        :param directly_train:
        :param backbone:
        :param method:
        :param model_path:
        :param augmentation:
        :param gpu_fraction:
        :return:
        """

        setup_gpu(gpu)
        import tensorflow as tf
        import keras.backend.tensorflow_backend as KTF
        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = gpu_fraction
        session = tf.Session(config=config)
        KTF.set_session(session)
        if method.lower() == 'retinanet' and directly_train and self.base.lower() == 'detection' and not self.yolo:
            args = parse_args()
            args.dataset_type = 'csv'
            args.batch_size = self.batch_size
            args.backbone = backbone
            args.gpu = gpu
            args.annotations = self.train_path
            args.classes = self.cls_path
            args.val_annotations = self.val_path
            args.evaluation = False
            args.snapshot_path = 'saved_models/detection'
            args.random_transform = True
            args.lr = 1e-3
            args.no_resize = True
            args.augmentation = augmentation
            args.width = self.resized_shape[1]
            args.height = self.resized_shape[0]
            args.compute_val_loss = True
            if model_path is None:
                # args.lr = 1e-3
                # args.freeze_backbone = True
                # args.snapshot = None
                # args.weights = None
                # args.imagenet_weights = True
                # args.epochs = 50
                # detection_main(args, train_steps=self.train_steps, val_steps=self.val_steps)

                args.lr = 1e-4
                args.freeze_backbone = False
                args.epochs = 100
                args.snapshot = '{dir}/retinanet_{backbone}_{dataset_type}_{height}x{width}.h5'.format(
                    dir=args.snapshot_path,
                    backbone=args.backbone,
                    dataset_type=args.dataset_type,
                    height=args.height, width=args.width)
                detection_main(args, train_steps=self.train_steps, val_steps=self.val_steps)
            else:
                args.lr = 1e-5
                args.epochs = 100
                args.freeze_backbone = False
                args.snapshot = model_path
                detection_main(args, train_steps=self.train_steps, val_steps=self.val_steps)

        elif method.lower() == 'efficientdet' and directly_train and self.base.lower() == 'detection' and not self.yolo:

            args = efficientdet_parse_args()
            args.dataset_type = 'csv'
            args.batch_size = self.batch_size
            args.backbone = backbone
            args.gpu = gpu
            args.annotations_path = self.train_path
            args.classes_path = self.cls_path
            args.val_annotations_path = self.val_path
            args.evaluation = False
            args.snapshot_path = 'saved_models/detection'
            args.detect_quadrangle = False
            args.detect_text = False
            args.no_resize = True
            args.augmentation = augmentation
            args.width = self.resized_shape[1]
            args.height = self.resized_shape[0]
            args.compute_val_loss = True
            args.random_transform = True
            args.phi = int(backbone[-1])

            if model_path is None:
                # step-1:
                args.lr = 1e-3
                args.freeze_backbone = True
                args.snapshot = 'imagenet'
                args.epochs = 50
                efficientdet_train(args, train_steps=self.train_steps, val_steps=self.val_steps)

                # step-2
                args.lr = 1e-4
                args.freeze_bn = True
                args.epochs = 100
                args.snapshot = '{dir}/efficientdet_{backbone}_{dataset_type}_{height}x{width}.h5'.format(
                    dir=args.snapshot_path,
                    backbone=args.backbone,
                    dataset_type=args.dataset_type,
                    height=args.height, width=args.width)
                efficientdet_train(args, train_steps=self.train_steps, val_steps=self.val_steps)

            else:
                args.lr = 1e-5
                args.freeze_bn = True
                args.epochs = 100
                args.snapshot = model_path
                efficientdet_train(args, train_steps=self.train_steps, val_steps=self.val_steps)

        elif self.base.lower() == 'segmentation':
            fine_tune = True if model_path else False
            self.app.train(model_name=method, backbone=backbone,
                           fine_tune=fine_tune, model_path=model_path,
                           opt='adam', lr=1e-3, shape=self.resized_shape)

        elif self.yolo:
            self.app.train_yolov3(gpu_id=1)

    def prediction(self,
                   gpu_id=0,
                   directly_train=True,
                   backbone='resnet50',
                   method='retinanet',
                   resized=True,
                   preprocess=True,
                   test_dir='data/personai_icartoonface_detval',
                   model_path='/data/shuai_li/FaceTask/saved_models/detection/resnet50_csv_16.h5',
                   show=False, write_prediction=True, gpu_fraction=0.2,test_on_val=False):
        """

        :param gpu_id:
        :param directly_train:
        :param backbone:
        :param method:
        :param resized:
        :param preprocess:
        :param test_dir:
        :param model_path:
        :param show:
        :param write_prediction:
        :param gpu_fraction:
        :param test_on_val:
        :return:
        """
        name = model_path.split('/')[-1].replace('.h5', '')
        full_name = [method,name]
        if directly_train:
            full_name.append('directly-train')
        if resized:
            full_name.append('resized-by-myself')
        if preprocess:
            full_name.append('preprocessed')
        if test_on_val:
            full_name.append('val-for-mAP-compute')

        setup_gpu(gpu_id)
        import tensorflow as tf
        import keras.backend.tensorflow_backend as KTF
        config = tf.ConfigProto()
        config.gpu_options.per_process_gpu_memory_fraction = gpu_fraction
        session = tf.Session(config=config)
        KTF.set_session(session)
        res_name = 'predictions/{}.csv'.format('_'.join(full_name))
        if directly_train and method.lower() == 'retinanet' and self.base.lower() == 'detection':
            if write_prediction:
                result = open('predictions/{}.csv'.format('_'.join(full_name)), 'w', encoding='utf-8')
            model = models.load_model(model_path, backbone_name=backbone)
            model = models.convert_model(model, nms=True, class_specific_filter=True)
            model.summary()
            label_to_name = {0: 'face'}
            files = os.listdir(test_dir)
            for img_name in tqdm.tqdm(files, total=len(files)):
                img_path = os.path.join(test_dir, img_name)
                image = cv2.imread(img_path)
                original_shape = image.shape
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                draw = image.copy()
                if resized:
                    image = cv2.resize(image, (720, 480), interpolation=cv2.INTER_AREA)
                if preprocess:
                    image = preprocess_image(image)
                scale = 1
                if not resized:
                    image, scale = resize_image(image)
                boxes, scores, labels = model.predict_on_batch(np.expand_dims(image, axis=0))
                if not resized:
                    boxes /= scale
                if resized:
                    h_rate = original_shape[0] / 480
                    w_rate = original_shape[1] / 720
                    boxes[:, 0] *= w_rate
                    boxes[:, 2] *= w_rate
                    boxes[:, 1] *= h_rate
                    boxes[:, 3] *= h_rate
                for box, score, label in zip(boxes[0], scores[0], labels[0]):
                    if score < 0.5:
                        break
                    color = label_color(label)
                    b = box.astype(int)
                    if write_prediction:
                        result.write('{},{},{},{},{},{},{}\n'.format(
                            img_name, b[0], b[1], b[2], b[3], 'face', score
                        ))
                    if show:
                        draw_box(draw, b, color=color)
                        caption = "{} {:.3f}".format(label_to_name[label], score)
                        draw_caption(draw, b, caption)
                if show:
                    plt.figure(figsize=(15, 15))
                    plt.axis('off')
                    plt.imshow(draw)
                    plt.show()
        elif self.base.lower() == 'segmentation':
            res_name = self.app.predict(
                model_path=model_path,
                test_img_dir=test_dir,
                shape=self.resized_shape,
            )
        elif directly_train and method.lower() == 'efficientdet' and self.base.lower() == 'detection':
            from libs.EfficientDet.model import efficientdet
            from libs.EfficientDet.utils import preprocess_image as pre_process, postprocess_boxes
            from libs.EfficientDet.utils.draw_boxes import draw_boxes
            if write_prediction:
                result = open('predictions/{}.csv'.format('_'.join(full_name)), 'w', encoding='utf-8')
            phi = int(backbone[-1])
            image_sizes = (512, 640, 768, 896, 1024, 1280, 1408)
            image_size = image_sizes[phi]
            classes = {0: 'face'}
            score_threshold = 0.2
            _, model = efficientdet(phi=phi,
                                    weighted_bifpn=False,
                                    num_classes=1,
                                    score_threshold=score_threshold)
            model.load_weights(model_path, by_name=True)
            for img_name in tqdm.tqdm(os.listdir(test_dir), total=len(os.listdir(test_dir))):
                img_path = os.path.join(test_dir, img_name)
                image = cv2.imread(img_path)
                image = image[:, :, ::-1]
                src_image = image.copy()
                h, w = image.shape[:2]
                if resized:
                    h_rate = h / self.resized_shape[0]
                    w_rate = w / self.resized_shape[1]
                    image = cv2.resize(image, (self.resized_shape[1], self.resized_shape[0]),
                                       interpolation=cv2.INTER_AREA)
                    image, scale = pre_process(image, image_size)
                    boxes, scores, labels = model.predict_on_batch([np.expand_dims(image, axis=0)])
                    boxes, scores, labels = np.squeeze(boxes), np.squeeze(scores), np.squeeze(labels)
                    boxes = postprocess_boxes(boxes=boxes, scale=scale, height=h, width=w)
                    boxes[:, 0] *= w_rate
                    boxes[:, 2] *= w_rate
                    boxes[:, 1] *= h_rate
                    boxes[:, 3] *= h_rate
                    boxes = boxes.astype(dtype='int8')
                    indices = np.where(scores[:] > score_threshold)[0]
                    boxes = boxes[indices]
                    labels = labels[indices]
                    if write_prediction:
                        for b, l, s in zip(boxes, labels, scores):
                            xmin, ymin, xmax, ymax = list(map(int, b))
                            result.write('{},{},{},{},{},{},{}\n'.format(
                                img_name, xmin, ymin, xmax, ymax, 'face', s
                            ))
                    if show:
                        draw_boxes(src_image, boxes, scores, labels, [[255, 0, 0]], classes)
                        plt.figure(figsize=(15, 15))
                        plt.axis('off')
                        plt.imshow(src_image)
                        plt.show()

                else:
                    image, scale = pre_process(image, image_size)
                    boxes, scores, labels = model.predict_on_batch([np.expand_dims(image, axis=0)])
                    boxes, scores, labels = np.squeeze(boxes), np.squeeze(scores), np.squeeze(labels)
                    boxes = postprocess_boxes(boxes=boxes, scale=scale, height=h, width=w)
                    indices = np.where(scores[:] > score_threshold)[0]
                    boxes = boxes[indices]
                    labels = labels[indices]
                    if write_prediction:
                        for b, l, s in zip(boxes, labels, scores):
                            xmin, ymin, xmax, ymax = list(map(int, b))
                            result.write('{},{},{},{},{},{},{}\n'.format(
                                img_name, xmin, ymin, xmax, ymax, 'face', s
                            ))
                    if show:
                        draw_boxes(src_image, boxes, scores, labels, [[255, 0, 0]], classes)
                        plt.figure(figsize=(15, 15))
                        plt.axis('off')
                        plt.imshow(src_image)
                        plt.show()

        return res_name

    def compute_mAP_on_val_data(self,
                                gpu_id=0, directly_train=True,
                                backbone='resnet50',method='retinanet',
                                resized=True,preprocess=True,
                                model_path='/data/shuai_li/FaceTask/saved_models/detection/resnet50_csv_16.h5',
                                show=False, write_prediction=True, gpu_fraction=0.2):
        pass

if __name__ == '__main__':
    # (240,360) (480,720) (512,768) (720,1080) multi-scale training
    app = Detection(img_dir='/data/shuai_li/FaceTask/data/personai_icartoonface_dettrain/icartoonface_dettrain',
                    label_csv_path='/data/shuai_li/FaceTask/data/personai_icartoonface_dettrain_anno_updatedv1.0.csv',
                    batch_size=10, resized_shape=(240, 360), base='detection')
    # app.train_model(gpu=3, directly_train=True, method='retinanet', backbone='resnet152',
    #                 model_path=None, augmentation=True, gpu_fraction=0.1)  # the last ms, use augmentation
    app.prediction(preprocess=True,resized=False,show=False,write_prediction=True,backbone='b0',method='efficientdet',gpu_id=0,
                   gpu_fraction=0.2,
                   model_path='/data/shuai_li/FaceTask/saved_models/detection/efficientdet_b0_csv_240x360_0011_0.62161.h5')


#!/usr/bin/python3

import sys
import cv2
import flask
import numpy as np
import tensorflow as tf
from flask import request, render_template, Response

app = flask.Flask(__name__)
app.config['DEBUG'] = True

model_dir = './models/'
age_labels = {0: 'Male', 1: 'Female'}

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict')
def predict():
    return render_template('predict.html')

@app.route('/video_feed')
def video_feed():
    video = cv2.VideoCapture(0)

    model_arch = 'vgg16'

    model = create_model(model_arch)
    loaded_model = load_model(model, model_dir+f'model_weights_{model_arch}.hdf5')

    return Response(gen_video(loaded_model, video), mimetype="multipart/x-mixed-replace; boundary=frame")

def gen_video(loaded_model, cam):
    clf = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')

    if not cam.isOpened():
        print("Cannot open camera exiting......")
        sys.exit()

    while True:
        success, frame = cam.read()

        if not success:
            print("Can't receive frame. Exiting ...")
            break

        frame = cv2.flip(frame, 1, 1)
        faces = clf.detectMultiScale(frame)

        for face in faces:
            (x, y, w, h) = [v for v in face]
            face_img = frame[y:y+h, x:x+w]
            face_img = cv2.resize(face_img, (48, 48)) / 255.

            result = loaded_model.predict(np.expand_dims(face_img, axis=0))
            cv2.rectangle(frame, (x, y), (x+w, y+h), (0,255,0))
            label = f'Age is {np.squeeze(result[0]).round()} and Gender is {age_labels[np.squeeze(result[1]).round()]}'
            cv2.putText(frame, str(label), (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

        _, jpeg = cv2.imencode('.jpg', frame)
        frame = jpeg.tobytes()

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

def create_model(model_arch):
    model = tf.keras.models.load_model(model_dir+f'model_{model_arch}.keras')

    return model

def load_model(model, fname):
    model.load_weights(fname)

    return model

if __name__ == '__main__':
    app.run()


from keras.applications.xception import Xception
from keras.applications.inception_v3 import InceptionV3
from keras.applications.vgg19 import VGG19
from keras.applications.vgg16 import VGG16
from keras.applications.resnet50 import ResNet50
from keras.preprocessing.image import ImageDataGenerator
from keras.models import Model, load_model
from keras.layers import Dense, GlobalAveragePooling2D, Input
from keras.callbacks import EarlyStopping, ModelCheckpoint
from keras.optimizers import SGD
import pandas as pd
import numpy as np

class FineTuneCNN:
    """Generic class that uses Xception, InceptionV3, VGG19, VGG16, or ResNet50 model pre-trained on imagenet weights
    to train, test, and predict image classifications.

    Attributes:
        model_name (str) : name of the model architecture
        dimensions (tup) : target dimensions to resize images to
        model_arch (keras.models.Model) : Xception, InceptionV3, VGG19, VGG16, or ResNet50 model
        num_classes (int) : number of classes for classification
        lr (float) : learning rate for the optimizer
        train_after_layer (int) : the layer after which you want to start training the model
        train_path (str) : path to the training data directory; must be nested by class
        num_training_samples (int) : number of training data samples
        validation_path (str) : path to the validation data directory; must be nested by class
        num_validation_samples (int) : number of validation data samples
        batch_size (int) : number of data samples to train on per batch
        num_epochs (int) : number epochs to train
        aug (bool) : True to augment the data, default is False
        trn_gen (keras.preprocessing.image.ImageDataGenerator) : train data image generator
        val_gen (keras.preprocessing.image.ImageDataGenerator) : validation data image generator
        test_gen (keras.preprocessing.image.ImageDataGenerator) : test data image generator
        early_stopping (bool) : True to stop training early based on the validation loss, default is True
        early_stop (keras.callbacks.Callback) : EarlyStopping callback from Keras
        model (keras.models.Model) : the instantiated model to be trained, tested, etc.
        test_path (str) : path to the testing data directory; must be nested by class
        num_test_samples (int) : number of testing data samples
        inference_path (str) : path to the inference data directory; must be nested within one folder
        num_inference_samples (int) : number of inference data samples
        checkpoint (keras.callbacks.Callback) : ModelCheckpoint callback from Keras
        model_path (str) : saved model path
    """

    def __init__(self, model_name=None, num_classes=None, lr=None):
        """Defines the target image size for the model based on the model architecture.
        """
        self.model_name = model_name
        assert model_name in ["VGG16","VGG19","ResNet50","Xception","InceptionV3"]
        if self.model_name in ["VGG16","VGG19","ResNet50"]:
            self.dimensions=(224,224)
        else:
            self.dimensions=(299,299)

    def __preprocess(self, x):
        """Preprocesses the image data based on the model architecture

        Returns:
            numpy.array : Normalized image data
        """
        if self.model_name in ["VGG16","VGG19","ResNet50"]:
            x[:, :, 0] -= 103.939
            x[:, :, 1] -= 116.779
            x[:, :, 2] -= 123.68
            # 'RGB'->'BGR'
            x = x[:, :, ::-1]
            return x
        else:
            x /= 255.
            x -= 0.5
            x *= 2.
            return x

    def build(self, model_arch=None, lr=None, num_classes=None, train_after_layer=-1):
        """Build the model

        Args:
            model_arch (keras.models.Model) : Xception, InceptionV3, VGG19, VGG16, or ResNet50 model
            lr (float) : learning rate for the optimizer
            num_classes (int) : number of classes for classification
            train_after_layer (int) : the layer after which you want to start training the model
        """
        self.model_arch = model_arch
        self.lr = lr
        self.num_classes = num_classes
        self.train_after_layer = train_after_layer

        # create the base pre-trained model
        base_model = self.model_arch
        x = base_model.output
        x = GlobalAveragePooling2D()(x)
        # and a logistic layer
        predictions = Dense(self.num_classes, activation='softmax')(x)

        # this is the model we will train
        self.model = Model(input=base_model.input, output=predictions)

        for layer in self.model.layers[:self.train_after_layer]:
           layer.trainable = False
        for layer in self.model.layers[self.train_after_layer:]:
           layer.trainable = True

        self.model.compile(optimizer=SGD(lr=self.lr, momentum=0.9), loss='categorical_crossentropy',metrics=['accuracy'])
        print 'Model built.'

    def train(self, train_path=None, num_training_samples=None, validation_path=None, num_validation_samples=None,
              batch_size=0, num_epochs=0, early_stopping=False, aug=False, checkpoint=None, model_path=None):
        """Fits the model

        Args:
            train_path (str) : path to the training data directory; must be nested by class
            num_training_samples (int) : number of training data samples
            validation_path (str) : path to the validation data directory; must be nested by class
            num_validation_samples (int) : number of validation data samples
            batch_size (int) : number of data samples to train on per batch
            num_epochs (int) : number epochs to train
            early_stopping (bool) : True to stop training early based on the validation loss, default is True
            aug (bool) : True to augment the data, default is False
            checkpoint (keras.callbacks.Callback) : ModelCheckpoint callback from Keras
        """
        self.train_path = train_path # training data needs to be within folders per class label
        self.num_training_samples = num_training_samples
        self.validation_path = validation_path
        self.num_validation_samples = num_validation_samples
        self.batch_size = batch_size
        self.num_epochs = num_epochs

        self.aug = aug
        if self.aug:
            train_datagen = ImageDataGenerator(rotation_range=10, width_shift_range=0.05, zoom_range=0.05,
                                               height_shift_range=0.05, shear_range=0.05,
                                               channel_shift_range=.1,horizontal_flip=True,
                                               preprocessing_function=self.__preprocess)
        else:
            train_datagen = ImageDataGenerator(preprocessing_function=self.__preprocess)

        self.trn_gen = train_datagen.flow_from_directory(self.train_path, target_size=self.dimensions,
                                                         batch_size=self.batch_size, class_mode='categorical', shuffle=True)
        if validation_path:
            val_datagen = ImageDataGenerator(preprocessing_function=self.__preprocess)

            self.val_gen = val_datagen.flow_from_directory(self.validation_path, target_size=self.dimensions,
                                                             batch_size=self.batch_size,
                                                             class_mode='categorical', shuffle=True)
        else:
            self.val_gen = None

        self.early_stopping = early_stopping
        if self.early_stopping:
            self.early_stop = EarlyStopping(monitor='val_loss', min_delta=0.10, patience=0, verbose=0, mode='auto')
            callback_list = [self.early_stop]
        else:
            callback_list = []

        self.checkpoint = checkpoint
        if self.checkpoint:
            callback_list.append(self.checkpoint)

        self.model.fit_generator(self.trn_gen, samples_per_epoch=self.num_training_samples,
                                 validation_data=self.val_gen, nb_val_samples=self.num_validation_samples,
                                 nb_epoch=self.num_epochs, verbose=1, callbacks=callback_list)

    def test(self, test_path=None, num_test_samples=None, batch_size=32):
        """Evaluates the model's performance on loss and accuracy

        Args:
            test_path (str) : path to the testing data directory; must be nested by class
            num_test_samples (int) : number of testing data samples
            batch_size (int) : number of data samples per batch
        Returns:
            float : loss
            float : accuracy
            numpy.array : ground truth class labels

        """
        self.test_path = test_path # images need to be nested in a folder in this directory
        self.num_test_samples = num_test_samples
        self.batch_size = batch_size

        test_datagen = ImageDataGenerator(preprocessing_function=self.__preprocess)
        self.test_gen = test_datagen.flow_from_directory(self.test_path, target_size=self.dimensions,
                                                         batch_size=self.batch_size,class_mode="categorical",
                                                         shuffle=False)

        loss,accuracy = self.model.evaluate_generator(self.test_gen, val_samples=self.num_test_samples)

        return loss, accuracy, self.test_gen.classes

    def load_weights(self, model_path=None):
        """Loads weights for a pre-trained model.

        Args:
            model_path (str) : saved model path
        """
        self.model = load_model(model_path)

    def inference(self, inference_path=None, num_inference_samples=None, batch_size=0, class_labels=None):
        """Predicts the probability of an image belonging to a certain class

        Args:
            inference_path (str) : path to the inference data directory; must be nested within one folder
            num_inference_samples (int) : number of inference data samples
            batch_size (int) : number of data samples per batch
            class_labels (list) : ordered list of class labels
        Returns:
            pandas.DataFrame : a dataframe of the images and their softmax probabilities
        """
        self.inference_path = inference_path # images need to be nested in a folder in this directory
        self.num_inference_samples = num_inference_samples
        self.batch_size = batch_size

        pred_datagen = ImageDataGenerator(preprocessing_function=self.__preprocess)
        pred_gen = pred_datagen.flow_from_directory(self.inference_path, target_size=self.dimensions, batch_size=self.batch_size,
                                                    class_mode=None, shuffle=False)

        predictions, filenames = self.model.predict_generator(pred_gen, val_samples=self.num_inference_samples), pred_gen.filenames

        prediction_list = []
        for i in range(len(filenames)):
            image = [filenames[i].split("/")[-1]]
            preds = predictions[i,:].tolist()
            prediction_list.append(image + preds)

        try:
            class_indices = self.trn_gen.class_indices
        except:
            class_indices = {value: index for index, value in enumerate(class_labels)}
        class_labels_ordered = [""]*len(class_indices)
        for label,index in class_indices.iteritems():
            class_labels_ordered[index] = label
        class_labels_ordered = ["image"] + class_labels_ordered

        inference_df = pd.DataFrame(prediction_list,columns=class_labels_ordered)

        return inference_df