#detect_video.py
import json
import trt_pose.coco
import trt_pose.models
import torch
import torch2trt
from torch2trt import TRTModule
import time, sys
import cv2
import torchvision.transforms as transforms
import PIL.Image, PIL.ImageDraw
from trt_pose.draw_objects import DrawObjects
from trt_pose.parse_objects import ParseObjects
import argparse
import os.path

#body_index replaces kpoint list
body_index = ["nose",
"L_eye",
"R_eye",
"L_ear",
"R_ear",
"L_sd",
"R_sd",
"L_elbow",
"R_elbow",
"L_wrist",
"R_wrist",
"L_hip",
"R_hip",
"L_knee",
"R_knee",
"L_ankle",
"R_ankle",
"neck"]

def get_keypoint(humans, hnum, peaks):
    #check invalid human index
    kpoint = []
    human = humans[0][hnum]
    C = human.shape[0]
    for j in range(C):
        k = int(human[j])
        if k >= 0:
            peak = peaks[0][j][k]   # peak[1]:width, peak[0]:height
            peak = (j, float(peak[0]), float(peak[1]))
            kpoint.append(peak)
            #print('index:%d : success [%5.3f, %5.3f]'%(j, peak[1], peak[2]) )
        else:
            peak = (j, None, None)
            kpoint.append(peak)
            #print('index:%d : None %d'%(j, k) )
    return kpoint


parser = argparse.ArgumentParser(description='TensorRT pose estimation run')
parser.add_argument('--model', type=str, default='resnet', help = 'resnet or densenet' )
args = parser.parse_args()

with open('human_pose.json', 'r') as f:
    human_pose = json.load(f)

topology = trt_pose.coco.coco_category_to_topology(human_pose)

num_parts = len(human_pose['keypoints'])
num_links = len(human_pose['skeleton'])


if 'resnet' in args.model:
    print('------ model = resnet--------')
    MODEL_WEIGHTS = 'resnet18_baseline_att_224x224_A_epoch_249.pth'
    OPTIMIZED_MODEL = 'resnet18_baseline_att_224x224_A_epoch_249_trt.pth'
    model = trt_pose.models.resnet18_baseline_att(num_parts, 2 * num_links).cuda().eval()
    WIDTH = 224
    HEIGHT = 224

else:
    print('------ model = densenet--------')
    MODEL_WEIGHTS = 'densenet121_baseline_att_256x256_B_epoch_160.pth'
    OPTIMIZED_MODEL = 'densenet121_baseline_att_256x256_B_epoch_160_trt.pth'
    model = trt_pose.models.densenet121_baseline_att(num_parts, 2 * num_links).cuda().eval()
    WIDTH = 256
    HEIGHT = 256

data = torch.zeros((1, 3, HEIGHT, WIDTH)).cuda()
if os.path.exists(OPTIMIZED_MODEL) == False:
    model.load_state_dict(torch.load(MODEL_WEIGHTS))
    model_trt = torch2trt.torch2trt(model, [data], fp16_mode=True, max_workspace_size=1<<25)
    torch.save(model_trt.state_dict(), OPTIMIZED_MODEL)

model_trt = TRTModule()
model_trt.load_state_dict(torch.load(OPTIMIZED_MODEL))

t0 = time.time()
torch.cuda.current_stream().synchronize()
for i in range(50):
    y = model_trt(data)
torch.cuda.current_stream().synchronize()
t1 = time.time()

print(50.0 / (t1 - t0))

mean = torch.Tensor([0.485, 0.456, 0.406]).cuda()
std = torch.Tensor([0.229, 0.224, 0.225]).cuda()
device = torch.device('cuda')

def preprocess(image):
    global device
    device = torch.device('cuda')
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    image = PIL.Image.fromarray(image)
    image = transforms.functional.to_tensor(image).to(device)
    image.sub_(mean[:, None, None]).div_(std[:, None, None])
    return image[None, ...]

#For draw line 0(nose)->17(neck), 1(L_eye)->3(L_ear)...
line_li = [[0,17],[1,3],[0,1],[2,4],[0,2],[5,11],[6,12],[7,9],[5,7],[6,5],[8,6],[10,8],[13,15],[11,13],[12,11],[14,12],[16,14]]

def draw_keypoint(src,counts,objects,peaks,src_w,src_h):
    X_compress = src_w / WIDTH * 1.9
    Y_compress =  src_h / HEIGHT / 2.0
    color = (0,255,0)
    for i in range(counts[0]):
        keypoints = get_keypoint(objects, i, peaks)
        for j in range(len(keypoints)):
            if keypoints[j][1]:
                x = round(keypoints[j][2] * WIDTH * X_compress)
                y = round(keypoints[j][1] * HEIGHT * Y_compress)
                cv2.circle(src, (x, y), 3, color, 2)
                cv2.putText(src , "%s" % body_index[keypoints[j][0]], (x + 5, y),  cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 1)
                cv2.circle(src, (x, y), 3, color, 2)
        for x in line_li:
            if keypoints[x[0]][1] != None and keypoints[x[1]][1] != None:
                x0 = round(keypoints[x[0]][2]*WIDTH*X_compress)
                y0 = round(keypoints[x[0]][1]*HEIGHT*Y_compress)
                x1 = round(keypoints[x[1]][2]*WIDTH*X_compress)
                y1 = round(keypoints[x[1]][1]*HEIGHT*Y_compress)
                cv2.line(src, (x0, y0), (x1, y1), color, 2)
    return src


def excute(img, src,t):
    src_w, src_h, tmp = src.shape
    data = preprocess(img)
    cmap, paf = model_trt(data)
    cmap, paf = cmap.detach().cpu(), paf.detach().cpu()
    counts, objects, peaks = parse_objects(cmap, paf)#, cmap_threshold=0.15, link_threshold=0.15)
    fps = 1.0 / (time.time()-t)
    src = draw_keypoint(src,counts,objects,peaks,src_w,src_h)
    print(fps)
    cv2.putText(src , "FPS: %f" % (fps), (50, 50),  cv2.FONT_HERSHEY_SIMPLEX, int(src_h/8), (0, 255, 0), 1)
    return src

cap = cv2.VideoCapture("testvideo.mp4")#video file path
ret_val, img = cap.read()
fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')#mp4 codek
count = 0

parse_objects = ParseObjects(topology)
draw_objects = DrawObjects(topology)

k=1
while True:
    t = time.time()
    ret,frame = cap.read()
    try:
        img = cv2.resize(frame, dsize=(WIDTH, HEIGHT), interpolation=cv2.INTER_AREA)
    except:
        break
    frame = excute(img, frame,t)
    cv2.imwrite("images/"+str(k)+".jpeg",frame)
    k+=1

cv2.destroyAllWindows()
cap.release()