File size: 3,075 Bytes
eea5f0e | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """Resample HF single-date datasets to ~0.55-0.65 m/px and pack into npy arrays.
FLAIR (0.2 m, 512^2) -> 192^2 (~0.53 m/px). classes: 1 building, 18 greenhouse -> building;
6 conifer, 7 deciduous, 16 mixed, 17 ligneous -> tree; 2,4,10,11,12,15 -> ground donors.
OAM-TCD (0.1 m, 2048^2) -> 320^2 (0.64 m/px). annotation > 0 -> tree.
"""
import glob
import io
import os
from concurrent.futures import ProcessPoolExecutor
import cv2
import numpy as np
import pyarrow.parquet as pq
import tifffile
from PIL import Image
OUT = "/workspace/prep"
os.makedirs(OUT, exist_ok=True)
BUILD = [1, 18]
TREE = [6, 7, 16, 17]
GROUND = [2, 4, 10, 11, 12, 15]
def flair_one(img_path):
msk_path = img_path.replace("/aerial/", "/labels/").replace("IMG_", "MSK_")
if not os.path.exists(msk_path):
cand = glob.glob(os.path.join(os.path.dirname(img_path).replace("img", "msk"), "MSK_" + img_path.split("IMG_")[-1]))
if not cand:
return None
msk_path = cand[0]
im = tifffile.imread(img_path)
if im.shape[0] in (4, 5):
im = np.transpose(im, (1, 2, 0))
rgb = np.ascontiguousarray(im[..., :3]).astype(np.uint8)
m = tifffile.imread(msk_path)
if m.ndim == 3:
m = m[0] if m.shape[0] < m.shape[-1] else m[..., 0]
rgb = cv2.resize(rgb, (192, 192), interpolation=cv2.INTER_AREA)
lab = np.zeros(m.shape, np.uint8)
lab[np.isin(m, TREE)] = 2
lab[np.isin(m, GROUND)] = 3
lab[np.isin(m, BUILD)] = 1
lab = cv2.resize(lab, (192, 192), interpolation=cv2.INTER_NEAREST)
return rgb, lab
def main():
imgs = sorted(glob.glob("/workspace/data/flair/unz/**/IMG_*.tif", recursive=True))
print("flair imgs", len(imgs), flush=True)
with ProcessPoolExecutor(48) as ex:
res = [r for r in ex.map(flair_one, imgs, chunksize=32) if r is not None]
fx = np.stack([r[0] for r in res]); fy = np.stack([r[1] for r in res])
np.save(f"{OUT}/flair_x.npy", fx); np.save(f"{OUT}/flair_y.npy", fy)
print("flair", fx.shape, "building frac", (fy == 1).mean(), "tree frac", (fy == 2).mean(), flush=True)
xs, ys = [], []
for f in sorted(glob.glob("/workspace/data/tcd/data/train-*.parquet")):
t = pq.read_table(f, columns=["image", "annotation"]).to_pylist()
for row in t:
im = np.asarray(Image.open(io.BytesIO(row["image"]["bytes"])).convert("RGB"))
an = np.asarray(Image.open(io.BytesIO(row["annotation"]["bytes"])))
if an.ndim == 3:
an = an[..., 0]
s = 320 * max(im.shape[:2]) // 2048
xs.append(cv2.resize(im, (s, s), interpolation=cv2.INTER_AREA) if s == 320 else cv2.resize(im, (320, 320), interpolation=cv2.INTER_AREA))
ys.append(cv2.resize(((an > 0) * 2).astype(np.uint8), (320, 320), interpolation=cv2.INTER_NEAREST))
print(f, len(xs), flush=True)
tx = np.stack(xs); ty = np.stack(ys)
np.save(f"{OUT}/tcd_x.npy", tx); np.save(f"{OUT}/tcd_y.npy", ty)
print("tcd", tx.shape, "tree frac", (ty == 2).mean(), flush=True)
if __name__ == "__main__":
main()
|