github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add ability to pull control images from same folder group

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-08 05:21:31 -06:00
parent 7bceec3b07
commit dd08579eda
2 changed files with 29 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
toolkit/config_modules.py
View File

@@ -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

28
toolkit/dataloader_mixins.py
View File

@@ -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: