File size: 4,433 Bytes
9d11120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eea32c6
9d11120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eea32c6
9d11120
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
eea32c6
 
 
 
 
 
 
 
 
 
 
9d11120
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
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("data"):
        gdown.download_folder("https://drive.google.com/drive/folders/1nQ1Cb_tBEhWxy183t-mIcVH7AhAfa6NO?usp=drive_link",
                              use_cookies=False)
    gaze_model_path = 'data/head_pose_estimation'
    gaze_model = tf.keras.models.load_model(gaze_model_path, custom_objects={"tf": tf})
    path_to_model = '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()