Added Ptv Image Folder Dataset node.

modules/pytorch_wrapper/ptv_image_folder_dataset.py

@@ -0,0 +1,73 @@
+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,)