diff --git a/Pytorch_Retinaface/pytorch_retinaface.py b/Pytorch_Retinaface/pytorch_retinaface.py index 69832c6..97364dc 100644 --- a/Pytorch_Retinaface/pytorch_retinaface.py +++ b/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 diff --git a/__init__.py b/__init__.py index ea78c2e..d9446b4 100644 --- a/__init__.py +++ b/__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 = [], []