added review changes

deepface/models/facial_recognition/Buffalo_L.py

@@ -16,9 +16,7 @@ class Buffalo_L(FacialRecognition):
         self.load_model()
 
     def load_model(self):
-        """
-        Load the InsightFace Buffalo_L recognition model.
-        """
+        """Load the InsightFace Buffalo_L recognition model."""
         try:
             from insightface.model_zoo import get_model
         except Exception as err:
@@ -29,76 +27,71 @@ class Buffalo_L(FacialRecognition):
                 "albumentations"
             ) from err
 
-        # Define the model filename and subdirectory
         sub_dir = "buffalo_l"
-        model_file = "webface_r50.onnx"  # Corrected from w600k_r50.onnx per serengil's comment
+        model_file = "webface_r50.onnx"
         model_rel_path = os.path.join(sub_dir, model_file)
-
-        # Get the DeepFace home directory and construct weights path
         home = folder_utils.get_deepface_home()
         weights_dir = os.path.join(home, ".deepface", "weights")
         buffalo_l_dir = os.path.join(weights_dir, sub_dir)
 
-        # Ensure the buffalo_l subdirectory exists
         if not os.path.exists(buffalo_l_dir):
             os.makedirs(buffalo_l_dir, exist_ok=True)
             logger.info(f"Created directory: {buffalo_l_dir}")
 
-        # Download the model weights if not already present
         weights_path = weight_utils.download_weights_if_necessary(
             file_name=model_rel_path,
             source_url="https://drive.google.com/uc?export=download&confirm=pbef&id=1N0GL-8ehw_bz2eZQWz2b0A5XBdXdxZhg" #pylint: disable=line-too-long
         )
 
-        # Verify the model file exists
         if not os.path.exists(weights_path):
             raise FileNotFoundError(f"Model file not found at: {weights_path}")
         else:
             logger.debug(f"Model file found at: {weights_path}")
 
-        # Load the model using the full path
-        self.model = get_model(weights_path)  # Updated per serengil's feedback
+        self.model = get_model(weights_path)
         self.model.prepare(ctx_id=-1, input_size=self.input_shape)
 
     def preprocess(self, img: np.ndarray) -> np.ndarray:
         """
-        Preprocess the input image for the Buffalo_L model.
+        Preprocess the input image or batch of images.
 
         Args:
-            img: Input image as a numpy array of shape (112, 112, 3) or (1, 112, 112, 3).
+            img: Input image or batch with shape (112, 112, 3) 
+            or (batch_size, 112, 112, 3).
 
         Returns:
-            Preprocessed image as numpy array.
+            Preprocessed image(s) with RGB converted to BGR.
         """
-        # Ensure input is a single image
-        if len(img.shape) == 4:
-            if img.shape[0] == 1:
-                img = img[0]  # Squeeze batch dimension if it's a single-image batch
-            else:
-                raise ValueError("Buffalo_L model expects a single image, not a batch.")
-        elif len(img.shape) != 3 or img.shape != (112, 112, 3):
-            raise ValueError("Input image must have shape (112, 112, 3).")
+        if len(img.shape) == 3:
+            img = np.expand_dims(img, axis=0)  # Convert single image to batch of 1
+        elif len(img.shape) != 4:
+            raise ValueError("Input must have shape (112, 112, 3) or (batch_size, 112, 112, 3).")
 
-        # Convert RGB to BGR as required by InsightFace
-        img = img[:, :, ::-1]
+        # Convert RGB to BGR for the entire batch
+        img = img[:, :, :, ::-1]
         return img
 
     def forward(self, img: np.ndarray) -> Union[List[float], List[List[float]]]:
         """
-        Extract face embedding from a pre-cropped face image.
-        Args:
-            img: Preprocessed face image with shape (1, 112, 112, 3) 
-            or batch (batch_size, 112, 112, 3)
-        Returns:
-            Face embedding as a list of floats or list of lists of floats
-        """
-        img = self.preprocess(img)
-        embedding = self.model.get_feat(img)
-        if isinstance(embedding, np.ndarray) and len(embedding.shape) > 1:
-            embedding = embedding.flatten()
-        elif isinstance(embedding, list):
-            embedding = np.array(embedding).flatten()
-        else:
-            raise ValueError(f"Unexpected embedding type: {type(embedding)}")
-        return embedding.tolist()
+        Extract facial embedding(s) from the input image or batch of images.
 
+        Args:
+            img: Input image or batch with shape (112, 112, 3) 
+            or (batch_size, 112, 112, 3).
+
+        Returns:
+            Embedding as a list of floats (single image) 
+            or list of lists of floats (batch).
+        """
+        # Preprocess the input (single image or batch)
+        img = self.preprocess(img)
+        batch_size = img.shape[0]
+
+        # Handle both single images and batches
+        embeddings = []
+        for i in range(batch_size):
+            embedding = self.model.get_feat(img[i])
+            embeddings.append(embedding.flatten().tolist())
+
+        # Return single embedding if batch_size is 1, otherwise return list of embeddings
+        return embeddings[0] if batch_size == 1 else embeddings