mirror of
https://github.com/serengil/deepface.git
synced 2025-06-07 12:05:22 +00:00
added review changes
This commit is contained in:
parent
2dbae0edca
commit
655a221650
@ -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()
|
Loading…
x
Reference in New Issue
Block a user