Spaces:
Runtime error
Runtime error
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() |