Spaces:
Runtime error
Runtime error
File size: 14,485 Bytes
3424266 |
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 |
# Copyright 2024 EPFL and Apple Inc.
#
# 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.
import os
import os.path
import pickle
import random
from copy import deepcopy
from typing import Any, Callable, Dict, List, Optional, Tuple, cast
import numpy as np
from torchvision.datasets.vision import VisionDataset
from fourm.data.modality_transforms import AbstractTransform, get_transform_key
IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp', '.jpx', '.npy', '.npz')
UNIFIED_EXTENSIONS = IMG_EXTENSIONS + ('.json', '.txt', '.json.gz')
def has_file_allowed_extension(filename: str, extensions: Tuple[str, ...]) -> bool:
"""Checks if a file is an allowed extension.
Args:
filename (string): path to a file
extensions (tuple of strings): extensions to consider (lowercase)
Returns:
bool: True if the filename ends with one of given extensions
"""
return filename.lower().endswith(extensions)
def is_image_file(filename: str) -> bool:
"""Checks if a file is an allowed image extension.
Args:
filename (string): path to a file
Returns:
bool: True if the filename ends with a known image extension
"""
return has_file_allowed_extension(filename, IMG_EXTENSIONS)
def make_dataset(
directory: str,
class_to_idx: Dict[str, int],
extensions: Optional[Tuple[str, ...]] = None,
is_valid_file: Optional[Callable[[str], bool]] = None,
cache_path: Optional[str] = None,
) -> List[Tuple[str, int]]:
if cache_path is not None and os.path.exists(cache_path):
# Load cached file paths from disk if it exists
with open(cache_path, 'rb') as f:
return pickle.load(f)
instances = []
directory = os.path.expanduser(directory)
both_none = extensions is None and is_valid_file is None
both_something = extensions is not None and is_valid_file is not None
if both_none or both_something:
raise ValueError("Both extensions and is_valid_file cannot be None or not None at the same time")
if extensions is not None:
def is_valid_file(x: str) -> bool:
return has_file_allowed_extension(x, cast(Tuple[str, ...], extensions))
is_valid_file = cast(Callable[[str], bool], is_valid_file)
for target_class in sorted(class_to_idx.keys()):
class_index = class_to_idx[target_class]
target_dir = os.path.join(directory, target_class)
if not os.path.isdir(target_dir):
continue
for root, _, fnames in sorted(os.walk(target_dir, followlinks=True)):
for fname in sorted(fnames):
path = os.path.join(root, fname)
if is_valid_file(path):
item = path, class_index
instances.append(item)
if cache_path is not None:
# Cache all file paths s.t. setting up the dataloader is instant in the future
os.makedirs(os.path.dirname(cache_path), exist_ok=True)
with open(cache_path, 'wb') as f:
pickle.dump(instances, f)
return instances
class DatasetFolder(VisionDataset):
"""A generic data loader where the samples are arranged in this way: ::
root/class_x/xxx.ext
root/class_x/xxy.ext
root/class_x/xxz.ext
root/class_y/123.ext
root/class_y/nsdf3.ext
root/class_y/asd932_.ext
Args:
root (string): Root directory path.
loader (callable): A function to load a sample given its path.
extensions (tuple[string]): A list of allowed extensions.
both extensions and is_valid_file should not be passed.
transform (callable, optional): A function/transform that takes in
a sample and returns a transformed version.
E.g, ``transforms.RandomCrop`` for images.
target_transform (callable, optional): A function/transform that takes
in the target and transforms it.
is_valid_file (callable, optional): A function that takes path of a file
and check if the file is a valid file (used to check of corrupt logs)
both extensions and is_valid_file should not be passed.
Attributes:
classes (list): List of the class names sorted alphabetically.
class_to_idx (dict): Dict with items (class_name, class_index).
samples (list): List of (sample path, class_index) tuples
targets (list): The class_index value for each image in the dataset
"""
def __init__(
self,
root: str,
loader: Callable[[str], Any],
extensions: Optional[Tuple[str, ...]] = None,
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
is_valid_file: Optional[Callable[[str], bool]] = None,
) -> None:
super(DatasetFolder, self).__init__(root, transform=transform,
target_transform=target_transform)
classes, class_to_idx = self._find_classes(self.root)
samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file)
if len(samples) == 0:
msg = "Found 0 logs in subfolders of: {}\n".format(self.root)
if extensions is not None:
msg += "Supported extensions are: {}".format(",".join(extensions))
raise RuntimeError(msg)
self.loader = loader
self.extensions = extensions
self.classes = classes
self.class_to_idx = class_to_idx
self.samples = samples
self.targets = [s[1] for s in samples]
def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
"""
Finds the class folders in a dataset.
Args:
dir (string): Root directory path.
Returns:
tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.
Ensures:
No class is a subdirectory of another.
"""
classes = [d.name for d in os.scandir(dir) if d.is_dir()]
classes.sort()
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
return classes, class_to_idx
def __getitem__(self, index: int) -> Tuple[Any, Any]:
"""
Args:
index (int): Index
Returns:
tuple: (sample, target) where target is class_index of the target class.
"""
while True:
try:
path, target = self.samples[index]
sample = self.loader(path)
break
except Exception as e:
print(e)
index = random.randint(0, len(self.samples) - 1)
if self.transform is not None:
sample = self.transform(sample)
if self.target_transform is not None:
target = self.target_transform(target)
return sample, target
def __len__(self) -> int:
return len(self.samples)
class MultiModalDatasetFolder(VisionDataset):
"""A generic multi-modal dataset loader where the samples are arranged in this way: ::
root/modality_a/class_x/xxx.ext
root/modality_a/class_y/xxy.ext
root/modality_a/class_z/xxz.ext
root/modality_b/class_x/xxx.ext
root/modality_b/class_y/xxy.ext
root/modality_b/class_z/xxz.ext
Args:
root (string): Root directory path.
modalities (list): List of modalities as strings
modality_paths (dict): Dict of paths to modalities
modality_transforms (dict): Dict of transforms for each modality
loader (callable): A function to load a sample given its path.
transform (callable, optional): A function/transform that takes in
a sample and returns a transformed version.
E.g, ``transforms.RandomCrop`` for images.
target_transform (callable, optional): A function/transform that takes
in the target and transforms it.
is_valid_file (callable, optional): A function that takes path of a file
and check if the file is a valid file (used to check of corrupt logs)
both extensions and is_valid_file should not be passed.
max_samples (int, optional): Maximum number of samples to load. If None, all samples are loaded.
pre_shuffle (bool, optional): Whether to shuffle the sample during the init.
return_paths (bool, optional): Whether to return the paths of the samples.
cache (bool, optional): Whether to cache the samples in memory. If True, the samples are loaded only once and then cached in memory.
Attributes:
classes (list): List of the class names sorted alphabetically.
class_to_idx (dict): Dict with items (class_name, class_index).
samples (list): List of (sample path, class_index) tuples
targets (list): The class_index value for each image in the dataset
"""
def __init__(
self,
root: str,
modalities: List[str],
modality_paths: Dict[str, str],
modality_transforms: Dict[str, AbstractTransform],
transform: Optional[Callable] = None,
target_transform: Optional[Callable] = None,
is_valid_file: Optional[Callable[[str], bool]] = None,
max_samples: Optional[int] = None,
pre_shuffle: bool = False,
cache: bool = False,
return_path: bool = False,
) -> None:
super(MultiModalDatasetFolder, self).__init__(root, transform=transform, target_transform=target_transform)
self.modalities = modalities
# If modality_paths is not provided, use the default paths
self.modality_paths = modality_paths
for mod in self.modalities:
if mod not in self.modality_paths:
modality_paths[mod] = mod
self.modality_transforms = modality_transforms
self.return_path = return_path
classes, class_to_idx = self._find_classes(os.path.join(self.root, list(self.modality_paths.values())[0]))
extensions = UNIFIED_EXTENSIONS if is_valid_file is None else None
samples = {
mod: make_dataset(
os.path.join(self.root, f'{self.modality_paths[mod]}'),
class_to_idx,
extensions,
is_valid_file,
cache_path=os.path.join(self.root, 'dataloader_cache', f'{self.modality_paths[mod]}.pkl') if cache else None)
for mod in self.modalities
}
for mod, mod_samples in samples.items():
if len(mod_samples) == 0:
msg = "Found 0 logs in subfolders of: {}\n".format(os.path.join(self.root, f'{self.modality_paths[mod]}'))
if extensions is not None:
msg += "Supported extensions are: {}".format(",".join(extensions))
raise RuntimeError(msg)
self.extensions = extensions
self.classes = classes
self.class_to_idx = class_to_idx
self.samples = samples
# Select random subset of dataset if so specified
if isinstance(max_samples, int):
total_samples = len(list(self.samples.values())[0])
np.random.seed(0)
permutation = np.random.permutation(total_samples)
for task in samples:
self.samples[task] = [self.samples[task][i] for i in permutation][:max_samples]
if pre_shuffle:
total_samples = len(list(self.samples.values())[0])
np.random.seed(100)
permutation = np.random.permutation(total_samples)
for task in samples:
self.samples[task] = [self.samples[task][i] for i in permutation]
self.cache = {}
self.imgs = self.samples
def _find_classes(self, dir: str) -> Tuple[List[str], Dict[str, int]]:
"""
Finds the class folders in a dataset.
Args:
dir (string): Root directory path.
Returns:
tuple: (classes, class_to_idx) where classes are relative to (dir), and class_to_idx is a dictionary.
Ensures:
No class is a subdirectory of another.
"""
classes = [d.name for d in os.scandir(dir) if d.is_dir()]
classes.sort()
class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
return classes, class_to_idx
def get_class_and_file(self, path: str) -> Tuple[str, str]:
""" Extracts the class and file name from a path. """
class_id, file_name = path.split('/')[-2:]
file_name = file_name.split('.')[0]
return class_id, file_name
def __getitem__(self, index: int) -> Tuple[Any, Any]:
"""
Args:
index (int): Index
Returns:
tuple: (sample, target) where target is class_index of the target class.
"""
if index in self.cache:
sample_dict, target = deepcopy(self.cache[index])
else:
sample_dict = {}
for mod in self.modalities:
path, target = self.samples[mod][index]
sample = self.modality_transforms[get_transform_key(mod)].load(path)
sample_dict[mod] = sample
# self.cache[index] = deepcopy((sample_dict, target))
if self.transform is not None:
sample_dict = self.transform(sample_dict)
if self.target_transform is not None:
target = self.target_transform(target)
sample_dict['class_idx'] = target
if self.return_path and not index in self.cache:
class_id, file_name = self.get_class_and_file(path)
sample_dict['class_id'] = class_id
sample_dict['file_name'] = file_name
return sample_dict
def __len__(self) -> int:
return len(list(self.samples.values())[0])
|