widerface.py 8.1 KB
Newer Older
Josh Bradley's avatar
Josh Bradley committed
1
2
import os
from os.path import abspath, expanduser
3
4
from pathlib import Path

5
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
6
7
8
9

import torch
from PIL import Image

10
from .utils import download_and_extract_archive, download_file_from_google_drive, extract_archive, verify_str_arg
Josh Bradley's avatar
Josh Bradley committed
11
12
13
14
15
16
17
from .vision import VisionDataset


class WIDERFace(VisionDataset):
    """`WIDERFace <http://shuoyang1213.me/WIDERFACE/>`_ Dataset.

    Args:
18
        root (str or ``pathlib.Path``): Root directory where images and annotations are downloaded to.
Josh Bradley's avatar
Josh Bradley committed
19
            Expects the following folder structure if download=False:
20
21
22

            .. code::

Josh Bradley's avatar
Josh Bradley committed
23
24
25
26
27
28
29
30
                <root>
                    └── widerface
                        ├── wider_face_split ('wider_face_split.zip' if compressed)
                        ├── WIDER_train ('WIDER_train.zip' if compressed)
                        ├── WIDER_val ('WIDER_val.zip' if compressed)
                        └── WIDER_test ('WIDER_test.zip' if compressed)
        split (string): The dataset split to use. One of {``train``, ``val``, ``test``}.
            Defaults to ``train``.
anthony-cabacungan's avatar
anthony-cabacungan committed
31
        transform (callable, optional): A function/transform that takes in a PIL image
Josh Bradley's avatar
Josh Bradley committed
32
33
34
35
36
37
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.
38

39
40
41
42
            .. warning::

                To download the dataset `gdown <https://github.com/wkentaro/gdown>`_ is required.

Josh Bradley's avatar
Josh Bradley committed
43
44
45
46
    """

    BASE_FOLDER = "widerface"
    FILE_LIST = [
47
48
49
50
        # File ID                             MD5 Hash                            Filename
        ("15hGDLhsx8bLgLcIRD5DhYt5iBxnjNF1M", "3fedf70df600953d25982bcd13d91ba2", "WIDER_train.zip"),
        ("1GUCogbp16PMGa39thoMMeWxp7Rp5oM8Q", "dfa7d7e790efa35df3788964cf0bbaea", "WIDER_val.zip"),
        ("1HIfDbVEWKmsYKJZm4lchTBDLW5N7dY5T", "e5d8f4248ed24c334bbd12f49c29dd40", "WIDER_test.zip"),
Josh Bradley's avatar
Josh Bradley committed
51
52
    ]
    ANNOTATIONS_FILE = (
53
        "http://shuoyang1213.me/WIDERFACE/support/bbx_annotation/wider_face_split.zip",
Josh Bradley's avatar
Josh Bradley committed
54
        "0e3767bcf0e326556d407bf5bff5d27c",
55
        "wider_face_split.zip",
Josh Bradley's avatar
Josh Bradley committed
56
57
58
    )

    def __init__(
59
        self,
60
        root: Union[str, Path],
61
62
63
64
        split: str = "train",
        transform: Optional[Callable] = None,
        target_transform: Optional[Callable] = None,
        download: bool = False,
Josh Bradley's avatar
Josh Bradley committed
65
    ) -> None:
66
        super().__init__(
67
68
            root=os.path.join(root, self.BASE_FOLDER), transform=transform, target_transform=target_transform
        )
Josh Bradley's avatar
Josh Bradley committed
69
70
71
72
73
74
75
        # check arguments
        self.split = verify_str_arg(split, "split", ("train", "val", "test"))

        if download:
            self.download()

        if not self._check_integrity():
76
            raise RuntimeError("Dataset not found or corrupted. You can use download=True to download and prepare it")
Josh Bradley's avatar
Josh Bradley committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

        self.img_info: List[Dict[str, Union[str, Dict[str, torch.Tensor]]]] = []
        if self.split in ("train", "val"):
            self.parse_train_val_annotations_file()
        else:
            self.parse_test_annotations_file()

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is a dict of annotations for all faces in the image.
            target=None for the test split.
        """

        # stay consistent with other datasets and return a PIL Image
        img = Image.open(self.img_info[index]["img_path"])

        if self.transform is not None:
            img = self.transform(img)

        target = None if self.split == "test" else self.img_info[index]["annotations"]
        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self) -> int:
        return len(self.img_info)

    def extra_repr(self) -> str:
        lines = ["Split: {split}"]
111
        return "\n".join(lines).format(**self.__dict__)
Josh Bradley's avatar
Josh Bradley committed
112
113
114
115
116

    def parse_train_val_annotations_file(self) -> None:
        filename = "wider_face_train_bbx_gt.txt" if self.split == "train" else "wider_face_val_bbx_gt.txt"
        filepath = os.path.join(self.root, "wider_face_split", filename)

117
        with open(filepath) as f:
Josh Bradley's avatar
Josh Bradley committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
            lines = f.readlines()
            file_name_line, num_boxes_line, box_annotation_line = True, False, False
            num_boxes, box_counter = 0, 0
            labels = []
            for line in lines:
                line = line.rstrip()
                if file_name_line:
                    img_path = os.path.join(self.root, "WIDER_" + self.split, "images", line)
                    img_path = abspath(expanduser(img_path))
                    file_name_line = False
                    num_boxes_line = True
                elif num_boxes_line:
                    num_boxes = int(line)
                    num_boxes_line = False
                    box_annotation_line = True
                elif box_annotation_line:
                    box_counter += 1
                    line_split = line.split(" ")
                    line_values = [int(x) for x in line_split]
                    labels.append(line_values)
                    if box_counter >= num_boxes:
                        box_annotation_line = False
                        file_name_line = True
                        labels_tensor = torch.tensor(labels)
142
143
144
145
                        self.img_info.append(
                            {
                                "img_path": img_path,
                                "annotations": {
146
147
148
149
150
151
152
                                    "bbox": labels_tensor[:, 0:4].clone(),  # x, y, width, height
                                    "blur": labels_tensor[:, 4].clone(),
                                    "expression": labels_tensor[:, 5].clone(),
                                    "illumination": labels_tensor[:, 6].clone(),
                                    "occlusion": labels_tensor[:, 7].clone(),
                                    "pose": labels_tensor[:, 8].clone(),
                                    "invalid": labels_tensor[:, 9].clone(),
153
154
155
                                },
                            }
                        )
Josh Bradley's avatar
Josh Bradley committed
156
157
158
                        box_counter = 0
                        labels.clear()
                else:
159
                    raise RuntimeError(f"Error parsing annotation file {filepath}")
Josh Bradley's avatar
Josh Bradley committed
160
161
162
163

    def parse_test_annotations_file(self) -> None:
        filepath = os.path.join(self.root, "wider_face_split", "wider_face_test_filelist.txt")
        filepath = abspath(expanduser(filepath))
164
        with open(filepath) as f:
Josh Bradley's avatar
Josh Bradley committed
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
            lines = f.readlines()
            for line in lines:
                line = line.rstrip()
                img_path = os.path.join(self.root, "WIDER_test", "images", line)
                img_path = abspath(expanduser(img_path))
                self.img_info.append({"img_path": img_path})

    def _check_integrity(self) -> bool:
        # Allow original archive to be deleted (zip). Only need the extracted images
        all_files = self.FILE_LIST.copy()
        all_files.append(self.ANNOTATIONS_FILE)
        for (_, md5, filename) in all_files:
            file, ext = os.path.splitext(filename)
            extracted_dir = os.path.join(self.root, file)
            if not os.path.exists(extracted_dir):
                return False
        return True

    def download(self) -> None:
        if self._check_integrity():
185
            print("Files already downloaded and verified")
Josh Bradley's avatar
Josh Bradley committed
186
187
188
189
190
191
192
193
194
            return

        # download and extract image data
        for (file_id, md5, filename) in self.FILE_LIST:
            download_file_from_google_drive(file_id, self.root, filename, md5)
            filepath = os.path.join(self.root, filename)
            extract_archive(filepath)

        # download and extract annotation files
195
196
197
        download_and_extract_archive(
            url=self.ANNOTATIONS_FILE[0], download_root=self.root, md5=self.ANNOTATIONS_FILE[1]
        )