Skip to content

Commit 88b5580

Browse files
committed
target size arg added into extract faces
1 parent 24f57ec commit 88b5580

File tree

3 files changed

+93
-3
lines changed

3 files changed

+93
-3
lines changed

retinaface/RetinaFace.py

+20-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import os
22
import warnings
33
import logging
4-
from typing import Union, Any, Optional, Dict
4+
from typing import Union, Any, Optional, Dict, Tuple, List
55

66
# this has to be set before importing tf
77
os.environ["TF_USE_LEGACY_KERAS"] = "1"
@@ -220,7 +220,9 @@ def extract_faces(
220220
align: bool = True,
221221
allow_upscaling: bool = True,
222222
expand_face_area: int = 0,
223-
) -> list:
223+
target_size: Optional[Tuple[int, int]] = None,
224+
min_max_norm: bool = True,
225+
) -> List[np.ndarray]:
224226
"""
225227
Extract detected and aligned faces
226228
Args:
@@ -230,6 +232,13 @@ def extract_faces(
230232
align (bool): enable or disable alignment
231233
allow_upscaling (bool): allowing up-scaling
232234
expand_face_area (int): expand detected facial area with a percentage
235+
target_size (optional tuple): resize the image by padding it with black pixels
236+
to fit the specified dimensions. default is None
237+
min_max_norm (bool): set this to True if you want to normalize image in [0, 1].
238+
this is only running when target_size is not none.
239+
for instance, matplotlib expects inputs in this scale. (default is True)
240+
Returns:
241+
result (List[np.ndarray]): list of extracted faces
233242
"""
234243
resp = []
235244

@@ -289,6 +298,14 @@ def extract_faces(
289298
int(rotated_y1) : int(rotated_y2), int(rotated_x1) : int(rotated_x2)
290299
]
291300

292-
resp.append(facial_img[:, :, ::-1])
301+
if target_size is not None:
302+
facial_img = postprocess.resize_image(
303+
img=facial_img, target_size=target_size, min_max_norm=min_max_norm
304+
)
305+
306+
# to rgb
307+
facial_img = facial_img[:, :, ::-1]
308+
309+
resp.append(facial_img)
293310

294311
return resp

retinaface/commons/postprocess.py

+63
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,18 @@
1+
# built-in dependencies
12
import math
23
from typing import Union, Tuple
4+
5+
# 3rd party dependencies
36
import numpy as np
47
from PIL import Image
8+
import cv2
9+
import tensorflow as tf
10+
11+
tf_major_version = int(tf.__version__.split(".", maxsplit=1)[0])
12+
if tf_major_version == 1:
13+
from keras.preprocessing import image
14+
else:
15+
from tensorflow.keras.preprocessing import image
516

617

718
# pylint: disable=unused-argument
@@ -143,6 +154,58 @@ def rotate_facial_area(
143154
return (x1, y1, x2, y2)
144155

145156

157+
def resize_image(
158+
img: np.ndarray, target_size: Tuple[int, int], min_max_norm: bool = True
159+
) -> np.ndarray:
160+
"""
161+
Resize an image to expected size of a ml model with adding black pixels.
162+
Ref: github.com/serengil/deepface/blob/master/deepface/modules/preprocessing.py
163+
Args:
164+
img (np.ndarray): pre-loaded image as numpy array
165+
target_size (tuple): input shape of ml model
166+
min_max_norm (bool): set this to True if you want to normalize image in [0, 1].
167+
this is only running when target_size is not none.
168+
for instance, matplotlib expects inputs in this scale. (default is True)
169+
Returns:
170+
img (np.ndarray): resized input image
171+
"""
172+
factor_0 = target_size[0] / img.shape[0]
173+
factor_1 = target_size[1] / img.shape[1]
174+
factor = min(factor_0, factor_1)
175+
176+
dsize = (
177+
int(img.shape[1] * factor),
178+
int(img.shape[0] * factor),
179+
)
180+
img = cv2.resize(img, dsize)
181+
182+
diff_0 = target_size[0] - img.shape[0]
183+
diff_1 = target_size[1] - img.shape[1]
184+
185+
# Put the base image in the middle of the padded image
186+
img = np.pad(
187+
img,
188+
(
189+
(diff_0 // 2, diff_0 - diff_0 // 2),
190+
(diff_1 // 2, diff_1 - diff_1 // 2),
191+
(0, 0),
192+
),
193+
"constant",
194+
)
195+
196+
# double check: if target image is not still the same size with target.
197+
if img.shape[0:2] != target_size:
198+
img = cv2.resize(img, target_size)
199+
200+
# make it 4-dimensional how ML models expect
201+
img = image.img_to_array(img)
202+
203+
if min_max_norm is True and img.max() > 1:
204+
img = (img.astype(np.float32) / 255.0).astype(np.float32)
205+
206+
return img
207+
208+
146209
def bbox_pred(boxes, box_deltas):
147210
"""
148211
This function is copied from the following code snippet:

tests/test_actions.py

+10
Original file line numberDiff line numberDiff line change
@@ -107,3 +107,13 @@ def test_different_expanding_ratios():
107107
plt.imshow(face)
108108
plt.axis("off")
109109
plt.show()
110+
111+
112+
def test_resize():
113+
faces = RetinaFace.extract_faces(img_path="tests/dataset/img11.jpg", target_size=(224, 224))
114+
for face in faces:
115+
assert face.shape == (224, 224, 3)
116+
if do_plotting is True:
117+
plt.imshow(face)
118+
plt.show()
119+
logger.info("✅ resize test done")

0 commit comments

Comments
 (0)