Spaces:
Sleeping
Sleeping
| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import logging | |
| import numpy as np | |
| import os | |
| from typing import Optional, Callable, Set | |
| import torch | |
| from torchvision.datasets.vision import VisionDataset | |
| from torchvision.transforms import ToTensor | |
| from fairseq.data import FairseqDataset | |
| logger = logging.getLogger(__name__) | |
| class ImageDataset(FairseqDataset, VisionDataset): | |
| def __init__( | |
| self, | |
| root: str, | |
| extensions: Set[str], | |
| load_classes: bool, | |
| transform: Optional[Callable] = None, | |
| shuffle=True, | |
| ): | |
| FairseqDataset.__init__(self) | |
| VisionDataset.__init__(self, root=root, transform=transform) | |
| self.shuffle = shuffle | |
| self.tensor_transform = ToTensor() | |
| self.classes = None | |
| self.labels = None | |
| if load_classes: | |
| classes = [d.name for d in os.scandir(root) if d.is_dir()] | |
| classes.sort() | |
| self.classes = {cls_name: i for i, cls_name in enumerate(classes)} | |
| logger.info(f"loaded {len(self.classes)} classes") | |
| self.labels = [] | |
| def walk_path(root_path): | |
| for root, _, fnames in sorted(os.walk(root_path, followlinks=True)): | |
| for fname in sorted(fnames): | |
| fname_ext = os.path.splitext(fname) | |
| if fname_ext[-1].lower() not in extensions: | |
| continue | |
| path = os.path.join(root, fname) | |
| yield path | |
| logger.info(f"finding images in {root}") | |
| if self.classes is not None: | |
| self.files = [] | |
| self.labels = [] | |
| for c, i in self.classes.items(): | |
| for f in walk_path(os.path.join(root, c)): | |
| self.files.append(f) | |
| self.labels.append(i) | |
| else: | |
| self.files = [f for f in walk_path(root)] | |
| logger.info(f"loaded {len(self.files)} examples") | |
| def __getitem__(self, index): | |
| from PIL import Image | |
| fpath = self.files[index] | |
| with open(fpath, "rb") as f: | |
| img = Image.open(f).convert("RGB") | |
| if self.transform is None: | |
| img = self.tensor_transform(img) | |
| else: | |
| img = self.transform(img) | |
| assert torch.is_tensor(img) | |
| res = {"id": index, "img": img} | |
| if self.labels is not None: | |
| res["label"] = self.labels[index] | |
| return res | |
| def __len__(self): | |
| return len(self.files) | |
| def collater(self, samples): | |
| if len(samples) == 0: | |
| return {} | |
| collated_img = torch.stack([s["img"] for s in samples], dim=0) | |
| res = { | |
| "id": torch.LongTensor([s["id"] for s in samples]), | |
| "net_input": { | |
| "img": collated_img, | |
| }, | |
| } | |
| if "label" in samples[0]: | |
| res["net_input"]["label"] = torch.LongTensor([s["label"] for s in samples]) | |
| return res | |
| def num_tokens(self, index): | |
| return 1 | |
| def size(self, index): | |
| return 1 | |
| def ordered_indices(self): | |
| """Return an ordered list of indices. Batches will be constructed based | |
| on this order.""" | |
| if self.shuffle: | |
| order = [np.random.permutation(len(self))] | |
| else: | |
| order = [np.arange(len(self))] | |
| return order[0] | |