Spaces:
Runtime error
Runtime error
| # 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 torch | |
| from . import BaseWrapperDataset | |
| class ColorizeDataset(BaseWrapperDataset): | |
| """ Adds 'colors' property to net input that is obtained from the provided color getter for use by models """ | |
| def __init__(self, dataset, color_getter): | |
| super().__init__(dataset) | |
| self.color_getter = color_getter | |
| def collater(self, samples): | |
| base_collate = super().collater(samples) | |
| if len(base_collate) > 0: | |
| base_collate["net_input"]["colors"] = torch.tensor( | |
| list(self.color_getter(self.dataset, s["id"]) for s in samples), | |
| dtype=torch.long, | |
| ) | |
| return base_collate | |