From 71eeb13d4204511b8c3b3f331fa6f2abf4dc81eb Mon Sep 17 00:00:00 2001
From: Stanislav Pidhorskyi <stanislav@podgorskiy.com>
Date: Thu, 23 Apr 2020 14:03:18 -0400
Subject: [PATCH] Added names for classifiers. Added comment about celeb-hq

---
 .../prepare_celeba_hq_tfrecords.py            |  3 ++
 principal_directions/classifier.py            | 43 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/dataset_preparation/prepare_celeba_hq_tfrecords.py b/dataset_preparation/prepare_celeba_hq_tfrecords.py
index 18960276..1a253e6a 100644
--- a/dataset_preparation/prepare_celeba_hq_tfrecords.py
+++ b/dataset_preparation/prepare_celeba_hq_tfrecords.py
@@ -28,6 +28,9 @@ def prepare_celeba(cfg, logger, train=True):
     os.makedirs(directory, exist_ok=True)
 
     images = []
+    # The official way of generating CelebA-HQ can be challenging.
+    # Please refer to this page: https://github.com/suvojit-0x55aa/celebA-HQ-dataset-download
+    # You can get pre-generated dataset from: https://drive.google.com/drive/folders/11Vz0fqHS2rXDb5pprgTjpD7S2BAJhi1P
     source_path = '/data/datasets/celeba-hq/data1024x1024'
     for filename in tqdm.tqdm(os.listdir(source_path)):
         images.append((int(filename[:-4]), filename))
diff --git a/principal_directions/classifier.py b/principal_directions/classifier.py
index 682869ad..b9bcca6c 100644
--- a/principal_directions/classifier.py
+++ b/principal_directions/classifier.py
@@ -67,6 +67,49 @@ def load_pkl(file_or_url):
     'https://drive.google.com/uc?id=1SF_mJIdyGINXoV-I6IAxHB_k5dxiF6M-',  # celebahq-classifier-39-wearing-necktie.pkl
 ]
 
+classifier_id_name = {
+    0: "gender",
+    1: "smiling",
+    2: "attractive",
+    3: "wavy-hair",
+    4: "young",
+    5: "5-o-clock-shadow",
+    6: "arched-eyebrows",
+    7: "bags-under-eyes",
+    8: "bald",
+    9: "bangs",
+    10: "big-lips",
+    11: "big-nose",
+    12: "black-hair",
+    13: "blond-hair",
+    14: "blurry",
+    15: "brown-hair",
+    16: "bushy-eyebrows",
+    17: "chubby",
+    18: "double-chin",
+    19: "eyeglasses",
+    20: "goatee",
+    21: "gray-hair",
+    22: "heavy-makeup",
+    23: "high-cheekbones",
+    24: "mouth-slightly-open",
+    25: "mustache",
+    26: "narrow-eyes",
+    27: "no-beard",
+    28: "oval-face",
+    29: "pale-skin",
+    30: "pointy-nose",
+    31: "receding-hairline",
+    32: "rosy-cheeks",
+    33: "sideburns",
+    34: "straight-hair",
+    35: "wearing-earrings",
+    36: "wearing-hat",
+    37: "wearing-lipstick",
+    38: "wearing-necklace",
+    39: "wearing-necktie",
+}
+
 
 def make_classifier(attrib_idx):
     classifier = load_pkl(classifier_urls[attrib_idx])