added review changes

This commit is contained in:
Raghucharan16 2025-02-26 23:41:16 +05:30
parent 2dbae0edca
commit 655a221650

View File

@ -1,60 +1,59 @@
import os import os
import numpy as np import numpy as np
from deepface.commons import weight_utils from typing import List
from deepface.commons.logger import Logger
from deepface.models.FacialRecognition import FacialRecognition
logger = Logger()
try: try:
from insightface.model_zoo import get_model from insightface.model_zoo import get_model
except ModuleNotFoundError as err: except ModuleNotFoundError as err:
raise ModuleNotFoundError( 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" "You can install it with: pip install insightface>=0.7.3"
) from err ) 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): class Buffalo_L(FacialRecognition):
def __init__(self): def __init__(self):
self.model = None self.model = None
self.input_shape = (112, 112) # Buffalo_L expects 112x112 self.input_shape = (112, 112)
self.output_shape = 512 # Embedding size self.output_shape = 512
self.load_model() self.load_model()
def load_model(self): def load_model(self):
""" """
Load the InsightFace Buffalo_L recognition model. Load the InsightFace Buffalo_L recognition model.
""" """
# Define the model filename and subdirectory # Define the relative model path
sub_dir = "buffalo_l" model_rel_path = os.path.join("buffalo_l", "w600k_r50.onnx")
model_file = "w600k_r50.onnx"
model_rel_path = os.path.join(sub_dir, model_file)
# Define the weights directory and ensure the buffalo_l subdirectory exists # Get the DeepFace weights directory
weights_dir = os.path.join(os.path.expanduser("~"), ".deepface", "weights") home = folder_utils.get_deepface_home()
buffalo_l_dir = os.path.join(weights_dir, sub_dir) 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): if not os.path.exists(buffalo_l_dir):
os.makedirs(buffalo_l_dir, exist_ok=True) 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( weights_path = weight_utils.download_weights_if_necessary(
file_name=model_rel_path, file_name=model_rel_path,
source_url="https://drive.google.com/uc?export=download&confirm=pbef&id=1N0GL-8ehw_bz2eZQWz2b0A5XBdXdxZhg" 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 # Verify the model file exists
expected_model_path = os.path.join(buffalo_l_dir, model_file) if os.path.exists(weights_path):
if os.path.exists(expected_model_path): logger.debug(f"Model file found at: {weights_path}")
print("Model file found at expected location:", expected_model_path)
else: 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 # Load the model using the full absolute path
full_model_path = os.path.join(buffalo_l_dir, model_file) self.model = get_model(weights_path)
print("Full model path:", full_model_path)
self.model = get_model(full_model_path)
self.model.prepare(ctx_id=-1, input_size=self.input_shape) self.model.prepare(ctx_id=-1, input_size=self.input_shape)
def preprocess(self, img: np.ndarray) -> np.ndarray: def preprocess(self, img: np.ndarray) -> np.ndarray:
@ -65,15 +64,14 @@ class Buffalo_L(FacialRecognition):
Returns: Returns:
Preprocessed image as numpy array Preprocessed image as numpy array
""" """
if len(img.shape) == 4: # (1, 112, 112, 3) if len(img.shape) == 4:
img = img[0] # Remove batch dimension img = img[0]
if img.max() <= 1.0: # If normalized to [0, 1] if img.max() <= 1.0:
img = (img * 255).astype(np.uint8) img = (img * 255).astype(np.uint8)
# Convert RGB to BGR (DeepFace outputs RGB, InsightFace expects BGR)
img = img[:, :, ::-1] img = img[:, :, ::-1]
return img 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. Extract face embedding from a pre-cropped face image.
Args: Args:
@ -84,7 +82,6 @@ class Buffalo_L(FacialRecognition):
img = self.preprocess(img) img = self.preprocess(img)
embedding = self.model.get_feat(img) embedding = self.model.get_feat(img)
# Handle different embedding formats
if isinstance(embedding, np.ndarray) and len(embedding.shape) > 1: if isinstance(embedding, np.ndarray) and len(embedding.shape) > 1:
embedding = embedding.flatten() embedding = embedding.flatten()
elif isinstance(embedding, list): elif isinstance(embedding, list):
@ -92,4 +89,4 @@ class Buffalo_L(FacialRecognition):
else: else:
raise ValueError(f"Unexpected embedding type: {type(embedding)}") raise ValueError(f"Unexpected embedding type: {type(embedding)}")
return embedding.tolist() # Convert to list per FacialRecognition spec return embedding.tolist()