Update app.py

app.py

@@ -48,21 +48,26 @@ def detect_objects(img, thr):
     # logger.info(f"Number of objs: {len(objs)}")
 
     # draw img
-    img = PersonDetectorAttrib.visualize(img, objs)
+    img, class_names = PersonDetectorAttrib.visualize(img, objs)
+    output_text = ""
+    for i, obj in enumerate(objs):
+        bbox = obj['points']
+        output_text += class_names[i] + "\n"
 
     # return img
-    return img
+    return img, output_text
 
 
 def detect_poses(img, thr):
     # infer img
     objs = person_detector.infer(img, threshold=thr)
-    objs = [obj for obj in objs if obj['confidence'] > thr]
+    objs = [obj for obj in objs if obj['confidence'] > 0.5]
+    output_text = ""
     for i, obj in enumerate(objs):
         bbox = obj['points']
         # crop img
         img_crop = img[bbox[1]:bbox[3], bbox[0]:bbox[2]]
-        objs_pose = pose_detector.infer(img_crop, 0)
+        objs_pose = pose_detector.infer(img_crop, threshold=0.0)
        # each_point in objs_pose['points'], add bbox[0] and bbox[1] to each_point
         for each_point in objs_pose:
             each_point['point'][0] += bbox[0]
@@ -71,37 +76,48 @@ def detect_poses(img, thr):
         objs_pose = [each_point['point']
                      for each_point in objs_pose]
         display_frame_pose(img, [bbox], [objs_pose])
-        # import ipdb; ipdb.set_trace()
-
-        # logger.info(f"Number of objs: {len(objs)}")
-
-        # draw img
-        # img = PoseDetector.visualize(img, objs_pose)
+        # draw bbox
+        img = cv2.rectangle(
+            img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 0, 255), 2)
+        # put bbox ID and pose to output_text
+        output_text += f"ID-{i}: " + "\t".join(
+            [f"point-{i}-[x-{int(point[0])},y-{int(point[1])}]" for i, point in enumerate(objs_pose)]) + "\n"
+        # get text size of output_text
+        text_size = cv2.getTextSize(
+            f"ID-{i}", cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
+        # put rectangle background to output_text at top left bbox
+        img = cv2.rectangle(img, (bbox[0], bbox[1]-text_size[0][1]-10),
+                            (bbox[0]+text_size[0][0], bbox[1]), (0, 0, 0), cv2.FILLED)
+        # putText output_text to image
+        img = cv2.putText(img, f"ID-{i}", (bbox[0], bbox[1]-10),
+                          cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
 
     # return img
-    return img
+    return img, output_text
 
 
 def detect_ages_gender(img, thr):
     # infer img
     objs = person_detector.infer(img, threshold=thr)
-    objs = [obj for obj in objs if obj['confidence'] > thr]
+    objs = [obj for obj in objs if obj['confidence'] > 0.5]
 
     width = img.shape[1]
     font_size = width/1284
     thickness = int((width/1284)*4)
 
+    output_text = ""
     for i, obj in enumerate(objs):
         bbox = obj['points']
         # crop img
         img_crop = img[bbox[1]:bbox[3], bbox[0]:bbox[2]]
-        # img_crop = img_crop[:,:,::-1]
-        
-        objs_ag = ages_genders_detector.infer([img_crop])
+
+        objs_ag = ages_genders_detector.infer(
+            [img_crop])
         # draw img
 
         # show bbox and age_gender to image
-        ages_genders = objs_ag['ages'][0] + "_"+objs_ag['genders'][0]
+        ages_genders = f"ID-{i} :" + \
+            objs_ag['ages'][0] + "_"+objs_ag['genders'][0]
 
         # draw bbox = [x1,y1,x2,y2]
         (text_width, text_height) = cv2.getTextSize(
@@ -121,9 +137,10 @@ def detect_ages_gender(img, thr):
         # putText ages_genders to bbox
         cv2.putText(img, ages_genders, (bbox[0], bbox[1]-2),
                     cv2.FONT_HERSHEY_SIMPLEX, font_size, (0, 0, 255), thickness, cv2.LINE_AA)
+        output_text += ages_genders
 
     # return img
-    return img
+    return img, output_text
 
 
 # ------------------------------------------------------------------------------
@@ -141,6 +158,7 @@ if __name__ == "__main__":
         outputs=[
             gr.Image(type="numpy", label="Processed Image",
                      width=800, height=480),
+            gr.Textbox(label="Person Detection info:", default="")
         ],
         title="Model Testing",
         description="Drag an image onto the box",
@@ -158,6 +176,7 @@ if __name__ == "__main__":
         outputs=[
             gr.Image(type="numpy", label="Processed Image",
                      width=800, height=480),
+            gr.Textbox(label="Human pose detection info:", default=""),
         ],
         title="Model Testing",
         description="Drag an image onto the box",
@@ -176,6 +195,7 @@ if __name__ == "__main__":
         outputs=[
             gr.Image(type="numpy", label="Processed Image",
                      width=800, height=480),
+            gr.Textbox(label="Age Gender detection info:", default=""),
         ],
         title="Model Testing",
         description="Drag an image onto the box",