All

GANimation/data/custom_dataset_data_loader.py

@@ -0,0 +1,25 @@
+import torch.utils.data
+from data.dataset import DatasetFactory
+
+
+class CustomDatasetDataLoader:
+    def __init__(self, opt, is_for_train=True):
+        self._opt = opt
+        self._is_for_train = is_for_train
+        self._num_threds = opt.n_threads_train if is_for_train else opt.n_threads_test
+        self._create_dataset()
+
+    def _create_dataset(self):
+        self._dataset = DatasetFactory.get_by_name(self._opt.dataset_mode, self._opt, self._is_for_train)
+        self._dataloader = torch.utils.data.DataLoader(
+            self._dataset,
+            batch_size=self._opt.batch_size,
+            shuffle=not self._opt.serial_batches,
+            num_workers=int(self._num_threds),
+            drop_last=True)
+
+    def load_data(self):
+        return self._dataloader
+
+    def __len__(self):
+        return len(self._dataset)

GANimation/data/dataset.py

@@ -0,0 +1,68 @@
+import torch.utils.data as data
+from PIL import Image
+import torchvision.transforms as transforms
+import os
+import os.path
+
+
+class DatasetFactory:
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_by_name(dataset_name, opt, is_for_train):
+        if dataset_name == 'aus':
+            from data.dataset_aus import AusDataset
+            dataset = AusDataset(opt, is_for_train)
+        else:
+            raise ValueError("Dataset [%s] not recognized." % dataset_name)
+
+        print('Dataset {} was created'.format(dataset.name))
+        return dataset
+
+
+class DatasetBase(data.Dataset):
+    def __init__(self, opt, is_for_train):
+        super(DatasetBase, self).__init__()
+        self._name = 'BaseDataset'
+        self._root = None
+        self._opt = opt
+        self._is_for_train = is_for_train
+        self._create_transform()
+
+        self._IMG_EXTENSIONS = [
+            '.jpg', '.JPG', '.jpeg', '.JPEG',
+            '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
+        ]
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def path(self):
+        return self._root
+
+    def _create_transform(self):
+        self._transform = transforms.Compose([])
+
+    def get_transform(self):
+        return self._transform
+
+    def _is_image_file(self, filename):
+        return any(filename.endswith(extension) for extension in self._IMG_EXTENSIONS)
+
+    def _is_csv_file(self, filename):
+        return filename.endswith('.csv')
+
+    def _get_all_files_in_subfolders(self, dir, is_file):
+        images = []
+        assert os.path.isdir(dir), '%s is not a valid directory' % dir
+
+        for root, _, fnames in sorted(os.walk(dir)):
+            for fname in fnames:
+                if is_file(fname):
+                    path = os.path.join(root, fname)
+                    images.append(path)
+
+        return images

GANimation/data/dataset_aus.py

@@ -0,0 +1,117 @@
+import os.path
+import torchvision.transforms as transforms
+from data.dataset import DatasetBase
+from PIL import Image
+import random
+import numpy as np
+import pickle
+from utils import cv_utils
+
+
+class AusDataset(DatasetBase):
+    def __init__(self, opt, is_for_train):
+        super(AusDataset, self).__init__(opt, is_for_train)
+        self._name = 'AusDataset'
+
+        # read dataset
+        self._read_dataset_paths()
+
+    def __getitem__(self, index):
+        assert (index < self._dataset_size)
+
+        # start_time = time.time()
+        real_img = None
+        real_cond = None
+        while real_img is None or real_cond is None:
+            # if sample randomly: overwrite index
+            if not self._opt.serial_batches:
+                index = random.randint(0, self._dataset_size - 1)
+
+            # get sample data
+            sample_id = self._ids[index]
+
+            real_img, real_img_path = self._get_img_by_id(sample_id)
+            real_cond = self._get_cond_by_id(sample_id)
+
+            if real_img is None:
+                print 'error reading image %s, skipping sample' % sample_id
+            if real_cond is None:
+                print 'error reading aus %s, skipping sample' % sample_id
+
+        desired_cond = self._generate_random_cond()
+
+        # transform data
+        img = self._transform(Image.fromarray(real_img))
+
+        # pack data
+        sample = {'real_img': img,
+                  'real_cond': real_cond,
+                  'desired_cond': desired_cond,
+                  'sample_id': sample_id,
+                  'real_img_path': real_img_path
+                  }
+
+        # print (time.time() - start_time)
+
+        return sample
+
+    def __len__(self):
+        return self._dataset_size
+
+    def _read_dataset_paths(self):
+        self._root = self._opt.data_dir
+        self._imgs_dir = os.path.join(self._root, self._opt.images_folder)
+
+        # read ids
+        use_ids_filename = self._opt.train_ids_file if self._is_for_train else self._opt.test_ids_file
+        use_ids_filepath = os.path.join(self._root, use_ids_filename)
+        self._ids = self._read_ids(use_ids_filepath)
+
+        # read aus
+        conds_filepath = os.path.join(self._root, self._opt.aus_file)
+        self._conds = self._read_conds(conds_filepath)
+
+        self._ids = list(set(self._ids).intersection(set(self._conds.keys())))
+
+        # dataset size
+        self._dataset_size = len(self._ids)
+
+    def _create_transform(self):
+        if self._is_for_train:
+            transform_list = [transforms.RandomHorizontalFlip(),
+                              transforms.ToTensor(),
+                              transforms.Normalize(mean=[0.5, 0.5, 0.5],
+                                                   std=[0.5, 0.5, 0.5]),
+                              ]
+        else:
+            transform_list = [transforms.ToTensor(),
+                              transforms.Normalize(mean=[0.5, 0.5, 0.5],
+                                                   std=[0.5, 0.5, 0.5]),
+                              ]
+        self._transform = transforms.Compose(transform_list)
+
+    def _read_ids(self, file_path):
+        ids = np.loadtxt(file_path, delimiter='\t', dtype=np.str)
+        return [id[:-4] for id in ids]
+
+    def _read_conds(self, file_path):
+        with open(file_path, 'rb') as f:
+            return pickle.load(f)
+
+    def _get_cond_by_id(self, id):
+        if id in self._conds:
+            return self._conds[id]/5.0
+        else:
+            return None
+
+    def _get_img_by_id(self, id):
+        filepath = os.path.join(self._imgs_dir, id+'.jpg')
+        return cv_utils.read_cv2_img(filepath), filepath
+
+    def _generate_random_cond(self):
+        cond = None
+        while cond is None:
+            rand_sample_id = self._ids[random.randint(0, self._dataset_size - 1)]
+            cond = self._get_cond_by_id(rand_sample_id)
+            cond += np.random.uniform(-0.1, 0.1, cond.shape)
+        return cond

GANimation/data/prepare_au_annotations.py

@@ -0,0 +1,39 @@
+import numpy as np
+import os
+from tqdm import tqdm
+import argparse
+import glob
+import re
+import pickle
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-ia', '--input_aus_filesdir', type=str, help='Dir with imgs aus files')
+parser.add_argument('-op', '--output_path', type=str, help='Output path')
+args = parser.parse_args()
+
+def get_data(filepaths):
+    data = dict()
+    for filepath in tqdm(filepaths):
+        content = np.loadtxt(filepath, delimiter=', ', skiprows=1)
+        data[os.path.basename(filepath[:-4])] = content[2:19]
+
+    return data
+
+def save_dict(data, name):
+    with open(name + '.pkl', 'wb') as f:
+        pickle.dump(data, f, pickle.HIGHEST_PROTOCOL)
+
+def main():
+    filepaths = glob.glob(os.path.join(args.input_aus_filesdir, '*.csv'))
+    filepaths.sort()
+
+    # create aus file
+    data = get_data(filepaths)
+
+    if not os.path.isdir(args.output_path):
+        os.makedirs(args.output_path)
+    save_dict(data, os.path.join(args.output_path, "aus"))
+
+
+if __name__ == '__main__':
+    main()