125 lines
4.2 KiB
Python
125 lines
4.2 KiB
Python
import inspect
|
|
from typing import Any, Dict, Optional, Callable, Tuple
|
|
import torch
|
|
import torchvision.datasets as datasets
|
|
import ast
|
|
|
|
from .utils import get_dataset_full_path
|
|
|
|
class PtvDataset:
|
|
"""
|
|
A Torchvision Dataset class wrapper.
|
|
|
|
Parameters
|
|
|
|
**name**
|
|
Specify a dataset class name such as MNIST, FashionMNIST in string.
|
|
This will be converted to a dataset class internally.
|
|
|
|
**download**
|
|
Set to True if you want to download.
|
|
|
|
**root**
|
|
Specify the root directory of the downloaded dataset relative to the
|
|
dataset directory of this extension, or specify the absolute path.
|
|
If you leave as a blank, the dataset will be downloaded under the `datasets` directory
|
|
of this extension.
|
|
|
|
**transform**
|
|
Data transforms e.g. transform to PyTorch tensor, normalize image.
|
|
Plug in the PTV Transforms node that contains Torchvision Transforms functionality.
|
|
|
|
**parameters**
|
|
Specify other parameters in Python dict format.
|
|
For example, if you want to set Train parameter to True to download the train set, specify:
|
|
```
|
|
{"train": True}
|
|
```
|
|
or if you want to download the test set, specify:
|
|
```
|
|
{"train": False}
|
|
```
|
|
|
|
category: PyTorch wrapper - Training
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
"""
|
|
Defines the required input types.
|
|
|
|
Returns:
|
|
Dict[str, Any]: A dictionary of required input types.
|
|
"""
|
|
return {
|
|
"required": {
|
|
"name": ("STRING", {"default": "", "multiline": False}),
|
|
"download": ("BOOLEAN", {"default": True}),
|
|
"root": ("STRING", {"default": "", "multiline": False}),
|
|
},
|
|
"optional": {
|
|
"transform": ("PTVTRANSFORM", {}),
|
|
"parameters": ("STRING", {"default": "", "multiline": True})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str] = ("PTVDATASET",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Training"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kw):
|
|
return float("NaN")
|
|
|
|
def f(self,
|
|
name: str,
|
|
download: bool,
|
|
root: str,
|
|
parameters: str,
|
|
transform: Optional[Callable] = None) -> Tuple:
|
|
"""
|
|
Loads a dataset from Torchvision with specified parameters.
|
|
|
|
Args:
|
|
name (str): Name of the dataset class (e.g., MNIST, FashionMNIST).
|
|
download (bool): Whether to download the dataset.
|
|
root (str): Root directory for storing the dataset.
|
|
parameters (str): Additional dataset parameters in stringified dictionary format.
|
|
transform (Optional[Callable]): Transformations to apply to the dataset.
|
|
|
|
Returns:
|
|
Tuple: A tuple containing the dataset instance.
|
|
"""
|
|
|
|
with torch.inference_mode(False):
|
|
param_dict = ast.literal_eval(parameters)
|
|
|
|
# Converts the datasets class name to dataset class itself
|
|
try:
|
|
dataset_class = getattr(datasets, name)
|
|
except AttributeError:
|
|
raise ValueError(f"Dataset '{name}' is not available in torchvision.datasets")
|
|
|
|
# Get valid parameters for the dataset class
|
|
valid_params = inspect.signature(dataset_class).parameters
|
|
|
|
# Filter only parameters that match the dataset's constructor
|
|
filtered_params = {k: v for k, v in param_dict.items() if k in valid_params}
|
|
|
|
if root:
|
|
dataset_path = get_dataset_full_path(root)
|
|
else: # Use default directory
|
|
dataset_path = get_dataset_full_path("")
|
|
|
|
if transform:
|
|
dc = dataset_class(download=download,
|
|
root=dataset_path,
|
|
transform=transform,
|
|
**filtered_params)
|
|
else:
|
|
dc = dataset_class(download=download,
|
|
root=dataset_path,
|
|
**filtered_params)
|
|
|
|
return (dc,)
|