microsoft / torchgeo

TorchGeo: datasets, samplers, transforms, and pre-trained models for geospatial data
https://www.osgeo.org/projects/torchgeo/
MIT License
2.32k stars 298 forks source link

VHR10 plot fails for negative split #2094

Closed robmarkcole closed 1 month ago

robmarkcole commented 1 month ago

Description

As title - there is a bug since if there are no annotations (true for the negative split) the sample keys change

Steps to reproduce

from torchgeo.datasets import VHR10

# Initialize the dataset
dataset = VHR10(
    root="data", 
    split="negative",
    download=False, 
    checksum=False
)

sample = dataset[0]
fig = dataset.plot(sample)

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[5], line 2
      1 sample = dataset[0]
----> 2 fig = dataset.plot(sample)

File ~/torchgeo/torchgeo/datasets/vhr10.py:386, in VHR10.plot(self, sample, show_titles, suptitle, show_feats, box_alpha, mask_alpha)
    384 assert show_feats in {'boxes', 'masks', 'both'}
    385 image = percentile_normalization(sample['image'].permute(1, 2, 0).numpy())
--> 386 boxes = sample['boxes'].cpu().numpy()
    387 labels = sample['labels'].cpu().numpy()
    389 if self.split == 'negative':

KeyError: 'boxes'

Version

master

ashnair1 commented 1 month ago

This is due to changes introduced in your PR - Refer to https://github.com/microsoft/torchgeo/pull/2092/files#r1617472554