Patch.py 2.88 KB
Newer Older
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
1
2
from bob.bio.base.annotator.FailSafe import translate_kwargs
from sklearn.base import BaseEstimator, TransformerMixin
3
from ..utils import extract_patches
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
4
from bob.bio.video import VideoLikeContainer
5
6
7
from collections import OrderedDict


Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
8
9
class ImagePatches(TransformerMixin, BaseEstimator):
    """Extracts patches of images and returns it in a VideoLikeContainer. You need
10
11
12
13
    to wrap the further blocks (extractor and algorithm) that come after this
    in bob.bio.video wrappers.
    """

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
14
15
16
    def __init__(
        self, block_size, block_overlap=(0, 0), n_random_patches=None, **kwargs
    ):
17
18
19
20
21
        super(ImagePatches, self).__init__(**kwargs)
        self.block_size = block_size
        self.block_overlap = block_overlap
        self.n_random_patches = n_random_patches

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
22
23
    def transform(self, images):
        return [self.transform_one_image(img) for img in images]
24

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
25
    def transform_one_image(self, image):
26

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
27
28
29
30
        patches = extract_patches(
            image, self.block_size, self.block_overlap, self.n_random_patches
        )
        vc = VideoLikeContainer(patches, range(len(patches)))
31

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
32
        return vc
33
34


Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
35
class VideoPatches(TransformerMixin, BaseEstimator):
36
    """Extracts patches of images from video containers and returns it in a
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
37
    VideoLikeContainer.
38
39
    """

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
40
41
42
43
44
45
46
47
48
    def __init__(
        self,
        face_cropper,
        block_size,
        block_overlap=(0, 0),
        n_random_patches=None,
        normalizer=None,
        **kwargs,
    ):
49
        super(VideoPatches, self).__init__(**kwargs)
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
50
        self.face_cropper = face_cropper
51
52
53
54
55
        self.block_size = block_size
        self.block_overlap = block_overlap
        self.n_random_patches = n_random_patches
        self.normalizer = normalizer

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
56
57
58
    def transform(self, videos, annotations=None):
        kwargs = translate_kwargs(dict(annotations=annotations), len(videos))
        return [self.transform_one_video(vid, **kw) for vid, kw in zip(videos, kwargs)]
59

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
60
61
    def transform_one_video(self, frames, annotations=None):
        annotations = annotations or {}
62
63
64
        if self.normalizer is not None:
            annotations = OrderedDict(self.normalizer(annotations))

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
65
66
        all_patches = []
        for frame, index in zip(frames, frames.indices):
67
68
69

            # if annotations are given, and if particular frame annotations are
            # not missing we take them:
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
70
            annots = annotations.get(str(index))
71

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
72
73
            # preprocess image (by default: crops a face)
            preprocessed = self.face_cropper(frame, annots)
74
75
76
77
78
            if preprocessed is None:
                continue

            # extract patches
            patches = extract_patches(
Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
79
80
81
82
83
                preprocessed, self.block_size, self.block_overlap, self.n_random_patches
            )
            all_patches.extend(patches)

        vc = VideoLikeContainer(all_patches, range(len(all_patches)))
84

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
85
        if not len(vc):
86
87
            return None

Amir MOHAMMADI's avatar
Amir MOHAMMADI committed
88
        return vc