"""
Inference of pose estimation model on webcam feed, ROIs are static
"""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
import cv2
import numpy as np

from modules.inference_utils import init_interhand_estimator, estimate_interhand, draw_pose
from modules.inference_utils import init_face_estimator, estimate_face, draw_keypoints
from modules.inference_utils import init_hand_estimator, estimate_hands


def interhand():
    """Infer Interhand model on webcam feed
    """
    # webcam feed
    cam = cv2.VideoCapture(0)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
    # get inference model
    model, skeleton = init_interhand_estimator()
    # static ROI
    bbox = [200, 340, 256, 256] # xmin, ymin, width, height
    while cv2.waitKey(1) != 27:
        # grab frame
        _, img = cam.read()
        # estimate keypoints
        handness, joints, relative_depth = estimate_interhand(model, img, [bbox])
        # visualize results
        img = cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), (0, 0, 255), (2))
        img = draw_pose(img, joints.astype(np.int32), handness > 0.5, skeleton)
        cv2.imshow('img', img)
    cv2.destroyAllWindows()


def face():
    """Infer face pose estimation model on webcam feed
    """
    # webcam feed
    cam = cv2.VideoCapture(0)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
    # get inference model
    model = init_face_estimator('efficientnet-b1')
    # static ROI
    bbox = [100, 240, 500, 640] # xmin, ymin, xmax, ymax
    while cv2.waitKey(1) != 27:
        # grab frame
        _, img = cam.read()
        # estimate keypoints
        kaypoints = estimate_face(model, img, [bbox])
        # visualize results
        img = cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), (2))
        img = draw_keypoints(img, kaypoints.astype(np.int32))
        cv2.imshow('img', img)
    cv2.destroyAllWindows()


def hands():
    """Infer hand pose estimation model on webcam feed
    """
    # webcam feed
    cam = cv2.VideoCapture(0)
    cam.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
    # get inference model
    model, skeleton = init_hand_estimator('efficientnet-b1')
    # static ROI
    bbox = [200, 340, 256, 256] # xmin, ymin, width, height
    while cv2.waitKey(1) != 27:
        # grab frame
        _, img = cam.read()
        # estimate keypoints
        handness, keypoints = estimate_hands(model, img, [bbox])
        # visualize results
        img = cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[0]+bbox[2], bbox[1]+bbox[3]), (0, 0, 255), (2))
        img = draw_pose(img, keypoints, handness, skeleton)
        cv2.imshow('img', img)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    # interhand()
    # face()
    hands()
