From 2b4b9a97d403d89b92f48581810a77c7d5cebd8b Mon Sep 17 00:00:00 2001 From: Northumber Date: Thu, 22 May 2025 18:20:30 +0200 Subject: [PATCH] fix load image batch with indexes --- nodes/from_dir_with_index_batch.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/nodes/from_dir_with_index_batch.py b/nodes/from_dir_with_index_batch.py index 6344898..644106e 100644 --- a/nodes/from_dir_with_index_batch.py +++ b/nodes/from_dir_with_index_batch.py @@ -62,9 +62,10 @@ class LoadImagesFromDirByIndexBatch: mask_np = np.array(i.getchannel('A')).astype(np.float32) / 255.0 mask_tensor = 1. - torch.from_numpy(mask_np) else: - mask_tensor = torch.zeros((image_tensor.shape[2], image_tensor.shape[3]), dtype=torch.float32) + # use image_tensor's height and width + mask_tensor = torch.zeros((image_tensor.shape[1], image_tensor.shape[2]), dtype=torch.float32) images.append(image_tensor) masks.append(mask_tensor) - image_batch = torch.cat(images, dim=0) if len(images) > 1 else images[0] - mask_batch = torch.stack(masks, dim=0) if len(masks) > 1 else masks[0] + image_batch = torch.cat(images, dim=0) if images else torch.zeros(0) + mask_batch = torch.stack(masks, dim=0) if masks else torch.zeros(0) return (image_batch, mask_batch, len(images)) \ No newline at end of file