From fa4044adae2ba84cf4b3d07916416717361a7a61 Mon Sep 17 00:00:00 2001 From: h-alice Date: Mon, 13 Jan 2025 23:14:40 +0800 Subject: [PATCH] patch: Greyscale image prediction condition. --- deepface/models/Demography.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/deepface/models/Demography.py b/deepface/models/Demography.py index 329a156..5fcc431 100644 --- a/deepface/models/Demography.py +++ b/deepface/models/Demography.py @@ -38,14 +38,15 @@ class Demography(ABC): if not self.model_name: # Check if called from derived class raise NotImplementedError("no model selected") assert img_batch.ndim == 4, "expected 4-dimensional tensor input" - # Single image - if img_batch.shape[0] == 1: + + if img_batch.shape[-1] != 3: # Handle grayscale image, check last dimension. # Check if grayscale by checking last dimension, if not 3, it is grayscale. - if img_batch.shape[-1] != 3: - # Remove batch dimension - img_batch = img_batch.squeeze(0) + img_batch = img_batch.squeeze(0) # Remove batch dimension + + if img_batch.shape[0] == 1: # Single image # Predict with legacy method. return self.model(img_batch, training=False).numpy()[0, :] + # Batch of images # Predict with batch prediction return self.model.predict_on_batch(img_batch)