TorchGeo is a PyTorch domain library similar to torchvision, specialized in geospatial data, offering datasets, samplers, transformations, and pre-trained models dedicated to this field. This guide introduces how to easily download the NWPU VHR-10 dataset from TorchGeo.

The VHR-10 dataset is a collection of Very High Resolution remote sensing images of 10 classes, provided by Northwestern Polytechnical University (NWPU) in China.

Comprising a total of 800 VHR optical remote sensing images, 715 of these are color images acquired from Google Earth with spatial resolutions between 0.5 to 2 meters. The remaining 85 are color-infrared (CIR) images pan-sharpened, obtained from the Vaihingen dataset with a spatial resolution of 0.08 meters.

The dataset is divided into two sets:

Positive image set (650 images): Images containing at least one object of interest.
Negative image set (150 images): Images with no objects of interest.

The positive image set includes objects from the following ten classes:

Airplanes (757)
Ships (302)
Storage tanks (655)
Baseball diamonds (390)
Tennis courts (524)
Basketball courts (159)
Ground track fields (163)
Harbors (224)
Bridges (124)
Vehicles (477)

It includes object detection bounding boxes and instance segmentation masks.

The TorchGeo documentation guides the use of additional libraries for utilizing this dataset as follows:

Since the dataset is stored in RAR files, [rarfile](https://pypi.org/project/rarfile/) is needed for extraction.
[pycocotools ](https://pypi.org/project/pycocotools/)is required to load the annotations.json file for the “positive” image set.

However, there are additional considerations beyond the libraries mentioned that need to be taken into account. In this context, I will share information obtained from the How to Download the VHR Dataset discussion.

Basic Code for Downloading the VHR-10 Dataset

TorchGeo installs only a core set of dependencies with pip install torchgeo to keep the installation relatively lightweight. For a full installation that includes the optional set of dependencies, you can use pip install torchgeo[datasets]. If you opt for the latter, rarfile and pycocotools, which are required for using the VHR-10 dataset, will be installed automatically. For more information, you might want to refer to the Required vs. optional dependencies discussion.

pip install torchgeo: Installs the "Required" dependency set.
pip install torchgeo[datasets]: Installs the full package including "Optional" dependency set.

%pip install -q -U torchgeo[datasets]

Import the necessary libraries and check the version of TorchGeo.

import torchgeo
from torchgeo.datasets import VHR10
import matplotlib.pyplot as plt

torchgeo.version

'0.5.2'

The code below downloads the VHR10 dataset to the location data/VHR10/ with the 'positive' split and applies transformations using the given preprocessing function. The role of the preprocess function is to convert each image into a floating-point type and normalize the values between 0 and 1 by dividing by 255. This process helps in preparing the data in a form that is suitable for model training.

def preprocess(sample):
sample["image"] = sample["image"].float() / 255.0
return sample

ds = VHR10(
root="data/VHR10/",
split="positive",
transforms=preprocess,
download=True,
checksum=True,
)

Error Type 1: gdown Module Not Installed

TorchGeo downloads the NWPU VHR-10 dataset.rar file from Google Drive. If the gdown module is not installed on your system, you may encounter the ModuleNotFoundError: No module named 'gdown' error. Therefore, execute the pip install gdown command to install the gdown module additionally.

%pip install -q -U gdown

Error Type 2: unrar Not Installed

The rarfile package requires unrar to decompress files in the rar format. If unrar is not installed on your system, you may encounter the RarCannotExec: Cannot find working tool or RarCannotExec: Unrar not installed error. In such cases, unrar needs to be installed directly. I resolved the issue by installing the Complete package file on Windows and added 'C:\Program Files (x86)\GnuWin32\bin' to the PATH in the environment variables.
Error Type 3: Google Drive File URL Retrieval Error

While installing the gdown module is unnecessary on Google Colab, you may encounter another error: FileURLRetrievalError: Cannot retrieve the public link of the file. You may need to change the permission to 'Anyone with the link', or have had many accesses.. The issue of downloading datasets from Google Drive has been a topic of discussion for a long time. In Colab, you can bypass this error by first downloading the NWPU VHR-10 dataset.rar file using the modified code below, then running the original code again. This approach allows you to use the VHR-10 dataset in Colab without issues.

import os, gdown

os.makedirs('data/VHR10/', exist_ok=True)

url = 'https://drive.google.com/uc?id=1--foZ3dV5OCsqXQXT84UeKtrAqc5CkAE'
output_path = 'data/VHR10/NWPU VHR-10 dataset.rar'
gdown.download(url, output_path, quiet=False)

Exploring the VHR-10 Dataset

Let’s now take a look at the positive image set. There are a total of 650 images.

print(f"VHR-10 dataset: {len(ds)}")

VHR-10 dataset: 650

To access the first item in the ds dataset and check the dimensions of the image corresponding to the "image" key, you can use the shape attribute. The shape attribute returns the dimensions of the image in the form of (number of channels, height, width).

ds[0]["image"].shape

torch.Size([3, 808, 958])

To open and display the sixth item’s image from the dataset, you’ll need to change its dimensions from (number of channels, height, width) to (height, width, number of channels).

image = ds[5]["image"].permute(1, 2, 0)
plt.imshow(image)
plt.show()

The VHR10 dataset in TorchGeo features a plot method that visualizes annotations included in the annotations.json. These annotations provide information for object detection and instance segmentation, including visual elements like bounding boxes and masks. This allows for direct visualization and verification of images marked with ground truth.

You can experience an example of object detection training in TorchGeo through torchgeo_object_detection_example.ipynb provided by Caleb Robinson from Microsoft AI for Earth.

ds.plot(ds[5])
plt.savefig('ground_truth.png', bbox_inches='tight')
plt.show()

Github code