refactored drawing part for small images, and added one useful flag: only face kpts print

app.py

@@ -11,7 +11,7 @@ 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, \
+from utils.img_util import resize_preserving_ar, percentage_to_pixel, draw_key_points_pose, \
     visualize_vector
 
 
@@ -25,7 +25,7 @@ def load_image(camera, ):
         return False, None
 
 
-def demo_play(img, laeo=True, rgb=False, show_keypoints=True):
+def demo_play(img, laeo=True, rgb=False, show_keypoints=True, only_face=False):
     # webcam in use
 
     # gpus = tf.config.list_physical_devices('GPU')
@@ -41,7 +41,7 @@ def demo_play(img, laeo=True, rgb=False, show_keypoints=True):
     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)
+    # 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'],
@@ -71,7 +71,7 @@ def demo_play(img, laeo=True, rgb=False, show_keypoints=True):
 
     if show_keypoints:
         for i in range(len(det)):
-            img = draw_key_points_pose(img, kpt[i])
+            img = draw_key_points_pose(img, kpt[i], only_face=only_face)
 
     # call LAEO
     clip_uncertainty = 0.5
@@ -87,14 +87,13 @@ def demo_play(img, laeo=True, rgb=False, show_keypoints=True):
         green = round((max(interaction_matrix[index, :])) * 255)
         colour = (0, green, 0)
         if green < 40:
-            colour = (0, 0, 255)
+            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
 
 
-
 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",
@@ -120,19 +119,20 @@ if __name__=='__main__':
     print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
 
     function_to_call = demo_play
-    outputs = "image"
+    # outputs = gr.Image(shape=(512, 512))
     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=outputs,
+            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.",
@@ -144,8 +144,10 @@ if __name__=='__main__':
                     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=outputs,
+            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.",

utils/img_util.py

@@ -265,15 +265,23 @@ def visualize_vector(image, center, unit_vector, title="", color=(0, 0, 255)):
     Returns:
         :result (numpy.ndarray): The image with the vectors drawn
     """
+
+    thickness = image.shape[0] // 100
+    if thickness == 0 or thickness == 1:
+        thickness = 1
+    # if image.shape[0] > 150 or image.shape[1] > 150:
+    #
+    # else:
+    #     thickness = 1
     unit_vector_draw = [unit_vector[0] * image.shape[0]*0.15, unit_vector[1] * image.shape[0]*0.15]
     point = [center[0] + unit_vector_draw[0], center[1] + unit_vector_draw[1]]
 
-    result = cv2.arrowedLine(image, (int(center[0]), int(center[1])), (int(point[0]), int(point[1])), color, thickness=4, tipLength=0.3)
+    result = cv2.arrowedLine(image, (int(center[0]), int(center[1])), (int(point[0]), int(point[1])), color, thickness=thickness, tipLength=0.3)
 
     return result
 
 
-def draw_key_points_pose(image, kpt, openpose=False):
+def draw_key_points_pose(image, kpt, openpose=False, only_face=False):
     """
     Draw the key points and the lines connecting them; it expects the output of CenterNet (not OpenPose format)
 
@@ -286,6 +294,13 @@ def draw_key_points_pose(image, kpt, openpose=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
+    if thickness == 1:
+        thickness = -1
+
+
     parts = body_parts_openpose if openpose else body_parts
     kpt_score = None
     threshold = 0.4
@@ -307,21 +322,29 @@ def draw_key_points_pose(image, kpt, openpose=False):
             color = color_pose["light_orange"]#LEar
         if j == face_pts[4]:
             color = color_pose["yellow"]# REar
-        if openpose:
-            cv2.circle(image, (int(kpt[j][0]), int(kpt[j][1])), 1, color, 2)
-        else:
-            cv2.circle(image, (int(kpt[j][1]), int(kpt[j][0])), 1, color, 2)
+        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)
+            else:
+                cv2.circle(image, (int(kpt[j][1]), int(kpt[j][0])), 1, color, thickness=thickness)
+        elif not only_face:
+            if openpose:
+                cv2.circle(image, (int(kpt[j][0]), int(kpt[j][1])), 1, color, thickness=thickness)
+            else:
+                cv2.circle(image, (int(kpt[j][1]), int(kpt[j][0])), 1, color, thickness=thickness)
+
         # cv2.putText(img, pose_id_part[i], (int(kpts[j][i, 1] * img.shape[1]), int(kpts[j][i, 0] * img.shape[0])), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1, cv2.LINE_AA)
 
-    for part in parts:
-        if int(kpt[part[0]][1]) != 0 and int(kpt[part[0]][0]) != 0 and int(kpt[part[1]][1]) != 0 and int(
-                kpt[part[1]][0]) != 0:
+    if not only_face:
+        for part in parts:
+            if int(kpt[part[0]][1]) != 0 and int(kpt[part[0]][0]) != 0 and int(kpt[part[1]][1]) != 0 and int(
+                    kpt[part[1]][0]) != 0:
 
-            if openpose:
-                cv2.line(overlay, (int(kpt[part[0]][0]), int(kpt[part[0]][1])), (int(kpt[part[1]][0]), int(kpt[part[1]][1])), (255, 255, 255), 2)
-            else:
-                cv2.line(overlay, (int(kpt[part[0]][1]), int(kpt[part[0]][0])),
-                         (int(kpt[part[1]][1]), int(kpt[part[1]][0])), (255, 255, 255), 2)
+                if openpose:
+                    cv2.line(overlay, (int(kpt[part[0]][0]), int(kpt[part[0]][1])), (int(kpt[part[1]][0]), int(kpt[part[1]][1])), (255, 255, 255), 2)
+                else:
+                    cv2.line(overlay, (int(kpt[part[0]][1]), int(kpt[part[0]][0])),
+                             (int(kpt[part[1]][1]), int(kpt[part[1]][0])), (255, 255, 255), 2)
 
     alpha = 0.4
     image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)