diff --git a/compass_image_loader.py b/compass_image_loader.py index b0be7ec..301867b 100644 --- a/compass_image_loader.py +++ b/compass_image_loader.py @@ -136,6 +136,65 @@ class CompassImageLoader: FUNCTION = "load_images" OUTPUT_NODE = False + def load_images( + self, directory, direction, modality, frame=None, width=0, height=0, + direction_in=None + ): + if direction_in and direction_in.strip(): + resolved_direction = direction_in.strip() + elif direction and direction.strip(): + resolved_direction = direction.strip() + else: + resolved_direction = "" + + base_dir = folder_paths.get_input_directory() + target_dir = _resolve_target_dir(base_dir, directory, resolved_direction) + modality_path = os.path.join(target_dir, modality) + + if not os.path.isdir(modality_path): + raise RuntimeError(f"Compass directory not found: {modality_path}") + + files = _list_image_files(modality_path) + if not files: + raise RuntimeError(f"No images found in: {modality_path}") + + if frame is None or str(frame).strip() == "": + selected_files = files + output_path = modality_path + else: + try: + index = int(str(frame).strip()) + except (ValueError, TypeError): + raise RuntimeError( + f"Invalid frame number: '{frame}'. Must be an integer." + ) + + if index < 0 or index >= len(files): + raise RuntimeError( + f"Frame index {index} out of bounds. " + f"Found {len(files)} images in {modality_path}." + ) + + selected_files = [files[index]] + output_path = os.path.join(modality_path, files[index]) + + tensors = [] + final_w, final_h = 0, 0 + + for filename in selected_files: + filepath = os.path.join(modality_path, filename) + image = Image.open(filepath).convert("RGB") + image, final_w, final_h = _resize_image(image, width, height) + + np_arr = np.array(image).astype(np.float32) / 255.0 + tensors.append(torch.from_numpy(np_arr)[None,]) + + image_batch = ( + tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0) + ) + + return (image_batch, output_path, resolved_direction, final_w, final_h, len(selected_files)) + @classmethod def IS_CHANGED( cls, directory, direction, modality, frame=None, width=0, height=0,