From 655a22165039a1acde18fea95ea9c4a522ef788b Mon Sep 17 00:00:00 2001 From: Raghucharan16 Date: Wed, 26 Feb 2025 23:41:16 +0530 Subject: [PATCH] added review changes --- .../models/facial_recognition/Buffalo_L.py | 65 +++++++++---------- 1 file changed, 31 insertions(+), 34 deletions(-) diff --git a/deepface/models/facial_recognition/Buffalo_L.py b/deepface/models/facial_recognition/Buffalo_L.py index 9b6b8d2..27e4e8a 100644 --- a/deepface/models/facial_recognition/Buffalo_L.py +++ b/deepface/models/facial_recognition/Buffalo_L.py @@ -1,60 +1,59 @@ import os import numpy as np -from deepface.commons import weight_utils -from deepface.commons.logger import Logger -from deepface.models.FacialRecognition import FacialRecognition - -logger = Logger() +from typing import List try: from insightface.model_zoo import get_model except ModuleNotFoundError as err: raise ModuleNotFoundError( - "InsightFace is an optional dependency for the Buffalo_L model. " + "InsightFace is an optional dependency for the Buffalo_L model." "You can install it with: pip install insightface>=0.7.3" ) from err +from deepface.commons import weight_utils, folder_utils +from deepface.commons.logger import Logger +from deepface.models.FacialRecognition import FacialRecognition + +logger = Logger() + class Buffalo_L(FacialRecognition): def __init__(self): self.model = None - self.input_shape = (112, 112) # Buffalo_L expects 112x112 - self.output_shape = 512 # Embedding size + self.input_shape = (112, 112) + self.output_shape = 512 self.load_model() def load_model(self): """ Load the InsightFace Buffalo_L recognition model. """ - # Define the model filename and subdirectory - sub_dir = "buffalo_l" - model_file = "w600k_r50.onnx" - model_rel_path = os.path.join(sub_dir, model_file) + # Define the relative model path + model_rel_path = os.path.join("buffalo_l", "w600k_r50.onnx") - # Define the weights directory and ensure the buffalo_l subdirectory exists - weights_dir = os.path.join(os.path.expanduser("~"), ".deepface", "weights") - buffalo_l_dir = os.path.join(weights_dir, sub_dir) + # Get the DeepFace weights directory + home = folder_utils.get_deepface_home() + weights_dir = os.path.join(home, ".deepface", "weights") + buffalo_l_dir = os.path.join(weights_dir, "buffalo_l") + + # Ensure the buffalo_l subdirectory exists if not os.path.exists(buffalo_l_dir): os.makedirs(buffalo_l_dir, exist_ok=True) - print("Created directory:", buffalo_l_dir) + logger.info(f"Created directory: {buffalo_l_dir}") - # Download the model weights if not already present + # Download the model weights 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" ) - print("Downloaded model path:", weights_path) - # Verify that the model file exists at the expected location - expected_model_path = os.path.join(buffalo_l_dir, model_file) - if os.path.exists(expected_model_path): - print("Model file found at expected location:", expected_model_path) + # Verify the model file exists + if os.path.exists(weights_path): + logger.debug(f"Model file found at: {weights_path}") else: - print("Model file NOT found at expected location:", expected_model_path) + logger.debug(f"Model file NOT found at: {weights_path}") - # Use the full absolute path for loading the model - full_model_path = os.path.join(buffalo_l_dir, model_file) - print("Full model path:", full_model_path) - self.model = get_model(full_model_path) + # Load the model using the full absolute 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: @@ -65,15 +64,14 @@ class Buffalo_L(FacialRecognition): Returns: Preprocessed image as numpy array """ - if len(img.shape) == 4: # (1, 112, 112, 3) - img = img[0] # Remove batch dimension - if img.max() <= 1.0: # If normalized to [0, 1] + if len(img.shape) == 4: + img = img[0] + if img.max() <= 1.0: img = (img * 255).astype(np.uint8) - # Convert RGB to BGR (DeepFace outputs RGB, InsightFace expects BGR) img = img[:, :, ::-1] return img - def forward(self, img: np.ndarray) -> list[float]: + def forward(self, img: np.ndarray) -> List[float]: """ Extract face embedding from a pre-cropped face image. Args: @@ -84,7 +82,6 @@ class Buffalo_L(FacialRecognition): img = self.preprocess(img) embedding = self.model.get_feat(img) - # Handle different embedding formats if isinstance(embedding, np.ndarray) and len(embedding.shape) > 1: embedding = embedding.flatten() elif isinstance(embedding, list): @@ -92,4 +89,4 @@ class Buffalo_L(FacialRecognition): else: raise ValueError(f"Unexpected embedding type: {type(embedding)}") - return embedding.tolist() # Convert to list per FacialRecognition spec \ No newline at end of file + return embedding.tolist() \ No newline at end of file