Complete refqctoring using Blocks and making a nicer web page

app.py

@@ -12,7 +12,60 @@ 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, percentage_to_pixel, draw_key_points_pose, \
-    visualize_vector
+    visualize_vector, draw_axis, draw_axis_3d, draw_cones
+
+# <a href="https://malga.unige.it/" target="_blank"><nobr>Lab MaLGa UniGe</nobr></a>
+WEBSITE = """
+<div class="embed_hidden">
+<h1 style='text-align: center'>Head Pose Estimation and LAEO computation </h1>
+<h2 style='text-align: center'>
+<a target="_blank" href="https://github.com/Malga-Vision/LAEO_demo"> <nobr> Code for LAEO </nobr></a>
+<br>
+<a target="_blank" href="https://github.com/Malga-Vision/HHP-Net/tree/master"> <nobr> Code for HPE </nobr></a>
+
+</h2>
+<h2 style='text-align: center'>
+<nobr><a href="https://github.com/Malga-Vision" target="_blank"><nobr>MaLGa Vision GitHub</nobr></a> &emsp;</nobr>
+</h2>
+
+<h3 style="text-align:center;">
+<a href="https://fede1995.github.io/" target="_blank"><nobr>Federico FT</nobr></a> &emsp;
+</h3>
+
+<h2> Description </h2>
+<p>
+This space illustrates a method for Head Pose Estimation and also LAEO detection. The code is based on experiments and research carried out at the UNiversity of Genoa (Italy) in the MaLGa Laboratory. 
+This demo has been set up by Federico Figari Tomenotti.
+DISCLAIMER: does not work properly on smartphones and sometimes on Safari web browser.
+</p>
+<h2> Usage </h2>
+<p>
+The flags allow the user to choose what to display on the result image, and to change the sensitivity for the person detection algorithm.
+The Head Pose orientation can be described only as one vector (arrow) or a triplet of angles: yaw, pitch and roll projected on the image plane.
+The uncertainty result is the mean of the uncertainty compute on the three angles.
+The run botton is needed to run the demo on an image after changing flag settings.
+For every detailed explanation on the algorithms refer to the paper which will be out soon.
+</p>
+
+</div>
+"""
+
+WEBSITE_citation = """
+<h2 style='text-align: center'>
+Citation
+</h2>
+
+If you find this code useful for your research, please use the following BibTeX entry.
+
+``` 
+@inproceedings{cantarini2022hhp,
+  title={HHP-Net: A light Heteroscedastic neural network for Head Pose estimation with uncertainty},
+  author={Cantarini, Giorgio and Tomenotti, Federico Figari and Noceti, Nicoletta and Odone, Francesca},
+  booktitle={Proceedings of the IEEE/CVF Winter Conference on applications of computer vision},
+  pages={3521--3530},
+  year={2022}
+}
+```"""
 
 
 def load_image(camera, ):
@@ -25,20 +78,26 @@ def load_image(camera, ):
         return False, None
 
 
-def demo_play(img, laeo=True, rgb=False, show_keypoints=True, only_face=False):
+def demo_play(img, laeo=True, rgb=False, show_keypoints=True, only_face=False, Head_Pose_representation='Vector', detection_threshold=0.45):
     # 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)
 
+    if not rgb:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # covert at grey scale
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) # it is still grey scale but with 3 channels to add the colours of the points and lines
+        # img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+    else: # if RGB
+        # img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        pass
+
     print('inference centernet')
-    detections, elapsed_time = detect(model, img_resized, min_score_thresh,
+    detections, elapsed_time = detect(model, img_resized, detection_threshold,
                                       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)
@@ -61,16 +120,19 @@ def demo_play(img, laeo=True, rgb=False, show_keypoints=True, only_face=False):
         #                    size=50)
 
         people_list.append({'yaw'      : yaw[0].numpy()[0],
-                            'yaw_u'    : 0,
+                            'yaw_u'    : yaw[0].numpy()[1],
                             'pitch'    : pitch[0].numpy()[0],
-                            'pitch_u'  : 0,
+                            'pitch_u'  : pitch[0].numpy()[1],
                             'roll'     : roll[0].numpy()[0],
-                            'roll_u'   : 0,
+                            'roll_u'   : roll[0].numpy()[1],
                             'center_xy': [tdx, tdy]
                             })
 
+
     if show_keypoints:
+        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
         for i in range(len(det)):
+            # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
             img = draw_key_points_pose(img, kpt[i], only_face=only_face)
 
     # call LAEO
@@ -83,15 +145,34 @@ def demo_play(img, laeo=True, rgb=False, show_keypoints=True, only_face=False):
         interaction_matrix = np.zeros((len(people_list), len(people_list)))
     # coloured arrow print per person
 
+    print(f'Head pose representation: {Head_Pose_representation}')
+    def visualise_hpe(yaw, pitch, roll, image=None, tdx=None, tdy=None, size=50, yaw_uncertainty=-1, pitch_uncertainty=-1, roll_uncertainty=-1, openpose=False, title="", color=(255, 0, 0)):
+        if str(Head_Pose_representation).lower() == 'vector':
+            vector = project_ypr_in2d(person['yaw'], person['pitch'], person['roll'])
+            image = visualize_vector(image, [tdx, tdy], vector, title=title, color=color)
+            return image
+        elif str(Head_Pose_representation).lower() == 'axis':
+           image = draw_axis_3d(yaw, pitch, roll, image=image, tdx=tdx, tdy=tdy, size=size)
+           return image
+        elif str(Head_Pose_representation).lower() == 'cone':
+            _, image = draw_cones(yaw, pitch, roll, unc_yaw=yaw_uncertainty, unc_pitch=pitch_uncertainty, unc_roll=roll_uncertainty, image=image, tdx=tdx, tdy=tdy, size=size)
+            return image
+        else:
+            return image
+
+
     for index, person in enumerate(people_list):
         green = round((max(interaction_matrix[index, :])) * 255)
         colour = (0, green, 0)
         if green < 40:
             colour = (255, 0, 0)
-        vector = project_ypr_in2d(person['yaw'], person['pitch'], person['roll'])
-        img = visualize_vector(img, person['center_xy'], vector, title="",
-                               color=colour)
-    return img
+        img = visualise_hpe(person['yaw'], person['pitch'], person['roll'], image=img, tdx=person['center_xy'][0], tdy=person['center_xy'][1], size=50, yaw_uncertainty=person['yaw_u'], pitch_uncertainty=person['pitch_u'], roll_uncertainty=person['roll_u'], title="", color=colour)
+        # vector = project_ypr_in2d(person['yaw'], person['pitch'], person['roll'])
+        # img = visualize_vector(img, person['center_xy'], vector, title="",
+        #                        color=colour)
+    uncertainty_mean = [i['yaw_u'] + i['pitch_u'] + i['roll_u'] for i in people_list]
+    uncertainty_mean_str =  ''.join([str(round(i, 2)) + ' ' for i in uncertainty_mean])
+    return img, uncertainty_mean_str
 
 
 if __name__=='__main__':
@@ -114,7 +195,7 @@ if __name__=='__main__':
 
     input_shape_od_model = (512, 512)
     # params
-    min_score_thresh, max_boxes_to_draw, min_distance = .45, 50, 1.5
+    min_score_thresh, max_boxes_to_draw, min_distance = .25, 50, 1.5
 
     print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
 
@@ -123,37 +204,65 @@ if __name__=='__main__':
     live = True
     title = "Head Pose Estimation and LAEO"
 
-    demo_webcam = gr.Interface(
-            fn=function_to_call,
-            inputs=[gr.Image(source="webcam"),  # with no streaming-> acquire images
-                    gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO"),
-                    gr.Checkbox(value=False, label="rgb", info="Display output on W/B image"),
-                    gr.Checkbox(value=True, label="show_keypoints", info="Display keypoints on image"),
-                    gr.Checkbox(value=True, label="show_keypoints_only_face",
-                                info="Display only face keypoints on image"),
-                    ],
-            outputs="image",
-            live=live,
-            title=title,
-            description="This is a demo developed by Federico Figari T. at MaLGa Lab, University of Genoa, Italy. You can choose to have only the Head Pose Estimation or also the LAEO computation (more than 1 person should be in the image). You need to take a picture and the algorithm will calculate the Head Pose and will be showed as an arrow on your face. LAEO, instead is showed colouring the arrow in green.",
-            )
-
-    demo_upload = gr.Interface(
-            fn=function_to_call,
-            inputs=[gr.Image(source="upload", ),  # with no streaming-> acquire images
-                    gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO"),
-                    gr.Checkbox(value=False, label="rgb", info="Display output on W/B image"),
-                    gr.Checkbox(value=True, label="show_keypoints", info="Display keypoints on image"),
-                    gr.Checkbox(value=True, label="show_keypoints_only_face",
-                                info="Display only face keypoints on image"),
-                    ],
-            outputs=gr.Image(height=238, width=585),  # shape=gr.Image(source="upload", ).shape
-            live=live,
-            title=title,
-            description="This is a demo developed by Federico Figari T. at MaLGa Lab, University of Genoa, Italy. You can choose to have only the Head Pose Estimation or also the LAEO computation (more than 1 person should be in the image). You need to upload an image and the algorithm will calculate the Head Pose and will be showed as an arrow on your face. LAEO, instead is showed colouring the arrow in green.",
-            examples=[["LAEO_demo_data/examples/1.jpg"], ["LAEO_demo_data/examples/20.jpg"]]
-            )
-
-    demo_tabbed = gr.TabbedInterface([demo_webcam, demo_upload], ["Demo from webcam", "Demo from upload"])
-
-    demo_tabbed.launch()
+    print(os.getcwd())
+
+    with gr.Blocks() as demo:
+        gr.Markdown(WEBSITE)
+        with gr.Tab("demo_webcam"):
+            with gr.Row():
+                with gr.Column():
+                    input_img = gr.Image(label="Input Image", source="webcam")
+                    laeo = gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO")
+                    rgb = gr.Checkbox(value=False, label="rgb", info="Display output on W/B image")
+                    show_keypoints = gr.Checkbox(value=True, label="show_keypoints", info="Display keypoints on image")
+                    show_keypoints_only_face = gr.Checkbox(value=True, label="show_keypoints_only_face",
+                                                              info="Display only face keypoints on image")
+                    Head_Pose_representation = gr.Radio(["Vector", "Axis", "None"], label="Head_Pose_representation",
+                                                        info="Which representation to show", value="Vector")
+                    detection_threshold = gr.Slider(0.01, 1, value=0.45, step=0.01, interactive=True,
+                                                    label="detection_threshold", info="Choose in [0, 1]")
+
+                    button = gr.Button(label="RUN", type="default")
+                with gr.Column():
+                    outputs = gr.Image(label="Output Image", shape=(512, 512))
+                    uncert = gr.Label(label="Uncertainty", value="0.0")
+
+        input_img.change(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
+                                                       Head_Pose_representation, detection_threshold], outputs=[outputs, uncert])
+        button.click(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
+                                                       Head_Pose_representation, detection_threshold], outputs=[outputs, uncert])
+
+        with gr.Tab("demo_upload"):
+            with gr.Row():
+                with gr.Column():
+                    input_img = gr.Image(label="Input Image", source="upload")
+                    laeo = gr.Checkbox(value=True, label="LAEO", info="Compute and display LAEO")
+                    rgb = gr.Checkbox(value=False, label="rgb", info="Display output on W/B image")
+                    show_keypoints = gr.Checkbox(value=True, label="show_keypoints", info="Display keypoints on image")
+                    show_keypoints_only_face = gr.Checkbox(value=True, label="show_keypoints_only_face",
+                                                           info="Display only face keypoints on image")
+                    Head_Pose_representation = gr.Radio(["Vector", "Axis", "None"],
+                                                        label="Head_Pose_representation",
+                                                        info="Which representation to show", value="Vector")
+                    detection_threshold = gr.Slider(0.01, 1, value=0.45, step=0.01, interactive=True,
+                                                    label="detection_threshold", info="Choose in [0, 1]")
+
+                    button = gr.Button(label="RUN", type="default")
+                with gr.Column():
+                    outputs = gr.Image(height=238, width=585, label="Output Image")
+                    uncert = gr.Label(label="Uncertainty", value="0.0")
+                    examples_text =gr.Markdown("## Image Examples")
+                    examples = gr.Examples([["LAEO_demo_data/examples/1.jpg"], ["LAEO_demo_data/examples/300wlp_0.png"],
+                                ["LAEO_demo_data/examples/AWFL_2.jpg"],
+                                ["LAEO_demo_data/examples/BIWI_3.png"]], inputs=input_img,)  # add all other flags
+
+        input_img.change(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
+                                                   Head_Pose_representation, detection_threshold],
+                         outputs=[outputs, uncert])
+        button.click(function_to_call, inputs=[input_img, laeo, rgb, show_keypoints, show_keypoints_only_face,
+                                               Head_Pose_representation, detection_threshold],
+                     outputs=[outputs, uncert])
+
+        gr.Markdown(WEBSITE_citation)
+
+    demo.launch()

utils/img_util.py

@@ -3,7 +3,11 @@ import os
 import json
 import numpy as np
 from math import cos, sin, pi
-from utils.labels import coco_category_index, rgb_colors, color_pose, color_pose_normalized, pose_id_part, face_category_index, body_parts_openpose, body_parts, face_points, face_points_openpose, pose_id_part_zedcam, face_points_zedcam, body_parts_zedcam
+from utils.labels import coco_category_index, rgb_colors, color_pose, color_pose_normalized, pose_id_part, \
+    face_category_index, body_parts_openpose, body_parts, face_points, face_points_openpose, pose_id_part_zedcam, \
+    face_points_zedcam, body_parts_zedcam, color_pose_rgb
+
+
 # from src.utils.my_utils import fit_plane_least_square  # , retrieve_line_from_two_points
 
 
@@ -294,6 +298,8 @@ def draw_key_points_pose(image, kpt, openpose=False, only_face=False):
         :img (numpy.ndarray): The image with the drawings of lines and key points
     """
 
+
+
     thickness = max (image.shape[0] // 100, image.shape[1] // 100)
     if thickness == 0:
         thickness = 1
@@ -311,17 +317,17 @@ def draw_key_points_pose(image, kpt, openpose=False, only_face=False):
 
     for j in range(len(kpt)):
         # 0 nose, 1/2 left/right eye, 3/4 left/right ear
-        color = color_pose["blue"]
+        color = color_pose_rgb["blue"]
         if j == face_pts[0]:
-            color = color_pose["purple"]# naso
+            color = color_pose_rgb["purple"]# naso
         if j == face_pts[1]:
-            color = color_pose["green"]#["light_pink"]#Leye
+            color = color_pose_rgb["green"]#["light_pink"]#Leye
         if j == face_pts[2]:
-            color = color_pose["dark_pink"]#Reye
+            color = color_pose_rgb["dark_pink"]#Reye
         if j == face_pts[3]:
-            color = color_pose["light_orange"]#LEar
+            color = color_pose_rgb["light_orange"]#LEar
         if j == face_pts[4]:
-            color = color_pose["yellow"]# REar
+            color = color_pose_rgb["yellow"]# REar
         if only_face and j in face_pts:
             if openpose:
                 cv2.circle(image, (int(kpt[j][0]), int(kpt[j][1])), 1, color, thickness=thickness)
@@ -652,7 +658,7 @@ def draw_cones(yaw, pitch, roll, unc_yaw, unc_pitch, unc_roll, image=None, tdx=N
     list_projection_xy = [sin(yaw), -cos(yaw) * sin(pitch)]
     return list_projection_xy, image
 
-def draw_axis_3d(yaw, pitch, roll, image=None, tdx=None, tdy=None, size=50, yaw_uncertainty=-1, pitch_uncertainty=-1, roll_uncertainty=-1):
+def draw_axis_3d(yaw, pitch, roll, image=None, tdx=None, tdy=None, size=50, yaw_uncertainty=-1, pitch_uncertainty=-1, roll_uncertainty=-1, openpose=False):
     """
     Draw yaw pitch and roll axis on the image if passed as input and returns the vector containing the projection of the vector on the image plane
     Args:
@@ -670,6 +676,11 @@ def draw_axis_3d(yaw, pitch, roll, image=None, tdx=None, tdy=None, size=50, yaw_
     Returns:
         :list_projection_xy (list): list containing the unit vector [x, y, z]
     """
+    if openpose:
+        temp= tdy
+        tdy = tdx
+        tdx = tdy
+
     pitch = pitch * np.pi / 180
     yaw = -(yaw * np.pi / 180)
     roll = roll * np.pi / 180
@@ -692,8 +703,18 @@ def draw_axis_3d(yaw, pitch, roll, image=None, tdx=None, tdy=None, size=50, yaw_
     x3 = size * (sin(yaw)) + tdx
     y3 = size * (-cos(yaw) * sin(pitch)) + tdy
     z3 = size * (cos(pitch) * cos(yaw)) + tdy
+
+    red_rgb = (255, 0, 0)
+    green_rgb = (0, 255, 0)
+    blue_rgb = (0, 0, 255)
+
+    red_bgr = (0, 0, 255)
+    green_bgr = (0, 255, 0)
+    blue_bgr = (255, 0, 0)
+
+
     if image is not None:
-        cv2.line(image, (int(tdx), int(tdy)), (int(x1), int(y1)), (0, 0, 255), 2)
-        cv2.line(image, (int(tdx), int(tdy)), (int(x2), int(y2)), (0, 255, 0), 2)
-        cv2.line(image, (int(tdx), int(tdy)), (int(x3), int(y3)), (255, 0, 0), 2)
+        cv2.line(image, (int(tdx), int(tdy)), (int(x1), int(y1)), red_rgb, 3)
+        cv2.line(image, (int(tdx), int(tdy)), (int(x2), int(y2)), green_rgb, 3)
+        cv2.line(image, (int(tdx), int(tdy)), (int(x3), int(y3)), blue_rgb, 3)
     return image

utils/labels.py

@@ -174,17 +174,29 @@ rgb_colors = {
     90: (186, 85, 211),
 }
 
-color_pose = {
+color_pose = { # BGR
     "purple": (255, 0, 100),
     "light_pink": (80, 0, 255),
     "dark_pink": (220, 0, 255),
-    "light_orange": (0, 80, 255),
+    "light_orange": (255, 80, 0),
     "dark_orange": (255, 220, 0.),
     "yellow": (0, 220, 255),
     "blue": (255, 0, 0),
     "green": (0,255,0),
 }
 
+color_pose_rgb= { # RGB
+    "purple": (100, 0, 255),
+    "light_pink": (255, 0, 80),
+    "dark_pink": (255, 0, 220),
+    "light_orange": (0, 80, 255),
+    "dark_orange": (0, 220, 255.),
+    "yellow": (255, 220, 0),
+    "blue": (0, 0, 255),
+    "green": (0,255,0),
+}
+
+
 color_pose_normalized = {
     "purple": (100/255., 0/255., 255/255.),
     "light_pink": (255/255., 0/255., 80/255.),