from keras import backend as K
from keras.models import load_model
from keras.preprocessing import image
from keras.optimizers import Adam
from imageio import imread
import numpy as np
import cv2
import copy
from matplotlib import pyplot as plt

from models.keras_ssd300 import ssd_300
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization

from ssd_encoder_decoder.ssd_output_decoder import decode_detections, decode_detections_fast

from data_generator.object_detection_2d_data_generator import DataGenerator
from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels
from data_generator.object_detection_2d_geometric_ops import Resize
from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms

img_height = 300
img_width = 300
confidence_threshold = 0.5
MATCH_OVERLAP_THRESHOLD = 0.3

K.clear_session() # Clear previous models from memory.

model = ssd_300(image_size=(img_height, img_width, 3),
                n_classes=20,
                mode='inference',
                l2_regularization=0.0005,
                scales=[0.1, 0.2, 0.37, 0.54, 0.71, 0.88, 1.05], # The scales for MS COCO are [0.07, 0.15, 0.33, 0.51, 0.69, 0.87, 1.05]
                aspect_ratios_per_layer=[[1.0, 2.0, 0.5],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                         [1.0, 2.0, 0.5],
                                         [1.0, 2.0, 0.5]],
                two_boxes_for_ar1=True,
                steps=[8, 16, 32, 64, 100, 300],
                offsets=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
                clip_boxes=False,
                variances=[0.1, 0.1, 0.2, 0.2],
                normalize_coords=True,
                subtract_mean=[123, 117, 104],
                swap_channels=[2, 1, 0],
                confidence_thresh=0.5,
                iou_threshold=0.45,
                top_k=200,
                nms_max_output_size=400)

# 2: Load the trained weights into the model.

# TODO: Set the path of the trained weights.
weights_path = './models/VGG_VOC0712_SSD_300x300_iter_120000.h5'

model.load_weights(weights_path, by_name=True)

# 3: Compile the model so that Keras won't complain the next time you load it.

adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)

ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0)

model.compile(optimizer=adam, loss=ssd_loss.compute_loss)

orig_images = [] # Store the images here.
input_images = [] # Store resized versions of the images here.

# We'll only load one image in this example.
img_path = 'examples/fish_bike.jpg'

orig_images.append(imread(img_path))
img = image.load_img(img_path, target_size=(img_height, img_width))
img = image.img_to_array(img)
input_images.append(img)
input_images = np.array(input_images)
print(input_images[0][1][1])
y_pred = model.predict(input_images)
confidence_threshold = 0.5
y_pred_thresh = [y_pred[k][y_pred[k,:,1] > confidence_threshold] for k in range(y_pred.shape[0])]
np.set_printoptions(precision=2, suppress=True, linewidth=90)
print("Predicted boxes:\n")
print('   class   conf xmin   ymin   xmax   ymax')
print(y_pred_thresh[0])

frame=cv2.imread(img_path)
frame=cv2.resize(frame, (img_width, img_height))
newframe = (frame[...,::-1].astype(np.float32))
newframe=newframe[np.newaxis, :]
y_pred = model.predict(newframe)
y_pred_thresh = [y_pred[k][y_pred[k,:,1] > confidence_threshold] for k in range(y_pred.shape[0])]
print("Predicted boxes:\n")
print('   class   conf xmin   ymin   xmax   ymax')
print(y_pred_thresh[0])

class BBox:
    def __init__(self,x,y,w,h,score,label):
        self.x=int(x)
        self.y=int(y)
        self.w=int(w)
        self.h=int(h)
        self.score=score
        self.label=label

    def size(self):
        return self.w * self.h

def boxOverlap(b1, b2):
	xmin = b1.x if (b1.x < b2.x ) else b2.x
	xmax = b1.x+b1.w if (b1.x+b1.w > b2.x+b2.w) else b2.x+b2.w
	ymin = b1.y if (b1.y < b2.y ) else b2.y
	ymax = b1.y+b1.h if (b1.y+b1.h > b2.y+b2.h) else b2.y+b2.h

	wid1 = int(b1.w)
	hei1 = int(b1.h)
	wid2 = int(b2.w)
	hei2 = int(b2.h)

	wid = int(xmax - xmin)
	hei = int(ymax - ymin)

	if (wid >= wid1 + wid2):
		return -1
	if (hei >= hei1 + hei2):
		return -1
	areaOR = float(np.sqrt(wid1 * hei1) + np.sqrt(wid2 * hei2))
	areaAnd = float(np.sqrt((wid1 + wid2 - wid) * (hei1 + hei2 - hei)))
	return areaAnd / (areaOR - areaAnd)

def creatTrackerList(bbox_list_track):
	trackercount=len(bbox_list_track)
	trackerlist=[]
	for i in range(0,trackercount):
		tracker=cv2.TrackerKCF_create()
		trackerlist.append(tracker)
	# Initialize tracker with frame and bounding box
	for (tracker,bbox) in zip(trackerlist,bbox_list_track):
		ok = tracker.init(frame,(bbox.x,bbox.y,bbox.w,bbox.h))
	return trackerlist

def detecting(frame):
	ori_width=frame.shape[1]
	ori_height=frame.shape[0]
	frame=cv2.resize(frame, (img_width, img_height))
	newframe = (frame[...,::-1].astype(np.float32))
	newframe=newframe[np.newaxis, :]
	y_pred = model.predict(newframe)
	y_pred_thresh = [y_pred[k][y_pred[k,:,1] > confidence_threshold] for k in range(y_pred.shape[0])]
	bbox_list_detect=[]
	for box in y_pred_thresh[0]:
	    # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
	    xmin = box[2] * ori_width/ img_width
	    ymin = box[3] * ori_height/ img_height
	    xmax = box[4] * ori_width/ img_width
	    ymax = box[5] * ori_height/ img_height
	    color = colors[int(box[0])]
	    bbox = BBox(xmin,ymin,xmax-xmin+1,ymax-ymin+1,box[1],box[0])
	    bbox_list_detect.append(bbox)
	return bbox_list_detect

def tracking(frame,trackerlist,bbox_list_track):
	for (tracker,bbox) in zip(trackerlist,bbox_list_track):
		# Update tracker
		ok, newbbox = tracker.update(frame)
		if not ok:
			break
		bbox.x = newbbox[0]
		bbox.y = newbbox[1]
		bbox.w = newbbox[2]
		bbox.h = newbbox[3]
	return bbox_list_track

def drawFrame(frame,bbox_list):
	colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
	for bbox in bbox_list:
		color = colors[int(bbox.label)]
		color = [colorname * 255 for colorname in color]
		display_txt = '%s: %.2f'%(classes[int(bbox.label)], bbox.score)
		cv2.rectangle(frame,(int(bbox.x),int(bbox.y)),(int(bbox.x)+int(bbox.w),int(bbox.y)+int(bbox.h)),color,2)
		cv2.putText(frame,display_txt,(int(bbox.x),int(bbox.y)), cv2.FONT_HERSHEY_SIMPLEX, 1,color,1,cv2.LINE_AA)
	return frame

def merging(bbox_list_detect,bbox_list_track):
	bbox_list_merge = []
	matchD2T = dict.fromkeys(range(len(bbox_list_detect)), -1)
	matchT2D = dict.fromkeys(range(len(bbox_list_track)), -1)
	bestT2D = dict.fromkeys(range(len(bbox_list_track)), -1)
	# Calculate IOU and cheak match pairs
	for i in range(0,len(bbox_list_track)):
		for j in range(0,len(bbox_list_detect)):
			d = boxOverlap(bbox_list_track[i],bbox_list_detect[j])
			if((d > MATCH_OVERLAP_THRESHOLD) and(bbox_list_track[i].label==bbox_list_detect[j].label)):
				matchD2T[j]=i;
				if((matchT2D[i] == -1) or (d > bestT2D[i])):
					matchT2D[i] = j;
					bestT2D[i] = d;
	# Update old vehicles from best overlap
	for i in range(0,len(bbox_list_track)):
		bbox = bbox_list_track[i];
		if (matchT2D[i] != -1):
			bbox = bbox_list_detect[matchT2D[i]];
		bbox_list_merge.append(bbox)
	# Create new object
	for i in range(0,len(bbox_list_detect)):
		if (matchD2T[i] == -1):
			bbox = bbox_list_detect[i];
			bbox_list_merge.append(bbox)
	return bbox_list_merge

colors = plt.cm.hsv(np.linspace(0, 1, 21)).tolist()
classes = ['background',
           'aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat',
           'chair', 'cow', 'diningtable', 'dog',
           'horse', 'motorbike', 'person', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor']

video = cv2.VideoCapture("examples/videos/4.mp4")
if not video.isOpened():
	print ("Could not open video")
	sys.exit()
fps = video.get(cv2.CAP_PROP_FPS)
print ("Frames per second using video.get(cv2.CAP_PROP_FPS) : {0}".format(fps))
video_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
video_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out = cv2.VideoWriter('examples/videos/4_out.mp4',fourcc, fps, (video_width,video_height))

framecount=0
while(True):
	print ('Framecount'+str(framecount))
	ok, frame = video.read()
	if not ok:
		break
	trackerlist = []
	if(framecount % 4 == 0):
		bbox_list_detect = detecting(frame)
		frame = drawFrame(frame,bbox_list_detect)
		out.write(frame)
		bbox_list_track = copy.deepcopy(bbox_list_detect)
		trackerlist = creatTrackerList(bbox_list_track)
	else:
		bbox_list_track = tracking(frame,trackerlist,bbox_list_track)
		bbox_list_detect = detecting(frame)
		bbox_list_track = merging(bbox_list_detect,bbox_list_track)
		trackerlist = creatTrackerList(bbox_list_track)
		frame = drawFrame(frame,bbox_list_track)
		out.write(frame)
	framecount+=1

print ("done")
out.release()
cv2.destroyAllWindows()

# for box in y_pred_thresh[0]:
#     # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
#     xmin = box[2] * orig_images[0].shape[1] / img_width
#     ymin = box[3] * orig_images[0].shape[0] / img_height
#     xmax = box[4] * orig_images[0].shape[1] / img_width
#     ymax = box[5] * orig_images[0].shape[0] / img_height
#     color = colors[int(box[0])]
#     label = '{}: {:.2f}'.format(classes[int(box[0])], box[1])
#     current_axis.add_patch(plt.Rectangle((xmin, ymin), xmax-xmin, ymax-ymin, color=color, fill=False, linewidth=2))
#     current_axis.text(xmin, ymin, label, size='x-large', color='white', bbox={'facecolor':color, 'alpha':1.0})


import tensorflow as tf
model = tf.keras.models.load_model('dogsVscats_tensorboard.model')
import matplotlib.pyplot as plt
import cv2
import matplotlib.image as mpimg
import numpy as np
CATOGRTIES = ['Dog','Cat']
def prepare(filepath):
	IMG_SIZE=100
	img_array = cv2.imread(filepath,cv2.IMREAD_GRAYSCALE)
	new_array = cv2.resize(img_array,(IMG_SIZE,IMG_SIZE))
	img=mpimg.imread(filepath)
	imgplot = plt.imshow(img)
	return new_array.reshape(-1,IMG_SIZE, IMG_SIZE, 1)

prediction = model.predict([prepare("dog_google.JPG")])
plt.title("Neural Network prediction is {}".format(CATOGRTIES[int(prediction[0][0])]))
plt.show()
print(CATOGRTIES[int(prediction[0][0])])


"""Fake eye contact."""
import argparse
import cv2
import dlib
import fcntl
import math
import numpy as np
import lib.v4l2 as v4l2
from threading import Thread
from time import perf_counter

MODEL_PATH = 'shape_predictor_68_face_landmarks.dat'

RIGHT_EYE_RANGE = (36, 42)
LEFT_EYE_RANGE = (42, 48)

EYE_SCALE = 1.5
EYE_BLUR = 21

DEFAULT_SCALE = 320

SHOW_FPS_EVEVRY = 10

def centers2(points):
    """Return eye centers and averall center."""
    centers = [p.mean(axis=0) for p in points]
    center = np.array(centers).mean(axis=0)
    return center, centers

def scale_box2d(box, scale):
    """Scale box."""
    h, w = box[1]
    return (box[0], (h * scale, w * scale), box[2])

def to_polar(p1, p2):
    """Get a length and angel in degrees."""
    dx = p2[0] - p1[0]
    dy = p2[1] - p1[1]

    return math.hypot(dx, dy), math.degrees(math.atan2(dy, dx))

class EyeContact:
    """EyeContact."""

    def __init__(self, img, scale=DEFAULT_SCALE):
        """Init the eyes replacer."""
        self.img = img
        self.mask = np.zeros(img.shape[:2], dtype=np.uint8)
        self.scale = scale
        # dlib
        self.detector = dlib.get_frontal_face_detector()
        self.predictor = dlib.shape_predictor(MODEL_PATH)

        elipses = []
        eyes = self.detect_eyes(img)
        for eye in eyes:
            ellipse = cv2.fitEllipse(eye)
            mask_ellipse = scale_box2d(ellipse, EYE_SCALE)
            elipses.append(mask_ellipse)
            cv2.ellipse(self.mask, mask_ellipse, 255, -1)

        self.mask = cv2.GaussianBlur(self.mask, (EYE_BLUR, EYE_BLUR), 0)

        # crop everything
        points = list(
            cv2.boxPoints(elipses[0])) + list(cv2.boxPoints(elipses[1]))
        (left, top, self.w, self.h) = cv2.boundingRect(np.array(points))
        left, top = left - EYE_BLUR, top - EYE_BLUR
        self.w, self.h = self.w + EYE_BLUR * 2, self.h + EYE_BLUR * 2
        self.img = self.img[top:top + self.h, left: left + self.w]
        self.mask = self.mask[top:top + self.h, left: left + self.w]

        # calculate the rest of params
        self.center, centers = centers2(eyes)

        self.length, self.angle = to_polar(centers[0], centers[1])
        self.center[0] -= left
        self.center[1] -= top

    def detect_eyes(self, img):
        """Detect eyes on image."""
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        h, w = gray.shape
        scale = w // self.scale
        gray = cv2.resize(gray, (self.scale, h // scale))

        rects = self.detector(gray, 1)
        # we assume only one face
        if len(rects) != 1:
            return None

        shape = self.predictor(gray, rects[0])
        shape = [(shape.part(i).x, shape.part(i).y) for i in range(0, 68)]
        shape = np.array(shape) * scale

        return [
            shape[RIGHT_EYE_RANGE[0]:RIGHT_EYE_RANGE[1]],
            shape[LEFT_EYE_RANGE[0]:LEFT_EYE_RANGE[1]]
        ]

    def open_eyes(self, img):
        """Replace eyes with opened."""
        eyes = self.detect_eyes(img)
        if eyes is None:
            return img

        center, centers = centers2(eyes)
        length, angle = to_polar(centers[0], centers[1])

        w, h = self.w, self.h
        m = cv2.getRotationMatrix2D(
            tuple(self.center),
            self.angle - angle,
            length / self.length)
        patch = cv2.warpAffine(self.img, m, (w, h))
        patch_mask = cv2.warpAffine(self.mask, m, (w, h))
        patch_mask = cv2.merge(3 * [patch_mask])

        x = int(round(center[0] - self.center[0]))
        y = int(round(center[1] - self.center[1]))
        patched = patch.astype(float) * patch_mask / 255
        x1, y1 = x + w, y + h
        if x < 0 or y < 0 or y1 > img.shape[1] or x1 > img.shape[0]:
            return img

        patched += img[y:y1, x:x1].astype(float) * (255 - patch_mask) / 255
        opn = img.copy()
        opn[y:y1, x:x1] = patched.astype(np.uint8)

        return opn

class CapStream:
    """Launch the thread to read cam async."""

    def __init__(self, cap):
        """Read first frame."""
        self.cap = cap
        (self.ret, self.frame) = self.cap.read()
        self.stopped = False
        self.count = 1

    def start(self):
        """Start the thread."""
        Thread(target=self.update, args=()).start()
        return self

    def update(self):
        """Loop infinitely until the thread is stopped."""
        while True:
            if self.stopped:
                return
            (ret, frame) = self.cap.read()
            if ret is not True:
                continue
            self.frame = frame
            self.count += 1

    def read(self, count):
        """Return the frame if there are new."""
        if count == self.count:
            return count, None

        return self.count, self.frame

    def stop(self):
        """Stop."""
        self.stopped = True

def _init_pair(source, dest):
    # Grab the webcam feed and get the dimensions of a frame
    cap_source = cv2.VideoCapture(source)
    width = int(cap_source.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap_source.get(cv2.CAP_PROP_FRAME_HEIGHT))
    channels = 3
    print('Source size {}x{}'.format(width, height))

    # Set up the formatting of our loopback device - boilerplate
    format = v4l2.v4l2_format()
    format.type = v4l2.V4L2_BUF_TYPE_VIDEO_OUTPUT

    dest_cap = open(dest, 'wb')
    if (fcntl.ioctl(dest_cap, v4l2.VIDIOC_G_FMT, format) == -1):
        print("Unable to get video format data.")
        return -1, None, None

    format.fmt.pix.field = v4l2.V4L2_FIELD_NONE
    format.fmt.pix.pixelformat = v4l2.V4L2_PIX_FMT_YUV420
    format.fmt.pix.width = width
    format.fmt.pix.height = height
    # format.fmt.pix.bytesperline = width * channels
    format.fmt.pix.sizeimage = width * height * channels // 2

    result = fcntl.ioctl(dest_cap, v4l2.VIDIOC_S_FMT, format)

    return result, cap_source, dest_cap

def _main_loop(eye_contact, stream, dest):
    count = 0
    while True:
        count, frame = stream.read(count)
        if frame is None:
            continue
        s = perf_counter()
        frame = eye_contact.open_eyes(frame)
        e = perf_counter()
        fps = round(1. / (e - s), 1)
        if count % SHOW_FPS_EVEVRY == 0:
            print('Eye opener FPS {}'.format(fps), end='\r')  # noqa
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2YUV_I420)
        dest.write(frame)

if __name__ == '__main__':
    desc = """Eye contact

    Sample usage:
        python3 eye_contact.py open.jpg /dev/video0 /dev/video1
    """

    p = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description=desc)

    p.add_argument('open', type=str, help='open eye frame')
    p.add_argument('source', type=str, help='video device ex. /dev/video0')
    p.add_argument('dest', type=str, help='video device destination')
    p.add_argument(
        '--scale', type=int, default=DEFAULT_SCALE,
        help='scale image on face detection')
    args = p.parse_args()

    r, src, dest = _init_pair(args.source, args.dest)
    print('Started ({} == 0). Press ctrl+c to exit.'.format(r))

    open_img = cv2.imread(args.open)
    eye_contact = EyeContact(open_img, args.scale)

    stream = CapStream(src).start()
    try:
        _main_loop(eye_contact, stream, dest)
    except KeyboardInterrupt:
        print('Exiting')

    stream.stop()
    src.release()
    dest.close()


# -*- coding:utf-8 -*-

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

import os
import sys
import tensorflow as tf
import time
import cv2
import numpy as np
import json
import math
from tqdm import tqdm
import argparse
from multiprocessing import Queue, Process
sys.path.append("../")

from data.io.image_preprocess import short_side_resize_for_inference_data
from libs.networks import build_whole_network
from help_utils import tools
from libs.label_name_dict.label_dict import *

def worker(gpu_id, images, det_net, eval_data, result_queue):
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    # 1. preprocess img
    img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3])  # is RGB. not BGR
    img_batch = tf.cast(img_plac, tf.float32)

    img_batch = short_side_resize_for_inference_data(img_tensor=img_batch,
                                                     target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
                                                     length_limitation=cfgs.IMG_MAX_LENGTH)
    if cfgs.NET_NAME in ['resnet152_v1d', 'resnet101_v1d', 'resnet50_v1d']:
        img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD)
    else:
        img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)

    img_batch = tf.expand_dims(img_batch, axis=0)

    detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
        input_img_batch=img_batch,
        gtboxes_batch=None)

    init_op = tf.group(
        tf.global_variables_initializer(),
        tf.local_variables_initializer()
    )

    restorer, restore_ckpt = det_net.get_restorer()

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        sess.run(init_op)
        if not restorer is None:
            restorer.restore(sess, restore_ckpt)
            print('restore model %d ...' % gpu_id)
        for a_img in images:
            raw_img = cv2.imread(os.path.join(eval_data, a_img['file_name']))
            raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]

            resized_img, detected_boxes, detected_scores, detected_categories = \
                sess.run(
                    [img_batch, detection_boxes, detection_scores, detection_category],
                    feed_dict={img_plac: raw_img[:, :, ::-1]}
                )

            resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]
            scales = [raw_w / resized_w, raw_h / resized_h]
            result_dict = {'scales': scales, 'boxes': detected_boxes,
                           'scores': detected_scores, 'labels': detected_categories,
                           'image_id': a_img['id']}
            result_queue.put_nowait(result_dict)

def test_coco(det_net, real_test_img_list, eval_data, gpu_ids):

    save_path = os.path.join('./eval_coco', cfgs.VERSION)
    tools.mkdir(save_path)
    fw_json_dt = open(os.path.join(save_path, 'coco_test-dev.json'), 'w')
    coco_det = []

    nr_records = len(real_test_img_list)
    pbar = tqdm(total=nr_records)
    gpu_num = len(gpu_ids.strip().split(','))

    nr_image = math.ceil(nr_records / gpu_num)
    result_queue = Queue(500)
    procs = []

    for i in range(gpu_num):
        start = i * nr_image
        end = min(start + nr_image, nr_records)
        split_records = real_test_img_list[start:end]
        proc = Process(target=worker, args=(int(gpu_ids.strip().split(',')[i]), split_records, det_net, eval_data, result_queue))
        print('process:%d, start:%d, end:%d' % (i, start, end))
        proc.start()
        procs.append(proc)

    for i in range(nr_records):
        res = result_queue.get()

        xmin, ymin, xmax, ymax = res['boxes'][:, 0], res['boxes'][:, 1], \
                                 res['boxes'][:, 2], res['boxes'][:, 3]

        xmin = xmin * res['scales'][0]
        xmax = xmax * res['scales'][0]

        ymin = ymin * res['scales'][1]
        ymax = ymax * res['scales'][1]

        boxes = np.transpose(np.stack([xmin, ymin, xmax-xmin, ymax-ymin]))

        for j, box in enumerate(boxes):
            coco_det.append({'bbox': [float(box[0]), float(box[1]), float(box[2]), float(box[3])],
                             'score': float(res['scores'][j]), 'image_id': res['image_id'],
                             'category_id': int(classes_originID[LABEL_NAME_MAP[res['labels'][j]]])})

        pbar.set_description("Test image %s" % res['image_id'])

        pbar.update(1)

    for p in procs:
        p.join()

    json.dump(coco_det, fw_json_dt)
    fw_json_dt.close()

def eval(num_imgs, eval_data, json_file, gpu_ids):

    with open(json_file) as f:
        test_img_list = json.load(f)['images']

    if num_imgs == np.inf:
        real_test_img_list = test_img_list
    else:
        real_test_img_list = test_img_list[: num_imgs]

    fcos = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME,
                                                is_training=False)
    test_coco(det_net=fcos, real_test_img_list=real_test_img_list, eval_data=eval_data, gpu_ids=gpu_ids)

def parse_args():

    parser = argparse.ArgumentParser('evaluate the result with Pascal2007 stdand')

    parser.add_argument('--eval_data', dest='eval_data',
                        help='evaluate imgs dir ',
                        default='/data/COCO/test2017', type=str)
    parser.add_argument('--json_file', dest='json_file',
                        help='test-dev json file',
                        default='image_info_test-dev2017.json', type=str)
    parser.add_argument('--showbox', dest='showbox',
                        help='whether show detecion results when evaluation',
                        default=False, type=bool)
    parser.add_argument('--gpus', dest='gpus',
                        help='gpu id',
                        default='0,1,2,3,4,5,6,7', type=str)
    parser.add_argument('--eval_num', dest='eval_num',
                        help='the num of eval imgs',
                        default=np.inf, type=int)
    args = parser.parse_args()
    return args

if __name__ == '__main__':

    args = parse_args()
    print(20*"--")
    print(args)
    print(20*"--")
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
    eval(np.inf,  # use np.inf to test all the imgs. use 10 to test 10 imgs.
         eval_data=args.eval_data,
         json_file=args.eval_gt,
         gpu_ids=args.gpus)

    # os.environ["CUDA_VISIBLE_DEVICES"] = '0'
    # eval(np.inf,  # use np.inf to test all the imgs. use 10 to test 10 imgs.
    #      eval_data='/data/COCO/test2017',
    #      json_file='/data/COCO/annotation/image_info_test-dev2017.json',
    #      gpu_ids='0,1,2,3')

    # cocoval('./eval_coco/FPN_Res101_20190108_v1/coco_res.json',
    #         '/unsullied/sharefs/_research_detection/GeneralDetection/COCO/data/MSCOCO/instances_minival2014.json')


"""
This is a torch implementation of DQN BreakoutDeterministic-v4

Revision compared to previous version:
1. convert data into tensors before we append them to memory.
2. add save_model module

Conclusion: It doesn't work. If you put too many things into GPU, it will OOM.
"""
import os
from time import time

import numpy as np
import gym
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.tensorboard import SummaryWriter
from collections import deque, defaultdict
from itertools import count
from contextlib import contextmanager

@contextmanager
def timeit():
    t1 = time()
    yield
    print("Total time cost is {:.0f}".format(time() - t1))

class Memory(object):
    def __init__(self, maxlen=100000):
        # TODO
        self.state = np.zeros(maxlen).astype(object)
        self.action = np.zeros(maxlen).astype(object)
        self.next_state = np.zeros(maxlen).astype(object)
        self.reward = np.zeros(maxlen).astype(object)
        self.done = np.zeros(maxlen).astype(object)
        self.ind = 0
        self.maxlen = maxlen

    @property
    def size(self):
        return min(self.ind, self.maxlen)

    def add_memory(self, state, action, next_state, reward, done):
        curr_ind = self.ind % self.maxlen
        self.state[curr_ind] = state
        self.action[curr_ind] = action
        self.next_state[curr_ind] = next_state
        self.reward[curr_ind] = reward
        self.done[curr_ind] = done
        self.ind += 1

    def get_memory(self, size):
        n_total = self.ind if self.ind < self.maxlen - 1 else self.maxlen - 1
        inds = np.random.choice(np.arange(n_total), size)
        return [self.state[inds],
                self.action[inds],
                self.next_state[inds],
                self.reward[inds],
                self.done[inds]]

class Net(nn.Module):
    def __init__(self, env):
        super(Net, self).__init__()
        obs_shape = env.observation_space.shape
        act_n = env.action_space.n - 1
        self.conv1 = nn.Conv2d(obs_shape[2], 32, 8, 4)
        self.conv2 = nn.Conv2d(32, 64, 4, 2)
        self.conv3 = nn.Conv2d(64, 64, 3, 1)
        self.dense1 = nn.Linear(22528, 512)
        self.dense2 = nn.Linear(512, act_n)

    def forward(self, x):
        # normalization 0-256 -> 0-1
        x = x / 256
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.conv3(x)
        x = F.relu(x)
        x = x.view(x.shape[0], -1)
        x = self.dense1(x)
        x = F.relu(x)
        x = self.dense2(x)
        return x

class DQNAgent(object):
    def __init__(self, net, memory, gamma):
        self.net = net
        self.memory = memory
        self.gamma = gamma
        self.optimizer = optim.Adam(self.net.parameters())

    def greedy_action(self, state):
        output = self.net.forward(state).detach().numpy()
        raw_action = output.argmax()
        return raw_action

    # don't use 1
    def eps_greedy_action(self, state, eps=0.9):
        output = self.net.forward(state).detach().cpu().numpy()[0]
        # probs_output = (1 - eps) * np.ones(output.shape) / (output.shape[1] - 1)
        # probs_output[(np.arange(output.shape[0]), output.argmax(axis=1))] = eps
        # return [np.random.choice(np.arange(output.shape[1]), p=p) for p in probs_output]\
        probs_output = (1 - eps) * np.ones(output.shape) / (output.shape[0] - 1)
        probs_output[output.argmax()] = eps
        raw_action = np.random.choice(np.arange(output.shape[0]), p=probs_output)
        return raw_action

    def update(self, sampled_memory, counter=None, writer=None):
        states = torch.cat(tuple(sampled_memory[0]), axis=0)
        actions = torch.cat(tuple(sampled_memory[1]), axis=0)
        next_states = torch.cat(tuple(sampled_memory[2]), axis=0)
        rewards = torch.cat(tuple(sampled_memory[3]), axis=0)
        dones = torch.cat(tuple(sampled_memory[4]), axis=0)

        self.optimizer.zero_grad()
        pred_next_state_vals = self.net.forward(next_states)
        pred_state_vals = self.net.forward(states)
        td_target = rewards + \
                    self.gamma * torch.max(pred_next_state_vals, dim=1, keepdims=True).values * (~dones).to(torch.float32)
        td_error = td_target - torch.gather(pred_state_vals, 1, actions)
        loss = torch.mean(td_error ** 2)
        if writer:
            writer.add_scalar("loss", loss, counter)
        loss.backward()
        self.optimizer.step()

def train():
    res = []
    n_episodes = 20000
    train_every_n = 50
    frame_skipping = 2
    batch_size = 32
    memory = Memory(100000)
    exp_name = "200131_5"
    eps_start = 0.75
    eps_end = 0.95

    env = gym.make("BreakoutDeterministic-v4")
    net = Net(env).cuda()
    dqn = DQNAgent(net, memory, 0.9)
    # to clear logs
    os.system('sudo rm -rf logs_{}'.format(exp_name))

    writer = SummaryWriter("logs_{}".format(exp_name))

    global_counter = 0

    for i_episode in range(n_episodes):
        t1 = time()
        state = env.reset()
        state_tensor = torch.from_numpy(state[np.newaxis, :]).cuda().permute(0, 3, 2, 1).to(torch.float32)
        next_state = None
        total_rewards = 0
        action_dict = defaultdict(int)
        for t in count():
            curr_eps = eps_start + (eps_end - eps_start) / n_episodes * i_episode
            raw_action = dqn.eps_greedy_action(state_tensor, eps=curr_eps)
            action_dict[raw_action] += 1
            action = raw_action if raw_action == 0 else raw_action + 1  # cuz we ignore action 1
            # frame skipping
            curr_rewards = 0
            for _ in range(frame_skipping):
                # env.render()
                next_state, reward, done, _ = env.step(action)
                # start new life
                if np.all(next_state == state):
                    next_state, reward, done, _ = env.step(1)
                curr_rewards += reward
                if done: break

            total_rewards += curr_rewards
            memory.add_memory(state_tensor,
                              torch.from_numpy(np.array([[raw_action]])).cuda(),
                              torch.from_numpy(next_state[np.newaxis, :]).cuda().permute(0, 3, 2, 1).to(torch.float32),
                              torch.from_numpy(np.array([[curr_rewards]])).cuda(),
                              torch.from_numpy(np.array([[done]])).cuda())

            if memory.size >= batch_size and global_counter % train_every_n == 0 :
                sampled_memory = memory.get_memory(batch_size)
                dqn.update(sampled_memory, counter=global_counter, writer=writer)

            if done:
                break
            state = next_state
            global_counter += 1
        res.append([i_episode, total_rewards])
        writer.add_scalar("total_rewards", total_rewards, i_episode)
        print("In episode {}, total rewards is {}, {}, time {:.2f}".format(
                                                              i_episode,
                                                              total_rewards,
                                                              [[i, action_dict[i]] for i in range(3)],
                                                              time() - t1))
    with open("result_{}.csv".format(exp_name), "w") as f:
        for i_episode, total_rewards in res:
            f.write("{},{}\n".format(i_episode, total_rewards))

    # save model
    torch.save(net.state_dict(), "model/{}_{}".format(exp_name, int(time.time())))

if __name__ == "__main__":
    with timeit():
        train()


#!/usr/bin/env python
from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import numpy as np
import torch

from torch.nn.modules.module import Module
from ..functions.roi_target import RoITargetFunction

class RoITarget(Module):
    def __init__(self, fg_overlap, box_encoding='fastrcnn'):
        super(RoITarget, self).__init__()
        self.fg_overlap = fg_overlap
        self.box_encoding = box_encoding

    def forward(self, rois, roi_batch_inds, gt_boxes, gt_batch_inds):
        # gt_boxes = gt_boxes[:, :4]
        # classes = gt_boxes[:, 4].long()
        rois = rois[:, :4]
        gt_boxes = gt_boxes.contiguous()
        gt_batch_inds = gt_batch_inds.contiguous()
        rois = rois.contiguous()
        labels, deltas, bbwgts = RoITargetFunction(self.fg_overlap)(rois, roi_batch_inds, gt_boxes, gt_batch_inds)

        return labels, deltas, bbwgts


import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import os

def read_image(image):
    return mpimg.imread(image)

def format_image(image):
    return tf.image.resize(image[tf.newaxis, ...], [224, 224]) / 255.0

def get_category(img):
    """Write a Function to Predict the Class Name

    Args:
        img [jpg]: image file

    Returns:
        [str]: Prediction
    """

    path = 'static/model/'
    tflite_model_file = 'converted_model.tflite'

    # Load TFLite model and allocate tensors.
    with open(path + tflite_model_file, 'rb') as fid:
        tflite_model = fid.read()

    # Interpreter interface for TensorFlow Lite Models.
    interpreter = tf.lite.Interpreter(model_content=tflite_model)
    interpreter.allocate_tensors()

    # Gets model input and output details.
    input_index = interpreter.get_input_details()[0]["index"]
    output_index = interpreter.get_output_details()[0]["index"]
    input_img = read_image(img)
    format_img = format_image(input_img)
    # Sets the value of the input tensor
    interpreter.set_tensor(input_index, format_img)
    # Invoke the interpreter.
    interpreter.invoke()

    predictions_array = interpreter.get_tensor(output_index)
    predicted_label = np.argmax(predictions_array)

    class_names = ['rock', 'paper', 'scissors']

    return class_names[predicted_label]

def plot_category(img, current_time):
    """Plot the input image

    Args:
        img [jpg]: image file
    """
    read_img = mpimg.imread(img)
    ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
    file_path = os.path.join(ROOT_DIR + f'/static/images/output_{current_time}.png')
    print(file_path)

    if os.path.exists(file_path):
        os.remove(file_path)

    plt.imsave(file_path, read_img)