#!/usr/bin/env python
from __future__ import division, print_function

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("trajlaps")

import signal
from picamera import PiCamera

import time
import numpy
import threading
try:
    from Queue import Queue
except:
    from queue import Queue
import datetime
import os
import sys
import math
from fractions import Fraction
import json
from collections import namedtuple, OrderedDict, deque
from argparse import ArgumentParser
import servo
from accelero import Accelerometer
from exposure import lens
from camera import Camera, Saver, Analyzer, Frame

Position = namedtuple("Position", ("pan", "tilt"))

trajectory={
"delay": 20,
"avg_wb":200,
"avg_ev":6,
"trajectory": [
 {"tilt":30,
  "pan":-50,
  "stay": 10,
  "move": 10},
 {"tilt":30,
  "pan":70,
  "stay": 10,
  "move": 21000},
 {"tilt":30,
  "pan":-50,
  "stay": 10,
  "move": 1000},
 ],
"histo_ev": [9,9,9,9],
"wb_red": [1.50390625,1.50390625,1.50390625,1.50390625],
"wb_blue": [1.56640625,1.56640625,1.56640625,1.56640625]
}

class DummyAccelero(object):
    """Dummy accelerometer"""
    def pause(self, wait=True):
        pass
    def resume(self):
        pass

class Trajectory(object):
    def __init__(self, delay_off=1, accelero=None, camera=None, config=None, json_file=None):
        self.start_time = time.time()
        self.accelero = accelero or DummyAccelero()
        self.camera = camera or DummyAccelero()
        self.delay_off = delay_off
        self.config = config or []
        self.servo_tilt = servo.tilt
        self.servo_pan = servo.pan
        self.default_pos = Position(0, 0)
        if json_file:
            self.load_config(json_file)
        self.goto_pos(self.default_pos)

    @property
    def duration(self):
        return sum( i.get("stay") + i.get("move") for i in self.config)

    def __repr__(self):
        return "Trajectory over %i points lasting %.3fs"%(len(self.config), self.duration)

    def load_config(self, filename):
        with open(filename) as jf:
            data = jf.read()
        config = json.loads(data)
        if "trajectory" in config:
           self.set_config(config["trajectory"])
        else:
           self.set_config(config)

    def get_config(self):
        return self.config

    def set_config(self, config):
        self.config = list(config)
        logger.info("Config set to :%s ", self.config)

    def goto(self, when):
        """move the camera to the position it need to be at a given timestamp"""
        pos = self.calc_pos(when)
        logger.info(pos)
        threading.Thread(target=self.goto_pos,args=(pos,)).start()
        return pos

    def goto_pos(self, pos):
        """Move the camera to the given position
        Operates usually in a separated thread
		"""
        pan, tilt = pos
        self.camera.pause(wait=False)
        self.accelero.pause()
        self.camera.pause()
        self.servo_tilt.move(tilt)
        self.servo_pan.move(pan)
        time.sleep(self.delay_off)
        #self.servo_tilt.off()
        #self.servo_pan.off()
        self.accelero.resume()
        self.camera.resume()
        res = OrderedDict([("pan", self.servo_pan.get_config()),
                           ("tilt", self.servo_tilt.get_config())])
        return res

    def calc_pos(self, when):
        """Calculate the position it need to be at a given timestamp"""
        next_pos = last_pos = self.default_pos
        last_timestamp = remaining = when
        if remaining <= 0:
            return last_pos
        #print(self.config)
        for point in self.config:
            next_pos = Position(point.get("pan",0), point.get("tilt",0))
            #print(last_pos, next_pos)
            remaining -= point.get("move")
            if remaining < 0:
                break
            remaining -= point.get("stay")
            if remaining < 0:
                return next_pos
            last_timestamp = remaining
            last_pos = next_pos
        else:
            #we ran out of points: stay on the last
            return next_pos
        delta = last_timestamp - remaining
        delta_p = next_pos.pan - last_pos.pan
        delta_t = next_pos.tilt - last_pos.tilt
        adv = last_timestamp / delta
        tilt = last_pos.tilt + adv * delta_t
        pan = last_pos.pan + adv * delta_p
        #print(last_timestamp, remaining, delta, adv)
        return Position(pan,tilt)

class TimeLaps(threading.Thread):
    def __init__(self, resolution=(3280, 2464), framerate=1, delay=20,
                 folder=".", avg_awb=200, avg_ev=25, config_file="parameters.json"):
        threading.Thread.__init__(self, name="TimeLaps")
        self.storage = {}
        self.storage_maxlen = 10
        self.pool_of_analyzers = []
        self.pool_size_analyzer = 2
        self.pool_of_savers = []
        self.pool_size_savers = 2
        self.do_analysis = False#True
        self.camera_queue = Queue()
        self.analysis_queue = Queue()
        self.saving_queue = Queue()
        self.config_queue = Queue()
        self.quit_event = threading.Event()
        self.delay = delay
        self.folder = os.path.abspath(folder)
        self.avg_wb = avg_awb # time laps over which average gains
        self.avg_ev = avg_ev # time laps over which average speed
        self.config_file = os.path.abspath(config_file)
        self.position = Position(0,0)
        self.start_time = self.last_img = time.time()
        self.next_img = self.last_img + 2 * self.delay
        signal.signal(signal.SIGINT, self.quit)
        self.accelero = Accelerometer()
        self.accelero.start()
        self.folder = folder
        self.servo_status = None

        self.camera = Camera(resolution=resolution,
                             framerate = framerate,
                             avg_ev=avg_ev,
                             avg_wb=avg_awb,
                             histo_ev=None,
                             wb_red=None,
                             wb_blue=None,
                             queue=self.camera_queue,
                             config_queue=self.config_queue,
                             quit_event=self.quit_event,
                             )
        self.camera.warm_up()
        self.trajectory = Trajectory(accelero=self.accelero, camera=self.camera)
        self.load_config(config_file)

        for i in range(self.pool_size_savers):
            saver = Saver(folder=self.folder,
                          queue=self.saving_queue,
                          quit_event=self.quit_event)
            saver.start()
            self.pool_of_savers.append(saver)

        for i in range(self.pool_size_analyzer):
            analyzer = Analyzer(frame_queue=self.analysis_queue,
                                 config_queue=self.config_queue,
                                 quit_event=self.quit_event)
            analyzer.start()
            self.pool_of_analyzers.append(analyzer)
        self.position = self.trajectory.goto(self.delay)
        self.camera.start()

    def __del__(self):
        self.quit_event.set()
        self.camera = None
        self.trajectory = None
        self.accelero = None

    def quit(self, *arg, **kwarg):
        """called with SIGINT: clean up and quit gracefully"""
        self.quit_event.set()

    def load_config(self, config_file=None):
        if config_file is None:
            config_file=self.config_file

        if os.path.isfile(config_file):
            with open(self.config_file) as jsonfile:
                dico = json.load(jsonfile)
            #print(dico)
            if "trajectory" in dico:
				self.trajectory.set_config(dico["trajectory"])
            if "camera" in dico:
                self.camera.set_config(dico["camera"])
            self.delay = dico.get("delay", self.delay)
            self.folder = dico.get("folder", self.folder)
            self.do_analysis = dico.get("do_analysis", self.do_analysis)
            self.camera.set_analysis(self.do_analysis)

    def save_config(self, index=None):
        if index is not None:
            fname = os.path.splitext(self.config_file)
            fname = fname[0] + "_%05i"%index + fname[1]
        else:
            fname = self.config_file
        logger.info("Save config to %s", fname)
        camera_config = self.camera.get_config()
        dico = OrderedDict([
                            ("delay", self.delay),
                            ("folder", self.folder),
                            ("do_analysis", self.do_analysis),
                            ("trajectory", self.trajectory.config),
                            ("camera", camera_config),
                            ])
        with open(fname,"w") as jsonfile:
            jsonfile.write(json.dumps(dico, indent=4))

    def run(self):
        "Actually does the timelaps"
        self.camera.set_analysis(self.do_analysis)
        while not self.quit_event.is_set():
            frame = self.camera_queue.get()
            frame.position = self.position
            frame.servo_status = self.servo_status
            if self.do_analysis:
                self.analysis_queue.put(frame)
            if self.position not in self.storage:
                self.storage[self.position] = deque(maxlen=self.storage_maxlen)
            self.storage[self.position].append(frame)
            if time.time() >= tl.next_img:
                frame.gravity = self.accelero.get()
                self.saving_queue.put(self.storage.pop(self.position))
                self.next_img = frame.timestamp + self.delay
                next_pos = self.trajectory.calc_pos(self.next_img - self.start_time)
                if next_pos != self.position:
                    self.servo_status = self.trajectory.goto_pos(next_pos)
                    self.position = next_pos
            self.camera_queue.task_done()
            logger.info("Frame #%04i. Length of queues: camera %i, analysis %i saving %i config %i",
                            frame.index,
                            self.camera_queue.qsize(),
                            self.analysis_queue.qsize(),
                            self.saving_queue.qsize(),
                            self.config_queue.qsize()
                            )
            if frame.index % 100 == 0:
                self.save_config(frame.index)

if __name__ == "__main__":
    try:
        from rfoo.utils import rconsole
    except:
        pass
    else:
        rconsole.spawn_server()

    parser = ArgumentParser("trajlaps",
                            description="TimeLaps over a trajectory")
    parser.add_argument("-j", "--json", help="config file")
    parser.add_argument("-d", "--debug", help="debug", default=False, action="store_true")

    args = parser.parse_args()
    if args.debug:
        logging.root.setLevel(logging.DEBUG)
    tl = TimeLaps(config_file=args.json, framerate=0.1)
    tl.run()


# load packages
import random
import yaml
from munch import Munch
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchaudio
import librosa
import os
import imageio

from Utils.ASR.models import ASRCNN
from Utils.JDC.model import JDCNet
from models import Generator, MappingNetwork, StyleEncoder, LmkEncoderPCA, MotDecoderPCATR, DiscriminatorMotPCA

# %matplotlib inline
emo_label = ['angry',  'contempt',  'disgusted',  'fear',  'happy',  'neutral',  'sad',  'surprised']
# speakers = [225,228,229,230,231,233,236,239,240,244,226,227,232,243,254,256,258,259,270,273]

to_mel = torchaudio.transforms.MelSpectrogram(
    n_mels=80, n_fft=2048, win_length=1200, hop_length=800)
mean, std = -4, 4

def preprocess(wave):
    wave_tensor = torch.from_numpy(wave).float()
#     print(wave_tensor.shape)
    mel_tensor = to_mel(wave_tensor)
#     print(mel_tensor.shape)
    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
    return mel_tensor

def build_model(model_params):
    args = Munch(model_params)
    generator = Generator(args.dim_in, args.style_dim, args.max_conv_dim, w_hpf=args.w_hpf, F0_channel=args.F0_channel)
    mapping_network = MappingNetwork(args.latent_dim, args.style_dim, args.num_domains, hidden_dim=args.max_conv_dim)
    style_encoder = StyleEncoder(args.dim_in, args.style_dim, args.num_domains, args.max_conv_dim)
    discriminator = DiscriminatorMotPCA(args.num_domains)

    lmk_encoder = LmkEncoderPCA()
    mot_decoder = MotDecoderPCATR()

    nets_ema = Munch(generator=generator,
                     mapping_network=mapping_network,
                     style_encoder=style_encoder)

    nets_lmk = Munch(lmk_encoder=lmk_encoder,
                     mot_decoder=mot_decoder)

    return nets_ema, nets_lmk

def compute_style(speaker_dicts):
    reference_embeddings = {}
    for key, (path, speaker) in speaker_dicts.items():
        if path == "":
            label = torch.LongTensor([speaker]).to('cuda')
            latent_dim = starganv2.mapping_network.shared[0].in_features
            ref = starganv2.mapping_network(torch.randn(1, latent_dim).to('cuda'), label)
        else:
            wave, sr = librosa.load(path, sr=24000)
            audio, index = librosa.effects.trim(wave, top_db=30)
            if sr != 24000:
                wave = librosa.resample(wave, sr, 24000)
            mel_tensor = preprocess(wave).to('cuda')
            extend_mel = torch.zeros((1, 80, 192)).float().to('cuda')
            mel_size = mel_tensor.size(2)
            extend_mel[0, :, :mel_size] = mel_tensor

            with torch.no_grad():
                label = torch.LongTensor([speaker])
                ref = starganv2.style_encoder(extend_mel.unsqueeze(1), label)
        reference_embeddings[key] = (ref, label)

    return reference_embeddings

#visualize mlmks
from functions import cv_draw_landmark_pca
def visual_landmarks(fl, writer):
    # fl = []
    # fl.append(initlm[:,:,:2])
    # # calculate fake landmarks
    # for i in range(fakemotion.shape[1]):
    #     fl.append(fl[i] + fakemotion[:,i,:,:])
    # fl = torch.cat(fl, 0).cpu().numpy()

    heatmap = 255*np.ones((256, 256, 3), dtype=np.uint8)
    fl[:, :, 0:2] = fl[:, :, 0:2] - [-0.5,-0.4]
    fl[:,:,0:2] = fl[:,:,0:2]*200
    fl = np.transpose(fl, (0, 2, 1))
    for l in fl:
        img_draw = cv_draw_landmark_pca(heatmap, l)
        writer.append_data(img_draw[:, :, ::-1])
    writer.close()
    return fl

def pca2lmk(pca, U, mean):
    b, l, d = pca.shape
    pca = pca.reshape(b*l, d)
    lmk = torch.mm(pca, U.t())
    lmk = lmk + mean.expand_as(lmk)
    lmk = lmk.reshape(b, l, -1).reshape(b, l, 468, 2)
    return lmk

# load F0 model

F0_model = JDCNet(num_class=1, seq_len=32)
params = torch.load("Utils/JDC/bst.t7", map_location='cpu')['net']
F0_model.load_state_dict(params)
# _ = F0_model.eval()
F0_model = F0_model.to('cuda')

# load vocoder
from parallel_wavegan.utils import load_model
vocoder = load_model("Vocoder/checkpoint-400000steps.pkl").to('cuda').eval()
vocoder.remove_weight_norm()
_ = vocoder.eval()

# load starganv2

# model_path = '/home/gy/gy/benchmark/A2L_StarG/Models/A2L/epoch_00750.pth'
# model_mot_path = '/home/gy/gy/benchmark/A2L_StarG2/Models/A2L_mot/epoch_00750.pth'
model_mot_path = '/home/gy/code/A2L_StarG2/Models/A2L_mot_M30_tran2/epoch_00800.pth'
wav_path = '/data3/shared/MEAD/M003/wav/neutral/level_1/001.wav'
# wav_path = '/data3/shared/MEAD/M031/wav/surprised/level_3/023.wav'

mean_mead = torch.from_numpy(np.load('./PCA/mean_mead.npy').astype(np.float32)).cuda()
U = torch.from_numpy(np.load('./PCA/U_mead.npy').astype(np.float32))[:,:32].cuda()

# model_mot_path = '/home/gy/gy/benchmark/A2L_StarG2/Models/A2L_mot_adv/epoch_01000.pth'

with open('Models/A2L_mot_M30_tran/config_lmk_M30_lstm2.yml') as f:
    starganv2_config = yaml.safe_load(f)
starganv2, models_mot = build_model(model_params=starganv2_config["model_params"])
# params = torch.load(model_path, map_location='cpu')
# params = params['model_ema']

params_mot = torch.load(model_mot_path, map_location='cpu')

params = params_mot['model']
params_mot = params_mot['model_mot']

_ = [starganv2[key].load_state_dict(params[key]) for key in starganv2]
# _ = [starganv2[key].eval() for key in starganv2]
starganv2.style_encoder = starganv2.style_encoder.to('cuda')
starganv2.mapping_network = starganv2.mapping_network.to('cuda')
starganv2.generator = starganv2.generator.to('cuda')

_ = [models_mot[key].load_state_dict(params_mot[key]) for key in models_mot]
models_mot.lmk_encoder = models_mot.lmk_encoder.to('cuda')
models_mot.mot_decoder = models_mot.mot_decoder.to('cuda')

source_emo = wav_path.split('/')[6]
bname = os.path.basename(wav_path)
dsp = wav_path.replace('wav/', 'video/front/')

example_landmark_path = '/data3/shared/other_images/init_leo.npy'
example_name = example_landmark_path.split('/')[-1][5:-4]
example_landmark = torch.from_numpy(np.load(example_landmark_path).astype(np.float32)).to('cuda')[:,:,:2]
print('example_landmark_shape: ', example_landmark.shape)

audio, source_sr = librosa.load(wav_path, sr=24000)
audio.dtype = np.float32

speaker_dicts = {}

speaker_dicts[f'happy'] = ('', emo_label.index('happy'))
speaker_dicts[f'sad'] = ('', emo_label.index('sad'))
speaker_dicts[f'disgusted'] = ('', emo_label.index('disgusted'))
speaker_dicts[f'contempt'] = ('', emo_label.index('contempt'))
speaker_dicts[f'surprised'] = ('', emo_label.index('surprised'))
speaker_dicts[f'angry'] = ('', emo_label.index('angry'))
speaker_dicts[f'fear'] = ('', emo_label.index('fear'))

reference_embeddings = compute_style(speaker_dicts)

# conversion
import time
start = time.time()

source = preprocess(audio).to('cuda')
extend_source = torch.zeros((1, 80, 96)).float().to('cuda')
mel_size = min(96, source.shape[2])
print('mel_size: ', mel_size)
extend_source[0, :, :mel_size] = source[:,:,:mel_size]
# extend_source = source[:,:,:mel_size]
print('source: ', extend_source.shape, source.shape)
length_mots = [mel_size]
print('length_mots: ', length_mots)
keys = []
converted_samples = {}
reconstructed_samples = {}
converted_mels = {}

for key, (ref, _) in reference_embeddings.items():
    with torch.no_grad():
        f0_feat = F0_model.get_feature_GAN(extend_source.unsqueeze(1))
        print('f0_feat: ', f0_feat.shape)
        out, audio_feature = starganv2.generator(extend_source.unsqueeze(1), ref, masks=None, F0=f0_feat)
        # assert(0)
        lm_feature = models_mot.lmk_encoder(example_landmark)
        print(lm_feature.shape, audio_feature.shape, length_mots)
        fake_pca = models_mot.mot_decoder(lm_feature, audio_feature, length_mots)
        print('fake_pca: ', fake_pca.shape)
        b, length, _ = fake_pca.shape
        fake_pca = fake_pca

        lmk = pca2lmk(fake_pca, U, mean_mead)

        print('fake_pca: ', fake_pca.shape)
        print('lmk: ', lmk.shape)

        # print(lmk.shape)
        writer = imageio.get_writer(f'/data3/shared/other_images/lmk_test/{example_name}_{bname[:-4]}_to_{key}_fake.mp4', fps=30)
        fake_fl = visual_landmarks(lmk[0,:mel_size,:].cpu().numpy(), writer)

        np.save(f'/data3/shared/other_images/lmk_test/{example_name}_{bname[:-4]}_to_{key}_fake', lmk[0,:mel_size,:].cpu().numpy())

        # assert(0)

        # c = out.transpose(-1, -2).squeeze().to('cuda')
        # y_out = vocoder.inference(c)
        # y_out = y_out.view(-1).cpu()
        # print('y_out shape:', y_out.shape)
        # if key not in speaker_dicts or speaker_dicts[key][0] == "":
        #     recon = None
        # else:
        #     wave, sr = librosa.load(speaker_dicts[key][0], sr=24000)
        #     mel = preprocess(wave)
        #     c = mel.transpose(-1, -2).squeeze().to('cuda')
        #     recon = vocoder.inference(c)
        #     recon = recon.view(-1).cpu().numpy()

    # converted_samples[key] = y_out.numpy()
    # reconstructed_samples[key] = recon

#     converted_mels[key] = out

#     keys.append(key)
# end = time.time()
# print('total processing time: %.3f sec' % (end - start) )

# import IPython.display as ipd
# for key, wave in converted_samples.items():
#     with open(f'./audio_result/converted_001_{key}.wav', 'wb') as f:
#         print('Converted: %s' % key)
#         f.write(ipd.Audio(wave, rate=24000).data)
#         print('Reference (vocoder): %s' % key)
#         if reconstructed_samples[key] is not None:
#             with open(f'recon_{key}.wav', 'wb') as f:
#                 f.write(ipd.Audio(reconstructed_samples[key], rate=24000).data)

# print('Original (vocoder):')
# wave, sr = librosa.load(wav_path, sr=24000)
# mel = preprocess(wave)
# c = mel.transpose(-1, -2).squeeze().to('cuda')
# with torch.no_grad():
#     recon = vocoder.inference(c)
#     recon = recon.view(-1).cpu().numpy()
# with open(f'./audio_result/source_recon.wav', 'wb') as f:
#     f.write(ipd.Audio(recon, rate=24000).data)
# print('Original:')
# with open(f'./audio_result/source.wav', 'wb') as f:
#     f.write(ipd.Audio(wav_path, rate=24000).data)


from math import ceil
import asyncio
import os
from PIL import Image
from numpy import array
from ctypes import *
from . import Arduino
from . import MAILBOX_OFFSET
from . import MAILBOX_PY2IOP_CMD_OFFSET

ARDUINO_CAM_PROGRAM = "arduino_cam.bin"
TEST = 0x1
TEST_2 = 0x2

class Arduino_CAM(object):
    '''
        Testing mailbox/circular buffer
    '''
    def __init__(self, mb_info):
        """Return a new instance of an Arduino_mailboxTest object.

        Parameters
        ----------
        mb_info : dict
            A dictionary storing Microblaze information, such as the
            IP name and the reset name.

        """
        self.microblaze = Arduino(mb_info, ARDUINO_CAM_PROGRAM)

    def test(self):
        self.microblaze.write_blocking_command(TEST)
        return self.microblaze.read_mailbox(0,1)

    def test_2(self):
        self.microblaze.write_blocking_command(TEST_2)
        return self.microblaze.read_mailbox(0,8)


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

''' Class to keep track of the seeds,
    i.e. all the anions, electrons and avalanches.

    This class is used to modify and get the properties of the seeds
    as well as classify, move and multiply them.
'''

# General imports
import numpy as np
import logging

# Import from project files
from .streamer_head import SHList

# settings
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())

# idea: refactor to a general class for particles
#       create a new class to manage the particles

class Seeds(object):
    def __init__(self, mu_e, mu_ion, efield_ion, efield_crit, time_step,
                 Q_crit, efield_dtype, micro_step_no
                 ):
        # initiate the seeds object without adding any seeds

        # define the logger
        self.logger = logging.getLogger(__name__ + '.Seeds')

        # set the dtype for electric field calculation
        if efield_dtype in ['sp', 'float32',]:
            self.efield_dtype = np.float32
        elif efield_dtype in ['dp', 'float64',]:
            self.efield_dtype = np.float64
        else:
            self.efield_dtype = np.float32
            msg = 'Can not use dtype "{}". Using "{}"" instead.'
            self.logger.warning(msg.format(efield_dtype, self.efield_dtype))

        # set the properties of the seeds
        self.mu_e         = mu_e        # electron mobility
        self.mu_ion       = mu_ion      # cation mobility
        self.efield_ion   = efield_ion  # electric field for ion detachment
        self.efield_crit  = efield_crit  # electric field for multiplication
        self.time_step    = time_step   # simulation max time step
        self.dt           = time_step   # simulation actual time step
        self.micro_step_no = micro_step_no   # micro time steps
        self.Q_crit       = Q_crit      # avalanche to streamer criterion

        # initialize the seed variables, start with no seeds
        self.pos = np.array([]).reshape(3, 0)   # [x, y, z] position
        self.Q = np.array([])                   # avalanche size
        self.clear()                            # initialize other variables

        # log initiation
        self.logger.debug('Initated Seeds')
        self.logger.log(5, 'Seeds.__dict__')
        for k, v in self.__dict__.items():
            self.logger.log(5, '  "{}": {}'.format(k, v))

    def clear(self):
        # Clear variables

        # rename for better readability below
        no = self.no
        nz = np.zeros  # return zeroes
        nf = lambda shape: np.zeros(shape, dtype=bool)  # return false shape

        # Should be explicitly reset
        self.no_added       = 0
        self.no_removed     = 0
        self.pos_to_append  = nz((3, 0))
        self.is_to_remove   = nz(no, dtype=bool)

        # Should be updated each iteration
        self.ds             = nz((3, no))   # last movement
        self.ds_abs         = nz(no)        # movement, absolute
        self.ds_max         = 0             # movement, maximum
        self.dQ             = nz(no)        # change in charge
        self.mu             = nz(no)        # mobility
        self.e_dir          = nz((3, no))   # electric field, unit vector
        self.e_vec          = nz((3, no))   # electric field, vector
        self.e_str          = nz(no)        # electric field, strength
        # bools
        self.is_in_streamer = nf(no)        # inside streamer
        self.is_critical    = nf(no)        # avalanche termination
        self.is_behind_roi  = nf(no)        # behind ROI
        self.is_avalanche   = nf(no)        # avalanche seed
        self.is_ion         = nf(no)        # electron seed
        self.is_electron    = nf(no)        # ion seed

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       DERIVED ATTIBUTES                                         #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # The user need to ensure that relevant variables are updated
    # before getting these properties

    Q_max = property(lambda self: self.Q.max())
    dQ_max = property(lambda self: self.dQ.max())
    charge = property(lambda self: np.exp(self.Q).sum())
    charge_gen = property(lambda self: (self.dQ * np.exp(self.Q)).sum())
    charge_rem = property(lambda self: np.exp(self.Q_to_remove).sum())

    no = property(lambda self: self.pos.shape[1])  # total number
    no_to_append = property(lambda self: self.pos_to_append.shape[1])

    no_ion = property(lambda self: self.is_ion.sum())
    no_electron = property(lambda self: self.is_electron.sum())
    no_avalanche = property(lambda self: self.is_avalanche.sum())
    no_critical = property(lambda self: self.is_critical.sum())
    no_behind_roi = property(lambda self: self.is_behind_roi.sum())
    no_in_streamer = property(lambda self: self.is_in_streamer.sum())
    no_to_remove = property(lambda self: self.is_to_remove.sum())

    pos_ion = property(lambda self: self.pos[:, self.is_ion])
    pos_electron = property(lambda self: self.pos[:, self.is_electron])
    pos_avalanche = property(lambda self: self.pos[:, self.is_avalanche])
    pos_critical = property(lambda self: self.pos[:, self.is_critical])
    pos_behind_roi = property(lambda self: self.pos[:, self.is_behind_roi])
    pos_in_streamer = property(lambda self: self.pos[:, self.is_in_streamer])
    pos_to_remove = property(lambda self: self.pos[:, self.is_to_remove])

    Q_ion = property(lambda self: self.Q[self.is_ion])
    Q_electron = property(lambda self: self.Q[self.is_electron])
    Q_avalanche = property(lambda self: self.Q[self.is_avalanche])
    Q_critical = property(lambda self: self.Q[self.is_critical])
    Q_behind_roi = property(lambda self: self.Q[self.is_behind_roi])
    Q_in_streamer = property(lambda self: self.Q[self.is_in_streamer])
    Q_to_remove = property(lambda self: self.Q[self.is_to_remove])

    dQ_ion = property(lambda self: self.dQ[self.is_ion])
    dQ_electron = property(lambda self: self.dQ[self.is_electron])
    dQ_avalanche = property(lambda self: self.dQ[self.is_avalanche])
    dQ_critical = property(lambda self: self.dQ[self.is_critical])
    dQ_behind_roi = property(lambda self: self.dQ[self.is_behind_roi])
    dQ_in_streamer = property(lambda self: self.dQ[self.is_in_streamer])
    dQ_to_remove = property(lambda self: self.dQ[self.is_to_remove])

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       UPDATE                                                    #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # Updating of bools should be called explicitly
    # since they may be expensive to calculate
    # or it has to be done at specific times
    # Maintain a comment to when they are updated!

    # `is_ion` is set during movement
    def update_is_ion(self, e_str=None):
        e_str = self.e_str if (e_str is None) else e_str
        self.is_ion = (e_str < self.efield_ion)
        return self.is_ion

    # ´is_electron´ is set during movement
    def update_is_electron(self, e_str=None):
        e_str = self.e_str if (e_str is None) else e_str
        self.is_electron = (e_str >= self.efield_ion)
        return self.is_electron

    # ´is_avalanche´ is set during movement
    def update_is_avalanche(self, e_str=None):
        e_str = self.e_str if (e_str is None) else e_str
        self.is_avalanche = (e_str >= self.efield_crit)
        return self.is_avalanche

    # Invoked in main simulation loop
    def update_is_critical(self, Q=None):
        Q = self.Q if (Q is None) else Q
        self.is_critical = (Q >= self.Q_crit)
        return self.is_critical

    # Invoked in main simulation loop
    def update_is_behind_roi(self, roi):
        self.is_behind_roi = roi.is_pos_behind(self.pos)
        return self.is_behind_roi

    # ´is_in_streamer´ is set during movement
    def update_is_in_streamer(self, heads):
        self.is_in_streamer = SHList(heads).is_inside(self.pos)
        return self.is_in_streamer

    def update_mu(self):
        self.mu = np.ones(self.no) * self.mu_ion
        self.mu[self.is_electron] = self.mu_e

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       OTHER                                                     #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def eprop(self, heads, pos=None):
        ''' Return `edir`, `estr`, `evec`, and `is_inside`.
            Convenience function to ensure correct dtype for seeds.
        '''
        if pos is None:
            pos = self.pos
        return SHList(heads).eprop(pos, dtype=self.efield_dtype)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       MOVE & MULTIPLY                                           #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def _update_ds(self, e_str, e_vec):
        ''' Calculate movement. Update `ds`.
            Also update `is_electron`, `is_ion`, `mu`.
        '''
        self.is_electron = (e_str > self.efield_ion)
        self.is_ion = ~self.is_electron  # note: need to be somewhere
        self.mu = self.mu_ion * np.ones(self.no)
        self.mu[self.is_electron] = self.mu_e
        self.ds = e_vec * self.mu * self.dt

    def _update_ds_abs(self, thresh):
        ''' Calculate the max movement. Update `ds_abs` and `ds_max`.
            Give warning if any movement exceeds `thresh`.
        '''
        self.ds_abs = np.linalg.norm(self.ds, axis=0)
        self.ds_max = self.ds_abs.max()
        if self.ds_max > thresh:
            mask = (self.ds_abs > thresh)
            msg = 'Warning! {}x large movement!'.format(mask.sum())
            self.logger.warning(msg)
            self.logger.debug('Movements {}'.format(self.ds_abs[mask]))

    # note: previously used function
    def move_multiply_single_step(self, heads, calc_alpha):
        ''' Move all seeds and multiply all avalanches, single step. '''

        # treat seeds below the plane as if they were at the plane
        pos = self.pos.copy()
        mask = (self.pos[2, :] < 0)
        pos[2, mask] = 0

        # Find the field at electrons!
        self.e_dir, self.e_str, self.e_vec, self.is_in_streamer = (
            self.eprop(heads, pos))

        # Move seeds
        self._update_ds(self.e_str, self.e_vec)
        self._update_ds_abs(thresh=heads[0].rp)
        self.pos += self.ds

        # Update electron numbers
        self.dQ = calc_alpha(self.e_str) * self.ds_abs
        self.dQ[self.e_str < self.efield_crit] = 0
        self.dQ[self.pos[2, :] <= 0] = 0
        self.Q += self.dQ

    def move_multiply(self, heads, calc_alpha):
        ''' Move all seeds, multiply electron in avalanches.

        Parameters
        ----------
        heads :     SHList
                    list of streamer heads used for field calculation
        calc_alpha : CalcAlpha
                     class instance for calculating electron multiplication

        Notes
        -----
        Seeds are classified by the electric field strength.
        Avalanches are moved in an inner loop.
        Critical avalanches and collisions with the streamer
        are saved and these break the inner loop.
        Other seeds are moved according to the time spent in the inner loop.
        Long seed movements give warnings.
        Positions and charge numbers are updated.
        '''

        # treat seeds below the plane as if they were at the plane
        # this implies that seeds below plane will move straight upwards
        pos = self.pos.copy()
        mask = (self.pos[2, :] < 0)
        pos[2, mask] = 0

        # find the field
        self.e_dir, self.e_str, self.e_vec, self.is_in_streamer = (
            self.eprop(heads, pos))

        # classify the seeds
        self.update_is_electron()   # uses self.e_str
        self.update_is_ion()        # uses self.e_str
        self.update_is_avalanche()  # uses self.e_str

        # set all mobilities
        self.update_mu()            # uses self.is_electron

        # extract information about just the avalanches
        is_avalanche = self.is_avalanche
        pos_avalanche = self.pos[:, is_avalanche]
        Q_avalanche = self.Q[is_avalanche]
        e_str = self.e_str[is_avalanche]
        e_vec = self.e_vec[:, is_avalanche]
        e_dir = self.e_dir[:, is_avalanche]
        is_in_streamer = self.is_in_streamer

        time_spent = 0

        def _get_ds(e_vec, mu=None, time_step=None):
            # return the change in avalanche positions
            if mu is None:
                mu = self.mu
            if time_step is None:
                time_step = self.time_step
            ds = e_vec * mu * time_step
            ds_abs = np.linalg.norm(ds, axis=0)
            return ds, ds_abs

        def _warn_large_ds(ds_abs, threshold, what):
            ds_max = ds_abs.max()
            if (ds_max < threshold):
                return ds_max
            mask = (ds_abs > threshold)
            msg = f'Warning! {mask.sum()}x large {what} movement!'
            self.logger.warning(msg)
            self.logger.info(f'Max movement: {(ds_max * 1e6):.2g} um')
            self.logger.log(5, f'Movements {ds_abs[mask]}')
            return ds_max

        # skip this part when it is not needed
        if any(is_avalanche):

            # do one iteration with the field already calculated
            time_spent += self.time_step        # update the time
            ds, ds_abs = _get_ds(e_vec, mu=self.mu_e) # find the movements
            _warn_large_ds(ds_abs, threshold=heads[0].rp/10, what='avalanches')
            pos_avalanche += ds                 # update positions
            dQ = calc_alpha(e_str) * ds_abs     # get avalanche change
            dQ[pos_avalanche[2, :] < 0] = 0     # no multiplication for z<0
            Q_avalanche += dQ                   # update avalanche electrons

            # do the rest of the iterations
            for i in range(self.micro_step_no - 1):  # one is already done
                # end when critical avalanche is reached
                if any(Q_avalanche > self.Q_crit):
                    logger.debug('Critical avalanche!')
                    break
                # end on collision with streamer
                if any(is_in_streamer):
                    if i != 0:  # not needed (or allowed) the first loop
                        self.is_in_streamer[is_avalanche] = is_in_streamer
                    logger.debug('Collision with streamer!')
                    break

                # treat seeds below the plane as if they were at the plane
                pos_avalanche_tmp = pos_avalanche.copy()
                mask = (pos_avalanche[2, :] < 0)
                pos_avalanche_tmp[2, mask] = 0

                # find the field
                e_dir, e_str, e_vec, is_in_streamer = (
                    self.eprop(heads, pos_avalanche_tmp))

                time_spent += self.time_step        # update the time
                ds, ds_abs = _get_ds(e_vec, mu=self.mu_e) # find the movements
                _warn_large_ds(
                    ds_abs, threshold=heads[0].rp/10, what='avalanches')
                pos_avalanche += ds                 # update positions
                dQ = calc_alpha(e_str) * ds_abs     # get avalanche change
                dQ[pos_avalanche[2, :] < 0] = 0     # no multiplication for z<0
                Q_avalanche += dQ                   # update avalanche electrons

        else:  # no avalanches, should be a rare event
            time_spent = self.time_step * self.micro_step_no
            logger.debug('No avalanches!')

        # update avalanche numbers
        self.is_avalanche = is_avalanche  # should already be correct
        self.dQ[is_avalanche] = Q_avalanche - self.Q[is_avalanche]
        self.Q[is_avalanche] = Q_avalanche

        # get seed movements
        self.e_vec[:, is_avalanche] = 0  # do not move avalanches here
        ds, ds_abs = _get_ds(self.e_vec, time_step=time_spent) # find movements
        _warn_large_ds(ds_abs, threshold=heads[0].rp, what='seed')

        # update positions
        pos_old = self.pos.copy()   # save old position for ds calculation
        self.pos += ds              # add ds for all seeds
        self.pos[:, is_avalanche] = pos_avalanche  # set these explicitly

        # store actual movements
        self.ds = self.pos - pos_old
        self.ds_abs = np.linalg.norm(self.ds, axis=0)
        self.ds_max = self.ds_abs.max()

        # store the time spent here
        self.dt = time_spent

        # log results, as required
        logger.log(5, f'Number of avalanches, {is_avalanche.sum()}')
        logger.log(5, f'Time spent, {(time_spent * 1e9):.3f} ns')
        logger.log(5, f'Maximum movement, {(self.ds_max * 1e6):.3f} um')
        if any(is_avalanche):
            msg = f'Maximum avalanche growth, {(self.dQ.max()):.3f}'
        else:
            msg = f'Maximum avalanche growth, 0'
        logger.log(5, msg)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       APPEND/REMOVE                                             #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def append(self, pos):
        ''' Append the given pos, and update vars as required.'''
        if not pos.size:
            self.logger.log(5, 'Tried to add zero seeds.')
            return
        assert pos.shape[0] == 3, 'pos.shape[0] == 3'
        Q = np.zeros(pos.shape[1])
        self.pos = np.hstack((self.pos, pos))
        self.Q = np.hstack((self.Q, Q))
        self.no_added += pos.shape[1]

    def remove(self, idx):
        ''' Remove pos and Q at the given indices.'''
        if idx.sum() == 0:
            self.logger.log(5, 'Tried to remove zero seeds.')
            return
        self.pos = self.pos[:, idx == 0]
        self.Q = self.Q[idx == 0]
        self.no_removed += idx.sum()

    def append_at_end(self, pos):
        ''' Appends the given pos to pos_to_append.'''
        assert pos.shape[0] == 3, 'pos.shape[0] == 3'
        pos = np.array(pos).reshape(3, -1)
        self.pos_to_append = np.hstack((self.pos_to_append, pos))

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #       CLEAN                                                     #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    def clean(self):

        # Remove seeds and reset var
        self.remove(self.is_to_remove)

        # Append seeds and reset var
        self.append(self.pos_to_append)

        # Clear other variables
        self.clear()

#


import time
from absl import app, flags, logging
from absl.flags import FLAGS

import cv2
import numpy as np
import tensorflow as tf

import core.utils as utils
from core.config import cfg
from core.yolov4 import YOLOv4, YOLOv3, YOLOv3_tiny, YOLOv4_tiny, decode

from tensorflow.python.saved_model import tag_constants

flags.DEFINE_string('framework', 'tf', '(tf, tflite, trt')
flags.DEFINE_string('weights', './weights/yolov4/yolov4_model_3.weights',
                    'path to weights file')
flags.DEFINE_integer('size', 512, 'resize images to')
flags.DEFINE_boolean('tiny', False, 'yolo or yolo-tiny')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')
flags.DEFINE_string('video', './data/test/video1.mp4', 'path to input video')
flags.DEFINE_string('output', None, 'path to output video')
flags.DEFINE_string('output_format', 'XVID',
                    'codec used in VideoWriter when saving video to file')

def main(_argv):
    physical_devices = tf.config.experimental.list_physical_devices('GPU')
    for physical_device in physical_devices:
        tf.config.experimental.set_memory_growth(physical_device, True)

    if FLAGS.tiny:
        STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
        XYSCALE = cfg.YOLO.XYSCALE_TINY
        if FLAGS.model == 'yolov4':
            ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
        else:
            ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY_V3, FLAGS.tiny)
    else:
        STRIDES = np.array(cfg.YOLO.STRIDES)
        XYSCALE = cfg.YOLO.XYSCALE
        if FLAGS.model == 'yolov4':
            ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
        else:
            ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)

    CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)
    NUM_CLASSES = len(CLASSES)
    input_size = FLAGS.size
    try:
        vid = cv2.VideoCapture(int(FLAGS.video))
    except:
        vid = cv2.VideoCapture(FLAGS.video)

    times = []
    if FLAGS.output:
        width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = int(vid.get(cv2.CAP_PROP_FPS))
        codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
        out = cv2.VideoWriter(FLAGS.output, codec, fps,
                              (width, height))  # TODO: switch to get vertical

    if FLAGS.framework == 'tf':
        input_layer = tf.keras.layers.Input([input_size, input_size, 3])
        if FLAGS.tiny:
            if FLAGS.model == 'yolov3':
                feature_maps = YOLOv3_tiny(input_layer, NUM_CLASSES)
            else:
                feature_maps = YOLOv4_tiny(input_layer, NUM_CLASSES)
            bbox_tensors = []
            for i, fm in enumerate(feature_maps):
                bbox_tensor = decode(fm, NUM_CLASSES, i)
                bbox_tensors.append(bbox_tensor)
            model = tf.keras.Model(input_layer, bbox_tensors)
            utils.load_weights_tiny(model, FLAGS.weights, FLAGS.model)
        else:
            if FLAGS.model == 'yolov3':
                feature_maps = YOLOv3(input_layer, NUM_CLASSES)
                bbox_tensors = []
                for i, fm in enumerate(feature_maps):
                    bbox_tensor = decode(fm, NUM_CLASSES, i)
                    bbox_tensors.append(bbox_tensor)
                model = tf.keras.Model(input_layer, bbox_tensors)
                utils.load_weights_v3(model, FLAGS.weights)
            elif FLAGS.model == 'yolov4':
                feature_maps = YOLOv4(input_layer, NUM_CLASSES)
                bbox_tensors = []
                for i, fm in enumerate(feature_maps):
                    bbox_tensor = decode(fm, NUM_CLASSES, i)
                    bbox_tensors.append(bbox_tensor)
                model = tf.keras.Model(input_layer, bbox_tensors)
                if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
                    utils.load_weights(model, FLAGS.weights)
                else:
                    model.load_weights(FLAGS.weights).expect_partial()
        model.summary()
    elif FLAGS.framework == 'tflite':
        interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
        interpreter.allocate_tensors()
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()
    elif FLAGS.framework == 'trt':
        saved_model_loaded = tf.saved_model.load(
            FLAGS.weights, tags=[tag_constants.SERVING])
        infer = saved_model_loaded.signatures['serving_default']

    logging.info("Model loaded!")
    while True:
        return_value, frame = vid.read()
        # frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)  # TODO: here
        if not return_value:
            logging.warning("Empty Frame")
            break

        frame_size = frame.shape[:2]
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

        image_data = utils.image_preprocess(
            np.copy(frame), [input_size, input_size])
        image_data = image_data[np.newaxis, ...].astype(np.float32)

        prev_time = time.time()
        if FLAGS.framework == 'tf':
            pred_bbox = model.predict(image_data)
        elif FLAGS.framework == 'tflite':
            interpreter.set_tensor(input_details[0]['index'], image_data)
            interpreter.invoke()
            pred_bbox = [interpreter.get_tensor(
                output_details[i]['index']) for i in range(len(output_details))]
        elif FLAGS.framework == 'trt':
            batched_input = tf.constant(image_data)
            pred_bbox = []
            result = infer(batched_input)
            for _, value in result.items():
                value = value.numpy()
                pred_bbox.append(value)
        curr_time = time.time()
        times.append(curr_time - prev_time)
        times = times[-20:]

        if FLAGS.model == 'yolov4':
            pred_bbox = utils.postprocess_bbbox(
                pred_bbox, ANCHORS, STRIDES, XYSCALE)
        else:
            pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
        bboxes = utils.postprocess_boxes(
            pred_bbox, frame_size, input_size, 0.5)  # 0.25
        bboxes = utils.nms(bboxes, 0.213, method='nms')  # 0.213

        image = utils.draw_bbox(frame, bboxes, classes=CLASSES)
        image = cv2.putText(image, "Time: {:.2f}ms".format(sum(times)/len(times)*1000), (0, 24),  # 24
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)  # 0.7

        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        cv2.namedWindow("Detections", cv2.WINDOW_AUTOSIZE)
        cv2.imshow("Detections", image)

        if FLAGS.output:
            out.write(image)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    vid.release()
    if FLAGS.output:
        out.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import numpy as np
import tqdm
import json
from AFSD.common import videotransforms
from AFSD.common.thumos_dataset import get_video_info, get_class_index_map
from AFSD.thumos14.BDNet import BDNet, DirichletLayer
from AFSD.common.segment_utils import softnms_v2
from AFSD.common.config import config

def get_path(input_path):
    if os.path.lexists(input_path):
        fullpath = os.path.realpath(input_path) if os.path.islink(input_path) else input_path
        real_name = fullpath.split('/')[-1]
        real_full_path = os.path.join(os.path.dirname(input_path), real_name)
    else:
        raise FileNotFoundError
    return real_full_path

def build_model(fusion=False):
    net, flow_net = None, None
    if fusion:
        rgb_net = BDNet(in_channels=3, training=False, use_edl=cfg.use_edl, use_rpl=cfg.use_rpl)
        flow_net = BDNet(in_channels=2, training=False, use_edl=cfg.use_edl, use_rpl=cfg.use_rpl)
        rgb_checkpoint_path = get_path(config['testing'].get('rgb_checkpoint_path',
                                            './models/thumos14/checkpoint-15.ckpt'))
        flow_checkpoint_path = get_path(config['testing'].get('flow_checkpoint_path',
                                             './models/thumos14_flow/checkpoint-16.ckpt'))
        rgb_net.load_state_dict(torch.load(rgb_checkpoint_path))
        flow_net.load_state_dict(torch.load(flow_checkpoint_path))
        rgb_net.eval().cuda()
        flow_net.eval().cuda()
        net = rgb_net
    else:
        net = BDNet(in_channels=config['model']['in_channels'],
                    training=False, use_edl=cfg.use_edl, use_rpl=cfg.use_rpl)
        checkpoint_path = get_path(config['testing']['checkpoint_path'])
        net.load_state_dict(torch.load(checkpoint_path))
        net.eval().cuda()
    return net, flow_net

def get_offsets(video_infos, video_name, clip_length, stride):
    sample_count = video_infos[video_name]['sample_count']
    if sample_count < clip_length:
        offsetlist = [0]
    else:
        offsetlist = list(range(0, sample_count - clip_length + 1, stride))
        if (sample_count - clip_length) % stride:
            offsetlist += [sample_count - clip_length]
    return offsetlist

def prepare_data(data_path, video_name, centor_crop):
    data = np.load(os.path.join(data_path, video_name + '.npy'))
    data = np.transpose(data, [3, 0, 1, 2])
    data = centor_crop(data)
    data = torch.from_numpy(data).cuda()
    return data

def prepare_clip(data, offset, clip_length):
    clip = data[:, offset: offset + clip_length]
    clip = clip.float()
    clip = (clip / 255.0) * 2.0 - 1.0
    if clip.size(1) < clip_length:
        tmp = torch.zeros([clip.size(0), clip_length - clip.size(1),
                            96, 96]).float().cuda()
        clip = torch.cat([clip, tmp], dim=1)
    clip = clip.unsqueeze(0)
    return clip

def parse_output(output_dict, flow_output_dict=None, fusion=False, use_edl=False, os_head=False, use_gcpl=False):
    act, prop_act, unct, prop_unct = None, None, None, None
    loc, conf, priors = output_dict['loc'][0], output_dict['conf'][0], output_dict['priors']
    prop_loc, prop_conf = output_dict['prop_loc'][0], output_dict['prop_conf'][0]
    if os_head:
        act, prop_act = output_dict['act'][0].squeeze(-1), output_dict['prop_act'][0].squeeze(-1)
    if use_gcpl:
        conf = -conf
        prop_conf = -prop_conf
    unct = output_dict['unct'][0] if use_edl else None
    prop_unct = output_dict['prop_unct'][0] if use_edl else None
    center = output_dict['center'][0]
    if fusion:
        flow_loc, flow_conf, priors = flow_output_dict['loc'][0], flow_output_dict['conf'][0], flow_output_dict['priors']
        flow_prop_loc, flow_prop_conf = flow_output_dict['prop_loc'][0], flow_output_dict['prop_conf'][0]
        flow_center = flow_output_dict['center'][0]
        # fusion by average
        loc = (loc + flow_loc) / 2.0
        prop_loc = (prop_loc + flow_prop_loc) / 2.0
        conf = (conf + flow_conf) / 2.0
        prop_conf = (prop_conf + flow_prop_conf) / 2.0
        center = (center + flow_center) / 2.0
        if os_head:
            flow_act, flow_prop_act = flow_output_dict['act'][0], flow_output_dict['prop_act'][0]
            act = (act + flow_act) / 2.0
            prop_act = (prop_act + flow_prop_act) / 2.0
        if use_edl:
            flow_prop_unct, flow_unct = flow_output_dict['prop_unct'][0], flow_output_dict['unct'][0]
            unct = (unct + flow_unct) / 2.0
            prop_unct = (prop_unct + flow_prop_unct) / 2.0
    return loc, conf, prop_loc, prop_conf, center, priors, unct, prop_unct, act, prop_act

def decode_predictions(loc, prop_loc, priors, conf, prop_conf, unct, prop_unct, act, prop_act, center, \
                        offset, sample_fps, clip_length, num_classes, score_func=nn.Softmax(dim=-1), use_edl=False, os_head=False):
    pre_loc_w = loc[:, :1] + loc[:, 1:]
    loc = 0.5 * pre_loc_w * prop_loc + loc
    segments = torch.cat(
        [priors[:, :1] * clip_length - loc[:, :1],
            priors[:, :1] * clip_length + loc[:, 1:]], dim=-1)
    segments.clamp_(min=0, max=clip_length)
    decoded_segments = (segments + offset) / sample_fps

    # compute uncertainty
    uncertainty = (unct + prop_unct) / 2.0 if use_edl else None

    actionness = None
    if os_head:
        # compute actionness
        act_score = act.sigmoid()
        prop_act_score = prop_act.sigmoid()
        actionness = (act_score + prop_act_score) / 2.0

    conf = score_func(conf)
    prop_conf = score_func(prop_conf)
    center = center.sigmoid()

    conf = (conf + prop_conf) / 2.0
    conf = conf * center * actionness.unsqueeze(-1) if os_head else conf * center
    conf = conf.view(-1, num_classes).transpose(1, 0)
    conf_scores = conf.clone()
    return decoded_segments, conf_scores, uncertainty, actionness

def filtering(decoded_segments, conf_score_cls, uncertainty, actionness, conf_thresh, use_edl=False, os_head=False):
    c_mask = conf_score_cls > conf_thresh
    if os_head:
        c_mask = c_mask & (actionness > 0.5)
    scores = conf_score_cls[c_mask]
    if scores.size(0) == 0:
        return None
    # masking segments
    l_mask = c_mask.unsqueeze(1).expand_as(decoded_segments)
    segments = decoded_segments[l_mask].view(-1, 2)
    segments = torch.cat([segments, scores.unsqueeze(1)], -1)  # (N, 3)
    if use_edl:
        # masking uncertainties
        uncertain_scores = uncertainty[c_mask]
        segments = torch.cat([segments, uncertain_scores.unsqueeze(1)], -1)  # (N, 4)
    if os_head:
        # masking actionness
        act_scores = actionness[c_mask]
        segments = torch.cat([segments, act_scores.unsqueeze(1)], -1)  # (N, 4) or (N, 5)
    return segments

def get_video_detections(output, idx_to_class, num_classes, top_k, nms_sigma, use_edl=False, os_head=False, cls_rng=None):
    res_dim = 3
    res_dim = res_dim + 1 if use_edl else res_dim  # 3 or 4
    res_dim = res_dim + 1 if os_head else res_dim  # 3 or 4 or 5
    res = torch.zeros(num_classes, top_k, res_dim)
    sum_count = 0
    for cl in cls_rng:  # from 1 to K+1 by default, or 0 to K for os_head
        if len(output[cl]) == 0:
            continue
        tmp = torch.cat(output[cl], 0)
        tmp, count = softnms_v2(tmp, sigma=nms_sigma, top_k=top_k, score_threshold=0.001, use_edl=use_edl, os_head=os_head)
        res[cl, :count] = tmp
        sum_count += count

    sum_count = min(sum_count, top_k)
    flt = res.contiguous().view(-1, res_dim)
    flt = flt.view(num_classes, -1, res_dim)
    proposal_list = []
    for cl in cls_rng:  # from 1 to K+1 by default, or 0 to K for os_head
        cl_idx = cl + 1 if os_head else cl
        class_name = idx_to_class[cl_idx]  # 1 to K
        tmp = flt[cl].contiguous()
        tmp = tmp[(tmp[:, 2] > 0).unsqueeze(-1).expand_as(tmp)].view(-1, res_dim)
        if tmp.size(0) == 0:
            continue
        tmp = tmp.detach().cpu().numpy()
        for i in range(tmp.shape[0]):
            tmp_proposal = {}
            tmp_proposal['label'] = class_name
            tmp_proposal['score'] = float(tmp[i, 2])
            tmp_proposal['segment'] = [float(tmp[i, 0]),
                                        float(tmp[i, 1])]
            tmp_proposal['uncertainty'] = float(tmp[i, 3]) if use_edl else 0.0
            tmp_proposal['actionness'] = float(tmp[i, 4]) if os_head else 0.0
            proposal_list.append(tmp_proposal)
    return proposal_list

def test(cfg):
    # get data
    video_infos = get_video_info(config['dataset']['testing']['video_info_path'])
    _, idx_to_class = get_class_index_map(config['dataset']['class_info_path'])
    npy_data_path = cfg.rgb_data_path if cfg.fusion else config['dataset']['testing']['video_data_path']
    class_range = range(1, cfg.num_classes) if not cfg.os_head else range(0, cfg.num_classes)

    # prepare model
    net, flow_net = build_model(fusion=cfg.fusion)
    out_layer = DirichletLayer(evidence=cfg.evidence, dim=-1) if cfg.use_edl else nn.Softmax(dim=-1)

    centor_crop = videotransforms.CenterCrop(config['dataset']['testing']['crop_size'])
    result_dict = {}
    for video_name in tqdm.tqdm(list(video_infos.keys()), ncols=0):
        # get the clip offsets
        offsetlist = get_offsets(video_infos, video_name, cfg.clip_length, cfg.stride)
        sample_fps = video_infos[video_name]['sample_fps']

        # load data
        data = prepare_data(npy_data_path, video_name, centor_crop)
        flow_data = prepare_data(cfg.flow_data_path, video_name, centor_crop) if cfg.fusion else None

        output = [[] for cl in range(cfg.num_classes)]
        output_dict_all = []
        for offset in offsetlist:
            # prepare clip of a video
            clip = prepare_clip(data, offset, cfg.clip_length)
            flow_clip = prepare_clip(flow_data, offset, cfg.clip_length) if cfg.fusion else None
            # run inference
            with torch.no_grad():
                output_dict = net(clip)
                flow_output_dict = flow_net(flow_clip) if cfg.fusion else None
            output_dict_all.append((output_dict, flow_output_dict, offset))

        # post-processing
        for (output_dict, flow_output_dict, offset) in output_dict_all:
            loc, conf, prop_loc, prop_conf, center, priors, unct, prop_unct, act, prop_act = parse_output(output_dict, flow_output_dict, \
                fusion=cfg.fusion, use_edl=cfg.use_edl, os_head=cfg.os_head, use_gcpl=cfg.use_gcpl)

            decoded_segments, conf_scores, uncertainty, actionness = decode_predictions(loc, prop_loc, priors, conf, prop_conf, unct, prop_unct, act, prop_act, \
                                                                center, offset, sample_fps, cfg.clip_length, cfg.num_classes, score_func=out_layer, use_edl=cfg.use_edl, os_head=cfg.os_head)
            # filtering out clip-level predictions with low confidence
            for cl in class_range:  # from 1 to K+1 by default, or 0 to K for os_head
                segments = filtering(decoded_segments, conf_scores[cl], uncertainty, actionness, cfg.conf_thresh, use_edl=cfg.use_edl, os_head=cfg.os_head)  # (N,5)
                if segments is None:
                    continue
                output[cl].append(segments)

        # get final detection results for each video
        result_dict[video_name] = get_video_detections(output, idx_to_class, cfg.num_classes, cfg.top_k, cfg.nms_sigma, use_edl=cfg.use_edl, os_head=cfg.os_head, cls_rng=class_range)

    output_dict = {"version": "THUMOS14", "results": dict(result_dict), "external_data": {}}
    with open(os.path.join(cfg.output_path, cfg.json_name), "w") as out:
        json.dump(output_dict, out)

def get_basic_config(config):
    class cfg: pass
    cfg.num_classes = config['dataset']['num_classes']
    cfg.conf_thresh = config['testing']['conf_thresh']
    cfg.top_k = config['testing']['top_k']
    cfg.nms_thresh = config['testing']['nms_thresh']
    cfg.nms_sigma = config['testing']['nms_sigma']
    cfg.clip_length = config['dataset']['testing']['clip_length']
    cfg.stride = config['dataset']['testing']['clip_stride']
    cfg.use_rpl = config['model']['use_rpl'] if 'use_rpl' in config['model'] else False
    cfg.use_gcpl = config['training']['rpl_config']['gcpl'] if cfg.use_rpl and 'gcpl' in config['training']['rpl_config'] else False
    cfg.use_edl = config['model']['use_edl'] if 'use_edl' in config['model'] else False
    if cfg.use_edl:
        cfg.evidence = config['model']['evidence']
    cfg.os_head = config['model']['os_head'] if 'os_head' in config['model'] else False
    if cfg.os_head:
        cfg.num_classes = cfg.num_classes - 1

    cfg.json_name = config['testing']['output_json']
    cfg.output_path = config['testing']['output_path']
    if not os.path.exists(cfg.output_path):
        os.makedirs(cfg.output_path)
    cfg.fusion = config['testing']['fusion']

    # getting path for fusion
    cfg.rgb_data_path = config['testing'].get('rgb_data_path',
                                        './datasets/thumos14/test_npy/')
    cfg.flow_data_path = config['testing'].get('flow_data_path',
                                        './datasets/thumos14/test_flow_npy/')
    return cfg

if __name__ == '__main__':

    cfg = get_basic_config(config)

    test(cfg)


import pandas as pd
import numpy as np

from nltk.corpus import stopwords
from nltk.tokenize import TweetTokenizer
from sklearn.svm import SVC
from sklearn.feature_extraction.text import TfidfVectorizer
from langdetect import detect
from datetime import datetime

from .tweet_cleaner import TweetCleaner

class SentimentClassifier:
    LANGS = {
        'hu': 'hungarian',
        'sv': 'swedish',
        'kk': 'kazakh',
        'fi': 'finnish',
        'no': 'norwegian',
        'ar': 'arabic',
        'id': 'indonesian',
        'pt': 'portuguese',
        'tr': 'turkish',
        'az': 'azerbaijani',
        'sl': 'slovene',
        'es': 'spanish',
        'da': 'danish',
        'ne': 'nepali',
        'ro': 'romanian',
        'el': 'greek',
        'nl': 'dutch',
        'tg': 'tajik',
        'de': 'german',
        'en': 'english',
        'ru': 'russian',
        'fr': 'french',
        'it': 'italian',
    }

    ORGS = {
        'apple': 0.1,
        'google': 0.2,
        'microsoft': 0.3,
        'twitter': 0.4
    }

    def __init__(self, normalization_lexicon_path: str, ngram_range=(1, 1), sentiment_rate_range=(1,5), start_day_hour=8):
        self._start_day_hour = start_day_hour
        self._sentiment_rate_range = sentiment_rate_range
        self._clf_sentiment = SVC(gamma="scale", probability=True)
        self._clf_organization = SVC(gamma="scale")
        self._cleaner = TweetCleaner()
        self._lexicon = {}
        self._tknzr = TweetTokenizer()
        self._vectorizer = TfidfVectorizer(ngram_range=ngram_range, sublinear_tf=True)
        with open(normalization_lexicon_path, "r", encoding="utf-8") as lexicon_file:
            for line in lexicon_file:
                informal, formal = line.split("\t")
                self._lexicon[informal] = formal

    def fit(self, tweets_dataset):
        corpus = []
        sentiment_labels = []
        organization_features = []
        temporal_features = []
        organization_labels = []
        for idx, tweet in tweets_dataset.iterrows():
            preprocessed_tweet = self._preprocess(tweet['TweetText'], discard_short_tweets=True)
            if preprocessed_tweet is not None:
                corpus.append(preprocessed_tweet)
                sentiment_labels.append(tweets_dataset['Sentiment'][idx])
                org = tweets_dataset['Topic'][idx]
                organization_labels.append(org)
                organization_features.append(self.ORGS[org])
                temporal_features.append(self._calc_temporal_feature(tweets_dataset['TweetDate'][idx]))
        features = self._vectorizer.fit_transform(corpus)
        features_sentiment = pd.DataFrame(features.todense()).join(pd.DataFrame({'Topic': organization_features}))
        features_sentiment = features_sentiment.join(pd.DataFrame({'TweetDate': temporal_features}))
        self._clf_sentiment.fit(features_sentiment, sentiment_labels)
        self._clf_organization.fit(features, organization_labels)

    def predict_sentiment(self, X, get_rate=False):
        features = self._prepare_for_prediction(X)
        features = pd.DataFrame(features.todense()).join(X['Topic'].apply(lambda x: self.ORGS[x]).reset_index(drop=True))
        features = features.join(X['TweetDate'].apply(self._calc_temporal_feature).reset_index(drop=True))
        if get_rate:
            probs = self._clf_sentiment.predict_proba(features)
            return list(map(self._calc_sentiment_rate_from_probs, probs))
        return self._clf_sentiment.predict(features)

    def predict_organization(self, X):
        features = self._prepare_for_prediction(X)
        return self._clf_organization.predict(features)

    def _calc_sentiment_rate_from_probs(self, probs):
        p_negative = probs[np.where(self._clf_sentiment.classes_ == 'negative')[0][0]]
        p_neutral = probs[np.where(self._clf_sentiment.classes_ == 'neutral')[0][0]]
        p_positive = probs[np.where(self._clf_sentiment.classes_ == 'positive')[0][0]]
        min_rate, max_rate = self._sentiment_rate_range
        middle_rate = (min_rate + max_rate) / 2.0
        return p_negative * min_rate + p_neutral * middle_rate + p_positive * max_rate

    def _calc_temporal_feature(self, str_time):
        t = datetime.strptime(str_time, '%a %b %d %H:%M:%S %z %Y')
        feature = (t.hour - self._start_day_hour) % 24 + 1  # from 1 to 24
        return feature / (24 * 2.0)                         # from 0.02 to 0.5 (normalized)

    def _prepare_for_prediction(self, X):
        corpus = [self._preprocess(tweet['TweetText']) for idx, tweet in X.iterrows()]
        features = self._vectorizer.transform(corpus)
        return features

    def _preprocess(self, text: str, discard_short_tweets=False):
        text = self._cleaner.clean_up(text)

        word_tokens = self._tknzr.tokenize(text)

        try:
            lang = detect(text)
        except:
            lang = 'en'
        stop_words = set(stopwords.words(self.LANGS.get(lang, 'english')))
        word_tokens = [word for word in word_tokens if word not in stop_words]

        word_tokens = [self._lexicon[word] if word in self._lexicon else word for word in word_tokens]

        if discard_short_tweets and len(word_tokens) < 3:
            return None

        return " ".join(word_tokens)


import torch
from polygon import Polygon, Bbox
import json
import socket
import pickle
import cv2
import struct
import socketserver
import datetime
import argparse
from polygon import Bbox, Keypoints
import utils
from PIL import Image
import os
import time
from detection.yolov3.batch_detect import model
from detection.yolov3.utils import do_detect
from detection.yolov3.image import letterbox_image
import numpy as np
import openpifpaf.network
import openpifpaf
import glob
from config import OUTPUT_DIR
from class_names import DEFAULT_CLASS_NAMES
import math
import pandas as pd

class Detection:
    def __init__(self, class_id=0, track_id=0, polygon=Polygon(), bbox=Bbox(), keypoints=Keypoints()):
        self.class_id = class_id
        self.track_id = track_id
        self.polygon = polygon
        self.bbox = bbox
        self.keypoints = keypoints

    @staticmethod
    def from_json(data):
        return Detection(data["class_id"], data["track_id"],
                         Polygon(data["polygon"]), Bbox(*data["bbox"]), Keypoints(data["keypoints"]))

    def to_json(self):
        return {
            "track_id": self.track_id,
            "class_id": self.class_id,
            "polygon": self.polygon.to_json(),
            "bbox": self.bbox.to_json(),
            "keypoints": self.keypoints.to_json()
        }

    def to_dict(self):
        return {"track_id": self.track_id, "class_id": self.class_id, **self.bbox.to_dict(),
                "polygon": self.polygon.to_str(), "kp": self.keypoints.to_str()}

    @staticmethod
    def from_df(row):
        bbox = Bbox(row.x, row.y, row.w, row.h)
        return Detection(row.class_id, row.track_id, Polygon.from_str(row.polygon), bbox, Keypoints.from_str(row.kp))

    def copy(self):
        return Detection(self.class_id, self.track_id, self.polygon.copy(), self.bbox.copy(), self.keypoints.copy())

    def __repr__(self):
        return "Detection(class_id={}, track_id={}, bbox={}, polygon={}, keypoints={})".format(self.class_id, self.track_id,
                                                                                self.bbox, self.polygon, self.keypoints)

class Detector:
    def __init__(self):
        self.use_cuda = torch.cuda.is_available()
        self.device = torch.device('cuda' if self.use_cuda else 'cpu')
        torch.backends.cudnn.benchmark = True

    def detect_single_image(self, image_path, crop_area=None):
        """
        Output:
            detections ([Detection])
        """
        raise NotImplementedError

    def detect_batch(self, images, crop_area=None):
        for image_path in images:
            yield self.detect_single_image(image_path, crop_area)

    def terminate(self):
        pass

class YOLODetector(Detector):
    CONFIG_FILE = "detection/yolov3/yolo_cfg/topview-6-predict.cfg"
    WEIGHTS_FILE = "detection/yolov3/weights/topview-final-6.weights"

    def __init__(self):
        super().__init__()

        self.model = model(self.CONFIG_FILE, self.WEIGHTS_FILE)
        self.use_cuda = torch.cuda.is_available()
        if self.use_cuda:
            self.model.cuda()

    def correct_yolo_boxes(self, boxes, im_w, im_h, net_w, net_h, cropping_area=None):
        im_w, im_h = float(im_w), float(im_h)
        net_w, net_h = float(net_w), float(net_h)

        if cropping_area is None:
            area_w, area_h = im_w, im_h
            factor_w, factor_h = 1, 1
            offset_w, offset_h = 0, 0
        else:
            area_w, area_h = cropping_area[2] - cropping_area[0], cropping_area[3] - cropping_area[1]
            factor_w, factor_h = area_w / im_w, area_h / im_h
            offset_w, offset_h = cropping_area[0] / im_w, cropping_area[1] / im_h

        if net_w/area_w < net_h/area_h:
            new_w = net_w
            new_h = (area_h * net_w)/area_w
        else:
            new_w = (area_w * net_h)/area_h
            new_h = net_h

        xo, xs = (net_w - new_w)/(2*net_w), net_w/new_w
        yo, ys = (net_h - new_h)/(2*net_h), net_h/new_h

        for i in range(len(boxes)):
            b = boxes[i]
            b[0] = (b[0] - xo) * xs * factor_w + offset_w
            b[1] = (b[1] - yo) * ys * factor_h + offset_h
            b[2] *= xs * factor_w
            b[3] *= ys * factor_h

    def yolo_to_coco_class(self, class_id):
        return class_id

    def detect_single_image(self, image_path, crop_area=None):
        img = Image.open(image_path).convert('RGB')
        W, H = img.size

        # Cropping
        if crop_area is None:
            crop_area = (0, 0, W, H)
        x_crop, y_crop, w_crop, h_crop = crop_area

        margin = 0.03 * max(W, H)  # 3% margin

        detections = []

        # Number of repeated cropping areas to span the entire image
        n_left = math.ceil(x_crop / w_crop)
        n_right = math.ceil((W - (x_crop + w_crop)) / w_crop)
        n_top = math.ceil(y_crop / h_crop)
        n_bottom = math.ceil((H - (y_crop + h_crop)) / h_crop)

        for i in range(-n_top, 1 + n_bottom):
            for j in range(-n_left, 1 + n_right):
                crop_area = (x_crop + j * w_crop - margin, y_crop + i * h_crop - margin,
                             x_crop + j * w_crop + w_crop + margin, y_crop + i * h_crop + h_crop + margin)
                print(i, j, crop_area)

                img_crop = img.crop(crop_area)

                sized = letterbox_image(img_crop, self.model.width, self.model.height)
                start = time.time()

                with torch.no_grad():
                    boxes = do_detect(self.model, sized, 0.5, 0.4, self.use_cuda)

                self.correct_yolo_boxes(boxes, W, H, self.model.width, self.model.height, crop_area)
                finish = time.time()

                print('%s: Predicted in %f seconds.' % (image_path, (finish - start)))

                detections.extend([Detection(class_id=self.yolo_to_coco_class(b[6].item()), track_id=i,
                                        bbox=Bbox((b[0] - b[2] / 2).item() * W, (b[1] - b[3] / 2).item() * H,
                                                  b[2].item() * W, b[3].item() * H)) for i, b in enumerate(boxes)])

        return detections

class OpenPifPafDetector(Detector):

    def __init__(self):
        super().__init__()

        parser = argparse.ArgumentParser(
            description=__doc__,
            formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        )
        openpifpaf.network.nets.cli(parser)
        openpifpaf.decoder.cli(parser, force_complete_pose=False, instance_threshold=0.2)
        parser.add_argument('--batch-size', default=1, type=int,
                            help='processing batch size')
        parser.add_argument('--loader-workers', default=2, type=int,
                            help='number of workers for data loading')
        self.args = parser.parse_args([])

        # add args.device
        self.args.device = torch.device('cpu')
        self.args.pin_memory = False
        if torch.cuda.is_available():
            self.args.device = torch.device('cuda')
            self.args.pin_memory = True

        model, _ = openpifpaf.network.nets.factory_from_args(self.args)
        model = model.to(self.args.device)
        self.processor = openpifpaf.decoder.factory_from_args(self.args, model)

    def detect(self, image_paths, image_tensors, processed_images_cpu):
        with torch.no_grad():
            images = image_tensors.permute(0, 2, 3, 1)

            processed_images = processed_images_cpu.to(self.args.device, non_blocking=True)
            fields_batch = self.processor.fields(processed_images)
            pred_batch = self.processor.annotations_batch(fields_batch, debug_images=processed_images_cpu)

            _, image, processed_image_cpu, pred = image_paths[0], images[0], processed_images_cpu[0], pred_batch[0]

            self.processor.set_cpu_image(image, processed_image_cpu)
            keypoint_sets, scores = self.processor.keypoint_sets_from_annotations(pred)  # keypoints shape: (nb_detections, 17, 3)

            keypoints = keypoint_sets.reshape(-1, 3 * 17)  # keypoints shape: (nb_detections, 3*17)
            print(keypoints)

            detections = [Detection(track_id=i, class_id=5, keypoints=Keypoints(kp), bbox=utils.keypoints_to_bbox(kp)) for i, kp in enumerate(keypoints)]

            return detections

    def detect_single_image(self, image_path, crop_area=None):
        data = openpifpaf.datasets.ImageList([image_path])
        data_loader = torch.utils.data.DataLoader(
            data, batch_size=self.args.batch_size, shuffle=False,
            pin_memory=self.args.pin_memory, num_workers=self.args.loader_workers)

        image_paths, image_tensors, processed_images_cpu = next(iter(data_loader))
        return self.detect(image_paths, image_tensors, processed_images_cpu)

    def detect_batch(self, images, crop_area=None):
        data = openpifpaf.datasets.ImageList(images)
        data_loader = torch.utils.data.DataLoader(
            data, batch_size=self.args.batch_size, shuffle=False,
            pin_memory=self.args.pin_memory, num_workers=self.args.loader_workers)

        for image_paths, image_tensors, processed_images_cpu in data_loader:
            yield self.detect(image_paths, image_tensors, processed_images_cpu)

class DetectionHandler():
    PAYLOAD_SIZE = struct.calcsize(">L")

    def __init__(self, detector, crop_area, sequence):
        self.detector = detector
        self.crop_area = crop_area

        if not os.path.exists(sequence):
            write_running_info(error="No such file or directory: {}".format(sequence))
            return

        file_names = sorted(glob.glob('{}/*.jpg'.format(sequence)), key=utils.natural_sort_key)
        video_name = os.path.basename(sequence)
        nb_frames = len(file_names)

        track_info = TrackInfo(video_name)
        start_time = datetime.datetime.now()

        for frame, detections in enumerate(detector.detect_batch(file_names, crop_area)):
            file_path = file_names[frame]
            base = os.path.basename(file_path)
            file_name = os.path.splitext(base)[0]

            track_info.write_detections(file_name, detections)
            write_running_info(video_name, start_time, frame, nb_frames)

        torch.cuda.empty_cache()
        track_info.save_to_disk()

    def write_running_info(video_name="", start_time="", frame=0, nb_frames=0, error=None):
        data = {
            "video_name": video_name,
            "current_frame": frame + 1,
            "total_frame": nb_frames,
            "start_time": str(start_time),
            "last_update": str(datetime.datetime.now()),
        }

        if error is not None:
            data["error"] = error

        file_name = os.path.join(OUTPUT_DIR, "running_info.json")

        with open(file_name, "w") as f:
            json.dump(data, f)


#! /bin/env python

# Copyright 2015 Peter Andrews @ CSHL

import sys
import mumdex
import signal
import numpy as np

def handler(signum, frame):
    print >> sys.stderr, "pipe closed in candidate_finder.py"
    sys.exit(1)
signal.signal(signal.SIGPIPE, handler)

if len(sys.argv) == 5:
    whole_genome = bool(int(sys.argv[1]))
    chrom = sys.argv[2]
    position = int(sys.argv[3])
    n_loci = int(sys.argv[4])
else:
    print >> sys.stderr, \
        "usage: candidate_finder.py whole_genome chromosome position n_loci"
    exit(1)

mums_dir="/data/safe/paa/analysis/mums"
if whole_genome:
    pop = mumdex.Population(mums_dir + "/wg-families.txt")
    samples_dir = mums_dir + "/wg-output/samples/"
else:
    pop = mumdex.Population(mums_dir + "/families.txt")
    samples_dir = mums_dir + "/output/samples/"

# ref = mumdex.Reference(mums_dir + "/hg19/chrAll.fa")
counts = mumdex.PopulationCounts(chrom, position, n_loci,
                                 whole_genome=whole_genome)

ignore_zeros = True
max_families = 5
min_coverage = 20
min_anchor_count = 10
min_well_covered = 0.5
max_parents_frac = 0.05
min_support = 20

good_anchors = set()
refs = counts.refs
anchors = counts.anchors
zero = counts.zero
edge = counts.edge
supports = counts.supports
positions = set()
n_cand = 0

is_parent = np.zeros((counts.samples.size), np.bool)
for sample in xrange(0, counts.samples.size):
    if counts.samples[sample]['rtp'] == "mother" or \
            counts.samples[sample]['rtp'] == "father":
        is_parent[sample] = True
    # print counts.samples[sample]['rtp'], is_parent[sample]

for position in xrange(0, counts.positions.size):
    # print position
    for out in xrange(0, 2):
        cand_samples = set()
        cand_families = set()
        n_samples_high = 0
        n_samples_low = 0
        n_parents_high = 0
        n_parents_low = 0
        max_support = 0
        for sample in xrange(0, counts.samples.size):
            count = int(anchors[out][sample][position])
            coverage = count + int(refs[out][sample][position])
            if ignore_zeros:
                count -= int(zero[out][sample][position]) + \
                    int(edge[out][sample][position])

            if coverage > min_coverage:
                if count >= min_anchor_count:
                    n_samples_high += 1
                    if is_parent[sample]:
                        n_parents_high += 1
                    cand_samples.add(sample)
                    cand_families.add(pop.sample_family(sample))
                else:
                    n_samples_low += 1
                    if is_parent[sample]:
                        n_parents_low += 1

            if max_support < supports[out][sample][position]:
                    max_support = supports[out][sample][position]
        parents_ok_coverage = n_parents_high + n_parents_low
        if parents_ok_coverage > 0:
            parent_frac = float(n_parents_high) / \
                (n_parents_high + n_parents_low)
            covered_frac = float(n_samples_high + n_samples_low) / \
                counts.samples.size
            if (n_samples_high >= 1 and
                len(cand_families) <= max_families and
                max_support > min_support and
                parent_frac < max_parents_frac and
                covered_frac > min_well_covered):
                posinfo = counts.positions[position]
                print posinfo[0], posinfo[1], out, \
                    len(cand_samples), len(cand_families), \
                    parent_frac, covered_frac, "s",
                for sample in cand_samples:
                    print sample,
                print "f",
                for family in cand_families:
                    print family,
                print
                sys.stdout.flush()
                positions.add((position, out))
                n_cand += 1

print >> sys.stderr, n_cand, "candidates"

exit()

all_info = dict()
for sample in xrange(0, counts.samples.size):
    print "\rReading bridges for sample", sample, "of", counts.samples.size,
    sys.stdout.flush()
    sample_info = dict()
    mumdex_name = pop.mumdex_name(samples_dir, sample)
    mums = mumdex.MUMdex(mumdex_name)
    for pos_key in positions:
        (position, out) = pos_key
        bridge_info = mums.bridges(
            ref.index(counts.positions[position][0]),
            counts.positions[position][1], out)
        sample_info[pos_key] = bridge_info
    all_info[sample] = sample_info
print


# BSD 3-Clause License
#
# Copyright (c) 2016-19, University of Liverpool
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# * Neither the name of the copyright holder nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
Parser module specific to CCMpred predictions
"""

from __future__ import division

__author__ = "Felix Simkovic"
__date__ = "03 Aug 2016"
__version__ = "0.1"

import numpy as np
import sys

from conkit.io._parser import ContactFileParser
from conkit.core.contact import Contact
from conkit.core.contactmap import ContactMap
from conkit.core.contactfile import ContactFile

class CCMpredParser(ContactFileParser):
    """
    Class to parse a CCMpred contact matrix

    """

    def __init__(self):
        super(CCMpredParser, self).__init__()

    def read(self, f_handle, f_id="ccmpred"):
        """Read a contact file

        Parameters
        ----------
        f_handle
           Open file handle [read permissions]
        f_id : str, optional
           Unique contact file identifier

        Returns
        -------
        :obj:`~conkit.core.contactfile.ContactFile`

        """
        contact_file = ContactFile(f_id)
        contact_file.method = "Contact map predicted using CCMpred"
        contact_map = ContactMap("map_1")
        contact_file.add(contact_map)

        # Bits ripped from Stefan Seemayer's script shipped with CCMpred
        mat = np.loadtxt(f_handle)
        if mat.size > 0:
            raw_contacts = self._get_contact_pairs(mat)
            for res1_seq, res2_seq, raw_score in zip(raw_contacts[0], raw_contacts[1], mat[raw_contacts]):
                if res1_seq > res2_seq:
                    continue
                # Matrix starts count at 0 so increment numbers by one straight away
                contact = Contact(int(res1_seq + 1), int(res2_seq + 1), float(raw_score))
                contact_map.add(contact)

        return contact_file

    def _get_contact_pairs(self, mat):
        """Get all contact pairs in the matrix

        Parameters
        ----------
        mat : :obj:`~numpy.ndarray`
           A :mod:`numpy` matrix

        Returns
        -------
        list
           A list of contact pairs

        """
        contacts = mat.argsort(axis=None)[::-1]
        contacts = (contacts % mat.shape[0]).astype(np.uint16), np.floor(contacts / mat.shape[0]).astype(np.uint16)
        return contacts

    def write(self, f_handle, hierarchy):
        """Write a contact file instance to to file

        Parameters
        ----------
        f_handle
           Open file handle [write permissions]
        hierarchy : :obj:`~conkit.core.contactfile.ContactFile`, :obj:`~conkit.core.contactmap.ContactMap`
                    or :obj:`~conkit.core.contact.Contact`

        Raises
        ------
        :exc:`RuntimeError`
           More than one contact map in the hierarchy
        :exc:`TypeError`
           Python3 requires f_handle to be in `wb` or `ab` mode

        """
        # Python3 support requires bytes mode
        if sys.version_info.major == 3 and not (f_handle.mode == "wb" or f_handle.mode == "ab"):
            raise TypeError("Python3 requires f_handle to be in 'wb' or 'ab' mode")

        # Double check the type of hierarchy and reconstruct if necessary
        contact_file = self._reconstruct(hierarchy)

        if len(contact_file) > 1:
            raise RuntimeError("More than one contact map provided")

        for contact_map in contact_file:
            len_mat = max([c.res1_seq for c in contact_map] + [c.res2_seq for c in contact_map])
            mat = np.zeros((len_mat, len_mat), np.float64)

            for contact in contact_map:
                mat[contact.res1_seq - 1][contact.res2_seq - 1] = contact.raw_score
                mat[contact.res2_seq - 1][contact.res1_seq - 1] = contact.raw_score

            np.savetxt(f_handle, mat, delimiter="\t")

        return