File size: 17,571 Bytes
ce190ee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
"""Data-loading functions in order to create a Dataset and DataLoaders.
Transforms for loaders are in transforms.py
"""

import json
import os
from pathlib import Path

import numpy as np
import torch
import yaml
from imageio import imread
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

from climategan.transforms import get_transforms
from climategan.tutils import get_normalized_depth_t
from climategan.utils import env_to_path, is_image_file

classes_dict = {
    "s": {  # unity
        0: [0, 0, 255, 255],  # Water
        1: [55, 55, 55, 255],  # Ground
        2: [0, 255, 255, 255],  # Building
        3: [255, 212, 0, 255],  # Traffic items
        4: [0, 255, 0, 255],  # Vegetation
        5: [255, 97, 0, 255],  # Terrain
        6: [255, 0, 0, 255],  # Car
        7: [60, 180, 60, 255],  # Trees
        8: [255, 0, 255, 255],  # Person
        9: [0, 0, 0, 255],  # Sky
        10: [255, 255, 255, 255],  # Default
    },
    "r": {  # deeplab v2
        0: [0, 0, 255, 255],  # Water
        1: [55, 55, 55, 255],  # Ground
        2: [0, 255, 255, 255],  # Building
        3: [255, 212, 0, 255],  # Traffic items
        4: [0, 255, 0, 255],  # Vegetation
        5: [255, 97, 0, 255],  # Terrain
        6: [255, 0, 0, 255],  # Car
        7: [60, 180, 60, 255],  # Trees
        8: [220, 20, 60, 255],  # Person
        9: [8, 19, 49, 255],  # Sky
        10: [0, 80, 100, 255],  # Default
    },
    "kitti": {
        0: [210, 0, 200],  # Terrain
        1: [90, 200, 255],  # Sky
        2: [0, 199, 0],  # Tree
        3: [90, 240, 0],  # Vegetation
        4: [140, 140, 140],  # Building
        5: [100, 60, 100],  # Road
        6: [250, 100, 255],  # GuardRail
        7: [255, 255, 0],  # TrafficSign
        8: [200, 200, 0],  # TrafficLight
        9: [255, 130, 0],  # Pole
        10: [80, 80, 80],  # Misc
        11: [160, 60, 60],  # Truck
        12: [255, 127, 80],  # Car
        13: [0, 139, 139],  # Van
        14: [0, 0, 0],  # Undefined
    },
    "flood": {
        0: [255, 0, 0],  # Cannot flood
        1: [0, 0, 255],  # Must flood
        2: [0, 0, 0],  # May flood
    },
}

kitti_mapping = {
    0: 5,  # Terrain -> Terrain
    1: 9,  # Sky -> Sky
    2: 7,  # Tree -> Trees
    3: 4,  # Vegetation ->  Vegetation
    4: 2,  # Building -> Building
    5: 1,  # Road -> Ground
    6: 3,  # GuardRail -> Traffic items
    7: 3,  # TrafficSign ->  Traffic items
    8: 3,  # TrafficLight -> Traffic items
    9: 3,  # Pole -> Traffic items
    10: 10,  # Misc -> default
    11: 6,  # Truck -> Car
    12: 6,  # Car -> Car
    13: 6,  # Van -> Car
    14: 10,  # Undefined -> Default
}


def encode_exact_segmap(seg, classes_dict, default_value=14):
    """
    When the mapping (rgb -> label) is known to be exact (no approximative rgb values)
    maps rgb image to segmap labels

    Args:
        seg (np.ndarray): H x W x 3 RGB image
        classes_dict (dict): Mapping {class: rgb value}
        default_value (int, optional): Value for unknown label. Defaults to 14.

    Returns:
        np.ndarray: Segmap as labels, not RGB
    """
    out = np.ones((seg.shape[0], seg.shape[1])) * default_value
    for cindex, cvalue in classes_dict.items():
        out[np.where((seg == cvalue).all(-1))] = cindex
    return out


def merge_labels(labels, mapping, default_value=14):
    """
    Maps labels from a source domain to labels of a target domain,
    typically kitti -> climategan

    Args:
        labels (np.ndarray): input segmap labels
        mapping (dict): source_label -> target_label
        default_value (int, optional): Unknown label. Defaults to 14.

    Returns:
        np.ndarray: Adapted labels
    """
    out = np.ones_like(labels) * default_value
    for source, target in mapping.items():
        out[labels == source] = target
    return out


def process_kitti_seg(path, kitti_classes, merge_map, default=14):
    """
    Processes a path to produce a 1 x 1 x H x W torch segmap

    %timeit process_kitti_seg(path, classes_dict, mapping, default=14)
    326 ms ± 118 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

    Args:
        path (str | pathlib.Path): Segmap RBG path
        kitti_classes (dict): Kitti map label -> rgb
        merge_map (dict): map kitti_label -> climategan_label
        default (int, optional): Unknown kitti label. Defaults to 14.

    Returns:
        torch.Tensor: 1 x 1 x H x W torch segmap
    """
    seg = imread(path)
    labels = encode_exact_segmap(seg, kitti_classes, default_value=default)
    merged = merge_labels(labels, merge_map, default_value=default)
    return torch.tensor(merged).unsqueeze(0).unsqueeze(0)


def decode_segmap_merged_labels(tensor, domain, is_target, nc=11):
    """Creates a label colormap for classes used in Unity segmentation benchmark.
    Arguments:
        tensor -- segmented image of size (1) x (nc) x (H) x (W)
            if prediction, or size (1) x (1) x (H) x (W) if target
    Returns:
        RGB tensor of size (1) x (3) x (H) x (W)
    #"""

    if is_target:  # Target is size 1 x 1 x H x W
        idx = tensor.squeeze(0).squeeze(0)
    else:  # Prediction is size 1 x nc x H x W
        idx = torch.argmax(tensor.squeeze(0), dim=0)

    indexer = torch.tensor(list(classes_dict[domain].values()))[:, :3]
    return indexer[idx.long()].permute(2, 0, 1).to(torch.float32).unsqueeze(0)


def decode_segmap_cityscapes_labels(image, nc=19):
    """Creates a label colormap used in CITYSCAPES segmentation benchmark.
    Arguments:
        image {array} -- segmented image
        (array of image size containing class at each pixel)
    Returns:
        array of size 3*nc -- A colormap for visualizing segmentation results.
    """
    colormap = np.zeros((19, 3), dtype=np.uint8)
    colormap[0] = [128, 64, 128]
    colormap[1] = [244, 35, 232]
    colormap[2] = [70, 70, 70]
    colormap[3] = [102, 102, 156]
    colormap[4] = [190, 153, 153]
    colormap[5] = [153, 153, 153]
    colormap[6] = [250, 170, 30]
    colormap[7] = [220, 220, 0]
    colormap[8] = [107, 142, 35]
    colormap[9] = [152, 251, 152]
    colormap[10] = [70, 130, 180]
    colormap[11] = [220, 20, 60]
    colormap[12] = [255, 0, 0]
    colormap[13] = [0, 0, 142]
    colormap[14] = [0, 0, 70]
    colormap[15] = [0, 60, 100]
    colormap[16] = [0, 80, 100]
    colormap[17] = [0, 0, 230]
    colormap[18] = [119, 11, 32]

    r = np.zeros_like(image).astype(np.uint8)
    g = np.zeros_like(image).astype(np.uint8)
    b = np.zeros_like(image).astype(np.uint8)

    for col in range(nc):
        idx = image == col
        r[idx] = colormap[col, 0]
        g[idx] = colormap[col, 1]
        b[idx] = colormap[col, 2]

    rgb = np.stack([r, g, b], axis=2)
    return rgb


def find_closest_class(pixel, dict_classes):
    """Takes a pixel as input and finds the closest known pixel value corresponding
    to a class in dict_classes

    Arguments:
        pixel -- tuple pixel (R,G,B,A)
    Returns:
        tuple pixel (R,G,B,A) corresponding to a key (a class) in dict_classes
    """
    min_dist = float("inf")
    closest_pixel = None
    for pixel_value in dict_classes.keys():
        dist = np.sqrt(np.sum(np.square(np.subtract(pixel, pixel_value))))
        if dist < min_dist:
            min_dist = dist
            closest_pixel = pixel_value
    return closest_pixel


def encode_segmap(arr, domain):
    """Change a segmentation RGBA array to a segmentation array
                            with each pixel being the index of the class
    Arguments:
        numpy array -- segmented image of size (H) x (W) x (4 RGBA values)
    Returns:
        numpy array of size (1) x (H) x (W) with each pixel being the index of the class
    """
    new_arr = np.zeros((1, arr.shape[0], arr.shape[1]))
    dict_classes = {
        tuple(rgba_value): class_id
        for (class_id, rgba_value) in classes_dict[domain].items()
    }
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            pixel_rgba = tuple(arr[i, j, :])
            if pixel_rgba in dict_classes.keys():
                new_arr[0, i, j] = dict_classes[pixel_rgba]
            else:
                pixel_rgba_closest = find_closest_class(pixel_rgba, dict_classes)
                new_arr[0, i, j] = dict_classes[pixel_rgba_closest]
    return new_arr


def encode_mask_label(arr, domain):
    """Change a segmentation RGBA array to a segmentation array
                            with each pixel being the index of the class
    Arguments:
        numpy array -- segmented image of size (H) x (W) x (3 RGB values)
    Returns:
        numpy array of size (1) x (H) x (W) with each pixel being the index of the class
    """
    diff = np.zeros((len(classes_dict[domain].keys()), arr.shape[0], arr.shape[1]))
    for cindex, cvalue in classes_dict[domain].items():
        diff[cindex, :, :] = np.sqrt(
            np.sum(
                np.square(arr - np.tile(cvalue, (arr.shape[0], arr.shape[1], 1))),
                axis=2,
            )
        )
    return np.expand_dims(np.argmin(diff, axis=0), axis=0)


def transform_segmap_image_to_tensor(path, domain):
    """
    Transforms a segmentation image to a tensor of size (1) x (1) x (H) x (W)
    with each pixel being the index of the class
    """
    arr = np.array(Image.open(path).convert("RGBA"))
    arr = encode_segmap(arr, domain)
    arr = torch.from_numpy(arr).float()
    arr = arr.unsqueeze(0)
    return arr


def save_segmap_tensors(path_to_json, path_to_dir, domain):
    """
    Loads the segmentation images mentionned in a json file, transforms them to
    tensors and save the tensors in the wanted directory

    Args:
        path_to_json: complete path to the json file where to find the original data
        path_to_dir: path to the directory where to save the tensors as tensor_name.pt
        domain: domain of the images ("r" or "s")

    e.g:
        save_tensors(
            "/network/tmp1/ccai/data/climategan/seg/train_s.json",
            "/network/tmp1/ccai/data/munit_dataset/simdata/Unity11K_res640/Seg_tensors/",
            "s",
        )
    """
    ims_list = None
    if path_to_json:
        path_to_json = Path(path_to_json).resolve()
        with open(path_to_json, "r") as f:
            ims_list = yaml.safe_load(f)

    assert ims_list is not None

    for im_dict in ims_list:
        for task_name, path in im_dict.items():
            if task_name == "s":
                file_name = os.path.splitext(path)[0]  # remove extension
                file_name = file_name.rsplit("/", 1)[-1]  # keep only the file_name
                tensor = transform_segmap_image_to_tensor(path, domain)
                torch.save(tensor, path_to_dir + file_name + ".pt")


def pil_image_loader(path, task):
    if Path(path).suffix == ".npy":
        arr = np.load(path).astype(np.uint8)
    elif is_image_file(path):
        # arr = imread(path).astype(np.uint8)
        arr = np.array(Image.open(path).convert("RGB"))
    else:
        raise ValueError("Unknown data type {}".format(path))

    # Convert from RGBA to RGB for images
    if len(arr.shape) == 3 and arr.shape[-1] == 4:
        arr = arr[:, :, 0:3]

    if task == "m":
        arr[arr != 0] = 1
        # Make sure mask is single-channel
        if len(arr.shape) >= 3:
            arr = arr[:, :, 0]

    # assert len(arr.shape) == 3, (path, task, arr.shape)

    return Image.fromarray(arr)


def tensor_loader(path, task, domain, opts):
    """load data as tensors
    Args:
        path (str): path to data
        task (str)
        domain (str)
    Returns:
        [Tensor]: 1 x C x H x W
    """
    if task == "s":
        if domain == "kitti":
            return process_kitti_seg(
                path, classes_dict["kitti"], kitti_mapping, default=14
            )
        return torch.load(path)
    elif task == "d":
        if Path(path).suffix == ".npy":
            arr = np.load(path)
        else:
            arr = imread(path)  # .astype(np.uint8) /!\ kitti is np.uint16
        tensor = torch.from_numpy(arr.astype(np.float32))
        tensor = get_normalized_depth_t(
            tensor,
            domain,
            normalize="d" in opts.train.pseudo.tasks,
            log=opts.gen.d.classify.enable,
        )
        tensor = tensor.unsqueeze(0)
        return tensor

    elif Path(path).suffix == ".npy":
        arr = np.load(path).astype(np.float32)
    elif is_image_file(path):
        arr = imread(path).astype(np.float32)
    else:
        raise ValueError("Unknown data type {}".format(path))

    # Convert from RGBA to RGB for images
    if len(arr.shape) == 3 and arr.shape[-1] == 4:
        arr = arr[:, :, 0:3]

    if task == "x":
        arr -= arr.min()
        arr /= arr.max()
        arr = np.moveaxis(arr, 2, 0)
    elif task == "s":
        arr = np.moveaxis(arr, 2, 0)
    elif task == "m":
        if arr.max() > 127:
            arr = (arr > 127).astype(arr.dtype)
        # Make sure mask is single-channel
        if len(arr.shape) >= 3:
            arr = arr[:, :, 0]
        arr = np.expand_dims(arr, 0)

    return torch.from_numpy(arr).unsqueeze(0)


class OmniListDataset(Dataset):
    def __init__(self, mode, domain, opts, transform=None):

        self.opts = opts
        self.domain = domain
        self.mode = mode
        self.tasks = set(opts.tasks)
        self.tasks.add("x")
        if "p" in self.tasks:
            self.tasks.add("m")

        file_list_path = Path(opts.data.files[mode][domain])
        if "/" not in str(file_list_path):
            file_list_path = Path(opts.data.files.base) / Path(
                opts.data.files[mode][domain]
            )

        if file_list_path.suffix == ".json":
            self.samples_paths = self.json_load(file_list_path)
        elif file_list_path.suffix in {".yaml", ".yml"}:
            self.samples_paths = self.yaml_load(file_list_path)
        else:
            raise ValueError("Unknown file list type in {}".format(file_list_path))

        if opts.data.max_samples and opts.data.max_samples != -1:
            assert isinstance(opts.data.max_samples, int)
            self.samples_paths = self.samples_paths[: opts.data.max_samples]

        self.filter_samples()
        if opts.data.check_samples:
            print(f"Checking samples ({mode}, {domain})")
            self.check_samples()
        self.file_list_path = str(file_list_path)
        self.transform = transform

    def filter_samples(self):
        """
        Filter out data which is not required for the model's tasks
        as defined in opts.tasks
        """
        self.samples_paths = [
            {k: v for k, v in s.items() if k in self.tasks} for s in self.samples_paths
        ]

    def __getitem__(self, i):
        """Return an item in the dataset with fields:
        {
            data: transform({
                domains: values
            }),
            paths: [{task: path}],
            domain: [domain],
            mode: [train|val]
        }
        Args:
            i (int): index of item to retrieve
        Returns:
            dict: dataset item where tensors of data are in item["data"] which is a dict
                  {task: tensor}
        """
        paths = self.samples_paths[i]

        # always apply transforms,
        # if no transform is specified, ToTensor and Normalize will be applied

        item = {
            "data": self.transform(
                {
                    task: tensor_loader(
                        env_to_path(path),
                        task,
                        self.domain,
                        self.opts,
                    )
                    for task, path in paths.items()
                }
            ),
            "paths": paths,
            "domain": self.domain if self.domain != "kitti" else "s",
            "mode": self.mode,
        }

        return item

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

    def json_load(self, file_path):
        with open(file_path, "r") as f:
            return json.load(f)

    def yaml_load(self, file_path):
        with open(file_path, "r") as f:
            return yaml.safe_load(f)

    def check_samples(self):
        """Checks that every file listed in samples_paths actually
        exist on the file-system
        """
        for s in self.samples_paths:
            for k, v in s.items():
                assert Path(v).exists(), f"{k} {v} does not exist"


def get_loader(mode, domain, opts):
    if (
        domain != "kitti"
        or not opts.train.kitti.pretrain
        or not opts.train.kitti.batch_size
    ):
        batch_size = opts.data.loaders.get("batch_size", 4)
    else:
        batch_size = opts.train.kitti.get("batch_size", 4)

    return DataLoader(
        OmniListDataset(
            mode,
            domain,
            opts,
            transform=transforms.Compose(get_transforms(opts, mode, domain)),
        ),
        batch_size=batch_size,
        shuffle=True,
        num_workers=opts.data.loaders.get("num_workers", 8),
        pin_memory=True,  # faster transfer to gpu
        drop_last=True,  # avoids batchnorm pbs if last batch has size 1
    )


def get_all_loaders(opts):
    loaders = {}
    for mode in ["train", "val"]:
        loaders[mode] = {}
        for domain in opts.domains:
            if mode in opts.data.files:
                if domain in opts.data.files[mode]:
                    loaders[mode][domain] = get_loader(mode, domain, opts)
    return loaders