data_loader.py

"""
Description : Set DataSet module for lip images
"""
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import os
import glob
from mxnet import nd
import mxnet.gluon.data.dataset as dataset
from mxnet.gluon.data.vision.datasets import image
from utils.align import Align

# pylint: disable=too-many-instance-attributes, too-many-arguments
class LipsDataset(dataset.Dataset):
    """
    Description : DataSet class for lip images
    """
    def __init__(self, root, align_root, flag=1,
                 mode='train', transform=None, seq_len=75):
        assert mode in ['train', 'valid']
        self._root = os.path.expanduser(root)
        self._align_root = align_root
        self._flag = flag
        self._transform = transform
        self._exts = ['.jpg', '.jpeg', '.png']
        self._seq_len = seq_len
        self._mode = mode
        self._list_images(self._root)

    def _list_images(self, root):
        """
        Description : generate list for lip images
        """
        self.labels = []
        self.items = []

        valid_unseen_sub_idx = [1, 2, 20, 22]
        skip_sub_idx = [21]

        if self._mode == 'train':
            sub_idx = ['s' + str(i) for i in range(1, 35) \
                             if i not in valid_unseen_sub_idx + skip_sub_idx]
        elif self._mode == 'valid':
            sub_idx = ['s' + str(i) for i in valid_unseen_sub_idx]

        folder_path = []
        for i in sub_idx:
            folder_path.extend(glob.glob(os.path.join(root, i, "*")))

        for folder in folder_path:
            filename = glob.glob(os.path.join(folder, "*"))
            if len(filename) != self._seq_len:
                continue
            filename.sort()
            label = os.path.split(folder)[-1]
            self.items.append((filename, label))

    def align_generation(self, file_nm, padding=75):
        """
        Description : Align to lip position
        """
        align = Align(self._align_root + '/' + file_nm + '.align')
        return nd.array(align.sentence(padding))

    def __getitem__(self, idx):
        img = list()
        for image_name in self.items[idx][0]:
            tmp_img = image.imread(image_name, self._flag)
            if self._transform is not None:
                tmp_img = self._transform(tmp_img)
            img.append(tmp_img)
        img = nd.stack(*img)
        img = nd.transpose(img, (1, 0, 2, 3))
        label = self.align_generation(self.items[idx][1],
                                      padding=self._seq_len)
        return img, label

    def __len__(self):
        return len(self.items)