import numpy as np
import cv2
import os
import time
from numpy import array

# some variables
weightsPath = './yolov3-tiny.weights'
configPath = './yolov3-tiny.cfg'
labelsPath = './obj.names'
LABELS = open(labelsPath).read().strip().split("\n")
colors = [(255, 255, 0), (255, 0, 255), (0, 255, 255), (0, 255, 0), (255, 0, 255)]
min_score = 0.3

# read darknet weights using opencv.dnn module
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)

# read video using opencv
cap = cv2.VideoCapture('./MY_TEST/8.h264') 

# loop for inference
while True:
    boxes = []
    confidences = []
    classIDs = []
    start = time.time()

    ret, frame = cap.read()
    frame = cv2.resize(frame, (744, 416), interpolation=cv2.INTER_CUBIC)
    image = frame
    (H, W) = image.shape[0: 2]


    # get output layer names
    ln = net.getLayerNames()
    out = net.getUnconnectedOutLayers()
    x = []
    for i in out:
        x.append(ln[i[0]-1])
    ln = x

    # create input data package with current frame
    blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), swapRB=True, crop=False)
    # set as input
    net.setInput(blob)

    # run!
    layerOutputs = net.forward(ln)

    # post-process
    # parsing the output and run nms
    for output in layerOutputs:
        for detection in output:
            scores = detection[5:]
            # class id
            classID = np.argmax(scores)
            # get score by classid
            score = scores[classID]

            # ignore if score is too low
            if score >= min_score:
                box = detection[0:4] * np.array([W, H, W, H])
                (centerX, centerY, width, height)= box.astype("int")

                x = int(centerX - (width / 2))
                y = int(centerY - (height / 2))
                
                boxes.append([x, y, int(width), int(height)])
                confidences.append(float(score))
                classIDs.append(classID)
    
    # run nms using opencv.dnn module
    idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.2, 0.3)
    
    # render on image
    idxs = array(idxs)
    box_seq = idxs.flatten()
    if len(idxs) > 0:
        for seq in box_seq:
            (x, y) = (boxes[seq][0], boxes[seq][1])
            (w, h) = (boxes[seq][2], boxes[seq][3])

            # draw what you want
            color = colors[classIDs[seq]]
            cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
            text = "{}: {:.3f}".format(LABELS[classIDs[seq]], confidences[seq])
            cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.3, color, 1)  

    cv2.namedWindow('Image', cv2.WINDOW_NORMAL)
    cv2.imshow("Image", image)

    # print fps 
    stop = time.time()
    fps = 1/(stop - start)
    print('fps>>> :', fps)

    # normal codes when displaying video
    c = cv2.waitKey(1) & 0xff
    if c == 27:
        cap.release()
        break

cv2.destroyAllWindows()