diff --git a/inference_video.py b/inference_video.py index 6e164d11..59d5f63d 100644 --- 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 index f87c98c6..16154f1f 100644 --- 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 \ No newline at end of file