Spaces:
Runtime error
Runtime error
| from scrfd import SCRFD, Threshold, Face | |
| from PIL.Image import Image | |
| import numpy as np | |
| from typing import List, Union | |
| from pathlib import Path | |
| import os | |
| from schemas.vision_schemas import FaceMainPoints, FaceDetector | |
| from utils.utils import open_image | |
| class SCRFDFaceDetector(FaceDetector): | |
| """ | |
| A face detector using the SCRFD model. | |
| """ | |
| def __init__(self, model_path: str, threshold_probability: float, nms:float): | |
| """ | |
| Initializes the face detector. | |
| Args: | |
| model_path: The path to the SCRFD model. | |
| threshold_probability: The probability threshold for face detection. | |
| nms: The NMS threshold for face detection. | |
| """ | |
| self.model = SCRFD.from_path(model_path) | |
| self.threshold = Threshold(probability=threshold_probability, nms=nms) | |
| def convert_face_to_face_main_points(self, face: Face) -> FaceMainPoints: | |
| """ | |
| Converts a Face object to a FaceMainPoints object. | |
| Args: | |
| face: The face to convert. | |
| Returns: | |
| A FaceMainPoints object. | |
| """ | |
| return FaceMainPoints( | |
| box_start_point=(face.bbox.upper_left.x, face.bbox.upper_left.y,), | |
| box_end_point=(face.bbox.lower_right.x, face.bbox.lower_right.y), | |
| box_probabilty_score=face.probability, | |
| left_eye=(face.keypoints.left_eye.x, face.keypoints.left_eye.y), | |
| right_eye=(face.keypoints.right_eye.x, face.keypoints.right_eye.y), | |
| nose=(face.keypoints.nose.x, face.keypoints.nose.y), | |
| left_mouth=(face.keypoints.left_mouth.x, face.keypoints.left_mouth.y), | |
| right_mouth=(face.keypoints.right_mouth.x, face.keypoints.right_mouth.y) | |
| ) | |
| def detect(self, image: Union[str, Path, np.ndarray, Image]) -> List[FaceMainPoints]: | |
| """ | |
| Detect faces in an image and return their main points. | |
| Args: | |
| image: Input image that can be: | |
| - str: Path to image file | |
| - Path: Path object to image file | |
| - np.ndarray: Image array (BGR or RGB) | |
| - PILImage.Image: PIL Image object | |
| Returns: | |
| List[FaceMainPoints]: List of face detections, each containing: | |
| - Bounding box coordinates | |
| - Key facial landmarks (eyes, nose, mouth) | |
| - Confidence score | |
| Note: | |
| - Uses a threshold set during initialization for detection confidence | |
| - Returns an empty list if no faces are detected | |
| - Image is automatically converted to RGB format | |
| """ | |
| image = open_image(image).convert("RGB") | |
| extracted_faces = self.model.detect(image, threshold=self.threshold) | |
| faces = [self.convert_face_to_face_main_points(extracted_face) for extracted_face in extracted_faces] | |
| return faces | |
| def extract_face_from_image(self, img, face_main_points:FaceMainPoints): | |
| """ | |
| Extract a face from an image using FaceMainPoints coordinates. | |
| Args: | |
| img: Input PIL Image | |
| face_main_points: FaceMainPoints object containing face coordinates | |
| Returns: | |
| PIL.Image: Cropped face image | |
| """ | |
| start_x, start_y = face_main_points.box_start_point | |
| end_x, end_y = face_main_points.box_end_point | |
| return img.crop((start_x, start_y, end_x, end_y)) | |
| def cut_extracted_faces(self, image: Union[str, Path, np.ndarray, Image], | |
| save_path: str = None) -> List[Image]: | |
| """ | |
| Detect and extract all faces from an image, optionally saving them. | |
| Args: | |
| image: Input image that can be: | |
| - str: Path to image file | |
| - Path: Path object to image file | |
| - np.ndarray: Image array (BGR or RGB) | |
| - PILImage.Image: PIL Image object | |
| save_path: Optional path to save extracted faces | |
| If provided, faces will be saved as 'face_{index}.jpg' in this directory | |
| Returns: | |
| List[PIL.Image]: List of cropped face images | |
| """ | |
| # Convert input to PIL Image | |
| img = open_image(image) | |
| # Detect faces and get their main points | |
| faces_information = self.detect(img) | |
| if not faces_information: | |
| return [] # Return empty list if no faces detected | |
| # Extract and optionally save each face | |
| extracted_faces = [] | |
| for i, face_info in enumerate(faces_information): | |
| face_image = self.extract_face_from_image(img, face_info) | |
| extracted_faces.append(face_image) | |
| # Save face if save_path is provided | |
| if save_path: | |
| os.makedirs(save_path, exist_ok=True) | |
| face_path = os.path.join(save_path, f'face_{i}.jpg') | |
| face_image.save(face_path) | |
| return extracted_faces | |