Annotation changes

sovit-123 · Sep 23, 2022 · 6e9c190 · 6e9c190
1 parent 9892a40
commit 6e9c190
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 7 deletions.
diff --git a/inference_video.py b/inference_video.py
@@ -10,7 +10,7 @@
 
 from models.create_fasterrcnn_model import create_model
 from utils.general import set_infer_dir
-from utils.annotations import inference_annotations
+from utils.annotations import inference_annotations, annotate_fps
 from utils.transforms import infer_transforms
 from torchvision import transforms as transforms
 
@@ -167,10 +167,7 @@ def main(args):
                     outputs, detection_threshold, CLASSES,
                     COLORS, frame
                 )
-            cv2.putText(frame, f"{fps:.1f} FPS", 
-                        (15, 25),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 
-                        2, lineType=cv2.LINE_AA)
+            frame = annotate_fps(frame, fps)
 
             final_end_time = time.time()
             forward_and_annot_time = final_end_time - start_time

diff --git a/utils/annotations.py b/utils/annotations.py
@@ -41,11 +41,52 @@ def inference_annotations(
         cv2.putText(
             orig_image, 
             class_name, 
-            (p1[0], p1[1] - 5),
+            (p1[0], p1[1] - 5 if outside else p1[1] + h + 2),
             cv2.FONT_HERSHEY_SIMPLEX, 
-            fontScale=lw/3, 
+            fontScale=lw / 3, 
             color=(255, 255, 255), 
             thickness=tf, 
             lineType=cv2.LINE_AA
         )
+    return orig_image
+
+def draw_text(
+        img,
+        text,
+        font=cv2.FONT_HERSHEY_SIMPLEX,
+        pos=(0, 0),
+        font_scale=1,
+        font_thickness=2,
+        text_color=(0, 255, 0),
+        text_color_bg=(0, 0, 0),
+    ):
+        offset = (5, 5)
+        x, y = pos
+        text_size, _ = cv2.getTextSize(text, font, font_scale, font_thickness)
+        text_w, text_h = text_size
+        rec_start = tuple(x - y for x, y in zip(pos, offset))
+        rec_end = tuple(x + y for x, y in zip((x + text_w, y + text_h), offset))
+        cv2.rectangle(img, rec_start, rec_end, text_color_bg, -1)
+        cv2.putText(
+            img,
+            text,
+            (x, int(y + text_h + font_scale - 1)),
+            font,
+            font_scale,
+            text_color,
+            font_thickness,
+            cv2.LINE_AA,
+        )
+        return img
+
+def annotate_fps(orig_image, fps_text):
+    draw_text(
+        orig_image,
+        f"FPS: {fps_text:0.1f}",
+        pos=(20, 20),
+        font_scale=1.0,
+        text_color=(204, 85, 17),
+        text_color_bg=(255, 255, 255),
+        font_thickness=2,
+    )
     return orig_image