diff --git a/README.md b/README.md index ac1d8e3..62fd4d1 100644 --- a/README.md +++ b/README.md @@ -1,47 +1,58 @@ -# Retinaface Face Detection (Pytorch) +# RetinaFace Face Detection *Authors: David Wang* + + ## Desription -This opertator detects faces in the images by using RetinaFace Detector[1]. It will returns the bounding box positions and the confidence scores of detected faces. This repo is a adopataion from [2]. +This operator detects faces in the images by using RetinaFace Detector[1]. It will return the bounding box positions and the confidence scores of detected faces. This repo is a adapataion from [2]. +## Code Example + +Load an image from path './dog.jpg' +and use the pretrained RetinaFace to generate face bounding boxes. + + *Write the pipeline in simplified style*: ```python from towhee import ops -model = ops.face_detection.retinaface() -embedding = model(img) +bboxes = dc.glob('nmb46.jpg') \ + .image_decode.cv2() \ + .face_detection.retinaface() \ + .to_list() ``` +*Write a same pipeline with explicit inputs/outputs name specifications:* + +```python +from towhee import dc +dc.glob['path']('./dog.jpg') \ + .image_decode.cv2['path', 'img']() \ + .image_embedding.timm['img', 'vec'](model_name='resnet50') \ + .face_detection.retinaface() \ + .to_list() +``` ## Factory Constructor -Create the operator via the following factory method +Create the operator via the following factory method. ***ops.face_detection.retinaface()*** - ## Interface -A face detection operator takes an image as input. it generates the bounding box position and confidence score back to ndarray. - -**Args:** - -​ ***framework*** - -​ the framework of the model - -​ supported types: `str`, default is 'pytorch' +A face detection operator takes an image as input. it generates the bounding box positions and confidence scores back to ndarray. **Parameters:** -​ ***image*** +​ ***image***: *numpy.ndarray.* ​ the image to detect faces. ​ supported types: numpy.ndarray -**Returns:**: +**Returns:** *numpy.ndarray* @@ -51,33 +62,6 @@ A face detection operator takes an image as input. it generates the bounding box ​ The detected face bounding boxes confident scores. -## Code Example - -get detected face bounding boxes from './img1.jpg'. - - *Write the pipeline in simplified style*: - -```python -import towhee.DataCollection as dc - -dc.glob('./img1.jpg') - .image_decode.cv2() - .face_detection.retinaface() - .to_list() -``` - -*Write a same pipeline with explicit inputs/outputs name specifications:* - -```python -import towhee.DataCollection as dc - -dc.glob['path']('./img1.jpg') - .image_decode.cv2['path', 'img']() - .face_detection.retinaface() - .to_list() -``` - - ## Reference [1]. https://arxiv.org/abs/1905.00641 diff --git a/__init__.py b/__init__.py index 98ce600..feafe9d 100644 --- a/__init__.py +++ b/__init__.py @@ -14,6 +14,6 @@ from .retinaface import Retinaface -def retinaface(): - return Retinaface() +def retinaface(framework: str = 'pytorch'): + return Retinaface(framework) diff --git a/retinaface.py b/retinaface.py index aab7280..889b5ce 100644 --- a/retinaface.py +++ b/retinaface.py @@ -25,25 +25,23 @@ import numpy from towhee import register from towhee.operator import Operator from towhee.types.image_utils import to_pil -from towhee._types import Image -from towhee.types import arg, to_image_color +from towhee._types import Image +from towhee.types.arg import arg, to_image_color -from timm.data import resolve_data_config -from timm.data.transforms_factory import create_transform -@register(output_schema=['bboxes', 'scores']) +@register(output_schema=['bbox', 'score']) class Retinaface(Operator): """ Retinaface """ - def __init__(self, framework: str = 'pytorch') -> None: + def __init__(self) -> None: super().__init__() sys.path.append(str(Path(__file__).parent)) from retinaface_impl import Model self.model = Model() @arg(1, to_image_color('RGB') ) - def __call__(self, image: 'towhee._types.Image'): + def __call__(self, image: Image): img = torch.FloatTensor(numpy.asarray(to_pil(image))) bboxes, keypoints = self.model(img) bboxes = bboxes.cpu().detach().numpy()