#!/usr/bin/env python

"""
Train an ngram model

    $ python -m ngram.train -o outdir [-t poolsize] [-k key] [-u username] [-p password]

where:
    - `outdir` is where the trained models will be saved
    - `poolsize` is the number of databases
    - `key` is the key to save the models in
    - `username` is the database username
    - `password` is the database password
"""

import numpy as np
import sys, os, re, music21
from optparse import OptionParser
from multiprocessing import Process

from collections import deque

from sqlalchemy import desc, asc

from db import Song, Track, Note, get_sessions
from ngram_helper import key_transpose_pitch
from exceptions import InvalidKeySignature

NUM_NOTES = 128

class RomanTrainer(object):
    """
    A RomanTrainer is the model trainer

    1. for a given process / database, and
    2. for a given roman numeral.
    """
    def __init__(self,p_id,rt_id,counts,options):
        """
        Initialize the RomanTrainer
        Args:
            p_id: process id
            rt_id: roman numeral id
            counts: counts matrix
            options: options passed into script
        """
        self.p_id = p_id
        self.rt_id = rt_id

        self.counts = counts
        self.triple = deque()
        self.options = options

        # assume the user has specified a major key
        self.dest_key = (music21.key.Key(options.key).sharps,0)

    def transposed_triple(self):
        """
        Transpose a triple into the appropriate key

        Returns:
            int[]: the transposed triple
        """
        res = []
        notes = list(self.triple)
        for note in notes:
            src_key = (note.track.key_sig_top,note.track.key_sig_bottom)
            res.append(key_transpose_pitch(note.pitch,src_key,self.dest_key))
	return res

    def train(self,note):
        """
        Train this RomanTrained on a given note

        Args:
            note: the note to train on
        """
        self.triple.append(note)

        if len(self.triple) > 3:
            # remove the old note
            old_note = self.triple.popleft()
            try:
                # increment the matrix, where appropriate
                np.add.at(self.counts, tuple(self.transposed_triple()), 1)
            except InvalidKeySignature, e:
                # remove the bad note, append the old note.
                self.triple.pop()
                self.triple.appendleft(old_note)

    def write(self):
        """
        Write the numpy counts matrix out to file.
        """
        with open(os.path.join(self.options.outdir,str(self.p_id),str(self.rt_id) + ".npy"), 'w') as outfile:
            np.save(outfile, self.counts)

class TrackTrainer(Process):
    """
    Separate process to train ngram models, all music sourcing from one database
    """
    def __init__(self,p_id,session,options):
        """
        Initialize the TrackTrainer process
        Args:
            p_id: process id
            session: the database session to load from
            options (dict): options passed to script
        """
        Process.__init__(self)
        self.session = session
        self.options = options

        self.rts = []
        matrix_size = (NUM_NOTES, NUM_NOTES, NUM_NOTES)

        # construct the roman trainers
        for i in xrange(7):
            rt = RomanTrainer(p_id,i + 1,np.zeros(matrix_size, dtype=np.int16),options)
            self.rts.append(rt)

    def run(self):
        """
        Start the process, training on each track separately
        """
        # iterate through all the tracks
        for trk in self.session.query(Track).all():
            self.train(trk)

        # write all the rts
        for rt in self.rts:
            rt.write()

    def train(self,trk):
        """
        Train the ngram model on a specific track

        Args:
            trk: the track on which to train
        """
        print os.path.basename(trk.song.title), ":", trk.instr_name

        # skip percurssion tracks
        regexp = re.compile(r'drum|cymbal', re.IGNORECASE)
        if trk.channel == 9 or regexp.search(trk.instr_name) is not None:
            # print 'skipped percussion track'
            return

        # skip bass tracks
        regexp = re.compile(r'bass', re.IGNORECASE)
        if (trk.channel >= 32 and trk.channel <= 39) or regexp.search(trk.instr_name) is not None:
            # print 'skipped bass track'
            return

        # and through all the notes in a track
        for note in trk.notes:
            if note.pitch < 0 or note.pitch >= NUM_NOTES:
                pass

            # train using the appropriate rt
            if note.roman:
                self.rts[note.roman-1].train(note)

def main():
    parser = OptionParser()

    parser.add_option("-o", "--outdir", dest="outdir")
    parser.add_option("-t", "--poolsize", dest="pool_size", default=8, type="int")
    parser.add_option("-k", "--key", dest="key", default="C")
    parser.add_option("-u", "--username", dest="db_username", default="postgres")
    parser.add_option("-p", "--password", dest="db_password", default="postgres")

    (options, args) = parser.parse_args()

    # make the process output directory if not there already
    for p_id in xrange(options.pool_size):
	print options.outdir
    	pt = os.path.join(options.outdir,str(p_id) + "/")
	print pt
        if not os.path.exists(pt):
            os.mkdir(pt)

    sessions = get_sessions(options.pool_size,options.db_username,options.db_password)
    processes = []

    # construct and start the threads
    for i in xrange(options.pool_size):
        p = TrackTrainer(str(i),sessions[i],options)
        processes.append(p)
        p.start()

    # wait for processes to complete
    for p in processes:
        p.join()

    # construct cumulative counts matrices
    matrix_size = (NUM_NOTES, NUM_NOTES, NUM_NOTES)
    cumulative_counts = []

    for i in xrange(7):
        cumulative_counts.append(np.zeros(matrix_size, dtype=np.int16))

    for p_id in xrange(options.pool_size):
        for rt_id in xrange(7):
            with open(os.path.join(options.outdir + "/",str(p_id) + "/",str(rt_id + 1) + ".npy")) as f:
                counts = np.load(f)
                cumulative_counts[rt_id] = np.add(cumulative_counts[rt_id],counts)

    for i in xrange(7):
        with open(os.path.join(options.outdir + "/",str(i+1) + ".npy"), "w") as f:
            np.save(f,cumulative_counts[i])

if __name__ == '__main__':
    main()


#!/usr/bin/env python
# encoding: utf-8
"""
color_names.py

Created by Matthew Loper on 2012-05-12.
Copyright (c) 2012 MPI. All rights reserved.
"""

import numpy as np

name_to_rgb = {
    'snow' : np.array([1.00, 0.98, 0.98]),
    'ghost white' : np.array([0.97, 0.97, 1.00]),
    'GhostWhite' : np.array([0.97, 0.97, 1.00]),
    'white smoke' : np.array([0.96, 0.96, 0.96]),
    'WhiteSmoke' : np.array([0.96, 0.96, 0.96]),
    'gainsboro' : np.array([0.86, 0.86, 0.86]),
    'floral white' : np.array([1.00, 0.98, 0.94]),
    'FloralWhite' : np.array([1.00, 0.98, 0.94]),
    'old lace' : np.array([0.99, 0.96, 0.90]),
    'OldLace' : np.array([0.99, 0.96, 0.90]),
    'linen' : np.array([0.98, 0.94, 0.90]),
    'antique white' : np.array([0.98, 0.92, 0.84]),
    'AntiqueWhite' : np.array([0.98, 0.92, 0.84]),
    'papaya whip' : np.array([1.00, 0.94, 0.84]),
    'PapayaWhip' : np.array([1.00, 0.94, 0.84]),
    'blanched almond' : np.array([1.00, 0.92, 0.80]),
    'BlanchedAlmond' : np.array([1.00, 0.92, 0.80]),
    'bisque' : np.array([1.00, 0.89, 0.77]),
    'peach puff' : np.array([1.00, 0.85, 0.73]),
    'PeachPuff' : np.array([1.00, 0.85, 0.73]),
    'navajo white' : np.array([1.00, 0.87, 0.68]),
    'NavajoWhite' : np.array([1.00, 0.87, 0.68]),
    'moccasin' : np.array([1.00, 0.89, 0.71]),
    'cornsilk' : np.array([1.00, 0.97, 0.86]),
    'ivory' : np.array([1.00, 1.00, 0.94]),
    'lemon chiffon' : np.array([1.00, 0.98, 0.80]),
    'LemonChiffon' : np.array([1.00, 0.98, 0.80]),
    'seashell' : np.array([1.00, 0.96, 0.93]),
    'honeydew' : np.array([0.94, 1.00, 0.94]),
    'mint cream' : np.array([0.96, 1.00, 0.98]),
    'MintCream' : np.array([0.96, 1.00, 0.98]),
    'azure' : np.array([0.94, 1.00, 1.00]),
    'alice blue' : np.array([0.94, 0.97, 1.00]),
    'AliceBlue' : np.array([0.94, 0.97, 1.00]),
    'lavender' : np.array([0.90, 0.90, 0.98]),
    'lavender blush' : np.array([1.00, 0.94, 0.96]),
    'LavenderBlush' : np.array([1.00, 0.94, 0.96]),
    'misty rose' : np.array([1.00, 0.89, 0.88]),
    'MistyRose' : np.array([1.00, 0.89, 0.88]),
    'white' : np.array([1.00, 1.00, 1.00]),
    'black' : np.array([0.00, 0.00, 0.00]),
    'dark slate gray' : np.array([0.18, 0.31, 0.31]),
    'DarkSlateGray' : np.array([0.18, 0.31, 0.31]),
    'dark slate grey' : np.array([0.18, 0.31, 0.31]),
    'DarkSlateGrey' : np.array([0.18, 0.31, 0.31]),
    'dim gray' : np.array([0.41, 0.41, 0.41]),
    'DimGray' : np.array([0.41, 0.41, 0.41]),
    'dim grey' : np.array([0.41, 0.41, 0.41]),
    'DimGrey' : np.array([0.41, 0.41, 0.41]),
    'slate gray' : np.array([0.44, 0.50, 0.56]),
    'SlateGray' : np.array([0.44, 0.50, 0.56]),
    'slate grey' : np.array([0.44, 0.50, 0.56]),
    'SlateGrey' : np.array([0.44, 0.50, 0.56]),
    'light slate gray' : np.array([0.47, 0.53, 0.60]),
    'LightSlateGray' : np.array([0.47, 0.53, 0.60]),
    'light slate grey' : np.array([0.47, 0.53, 0.60]),
    'LightSlateGrey' : np.array([0.47, 0.53, 0.60]),
    'gray' : np.array([0.75, 0.75, 0.75]),
    'grey' : np.array([0.75, 0.75, 0.75]),
    'light grey' : np.array([0.83, 0.83, 0.83]),
    'LightGrey' : np.array([0.83, 0.83, 0.83]),
    'light gray' : np.array([0.83, 0.83, 0.83]),
    'LightGray' : np.array([0.83, 0.83, 0.83]),
    'midnight blue' : np.array([0.10, 0.10, 0.44]),
    'MidnightBlue' : np.array([0.10, 0.10, 0.44]),
    'navy' : np.array([0.00, 0.00, 0.50]),
    'navy blue' : np.array([0.00, 0.00, 0.50]),
    'NavyBlue' : np.array([0.00, 0.00, 0.50]),
    'cornflower blue' : np.array([0.39, 0.58, 0.93]),
    'CornflowerBlue' : np.array([0.39, 0.58, 0.93]),
    'dark slate blue' : np.array([0.28, 0.24, 0.55]),
    'DarkSlateBlue' : np.array([0.28, 0.24, 0.55]),
    'slate blue' : np.array([0.42, 0.35, 0.80]),
    'SlateBlue' : np.array([0.42, 0.35, 0.80]),
    'medium slate blue' : np.array([0.48, 0.41, 0.93]),
    'MediumSlateBlue' : np.array([0.48, 0.41, 0.93]),
    'light slate blue' : np.array([0.52, 0.44, 1.00]),
    'LightSlateBlue' : np.array([0.52, 0.44, 1.00]),
    'medium blue' : np.array([0.00, 0.00, 0.80]),
    'MediumBlue' : np.array([0.00, 0.00, 0.80]),
    'royal blue' : np.array([0.25, 0.41, 0.88]),
    'RoyalBlue' : np.array([0.25, 0.41, 0.88]),
    'blue' : np.array([0.00, 0.00, 1.00]),
    'dodger blue' : np.array([0.12, 0.56, 1.00]),
    'DodgerBlue' : np.array([0.12, 0.56, 1.00]),
    'deep sky blue' : np.array([0.00, 0.75, 1.00]),
    'DeepSkyBlue' : np.array([0.00, 0.75, 1.00]),
    'sky blue' : np.array([0.53, 0.81, 0.92]),
    'SkyBlue' : np.array([0.53, 0.81, 0.92]),
    'light sky blue' : np.array([0.53, 0.81, 0.98]),
    'LightSkyBlue' : np.array([0.53, 0.81, 0.98]),
    'steel blue' : np.array([0.27, 0.51, 0.71]),
    'SteelBlue' : np.array([0.27, 0.51, 0.71]),
    'light steel blue' : np.array([0.69, 0.77, 0.87]),
    'LightSteelBlue' : np.array([0.69, 0.77, 0.87]),
    'light blue' : np.array([0.68, 0.85, 0.90]),
    'LightBlue' : np.array([0.68, 0.85, 0.90]),
    'powder blue' : np.array([0.69, 0.88, 0.90]),
    'PowderBlue' : np.array([0.69, 0.88, 0.90]),
    'pale turquoise' : np.array([0.69, 0.93, 0.93]),
    'PaleTurquoise' : np.array([0.69, 0.93, 0.93]),
    'dark turquoise' : np.array([0.00, 0.81, 0.82]),
    'DarkTurquoise' : np.array([0.00, 0.81, 0.82]),
    'medium turquoise' : np.array([0.28, 0.82, 0.80]),
    'MediumTurquoise' : np.array([0.28, 0.82, 0.80]),
    'turquoise' : np.array([0.25, 0.88, 0.82]),
    'cyan' : np.array([0.00, 1.00, 1.00]),
    'light cyan' : np.array([0.88, 1.00, 1.00]),
    'LightCyan' : np.array([0.88, 1.00, 1.00]),
    'cadet blue' : np.array([0.37, 0.62, 0.63]),
    'CadetBlue' : np.array([0.37, 0.62, 0.63]),
    'medium aquamarine' : np.array([0.40, 0.80, 0.67]),
    'MediumAquamarine' : np.array([0.40, 0.80, 0.67]),
    'aquamarine' : np.array([0.50, 1.00, 0.83]),
    'dark green' : np.array([0.00, 0.39, 0.00]),
    'DarkGreen' : np.array([0.00, 0.39, 0.00]),
    'dark olive green' : np.array([0.33, 0.42, 0.18]),
    'DarkOliveGreen' : np.array([0.33, 0.42, 0.18]),
    'dark sea green' : np.array([0.56, 0.74, 0.56]),
    'DarkSeaGreen' : np.array([0.56, 0.74, 0.56]),
    'sea green' : np.array([0.18, 0.55, 0.34]),
    'SeaGreen' : np.array([0.18, 0.55, 0.34]),
    'medium sea green' : np.array([0.24, 0.70, 0.44]),
    'MediumSeaGreen' : np.array([0.24, 0.70, 0.44]),
    'light sea green' : np.array([0.13, 0.70, 0.67]),
    'LightSeaGreen' : np.array([0.13, 0.70, 0.67]),
    'pale green' : np.array([0.60, 0.98, 0.60]),
    'PaleGreen' : np.array([0.60, 0.98, 0.60]),
    'spring green' : np.array([0.00, 1.00, 0.50]),
    'SpringGreen' : np.array([0.00, 1.00, 0.50]),
    'lawn green' : np.array([0.49, 0.99, 0.00]),
    'LawnGreen' : np.array([0.49, 0.99, 0.00]),
    'green' : np.array([0.00, 1.00, 0.00]),
    'chartreuse' : np.array([0.50, 1.00, 0.00]),
    'medium spring green' : np.array([0.00, 0.98, 0.60]),
    'MediumSpringGreen' : np.array([0.00, 0.98, 0.60]),
    'green yellow' : np.array([0.68, 1.00, 0.18]),
    'GreenYellow' : np.array([0.68, 1.00, 0.18]),
    'lime green' : np.array([0.20, 0.80, 0.20]),
    'LimeGreen' : np.array([0.20, 0.80, 0.20]),
    'yellow green' : np.array([0.60, 0.80, 0.20]),
    'YellowGreen' : np.array([0.60, 0.80, 0.20]),
    'forest green' : np.array([0.13, 0.55, 0.13]),
    'ForestGreen' : np.array([0.13, 0.55, 0.13]),
    'olive drab' : np.array([0.42, 0.56, 0.14]),
    'OliveDrab' : np.array([0.42, 0.56, 0.14]),
    'dark khaki' : np.array([0.74, 0.72, 0.42]),
    'DarkKhaki' : np.array([0.74, 0.72, 0.42]),
    'khaki' : np.array([0.94, 0.90, 0.55]),
    'pale goldenrod' : np.array([0.93, 0.91, 0.67]),
    'PaleGoldenrod' : np.array([0.93, 0.91, 0.67]),
    'light goldenrod yellow' : np.array([0.98, 0.98, 0.82]),
    'LightGoldenrodYellow' : np.array([0.98, 0.98, 0.82]),
    'light yellow' : np.array([1.00, 1.00, 0.88]),
    'LightYellow' : np.array([1.00, 1.00, 0.88]),
    'yellow' : np.array([1.00, 1.00, 0.00]),
    'gold' : np.array([1.00, 0.84, 0.00]),
    'light goldenrod' : np.array([0.93, 0.87, 0.51]),
    'LightGoldenrod' : np.array([0.93, 0.87, 0.51]),
    'goldenrod' : np.array([0.85, 0.65, 0.13]),
    'dark goldenrod' : np.array([0.72, 0.53, 0.04]),
    'DarkGoldenrod' : np.array([0.72, 0.53, 0.04]),
    'rosy brown' : np.array([0.74, 0.56, 0.56]),
    'RosyBrown' : np.array([0.74, 0.56, 0.56]),
    'indian red' : np.array([0.80, 0.36, 0.36]),
    'IndianRed' : np.array([0.80, 0.36, 0.36]),
    'saddle brown' : np.array([0.55, 0.27, 0.07]),
    'SaddleBrown' : np.array([0.55, 0.27, 0.07]),
    'sienna' : np.array([0.63, 0.32, 0.18]),
    'peru' : np.array([0.80, 0.52, 0.25]),
    'burlywood' : np.array([0.87, 0.72, 0.53]),
    'beige' : np.array([0.96, 0.96, 0.86]),
    'wheat' : np.array([0.96, 0.87, 0.70]),
    'sandy brown' : np.array([0.96, 0.64, 0.38]),
    'SandyBrown' : np.array([0.96, 0.64, 0.38]),
    'tan' : np.array([0.82, 0.71, 0.55]),
    'chocolate' : np.array([0.82, 0.41, 0.12]),
    'firebrick' : np.array([0.70, 0.13, 0.13]),
    'brown' : np.array([0.65, 0.16, 0.16]),
    'dark salmon' : np.array([0.91, 0.59, 0.48]),
    'DarkSalmon' : np.array([0.91, 0.59, 0.48]),
    'salmon' : np.array([0.98, 0.50, 0.45]),
    'light salmon' : np.array([1.00, 0.63, 0.48]),
    'LightSalmon' : np.array([1.00, 0.63, 0.48]),
    'orange' : np.array([1.00, 0.65, 0.00]),
    'dark orange' : np.array([1.00, 0.55, 0.00]),
    'DarkOrange' : np.array([1.00, 0.55, 0.00]),
    'coral' : np.array([1.00, 0.50, 0.31]),
    'light coral' : np.array([0.94, 0.50, 0.50]),
    'LightCoral' : np.array([0.94, 0.50, 0.50]),
    'tomato' : np.array([1.00, 0.39, 0.28]),
    'orange red' : np.array([1.00, 0.27, 0.00]),
    'OrangeRed' : np.array([1.00, 0.27, 0.00]),
    'red' : np.array([1.00, 0.00, 0.00]),
    'hot pink' : np.array([1.00, 0.41, 0.71]),
    'HotPink' : np.array([1.00, 0.41, 0.71]),
    'deep pink' : np.array([1.00, 0.08, 0.58]),
    'DeepPink' : np.array([1.00, 0.08, 0.58]),
    'pink' : np.array([1.00, 0.75, 0.80]),
    'light pink' : np.array([1.00, 0.71, 0.76]),
    'LightPink' : np.array([1.00, 0.71, 0.76]),
    'pale violet red' : np.array([0.86, 0.44, 0.58]),
    'PaleVioletRed' : np.array([0.86, 0.44, 0.58]),
    'maroon' : np.array([0.69, 0.19, 0.38]),
    'medium violet red' : np.array([0.78, 0.08, 0.52]),
    'MediumVioletRed' : np.array([0.78, 0.08, 0.52]),
    'violet red' : np.array([0.82, 0.13, 0.56]),
    'VioletRed' : np.array([0.82, 0.13, 0.56]),
    'magenta' : np.array([1.00, 0.00, 1.00]),
    'violet' : np.array([0.93, 0.51, 0.93]),
    'plum' : np.array([0.87, 0.63, 0.87]),
    'orchid' : np.array([0.85, 0.44, 0.84]),
    'medium orchid' : np.array([0.73, 0.33, 0.83]),
    'MediumOrchid' : np.array([0.73, 0.33, 0.83]),
    'dark orchid' : np.array([0.60, 0.20, 0.80]),
    'DarkOrchid' : np.array([0.60, 0.20, 0.80]),
    'dark violet' : np.array([0.58, 0.00, 0.83]),
    'DarkViolet' : np.array([0.58, 0.00, 0.83]),
    'blue violet' : np.array([0.54, 0.17, 0.89]),
    'BlueViolet' : np.array([0.54, 0.17, 0.89]),
    'purple' : np.array([0.63, 0.13, 0.94]),
    'medium purple' : np.array([0.58, 0.44, 0.86]),
    'MediumPurple' : np.array([0.58, 0.44, 0.86]),
    'thistle' : np.array([0.85, 0.75, 0.85]),
    'snow1' : np.array([1.00, 0.98, 0.98]),
    'snow2' : np.array([0.93, 0.91, 0.91]),
    'snow3' : np.array([0.80, 0.79, 0.79]),
    'snow4' : np.array([0.55, 0.54, 0.54]),
    'seashell1' : np.array([1.00, 0.96, 0.93]),
    'seashell2' : np.array([0.93, 0.90, 0.87]),
    'seashell3' : np.array([0.80, 0.77, 0.75]),
    'seashell4' : np.array([0.55, 0.53, 0.51]),
    'AntiqueWhite1' : np.array([1.00, 0.94, 0.86]),
    'AntiqueWhite2' : np.array([0.93, 0.87, 0.80]),
    'AntiqueWhite3' : np.array([0.80, 0.75, 0.69]),
    'AntiqueWhite4' : np.array([0.55, 0.51, 0.47]),
    'bisque1' : np.array([1.00, 0.89, 0.77]),
    'bisque2' : np.array([0.93, 0.84, 0.72]),
    'bisque3' : np.array([0.80, 0.72, 0.62]),
    'bisque4' : np.array([0.55, 0.49, 0.42]),
    'PeachPuff1' : np.array([1.00, 0.85, 0.73]),
    'PeachPuff2' : np.array([0.93, 0.80, 0.68]),
    'PeachPuff3' : np.array([0.80, 0.69, 0.58]),
    'PeachPuff4' : np.array([0.55, 0.47, 0.40]),
    'NavajoWhite1' : np.array([1.00, 0.87, 0.68]),
    'NavajoWhite2' : np.array([0.93, 0.81, 0.63]),
    'NavajoWhite3' : np.array([0.80, 0.70, 0.55]),
    'NavajoWhite4' : np.array([0.55, 0.47, 0.37]),
    'LemonChiffon1' : np.array([1.00, 0.98, 0.80]),
    'LemonChiffon2' : np.array([0.93, 0.91, 0.75]),
    'LemonChiffon3' : np.array([0.80, 0.79, 0.65]),
    'LemonChiffon4' : np.array([0.55, 0.54, 0.44]),
    'cornsilk1' : np.array([1.00, 0.97, 0.86]),
    'cornsilk2' : np.array([0.93, 0.91, 0.80]),
    'cornsilk3' : np.array([0.80, 0.78, 0.69]),
    'cornsilk4' : np.array([0.55, 0.53, 0.47]),
    'ivory1' : np.array([1.00, 1.00, 0.94]),
    'ivory2' : np.array([0.93, 0.93, 0.88]),
    'ivory3' : np.array([0.80, 0.80, 0.76]),
    'ivory4' : np.array([0.55, 0.55, 0.51]),
    'honeydew1' : np.array([0.94, 1.00, 0.94]),
    'honeydew2' : np.array([0.88, 0.93, 0.88]),
    'honeydew3' : np.array([0.76, 0.80, 0.76]),
    'honeydew4' : np.array([0.51, 0.55, 0.51]),
    'LavenderBlush1' : np.array([1.00, 0.94, 0.96]),
    'LavenderBlush2' : np.array([0.93, 0.88, 0.90]),
    'LavenderBlush3' : np.array([0.80, 0.76, 0.77]),
    'LavenderBlush4' : np.array([0.55, 0.51, 0.53]),
    'MistyRose1' : np.array([1.00, 0.89, 0.88]),
    'MistyRose2' : np.array([0.93, 0.84, 0.82]),
    'MistyRose3' : np.array([0.80, 0.72, 0.71]),
    'MistyRose4' : np.array([0.55, 0.49, 0.48]),
    'azure1' : np.array([0.94, 1.00, 1.00]),
    'azure2' : np.array([0.88, 0.93, 0.93]),
    'azure3' : np.array([0.76, 0.80, 0.80]),
    'azure4' : np.array([0.51, 0.55, 0.55]),
    'SlateBlue1' : np.array([0.51, 0.44, 1.00]),
    'SlateBlue2' : np.array([0.48, 0.40, 0.93]),
    'SlateBlue3' : np.array([0.41, 0.35, 0.80]),
    'SlateBlue4' : np.array([0.28, 0.24, 0.55]),
    'RoyalBlue1' : np.array([0.28, 0.46, 1.00]),
    'RoyalBlue2' : np.array([0.26, 0.43, 0.93]),
    'RoyalBlue3' : np.array([0.23, 0.37, 0.80]),
    'RoyalBlue4' : np.array([0.15, 0.25, 0.55]),
    'blue1' : np.array([0.00, 0.00, 1.00]),
    'blue2' : np.array([0.00, 0.00, 0.93]),
    'blue3' : np.array([0.00, 0.00, 0.80]),
    'blue4' : np.array([0.00, 0.00, 0.55]),
    'DodgerBlue1' : np.array([0.12, 0.56, 1.00]),
    'DodgerBlue2' : np.array([0.11, 0.53, 0.93]),
    'DodgerBlue3' : np.array([0.09, 0.45, 0.80]),
    'DodgerBlue4' : np.array([0.06, 0.31, 0.55]),
    'SteelBlue1' : np.array([0.39, 0.72, 1.00]),
    'SteelBlue2' : np.array([0.36, 0.67, 0.93]),
    'SteelBlue3' : np.array([0.31, 0.58, 0.80]),
    'SteelBlue4' : np.array([0.21, 0.39, 0.55]),
    'DeepSkyBlue1' : np.array([0.00, 0.75, 1.00]),
    'DeepSkyBlue2' : np.array([0.00, 0.70, 0.93]),
    'DeepSkyBlue3' : np.array([0.00, 0.60, 0.80]),
    'DeepSkyBlue4' : np.array([0.00, 0.41, 0.55]),
    'SkyBlue1' : np.array([0.53, 0.81, 1.00]),
    'SkyBlue2' : np.array([0.49, 0.75, 0.93]),
    'SkyBlue3' : np.array([0.42, 0.65, 0.80]),
    'SkyBlue4' : np.array([0.29, 0.44, 0.55]),
    'LightSkyBlue1' : np.array([0.69, 0.89, 1.00]),
    'LightSkyBlue2' : np.array([0.64, 0.83, 0.93]),
    'LightSkyBlue3' : np.array([0.55, 0.71, 0.80]),
    'LightSkyBlue4' : np.array([0.38, 0.48, 0.55]),
    'SlateGray1' : np.array([0.78, 0.89, 1.00]),
    'SlateGray2' : np.array([0.73, 0.83, 0.93]),
    'SlateGray3' : np.array([0.62, 0.71, 0.80]),
    'SlateGray4' : np.array([0.42, 0.48, 0.55]),
    'LightSteelBlue1' : np.array([0.79, 0.88, 1.00]),
    'LightSteelBlue2' : np.array([0.74, 0.82, 0.93]),
    'LightSteelBlue3' : np.array([0.64, 0.71, 0.80]),
    'LightSteelBlue4' : np.array([0.43, 0.48, 0.55]),
    'LightBlue1' : np.array([0.75, 0.94, 1.00]),
    'LightBlue2' : np.array([0.70, 0.87, 0.93]),
    'LightBlue3' : np.array([0.60, 0.75, 0.80]),
    'LightBlue4' : np.array([0.41, 0.51, 0.55]),
    'LightCyan1' : np.array([0.88, 1.00, 1.00]),
    'LightCyan2' : np.array([0.82, 0.93, 0.93]),
    'LightCyan3' : np.array([0.71, 0.80, 0.80]),
    'LightCyan4' : np.array([0.48, 0.55, 0.55]),
    'PaleTurquoise1' : np.array([0.73, 1.00, 1.00]),
    'PaleTurquoise2' : np.array([0.68, 0.93, 0.93]),
    'PaleTurquoise3' : np.array([0.59, 0.80, 0.80]),
    'PaleTurquoise4' : np.array([0.40, 0.55, 0.55]),
    'CadetBlue1' : np.array([0.60, 0.96, 1.00]),
    'CadetBlue2' : np.array([0.56, 0.90, 0.93]),
    'CadetBlue3' : np.array([0.48, 0.77, 0.80]),
    'CadetBlue4' : np.array([0.33, 0.53, 0.55]),
    'turquoise1' : np.array([0.00, 0.96, 1.00]),
    'turquoise2' : np.array([0.00, 0.90, 0.93]),
    'turquoise3' : np.array([0.00, 0.77, 0.80]),
    'turquoise4' : np.array([0.00, 0.53, 0.55]),
    'cyan1' : np.array([0.00, 1.00, 1.00]),
    'cyan2' : np.array([0.00, 0.93, 0.93]),
    'cyan3' : np.array([0.00, 0.80, 0.80]),
    'cyan4' : np.array([0.00, 0.55, 0.55]),
    'DarkSlateGray1' : np.array([0.59, 1.00, 1.00]),
    'DarkSlateGray2' : np.array([0.55, 0.93, 0.93]),
    'DarkSlateGray3' : np.array([0.47, 0.80, 0.80]),
    'DarkSlateGray4' : np.array([0.32, 0.55, 0.55]),
    'aquamarine1' : np.array([0.50, 1.00, 0.83]),
    'aquamarine2' : np.array([0.46, 0.93, 0.78]),
    'aquamarine3' : np.array([0.40, 0.80, 0.67]),
    'aquamarine4' : np.array([0.27, 0.55, 0.45]),
    'DarkSeaGreen1' : np.array([0.76, 1.00, 0.76]),
    'DarkSeaGreen2' : np.array([0.71, 0.93, 0.71]),
    'DarkSeaGreen3' : np.array([0.61, 0.80, 0.61]),
    'DarkSeaGreen4' : np.array([0.41, 0.55, 0.41]),
    'SeaGreen1' : np.array([0.33, 1.00, 0.62]),
    'SeaGreen2' : np.array([0.31, 0.93, 0.58]),
    'SeaGreen3' : np.array([0.26, 0.80, 0.50]),
    'SeaGreen4' : np.array([0.18, 0.55, 0.34]),
    'PaleGreen1' : np.array([0.60, 1.00, 0.60]),
    'PaleGreen2' : np.array([0.56, 0.93, 0.56]),
    'PaleGreen3' : np.array([0.49, 0.80, 0.49]),
    'PaleGreen4' : np.array([0.33, 0.55, 0.33]),
    'SpringGreen1' : np.array([0.00, 1.00, 0.50]),
    'SpringGreen2' : np.array([0.00, 0.93, 0.46]),
    'SpringGreen3' : np.array([0.00, 0.80, 0.40]),
    'SpringGreen4' : np.array([0.00, 0.55, 0.27]),
    'green1' : np.array([0.00, 1.00, 0.00]),
    'green2' : np.array([0.00, 0.93, 0.00]),
    'green3' : np.array([0.00, 0.80, 0.00]),
    'green4' : np.array([0.00, 0.55, 0.00]),
    'chartreuse1' : np.array([0.50, 1.00, 0.00]),
    'chartreuse2' : np.array([0.46, 0.93, 0.00]),
    'chartreuse3' : np.array([0.40, 0.80, 0.00]),
    'chartreuse4' : np.array([0.27, 0.55, 0.00]),
    'OliveDrab1' : np.array([0.75, 1.00, 0.24]),
    'OliveDrab2' : np.array([0.70, 0.93, 0.23]),
    'OliveDrab3' : np.array([0.60, 0.80, 0.20]),
    'OliveDrab4' : np.array([0.41, 0.55, 0.13]),
    'DarkOliveGreen1' : np.array([0.79, 1.00, 0.44]),
    'DarkOliveGreen2' : np.array([0.74, 0.93, 0.41]),
    'DarkOliveGreen3' : np.array([0.64, 0.80, 0.35]),
    'DarkOliveGreen4' : np.array([0.43, 0.55, 0.24]),
    'khaki1' : np.array([1.00, 0.96, 0.56]),
    'khaki2' : np.array([0.93, 0.90, 0.52]),
    'khaki3' : np.array([0.80, 0.78, 0.45]),
    'khaki4' : np.array([0.55, 0.53, 0.31]),
    'LightGoldenrod1' : np.array([1.00, 0.93, 0.55]),
    'LightGoldenrod2' : np.array([0.93, 0.86, 0.51]),
    'LightGoldenrod3' : np.array([0.80, 0.75, 0.44]),
    'LightGoldenrod4' : np.array([0.55, 0.51, 0.30]),
    'LightYellow1' : np.array([1.00, 1.00, 0.88]),
    'LightYellow2' : np.array([0.93, 0.93, 0.82]),
    'LightYellow3' : np.array([0.80, 0.80, 0.71]),
    'LightYellow4' : np.array([0.55, 0.55, 0.48]),
    'yellow1' : np.array([1.00, 1.00, 0.00]),
    'yellow2' : np.array([0.93, 0.93, 0.00]),
    'yellow3' : np.array([0.80, 0.80, 0.00]),
    'yellow4' : np.array([0.55, 0.55, 0.00]),
    'gold1' : np.array([1.00, 0.84, 0.00]),
    'gold2' : np.array([0.93, 0.79, 0.00]),
    'gold3' : np.array([0.80, 0.68, 0.00]),
    'gold4' : np.array([0.55, 0.46, 0.00]),
    'goldenrod1' : np.array([1.00, 0.76, 0.15]),
    'goldenrod2' : np.array([0.93, 0.71, 0.13]),
    'goldenrod3' : np.array([0.80, 0.61, 0.11]),
    'goldenrod4' : np.array([0.55, 0.41, 0.08]),
    'DarkGoldenrod1' : np.array([1.00, 0.73, 0.06]),
    'DarkGoldenrod2' : np.array([0.93, 0.68, 0.05]),
    'DarkGoldenrod3' : np.array([0.80, 0.58, 0.05]),
    'DarkGoldenrod4' : np.array([0.55, 0.40, 0.03]),
    'RosyBrown1' : np.array([1.00, 0.76, 0.76]),
    'RosyBrown2' : np.array([0.93, 0.71, 0.71]),
    'RosyBrown3' : np.array([0.80, 0.61, 0.61]),
    'RosyBrown4' : np.array([0.55, 0.41, 0.41]),
    'IndianRed1' : np.array([1.00, 0.42, 0.42]),
    'IndianRed2' : np.array([0.93, 0.39, 0.39]),
    'IndianRed3' : np.array([0.80, 0.33, 0.33]),
    'IndianRed4' : np.array([0.55, 0.23, 0.23]),
    'sienna1' : np.array([1.00, 0.51, 0.28]),
    'sienna2' : np.array([0.93, 0.47, 0.26]),
    'sienna3' : np.array([0.80, 0.41, 0.22]),
    'sienna4' : np.array([0.55, 0.28, 0.15]),
    'burlywood1' : np.array([1.00, 0.83, 0.61]),
    'burlywood2' : np.array([0.93, 0.77, 0.57]),
    'burlywood3' : np.array([0.80, 0.67, 0.49]),
    'burlywood4' : np.array([0.55, 0.45, 0.33]),
    'wheat1' : np.array([1.00, 0.91, 0.73]),
    'wheat2' : np.array([0.93, 0.85, 0.68]),
    'wheat3' : np.array([0.80, 0.73, 0.59]),
    'wheat4' : np.array([0.55, 0.49, 0.40]),
    'tan1' : np.array([1.00, 0.65, 0.31]),
    'tan2' : np.array([0.93, 0.60, 0.29]),
    'tan3' : np.array([0.80, 0.52, 0.25]),
    'tan4' : np.array([0.55, 0.35, 0.17]),
    'chocolate1' : np.array([1.00, 0.50, 0.14]),
    'chocolate2' : np.array([0.93, 0.46, 0.13]),
    'chocolate3' : np.array([0.80, 0.40, 0.11]),
    'chocolate4' : np.array([0.55, 0.27, 0.07]),
    'firebrick1' : np.array([1.00, 0.19, 0.19]),
    'firebrick2' : np.array([0.93, 0.17, 0.17]),
    'firebrick3' : np.array([0.80, 0.15, 0.15]),
    'firebrick4' : np.array([0.55, 0.10, 0.10]),
    'brown1' : np.array([1.00, 0.25, 0.25]),
    'brown2' : np.array([0.93, 0.23, 0.23]),
    'brown3' : np.array([0.80, 0.20, 0.20]),
    'brown4' : np.array([0.55, 0.14, 0.14]),
    'salmon1' : np.array([1.00, 0.55, 0.41]),
    'salmon2' : np.array([0.93, 0.51, 0.38]),
    'salmon3' : np.array([0.80, 0.44, 0.33]),
    'salmon4' : np.array([0.55, 0.30, 0.22]),
    'LightSalmon1' : np.array([1.00, 0.63, 0.48]),
    'LightSalmon2' : np.array([0.93, 0.58, 0.45]),
    'LightSalmon3' : np.array([0.80, 0.51, 0.38]),
    'LightSalmon4' : np.array([0.55, 0.34, 0.26]),
    'orange1' : np.array([1.00, 0.65, 0.00]),
    'orange2' : np.array([0.93, 0.60, 0.00]),
    'orange3' : np.array([0.80, 0.52, 0.00]),
    'orange4' : np.array([0.55, 0.35, 0.00]),
    'DarkOrange1' : np.array([1.00, 0.50, 0.00]),
    'DarkOrange2' : np.array([0.93, 0.46, 0.00]),
    'DarkOrange3' : np.array([0.80, 0.40, 0.00]),
    'DarkOrange4' : np.array([0.55, 0.27, 0.00]),
    'coral1' : np.array([1.00, 0.45, 0.34]),
    'coral2' : np.array([0.93, 0.42, 0.31]),
    'coral3' : np.array([0.80, 0.36, 0.27]),
    'coral4' : np.array([0.55, 0.24, 0.18]),
    'tomato1' : np.array([1.00, 0.39, 0.28]),
    'tomato2' : np.array([0.93, 0.36, 0.26]),
    'tomato3' : np.array([0.80, 0.31, 0.22]),
    'tomato4' : np.array([0.55, 0.21, 0.15]),
    'OrangeRed1' : np.array([1.00, 0.27, 0.00]),
    'OrangeRed2' : np.array([0.93, 0.25, 0.00]),
    'OrangeRed3' : np.array([0.80, 0.22, 0.00]),
    'OrangeRed4' : np.array([0.55, 0.15, 0.00]),
    'red1' : np.array([1.00, 0.00, 0.00]),
    'red2' : np.array([0.93, 0.00, 0.00]),
    'red3' : np.array([0.80, 0.00, 0.00]),
    'red4' : np.array([0.55, 0.00, 0.00]),
    'DeepPink1' : np.array([1.00, 0.08, 0.58]),
    'DeepPink2' : np.array([0.93, 0.07, 0.54]),
    'DeepPink3' : np.array([0.80, 0.06, 0.46]),
    'DeepPink4' : np.array([0.55, 0.04, 0.31]),
    'HotPink1' : np.array([1.00, 0.43, 0.71]),
    'HotPink2' : np.array([0.93, 0.42, 0.65]),
    'HotPink3' : np.array([0.80, 0.38, 0.56]),
    'HotPink4' : np.array([0.55, 0.23, 0.38]),
    'pink1' : np.array([1.00, 0.71, 0.77]),
    'pink2' : np.array([0.93, 0.66, 0.72]),
    'pink3' : np.array([0.80, 0.57, 0.62]),
    'pink4' : np.array([0.55, 0.39, 0.42]),
    'LightPink1' : np.array([1.00, 0.68, 0.73]),
    'LightPink2' : np.array([0.93, 0.64, 0.68]),
    'LightPink3' : np.array([0.80, 0.55, 0.58]),
    'LightPink4' : np.array([0.55, 0.37, 0.40]),
    'PaleVioletRed1' : np.array([1.00, 0.51, 0.67]),
    'PaleVioletRed2' : np.array([0.93, 0.47, 0.62]),
    'PaleVioletRed3' : np.array([0.80, 0.41, 0.54]),
    'PaleVioletRed4' : np.array([0.55, 0.28, 0.36]),
    'maroon1' : np.array([1.00, 0.20, 0.70]),
    'maroon2' : np.array([0.93, 0.19, 0.65]),
    'maroon3' : np.array([0.80, 0.16, 0.56]),
    'maroon4' : np.array([0.55, 0.11, 0.38]),
    'VioletRed1' : np.array([1.00, 0.24, 0.59]),
    'VioletRed2' : np.array([0.93, 0.23, 0.55]),
    'VioletRed3' : np.array([0.80, 0.20, 0.47]),
    'VioletRed4' : np.array([0.55, 0.13, 0.32]),
    'magenta1' : np.array([1.00, 0.00, 1.00]),
    'magenta2' : np.array([0.93, 0.00, 0.93]),
    'magenta3' : np.array([0.80, 0.00, 0.80]),
    'magenta4' : np.array([0.55, 0.00, 0.55]),
    'orchid1' : np.array([1.00, 0.51, 0.98]),
    'orchid2' : np.array([0.93, 0.48, 0.91]),
    'orchid3' : np.array([0.80, 0.41, 0.79]),
    'orchid4' : np.array([0.55, 0.28, 0.54]),
    'plum1' : np.array([1.00, 0.73, 1.00]),
    'plum2' : np.array([0.93, 0.68, 0.93]),
    'plum3' : np.array([0.80, 0.59, 0.80]),
    'plum4' : np.array([0.55, 0.40, 0.55]),
    'MediumOrchid1' : np.array([0.88, 0.40, 1.00]),
    'MediumOrchid2' : np.array([0.82, 0.37, 0.93]),
    'MediumOrchid3' : np.array([0.71, 0.32, 0.80]),
    'MediumOrchid4' : np.array([0.48, 0.22, 0.55]),
    'DarkOrchid1' : np.array([0.75, 0.24, 1.00]),
    'DarkOrchid2' : np.array([0.70, 0.23, 0.93]),
    'DarkOrchid3' : np.array([0.60, 0.20, 0.80]),
    'DarkOrchid4' : np.array([0.41, 0.13, 0.55]),
    'purple1' : np.array([0.61, 0.19, 1.00]),
    'purple2' : np.array([0.57, 0.17, 0.93]),
    'purple3' : np.array([0.49, 0.15, 0.80]),
    'purple4' : np.array([0.33, 0.10, 0.55]),
    'MediumPurple1' : np.array([0.67, 0.51, 1.00]),
    'MediumPurple2' : np.array([0.62, 0.47, 0.93]),
    'MediumPurple3' : np.array([0.54, 0.41, 0.80]),
    'MediumPurple4' : np.array([0.36, 0.28, 0.55]),
    'thistle1' : np.array([1.00, 0.88, 1.00]),
    'thistle2' : np.array([0.93, 0.82, 0.93]),
    'thistle3' : np.array([0.80, 0.71, 0.80]),
    'thistle4' : np.array([0.55, 0.48, 0.55]),
    'gray0' : np.array([0.00, 0.00, 0.00]),
    'grey0' : np.array([0.00, 0.00, 0.00]),
    'gray1' : np.array([0.01, 0.01, 0.01]),
    'grey1' : np.array([0.01, 0.01, 0.01]),
    'gray2' : np.array([0.02, 0.02, 0.02]),
    'grey2' : np.array([0.02, 0.02, 0.02]),
    'gray3' : np.array([0.03, 0.03, 0.03]),
    'grey3' : np.array([0.03, 0.03, 0.03]),
    'gray4' : np.array([0.04, 0.04, 0.04]),
    'grey4' : np.array([0.04, 0.04, 0.04]),
    'gray5' : np.array([0.05, 0.05, 0.05]),
    'grey5' : np.array([0.05, 0.05, 0.05]),
    'gray6' : np.array([0.06, 0.06, 0.06]),
    'grey6' : np.array([0.06, 0.06, 0.06]),
    'gray7' : np.array([0.07, 0.07, 0.07]),
    'grey7' : np.array([0.07, 0.07, 0.07]),
    'gray8' : np.array([0.08, 0.08, 0.08]),
    'grey8' : np.array([0.08, 0.08, 0.08]),
    'gray9' : np.array([0.09, 0.09, 0.09]),
    'grey9' : np.array([0.09, 0.09, 0.09]),
    'gray10' : np.array([0.10, 0.10, 0.10]),
    'grey10' : np.array([0.10, 0.10, 0.10]),
    'gray11' : np.array([0.11, 0.11, 0.11]),
    'grey11' : np.array([0.11, 0.11, 0.11]),
    'gray12' : np.array([0.12, 0.12, 0.12]),
    'grey12' : np.array([0.12, 0.12, 0.12]),
    'gray13' : np.array([0.13, 0.13, 0.13]),
    'grey13' : np.array([0.13, 0.13, 0.13]),
    'gray14' : np.array([0.14, 0.14, 0.14]),
    'grey14' : np.array([0.14, 0.14, 0.14]),
    'gray15' : np.array([0.15, 0.15, 0.15]),
    'grey15' : np.array([0.15, 0.15, 0.15]),
    'gray16' : np.array([0.16, 0.16, 0.16]),
    'grey16' : np.array([0.16, 0.16, 0.16]),
    'gray17' : np.array([0.17, 0.17, 0.17]),
    'grey17' : np.array([0.17, 0.17, 0.17]),
    'gray18' : np.array([0.18, 0.18, 0.18]),
    'grey18' : np.array([0.18, 0.18, 0.18]),
    'gray19' : np.array([0.19, 0.19, 0.19]),
    'grey19' : np.array([0.19, 0.19, 0.19]),
    'gray20' : np.array([0.20, 0.20, 0.20]),
    'grey20' : np.array([0.20, 0.20, 0.20]),
    'gray21' : np.array([0.21, 0.21, 0.21]),
    'grey21' : np.array([0.21, 0.21, 0.21]),
    'gray22' : np.array([0.22, 0.22, 0.22]),
    'grey22' : np.array([0.22, 0.22, 0.22]),
    'gray23' : np.array([0.23, 0.23, 0.23]),
    'grey23' : np.array([0.23, 0.23, 0.23]),
    'gray24' : np.array([0.24, 0.24, 0.24]),
    'grey24' : np.array([0.24, 0.24, 0.24]),
    'gray25' : np.array([0.25, 0.25, 0.25]),
    'grey25' : np.array([0.25, 0.25, 0.25]),
    'gray26' : np.array([0.26, 0.26, 0.26]),
    'grey26' : np.array([0.26, 0.26, 0.26]),
    'gray27' : np.array([0.27, 0.27, 0.27]),
    'grey27' : np.array([0.27, 0.27, 0.27]),
    'gray28' : np.array([0.28, 0.28, 0.28]),
    'grey28' : np.array([0.28, 0.28, 0.28]),
    'gray29' : np.array([0.29, 0.29, 0.29]),
    'grey29' : np.array([0.29, 0.29, 0.29]),
    'gray30' : np.array([0.30, 0.30, 0.30]),
    'grey30' : np.array([0.30, 0.30, 0.30]),
    'gray31' : np.array([0.31, 0.31, 0.31]),
    'grey31' : np.array([0.31, 0.31, 0.31]),
    'gray32' : np.array([0.32, 0.32, 0.32]),
    'grey32' : np.array([0.32, 0.32, 0.32]),
    'gray33' : np.array([0.33, 0.33, 0.33]),
    'grey33' : np.array([0.33, 0.33, 0.33]),
    'gray34' : np.array([0.34, 0.34, 0.34]),
    'grey34' : np.array([0.34, 0.34, 0.34]),
    'gray35' : np.array([0.35, 0.35, 0.35]),
    'grey35' : np.array([0.35, 0.35, 0.35]),
    'gray36' : np.array([0.36, 0.36, 0.36]),
    'grey36' : np.array([0.36, 0.36, 0.36]),
    'gray37' : np.array([0.37, 0.37, 0.37]),
    'grey37' : np.array([0.37, 0.37, 0.37]),
    'gray38' : np.array([0.38, 0.38, 0.38]),
    'grey38' : np.array([0.38, 0.38, 0.38]),
    'gray39' : np.array([0.39, 0.39, 0.39]),
    'grey39' : np.array([0.39, 0.39, 0.39]),
    'gray40' : np.array([0.40, 0.40, 0.40]),
    'grey40' : np.array([0.40, 0.40, 0.40]),
    'gray41' : np.array([0.41, 0.41, 0.41]),
    'grey41' : np.array([0.41, 0.41, 0.41]),
    'gray42' : np.array([0.42, 0.42, 0.42]),
    'grey42' : np.array([0.42, 0.42, 0.42]),
    'gray43' : np.array([0.43, 0.43, 0.43]),
    'grey43' : np.array([0.43, 0.43, 0.43]),
    'gray44' : np.array([0.44, 0.44, 0.44]),
    'grey44' : np.array([0.44, 0.44, 0.44]),
    'gray45' : np.array([0.45, 0.45, 0.45]),
    'grey45' : np.array([0.45, 0.45, 0.45]),
    'gray46' : np.array([0.46, 0.46, 0.46]),
    'grey46' : np.array([0.46, 0.46, 0.46]),
    'gray47' : np.array([0.47, 0.47, 0.47]),
    'grey47' : np.array([0.47, 0.47, 0.47]),
    'gray48' : np.array([0.48, 0.48, 0.48]),
    'grey48' : np.array([0.48, 0.48, 0.48]),
    'gray49' : np.array([0.49, 0.49, 0.49]),
    'grey49' : np.array([0.49, 0.49, 0.49]),
    'gray50' : np.array([0.50, 0.50, 0.50]),
    'grey50' : np.array([0.50, 0.50, 0.50]),
    'gray51' : np.array([0.51, 0.51, 0.51]),
    'grey51' : np.array([0.51, 0.51, 0.51]),
    'gray52' : np.array([0.52, 0.52, 0.52]),
    'grey52' : np.array([0.52, 0.52, 0.52]),
    'gray53' : np.array([0.53, 0.53, 0.53]),
    'grey53' : np.array([0.53, 0.53, 0.53]),
    'gray54' : np.array([0.54, 0.54, 0.54]),
    'grey54' : np.array([0.54, 0.54, 0.54]),
    'gray55' : np.array([0.55, 0.55, 0.55]),
    'grey55' : np.array([0.55, 0.55, 0.55]),
    'gray56' : np.array([0.56, 0.56, 0.56]),
    'grey56' : np.array([0.56, 0.56, 0.56]),
    'gray57' : np.array([0.57, 0.57, 0.57]),
    'grey57' : np.array([0.57, 0.57, 0.57]),
    'gray58' : np.array([0.58, 0.58, 0.58]),
    'grey58' : np.array([0.58, 0.58, 0.58]),
    'gray59' : np.array([0.59, 0.59, 0.59]),
    'grey59' : np.array([0.59, 0.59, 0.59]),
    'gray60' : np.array([0.60, 0.60, 0.60]),
    'grey60' : np.array([0.60, 0.60, 0.60]),
    'gray61' : np.array([0.61, 0.61, 0.61]),
    'grey61' : np.array([0.61, 0.61, 0.61]),
    'gray62' : np.array([0.62, 0.62, 0.62]),
    'grey62' : np.array([0.62, 0.62, 0.62]),
    'gray63' : np.array([0.63, 0.63, 0.63]),
    'grey63' : np.array([0.63, 0.63, 0.63]),
    'gray64' : np.array([0.64, 0.64, 0.64]),
    'grey64' : np.array([0.64, 0.64, 0.64]),
    'gray65' : np.array([0.65, 0.65, 0.65]),
    'grey65' : np.array([0.65, 0.65, 0.65]),
    'gray66' : np.array([0.66, 0.66, 0.66]),
    'grey66' : np.array([0.66, 0.66, 0.66]),
    'gray67' : np.array([0.67, 0.67, 0.67]),
    'grey67' : np.array([0.67, 0.67, 0.67]),
    'gray68' : np.array([0.68, 0.68, 0.68]),
    'grey68' : np.array([0.68, 0.68, 0.68]),
    'gray69' : np.array([0.69, 0.69, 0.69]),
    'grey69' : np.array([0.69, 0.69, 0.69]),
    'gray70' : np.array([0.70, 0.70, 0.70]),
    'grey70' : np.array([0.70, 0.70, 0.70]),
    'gray71' : np.array([0.71, 0.71, 0.71]),
    'grey71' : np.array([0.71, 0.71, 0.71]),
    'gray72' : np.array([0.72, 0.72, 0.72]),
    'grey72' : np.array([0.72, 0.72, 0.72]),
    'gray73' : np.array([0.73, 0.73, 0.73]),
    'grey73' : np.array([0.73, 0.73, 0.73]),
    'gray74' : np.array([0.74, 0.74, 0.74]),
    'grey74' : np.array([0.74, 0.74, 0.74]),
    'gray75' : np.array([0.75, 0.75, 0.75]),
    'grey75' : np.array([0.75, 0.75, 0.75]),
    'gray76' : np.array([0.76, 0.76, 0.76]),
    'grey76' : np.array([0.76, 0.76, 0.76]),
    'gray77' : np.array([0.77, 0.77, 0.77]),
    'grey77' : np.array([0.77, 0.77, 0.77]),
    'gray78' : np.array([0.78, 0.78, 0.78]),
    'grey78' : np.array([0.78, 0.78, 0.78]),
    'gray79' : np.array([0.79, 0.79, 0.79]),
    'grey79' : np.array([0.79, 0.79, 0.79]),
    'gray80' : np.array([0.80, 0.80, 0.80]),
    'grey80' : np.array([0.80, 0.80, 0.80]),
    'gray81' : np.array([0.81, 0.81, 0.81]),
    'grey81' : np.array([0.81, 0.81, 0.81]),
    'gray82' : np.array([0.82, 0.82, 0.82]),
    'grey82' : np.array([0.82, 0.82, 0.82]),
    'gray83' : np.array([0.83, 0.83, 0.83]),
    'grey83' : np.array([0.83, 0.83, 0.83]),
    'gray84' : np.array([0.84, 0.84, 0.84]),
    'grey84' : np.array([0.84, 0.84, 0.84]),
    'gray85' : np.array([0.85, 0.85, 0.85]),
    'grey85' : np.array([0.85, 0.85, 0.85]),
    'gray86' : np.array([0.86, 0.86, 0.86]),
    'grey86' : np.array([0.86, 0.86, 0.86]),
    'gray87' : np.array([0.87, 0.87, 0.87]),
    'grey87' : np.array([0.87, 0.87, 0.87]),
    'gray88' : np.array([0.88, 0.88, 0.88]),
    'grey88' : np.array([0.88, 0.88, 0.88]),
    'gray89' : np.array([0.89, 0.89, 0.89]),
    'grey89' : np.array([0.89, 0.89, 0.89]),
    'gray90' : np.array([0.90, 0.90, 0.90]),
    'grey90' : np.array([0.90, 0.90, 0.90]),
    'gray91' : np.array([0.91, 0.91, 0.91]),
    'grey91' : np.array([0.91, 0.91, 0.91]),
    'gray92' : np.array([0.92, 0.92, 0.92]),
    'grey92' : np.array([0.92, 0.92, 0.92]),
    'gray93' : np.array([0.93, 0.93, 0.93]),
    'grey93' : np.array([0.93, 0.93, 0.93]),
    'gray94' : np.array([0.94, 0.94, 0.94]),
    'grey94' : np.array([0.94, 0.94, 0.94]),
    'gray95' : np.array([0.95, 0.95, 0.95]),
    'grey95' : np.array([0.95, 0.95, 0.95]),
    'gray96' : np.array([0.96, 0.96, 0.96]),
    'grey96' : np.array([0.96, 0.96, 0.96]),
    'gray97' : np.array([0.97, 0.97, 0.97]),
    'grey97' : np.array([0.97, 0.97, 0.97]),
    'gray98' : np.array([0.98, 0.98, 0.98]),
    'grey98' : np.array([0.98, 0.98, 0.98]),
    'gray99' : np.array([0.99, 0.99, 0.99]),
    'grey99' : np.array([0.99, 0.99, 0.99]),
    'gray100' : np.array([1.00, 1.00, 1.00]),
    'grey100' : np.array([1.00, 1.00, 1.00]),
    'dark grey' : np.array([0.66, 0.66, 0.66]),
    'DarkGrey' : np.array([0.66, 0.66, 0.66]),
    'dark gray' : np.array([0.66, 0.66, 0.66]),
    'DarkGray' : np.array([0.66, 0.66, 0.66]),
    'dark blue' : np.array([0.00, 0.00, 0.55]),
    'DarkBlue' : np.array([0.00, 0.00, 0.55]),
    'dark cyan' : np.array([0.00, 0.55, 0.55]),
    'DarkCyan' : np.array([0.00, 0.55, 0.55]),
    'dark magenta' : np.array([0.55, 0.00, 0.55]),
    'DarkMagenta' : np.array([0.55, 0.00, 0.55]),
    'dark red' : np.array([0.55, 0.00, 0.00]),
    'DarkRed' : np.array([0.55, 0.00, 0.00]),
    'light green' : np.array([0.56, 0.93, 0.56]),
    'LightGreen' : np.array([0.56, 0.93, 0.56])
}

def main():
    """ Generates code for name_to_rgb dict, assuming an rgb.txt file available (in X11 format)."""
    import re
    with open('rgb.txt') as fp:
        line = fp.readline()
        while line:
            reg = re.match(r'\s*(\d+)\s*(\d+)\s*(\d+)\s*(\w.*\w).*', line)
            if reg:
                r = int(reg.group(1)) / 255.
                g = int(reg.group(2)) / 255.
                b = int(reg.group(3)) / 255.
                d = reg.group(4)
                print("'%s' : np.array([%.2f, %.2f, %.2f])," % (d, r, g, b))
            line = fp.readline()

if __name__ == '__main__':
    main()


import numpy as np
import dent.assets
import procgen.Rock
import dent.Shaders as Shaders
import dent.transforms
import OpenGL.GL as gl

shader = Shaders.getShader('rock', instance=True)

class Rock(object):
  def __init__(self, instance_buffer, instance_count, detail):
    self.position = np.array([0,0.5,0],dtype=float)

    self.shader = shader

    self.meshdata, self.meshindices = dent.assets.getAsset('rock'+str(detail),
        procgen.Rock.get_rock_mesh, (1, detail))

    self.instance_template = np.zeros(0, dtype=[("model", np.float32, (4, 4))])

    self.renderID = self.shader.setData(self.meshdata,
                                        self.meshindices,
                                        self.instance_template,
                                        instance_buffer)
    self.instance_count = instance_count

  def display(self):
    if self.renderID is None:
      return
    self.shader.draw(gl.GL_TRIANGLES, self.renderID, self.instance_count)


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

# # 剽窃检测模型
#
# 创建了训练和测试数据后，可以定义和训练模型了。在此 notebook 中，你的目标是训练一个二元分类模型，它会根据你提供的特征学习将答案文件标为剽窃文件或非剽窃文件。
#
# 此任务将分成以下几个步骤：
#
# * 将数据上传到 S3。
# * 定义一个二元分类模型和训练脚本。
# * 训练和部署模型。
# * 评估部署的分类器并回答关于所采用方法的一些问题。
#
# 要完成此 notebook，你需要完成此 notebook 中的所有练习并回答所有问题。
# > 所有任务将清晰地标为**练习**，问题都标为**问题**。
#
# 你可以尝试不同的分类模型，并选择一个在此数据集上效果最佳的模型。
#
# ---

# ## 将数据上传到 S3
#
# 在上个 notebook 中，你应该使用给定剽窃/非剽窃文本数据语料库的特征和类别标签创建了两个文件：`training.csv` 和 `test.csv` 文件。
#
# >以下单元格将加载一些 AWS SageMaker 库并创建一个默认存储桶。创建此存储桶后，你可以将本地存储的数据上传到 S3。
#
# 将训练和测试 `.csv` 特征文件保存到本地。你可以在 SageMaker 中运行第二个 notebook“2_Plagiarism_Feature_Engineering”，或者使用 Jupyter Lab 中的上传图标手动将文件上传到此 notebook。然后，你可以使用 `sagemaker_session.upload_data` 将本地文件上传到 S3，并直接指向训练数据的存储位置。

# In[1]:

import pandas as pd
import boto3
import sagemaker

# In[33]:

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
# session and role
sagemaker_session = sagemaker.Session()
role = sagemaker.get_execution_role()

# create an S3 bucket
bucket = sagemaker_session.default_bucket()

# ## 练习：将训练数据上传到 S3
#
# 指定在其中保存了 `train.csv` 文件的 `data_dir`。指定一个描述性 `prefix`，指出数据将上传到默认 S3 存储桶的什么位置。最后，通过调用 `sagemaker_session.upload_data` 并传入必要的参数，创建一个指向训练数据的指针。建议参考 [Session 文档](https://sagemaker.readthedocs.io/en/stable/session.html#sagemaker.session.Session.upload_data)或之前的 SageMaker 代码示例。
#
# 你需要上传整个目录。之后，训练脚本将仅访问 `train.csv` 文件。

# In[34]:

# should be the name of directory you created to save your features data
data_dir =  'plagiarism_data'

# set prefix, a descriptive name for a directory
prefix = 'plagiarism_detection'

# upload all data to S3
input_data = sagemaker_session.upload_data(path= data_dir, bucket=bucket, key_prefix = prefix)

# ### 测试单元格
#
# 测试数据是否已成功上传。以下单元格将输出 S3 存储桶中的内容，如果为空，将抛出错误。你应该看到 `data_dir` 的内容，或许还有一些检查点。如果你看到其中列出了任何其他文件，那么你也许有一些旧的模型文件，你可以通过 S3 控制台删除这些旧文件（不过多余的文件应该不会影响在此 notebook 中开发的模型的性能）。

# In[35]:

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
# confirm that data is in S3 bucket
empty_check = []
for obj in boto3.resource('s3').Bucket(bucket).objects.all():
    empty_check.append(obj.key)
    print(obj.key)

assert len(empty_check) !=0, 'S3 bucket is empty.'
print('Test passed!')

# ---
#
# # 建模
#
# 上传训练数据后，下面定义并训练模型。
#
# 你可以决定创建什么类型的模型。对于二元分类任务，你可以选择采用以下三种方法之一：
# * 使用内置的分类算法，例如 LinearLearner。
# * 定义自定义 Scikit-learn 分类器，可以在[此处](https://scikit-learn.org/stable/auto_examples/classification/plot_classifier_comparison.html)找到各个模型的比较情况。
# * 定义自定义 PyTorch 神经网络分类器。
#
# 你需要测试各种模型并选择最佳模型。我们将根据最终模型的准确率对你的项目打分。
#
# ---
#
# ## 练习：完成训练脚本
#
# 为了实现自定义分类器，你需要完成 `train.py` 脚本。我们提供了文件夹 `source_sklearn` 和 `source_pytorch`，其中分别包含自定义 Scikit-learn 模型和 PyTorch 模型的起始代码。每个目录都有一个 `train.py` 训练脚本。要完成此项目，**你只需完成其中一个脚本**，即负责训练最终模型的脚本。
#
# 典型的训练脚本会：
# * 从指定的目录加载训练数据
# * 解析所有的训练和模型超参数（例如神经网络中的节点数，训练周期，等等）
# * 实例化你设计的模型，并采用指定的超参数
# * 训练该模型
# * 最后，保存模型，以便之后托管/部署模型
#
# ### 定义和训练模型
# 我们已经提供了大部分训练脚本。几乎所有任务都位于 `if __name__ == '__main__':` 部分。为了完成 `train.py` 文件，你需要：
# 1. 导入所需的任何额外库
# 2. 使用 `parser.add_argument` 定义任何其他模型训练超参数
# 2. 在 `if __name__ == '__main__':` 部分定义模型
# 3. 在此部分训练模型
#
# 你可以在下面使用 `!pygmentize` 显示现有的 `train.py` 文件。请通读代码，所有任务都标有 `TODO` 注释。
#
# **注意：如果你选择创建自定义 PyTorch 模型，需要在 `model.py` 文件中定义模型**，并且我们提供了 `predict.py` 文件。如果你选择使用 Scikit-learn，则只需 `train.py` 文件；你可以从 `sklearn` 库中导入一个分类器。

# In[36]:

# directory can be changed to: source_sklearn or source_pytorch
get_ipython().system('pygmentize source_sklearn/train.py')

# ### 提供的代码
#
# 如果你阅读了上述代码，就会发现起始代码包含：
# * 模型加载 (`model_fn`) 和保存代码
# * 获取 SageMaker 的默认超参数
# * 按照名称 `train.csv` 加载训练数据，并提取特征和标签 `train_x` 和 `train_y`
#
# 如果你想详细了解如何通过 [joblib for sklearn](https://scikit-learn.org/stable/modules/model_persistence.html) 或 [torch.save](https://pytorch.org/tutorials/beginner/saving_loading_models.html) 保存模型，请点击提供的链接。

# ---
# # 创建评估器
#
# 在 SageMaker 中构建自定义模型时，必须指定入口点。入口点是一个 Python 文件，当模型被训练时，该文件将执行，即你在上面指定的 `train.py` 函数。要在 SageMaker 中运行自定义训练脚本，你需要构建评估器并指定相应的构造函数参数：
#
# * **entry_point**：SageMaker 训练模型和预测时运行的 Python 脚本的路径。
# * **source_dir**：训练脚本目录 `source_sklearn` 或 `source_pytorch` 的路径。
# * **entry_point**：SageMaker 训练模型和预测时运行的 Python 脚本的路径。
# * **source_dir**：训练脚本目录 `train_sklearn` 或 `train_pytorch` 的路径。
# * **entry_point**：SageMaker 训练模型时运行的 Python 脚本的路径。
# * **source_dir**：训练脚本目录 `train_sklearn` 或 `train_pytorch` 的路径。
# * **role**：角色 ARN，在上面已指定。
# * **train_instance_count**：训练实例的数量（应该保留为 1）。
# * **train_instance_type**：SageMaker 训练实例的类型。注意，因为 Scikit-learn 不提供 GPU 训练原生支持，所以 Sagemaker Scikit-learn 目前不支持在 GPU 实例上训练模型。
# * **sagemaker_session**：在 Sagemaker 中训练时使用的会话。
# * **hyperparameters**（可选）：作为超参数传递给训练函数的字典 `{'name':value, ..}`。
#
# 注意：对于 PyTorch 模型，还有一个可选参数 **framework_version**，你可以将其设为最新的 PyTorch 版本 `1.0`。
#
# ## 练习：定义 Scikit-learn 或 PyTorch 评估器
#
# 你可以使用以下命令之一导入一个评估器：
# ```
# from sagemaker.sklearn.estimator import SKLearn
# ```
# ```
# from sagemaker.pytorch import PyTorch
# ```

# In[37]:

# your import and estimator code, here
from sagemaker.sklearn.estimator import SKLearn

estimator = SKLearn(entry_point="train.py",
                    source_dir="source_sklearn",
                    role=role,
                    train_instance_count=1,
                    train_instance_type='ml.c4.xlarge')

# ## 练习：训练评估器
#
# 使用在 S3 中存储的训练数据训练评估器。代码应该创建一个训练作业，你可以在 SageMaker 控制台中监控该作业。

# In[38]:

get_ipython().run_cell_magic('time', '', "\n# Train your estimator on S3 training data\n\nestimator.fit({'train': input_data})")

# ## 练习：部署训练过的模型
#
# 训练之后，部署模型以创建 `predictor`。如果你使用的是 PyTorch 模型，你需要创建一个训练过的 `PyTorchModel`，它会接受训练过的 `<model>.model_data` 作为输入参数，并指向提供的 `source_pytorch/predict.py` 文件作为入口点。
#
# 为了部署训练过的模型，你需要使用 `<model>.deploy`，它接受两个参数：
# * **initial_instance_count**：部署实例的数量 (1)。
# * **instance_type**：部署 SageMaker 实例的类型。
#
# 注意：如果你遇到实例错误，可能是因为你选择了错误的训练或部署实例类型。建议参考之前的练习代码，看看我们使用了哪种类型的实例。

# In[39]:

get_ipython().run_cell_magic('time', '', "\n# uncomment, if needed\n# from sagemaker.pytorch import PyTorchModel\n\n\n# deploy your model to create a predictor\npredictor = estimator.deploy(initial_instance_count=1, instance_type='ml.t2.medium')")

# ---
# # 评估模型
#
# 模型部署后，你可以将模型应用到测试数据上，看看模型的效果如何。
#
# 下面提供的单元格会读入测试数据，并假设它存储在本地 `data_dir` 目录下，名称为 `test.csv`。标签和特征是从 `.csv` 文件提取的。

# In[40]:

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
import os

# read in test data, assuming it is stored locally
test_data = pd.read_csv(os.path.join(data_dir, "test.csv"), header=None, names=None)

# labels are in the first column
test_y = test_data.iloc[:,0]
test_x = test_data.iloc[:,1:]

# ## 练习：确定模型的准确率
#
# 使用部署的 `predictor` 为测试数据预测类别标签。将这些标签与真实标签 `test_y` 进行比较，并计算 0-1 之间的准确率，表示模型分类正确的测试数据所占的比例。你可以使用 [sklearn.metrics](https://scikit-learn.org/stable/modules/classes.html#module-sklearn.metrics) 计算准确率。
#
# **要通过此项目，你的模型测试准确率应至少达到 90%。**

# In[41]:

# First: generate predicted, class labels
test_y_preds = predictor.predict(test_x)

"""
DON'T MODIFY ANYTHING IN THIS CELL THAT IS BELOW THIS LINE
"""
# test that your model generates the correct number of labels
assert len(test_y_preds)==len(test_y), 'Unexpected number of predictions.'
print('Test passed!')

# In[42]:

from sklearn.metrics import accuracy_score

# Second: calculate the test accuracy
accuracy = accuracy_score(test_y, test_y_preds)

print(accuracy)

## print out the array of predicted and true labels, if you want
print('\nPredicted class labels: ')
print(test_y_preds)
print('\nTrue class labels: ')
print(test_y.values)

# ### 问题 1：你的模型生成了多少个假正例和假负例？为何会这样？

# **回答：**
#
# 假正例（False Positive）指的是预测为1，实际为0的样本；假负例（False Negative）指的是预测为0，实际为1的样本。我的模型生成了0个假正例，0个假负例。这可能是因为数据量较小，模型能够很好的拟合。

# ### 问题 2：你是如何决定要使用什么类型的模型？

# **回答：**
#
# 因为这个是一个分类模型，因此可以选择决策树，支持向量机，逻辑回归等模型。在这里数据集很小，因此决定采用支持向量机分类。

# ----
# ## 练习：清理资源
#
# 评估完模型后，记得**删除模型端点**。你可以通过调用 `.delete_endpoint()` 删除端点。你需要在此 notebook 中演示端点已删除。你可以从 AWS 控制台删除任何其他资源，并且在下面找到更多关于删除所有资源的说明。

# In[43]:

# uncomment and fill in the line below!
# <name_of_deployed_predictor>.delete_endpoint()

predictor.delete_endpoint()

# ### 删除 S3 存储桶
#
# 完全训练和测试完模型后，你可以删除整个 S3 存储桶。如果你在训练模型之前删除存储桶，需要重新创建 S3 存储桶并将训练数据再次上传到存储桶中。

# In[44]:

# deleting bucket, uncomment lines below

bucket_to_delete = boto3.resource('s3').Bucket(bucket)
bucket_to_delete.objects.all().delete()

# ### 删除所有模型和实例
#
# 当你完全处理完模型，并且**不**需要重新访问此 notebook，你可以根据[这些说明](https://docs.aws.amazon.com/sagemaker/latest/dg/ex1-cleanup.html)删除所有 SageMaker notebook 实例和模型。在删除此 notebook 实例之前，建议至少下载一个副本并将其保存到本地。

# ---
# ## 后续改进建议
#
# 此项目还有很多改进或扩展方式，可以拓宽你的知识面，或让此项目更独特。下面是一些建议：
# * 训练分类器预测剽窃类别 (1-3)，而不仅仅是剽窃 (1) 或非剽窃 (0)。
# * 利用其他更大型的数据集检测此模型能否扩展到其他类型的剽窃行为。
# * 利用语言或字符级分析寻找不同（及更多）相似性特征。
# * 编写完整的管道函数，它会接受原文和提交的文本，并将提交的文本分类为剽窃或非剽窃文本。
# * 使用 API Gateway 和 lambda 函数将模型部署到网络应用上。
#
# 这些都只是扩展项目的建议。如果你完成了此 notebook 中的所有练习，你已经完成了一个真实的应用，可以提交项目了。棒棒哒！


import numpy as np
from tensorboardX import SummaryWriter
from MyDQN import vrep
import time
import random
import cv2 as cv

image_pix = 224  # 输入图像的维度 image_pix * image_pix 灰度图

class EnvGrasp(object):
    def __init__(self):
        self.total_success = 0
        self.total_try = 0
        self.writer = SummaryWriter('./logs_grasp_move_rate')
        self.object_work_space = np.asarray([[-0.7, -0.3], [-0.2, 0.2], [0.01, 0.31]])

        self.work_space = np.asarray([[-0.75, -0.25], [-0.25, 0.25], [0.01, 0.31]])

        vrep.simxFinish(-1)
        self.sim_client = vrep.simxStart('127.0.0.1', 19997, True, True, 5000, 5)  # Connect to V-REP on port 19997

        self.current_state = []     # 2d

        self.target = []    # 3d
        self.target_orientation = []
        self.target_handle = 0

        self.current_handle = 0
        self.current_position = [0, 0, 0]   # 3d
        self.current_orientation = []

        self.pre_distance = 0
        self.pre_orientation_distance = 0

        self.correct_count = 0

        if self.sim_client == -1:
            print('Failed to connect to simulation (V-REP remote API server). Exiting.')
            exit()
        else:
            print('Connected to simulation.')

    def get_state(self):
        sim_ret, self.current_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target',
                                                                vrep.simx_opmode_blocking)
        sim_ret, self.current_position = vrep.simxGetObjectPosition(self.sim_client, self.current_handle, -1,
                                                                    vrep.simx_opmode_blocking)
        sim_ret, self.current_orientation = vrep.simxGetObjectOrientation(self.sim_client, self.current_handle, -1, vrep.simx_opmode_blocking)
        sim_ret, self.target_handle = vrep.simxGetObjectHandle(self.sim_client, 'shape001', vrep.simx_opmode_blocking)
        sim_ret, self.target = vrep.simxGetObjectPosition(self.sim_client, self.target_handle, -1, vrep.simx_opmode_blocking)

        sim_ret, self.target_orientation = vrep.simxGetObjectOrientation(self.sim_client, self.target_handle, -1, vrep.simx_opmode_blocking)

        self.target = np.asarray(self.target)
        self.current_position = np.asarray(self.current_position)

        if self.target_orientation[1] < 0:
            ratio = -1
        else:
            ratio = 1

        orientation_dis = (self.current_orientation[1] - ratio * (np.pi/2 - abs(self.target_orientation[1]))) / (np.pi)
        self.current_state = self.current_position - self.target
        self.current_state = np.append(self.current_state, orientation_dis)
        # print(orientation_dis)
        # print('current position : ', self.current_position, 'target position : ', self.target, 'target orientation : ', self.target_orientation, 'current orientation : ', self.current_orientation)
        # print('current state : ', self.current_state)
        frame = self.get_sensor_data()
        # pos = self.current_position
        pos = np.append(self.current_position, self.current_orientation[1])
        return self.current_state, frame, pos

    def get_reward(self, origin_action):
        done = 0
        if self.current_position[0] < self.work_space[0][0] or self.current_position[0] > self.work_space[0][1] \
                        or self.current_position[1] < self.work_space[1][0] or self.current_position[1] > self.work_space[1][1] \
                        or self.current_position[2] < self.work_space[2][0] or self.current_position[2] > self.work_space[2][1]:
            print('arm is out of workspace!')
            done = 1
            self.correct_count = 0
            return [-1, -1], done
        # print('self.current_state : ', self.current_state)

        cs = self.current_state[0:2]
        distance = np.linalg.norm(cs)
        cs = np.append(cs, self.current_state[-1])

        orientation_distance = np.linalg.norm(self.current_state[-1])
        # print('distance : ', distance, ' predistance : ', self.pre_distance)
        # reward = -distance

        # distance = np.linalg.norm(origin_action[0:2])
        # print('distance : ', distance)

        reward = [-distance, -orientation_distance/1.5]
        # print(distance)
        total_distance = np.linalg.norm(cs)
        if distance < self.pre_distance:
            reward[0] += 0.3
            # reward = 0
        else:
            pass
        if orientation_distance < self.pre_orientation_distance:
            reward[1] += 0.3
        else:
            reward[1] -= 0.3
        if distance < 0.015:
            # print('distance reached !')
            reward[0] += 0.5
        if orientation_distance < 0.015:
            print('orientation reached !')
            reward[1] += 0.5
        if total_distance < 0.015:
            self.correct_count += 1
            print('reached !')
            # reward[0] += 0.5
            # reward[1] += 0.5
            # reward = 0
        else:
            self.correct_count = 0
        if self.correct_count >= 5:
            # grasp
            success = self.grasp()
            if success:
                print('grasp success ! ')
                reward[0] = 1
                reward[1] = 1
            else:
                print('grasp failed !')
                reward[0] = -1
                reward[1] = -1
                self.correct_count = 0
            done = 1
        self.pre_distance = distance
        self.pre_orientation_distance = orientation_distance
        return reward, done

    def reset(self):
        sim_ret, self.UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target',
                                                                   vrep.simx_opmode_blocking)
        vrep.simxStopSimulation(self.sim_client, vrep.simx_opmode_blocking)
        vrep.simxStartSimulation(self.sim_client, vrep.simx_opmode_blocking)
        time.sleep(1)
        sim_ret, self.RG2_tip_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_tip', vrep.simx_opmode_blocking)
        sim_ret, gripper_position = vrep.simxGetObjectPosition(self.sim_client, self.RG2_tip_handle, -1,
                                                               vrep.simx_opmode_blocking)

        time.sleep(1)
        vrep.simxSetObjectPosition(self.sim_client, self.UR5_target_handle, -1, (-0.47, 0, 0.254),
                                   vrep.simx_opmode_blocking)
        vrep.simxSetObjectOrientation(self.sim_client, self.UR5_target_handle, -1, (np.pi / 2, 0, np.pi / 2),
                                      vrep.simx_opmode_blocking)
        while gripper_position[2] > 0.4:  # V-REP bug requiring multiple starts and stops to restart
            self.open_griper()
            vrep.simxStopSimulation(self.sim_client, vrep.simx_opmode_blocking)
            vrep.simxStartSimulation(self.sim_client, vrep.simx_opmode_blocking)
            time.sleep(1)
            sim_ret, gripper_position = vrep.simxGetObjectPosition(self.sim_client, self.RG2_tip_handle, -1,
                                                                   vrep.simx_opmode_blocking)
            vrep.simxSetObjectPosition(self.sim_client, self.UR5_target_handle, -1, (-0.47, 0, 0.254),
                                       vrep.simx_opmode_blocking)
            vrep.simxSetObjectOrientation(self.sim_client, self.UR5_target_handle, -1, (np.pi / 2, 0, np.pi / 2),
                                          vrep.simx_opmode_blocking)
        # print('started!')

        target_x = random.random() * (self.object_work_space[0][1] - self.object_work_space[0][0]) + \
                   self.object_work_space[0][0]
        target_y = random.random() * (self.object_work_space[1][1] - self.object_work_space[1][0]) + \
                   self.object_work_space[1][0]
        # target_z = random.random() * (self.work_space[2][1] - self.work_space[2][0]) + self.work_space[2][0]
        # target_x = -0.5283
        # target_y = 0.1640

        target_z = 0.025
        self.open_griper()
        self.move_position([target_x + 0.47, target_y, 0.03 - 0.254])
        self.target = [target_x, target_y, target_z]

        orientation_y = -random.random() * np.pi + np.pi/2

        # print()
        # vrep.simxPauseCommunication(self.sim_client, False)
        object_orientation = [-np.pi/2, orientation_y, -np.pi/2]
        curr_mesh_file = '/home/chen/stl-model/surface_car.obj'
        curr_shape_name = 'shape001'
        object_color = [0.4, 0.7, 1]
        ret_resp, ret_ints, ret_floats, ret_strings, ret_buffer = vrep.simxCallScriptFunction(self.sim_client,
                                                                                'remoteApiCommandServer',
                                                                                vrep.sim_scripttype_childscript,
                                                                                'importShape',
                                                                                [0, 0, 255, 0],
                                                                                self.target + object_orientation + object_color,
                                                                                [curr_mesh_file,
                                                                                curr_shape_name],
                                                                                bytearray(),
                                                                                vrep.simx_opmode_blocking)
        self.target = np.asarray(self.target)
        self.open_griper()
        sim_ret, self.current_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target',
                                                                vrep.simx_opmode_blocking)
        sim_ret, self.current_position = vrep.simxGetObjectPosition(self.sim_client, self.current_handle, -1,
                                                                    vrep.simx_opmode_blocking)
        if orientation_y < 0:
            ratio = -1
        else:
            ratio = 1
        # target_state = np.append(self.target[0:2], ratio * (np.pi/2 - abs(orientation_y))/(np.pi))
        # current_state = np.append(self.current_position[0:2], 0)
        self.pre_distance = np.linalg.norm(self.target[0:2] - self.current_position[0:2])
        self.pre_orientation_distance = np.linalg.norm(ratio * (np.pi/2 - abs(orientation_y))/(np.pi))
        self.current_state, frame, pos = self.get_state()
        # test_pos = np.append([0, 0, 0], orientation_y - np.pi/2)
        # self.move_to(test_pos)
        cs = self.current_state[0:2]

        cs = np.append(cs, self.current_state[-1])

        return cs, frame, pos

    def open_griper(self):
        gripper_motor_velocity = 0.5
        gripper_motor_force = 20
        sim_ret, RG2_gripper_handle = vrep.simxGetObjectHandle(self.sim_client, 'RG2_openCloseJoint',
                                                               vrep.simx_opmode_blocking)
        sim_ret, gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                    vrep.simx_opmode_blocking)
        vrep.simxSetJointForce(self.sim_client, RG2_gripper_handle, gripper_motor_force, vrep.simx_opmode_blocking)
        vrep.simxSetJointTargetVelocity(self.sim_client, RG2_gripper_handle, gripper_motor_velocity,
                                        vrep.simx_opmode_blocking)
        gripper_fully_opened = False
        while gripper_joint_position < 0.0536:  # Block until gripper is fully open
            sim_ret, new_gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                        vrep.simx_opmode_blocking)
            if new_gripper_joint_position <= gripper_joint_position:
                return gripper_fully_opened
            gripper_joint_position = new_gripper_joint_position
        gripper_fully_opened = True
        return gripper_fully_opened

    def move_to(self, tool_position):
        sim_ret, UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        # UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client,'UR5_target',vrep.simx_opmode_blocking)
        sim_ret, UR5_target_position = vrep.simxGetObjectPosition(self.sim_client, UR5_target_handle, -1,
                                                                  vrep.simx_opmode_blocking)
        sim_ret, UR5_target_orientation = vrep.simxGetObjectOrientation(self.sim_client, UR5_target_handle, -1, vrep.simx_opmode_blocking)
        # print(UR5_target_position)
        tool_orientation = tool_position[-1]
        tool_position = tool_position[0:3]
        if UR5_target_position[2] <= 0.031:
            tool_position[2] = 0
        move_direction = np.asarray(tool_position)
        move_magnitude = np.linalg.norm(move_direction)
        if move_magnitude == 0 or not move_magnitude == move_magnitude:
            move_step = [0, 0, 0]
            num_move_steps = 0
            # print('magnitude error~!', move_magnitude, tool_position)
        else:
            move_step = 0.005 * move_direction / move_magnitude
            num_move_steps = int(np.floor(move_magnitude / 0.005))

        rotation_step = 0.05 if (tool_orientation > 0) else -0.05
        num_rotation_steps = int(np.floor(tool_orientation / rotation_step))

        # print('move direction : ', move_direction, 'move magnitude : ', move_magnitude, 'move step : ', move_step, 'num : ', num_move_steps)

        for step_iter in range(max(num_move_steps, num_rotation_steps)):
            vrep.simxSetObjectPosition(self.sim_client, UR5_target_handle, -1, (UR5_target_position[0] +
                                                                                move_step[0] * min(step_iter,
                                                                                num_move_steps),
                                                                                UR5_target_position[1] +
                                                                                move_step[1] * min(step_iter,
                                                                                num_move_steps),
                                                                                UR5_target_position[2] +
                                                                                move_step[2] * min(step_iter,
                                                                                num_move_steps)),
                                       vrep.simx_opmode_blocking)
            vrep.simxSetObjectOrientation(self.sim_client, UR5_target_handle, -1,
                                          (np.pi/2, UR5_target_orientation[1] + rotation_step * min(step_iter, num_rotation_steps), np.pi/2),
                                          vrep.simx_opmode_blocking)

    def move_position(self, tool_position):
        sim_ret, UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        # UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client,'UR5_target',vrep.simx_opmode_blocking)
        sim_ret, UR5_target_position = vrep.simxGetObjectPosition(self.sim_client, UR5_target_handle, -1,
                                                                  vrep.simx_opmode_blocking)
        if UR5_target_position[2] <= 0.031:
            print(UR5_target_position[2])
            tool_position[2] = 0
        move_direction = np.asarray(tool_position)
        move_magnitude = np.linalg.norm(move_direction)
        if move_magnitude == 0 or not move_magnitude == move_magnitude:
            move_step = [0, 0, 0]
            num_move_steps = 0
            # print('magnitude error~!', move_magnitude, tool_position)
        else:
            move_step = 0.005 * move_direction / move_magnitude
            num_move_steps = int(np.floor(move_magnitude / 0.005))
        # print(move_step)
        for step_iter in range(num_move_steps):
            vrep.simxSetObjectPosition(self.sim_client, UR5_target_handle, -1, (UR5_target_position[0] +
                                                                                move_step[0] * min(step_iter,
                                                                                num_move_steps),
                                                                                UR5_target_position[1] +
                                                                                move_step[1] * min(step_iter,
                                                                                num_move_steps),
                                                                                UR5_target_position[2] +
                                                                                move_step[2] * min(step_iter,
                                                                                num_move_steps)),
                                       vrep.simx_opmode_blocking)

    def move_orientation(self, tool_orientation):
        sim_ret, UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        sim_ret, UR5_target_orientation = vrep.simxGetObjectOrientation(self.sim_client, UR5_target_handle, -1,
                                                                        vrep.simx_opmode_blocking)
        rotation_step = 0.05 if (tool_orientation > 0) else -0.05
        num_rotation_steps = int(np.floor(tool_orientation / rotation_step))
        for step_iter in range(num_rotation_steps):
            vrep.simxSetObjectOrientation(self.sim_client, UR5_target_handle, -1,
                                          (np.pi/2, UR5_target_orientation[1] + rotation_step * min(step_iter, num_rotation_steps), np.pi/2),
                                          vrep.simx_opmode_blocking)

    def move_down(self, direction=1):
        sim_ret, UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        # UR5_target_handle = vrep.simxGetObjectHandle(self.sim_client,'UR5_target',vrep.simx_opmode_blocking)
        sim_ret, UR5_target_position = vrep.simxGetObjectPosition(self.sim_client, UR5_target_handle, -1,
                                                                  vrep.simx_opmode_blocking)
        # print(UR5_target_position)
        # time.sleep(200)
        if direction == 1:
            move_direction = np.asarray([0, 0, 0.01 - UR5_target_position[2]])
            move_magnitude = UR5_target_position[2] - 0.01
        else:
            move_direction = np.asarray([0, 0, 0.301 - UR5_target_position[2]])
            move_magnitude = 0.301 - UR5_target_position[2]
        if move_magnitude == 0 or not move_magnitude == move_magnitude:
            move_step = [0, 0, 0]
            num_move_steps = 0
            print('magnitude error~!', move_magnitude)
        else:
            move_step = 0.005 * move_direction / move_magnitude
            num_move_steps = int(np.floor(move_magnitude / 0.005))
            # print('move_direction : ', move_direction, 'move_magnitude : ', move_magnitude)
            # print('move step : ', move_step, 'num : ', num_move_steps)

        for step_iter in range(num_move_steps):
            vrep.simxSetObjectPosition(self.sim_client, UR5_target_handle, -1, (UR5_target_position[0],
                                                                                UR5_target_position[1],
                                                                                UR5_target_position[2] +
                                                                                move_step[2] * (step_iter+1)),
                                                                                vrep.simx_opmode_blocking)

    def close_gripper(self):
        gripper_motor_velocity = -0.5
        gripper_motor_force = 100
        sim_ret, RG2_gripper_handle = vrep.simxGetObjectHandle(self.sim_client, 'RG2_openCloseJoint',
                                                               vrep.simx_opmode_blocking)
        sim_ret, gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                    vrep.simx_opmode_blocking)
        vrep.simxSetJointForce(self.sim_client, RG2_gripper_handle, gripper_motor_force, vrep.simx_opmode_blocking)
        vrep.simxSetJointTargetVelocity(self.sim_client, RG2_gripper_handle, gripper_motor_velocity,
                                        vrep.simx_opmode_blocking)
        gripper_fully_closed = False
        while gripper_joint_position > -0.047:  # Block until gripper is fully closed
            sim_ret, new_gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                            vrep.simx_opmode_blocking)
            # print(gripper_joint_position)
            if new_gripper_joint_position >= gripper_joint_position:
                return gripper_fully_closed
            gripper_joint_position = new_gripper_joint_position
        gripper_fully_closed = True
        return gripper_fully_closed

    def grasp(self):
        self.total_try += 1

        self.move_down()
        # time.sleep(0.5)
        self.close_gripper()
        self.move_down(-1)  # move up
        time.sleep(1)
        sim_ret, RG2_gripper_handle = vrep.simxGetObjectHandle(self.sim_client, 'RG2_openCloseJoint',
                                                               vrep.simx_opmode_blocking)
        sim_ret, gripper_joint_position = vrep.simxGetJointPosition(self.sim_client, RG2_gripper_handle,
                                                                    vrep.simx_opmode_blocking)
        if 0.055 >= gripper_joint_position >= -0.04:
            success = True
        else:
            success = False
        if success:
            self.total_success += 1

        accuracy = float(self.total_success) / self.total_try

        self.writer.add_scalar('grasp success rate', accuracy, self.total_try)

        return success

    def get_sensor_data(self, is_save=False):
        sim_ret, cam_handle = vrep.simxGetObjectHandle(self.sim_client, 'Vision_sensor_hand', vrep.simx_opmode_blocking)

        # Get color image from simulation
        sim_ret, resolution, raw_image = vrep.simxGetVisionSensorImage(self.sim_client, cam_handle, 0,
                                                                       vrep.simx_opmode_blocking)
        color_img = np.asarray(raw_image)
        # print(color_img.shape)
        color_img.shape = (resolution[1], resolution[0], 3)
        # print(color_img)
        color_img = color_img.astype(np.float) / 255

        color_img[color_img < 0] += 1

        # color_img *= 255
        color_img = np.fliplr(color_img)

        test_show_image = color_img * 255
        test_show_image = test_show_image.astype(np.uint8)

        test_show_image = cv.cvtColor(test_show_image, cv.COLOR_BGR2GRAY)

        resize_color = cv.resize(test_show_image, (image_pix, image_pix))
        # print(resize_color.shape)
        # cropped_color = resize_color[84:132, 78:126]
        cropped_color = resize_color

        if is_save:
            # print(test_show_image.shape)
            # test_show_image = Image.fromarray(test_show_image)
            # test_show_image.show()
            cv.imwrite('/home/chen/PycharmProjects/Reinforcement/VisualGrasp/image/aob'
                       + str(self.target_orientation[0]) + '_'
                       + str(self.target_orientation[1]) + '_'
                       + str(self.target_orientation[2]) + '_' + '.png', cropped_color)
            cv.imwrite('/home/chen/PycharmProjects/Reinforcement/VisualGrasp/image/aob'
                       + str(self.target_orientation[0]) + '_'
                       + str(self.target_orientation[1]) + '_'
                       + str(self.target_orientation[2]) + '_big' + '.png', resize_color)

        # print(cropped_color.shape)
        cropped_color = np.asarray(cropped_color)
        # print(cropped_color.shape)
        # cv.imwrite('t.png', resize_color)
        cropped_color = cropped_color.astype(np.float) / 255.
        # print(cropped_color.shape)
        # np.set_printoptions(threshold=np.inf, suppress=True)
        # f = open('test.txt', 'w')
        # print(resize_color, file=f)
        # f.close()
        # print('ok')
        # time.sleep(5)
        # print(cropped_color.shape)
        return cropped_color

    def get_target_orientation(self):
        sim_ret, self.target_handle = vrep.simxGetObjectHandle(self.sim_client, 'shape001', vrep.simx_opmode_blocking)
        sim_ret, self.target_orientation = vrep.simxGetObjectOrientation(self.sim_client, self.target_handle, -1,
                                                                         vrep.simx_opmode_blocking)
        return self.target_orientation[1]

    def step(self, action):
        # action : [x_ratio, y_ratio] ~ [-1, 1]
        # origin_action = action
        move_magnitude = 0.05
        rotation_magnitude = np.pi/6

        move = action[0:2]
        move = np.append(move, 0)
        rotation = action[-1]
        # print('rotation : ', rotation)
        rotation = rotation * rotation_magnitude
        move = move * move_magnitude

        self.move_position(move)
        self.move_orientation(rotation)

        # action = np.append(move, rotation)
        # self.move_to(action)
        self.current_state, frame, pos = self.get_state()
        reward, done = self.get_reward(action)
        cs = self.current_state[0:2]
        cs = np.append(cs, self.current_state[-1])
        return reward, cs, done, frame, pos


import numpy as np
import numpy.random as npr
import os

data_dir = 'M:/MTCNN/data'
#anno_file = os.path.join(data_dir, "anno.txt")

size = 12

if size == 12:
    net = "PNet"
elif size == 24:
    net = "RNet"
elif size == 48:
    net = "ONet"

with open(os.path.join(data_dir, '%s/pos_%s.txt' % (size, size)), 'r') as f:
    pos = f.readlines()

with open(os.path.join(data_dir, '%s/neg_%s.txt' % (size, size)), 'r') as f:
    neg = f.readlines()

with open(os.path.join(data_dir, '%s/part_%s.txt' % (size, size)), 'r') as f:
    part = f.readlines()

with open(os.path.join(data_dir,'%s/landmark_%s_aug.txt' %(size,size)), 'r') as f:
    landmark = f.readlines()

dir_path = os.path.join(data_dir, 'imglists')
if not os.path.exists(dir_path):
    os.makedirs(dir_path)
if not os.path.exists(os.path.join(dir_path, "%s" %(net))):
    os.makedirs(os.path.join(dir_path, "%s" %(net)))
with open(os.path.join(dir_path, "%s" %(net),"train_%s_landmark.txt" % (net)), "w") as f:
    nums = [len(neg), len(pos), len(part)]
    ratio = [3, 1, 1]
    #base_num = min(nums)
    base_num = 250000
    print(len(neg), len(pos), len(part), base_num)

    #shuffle the order of the initial data
    #if negative examples are more than 750k then only choose 750k
    if len(neg) > base_num * 3:
        neg_keep = npr.choice(len(neg), size=base_num * 3, replace=True)
    else:
        neg_keep = npr.choice(len(neg), size=len(neg), replace=True)
    pos_keep = npr.choice(len(pos), size=base_num, replace=True)
    part_keep = npr.choice(len(part), size=base_num, replace=True)
    print(len(neg_keep), len(pos_keep), len(part_keep))

    # write the data according to the shuffled order
    for i in pos_keep:
        f.write(pos[i])
    for i in neg_keep:
        f.write(neg[i])
    for i in part_keep:
        f.write(part[i])
    for item in landmark:
        f.write(item)


import os
from tqdm import trange
from PIL import Image, ImageOps, ImageFilter
import numpy as np
import torch

from .base import BaseDataset

"""
NUM_CHANNEL = 91
[] background
[5] airplane
[2] bicycle
[16] bird
[9] boat
[44] bottle
[6] bus
[3] car
[17] cat
[62] chair
[21] cow
[67] dining table
[18] dog
[19] horse
[4] motorcycle
[1] person
[64] potted plant
[20] sheep
[63] couch
[7] train
[72] tv
"""
CAT_LIST = [0, 5, 2, 16, 9, 44, 6, 3, 17, 62, 21, 67, 18, 19, 4,
    1, 64, 20, 63, 7, 72]

class COCOSegmentation(BaseDataset):
    def __init__(self, root=os.path.expanduser('~/.encoding/data'), split='train',
                 mode=None, transform=None, target_transform=None):
        super(COCOSegmentation, self).__init__(
            root, split, mode, transform, target_transform)
        from pycocotools.coco import COCO
        from pycocotools import mask
        if mode == 'train':
            print('train set')
            ann_file = os.path.join(root, 'coco/annotations/instances_train2014.json')
            ids_file = os.path.join(root, 'coco/annotations/train_ids.pth')
            root = os.path.join(root, 'coco/train2014')
        else:
            print('val set')
            ann_file = os.path.join(root, 'coco/annotations/instances_val2014.json')
            ids_file = os.path.join(root, 'coco/annotations/val_ids.pth')
            root = os.path.join(root, 'coco/val2014')
        self.coco = COCO(ann_file)
        self.coco_mask = mask
        if os.path.exists(ids_file):
            self.ids = torch.load(ids_file)
        else:
            ids = list(self.coco.imgs.keys())
            self.ids = self._preprocess(ids, ids_file)
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        coco = self.coco
        img_id = self.ids[index]
        img_metadata = coco.loadImgs(img_id)[0]
        path = img_metadata['file_name']
        img = Image.open(os.path.join(self.root, path)).convert('RGB')
        cocotarget = coco.loadAnns(coco.getAnnIds(imgIds=img_id))
        mask = Image.fromarray(self._gen_seg_mask(cocotarget, img_metadata['height'],
                               img_metadata['width']))
        # synchrosized transform
        if self.mode == 'train':
            img, mask = self._sync_transform(img, mask)
        elif self.mode == 'val':
            img, mask = self._val_sync_transform(img, mask)
        else:
            assert self.mode == 'testval'
            mask = self._mask_transform(mask)
        # general resize, normalize and toTensor
        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            mask = self.target_transform(mask)
        return img, mask

    def __len__(self):
        return len(self.ids)

    def _gen_seg_mask(self, target, h, w):
        mask = np.zeros((h, w), dtype=np.uint8)
        coco_mask = self.coco_mask
        for instance in target:
            rle = coco_mask.frPyObjects(instance['segmentation'], h, w)
            m = coco_mask.decode(rle)
            cat = instance['category_id']
            if cat in CAT_LIST:
                c = CAT_LIST.index(cat)
            else:
                continue
            if len(m.shape) < 3:
                mask[:, :] += (mask == 0) * (m * c)
            else:
                mask[:, :] += (mask == 0) * (((np.sum(m, axis=2)) > 0) * c).astype(np.uint8)
        return mask

    def _preprocess(self, ids, ids_file):
        print("Preprocessing mask, this will take a while." + \
            "But don't worry, it only run once for each split.")
        tbar = trange(len(ids))
        new_ids = []
        for i in tbar:
            img_id = ids[i]
            cocotarget = self.coco.loadAnns(self.coco.getAnnIds(imgIds=img_id))
            img_metadata = self.coco.loadImgs(img_id)[0]
            mask = self._gen_seg_mask(cocotarget, img_metadata['height'],
                                      img_metadata['width'])
            # more than 1k pixels
            if (mask > 0).sum() > 1000:
                new_ids.append(img_id)
            tbar.set_description('Doing: {}/{}, got {} qualified images'.\
                format(i, len(ids), len(new_ids)))
        print('Found number of qualified images: ', len(new_ids))
        torch.save(new_ids, ids_file)
        return new_ids


#!/usr/bin/env python
import cv2
import numpy as np
import os
import copy
import rospkg

class SignalDetecor:

    #generating diff image
    def diffimg(self, after, before):

        blue_before  = before[:,:,0]
        green_before = before[:,:,1]
        red_before   = before[:,:,2]

        blue_after  = after[:,:,0]
        green_after = after[:,:,1]
        red_after   = after[:,:,2]

        blue_final  = blue_after.astype(np.int32) - blue_before.astype(np.int32) + 255
        green_final = green_after.astype(np.int32)- green_before.astype(np.int32) + 255
        red_final   = red_after.astype(np.int32) - red_before.astype(np.int32) + 255

        output=np.zeros((after.shape[0],after.shape[1],3))
        output[:,:,0]=blue_final
        output[:,:,1]=green_final
        output[:,:,2]=red_final

        return (output/2.0).astype(np.uint8)

    def __init__(self,template_path):

        print template_path
        # load template images
        self.temp0 = cv2.imread(template_path+'red2green.png')
        self.temp1 = cv2.imread(template_path+'red2black.png')
        self.temp2 = cv2.imread(template_path+'black2green.png')
        self.temp3 = cv2.imread(template_path+'green2red.png')
        # self.temp0 = cv2.imread('templates/red2green.png')
        # self.temp1 = cv2.imread('templates/red2black.png')
        # self.temp2 = cv2.imread('templates/black2green.png')
        # self.temp3 = cv2.imread('templates/green2red.png')

        self.temp0_h = self.temp0.shape[0]
        self.temp0_w = self.temp0.shape[1]
        self.temp3_h = self.temp3.shape[0]
        self.temp3_w = self.temp3.shape[1]

        self.frame4=0
        self.frame3=0
        self.frame2=0
        self.frame1=0

        self.ok_desu=0
        self.output=np.zeros((600,600,3))

        self.green_cnt=0
        self.red_cnt=0

        # threshold of template matching
        self.green_th = 0.65 #stable:0.65
        self.red_th = 0.7 #stable:0.7

        # number of detections
        self.nd = 3

        # signal states (0:unkown, 1:green, 2:red)
        self.state = 0

    def __call__(self, frame):

        new_frame = frame
        display_frame = copy.deepcopy(new_frame)

        if self.ok_desu>4:

            output=self.diffimg(new_frame,self.frame4)
            result0=cv2.matchTemplate(output,self.temp0,method=cv2.TM_CCOEFF_NORMED)
            result1=cv2.matchTemplate(output,self.temp1,method=cv2.TM_CCOEFF_NORMED)
            result2=cv2.matchTemplate(output,self.temp2,method=cv2.TM_CCOEFF_NORMED)
            result3=cv2.matchTemplate(output,self.temp3,method=cv2.TM_CCOEFF_NORMED)

            mmlr0=cv2.minMaxLoc(result0)
            mmlr1=cv2.minMaxLoc(result1)
            mmlr2=cv2.minMaxLoc(result2)
            mmlr3=cv2.minMaxLoc(result3)

            # condition1: RED to GREEN
            if mmlr0[1]>self.green_th and mmlr0[1] > mmlr2[1]:
                self.green_cnt+=1
            else:
                self.green_cnt=0

            if self.green_cnt==self.nd:
                self.state = 1
                print "The signal turned from RED to GREEN !"
                cv2.rectangle(display_frame,
                              (mmlr0[3][0],mmlr0[3][1]),
                              (mmlr0[3][0]+temp0_w,mmlr0[3][1]+temp0_h),
                              (0,255,0),3)
                self.green_cnt=0

            # condition2: GREEN to RED
            if mmlr3[1]>self.red_th and mmlr3[1] > mmlr1[1]:
                self.red_cnt+=1
            else:
                self.red_cnt=0

            if self.red_cnt==self.nd:
                self.state = 2
                print "The signal turned from GREEN to RED !"
                cv2.rectangle(display_frame,
                              (mmlr3[3][0],mmlr3[3][1]),
                              (mmlr3[3][0]+temp3_w,mmlr3[3][1]+temp3_h),
                              (0,0,255),3)
                self.red_cnt = 0

            # display the signal state
            if self.state == 0:
                signal_state = "UNKOWN"
                color = (0,0,0)
            elif self.state == 1:
                signal_state = "GREEN"
                color = (0,255,0)
            else:
                signal_state = "RED"
                color = (0,0,255)

            # message = "signal state "+str(self.state)+": "+signal_state
            #print message
            # cv2.putText(display_frame, signal_state, (10, 25),
            #             cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

            # display the resulting frame
            # cv2.imshow('frame',display_frame)
            # cv2.imshow('diff4',output)
            # cv2.waitKey(3)

        self.frame4=self.frame3
        self.frame3=self.frame2
        self.frame2=self.frame1
        self.frame1=new_frame
        self.ok_desu+=1

        return self.state


#! /usr/bin/env python
"""
Copyright (c) 2019 - 2022 Jolle Jolles <j.w.jolles@gmail.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at:

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from threading import Thread
import os
import cv2
import time
import numpy as np

from pythutils.sysutils import isrpi
from pythutils.mediautils import *

class VideoIn:

    def __init__(self, system = "auto", vidsize = 0.2, framerate = 32,
                 crop = False, rotation = 0, maxres = None):

        """Opens a video stream from native camera, webcam or rpi camera"""

        if system == "auto":
            self.cam = "rpi" if isrpi() else 0
        elif system in ["rpi",0,1,2]:
            self.cam = system
        else:
            self.cam = 0

        self.crop = crop

        if self.cam == "rpi":
            from picamera.array import PiRGBArray
            from picamera import PiCamera

            if isinstance(maxres, tuple):
                self.maxres = maxres
            elif maxres == "v2":
                self.maxres = (3264,2464)
            elif maxres == "hq":
                self.maxres = (4056,3040)
            else:
                self.maxres = (2592,1952)
            self.res = (self.maxres[0]*vidsize, self.maxres[1]*vidsize)
            self.res = picamconv(self.res)
            self.camera = PiCamera()
            self.camera.resolution = self.res
            self.camera.framerate = framerate
            self.camera.rotation = rotation
            self.rawCapture = PiRGBArray(self.camera, size=self.res)
            self.stream = self.camera.capture_continuous(self.rawCapture,
                          format="bgr", use_video_port=True)

        else:
            self.stream = cv2.VideoCapture(self.cam)
            self.stream.set(3, 4000)
            self.stream.set(4, 4000)
            self.maxres = (int(self.stream.get(3)), int(self.stream.get(4)))
            self.stream.set(3, int(self.maxres[0]*vidsize))
            self.stream.set(4, int(self.maxres[1]*vidsize))
            self.res = (int(self.stream.get(3)), int(self.stream.get(4)))

        self.stopped = False

    def start(self):
        Thread(target=self.update, args=()).start()
        time.sleep(2)
        return self

    def update(self):
        if self.cam == "rpi":
            for f in self.stream:
                self.frame = f.array
                self.rawCapture.truncate(0)
                if self.stopped:
                    self.stream.close()
                    self.rawCapture.close()
                    self.camera.close()
                    return
        else:
            while True:
                _, self.frame = self.stream.read()
                if self.stopped:
                    self.stream.release()
                    return

    def read(self):
        if not hasattr(self, 'frame'):
            time.sleep(4)
        if self.crop:
            self.frame = crop(self.frame, self.crop[0], self.crop[1])
        return self.frame

    def img(self):
        w,h = self.maxres
        if self.cam == "rpi":
            self.camera.resolution = self.maxres
            self.image = np.empty((h * w * 3,), dtype=np.uint8)
            time.sleep(1)
            self.camera.capture(self.image, 'bgr')
            self.image = self.image.reshape((h, w, 3))
            self.stream.close()
            self.rawCapture.close()
            self.camera.close()
        else:
            self.stream.release()
            self.stream = cv2.VideoCapture(self.cam)
            time.sleep(1)
            self.stream.set(3, w)
            self.stream.set(4, h)
            _, self.image = self.stream.read()
            self.stream.release()

        if self.crop:
            zoom = roi_to_zoom(self.crop, self.res)
            (rx1,ry1),(rx2,ry2) = zoom_to_roi(zoom, self.maxres)
            fixx, fixy = fix_vidshape(self.res, self.maxres)
            if fixx > 100 or fixy > 100:
                rx1 = rx1+int(((self.maxres[0]/2.)-rx1)/(self.maxres[0]/2.)*fixx)
                ry1 = ry1+int(((self.maxres[1]/2.)-ry1)/(self.maxres[1]/2.)*fixy)
                rx2 = rx2+int(((self.maxres[0]/2.)-rx2)/(self.maxres[0]/2.)*fixx)
                ry2 = ry2+int(((self.maxres[1]/2.)-ry2)/(self.maxres[1]/2.)*fixy)
            self.roil = ((rx1,ry1),(rx2,ry2))
            self.roiw = self.roil[1][0] - self.roil[0][0]
            self.roih = self.roil[1][1] - self.roil[0][1]
            self.image = crop(self.image, self.roil[0], self.roil[1])

        return self.image

    def stop(self):
        self.stopped = True


"""Evaluate a cost model on a network with dataset simulator"""
import argparse
import os
import pickle

import numpy as np

import tvm
from tvm import auto_scheduler
from tvm.auto_scheduler.dataset import LearningTask
from tvm.auto_scheduler.cost_model.xgb_model import XGBModelInternal

from common import MEASURE_RECORD_FOLDER, NETWORK_INFO_FOLDER

def eval_cost_model_on_weighted_tasks(model, eval_task_dict, eval_dataset, top_ks):
    """Evaluate a cost model on weighted tasks"""
    preds_dict = model.predict(eval_dataset)

    best_latency = 0
    latencies = [0] * len(top_ks)
    for task, weight in eval_task_dict.items():
        if task not in eval_dataset.throughputs:
            print(f"Warning: cannot find {task.workload_key} in the eval_dataset. Skipped.")
            continue

        preds = preds_dict[task]
        labels, min_latency = eval_dataset.throughputs[task], eval_dataset.min_latency[task]

        real_values = labels[np.argsort(-preds)]
        real_latency = min_latency / np.maximum(real_values, 1e-5)

        for i, top_k in enumerate(top_ks):
            latencies[i] += np.min(real_latency[:top_k]) * weight
        best_latency += min_latency * weight

    return latencies, best_latency

def eval_cost_model_on_network(model, network, target, top_ks):
    # Read tasks of the network
    target = tvm.target.Target(target)
    network_task_key = (network_key, str(target.kind))
    task_info_filename = f"dataset/network_info/{network_task_key}.task.pkl"
    tasks, task_weights = pickle.load(open(task_info_filename, "rb"))
    network_task_key2 = (network_key, str(target))

    dataset_file = f".dataset_cache/{network_task_key2}.network.pkl"
    if not os.path.exists(dataset_file):
        # get file names of these tasks
        filenames = []
        for task in tasks:
            task_key = (task.workload_key, str(task.target.kind))
            filename = f"dataset/measure_records/{target.model}/{task_key}.json"
            filenames.append(filename)

        # make a dataset
        auto_scheduler.dataset.make_dataset_from_log_file(
            filenames, dataset_file, min_sample_size=0)

    dataset = pickle.load(open(dataset_file, "rb"))
    target = dataset.tasks()[0].target

    learning_tasks = [LearningTask(t.workload_key, target) for t in tasks]

    task_dict = {task: weight for task, weight in zip(learning_tasks, task_weights)}
    return eval_cost_model_on_weighted_tasks(model, task_dict, dataset, top_ks)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-file", type=str)
    args= parser.parse_args()

    model_file = args.model_file
    network_keys = [
        ("resnet_50", [(1, 3, 224,224)]),
        ("mobilenet_v2", [(1, 3, 224,224)]),
        ("mobilenet_v3", [(1, 3, 224,224)]),
        ("bert_base", [(1, 128)]),
    ]
    target = "llvm -model=e5-2666"

    model = XGBModelInternal()
    model.load(model_file)

    top_ks = [1, 5]
    for network_key in network_keys:
        latencies, best_latency = eval_cost_model_on_network(model, network_key, target, top_ks)
        for top_k, latency in zip(top_ks, latencies):
            print(f"Network: {network_key}\tTop-{top_k} score: {best_latency / latency}")