Source code for tensorbay.opendataset.SCUT_FBP5500.loader

#!/usr/bin/env python3
#
# Copytright 2021 Graviti. Licensed under MIT License.
#
# pylint: disable=invalid-name

"""Dataloader of SCUT_FBP5500 dataset."""

import csv
import os
import struct
from itertools import islice

from tensorbay.dataset import Data, Dataset
from tensorbay.geometry.keypoint import Keypoint2D
from tensorbay.label import LabeledKeypoints2D
from tensorbay.utility import chunked

_DATASET_NAME = "SCUT_FBP5500"
_SEGMENT_NAMES = ("train", "test")
_CATEGORY_NAMES = {
    "AF": "Asian females",
    "AM": "Asian males",
    "CF": "Caucasian females",
    "CM": "Caucasian males",
}


[docs]def SCUT_FBP5500(path: str) -> Dataset: """`SCUT_FBP5500 <https://github.com/HCIILAB/SCUT-FBP5500-Database-Release>`_ dataset. The folder structure should be like:: <path> facial landmark/ <file_name>.pts ... Images/ <file_name>.jpg ... train_test_files/ split_of_60%training and 40%testing/ test.txt train.txt ... ... Arguments: path: The root directory of the dataset. Returns: Loaded :class:`~tensorbay.dataset.dataset.Dataset` instance. """ root_path = os.path.abspath(os.path.expanduser(path)) segment_path = os.path.join( root_path, "train_test_files", "split_of_60%training and 40%testing" ) image_path = os.path.join(root_path, "Images") label_path = os.path.join(root_path, "facial landmark") dataset = Dataset(_DATASET_NAME) dataset.load_catalog(os.path.join(os.path.dirname(__file__), "catalog.json")) for segment_name in _SEGMENT_NAMES: segment = dataset.create_segment(segment_name) with open(os.path.join(segment_path, f"{segment_name}.txt"), encoding="utf-8") as fp: reader = csv.reader(fp, delimiter=" ") for filename, beauty_score in reader: if filename == "CM152.jpg": # This image lacks 2D keypoints. continue segment.append(_get_data(filename, beauty_score, image_path, label_path)) return dataset
def _get_data(filename: str, beauty_score: str, image_path: str, label_path: str) -> Data: stem = os.path.splitext(os.path.basename(filename))[0] data = Data(os.path.join(image_path, filename)) keypoints2d = LabeledKeypoints2D() keypoints2d.attributes = {"beauty_score": float(beauty_score)} keypoints2d.category = _CATEGORY_NAMES[stem[:2]] with open(os.path.join(label_path, f"{stem}.pts"), "rb") as fp: points = struct.unpack("i172f", fp.read()) for x, y in chunked(islice(points, 1, None), 2): keypoints2d.append(Keypoint2D(float(x), float(y))) data.label.keypoints2d = [keypoints2d] return data