fix aspect ratio

Pytorch_Retinaface/pytorch_retinaface.py

@@ -1,5 +1,6 @@
 import os
 import time
+import math
 import torch
 import torch.backends.cudnn as cudnn
 import numpy as np
@@ -81,10 +82,13 @@ class Pytorch_RetinaFace:
             face_width = x2 - x1
             face_height = y2 - y1
 
-            # New height and width based on scale factor
-            new_face_height = int(face_height * scale_factor)
-            # new_face_width = int(new_face_height * (face_width / face_height))
-            new_face_width = int(new_face_height / aspect_ratio)
+            default_area = face_width * face_height
+            default_side = math.sqrt(default_area)
+            default_area *= scale_factor
+
+            # New height and width based on aspect_ratio
+            new_face_width = int(default_side * math.sqrt(aspect_ratio))
+            new_face_height = int(default_side / math.sqrt(aspect_ratio))
 
             # Center coordinates of the detected face
             center_x = x1 + face_width // 2

__init__.py

@@ -42,11 +42,14 @@ class AutoCropFaces:
                     "max": 1000,
                     "step": 1,
                 }),
-                "aspect_ratio": ("FLOAT", {
-                    "default": 1, 
-                    "min": 0.2,
-                    "max": 5,
-                    "step": 0.1,
+                # "aspect_ratio": ("FLOAT", {
+                #     "default": 1, 
+                #     "min": 0.2,
+                #     "max": 5,
+                #     "step": 0.1,
+                # }),
+                "aspect_ratio": (["9:16", "2:3", "3:4", "4:5", "1:1", "5:4", "4:3", "3:2", "16:9"], {
+                    "default": "1:1",
                 }),
             },
         }
@@ -58,6 +61,10 @@ class AutoCropFaces:
 
     CATEGORY = "Faces"
 
+    def aspect_ratio_string_to_float(self, str_aspect_ratio="1:1"):
+        a, b = map(float, str_aspect_ratio.split(':'))
+        return a / b
+
     def auto_crop_faces_in_image (self, image, max_number_of_faces, scale_factor, shift_factor, aspect_ratio, method='lanczos'): 
         image_255 = image * 255
         rf = Pytorch_RetinaFace(top_k=50, keep_top_k=max_number_of_faces)
@@ -78,6 +85,9 @@ class AutoCropFaces:
         "aspect_ratio" - When we crop, you can have it crop down at a particular aspect ratio.
         "method" - Scaling pixel sampling interpolation method.
         """
+        
+        # Turn aspect ratio to float value
+        aspect_ratio = self.aspect_ratio_string_to_float(aspect_ratio)
 
         selected_faces, detected_cropped_faces = [], []
         selected_crop_data, detected_crop_data = [], []