import gdown import gradio as gr import logging import os import cv2 import numpy as np import tensorflow as tf from ai.detection import detect from laeo_per_frame.interaction_per_frame_uncertainty import LAEO_computation from utils.hpe import hpe, project_ypr_in2d from utils.img_util import resize_preserving_ar, draw_detections, percentage_to_pixel, draw_key_points_pose, \ visualize_vector def load_image(camera, ): # Capture the video frame by frame try: ret, frame = camera.read() return True, frame except: logging.Logger('Error reading frame') return False, None def demo_play(img, laeo=True, rgb=False): # webcam in use # gpus = tf.config.list_physical_devices('GPU') # img = np.array(frame) if not rgb: img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) img_resized, new_old_shape = resize_preserving_ar(img, input_shape_od_model) print('inference centernet') detections, elapsed_time = detect(model, img_resized, min_score_thresh, new_old_shape) # detection classes boxes scores # probably to draw on resized img_with_detections = draw_detections(img_resized, detections, max_boxes_to_draw, None, None, None) # cv2.imshow("aa", img_with_detections) det, kpt = percentage_to_pixel(img.shape, detections['detection_boxes'], detections['detection_scores'], detections['detection_keypoints'], detections['detection_keypoint_scores']) # center_xy, yaw, pitch, roll = head_pose_estimation(kpt, 'centernet', gaze_model=gaze_model) # _________ extract hpe and print to img people_list = [] print('inferece hpe') for j, kpt_person in enumerate(kpt): yaw, pitch, roll, tdx, tdy = hpe(gaze_model, kpt_person, detector='centernet') # img = draw_axis_3d(yaw[0].numpy()[0], pitch[0].numpy()[0], roll[0].numpy()[0], image=img, tdx=tdx, tdy=tdy, # size=50) people_list.append({'yaw' : yaw[0].numpy()[0], 'yaw_u' : 0, 'pitch' : pitch[0].numpy()[0], 'pitch_u' : 0, 'roll' : roll[0].numpy()[0], 'roll_u' : 0, 'center_xy': [tdx, tdy] }) for i in range(len(det)): img = draw_key_points_pose(img, kpt[i]) # call LAEO clip_uncertainty = 0.5 binarize_uncertainty = False if laeo: interaction_matrix = LAEO_computation(people_list, clipping_value=clip_uncertainty, clip=binarize_uncertainty) else: interaction_matrix = np.zeros((len(people_list), len(people_list))) # coloured arrow print per person for index, person in enumerate(people_list): green = round((max(interaction_matrix[index, :])) * 255) colour = (0, green, 0) if green < 40: colour = (0, 0, 255) vector = project_ypr_in2d(person['yaw'], person['pitch'], person['roll']) img = visualize_vector(img, person['center_xy'], vector, title="", color=colour) return img demo = gr.Interface( fn=demo_play, inputs=[gr.Image(source="webcam", streaming=True), gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO"), gr.Checkbox(value=True, label="rgb", info="Display output on W/B image"), ], outputs="image", live=True ) if __name__=='__main__': if not os.path.exists("LAEO_demo_data"): gdown.download_folder("https://drive.google.com/drive/folders/1nQ1Cb_tBEhWxy183t-mIcVH7AhAfa6NO?usp=drive_link", use_cookies=False) gaze_model_path = 'LAEO_demo_data/head_pose_estimation' gaze_model = tf.keras.models.load_model(gaze_model_path, custom_objects={"tf": tf}) path_to_model = 'LAEO_demo_data/keypoint_detector/centernet_hg104_512x512_kpts_coco17_tpu-32' model = tf.saved_model.load(os.path.join(path_to_model, 'saved_model')) input_shape_od_model = (512, 512) # params min_score_thresh, max_boxes_to_draw, min_distance = .45, 50, 1.5 print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU'))) demo.launch()