From 527a89c9318079ef260ba7a511efa8ffe2ff5441 Mon Sep 17 00:00:00 2001
From: raghucharan16 <narravenkataraghucharan@gmail.com>
Date: Sat, 1 Mar 2025 12:18:09 +0530
Subject: [PATCH] review chanages

---
 .../models/facial_recognition/Buffalo_L.py    | 52 +++++++++++--------
 requirements_additional.txt                   |  3 +-
 2 files changed, 31 insertions(+), 24 deletions(-)

diff --git a/deepface/models/facial_recognition/Buffalo_L.py b/deepface/models/facial_recognition/Buffalo_L.py
index bf44379..7e4b0d5 100644
--- a/deepface/models/facial_recognition/Buffalo_L.py
+++ b/deepface/models/facial_recognition/Buffalo_L.py
@@ -1,5 +1,5 @@
 import os
-from typing import List
+from typing import List, Union
 import numpy as np
 
 from deepface.commons import weight_utils, folder_utils
@@ -25,12 +25,12 @@ class Buffalo_L(FacialRecognition):
             raise ModuleNotFoundError(
                 "InsightFace and its dependencies are optional for the Buffalo_L model. "
                 "Please install them with: "
-                "pip install insightface>=0.7.3 onnxruntime>=1.9.0 typing-extensions pydantic"
+                "pip install insightface>=0.7.3 onnxruntime>=1.9.0 typing-extensions pydantic albumentations"
             ) from err
 
         # Define the model filename and subdirectory
         sub_dir = "buffalo_l"
-        model_file = "w600k_r50.onnx"
+        model_file = "webface_r50.onnx"
         model_rel_path = os.path.join(sub_dir, model_file)
 
         # Get the DeepFace home directory and construct weights path
@@ -46,7 +46,7 @@ class Buffalo_L(FacialRecognition):
         # 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
+            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
@@ -55,20 +55,27 @@ class Buffalo_L(FacialRecognition):
         else:
             raise FileNotFoundError(f"Model file not found at: {weights_path}")
 
-        # Load the model
-        self.model = get_model(model_file, root=buffalo_l_dir)
+        # Load the model using the full path
+        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 image to match InsightFace recognition model expectations.
+        Preprocess the image or batch of images to match InsightFace recognition model expectations.
         Args:
-            img: Image in shape (1, 112, 112, 3) or (112, 112, 3)
+            img: Image in shape (1, 112, 112, 3) or (112, 112, 3) or batch (batch_size, 112, 112, 3)
         Returns:
-            Preprocessed image as numpy array
+            Preprocessed image or batch as numpy array
         """
-        if len(img.shape) == 4:
-            img = img[0]
+        if len(img.shape) == 4:  # Batch of images
+            preprocessed_imgs = []
+            for i in range(img.shape[0]):
+                single_img = img[i]
+                if single_img.max() <= 1.0:
+                    single_img = (single_img * 255).astype(np.uint8)
+                single_img = single_img[:, :, ::-1]  # Convert RGB to BGR
+                preprocessed_imgs.append(single_img)
+            return np.array(preprocessed_imgs)
         if len(img.shape) != 3:
             raise ValueError(
                 f"Expected image to be 3D after preprocessing, but got shape: {img.shape}")
@@ -77,21 +84,20 @@ class Buffalo_L(FacialRecognition):
         img = img[:, :, ::-1]  # Convert RGB to BGR
         return img
 
-    def forward(self, img: np.ndarray) -> List[float]:
+    def forward(self, img: np.ndarray) -> Union[List[float], List[List[float]]]:
         """
-        Extract face embedding from a pre-cropped face image.
+        Extract face embedding from a pre-cropped face image or batch of images.
         Args:
-            img: Preprocessed face image with shape (1, 112, 112, 3)
+            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
+            Face embedding as a list of floats (single image) or list of lists of floats (batch)
         """
         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()
+        if len(img.shape) == 4:  # Batch
+            embeddings = self.model.get_feat(img)
+            return [embedding.tolist() for embedding in embeddings]
+        elif len(img.shape) == 3:  # Single image
+            embedding = self.model.get_feat(np.expand_dims(img, axis=0))[0]
+            return embedding.tolist()
         else:
-            raise ValueError(f"Unexpected embedding type: {type(embedding)}")
-        return embedding.tolist()
-    
\ No newline at end of file
+            raise ValueError(f"Unexpected embedding type after preprocessing: {img.shape}")
\ No newline at end of file
diff --git a/requirements_additional.txt b/requirements_additional.txt
index dd30b95..cc38a19 100644
--- a/requirements_additional.txt
+++ b/requirements_additional.txt
@@ -8,4 +8,5 @@ insightface>=0.7.3
 onnxruntime>=1.9.0
 tf-keras
 typing-extensions
-pydantic
\ No newline at end of file
+pydantic
+albumentations
\ No newline at end of file