File size: 2,887 Bytes
d9bb75c | 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 | # Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This software may be used and distributed in accordance with
# the terms of the DINOv3 License Agreement.
import os
from enum import Enum
from typing import Any, Callable, List, Optional, Tuple, Union
from PIL import Image
from .decoders import Decoder, DenseTargetDecoder, ImageDataDecoder
from .extended import ExtendedVisionDataset
class _Split(Enum):
TRAIN = "train"
VAL = "val"
@property
def dirname(self) -> str:
return {
_Split.TRAIN: "training",
_Split.VAL: "validation",
}[self]
def _file_to_segmentation_path(file_name: str, segm_base_path: str) -> str:
file_name_noext = os.path.splitext(file_name)[0]
return os.path.join(segm_base_path, file_name_noext + ".png")
def _load_segmentation(root: str, split_file_names: List[str]):
segm_base_path = "annotations"
segmentation_paths = [_file_to_segmentation_path(file_name, segm_base_path) for file_name in split_file_names]
return segmentation_paths
def _load_file_paths(root: str, split: _Split) -> Tuple[List[str], List[str]]:
with open(os.path.join(root, f"ADE20K_object150_{split.value}.txt")) as f:
split_file_names = sorted(f.read().strip().split("\n"))
all_segmentation_paths = _load_segmentation(root, split_file_names)
file_names = [os.path.join("images", el) for el in split_file_names]
return file_names, all_segmentation_paths
class ADE20K(ExtendedVisionDataset):
Split = Union[_Split]
Labels = Union[Image.Image]
def __init__(
self,
split: "ADE20K.Split",
root: Optional[str] = None,
transforms: Optional[Callable] = None,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
image_decoder: Decoder = ImageDataDecoder,
target_decoder: Decoder = DenseTargetDecoder,
) -> None:
super().__init__(
root=root,
transforms=transforms,
transform=transform,
target_transform=target_transform,
image_decoder=image_decoder,
target_decoder=target_decoder,
)
self.image_paths, self.target_paths = _load_file_paths(root, split)
def get_image_data(self, index: int) -> bytes:
image_relpath = self.image_paths[index]
image_full_path = os.path.join(self.root, image_relpath)
with open(image_full_path, mode="rb") as f:
image_data = f.read()
return image_data
def get_target(self, index: int) -> Any:
target_relpath = self.target_paths[index]
target_full_path = os.path.join(self.root, target_relpath)
with open(target_full_path, mode="rb") as f:
target_data = f.read()
return target_data
def __len__(self) -> int:
return len(self.image_paths)
|