diff --git a/README.md b/README.md index 5876f78..3adedf6 100644 --- a/README.md +++ b/README.md @@ -6,12 +6,11 @@ ## Desription -This operator detects faces in the images by using [RetinaFace](https://arxiv.org/abs/1905.00641) Detector. It will return the bounding box positions and the confidence scores of detected faces. This repo is a adapataion from [repo](https://github.com/biubug6/Pytorch_Retinaface). +This operator detects faces in the images by using [RetinaFace](https://arxiv.org/abs/1905.00641) Detector. It will return the bounding box positions and the confidence scores of detected faces. This repo is an adaptaion from [biubug6/Pytorch_Retinaface](https://github.com/biubug6/Pytorch_Retinaface). ## Code Example -Load an image from path './dog.jpg' -and use the pretrained RetinaFace to generate face bounding boxes. +Load an image from path './turing.png' and use the pretrained RetinaFace model to generate face bounding boxes and confidence scores. *Write the pipeline in simplified style*: @@ -34,7 +33,8 @@ import towhee towhee.glob['path']('turing.png') \ .image_decode.cv2['path', 'img']() \ .face_detection.retinaface['img', ('bbox','score')]() \ - .select('img', 'bbox', 'score').show() + .select('img', 'bbox', 'score') \ + .show() ``` result2 @@ -43,15 +43,15 @@ towhee.glob['path']('turing.png') \ Create the operator via the following factory method. -***ops.face_detection.retinaface()*** +***face_detection.retinaface()*** ## Interface -A face detection operator takes an image as input. it generates the bounding box positions and confidence scores back to ndarray. +A face detection operator takes an image as input. It generates the bounding box positions and confidence scores back to ndarray. **Parameters:** -​ ***image***: *numpy.ndarray.* +​ ***img***: *numpy.ndarray.* ​ the image to detect faces. @@ -59,11 +59,11 @@ A face detection operator takes an image as input. it generates the bounding box **Returns:** -*numpy.ndarray* +*List[(int, int, int, int)]* ​ The detected face bounding boxes. -*numpy.ndarray* +*List[float]* ​ The detected face bounding boxes confident scores. diff --git a/result1.png b/result1.png index f9e17d7..e29b9f8 100644 Binary files a/result1.png and b/result1.png differ diff --git a/result2.png b/result2.png index 90a6212..dae2e9f 100644 Binary files a/result2.png and b/result2.png differ diff --git a/retinaface.py b/retinaface.py index 889b5ce..11561af 100644 --- a/retinaface.py +++ b/retinaface.py @@ -29,7 +29,7 @@ from towhee._types import Image from towhee.types.arg import arg, to_image_color -@register(output_schema=['bbox', 'score']) +@register(output_schema=['box', 'score']) class Retinaface(Operator): """ Retinaface @@ -41,8 +41,11 @@ class Retinaface(Operator): self.model = Model() @arg(1, to_image_color('RGB') ) - def __call__(self, image: Image): - img = torch.FloatTensor(numpy.asarray(to_pil(image))) + def __call__(self, img: Image): + img = torch.FloatTensor(numpy.asarray(to_pil(img))) bboxes, keypoints = self.model(img) bboxes = bboxes.cpu().detach().numpy() - return bboxes[:,:4].astype(int), bboxes[:,4] + bl, sl = bboxes[:,:4].astype(int), bboxes[:,4] + rbboxes = [tuple(box) for box in bl] + return rbboxes, sl.tolist() +