#! /usr/bin/env python
from __future__ import print_function
import os
import sys
import io
import csv
from optparse import OptionParser

import numpy as np
import tensorflow as tf
from flask import Flask, jsonify, render_template, request
from tensorflow.contrib import learn

import data_helpers
from flask_restplus import Api, Resource, fields, marshal_with, reqparse
from text_cnn import TextCNNModel

from werkzeug.datastructures import FileStorage

# Create event type model
def __load_models():
    mod = "models/event-types/model.ckpt"
    vocab = "models/event-types/vocabulary.voc"
    data_helpers.merge_model_file(mod, remove=True)
    type_classifier = TextCNNModel.restore(
        mod, vocab, sequence_length=32, num_classes=14, vocab_size=44345)
    type_classifier.labels = ["bombings", "collapse", "crash", "derailment", "earthquake",
                              "explosion", "fire", "floods", "haze", "meteorite",
                              "none", "shootings", "typhoon", "wildfire"]

    # Create related event type model:
    mod = "models/event-related/model.ckpt"
    vocab = "models/event-related/vocabulary.voc"
    data_helpers.merge_model_file(mod, remove=True)
    related_classifier = TextCNNModel.restore(
        mod, vocab, sequence_length=43, num_classes=2, vocab_size=87420)
    related_classifier.labels = ["non-related", "related"]

    # Create related info type model:
    mod = "models/info-types/model.ckpt"
    vocab = "models/info-types/vocabulary.voc"
    data_helpers.merge_model_file(mod, remove=True)
    info_classifier = TextCNNModel.restore(
        mod, vocab, sequence_length=32, num_classes=8, vocab_size=44345)
    info_classifier.labels = ["affected_individuals", "caution_and_advice", "donations_and_volunteering",
                              "infrastructure_and_utilities", "not_applicable", "not_labeled",
                              "other_useful_information", "sympathy_and_support"]

    return (type_classifier, related_classifier, info_classifier)

def main(argv):

    # Parse command line arguments:
    parser = OptionParser()
    parser.add_option('-p', '--port', type='int', dest='port', default=80,
                      help="the API port for serving CREES [default: %default]")

    parser.add_option('-n', '--namespace', type='string', dest='api_namespace', default="comrades",
                      help="the API namespace for CREES [default: %default]")

    parser.add_option('-d', '--debug', action="store_true", dest='debug', default=False,
                      help="show debugging information")

    (options, args) = parser.parse_args()

    # Use environment variables if available:
    if os.environ.get('CREES_PORT') is not None and options.port is 80:
        options.port = int(os.environ.get('CREES_PORT'))
    if os.environ.get('CREES_NAMESPACE') is not None and options.api_namespace is 'comrades':
        options.api_namespace = str(os.environ.get('CREES_NAMESPACE'))

    print(' * Using port: '+str(options.port))
    print(' * Using namespace: /'+str(options.api_namespace))

    # Web App:
    app = Flask(__name__, static_url_path='')
    app.config['RESTPLUS_MASK_SWAGGER'] = False
    app.config['APPLICATION_ROOT'] = '/' + options.api_namespace
    app.config.SWAGGER_UI_DOC_EXPANSION = 'list'
    api = Api(app, version='0.3', title='COMRADES Event API',
              description='A set of tools for analysing short textual documents (e.g. tweets).',
              doc='/' + options.api_namespace + '/',
              endpoint='/' + options.api_namespace
              )

    ns = api.namespace(options.api_namespace,
                       description='Event detection tools.')

    # Load models:
    type_classifier, related_classifier, info_classifier = __load_models()

    # Routes:
    @ns.route('/')
    class RootController(Resource):
        def get(self):
            return app.send_static_file('index.html')

    @ns.route('/events/eventType')
    class EventClassifierController(Resource):
        """
        Performs event detection on short piece of text (e.g. tweets).
        """

        # API arguments:
        get_arguments = reqparse.RequestParser()
        get_arguments.add_argument(
            'text', required=True, help="The text to be analysed.")

        # Output arguments:
        model = api.model('Category', {
            'input': fields.String,
            'label': fields.String,
            'classifier': fields.String,
            'version': fields.Float
        })

        @api.doc(id='textEventType')
        @api.expect(get_arguments, validate=True)
        @api.marshal_with(model, description='Event type')
        def get(self):
            """Obtains the type of event associated with a post."""
            args = self.get_arguments.parse_args()
            text = args['text']
            results = type_classifier.predict(text)

            return {'input': text, 'label': results, 'classifier': "CNN", 'version': 0.3}

        post_arguments = reqparse.RequestParser()
        post_arguments.add_argument(
            'texts', required=True, type=list, location='json', help="The JSON array containning the strings to be analysed.")

        # Output arguments:
        model2_inner = api.model('Category', {
            'input': fields.String,
            'label': fields.String,
        })
        model2 = api.model('Categories', {
            'labels': fields.List(fields.Nested(model2_inner)),
            'classifier': fields.String,
            'version': fields.Float
        })

        @api.doc(id='textEventType')
        @api.expect(post_arguments, validate=False)
        @api.marshal_with(model2, description='Event Type')
        def post(self):
            """Identifies the type of events associated with multiple posts."""
            data = request.get_json()

            resp = []
            for text in data:
                results = type_classifier.predict(text)
                resp.append({'input': text, 'label': results})

            return {'labels': resp, 'classifier': "CNN", 'version': 0.3}

    @ns.route('/events/infoType')
    class InfoTypeClassifierController(Resource):
        """
        Identifies the type of information associated with a post.
        """

        # API arguments:
        get_arguments = reqparse.RequestParser()
        get_arguments.add_argument(
            'text', required=True, help="The text to be analysed.")

        # Output arguments:
        model = api.model('Category', {
            'input': fields.String,
            'label': fields.String,
            'classifier': fields.String,
            'version': fields.Float
        })

        @api.doc(id='textInfoType')
        @api.expect(get_arguments, validate=True)
        @api.marshal_with(model, description='Info type')
        def get(self):
            """Identifies the type of information associated with a post."""
            args = self.get_arguments.parse_args()
            text = args['text']
            results = info_classifier.predict(text)

            return {'input': text, 'label': results, 'classifier': "CNN", 'version': 0.3}

        post_arguments = reqparse.RequestParser()
        post_arguments.add_argument(
            'texts', required=True, type=list, location='json', help="The JSON array containning the strings to be analysed.")

        # Output arguments:
        model2_inner = api.model('Category', {
            'input': fields.String,
            'label': fields.String,
        })
        model2 = api.model('Categories', {
            'labels': fields.List(fields.Nested(model2_inner)),
            'classifier': fields.String,
            'version': fields.Float
        })

        @api.doc(id='textInfoType')
        @api.expect(post_arguments, validate=False)
        @api.marshal_with(model2, description='Info Type')
        def post(self):
            """Identifies the type of information associated with multiple posts'"""
            data = request.get_json()

            resp = []
            for text in data:
                results = info_classifier.predict(text)
                resp.append({'input': text, 'label': results})

            return {'labels': resp, 'classifier': "CNN", 'version': 0.3}

    @ns.route('/events/eventRelated')
    class RelatedClassifierController(Resource):
        """
        Identifies if a post is talking about an event.
        """

        # API arguments:
        get_arguments = reqparse.RequestParser()
        get_arguments.add_argument(
            'text', required=True, help="The text to be analysed.")

        # Output arguments:
        model = api.model('Category', {
            'input': fields.String,
            'label': fields.String,
            'classifier': fields.String,
            'version': fields.Float
        })

        @api.doc(id='textEventRelated')
        @api.expect(get_arguments, validate=True)
        @api.marshal_with(model, description='Event relation')
        def get(self):
            """Identifies if a post is talking about an event."""
            args = self.get_arguments.parse_args()
            text = args['text']
            results = related_classifier.predict(text)

            return {'input': text, 'label': results, 'classifier': "CNN", 'version': 0.3}

        post_arguments = reqparse.RequestParser()
        post_arguments.add_argument(
            'texts', required=True, type=list, location='json', help="The JSON array containning the strings to be analysed.")

        # Output arguments:
        model2_inner = api.model('Category', {
            'input': fields.String,
            'label': fields.String,
        })
        model2 = api.model('Categories', {
            'labels': fields.List(fields.Nested(model2_inner)),
            'classifier': fields.String,
            'version': fields.Float
        })

        @api.doc(id='textEventRelated')
        @api.expect(post_arguments, validate=False)
        @api.marshal_with(model2, description='Event relations')
        def post(self):
            """Identifies if multiple posts are talking about an event."""
            data = request.get_json()

            resp = []
            for text in data:
                results = related_classifier.predict(text)
                resp.append({'input': text, 'label': results})

            return {'labels': resp, 'classifier': "CNN", 'version': 0.3}

    # Start App:
    app.run(host='0.0.0.0', port=options.port, debug=options.debug)

if __name__ == "__main__":
    main(sys.argv)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright © 2017 Matthew Stone <mstone5@mgh.harvard.edu>
# Distributed under terms of the MIT license.

"""

"""

import numpy as np
import scipy.stats as ss
import pandas as pd
from .pesr_test import PESRTest, PESRTestRunner

class SRTest(PESRTest):
    def __init__(self, countfile, common=False, window=50, medians=None):
        self.countfile = countfile
        self.window = window
        self.common = common

        super().__init__(medians, common)

    def test_record_rec(self, record, called, background):
        # Test SR support at all coordinates within window of start/end
        results = []
        for coord in 'posA posB'.split():
            result = self._test_coord(record, coord, called, background)
            result['coord'] = coord
            results.append(result)
        results = pd.concat(results, ignore_index=True)

        # Add test for sum of posA and posB
        total = self._test_total(results)
        results = pd.concat([results, total], ignore_index=True)

        # Clean up columns
        results['name'] = record.id
        results['bg_frac'] = results.called / \
            (results.background + results.called)
        results['bg_frac'] = results.bg_frac.fillna(0)
        cols = 'name coord pos log_pval called background bg_frac'.split()
        return results[cols]

    def test_record(self, record, called, background):
        # Test SR support at all coordinates within window of start/end
        resultA = self._test_coord(
            record, 'posA', record.chrom, called, background)
        resultA['coord'] = 'posA'

        resultB = self._test_coord(
            record, 'posB', record.info['CHR2'], called, background)
        resultB['coord'] = 'posB'

        if (record.chrom == record.info['CHR2']) and (int(resultA.pos) >= int(resultB.pos)):
            rec = 1
            while True:
                result = self._test_coord_V2(
                    record, 'posB', record.chrom, called, background, rec)
                if int(result.pos) > int(resultA.pos.iloc[0]):
                    result['coord'] = 'posB'
                    resultB = result
                    break
                rec = rec + 1

        results = pd.concat([resultA, resultB], ignore_index=True)
        # Add test for sum of posA and posB
        total = self._test_total(results)
        results = pd.concat([results, total], ignore_index=True)

        # Clean up columns
        results['name'] = record.id
        results['bg_frac'] = results.called / \
            (results.background + results.called)
        results['bg_frac'] = results.bg_frac.fillna(0)
        cols = 'name coord pos log_pval called background bg_frac'.split()

        return results[cols]

    def _test_coord_V2(self, record, coord, chrom, samples, background, optimal=0):
        """Test enrichment at all positions within window"""
        if coord == 'posA':
            coord, strand = record.pos, record.info['STRANDS'][0]
        else:
            coord, strand = record.stop, record.info['STRANDS'][1]

        # Run SR test at each position
        results = []
        for pos in range(coord - self.window, coord + self.window + 1):
            result = self.test(chrom, pos, strand, samples, background)
            result = result.to_frame().transpose()
            result['pos'] = pos
            result['dist'] = np.abs(pos - coord)
            results.append(result)

        results = pd.concat(results, ignore_index=True)

        # Choose most significant position, using distance to predicted
        # breakpoint as tiebreaker
        results = results.sort_values(['log_pval', 'dist'], ascending=False)
        if optimal < len(results):
            best = results.iloc[optimal].to_frame().transpose()
        else:
            best = results.iloc[-1].to_frame().transpose()
            best['log_pval'] = 0
            best['pos'] = record.stop + self.window + 1
        return best

    def test(self, chrom, pos, strand, called, background):
        """
        Test enrichment of clipped reads in a set of samples at a given coord.

        Arguments
        ---------
        chrom : str
        pos : int
        strand : str
        called : list of str
            List of called samples to test
        background : list of str
            List of samples to use as background

        Returns
        -------
        called_median : float
        background_median : float
        log_pval : float
            Negative log10 p-value
        """

        # Load split counts.
        counts = self.load_counts(chrom, pos, strand)
        counts = self.normalize_counts(counts)

        return super().test(counts, called, background)

    def load_counts(self, chrom, pos, strand):
        """Load pandas DataFrame from tabixfile"""

        if pos > 0:
            region = '{0}:{1}-{1}'.format(chrom, pos)
            try:
                lines = self.countfile.fetch(region)
            except ValueError:
                lines = []
        else:
            lines = []
        #  counts = io.StringIO('\n'.join([l for l in lines]))

        cols = 'chrom pos clip count sample'.split()
        #  dtypes = dict(chrom=str, pos=int, clip=str, count=int, sample=str)

        counts = pd.DataFrame.from_records(
            [l[:5] for l in lines], columns=cols)
        counts['count'] = counts['count'].astype(int)

        # Restrict to splits in orientation of interest
        clip = 'right' if strand == '+' else 'left'
        counts = counts.loc[counts['clip'] == clip].copy()

        return counts

    def _test_total(self, results):
        """Test enrichment of posA+posB"""
        total = results['called background'.split()].sum()
        pval = ss.poisson.cdf(total.background, total.called)
        total['log_pval'] = np.abs(np.log10(pval))

        # format and add dummy metadata
        total = total.to_frame().transpose()
        total['coord'] = 'sum'
        total['pos'] = 0

        return total

    def _test_coord(self, record, coord, chrom, samples, background):
        """Test enrichment at all positions within window"""
        if coord == 'posA':
            coord, strand = record.pos, record.info['STRANDS'][0]
        else:
            coord, strand = record.stop, record.info['STRANDS'][1]

        # Run SR test at each position
        results = []
        for pos in range(coord - self.window, coord + self.window + 1):
            result = self.test(chrom, pos, strand, samples, background)
            result = result.to_frame().transpose()
            result['pos'] = pos
            result['dist'] = np.abs(pos - coord)
            results.append(result)

        results = pd.concat(results, ignore_index=True)

        # Choose most significant position, using distance to predicted
        # breakpoint as tiebreaker
        results = results.sort_values(['log_pval', 'dist'], ascending=False)
        best = results.iloc[0].to_frame().transpose()

        return best

class SRTestRunner(PESRTestRunner):
    def __init__(self, vcf, countfile, fout, n_background=160, common=False, window=100,
                 whitelist=None, blacklist=None, medians=None, log=False):
        """
        vcf : pysam.VariantFile
        countfile : pysam.TabixFile
        fout : writable file
        n_background : int
        window : int
        whitelist : list of str
        blacklist : list of str
        """
        self.srtest = SRTest(countfile, common, window, medians=medians)
        self.fout = fout

        super().__init__(vcf, common, n_background, whitelist, blacklist, log)

    def test_record(self, record):
        called, background = self.choose_background(record)
        counts = self.srtest.test_record(record, called, background)
        counts = counts.rename(columns={'called': 'called_median',
                                        'background': 'bg_median'})
        counts.to_csv(self.fout, header=False, index=False, sep='\t')


# Standard library
import os
import pickle
from glob import glob

# Third-party
import numpy as np
from astropy import units as u
from astropy.table import Table

# Project
from . import data_dir

__all__ = ['phot_system_list',
           'FilterSystem',
           'get_filter_names',
           'get_filter_properties',
           'phot_system_lookup',
           'load_zero_point_converter']

# list of photometric systems with pre-calculated filter properties
phot_system_list = [
    'HST_WFC3', 'HST_ACSWF', 'SDSSugriz', 'CFHTugriz', 'DECam', 'HSC',
    'JWST', 'LSST', 'UBVRIplus', 'UKIDSS', 'WFIRST'
]

class FilterSystem(object):
    """
    Class for calculating filter parameters from throughput curves.
    The parameter definitions are taken from
    `Fukugita et al. (1996) AJ 111, 1748.
    <https://ui.adsabs.harvard.edu/abs/1996AJ....111.1748F/abstract>`_

    Parameters
    ----------
    filter_curve_files : list
        List of the file names of the filter curves.
    filter_names : str
        Names of filters. Must be the same length as ``filter_curve_files``.
    **kwargs
        Optional argumnets for `~numpy.loadtxt`, which is used to load the
        filter curve files.
    """

    def __init__(self, filter_curve_files, filter_names, **kwargs):

        self.filter_names = filter_names
        for fn, name in zip(filter_curve_files, filter_names):
            data = np.loadtxt(fn, **kwargs)
            table = Table(data=data, names=['wave', 'trans'])
            setattr(self, name, table)
            self.filter_names.append(name)

    def _get_trans(self, bandpass):
        """
        Get the filter throughput curve for the given bandpass.
        """
        lam = getattr(self, bandpass)['wave']
        trans = getattr(self, bandpass)['trans']
        return lam, trans

    def effective_throughput(self, bandpass):
        """
        Calculate the effective throughput for the given bandpass.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        theff : float
            Effective effective throughput.
        """
        lam, trans = self._get_trans(bandpass)
        theff = np.trapz(trans, np.log(lam))
        return theff

    def lam_eff(self, bandpass):
        """
        Calculate effective wavelength for the given bandpass.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
            leff : float
                The effective wavelength.
        """
        lam, trans = self._get_trans(bandpass)
        log_leff = np.trapz(np.log(lam) * trans, np.log(lam))
        log_leff /= np.trapz(trans, np.log(lam))
        leff =  np.exp(log_leff) * u.angstrom
        return leff

    def lam_pivot(self, bandpass):
        """
        Calculate pivot wavelength for the given bandpass.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
            lpivot : float
                The pivot wavelength.
        """
        lam, trans = self._get_trans(bandpass)
        lpivot = np.trapz(lam * trans, lam)
        lpivot /= np.trapz(trans, np.log(lam))
        lpivot = np.sqrt(lpivot) * u.angstrom
        return lpivot

    def dlam(self, bandpass):
        """
        Calculate the bandpass width.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
            width : float
                The bandpass width.
        """
        lam, trans = self._get_trans(bandpass)
        norm = np.trapz(trans, lam)
        width = (norm / trans.max()) * u.angstrom
        return width

class ZeroPointConverter(object):
    """
    Help class for converting between AB, ST, and Vega magnitudes & colors.

    Parameters
    ----------
    zpt_table : `~astropy.table.Table`
        MIST zero point table, which can be `downloaded here
        <http://waps.cfa.harvard.edu/MIST/BC_tables/zeropoints.txt>`_.
    """

    def __init__(self, zpt_table):
        for f, system, v_to_st, v_to_ab in zpt_table:
            setattr(self, f, [system, v_to_st, v_to_ab])
        self.zpt_table = zpt_table

    def to_vega(self, bandpass):
        """
        Convert to Vega magnitudes.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        system, v_to_st, v_to_ab = getattr(self, bandpass)
        if system == 'Vega':
            zpt_convert = 0.
        elif system == 'AB':
            zpt_convert = -v_to_ab
        return zpt_convert

    def to_ab(self, bandpass):
        """
        Convert to AB magnitudes.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        system, v_to_st, v_to_ab = getattr(self, bandpass)
        if system == 'AB':
            zpt_convert = 0.
        elif system == 'Vega':
            zpt_convert = v_to_ab
        return zpt_convert

    def to_st(self, bandpass):
        """
        Convert to ST magnitudes.

        Parameters
        ----------
        bandpass : str
            The name of the filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        system, v_to_st, v_to_ab = getattr(self, bandpass)
        if system == 'AB':
            zpt_convert = v_to_st - v_to_ab
        elif system == 'Vega':
            zpt_convert = v_to_st
        return zpt_convert

    def color_to_vega(self, blue, red):
        """
        Convert to Vega colors.

        Parameters
        ----------
        blue : str
            The name of the blue filter.
        red : str
            The name of the red filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        blue_convert = self.to_vega(blue)
        red_convert = self.to_vega(red)
        return blue_convert - red_convert

    def color_to_ab(self, blue, red):
        """
        Convert to AB colors.

        Parameters
        ----------
        blue : str
            The name of the blue filter.
        red : str
            The name of the red filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        blue_convert = self.to_ab(blue)
        red_convert = self.to_ab(red)
        return blue_convert - red_convert

    def color_to_st(self, blue, red):
        """
        Convert to ST colors.

        Parameters
        ----------
        blue : str
            The name of the blue filter.
        red : str
            The name of the red filter.

        Returns
        -------
        zpt_convert : float
            Zero point conversion magnitude.
        """
        blue_convert = self.to_st(blue)
        red_convert = self.to_st(red)
        return blue_convert - red_convert

def get_filter_names(phot_system=None):
    """
    Get MIST photometric systems and filter names.

    Parameters
    ----------
    phot_system : str or None
        The desired photometric system.

    Returns
    -------
    mist_filter_names : list or dict
        If ``phot_system`` given, then a list of filter names is returned. If
        ``phot_system`` is ``None``, then a dict of all photometric systems
        and filters is returned.
    """
    fn = os.path.join(data_dir, 'mist_filter_names.pkl')
    pickle_in = open(fn, 'rb')
    mist_filter_names = pickle.load(pickle_in)
    pickle_in.close()
    if phot_system is not None:
        if type(phot_system) == str:
            phot_system = [phot_system]
        names = []
        for p in phot_system:
            names.extend(mist_filter_names[p])
        mist_filter_names = names
    return mist_filter_names

def get_filter_properties():
    """

    """
    pkl_fn = os.path.join(data_dir, 'filter_properties.pkl')
    with open(pkl_fn, 'rb') as pkl_file:
        data = pickle.load(pkl_file)
        props = Table(data)
        props['dlam'] *= u.angstrom
        props['lam_eff'] *= u.angstrom
    return props

def phot_system_lookup(filter_name=None):
    """
    Lookup the photometric system name associated with a given filter name.

    Parameters
    ----------
    filter_name : str, optional
        Filter name to lookup.

    Returns
    -------
    lookup : dict or str
        If ``filter_name`` is ``None``, a dictionary with the filter names as
        keywords and the photometric systems as values. If ``filter_name`` is
        not ``None``, its photometric system is returned.
    """
    fn = os.path.join(data_dir, 'phot_system_lookup.pkl')
    pickle_in = open(fn, 'rb')
    lookup = pickle.load(pickle_in)
    pickle_in.close()
    if filter_name is not None:
        lookup = lookup[filter_name]
    return lookup

def load_zero_point_converter():
    """
    Create and return a `~artpop.filters.ZeroPointConverter` object.
    """
    from astropy.io import ascii
    fn = os.path.join(data_dir, 'zeropoints.txt')
    table= ascii.read(fn)
    return ZeroPointConverter(table)


import sys
sys.path.append('build')

from MatterSim import Simulator
import math
import cv2
import json
import numpy as np

sim = Simulator()
sim.setCameraResolution(500, 300)
sim.setCameraVFOV(math.radians(60))
sim.setElevationLimits(math.radians(-40),math.radians(50))
sim.initialize()

with open("src/test/rendertest_spec.json") as f:
    spec = json.load(f)
    for tc in spec[:1]:
        sim.newEpisode(tc["scanId"], tc["viewpointId"], tc["heading"], tc["elevation"])
        state = sim.getState()
        im = np.array(state.rgb, copy=False)
        imgfile = tc["reference_image"]
        cv2.imwrite("sim_imgs/"+imgfile, im);
        cv2.imshow('rendering', im)
        cv2.waitKey(0)


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import argparse
import os
import sys
import math
import cv2
import numpy as np
import align.detect_face
import facenet
from facenet import crop
import tensorflow as tf
import pickle
from PIL import Image,ImageDraw,ImageFont
import logging
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s -   %(message)s',
                           datefmt = '%m/%d/%Y %H:%M:%S',
                           level = logging.INFO)
logger = logging.getLogger(__name__)

# 加载opencv人脸识别模型
face_cascade = cv2.CascadeClassifier('D:\opencv\sources\data\haarcascades\haarcascade_frontalface_default.xml')

def opencvDetect(img):
    faces = face_cascade.detectMultiScale(img,1.3,5)#faces返回人脸位置坐标
    if len(faces) > 0 :
        for (x,y,w,h) in faces:
            cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)#画矩形框标记人脸,框大点
    return faces

# 加载mtcnn
minsize = 20  # minimum size of face
threshold = [0.6, 0.7, 0.7]  # three steps's threshold
factor = 0.709  # scale factor
logger.info('加载MTCNN模型')
with tf.Graph().as_default():
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=1.0)
    sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False))
    with sess.as_default():
        pnet, rnet, onet = align.detect_face.create_mtcnn(sess, None)
logger.info('加载完成')

def mtcnnDetect(img):
    bounding_boxes, _ = align.detect_face.detect_face(img, minsize, pnet, rnet, onet, threshold, factor)
    return bounding_boxes

def main(args):
    # 加载svm
    classifier_filename_exp = os.path.expanduser(args.classifier_filename)
    logger.info('加载分类器"%s"' % classifier_filename_exp)
    with open(classifier_filename_exp, 'rb') as infile:
        (classifier_model, class_names) = pickle.load(infile)
        logger.info('分类器识别类别有:{}'.format(class_names))
    logger.info('加载完成')

    with tf.Graph().as_default():
        with tf.Session() as sess:
            logger.info('加载Facenet进行人脸识别')
            facenet.load_model(args.model)
            # Get input and output tensors
            images_placeholder = tf.get_default_graph().get_tensor_by_name("input:0")
            embeddings = tf.get_default_graph().get_tensor_by_name("embeddings:0")
            phase_train_placeholder = tf.get_default_graph().get_tensor_by_name("phase_train:0")
            embedding_size = embeddings.get_shape()[1]
            logger.info('加载完成')
            cap = cv2.VideoCapture(0)
            num = 1
            while(1):
                key = cv2.waitKey(1)
                # get a frame
                ret, frame = cap.read()
                if(num%args.interval == 0):
                    logger.info('识别第%d帧' %num)
                    if frame.ndim == 3:
                        gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
                    else:
                        gray = frame
                    if gray.ndim == 2:
                        img = facenet.to_rgb(gray)

                    bounding_boxes = mtcnnDetect(img)
                    nrof_faces = bounding_boxes.shape[0]  # number of faces
                    crop_faces=np.zeros((len(bounding_boxes), args.image_size, args.image_size, 3))
                    label=[]
                    if nrof_faces < 1:
                        logger.info('未检测到人脸')
                        pass
                    else:
                        logger.info('检测到%d个人脸'%nrof_faces)
                        for i,face_position in enumerate(bounding_boxes):
                            face_position = face_position.astype(int)
                            cv2.rectangle(frame, (face_position[0],
                                                  face_position[1]),
                                          (face_position[2], face_position[3]),
                                          (0, 255, 0), 2)
                            #将图像reshape成输入进facenet模型的结构
                            face = frame[face_position[1]:face_position[3],face_position[0]:face_position[2], ]
                            face = cv2.resize(face, (args.image_size, args.image_size), interpolation=cv2.INTER_CUBIC)#crop shape 96*96*3
                            data = face.reshape(-1, args.image_size, args.image_size, 3) #reshape image
                            feed_dict = {images_placeholder: data, phase_train_placeholder: False}
                            #得到图像embedding
                            emb_array = sess.run(embeddings,feed_dict=feed_dict)
                            #embedding输入到svm分类
                            predictions = classifier_model.predict_proba(emb_array)
                            best_class_indices = np.argmax(predictions, axis=1)
                            best_class_probabilities = predictions[np.arange(len(best_class_indices)), best_class_indices]
                            for i in range(len(best_class_indices)):
                                label.append(class_names[best_class_indices[i]])
                            results = set(zip(label,best_class_probabilities))
                            logger.info('识别出:{}'.format(results))

                            cv2.putText(frame, 'detected:{}'.format(results), (50, 100),
                                        cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (255, 0, 0),
                                        thickness=2, lineType=2)

                cv2.imshow("capture", frame)
                num += 1 #帧数+1
                # 键盘输入q则退出,ord返回ascii码
                if key == ord('q'):
                    break

    cap.release()
    cv2.destroyAllWindows()

def parse_arguments(argv):
    parser = argparse.ArgumentParser()
    parser.add_argument('classifier_filename',
                        help='Classifier model file name as a pickle (.pkl) file. ' +
                             'For training this is the output and for classification this is an input.')
    parser.add_argument('model', type=str,
                        help='Could be either a directory containing the meta_file and ckpt_file or a model protobuf (.pb) file')
    parser.add_argument('--interval',type=int,help='Frame by second',default=1)

    parser.add_argument('--image_size', type=int,
                        help='Image size (height, width) in pixels.', default=160)

    return parser.parse_args(argv)

if __name__ == '__main__':
    main(parse_arguments(sys.argv[1:]))


import base64
import io
import zmq
import cv2
import os
import numpy as np
import tensorflow as tf

class CarDetector(object):
    def __init__(self, frozen_graph):

#Tensorflow localization/detection model
# Single-shot-dectection with mobile net architecture trained on COCO dataset
# setup tensorflow graph

        self.detection_graph = tf.Graph()

# configuration for possible GPU use
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

# load frozen tensorflow detection model and initialize
# the tensorflow graph

        with tf.gfile.GFile(frozen_graph, "rb") as f:
            self.graph_def = tf.GraphDef()
            self.graph_def.ParseFromString(f.read())

# import the graph_def into a new Graph and returns it
        with self.detection_graph.as_default():
            tf.import_graph_def(self.graph_def, name="")

            self.sess = tf.Session(graph=self.detection_graph, config=config)
            self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
# Each box represents a part of the image where a particular object was detected.
            self.boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represent how level of confidence for each of the objects.
# Score is shown on the result image, together with the class label.
            self.scores =self.detection_graph.get_tensor_by_name('detection_scores:0')
            self.classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
            self.num_detections =self.detection_graph.get_tensor_by_name('num_detections:0')
        self.image_width = 640
        self.image_height = 480

    def detect(self, image):
# Detect the results without output
        """
        Args:
            image: camera image

        Returns:
            list of bounding boxes: coordinates [y_up, x_left, y_down, x_right]

        """
        self.image_width = image.shape[0]
        self.image_height = image.shape[1]

        with self.detection_graph.as_default():
            image_expanded = np.expand_dims(image, axis=0)
            (boxes, scores, classes, num_detections) = self.sess.run(
                [self.boxes, self.scores, self.classes, self.num_detections],
                feed_dict={self.image_tensor: image_expanded})

            self.out_boxes = np.squeeze(boxes)
            self.out_classes = np.squeeze(classes).astype("int")
            self.out_scores = np.squeeze(scores)
            self.out_number = len(self.out_classes)

    def box_isvalid(self, box, score, category, confidence = 0.3, size_limit = 0.5, ratio_limit = 0.1):
# Examine whether a detected box is valid (good confidence, reasonable shape and size, etc).
        h = box[2] - box[0]
        w = box[3] - box[1]
        ratio = h/(w+0.01)
        return ratio>ratio_limit and ratio<1/ratio_limit and h<size_limit and w<size_limit and category<9 and score>confidence

    def output(self, confidence =0.3, quiet = True):
# Examine the detected results and select the desired output
        self.car_boxes = []
        category_index={1: {'id': 1, 'name': u'person'},
                        2: {'id': 2, 'name': u'bicycle'},
                        3: {'id': 3, 'name': u'car'},
                        4: {'id': 4, 'name': u'motorcycle'},
                        5: {'id': 5, 'name': u'airplane'},
                        6: {'id': 6, 'name': u'bus'},
                        7: {'id': 7, 'name': u'train'},
                        8: {'id': 8, 'name': u'truck'},
                        9: {'id': 9, 'name': u'boat'},
                        10: {'id': 10, 'name': u'traffic light'},
                        11: {'id': 11, 'name': u'fire hydrant'},
                        13: {'id': 13, 'name': u'stop sign'},
                        14: {'id': 14, 'name': u'parking meter'}}
# The ID for car in COCO data set is 3
#idx_vec = [i for i, v in enumerate(cls) if ((v==3) and (scores[i]>0.3))]
        for i in range(self.out_number):
            category = self.out_classes[i]
            score = self.out_scores[i]
            box = self.out_boxes[i]
            if self.box_isvalid(box, score, category):
                box[0] = self.image_width*box[0]
                box[2] = self.image_width*box[2]
                box[1] = self.image_height*box[1]
                box[3] = self.image_height*box[3]
                box = box.astype("int")
                self.car_boxes.append([box[0],box[1],box[2],box[3],score,category])
                if not quiet:
                    print(box, ', confidence: ', self.out_scores[i], ', category: ',category_index[category]['name'])
        if len(self.car_boxes) ==0 and not quiet:
            print('no detection!')

        return self.car_boxes

    def output_as_str(self, confidence =0.3, quiet = True):
# As title - output the results as string, to be sent as ZeroMQ message
        self.output(confidence = confidence, quiet = quiet)
        return [[str(item) for item in box] for box in self.car_boxes]

def DL_init(frozen_graph = "/home/ubuntu/frozen_inference_graph.pb", dummy = cv2.imread("/home/ubuntu/dummy.jpg")):
# Initialize the SSD detector to save time.
    det = CarDetector(frozen_graph)
    det.detect(dummy)
    return det

def overlay_boxes(boxes,image):
# Overlay the detected box on the image. It directly modifies the raw images
# so it should only be used for showing the graph.
    for box in boxes:
        x1, y1, x2, y2 = box[:4]
        x1, y1, x2, y2 = int(x1),int(y1),int(x2),int(y2)
        image[x1:x1+4,y1:y2+1,:] =0
        image[x2:x2+4,y1:y2+1,:] =0
        image[x1:x2+1,y1:y1+4,:] =0
        image[x1:x2+1,y2:y2+4,:] = 0
        image[x1:x1+3,y1:y2+1,0] = 255
        image[x2:x2+3,y1:y2+1,0] = 255
        image[x1:x2+1,y1:y1+3,0] = 255
        image[x1:x2+1,y2:y2+3,0] = 255

def worker():
# Ingest the data from ZeroMQ queue and apply SSD model. The results will be sent as meta data.
    producer_ip = "tcp://10.0.0.9:5557"
    consumer_ip = "tcp:/10.0.0.14:5558"
    context = zmq.Context()
    detector = DL_init()

    receiver = context.socket(zmq.PULL)
    receiver.connect(producer_ip)
    sender = context.socket(zmq.PUSH)
    sender.connect(consumer_ip)
    print('Ingest images from ', producer_ip)
    print('Send detection results to ', consumer_ip)

    # recieve work

    while True:
        meta_data = receiver.recv_json()
        message = receiver.recv()

        b64_string = base64.b64decode(message)
        img = cv2.imdecode(np.frombuffer(b64_string, np.uint8),cv2.IMREAD_ANYCOLOR)
        detector.detect(img)
        boxes = detector.output_as_str()

        new_meta = {'camera' : meta_data['camera'],
                    'counter':meta_data['counter'],
                    'time': meta_data['time'],
                    'boxes': boxes}

        sender.send_json(new_meta, flags=zmq.SNDMORE)

        overlay_boxes(boxes,img)
        img_bytes = base64.b64encode(cv2.imencode('.jpg',img)[1].tobytes())
# Draw the detected box onto the image to save the time on the front end.

        sender.send(img_bytes)
        time.sleep(.01)

if __name__=="__main__":
    worker()


### General imports ###
from __future__ import division
import numpy as np
import pandas as pd
import time
from time import sleep
import re
import os
import requests
import argparse
from collections import OrderedDict

### Image processing ###
import cv2
from scipy.ndimage import zoom
from scipy.spatial import distance
import imutils
from scipy import ndimage
import dlib
from imutils import face_utils

### Model ###
from tensorflow.keras.models import load_model
from tensorflow.keras import backend as K

################################################################################
########## Change this depending on where your recordings are located ##########
################################################################################
rec_dir = 'recordings/'
################################################################################

def getVideoEmotions(input_video_filepath,output_name):
    # Start video capute. 0 = Webcam, 1 = Video file, -1 = Webcam for Web
    video_capture = cv2.VideoCapture(input_video_filepath)

    # Image shape
    shape_x = 48
    shape_y = 48
    input_shape = (shape_x, shape_y, 1)

    # We have 7 emotions
    nClasses = 7

    # Timer until the end of the recording
    curr_frame_num = 0

    # Initiate Landmarks
    (lStart, lEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eye"]
    (rStart, rEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eye"]

    (nStart, nEnd) = face_utils.FACIAL_LANDMARKS_IDXS["nose"]
    (mStart, mEnd) = face_utils.FACIAL_LANDMARKS_IDXS["mouth"]
    (jStart, jEnd) = face_utils.FACIAL_LANDMARKS_IDXS["jaw"]

    (eblStart, eblEnd) = face_utils.FACIAL_LANDMARKS_IDXS["left_eyebrow"]
    (ebrStart, ebrEnd) = face_utils.FACIAL_LANDMARKS_IDXS["right_eyebrow"]

    # Load the pre-trained X-Ception model
    model = load_model('Models/video.h5')

    # Load the face detector
    face_detect = dlib.get_frontal_face_detector()

    # Load the facial landmarks predictor
    predictor_landmarks  = dlib.shape_predictor("Models/face_landmarks.dat")

    # Prediction vector
    predictions = []

    # Timer for length of video
    max_frame_num = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = video_capture.get(cv2.CAP_PROP_FPS)
    frames_btwn = int(fps/3) # measure emotion every 1/3 of a sec

    # angry_0 = []
    # disgust_1 = []
    # fear_2 = []
    # happy_3 = []
    # sad_4 = []
    # surprise_5 = []
    # neutral_6 = []

    # Initialize arrays for saving predictions
    emotions = []
    face_indices = []
    timestamps = []

    # Analyze video until the end
    curr_frame_num = 0
    iter_percent = 0 # for printing
    while curr_frame_num < max_frame_num:
        # Set the frame to be read
        video_capture.set(cv2.CAP_PROP_POS_FRAMES, curr_frame_num)

        # Capture the frame number set above (frame here means image)
        ret, frame = video_capture.read()

        # Face index, face by face
        face_index = 0

        # Image to gray scale
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

        # All faces detected
        rects = face_detect(gray, 1)

        # For each detected face
        for (i, rect) in enumerate(rects):
            # Identify face coordinates
            (x, y, w, h) = face_utils.rect_to_bb(rect)
            face = gray[y:y+h,x:x+w]

            # Zoom on extracted face (if face extracted)
            if face.shape[0] and face.shape[1]:
                face = zoom(face, (shape_x / face.shape[0],shape_y / face.shape[1]))

                # Cast type float
                face = face.astype(np.float32)

                # Scale the face
                face /= float(face.max())
                face = np.reshape(face.flatten(), (1, 48, 48, 1))

                # Make Emotion prediction on the face, outputs probabilities
                prediction = model.predict(face)

                # Most likely emotion (as an index)
                prediction_result = np.argmax(prediction)

                # Convert emotion index to an emotion (string)
                emotion = ''
                if (prediction_result == 0):
                    emotion = 'Angry'
                elif (prediction_result == 1):
                    emotion = 'Disgust'
                elif (prediction_result == 2):
                    emotion = 'Fear'
                elif (prediction_result == 3):
                    emotion = 'Happy'
                elif (prediction_result == 4):
                    emotion = 'Sad'
                elif (prediction_result == 5):
                    emotion = 'Surprise'
                elif (prediction_result == 6):
                    emotion = 'Neutral'
                else:
                    emotion = 'Unknown emotion'
            else:
                emotion = 'No face extracted'
            # save results for later
            emotions.append(emotion)
            face_indices.append(face_index)
            timestamps.append(curr_frame_num/fps)

        # every so often, show percent done
        percent_done = curr_frame_num/max_frame_num*100
        if (percent_done > iter_percent):
            print('current frame: %.0f' % curr_frame_num)
            print('percent done: %.1f%%' % percent_done)
            iter_percent += 20

        # increment frame
        curr_frame_num += frames_btwn

    video_capture.release()

    # # Export predicted emotions to .csv format
    df = pd.DataFrame({'EMOTION': emotions, 'FACE_INDEX': face_indices, 'TIMESTAMP_SEC': timestamps})
    df.to_csv(os.path.join('output', 'video_emotions_' + output_name + '.csv'), sep=',', index=False)

    print('🎉 Done! 🎉')
    print('See the output file:')
    print('output/' + 'video_emotions_' + output_name + '.csv')

if __name__ == '__main__':
    # Loop through specific files and analyze their video
    files_in_dir = [f for f in os.listdir(rec_dir) if os.path.isfile(os.path.join(rec_dir, f))]
    i = 0
    for f in files_in_dir:
        if f.split('.')[1] == 'avi' or f.split('.')[1] == 'mp4':
            input_video_filepath = os.path.join(rec_dir,f)
            print(f'Reading from {input_video_filepath}')
            output_name = f.split('.')[0]

            getVideoEmotions(input_video_filepath, output_name)

        i += 1
        print(f"""Number of files to go: {len(files_in_dir) - i}
            Percent files done: {i/len(files_in_dir)*100}\n""")


import numpy as np
import theano.tensor as T
import lasagne
from lasagne.layers import InputLayer, DropoutLayer, ReshapeLayer, \
        NonlinearityLayer, DimshuffleLayer, ConcatLayer
from lasagne.layers import batch_norm, BatchNormLayer
from lasagne.layers import Pool1DLayer as PoolLayer
from lasagne.layers import Conv1DLayer as ConvLayer
from lasagne.layers import Upscale1DLayer as UpscaleLayer
from lasagne.layers import PadLayer, DenseLayer
from lasagne.layers import ElemwiseSumLayer, ElemwiseMergeLayer
from lasagne.nonlinearities import softmax, linear, rectify

def conv_bn_relu(net, incoming_layer, depth, num_filters, filter_size,
                    pad = 'same', name=''):
    net[name+'conv'+str(depth)] = ConvLayer(net[incoming_layer],
                num_filters = num_filters,
                filter_size = filter_size,
                pad = pad,
                nonlinearity=None)
    net[name+'bn'+str(depth)] = BatchNormLayer(net['conv'+str(depth)])
    net[name+'relu'+str(depth)] = NonlinearityLayer( net['bn'+str(depth)],
                                    nonlinearity = rectify)
    incoming_layer = name+'relu'+str(depth)

    return incoming_layer

def bn_relu_conv(net, incoming_layer, depth, num_filters, filter_size,
                    pad = 'same', name=''):

    net[name+'bn'+str(depth)] = BatchNormLayer(net[incoming_layer])
    net[name+'relu'+str(depth)] = NonlinearityLayer( net['bn'+str(depth)],
                                    nonlinearity = rectify)
    net[name+'conv'+str(depth)] = ConvLayer(net['relu'+str(depth)],
                num_filters = num_filters,
                filter_size = filter_size,
                pad = pad,
                nonlinearity=None)
    incoming_layer = name+'conv'+str(depth)

    return incoming_layer

def build_region_model(input_var,
        filter_size=[25],
        n_filters = 64,
        n_classes = 41,
        poolings = [0,2,0,2,0],
        hidden_size = 512,
        n_hidden = 2,
        last_filter_size = 1,
        nb_in_channels = 1,
        block = 'bn_relu_conv',
        out_nonlin = softmax):
    '''
    Parameters:
    -----------
    input_var : theano tensor
    filter_size : list of odd int (to fit with same padding),
                size of filter_size list determines the number of
                convLayer to Concatenate
    n_filters : int, number of filters for each convLayer
    n_classes : int
    poolings : list of int, represents the structure of the network, each
            number corresponds to a layer of the model, indicating the pooling
            stride. depth = len(poolings)
    hidden_size: number of units in fully connected layers
    n_hidden: number of fully connected layers
    last_filter_size : int, must be set to 1 (the older version had
            a last_filter_size of 3, that was an error
            the argument is there to be able to reassign weights correctly
            when testing)
    out_nonlin : default=softmax, non linearity function
    '''

    net = {}

    net['input'] = InputLayer((None, nb_in_channels, 200), input_var)
    incoming_layer = 'input'

    depth = len(poolings)

    for d in range(depth):
        if block == 'bn_relu_conv':
            incoming_layer = bn_relu_conv(net, incoming_layer, depth = d,
                            num_filters= n_filters, filter_size=filter_size[0])
        elif block == 'conv_bn_relu':
            incoming_layer = conv_bn_relu(net, incoming_layer, depth = d,
                            num_filters= n_filters, filter_size=filter_size[0])

        if poolings[d]:
            net['pool'+str(depth)] = PoolLayer(net[incoming_layer], poolings[d])
            incoming_layer = 'pool'+str(depth)

    for i in range(n_hidden):
        name = 'dense_hidden'+str(i)
        net[name] = DenseLayer(net[incoming_layer], num_units=hidden_size)
        incoming_layer = name

    #Output layer
    net['dense_out'] = DenseLayer(net[incoming_layer], num_units=n_classes)
    incoming_layer = 'dense_out'

    net['last_layer'] = NonlinearityLayer(net[incoming_layer],
                    nonlinearity = out_nonlin)
    incoming_layer = 'last_layer'

    return [net[l] for l in ['last_layer']], net

def build_simple_join_model(input_var,
        filter_size=[25],
        n_filters = 64,
        n_filters_reg = 64,
        n_classes = 6,
        n_regions = 6,
        depth = 1,
        region_depth = 3,
        poolings = [0,2,0,2,0],
        hidden_size = 512,
        n_hidden = 2,
        last_filter_size = 1,
        nb_in_channels = 1,
        block = 'bn_relu_conv',
        #bn_relu_conv = False, #unused for now
        out_nonlin = softmax):
    '''
    Parameters:
    -----------
    input_var : theano tensor
    filter_size : list of odd int (to fit with same padding),
                size of filter_size list determines the number of
                convLayer to Concatenate
    n_filters : int, number of filters for each convLayer
    n_filters_reg : int, number of filters for each convLayer in the region
                classification network
    n_classes : int
    depth : int, number of stacked convolution before concatenation
    region_depth: int, index of layer that the region classification network
                will take as input.
    poolings : list of int, represents the structure of the region
            classification network, each number corresponds to a layer of the
            model, indicating the pooling stride. depth = len(poolings)
    hidden_size: number of units in fully connected layers of the class model
    n_hidden: number of fully connected layers in the class model
    last_filter_size : int, must be set to 1 (the older version had
            a last_filter_size of 3, that was an error
            the argument is there to be able to reassign weights correctly
            when testing)
    out_nonlin : default=softmax, non linearity function
    '''

    net = {}

    net['input'] = InputLayer((None, nb_in_channels, 200), input_var)
    incoming_layer = 'input'
    #incoming_layer = 'input'

    #Convolution layers
    for d in range(depth):
        if block == 'bn_relu_conv':
            incoming_layer = bn_relu_conv(net, incoming_layer, depth = d,
                            num_filters= n_filters, filter_size=filter_size[0])

        elif block == 'conv_bn_relu':
            incoming_layer = conv_bn_relu(net, incoming_layer, depth = d,
                            num_filters= n_filters, filter_size=filter_size[0])

        if d == region_depth:
            # this is the layer that the region classification network will
            # take as input
            region_incoming_layer = incoming_layer

    #Output layer
    net['final_conv'] = ConvLayer(net[incoming_layer],
                    num_filters = n_classes,
                    filter_size = last_filter_size,
                    pad='same')
    incoming_layer = 'final_conv'

    #DimshuffleLayer and ReshapeLayer to fit the softmax implementation
    #(it needs a 1D or 2D tensor, not a 3D tensor)
    net['final_dimshuffle'] = DimshuffleLayer(net[incoming_layer], (0,2,1))
    incoming_layer = 'final_dimshuffle'

    layerSize = lasagne.layers.get_output(net[incoming_layer]).shape
    net['final_reshape'] = ReshapeLayer(net[incoming_layer],
                                (T.prod(layerSize[0:2]),layerSize[2]))
                                # (200*batch_size,n_classes))
    incoming_layer = 'final_reshape'

    #This is the layer that computes the prediction
    net['last_layer'] = NonlinearityLayer(net[incoming_layer],
                    nonlinearity = out_nonlin)
    incoming_layer = 'last_layer'

    #Layers needed to visualize the prediction of the network
    net['probs_reshape'] = ReshapeLayer(net[incoming_layer],
                    (layerSize[0], layerSize[1], n_classes))
    incoming_layer = 'probs_reshape'

    net['probs_dimshuffle'] = DimshuffleLayer(net[incoming_layer], (0,2,1))

    # Region classifier net starting from a given layer in the segmentation
    # network stored in region_incoming_layer
    for d in range(len(poolings)):
        if block == 'bn_relu_conv':
            region_incoming_layer = bn_relu_conv(net, region_incoming_layer, depth = d,
                            num_filters= n_filters_reg, filter_size=filter_size[0], name='reg_')
        elif block == 'conv_bn_relu':
            region_incoming_layer = conv_bn_relu(net, region_incoming_layer, depth = d,
                            num_filters= n_filters_reg, filter_size=filter_size[0], name='reg_')

        if poolings[d]:
            net['pool'+str(depth)] = PoolLayer(net[region_incoming_layer], poolings[d])
            region_incoming_layer = 'pool'+str(depth)

    for i in range(n_hidden):
        name = 'dense_hidden'+str(i)
        net[name] = DenseLayer(net[region_incoming_layer], num_units=hidden_size)
        region_incoming_layer = name

    #Output layer
    net['dense_out'] = DenseLayer(net[region_incoming_layer], num_units=n_regions)
    region_incoming_layer = 'dense_out'

    net['last_layer_reg'] = NonlinearityLayer(net[region_incoming_layer],
                    nonlinearity = out_nonlin)
    region_incoming_layer = 'last_layer_reg'

    # [net[l] for l in ['last_layer']] : used to directly compute the output
    #                       of the network
    # net : dictionary containing each layer {name : Layer instance}
    return [net[l] for l in ['last_layer', 'last_layer_reg']], net


import numpy as np
import matplotlib
# matplotlib.use("pdf")
import matplotlib.pyplot as plt
from pygimli.viewer.mpl import underlayBKGMap
from saem import CSEMData
from pyproj import Proj

txpos = np.array([[582000, 584000],
                  [5744000, 5742000]]).T
ground = 400  #  altitude of ground
# %% read in the old data and take every second data point (12->24m)
self = CSEMData(datafile="L*.mat", txPos=txpos, txalt=ground)
self.showField("alt")
self.generateDataPDF()
sdfsdfsdf
# self.filter(every=2)
# %%
# from hdf5storage import loadmat
from scipy.io import loadmat

matfile = "L1_Source_LIAGMUN_Tx02_Ncyc16_ltsregress.mat"
bla = loadmat(matfile)["ztfs"][0][0]
# %%
matfile = matfile.replace("L1", "L2")
bla2 = loadmat(matfile)["ztfs"][0][0]
# %%
for i in range(20):
    bla[i] = np.concatenate((bla[i], bla2[i]), axis=-1)
# %%
sdfsdfsdf
# %%
f = 1./bla[12]
# time = bla[16]
gx, gy, gz = bla[17]
utm = Proj(proj='utm', zone=32, ellps='WGS84')
ggx, ggy = utm(gx, gy)
z = bla[18][0]
y, x = bla[19]
DATA = np.squeeze(bla[13])
ERR = np.squeeze(bla[14])
DN = np.squeeze(bla[15])
# %%
fig, ax = plt.subplots()
ax.plot(x, y, ".-")  # , ggx, ggy, ".-")
ax.set_aspect(1.0)
xl = ax.get_xlim()
ax.set_xlim(xl[0]-2100, xl[1]+2100)
underlayBKGMap(ax, mode="DOP")  # , uuid='139409f9-5953-1cf5-f19c-75262cfa13d6')
#fig.savefig("L1.pdf", bbox_inches="tight", dip=300)
fig.savefig("L1.pdf", dpi=200)
# %%
fig, ax = plt.subplots()
ax.plot(y, gz, "x-", y, bla[18][0], "+-")


import math
import os
from collections import OrderedDict
from pathlib import Path
from itertools import permutations
import math

import cv2
import numpy as np
from numpy.core.fromnumeric import amin

import craft_text_detector.file_utils as file_utils
import craft_text_detector.torch_utils as torch_utils

CRAFT_GDRIVE_URL = "https://drive.google.com/uc?id=1bupFXqT-VU6Jjeul13XP7yx2Sg5IHr4J"
REFINENET_GDRIVE_URL = (
    "https://drive.google.com/uc?id=1xcE9qpJXp4ofINwXWVhhQIh9S8Z7cuGj"
)

# unwarp corodinates
def warpCoord(Minv, pt):
    out = np.matmul(Minv, (pt[0], pt[1], 1))
    return np.array([out[0] / out[2], out[1] / out[2]])

def copyStateDict(state_dict):
    if list(state_dict.keys())[0].startswith("module"):
        start_idx = 1
    else:
        start_idx = 0
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = ".".join(k.split(".")[start_idx:])
        new_state_dict[name] = v
    return new_state_dict

def load_craftnet_model(cuda: bool = False):
    '''
    Loads and returns the craftnet model with the state value as defined in the
    state dictionary
    Args:
        cuda (bool): if True, to use GPU for compute else cpu
    Returns
        A pytorch model object
    '''
    # get craft net path
    home_path = str(Path.home())
    weight_path = os.path.join(
        home_path, ".craft_text_detector", "weights", "craft_mlt_25k.pth"
    )
    # load craft net
    from craft_text_detector.models.craftnet import CraftNet

    craft_net = CraftNet()  # initialize

    # check if weights are already downloaded, if not download
    url = CRAFT_GDRIVE_URL
    if os.path.isfile(weight_path) is not True:
        print("Craft text detector weight will be downloaded to {}".format(weight_path))

        file_utils.download(url=url, save_path=weight_path)

    # arange device
    if cuda:
        craft_net.load_state_dict(copyStateDict(torch_utils.load(weight_path)))

        craft_net = craft_net.cuda()
        craft_net = torch_utils.DataParallel(craft_net)
        torch_utils.cudnn_benchmark = False
    else:
        craft_net.load_state_dict(
            copyStateDict(torch_utils.load(weight_path, map_location="cpu"))
        )
    craft_net.eval()
    return craft_net

def load_refinenet_model(cuda: bool = False):
    '''
    Load a refine net model, the refine net is used to make better
    heatmaps based on the outputs and intermediate features of craft.
    It is also trained using the same semisupervised objective function.
    Args:
        cuda (bool): if True, to use GPU for compute else cpu
    Returns
        A pytorch model object
    '''
    # get refine net path
    home_path = str(Path.home())
    weight_path = os.path.join(
        home_path, ".craft_text_detector", "weights", "craft_refiner_CTW1500.pth"
    )
    # load refine net
    from craft_text_detector.models.refinenet import RefineNet

    refine_net = RefineNet()  # initialize

    # check if weights are already downloaded, if not download
    url = REFINENET_GDRIVE_URL
    if os.path.isfile(weight_path) is not True:
        print("Craft text refiner weight will be downloaded to {}".format(weight_path))

        file_utils.download(url=url, save_path=weight_path)

    # arange device
    if cuda:
        refine_net.load_state_dict(copyStateDict(torch_utils.load(weight_path)))

        refine_net = refine_net.cuda()
        refine_net = torch_utils.DataParallel(refine_net)
        torch_utils.cudnn_benchmark = False
    else:
        refine_net.load_state_dict(
            copyStateDict(torch_utils.load(weight_path, map_location="cpu"))
        )
    refine_net.eval()
    return refine_net

def getDetBoxes_core(textmap, linkmap, text_threshold, link_threshold, low_text):
    '''
    Args:
        text_map: a 2d np float array containing the probability distribution of a
            pixel being a character region
        linkmap: a 2d np float array containing the probability distribution of a
            pixel being an inter character spacing
        text_threshold: text confidence threshold
        link_threshold: link confidence threshold
        low_text: text low-bound score
    Output:
        det: a list of predicted boxes
        labels: a list of keys of each signle line word segmentation
        mapper: a map between the detected boxes and all possible labels in the heatmaps
        num_characters: estimate of number of characters in each single line text crop
        avg_character_sizes: average size of characters in the text crops
        word_id: next text crop adjacency vector, -1 when for a given crop index has no consecutive
            multiline text element, equal to the next line text crop index when it has one.

    '''
    def compare_boxes(box1, box2, slope_thresh = 0.1, min_edge_distance = 10):
        '''
        Given 2 4X2 numpy array containing [[x,y]..] points
        Returns True if they are next to each other else False
        Args:
            box1 (np.ndarray): a 4X2 nummy float32 array containing the 4 coordinates in
                TL->TR->BR->BL sequence in (x,y) order, This box is assumed to be the one on top.
            box2 (np.ndarray): a 4X2 nummy float32 array containing the 4 coordinates in
                TL->TR->BR->BL sequence in (x,y) order, This box is assumed to be the one at the bottom.
            slope_thresh (float): minimum difference between any 2 edges to be considered parallell
            min_edge_distance (float): minimum distance between the mid points of any 2 edges to be
                considered adjacent/overlapping/"intersecting"

        '''
        boxcentroid1 = np.mean(box1, axis=0)
        boxcentroid2 = np.mean(box2, axis=0)

        # check if box1 is above box2 spatially
        if boxcentroid1[1] < boxcentroid2[1]:
            for i, j in permutations(range(4), 2):
                for k,l in permutations(range(4), 2):
                    # compute slope and mid points of any 2 edges
                    #  and compare their differences
                    slope1 = (box1[i, 1]-box1[j, 1])/(box2[i, 0]-box2[j, 0]\
                        if (box2[i, 0]-box2[j, 0]) != 0 else 0.001)
                    slope2 = (box1[k, 1]-box1[l, 1])/(box2[k, 0]-box2[l, 0]\
                        if (box2[k, 0]-box2[l, 0]) != 0 else 0.001)

                    linecentroid1 = (box1[i,:] + box1[j,:])/2
                    linecentroid2 = (box2[j,:] + box2[k,:])/2

                    if abs(slope1 - slope2) < slope_thresh and\
                        np.sum(abs((linecentroid1 - linecentroid2))) < min_edge_distance:
                        return True
        else:
            return False

    def d(cord1, cord2):
        return np.linalg.norm(cord1-cord2)

    def sort_box(box):
        '''
        Returns the bounding box sorted with BL->BR->TR->TL sequence
        given a 2d 4X2 list of vertices of a cyclic bbox
        Args (np.ndarray): a 4X2 np float32 array with the coordinated
            of the text quadrilateral in cyclic order
        Returns:
            a 4X2 np float32 array with the coordinates sorted in
            TL->TR->BR->BL sequence
        '''
        aymin = np.argmin(box[:,1])
        tl=0
        if d(box[aymin,:], box[(aymin+1)%4,:]) > d(box[(aymin+2)%4, :], box[(aymin+1)%4,:]):
            # aymin is top left
            tl = np.argmin(box[:,1])
        else:
            tl = (np.argmin(box[:,1]) -1)%4
            # axmin is top right
        box = box[tl:,:].tolist() + box[:tl,:].tolist()
        return np.array(box).astype(np.float32)

    # prepare data
    linkmap = linkmap.copy()
    textmap = textmap.copy()
    img_h, img_w = textmap.shape

    """ labeling method """
    ret, text_score = cv2.threshold(textmap, low_text, 1, 0)
    ret, link_score = cv2.threshold(linkmap, link_threshold, 1, 0)

    text_score_comb = np.clip(text_score + link_score, 0, 1)
    nLabels, labels, stats, centroids = cv2.connectedComponentsWithStats(
        text_score_comb.astype(np.uint8), connectivity=4)

    det = []
    mapper = []
    num_characters = []
    avg_character_sizes = []
    for k in range(1, nLabels):
        # size filtering
        size = stats[k, cv2.CC_STAT_AREA]
        if size < 10:
            continue

        # thresholding
        if np.max(textmap[labels == k]) < text_threshold:
            continue

        # make segmentation map
        segmap = np.zeros(textmap.shape, dtype=np.uint8)
        segmap[labels == k] = 255

        # remove link area
        segmap[np.logical_and(link_score == 1, text_score == 0)] = 0

        # count number of characters in word segementation
        word_chars, word_label, word_stats, word_centroid = cv2.connectedComponentsWithStats(
            segmap.astype(np.uint8), connectivity=4)
        num_characters.append(word_chars)

        avg_character_size = np.mean(word_stats[:, 4])
        avg_character_sizes.append(avg_character_size)

        x, y = stats[k, cv2.CC_STAT_LEFT], stats[k, cv2.CC_STAT_TOP]
        w, h = stats[k, cv2.CC_STAT_WIDTH], stats[k, cv2.CC_STAT_HEIGHT]
        niter = int(math.sqrt(size * min(w, h) / (w * h)) * 2)
        sx, ex, sy, ey = (x - niter, x + w + niter + 1, y - niter, y + h + niter + 1)
        # boundary check
        if sx < 0:
            sx = 0
        if sy < 0:
            sy = 0
        if ex >= img_w:
            ex = img_w
        if ey >= img_h:
            ey = img_h
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1 + niter, 1 + niter))
        segmap[sy:ey, sx:ex] = cv2.dilate(segmap[sy:ey, sx:ex], kernel)

        # make box
        np_temp = np.roll(np.array(np.where(segmap != 0)), 1, axis=0)
        np_contours = np_temp.transpose().reshape(-1, 2)
        rectangle = cv2.minAreaRect(np_contours)
        box = cv2.boxPoints(rectangle)

        # align diamond-shape
        w, h = np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[1] - box[2])
        box_ratio = max(w, h) / (min(w, h) + 1e-5)
        if abs(1 - box_ratio) <= 0.1:
            l, r = min(np_contours[:, 0]), max(np_contours[:, 0])
            t, b = min(np_contours[:, 1]), max(np_contours[:, 1])
            box = np.array([[l, t], [r, t], [r, b], [l, b]], dtype=np.float32)

        # make clock-wise order
        startidx = box.sum(axis=1).argmin()
        box = np.roll(box, 4 - startidx, 0)
        box = np.array(box)

        det.append(sort_box(box))
        mapper.append(k)

    word_id = [-1]*len(det)
    for id1, (box1, size1) in enumerate(zip(det, avg_character_sizes)):
        for id2, (box2, size2) in enumerate(zip(det, avg_character_sizes)):
            if id1 == id2:
                pass
            else:
                if compare_boxes(box1=box1, box2=box2, min_edge_distance=30):
                    word_id[id1] = id2
    return det, labels, mapper, num_characters, avg_character_sizes, word_id

def getPoly_core(boxes, labels, mapper, linkmap):
    # configs
    num_cp = 5
    max_len_ratio = 0.7
    expand_ratio = 1.45
    max_r = 2.0
    step_r = 0.2

    polys = []
    for k, box in enumerate(boxes):
        # size filter for small instance
        w, h = (
            int(np.linalg.norm(box[0] - box[1]) + 1),
            int(np.linalg.norm(box[1] - box[2]) + 1),
        )
        if w < 10 or h < 10:
            polys.append(None)
            continue

        # warp image
        tar = np.float32([[0, 0], [w, 0], [w, h], [0, h]])
        M = cv2.getPerspectiveTransform(box, tar)
        word_label = cv2.warpPerspective(labels, M, (w, h), flags=cv2.INTER_NEAREST)
        try:
            Minv = np.linalg.inv(M)
        except:
            polys.append(None)
            continue

        # binarization for selected label
        cur_label = mapper[k]
        word_label[word_label != cur_label] = 0
        word_label[word_label > 0] = 1

        """ Polygon generation """
        # find top/bottom contours
        cp = []
        max_len = -1
        for i in range(w):
            region = np.where(word_label[:, i] != 0)[0]
            if len(region) < 2:
                continue
            cp.append((i, region[0], region[-1]))
            length = region[-1] - region[0] + 1
            if length > max_len:
                max_len = length

        # pass if max_len is similar to h
        if h * max_len_ratio < max_len:
            polys.append(None)
            continue

        # get pivot points with fixed length
        tot_seg = num_cp * 2 + 1
        seg_w = w / tot_seg  # segment width
        pp = [None] * num_cp  # init pivot points
        cp_section = [[0, 0]] * tot_seg
        seg_height = [0] * num_cp
        seg_num = 0
        num_sec = 0
        prev_h = -1
        for i in range(0, len(cp)):
            (x, sy, ey) = cp[i]
            if (seg_num + 1) * seg_w <= x and seg_num <= tot_seg:
                # average previous segment
                if num_sec == 0:
                    break
                cp_section[seg_num] = [
                    cp_section[seg_num][0] / num_sec,
                    cp_section[seg_num][1] / num_sec,
                ]
                num_sec = 0

                # reset variables
                seg_num += 1
                prev_h = -1

            # accumulate center points
            cy = (sy + ey) * 0.5
            cur_h = ey - sy + 1
            cp_section[seg_num] = [
                cp_section[seg_num][0] + x,
                cp_section[seg_num][1] + cy,
            ]
            num_sec += 1

            if seg_num % 2 == 0:
                continue  # No polygon area

            if prev_h < cur_h:
                pp[int((seg_num - 1) / 2)] = (x, cy)
                seg_height[int((seg_num - 1) / 2)] = cur_h
                prev_h = cur_h

        # processing last segment
        if num_sec != 0:
            cp_section[-1] = [cp_section[-1][0] / num_sec, cp_section[-1][1] / num_sec]

        # pass if num of pivots is not sufficient or segment widh
        # is smaller than character height
        if None in pp or seg_w < np.max(seg_height) * 0.25:
            polys.append(None)
            continue

        # calc median maximum of pivot points
        half_char_h = np.median(seg_height) * expand_ratio / 2

        # calc gradiant and apply to make horizontal pivots
        new_pp = []
        for i, (x, cy) in enumerate(pp):
            dx = cp_section[i * 2 + 2][0] - cp_section[i * 2][0]
            dy = cp_section[i * 2 + 2][1] - cp_section[i * 2][1]
            if dx == 0:  # gradient if zero
                new_pp.append([x, cy - half_char_h, x, cy + half_char_h])
                continue
            rad = -math.atan2(dy, dx)
            c, s = half_char_h * math.cos(rad), half_char_h * math.sin(rad)
            new_pp.append([x - s, cy - c, x + s, cy + c])

        # get edge points to cover character heatmaps
        isSppFound, isEppFound = False, False
        grad_s = (pp[1][1] - pp[0][1]) / (pp[1][0] - pp[0][0]) + (
            pp[2][1] - pp[1][1]
        ) / (pp[2][0] - pp[1][0])
        grad_e = (pp[-2][1] - pp[-1][1]) / (pp[-2][0] - pp[-1][0]) + (
            pp[-3][1] - pp[-2][1]
        ) / (pp[-3][0] - pp[-2][0])
        for r in np.arange(0.5, max_r, step_r):
            dx = 2 * half_char_h * r
            if not isSppFound:
                line_img = np.zeros(word_label.shape, dtype=np.uint8)
                dy = grad_s * dx
                p = np.array(new_pp[0]) - np.array([dx, dy, dx, dy])
                cv2.line(
                    line_img,
                    (int(p[0]), int(p[1])),
                    (int(p[2]), int(p[3])),
                    1,
                    thickness=1,
                )
                if (
                    np.sum(np.logical_and(word_label, line_img)) == 0
                    or r + 2 * step_r >= max_r
                ):
                    spp = p
                    isSppFound = True
            if not isEppFound:
                line_img = np.zeros(word_label.shape, dtype=np.uint8)
                dy = grad_e * dx
                p = np.array(new_pp[-1]) + np.array([dx, dy, dx, dy])
                cv2.line(
                    line_img,
                    (int(p[0]), int(p[1])),
                    (int(p[2]), int(p[3])),
                    1,
                    thickness=1,
                )
                if (
                    np.sum(np.logical_and(word_label, line_img)) == 0
                    or r + 2 * step_r >= max_r
                ):
                    epp = p
                    isEppFound = True
            if isSppFound and isEppFound:
                break

        # pass if boundary of polygon is not found
        if not (isSppFound and isEppFound):
            polys.append(None)
            continue

        # make final polygon
        poly = []
        poly.append(warpCoord(Minv, (spp[0], spp[1])))
        for p in new_pp:
            poly.append(warpCoord(Minv, (p[0], p[1])))
        poly.append(warpCoord(Minv, (epp[0], epp[1])))
        poly.append(warpCoord(Minv, (epp[2], epp[3])))
        for p in reversed(new_pp):
            poly.append(warpCoord(Minv, (p[2], p[3])))
        poly.append(warpCoord(Minv, (spp[2], spp[3])))

        # add to final result
        polys.append(np.array(poly))

    return polys

def getDetBoxes(textmap, linkmap, text_threshold, link_threshold, low_text, poly=False):
    boxes, labels, mapper, num_characters, avg_character_size, word_id = getDetBoxes_core(
        textmap, linkmap, text_threshold, link_threshold, low_text
    )
    '''
    Returns the detection boxes and predictions of CRAFT as a dictionary
    Args:
        text_map: a 2d np float array containing the probability distribution of a
            pixel being a character region
        linkmap: a 2d np float array containing the probability distribution of a
            pixel being an inter character spacing
        text_threshold: text confidence threshold
        link_threshold: link confidence threshold
        low_text: text low-bound score
        poly (bool): whether to return the polyline seg or not
    '''
    if poly:
        polys = getPoly_core(boxes, labels, mapper, linkmap)
    else:
        polys = [None] * len(boxes)
    return_dict = {"boxes": boxes, "labels": labels, "mapper": mapper, "num_characters": num_characters,\
    "average_character_size": avg_character_size,"word_id": word_id, "polys": polys}
    return return_dict

def adjustResultCoordinates(polys, ratio_w, ratio_h, ratio_net=2):
    if len(polys) > 0:
        polys = np.array(polys)
        for k in range(len(polys)):
            if polys[k] is not None:
                polys[k] *= (ratio_w * ratio_net, ratio_h * ratio_net)
    return polys


import numpy as np

from .modeling.feature.feature_definition import FeatureDefinition
from .modeling.utility import p, to_key_value_columns

class OnlineFeatureExtractor(FeatureDefinition):
    def reset(self, initial_size=1):
        self.n_samples = 0
        self.X = np.zeros([initial_size, self.n_features], dtype=int)
        self.y = np.zeros(initial_size, dtype=int)
        self.sample = np.zeros(self.n_features, dtype=int)
        self.tokens = []

        self.index = []
        self.completions = []
        self.current_completion_start_index = 0
        self.last_token = None
        # self.stack_context_info = self.get_stack_context_info(None)

    def fill_preprocessor_context(self, line_content, line, doc):
        for f in FeatureDefinition.preprocessors:
            f(line_content=line_content, line=line, ch=0, doc=doc, context=self.context)

    def extract_online(self, completions, line_content, line, ch, doc, call_signatures):
        self.reset(len(completions))

        # context features
        completion = completions[0]
        # self.stack_context_info = self.get_stack_context_info(completion)
        completion_data_type = 'unknown'
        if completion.type == 'instance':
            definitions = completion.infer()
            if definitions:
                completion_data_type = definitions[0].name

        for f in FeatureDefinition.preprocessors:
            f(
                line_content=line_content[:ch],  # TODO: should slice at ch or not?
                line=line,
                ch=ch - 1,
                doc=doc,
                context=self.context)
        for i, f in enumerate(FeatureDefinition.context_features.values()):
            feature = f(
                line_content=line_content[:ch],
                line=line,
                ch=ch,
                doc=doc,
                call_signatures=call_signatures,
                completion_data_type=completion_data_type,
                context=self.context
                # stack_context_info=self.stack_context_info
            )
            self.sample[i + len(FeatureDefinition.completion_features)] = feature

        # completion features
        for completion in completions:
            for i, f in enumerate(FeatureDefinition.completion_features.values()):
                self.sample[i] = f(
                    completion=completion,
                    line_content=line_content[:ch],
                    line=line,
                    ch=ch,
                    doc=doc,
                    call_signatures=call_signatures,
                    completion_data_type=completion_data_type,
                    context=self.context
                    # stack_context_info=self.stack_context_info
                )
            self.X[self.n_samples] = self.sample
            self.n_samples += 1
            p('=' * 40 + f' {completion.name} ' + '=' * 40)
            p(to_key_value_columns(self.name_to_feature_index.keys(), self.sample))
        self.normalize_feature()


#! /usr/bin/env python
# coding=utf-8

import tensorflow as tf
from tensorflow import keras
import numpy as np
import cv2
import os
from tensorflow.python.saved_model import tag_constants

class AIE:
    """create an instance that provides inference of stage_1 and stage_2 models"""

    def __init__(self):
        print("Initializing...")
        gpus = tf.config.list_physical_devices('GPU')
        if gpus:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            os.environ['TF_GPU_THREAD_MODE'] = 'gpu_private'
        self.input_size = 416
        self.target_size = 96
        self.id_to_label = ('cycling', 'jogging', 'sitting', 'other')
        self.id_to_color = ((255, 255, 0), (0, 255, 255), (255, 0, 255), (128, 128, 128))
        self.skeleton = (
            (0, 1), (0, 2), (1, 2), (1, 3), (2, 4), (3, 5), (4, 6), (1, 7), (2, 8), (7, 8), (7, 9), (8, 10), (9, 11),
            (10, 12))
        self.edge_colors = ((255, 0, 0),
                            (255, 135, 0),
                            (255, 211, 0),
                            (222, 255, 10),
                            (161, 255, 10),
                            (10, 255, 153),
                            (10, 239, 255),
                            (20, 125, 245),
                            (88, 10, 255),
                            (239, 195, 5),
                            (116, 255, 82),
                            (15, 190, 199),
                            (105, 68, 250),
                            (239, 195, 5))
        print("Loading models...")
        self.model_human_detector = tf.saved_model.load("./aieAPI/human_detector", tags=[tag_constants.SERVING])
        self.infer_human_detector = self.model_human_detector.signatures['serving_default']
        self.model_classifier = tf.saved_model.load("./aieAPI/classifier", tags=[tag_constants.SERVING])
        self.infer_classifier = self.model_classifier.signatures['serving_default']
        self.model_keypoint_estimator = tf.saved_model.load("./aieAPI/keypoint_estimator", tags=[tag_constants.SERVING])
        self.infer_keypoint_estimator = self.model_keypoint_estimator.signatures['serving_default']
        print("Allocating memory...")
        _ = self.detect(
            cv2.cvtColor(cv2.imread("./aieAPI/warm_up/warm_up.jpg"), cv2.COLOR_BGR2RGB),
            show_bbox=False, show_skeleton=False, show_activity=False)
        print("Ready√")

    def get_activity(self, logits):
        activity_id = np.argmax(logits)
        if logits[activity_id] < 0.6:
            activity_id = 3
        return activity_id, self.id_to_label[activity_id], self.id_to_color[activity_id]

    def detect(self, original_image, show_bbox=True, show_skeleton=True, show_activity=True):
        """
        Parameters:
                original_image: RGB image (numpy array) of any size, dtype=uint8 ~ [0, 255]
                show_bbox: whether to draw bounding boxes on targets
                show_skeleton: whether to draw skeletons on targets
                show_activity: whether to display activity labels on targets

        Returns:
                output_image: RGB image (numpy array) with skeletons and activity labels, dtype=uint8 ~ [0, 255]
                count of activities: [<cycling>, <jogging>, <sitting>, <unknown>]
        """
        output_image = original_image.copy()
        H, W, _ = output_image.shape
        resized_image = cv2.resize(output_image, (self.input_size, self.input_size), cv2.INTER_CUBIC) / 255.
        resized_image = tf.expand_dims(resized_image.astype(np.float32), axis=0)

        pred_bbox = self.infer_human_detector(resized_image)['tf.concat_12'].numpy().reshape(-1, 84)
        boxes = pred_bbox[:, 0:4]
        scores = pred_bbox[:, 4:]
        human_indices = np.argwhere(np.argmax(scores, axis=-1) == 0)
        boxes = boxes[human_indices, :].reshape(-1, 4)
        scores = scores[human_indices, 0].reshape(-1)
        selected_indices = tf.image.non_max_suppression(
            boxes=boxes,
            scores=scores,
            max_output_size=20,
            iou_threshold=0.45,
            score_threshold=0.25)
        boxes = np.rint(np.clip(boxes[selected_indices.numpy(), :], 0., 1.) * [H, W, H, W]).astype(np.int32)

        activity_count = [0, 0, 0, 0]
        for y_min, x_min, y_max, x_max in boxes:
            line_px = min(5, max(1, round(100 * (y_max - y_min) * (x_max - x_min) / H / W)))
            target = keras.applications.mobilenet_v2.preprocess_input(tf.expand_dims(cv2.resize(
                output_image[y_min:y_max, x_min:x_max, :], (self.target_size, self.target_size),
                cv2.INTER_CUBIC).astype(np.float32), axis=0))
            keypoints = self.infer_keypoint_estimator(target)['regression'].numpy().reshape(-1, 2)
            activity_id, activity, label_color = self.get_activity(
                np.squeeze(self.infer_classifier(target)['classification'].numpy()))

            activity_count[activity_id] += 1
            Xs = np.rint(keypoints[:, 0] * (x_max - x_min) + x_min).astype(np.int32)
            Ys = np.rint(keypoints[:, 1] * (y_max - y_min) + y_min).astype(np.int32)

            if show_bbox:
                _ = cv2.rectangle(output_image, (x_min, y_min), (x_max, y_max), label_color, line_px)

            if show_skeleton:
                for index, edge_color in zip(self.skeleton, self.edge_colors):
                    _ = cv2.line(output_image, (Xs[index[0]], Ys[index[0]]), (Xs[index[1]], Ys[index[1]]), edge_color,
                                 line_px, cv2.LINE_AA)

            if show_activity:
                text_x, text_y = (Xs[0], Ys[0])
                (text_w, text_h), text_bl = cv2.getTextSize(activity, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                _ = cv2.rectangle(output_image, (text_x, text_y - text_h), (text_x + text_w, text_y + text_bl),
                                  (0, 0, 0), -1)
                _ = cv2.putText(output_image, activity, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, label_color, 1,
                                cv2.LINE_AA)

        return output_image, activity_count

def call_aie(original_frame, fps, aie):
    original_frame = cv2.cvtColor(original_frame, cv2.COLOR_BGR2RGB)
    # original_frame = cv2.resize(original_frame, (1100, 700))
    original_frame, _ = aie.detect(original_frame, show_bbox=False, show_skeleton=True, show_activity=True)
    _ = cv2.putText(original_frame, "FPS: %.1f" % fps, (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1,
                    cv2.LINE_AA)
    return original_frame


import numpy as np
import unittest
import gnas

class TestGenetic(unittest.TestCase):

    def test_search_cnn_space(self):
        ss = gnas.get_gnas_cnn_search_space(5, 1, gnas.SearchSpaceType.CNNSingleCell)
        ga = gnas.genetic_algorithm_searcher(ss, population_size=20, generation_size=20, p_cross_over=0.8)
        for i in range(30):
            for i, ind in enumerate(ga.get_current_generation()):
                ga.sample_child()
                ga.update_current_individual_fitness(ind, 0 + np.random.rand(1))
            ga.update_population()
            self.assertTrue(len(ga.max_dict) <= 200)
            self.assertTrue(len(ga.generation))

if __name__ == '__main__':
    unittest.main()


import io
import time
import logging
from urllib.request import urlopen
from urllib.error import HTTPError
from importlib import import_module

from PIL import Image
import numpy as np
import tensorflow as tf
from tensorflow.python.framework import errors
from tensorflow.python.lib.io import file_io

# This class is modified from
# https://github.com/GoogleCloudPlatform/cloudml-samples/blob/master/flowers/trainer/preprocess.py
# - supports inception v4 model
class EmbeddingsGraph(object):
  """Builds a graph and uses it to extract embeddings from images.
  """

  # These constants are set by Inception v3's expectations.
  WIDTH = 299
  HEIGHT = 299
  CHANNELS = 3

  def __init__(self, tf_session, version='inception_v3'):
    self.version = version
    self.tf_session = tf_session
    # input_jpeg is the tensor that contains raw image bytes.
    # It is used to feed image bytes and obtain embeddings.
    self.input_jpeg, self.embedding = self.build_graph()

    init_op = tf.global_variables_initializer()
    self.tf_session.run(init_op)
    self.restore_from_checkpoint('models/%s.ckpt' % self.version)

  def dim(self):
    return self.embedding.shape[-1].value

  def build_graph(self):
    """Forms the core by building a wrapper around the inception graph.

      Here we add the necessary input & output tensors, to decode jpegs,
      serialize embeddings, restore from checkpoint etc.

      To use other Inception models modify this file. Note that to use other
      models beside Inception, you should make sure input_shape matches
      their input. Resizing or other modifications may be necessary as well.
      See tensorflow/contrib/slim/python/slim/nets/inception_v3.py for
      details about InceptionV3.

    Returns:
      input_jpeg: A tensor containing raw image bytes as the input layer.
      embedding: The embeddings tensor, that will be materialized later.
    """

    input_jpeg = tf.placeholder(tf.string, shape=None)
    image = tf.image.decode_jpeg(input_jpeg, channels=self.CHANNELS)

    # Note resize expects a batch_size, but we are feeding a single image.
    # So we have to expand then squeeze.  Resize returns float32 in the
    # range [0, uint8_max]
    image = tf.expand_dims(image, 0)

    # convert_image_dtype also scales [0, uint8_max] -> [0 ,1).
    image = tf.image.convert_image_dtype(image, dtype=tf.float32)
    image = tf.image.resize_bilinear(
        image, [self.HEIGHT, self.WIDTH], align_corners=False)

    # Then rescale range to [-1, 1) for Inception.
    image = tf.subtract(image, 0.5)
    inception_input = tf.multiply(image, 2.0)

    embedding = self.get_pre_logits(inception_input)
    return input_jpeg, embedding

  def get_pre_logits(self, inception_input):
    # import model
    if self.version == 'inception_v3':
        from tensorflow.contrib.slim.python.slim.nets import inception_v3 as inception
        inception = import_module('tensorflow.contrib.slim.python.slim.nets.inception_v3')
    else:
        inception = import_module('nets.inception_v4')

    # Build Inception layers, which expect a tensor of type float from [-1, 1)
    # and shape [batch_size, height, width, channels].
    arg_scope = getattr(inception, '%s_arg_scope' % self.version)
    model = getattr(inception, self.version)
    with tf.contrib.slim.arg_scope(arg_scope()):
      _, end_points = model(inception_input, is_training=False)

    pre_logits_name = (self.version == 'inception_v3' and 'PreLogits' or 'PreLogitsFlatten')
    return end_points[pre_logits_name]

  def restore_from_checkpoint(self, checkpoint_path):
    """To restore inception model variables from the checkpoint file.

       Some variables might be missing in the checkpoint file, so it only
       loads the ones that are avialable, assuming the rest would be
       initialized later.
    Args:
      checkpoint_path: Path to the checkpoint file for the Inception graph.
    """
    # Get all variables to restore. Exclude Logits and AuxLogits because they
    # depend on the input data and we do not need to intialize them from
    # checkpoint.
    scope = ''.join([x.capitalize() for x in self.version.split('_')])
    all_vars = tf.contrib.slim.get_variables_to_restore(
            exclude=['%s/AuxLogits' % scope, '%s/Logits' % scope, 'global_step'])

    saver = tf.train.Saver(all_vars)
    saver.restore(self.tf_session, checkpoint_path)

  def calculate_embedding(self, batch_image_bytes):
    """Get the embeddings for a given JPEG image.

    Args:
      batch_image_bytes: As if returned from [ff.read() for ff in file_list].

    Returns:
      The Inception embeddings (bottleneck layer output)
    """
    return self.tf_session.run(
        self.embedding, feed_dict={self.input_jpeg: batch_image_bytes})

class ImageEmbeddingService:
    def __init__(self, version):
        tf_session = tf.InteractiveSession()
        self.graph = EmbeddingsGraph(tf_session, version)

    def dim(self):
        return self.graph.dim()

    def get_image_bytes(self, uri):
        def _open_file_read_binary(uri):
            try:
                return file_io.FileIO(uri, mode='rb')
            except errors.InvalidArgumentError:
                return file_io.FileIO(uri, mode='r')

        try:
          if uri.startswith('https://'):
              image_bytes = urlopen(uri, timeout=5).read()
          else:
              with _open_file_read_binary(uri) as f:
                  image_bytes = f.read()
          img = Image.open(io.BytesIO(image_bytes)).convert('RGB')
        except HTTPError as e:
          logging.warning('%s(%s) uri: %s', e.msg, e.code, uri)
          if e.code == 403:
            return
          raise e

        # A variety of different calling libraries throw different exceptions here.
        # They all correspond to an unreadable file so we treat them equivalently.
        except Exception as e:  # pylint: disable=broad-except
          logging.exception('Error processing image %s: %s', uri, str(e))
          raise e

        # Convert to desired format and output.
        output = io.BytesIO()
        img.save(output, 'jpeg')
        image_bytes = output.getvalue()
        return image_bytes

    def get_embedding(self, url):
        image_bytes = self.get_image_bytes(url)
        if image_bytes is None:
            return
        embedding = self.graph.calculate_embedding(image_bytes)
        return np.squeeze(embedding)

if __name__ == '__main__':
    url = ''
    service = ImageEmbeddingService()
    embedding = service.get_embedding(url)
    print(embedding)
    print(embedding.shape)