File size: 2,065 Bytes
8acf2cd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | """HOT building tiles (OAM imagery CC BY 4.0, OSM labels ODbL): mosaic 2x2 adjacent z19 tiles (~0.3 m) -> 512 -> 256 (~0.6 m).
Outputs prep/hot_x.npy (M,256,256,3), prep/hot_y.npy (M,256,256) with 1 = building."""
import io, sys, os
from collections import defaultdict
import numpy as np, cv2, pyarrow.parquet as pq
from PIL import Image
from huggingface_hub import hf_hub_download
NSH = int(sys.argv[1]) if len(sys.argv) > 1 else 4
tiles = {}
for i in range(NSH):
p = hf_hub_download("hotosm/vhr-building-segmentation", f"data/train-{i:05d}-of-00011.parquet", repo_type="dataset", local_dir="/workspace/hfd")
t = pq.read_table(p, columns=["image", "mask", "tile_x", "tile_y", "project_id"]).to_pylist()
for r in t:
tiles[(r["project_id"], r["tile_x"], r["tile_y"])] = (r["image"]["bytes"], r["mask"]["bytes"])
os.remove(p); print("shard", i, len(tiles), flush=True)
blocks = defaultdict(dict)
for (pid, x, y) in tiles: blocks[(pid, x // 2, y // 2)][(x % 2, y % 2)] = (pid, x, y)
X, Y = [], []
def dec(b, gray=False):
im = Image.open(io.BytesIO(b)); a = np.asarray(im.convert("L" if gray else "RGB"))
return a
for k, v in blocks.items():
if len(v) < 4: continue
try:
img = np.zeros((512, 512, 3), np.uint8); msk = np.zeros((512, 512), np.uint8)
for (dx, dy), key in v.items():
ib, mb = tiles[key]; a = dec(ib); m = dec(mb, True)
if a.shape[:2] != (256, 256) or m.shape != (256, 256): raise ValueError
img[dy * 256:(dy + 1) * 256, dx * 256:(dx + 1) * 256] = a; msk[dy * 256:(dy + 1) * 256, dx * 256:(dx + 1) * 256] = (m > 0)
except Exception:
continue
if msk.mean() < 0.01 or (img.max(2) == 0).mean() > 0.05: continue
X.append(cv2.resize(img, (256, 256), interpolation=cv2.INTER_AREA)); Y.append(cv2.resize(msk, (256, 256), interpolation=cv2.INTER_NEAREST))
X, Y = np.stack(X), np.stack(Y)
np.save("/workspace/prep/hot_x.npy", X); np.save("/workspace/prep/hot_y.npy", Y)
print("HOT_DONE", X.shape, "building frac", round(float(Y.mean()), 4), flush=True)
|