# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""ADE20K dataset."""

import os
import numpy as np
from PIL import Image
import torch


ADE_CLASSES = [
    'wall',
    'building, edifice',
    'sky',
    'floor, flooring',
    'tree',
    'ceiling',
    'road, route',
    'bed',
    'windowpane, window',
    'grass',
    'cabinet',
    'sidewalk, pavement',
    'person, individual, someone, somebody, mortal, soul',
    'earth, ground',
    'door, double, door',
    'table',
    'mountain, mount',
    'plant, flora, plant, life',
    'curtain, drape, drapery, mantle, pall',
    'chair',
    'car, auto, automobile, machine, motorcar',
    'water',
    'painting, picture',
    'sofa, couch, lounge',
    'shelf',
    'house',
    'sea',
    'mirror',
    'rug, carpet, carpeting',
    'field',
    'armchair',
    'seat',
    'fence, fencing',
    'desk',
    'rock, stone',
    'wardrobe, closet, press',
    'lamp',
    'bathtub, bathing, tub, bath, tub',
    'railing, rail',
    'cushion',
    'base, pedestal, stand',
    'box',
    'column, pillar',
    'signboard, sign',
    'chest, of, drawers, chest, bureau, dresser',
    'counter',
    'sand',
    'sink',
    'skyscraper',
    'fireplace, hearth, open, fireplace',
    'refrigerator, icebox',
    'grandstand, covered, stand',
    'path',
    'stairs, steps',
    'runway',
    'case, display, case, showcase, vitrine',
    'pool, table, billiard, table, snooker, table',
    'pillow',
    'screen, door, screen',
    'stairway, staircase',
    'river',
    'bridge, span',
    'bookcase',
    'blind, screen',
    'coffee, table, cocktail, table',
    'toilet, can, commode, crapper, pot, potty, stool, throne',
    'flower',
    'book',
    'hill',
    'bench',
    'countertop',
    'stove, kitchen, stove, range, kitchen, range, cooking, stove',
    'palm, palm, tree',
    'kitchen, island',
    (
        'computer, computing, machine, computing, device, data, processor,'
        ' electronic, computer, information, processing, system'
    ),
    'swivel, chair',
    'boat',
    'bar',
    'arcade, machine',
    'hovel, hut, hutch, shack, shanty',
    (
        'bus, autobus, coach, charabanc, double-decker, jitney, motorbus,'
        ' motorcoach, omnibus, passenger, vehicle'
    ),
    'towel',
    'light, light, source',
    'truck, motortruck',
    'tower',
    'chandelier, pendant, pendent',
    'awning, sunshade, sunblind',
    'streetlight, street, lamp',
    'booth, cubicle, stall, kiosk',
    (
        'television, television, receiver, television, set, tv, tv, set, idiot,'
        ' box, boob, tube, telly, goggle, box'
    ),
    'airplane, aeroplane, plane',
    'dirt, track',
    'apparel, wearing, apparel, dress, clothes',
    'pole',
    'land, ground, soil',
    'bannister, banister, balustrade, balusters, handrail',
    'escalator, moving, staircase, moving, stairway',
    'ottoman, pouf, pouffe, puff, hassock',
    'bottle',
    'buffet, counter, sideboard',
    'poster, posting, placard, notice, bill, card',
    'stage',
    'van',
    'ship',
    'fountain',
    'conveyer, belt, conveyor, belt, conveyer, conveyor, transporter',
    'canopy',
    'washer, automatic, washer, washing, machine',
    'plaything, toy',
    'swimming, pool, swimming, bath, natatorium',
    'stool',
    'barrel, cask',
    'basket, handbasket',
    'waterfall, falls',
    'tent, collapsible, shelter',
    'bag',
    'minibike, motorbike',
    'cradle',
    'oven',
    'ball',
    'food, solid, food',
    'step, stair',
    'tank, storage, tank',
    'trade, name, brand, name, brand, marque',
    'microwave, microwave, oven',
    'pot, flowerpot',
    'animal, animate, being, beast, brute, creature, fauna',
    'bicycle, bike, wheel, cycle',
    'lake',
    'dishwasher, dish, washer, dishwashing, machine',
    'screen, silver, screen, projection, screen',
    'blanket, cover',
    'sculpture',
    'hood, exhaust, hood',
    'sconce',
    'vase',
    'traffic, light, traffic, signal, stoplight',
    'tray',
    (
        'ashcan, trash, can, garbage, can, wastebin, ash, bin, ash-bin, ashbin,'
        ' dustbin, trash, barrel, trash, bin'
    ),
    'fan',
    'pier, wharf, wharfage, dock',
    'crt, screen',
    'plate',
    'monitor, monitoring, device',
    'bulletin, board, notice, board',
    'shower',
    'radiator',
    'glass, drinking, glass',
    'clock',
    'flag',
]


ADE_STUFF_CLASS = [
    'wall',
    'sky',
    'floor, flooring',
    'tree',
    'ceiling',
    'road, route',
    'grass',
    'earth, ground',
    'mountain, mount',
    'plant, flora, plant, life',
    'water',
    'sea',
    'field',
    'sand',
    'skyscraper',
    'path',
    'river',
    'bridge, span',
    'flower',
    'hill',
    'land, ground, soil',
    'dirt, track',
    'apparel, wearing, apparel, dress, clothes',
    'lake',
    'waterfall, falls',
]

ADE_THING_CLASS = [
    'building, edifice',
    'bed',
    'windowpane, window',
    'cabinet',
    'sidewalk, pavement',
    'person, individual, someone, somebody, mortal, soul',
    'door, double, door',
    'table',
    'curtain, drape, drapery, mantle, pall',
    'chair',
    'car, auto, automobile, machine, motorcar',
    'painting, picture',
    'sofa, couch, lounge',
    'shelf',
    'house',
    'mirror',
    'rug, carpet, carpeting',
    'armchair',
    'seat',
    'fence, fencing',
    'desk',
    'rock, stone',
    'wardrobe, closet, press',
    'lamp',
    'bathtub, bathing, tub, bath, tub',
    'railing, rail',
    'cushion',
    'base, pedestal, stand',
    'box',
    'column, pillar',
    'signboard, sign',
    'chest, of, drawers, chest, bureau, dresser',
    'counter',
    'sink',
    'fireplace, hearth, open, fireplace',
    'refrigerator, icebox',
    'grandstand, covered, stand',
    'stairs, steps',
    'runway',
    'case, display, case, showcase, vitrine',
    'pool, table, billiard, table, snooker, table',
    'pillow',
    'screen, door, screen',
    'stairway, staircase',
    'bookcase',
    'blind, screen',
    'coffee, table, cocktail, table',
    'toilet, can, commode, crapper, pot, potty, stool, throne',
    'book',
    'bench',
    'countertop',
    'stove, kitchen, stove, range, kitchen, range, cooking, stove',
    'palm, palm, tree',
    'kitchen, island',
    (
        'computer, computing, machine, computing, device, data, processor,'
        ' electronic, computer, information, processing, system'
    ),
    'swivel, chair',
    'boat',
    'bar',
    'arcade, machine',
    'hovel, hut, hutch, shack, shanty',
    (
        'bus, autobus, coach, charabanc, double-decker, jitney, motorbus,'
        ' motorcoach, omnibus, passenger, vehicle'
    ),
    'towel',
    'light, light, source',
    'truck, motortruck',
    'tower',
    'chandelier, pendant, pendent',
    'awning, sunshade, sunblind',
    'streetlight, street, lamp',
    'booth, cubicle, stall, kiosk',
    (
        'television, television, receiver, television, set, tv, tv, set, idiot,'
        ' box, boob, tube, telly, goggle, box'
    ),
    'airplane, aeroplane, plane',
    'pole',
    'bannister, banister, balustrade, balusters, handrail',
    'escalator, moving, staircase, moving, stairway',
    'ottoman, pouf, pouffe, puff, hassock',
    'bottle',
    'buffet, counter, sideboard',
    'poster, posting, placard, notice, bill, card',
    'stage',
    'van',
    'ship',
    'fountain',
    'conveyer, belt, conveyor, belt, conveyer, conveyor, transporter',
    'canopy',
    'washer, automatic, washer, washing, machine',
    'plaything, toy',
    'swimming, pool, swimming, bath, natatorium',
    'stool',
    'barrel, cask',
    'basket, handbasket',
    'tent, collapsible, shelter',
    'bag',
    'minibike, motorbike',
    'cradle',
    'oven',
    'ball',
    'food, solid, food',
    'step, stair',
    'tank, storage, tank',
    'trade, name, brand, name, brand, marque',
    'microwave, microwave, oven',
    'pot, flowerpot',
    'animal, animate, being, beast, brute, creature, fauna',
    'bicycle, bike, wheel, cycle',
    'dishwasher, dish, washer, dishwashing, machine',
    'screen, silver, screen, projection, screen',
    'blanket, cover',
    'sculpture',
    'hood, exhaust, hood',
    'sconce',
    'vase',
    'traffic, light, traffic, signal, stoplight',
    'tray',
    (
        'ashcan, trash, can, garbage, can, wastebin, ash, bin, ash-bin, ashbin,'
        ' dustbin, trash, barrel, trash, bin'
    ),
    'fan',
    'pier, wharf, wharfage, dock',
    'crt, screen',
    'plate',
    'monitor, monitoring, device',
    'bulletin, board, notice, board',
    'shower',
    'radiator',
    'glass, drinking, glass',
    'clock',
    'flag',
]


ADE_STUFF_CLASS_ID = [
    0,
    2,
    3,
    4,
    5,
    6,
    9,
    13,
    16,
    17,
    21,
    26,
    29,
    46,
    48,
    52,
    60,
    61,
    66,
    68,
    94,
    91,
    92,
    128,
    113,
]

ADE_THING_CLASS_ID = [
    1,
    7,
    8,
    10,
    11,
    12,
    14,
    15,
    18,
    19,
    20,
    22,
    23,
    24,
    25,
    27,
    28,
    30,
    31,
    32,
    33,
    34,
    35,
    36,
    37,
    38,
    39,
    40,
    41,
    42,
    43,
    44,
    45,
    47,
    49,
    50,
    51,
    53,
    54,
    55,
    56,
    57,
    58,
    59,
    62,
    63,
    64,
    65,
    67,
    69,
    70,
    71,
    72,
    73,
    74,
    75,
    76,
    77,
    78,
    79,
    80,
    81,
    82,
    83,
    84,
    85,
    86,
    87,
    88,
    89,
    90,
    93,
    95,
    96,
    97,
    98,
    99,
    100,
    101,
    102,
    103,
    104,
    105,
    106,
    107,
    108,
    109,
    110,
    111,
    112,
    114,
    115,
    116,
    117,
    118,
    119,
    120,
    121,
    122,
    123,
    124,
    125,
    126,
    127,
    129,
    130,
    131,
    132,
    133,
    134,
    135,
    136,
    137,
    138,
    139,
    140,
    141,
    142,
    143,
    144,
    145,
    146,
    147,
    148,
    149,
]


class ADEDataset(torch.utils.data.Dataset):
  """ADE dataset."""

  def __init__(self, root, split='validation', transform=None):
    """Construct ADE dataset.

    Args:
      root (string): Root directory where images are downloaded.
      split (string): The split of the dataset.
      transform (callable, optional): Optional transform to be applied on a
        sample.
    """
    self.root = root
    self.image_dir = os.path.join(root, 'images', split)
    self.ann_dir = os.path.join(root, 'annotations', split)
    self.images = os.listdir(self.image_dir)
    self.transform = transform

  def __getitem__(self, index):
    img_path = os.path.join(self.image_dir, self.images[index])
    img = Image.open(img_path).convert('RGB')
    img = np.asarray(img)
    idx = self.images[index].split('.')[0]
    ann_path = os.path.join(self.ann_dir, f'{idx}.png')
    ann = np.asarray(Image.open(ann_path), dtype=np.int32)
    return img, img_path, ann, idx

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