# Controller code for Digital Sky lighting project.
# This code is based upon the low level python wrapper for rpi_ws281x library
# which was produced by Tony DiCola (tony@tonydicory0x000080,a.com), Jeremy Garff (jer@jers.net)
# This code will animate a number of WS281x LEDs, and a number of LED Strips driven of WS2811 ICs on the same neopixel bus.
import sys,time, urllib.request, urllib.parse, urllib.error, traceback, random
from PIL import Image
from numpy import array, bitwise_xor, clip, greater, dstack, full, uint8, maximum
from neopixel import *
from filters.make_it_red import make_it_red

# LED strip configuration:
POWER_PIN      = 15      # GPIO pin which controlls the 36V power supply.
LED_PIN        = 18      # GPIO pin connected to the pixels (must support PWM!).
LED_FREQ_HZ    = 800000  # LED signal frequency in hertz (usually 800khz)
LED_DMA        = 10      # DMA channel to use for generating signal (try 10)
LED_BRIGHTNESS = 255     # Set to 0 for darkest and 255 for brightest
LED_INVERT     = False   # True to invert the signal (when using NPN transistor level shift)
LED_CHANNEL    = 0
LED_STRIP      = ws.SK6812W_STRIP

class LED_Control():
    def __init__(self, datastore):
        self.datastore = datastore
        self.strip = Adafruit_NeoPixel(self.datastore.LED_COUNT, LED_PIN, LED_FREQ_HZ, LED_DMA, LED_INVERT, LED_BRIGHTNESS, LED_CHANNEL, LED_STRIP)
        self.strip.begin()

    def service_leds(self):
        try:
            if self.datastore.power != 0:
                # Set up an empty blank row of Strip pixels
                self.datastore.strips=full((self.datastore.LED_COUNT,4),0)
                for animation in self.datastore.strip_animations:
                    self.datastore.strips=maximum(self.datastore.strips, animation.emit_row())
                # Set up an empty blank row of RGBW pixels
                rowdata=full((self.datastore.LED_COUNT,4),0)
                for animation in self.datastore.animations:
                    rowdata=maximum(rowdata, animation.emit_row())
                # Scale RGBW elements individually (colour tint)
                rowdata=rowdata*self.datastore.rgbw_brightness
                # Then scale everything by master_brightness
                rowdata=rowdata*float(self.datastore.master_brightness)
                rowdata=uint8(rowdata)
                # Update each LED color in the buffer.
                for i in range(self.strip.numPixels()):
                    if i % self.datastore.LAMP_LENGTH < self.datastore.STRIP_LEDS:
                        #self.strip.setPixelColor(i, Color(ib,ww,nw,dw))

                         self.strip._led_data[i]=(int(self.datastore.strips[i][3]) << 24) | \
                                                 (int(self.datastore.strips[i][0]) << 16) | \
                                                 (int(self.datastore.strips[i][1]) << 8 ) | \
                                                  int(self.datastore.strips[i][2])
                    else:
                    # Set the LED color buffer value.
                        #self.strip.setPixelColor(i, Color(r,g,b,w))
                        self.strip._led_data[i]=(int(rowdata[i][3]) << 24) | \
                                                (int(rowdata[i][0]) << 16) | \
                                                (int(rowdata[i][1]) << 8 ) | \
                                                 int(rowdata[i][2])
                # Send the LED color data to the hardware.
                self.strip.show()
        except Exception as err:
            print((self.datastore.strips))
            print(err)
            traceback.print_exc(file=sys.stdout)

def Color(red, green, blue, white = 0):
	"""Convert the provided red, green, blue color to a 24-bit color value.
	Each color component should be a value 0-255 where 0 is the lowest intensity
	and 255 is the highest intensity.
	"""
	return (white << 24) | (red << 16)| (green << 8) | blue


# import the necessary packages
from PIL import Image
import os, ntpath
from random import shuffle
import flask
from flask import send_file, request, Response, jsonify
import io
import json
import constants as P

# KERAS
from keras.models import load_model
import keras.backend as K

# initialize our Flask application and the Keras model
app = flask.Flask(__name__)

##########################################
##############  KERAS MODEL  #############

import tensorflow.compat.v1 as tf
from tensorflow.compat.v1.keras.models import load_model
import tensorflow.compat.v1.keras.backend as K
from io import StringIO
import base64, uuid
import numpy as np

tf.disable_v2_behavior()

print("Loading pix2pix model: " + P.BASE_DIR + P.MODEL + "..")
model = load_model(P.BASE_DIR + P.MODEL, compile=False)
print("Loading model.. DONE!")

graph = tf.get_default_graph()

@app.route("/magic", methods=["POST"])
def magic():
    image_uuid = uuid.uuid4().hex
    filename_orig = P.BASE_DIR + P.TEMP_DIR + image_uuid + '_orig.png'
    filename_segm = P.BASE_DIR + P.TEMP_DIR + image_uuid + '_segm.png'
    filename_pred = P.BASE_DIR + P.TEMP_DIR + image_uuid + '_pred.png'

    print ('MAGIC BLACKBOARD REQUEST: ' + filename_orig)
    with open(filename_orig,"wb+") as f:
        f.write(base64.b64decode(request.form['original']))
    with open(filename_segm,"wb+") as f:
        f.write(base64.b64decode(request.form['segmentation']))

    # OPEN AND PREPROCESS IMAGES
    img_org = Image.open(filename_orig)
    # WHITE BACKGROUND + PAINT CAUSE IT IS RGBA AND WE NEED ONLY RGB
    paint_sgm = Image.open(filename_segm)
    img_sgm = Image.new("RGB", paint_sgm.size, (255, 255, 255))
    img_sgm.paste(paint_sgm, mask=paint_sgm.split()[3])

    # DELETE TEMP FILE
    if not P.SAVE_DATA:
        os.unlink(filename_orig)
        os.unlink(filename_segm)

    # PREPARE IMAGES FOR MODEL INGESTION
    img_arr_org = np.array(img_org.resize(P.MODEL_SIZE_ORG, Image.BILINEAR))
    img_arr_org = np.asarray(img_arr_org/127.5 - 1)
    img_arr_org = np.expand_dims(img_arr_org, axis=0)

    img_arr_sgm = np.array(img_sgm.resize(P.MODEL_SIZE_SGM, Image.BILINEAR))
    img_arr_sgm = np.asarray(img_arr_sgm/127.5 - 1)
    img_arr_sgm = np.expand_dims(img_arr_sgm, axis=0)

    # DO THE MAGIC
    global graph
    with graph.as_default():
        out = model.predict([img_arr_sgm, img_arr_org])
    out_img = Image.fromarray(np.uint8((out[0]+1)*127.5))

    if P.SAVE_DATA:
        out_img.save(filename_pred)

    img_io = io.BytesIO()
    out_img.save(img_io, 'PNG')
    img_io.seek(0)
    return send_file(img_io, mimetype='image/png')

############# KERA MODEL #################
##########################################

@app.route('/')
def homepage():
    return "NOTHING TO DO HERE"

def start_api():
    app.run(host='0.0.0.0')


#!/usr/bin python3
""" Stats functions for the GUI """

import logging
import time
import os
import warnings

from math import ceil, sqrt

import numpy as np
import tensorflow as tf
from lib.Serializer import JSONSerializer

logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

def convert_time(timestamp):
    """ Convert time stamp to total hours, minutes and seconds """
    hrs = int(timestamp // 3600)
    if hrs < 10:
        hrs = "{0:02d}".format(hrs)
    mins = "{0:02d}".format((int(timestamp % 3600) // 60))
    secs = "{0:02d}".format((int(timestamp % 3600) % 60))
    return hrs, mins, secs

class TensorBoardLogs():
    """ Parse and return data from TensorBoard logs """
    def __init__(self, logs_folder):
        self.folder_base = logs_folder
        self.log_filenames = self.set_log_filenames()

    def set_log_filenames(self):
        """ Set the TensorBoard log filenames for all existing sessions """
        logger.debug("Loading log filenames. base_dir: '%s'", self.folder_base)
        log_filenames = dict()
        for dirpath, _, filenames in os.walk(self.folder_base):
            if not any(filename.startswith("events.out.tfevents") for filename in filenames):
                continue
            logfiles = [filename for filename in filenames
                        if filename.startswith("events.out.tfevents")]
            # Take the last logfile, in case of previous crash
            logfile = os.path.join(dirpath, sorted(logfiles)[-1])
            side, session = os.path.split(dirpath)
            side = os.path.split(side)[1]
            session = int(session[session.rfind("_") + 1:])
            log_filenames.setdefault(session, dict())[side] = logfile
        logger.debug("logfiles: %s", log_filenames)
        return log_filenames

    def get_loss(self, side=None, session=None):
        """ Read the loss from the TensorBoard logs
            Specify a side or a session or leave at None for all
        """
        logger.debug("Getting loss: (side: %s, session: %s)", side, session)
        all_loss = dict()
        for sess, sides in self.log_filenames.items():
            if session is not None and sess != session:
                logger.debug("Skipping session: %s", sess)
                continue
            loss = dict()
            for sde, logfile in sides.items():
                if side is not None and sde != side:
                    logger.debug("Skipping side: %s", sde)
                    continue
                for event in tf.train.summary_iterator(logfile):
                    for summary in event.summary.value:
                        if "loss" not in summary.tag:
                            continue
                        tag = summary.tag.replace("batch_", "")
                        loss.setdefault(tag,
                                        dict()).setdefault(sde,
                                                           list()).append(summary.simple_value)
            all_loss[sess] = loss
        return all_loss

    def get_timestamps(self, session=None):
        """ Read the timestamps from the TensorBoard logs
            Specify a session or leave at None for all
            NB: For all intents and purposes timestamps are the same for
                both sides, so just read from one side """
        logger.debug("Getting timestamps")
        all_timestamps = dict()
        for sess, sides in self.log_filenames.items():
            if session is not None and sess != session:
                logger.debug("Skipping sessions: %s", sess)
                continue
            for logfile in sides.values():
                timestamps = [event.wall_time
                              for event in tf.train.summary_iterator(logfile)
                              if event.summary.value]
                logger.debug("Total timestamps for session %s: %s", sess, len(timestamps))
                all_timestamps[sess] = timestamps
                break  # break after first file read
        return all_timestamps

class Session():
    """ The Loaded or current training session """
    def __init__(self, model_dir=None, model_name=None):
        logger.debug("Initializing %s: (model_dir: %s, model_name: %s)",
                     self.__class__.__name__, model_dir, model_name)
        self.serializer = JSONSerializer
        self.state = None
        self.modeldir = model_dir  # Set and reset by wrapper for training sessions
        self.modelname = model_name  # Set and reset by wrapper for training sessions
        self.tb_logs = None
        self.initialized = False
        self.session_id = None  # Set to specific session_id or current training session
        self.summary = SessionsSummary(self)
        logger.debug("Initialized %s", self.__class__.__name__)

    @property
    def batchsize(self):
        """ Return the session batchsize """
        return self.session["batchsize"]

    @property
    def config(self):
        """ Return config and other information """
        retval = {key: val for key, val in self.state["config"]}
        retval["training_size"] = self.state["training_size"]
        retval["input_size"] = [val[0] for key, val in self.state["inputs"].items()
                                if key.startswith("face")][0]
        return retval

    @property
    def full_summary(self):
        """ Retun all sessions summary data"""
        return self.summary.compile_stats()

    @property
    def iterations(self):
        """ Return session iterations """
        return self.session["iterations"]

    @property
    def logging_disabled(self):
        """ Return whether logging is disabled for this session """
        return self.session["no_logs"] or self.session["pingpong"]

    @property
    def loss(self):
        """ Return loss from logs for current session """
        loss_dict = self.tb_logs.get_loss(session=self.session_id)[self.session_id]
        return loss_dict

    @property
    def loss_keys(self):
        """ Return list of unique session loss keys """
        if self.session_id is None:
            loss_keys = self.total_loss_keys
        else:
            loss_keys = set(loss_key for side_keys in self.session["loss_names"].values()
                            for loss_key in side_keys)
        return list(loss_keys)

    @property
    def lowest_loss(self):
        """ Return the lowest average loss per save iteration seen """
        return self.state["lowest_avg_loss"]

    @property
    def session(self):
        """ Return current session dictionary """
        return self.state["sessions"].get(str(self.session_id), dict())

    @property
    def session_ids(self):
        """ Return sorted list of all existing session ids in the state file """
        return sorted([int(key) for key in self.state["sessions"].keys()])

    @property
    def timestamps(self):
        """ Return timestamps from logs for current session """
        ts_dict = self.tb_logs.get_timestamps(session=self.session_id)
        return ts_dict[self.session_id]

    @property
    def total_batchsize(self):
        """ Return all session batch sizes """
        return {int(sess_id): sess["batchsize"]
                for sess_id, sess in self.state["sessions"].items()}

    @property
    def total_iterations(self):
        """ Return session iterations """
        return self.state["iterations"]

    @property
    def total_loss(self):
        """ Return collated loss for all session """
        loss_dict = dict()
        all_loss = self.tb_logs.get_loss()
        for key in sorted(int(idx) for idx in all_loss.keys()):
            for loss_key, side_loss in all_loss[key].items():
                for side, loss in side_loss.items():
                    loss_dict.setdefault(loss_key, dict()).setdefault(side, list()).extend(loss)
        return loss_dict

    @property
    def total_loss_keys(self):
        """ Return list of unique session loss keys across all sessions """
        loss_keys = set(loss_key
                        for session in self.state["sessions"].values()
                        for loss_keys in session["loss_names"].values()
                        for loss_key in loss_keys)
        return list(loss_keys)

    @property
    def total_timestamps(self):
        """ Return timestamps from logs seperated per session for all sessions """
        return self.tb_logs.get_timestamps()

    def initialize_session(self, is_training=False, session_id=None):
        """ Initialize the training session """
        logger.debug("Initializing session: (is_training: %s, session_id: %s)",
                     is_training, session_id)
        self.load_state_file()
        self.tb_logs = TensorBoardLogs(os.path.join(self.modeldir,
                                                    "{}_logs".format(self.modelname)))
        if is_training:
            self.session_id = max(int(key) for key in self.state["sessions"].keys())
        else:
            self.session_id = session_id
        self.initialized = True
        logger.debug("Initialized session. Session_ID: %s", self.session_id)

    def load_state_file(self):
        """ Load the current state file """
        state_file = os.path.join(self.modeldir, "{}_state.json".format(self.modelname))
        logger.debug("Loading State: '%s'", state_file)
        try:
            with open(state_file, "rb") as inp:
                state = self.serializer.unmarshal(inp.read().decode("utf-8"))
                self.state = state
                logger.debug("Loaded state: %s", state)
        except IOError as err:
            logger.warning("Unable to load state file. Graphing disabled: %s", str(err))

    def get_iterations_for_session(self, session_id):
        """ Return the number of iterations for the given session id """
        session = self.state["sessions"].get(str(session_id), None)
        if session is None:
            logger.warning("No session data found for session id: %s", session_id)
            return 0
        return session["iterations"]

class SessionsSummary():
    """ Calculations for analysis summary stats """

    def __init__(self, session):
        logger.debug("Initializing %s: (session: %s)", self.__class__.__name__, session)
        self.session = session
        logger.debug("Initialized %s", self.__class__.__name__)

    @property
    def time_stats(self):
        """ Return session time stats """
        ts_data = self.session.tb_logs.get_timestamps()
        time_stats = {sess_id: {"start_time": min(timestamps) if timestamps else 0,
                                "end_time": max(timestamps) if timestamps else 0,
                                "datapoints": len(timestamps) if timestamps else 0}
                      for sess_id, timestamps in ts_data.items()}
        return time_stats

    @property
    def sessions_stats(self):
        """ Return compiled stats """
        compiled = list()
        for sess_idx, ts_data in self.time_stats.items():
            logger.debug("Compiling session ID: %s", sess_idx)
            if self.session.state is None:
                logger.debug("Session state dict doesn't exist. Most likely task has been "
                             "terminated during compilation")
                return None
            iterations = self.session.get_iterations_for_session(sess_idx)
            elapsed = ts_data["end_time"] - ts_data["start_time"]
            batchsize = self.session.total_batchsize.get(sess_idx, 0)
            compiled.append({"session": sess_idx,
                             "start": ts_data["start_time"],
                             "end": ts_data["end_time"],
                             "elapsed": elapsed,
                             "rate": (batchsize * iterations) / elapsed if elapsed != 0 else 0,
                             "batch": batchsize,
                             "iterations": iterations})
        compiled = sorted(compiled, key=lambda k: k["session"])
        return compiled

    def compile_stats(self):
        """ Compile sessions stats with totals, format and return """
        logger.debug("Compiling sessions summary data")
        compiled_stats = self.sessions_stats
        if compiled_stats is None:
            return compiled_stats
        logger.debug("sessions_stats: %s", compiled_stats)
        total_stats = self.total_stats(compiled_stats)
        compiled_stats.append(total_stats)
        compiled_stats = self.format_stats(compiled_stats)
        logger.debug("Final stats: %s", compiled_stats)
        return compiled_stats

    @staticmethod
    def total_stats(sessions_stats):
        """ Return total stats """
        logger.debug("Compiling Totals")
        elapsed = 0
        rate = 0
        batchset = set()
        iterations = 0
        total_summaries = len(sessions_stats)
        for idx, summary in enumerate(sessions_stats):
            if idx == 0:
                starttime = summary["start"]
            if idx == total_summaries - 1:
                endtime = summary["end"]
            elapsed += summary["elapsed"]
            rate += summary["rate"]
            batchset.add(summary["batch"])
            iterations += summary["iterations"]
        batch = ",".join(str(bs) for bs in batchset)
        totals = {"session": "Total",
                  "start": starttime,
                  "end": endtime,
                  "elapsed": elapsed,
                  "rate": rate / total_summaries,
                  "batch": batch,
                  "iterations": iterations}
        logger.debug(totals)
        return totals

    @staticmethod
    def format_stats(compiled_stats):
        """ Format for display """
        logger.debug("Formatting stats")
        for summary in compiled_stats:
            hrs, mins, secs = convert_time(summary["elapsed"])
            summary["start"] = time.strftime("%x %X", time.gmtime(summary["start"]))
            summary["end"] = time.strftime("%x %X", time.gmtime(summary["end"]))
            summary["elapsed"] = "{}:{}:{}".format(hrs, mins, secs)
            summary["rate"] = "{0:.1f}".format(summary["rate"])
        return compiled_stats

class Calculations():
    """ Class to pull raw data for given session(s) and perform calculations """
    def __init__(self, session, display="loss", loss_keys=["loss"], selections=["raw"],
                 avg_samples=500, smooth_amount=0.90, flatten_outliers=False, is_totals=False):
        logger.debug("Initializing %s: (session: %s, display: %s, loss_keys: %s, selections: %s, "
                     "avg_samples: %s, smooth_amount: %s, flatten_outliers: %s, is_totals: %s",
                     self.__class__.__name__, session, display, loss_keys, selections, avg_samples,
                     smooth_amount, flatten_outliers, is_totals)

        warnings.simplefilter("ignore", np.RankWarning)

        self.session = session
        self.display = display
        self.loss_keys = loss_keys
        self.selections = selections
        self.is_totals = is_totals
        self.args = {"avg_samples": avg_samples,
                     "smooth_amount": smooth_amount,
                     "flatten_outliers": flatten_outliers}
        self.iterations = 0
        self.stats = None
        self.refresh()
        logger.debug("Initialized %s", self.__class__.__name__)

    def refresh(self):
        """ Refresh the stats """
        logger.debug("Refreshing")
        if not self.session.initialized:
            logger.warning("Session data is not initialized. Not refreshing")
            return None
        self.iterations = 0
        self.stats = self.get_raw()
        self.get_calculations()
        self.remove_raw()
        logger.debug("Refreshed")
        return self

    def get_raw(self):
        """ Add raw data to stats dict """
        logger.debug("Getting Raw Data")

        raw = dict()
        iterations = set()
        if self.display.lower() == "loss":
            loss_dict = self.session.total_loss if self.is_totals else self.session.loss
            for loss_name, side_loss in loss_dict.items():
                if loss_name not in self.loss_keys:
                    continue
                for side, loss in side_loss.items():
                    if self.args["flatten_outliers"]:
                        loss = self.flatten_outliers(loss)
                    iterations.add(len(loss))
                    raw["raw_{}_{}".format(loss_name, side)] = loss

            self.iterations = 0 if not iterations else min(iterations)
            if len(iterations) > 1:
                # Crop all losses to the same number of items
                if self.iterations == 0:
                    raw = {lossname: list() for lossname in raw.keys()}
                else:
                    raw = {lossname: loss[:self.iterations] for lossname, loss in raw.items()}

        else:  # Rate calulation
            data = self.calc_rate_total() if self.is_totals else self.calc_rate()
            if self.args["flatten_outliers"]:
                data = self.flatten_outliers(data)
            self.iterations = len(data)
            raw = {"raw_rate": data}

        logger.debug("Got Raw Data")
        return raw

    def remove_raw(self):
        """ Remove raw values from stats if not requested """
        if "raw" in self.selections:
            return
        logger.debug("Removing Raw Data from output")
        for key in list(self.stats.keys()):
            if key.startswith("raw"):
                del self.stats[key]
        logger.debug("Removed Raw Data from output")

    def calc_rate(self):
        """ Calculate rate per iteration """
        logger.debug("Calculating rate")
        batchsize = self.session.batchsize
        timestamps = self.session.timestamps
        iterations = range(len(timestamps) - 1)
        rate = [batchsize / (timestamps[i + 1] - timestamps[i]) for i in iterations]
        logger.debug("Calculated rate: Item_count: %s", len(rate))
        return rate

    def calc_rate_total(self):
        """ Calculate rate per iteration
            NB: For totals, gaps between sessions can be large
            so time difference has to be reset for each session's
            rate calculation """
        logger.debug("Calculating totals rate")
        batchsizes = self.session.total_batchsize
        total_timestamps = self.session.total_timestamps
        rate = list()
        for sess_id in sorted(total_timestamps.keys()):
            batchsize = batchsizes[sess_id]
            timestamps = total_timestamps[sess_id]
            iterations = range(len(timestamps) - 1)
            print("===========\n")
            print(timestamps[:100])
            print([batchsize / (timestamps[i + 1] - timestamps[i]) for i in iterations][:100])
            rate.extend([batchsize / (timestamps[i + 1] - timestamps[i]) for i in iterations])
        logger.debug("Calculated totals rate: Item_count: %s", len(rate))
        return rate

    @staticmethod
    def flatten_outliers(data):
        """ Remove the outliers from a provided list """
        logger.debug("Flattening outliers")
        retdata = list()
        samples = len(data)
        mean = (sum(data) / samples)
        limit = sqrt(sum([(item - mean)**2 for item in data]) / samples)
        logger.debug("samples: %s, mean: %s, limit: %s", samples, mean, limit)

        for idx, item in enumerate(data):
            if (mean - limit) <= item <= (mean + limit):
                retdata.append(item)
            else:
                logger.trace("Item idx: %s, value: %s flattened to %s", idx, item, mean)
                retdata.append(mean)
        logger.debug("Flattened outliers")
        return retdata

    def get_calculations(self):
        """ Perform the required calculations """
        for selection in self.selections:
            if selection == "raw":
                continue
            logger.debug("Calculating: %s", selection)
            method = getattr(self, "calc_{}".format(selection))
            raw_keys = [key for key in self.stats.keys() if key.startswith("raw_")]
            for key in raw_keys:
                selected_key = "{}_{}".format(selection, key.replace("raw_", ""))
                self.stats[selected_key] = method(self.stats[key])

    def calc_avg(self, data):
        """ Calculate rolling average """
        logger.debug("Calculating Average")
        avgs = list()
        presample = ceil(self.args["avg_samples"] / 2)
        postsample = self.args["avg_samples"] - presample
        datapoints = len(data)

        if datapoints <= (self.args["avg_samples"] * 2):
            logger.info("Not enough data to compile rolling average")
            return avgs

        for idx in range(0, datapoints):
            if idx < presample or idx >= datapoints - postsample:
                avgs.append(None)
                continue
            else:
                avg = sum(data[idx - presample:idx + postsample]) \
                        / self.args["avg_samples"]
                avgs.append(avg)
        logger.debug("Calculated Average")
        return avgs

    def calc_smoothed(self, data):
        """ Smooth the data """
        last = data[0]  # First value in the plot (first timestep)
        weight = self.args["smooth_amount"]
        smoothed = list()
        for point in data:
            smoothed_val = last * weight + (1 - weight) * point  # Calculate smoothed value
            smoothed.append(smoothed_val)                        # Save it
            last = smoothed_val                                  # Anchor the last smoothed value

        return smoothed

    @staticmethod
    def calc_trend(data):
        """ Compile trend data """
        logger.debug("Calculating Trend")
        points = len(data)
        if points < 10:
            dummy = [None for i in range(points)]
            return dummy
        x_range = range(points)
        fit = np.polyfit(x_range, data, 3)
        poly = np.poly1d(fit)
        trend = poly(x_range)
        logger.debug("Calculated Trend")
        return trend


import numpy as np
from PyQt5 import QtCore
import pyqtgraph as pg

from acconeer_utils.clients import configs
from acconeer_utils import example_utils

import logging

log = logging.getLogger("acconeer_utils.examples.gui_iq")

def get_processing_config():
    return {
        "image_buffer": {
            "name": "Image history",
            "value": 100,
            "limits": [10, 16384],
            "type": int,
            "text": None,
        },
        "averaging": {
            "name": "Averaging",
            "value": 0,
            "limits": [0, 0.9999],
            "type": float,
            "text": None,
        },
    }

def get_sensor_config():
    return configs.IQServiceConfig()

class IQProcessor:
    def __init__(self, sensor_config, processing_config, session_info):
        self.data_processing = processing_config["processing_handle"]
        self.sensor_config = sensor_config
        self.mode = self.sensor_config.mode
        self.start_x = self.sensor_config.range_interval[0]
        self.stop_x = self.sensor_config.range_interval[1]
        self.sweep = 0
        self.time_filter = processing_config["averaging"]["value"]

        if "create_clutter" in processing_config:
            self.create_cl = processing_config["create_clutter"]
            self.use_cl = processing_config["use_clutter"]
            self.cl_file = processing_config["clutter_file"]
            self.sweeps = processing_config["sweeps_requested"]
        else:
            self.create_cl = None
            self.use_cl = None
            self.cl_file = None
            self.sweeps = -1

        self.rate = 1/self.sensor_config.sweep_rate

        self.image_buffer = processing_config["image_buffer"]["value"]

        if self.sweeps < 0:
            self.sweeps = self.image_buffer

        if self.create_cl:
            self.sweeps = max(self.sweeps, 100)

    def update_processing_config(self, processing_config):
        self.use_cl = processing_config["use_clutter"]

    def process(self, sweep):
        snr = {}
        peak_data = {}

        if len(sweep.shape) == 1:
            sweep = np.expand_dims(sweep, 0)

        if self.sweep == 0:
            self.num_sensors, self.data_len = sweep.shape
            self.env_x_mm = np.linspace(self.start_x, self.stop_x, self.data_len) * 1000

            self.cl_empty = np.zeros(self.data_len)
            self.last_iq = np.zeros((self.num_sensors, self.data_len))

            if self.num_sensors == 1:
                self.cl, self.cl_iq, self.n_std_avg = \
                    self.data_processing.load_clutter_data(self.data_len, self.cl_file)
            else:
                self.cl = self.cl_iq = self.n_std_avg = self.cl_empty
                self.create_cl = False
                if self.use_cl:
                    print("Background not supported for multiple sensors!")
                    self.use_cl = False

            self.hist_env = np.zeros(
                (self.num_sensors, len(self.env_x_mm), self.image_buffer)
                )
            self.peak_history = np.zeros(
                (self.num_sensors, self.image_buffer),
                dtype="float"
                )

        iq = sweep.copy()

        for s in range(self.num_sensors):
            if self.use_cl:
                try:
                    iq[s, :] = iq[s, :] - self.cl_iq
                except Exception as e:
                    log.error("Background has wrong format!\n{}".format(e))
                    self.use_cl = False
                    self.cl = self.cl_iq = np.zeros(len(self.cl))

            time_filter = self.time_filter
            if time_filter >= 0:
                if self.sweep < np.ceil(1.0 / (1.0 - self.time_filter) - 1):
                    time_filter = min(1.0 - 1.0 / (self.sweep + 1), self.time_filter)
                if self.sweep:
                    iq[s, :] = (1 - time_filter) * iq[s, :] + time_filter * self.last_iq[s, :]
                self.last_iq[s, :] = iq[s, :].copy()

        env = np.abs(iq)

        if self.create_cl:
            if self.sweep == 0:
                self.cl = np.zeros((self.sweeps, self.data_len))
                self.cl_iq = np.zeros((self.sweeps, self.data_len), dtype="complex")
            self.cl[self.sweep, :] = env[0, :]
            self.cl_iq[self.sweep, :] = iq[0, :]

        peak_data = {
            'peak_idx': np.zeros(self.num_sensors),
            'peak_mm': np.zeros(self.num_sensors),
            'peak_amp': np.zeros(self.num_sensors),
            'snr': np.zeros(self.num_sensors),
            }
        env_max = np.zeros(self.num_sensors)
        phase = np.zeros((self.num_sensors, self.data_len))

        hist_plots = []
        for s in range(self.num_sensors):
            peak_idx = np.argmax(env[s, :])
            peak_mm = self.env_x_mm[peak_idx]
            if peak_mm < self.start_x * 1000:
                peak_mm = self.stop_x * 1000
            peak_data['peak_mm'][s] = peak_mm
            peak_data['peak_idx'][s] = peak_idx
            env_max[s] = np.max(env[s, :])
            peak_data['peak_amp'][s] = env_max[s]

            snr = None
            signal = env[s, peak_idx]
            if self.use_cl and self.n_std_avg[peak_idx] > 0:
                noise = self.n_std_avg[peak_idx]
                snr = 20*np.log10(signal / noise)
            else:
                # Simple noise estimate: noise ~ mean(envelope)
                noise = np.mean(env[s, :])
                snr = 20*np.log10(signal / noise)
            peak_data["snr"][s] = snr

            p = np.angle(iq[s, :])
            p /= np.max(np.abs(p))

            phase[s, :] = p

            hist_plots.append(np.flip(self.peak_history[s, :], axis=0))
            self.peak_history[s, :] = np.roll(self.peak_history[s, :], 1)
            self.peak_history[s, 0] = peak_mm

            self.hist_env[s, :, :] = np.roll(self.hist_env[s, :, :], 1, axis=1)
            self.hist_env[s, :, 0] = env[s, :]

        cl = self.cl
        cl_iq = self.cl_iq
        if self.create_cl:
            cl = self.cl[self.sweep, :]
            cl_iq = self.cl_iq[self.sweep, :]
        elif not self.use_cl:
            cl = self.cl_empty
            cl_iq = self.cl_empty

        plot_data = {
            "iq_data": iq,
            "env_ampl": env,
            "phase": phase,
            "env_clutter": cl,
            "iq_clutter": cl_iq,
            "clutter_raw": self.cl_iq,
            "env_max": env_max,
            "n_std_avg": self.n_std_avg,
            "hist_plots": hist_plots,
            "hist_env": self.hist_env,
            "sensor_config": self.sensor_config,
            "peaks": peak_data,
            "x_mm": self.env_x_mm,
            "cl_file": self.cl_file,
            "sweep": self.sweep,
            "num_sensors": self.num_sensors,
        }

        self.sweep += 1

        return plot_data

class PGUpdater:
    def __init__(self, sensor_config, processing_config, session_info):
        self.sensor_config = sensor_config
        self.num_sensors = len(sensor_config.sensor)

    def setup(self, win):
        win.setWindowTitle("Acconeer IQ mode example")
        self.envelope_plot_window = win.addPlot(row=0, col=0, title="Envelope",
                                                colspan=self.num_sensors)
        self.envelope_plot_window.showGrid(x=True, y=True)
        self.envelope_plot_window.addLegend(offset=(-10, 10))
        self.envelope_plot_window.setYRange(0, 1)
        self.envelope_plot_window.setLabel("left", "Amplitude")
        self.envelope_plot_window.setLabel("bottom", "Distance (mm)")

        self.peak_text = pg.TextItem(text="", color=(1, 1, 1), anchor=(0, 1), fill="#f0f0f0")
        self.peak_text.setZValue(3)
        self.envelope_plot_window.addItem(self.peak_text)

        self.phase_plot_window = win.addPlot(row=1, col=0, title="Phase", colspan=self.num_sensors)
        self.phase_plot_window.showGrid(x=True, y=True)
        self.phase_plot_window.addLegend(offset=(-10, 10))
        self.phase_plot_window.setLabel("left", "Normalized phase")
        self.phase_plot_window.setLabel("bottom", "Distance (mm)")

        self.envelope_plots = []
        self.peak_vlines = []
        self.clutter_plots = []
        self.phase_plots = []
        self.hist_plot_images = []
        self.hist_plots = []
        self.hist_plot_peaks = []

        lut = example_utils.pg_mpl_cmap("viridis")
        hist_pen = example_utils.pg_pen_cycler(1)

        for s in range(self.num_sensors):
            legend_text = "Sensor {}".format(self.sensor_config.sensor[s])
            pen = example_utils.pg_pen_cycler(s+1)
            self.envelope_plots.append(
                self.envelope_plot_window.plot(range(10), np.zeros(10), pen=pen, name=legend_text)
                )
            self.peak_vlines.append(
                pg.InfiniteLine(pos=0, angle=90, pen=pg.mkPen(width=2, style=QtCore.Qt.DotLine))
                )
            self.envelope_plot_window.addItem(self.peak_vlines[s])

            self.phase_plots.append(
                self.phase_plot_window.plot(range(10), np.arange(10)*0, pen=pen, name=legend_text)
                    )

            hist_title = "Envelope history"
            if self.num_sensors == 1:
                pen = pg.mkPen(0.2, width=2, style=QtCore.Qt.DotLine)
                self.clutter_plots.append(
                    self.envelope_plot_window.plot(
                        range(10),
                        np.zeros(10),
                        pen=pen,
                        name="Background")
                    )
                self.clutter_plots[0].setZValue(2)
            else:
                hist_title = "History {}".format(legend_text)

            self.hist_plot_images.append(win.addPlot(row=2, col=s,
                                                     title=hist_title))
            self.hist_plot_images[s].setLabel("left", "Distance (mm)")
            self.hist_plot_images[s].setLabel("bottom", "Time (s)")
            self.hist_plots.append(pg.ImageItem())
            self.hist_plots[s].setAutoDownsample(True)
            self.hist_plots[s].setLookupTable(lut)
            self.hist_plot_images[s].addItem(self.hist_plots[s])

            self.hist_plot_peaks.append(
                self.hist_plot_images[s].plot(range(10), np.zeros(10), pen=hist_pen))

    def update(self, data):
        xstart = data["x_mm"][0]
        xend = data["x_mm"][-1]
        xdim = data["hist_env"].shape[1]

        num_sensors = len(data["env_ampl"])
        if self.num_sensors < num_sensors:
            num_sensors = self.num_sensors

        if data["sweep"] <= 2:
            self.env_plot_max_y = np.zeros(self.num_sensors)
            self.envelope_plot_window.setXRange(xstart, xend)
            self.phase_plot_window.setXRange(xstart, xend)
            self.phase_plot_window.setYRange(-1.1, 1.1)

            self.smooth_envelope = example_utils.SmoothMax(
                int(self.sensor_config.sweep_rate),
                tau_decay=1,
                tau_grow=0.2
                )

            for s in range(num_sensors):
                self.peak_text.setPos(xstart, 0)

                yax = self.hist_plot_images[s].getAxis("left")
                y = np.round(np.arange(0, xdim+xdim/9, xdim/9))
                labels = np.round(np.arange(xstart, xend+(xend-xstart)/9,
                                  (xend-xstart)/9))
                ticks = [list(zip(y, labels))]
                yax.setTicks(ticks)
                self.hist_plot_images[s].setYRange(0, xdim)

                s_buff = data["hist_env"].shape[2]
                t_buff = s_buff / data["sensor_config"].sweep_rate
                tax = self.hist_plot_images[s].getAxis("bottom")
                t = np.round(np.arange(0, s_buff + 1, s_buff / min(10 / self.num_sensors, s_buff)))
                labels = np.round(t / s_buff * t_buff, decimals=3)
                ticks = [list(zip(t, labels))]
                tax.setTicks(ticks)

        peaks = data["peaks"]
        peak_txt = "Peak: N/A"
        for s in range(num_sensors):
            sensor = self.sensor_config.sensor[s]
            if peaks:
                self.peak_vlines[s].setValue(peaks["peak_mm"][s])
                if s == 0:
                    peak_txt = ""
                if np.isfinite(peaks["snr"][s]):
                    peak_txt += "Peak S{}: {:.1f}mm, SNR: {:.1f}dB".format(
                        sensor, peaks["peak_mm"][s], peaks["snr"][s])
                else:
                    peak_txt += "Peak S{}: %.1fmm".format(sensor, peaks["peak_mm"][s])
                if s < num_sensors - 1:
                    peak_txt += "\n"

            max_val = max(np.max(data["env_clutter"] + data["env_ampl"][s]),
                          np.max(data["env_clutter"]))
            peak_line = np.flip((data["hist_plots"][s] - xstart) / (xend - xstart) * xdim, axis=0)

            self.envelope_plots[s].setData(data["x_mm"], data["env_ampl"][s] + data["env_clutter"])
            self.phase_plots[s].setData(data["x_mm"], data["phase"][s])

            self.envelope_plot_window.setYRange(0, self.smooth_envelope.update(max_val))

            if num_sensors == 1 and len(self.clutter_plots):
                self.clutter_plots[0].setData(data["x_mm"], data["env_clutter"])

            ymax_level = min(1.5 * np.max(np.max(data["hist_env"][s, :, :])),
                             self.env_plot_max_y[s])

            self.hist_plots[s].updateImage(data["hist_env"][s, :, :].T, levels=(0, ymax_level))
            self.hist_plot_peaks[s].setData(peak_line)
            self.hist_plot_peaks[s].setZValue(2)

            if max_val > self.env_plot_max_y[s]:
                self.env_plot_max_y[s] = 1.2 * max_val
        self.peak_text.setText(peak_txt, color=(1, 1, 1))


import sys                          # LIB: System
import os.path, stat                # LIB: Path, File status
import glob                         # LIB: Global
import cv2                          # LIB: OpenCV
import numpy as np                  # LIB: Numpy
import torch                        # LIB: Pytorch
import architecture as arch         # LIB: ERSGAN architecture
import subprocess                   # LIB: Call Subprocess
import pathlib                      # LIB: Pathlib
import time                         # LIB: Time
from PIL import Image               # LIB: PIL
from PIL import ImageEnhance        # LIB: PIL Enhancement
from PIL import ImageFilter         # LIB: PIL Filters
from os.path import splitext        # LIB: extension split

# Changeable flags
powertwo = True                     # check for and correct textures which are not power of two size
rtcwexcludes = True                 # exclude defined RTCW/ET folders and use standard settings there
alphaoptimize = True                # use gaussian blur, contrast and brightness (or not, if not needed)
usesharpen = True                   # sharpen the high resolution texture before resize to increase quality
autoconvert = True                  # convert the image to RGB if it is NOT RBG/RGBA!
skiptracemap = True                 # don't resize / include the tracemap
scalelightmaps = True               # resize Lightmaps (could look better, could look strange)
scalelarge = False                  # scale large images too (True = they are initially resized to a lower res)
testmode = False                    # in Testmode, a Lancosz method is used instead of the ESRGAN method
warnings = False                    # ignore (False) or show (True) warnings

# VRAM limits 8GB
largelimit = 2048*2048              # maximum texture scaling limit (stop scaling if texture is below this size)
vramlimit = 1024*512                # maximum size a texture can have before scaling that no CUDA error occurs
									# this depends on available VRAM size, 1024*512 ist for 8GB VRAM, an
									# approx. calculation of the VRAM usage is: width*height*8192 in Bytes
									# so 1024*512*8192 = ~ 4.2GB plus the VRAM already used by Windows/Apps

# Predefined Values
modelfactor = 4                     # the scale the selected model has been trained on (default is 4x)
allowed = [".png",".tga",".jpg"]    # allowed image file extensions to process (default: PNG, TGA, JPG)
scaling = Image.LANCZOS             # scaling method reducing too large images for next scale pass
finishing = Image.LANCZOS           # scaling method reducing the highres image to the desired resolution
target = 'cuda'                     # ESRGAN target device: 'cuda' for nVidia card (fast) or 'cpu' for ATI/CPU

# Predefined ESRGAN Models
model_path = 'models/cartoonpainted_400000.pth'       # default model
font_model_path='models/ReducedColorsAttempt.pth'     # font model

# RTCW exclude files (not implemented yet)
excludes = ["gfx/2d/backtile.jpg"]  # ET: gives a strange background texture in the loading screen, should be black

# create logfile
log=open("convert.log","w+")

# ignore warning
if(warnings==False):
	import warnings
	warnings.filterwarnings("ignore")

# function: write to logfile and output to stdout
def write_log(*args):
	line = ' '.join([str(a) for a in args])
	log.write(line+'\n')
	print(line)

# function: delete a single directory
def remove_empty_dir(path):
	try:
		if(os.rmdir(path)):
			write_log("Removed: " + path)
	except OSError:
		write_log("Not removed: " + path)
		pass

# function: delete a directory tree
def remove_empty_dirs(path):
	for root, dirnames, filenames in os.walk(path, topdown=False):
		for dirname in dirnames:
			remove_empty_dir(os.path.realpath(os.path.join(root, dirname)))

# function: upscale a PNG image
def upscale(im, device, model):
	img = cv2.cvtColor(np.array(im), cv2.COLOR_RGB2BGR)
	img = img * 1.0 / 255
	img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
	img_LR = img.unsqueeze(0)
	img_LR = img_LR.to(device)
	output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
	output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
	output = (output * 255.0).round()
	output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
	return output.astype(np.uint8)

# get the next power of two value of a value (texture size correction)
def poweroftwo(val):
	val=int(val)
	# check range from 2^0=1 to 2^15 = 32768 (should be large enough :-)
	for i in range(0,15):
		# value is below current power of two? found!
		if val<pow(2,i):
			# get previous power of two and next power of two
			mn=pow(2,i-1)
			mx=pow(2,i)
			# get delta between previous and next power (middle)
			delta=(mx-mn)/2
			# value above the middle: use higher power of two else use lower power of two
			if val>=(mn+delta):
				return mx
			else:
				return mn

def hms_string(sec_elapsed):
    h = int(sec_elapsed / (60 * 60))
    m = int((sec_elapsed % (60 * 60)) / 60)
    s = sec_elapsed % 60.
    return "{}:{:>02}:{:>05.2f}".format(h, m, s)

# "resize" the image until it fits in vram, return the width (aspect is known)
def fitimage(width, height, vramlimit):

	# reduce image size by factor 2
	while(width*height>vramlimit):
		width=int(width/2)
		height=int(height/2)

	write_log("  - Image resized to "+str(width)+"x"+str(height))

	return width

# commandline input
input = sys.argv[1]
factor = sys.argv[2]
maxsize = sys.argv[3]
blur = sys.argv[4]
contrast = sys.argv[5]
brightness = sys.argv[6]
sharpen = sys.argv[7]
jpegquality = sys.argv[8]

# set data types
input=str(input)
factor=int(factor)
maxsize=int(maxsize)
blur=int(blur)
contrast=float(contrast)
brightness=float(brightness)
sharpen=int(sharpen)
jpegquality=int(jpegquality)

# start
write_log("======================================================================")
write_log("RTCWHQ batch upscaling with ERSGAN started")
write_log("======================================================================")
write_log("Model:           {:s}".format(model_path))
write_log("Fontmodel:       {:s}".format(font_model_path))
write_log("Folder:          " + input)
write_log("Maximum scaling: " + str(factor) + "x")
write_log("Maximum size:    " + str(maxsize) + " Pixel")
write_log("Gaussian Blur:   " + str(blur) + " Pixel")
write_log("Contrast:        " + str(int(contrast*100)) + "%")
write_log("Brightness:      " + str(int(brightness*100)) + "%")
write_log("Sharpen:         " + str(sharpen) + " Pixel")
write_log("JPEG Quality:    " + str(jpegquality))
write_log("----------------------------------------------------------------------")

# count files to process
pngCounter = sum(1 for f in pathlib.Path(input).glob('**/*.png'))
tgaCounter = sum(1 for f in pathlib.Path(input).glob('**/*.tga'))
jpgCounter = sum(1 for f in pathlib.Path(input).glob('**/*.jpg'))
Counter=pngCounter+tgaCounter+jpgCounter
write_log("Found "+str(Counter)+" images: "+str(pngCounter)+" PNG, "+str(tgaCounter)+" TGA and "+str(jpgCounter)+" JPEG")

# init model
if(testmode==False):
	write_log("----------------------------------------------------------------------")
	write_log("Puny human is instructed to wait until the models are been prepared...")

	device = torch.device(target)

	# prepare model
	model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1, upsample_mode='upconv')
	model.load_state_dict(torch.load(model_path), strict=True)
	model.eval()

	for k, v in model.named_parameters():
		v.requires_grad = False
		model = model.to(device)

	write_log("Model ready.")

	# prepare font model
	fontmodel = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1, upsample_mode='upconv')
	fontmodel.load_state_dict(torch.load(font_model_path), strict=True)
	fontmodel.eval()

	for k, v in fontmodel.named_parameters():
		v.requires_grad = False
		fontmodel = fontmodel.to(device)

	write_log("FontModel ready. Let's go!")

starttime=time.time()
cnt=0
dcnt=0

# iterate through all subfolders
for dirName, subdirList, fileList in os.walk(input, topdown=False):

	for fname in fileList:

		dirName.replace("/","\\")
		fname.replace("/","\\")
		path=dirName+"/"+fname

		delete=False
		delreason=""

		# split filename, make it writable and convert extension to lowercase
		os.chmod(path ,stat.S_IWRITE)
		filename,ext=splitext(fname)
		ext=ext.lower()

		# only convert allowed file extensions
		if ext in allowed:

			# open image, add extension to filename and save it to PNG
			fullname=dirName + "\\" + filename + ext

			# try to load the image or throw error if there is a problem
			try:
				im = Image.open(fullname)
				loaded=True
				width, height = im.size
				width=int(width)
				height=int(height)
				imode=im.mode
				cnt+=1

			# image could not open? error!
			except IOError:
				loaded=False
				delete=True
				delreason="could not open image"
				im.close()

			# skip tracemap if set
			if(skiptracemap==True and "_tracemap" in filename):
				loaded=False
				delete=True
				delreason="tracemaps not allowed"
				im.close()

			# skip lightmaps if set
			if(scalelightmaps==False and "lm_0" in filename):
				loaded=Fals
				delete=True
				delreason="lightmaps not allowed"
				im.close()

			# loading successful? check colormode first!
			if(loaded==True):

				write_log("----------------------------------------------------------------------")
				write_log("IMAGE "+str(cnt)+" of "+str(Counter)+": "+fullname)
				write_log("----------------------------------------------------------------------")
				stime=time.time()

				add="unknown"
				if(imode=="L"):
					add="RGB 8bit Greyscale"

				if(imode=="P"):
					add="RGB 8bit Palette"

				if(imode=="RGB"):
					add="RGB 24bit"

				if(imode=="RGBA"):
					add="RGB 32bit with Alpha Channel"

				write_log("- Colormode: "+add)

				# image is NOT RGB(A)? then try to convert it
				if(autoconvert==True and imode != 'RGBA' and imode != 'RGB'):

					# 8bit color palette/greyscale? then convert to 24bit RGB
					if(imode=="P" or imode=="L"):
						im2=Image.new("RGB",im.size)
						im2.paste(im)
						im=im2
						im2=None
						write_log("- NOTICE: 8bit converted to 24bit RGB")
					else:
						# image was not valid, skip
						loaded=False
						delete=True
						delreason="no valid image"
						im.close()

			# still a valid image? then process it
			if(loaded==True):

				# optional: check and correct textures which are NOT power of two size (can cause errors otherwise)
				if(powertwo):

					twow=poweroftwo(width)
					twoh=poweroftwo(height)

					# width/height doesn't match with calculated power of two value? correct it = resize to next power of two
					if (twow != width) or (twoh != height):
						write_log("- NOTICE: Texture size "+str(width)+"x"+str(height)+" corrected - was NOT power of two!")
						width=twow
						height=twoh
						im.resize((width,height),Image.LANCZOS)

				# store original width for later use
				ow=width
				oh=height

				# if exists, get alpha channel first before we mess with the original image
				alpha=None
				if im.mode in ('RGBA', 'LA') or (im.mode == 'P' and 'transparency' in im.info):
					alpha=im.split()[-1]
					# convert to RGB to work with ESRGAN
					alpha = alpha.convert("RGB")

				scalepass=1

				write_log("- Resolution: "+str(width)+"x"+str(height)+" ("+str(factor)+"x = "+str(width*factor)+"x"+str(height*factor)+")")

				# initial check if image is already too large: reduce it before enlargement
				process=True
				if(scalelarge==True):
					if((width*height)>vramlimit):
						write_log("  - Image too large, must be resized first")

						aspect=width/height
						width=fitimage(width,height,vramlimit)
						height=int(width/aspect)

						im=im.resize((width,height),scaling)
						if(alpha):
							alpha=alpha.resize((width,height),Image.LANCZOS)
				# simple check if the image needs processing or not
				else:
					if(width*height>=maxsize*maxsize):
						write_log("- no processing: image is already High Resolution")
						process=False

				# assume we're rescaling the image
				rescale=True

				# process images if necessary
				if(process==True):

					# resize image until it is large enough (change the largelimit var to a lower value if it crashes)
					while(width*height<largelimit and rescale==True):

						write_log("- Scale Pass #"+str(scalepass))

						# -------------------------------------------------------------------------------
						# Image is too large to scale? reduce size by factor two until the size is valid
						# according to VRAM limit.
						#
						# Notice: ESRGAN uses a lot of GPU VRAM so there is a limitation for the input
						# depending on the available VRAM. An estimated maximum value for 8GB VRAM is
						# about 1024x512 = 524288 Pixels, so change the vramlimit variable if you have
						# more or less than 8GB VRAM. An approximate formula to calculate the size:
						#
						# VRAM needed = width*height*8192 in Bytes
						#
						# You must add to this value the VRAM already assigned by Windows and Apps, which
						# can already take 2-3GB on 8GB VRAM, consider this
						# -------------------------------------------------------------------------------
						if((width*height)>vramlimit):

							write_log("  - Image doesn't fit in VRAM, must be resized")

							aspect=width/height
							width=fitimage(width,height,vramlimit)
							height=int(width/aspect)

							im=im.resize((width,height),scaling)
							if(alpha):
								alpha=alpha.resize((width,height),Image.LANCZOS)

						# scale color image
						width=width*modelfactor
						height=height*modelfactor
						write_log("  - ERSGAN scales Colormap to "+str(width)+"x"+str(height))
						if(testmode==False):
							if(("font" or "hudchars") in dirName):
								im=Image.fromarray(upscale(im,device,fontmodel))
							else:
								im=Image.fromarray(upscale(im,device,model))
						else:
							im=im.resize((width*modelfactor,height*modelfactor),scaling)

						# scale alpha channel
						if(alpha):
							write_log("  - ERSGAN scales Alphamap to "+str(width)+"x"+str(height))
							if(testmode==False):
								if(("font" or "hudchars") in dirName):
									alpha=Image.fromarray(upscale(alpha,device,fontmodel))
								else:
									alpha=Image.fromarray(upscale(alpha,device,model))
							else:
								alpha=alpha.resize((width*modelfactor,height*modelfactor),scaling)

						# image has target size? don't rescale anymore
						if(width==(ow*factor) and height==(oh*factor)):
							rescale=False
						# otherwise do the next scale pass
						else:
							scalepass+=1

					# calculate final texture size
					nsw=int(ow*factor)
					nsh=int(oh*factor)
					ms=maxsize
					bl=blur
					co=contrast
					br=brightness
					sh=sharpen

					# rtcw/et excludes and default settings for specific folder names
					if(rtcwexcludes):

						# limit general font size to 1024 and use different values for blur/contrast/brightness/sharpen
						if("font" in dirName):
							ms=1024
							sh=4
							bl=1
							co=2.0
							br=-0.5
							write_log("- Font Texture found, limiting size to "+str(ms)+" Pixel")

						# limit ET HUD font size to 1024 and use different values for blur/contrast/brightness/sharpen
						if("hudchars" in filename):
							ms=1024
							sh=4
							bl=0
							co=4.0
							br=-0.5
							write_log("- Font Texture found, limiting size to "+str(ms)+" Pixel")

						# limit leveshots image size to 1024
						if("levelshots" in dirName):
							ms=512
							# but the survey map can still be large
							if("_cc") in filename:
								ms=maxsize
							write_log("- Levelshot Texture found, limiting size to "+str(ms)+" Pixel")

						# limit lightmaps image size to 1024
						if("maps" in dirName):
							ms=1024
							write_log("- Lightmap Texture found, limiting size to "+str(ms)+" Pixel")

						# dont' add contrast to the skies and user different blur value
						folders=["skies","sfx","liquids"]
						if(dirName in folders):
						#if(("skies" in dirName) or ("sfx" in dirName) or ("liquids" in dirName)):
							co=0.0
							br=0.0
							bl=2
							write_log("- Blurry Alpha Texture found - no contrast or brightness change!")

					# optional: sharpen filter
					if(usesharpen and sh!=0):

						# don't sharpen lightmaps!
						if ("maps" in dirName):
							write_log("- Lightmap texture found, no sharpen on lightmaps")
						else:
							write_log("- Colormap: Sharpen")
							im=ImageEnhance.Sharpness(im).enhance(sh)

					# if texture size is too large? reduce by factor
					if(nsw>ms) or (nsh>ms):
						f=min(ms/nsw,ms/nsh)
						sw=int(nsw*f)
						sh=int(nsh*f)
						write_log("- Colormap: scaled to "+str(sw)+"x"+str(sh)+" ("+str(nsw)+"x"+str(nsh)+" is too large)")
					# or use the calculated values
					else:
						sw=nsw
						sh=nsh
						write_log("- Colormap: scaled to "+str(sw)+"x"+str(sh))

					# scale colormap to desired resolution
					im=im.resize((sw,sh),finishing)

					# scae alphamap, if there is alpha
					if(alpha):

						# convert alphamap to 8bit greyscale to increase processing speed (it's greyscale only)
						alpha = alpha.convert("L")

						# only perform this if alpha optimizations are desired
						if(alphaoptimize):

							# apply gaussian blur
							if(bl!=0):
								write_log("- Alphamap: Gaussian Blur")
								alpha = alpha.filter(ImageFilter.GaussianBlur(bl))

							# apply brightness
							if(br!=0.0):
								write_log("- Alphamap: Brightness")
								alpha = ImageEnhance.Brightness(alpha).enhance(1.0+br)

							# apply contrast
							if(co!=0.0):
								write_log("- Alphamap: Contrast")
								alpha = ImageEnhance.Contrast(alpha).enhance(1.0+co)

						# scale alpha to desired resolution
						write_log("- Alphamap scaled to "+str(sw)+"x"+str(sh))
						alpha=alpha.resize((sw,sh),finishing)

						# merge alpha channel with RGB
						write_log("- Merging Colormap with Alphamap")
						im.putalpha(alpha.split()[-1])

					# save file
					write_log("- Replacing original Texture")
					im.save(fullname,quality=jpegquality,optimize=True,progressive=True)
					im.close()
					write_log("- Conversion completed in "+str(hms_string(time.time()-stime)))

		else:
			# remove all other files
			delete=True

		# delete files if flagged for deletion
		if(delete==True):
			write_log("----------------------------------------------------------------------")
			write_log("- DELETED "+path+" ("+delreason+")")
			os.remove(path)
			dcnt+=1

# finish
write_log("----------------------------------------------------------------------")
write_log("Removing empty directories...")
remove_empty_dirs(input)
write_log("Converted " + str(cnt) + " images and deleted "+str(dcnt)+" other files in "+str(hms_string(time.time()-starttime))+". Done.")
log.close()

# wait for input and exit
subprocess.call('timeout /T 5')


#! /usr/bin/env python
# coding=utf-8
# Copyright (c) 2021 Jiyang Chen
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

# File name   : visualize.py
# Author      : Jiyang Chen
# Created date: 2021-02-09
# Description : Test if YOLO can detect a person in the drone picture and report its precision.
#               This is used for the multi-objective project.

import os
import cv2
import time
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import numpy as np
import tensorflow as tf
import core.utils as utils
from PIL import Image
from core.yolov4 import YOLOv4, YOLOv3, YOLOv3_tiny, YOLOv4_tiny, decode
from core.yolov4 import filter_boxes
from core.config import cfg
from absl import app, flags, logging
from absl.flags import FLAGS
from tensorflow.compat.v1 import ConfigProto
from tensorflow.compat.v1 import InteractiveSession
from tensorflow.python.saved_model import tag_constants

physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

flags.DEFINE_string('framework', 'tf', '(tf, tflite, trt')
flags.DEFINE_string('weights', './checkpoints/yolov4-tiny-416',
                    'path to weights file')
flags.DEFINE_integer('size', 416, 'resize images to')
flags.DEFINE_boolean('tiny', True, 'yolo or yolo-tiny')
flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4')
flags.DEFINE_string('image', './data/kite.jpg', 'path to input image')
flags.DEFINE_string('output', 'result.png', 'path to output image')
flags.DEFINE_float('iou', 0.45, 'iou threshold')
flags.DEFINE_float('score', 0.25, 'score threshold')

def main(_argv):
    config = ConfigProto()
    config.gpu_options.allow_growth = True
    session = InteractiveSession(config=config)
    STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)

    input_size = FLAGS.size
    image_path = FLAGS.image

    image_folder = "./data/aerial_photos/"
    # Find all png files within the given directory, sorted numerically
    image_files = []
    file_names = os.listdir(image_folder)
    saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING])
    infer = saved_model_loaded.signatures['serving_default']

    for file in file_names:
        if ".JPG" in file:
            # image_files.append(os.path.join(image_folder, file))
            image_files.append(file)
    image_files.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))

    for image in image_files[:]:
        original_image = cv2.imread(os.path.join(image_folder, image))
        original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
        original_image_size = original_image.shape[:2]
        image_data = cv2.resize(original_image, (input_size, input_size))

        # image_data = utils.image_preprocess(np.copy(original_image), [input_size, input_size])
        partial_image = original_image[int(original_image_size[0]/3):int(original_image_size[0]/3*2),
                                       int(original_image_size[1]/3):int(original_image_size[1]/3*2), :]
        # partial_image = original_image[int(original_image_size[0]/5*2):int(original_image_size[0]/5*3),
        #                                int(original_image_size[1]/5*2):int(original_image_size[1]/5*3), :]
        image_data = cv2.resize(partial_image, (input_size, input_size))

        image_data = image_data / 255.
        # image_data = image_data[np.newaxis, ...].astype(np.float32)

        images_data = []
        for i in range(1):
            images_data.append(image_data)
        images_data = np.asarray(images_data).astype(np.float32)

        batch_data = tf.constant(images_data)

        pred_bbox = infer(batch_data)

        for key, value in pred_bbox.items():
            boxes = value[:, :, 0:4]
            pred_conf = value[:, :, 4:]

        boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
            boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)),
            scores=tf.reshape(
                pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])),
            max_output_size_per_class=50,
            max_total_size=50,
            iou_threshold=FLAGS.iou,
            score_threshold=FLAGS.score
        )
        pred_bbox = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()]
        # image = utils.draw_bbox(original_image, pred_bbox)

        detect_person = False
        _, out_scores, out_classes, num_boxes = pred_bbox
        # print("num_boxes is ", num_boxes)
        # print("out_class is ", out_classes)
        for i in range(num_boxes[0]):
            class_id = int(out_classes[0][i])
            # print(class_id)
            if class_id == 0:
                detect_person = True
                print('%s: %.2f' % (image, out_scores[0][i]))
                break

        if not detect_person:
            print('%s: %.2f' % (image, 0))

        # image = utils.draw_bbox(image_data*255, pred_bbox)
        # image = Image.fromarray(image.astype(np.uint8))
        # image.show()
        # image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
        # cv2.imwrite(FLAGS.output, image)

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


import time
import sys
import argparse

# import image and DL processing
import cv2
import numpy as np
import dlib
from random import randrange
# from edgetpu.detection.engine import DetectionEngine
from pycoral.adapters import common
from pycoral.adapters import detect
from pycoral.utils.edgetpu import make_interpreter
from scipy.interpolate import UnivariateSpline

from imutils.video import VideoStream
from PIL import Image, ImageDraw

# import local helper classes
from faceextractor import FaceDataExtractor
from recognizer import FaceRecognizer

# construct the argument parser and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-o", "--output", default=False, action="store_true",
	help="Display dalek PoV")
ap.add_argument("-f", "--face", type=float, default=0.7,
	help="Face detection certainty")
ap.add_argument("-r", "--recognize", type=float, default=0.7,
	help="Face recognition certainty")
args = vars(ap.parse_args())

print(args)

print("Loading face detection engine...")
interpreter = make_interpreter("/home/pi/coral-dalek/mobilenet_ssd_v2_face_quant_postprocess_edgetpu.tflite")
interpreter.allocate_tensors()

print("Loading face landmark detection engine...")
shape_pred = dlib.shape_predictor("./shape_predictor_5_face_landmarks.dat")
face_ext = FaceDataExtractor()
print("Loading face recognitn engine...")
facerec = dlib.face_recognition_model_v1("./dlib_face_recognition_resnet_model_v1.dat")
face_recog = FaceRecognizer()

# https://www.askaswiss.com/2016/02/how-to-manipulate-color-temperature-opencv-python.html

if args['output']:
    pov = 0
    overlay=[]
    overlay.append(cv2.imread('dalekpov-a.png'))
    overlay.append(cv2.imread('dalekpov-b.png'))
    overlay.append(cv2.imread('dalekpov-c.png'))

    def create_transform(x, y):
        spl = UnivariateSpline(x, y)
        return spl(range(256))

    inc_col = create_transform([0, 64, 128, 192, 256],[150, 175, 200, 225, 256])
    dec_col = create_transform([0, 64, 128, 192, 256],[28, 64, 90, 110, 128])

print("Starting video capture")

vc = cv2.VideoCapture(0)
if not vc.isOpened():
    print("Cannot open USB camera.")
    exit()

cap_width  = vc.get(cv2.CAP_PROP_FRAME_WIDTH)
cap_height = vc.get(cv2.CAP_PROP_FRAME_HEIGHT)
cap_fps = vc.get(cv2.CAP_PROP_FPS)
print(cap_width," x ", cap_height," @ ", cap_fps)

print("Entering main loop, press CTRL+C to exit...")
while True:
    try:
        ret, frame = vc.read()
        if not ret:
            print("No frame received from camera; exiting...")
            break
        # Convert frame from color_coverted = cv2.cvtColor(opencv_image, cv2.COLOR_BGR2RGB)
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        image = Image.fromarray(frame)
        _, scale = common.set_resized_input(
            interpreter, image.size, lambda size: image.resize(size, Image.ANTIALIAS))
        interpreter.invoke()
        face_box_list = detect.get_objects(interpreter, args['face'], scale)

        draw = ImageDraw.Draw(image)
        for face in face_box_list:
            bbox = face.bbox
            draw.rectangle([(bbox.xmin, bbox.ymin), (bbox.xmax, bbox.ymax)], outline='black')
            box = dlib.rectangle(left = bbox.xmin,
                                right = bbox.xmax,
                                top = bbox.ymin,
                                bottom = bbox.ymax)
            shape = shape_pred(frame, box)
            if shape:
                face_chip_img = dlib.get_face_chip(frame, shape)
                face_descriptor = facerec.compute_face_descriptor(face_chip_img)
                name = face_recog.recognize_face(face_descriptor, threshold = args['recognize'])
            if name:
                if output:
                    draw.text((bbox.xmin, bbox.ymin - 20), name, fill='black')
                else:
                    print(name)

        if args['output']:
            displayImage = np.asarray(image)
            blue, green, red = cv2.split(displayImage)
            red = cv2.LUT(red, dec_col).astype(np.uint8)
            blue = cv2.LUT(blue, dec_col).astype(np.uint8)
            green = cv2.LUT(green, inc_col).astype(np.uint8)
            displayImage = cv2.merge((red, green, blue))

            # displayImage = cv2.cvtColor(displayImage, cv2.COLOR_BGR2GRAY)
            if (randrange(10) > 6): pov = randrange(3)
            displayImage = cv2.addWeighted(displayImage,0.8,overlay[pov],0.2,0)
            cv2.imshow('Dalek Fry Eyestalk PoV', displayImage)
            if cv2.waitKey(1) == ord('q'):
                raise KeyboardInterrupt
    except KeyboardInterrupt:
        vc.release()
        cv2.destroyAllWindows()
        print("Program halted by CTRL+C")
        sys.exit(0)