#ipcam.py
import sys
import socket
import numpy as np

sys.path.append(r'C:\Users\dledb\Anaconda3\Lib\site-packages')

import cv2

###
# Configuration Step for darknet in python.
###
yolo_weight = r"C:\darknet\build\darknet\x64\yolov4.weights"
yolo_config = r"C:\darknet\build\darknet\x64\cfg\yolov4.cfg"
coco_labels = r"C:\darknet\build\darknet\x64\data\coco.names"
net = cv2.dnn.readNet(yolo_weight, yolo_config)

# Load coco object names file
classes = []
with open(coco_labels, "r") as f:
    classes = [line.strip() for line in f.readlines()]

layer_names = net.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
colors = np.random.uniform(0, 255, size=(len(classes), 3))

yolov4_width = 640
yolov4_height = 480

###
# End Darknet config
###

client_socket = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
host_ip = '192.168.4.1' # paste your server ip address here
port = 5000
print("Connecting to socket on {}:{}...".format(host_ip, port))
client_socket.connect((host_ip,port)) # a tuple
print("Connection Successful")

imgdata = None
data_buffer = b''
# cap = cv2.VideoCapture('rtsp://admin:123456@192.168.4.1/H264?ch=1&subtype=0')
# cap = cv2.VideoCapture(0)

while True:
    data_buffer+=(client_socket.recv(512))
    start_idx = data_buffer.find(b'\xff\xd8')
    end_idx = data_buffer.find(b'\xff\xd9')
    # print("found begin: ", a, " found end: ", b)
    i = None

    # At startup we might get an end before we get the first start, if
    # that is the case then throw away the data before start
    if end_idx > -1 and end_idx < start_idx:
        data_buffer = data_buffer[start_idx:]

    # We have a start and an end of the image in the buffer now
    if start_idx > -1 and end_idx > -1 and end_idx > start_idx:
        # Pick out the image to render ...
        imgdata = data_buffer[start_idx:end_idx + 2]
        # .. and remove it from the buffer
        data_buffer = data_buffer[end_idx + 2 :]
        print(imgdata)
        i = cv2.imdecode(np.fromstring(imgdata, dtype=np.uint8),cv2.IMREAD_GRAYSCALE)

        #print('About to start the Read command')
        img = i

        # gives frame 3 channels
        img = np.stack((img,)*3, axis=-1)

        height, width, channels = img.shape

        # Detecting objects
        blob = cv2.dnn.blobFromImage(img, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
        net.setInput(blob)
        outs = net.forward(output_layers)

        # Showing informations on the screen
        class_ids = []
        confidences = []
        boxes = []
        for out in outs:
            for detection in out:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                if confidence > 0.5:
                    # Object detected
                    center_x = int(detection[0] * width)
                    center_y = int(detection[1] * height)
                    w = int(detection[2] * width)
                    h = int(detection[3] * height)
                    # Rectangle coordinates
                    x = int(center_x - w / 2)
                    y = int(center_y - h / 2)
                    boxes.append([x, y, w, h])
                    confidences.append(float(confidence))
                    class_ids.append(class_id)

        indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)

        font = cv2.FONT_HERSHEY_DUPLEX
        for i in range(len(boxes)):
            if i in indexes:
                x, y, w, h = boxes[i]
                label = str(classes[class_ids[i]])
                confidence_label = int(confidences[i] * 100)
                color = colors[i]
                cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
                cv2.putText(img, f'{label, confidence_label}', (x-25, y + 75), font, 2, color, 2)

        cv2.imshow("Image", img)
        # Close video window by pressing 'x'
        if cv2.waitKey(1) & 0xFF == ord('x'):
            break

# cap.release()
client_socket.close()
cv2.destroyAllWindows()


import tensorflow as tf
import tensorflow_hub as hub
# import streamlit as st
# import matplotlib
# import matplotlib.pyplot as plt
import os
import pandas
import numpy as np
from PIL import Image, ImageOps

# What I have to do:

# 1. Get create_batches
# 2. Get unbatchify
# 3. download the model.h5 thing
# 4. Load it
# 5. Predict
# 6. Plot
# 7. Streamlit stuff

breednames = np.array(["A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q",
             "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "del", "nothing", "space"])
IMAGE_SIZE = 200
def bgr_to_rgb(input, name=None):
    """
    Convert a BGR image to RGB.
    Args:
      input: A 3-D (`[H, W, 3]`) or 4-D (`[N, H, W, 3]`) Tensor.
      name: A name for the operation (optional).
    Returns:
      A 3-D (`[H, W, 3]`) or 4-D (`[N, H, W, 3]`) Tensor.
    """
    bgr = tf.unstack(input, axis=-1)
    b, g, r = bgr[0], bgr[1], bgr[2]
    return tf.stack([r, g, b], axis=-1)

def process_image(img):
	data = np.ndarray(shape=(1, 200, 200, 3), dtype=np.float32)
	# image = tf.constant(img, dtype="float32")

	image = tf.image.resize(img, [200, 200])
	image = bgr_to_rgb(image)

	image_array = np.asarray(image)
	normalized_image_array = (image_array.astype(np.float32) / 127.0) -1

	data[0] = normalized_image_array
	return data

def preds_to_text(prediction_proba):
  return breednames[np.argmax(prediction_proba)]

def load_model(model_path):
  print(f'Loading model from: {model_path}...')
  model = tf.keras.models.load_model(model_path,
                                     custom_objects={"KerasLayer": hub.KerasLayer})
  return model

def predict_custom(image, model):

	custom_data = process_image(image)
	custom_preds = model.predict(custom_data)
	conf = f'{np.max(custom_preds[0])* 100:.2f}%'
	custom_preds_labels = preds_to_text(custom_preds)
	return custom_preds_labels, conf


from functools import reduce

# from keras import layers
# from keras import initializers
# from keras import models
# from keras_ import EfficientNetB0, EfficientNetB1, EfficientNetB2
# from keras_ import EfficientNetB3, EfficientNetB4, EfficientNetB5, EfficientNetB6

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers as KL
from tensorflow.keras import initializers
from tensorflow.keras import models
from tensorflow.keras import backend as K
from tensorflow.keras.optimizers import Adam, SGD
from tfkeras import EfficientNetB0, EfficientNetB1, EfficientNetB2
from tfkeras import EfficientNetB3, EfficientNetB4, EfficientNetB5, EfficientNetB6

from layers import ClipBoxes, RegressBoxes, FilterDetections, wBiFPNAdd, BatchNormalization
from layers import DetectionLayer
from initializers import PriorProbability
from utils.anchors import anchors_for_shape, anchor_targets_bbox, AnchorParameters
import numpy as np
from utils.graph_funcs import trim_zeros_graph, overlaps_graph, batch_slice
from utils.graph_funcs import norm_boxes_graph, denorm_boxes_graph
from losses import focal, smooth_l1

from utils.helper import norm_boxes
import os

w_bifpns = [64, 88, 112, 160, 224, 288, 384]
d_bifpns = [3, 4, 5, 6, 7, 7, 8]
d_heads = [3, 3, 3, 4, 4, 4, 5]
image_sizes = [512, 640, 768, 896, 1024, 1280, 1408]
backbones = [EfficientNetB0, EfficientNetB1, EfficientNetB2,
             EfficientNetB3, EfficientNetB4, EfficientNetB5, EfficientNetB6]

MOMENTUM = 0.997
EPSILON = 1e-4
MAX_INSTANCES = 10
TRAIN_ROIS_PER_IMAGE = 100
DETECTION_MAX_INSTANCES = 100
# Minimum probability value to accept a detected instance
# ROIs below this threshold are skipped
DETECTION_MIN_CONFIDENCE = 0.7

# Non-maximum suppression threshold for detection
DETECTION_NMS_THRESHOLD = 0.3

def SeparableConvBlock(num_channels, kernel_size, strides, name, freeze_bn=False):
    f1 = KL.SeparableConv2D(num_channels, kernel_size=kernel_size, strides=strides, padding='same',
                                use_bias=True, name=f'{name}/conv')
    f2 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON, name=f'{name}/bn')
    # f2 = BatchNormalization(freeze=freeze_bn, name=f'{name}/bn')
    return reduce(lambda f, g: lambda *args, **kwargs: g(f(*args, **kwargs)), (f1, f2))

def ConvBlock(num_channels, kernel_size, strides, name, freeze_bn=False):
    f1 = KL.Conv2D(num_channels, kernel_size=kernel_size, strides=strides, padding='same',
                       use_bias=True, name='{}_conv'.format(name))
    f2 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON, name='{}_bn'.format(name))
    # f2 = BatchNormalization(freeze=freeze_bn, name='{}_bn'.format(name))
    f3 = KL.ReLU(name='{}_relu'.format(name))
    return reduce(lambda f, g: lambda *args, **kwargs: g(f(*args, **kwargs)), (f1, f2, f3))

def build_wBiFPN(features, num_channels, id, freeze_bn=False):
    if id == 0:
        _, _, C3, C4, C5 = features
        P3_in = C3
        P4_in = C4
        P5_in = C5
        P6_in = KL.Conv2D(num_channels, kernel_size=1, padding='same', name='resample_p6/conv2d')(C5)
        P6_in = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON, name='resample_p6/bn')(P6_in)
        # P6_in = BatchNormalization(freeze=freeze_bn, name='resample_p6/bn')(P6_in)
        P6_in = KL.MaxPooling2D(pool_size=3, strides=2, padding='same', name='resample_p6/maxpool')(P6_in)
        P7_in = KL.MaxPooling2D(pool_size=3, strides=2, padding='same', name='resample_p7/maxpool')(P6_in)
        P7_U = KL.UpSampling2D()(P7_in)
        P6_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode0/add')([P6_in, P7_U])
        P6_td = KL.Activation(lambda x: tf.nn.swish(x))(P6_td)
        P6_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode0/op_after_combine5')(P6_td)
        P5_in_1 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/conv2d')(P5_in)
        P5_in_1 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/bn')(P5_in_1)
        # P5_in_1 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/bn')(P5_in_1)
        P6_U = KL.UpSampling2D()(P6_td)
        P5_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode1/add')([P5_in_1, P6_U])
        P5_td = KL.Activation(lambda x: tf.nn.swish(x))(P5_td)
        P5_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode1/op_after_combine6')(P5_td)
        P4_in_1 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/conv2d')(P4_in)
        P4_in_1 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/bn')(P4_in_1)
        # P4_in_1 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/bn')(P4_in_1)
        P5_U = KL.UpSampling2D()(P5_td)
        P4_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode2/add')([P4_in_1, P5_U])
        P4_td = KL.Activation(lambda x: tf.nn.swish(x))(P4_td)
        P4_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode2/op_after_combine7')(P4_td)
        P3_in = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                              name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/conv2d')(P3_in)
        P3_in = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                          name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/bn')(P3_in)
        # P3_in = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/bn')(P3_in)
        P4_U = KL.UpSampling2D()(P4_td)
        P3_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode3/add')([P3_in, P4_U])
        P3_out = KL.Activation(lambda x: tf.nn.swish(x))(P3_out)
        P3_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode3/op_after_combine8')(P3_out)
        P4_in_2 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/conv2d')(P4_in)
        P4_in_2 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/bn')(P4_in_2)
        # P4_in_2 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/bn')(P4_in_2)
        P3_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P3_out)
        P4_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode4/add')([P4_in_2, P4_td, P3_D])
        P4_out = KL.Activation(lambda x: tf.nn.swish(x))(P4_out)
        P4_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode4/op_after_combine9')(P4_out)

        P5_in_2 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/conv2d')(P5_in)
        P5_in_2 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/bn')(P5_in_2)
        # P5_in_2 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/bn')(P5_in_2)
        P4_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P4_out)
        P5_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode5/add')([P5_in_2, P5_td, P4_D])
        P5_out = KL.Activation(lambda x: tf.nn.swish(x))(P5_out)
        P5_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode5/op_after_combine10')(P5_out)

        P5_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P5_out)
        P6_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode6/add')([P6_in, P6_td, P5_D])
        P6_out = KL.Activation(lambda x: tf.nn.swish(x))(P6_out)
        P6_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode6/op_after_combine11')(P6_out)

        P6_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P6_out)
        P7_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode7/add')([P7_in, P6_D])
        P7_out = KL.Activation(lambda x: tf.nn.swish(x))(P7_out)
        P7_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode7/op_after_combine12')(P7_out)

    else:
        P3_in, P4_in, P5_in, P6_in, P7_in = features
        P7_U = KL.UpSampling2D()(P7_in)
        P6_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode0/add')([P6_in, P7_U])
        P6_td = KL.Activation(lambda x: tf.nn.swish(x))(P6_td)
        P6_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode0/op_after_combine5')(P6_td)
        P6_U = KL.UpSampling2D()(P6_td)
        P5_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode1/add')([P5_in, P6_U])
        P5_td = KL.Activation(lambda x: tf.nn.swish(x))(P5_td)
        P5_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode1/op_after_combine6')(P5_td)
        P5_U = KL.UpSampling2D()(P5_td)
        P4_td = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode2/add')([P4_in, P5_U])
        P4_td = KL.Activation(lambda x: tf.nn.swish(x))(P4_td)
        P4_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode2/op_after_combine7')(P4_td)
        P4_U = KL.UpSampling2D()(P4_td)
        P3_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode3/add')([P3_in, P4_U])
        P3_out = KL.Activation(lambda x: tf.nn.swish(x))(P3_out)
        P3_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode3/op_after_combine8')(P3_out)
        P3_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P3_out)
        P4_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode4/add')([P4_in, P4_td, P3_D])
        P4_out = KL.Activation(lambda x: tf.nn.swish(x))(P4_out)
        P4_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode4/op_after_combine9')(P4_out)

        P4_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P4_out)
        P5_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode5/add')([P5_in, P5_td, P4_D])
        P5_out = KL.Activation(lambda x: tf.nn.swish(x))(P5_out)
        P5_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode5/op_after_combine10')(P5_out)

        P5_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P5_out)
        P6_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode6/add')([P6_in, P6_td, P5_D])
        P6_out = KL.Activation(lambda x: tf.nn.swish(x))(P6_out)
        P6_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode6/op_after_combine11')(P6_out)

        P6_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P6_out)
        P7_out = wBiFPNAdd(name=f'fpn_cells/cell_{id}/fnode7/add')([P7_in, P6_D])
        P7_out = KL.Activation(lambda x: tf.nn.swish(x))(P7_out)
        P7_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode7/op_after_combine12')(P7_out)
    return P3_out, P4_td, P5_td, P6_td, P7_out

def build_BiFPN(features, num_channels, id, freeze_bn=False):
    if id == 0:
        _, _, C3, C4, C5 = features
        P3_in = C3
        P4_in = C4
        P5_in = C5
        P6_in = KL.Conv2D(num_channels, kernel_size=1, padding='same', name='resample_p6/conv2d')(C5)
        P6_in = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON, name='resample_p6/bn')(P6_in)
        # P6_in = BatchNormalization(freeze=freeze_bn, name='resample_p6/bn')(P6_in)
        P6_in = KL.MaxPooling2D(pool_size=3, strides=2, padding='same', name='resample_p6/maxpool')(P6_in)
        P7_in = KL.MaxPooling2D(pool_size=3, strides=2, padding='same', name='resample_p7/maxpool')(P6_in)
        P7_U = KL.UpSampling2D()(P7_in)
        P6_td = KL.Add(name=f'fpn_cells/cell_{id}/fnode0/add')([P6_in, P7_U])
        P6_td = KL.Activation(lambda x: tf.nn.swish(x))(P6_td)
        P6_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode0/op_after_combine5')(P6_td)
        P5_in_1 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/conv2d')(P5_in)
        P5_in_1 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/bn')(P5_in_1)
        # P5_in_1 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode1/resample_0_2_6/bn')(P5_in_1)
        P6_U = KL.UpSampling2D()(P6_td)
        P5_td = KL.Add(name=f'fpn_cells/cell_{id}/fnode1/add')([P5_in_1, P6_U])
        P5_td = KL.Activation(lambda x: tf.nn.swish(x))(P5_td)
        P5_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode1/op_after_combine6')(P5_td)
        P4_in_1 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/conv2d')(P4_in)
        P4_in_1 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/bn')(P4_in_1)
        # P4_in_1 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode2/resample_0_1_7/bn')(P4_in_1)
        P5_U = KL.UpSampling2D()(P5_td)
        P4_td = KL.Add(name=f'fpn_cells/cell_{id}/fnode2/add')([P4_in_1, P5_U])
        P4_td = KL.Activation(lambda x: tf.nn.swish(x))(P4_td)
        P4_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode2/op_after_combine7')(P4_td)
        P3_in = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                              name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/conv2d')(P3_in)
        P3_in = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                          name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/bn')(P3_in)
        # P3_in = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode3/resample_0_0_8/bn')(P3_in)
        P4_U = KL.UpSampling2D()(P4_td)
        P3_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode3/add')([P3_in, P4_U])
        P3_out = KL.Activation(lambda x: tf.nn.swish(x))(P3_out)
        P3_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode3/op_after_combine8')(P3_out)
        P4_in_2 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/conv2d')(P4_in)
        P4_in_2 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/bn')(P4_in_2)
        # P4_in_2 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode4/resample_0_1_9/bn')(P4_in_2)
        P3_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P3_out)
        P4_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode4/add')([P4_in_2, P4_td, P3_D])
        P4_out = KL.Activation(lambda x: tf.nn.swish(x))(P4_out)
        P4_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode4/op_after_combine9')(P4_out)

        P5_in_2 = KL.Conv2D(num_channels, kernel_size=1, padding='same',
                                name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/conv2d')(P5_in)
        P5_in_2 = KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON,
                                            name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/bn')(P5_in_2)
        # P5_in_2 = BatchNormalization(freeze=freeze_bn, name=f'fpn_cells/cell_{id}/fnode5/resample_0_2_10/bn')(P5_in_2)
        P4_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P4_out)
        P5_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode5/add')([P5_in_2, P5_td, P4_D])
        P5_out = KL.Activation(lambda x: tf.nn.swish(x))(P5_out)
        P5_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode5/op_after_combine10')(P5_out)

        P5_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P5_out)
        P6_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode6/add')([P6_in, P6_td, P5_D])
        P6_out = KL.Activation(lambda x: tf.nn.swish(x))(P6_out)
        P6_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode6/op_after_combine11')(P6_out)

        P6_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P6_out)
        P7_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode7/add')([P7_in, P6_D])
        P7_out = KL.Activation(lambda x: tf.nn.swish(x))(P7_out)
        P7_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode7/op_after_combine12')(P7_out)

    else:
        P3_in, P4_in, P5_in, P6_in, P7_in = features
        P7_U = KL.UpSampling2D()(P7_in)
        P6_td = KL.Add(name=f'fpn_cells/cell_{id}/fnode0/add')([P6_in, P7_U])
        P6_td = KL.Activation(lambda x: tf.nn.swish(x))(P6_td)
        P6_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode0/op_after_combine5')(P6_td)
        P6_U = KL.UpSampling2D()(P6_td)
        P5_td = KL.Add(name=f'fpn_cells/cell_{id}/fnode1/add')([P5_in, P6_U])
        P5_td = KL.Activation(lambda x: tf.nn.swish(x))(P5_td)
        P5_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode1/op_after_combine6')(P5_td)
        P5_U = KL.UpSampling2D()(P5_td)
        P4_td = KL.Add(name=f'fpn_cells/cell_{id}/fnode2/add')([P4_in, P5_U])
        P4_td = KL.Activation(lambda x: tf.nn.swish(x))(P4_td)
        P4_td = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                   name=f'fpn_cells/cell_{id}/fnode2/op_after_combine7')(P4_td)
        P4_U = KL.UpSampling2D()(P4_td)
        P3_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode3/add')([P3_in, P4_U])
        P3_out = KL.Activation(lambda x: tf.nn.swish(x))(P3_out)
        P3_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode3/op_after_combine8')(P3_out)
        P3_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P3_out)
        P4_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode4/add')([P4_in, P4_td, P3_D])
        P4_out = KL.Activation(lambda x: tf.nn.swish(x))(P4_out)
        P4_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode4/op_after_combine9')(P4_out)

        P4_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P4_out)
        P5_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode5/add')([P5_in, P5_td, P4_D])
        P5_out = KL.Activation(lambda x: tf.nn.swish(x))(P5_out)
        P5_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode5/op_after_combine10')(P5_out)

        P5_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P5_out)
        P6_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode6/add')([P6_in, P6_td, P5_D])
        P6_out = KL.Activation(lambda x: tf.nn.swish(x))(P6_out)
        P6_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode6/op_after_combine11')(P6_out)

        P6_D = KL.MaxPooling2D(pool_size=3, strides=2, padding='same')(P6_out)
        P7_out = KL.Add(name=f'fpn_cells/cell_{id}/fnode7/add')([P7_in, P6_D])
        P7_out = KL.Activation(lambda x: tf.nn.swish(x))(P7_out)
        P7_out = SeparableConvBlock(num_channels=num_channels, kernel_size=3, strides=1,
                                    name=f'fpn_cells/cell_{id}/fnode7/op_after_combine12')(P7_out)
    return P3_out, P4_td, P5_td, P6_td, P7_out

class BoxNet(KL.Layer):
    def __init__(self, width, depth, num_anchors=9, separable_conv=True, freeze_bn=False, detect_quadrangle=False, **kwargs):
        super(BoxNet, self).__init__(**kwargs)
        self.width = width
        self.depth = depth
        self.num_anchors = num_anchors
        self.separable_conv = separable_conv
        self.detect_quadrangle = detect_quadrangle
        num_values = 9 if detect_quadrangle else 4
        options = {
            'kernel_size': 3,
            'strides': 1,
            'padding': 'same',
            'bias_initializer': 'zeros',
        }
        if separable_conv:
            kernel_initializer = {
                'depthwise_initializer': initializers.VarianceScaling(),
                'pointwise_initializer': initializers.VarianceScaling(),
            }
            options.update(kernel_initializer)
            self.convs = [KL.SeparableConv2D(filters=width, name=f'{self.name}/box-{i}', **options) for i in
                          range(depth)]
            self.head = KL.SeparableConv2D(filters=num_anchors * num_values,
                                               name=f'{self.name}/box-predict', **options)
        else:
            kernel_initializer = {
                'kernel_initializer': initializers.RandomNormal(mean=0.0, stddev=0.01, seed=None)
            }
            options.update(kernel_initializer)
            self.convs = [KL.Conv2D(filters=width, name=f'{self.name}/box-{i}', **options) for i in range(depth)]
            self.head = KL.Conv2D(filters=num_anchors * num_values, name=f'{self.name}/box-predict', **options)
        self.bns = [
            [KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON, name=f'{self.name}/box-{i}-bn-{j}') for j in
             range(3, 8)]
            for i in range(depth)]
        # self.bns = [[BatchNormalization(freeze=freeze_bn, name=f'{self.name}/box-{i}-bn-{j}') for j in range(3, 8)]
        #             for i in range(depth)]
        self.relu = KL.Lambda(lambda x: tf.nn.swish(x))
        self.reshape = KL.Reshape((-1, num_values))
        self.level = 0

    def call(self, inputs, **kwargs):
        feature, level = inputs
        for i in range(self.depth):
            feature = self.convs[i](feature)
            feature = self.bns[i][self.level](feature)
            feature = self.relu(feature)
        outputs = self.head(feature)
        outputs = self.reshape(outputs)
        self.level += 1
        return outputs

class ClassNet(KL.Layer):
    def __init__(self, width, depth, num_classes=20, num_anchors=9, separable_conv=True, freeze_bn=False, **kwargs):
        super(ClassNet, self).__init__(**kwargs)
        self.width = width
        self.depth = depth
        self.num_classes = num_classes
        self.num_anchors = num_anchors
        self.separable_conv = separable_conv
        options = {
            'kernel_size': 3,
            'strides': 1,
            'padding': 'same',
        }
        if self.separable_conv:
            kernel_initializer = {
                'depthwise_initializer': initializers.VarianceScaling(),
                'pointwise_initializer': initializers.VarianceScaling(),
            }
            options.update(kernel_initializer)
            self.convs = [KL.SeparableConv2D(filters=width, bias_initializer='zeros', name=f'{self.name}/class-{i}',
                                                 **options)
                          for i in range(depth)]
            self.head = KL.SeparableConv2D(filters=num_classes * num_anchors,
                                               bias_initializer=PriorProbability(probability=0.01),
                                               name=f'{self.name}/class-predict', **options)
        else:
            kernel_initializer = {
                'kernel_initializer': initializers.RandomNormal(mean=0.0, stddev=0.01, seed=None)
            }
            options.update(kernel_initializer)
            self.convs = [KL.Conv2D(filters=width, bias_initializer='zeros', name=f'{self.name}/class-{i}',
                                        **options)
                          for i in range(depth)]
            self.head = KL.Conv2D(filters=num_classes * num_anchors,
                                      bias_initializer=PriorProbability(probability=0.01),
                                      name='class-predict', **options)
        self.bns = [
            [KL.BatchNormalization(momentum=MOMENTUM, epsilon=EPSILON, name=f'{self.name}/class-{i}-bn-{j}') for j
             in range(3, 8)]
            for i in range(depth)]
        # self.bns = [[BatchNormalization(freeze=freeze_bn, name=f'{self.name}/class-{i}-bn-{j}') for j in range(3, 8)]
        #             for i in range(depth)]
        self.relu = KL.Lambda(lambda x: tf.nn.swish(x))
        self.reshape = KL.Reshape((-1, num_classes))
        self.activation = KL.Activation('sigmoid')
        self.level = 0

    def call(self, inputs, **kwargs):
        feature, level = inputs
        for i in range(self.depth):
            feature = self.convs[i](feature)
            feature = self.bns[i][self.level](feature)
            feature = self.relu(feature)
        outputs = self.head(feature)
        outputs = self.reshape(outputs)
        outputs = self.activation(outputs)
        self.level += 1
        return outputs

def log2_graph(x):
    """Implementation of Log2. TF doesn't have a native implementation."""
    return tf.log(x) / tf.log(2.0)

class BatchNorm(KL.BatchNormalization):
    """Extends the Keras BatchNormalization class to allow a central place
    to make changes if needed.
    Batch normalization has a negative effect on training if batches are small
    so this layer is often frozen (via setting in Config class) and functions
    as linear layer.
    """
    def call(self, inputs, training=None):
        """
        Note about training values:
            None: Train BN layers. This is the normal mode
            False: Freeze BN layers. Good when batch size is small
            True: (don't use). Set layer in training mode even when making inferences
        """
        return super(self.__class__, self).call(inputs, training=training)

class ROIAlignLayer(KL.Layer):
    """Implements ROI Pooling on multiple levels of the feature pyramid.
    Params:
    - pool_shape: [pool_height, pool_width] of the output pooled regions. Usually [7, 7]
    Inputs:
    - boxes: [batch, num_boxes, (y1, x1, y2, x2)] in normalized
             coordinates. Possibly padded with zeros if not enough
             boxes to fill the array.
    - image_meta: [batch, (meta data)] Image details. See compose_image_meta()
    - feature_maps: List of feature maps from different levels of the pyramid.
                    Each is [batch, height, width, channels]
    Output:
    Pooled regions in the shape: [batch, num_boxes, pool_height, pool_width, channels].
    The width and height are those specific in the pool_shape in the layer
    constructor.
    """

    def __init__(self, pool_shape, **kwargs):
        super(ROIAlignLayer, self).__init__(**kwargs)
        self.pool_shape = tuple(pool_shape)

    def call(self, inputs):
        # Crop boxes [batch, num_boxes, (y1, x1, y2, x2)] in normalized coords
        boxes = inputs[0]

        # Feature Maps. List of feature maps from different level of the
        # feature pyramid. Each is [batch, height, width, channels]
        feature_maps = inputs[2]

        # Assign each ROI to a level in the pyramid based on the ROI area.
        y1, x1, y2, x2 = tf.split(boxes, 4, axis=2)
        h = y2 - y1
        w = x2 - x1
        # Use shape of first image. Images in a batch must have the same size.
        #image_shape = inputs[1]
        # Equation 1 in the Feature Pyramid Networks paper. Account for
        # the fact that our coordinates are normalized here.
        # e.g. a 224x224 ROI (in pixels) maps to P4
        #image_area = tf.cast(image_shape[0] * image_shape[1], tf.float32)
        image_area = tf.cast(inputs[1] * inputs[1], tf.float32)
        roi_level = log2_graph(tf.sqrt(h * w) / (224.0 / tf.sqrt(image_area)))
        roi_level = tf.minimum(5, tf.maximum(
            2, 4 + tf.cast(tf.round(roi_level), tf.int32)))
        roi_level = tf.squeeze(roi_level, 2)

        # Loop through levels and apply ROI pooling to each. P2 to P5.
        pooled = []
        box_to_level = []
        for i, level in enumerate(range(2, 6)):
            ix = tf.where(tf.equal(roi_level, level))
            level_boxes = tf.gather_nd(boxes, ix)

            # Box indices for crop_and_resize.
            box_indices = tf.cast(ix[:, 0], tf.int32)

            # Keep track of which box is mapped to which level
            box_to_level.append(ix)

            # Stop gradient propogation to ROI proposals
            level_boxes = tf.stop_gradient(level_boxes)
            box_indices = tf.stop_gradient(box_indices)
            print("box at level", level_boxes.shape)
            print(feature_maps)
            print("feature_map size", feature_maps[i])
            # Crop and Resize
            # From Mask R-CNN paper: "We sample four regular locations, so
            # that we can evaluate either max or average pooling. In fact,
            # interpolating only a single value at each bin center (without
            # pooling) is nearly as effective."
            #
            # Here we use the simplified approach of a single value per bin,
            # which is how it's done in tf.crop_and_resize()
            # Result: [batch * num_boxes, pool_height, pool_width, channels]
            pooled.append(tf.image.crop_and_resize(
                feature_maps[i], level_boxes, box_indices, self.pool_shape,
                method="bilinear"))

        # Pack pooled features into one tensor
        pooled = tf.concat(pooled, axis=0)

        # Pack box_to_level mapping into one array and add another
        # column representing the order of pooled boxes
        box_to_level = tf.concat(box_to_level, axis=0)
        box_range = tf.expand_dims(tf.range(tf.shape(box_to_level)[0]), 1)
        box_to_level = tf.concat([tf.cast(box_to_level, tf.int32), box_range],
                                 axis=1)

        # Rearrange pooled features to match the order of the original boxes
        # Sort box_to_level by batch then box index
        # TF doesn't have a way to sort by two columns, so merge them and sort.
        sorting_tensor = box_to_level[:, 0] * 100000 + box_to_level[:, 1]
        ix = tf.nn.top_k(sorting_tensor, k=tf.shape(
            box_to_level)[0]).indices[::-1]
        ix = tf.gather(box_to_level[:, 2], ix)
        pooled = tf.gather(pooled, ix)

        # Re-add the batch dimension
        shape = tf.concat([tf.shape(boxes)[:2], tf.shape(pooled)[1:]], axis=0)
        pooled = tf.reshape(pooled, shape)
        return pooled

    def compute_output_shape(self, input_shape):
        return input_shape[0][:2] + self.pool_shape + (input_shape[2][-1], )

def build_mask_target_graph(rois,
                            b_gt_boxes,
                            b_gt_masks,
                            gt_class_ids,
                            negative_overlap=0,
                            positive_overlap=0,
                            num_classes=13,
                            mask_shape=[28,28]):
    '''
    Inputs:
    proposals: [POST_NMS_ROIS_TRAINING, (y1, x1, y2, x2)] in normalized coordinates. Might
               be zero padded if there are not enough proposals.
    gt_class_ids: [MAX_GT_INSTANCES] int class IDs
    gt_boxes: [MAX_GT_INSTANCES, (y1, x1, y2, x2)] in normalized coordinates.
    gt_masks: [height, width, MAX_GT_INSTANCES] of boolean type.
    Returns: Target ROIs and corresponding class IDs, bounding box shifts,
    and masks.
    rois: [batch, TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)] in normalized coordinates
    class_ids: [batch, TRAIN_ROIS_PER_IMAGE]. Integer class IDs. Zero padded.
    masks: [batch, TRAIN_ROIS_PER_IMAGE, height, width]. Masks cropped to bbox
           boundaries and resized to neural network output size.
    Note: Returned arrays might be zero padded if not enough target ROIs.
    '''
    print("b_gt_boxes", b_gt_boxes.shape)
    rois, _ = trim_zeros_graph(rois, name="trim_rois")
    b_gt_boxes, non_zeros = trim_zeros_graph(b_gt_boxes, name='trim_gt_boxes')
    b_gt_masks = tf.gather(b_gt_masks, tf.where(non_zeros)[:, 0], axis=0,
                        name="trim_gt_masks")
    # skip handle coco crowd for now
    gt_class_ids = tf.boolean_mask(gt_class_ids, non_zeros,
                                   name="trim_gt_class_ids")
    overlaps = overlaps_graph(rois, b_gt_boxes)

    # Determine positive and negative ROIs
    roi_iou_max = tf.reduce_max(overlaps, axis=1)
    # 1. Positive ROIs are those with >= 0.5 IoU with a GT box
    positive_roi_bool = (roi_iou_max >= positive_overlap)
    positive_indices = tf.where(positive_roi_bool)[:, 0]
    # 2. Negative ROIs are those with < 0.5 with every GT box. Skip crowds.
    negative_indices = tf.where((roi_iou_max < negative_overlap))[:, 0]

    positive_count = int(TRAIN_ROIS_PER_IMAGE * 0.3)
    positive_indices = tf.random_shuffle(positive_indices)[:positive_count]
    positive_count = tf.shape(positive_indices)[0]
    # Negative ROIs. Add enough to maintain positive:negative ratio.
    #r = 1.0 / 0.3
    negative_count = int(TRAIN_ROIS_PER_IMAGE * 0.6)
    #negative_count = tf.cast(r * tf.cast(positive_count, tf.float32), tf.int32) - positive_count
    negative_indices = tf.random_shuffle(negative_indices)[:negative_count]

    positive_overlaps = tf.gather(overlaps, positive_indices)
    roi_gt_box_assignment = tf.cond(
        tf.greater(tf.shape(positive_overlaps)[1], 0),
        true_fn = lambda: tf.argmax(positive_overlaps, axis=1),
        false_fn = lambda: tf.cast(tf.constant([]),tf.int64)
    )

    # Permute masks to [N, height, width, 1]
    transposed_masks = tf.expand_dims(b_gt_masks, -1)
    # Pick the right mask for each ROI
    roi_masks = tf.gather(transposed_masks, roi_gt_box_assignment)
    roi_gt_class_ids = tf.gather(gt_class_ids, roi_gt_box_assignment)
    boxes = tf.gather(rois, positive_indices)
    box_ids = tf.range(0, tf.shape(roi_masks)[0])
    b_masks = tf.image.crop_and_resize(tf.cast(roi_masks, tf.float32), boxes,
                                    box_ids,
                                    mask_shape)
    b_masks = tf.round(tf.squeeze(b_masks, axis=-1))
    positive_rois = tf.gather(rois, positive_indices)
    negative_rois = tf.gather(rois, negative_indices)
    # Append negative ROIs and pad bbox deltas and masks that
    # are not used for negative ROIs with zeros.
    N = tf.shape(negative_rois)[0]
    NP = tf.shape(positive_rois)[0]
    P = tf.maximum(TRAIN_ROIS_PER_IMAGE - N - NP, 0)
    rois = tf.concat([positive_rois, negative_rois], axis=0)
    #print("before", rois.shape)
    rois = tf.pad(rois, [(0, P), (0, 0)])
    roi_gt_class_ids = tf.pad(roi_gt_class_ids, [(0, N + P)], mode='CONSTANT', constant_values=-1)
    b_masks = tf.pad(b_masks, [[0, N + P], (0, 0), (0, 0)])
    print("in function", N+P, rois.shape, b_masks.shape, roi_gt_class_ids.shape)
    return rois, roi_gt_class_ids, b_masks

class DetectionTargetLayer(KL.Layer):
    """Subsamples proposals and generates target box refinement, class_ids,
    and masks for each.
    Inputs:
    proposals: [batch, N, (y1, x1, y2, x2)] in normalized coordinates. Might
               be zero padded if there are not enough proposals.
    gt_class_ids: [batch, MAX_GT_INSTANCES] Integer class IDs.
    gt_boxes: [batch, MAX_GT_INSTANCES, (y1, x1, y2, x2)] in normalized
              coordinates.
    gt_masks: [batch, height, width, MAX_GT_INSTANCES] of boolean type
    Returns: Target ROIs and corresponding class IDs, bounding box shifts,
    and masks.
    rois: [batch, TRAIN_ROIS_PER_IMAGE, (y1, x1, y2, x2)] in normalized
          coordinates
    target_class_ids: [batch, TRAIN_ROIS_PER_IMAGE]. Integer class IDs.
    target_deltas: [batch, TRAIN_ROIS_PER_IMAGE, (dy, dx, log(dh), log(dw)]
    target_mask: [batch, TRAIN_ROIS_PER_IMAGE, height, width]
                 Masks cropped to bbox boundaries and resized to neural
                 network output size.
    Note: Returned arrays might be zero padded if not enough target ROIs.
    """

    def __init__(self, **kwargs):
        super(DetectionTargetLayer, self).__init__(**kwargs)
        #self.config = config

    def call(self, inputs):
        proposals = inputs[0]
        gt_class_ids = inputs[1]
        gt_boxes = inputs[2]
        gt_masks = inputs[3]

        print('shapes', proposals.shape, gt_class_ids.shape, gt_boxes.shape, gt_masks.shape)

        # Slice the batch and run a graph for each slice
        # TODO: Rename target_bbox to target_deltas for clarity
        names = ["rois", "target_class_ids", "target_mask"]
        outputs = batch_slice(
            [proposals, gt_class_ids, gt_boxes, gt_masks],
            lambda w, x, y, z: build_mask_target_graph(
                w, y, z, x),
            1, names=names)
        return outputs

    def compute_output_shape(self, input_shape):
        return [
            (None, TRAIN_ROIS_PER_IMAGE, 4),  # rois
            (None, TRAIN_ROIS_PER_IMAGE),  # class_ids
            (None, TRAIN_ROIS_PER_IMAGE, 28, 28) # masks
        ]

    def compute_mask(self, inputs, mask=None):
        return [None, None, None, None]

import cv2
def data_generator(dataset, shuffle=True, phi=0,
                   batch_size=1):
    '''
    '''
    b = 0  # batch item index
    image_index = -1
    image_ids = np.copy(dataset.image_ids)
    num_classes = dataset.num_classes()

    # Keras requires a generator to run indefinitely.
    while True:
        try:
            group=[]

            while b < batch_size:
                # Increment index to pick next image. Shuffle if at the start of an epoch.
                image_index = (image_index + 1) % len(image_ids)
                group.append(image_index)
                b+=1

            if shuffle and image_index == 0:
                np.random.shuffle(image_ids)

            images, outs, annotations_group = dataset.compute_inputs_targets(group)
            images = np.squeeze(images, axis=0)
            # Skip images that have no instances. This can happen in cases
            # where we train on a subset of classes and the image doesn't
            # have any of the classes we care about.
            #if not np.any(annotations_group['labels'] > 0):
            #    continue

            # first stage bbox Targets
            #outs = build_bbox_target(annotations_group, images, num_classes, image_size, detect_quadrangle=False)

            # mask targets [batch, max_instances, 28, 28] and second state bbox targets [batch, max_instaces, 4]
            masks_batch = np.zeros((batch_size, MAX_INSTANCES, images.shape[1], images.shape[2]), dtype=np.float32)
            boxes_batch = np.zeros((batch_size, MAX_INSTANCES, 4), dtype=np.float32)
            labels_batch = np.zeros((batch_size, MAX_INSTANCES), dtype=np.float32)
            for index, annotations in enumerate(annotations_group):
                # TODO: apply transformation to gt_masks
                gt_masks = annotations['masks']
                # for debug
                gt_boxes = annotations['bboxes']
                gt_boxes = norm_boxes(gt_boxes, (images.shape[1], images.shape[2]))
                '''
                for i,mask in enumerate(gt_masks):
                    cv2.imwrite('test_effects/post_mask_{}.jpg'.format(i), mask*255)
                    print("mask shape", mask.shape)'''
                gt_labels = annotations['labels']
                masks_batch[index, :gt_masks.shape[0], :, :] = gt_masks
                boxes_batch[index, :gt_boxes.shape[0], :] = gt_boxes
                labels_batch[index, :gt_labels.shape[0]] = gt_labels
            inputs = [images, outs[0], outs[1], masks_batch, boxes_batch, labels_batch]
            targets = []

            outputs = targets
            yield inputs, outputs
            # start a new batch
            b = 0

        except (GeneratorExit, KeyboardInterrupt):
            raise

def build_fpn_mask_graph(rois, feature_maps, image_size, num_classes,
                         pool_size, train_bn=True):
    """Builds the computation graph of the mask head of Feature Pyramid Network.
    rois: [batch, num_rois, (y1, x1, y2, x2)] Proposal boxes in normalized
          coordinates.
    feature_maps: List of feature maps from different layers of the pyramid,
                  [P1, P2, P3, P4, P5]. Each has a different resolution.
    image_meta: [batch, (meta data)] Image details. See compose_image_meta()
    pool_size: The width of the square feature map generated from ROI Pooling.
    num_classes: number of classes, which determines the depth of the results
    train_bn: Boolean. Train or freeze Batch Norm layers
    Returns: Masks [batch, num_rois, MASK_POOL_SIZE, MASK_POOL_SIZE, NUM_CLASSES]
    """
    # ROI Pooling
    # Shape: [batch, num_rois, MASK_POOL_SIZE, MASK_POOL_SIZE, channels]
    x = ROIAlignLayer([pool_size, pool_size],
                        name="roi_align_mask")((rois, image_size, feature_maps))
    # x [1, num_rois, 14, 14, 64]
    # Conv layers
    x = KL.TimeDistributed(KL.Conv2D(256, (3, 3), padding="same"),
                           name="mrcnn_mask_conv1")(x)
    x = KL.TimeDistributed(BatchNorm(),
                           name='mrcnn_mask_bn1')(x, training=train_bn)
    x = KL.Activation('relu')(x)

    x = KL.TimeDistributed(KL.Conv2D(256, (3, 3), padding="same"),
                           name="mrcnn_mask_conv2")(x)
    x = KL.TimeDistributed(BatchNorm(),
                           name='mrcnn_mask_bn2')(x, training=train_bn)
    x = KL.Activation('relu')(x)

    x = KL.TimeDistributed(KL.Conv2D(256, (3, 3), padding="same"),
                           name="mrcnn_mask_conv3")(x)
    x = KL.TimeDistributed(BatchNorm(),
                           name='mrcnn_mask_bn3')(x, training=train_bn)
    x = KL.Activation('relu')(x)

    x = KL.TimeDistributed(KL.Conv2D(256, (3, 3), padding="same"),
                           name="mrcnn_mask_conv4")(x)
    x = KL.TimeDistributed(BatchNorm(),
                           name='mrcnn_mask_bn4')(x, training=train_bn)
    x = KL.Activation('relu')(x)

    x = KL.TimeDistributed(KL.Conv2DTranspose(256, (2, 2), strides=2, activation="relu"),
                           name="mrcnn_mask_deconv")(x)
    x = KL.TimeDistributed(KL.Conv2D(num_classes, (1, 1), strides=1, activation="sigmoid"),
                           name='mask')(x)
    print("pool mask feature map size", x.shape)
    return x

def mrcnn_mask_loss_graph(target_masks, target_class_ids, pred_masks):
    """Mask binary cross-entropy loss for the masks head.
    target_masks: [batch, num_rois, height, width].
        A float32 tensor of values 0 or 1. Uses zero padding to fill array.
    target_class_ids: [batch, num_rois]. Integer class IDs. Zero padded.
    pred_masks: [batch, proposals, height, width, num_classes] float32 tensor
                with values from 0 to 1.
    """
    print("target_class_ids", target_class_ids.shape)
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids >= 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    # Compute binary cross entropy. If no positive ROIs, then return 0.
    # shape: [batch, roi, num_classes]
    loss = K.switch(tf.size(y_true) > 0,
                    K.binary_crossentropy(target=y_true, output=y_pred),
                    tf.constant(0.0))
    loss = K.mean(loss)
    return loss

class EfficientDetModel():
    def __init__(self, phi):
        assert phi in range(7)
        self.phi = phi
        self.input_size = image_sizes[phi]
        self.input_shape = (self.input_size, self.input_size, 3)
        self.w_bifpn = w_bifpns[phi]
        self.d_bifpn = d_bifpns[phi]
        self.w_head = self.w_bifpn
        self.d_head = d_heads[phi]
        self.backbone_cls = backbones[phi]
        self.model, self.p_model = self.efficientdet()

    def efficientdet(self, num_classes=13, num_anchors=9, weighted_bifpn=False, freeze_bn=False,
                    score_threshold=0.01, detect_quadrangle=False, anchor_parameters=None, separable_conv=True):
        image_input = KL.Input(self.input_shape, name="image_input")
        '''
        w_bifpn = w_bifpns[phi]
        d_bifpn = d_bifpns[phi]
        w_head = w_bifpn
        d_head = d_heads[phi]
        '''
        features = self.backbone_cls(input_tensor=image_input, freeze_bn=freeze_bn)
        if weighted_bifpn:
            fpn_features = features
            for i in range(self.d_bifpn):
                fpn_features = build_wBiFPN(fpn_features, self.w_bifpn, i, freeze_bn=freeze_bn)
        else:
            fpn_features = features
            for i in range(self.d_bifpn):
                fpn_features = build_BiFPN(fpn_features, self.w_bifpn, i, freeze_bn=freeze_bn)
        box_net = BoxNet(self.w_head, self.d_head, num_anchors=num_anchors, separable_conv=separable_conv, freeze_bn=freeze_bn,
                        detect_quadrangle=detect_quadrangle, name='box_net')
        class_net = ClassNet(self.w_head, self.d_head, num_classes=num_classes, num_anchors=num_anchors,
                            separable_conv=separable_conv, freeze_bn=freeze_bn, name='class_net')
        classification = [class_net([feature, i]) for i, feature in enumerate(fpn_features)]
        classification = KL.Concatenate(axis=1, name='classification')(classification)
        regression = [box_net([feature, i]) for i, feature in enumerate(fpn_features)]
        regression = KL.Concatenate(axis=1, name='regression')(regression)
        # regression has shpae [batch, num_detections, 4]
        print("classification shape", classification.shape)

        # apply predicted regression to anchors
        anchors = anchors_for_shape((self.input_size, self.input_size), anchor_params=anchor_parameters)
        anchors_input = np.expand_dims(anchors, axis=0)

        # build input layer for box detection targets
        box_target = KL.Input((anchors.shape[0], 5), name="box_target")
        class_target = KL.Input((anchors.shape[0], num_classes+1), name="class_target")

        boxes = RegressBoxes(name='boxes')([anchors_input, regression[..., :4]])
        boxes = ClipBoxes(name='clipped_boxes')([image_input, boxes])

        # filter detections (apply NMS / score threshold / select top-k)
        if detect_quadrangle:
            detections = FilterDetections(
                name='filtered_detections',
                score_threshold=score_threshold,
                detect_quadrangle=True
            )([boxes, classification, regression[..., 4:8], regression[..., 8]])
        else:
            detections = FilterDetections(
                name='filtered_detections',
                score_threshold=score_threshold
            )([boxes, classification])

        output_boxes = detections[0]
        norm_boxes = norm_boxes_graph(output_boxes, K.shape(image_input)[1:3])
        gt_masks = KL.Input((MAX_INSTANCES, self.input_shape[0], self.input_shape[1]), name="gt_masks")
        gt_boxes = KL.Input((MAX_INSTANCES, 4), name="gt_boxes")
        gt_labels = KL.Input((MAX_INSTANCES), name="gt_labels")

        rois, target_class_ids, target_mask = DetectionTargetLayer(name="proposal_targets")([
                    norm_boxes, gt_labels, gt_boxes, gt_masks])
        roi_masks = build_fpn_mask_graph(rois, fpn_features, self.input_size, 13, 14) #[batch, num_proposals,h,w,num_classes]
        #print("shape of roi_masks", roi_masks.shape)

        # build losses
        #output_rois = KL.Lambda(lambda x: x * 1, name="output_rois")(rois)

        # Losses
        regression_loss = KL.Lambda(lambda x: smooth_l1(*x), name="regression_loss")([box_target, regression])
        classification_loss = KL.Lambda(lambda x: focal(*x), name="classification_loss")([class_target, classification])
        mask_loss = KL.Lambda(lambda x: mrcnn_mask_loss_graph(*x), name="mrcnn_mask_loss")(
                                        [target_mask, target_class_ids, roi_masks])

        model = models.Model(inputs=[image_input, class_target, box_target, gt_masks, gt_boxes, gt_labels],
                             outputs=[classification, regression, roi_masks,
                                      regression_loss, classification_loss, mask_loss], name='efficientdet')
        prediction_model = models.Model(inputs=[image_input, gt_masks, gt_boxes, gt_labels],
                                        outputs=[detections, roi_masks], name='efficientdet_p')
        return model, prediction_model

    def compile(self,learning_rate=1e-3):
        # compile losses
        self.model._losses = []
        self.model._per_input_losses = {}
        loss_names = [
            "regression_loss", "classification_loss", "mrcnn_mask_loss"]
        for name in loss_names:
            layer = self.model.get_layer(name)
            if layer.output in self.model.losses:
                continue
            loss = (
                tf.reduce_mean(layer.output, keepdims=True))
            self.model.add_loss(loss)

        # Add L2 Regularization
        # Skip gamma and beta weights of batch normalization layers.
        '''
        reg_losses = [
            keras.regularizers.l2(self.config.WEIGHT_DECAY)(w) / tf.cast(tf.size(w), tf.float32)
            for w in self.keras_model.trainable_weights
            if 'gamma' not in w.name and 'beta' not in w.name]
        self.keras_model.add_loss(tf.add_n(reg_losses))'''

        # Compile
        self.model.compile(
            optimizer=Adam(lr=1e-3),
            loss=[None] * len(self.model.outputs))

        # Add metrics for losses
        for name in loss_names:
            if name in self.model.metrics_names:
                continue
            layer = self.model.get_layer(name)
            loss = (
                tf.reduce_mean(layer.output, keepdims=True))
            self.model.add_metric(loss, aggregation='mean', name=name)

    def create_callbacks(self, training_model, prediction_model, validation_generator, args):
        """
        Creates the callbacks to use during training.

        Args
            training_model: The model that is used for training.
            prediction_model: The model that should be used for validation.
            validation_generator: The generator for creating validation data.
            args: parseargs args object.

        Returns:
            A list of callbacks used for training.
        """
        callbacks = []

        tensorboard_callback = None

        if args.tensorboard_dir:
            if tf.version.VERSION > '2.0.0':
                file_writer = tf.summary.create_file_writer(args.tensorboard_dir)
                file_writer.set_as_default()
            tensorboard_callback = keras.callbacks.TensorBoard(
                log_dir=args.tensorboard_dir,
                histogram_freq=0,
                batch_size=args.batch_size,
                write_graph=True,
                write_grads=False,
                write_images=False,
                embeddings_freq=0,
                embeddings_layer_names=None,
                embeddings_metadata=None
            )
            callbacks.append(tensorboard_callback)

        if args.evaluation and validation_generator:
            if args.dataset_type == 'coco':
                from eval.coco import Evaluate
                # use prediction model for evaluation
                evaluation = Evaluate(validation_generator, prediction_model, tensorboard=tensorboard_callback)
            else:
                from eval.pascal import Evaluate
                evaluation = Evaluate(validation_generator, prediction_model, tensorboard=tensorboard_callback)
            callbacks.append(evaluation)

        # save the model
        if args.snapshots:
            # ensure directory created first; otherwise h5py will error after epoch.
            makedirs(args.snapshot_path)
            checkpoint = keras.callbacks.ModelCheckpoint(
                os.path.join(
                    args.snapshot_path,
                    f'{args.dataset_type}_{{epoch:02d}}_{{loss:.4f}}_{{val_loss:.4f}}.h5' if args.compute_val_loss
                    else f'{args.dataset_type}_{{epoch:02d}}_{{loss:.4f}}.h5'
                ),
                verbose=1,
                save_weights_only=True,
                # save_best_only=True,
                # monitor="mAP",
                # mode='max'
            )
            callbacks.append(checkpoint)

        # callbacks.append(keras.callbacks.ReduceLROnPlateau(
        #     monitor='loss',
        #     factor=0.1,
        #     patience=2,
        #     verbose=1,
        #     mode='auto',
        #     min_delta=0.0001,
        #     cooldown=0,
        #     min_lr=0
        # ))

        return callbacks

    def train(self, train_dataset, learning_rate, epochs,
              args, val_dataset=None, custom_callbacks=None):
        # Data generators
        train_generator = data_generator(train_dataset, shuffle=True,
                                         phi=self.phi, batch_size=1)
        '''
        val_generator = data_generator(val_dataset, shuffle=True,
                                       batch_size=1, phi=self.phi)'''

        # Create log_dir if it does not exist
        if not os.path.exists(args.log_dir):
            os.makedirs(args.log_dir)

        # Callbacks
        callbacks = [
            keras.callbacks.TensorBoard(log_dir=args.log_dir,
                                        histogram_freq=0, write_graph=True, write_images=False),
            keras.callbacks.ModelCheckpoint(os.path.join(
                                            args.checkpoint_path,
                                            f'{args.dataset_type}_{{epoch:02d}}_{{loss:.4f}}_{{val_loss:.4f}}.h5' if args.compute_val_loss
                                            else f'{args.dataset_type}_{{epoch:02d}}_{{loss:.4f}}.h5'),
                                            verbose=0, save_weights_only=True)
        ]

        # Add custom callbacks to the list
        if custom_callbacks:
            callbacks += custom_callbacks

        # Train
        log("\nStarting at epoch {}. LR={}\n".format(self.epoch, learning_rate))
        log("Checkpoint Path: {}".format(self.checkpoint_path))
        self.compile()
        #self.set_trainable(layers)
        #self.compile(learning_rate, self.config.LEARNING_MOMENTUM)

        # Work-around for Windows: Keras fails on Windows when using
        # multiprocessing workers. See discussion here:
        # https://github.com/matterport/Mask_RCNN/issues/13#issuecomment-353124009
        if os.name is 'nt':
            workers = 0
        else:
            workers = multiprocessing.cpu_count()

        self.model.fit_generator(
            train_generator,
            initial_epoch=self.epoch,
            epochs=epochs,
            steps_per_epoch=1000,
            callbacks=callbacks,
            validation_data=None,
            #validation_steps=None,
            max_queue_size=100,
            workers=workers,
            use_multiprocessing=True,
        )
        self.epoch = max(self.epoch, epochs)

if __name__ == '__main__':
    models = EfficientDetModel(0)
    models.compile()
    train_model = models.model
    from generators.coco_ins import CocoDataset
    from augmentor.mask_misc import MiscEffectMask
    misc_effect = MiscEffectMask()
    train_dataset = CocoDataset('data/sample_val', ['train', 'val'], misc_effect=misc_effect)
    train_generator = data_generator(train_dataset, shuffle=True,
                                         phi=0, batch_size=2)
    checkpoint = keras.callbacks.ModelCheckpoint(
            os.path.join(
                'checkpoints',
                'deepfashion.h5'
            ),
            verbose=1,
            save_weights_only=True,
            # save_best_only=True,
            # monitor="mAP",
            # mode='max'
        )
    train_model.fit_generator(train_generator, epochs=1, steps_per_epoch=1000, callbacks=[checkpoint])


# GAIL, DAC, PURL
# Off-PAC sasaki 2019
from stable_baselines3.common.off_policy_algorithm import OffPolicyAlgorithm
from stable_baselines3.common.on_policy_algorithm import OnPolicyAlgorithm
from stable_baselines3.common.buffers import DictReplayBuffer, ReplayBuffer, RolloutBuffer, DictRolloutBuffer
from stable_baselines3.common.preprocessing import get_action_dim, get_flattened_obs_dim
from stable_baselines3.common.type_aliases import (
    DictReplayBufferSamples,
    DictRolloutBufferSamples,
    ReplayBufferSamples,
    RolloutBufferSamples,
)
from sbil.demo.utils import get_demo_buffer, state_action, all_state_action
from sbil.utils import set_method, MLP, set_restore, save_torch, get_policy, get_features_extractor, load_torch
import sbil

from typing import Any, Dict, List, Optional, Tuple, Type, Union
from pathlib import Path
from types import MethodType
from operator import attrgetter
from functools import partial

import torch as th
from torch.nn.functional import logsigmoid, mse_loss, binary_cross_entropy_with_logits
import numpy as np
import gym

def discriminator_step(discriminator, buffer_sample, demo_buffer, learner, state_only, λ=1, center=1, η=None):
    batch_size = buffer_sample.actions.shape[0]
    extract_features = get_policy(learner).extract_features
    demo_sample = demo_buffer.sample(batch_size, learner._vec_normalize_env)
    # get state-action pairs
    sa = state_action(buffer_sample.observations, buffer_sample.actions, learner, state_only)
    demo_sa = state_action(demo_sample.observations, demo_sample.actions, learner, state_only)

    # regularization: zero-centered gradient penalty: https://openreview.net/pdf?id=ByxPYjC5KQ
    α = th.rand(batch_size,1).to(learner.device)
    inputs = th.autograd.Variable(sa*α + (1-α)*demo_sa, requires_grad=True).to(learner.device)
    outputs = discriminator(inputs)
    gradients = th.autograd.grad(
        outputs=outputs,
        inputs=inputs,
        grad_outputs=th.ones(outputs.size()).to(learner.device),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    penalty = λ * th.square(th.linalg.norm(gradients)-center)
    # discrimination loss
    if η is None: # positive negative
        input = discriminator(th.cat((demo_sa, sa), dim=0)).squeeze()
        target = th.cat((th.ones(batch_size), th.zeros(batch_size)), dim=0)
        weight = None
    else: # Positive unlabeled
        input = discriminator(th.cat((demo_sa, sa, demo_sa), dim=0)).squeeze()
        target = th.cat((th.ones(batch_size), th.zeros(batch_size*2)), dim=0)
        weight = th.cat((th.ones(batch_size)*η, th.ones(batch_size), -th.ones(batch_size)*η), dim=0).to(learner.device)
    loss = binary_cross_entropy_with_logits(input=input, target=target.to(learner.device), weight=weight)

    loss = loss.mean() + penalty.mean()
    discriminator.optimizer.zero_grad()
    loss.backward()
    discriminator.optimizer.step()

    return loss.item()

def train_off(self, gradient_steps, super_, discriminator, demo_buffer, state_only: bool = False, η=None, *args, **kwargs) -> None:
    for gradient_step in range(gradient_steps):
        replay_data = self.replay_buffer.sample(self.batch_size, self._vec_normalize_env)
        discriminator_step(discriminator, replay_data, demo_buffer, self, state_only, η=η)
    super_(gradient_steps, *args, **kwargs)

def train_on(self, super_, discriminator, demo_buffer, state_only: bool = False, η=None, *args, **kwargs) -> None:
    for epoch in range(getattr(self, "n_epochs", 1)):
        for rollout_data in self.rollout_buffer.get(batch_size=getattr(self, "batch_size", None)):
            discriminator_step(discriminator, rollout_data, demo_buffer, self, state_only, η=η)
    super_(*args, **kwargs)

def sample(self, batch_size, env, *args, super_, discriminator, learner, state_only: bool = False, **kwargs) -> Union[DictReplayBufferSamples, ReplayBufferSamples]:
    replay_data = super_(batch_size=batch_size, env=env, *args, **kwargs)
    sa = state_action(replay_data.observations, replay_data.actions, learner, state_only)
    with th.no_grad():
        replay_data.rewards[:] = discriminator(sa)

    return replay_data

def compute_returns_and_advantage(self, super_, discriminator, learner, state_only: bool = False, *args, **kwargs) -> None:
    sa = all_state_action(self, learner, state_only)
    with th.no_grad():
        out = discriminator(sa)
    self.rewards[:] = out.view(self.buffer_size, self.n_envs).cpu().numpy() # unstack
    super_(*args, **kwargs)

def adversarial(
    learner: Union[OnPolicyAlgorithm, OffPolicyAlgorithm],
    demo_buffer: Union[DictReplayBuffer, ReplayBuffer, str, Path],
    state_only: bool = False,
    net_arch: List[int] = [64, 64],
    η: Optional[float] = None,
    load: Optional[str] = None,
) -> Union[OnPolicyAlgorithm, OffPolicyAlgorithm]:
    """
    Adversarial imitation learning decorator: https://arxiv.org/pdf/1809.02925.pdf
    The discriminator is trained in learner.train
    The reward is modified in:
    - learner.replay_buffer._get_samples if OffPolicyAlgorithm
    - learner.rollout_buffer.compute_returns_and_advantage if OnPolicyAlgorithm

    :param learner: Stable baselines learner object
    :param demo_buffer: Demonstration replay buffer
    :param state_only: Use state only
        default is the concatenation of the state-action pair
    :param η: Positive class prior, usually set to 0.5 in papers. PURL is disabled by default.
        https://arxiv.org/abs/1911.00459
    :param load: Zip file containing state_dict of the discriminator to resume a training.
        It should be the same as the path used with .load(path)
    :return learner: Decorated learner
    """
    demo_buffer = get_demo_buffer(demo_buffer, learner)
    is_off = isinstance(learner, OffPolicyAlgorithm)
    is_on = not is_off

    buffer_name = "replay_buffer" if is_off else "rollout_buffer"
    set_restore(getattr(learner, buffer_name))
    set_restore(learner, lambda self: getattr(self, buffer_name).restore())

    # discriminator
    if load is None:
        learner.state_only = state_only
        learner.net_arch = net_arch
    p = get_features_extractor(learner)
    input_dim = p.features_extractor.features_dim
    if not learner.state_only:
        input_dim += get_flattened_obs_dim(learner.action_space)
    discriminator = MLP(
        input_dim=input_dim,
        output_dim=1,
        net_arch=learner.net_arch,
        spectral_norm=True,
    ).to(learner.device)
    modules = {'discriminator': discriminator, 'discriminator_optimizer': discriminator.optimizer}
    if load is not None:
        load_torch(load, modules=modules)
    set_method(learner, old="save", new=partial(save_torch, modules=modules))

    # modify train() to train the discriminator
    set_method(
        learner,
        old="train",
        new=train_off if is_off else train_on,
        discriminator=discriminator,
        demo_buffer=demo_buffer,
        state_only=state_only,
        η=η,
    )

    # overwrite set_method with additional arguments
    set_method_ = partial(
        set_method,
        discriminator=discriminator,
        state_only=state_only,
        learner=learner,
    )

    # modify the buffer to change the reward
    if is_off:
        set_method_(learner.replay_buffer, old="sample", new=sample)
    elif is_on:
        set_method_(learner.rollout_buffer, old="compute_returns_and_advantage", new=compute_returns_and_advantage)
    else:
        raise NotImplementedError()
    return learner


"""
Routines for handling ASPCAP bitmasks
"""

import numpy as np

class BitMask():
    '''
    Base class for bitmasks, define common methods

    At a minimum, a BitMask will have a set of name, level, descrip

    BitMask provides 3 methods:
        getname(val,level=level) : returns name(s) of all set bits
                                  (optionally, of requested level)
        getval(name) : returns value of bit with input name
        badval()     : returns value of all bits that are marked bad (level=1)
        warnval()    : returns value of all bits that are marked warn (level=2)
    '''
    def getname(self,val,level=0,strip=True):
        '''
        Given input value, returns names of all set bits, optionally of a given level
        '''
        strflag=''
        for ibit,name in enumerate(self.name) :
            if ( val & 2**ibit ) > 0 and ( level == 0 or self.level == level ) :
              strflag = strflag + name +','
        if strip : return strflag.strip(',')
        else : return strflag

    def getval(self,name) :
        """
        Get the numerical bit value of a given character name(s)
        """
        if type(name) is str :
            name = [name]
        bitval = 0
        for n in name :
            try:
                j=self.name.index(n.strip())
                bitval|=2**j
            except :
                print('WARNING: undefined name: ',n)
        return bitval

    def badval(self) :
        """
        Return bitmask value of all bits that indicate BAD in input bitmask
        """
        val=0
        for i,level in enumerate(self.level) :
            if level == 1 :
                val=val | 2**i
        return val

    def warnval(self) :
        """
        Return bitmask value of all bits that indicate BAD in input bitmask
        """
        val=0
        for i,level in enumerate(self.level) :
            if level == 2 :
                val=val | 2**i
        return val

class StarBitMask(BitMask):
    '''
    BitMask class for APOGEE star bitmask (APOGEE_STARFLAG)
    '''

    name=(['BAD_PIXELS','COMMISSIONING','BRIGHT_NEIGHBOR','VERY_BRIGHT_NEIGHBOR','LOW_SNR','','','',
          '','PERSIST_HIGH','PERSIST_MED','PERSIST_LOW','PERSIST_JUMP_POS','PERSIST_JUMP_NEG','','',
          'SUSPECT_RV_COMBINATION','SUSPECT_BROAD_LINES','BAD_RV_COMBINATION','RV_REJECT','RV_SUSPECT','MULTIPLE_SUSPECT','','',
          '','','','','','','',''])
    level=([1,0,0,1,0,0,0,0,
             0,0,0,0,0,0,0,0,
             0,0,1,0,0,0,0,0,
             0,0,0,0,0,0,0,0])

    descrip=([
     'Spectrum has many bad pixels (>20%):  BAD',
     'Commissioning data (MJD<55761), non-standard configuration, poor LSF: WARN',
     'Star has neighbor more than 10 times brighter: WARN',
     'Star has neighbor more than 100 times brighter: BAD',
     'Spectrum has low S/N (S/N<5)',
     '',
     '',
     '',
     '',
     'Spectrum has significant number (>20%) of pixels in high persistence region: WARN',
     'Spectrum has significant number (>20%) of pixels in medium persistence region: WARN',
     'Spectrum has significant number (>20%) of pixels in low persistence region: WARN',
     'Spectrum show obvious positive jump in blue chip: WARN',
     'Spectrum show obvious negative jump in blue chip: WARN',
     '',
     '',
     'RVs from synthetic template differ significantly (~2 km/s) from those from combined template: WARN',
     'Cross-correlation peak with template significantly broader than autocorrelation of template: WARN',
     'RVs from synthetic template differ very significatly (~10 km/s) from those from combined template: BAD',
     'Rejected visit because cross-correlation RV differs significantly from least squares RV',
     'Suspect visit (but used!) because cross-correlation RV differs slightly from least squares RV',
     'Suspect multiple components from Gaussian decomposition of cross-correlation',
     '',
     '',
     '',
     '',
     '',
     '',
     '',
     '',
     '',
     ''
     ])
    def persist(self) :
        '''
        Returns bitwise OR of all persistence bits
        '''
        return self.getval(['PERSIST_HIGH','PERSIST_MED','PERSIST_LOW',
                             'PERSIST_JUMP_POS','PERSIST_JUMP_NEG'] )

class AspcapBitMask(BitMask):
    '''
    BitMask class for APOGEE ASPCAP bitmask (APOGEE_ASPCAPFLAG)
    '''
    name=(['TEFF_WARN','LOGG_WARN','VMICRO_WARN','M_H_WARN','ALPHA_M_WARN','C_M_WARN','N_M_WARN','STAR_WARN',
          'CHI2_WARN','COLORTE_WARN','ROTATION_WARN','SN_WARN','SPEC_HOLE_WARN','ATMOS_HOLE_WARN','VSINI_WARN','',
          'TEFF_BAD','LOGG_BAD','VMICRO_BAD','M_H_BAD','ALPHA_M_BAD','C_M_BAD','N_M_BAD','STAR_BAD',
          'CHI2_BAD','COLORTE_BAD','ROTATION_BAD','SN_BAD','SPEC_HOLE_BAD','ATMOS_HOLE_BAD','VSINI_BAD','NO_ASPCAP_RESULT'])
    level=([2,2,0,0,0,0,0,2,
              2,2,2,2,2,2,0,0,
              1,1,0,0,0,0,0,1,
              1,1,1,1,1,2,0,1])

    descrip=([
     'WARNING on effective temperature (see PARAMFLAG[0] for details) ',
     'WARNING on log g (see PARAMFLAG[1] for details) ',
     'WARNING on vmicro (see PARAMFLAG[2] for details) ',
     'WARNING on metals (see PARAMFLAG[3] for details) ',
     'WARNING on [alpha/M] (see PARAMFLAG[4] for details) ',
     'WARNING on [C/M] (see PARAMFLAG[5] for details) ',
     'WARNING on [N/M] (see PARAMFLAG[6] for details) ',
     'WARNING overall for star: set if any of TEFF, LOGG, CHI2, COLORTE, ROTATION, SN warn are set ',
     'high chi^2 (> 2*median at ASPCAP temperature (WARN)',
     'effective temperature more than 500K from photometric temperature for dereddened color (WARN)',
     'Spectrum has broad lines, with possible bad effects: FWHM of cross-correlation of spectrum with best RV template relative to auto-correltion of template > 1.5 (WARN)',
     'S/N<70 (WARN)',
     'Grid point within 2 grid steps of hole-filled synthesis ',
     'Grid point within 2 grid steps of hole-filled atmosphere ',
     ' ',
     ' ',
     'BAD effective temperature (see PARAMFLAG[0] for details) ',
     'BAD log g (see PARAMFLAG[1] for details) ',
     'BAD vmicro (see PARAMFLAG[2] for details) ',
     'BAD metals (see PARAMFLAG[3] for details) ',
     'BAD [alpha/M] (see PARAMFLAG[4] for details) ',
     'BAD [C/M] (see PARAMFLAG[5] for details) ',
     'BAD [N/M] (see PARAMFLAG[6] for details) ',
     'BAD overall for star: set if any of TEFF, LOGG, CHI2, COLORTE, ROTATION, SN error are set, or any GRIDEDGE_BAD ',
     'high chi^2 (> 5*median at ASPCAP temperature (BAD)',
     'effective temperature more than 1000K from photometric temperature for dereddened color (BAD)',
     'Spectrum has broad lines, with possible bad effects: FWHM of cross-correlation of spectrum with best RV template relative to auto-correltion of template > 2 (BAD)',
     'S/N<50 (BAD)',
     'Grid point within 1 grid steps of hole-filled synthesis ',
     'Grid point within 1 grid steps of hole-filled atmosphere ',
     ' ',
     ' '
     ])

class ParamBitMask(BitMask):
    '''
    BitMask class for APOGEE ASPCAP bitmask (APOGEE_ASPCAPFLAG)
    '''
    name =['GRIDEDGE_BAD','CALRANGE_BAD','OTHER_BAD','FERRE_BAD','PARAM_MISMATCH_BAD','','','',
           'GRIDEDGE_WARN','CALRANGE_WARN','OTHER_WARN','FERRE_WARN','PARAM_MISMATCH_WARN','OPTICAL_WARN','ERR_WARN','FAINT_WARN',
           'PARAM_FIXED','RV_WARN','','','','','','',
           'LOGG_CAL_RC','LOGG_CAL_RGB','LOGG_CAL_MS','LOGG_CAL_RGB_MS','','','','']

    level=[1,1,1,1,1,0,0,0,
             0,0,0,0,0,0,0,0,
             0,0,0,0,0,0,0,0,
             0,0,0,0,0,0,0,0]

    descrip=[
     'Parameter within 1/8 grid spacing of grid edge ',
     'Parameter outside valid range of calibration determination ',
     'Other error condition ',
     'Failed solution in FERRE ',
     'Elemental abundance from window differs significantly from parameter abundance ',
     ' ',
     ' ',
     ' ',
     'Parameter within 1/2 grid spacing of grid edge ',
     'Parameter in possibly unreliable range of calibration determination ',
     'Other warning condition ',
     'FERRE warning (not implemented?) ',
     'Elemental abundance from window differs from parameter abundance ',
     'Comparison with optical abundances suggests problem ',
     'Large expected uncertainty or upper limit based on location in parameter space (Teff, [M/H], S/N) ',
     'Warning based on faint star/RV combination ',
     'Parameter set at fixed value, not fit',
     'RV puts important line off of chip ',
     ' ',
     ' ',
     ' ',
     ' ',
     ' ',
     ' ',
     'Use RC gravity calibration ',
     'Use RGB gravity calibration ',
     'Use MS gravity calibration ',
     'Use RBG/MS transition gravity calibration ',
     ' ',
     ' ',
     ' ',
     ' '
     ]

class PixelBitMask(BitMask) :
    '''
    BitMask class for APOGEE pixel bitmask (APOGEE_PIXMASK)
    '''
    name=(['BADPIX','CRPIX','SATPIX','UNFIXABLE','BADDARK','BADFLAT','BADERR','NOSKY',
          'LITTROW_GHOST','PERSIST_HIGH','PERSIST_MED','PERSIST_LOW','SIG_SKYLINE','SIG_TELLURIC','NOT_ENOUGH_PSF',''])

    level=([1,1,1,1,1,1,1,1,
             0,0,0,0,0,0,1,0])

    maskcontrib=([0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,
                 0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.0])

    descrip=([
     'Pixel marked as BAD in bad pixel mask or from strong persistence jump',
     'Pixel marked as cosmic ray in ap3d',
     'Pixel marked as saturated in ap3d',
     'Pixel marked as unfixable in ap3d',
     'Pixel marked as bad as determined from dark frame',
     'Pixel marked as bad as determined from flat frame',
     'Pixel set to have very high error (not used)',
     'No sky available for this pixel from sky fibers',
     'Pixel falls in Littrow ghost, may be affected',
     'Pixel falls in high persistence region, may be affected',
     'Pixel falls in medium persistence region, may be affected',
     'Pixel falls in low persistence region, may be affected',
     'Pixel falls near sky line that has significant flux compared with object',
     'Pixel falls near telluric line that has significant absorption',
     'Less than 50 percent PSF in good pixels',
     ''
    ])

class Apogee2Target1(BitMask) :
    '''
    BitMask class for APOGEE2_TARGET1
    '''

    name = ([ 'APOGEE2_ONEBIN_GT_0_5','APOGEE2_TWOBIN_0_5_TO_0_8','APOGEE2_TWOBIN_GT_0_8','APOGEE2_IRAC_DERED',
              'APOGEE2_WISE_DERED','APOGEE2_SFD_DERED','APOGEE2_NO_DERED','APOGEE2_WASH_GIANT',
              'APOGEE2_WASH_DWARF','APOGEE2_SCI_CLUSTER','APOGEE2_','APOGEE2_SHORT',
              'APOGEE2_MEDIUM','APOGEE2_LONG','APOGEE2_NORMAL_SAMPLE','APOGEE2_MANGA_LED',
              'APOGEE2_ONEBIN_GT_0_3','APOGEE2_WASH_NOCLASS','APOGEE2_STREAM_MEMBER','APOGEE2_STREAM_CANDIDATE',
              'APOGEE2_DSPH_MEMBER','APOGEE2_DSPH_CANDIDATE','APOGEE2_MAGCLOUD_MEMBER','APOGEE2_MAGCLOUD_CANDIDATE',
              'APOGEE2_RRLYR','APOGEE2_BULGE_RC','APOGEE2_SGR_DSPH','APOGEE2_APOKASC_GIANT',
              'APOGEE2_APOKASC_DWARF','APOGEE2_FAINT_EXTRA','APOGEE2_APOKASC',''])

class Apogee2Target2(BitMask) :
    '''
    BitMask class for APOGEE2_TARGET2
    '''
    name = ([ 'LIGHT_TRAP','APOGEE2_CCLOUD_AS4','APOGEE2_STANDARD_STAR','APOGEE2_RV_STANDARD',
              'APOGEE2_SKY','APOGEE2_EXTERNAL_CALIB','APOGEE2_INTERNAL_CALIB','APOGEE2_DISK_SUBSTRUCTURE_MEMBER',
              'APOGEE2_DISK_SUBSTRUCTURE_CANDIDATE','APOGEE2_TELLURIC','APOGEE2_CALIB_CLUSTER','APOGEE2_K2_PLANETHOST',
              'APOGEE2_TIDAL_BINARY','APOGEE2_LITERATURE_CALIB','APOGEE2_GES_OVERLAP','APOGEE2_ARGOS_OVERLAP',
              'APOGEE2_GAIA_OVERLAP','APOGEE2_GALAH_OVERLAP','APOGEE2_RAVE_OVERLAP','APOGEE2_COMMIS_SOUTH_SPEC',
              'APOGEE2_HALO_MEMBER','APOGEE2_HALO_CANDIDATE','APOGEE2_1M_TARGET','APOGEE2_MOD_BRIGHT_LIMIT',
              'APOGEE2_CIS','APOGEE2_CNTAC','APOGEE2_EXTERNAL','APOGEE2_CVZ_AS4_OBAF',
              'APOGEE2_CVZ_AS4_GI','APOGEE2_CVZ_AS4_CTL','APOGEE2_CVZ_AS4_GIANT',''])

class Apogee2Target3(BitMask) :
    '''
    BitMask class for APOGEE2_TARGET3
    '''
    name = ([ 'APOGEE2_KOI','APOGEE2_EB','APOGEE2_KOI_CONTROL','APOGEE2_MDWARF',
              'APOGEE2_SUBSTELLAR_COMPANIONS','APOGEE2_YOUNG_CLUSTER','APOGEE2_K2','APOGEE2_OBJECT',
              'APOGEE2_ANCILLARY','APOGEE2_MASSIVE_STAR','APOGEE2_QSO','APOGEE2_CEPHEID',
              'APOGEE2_LOW_AV_WINDOWS','APOGEE2_BE_STAR','APOGEE2_YOUNG_MOVING_GROUP','APOGEE2_NGC6791',
              'APOGEE2_LABEL_STAR','APOGEE2_FAINT_KEPLER_GIANTS','APOGEE2_W345','APOGEE2_MASSIVE_EVOLVED',
              'APOGEE2_REDDENING_TARGETS','APOGEE2_KEPLER_MDWARF_KOI','APOGEE2_AGB','APOGEE2_M33',
              'APOGEE2_ULTRACOOL','APOGEE2_DISTANT_SEGUE_GIANTS','APOGEE2_CEPHEID_MAPPING','APOGEE2_SA57',
              'APOGEE2_K2_MDWARF','APOGEE2_RVVAR','APOGEE2_M31','APOGEE2_'])

class ApogeeTarget1(BitMask) :
    '''
    BitMask class for APOGEE_TARGET1
    '''
    name = ([ 'APOGEE_FAINT','APOGEE_MEDIUM','APOGEE_BRIGHT','APOGEE_IRAC_DERED',
              'APOGEE_WISE_DERED','APOGEE_SFD_DERED','APOGEE_NO_DERED','APOGEE_WASH_GIANT',
              'APOGEE_WASH_DWARF','APOGEE_SCI_CLUSTER','APOGEE_EXTENDED','APOGEE_SHORT',
              'APOGEE_INTERMEDIATE','APOGEE_LONG','APOGEE_DO_NOT_OBSERVE','APOGEE_SERENDIPITOUS',
              'APOGEE_FIRST_LIGHT','APOGEE_ANCILLARY','APOGEE_M31_CLUSTER','APOGEE_MDWARF',
              'APOGEE_HIRES','APOGEE_OLD_STAR','APOGEE_DISK_RED_GIANT','APOGEE_KEPLER_EB',
              'APOGEE_GC_PAL1','APOGEE_MASSIVE_STAR','APOGEE_SGR_DSPH','APOGEE_KEPLER_SEISMO',
              'APOGEE_KEPLER_HOST','APOGEE_FAINT_EXTRA','APOGEE_SEGUE_OVERLAP',''])

class ApogeeTarget2(BitMask) :
    '''
    BitMask class for APOGEE_TARGET2
    '''
    name = ([ 'LIGHT_TRAP','APOGEE_FLUX_STANDARD','APOGEE_STANDARD_STAR','APOGEE_RV_STANDARD',
              'APOGEE_SKY','APOGEE_SKY_BAD','APOGEE_GUIDE_STAR','APOGEE_BUNDLE_HOLE',
              'APOGEE_TELLURIC_BAD','APOGEE_TELLURIC','APOGEE_CALIB_CLUSTER','APOGEE_BULGE_GIANT',
              'APOGEE_BULGE_SUPER_GIANT','APOGEE_EMBEDDEDCLUSTER_STAR','APOGEE_LONGBAR','APOGEE_EMISSION_STAR',
              'APOGEE_KEPLER_COOLDWARF','APOGEE_MIRCLUSTER_STAR','APOGEE_RV_MONITOR_IC348','APOGEE_RV_MONITOR_KEPLER',
              'APOGEE_GES_CALIBRATE','APOGEE_BULGE_RV_VERIFY','APOGEE_1MTARGET','APOGEE_',
              'APOGEE_','APOGEE_','APOGEE_','APOGEE_',
              'APOGEE_','APOGEE_','APOGEE_',''])

class ApogeeTarget3(BitMask) :
    '''
    BitMask class for APOGEE_TARGET3
    '''
    name = ([ '','','','','','','','',
              '','','','','','','','',
              '','','','','','','','',
              '','','','','','','',''])

def targflags(targ1,targ2,targ3,survey='apogee2') :

    if 'apogee2' in survey :
        mask1=Apogee2Target1()
        mask2=Apogee2Target2()
        mask3=Apogee2Target3()
        return ','.join([mask1.getname(targ1),mask2.getname(targ2),mask3.getname(targ3)]).strip(',')
    else :
        mask1=ApogeeTarget1()
        mask2=ApogeeTarget2()
        return ','.join([mask1.getname(targ1),mask2.getname(targ2)]).strip(',')


import socket
import numpy as np
import cv2 as cv

recevier = socket.socket(family=socket.AF_INET, type=socket.SOCK_DGRAM)
# socket 통신준비
recevier.bind(('192.168.16.24', 7778))
# 바인딩

# 480 * 640 * 3 / 20 = 46080
perlength = int( (480 * 640 * 3) / 20) # packet size
reallength = perlength + 1

array = list()
while True:
    message, address= recevier.recvfrom(reallength)
    # 통신 사이즈 설정

    str = message[1:46081]
    num_array = b''

    array[message[0]] =  str

    if message[0] == 19:
        for i in range(20):
            num_array += array[i]
        frame = np.fromstring(num_array, dtype=np.uint8)
        frame = frame.reshape(480,640,3)
        cv.imshow('video receiver',frame)

    pass


import argparse
from csv import DictReader
import glob
from multiprocessing import Pool, cpu_count
from os.path import splitext, basename
from pathlib import Path
import pickle
import numpy as np
from typing import Union
from utils import hparams as hp
from utils.display import *
from utils.dsp import *
from utils.files import get_files
from utils.paths import Paths
from utils.display import simple_table
from functools import partial

from 臺灣言語工具.解析整理.拆文分析器 import 拆文分析器
from 臺灣言語工具.語音合成 import 台灣話口語講法

def valid_n_workers(num):
    n = int(num)
    if n < 1:
        raise argparse.ArgumentTypeError(
            '%r must be an integer greater than 0' % num)
    return n

def wav_to_mel_quant(path: Path):
    y = load_wav(path)
    # === 1. wav normalization === #
    peak = np.abs(y).max()
    if hp.peak_norm or peak > 1.0:
        y /= peak

    # === 2. mel === #
    mel = melspectrogram(y)

    # === 3. quant === #
    if hp.voc_mode == 'RAW':
        if hp.mu_law:
            quant = encode_mu_law(y, mu=2**hp.bits)
        else:
            quant = float_2_label(y, bits=hp.bits)

    elif hp.voc_mode == 'MOL':
        quant = float_2_label(y, bits=16)

    return mel.astype(np.float32), quant.astype(np.int64)

class Preprocessor():
    def __init__(self, args):
        self.args = args
        self.config()

        # === args === #
        self.audio_extension = args.extension
        self.dataset_path = args.data_dir
        self.n_workers = max(1, args.num_workers)
        self.paths = Paths(args.save_dir, hp.voc_model_id, hp.tts_model_id)

    def config(self):
        hp.configure(self.args.hp_file)

    # === Dataset === #
    def suisiann(self, path: Union[str, Path], wav_files):
        csv_file = get_files(path, extension='.csv')

        assert len(csv_file) == 1

        u_tihleh = set()
        for sootsai in wav_files:
            u_tihleh.add(basename(sootsai))
        text_dict = {}

        with open(csv_file[0], encoding='utf-8') as f:
            for tsua in DictReader(f):
                mia = basename(tsua['音檔'])
                if mia in u_tihleh:
                    imtong = splitext(mia)[0]
                    hj = tsua['漢字']
                    lmj = tsua['羅馬字']
                    text = 台灣話口語講法(
                        拆文分析器.建立句物件(hj, lmj)
                    )
                    text_dict[imtong] = text

        return text_dict

    def process_wav(self, path: Path):
        wav_id = path.stem
        m, x = wav_to_mel_quant(path)
        np.save(self.paths.mel / f'{wav_id}.npy', m, allow_pickle=False)
        np.save(self.paths.quant / f'{wav_id}.npy', x, allow_pickle=False)
        return wav_id, m.shape[-1]

    def run(self):
        '''
            1. Load .wav files
            2. Process suisiann csv file, save as text_dict.pkl
            3. Multiprocess .wav files, save correspondings in ./mel and ./quant
            4. Save wav_id & mel length in ./dataset.pkl
        '''
        # === wav files === #
        wav_files = get_files(
            self.dataset_path, extension=self.audio_extension)
        print(
            f'\n{len(wav_files)} {self.audio_extension[1:]} files found in "{self.dataset_path}"\n')

        if len(wav_files) == 0:
            print(
                f'no wav file found in {self.dataset_path}, please check specify --data_dir correctly!.\n')

        # === process === #
        else:
            if not hp.ignore_tts:
                text_dict = self.suisiann(self.dataset_path, wav_files)
                with open(self.paths.data / 'text_dict.pkl', 'wb') as f:
                    pickle.dump(text_dict, f)

            simple_table([
                ('Sample Rate', hp.sample_rate),
                ('Bit Depth',   hp.bits),
                ('Mu Law',      hp.mu_law),
                ('Hop Length',  hp.hop_length),
                ('CPU Usage', f'{self.n_workers}/{cpu_count()}')
            ])

            pool = Pool(processes=self.n_workers)

            dataset = []

            #func = partial(process_wav, self.paths)
            for i, (item_id, length) in enumerate(pool.imap_unordered(self.process_wav, wav_files), 1):
                dataset += [(item_id, length)]
                bar = progbar(i, len(wav_files))
                message = f'{bar} {i}/{len(wav_files)} '
                stream(message)

            with open(self.paths.data / 'dataset.pkl', 'wb') as f:
                pickle.dump(dataset, f)

            print(
                '\n\nCompleted. Ready to run "python train_tacotron2.py" or "python train_wavernn.py". \n')

def main(args):
    P = Preprocessor(args)
    P.run()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Preprocessing for WaveRNN and Tacotron2')

    parser.add_argument("data_dir", type=str)
    parser.add_argument("save_dir", type=str)

    parser.add_argument(
        '--extension', '-e', metavar='EXT', default='.wav',
        help='File extension to search for in dataset folder'
    )

    parser.add_argument(
        '--num_workers', '-w', metavar='N', type=valid_n_workers,
        default=cpu_count() - 1,
        help='The number of worker threads to use for preprocessing'
    )

    parser.add_argument(
        '--hp_file', metavar='FILE', default='hparams.py',
        help='The file to use for the hyperparameters'
    )
    args = parser.parse_args()
    main(args)


# Copyright (C) 2020 THL A29 Limited, a Tencent company.
# All rights reserved.
# Licensed under the BSD 3-Clause License (the "License"); you may
# not use this file except in compliance with the License. You may
# obtain a copy of the License at
# https://opensource.org/licenses/BSD-3-Clause
# 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.
# See the AUTHORS file for names of contributors.

import torch
import torch.utils.dlpack as dlpack
from torch import Tensor, device, dtype, nn
from typing import Union
import numpy as np
from typing import Callable, Dict, List, Optional, Tuple

try:
    # `turbo_transformers_cxxd` is the name on debug mode
    import turbo_transformers.turbo_transformers_cxxd as cxx
except ImportError:
    import turbo_transformers.turbo_transformers_cxx as cxx
from .return_type import convert_returns_as_type, ReturnType

__all__ = [
    'try_convert', 'convert2tt_tensor', 'to_param_dict_convert_tt',
    'to_param_dict', 'create_empty_if_none', 'AnyTensor'
]

def convert2tt_tensor(t):
    return cxx.Tensor.from_dlpack(dlpack.to_dlpack(t))

def try_convert(t):
    if isinstance(t, torch.Tensor):
        return convert2tt_tensor(t)
    elif isinstance(t, np.ndarray):
        return convert2tt_tensor(torch.from_numpy(t))
    else:
        return t

def to_param_dict_convert_tt(torch_module: torch.nn.Module):
    return {
        k: convert2tt_tensor(v)
        for k, v in torch_module.named_parameters()
    }

def to_param_dict(torch_module: torch.nn.Module):
    return {k: v for k, v in torch_module.named_parameters()}

def create_empty_if_none(output):
    return output if output is not None else cxx.Tensor.create_empty()

def get_head_mask(head_mask,
                  num_hidden_layers: int,
                  is_attention_chunked: bool = False):
    """
    # Prepare head mask if needed
    # 1.0 in head_mask indicate we keep the head
    attention_probs has shape bsz x n_heads x N x N
    Arguments:
        head_mask: torch.Tensor or None: has shape [num_heads] or [num_hidden_layers x num_heads]
        num_hidden_layers: int
    Returns:
            Tensor of shape shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
            or list with [None] for each layer
    """
    if head_mask is not None:
        head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
        if is_attention_chunked is True:
            head_mask = head_mask.unsqueeze(-1)
    else:
        head_mask = [None] * num_hidden_layers

    return head_mask

def get_extended_attention_mask(attention_mask: Tensor, input_shape: Tuple,
                                device: device) -> Tensor:
    """Makes broadcastable attention mask and causal mask so that future and maked tokens are ignored.
    Arguments:
        attention_mask: torch.Tensor with 1 indicating tokens to ATTEND to
        input_shape: tuple, shape of input_ids
        device: torch.Device, usually self.device
    Returns:
        torch.Tensor with dtype of attention_mask.dtype
    """
    # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
    # ourselves in which case we just need to make it broadcastable to all heads.
    if attention_mask.dim() == 3:
        extended_attention_mask = attention_mask[:, None, :, :]
    elif attention_mask.dim() == 2:
        # Provided a padding mask of dimensions [batch_size, seq_length]
        # - if the model is a decoder, apply a causal mask in addition to the padding mask
        # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
        extended_attention_mask = attention_mask[:, None, None, :]
    else:
        raise ValueError(
            "Wrong shape for input_ids (shape {}) or attention_mask (shape {})"
            .format(input_shape, attention_mask.shape))

    # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
    # masked positions, this operation will create a tensor which is 0.0 for
    # positions we want to attend and -10000.0 for masked positions.
    # Since we are adding it to the raw scores before the softmax, this is
    # effectively the same as removing these entirely.
    extended_attention_mask = extended_attention_mask.to(
        dtype=torch.float32)  # fp16 compatibility
    extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
    return extended_attention_mask

AnyTensor = Union[cxx.Tensor, torch.Tensor]