Spaces:
Runtime error
Runtime error
File size: 5,171 Bytes
e0f2d0e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
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
|