End-to-end Reference Implementation

dataset.py — patches + DataLoader

Save to /content/drive/MyDrive/ceamls-project/notebooks/dataset.py.

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    ])
15def val_aug():
16    return A.Compose([ToTensorV2()])
17
18class XBDPatchDataset(Dataset):
19    def __init__(self, index_csv, patch_dir, transform=None, siamese=False):
20        self.idx = pd.read_csv(index_csv)
21        self.patch_dir = patch_dir
22        self.transform = transform
23        self.siamese = siamese
24    def __len__(self):
25        return len(self.idx)
26    def __getitem__(self, i):
27        row = self.idx.iloc[i]
28        post = np.load(f"{self.patch_dir}/{row['patch']}")
29        pre  = np.load(f"{self.patch_dir}/{row['patch'].replace('_post_','_pre_')}") \
30                if self.siamese else post.copy()
31        cls = CLASS_TO_ID[row['class']]
32        if self.transform:
33            post = self.transform(image=post)['image']
34            pre  = self.transform(image=pre)['image']
35        return {'pre': pre, 'post': post, 'label': torch.tensor(cls, dtype=torch.long)}
36
37def get_loaders(index_csv, patch_dir, batch=16, siamese=False):
38    full = XBDPatchDataset(index_csv, patch_dir, transform=train_aug(), siamese=siamese)
39    n = len(full); n_tr, n_va = int(0.8*n), int(0.1*n)
40    tr, va, te = random_split(full, [n_tr, n_va, n - n_tr - n_va],
41                              generator=torch.Generator().manual_seed(42))
42    return (DataLoader(tr, batch, shuffle=True,  num_workers=2, pin_memory=True),
43            DataLoader(va, batch, shuffle=False, num_workers=2, pin_memory=True),
44            DataLoader(te, batch, shuffle=False, num_workers=2, pin_memory=True))

metrics.py — OA, weighted F1, mIoU

python
1import torch
2from sklearn.metrics import f1_score, accuracy_score, jaccard_score
3
4@torch.no_grad()
5def evaluate(model, loader, device, siamese=False):
6    model.eval()
7    all_preds, all_lbls = [], []
8    for batch in loader:
9        if siamese:
10            logits = model(batch['pre'].to(device), batch['post'].to(device))
11        else:
12            logits = model(batch['post'].to(device))
13        all_preds.extend(logits.argmax(dim=1).cpu().tolist())
14        all_lbls.extend(batch['label'].tolist())
15    return {
16        'OA':           accuracy_score(all_lbls, all_preds),
17        'wF1':          f1_score(all_lbls, all_preds, average='weighted'),
18        'macroF1':      f1_score(all_lbls, all_preds, average='macro'),
19        'mIoU':         jaccard_score(all_lbls, all_preds, average='macro'),
20        'f1_destroyed': f1_score(all_lbls, all_preds, average=None)[3],
21    }

train_baseline.py — ResNet-50 + U-Net

python
1import torch, wandb, segmentation_models_pytorch as smp
2from dataset import get_loaders
3from metrics import evaluate
4
5device = 'cuda'
6wandb.init(project='ceamls-baseline')
7model = smp.Unet('resnet50', encoder_weights='imagenet', in_channels=3, classes=4).to(device)
8tr, va, te = get_loaders(IDX, PATCHES, batch=16)
9
10W = torch.tensor([0.07, 0.99, 0.65, 14.3]).to(device)
11criterion = torch.nn.CrossEntropyLoss(weight=W)
12optim     = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
13
14for epoch in range(30):
15    model.train()
16    for b in tr:
17        logits = model(b['post'].to(device))
18        loss = criterion(logits, b['label'].to(device))
19        optim.zero_grad(); loss.backward(); optim.step()
20        wandb.log({'train_loss': loss.item()})
21    wandb.log({'epoch': epoch, **evaluate(model, va, device)})
22torch.save(model.state_dict(), 'results/baseline.pt')

train_siamese.py — pre/post change detection

python
1import torch, torch.nn as nn, segmentation_models_pytorch as smp
2from dataset import get_loaders
3from metrics import evaluate
4
5class SiameseDamageNet(nn.Module):
6    def __init__(self, num_classes=4):
7        super().__init__()
8        self.shared = smp.encoders.get_encoder('resnet50', weights='imagenet')
9        self.head = nn.Sequential(
10            nn.AdaptiveAvgPool2d(1), nn.Flatten(),
11            nn.Linear(2048, 256), nn.ReLU(),
12            nn.Linear(256, num_classes),
13        )
14    def forward(self, pre, post):
15        return self.head(torch.abs(self.shared(pre)[-1] - self.shared(post)[-1]))
16
17model = SiameseDamageNet().cuda()
18tr, va, te = get_loaders(IDX, PATCHES, batch=8, siamese=True)
19optim = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
20W = torch.tensor([0.07, 0.99, 0.65, 14.3]).cuda()
21criterion = torch.nn.CrossEntropyLoss(weight=W)
22for epoch in range(25):
23    model.train()
24    for b in tr:
25        logits = model(b['pre'].cuda(), b['post'].cuda())
26        loss = criterion(logits, b['label'].cuda())
27        optim.zero_grad(); loss.backward(); optim.step()
28    print(epoch, evaluate(model, va, 'cuda', siamese=True))
29torch.save(model.state_dict(), 'results/siamese.pt')

export_onnx.py — for the FastAPI app

python
1import torch
2from train_siamese import SiameseDamageNet
3
4model = SiameseDamageNet().cuda()
5model.load_state_dict(torch.load('results/siamese.pt'))
6model.eval()
7torch.onnx.export(
8    model,
9    (torch.randn(1,3,224,224).cuda(), torch.randn(1,3,224,224).cuda()),
10    'results/model.onnx',
11    input_names=['pre','post'], output_names=['logits'],
12    dynamic_axes={'pre':{0:'B'},'post':{0:'B'},'logits':{0:'B'}},
13    opset_version=17,
14)
15print('Exported results/model.onnx')

main.py — FastAPI deployment app (Step 14)

This is the field-inspector web app that S4 deploys. A browser uploads a building photo or GeoTIFF → the API returns the ATC-20 placard.

python
1from fastapi import FastAPI, UploadFile, File
2from fastapi.responses import JSONResponse, HTMLResponse
3import onnxruntime, numpy as np, rasterio, cv2, tempfile, os
4
5app = FastAPI(title='CEAMLS Building Damage Classifier')
6sess = onnxruntime.InferenceSession('model.onnx')
7
8CLASS_NAMES  = {0:'No Damage', 1:'Minor Damage', 2:'Major Damage', 3:'Destroyed'}
9ATC20_TAGS   = {0:'GREEN - Safe', 1:'YELLOW - Restricted', 2:'YELLOW - Restricted', 3:'RED - Unsafe'}
10ATC20_COLORS = {0:'#16A34A', 1:'#CA8A04', 2:'#CA8A04', 3:'#DC2626'}
11
12@app.get('/', response_class=HTMLResponse)
13async def home():
14    with open('index.html') as f: return f.read()
15
16@app.post('/classify')
17async def classify(file: UploadFile = File(...)):
18    suffix = '.tif' if file.filename.endswith('.tif') else '.jpg'
19    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
20        tmp.write(await file.read()); tmp_path = tmp.name
21
22    if suffix == '.tif':
23        with rasterio.open(tmp_path) as src:
24            img = src.read([5,3,2]).astype(np.float32) / 65535.0
25            img = np.transpose(img, (1,2,0))
26    else:
27        img = cv2.imread(tmp_path).astype(np.float32) / 255.0
28        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
29    os.unlink(tmp_path)
30
31    img = cv2.resize(img, (224,224))
32    img = np.transpose(img, (2,0,1))[np.newaxis].astype(np.float32)
33
34    logits = sess.run(None, {'input': img})[0]
35    cls    = int(np.argmax(logits))
36    return JSONResponse({
37        'class':      CLASS_NAMES[cls],
38        'atc20_tag':  ATC20_TAGS[cls],
39        'color':      ATC20_COLORS[cls],
40        'confidence': round(float(np.max(logits)), 3),
41    })

python
1import nest_asyncio, uvicorn, shutil, sys
2nest_asyncio.apply()
3shutil.copy('results/model.onnx', 'model.onnx')
4shutil.copy('notebooks/index.html', 'index.html')
5sys.path.insert(0, 'notebooks')
6from main import app
7uvicorn.run(app, host='0.0.0.0', port=8000)
8# Colab prints a URL — open it in a new Chrome tab to test the app