"""
Hand and face detection and pose estimation inference on images
"""
import sys
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import cv2

from detection.modules.inference_utils import init_detector, detect, draw_boxes
from pose_estimation.modules.inference_utils import init_interhand_estimator, estimate_interhand, draw_pose
from pose_estimation.modules.inference_utils import init_face_estimator, estimate_face, draw_keypoints
from pose_estimation.modules.inference_utils import init_hand_estimator, estimate_hands


def main():
    base_path = os.path.dirname(os.path.realpath(__file__))
    # init models
    detector_cfg = 'yolov5m'
    detector = init_detector(detector_cfg)
    hand_estimator, hand_skeleton = init_hand_estimator('efficientnet-b1')
    face_estimator = init_face_estimator('efficientnet-b1')
    # list of images to process
    imgs = ['img_2.jpg']
    # for each image
    for img_name in imgs:
        # input path, input videos are expected to be in data/images/
        path = os.path.join(base_path, 'data', 'images', img_name)
        # read image
        img = cv2.imread(path)
        # detection inference
        boxes, scores, classes = detect(detector, img)
        hand_boxes = boxes[classes == 0, :]
        face_boxes = boxes[classes == 1, :]
        # hand and face pose estimation
        handness, hand_keypoints = estimate_hands(hand_estimator, img, hand_boxes, center_format=True)
        face_keypoints = estimate_face(face_estimator, img, face_boxes, center_format=True)
        # draw detections on one image and pose estimation on a copy of input image
        detection_img = img.copy()
        draw_boxes(detection_img, (boxes, scores, classes))
        cv2.imshow('detections', detection_img)
        
        draw_pose(img, hand_keypoints, handness, hand_skeleton)
        draw_keypoints(img, face_keypoints)
        # visualize and write images
        cv2.imshow('pose', img)
        if cv2.waitKey(0) == 27:
            break
        cv2.imwrite(os.path.join(base_path, 'data', 'images', f'detection_{img_name}'), detection_img)
        cv2.imwrite(os.path.join(base_path, 'data', 'images', f'pose_{img_name}'), img)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()
