bobiptinyface.py 1.01 KB
Newer Older
Yu LINGHU's avatar
Yu LINGHU committed
1
2
from . import Base

Xinyi ZHANG's avatar
Xinyi ZHANG committed
3

Yu LINGHU's avatar
Yu LINGHU committed
4
5
class BobIpTinyface(Base):
    """Annotator using tinyface in bob.ip.facedetect"""
Xinyi ZHANG's avatar
Xinyi ZHANG committed
6

7
    def __init__(self, prob_thresh=0.5, **kwargs):
Yu LINGHU's avatar
Yu LINGHU committed
8
        super(BobIpTinyface, self).__init__(**kwargs)
9
10
        import bob.ip.facedetect.tinyface

11
12
13
14
15
16
17
        self.detector = bob.ip.facedetect.tinyface.TinyFacesDetector(
            prob_thresh=prob_thresh
        )

    @property
    def prob_thresh(self):
        return self.detector.prob_thresh
Yu LINGHU's avatar
Yu LINGHU committed
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

    def annotate(self, image, **kwargs):
        """Annotates an image using tinyface

        Parameters
        ----------
        image : numpy.array
            An RGB image in Bob format.
        **kwargs
            Ignored.

        Returns
        -------
        dict
            Annotations with (topleft, bottomright) keys (or None).
        """

35
36
37
        # return the annotations for the first/largest face
        annotations = self.detector.detect(image)

Yu LINGHU's avatar
Yu LINGHU committed
38
        if annotations is not None:
Xinyi ZHANG's avatar
2    
Xinyi ZHANG committed
39
            return annotations[0]
Yu LINGHU's avatar
Yu LINGHU committed
40
        else:
Tiago de Freitas Pereira's avatar
Tiago de Freitas Pereira committed
41
            return None