import winreg, subprocess, gym, requests, psutil
import numpy as np

from time import sleep, time

class ExcavatorEnv(gym.Env):

    metadata = {'render.modes': ['human']}

    def __init__(self, id, policy,
        obs_dim=5,
        action_dim=4,
        http_url='http://127.0.0.1:5000'
    ):

        super(ExcavatorEnv, self).__init__()
        self.obs_dim = obs_dim
        self.action_dim = action_dim

        # id and counts

        self.env_id = id
        self.policy = policy
        self.episode_count = 0
        self.step_count = 0
        self.n_steps = 8
        self.step_count_max = 8
        self.backend_assigned = False
        self.dig_target = None
        self.emp_target = None
        self.emp_trajectory = None
        self.http_url = http_url

        # thresholds and coefficients

        self.dist_coeff = 0.5
        self.mass_coeff = 0.5
        self.mass_thr = 0.05
        self.x_min = np.array([-180.0, 3.9024162648733514, 13.252630737652677, 16.775050853637147])
        self.x_max = np.array([180.0, 812.0058600513476, 1011.7128949856826, 787.6024456729566])
        self.d_min = np.array([-0.05, -0.3, -0.5, -0.8])
        self.d_max = np.array([0.05, 0.3, 0.5, 0.8])
        self.m_max = 1000.0
        self.t_max = 60.0
        self.delay_max = 3.0

        # collision accounting

        self.collision_coeff = 0.0
        self.restart_on_collision = False

        # observation and action spaces

        self.observation_space = gym.spaces.Box(low=0, high=1, shape=(obs_dim,), dtype=np.float)
        self.action_space = gym.spaces.Box(low=-1, high=1, shape=(action_dim * self.n_steps,), dtype=np.float)
        self.debug = False

    def step(self, action, a_thr=3.0, x_thr=5.0):

        # current position

        jdata = self._post_target()
        x_current = (jdata['x'] - self.x_min) / (self.x_max - self.x_min)

        # generate full trajectory

        action = np.clip(action, self.action_space.low, self.action_space.high)  # action values are now in range -1..1
        action = (action - self.action_space.low) / (self.action_space.high - self.action_space.low)  # action values are now in range 0..1
        dig_deltas_std = action.reshape(self.n_steps, self.action_dim)
        dig_deltas = np.vstack([self.d_min + x * (self.d_max - self.d_min) for x in dig_deltas_std])  # action values are now in range d_min..d_max
        if self.policy == 'residual':
            dig_the_target = np.vstack([x + y for x, y in zip(dig_deltas, self.dig_trajectory)])  # 1..n_steps
        else:
            dig_the_target = np.vstack([x + self.dig_target for x in dig_deltas])  # 1..n_steps
        turn_to_the_target = np.hstack([dig_the_target[0, 0], x_current[1:]])  # 0
        prepare_for_the_turn_to_the_dumper = np.hstack([dig_the_target[-1, 0], self.emp_trajectory[0, 1:]])  # n_steps + 1
        trajectory = np.vstack([turn_to_the_target, dig_the_target, prepare_for_the_turn_to_the_dumper, self.emp_trajectory])  # n_steps + 2..n_steps + 2 + emp_trajectory_len

        # move along the trajectory

        x_dig = None
        n_collisions = 0
        mass_max = 0
        soil_grabbed = 0
        t_start = time()
        for i in range(trajectory.shape[0]):
            target = np.clip(self.x_min + (self.x_max - self.x_min) * trajectory[i, :], self.x_min, self.x_max)
            jdata = self._post_target(target)
            n_collisions += jdata['c']
            mass = jdata['m'] / self.m_max
            if i <= self.n_steps and mass > self.mass_thr and x_dig is None:
                x_dig = (jdata['x'] - self.x_min) / (self.x_max - self.x_min)
            if mass > mass_max:
                mass_max = mass
            in_target = np.zeros_like(target)
            in_angle_target = 0
            t_step_start = time()
            while not np.all(in_target):
                jdata = self._post_target()
                t_step_delta = time() - t_step_start
                if jdata is not None:
                    current = jdata['x']
                    dist_to_x = np.abs(np.array(current) - target)
                    if np.abs(current[0] - target[0]) < a_thr:
                        in_angle_target = 1
                    for j in range(self.action_dim):
                        if dist_to_x[j] < x_thr:
                            in_target[j] = 1
                if i in [0, self.n_steps + 2] and in_angle_target:
                    break
                elif i not in [0, self.n_steps + 2] and t_step_delta > self.delay_max:
                    break

            if i == self.n_steps + 2:
                soil_grabbed = mass
        state = self._construct_state(x_dig, mass_max)
        self.step_count += 1
        reward = self._calculate_reward(x_dig, soil_grabbed, n_collisions)
        if self.debug:
            print(self.step_count, reward, time() - t_start)
            print(self.env_id, self.step_count, state, reward, x_dig)
        if self.step_count == self.step_count_max:
            print('After working for {0} seconds, solver {1} stops as it has reached the maximum step count.'.format(time() - self.time_start, self.env_id))
            done = True
        elif n_collisions > 0 and self.restart_on_collision:
            print('Solver {0} will restart due to unexpected collision!'.format(self.env_id))
            done = True
        else:
            done = False
        return state, reward, done, {'r': reward, 'l': self.step_count}

    def reset(self, delay=1.0):

        self._request_targets()
        self._generate_dig_trajectory()
        self._generate_emp_trajectory()
        self.target_idx = 0

        # wait for the backend to start

        ready = False
        while not ready:
            jdata = self._post_target()
            if jdata is not None:
                if jdata['backend_running'] and jdata['x'] is not None and jdata['l'] is not None and jdata['m'] is not None:
                    ready = True
                else:
                    sleep(delay)

        state = self._construct_state()
        self.step_count = 0
        self.time_start = time()
        return state

    def render(self, mode='human', close=False):
        pass

    def _get_mode(self, new_mode=None, uri='mode'):
        url = '{0}/{1}'.format(self.http_url, uri)
        if new_mode is not None:
            jdata = requests.post(url, json={'id': self.env_id, 'mode': new_mode}).json()
        else:
            jdata = requests.get(url, json={'id': self.env_id}).json()
        return jdata['mode']

    def _request_targets(self, uri='targets', eps=1e-10):
        url = '{0}/{1}'.format(self.http_url, uri)
        target_list = [None, None]
        while None in target_list:
            try:
                j = requests.get(url, json={'id': self.env_id}).json()
                target_list = j['targets']
            except Exception as e:
                print('Exception when requesting target list:')
                print(e)
        self.dig_target = (np.array(target_list[0]) - self.x_min) / (self.x_max - self.x_min + eps)
        self.emp_target = (np.array(target_list[1]) - self.x_min) / (self.x_max - self.x_min + eps)

    def _generate_dig_trajectory(self, fname='data/dig.txt'):
        dig = []
        with open(fname) as f:
            lines = f.readlines()
        for line in lines:
            dig.append([float(x) for x in line.split(',')])
        dig = np.vstack(dig)
        dig_old_angle = np.mean(dig[:, 0])
        dig_new_angles = dig[:, 0] - dig_old_angle + self.dig_target[0]
        self.dig_trajectory = dig
        self.dig_trajectory[:, 0] = dig_new_angles

    def _generate_emp_trajectory(self, fname='data/emp.txt'):
        emp = []
        with open(fname) as f:
            lines = f.readlines()
        for line in lines:
            emp.append([float(x) for x in line.split(',')])
        emp = np.vstack(emp)
        emp_old_angle = np.mean(emp[:, 0])
        emp_new_angles = emp[:, 0] - emp_old_angle + self.emp_target[0]
        self.emp_trajectory = emp
        self.emp_trajectory[:, 0] = emp_new_angles

    def _post_target(self, target=None, uri='p_target'):
        url = '{0}/{1}'.format(self.http_url, uri)
        if target is not None:
            target = target.tolist()
        try:
            jdata = requests.post(url, json={'id': self.env_id, 'y': target}).json()
        except Exception as e:
            print('Exception when posting new p_target:')
            print(e)
            jdata = None
        return jdata

    def _construct_state(self, x_dig=None, soil_taken=None):
        state = np.zeros(self.obs_dim)
        if x_dig is not None:
            state[:self.action_dim] = x_dig - self.dig_target
        if soil_taken is not None:
            state[self.action_dim] = soil_taken
        return state

    def _calculate_reward(self, x_dig, m_grabbed, n_collisions):
        if x_dig is None or m_grabbed == 0:
            r = -1
        else:
            dist_to_target = np.linalg.norm(x_dig - self.dig_target)
            r = self.dist_coeff * (1 - dist_to_target) + self.mass_coeff * m_grabbed - self.collision_coeff * n_collisions
        return r


#!/usr/bin/env python

"""
Sonify a FITS image
"""

from __future__ import absolute_import, division, print_function

import os.path
import argparse
from glob import glob
import astropy.io.fits as pyfits
from skylib.sonification import sonify_image

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        'infile', metavar='"infile"', nargs='?', default='*.fits',
        help='comma-separated list of input FITS file names/masks')
    parser.add_argument(
        '-a', '--coordinates', metavar='rect|radial|circ',
        choices=('rect', 'radial', 'circ'),
        default='rect', help='sonification coordinates')
    parser.add_argument(
        '-b', '--barycenter', action='store_true',
        help='move origin to barycenter')
    parser.add_argument(
        '-c', '--min-connected', metavar='n', type=int, default=5,
        help='minimum number of pixels per connected group')
    parser.add_argument(
        '-d', '--debug', action='store_true',
        help='write intermediate images')
    parser.add_argument(
        '-i', '--noise-hi', metavar='%', type=float, default=99.9,
        help='hi percentile noise clipping value')
    parser.add_argument(
        '-l', '--noise-lo', metavar='%', type=float, default=50.0,
        help='lo percentile noise clipping value')
    parser.add_argument(
        '-m', '--sampling-rate', metavar='n', type=int, default=44100,
        help='sampling rate in beats per second')
    parser.add_argument(
        '-n', '--noise-volume', metavar='n', type=int, default=1000,
        help='max noise volume')
    parser.add_argument(
        '-o', '--tones', metavar='n', type=int, default=22,
        help='number of discrete tones')
    parser.add_argument(
        '-p', '--hi-clip', metavar='%', type=float, default=99.9,
        help='high percentile image clipping value')
    parser.add_argument(
        '-r', '--threshold', metavar='x', type=float,
        default=1.5, help='detection threshold in units of RMS')
    parser.add_argument(
        '-s', '--background-scale', metavar='n|f', type=float, default=1/64,
        help='background inhomogeneity scale, pixels or fraction of image size')
    parser.add_argument(
        '-t', '--tempo', metavar='x', type=float, default=100.0,
        help='rows per second')
    parser.add_argument(
        '-v', '--volume', metavar='n', type=int, default=16384,
        help='max sound volume')
    args = parser.parse_args()

    for infile in sum([glob(mask) for mask in args.infile.split(',')], []):
        print('Generating sound from image "{}"'.format(infile))
        with pyfits.open(infile, 'readonly', uint=True) as f:
            img = f[0].data.astype(float)
        outfile = os.path.splitext(infile)[0] + '_{}.wav'.format(args.coord)
        print('Writing waveform file "{}"'.format(outfile))
        sonify_image(
            img, outfile, coord=args.coord, barycenter=args.barycenter,
            tempo=args.tempo, sampling_rate=args.sampling_rate,
            num_tones=args.tones, volume=args.volume,
            noise_volume=args.noise_volume, bkg_scale=args.background_scale,
            threshold=args.threshold, min_connected=args.min_connected,
            hi_clip=args.hi_clip, noise_lo=args.noise_lo,
            noise_hi=args.noise_hi)

if __name__ == '__main__':
    main()


from rlgym.utils.gamestates import PhysicsObject
import numpy as np

class PhysicsWrapper(object):

    def __init__(self, phys_obj: PhysicsObject = None):
        """
        PhysicsWrapper constructor. Under most circumstances, users should not expect to instantiate their own PhysicsWrapper objects.

        :param phys_obj: PhysicsObject object from which values will be read.
        """
        if phys_obj is None:
            self.position: np.ndarray = np.asarray([0, 0, 93])
            self.linear_velocity: np.ndarray = np.zeros(3)
            self.angular_velocity: np.ndarray = np.zeros(3)
        else:
            self._read_from_physics_object(phys_obj)

    def _read_from_physics_object(self, phys_obj: PhysicsObject):
        """
        A function to modify PhysicsWrapper values from values in a PhysicsObject object.
        """
        self.position = phys_obj.position
        self.linear_velocity = phys_obj.linear_velocity
        self.angular_velocity = phys_obj.angular_velocity

    def set_pos(self, x: float = None, y: float = None, z: float = None):
        """
        Sets position.

        :param x: Float indicating x position value.
        :param y: Float indicating y position value.
        :param z: Float indicating z position value.
        """
        if x is not None:
            self.position[0] = x
        if y is not None:
            self.position[1] = y
        if z is not None:
            self.position[2] = z

    def set_lin_vel(self, x: float = None, y: float = None, z: float = None):
        """
        Sets linear velocity.

        :param x: Float indicating x velocity value.
        :param y: Float indicating y velocity value.
        :param z: Float indicating z velocity value.
        """
        if x is not None:
            self.linear_velocity[0] = x
        if y is not None:
            self.linear_velocity[1] = y
        if z is not None:
            self.linear_velocity[2] = z

    def set_ang_vel(self, x: float = None, y: float = None, z: float = None):
        """
        Sets angular velocity.

        :param x: Float indicating x angular velocity value.
        :param y: Float indicating y angular velocity value.
        :param z: Float indicating z angular velocity value.
        """
        if x is not None:
            self.angular_velocity[0] = x
        if y is not None:
            self.angular_velocity[1] = y
        if z is not None:
            self.angular_velocity[2] = z

    def _encode(self) -> list:
        """
        Function called by a StateWrapper to produce a state string.

        :return: String containing value data.
        """
        encoded = np.concatenate((self.position, self.linear_velocity, self.angular_velocity))
        return encoded.tolist()


print("Loading TRT Modules....")
import json
import trt_pose.coco
import trt_pose.models
import torch
import torch2trt
from torch2trt import TRTModule
import cv2
import torchvision.transforms as transforms
import PIL.Image
import numpy as np
#from jetcam.usb_camera import USBCamera

from trt_pose.draw_objects import DrawObjects
from trt_pose.parse_objects import ParseObjects

OPTIMIZED_MODEL = 'models/trt_resnet18_baseline_att_224x224_A_epoch_249.pth'

with open('human_pose.json', 'r') as f:
    human_pose = json.load(f)

topology = trt_pose.coco.coco_category_to_topology(human_pose)

print("Done")

print("Loading ROS Modules....")
import rospy
from sensor_msgs.msg import Image
from trt_bridge.msg import *
human_msg = Human()

rospy.init_node('Inference')
image_pub = rospy.Publisher("/inference/np_out",Image,queue_size=1)
human_pub = rospy.Publisher("/inference/humans", Human,queue_size=1)
print("Done")

num_parts = len(human_pose['keypoints'])
num_links = len(human_pose['skeleton'])
print("Loading Model...")
model_trt = TRTModule()
model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))
print("Done")
mean = torch.Tensor([0.485, 0.456, 0.406]).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).cuda()
device = torch.device('cuda')

def preprocess(image):
    global device
    device = torch.device('cuda')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = PIL.Image.fromarray(image)
    image = transforms.functional.to_tensor(image).to(device)
    image.sub_(mean[:, None, None]).div_(std[:, None, None])
    return image[None, ...]

parse = ParseObjects(topology)
draw = DrawObjects(topology)
#camera = USBCamera(width=224,height=224,capture_fps=30)
#camera.running = True

def publish_image(image):
    img = Image()
    (img.width,img.height,n) = image.shape
    img.encoding = "rgb8"
    img.is_bigendian = 0
    img.data = image.ravel().tolist()
    img.step = 3
    image_pub.publish(img)

def human_publish(people):
    for i in range(len(people)):
        human_msg.id = i
        for key in people[i]:
            part = keypoint()
            part.x,part.y = people[i][key]
            setattr(human_msg,key,part)
        human_pub.publish(human_msg)

print("Ready to Start Inference")
input("\t Press [Enter] to start")
def execute(image):

#    image = np.array(list(bytearray(msg.data)),dtype='uint8')
#    image = image.reshape(msg.width,msg.height,3)
    try:
        data = preprocess(image)
        cmap, paf = model_trt(data)
        cmap, paf = cmap.detach().cpu(), paf.detach().cpu()
        counts, objects, peaks = parse(cmap, paf)#, cmap_threshold=0.15, link_threshold=0.15)
        draw(image, counts, objects, peaks)
        publish_image(image)
        human_publish(get_points(counts,objects,peaks))
        rospy.loginfo("Inference complete {} person(s) detected".format(counts[0]))
    except KeyboardInterrupt:
        print("Shutting Down...")
        exit()

def get_points(counts,objects,peaks):
    height = 224
    width = 224
    k = topology.shape[0]
    count = counts[0]
    people = []
    for human in range(count):
        obj = objects[0][human]
        person = {}
        for key in range(obj.shape[0]):
            value = int(obj[key])
            if value >=0:
                peak = peaks[0][key][value]
                x,y = (round(float(peak[1])*width),round(float(peak[0])*height))
            else:
                x,y = -1,-1
            person[human_pose['keypoints'][key]]=(x,y)
        people.append(person)
    return people

cam = cv2.VideoCapture(0)
rate = rospy.Rate(30)
while not rospy.is_shutdown():
        ret,image = cam.read()
        image = cv2.resize(image,(224,224))
        execute(image)
        rate.sleep()


from typing import Dict, List, Optional, Any
from collections import namedtuple
import numpy as np
import torch
import torch.nn.functional as F

from ding.utils.data import default_collate, default_decollate
from ding.torch_utils import to_device
from haco.DIDrive_core.models import CILRSModel
from .base_carla_policy import BaseCarlaPolicy

class CILRSPolicy(BaseCarlaPolicy):
    """
    CILRS driving policy. It has a CILRS NN model which can handle
    observations from several environments by collating data into batch. It contains 2
    modes: `eval` and `learn`. The learn mode will calculate all losses, but will not
    back-propregate it. In `eval` mode, the output control signal will be postprocessed to
    standard control signal in Carla, and it can avoid stopping in the staring ticks.

    :Arguments:
        - cfg (Dict): Config Dict.

    :Interfaces:
        reset, forward
    """

    config = dict(
        cuda=True,
        max_throttle=0.75,
        model=dict(),
        learn=dict(
            epoches=200,
            lr=1e-4,
            batch_size=128,
            loss='l1',
            speed_weight=0.05,
            control_weights=[0.5, 0.45, 0.05],
        ),
    )

    def __init__(
            self,
            cfg: Dict,
    ) -> None:
        super().__init__(cfg, enable_field=[])
        self._enable_field = ['eval', 'learn']
        self._cuda = self._cfg.cuda
        self._max_throttle = self._cfg.max_throttle
        self._model = CILRSModel(**self._cfg.model)
        if self._cuda:
            self._model.cuda()

        for field in self._enable_field:
            getattr(self, '_init_' + field)()

    def _process_sensors(self, sensor: np.ndarray) -> np.ndarray:
        sensor = sensor[:, :, ::-1]  # BGR->RGB
        sensor = np.transpose(sensor, (2, 0, 1))
        sensor = sensor / 255.0

        return sensor

    def _process_model_outputs(self, data: Dict, output: List) -> List:
        action = []
        for i, d in enumerate(data.values()):
            control_pred = output[i][0]
            steer = control_pred[0] * 2 - 1.  # convert from [0,1] to [-1,1]
            throttle = min(control_pred[1], self._max_throttle)
            brake = control_pred[2]
            if d['tick'] < 20 and d['speed'] < 0.1:
                throttle = self._max_throttle
                brake = 0
            if brake < 0.05:
                brake = 0
            action.append({'steer': steer, 'throttle': throttle, 'brake': brake})
        return action

    def _reset_eval(self, data_id: Optional[List[int]] = None) -> None:
        """
        Reset policy of `eval` mode. It will change the NN model into 'eval' mode.

        :Arguments:
            - data_id (List[int], optional): List of env id to reset. Defaults to None.
        """
        self._model.eval()

    @torch.no_grad()
    def _forward_eval(self, data: Dict) -> Dict[str, Any]:
        """
        Running forward to get control signal of `eval` mode.

        :Arguments:
            - data (Dict): Input dict, with env id in keys and related observations in values,

        :Returns:
            Dict: Control and waypoints dict stored in values for each provided env id.
        """
        data_id = list(data.keys())

        new_data = dict()
        for id in data.keys():
            new_data[id] = dict()
            new_data[id]['rgb'] = self._process_sensors(data[id]['rgb'].numpy())
            new_data[id]['command'] = data[id]['command']
            new_data[id]['speed'] = data[id]['speed']

        new_data = default_collate(list(new_data.values()))
        if self._cuda:
            new_data = to_device(new_data, 'cuda')

        embedding = self._model.encode([new_data['rgb']])
        output = self._model(embedding, new_data['speed'], new_data['command'])
        if self._cuda:
            output = to_device(output, 'cpu')

        actions = default_decollate(output)
        actions = self._process_model_outputs(data, actions)
        return {i: {'action': d} for i, d in zip(data_id, actions)}

    def _init_learn(self) -> None:
        if self._cfg.learn.loss == 'l1':
            self._criterion = F.l1_loss
        elif self._cfg.policy.learn.loss == 'l2':
            self._criterion = F.mse_loss

    def _reset_learn(self, data_id: Optional[List[int]] = None) -> None:
        """
        Reset policy of `learn` mode. It will change the NN model into 'train' mode.

        :Arguments:
            - data_id (List[int], optional): List of env id to reset. Defaults to None.
        """
        self._model.train()

    def _forward_learn(self, data: Dict) -> Dict[str, Any]:
        """
        Running forward of `learn` mode to get loss.

        :Arguments:
            - data (Dict): Input dict, with env id in keys and related observations in values,

        :Returns:
            Dict: information about training loss.
        """
        if self._cuda:
            data = to_device(data, 'cuda')

        rgb = data['rgb']
        steer_gt, throttle_gt, brake_gt = data['steer'], data['throttle'], data['brake']
        speed = data['speed']
        command = data['command']
        embedding = self._model.encode([rgb])
        output = self._model(embedding, speed, command)
        control_pred, speed_pred = output
        steer_pred = control_pred[:, 0]
        throttle_pred = control_pred[:, 1]
        brake_pred = control_pred[:, 2]

        speed_loss = self._criterion(speed_pred.squeeze(), speed.squeeze()).mean() * self._cfg.learn.speed_weight
        steer_loss = self._criterion(steer_pred, steer_gt.squeeze()).mean() * self._cfg.learn.control_weights[0]
        throttle_loss = self._criterion(throttle_pred,
                                        throttle_gt.squeeze()).mean() * self._cfg.learn.control_weights[1]
        brake_loss = self._criterion(brake_pred, brake_gt.squeeze()).mean() * self._cfg.learn.control_weights[2]

        total_loss = speed_loss + steer_loss + throttle_loss + brake_loss

        return_info = {
            'total_loss': total_loss,
            'speed_loss': speed_loss,
            'steer_loss': steer_loss,
            'throttle_loss': throttle_loss,
            'brake_loss': brake_loss,
            # 'steer_mean': steer_pred.item().mean(),
            # 'throttle_mean': throttle_pred.item().mean(),
            # 'brake_mean': brake_pred.item().mean(),
        }

        return return_info


import time
import cv2
import numpy as np
import tensorflow as tf
import yolov3_tf2
import valohai
import os

from yolov3_tf2.models import (
    YoloV3
)
from yolov3_tf2.dataset import transform_images
from yolov3_tf2.utils import draw_outputs

params = {
    "size": 416,
}

inputs = {
    "classes": ".data/classes.txt",
    "model": ".data/model/model.tf.data-00000-of-00001",
    "image": "data/street.jpg"
}

valohai.prepare(step="detect", default_parameters=params, default_inputs=inputs)

def main():
    yolov3_tf2.YOLO_IOU_THRESHOLD = valohai.parameters('iou_threshold').value
    yolov3_tf2.YOLO_SCORE_THRESHOLD = valohai.parameters('score_threshold').value

    with open(valohai.inputs('classes').path(), "r") as f:
        num_classes = len([line.strip("\n") for line in f if line != "\n"])
        print(f"Number of classes used: {num_classes}")

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

    yolo = YoloV3(classes=num_classes)
    weights_path = os.path.join(os.path.dirname(valohai.inputs('model').path()), "model.tf")
    yolo.load_weights(weights_path).expect_partial()
    print('weights loaded')

    class_names = [c.strip() for c in open(valohai.inputs('classes').path()).readlines()]
    print(f'classes loaded ({class_names})')

    for path in valohai.inputs('image').paths():
        img_raw = tf.image.decode_image(open(path, 'rb').read(), channels=3)

        img = tf.expand_dims(img_raw, 0)
        img = transform_images(img, valohai.parameters('size').value)

        t1 = time.time()
        boxes, scores, classes, nums = yolo(img)
        t2 = time.time()
        print('time: {}'.format(t2 - t1))

        print('detections:')
        for i in range(nums[0]):
            print('\t{}, {}, {}'.format(class_names[int(classes[0][i])],
                                               np.array(scores[0][i]),
                                               np.array(boxes[0][i])))

        img = cv2.cvtColor(img_raw.numpy(), cv2.COLOR_RGB2BGR)
        img = draw_outputs(img, (boxes, scores, classes, nums), class_names)

        basename = os.path.basename(path)
        output_path = valohai.outputs().path(f'output_{basename}')
        cv2.imwrite(output_path, img)
        print('output saved to: {}'.format(output_path))

if __name__ == "__main__":
    main()


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
import os
import os.path as osp
from priv_config import cfg_priv

sys.path.append(cfg_priv.GLOBAL.PYTORCH_EVERYTHING_ROOT)

# Numerical libs
import cv2
import math
import numpy as np
import torch
import torch.backends.cudnn as cudnn

import pet4.models.imagenet as models
from pet4.utils.misc import weight_filler
from pet4.utils.transforms import bgr2rgb, normalize
from pet4.cls.core.config import cfg, merge_cfg_from_file

merge_cfg_from_file(osp.join(cfg_priv.GLOBAL.PYTORCH_EVERYTHING_ROOT, cfg_priv.MODULES.FKP.CFG))
# pprint.pprint(cfg)

cudnn.benchmark = True
_GRAY = [218, 227, 218]
_GREEN = [18, 127, 15]
_WHITE = [255, 255, 255]

extend_scale = 0.85
bbox_adapt_R = np.array([[0.355067475559597, 0.119262390148535, -0.355067475560677, -0.119262390148450],
                         [0.332986547478306, 0.246449822265788, -0.332986547479770, -0.246449822265689],
                         [-0.331071700057979, -0.146350441272887, 0.331071700059095, 0.146350441272787],
                         [-0.145853845175441, -0.399418765180059, 0.145853845176905, 0.399418765179917]])

class FaceKpts:
    def __init__(self, gpu_id=0):
        self.gpu_id = gpu_id
        torch.cuda.set_device(gpu_id)

        # Create model
        self.model = models.__dict__[cfg.MODEL.CONV_BODY]()
        print(self.model)

        # Load pre-train model
        model_weights = get_weights()
        pretrained_dict = torch.load(model_weights)
        # pretrained_dict = torch.load(model_weights, map_location='cpu')
        model_dict = self.model.state_dict()
        updated_dict, match_layers, mismatch_layers = weight_filler(pretrained_dict, model_dict)
        model_dict.update(updated_dict)
        self.model.load_state_dict(model_dict)
        print('==> Mismatch: ', mismatch_layers)

        self.model.cuda()
        self.model.eval()

    def __call__(self, img, cls_boxes_i, face_idx=1, min_thresh=0.8, min_abso_area=50*50, min_rela_area=0.01,
                 max_aspect_ratio=2.0):
        # img = bgr2rgb(img)
        normalized_im = normalize(np.asarray(img) / 255.0, mean=cfg.PIXEL_MEANS, std=cfg.PIXEL_STDS)
        img_area = img.shape[0] * img.shape[1]

        # boxes processing
        crop_imgs = []
        ldmk_bboxes = []
        bboxes = cls_boxes_i[face_idx]
        bboxes = boxes_filter(bboxes, img_area=img_area, min_thresh=min_thresh, min_abso_area=min_abso_area,
                              min_rela_area=min_rela_area, max_aspect_ratio=max_aspect_ratio)
        for i in range(len(bboxes)):
            bbox = bboxes[i]
            ldmk_bbox = bbox_adapt(np.asarray(bbox))
            ldmk_bbox = crop_bbox(ldmk_bbox)
            ldmk_bboxes.append(ldmk_bbox)
            crop_img = sample_img_ignore_bound(ldmk_bbox, normalized_im)
            crop_img = cv2.resize(crop_img, (120, 120), interpolation=cv2.INTER_CUBIC)
            crop_img = crop_img.astype(float)
            crop_imgs.append(crop_img)

        if len(crop_imgs) == 0:
            return None
        input_data = np.asarray(crop_imgs, dtype=np.float32)
        input_data = input_data.transpose(0, 3, 1, 2)

        with torch.no_grad():
            input_data = torch.from_numpy(input_data).float()
            input_data = input_data.cuda()
            outputs = self.model(input_data)
            feature = outputs.data.cpu().numpy()
            norm_pts84 = feature.reshape([-1, 2, cfg.MODEL.NUM_CLASSES // 2])

            all_pts = []
            for i in range(len(norm_pts84)):
                norm_pts = de_norm_kpt(ldmk_bboxes[i], norm_pts84[i])
                all_pts.append(norm_pts)

        return np.asarray(all_pts)

def get_weights(mode='ldmk10_84'):
    if os.path.exists(cfg.TEST.WEIGHTS):
        weights = cfg.TEST.WEIGHTS
    else:
        weights = os.path.join(cfg_priv.GLOBAL.PYTORCH_EVERYTHING_ROOT,
                               cfg.CKPT,
                               '{}.pth'.format(mode))
    return weights

def boxes_filter(boxes, img_area=800*1000, min_thresh=0.8, min_abso_area=50*50, min_rela_area=0.01,
                 max_aspect_ratio=2.0):
    boxes_fl = []
    for i in range(len(boxes)):
        box = boxes[i][:4]
        score = boxes[i][-1]
        box_area = (box[2] - box[0]) * (box[3] - box[1])
        box_aspect_ratio = max((box[2] - box[0]) / float(box[3] - box[1]), (box[3] - box[1]) / float(box[2] - box[0]))
        if score < min_thresh:
            print('min_thresh', score)
            continue
        if 0 < min_abso_area and box_area < min_abso_area:
            print('min_abso_area', box_area)
            continue
        if 0 < min_rela_area <= 1 and box_area / float(img_area) < min_rela_area:
            print('min_rela_area', box_area / float(img_area))
            continue
        if 1 <= max_aspect_ratio and box_aspect_ratio > max_aspect_ratio:
            print('max_aspect_ratio', box_aspect_ratio)
            continue
        boxes_fl.append(box)

    return np.asarray(boxes_fl)

def bbox_adapt(bbox_dt):
    # bbox norm
    centers = np.array([(bbox_dt[0] + bbox_dt[2]) / 2, (bbox_dt[1] + bbox_dt[3]) / 2])
    llengthes = ((bbox_dt[2] - bbox_dt[0]) + (bbox_dt[3] - bbox_dt[1])) / 4
    bbox_norm = np.array([bbox_dt[0] - centers[0], bbox_dt[1] - centers[1],
                          bbox_dt[2] - centers[0], bbox_dt[3] - centers[1]])
    bbox_norm = bbox_norm / llengthes

    # bbox adapt
    bbox_dt = np.transpose(np.dot(bbox_adapt_R, np.transpose(bbox_norm)))
    # bbox denorm
    bbox_dt = bbox_dt * llengthes
    bbox_dt = np.array([bbox_dt[0] + centers[0], bbox_dt[1] + centers[1],
                        bbox_dt[2] + centers[0], bbox_dt[3] + centers[1]])

    return bbox_dt

def crop_bbox(bbox_dt):
    half_width = bbox_dt[2] - bbox_dt[0]
    half_height = bbox_dt[3] - bbox_dt[1]
    llength = math.sqrt(half_width * half_width + half_height * half_height)
    llength = round(llength * extend_scale)
    centers = np.array([(bbox_dt[0] + bbox_dt[2]) / 2, (bbox_dt[1] + bbox_dt[3]) / 2])
    x1 = centers[0] - llength * 0.5
    y1 = centers[1] - llength * 0.43
    x2 = x1 + llength
    y2 = y1 + llength
    roi_box = np.array([x1, y1, x2, y2])

    return roi_box

# denorm pts5 via ldmk bbox
def de_norm_kpt(bbox, pt):
    pts_num = pt.shape[1]
    for i in range(pts_num):
        pt[0, i] = pt[0, i] * (bbox[2] - bbox[0]) + bbox[0]
        pt[1, i] = pt[1, i] * (bbox[3] - bbox[1]) + bbox[1]

    return pt

def sample_img_ignore_bound(roi, img):
    x1 = int(round(roi[0]))
    y1 = int(round(roi[1]))
    x2 = int(round(roi[2]))
    y2 = int(round(roi[3]))

    if (y2 - y1) > (x2 - x1):
        x2 += 1
    elif (y2 - y1) < (x2 - x1):
        y2 += 1

    img1 = np.zeros([y2 - y1, x2 - x1, 3])
    x1_ = max(x1, 0)
    y1_ = max(y1, 0)
    x2_ = min(x2, img.shape[1])
    y2_ = min(y2, img.shape[0])

    assert (y2 - y1) == (x2 - x1)

    img1[(y1_ - y1):(y2_ - y1), (x1_ - x1):(x2_ - x1), :] = img[y1_:y2_, x1_:x2_, :]
    return img1

def vis_kpts(img, kpts, color=(255, 255, 255), size=2):
    face_num, _, kpt_num = kpts.shape
    for i in range(face_num):
        for j in range(kpt_num):
            cv2.circle(img, (int(round(kpts[i, 0, j])), int(round(kpts[i, 1, j]))), size, color, -1)


from lxml import etree
import numpy
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
import operator
import pandas as pd
import requests
import json
from collections import defaultdict, Counter
import re

###GET FF-ICF PER EVENT TYPE###

def frames_naf_predicate(path_to_doc, frame_to_info, languages={'en'}):
    """Load a NAF file, extract the frames from their predicate layers and add them to a list."""
    doc_tree = etree.parse(path_to_doc) #parse the NAF file
    root = doc_tree.getroot()

    lang = root.get('{http://www.w3.org/XML/1998/namespace}lang')
    if lang not in languages:
        return []

    frames = []

    if root.find('srl') in root:
        for predicate in root.find('srl'):
            ext_ref_el = predicate.find('externalReferences/externalRef')
            uri = ext_ref_el.get('reference')
            label = frame_to_info[uri]['frame_label']
            frames.append(label) #append the frames to a list
    return frames

def frames_collection(collection, frame_to_info):
    """returns a list of frames extracted from a collection of NAF files."""
    collection_frames = []

    for file in collection: #iterate over the filepaths in a list
        for frame in frames_naf_predicate(file, frame_to_info): #iterate over the frames extracted from each NAF file
            collection_frames.append(frame) #append the frames to a list
    return collection_frames

def frames_collections(event_types, collection_of_collections, frame_to_info):
    """returns a dictionary with the event type as key and list of frames as value"""
    event_type_frames_dict = {}

    for event_type, collection in zip(event_types, collection_of_collections): #iterate over each event type and the corresponding list of sets of filepaths
        event_type_frames_dict[event_type] = frames_collection(collection, frame_to_info) #add each event type and the corresponding list of frames as key-value pairs to a dictionary
    return event_type_frames_dict

def contrastive_analysis(event_type_frames_dict):
    """returns a dictionary with event type as key and a sorted list of frames and their tf-idf values"""
    lists_frames = []

    for key in event_type_frames_dict: #iterate over the key:value (event type:list of frames) pairs
        values = event_type_frames_dict[key] #create a variable for each list of frames
        space = ' '
        space = space.join(values) #join the frames
        lists_frames.append(space) #append the string to a list

    vectorizer = CountVectorizer() #frame vocabulary
    lists_vector_data = vectorizer.fit_transform(lists_frames) #data structure that represents the instances through their vectors
    column_headers = vectorizer.get_feature_names() #frame vocabulary mapped to data columns
    tfidf_transformer = TfidfTransformer()
    lists_frames_tfidf = tfidf_transformer.fit_transform(lists_vector_data)
    tf_idf_array = lists_frames_tfidf.toarray() #apply tf-idf
    tf_idf_array_round = numpy.round(tf_idf_array, decimals=3)

    tf_idf_dict = {}

    for key, array in zip(event_type_frames_dict, tf_idf_array): #iterate over the event types and the list of corresponding arrays of tf-idf values
        frame_valuedict = {}
        for frame, value in zip(column_headers, array): #iterate over each frame and its corresponding value
            frame_valuedict[frame] = value #add the frame and value as key-value pair to a dictionary
        sorted_tuples = sorted(frame_valuedict.items(), key=operator.itemgetter(1), reverse=True) #convert the dictionary to a list of tuples sorted in descending order of the values
        tf_idf_dict[key] = sorted_tuples #add the event type and its list of tuples as a key-value pair to the tf_idfdict
    return tf_idf_dict

def output_tfidf_to_format(tf_idf_dict):
    """exports the output of the tf-idf analysis to an excel format"""
    headers = ['event type', 'rank', 'frame', 'tf-idf value', 'judgement']
    list_of_lists = []

    for key in tf_idf_dict: #iterate over the tf-idf dictionary
        for tupl, number in zip(tf_idf_dict[key][:40], range(1,41)): #iterate over n tuples and corresponding range of numbers
                one_row = [key, number] #create a list with the event type and the number of each tuple
                for element in tupl: #iterate over the frame and its tf-idf value of each tuple
                    one_row.append(element) #append both elements to the list
                one_row.append('') #append placeholder for validation
                list_of_lists.append(one_row) #append the list to a list of lists

    df = pd.DataFrame(list_of_lists, columns=headers) #turn the list into a table

    df.to_excel('tf_idf.xlsx', index=False) #export the table to an excel file

def capitalize_frame(frame):
    label = frame[0].upper() + frame[1:]
    return label

def validation_to_json(tf_idf_dict, output_path):
    """exports the tf-idf dictionary to json with validation of typical frames"""
    typical_frame_dict = {}

    for key, value in tf_idf_dict.items(): #iterate over the key/value pairs of the tf_idf dictionary
        typical = []
        other = []

        for frame in value[:10]: #iterate over the first n tuples in the list
            label = capitalize_frame(frame[0])
            typical.append(label) #append the frame of each tuple to a list
        for frame in value[10:]: #iterate over the rest of the tuples in the list
            label = capitalize_frame(frame[0])
            other.append(label) #append the frame of each tuple to another list
        validation_dict = {'typical': typical, 'other': other} #create dictionary with both lists as values
        typical_frame_dict[key] = validation_dict #add the dictionary to typical_frame_dict with event types as keys

    with open(output_path, 'w') as outfile:
        json.dump(typical_frame_dict, outfile, indent=4, sort_keys=True)

### CONVERT WIKIDATA IDENTIFIER TO ENTITY NAME ###

def get_entity_name(identifier):
    """returns name of the event type for a given identifier in wikidata"""
    r = requests.get(f"https://www.wikidata.org/entity/{identifier}.json")
    data = json.loads(r.text)
    return data["entities"][identifier]["labels"]["en"]["value"]

def get_entity_list(event_types):
    """returns a list with the wikidata identifiers converted to their entity name"""
    entity_list = []

    for identifier in event_types:
        entity = get_entity_name(identifier)
        entity_list.append(entity)
    return entity_list


import numpy as np
import tensorflow as tf
from PIL import Image
import cv2 as cv
from styx_msgs.msg import TrafficLight

class TLClassifier(object):
    def __init__(self):
        #TODO load classifier
        #pass
		self.is_site = FALSE
        self.image_result_show = False
		# load frozen tensorflow model into memory
        GRAPH_FILE = 'light_classification/model/frozen_inference_graph_sim.pb'
        self.detection_graph = load_graph(GRAPH_FILE)
        # create tensorflow session
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        self.session = tf.Session(graph=self.detection_graph, config=config)

		# `get_tensor_by_name` returns the Tensor with the associated name in the Graph.
        # extract the relevant tensors which reflect the outputs of the graph
		ops = self.detection_graph.get_operations()
        all_tensor_names = {output.name for op in ops for output in op.outputs}
        self.tensor_dict = {}
        for key in [
            'num_detections', 'detection_boxes', 'detection_scores',
            'detection_classes'
        ]:
            tensor_name = key + ':0'
            if tensor_name in all_tensor_names:
                self.tensor_dict[key] = self.detection_graph.get_tensor_by_name(
                tensor_name)

		# extract the relevant tensors which reflect the input of the graph
		# The input placeholder for the image.
        self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')

    def get_classification(self, image):
        """Determines the color of the traffic light in the image

        Args:
            image (cv::Mat): image containing the traffic light

        Returns:
            int: ID of traffic light color (specified in styx_msgs/TrafficLight)

        """
        #TODO implement light color prediction
        light_thresh = 0.4
        current_light = TrafficLight.UNKNOWN
        # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
        image_cp = image.copy()
        image_np_expanded = np.expand_dims(image_cp, axis=0)

        # Actual detection.
        output_dict = self.run_inference_for_single_image(image_np_expanded, self.detection_graph)

        # Show classified result image for test
        if self.image_result_show:
            img_width,img_height,depth = image.shape
            font = cv.FONT_HERSHEY_COMPLEX
            category_index = {1: {'id': 1, 'name': 'Green'}, 2: {'id': 2, 'name': 'Red'},
                    3: {'id': 3, 'name': 'Yellow'}, 4: {'id': 4, 'name': 'off'}}
            for i in range(output_dict['num_detections']):
                top_y = output_dict['detection_boxes'][i][0] * img_width
                top_x = output_dict['detection_boxes'][i][1] * img_height
                botom_y = output_dict['detection_boxes'][i][2] * img_width
                botom_x = output_dict['detection_boxes'][i][3] * img_height
                score = output_dict['detection_scores'][i]
                class_name = category_index[output_dict['detection_classes'][i]]['name']
                if score > light_thresh:
                    cv.rectangle(image_cp, (int(top_x), int(top_y)), (int(botom_x), int(botom_y)), (0, 0, 255), thickness=2)
                    cv.putText(image_cp, str(class_name) + ' ' + str(score), (int(top_x), int(top_y)), font, 0.8,(255, 255, 255), 1)
            cv.namedWindow('result', cv.WINDOW_NORMAL)
            cv.imshow('result', image_cp)
            cv.waitKey(1)

        # init red light score
        red_score_max = 0
        img_width,img_height,depth = image.shape
        for i in range(output_dict['num_detections']):
            if self.is_site or output_dict['detection_classes'][i]==2:        #if sim, just get red light max score
                if output_dict['detection_scores'][i] > red_score_max:        #if site, need get all traffic light max score
                    red_score_max = output_dict['detection_scores'][i]
                    top_y = output_dict['detection_boxes'][i][0] * img_width
                    top_x = output_dict['detection_boxes'][i][1] * img_height
                    botom_y = output_dict['detection_boxes'][i][2] * img_width
                    botom_x = output_dict['detection_boxes'][i][3] * img_height
        if self.is_site:       # if site mode, need classified the light color based image color
            if red_score_max > light_thresh:
                lightImg = image[int(top_y):int(botom_y),int(top_x):int(botom_x)]   # crop traffic image
                current_light = self.light_color(lightImg)        # get traffic light color
            else:
                current_light = TrafficLight.GREEN
        else:
            if red_score_max > light_thresh:
                current_light = TrafficLight.RED
            else:
                current_light = TrafficLight.GREEN

        if current_light == TrafficLight.RED:
            print('currend light is RED')
        else:
            print('currend light is GREEN')

        return current_light

    def load_graph(self, graph_file):
        """Loads a frozen inference graph"""
        graph = tf.Graph()
        with graph.as_default():
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(graph_file, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
        return graph

    def run_inference_for_single_image(self, image, graph):
        """Run detection and classification on a sample image.

        Args:
            image：
            graph：

        Returns:
        """
        output_dict = self.session.run(self.tensor_dict,
                                feed_dict={self.image_tensor: image})

        # all outputs are float32 numpy arrays, so convert types as appropriate
        output_dict['num_detections'] = int(output_dict['num_detections'][0])
        output_dict['detection_classes'] = output_dict['detection_classes'][0].astype(np.int64)
        output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
        output_dict['detection_scores'] = output_dict['detection_scores'][0]
        return output_dict

    def light_color(self,lightImg):
        """.

        Args:
            lightImg：

        Returns:
        """
        if self.image_result_show:
            cv.namedWindow('light_image', cv.WINDOW_NORMAL)
            cv.imshow('light_image', lightImg)
            cv.waitKey(1)

        lightImg = cv.cvtColor(lightImg, cv.COLOR_BGR2HSV)    # translate bgr to HSV

        height,width,depth = lightImg.shape

        height_start = int (height/8)      # split 3 part in height
        height_1st = int(height/3)
        height_2nd = int(height*2/3)
        height_end = int(height*7/8)

        lightImg_red = lightImg[height_start:height_1st,:,2]     # red part image
        lightImg_yellow = lightImg[height_1st:height_2nd,:,2]    # yellow part
        lightImg_green = lightImg[height_2nd:height_end,:,2]     # green part

        # Calculating the average of pixels each color part image
        red_height,red_width = lightImg_red.shape
        red_pixel_sum = np.sum(lightImg_red)/(red_height*red_width)

        yellow_height,yellow_width = lightImg_yellow.shape
        yellow_pixel_sum = np.sum(lightImg_yellow)/(yellow_height*yellow_width)

        green_height,green_width = lightImg_green.shape
        green_pixel_sum = np.sum(lightImg_green)/(green_height*green_width)

        print ("red is%f, yellow is %f, green is %f"%(red_pixel_sum,yellow_pixel_sum,green_pixel_sum))
        # red part is biggest indicate red light
        if red_pixel_sum > yellow_pixel_sum and red_pixel_sum > green_pixel_sum:
            current_light = TrafficLight.RED
        else:
            current_light = TrafficLight.GREEN
        return current_light