import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import numpy as np
import torch.optim as optim
from args import get_args
import random
import os
import pickle
from torch.optim.lr_scheduler import StepLR
import logging
from transformers import DistilBertTokenizer
from data.howto_loader import HowTo_Dataset, howto_collate_fn
from data.videotext_loader import (
    VideoText_Dataset,
    Youcook_Dataset,
    videotext_collate_fn,
)
from model.multimodal_transformer import MMT_VideoQA
from train.train_htm import train_mlmcm, eval_mlm, eval_retrieval

# args, logging
args = get_args()
assert args.checkpoint_dir
assert args.dataset == "howto100m"
if not (os.path.isdir(args.save_dir)):
    os.mkdir(args.save_dir)
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s %(levelname)-8s %(message)s"
)
logFormatter = logging.Formatter("%(asctime)s %(levelname)-8s %(message)s")
rootLogger = logging.getLogger()
fileHandler = logging.FileHandler(os.path.join(args.save_dir, "stdout.log"), "w+")
fileHandler.setFormatter(logFormatter)
rootLogger.addHandler(fileHandler)
logging.info(args)

# set random seeds
torch.manual_seed(args.seed)
np.random.seed(args.seed)
random.seed(args.seed)

# Model
model = MMT_VideoQA(
    feature_dim=args.feature_dim,
    word_dim=args.word_dim,
    N=args.n_layers,
    d_model=args.embd_dim,
    d_ff=args.ff_dim,
    h=args.n_heads,
    dropout=args.dropout,
    T=args.max_feats,
    Q=args.qmax_words,
    baseline=args.baseline,
    n_negs=args.n_negs,
)
model = nn.DataParallel(model)
model.cuda()

# Load captions, dataloaders
caption = pickle.load(open(args.caption_path, "rb"))
logging.info("Pickle loaded")
bert_tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")

dataset = HowTo_Dataset(
    csv_path=args.train_csv_path,
    caption=caption,
    features_path=args.features_path,
    min_time=args.min_time,
    max_time=args.max_feats,
    max_words=args.qmax_words,
    min_words=args.min_words,
    n_pair=args.n_pair,
    bert_tokenizer=bert_tokenizer,
)

dataset_size = len(dataset)
dataloader = DataLoader(
    dataset,
    batch_size=args.batch_size,
    num_workers=args.num_thread_reader,
    shuffle=True,
    batch_sampler=None,
    drop_last=True,
    collate_fn=howto_collate_fn,
)

youcook_dataset = Youcook_Dataset(
    data=args.youcook_val_path,
    max_words=args.qmax_words,
    bert_tokenizer=bert_tokenizer,
    max_feats=args.max_feats,
)
youcook_loader = DataLoader(
    youcook_dataset,
    batch_size=args.batch_size_val,
    num_workers=args.num_thread_reader,
    shuffle=False,
    collate_fn=videotext_collate_fn,
)

msrvtt_dataset = VideoText_Dataset(
    csv_path=args.msrvtt_test_csv_path,
    features_path=args.msrvtt_test_features_path,
    max_words=args.qmax_words,
    bert_tokenizer=bert_tokenizer,
    max_feats=args.max_feats,
)
msrvtt_loader = DataLoader(
    msrvtt_dataset,
    batch_size=args.batch_size_val,
    num_workers=args.num_thread_reader,
    shuffle=False,
    drop_last=False,
    collate_fn=videotext_collate_fn,
)

# Optimizer, Scheduler
params_for_optimization = list(p for p in model.parameters() if p.requires_grad)
optimizer = optim.Adam(params_for_optimization, lr=args.lr)
scheduler = StepLR(optimizer, step_size=len(dataloader), gamma=args.lr_decay)

# Train
for epoch in range(args.epochs):
    eval_mlm(model, youcook_loader, "YouCook2", epoch)
    eval_retrieval(model, youcook_loader, "YouCook2", epoch)
    eval_mlm(model, msrvtt_loader, "MSR-VTT", epoch)
    eval_retrieval(model, msrvtt_loader, "MSR-VTT", epoch)
    train_mlmcm(model, optimizer, dataloader, scheduler, epoch, args)
    torch.save(model.state_dict(), os.path.join(args.save_dir, f"e{epoch}.pth"))
eval_mlm(model, youcook_loader, "YouCook2", args.epochs)
eval_retrieval(model, youcook_loader, "YouCook2", args.epochs)
eval_mlm(model, msrvtt_loader, "MSR-VTT", args.epochs)
eval_retrieval(model, msrvtt_loader, "MSR-VTT", args.epochs)


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

"""
# TpgDraw モジュール

## 概要
TestPattern作成時の各種描画関数をまとめたもの。
"""

import os
import numpy as np
import transfer_functions as tf
import color_space as cs
import test_pattern_generator2 as tpg
from PIL import Image
from PIL import ImageFont
from PIL import ImageDraw
import colour
import imp
imp.reload(tpg)

class TpgDraw:
    """
    テストパターン用の各種描画を行う。
    """
    def __init__(self, draw_param, preview):
        # TpgControl から受け通るパラメータ
        self.preview = preview
        self.bg_color = draw_param['bg_color']
        self.fg_color = draw_param['fg_color']
        self.img_width = draw_param['img_width']
        self.img_height = draw_param['img_height']
        self.bit_depth = draw_param['bit_depth']
        self.img_max = (2 ** self.bit_depth) - 1
        self.transfer_function = draw_param['transfer_function']
        self.color_space = draw_param['color_space']
        self.white_point = draw_param['white_point']
        self.revision = draw_param['revision']

        self.convert_to_8bit_coef = 2 ** (self.bit_depth - 8)
        self.convert_from_10bit_coef = 2 ** (16 - self.bit_depth)

        # TpgDraw 内部パラメータ(細かい座標など)
        self.ramp_height_coef = 0.079  # range is [0.0:1.0]
        # self.ramp_st_pos_h_coef = 0.442  # range is [0.0:1.0]
        self.ramp_st_pos_h_coef = 0.03  # range is [0.0:1.0]
        self.ramp_st_pos_v_coef = 0.235  # range is [0.0:1.0]
        self.checker_8bit_st_pos_v_coef = self.ramp_st_pos_v_coef + 0.115
        self.checker_10bit_st_pos_v_coef\
            = self.checker_8bit_st_pos_v_coef + 0.115
        self.each_spec_text_size_coef = 0.02  # range is [0.0:1.0]
        self.outline_text_size_coef = 0.02  # range is [0.0:1.0]
        self.step_bar_width_coef = 0.95
        self.step_bar_height_coef = 0.2
        self.step_bar_st_pos_v_coef = 0.74
        self.step_bar_text_width = 0.3

        self.color_checker_size_coef = 0.060
        self.color_checker_padding_coef = 0.0035
        self.color_checker_st_pos_v_coef = 0.1
        # self.color_checker_st_pos_h_coef = 0.03
        self.color_checker_st_pos_h_coef = 0.59

        self.csf_img_size_coef = 0.017
        self.csf_img_st_pos_v_coef = 0.03
        self.csf_img_low_st_pos_h_coef = 0.8
        self.csf_img_high_st_pos_h_coef = 0.842

        self.ycbcr_img_st_pos_v_coef = 0.03
        self.ycbcr_img_st_pos_h_coef = 0.862
        self.ycbcr_img_size_coef = 0.017

        self.dot_img_st_pos_v_coef = 0.03
        self.dot_img_st_pos_h_coef = 0.90

        self.info_text_st_pos_h_coef = 0.3
        self.info_text_st_pos_v_coef = 0.043
        self.info_text_size_coef = 0.3

        self.info2_text_st_pos_h_coef = 0.83
        self.info2_text_st_pos_v_coef = 0.03
        self.info2_text_size_coef = 0.02

        self.step_bar_width_coef_type2 = 0.8
        self.step_bar_height_coef_type2 = 0.09 * self.step_bar_width_coef_type2
        self.step_bar_st_pos_h_coef_type2 = 0.02
        self.step_bar_st_pos_v_coef_type2 = 0.17
        self.step_bar_v_space_coef_type2 = 0.24
        self.step_bar_text_width_type2 = 0.15

        self.set_fg_code_value()
        self.set_bg_code_value()

    def set_bg_code_value(self):
        code_value = tf.oetf_from_luminance(self.bg_color,
                                            self.transfer_function)
        code_value = int(round(code_value * self.img_max))
        self.bg_code_value = code_value

    def get_bg_code_value(self):
        return self.bg_code_value

    def set_fg_code_value(self):
        code_value = tf.oetf_from_luminance(self.fg_color,
                                            self.transfer_function)
        code_value = int(round(code_value * self.img_max))
        self.fg_code_value = code_value

    def get_fg_code_value(self):
        return self.fg_code_value

    def preview_iamge(self, img, order='rgb'):
        tpg.preview_image(img, order)

    def draw_bg_color(self):
        """
        背景色を描く。
        nits で指定したビデオレベルを使用
        """
        code_value = self.get_bg_code_value()
        self.img *= code_value

    def draw_outline(self):
        """
        外枠として1pxの直線を描く。
        nits で指定したビデオレベルを使用
        """
        code_value = self.get_fg_code_value()

        st_h, st_v = (0, 0)
        ed_h, ed_v = (self.img_width - 1, self.img_height - 1)

        self.img[st_v, st_h:ed_h, :] = code_value
        self.img[ed_v, st_h:ed_h, :] = code_value
        self.img[st_v:ed_v, st_h, :] = code_value
        self.img[st_v:ed_v, ed_h, :] = code_value

    def get_each_spec_text_height_and_size(self):
        """
        各パーツ説明のテキストの高さ(px)とフォントサイズを吐く
        """
        font_size = self.img_height * self.each_spec_text_size_coef
        text_height = font_size / 72 * 96 * 1.1

        return int(text_height), int(font_size)

    def get_info_text_height_and_size(self):
        """
        冒頭の解説用テキストの情報を取得
        """
        font_size = self.img_height * self.info_text_size_coef
        text_height = font_size / 72 * 96 * 1.1

        return int(text_height), int(font_size)

    def get_info2_text_height_and_size(self):
        """
        冒頭の解説用テキストの情報を取得
        """
        font_size = self.img_height * self.info2_text_size_coef
        text_height = font_size / 72 * 96 * 1.1

        return int(text_height), int(font_size)

    def get_color_bar_text_font_size(self, text_height):
        """
        カラーバー横に表示する階調＋輝度のフォントサイズを取得する
        """
        font_size = int(text_height / 96 * 72)
        return font_size

    def get_text_st_pos_for_over_info(self, tp_pos, text_height):
        return (tp_pos[0], tp_pos[1] - text_height)

    def get_fg_color_for_pillow(self):
        """
        Pillow 用 に 8bit精度のFG COLORを算出する
        """
        text_video_level_8bit\
            = int(self.fg_code_value / self.convert_to_8bit_coef)
        fg_color = tuple([text_video_level_8bit for x in range(3)])
        return fg_color

    def convert_from_pillow_to_numpy(self, img):
        img = np.uint16(np.asarray(img)) * self.convert_to_8bit_coef

        return img

    def merge_text(self, txt_img, pos):
        """
        テキストを合成する作業の最後の部分。
        pos は テキストの (st_pos_h, st_pos_v) 。

        ## 個人的実装メモ
        今回はちゃんとアルファチャンネルを使った合成をしたかったが、
        PILは8bit, それ以外は 10～16bit により BG_COLOR に差が出るので断念。
        """
        st_pos_v = pos[1]
        ed_pos_v = pos[1] + txt_img.shape[0]
        st_pos_h = pos[0]
        ed_pos_h = pos[0] + txt_img.shape[1]

        # かなり汚い実装。0x00 で無いピクセルのインデックスを抽出し、
        # そのピクセルのみを元の画像に上書きするという処理をしている。
        text_index = txt_img > 0
        temp_img = self.img[st_pos_v:ed_pos_v, st_pos_h:ed_pos_h]
        temp_img[text_index] = txt_img[text_index]
        self.img[st_pos_v:ed_pos_v, st_pos_h:ed_pos_h] = temp_img

    def merge_each_spec_text(self, pos, font_size, text_img_size, text):
        """
        各パーツの説明テキストを合成。
        pos は テキストの (st_pos_h, st_pos_v) 。
        text_img_size = (size_h, size_v)

        ## 個人的実装メモ
        今回はちゃんとアルファチャンネルを使った合成をしたかったが、
        PILは8bit, それ以外は 10～16bit により BG_COLOR に差が出るので断念。
        """
        # テキストイメージ作成
        text_width = text_img_size[0]
        text_height = text_img_size[1]
        fg_color = self.get_fg_color_for_pillow()
        bg_coor = (0x00, 0x00, 0x00)
        txt_img = Image.new("RGB", (text_width, text_height), bg_coor)
        draw = ImageDraw.Draw(txt_img)
        font = ImageFont.truetype("./fonts/NotoSansMonoCJKjp-Regular.otf",
                                  font_size)
        draw.text((0, 0), text, font=font, fill=fg_color)
        txt_img = self.convert_from_pillow_to_numpy(txt_img)

        self.merge_text(txt_img, pos)

    def draw_10bit_ramp(self):
        """
        10bitのRampパターンを描く
        """
        global g_cuurent_pos_v

        # パラメータ計算
        width = (self.img_height // 1080) * 1024
        height = int(self.img_height * self.ramp_height_coef)
        ramp_st_pos_h = self.get_ramp_st_pos_h(width)
        ramp_st_pos_v = int(self.img_height * self.ramp_st_pos_v_coef)
        ramp_pos = (ramp_st_pos_h, ramp_st_pos_v)
        text_height, font_size = self.get_each_spec_text_height_and_size()
        text_sub = " 10bit gray ramp (0, 1, 2, ..., 1021, 1022, 1023 level)."
        text = "▼" + text_sub
        text_pos = self.get_text_st_pos_for_over_info(ramp_pos, text_height)

        # ramp パターン作成
        ramp_10bit = tpg.gen_step_gradation(width=width, height=height,
                                            step_num=1025,
                                            bit_depth=self.bit_depth,
                                            color=(1.0, 1.0, 1.0),
                                            direction='h')
        tpg.merge(self.img, ramp_10bit, pos=ramp_pos)
        self.merge_each_spec_text(text_pos, font_size,
                                  (width, text_height), text)

    def get_bit_depth_checker_grad_width(self):
        """
        画面中央のグラデーション(256～768)の幅を求める
        """
        if self.img_height == 1080:
            width = 2048
        elif self.img_height == 2160:
            width = 4096
        else:
            raise ValueError('invalid img_height')

        return width

    def get_bit_depth_checker_grad_st_ed(self):
        """
        8bit/10bit チェック用のグラデーションは長めに作ってあるので
        最後にトリミングが必要となる。
        トリミングポイントの st, ed を返す
        """
        grad_width = self.get_bit_depth_checker_grad_width()
        grad_st_h = grad_width // 4
        grad_ed_h = grad_st_h + (grad_width // 2)

        return grad_st_h, grad_ed_h

    def get_ramp_st_pos_h(self, width):
        return int(self.img_width * self.ramp_st_pos_h_coef)

    def draw_8bit_10bit_checker(self, bit_depth='8bit', pos_v_coef=0.5):
        """
        256～768 の 8bit/10bitのRampパターンを描く
        """
        # パラメータ計算
        width = self.get_bit_depth_checker_grad_width()
        height = int(self.img_height * self.ramp_height_coef)
        grad_st_pos_h, grad_ed_pos_h = self.get_bit_depth_checker_grad_st_ed()
        width_after_trim = grad_ed_pos_h - grad_st_pos_h
        ramp_st_pos_h = self.get_ramp_st_pos_h(width_after_trim)
        ramp_st_pos_v = int(self.img_height * pos_v_coef)
        ramp_pos = (ramp_st_pos_h, ramp_st_pos_v)
        text_height, font_size = self.get_each_spec_text_height_and_size()
        text_8 = "8bit gray ramp (256, 260, 264, ..., 756, 760, 764 level)."
        text_10 = "10bit gray ramp (256, 257, 258, ..., 765, 766, 767 level)."
        text_sub = text_8 if bit_depth == '8bit' else text_10
        text = "▼ " + text_sub
        text_pos = self.get_text_st_pos_for_over_info(ramp_pos, text_height)

        # ramp パターン作成
        step_num = 257 if bit_depth == '8bit' else 1025
        ramp = tpg.gen_step_gradation(width=width, height=height,
                                      step_num=step_num,
                                      bit_depth=self.bit_depth,
                                      color=(1.0, 1.0, 1.0),
                                      direction='h')
        tpg.merge(self.img, ramp[:, grad_st_pos_h:grad_ed_pos_h],
                  pos=ramp_pos)
        self.merge_each_spec_text(text_pos, font_size,
                                  (width_after_trim, text_height), text)

    def get_color_bar_st_pos_h(self, width):
        return (self.img_width - width) // 2

    def draw_wrgbmyc_color_bar(self):
        """
        階段状のカラーバーをプロットする
        """
        scale_step = 65
        color_list = [(1, 1, 1), (1, 1, 1), (1, 0, 0), (0, 1, 0), (0, 0, 1),
                      (1, 0, 1), (1, 1, 0), (0, 1, 1)]
        width = int(self.img_width * self.step_bar_width_coef)
        height = int(self.img_height * self.step_bar_height_coef)
        color_bar_st_pos_h = self.get_color_bar_st_pos_h(width)
        color_bar_st_pos_v = int(self.img_height * self.step_bar_st_pos_v_coef)
        st_pos = (color_bar_st_pos_h, color_bar_st_pos_v)

        bar_height_list = tpg.equal_devision(height, len(color_list))
        bar_img_list = []
        for color, bar_height in zip(color_list, bar_height_list):
            color_bar = tpg.gen_step_gradation(width=width, height=bar_height,
                                               step_num=scale_step,
                                               bit_depth=self.bit_depth,
                                               color=color, direction='h')
            bar_img_list.append(color_bar)
        color_bar = np.vstack(bar_img_list)
        tpg.merge(self.img, color_bar, st_pos)

        # ここからテキスト。あらかじめV方向で作っておき、最後に回転させる
        txt_img = self.get_video_level_text_img(scale_step, width)
        text_pos = self.get_text_st_pos_for_over_info(st_pos, txt_img.shape[0])
        self.merge_text(txt_img, text_pos)

        # 説明文を下に追加する
        text_pos_v = st_pos[1] + color_bar.shape[0]
        text_pos = (st_pos[0], text_pos_v)
        text_height, font_size = self.get_each_spec_text_height_and_size()
        level_text = " (0, 16, 32, ..., 992, 1008, 1023 Level)"
        text = "▲ WRGBMYC Color Gradation" + level_text
        self.merge_each_spec_text(text_pos, font_size,
                                  (width, text_height), text)

    def get_video_level_text_img(self, scale_step, width, type=1):
        """
        ステップカラーに付与する VideoLevel & Luminance 情報。
        最初は縦向きで作って、最後に横向きにする
        """
        fg_color = self.get_fg_color_for_pillow()
        text_height_list = tpg.equal_devision(width, scale_step)
        font_size = self.get_color_bar_text_font_size(width / scale_step)
        video_level = np.linspace(0, 2 ** self.bit_depth, scale_step)
        video_level[-1] -= 1
        video_level_float = video_level / self.img_max
        bright_list = tf.eotf_to_luminance(video_level_float,
                                           self.transfer_function)
        if type == 1:
            text_width = int(self.step_bar_text_width * self.img_height)
        elif type == 2:
            text_width = int(self.step_bar_text_width_type2 * self.img_height)
        else:
            raise ValueError('invalid tpg-type')
        txt_img = Image.new("RGB", (text_width, width), (0x00, 0x00, 0x00))
        draw = ImageDraw.Draw(txt_img)
        font = ImageFont.truetype("./fonts/NotoSansMonoCJKjp-Regular.otf",
                                  font_size)
        st_pos_h = 0
        st_pos_v = 0
        for idx in range(scale_step):
            pos = (st_pos_h, st_pos_v)
            if bright_list[idx] < 999.99999:
                text_data = " {:>4.0f},{:>6.1f} nit".format(video_level[idx],
                                                            bright_list[idx])
            else:
                text_data = " {:>4.0f},{:>5.0f}  nit".format(video_level[idx],
                                                             bright_list[idx])
            draw.text(pos, text_data, font=font, fill=fg_color)
            st_pos_v += text_height_list[idx]

        txt_img = self.convert_from_pillow_to_numpy(txt_img)
        txt_img = np.rot90(txt_img)

        return txt_img

    def get_color_checker_rgb_value(self):
        """
        24パターンの Color Checker の RGB値を得る
        """
        colour_checker_param = colour.COLOURCHECKERS.get('ColorChecker 2005')

        # 今回の処理では必要ないデータもあるので xyY と whitepoint だけ抽出
        # -------------------------------------------------------------
        _name, data, whitepoint = colour_checker_param
        temp_xyY = []
        for key in data.keys():
            temp_xyY.append(data[key])
        temp_xyY = np.array(temp_xyY)
        large_xyz = colour.models.xyY_to_XYZ(temp_xyY)

        rgb_white_point\
            = colour.colorimetry.ILLUMINANTS['cie_2_1931'][self.white_point]

        illuminant_XYZ = whitepoint   # ColorCheckerのオリジナルデータの白色点
        illuminant_RGB = rgb_white_point  # RGBの白色点を設定
        chromatic_adaptation_transform = 'CAT02'
        large_xyz_to_rgb_matrix\
            = cs.get_xyz_to_rgb_matrix(self.color_space.name)
        # large_xyz_to_rgb_matrix = self.color_space.XYZ_to_RGB_matrix
        rgb = colour.models.XYZ_to_RGB(large_xyz, illuminant_XYZ,
                                       illuminant_RGB,
                                       large_xyz_to_rgb_matrix,
                                       chromatic_adaptation_transform)

        # overflow, underflow check
        # -----------------------------
        rgb[rgb < 0.0] = 0.0
        rgb[rgb > 1.0] = 1.0

        point_100nits = 100 / tf.PEAK_LUMINANCE[self.transfer_function]
        rgb = tf.oetf(rgb * point_100nits, self.transfer_function)
        rgb = np.uint16(np.round(rgb * self.img_max))

        return rgb

    def draw_color_checker(self):
        # 基本パラメータ算出
        # --------------------------------------
        h_num = 6
        v_num = 4
        patch_width = int(self.color_checker_size_coef * self.img_width)
        patch_height = patch_width
        patch_space = int(self.color_checker_padding_coef * self.img_width)
        img_width = patch_width * h_num + patch_space * (h_num - 1)
        img_height = patch_height * v_num + patch_space * (v_num - 1)
        rgb = self.get_color_checker_rgb_value()
        st_pos_h = int(self.color_checker_st_pos_h_coef * self.img_width)
        st_pos_v = int(self.color_checker_st_pos_v_coef * self.img_height)
        pos = (st_pos_h, st_pos_v)
        text = "▼ Color Checker"
        text_height, font_size = self.get_each_spec_text_height_and_size()
        text_pos = self.get_text_st_pos_for_over_info(pos, text_height)

        # 24ループで1枚の画像に24パッチを描画
        # -------------------------------------------------
        img_all_patch = np.zeros((img_height, img_width, 3))
        for idx in range(h_num * v_num):
            v_idx = idx // h_num
            h_idx = (idx % h_num)
            patch = np.ones((patch_height, patch_width, 3))
            patch[:, :] = rgb[idx]
            st_h = (patch_width + patch_space) * h_idx
            st_v = (patch_height + patch_space) * v_idx
            img_all_patch[st_v:st_v+patch_height, st_h:st_h+patch_width]\
                = patch

        tpg.merge(self.img, img_all_patch, pos)
        self.merge_each_spec_text(text_pos, font_size,
                                  (img_width, text_height), text)

    def draw_info_text(self):
        text_list = ["■ Information",
                     "    OETF/EOTF: {}".format(self.transfer_function),
                     "    White Point: {}".format(self.white_point),
                     "    Gamut: {}".format(self.color_space.name),
                     "    Revision: {:02d}".format(self.revision)]
        text_height, font_size = self.get_each_spec_text_height_and_size()
        st_pos_h = int(self.info_text_st_pos_h_coef * self.img_width)
        st_pos_v = int(self.info_text_st_pos_v_coef * self.img_height)
        for text in text_list:
            pos = (st_pos_h, st_pos_v)
            text_img_size = (self.img_width - st_pos_h, text_height)
            self.merge_each_spec_text(pos, font_size, text_img_size, text)
            st_pos_v += text_height

    def draw_info_text_type2(self):
        text_list = ["■ Information",
                     "    OETF/EOTF: {}".format(self.transfer_function),
                     "    Revison: {:02d}".format(self.revision)]
        text_height, font_size = self.get_info2_text_height_and_size()
        st_pos_h = int(self.info2_text_st_pos_h_coef * self.img_width)
        st_pos_v = int(self.info2_text_st_pos_v_coef * self.img_height)
        for text in text_list:
            pos = (st_pos_h, st_pos_v)
            text_img_size = (self.img_width - st_pos_h, text_height)
            self.merge_each_spec_text(pos, font_size, text_img_size, text)
            st_pos_v += text_height

    def draw_dot_mesh(self):
        """
        dot mesh パターンを描画。
        """
        # UHD画像の場合はトリミングを行う
        triming = True if self.img_height != 1080 else False

        kind_num = 2 if self.img_height == 1080 else 3
        fg_color_list = [[1.0, 1.0, 1.0], [1.0, 0.0, 0.0],
                         [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]
        fg_color_list = np.uint16(np.array(fg_color_list) * self.img_max)
        bg_color = np.uint16(np.array([self.get_bg_code_value()
                                       for x in range(3)]))

        args = {'kind_num': kind_num, 'whole_repeat': 1,
                'bg_color': bg_color}
        dot_img_list = []
        for idx in range(4):
            args['fg_color'] = fg_color_list[idx]
            temp_img = tpg.complex_dot_pattern(**args)
            if triming:
                length = (temp_img.shape[0] // 3) * 2
                temp_img = temp_img[:length, :length, :]
            dot_img_list.append(temp_img)
        dot_img = np.hstack(dot_img_list)
        st_pos_h = int(self.dot_img_st_pos_h_coef * self.img_width)
        st_pos_v = int(self.dot_img_st_pos_v_coef * self.img_height)
        dot_img_pos = (st_pos_h, st_pos_v)
        tpg.merge(self.img, dot_img, pos=dot_img_pos)

    def draw_csf_pattern(self):
        """
        Limited Range 確認用のパターンを追加。
        """
        lv_0 = np.array([0, 0, 0], dtype=np.uint16)
        lv_64 = np.array([64, 64, 64], dtype=np.uint16)
        lv_940 = np.array([940, 940, 940], dtype=np.uint16)
        lv_1023 = np.array([1023, 1023, 1023], dtype=np.uint16)
        width = int(self.csf_img_size_coef * self.img_width)
        num = 2
        low_img = tpg.make_csf_color_image(width, width, lv_0, lv_64, num)
        high_img = tpg.make_csf_color_image(width, width, lv_940, lv_1023, num)

        high_img = tpg.make_tile_pattern(width=width, height=width,
                                         h_tile_num=4, v_tile_num=4,
                                         low_level=lv_940, high_level=lv_1023)
        low_img = tpg.make_tile_pattern(width=width, height=width,
                                        h_tile_num=4, v_tile_num=4,
                                        low_level=lv_0, high_level=lv_64)

        st_pos_h = int(self.csf_img_low_st_pos_h_coef * self.img_width)
        st_pos_v = int(self.csf_img_st_pos_v_coef * self.img_height)
        csf_img_pos = (st_pos_h, st_pos_v)
        tpg.merge(self.img, low_img, csf_img_pos)

        st_pos_h = int(self.csf_img_high_st_pos_h_coef * self.img_width)
        csf_img_pos = (st_pos_h, st_pos_v)
        tpg.merge(self.img, high_img, csf_img_pos)
        # tpg.preview_image(csf_img_pos)

    def draw_ycbcr_err_checker(self):
        width = int(self.ycbcr_img_size_coef * self.img_width)
        img = tpg.make_ycbcr_checker(height=width, v_tile_num=4)

        st_pos_h = int(self.ycbcr_img_st_pos_h_coef * self.img_width)
        st_pos_v = int(self.ycbcr_img_st_pos_v_coef * self.img_height)
        ycbcr_img_pos = (st_pos_h, st_pos_v)
        tpg.merge(self.img, img, ycbcr_img_pos)

    def draw_tpg_type1(self):
        """
        Color Checker付きの基本的なパターンを描画。
        """
        self.img = np.ones((self.img_height, self.img_width, 3),
                           dtype=np.uint16)
        self.draw_bg_color()
        self.draw_outline()
        self.draw_10bit_ramp()
        self.draw_8bit_10bit_checker('8bit', self.checker_8bit_st_pos_v_coef)
        self.draw_8bit_10bit_checker('10bit', self.checker_10bit_st_pos_v_coef)
        self.draw_wrgbmyc_color_bar()
        self.draw_color_checker()
        self.draw_dot_mesh()
        self.draw_csf_pattern()
        self.draw_ycbcr_err_checker()
        self.draw_info_text()

        if self.preview:
            self.preview_iamge(self.img / self.img_max)

        return self.img

    def draw_wrgbmyc_color_bar_type2(self):
        """
        階段状のカラーバーをプロットする
        """
        scale_step = 257
        color_list = [(1, 1, 1)]
        width = int(self.img_width * self.step_bar_width_coef_type2) * 4
        height = int(self.img_height * self.step_bar_height_coef_type2)
        vspace = int(self.img_height * self.step_bar_v_space_coef_type2)
        color_bar_st_pos_h = int(self.img_width
                                 * self.step_bar_st_pos_h_coef_type2)
        color_bar_st_pos_v = int(self.img_height
                                 * self.step_bar_st_pos_v_coef_type2)
        st_pos = (color_bar_st_pos_h, color_bar_st_pos_v)

        bar_height_list = tpg.equal_devision(height, len(color_list))
        bar_img_list = []
        for color, bar_height in zip(color_list, bar_height_list):
            color_bar = tpg.gen_step_gradation(width=width, height=bar_height,
                                               step_num=scale_step,
                                               bit_depth=self.bit_depth,
                                               color=color, direction='h')
            bar_img_list.append(color_bar)
        color_bar = np.vstack(bar_img_list)

        # テキストデータ作成
        txt_img = self.get_video_level_text_img(scale_step, width, type=2)

        # 最終データだけ削除（後で4分割できるように）
        width = int(round(width/257*256))
        color_bar = color_bar[:, :width, :]
        txt_img = txt_img[:, :width, :]

        for idx in range(4):
            img_pos_h = (width // 4) * idx
            img_merge_width = (width // 4)
            color_bar_temp = color_bar[:, img_pos_h:img_pos_h+img_merge_width]
            txt_img_temp = txt_img[:, img_pos_h:img_pos_h+img_merge_width]
            tpg.merge(self.img, color_bar_temp, st_pos)
            text_pos = self.get_text_st_pos_for_over_info(st_pos,
                                                          txt_img.shape[0])
            self.merge_text(txt_img_temp, text_pos)
            st_pos = (st_pos[0], st_pos[1] + vspace)

    def draw_tpg_type2(self):
        """
        グレースケールを4stepで描画した画像
        """
        self.img = np.ones((self.img_height, self.img_width, 3),
                           dtype=np.uint16)
        self.draw_bg_color()
        self.draw_outline()

        self.draw_info_text_type2()
        self.draw_wrgbmyc_color_bar_type2()
        if self.preview:
            self.preview_iamge(self.img / self.img_max)
        return self.img

if __name__ == '__main__':
    os.chdir(os.path.dirname(os.path.abspath(__file__)))


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

"""
Copyright 2020 Tianshu AI Platform. All Rights Reserved.

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.
=============================================================
"""
import numpy as np
import cv2

class Extractor:

    def __init__(
            self,
            model_name,
            load_path,
            gpu_ids='0',
            use_gpu=True,
            height=256,
            width=128,
            seed=1,
            cls=0):
        self.model_name = model_name
        self.load_path = load_path
        self.gpu_ids = gpu_ids
        self.use_gpu = use_gpu
        self.height = height
        self.width = width
        self.seed = seed

        winSize = (20, 20)
        blockSize = (10, 10)
        blockStride = (5, 5)
        cellSize = (5, 5)
        nbins = 9
        self.hog = cv2.HOGDescriptor(
            winSize, blockSize, blockStride, cellSize, nbins)

    def __call__(self, input, batch_size=10, feature_type=0):
        """
        :param input: detected images, numpy array
        feature_type = 0 表示提取reid特征，1 表示提取hog特征
        :return: image features extracted from input
        """
        if feature_type == 1:
            winStride = (20, 20)
            padding = (0, 0)
            if isinstance(input, list):
                if len(input) == 0:
                    return np.array([])
                features = []
                for ind in input:
                    ind_ = cv2.resize(
                        ind, (100, 75), interpolation=cv2.INTER_LINEAR)
                    extracted_feature = self.hog.compute(
                        ind_, winStride, padding)
                    extracted_feature = extracted_feature.T
                    features.append(extracted_feature)
            else:
                input_ = cv2.resize(
                    input, (100, 75), interpolation=cv2.INTER_LINEAR)
                features = self.hog.compute(input_, winStride, padding)
                features = features.T
            features = np.vstack(features)
            return features
        else:
            return np.array([])


#!/usr/bin/python3

import argparse
import sys
import pandas as pd
import numpy as np
import gzip
import ray

def main():

    parser = argparse.ArgumentParser(
        description=
"Simplistic taxonomic classifier based on BLAST 6 output table.",

        formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("-i", "--input", metavar="PATH",
                        type=str, help="Input table of BLAST6OUT.",
                        required=True)

    parser.add_argument("-f", "--fasta", metavar="PATH",
                        type=str, help="FASTA file of study ASVs.",
                        required=True)

    parser.add_argument("-t", "--taxa", metavar="PATH",
                        type=str, help="Reference taxonomy database.",
                        required=True)

    parser.add_argument("--min_prob", metavar="FLOAT",
                        type=float, help="Max probability for taxon to be unclassified (i.e. a taxon needs probability greater than this value to be called).",
                        required=False, default=0.9)

    parser.add_argument("-p", "--proc", metavar="INT",
                        type=int, help="Number of processes to run in parallel.",
                        required=False, default=1)

    args = parser.parse_args()

    ray.init(num_cpus=args.proc)

    print("Reading in files.", file=sys.stderr)

    seqs = read_fasta(args.fasta)

    # The taxonomic and match tables will be made into global variables to make them easier to run.
    matches = pd.read_csv(args.input, sep="\t", header=None,
                             names=['query', 'target', 'identity',
                             'align_length', 'num_mismatch', 'num_gap_opens',
                             'query_start', 'query_end', 'target_start',
                             'target_end', 'evalue', 'bitscore'])

    taxa = pd.read_csv(args.taxa, sep="\t", index_col=0)

    print("Started parsing taxa table.", file=sys.stderr)

    taxa[['Kingdom', 'Phylum', 'Class', 'Order', 'Family', 'Genus', 'Species']] = taxa['Taxon'].str.split(pat=';', expand=True)

    print("Done parsing taxa table.", file=sys.stderr)

    print("Started transforming identity.", file=sys.stderr)

    matches['exp_identity'] = 2 ** matches['identity'] / 2 ** 80

    print("Done transforming identity.", file=sys.stderr)

    print("\t".join(['Sequence', 'Kingdom', 'Phylum', 'Class', 'Order', 'Family', 'Genus', 'Species']))

    matches_mem = ray.put(matches)
    taxa_mem = ray.put(taxa)

    print("Classifying ASVs.", file=sys.stderr)

    assignments_raw = [assign_asv.remote(asv, matches_mem, taxa_mem, args.min_prob) for asv in seqs.keys()]

    for assignment in ray.get(assignments_raw):
        print(assignment)

@ray.remote
def assign_asv(asv_id, in_matches, in_taxa, prob_cutoff):
    '''Assign taxonomy to each ASV based hits to target sequences in BLAST6OUT
    table.'''

    taxa_max_cutoff = {}
    taxa_max_cutoff["Species"] = 99
    taxa_max_cutoff["Genus"] = 97
    taxa_max_cutoff["Family"] = 95
    taxa_max_cutoff["Order"] = 93
    taxa_max_cutoff["Class"] = 91
    taxa_max_cutoff["Phylum"] = 89
    taxa_max_cutoff["Kingdom"] = 80

    asv_matches = in_matches.loc[in_matches['query'] == asv_id]

    # If no matches then output as Unassigned and move on.
    if asv_matches.shape[0] == 0:
        return(asv_id + "\tUnassigned" + "\tUnclassified" * 6)

    # Otherwise weight all ASV matches by the sequence similarity, where
    # the weight value is given by 2^(similarity).
    # For each taxonomic level only accept a classification if >=90% of
    # the weight supports a call, otherwise leave that level blank.

    matching_ids = list(asv_matches['target'])

    matching_taxa = in_taxa.loc[matching_ids, :]

    max_identity = asv_matches['identity'].max()

    asv_taxonomy = ""

    last_unclassified = False

    for level in ['Kingdom', 'Phylum', 'Class', 'Order', 'Family', 'Genus', 'Species']:

        if max_identity < taxa_max_cutoff[level] or last_unclassified:
            asv_taxonomy = asv_taxonomy + "\t" + "Unclassified"
            continue

        asv_level_matches = asv_matches.copy()

        asv_level_matches[level] = list(matching_taxa.loc[asv_matches['target'], level])

        asv_level_matches['prob'] = asv_matches['exp_identity'] / asv_matches['exp_identity'].sum()

        asv_level_matches = asv_level_matches.loc[:, [level, 'prob']]

        asv_level_pivot = pd.pivot_table(asv_level_matches, index=level, values='prob', aggfunc=np.sum)

        max_prob = asv_level_pivot['prob'].max()

        if max_prob > prob_cutoff:
            asv_taxonomy = asv_taxonomy + "\t" + list(asv_level_pivot.loc[asv_level_pivot['prob'] == max_prob].index)[0]
        else:
            asv_taxonomy = asv_taxonomy + "\tUnclassified"
            last_unclassified = True

    return(asv_id + asv_taxonomy)

def read_fasta(filename):
    '''Read FASTA file into a dictionary - ids are keys and sequences are
    values.'''

    seq = {}

    name = None

    fasta_in = open(filename, "r")

    for line in fasta_in:

        if line[0] == ">":

            name = line[1:]

            name = name.rstrip("\r\n")

            seq[name] = ""

        else:

            line = line.rstrip("\r\n")

            seq[name] += line

    fasta_in.close()

    return seq

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
    poc 3
    PyInstallerで生成される実行ファイルのサイズ確認用
    ◇opencvを使用
"""

from pathlib import Path
# library
import numpy as np
import cv2

def read_file(file_name: str):
    p = Path(file_name)
    with p.open('rb') as file:
        return file.read()

class ImageData(object):
    """
        cv2.adaptiveThresholdの元データに使用するグレースケール画像を管理する。
    """
    def __init__(self, file_name: str):
        """
        :param file_name: 画像ファイル名
        """
        self.__file_name = file_name
        self.__gray_scale = ImageData.imread(file_name, cv2.IMREAD_GRAYSCALE)

    @staticmethod
    def imread(file_name: str, flags: int = cv2.IMREAD_COLOR):
        """
        Unicode Path/Filename for imread Not supported.
        @see https://github.com/opencv/opencv/issues/4292
        cv2.imread alternative = np.asarray & cv2.imdecode
        Unicode Path/Filename image file read.
        :param file_name:
        :param flags: cv2.IMREAD_COLOR
        :return: {Mat}image BGR
            FileNotFoundError image is None
        """
        image = None
        try:

            buffer = np.asarray(bytearray(read_file(file_name)), dtype=np.uint8)
            image = cv2.imdecode(buffer, flags)
        except FileNotFoundError as ex:
            # cv2.imread compatible
            pass
        return image

    @staticmethod
    def imwrite(file_name: str, image, params=None) -> bool:
        """
        Unicode Path/Filename for imwrite Not supported.
        cv2.imwrite alternative = cv2.imencode & numpy.ndarray.tofile
        :param file_name
        :param image imagedata
        :param params encode params
        """
        try:
            p = Path(file_name)
            ret_val, buf = cv2.imencode(p.suffix, image, params)
            assert ret_val, 'imencode failure'
            with p.open('wb') as out:
                buf.tofile(out)
            return True
        except IOError as ex:
            pass
        return False

    @staticmethod
    def encode2PNM(image, flags: int):
        """
        読み込んだ画像データをPGMまたはPPM形式にエンコードする。
        :param image:
        :param flags:
        :return:
        """
        ext = {cv2.IMREAD_GRAYSCALE: '.PGM', cv2.IMREAD_COLOR: '.PPM'}.get(flags)
        assert ext
        ret_val, enc_img = cv2.imencode(ext, image, None)
        assert ret_val, 'encode2PNM failure'
        return enc_img.tobytes()

    @property
    def file_name(self) -> str:
        """
            画像ファイル名
        """
        return self.__file_name

    @property
    def color(self):
        """
            カラー画像(np.array)
        """
        return ImageData.imread(self.file_name)

    @property
    def gray_scale(self):
        """
            グレースケール画像(np.array)
        """
        assert self.__gray_scale is not None, 'cannot open file as image.'
        return self.__gray_scale

def main():
    image_file = r'../images/kodim07.png'
    data = ImageData(image_file)
    print(type(data.color))
    cv2.imshow('image', data.color)
    cv2.waitKey()

if __name__ == "__main__":
    main()


"""This is the WSS module.
"""
import logging
import time

import numpy as np

import wsapi

logger = logging.getLogger("WSS")
logger.addHandler(logging.NullHandler())

class WSS:
    """
    This is a class for WaveShaper module.

    :var float frequency_start: Start frequency
    :var float frequency_end: End frequency
    :var float step: Frequency step
    :var float speed_of_light: Speed of light in m/s
    :var int time_sleep: Time needed to load the WaveShaper profile before returning the bandwidth and port attenuation
    :var str folder: Folder that contains the configuration files
    :var str configfile_1: Configuration file for the WaveShaper 1
    :var str configfile_2: Configuration file for the WaveShaper 2
    """
    frequency_start = 191.250
    frequency_end = 196.274
    step = 0.001
    speed_of_light = 299792.458
    time_sleep = 5
    folder = "C:/Users/CTTC/Desktop/agent-bvt/config/"
    configfile_1 = "SN042561.wsconfig"
    configfile_2 = "SN200162.wsconfig"

    def __init__(self, id, n, m):
        """
        The constructor for the WaveShaper class.
        Initialize the WaveShaper default parameters:

            - wavelength (float): Central frequency. The range of wavelength takes 1528'773 to 1566'723 nm.
            - attenuation (float): Port attenuation.
            - phase (float): Phase.
            - bandwidth (float): Bandwidth.

        :param id: identify the WaveShaper
        :type id: int
        :param n: max number of the input ports
        :type n: int
        :param m: max number of the output ports
        :type m: int
        """
        self.id = id
        if id == 1:
            self.filename = self.folder + self.configfile_1
        elif id == 2:
            self.filename = self.folder + self.configfile_2

        self.n = n
        self.m = m
        self.wavelength = np.ones(shape=(self.n, self.m), dtype=float)
        self.attenuation = 60 * np.ones(shape=(self.n, self.m), dtype=float)
        self.phase = np.zeros(shape=(self.n, self.m), dtype=float)
        self.bandwidth = np.zeros(shape=(self.n, self.m), dtype=float)
        self.open()

    def open(self):
        """
        Create and open the WaveShaper.
        """
        name = str(self.id)
        filename = self.filename
        try:
            wsapi.ws_create_waveshaper(name, filename)
            logger.debug("WaveShaper {} created with configuration file {}".format(name, filename))
            try:
                wsapi.ws_open_waveshaper(name)
                logger.debug("WaveShaper %s opened" % name)

            except Exception as e:
                logger.error("WaveShaper {} not opened, {}".format(name, e))

        except Exception as e:
            logger.error("WaveShaper {} not created, {}".format(name, e))

    def close(self):
        """
        Close and delete the WaveShaper.
        """
        name = str(self.id)
        try:
            wsapi.ws_close_waveshaper(name)
            logger.debug("WaveShaper %s closed" % name)
            try:
                wsapi.ws_delete_waveshaper(name)
                logger.debug("WaveShaper %s deleted" % name)

            except Exception as e:
                logger.error("WaveShaper {} not closed, {}".format(name, e))

        except Exception as e:
            logger.error("WaveShaper {} not deleted, {}".format(name, e))

    def execute(self):
        """
        Load the desired profile according to the WaveShaper wavelength, port attenuation,
        phase and bandwidth for the filter configuration.
        """
        wss_id = str(self.id)
        profiletext = ""
        freq = self.speed_of_light / self.wavelength
        startfreq = freq - self.bandwidth * 0.5 * 1e-3  # start frequency in THz
        stopfreq = freq + self.bandwidth * 0.5 * 1e-3  # stop frequency in THz

        try:
            for frequency in np.arange(self.frequency_start, self.frequency_end, self.step, dtype=float):
                i = 0
                for k in range(self.n):
                    if self.wavelength[k] > 1 and startfreq[k] < frequency < stopfreq[k]:
                        profiletext += "%.3f %.1f %.1f %d\n" % (frequency, self.attenuation[k], self.phase[k], k + 1)
                        i = 1
                if i == 0:
                    profiletext += "%.3f 60.0 0.0 0\n" % frequency

            wsapi.ws_load_profile(wss_id, profiletext)
            logger.debug("WaveShaper %s profile loaded" % wss_id)

        except Exception as error:
            logger.error("WaveShaper {} profile not loaded, result description: {}, "
                         "error: {}".format(wss_id, wsapi.ws_get_result_description(-1), error))

    def configuration(self, operations):
        """
        WaveShaper configuration:

            - Set the wavelength, port attenuation, phase and bandwidth for the filter configuration.
            - Load the desired profile according to the WaveShaper filter configuration.

        :param operations: operations to configure the filter of the WaveShaper
        :type operations: list
        """
        wss_id = str(self.id)
        try:
            for i in range(len(operations)):  # for each operation
                x = operations[i]['port_in'] - 1
                y = operations[i]['port_out'] - 1
                self.wavelength[x][y] = operations[i]['lambda0']
                self.attenuation[x][y] = operations[i]['att']
                self.phase[x][y] = operations[i]['phase']
                self.bandwidth[x][y] = operations[i]['bw']

            self.execute()
            time.sleep(self.time_sleep)
            self.close()
            logger.debug("WaveShaper %s configuration finished" % wss_id)

        except Exception as error:
            logger.error("WaveShaper {} filter configuration failed, error: {}".format(wss_id, error))
            raise error


"""
Copyright 2017-2018 Fizyr (https://fizyr.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.
"""

import numpy as np
import keras
from .compute_overlap import compute_overlap

class AnchorParameters:
    """ The parameteres that define how anchors are generated.

    Args
        sizes   : List of sizes to use. Each size corresponds to one feature level.
        strides : List of strides to use. Each stride correspond to one feature level.
        ratios  : List of ratios to use per location in a feature map.
        scales  : List of scales to use per location in a feature map.
    """
    def __init__(self, sizes, strides, ratios, scales):
        self.sizes   = sizes
        self.strides = strides
        self.ratios  = ratios
        self.scales  = scales

    def num_anchors(self):
        return len(self.ratios) * len(self.scales)

"""
The default anchor parameters.
"""
AnchorParameters.default = AnchorParameters(
    sizes   = [32, 64, 128, 256, 512],
    strides = [8, 16, 32, 64, 128],
    ratios  = np.array([0.5, 1, 2], keras.backend.floatx()),
    scales  = np.array([2 ** 0, 2 ** (1.0 / 3.0), 2 ** (2.0 / 3.0)], keras.backend.floatx()),
)

def anchor_targets_bbox(
    anchors,
    image_group,
    annotations_group,
    num_classes,
    negative_overlap=0.4,
    positive_overlap=0.5
):
    """ Generate anchor targets for bbox detection.

    Args
        anchors: np.array of annotations of shape (N, 4) for (x1, y1, x2, y2).
        image_group: List of BGR images.
        annotations_group: List of annotation dictionaries with each annotation containing 'labels' and 'bboxes' of an image.
        num_classes: Number of classes to predict.
        mask_shape: If the image is padded with zeros, mask_shape can be used to mark the relevant part of the image.
        negative_overlap: IoU overlap for negative anchors (all anchors with overlap < negative_overlap are negative).
        positive_overlap: IoU overlap or positive anchors (all anchors with overlap > positive_overlap are positive).

    Returns
        labels_batch: batch that contains labels & anchor states (np.array of shape (batch_size, N, num_classes + 1),
                      where N is the number of anchors for an image and the last column defines the anchor state (-1 for ignore, 0 for bg, 1 for fg).
        regression_batch: batch that contains bounding-box regression targets for an image & anchor states (np.array of shape (batch_size, N, 4 + 1),
                      where N is the number of anchors for an image, the first 4 columns define regression targets for (x1, y1, x2, y2) and the
                      last column defines anchor states (-1 for ignore, 0 for bg, 1 for fg).
    """

    assert(len(image_group) == len(annotations_group)), "The length of the images and annotations need to be equal."
    assert(len(annotations_group) > 0), "No data received to compute anchor targets for."
    for annotations in annotations_group:
        assert('bboxes' in annotations), "Annotations should contain bboxes."
        assert('labels' in annotations), "Annotations should contain labels."

    batch_size = len(image_group)

    regression_batch  = np.zeros((batch_size, anchors.shape[0], 4 + 1), dtype=keras.backend.floatx())
    labels_batch      = np.zeros((batch_size, anchors.shape[0], num_classes + 1), dtype=keras.backend.floatx())

    # compute labels and regression targets
    for index, (image, annotations) in enumerate(zip(image_group, annotations_group)):
        if annotations['bboxes'].shape[0]:
            # obtain indices of gt annotations with the greatest overlap
            positive_indices, ignore_indices, argmax_overlaps_inds = compute_gt_annotations(anchors, annotations['bboxes'], negative_overlap, positive_overlap)

            labels_batch[index, ignore_indices, -1]       = -1
            labels_batch[index, positive_indices, -1]     = 1

            regression_batch[index, ignore_indices, -1]   = -1
            regression_batch[index, positive_indices, -1] = 1

            # compute target class labels
            labels_batch[index, positive_indices, annotations['labels'][argmax_overlaps_inds[positive_indices]].astype(int)] = 1

            regression_batch[index, :, :-1] = bbox_transform(anchors, annotations['bboxes'][argmax_overlaps_inds, :])

        # ignore annotations outside of image
        if image.shape:
            anchors_centers = np.vstack([(anchors[:, 0] + anchors[:, 2]) / 2, (anchors[:, 1] + anchors[:, 3]) / 2]).T
            indices = np.logical_or(anchors_centers[:, 0] >= image.shape[1], anchors_centers[:, 1] >= image.shape[0])

            labels_batch[index, indices, -1]     = -1
            regression_batch[index, indices, -1] = -1

    return regression_batch, labels_batch

def compute_gt_annotations(
    anchors,
    annotations,
    negative_overlap=0.4,
    positive_overlap=0.5
):
    """ Obtain indices of gt annotations with the greatest overlap.

    Args
        anchors: np.array of annotations of shape (N, 4) for (x1, y1, x2, y2).
        annotations: np.array of shape (N, 5) for (x1, y1, x2, y2, label).
        negative_overlap: IoU overlap for negative anchors (all anchors with overlap < negative_overlap are negative).
        positive_overlap: IoU overlap or positive anchors (all anchors with overlap > positive_overlap are positive).

    Returns
        positive_indices: indices of positive anchors
        ignore_indices: indices of ignored anchors
        argmax_overlaps_inds: ordered overlaps indices
    """

    overlaps = compute_overlap(anchors.astype(np.float64), annotations.astype(np.float64))
    argmax_overlaps_inds = np.argmax(overlaps, axis=1)
    max_overlaps = overlaps[np.arange(overlaps.shape[0]), argmax_overlaps_inds]

    # assign "dont care" labels
    positive_indices = max_overlaps >= positive_overlap
    ignore_indices = (max_overlaps > negative_overlap) & ~positive_indices

    return positive_indices, ignore_indices, argmax_overlaps_inds

def layer_shapes(image_shape, model):
    """Compute layer shapes given input image shape and the model.

    Args
        image_shape: The shape of the image.
        model: The model to use for computing how the image shape is transformed in the pyramid.

    Returns
        A dictionary mapping layer names to image shapes.
    """
    shape = {
        model.layers[0].name: (None,) + image_shape,
    }

    for layer in model.layers[1:]:
        nodes = layer._inbound_nodes
        for node in nodes:
            inputs = [shape[lr.name] for lr in node.inbound_layers]
            if not inputs:
                continue
            shape[layer.name] = layer.compute_output_shape(inputs[0] if len(inputs) == 1 else inputs)

    return shape

def make_shapes_callback(model):
    """ Make a function for getting the shape of the pyramid levels.
    """
    def get_shapes(image_shape, pyramid_levels):
        shape = layer_shapes(image_shape, model)
        image_shapes = [shape["P{}".format(level)][1:3] for level in pyramid_levels]
        return image_shapes

    return get_shapes

def guess_shapes(image_shape, pyramid_levels):
    """Guess shapes based on pyramid levels.

    Args
         image_shape: The shape of the image.
         pyramid_levels: A list of what pyramid levels are used.

    Returns
        A list of image shapes at each pyramid level.
    """
    image_shape = np.array(image_shape[:2])
    image_shapes = [(image_shape + 2 ** x - 1) // (2 ** x) for x in pyramid_levels]
    return image_shapes

def anchors_for_shape(
    image_shape,
    pyramid_levels=None,
    anchor_params=None,
    shapes_callback=None,
):
    """ Generators anchors for a given shape.

    Args
        image_shape: The shape of the image.
        pyramid_levels: List of ints representing which pyramids to use (defaults to [3, 4, 5, 6, 7]).
        anchor_params: Struct containing anchor parameters. If None, default values are used.
        shapes_callback: Function to call for getting the shape of the image at different pyramid levels.

    Returns
        np.array of shape (N, 4) containing the (x1, y1, x2, y2) coordinates for the anchors.
    """

    if pyramid_levels is None:
        pyramid_levels = [3, 4, 5, 6, 7]

    if anchor_params is None:
        anchor_params = AnchorParameters.default

    if shapes_callback is None:
        shapes_callback = guess_shapes
    image_shapes = shapes_callback(image_shape, pyramid_levels)

    # compute anchors over all pyramid levels
    all_anchors = np.zeros((0, 4))
    for idx, p in enumerate(pyramid_levels):
        anchors = generate_anchors(
            base_size=anchor_params.sizes[idx],
            ratios=anchor_params.ratios,
            scales=anchor_params.scales
        )
        shifted_anchors = shift(image_shapes[idx], anchor_params.strides[idx], anchors)
        all_anchors     = np.append(all_anchors, shifted_anchors, axis=0)

    return all_anchors

def shift(shape, stride, anchors):
    """ Produce shifted anchors based on shape of the map and stride size.

    Args
        shape  : Shape to shift the anchors over.
        stride : Stride to shift the anchors with over the shape.
        anchors: The anchors to apply at each location.
    """

    # create a grid starting from half stride from the top left corner
    shift_x = (np.arange(0, shape[1]) + 0.5) * stride
    shift_y = (np.arange(0, shape[0]) + 0.5) * stride

    shift_x, shift_y = np.meshgrid(shift_x, shift_y)

    shifts = np.vstack((
        shift_x.ravel(), shift_y.ravel(),
        shift_x.ravel(), shift_y.ravel()
    )).transpose()

    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    # reshape to (K*A, 4) shifted anchors
    A = anchors.shape[0]
    K = shifts.shape[0]
    all_anchors = (anchors.reshape((1, A, 4)) + shifts.reshape((1, K, 4)).transpose((1, 0, 2)))
    all_anchors = all_anchors.reshape((K * A, 4))

    return all_anchors

def generate_anchors(base_size=16, ratios=None, scales=None):
    """
    Generate anchor (reference) windows by enumerating aspect ratios X
    scales w.r.t. a reference window.
    """

    if ratios is None:
        ratios = AnchorParameters.default.ratios

    if scales is None:
        scales = AnchorParameters.default.scales

    num_anchors = len(ratios) * len(scales)

    # initialize output anchors
    anchors = np.zeros((num_anchors, 4))

    # scale base_size
    anchors[:, 2:] = base_size * np.tile(scales, (2, len(ratios))).T

    # compute areas of anchors
    areas = anchors[:, 2] * anchors[:, 3]

    # correct for ratios
    anchors[:, 2] = np.sqrt(areas / np.repeat(ratios, len(scales)))
    anchors[:, 3] = anchors[:, 2] * np.repeat(ratios, len(scales))

    # transform from (x_ctr, y_ctr, w, h) -> (x1, y1, x2, y2)
    anchors[:, 0::2] -= np.tile(anchors[:, 2] * 0.5, (2, 1)).T
    anchors[:, 1::2] -= np.tile(anchors[:, 3] * 0.5, (2, 1)).T

    return anchors

def bbox_transform(anchors, gt_boxes, mean=None, std=None):
    """Compute bounding-box regression targets for an image."""

    # The Mean and std are calculated from COCO dataset.
    # Bounding box normalization was firstly introduced in the Fast R-CNN paper.
    # See https://github.com/fizyr/keras-retinanet/issues/1273#issuecomment-585828825  for more details
    if mean is None:
        mean = np.array([0, 0, 0, 0])
    if std is None:
        std = np.array([0.2, 0.2, 0.2, 0.2])

    if isinstance(mean, (list, tuple)):
        mean = np.array(mean)
    elif not isinstance(mean, np.ndarray):
        raise ValueError('Expected mean to be a np.ndarray, list or tuple. Received: {}'.format(type(mean)))

    if isinstance(std, (list, tuple)):
        std = np.array(std)
    elif not isinstance(std, np.ndarray):
        raise ValueError('Expected std to be a np.ndarray, list or tuple. Received: {}'.format(type(std)))

    anchor_widths  = anchors[:, 2] - anchors[:, 0]
    anchor_heights = anchors[:, 3] - anchors[:, 1]

    # According to the information provided by a keras-retinanet author, they got marginally better results using
    # the following way of bounding box parametrization.
    # See https://github.com/fizyr/keras-retinanet/issues/1273#issuecomment-585828825 for more details
    targets_dx1 = (gt_boxes[:, 0] - anchors[:, 0]) / anchor_widths
    targets_dy1 = (gt_boxes[:, 1] - anchors[:, 1]) / anchor_heights
    targets_dx2 = (gt_boxes[:, 2] - anchors[:, 2]) / anchor_widths
    targets_dy2 = (gt_boxes[:, 3] - anchors[:, 3]) / anchor_heights

    targets = np.stack((targets_dx1, targets_dy1, targets_dx2, targets_dy2))
    targets = targets.T

    targets = (targets - mean) / std

    return targets


#!/usr/bin/python2.7

import tensorflow as tf
import numpy as np
import argparse
from models.cnn import ModelCNN
from models.rnn import ModelRNN
from utils.base_batch_gen import Base_batch_generator
from utils.rnn_batch_gen import RNN_batch_generator
from utils.cnn_batch_gen import CNN_batch_generator
from utils.helper_functions import read_mapping_dict, encode_content, write_predictions, get_label_length_seq

parser = argparse.ArgumentParser()

parser.add_argument("--model", default="rnn", help="select model: [\"rnn\", \"cnn\"]")
parser.add_argument("--action", default="predict", help="select action: [\"train\", \"predict\"]")

parser.add_argument("--mapping_file", default="./data/mapping_bf.txt")
parser.add_argument("--vid_list_file", default="./data/test.split1.bundle")
parser.add_argument("--model_save_path", default="./save_dir/models/rnn")
parser.add_argument("--results_save_path", default="./save_dir/results/rnn")

parser.add_argument("--save_freq", type=int, default=1)
parser.add_argument("--batch_size", type=int, default=512)
parser.add_argument("--nEpochs", type=int, default=20)
parser.add_argument("--eval_epoch", type=int, default=20)

#RNN specific parameters
parser.add_argument("--rnn_size", type=int, default=256)
parser.add_argument("--num_layers", type=int, default=2)
parser.add_argument("--max_seq_sz", type=int, default=25)
parser.add_argument("--alpha", type=float, default=6, help="a scalar value used in normalizing the input length")
parser.add_argument("--n_iterations", type=int, default=10, help="number of training examples corresponding to each action segment for the rnn")

#CNN specific parameters
parser.add_argument("--nRows", type=int, default=128)
parser.add_argument("--sigma", type=int, default=3, help="sigma for the gaussian smoothing step")

#Test on GT or decoded input
parser.add_argument("--input_type", default="decoded", help="select input type: [\"decoded\", \"gt\"]")
parser.add_argument("--decoded_path", default="./data/decoded/split1")

################################################################################################################################################

args, unknown = parser.parse_known_args()

actions_dict = read_mapping_dict(args.mapping_file)
nClasses = len(actions_dict)

file_ptr = open(args.vid_list_file, 'r')
list_of_videos = file_ptr.read().split('\n')[1:-1]

################
# Training #####
################
if args.action == "train":
    model = None
    batch_gen = Base_batch_generator()

    if args.model == "rnn":
        model = ModelRNN(nClasses, args.rnn_size, args.max_seq_sz, args.num_layers)
        batch_gen = RNN_batch_generator(nClasses, args.n_iterations, args.max_seq_sz, actions_dict, args.alpha)
    elif args.model == "cnn":
        model = ModelCNN(args.nRows, nClasses)
        batch_gen = CNN_batch_generator(args.nRows, nClasses, actions_dict)

    batch_gen.read_data(list_of_videos)
    with tf.Session() as sess:
        model.train(sess, args.model_save_path, batch_gen, args.nEpochs, args.save_freq, args.batch_size)

##################
# Prediction #####
##################
elif args.action == "predict":
    pred_percentages = [.1, .2, .3, .5]
    obs_percentages = [.2, .3, .5]
    model_restore_path = args.model_save_path+"/epoch-"+str(args.eval_epoch)+"/model.ckpt"

    if args.model == "rnn":
        model = ModelRNN(nClasses, args.rnn_size, args.max_seq_sz, args.num_layers)
        for vid in list_of_videos:
            f_name = vid.split('/')[-1].split('.')[0]
            observed_content=[]
            vid_len = 0
            if args.input_type == "gt":
                file_ptr = open(vid, 'r')
                content = file_ptr.read().split('\n')[:-1]
                vid_len = len(content)

            for obs_p in obs_percentages:

                if args.input_type == "decoded":
                    file_ptr = open(args.decoded_path+"/obs"+str(obs_p)+"/"+f_name+'.txt', 'r')
                    observed_content = file_ptr.read().split('\n')[:-1]
                    vid_len = int(len(observed_content)/obs_p)
                elif args.input_type == "gt":
                    observed_content = content[:int(obs_p*vid_len)]
                T = (1.0/args.alpha)*vid_len

                for pred_p in pred_percentages:
                    pred_len = int(pred_p*vid_len)
                    output_len = pred_len + len(observed_content)

                    label_seq, length_seq = get_label_length_seq(observed_content)
                    with tf.Session() as sess:
                        label_seq, length_seq = model.predict(sess, model_restore_path, pred_len, label_seq, length_seq, actions_dict, T)

                    recognition = []
                    for i in range(len(label_seq)):
                        recognition = np.concatenate((recognition, [label_seq[i]]*int(length_seq[i])))
                    recognition = recognition[:output_len]
                    #write results to file
                    f_name = vid.split('/')[-1].split('.')[0]
                    path=args.results_save_path+"/obs"+str(obs_p)+"-pred"+str(pred_p)
                    write_predictions(path, f_name, recognition)

    elif args.model == "cnn":
        model = ModelCNN(args.nRows, nClasses)
        for vid in list_of_videos:
            f_name = vid.split('/')[-1].split('.')[0]
            observed_content=[]
            vid_len = 0
            if args.input_type == "gt":
                file_ptr = open(vid, 'r')
                content = file_ptr.read().split('\n')[:-1]
                vid_len = len(content)

            for obs_p in obs_percentages:
                if args.input_type == "decoded":
                    file_ptr = open(args.decoded_path+"/obs"+str(obs_p)+"/"+f_name+'.txt', 'r')
                    observed_content = file_ptr.read().split('\n')[:-1]
                    vid_len = int(len(observed_content)/obs_p)
                elif args.input_type == "gt":
                    observed_content = content[:int(obs_p*vid_len)]

                input_x = encode_content(observed_content, args.nRows, nClasses, actions_dict)
                input_x = [np.reshape(input_x, [args.nRows, nClasses, 1])]

                with tf.Session() as sess:
                    label_seq, length_seq = model.predict(sess, model_restore_path, input_x, args.sigma, actions_dict)

                recognition = []
                for i in range(len(label_seq)):
                    recognition = np.concatenate((recognition, [label_seq[i]]*int(0.5*vid_len*length_seq[i]/args.nRows)))
                recognition = np.concatenate((observed_content,recognition))
                diff = int((0.5+obs_p)*vid_len)-len(recognition)
                for i in range(diff):
                    recognition = np.concatenate((recognition, [label_seq[-1]]))
                #write results to file
                for pred_p in pred_percentages:
                    path=args.results_save_path+"/obs"+str(obs_p)+"-pred"+str(pred_p)
                    write_predictions(path, f_name, recognition[:int((pred_p+obs_p)*vid_len)])


#!/usr/bin/env python

import sys
import sqlite3
import Bio
from Bio import SeqIO
import pickle
import tagmatch
import pandas as pd
import numpy as np

db = sqlite3.connect(sys.argv[1])
cursor = db.cursor()
prec_tol = float(sys.argv[2])
gap_tol = float(sys.argv[3])
fixed_mods = [i for i in sys.argv[4].split(';') if i!= '']
variable_mods = [i for i in sys.argv[5].split(';') if i!= '']

limit = prec_tol + gap_tol + gap_tol

#cursor.execute("CREATE TABLE IF NOT EXISTS tags (scanid TEXT, tagid TEXT, seqstr TEXT, rec BYTEA, source TEXT, UNIQUE( tagid, seqstr, source ) )")

tag_df = pd.read_sql("SELECT DISTINCT tagid, seqstr from tags", db, columns=['tagid', 'seqstr'])
tag_df = tag_df.drop_duplicates()
print(len(tag_df))
load = []

def validate(table):
    global load
    tid = table['tagid']
    _ = tid.split(';')
    mw = np.round(float(_[2].split('mw=')[1]),4)
    ngap = np.round(float(_[3].split('ngap=')[1]),4)
    cgap = np.round(float(_[4].split('cgap=')[1]),4)
    seq = table['seqstr']
    for mass in tagmatch.peptide_mass(seq, fixed_modifications=fixed_mods, variable_modifications = variable_mods, nterm=True, cterm=True):
        total = mass + ngap + cgap
        diff = np.absolute(total - mw)
        if diff < limit:
            load.append((tid, seq))

tag_df.apply(validate, axis=1)
del tag_df

cursor.execute("CREATE TABLE IF NOT EXISTS valid (tagid TEXT, seqstr TEXT);")

db.commit()   # This creates table, makes sure that tagid, seqstr and source are unique

sql = ("DROP INDEX IF EXISTS index_seqs_valid;")
cursor.execute(sql)
db.commit()

sql = ("DROP INDEX IF EXISTS index_tagids_valid;")
cursor.execute(sql)
db.commit()

sql = ("DROP INDEX IF EXISTS index_tagids_seqs_valid;")
cursor.execute(sql)
db.commit()

sql = "INSERT OR IGNORE INTO valid VALUES(?, ?);"
cursor.executemany(sql, list(set(load)) )
db.commit()

sql = ("CREATE INDEX IF NOT EXISTS index_seqs_valid ON valid (seqstr);")
cursor.execute(sql)
db.commit()

sql = ("CREATE INDEX IF NOT EXISTS index_tagids_valid ON valid (tagid);")
cursor.execute(sql)
db.commit()

sql = ("CREATE INDEX IF NOT EXISTS index_tagids_seqs_valid ON valid (tagid, seqstr);")
cursor.execute(sql)
db.commit()
print(len(list(set(load))))


# -*- coding: utf-8 -*-
"""
Class definition of YOLO_v3 style detection model on image and video
"""

import colorsys
# import os
from timeit import default_timer as timer

import numpy as np
from keras import backend as K
from keras.models import load_model
# from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw

from yolo3.model import yolo_eval, yolo_body, tiny_yolo_body
from yolo3.utils import letterbox_image
import os

class YOLO(object):
    _defaults = {
        "model_path": 'model_data/yolo.h5',
        "anchors_path": 'model_data/yolo_anchors.txt',
        "classes_path": 'model_data/coco_classes.txt',
        "score": 0.3,
        "iou": 0.45,
        "model_image_size" : (416, 416),
        "gpu_num": 1,
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)  # set up default values
        self.__dict__.update(kwargs)  # and update with user overrides

        # get class names
        self.class_names = self._get_class()
        # get yolo anchor parameter
        self.anchors = self._get_anchors()

        # generate session, and loading model
        self.sess = K.get_session()
        self.boxes, self.scores, self.classes = self.generate()

    def _get_class(self):
        """
        Get class names from txt file
        for example person, traffic light, animal
        :return:
        """
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    def _get_anchors(self):
        """
        Get anchors parameter of yolo
        you can get information from Internet
        :return:
        """
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    def generate(self):
        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras model or weights must be a .h5 file.'

        # Load model
        self.yolo_model = load_model(model_path, compile=False)
        print('{} model, anchors, and classes loaded.'.format(model_path))

        # Generate colors for drawing bounding boxes.
        hsv_tuples = [(x / len(self.class_names), 1., 1.)
                      for x in range(len(self.class_names))]
        self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
        self.colors = list(
            map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                self.colors))
        np.random.seed(10101)  # Fixed seed for consistent colors across runs.
        np.random.shuffle(self.colors)  # Shuffle colors to decorrelate adjacent classes.
        np.random.seed(None)  # Reset seed to default.

        # Generate output tensor targets for filtered bounding boxes.
        self.input_image_shape = K.placeholder(shape=(2, ))

        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
                len(self.class_names), self.input_image_shape,
                score_threshold=self.score, iou_threshold=self.iou)

        return boxes, scores, classes

    def detect_image(self, image):
        """
        detect objects from image
        and draw the objects in image with boxs
        the resault is a new image
        :param image:
        :return:
        """
        if self.model_image_size != (None, None):
            assert self.model_image_size[0] % 32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1] % 32 == 0, 'Multiples of 32 required'
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        # print(image_data.shape)
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,
                self.input_image_shape: [image.size[1], image.size[0]],
                K.learning_phase(): 0
            })

        # print('Found {} boxes for {}'.format(len(out_boxes), 'img'))
        # font setting
        font = ImageFont.truetype(font='font/FiraMono-Medium.otf',
                    size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        thickness = (image.size[0] + image.size[1]) // 300

        for i, c in reversed(list(enumerate(out_classes))):
            predicted_class = self.class_names[c]
            box = out_boxes[i]
            score = out_scores[i]

            top, left, bottom, right = box
            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))

            label = '{} {:.2f}'.format(predicted_class, score)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)

            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            # My kingdom for a good redistributable image drawing library.
            for i in range(thickness):
                draw.rectangle(
                    [left + i, top + i, right - i, bottom - i],
                    outline=self.colors[c])
            draw.rectangle(
                [tuple(text_origin), tuple(text_origin + label_size)],
                fill=self.colors[c])
            draw.text(text_origin, label, fill=(0, 0, 0), font=font)
            del draw

        return image

    def detect_traffic_light(self, image):
        """
        detect traffic light from image
        the resault is the location of the traffic light
        :param image:
        :return:
        """
        if self.model_image_size != (None, None):
            assert self.model_image_size[0] % 32 == 0, 'Multiples of 32 required'
            assert self.model_image_size[1] % 32 == 0, 'Multiples of 32 required'
            boxed_image = letterbox_image(image, tuple(reversed(self.model_image_size)))
        else:
            new_image_size = (image.width - (image.width % 32),
                              image.height - (image.height % 32))
            boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')

        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,
                self.input_image_shape: [image.size[1], image.size[0]],
                K.learning_phase(): 0
            })

        traffic_light = None

        for i, c in reversed(list(enumerate(out_classes))):
            predicted_class = self.class_names[c]

            if predicted_class == 'traffic light':
                top, left, bottom, right = out_boxes[i]

                top = max(0, np.floor(top + 0.5).astype('int32'))
                left = max(0, np.floor(left + 0.5).astype('int32'))
                bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
                right = min(image.size[0], np.floor(right + 0.5).astype('int32'))

                traffic_light = ((left, top), (right, bottom))
                break

        return traffic_light

    def close_session(self):
        self.sess.close()

if __name__ == '__main__':
    yolo = YOLO()
    start = timer()  # start timer

    image = Image.open(r'./image_{0}'.format(0) + '.jpg')
    image = yolo.detect_image(image)
    image.save('./image_out_{0}'.format(0) + '.jpg')

    # ((left, top), (right, bottom)) = yolo.detect_traffic_light(image)

    end = timer()   # end timer
    print(end - start)

    # finish session
    yolo.close_session()
    # print((left, top), (right, bottom))