mirror of
https://github.com/serengil/deepface.git
synced 2025-06-05 19:15:23 +00:00
review changes
This commit is contained in:
parent
e19c7fcc1c
commit
2dbae0edca
@ -10,7 +10,7 @@ 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
|
||||
|
||||
@ -18,21 +18,43 @@ 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.output_shape = 512 # Embedding size
|
||||
self.load_model()
|
||||
|
||||
def load_model(self):
|
||||
"""
|
||||
Load the InsightFace Buffalo_L recognition model.
|
||||
"""
|
||||
# Use DeepFace's utility to download weights if necessary
|
||||
model_rel_path = os.path.join("insightface", "buffalo_l", "w600k_r50.onnx")
|
||||
# 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 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)
|
||||
if not os.path.exists(buffalo_l_dir):
|
||||
os.makedirs(buffalo_l_dir, exist_ok=True)
|
||||
print("Created directory:", buffalo_l_dir)
|
||||
|
||||
# Download the model weights if not already present
|
||||
weights_path = weight_utils.download_weights_if_necessary(
|
||||
file_name="webface_r50.onnx",
|
||||
file_name=model_rel_path,
|
||||
source_url="https://drive.google.com/uc?export=download&confirm=pbef&id=1N0GL-8ehw_bz2eZQWz2b0A5XBdXdxZhg"
|
||||
)
|
||||
# Load model from weights folder
|
||||
self.model = get_model("buffalo_l/w600k_r50.onnx", root=os.path.dirname(weights_path))
|
||||
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)
|
||||
else:
|
||||
print("Model file NOT found at expected location:", expected_model_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)
|
||||
self.model.prepare(ctx_id=-1, input_size=self.input_shape)
|
||||
|
||||
def preprocess(self, img: np.ndarray) -> np.ndarray:
|
||||
@ -47,7 +69,7 @@ class Buffalo_L(FacialRecognition):
|
||||
img = img[0] # Remove batch dimension
|
||||
if img.max() <= 1.0: # If normalized to [0, 1]
|
||||
img = (img * 255).astype(np.uint8)
|
||||
# Always convert RGB to BGR (DeepFace outputs RGB, InsightFace expects BGR)
|
||||
# Convert RGB to BGR (DeepFace outputs RGB, InsightFace expects BGR)
|
||||
img = img[:, :, ::-1]
|
||||
return img
|
||||
|
||||
|
Loading…
x
Reference in New Issue
Block a user