"""Split JL1 binary change into new_building / tree_removal with the Satlas (siam_v2) semantic heads."""
import numpy as np, torch
from model_v2 import SiamCD
pre = np.load("/workspace/prep/jl_pre.npy", mmap_mode="r"); post = np.load("/workspace/prep/jl_post.npy", mmap_mode="r"); chg = np.load("/workspace/prep/jl_chg.npy")
m = SiamCD(); m.load_state_dict(torch.load("/workspace/siam_v2.pt", map_location="cpu", weights_only=True)["state_dict"]); m = m.cuda().eval()
lb = np.zeros_like(chg); lt = np.zeros_like(chg)
with torch.no_grad():
    for s in range(0, len(chg), 16):
        a = torch.from_numpy(np.array(pre[s:s + 16])).permute(0, 3, 1, 2).cuda().float() / 255
        b = torch.from_numpy(np.array(post[s:s + 16])).permute(0, 3, 1, 2).cuda().float() / 255
        with torch.autocast("cuda", dtype=torch.float16):
            _, _, sp, sq = m(a, b)
        sp = sp.float().argmax(1).cpu().numpy(); sq = sq.float().argmax(1).cpu().numpy()
        c = chg[s:s + 16].astype(bool)
        lb[s:s + 16] = (c & (sq == 1)).astype(np.uint8)
        lt[s:s + 16] = (c & (sp == 2) & (sq != 2)).astype(np.uint8)
np.save("/workspace/prep/jl_lb.npy", lb); np.save("/workspace/prep/jl_lt.npy", lt)
n = lambda x: int((x.reshape(len(x), -1).sum(1) >= 20).sum())
print("PSEUDO_DONE pairs: change", n(chg), "new_building", n(lb), "tree_removal", n(lt), "| px frac b/t", round(float(lb.mean()), 4), round(float(lt.mean()), 4), flush=True)