From dd08579eda9268a4d200c7977914f6cd9f6d66b7 Mon Sep 17 00:00:00 2001
From: Jaret Burkett <jaretburkett@gmail.com>
Date: Wed, 8 Jul 2026 05:21:31 -0600
Subject: [PATCH] Add ability to pull control images from same folder group

---
 toolkit/config_modules.py    |  4 ++++
 toolkit/dataloader_mixins.py | 28 +++++++++++++++++++++++++---
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py
index 3a56984..9c9a04b 100644
--- a/toolkit/config_modules.py
+++ b/toolkit/config_modules.py
@@ -923,6 +923,10 @@ class DatasetConfig:
         self.flip_y: bool = kwargs.get('flip_y', False)
         self.augments: List[str] = kwargs.get('augments', [])
         self.control_path: Union[str,List[str]] = kwargs.get('control_path', None)  # depth maps, etc
+        # pull a random control image from the same folder as the image. Useful for folder grouped pairs.
+        self.control_from_same_folder: bool = kwargs.get('control_from_same_folder', False)
+        self.num_controls_from_same_folder: int = kwargs.get('num_controls_from_same_folder', 1)
+        
         if self.control_path == '':
             self.control_path = None
         
diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py
index 32861f3..cde15b5 100644
--- a/toolkit/dataloader_mixins.py
+++ b/toolkit/dataloader_mixins.py
@@ -996,11 +996,33 @@ class ControlFileItemDTOMixin:
                 # only do one
                 self.control_path = self.control_path[0]
 
+        if dataset_config.control_from_same_folder:
+            # assume we have them. We will pull them on load.
+            self.full_size_control_images = dataset_config.full_size_control_images
+            self.has_control_image = True
+
+    def get_new_control_paths(self: 'FileItemDTO'):
+        if self.dataset_config.control_from_same_folder:
+            # randomly grab image paths from the same folder as if they came from control_path
+            pool_folder = os.path.dirname(self.path)
+            # find all images in the folder
+            img_files = []
+            for ext in img_ext_list:
+                img_files += glob.glob(os.path.join(pool_folder, f'*{ext}'))
+            # remove the current image if len is greater than 1
+            if len(img_files) > 1:
+                img_files.remove(self.path)
+            num_controls = min(self.dataset_config.num_controls_from_same_folder, len(img_files))
+            # randomly grab them
+            return random.sample(img_files, num_controls)
+        else:
+            return self.control_path
+
     def load_control_image(self: 'FileItemDTO'):
         control_tensors = []
-        control_path_list = self.control_path
-        if not isinstance(self.control_path, list):
-            control_path_list = [self.control_path]
+        control_path_list = self.get_new_control_paths()
+        if not isinstance(control_path_list, list):
+            control_path_list = [control_path_list]
         
         for control_path in control_path_list:
             try: