[fix] lint

This commit is contained in:
NatLee 2025-01-13 22:27:11 +08:00
parent eb7b8411e8
commit 688fbe6b90
2 changed files with 25 additions and 21 deletions

View File

@ -28,24 +28,32 @@ class Demography(ABC):
And switch to batch prediction if receives batched images. And switch to batch prediction if receives batched images.
Args: Args:
img_batch: Batch of images as np.ndarray (n, x, y, c), with n >= 1, x = image width, y = image height, c = channel img_batch:
Or Single image as np.ndarray (1, x, y, c), with x = image width, y = image height and c = channel Batch of images as np.ndarray (n, x, y, c)
with n >= 1, x = image width, y = image height, c = channel
Or Single image as np.ndarray (1, x, y, c)
with x = image width, y = image height and c = channel
The channel dimension may be omitted if the image is grayscale. (For emotion model) The channel dimension may be omitted if the image is grayscale. (For emotion model)
""" """
if not self.model_name: # Check if called from derived class if not self.model_name: # Check if called from derived class
raise NotImplementedError("no model selected") raise NotImplementedError("no model selected")
assert img_batch.ndim == 4, "expected 4-dimensional tensor input" assert img_batch.ndim == 4, "expected 4-dimensional tensor input"
# Single image
if img_batch.shape[0] == 1:
# 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)
# 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)
if img_batch.shape[0] == 1: # Single image def _preprocess_batch_or_single_input(
if img_batch.shape[-1] != 3: # Check if grayscale by checking last dimension, if not 3, it is grayscale. self,
img_batch = img_batch.squeeze(0) # Remove batch dimension img: Union[np.ndarray, List[np.ndarray]]
predict_result = self.model(img_batch, training=False).numpy()[0, :] # Predict with legacy method. ) -> np.ndarray:
return predict_result
else: # Batch of images
return self.model.predict_on_batch(img_batch) # Predict with batch prediction
def _preprocess_batch_or_single_input(self, img: Union[np.ndarray, List[np.ndarray]]) -> np.ndarray:
""" """
Preprocess single or batch of images, return as 4-D numpy array. Preprocess single or batch of images, return as 4-D numpy array.
@ -56,15 +64,11 @@ class Demography(ABC):
Returns: Returns:
Four-dimensional numpy array (n, 224, 224, 3) Four-dimensional numpy array (n, 224, 224, 3)
""" """
image_batch = np.array(img) image_batch = np.array(img)
# Remove batch dimension in advance if exists # Remove batch dimension in advance if exists
image_batch = image_batch.squeeze() image_batch = image_batch.squeeze()
# Check input dimension # Check input dimension
if len(image_batch.shape) == 3: if len(image_batch.shape) == 3:
# Single image - add batch dimension # Single image - add batch dimension
image_batch = np.expand_dims(image_batch, axis=0) image_batch = np.expand_dims(image_batch, axis=0)
return image_batch return image_batch

View File

@ -168,9 +168,9 @@ def analyze(
model = modeling.build_model(task="facial_attribute", model_name=action.capitalize()) model = modeling.build_model(task="facial_attribute", model_name=action.capitalize())
predictions = model.predict(faces_array) predictions = model.predict(faces_array)
# If the model returns a single prediction, reshape it to match the number of faces # If the model returns a single prediction, reshape it to match the number of faces.
# Use number of faces and number of predictions shape to determine the correct shape of predictions # Determine the correct shape of predictions by using number of faces and predictions shape.
# For example, if there are 1 face to predict with Emotion model, reshape predictions to (1, 7) # Example: For 1 face with Emotion model, predictions will be reshaped to (1, 7).
if faces_array.shape[0] == 1 and len(predictions.shape) == 1: if faces_array.shape[0] == 1 and len(predictions.shape) == 1:
# For models like `Emotion`, which return a single prediction for a single face # For models like `Emotion`, which return a single prediction for a single face
predictions = predictions.reshape(1, -1) predictions = predictions.reshape(1, -1)