Problem with custom intersection dataset

[lcoandrade]

Hi there,
I'm working on a Kaggle notebook to perform building footprint extraction.

I'm using this dataset. That has 2 folders to train (gt and images). I'm using a RasterDataset.

Both gt and images have the same extents and CRS:
GT:

Width: 5000
Height: 5000
Transform: | 0.30, 0.00, 616500.00|
| 0.00,-0.30, 3345000.00|
| 0.00, 0.00, 1.00|
CRS: EPSG:26914

Images:

Width: 5000
Height: 5000
Transform: | 0.30, 0.00, 616500.00|
| 0.00,-0.30, 3345000.00|
| 0.00, 0.00, 1.00|
CRS: EPSG:26914

I create the datasets like this:

image_set = RasterDataset(
    paths=os.path.join(INPUT_DIR, TRAIN_DIR, IMG_DIR)
)

gt_set = RasterDataset(
    paths=os.path.join(INPUT_DIR, TRAIN_DIR, LABEL_DIR)
)

Printing the datasets, I get:

RasterDataset Dataset
    type: GeoDataset
    bbox: BoundingBox(minx=-6823569.7720385045, maxx=6776.0, miny=237532.76998160873, maxy=8895035.624662481, mint=0.0, maxt=9.223372036854776e+18)
    size: 180
RasterDataset Dataset
    type: GeoDataset
    bbox: BoundingBox(minx=-2165356.726339508, maxx=5727468.375156833, miny=3391632.848544845, maxy=11305853.12664395, mint=0.0, maxt=9.223372036854776e+18)
    size: 180

To check the dataset, I'm creating a data loader to use with visualization code:

dataset = image_set & gt_set
sampler = RandomGeoSampler(dataset, size=IMG_SIZE, length=MAX_WINDOWS)
dataloader = DataLoader(dataset, sampler=sampler, collate_fn=stack_samples)

This is my visualization code:

reverse_transform = transforms.Compose(
    [transforms.ToPILImage()]
)

# helper function for data visualization
def visualize(**images):
    """PLot images in one row."""
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(' '.join(name.split('_')).title())
        plt.imshow(image)
    plt.show()

for n, sample in enumerate(dataloader):
    if n == 5:
        break
        
    image, gt_mask = sample['image'], sample['mask']
             
    visualize(
        image=reverse_transform(image.squeeze()[:3]),
        ground_truth = gt_mask.squeeze(),
    )

Which gives me this error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[18], line 26
     20     break
     22 image, gt_mask = sample['image'], sample['mask']
     24 visualize(
     25     image=reverse_transform(image.squeeze()[:3]),
---> 26     ground_truth = gt_mask.squeeze(),
     27 )

AttributeError: 'list' object has no attribute 'squeeze'

Checking the samples (print(sample)), I get this:

defaultdict(<class 'list'>, {'crs': [CRS.from_epsg(31256)], 'bbox': [BoundingBox(minx=-5824671.499583863, maxx=-5824594.699583863, miny=6680402.841688732, maxy=6680479.6416887315, mint=0.0, maxt=9.223372036854776e+18)], 'image': tensor([[[[147.,  39.,  44.,  ..., 186., 182., 182.],
          [115.,  80.,  84.,  ..., 144., 153., 153.],
          [110., 129., 127.,  ..., 170., 164., 164.],
          ...,
          [153., 155., 157.,  ..., 239., 239., 239.],
          [148., 156., 157.,  ..., 219., 242., 242.],
          [153., 155., 157.,  ..., 137., 251., 251.]],

         [[161.,  38.,  41.,  ..., 189., 185., 185.],
          [128.,  84.,  93.,  ..., 150., 158., 158.],
          [121., 138., 134.,  ..., 168., 166., 166.],
          ...,
          [152., 155., 156.,  ..., 244., 244., 244.],
          [146., 156., 156.,  ..., 221., 245., 245.],
          [153., 155., 156.,  ..., 143., 253., 253.]],

         [[157.,  37.,  47.,  ..., 183., 183., 183.],
          [122.,  88.,  92.,  ..., 135., 146., 146.],
          [112., 135., 124.,  ..., 162., 158., 158.],
          ...,
          [134., 136., 138.,  ..., 238., 238., 238.],
          [129., 137., 138.,  ..., 217., 247., 247.],
          [134., 136., 138.,  ..., 138., 252., 252.]],

         [[255., 255., 255.,  ..., 255., 255., 255.],
          [255., 255., 255.,  ..., 255., 255., 255.],
          [255., 255., 255.,  ..., 255., 255., 255.],
          ...,
          [  0.,   0.,   0.,  ..., 255., 255., 255.],
          [  0.,   0.,   0.,  ..., 255., 255., 255.],
          [  0.,   0.,   0.,  ..., 255., 255., 255.]]]])})

There is no 'mask' in it.

How can I solve this?

[adamjstewart]

By default, RasterDataset assumes that all files are images, you have to tell it one of those is a mask. See https://torchgeo.readthedocs.io/en/stable/tutorials/custom_raster_dataset.html.

[lcoandrade]

Thanks!!!!