Resize in keeping aspect ratio in visualization

rtmo_demo.py

@@ -5,7 +5,7 @@ import cv2
 from pathlib import Path
 import argparse
 import os
-from rtmo_gpu import RTMO_GPU, draw_skeleton
+from rtmo_gpu import RTMO_GPU, draw_skeleton, resize_to_fit_screen
 
 if __name__ == "__main__":
 
@@ -50,6 +50,6 @@ if __name__ == "__main__":
                                     scores,
                                     kpt_thr=0.3,
                                     line_width=2)
-            img_show = cv2.resize(img_show, (788, 525))
+            img_show = resize_to_fit_screen(img_show, 720, 480)
             cv2.imshow(f'{model}', img_show)
             cv2.waitKey(10)

rtmo_demo_batch.py

@@ -4,7 +4,7 @@ import time
 import cv2
 from pathlib import Path
 import argparse
-from rtmo_gpu import RTMO_GPU_Batch, draw_skeleton  # Ensure to import RTMO_GPU_Batch
+from rtmo_gpu import RTMO_GPU_Batch, draw_skeleton, resize_to_fit_screen  # Ensure to import RTMO_GPU_Batch
 
 def process_video(video_path, body_estimator, batch_size=4):
     cap = cv2.VideoCapture(video_path)
@@ -33,7 +33,7 @@ def process_video(video_path, body_estimator, batch_size=4):
                 frame = batch_frames[i]
                 img_show = frame.copy()
                 img_show = draw_skeleton(img_show, keypoints, scores, kpt_thr=0.3, line_width=2)
-                img_show = cv2.resize(img_show, (788, 525))
+                img_show = resize_to_fit_screen(img_show, 720, 480)
                 cv2.imshow(f'{video_path}', img_show)
                 cv2.waitKey(10)
 
@@ -56,7 +56,7 @@ def process_video(video_path, body_estimator, batch_size=4):
             frame = batch_frames[i]
             img_show = frame.copy()
             img_show = draw_skeleton(img_show, keypoints, scores, kpt_thr=0.3, line_width=2)
-            img_show = cv2.resize(img_show, (720, 480))
+            img_show = resize_to_fit_screen(img_show, 720, 480)
             cv2.imshow(f'{video_path}', img_show)
             #cv2.waitKey(10)
 

rtmo_gpu.py

@@ -561,3 +561,22 @@ class RTMO_GPU_Batch(RTMO_GPU):
         outputs = self.inference(batch_img)
         keypoints, scores = self.postprocess_batch(outputs, ratios)
         return keypoints, scores
+
+def resize_to_fit_screen(image, screen_width, screen_height):
+    # Get the dimensions of the image
+    h, w = image.shape[:2]
+
+    # Calculate the aspect ratio of the image
+    aspect_ratio = w / h
+
+    # Determine the scaling factor
+    scale = min(screen_width / w, screen_height / h)
+
+    # Calculate the new dimensions
+    new_width = int(w * scale)
+    new_height = int(h * scale)
+
+    # Resize the image
+    resized_image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
+
+    return resized_image