diff --git a/deepface/models/facial_recognition/Buffalo_L.py b/deepface/models/facial_recognition/Buffalo_L.py index efe1524..ad59169 100644 --- a/deepface/models/facial_recognition/Buffalo_L.py +++ b/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. + Extract facial embedding(s) from the input image or batch of images. + Args: - img: Preprocessed face image with shape (1, 112, 112, 3) - or batch (batch_size, 112, 112, 3) + img: Input image or batch with shape (112, 112, 3) + or (batch_size, 112, 112, 3). + Returns: - Face embedding as a list of floats or list of lists of floats + 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) - 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() - \ No newline at end of file + 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 \ No newline at end of file