From 36fb512bec44284f0b02bed98f8c121f7bc54884 Mon Sep 17 00:00:00 2001 From: NatLee Date: Tue, 7 Jan 2025 04:55:14 +0800 Subject: [PATCH] [fix] handle between grayscale and RGB image for models --- deepface/models/Demography.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepface/models/Demography.py b/deepface/models/Demography.py index d0a00f1..7bba87f 100644 --- a/deepface/models/Demography.py +++ b/deepface/models/Demography.py @@ -36,7 +36,8 @@ class Demography(ABC): assert img_batch.ndim == 4, "expected 4-dimensional tensor input" if img_batch.shape[0] == 1: # Single image - img_batch = img_batch.squeeze(0) # Remove batch dimension + if img_batch.shape[-1] != 3: # Check if grayscale + img_batch = img_batch.squeeze(0) # Remove batch dimension predict_result = self.model(img_batch, training=False).numpy()[0, :] # Predict with legacy method. return predict_result else: # Batch of images