Interactive · Damage Assessment

Code Explorer · xBD Study

Every Python file the CEAMLS team writes for the post-disaster damage assessment paper. Pick a file, then click any line to see what it does and why it's there.

S1's deliverable. Reads xBD GeoTIFF patches, applies Albumentations augmentation, and serves train/val/test DataLoaders to S2 and S3.

notebooks/dataset.pyclick any line

python
1import torch, numpy as np, pandas as pd
2from torch.utils.data import Dataset, DataLoader, random_split
3import albumentations as A
4from albumentations.pytorch import ToTensorV2
5
6CLASS_TO_ID = {'no-damage':0, 'minor-damage':1, 'major-damage':2, 'destroyed':3}
7
8def train_aug():
9    return A.Compose([
10        A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.5),
11        A.RandomRotate90(p=0.5),
12        A.RandomBrightnessContrast(0.2, 0.2, p=0.5),
13        ToTensorV2(),
14    ])
15
16def val_aug():
17    return A.Compose([ToTensorV2()])  # NEVER augment validation
18
19class XBDPatchDataset(Dataset):
20    def __init__(self, index_csv, patch_dir, transform=None, siamese=False):
21        self.idx = pd.read_csv(index_csv)
22        self.patch_dir = patch_dir
23        self.transform = transform
24        self.siamese = siamese
25    def __len__(self):
26        return len(self.idx)
27    def __getitem__(self, i):
28        row = self.idx.iloc[i]
29        post = np.load(f"{self.patch_dir}/{row['patch']}")
30        pre  = np.load(f"{self.patch_dir}/{row['patch'].replace('_post_','_pre_')}") \
31                if self.siamese else post.copy()
32        cls = CLASS_TO_ID[row['class']]
33        if self.transform:
34            post = self.transform(image=post)['image']
35            pre  = self.transform(image=pre)['image']
36        return {'pre': pre, 'post': post, 'label': torch.tensor(cls, dtype=torch.long)}
37
38def get_loaders(index_csv, patch_dir, batch=16, siamese=False):
39    full = XBDPatchDataset(index_csv, patch_dir, transform=train_aug(), siamese=siamese)
40    n = len(full); n_tr, n_va = int(0.8*n), int(0.1*n)
41    tr, va, te = random_split(full, [n_tr, n_va, n - n_tr - n_va],
42                              generator=torch.Generator().manual_seed(42))
43    return (DataLoader(tr, batch, shuffle=True,  num_workers=2, pin_memory=True),
44            DataLoader(va, batch, shuffle=False, num_workers=2, pin_memory=True),
45            DataLoader(te, batch, shuffle=False, num_workers=2, pin_memory=True))

Line 1

What it does

Import PyTorch + NumPy + Pandas.

Why it's needed

Tensors, arrays, and the CSV index of patches.

Files in this study