manhdo/head_pose_estimation_tracking_app

1import cv22import pickle3import os4import argparse5import mediapipe as mp6import numpy as np7import glob8import time9import yaml10from PIL import Image11 12import streamlit as st13st.set_page_config(layout="wide")14 15from utils.drawing_utils import draw_all_informations16from utils.general import resize_img, get_cache_informations, parse_head_pose_informations17from utils.streamlit_options import default_UI18from utils.detection import detect_face_pose_informations_from_image19 20 21IMG_SUFFIX = ['jpeg', 'jpg', 'png']22 23 24def main():25    parser = argparse.ArgumentParser()26    # Pose position27    parser.add_argument('--head_pose_info', default='configs.yaml', help='path to head pose information')28 29    args = parser.parse_args()30 31    ## Streamlit options32    st.title("Head pose estimation tracking app 📷")33    left_col, right_col = st.columns([2, 6])34 35    with left_col:36        default_UI()37 38    img_size = st.session_state.img_size.split('x')39    img_size = [int(s) for s in img_size]40 41 42    if st.session_state.img_upload is not None:43        mp_face_mesh = mp.solutions.face_mesh44        mp_face_detection = mp.solutions.face_detection45        face_mesh = mp_face_mesh.FaceMesh(min_detection_confidence=.5, min_tracking_confidence=0.5)46 47        with open(args.head_pose_info, 'r') as f:48            head_pose_info = yaml.load(f, Loader=yaml.FullLoader)49 50        parse_head_pose_informations(head_pose_info, st.session_state)51 52        img_dict = {}  # List of images upload in RGB53        save_dict = {'img_size': img_size, 54                    'position_horizontal_thresholds': st.session_state.position_horizontal_thresholds,55                    'position_vertical_thresholds': st.session_state.position_vertical_thresholds}56        image_results = {}  # List of results for each image in RGB57        image_names = []  # List image names 58        face_detection_dict = {}  # Face detection information of each image59        face_pose_dict = {}  # Face pose information of each image60        chosen_rectangle_pos_dict = {}  # Positions where have face for each image61        face_direction_dict = {}  # direction of each face for each image62        face_position_dict = {}  # position of each face for each image63        face_coordinate_dict = {}  # coordinate of each face for each image64        face_area_dict = {}65        66        if st.session_state.using_local_cache:67            cache_dict = get_cache_informations(save_dict)68        else:69            cache_dict = {}70 71        for img_upload_file in st.session_state.img_upload:72            image = np.array(Image.open(img_upload_file))73            if len(image.shape) > 2 and image.shape[2] == 4:74                image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)75 76            image_name = img_upload_file.name.split('.')[0]77 78            image_names.append(image_name)79            img_dict[image_name] = image80 81            if image_name in cache_dict:82                face_detection_dict[image_name] = cache_dict[image_name]['face_detection']83                face_pose_dict[image_name] = cache_dict[image_name]['face_pose']84                chosen_rectangle_pos_dict[image_name] = cache_dict[image_name]['chosen_rectangle_pos']85                face_direction_dict[image_name] = cache_dict[image_name]['face_direction']86                face_position_dict[image_name] = cache_dict[image_name]['face_position']87                face_coordinate_dict[image_name] = cache_dict[image_name]['face_coordinate']88                face_area_dict[image_name] = cache_dict[image_name]['face_area']89        total_time = 090 91        with mp_face_detection.FaceDetection(92            model_selection=1, min_detection_confidence=0.5) as face_detection:93            for id, (image_name, image) in enumerate(img_dict.items()):94                if isinstance(image, str):95                    image = cv2.imread(image)96 97                # Check if the image has only 3 channels98                if image.shape[-1] == 1:99                    image = np.stack([image]*3, axis=-1)100 101                start = time.time()102 103                image, h_ratio, w_ratio, pad_top, pad_bot, pad_left, pad_right, img_size_no_pad = \104                                resize_img(image, img_size, return_all_infos=True)105                image_results[image_name] = image106 107                # To improve performance108                image.flags.writeable = False109 110                if not image_name in cache_dict:111                    face_detection_infos, face_direction_infos, face_position_infos, face_coordinate_infos, face_pose_infos, face_area_infos, chosen_rectangle_pos_list = \112                                detect_face_pose_informations_from_image(image, face_detection, face_mesh, img_size, (pad_top, pad_bot, pad_left, pad_right), img_size_no_pad, head_pose_info)113 114                    face_detection_dict[image_name] = face_detection_infos115                    face_direction_dict[image_name] = face_direction_infos116                    face_position_dict[image_name] = face_position_infos117                    face_coordinate_dict[image_name] = face_coordinate_infos118                    face_pose_dict[image_name] = face_pose_infos119                    chosen_rectangle_pos_dict[image_name] = chosen_rectangle_pos_list120                    face_area_dict[image_name] = face_area_infos121 122                end = time.time()123                cur_time = end - start124                total_time += cur_time125 126                draw_all_informations(image_results[image_name], face_detection_dict[image_name], face_direction_dict[image_name], face_position_dict[image_name],127                                    face_coordinate_dict[image_name], face_pose_dict[image_name], face_area_dict[image_name], chosen_rectangle_pos_dict[image_name],128                                    img_size, (pad_top, pad_bot, pad_left, pad_right), (w_ratio, h_ratio), st.session_state, head_pose_info)129 130                ## Save results to dict131                if st.session_state.using_local_cache:132                    save_dict[image_name] = {}133                    save_dict[image_name]['face_detection'] = face_detection_dict[image_name]134                    save_dict[image_name]['face_pose'] = face_pose_dict[image_name]135                    save_dict[image_name]['face_direction'] = face_direction_dict[image_name]136                    save_dict[image_name]['face_position'] = face_position_dict[image_name]137                    save_dict[image_name]['chosen_rectangle_pos'] = chosen_rectangle_pos_dict[image_name]138                    save_dict[image_name]['face_coordinate'] = face_coordinate_dict[image_name]139                    save_dict[image_name]['face_area'] = face_area_dict[image_name]140 141                    if image_name in cache_dict:  # Remove image from cache to save memory142                        del cache_dict[image_name]143 144        print(f"FPS: {len(image_results) / total_time}")145    ## Save results to local cache146    if st.session_state.using_local_cache:147        with open('cache.pkl', 'wb') as f:148            pickle.dump(save_dict, f, protocol=pickle.HIGHEST_PROTOCOL)149 150    image_results = [image_results[image_name] for image_name in image_names]151    with right_col:152        st.image(image_results, width=st.session_state.width_visual, caption=image_names)153 154 155if __name__ == '__main__':156    main()