74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
from typing import Any, Dict, Optional, Callable, Tuple
|
|
import torch
|
|
from torchvision.datasets import ImageFolder
|
|
|
|
from .utils import get_dataset_full_path
|
|
|
|
class PtvImageFolderDataset:
|
|
"""
|
|
A Torchvision ImageFolder Dataset class wrapper.
|
|
|
|
Parameters
|
|
|
|
**root**
|
|
Specify the root directory that contains subdirectories which of which contains images for each class.
|
|
The subdirectory name needs to be named after each class.
|
|
|
|
**transform**
|
|
Data transforms e.g. transform to PyTorch tensor, normalize image.
|
|
Plug in the PTV Transforms node that contains Torchvision Transforms functionality.
|
|
|
|
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": {
|
|
"root": ("STRING", {"default": "", "multiline": False})
|
|
},
|
|
"optional": {
|
|
"transform": ("PTVTRANSFORM", {})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES: Tuple[str] = ("PTVDATASET",)
|
|
FUNCTION: str = "f"
|
|
CATEGORY: str = "Training"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kw):
|
|
return float("NaN")
|
|
|
|
def f(self,
|
|
root: str,
|
|
transform: Optional[Callable] = None) -> Tuple:
|
|
"""
|
|
Loads a dataset from Torchvision with specified parameters.
|
|
|
|
Args:
|
|
root (str): Root directory for storing the dataset.
|
|
transform (Optional[Callable]): Transformations to apply to the dataset.
|
|
|
|
Returns:
|
|
Tuple: A tuple containing the dataset instance.
|
|
"""
|
|
|
|
with torch.inference_mode(False):
|
|
|
|
if root:
|
|
dataset_path = get_dataset_full_path(root)
|
|
else: # Use default directory
|
|
dataset_path = get_dataset_full_path("")
|
|
|
|
dc = ImageFolder(root=dataset_path,
|
|
transform=transform
|
|
)
|
|
return (dc,)
|