app.py

"""
Pixels, Points and Polygons (P3) — Building Vectorizer (image modality).

Self-contained ZeroGPU Gradio demo of the Pix2Poly image-only model from
"The P3 Dataset: Pixels, Points and Polygons for Multimodal Building
Vectorization" (Sulzer et al., 2025).  arXiv:2505.15379

The model takes a 224x224 aerial RGB tile (25 cm GSD) and autoregressively
predicts building outline polygons, which are drawn as a vector overlay.

The Pix2Poly image path (DINO ViT-S/8 encoder + transformer polygon decoder +
optimal-transport permutation head) is ported 1:1 from the original repo
(github.com/raphaelsulzer/pixelspointspolygons) with weights from
huggingface.co/rsi/PixelsPointsPolygons.  The LiDAR / fusion / HiSup / FFL
paths (which need Open3D-ML and custom CUDA ops) are intentionally omitted.
"""

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import timm
import gradio as gr

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as Patches

from scipy.optimize import linear_sum_assignment
from shapely.geometry import Polygon
from shapely.validation import make_valid
from huggingface_hub import hf_hub_download

MODEL_REPO = "rsi/PixelsPointsPolygons"
CKPT_PATH = "pix2poly/224/v4_image_vit_bs4x16/checkpoints/best_val_iou.pth"
BACKBONE_PATH = "backbones/dino_deitsmall8_pretrain.pth"

DEVICE = "cuda"
IN_SIZE = 224
OUT_FEATURE_DIM = 256        # model.decoder.in_feature_dim
NUM_PATCHES = 784            # (224 // 8) ** 2
NUM_BINS = 224               # tokenizer.num_bins == in_size
MAX_NUM_VERTICES = 192       # tokenizer.max_num_vertices
TOKEN_MODE = 2
SINKHORN_ITERATIONS = 100

IMAGE_MEAN = [0.0, 0.0, 0.0]
IMAGE_STD = [1.0, 1.0, 1.0]


# --------------------------------------------------------------------------- #
# Tokenizer (ported from pixelspointspolygons/models/pix2poly/tokenizer.py)
# --------------------------------------------------------------------------- #
class Tokenizer:
    def __init__(self):
        self.token_mode = TOKEN_MODE
        self.num_bins = NUM_BINS
        self.width = IN_SIZE
        self.height = IN_SIZE
        self.max_len = MAX_NUM_VERTICES * self.token_mode + 2

        self.BOS_code = self.num_bins
        self.EOS_code = self.BOS_code + 1
        self.PAD_code = self.EOS_code + 1
        self.pad_idx = self.PAD_code

        self.vocab_size = self.num_bins + 3
        self.generation_steps = MAX_NUM_VERTICES * self.token_mode + 1

    def dequantize(self, x: np.ndarray):
        return x.astype("float32") / (self.num_bins - 1)

    def decode(self, tokens: torch.Tensor):
        mask = tokens != self.PAD_code
        tokens = tokens[mask]
        tokens = tokens[1:-1]
        assert len(tokens) % self.token_mode == 0, "Invalid tokens!"
        coords = np.array(tokens).reshape(-1, self.token_mode)[:, :2]
        coords = self.dequantize(coords)
        if len(coords) > 0:
            coords[:, 0] = coords[:, 0] * self.width
            coords[:, 1] = coords[:, 1] * self.height
        return coords


# --------------------------------------------------------------------------- #
# Model (ported from models/pix2poly/model_pix2poly.py + vision_transformer/vit.py)
# --------------------------------------------------------------------------- #
def generate_square_subsequent_mask(sz, device):
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    mask = mask.float().masked_fill(mask == 0, float("-inf")).masked_fill(mask == 1, float(0.0))
    return mask


def create_mask(tgt, pad_idx):
    tgt_seq_len = tgt.size(1)
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len, device=tgt.device)
    tgt_padding_mask = (tgt == pad_idx).to(dtype=tgt_mask.dtype)
    return tgt_mask, tgt_padding_mask


class ScoreNet(nn.Module):
    def __init__(self, n_vertices, in_channels=512, token_mode=2):
        super().__init__()
        self.n_vertices = n_vertices
        self.in_channels = in_channels
        self.relu = nn.ReLU(inplace=True)
        self.conv1 = nn.Conv2d(in_channels, 256, kernel_size=1, stride=1, padding=0, bias=True)
        self.bn1 = nn.BatchNorm2d(256)
        self.conv2 = nn.Conv2d(256, 128, kernel_size=1, stride=1, padding=0, bias=True)
        self.bn2 = nn.BatchNorm2d(128)
        self.conv3 = nn.Conv2d(128, 64, kernel_size=1, stride=1, padding=0, bias=True)
        self.bn3 = nn.BatchNorm2d(64)
        self.conv4 = nn.Conv2d(64, 1, kernel_size=1, stride=1, padding=0, bias=True)
        self.token_mode = token_mode

    def forward(self, feats):
        feats = feats[:, 1:]
        feats = feats.unsqueeze(2)
        feats = feats.view(feats.size(0), feats.size(1) // self.token_mode, self.token_mode, feats.size(3))
        feats = torch.mean(feats, dim=2)

        x = torch.transpose(feats, 1, 2)
        x = x.unsqueeze(-1)
        x = x.repeat(1, 1, 1, self.n_vertices)
        t = torch.transpose(x, 2, 3)
        x = torch.cat((x, t), dim=1)

        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.relu(self.bn3(self.conv3(x)))
        x = self.conv4(x)
        return x[:, 0]


class Decoder(nn.Module):
    def __init__(self, vocab_size, encoder_len, dim, num_heads, num_layers, max_len, pad_idx):
        super().__init__()
        self.dim = dim
        self.max_len = max_len
        self.pad_idx = pad_idx

        self.embedding = nn.Embedding(vocab_size, dim)
        self.decoder_pos_embed = nn.Parameter(torch.randn(1, self.max_len - 1, dim) * 0.02)
        self.decoder_pos_drop = nn.Dropout(p=0.05)

        decoder_layer = nn.TransformerDecoderLayer(d_model=dim, nhead=num_heads)
        self.decoder = nn.TransformerDecoder(decoder_layer=decoder_layer, num_layers=num_layers)
        self.output = nn.Linear(dim, vocab_size)

        self.encoder_pos_embed = nn.Parameter(torch.randn(1, encoder_len, dim) * 0.02)
        self.encoder_pos_drop = nn.Dropout(p=0.05)

    def predict(self, encoder_out, tgt):
        length = tgt.size(1)
        padding = (
            torch.ones((tgt.size(0), self.max_len - length - 1), device=tgt.device)
            .fill_(self.pad_idx)
            .long()
        )
        tgt = torch.cat([tgt, padding], dim=1)
        tgt_mask, tgt_padding_mask = create_mask(tgt, self.pad_idx)
        tgt_embedding = self.embedding(tgt)
        tgt_embedding = self.decoder_pos_drop(tgt_embedding + self.decoder_pos_embed)

        encoder_out = self.encoder_pos_drop(encoder_out + self.encoder_pos_embed)

        encoder_out = encoder_out.transpose(0, 1)
        tgt_embedding = tgt_embedding.transpose(0, 1)

        preds = self.decoder(
            memory=encoder_out,
            tgt=tgt_embedding,
            tgt_mask=tgt_mask,
            tgt_key_padding_mask=tgt_padding_mask,
        )
        preds = preds.transpose(0, 1)
        return self.output(preds)[:, length - 1, :], preds


class ViT(nn.Module):
    def __init__(self, backbone_ckpt=None):
        super().__init__()
        self.vit = timm.create_model(
            model_name="vit_small_patch8_224.dino",
            num_classes=0,
            global_pool="",
        )
        if backbone_ckpt is not None:
            self.vit.load_state_dict(torch.load(backbone_ckpt, map_location="cpu"), strict=False)
        self.bottleneck = nn.AdaptiveAvgPool1d(OUT_FEATURE_DIM)

    def forward(self, x):
        x = self.vit(x)
        x = self.bottleneck(x[:, 1:, :])
        return x


class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.token_mode = TOKEN_MODE
        self.encoder = encoder
        self.decoder = decoder
        self.max_num_vertices = MAX_NUM_VERTICES
        self.sinkhorn_iterations = SINKHORN_ITERATIONS
        self.scorenet1 = ScoreNet(self.max_num_vertices, token_mode=self.token_mode)
        self.scorenet2 = ScoreNet(self.max_num_vertices, token_mode=self.token_mode)
        self.bin_score = torch.nn.Parameter(torch.tensor(1.0))
        self.bottleneck = nn.AdaptiveAvgPool1d(OUT_FEATURE_DIM)

    def predict(self, encoded_image, tgt):
        return self.decoder.predict(encoded_image, tgt)


def build_model(tokenizer, backbone_ckpt=None):
    encoder = ViT(backbone_ckpt=backbone_ckpt)
    decoder = Decoder(
        vocab_size=tokenizer.vocab_size,
        encoder_len=NUM_PATCHES,
        dim=OUT_FEATURE_DIM,
        num_heads=8,
        num_layers=6,
        max_len=tokenizer.max_len,
        pad_idx=tokenizer.pad_idx,
    )
    return EncoderDecoder(encoder, decoder)


# --------------------------------------------------------------------------- #
# smart state-dict loading (ported from misc/shared_utils.smart_load_state_dict)
# --------------------------------------------------------------------------- #
def smart_load_state_dict(model, checkpoint_state_dict, strict=True):
    from collections import OrderedDict
    from copy import deepcopy

    model_state_dict = model.state_dict()
    new_state_dict = OrderedDict()

    unmatched_model_keys = set(model_state_dict.keys())
    unmatched_checkpoint_keys = set(checkpoint_state_dict.keys())

    temp = deepcopy(checkpoint_state_dict)
    for k, v in checkpoint_state_dict.items():
        temp[k.replace("encoder.model.", "encoder.vit.")] = v
    checkpoint_state_dict = temp

    for ckpt_key in checkpoint_state_dict.keys():
        if ckpt_key in model_state_dict:
            new_state_dict[ckpt_key] = checkpoint_state_dict[ckpt_key]
            unmatched_model_keys.discard(ckpt_key)
            unmatched_checkpoint_keys.discard(ckpt_key)
        else:
            for model_key in model_state_dict.keys():
                if ckpt_key.endswith(model_key):
                    new_state_dict[model_key] = checkpoint_state_dict[ckpt_key]
                    unmatched_model_keys.discard(model_key)
                    unmatched_checkpoint_keys.discard(ckpt_key)
                    break
                if model_key.endswith(ckpt_key):
                    new_state_dict[model_key] = checkpoint_state_dict[ckpt_key]
                    unmatched_model_keys.discard(model_key)
                    unmatched_checkpoint_keys.discard(ckpt_key)
                    break

    print(f"[load] matched {len(model_state_dict) - len(unmatched_model_keys)}/{len(model_state_dict)} model keys")
    if unmatched_model_keys:
        print(f"[load] {len(unmatched_model_keys)} unmatched model keys, e.g. {list(unmatched_model_keys)[:5]}")
    model.load_state_dict(new_state_dict, strict=strict)
    return model


# --------------------------------------------------------------------------- #
# Load model at module scope, eagerly to CUDA (ZeroGPU pattern)
# --------------------------------------------------------------------------- #
print("Downloading weights...")
backbone_file = hf_hub_download(MODEL_REPO, BACKBONE_PATH)
ckpt_file = hf_hub_download(MODEL_REPO, CKPT_PATH)

tokenizer = Tokenizer()
model = build_model(tokenizer, backbone_ckpt=backbone_file)

print("Loading checkpoint...")
checkpoint = torch.load(ckpt_file, map_location="cpu", weights_only=False)
for k in list(checkpoint.keys()):
    if "_state_dict" in k:
        checkpoint[k.replace("_state_dict", "")] = checkpoint.pop(k)
model = smart_load_state_dict(model, checkpoint["model"], strict=True)
model.eval()
model.to(DEVICE)
print("Model ready.")


# --------------------------------------------------------------------------- #
# Inference helpers (ported from predict/predictor_pix2poly.py)
# --------------------------------------------------------------------------- #
def scores_to_permutations(scores):
    B, N, _ = scores.shape
    scores = scores.detach().cpu().numpy()
    perm = np.zeros_like(scores)
    for b in range(B):
        r, c = linear_sum_assignment(-scores[b])
        perm[b, r, c] = 1
    return torch.tensor(perm)


def permutations_to_polygons(perm, graph):
    B, N, _ = perm.shape
    device = perm.device

    def bubble_merge(poly):
        s = 0
        P = len(poly)
        while s < P:
            head = poly[s][-1]
            t = s + 1
            while t < P:
                tail = poly[t][0]
                if head == tail:
                    poly[s] = poly[s] + poly[t][1:]
                    del poly[t]
                    poly = bubble_merge(poly)
                    P = len(poly)
                t += 1
            s += 1
        return poly

    diag = torch.logical_not(perm[:, range(N), range(N)])
    batch = []
    for b in range(B):
        b_perm = perm[b]
        b_graph = graph[b]
        b_diag = diag[b]
        idx = torch.arange(N, device=perm.device)[b_diag]
        if idx.shape[0] > 0:
            b_perm = b_perm[idx, :]
            b_graph = b_graph[idx, :]
            b_perm = b_perm[:, idx]
            first = torch.arange(idx.shape[0]).unsqueeze(1).to(device=device)
            second = torch.argmax(b_perm, dim=1).unsqueeze(1)
            polygons_idx = torch.cat((first, second), dim=1).tolist()
            polygons_idx = bubble_merge(polygons_idx)
            batch_poly = []
            for p_idx in polygons_idx:
                batch_poly.append(b_graph[p_idx, :])
            batch.append(batch_poly)
        else:
            batch.append([])
    return batch


def postprocess(batch_preds):
    EOS_idxs = (batch_preds == tokenizer.EOS_code).float().argmax(dim=-1)
    invalid_idxs = ((EOS_idxs - 1) % tokenizer.token_mode != 0).nonzero().view(-1)
    EOS_idxs[invalid_idxs] = 0
    all_coords = []
    for i, EOS_idx in enumerate(EOS_idxs.tolist()):
        if EOS_idx == 0:
            all_coords.append(None)
            continue
        coords = tokenizer.decode(batch_preds[i, : EOS_idx + 1])
        all_coords.append(coords)
    return all_coords


def coord_and_perm_to_polygons(coord_preds, perm_preds):
    vertex_coords = postprocess(coord_preds)
    coords = []
    for i in range(len(vertex_coords)):
        if vertex_coords[i] is not None:
            coord = torch.from_numpy(vertex_coords[i])
        else:
            coord = torch.tensor([])
        padd = torch.ones((MAX_NUM_VERTICES - len(coord), 2)).fill_(tokenizer.pad_idx)
        coord = torch.cat([coord, padd], dim=0)
        coords.append(coord)

    batch_polygons = permutations_to_polygons(perm_preds, coords)
    batch_polygons_processed = []
    for pp in batch_polygons:
        polys = []
        for p in pp:
            p = torch.fliplr(p)
            p = p[p[:, 0] != tokenizer.pad_idx]
            if len(p) > 0:
                polys.append(p)
        batch_polygons_processed.append(polys)
    return batch_polygons_processed


def test_generate(model, x_images):
    batch_size = x_images.size(0)
    batch_preds = torch.ones((batch_size, 1), device=DEVICE).fill_(tokenizer.BOS_code).long()
    sample = lambda preds: torch.softmax(preds, dim=-1).argmax(dim=-1).view(-1, 1)

    with torch.no_grad():
        features = model.encoder(x_images)
        feats = None
        for _ in range(tokenizer.generation_steps):
            preds, feats = model.predict(features, batch_preds)
            preds = sample(preds)
            batch_preds = torch.cat([batch_preds, preds], dim=1)

        perm_preds = model.scorenet1(feats) + torch.transpose(model.scorenet2(feats), 1, 2)
        perm_preds = scores_to_permutations(perm_preds)

    return batch_preds.cpu(), perm_preds


def preprocess_image(image_np):
    """image_np: HxWx3 uint8 RGB. Returns normalized (1,3,224,224) float tensor and resized uint8 array."""
    from PIL import Image
    pil = Image.fromarray(image_np.astype(np.uint8)).convert("RGB")
    if pil.size != (IN_SIZE, IN_SIZE):
        pil = pil.resize((IN_SIZE, IN_SIZE), Image.BILINEAR)
    arr = np.asarray(pil).astype(np.uint8)

    t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(torch.float32) / 255.0
    mean = torch.tensor(IMAGE_MEAN).view(1, 3, 1, 1)
    std = torch.tensor(IMAGE_STD).view(1, 3, 1, 1)
    t = (t - mean) / std
    return t, arr


def render_overlay(image_np, polygons):
    """Draw the predicted building polygons over the input tile."""
    px = 1 / plt.rcParams["figure.dpi"]
    fig, ax = plt.subplots(1, 1, figsize=(800 * px, 800 * px))
    ax.axis("off")
    ax.imshow(image_np)

    shapely_polygons = []
    for poly in polygons:
        arr = poly.cpu().numpy() if isinstance(poly, torch.Tensor) else np.asarray(poly)
        if len(arr) < 3:
            continue
        try:
            sp = Polygon(arr)
            if not sp.is_valid:
                sp = make_valid(sp)
            shapely_polygons.append((sp, arr))
        except Exception:
            continue

    n_buildings = 0
    for sp, arr in shapely_polygons:
        ax.add_patch(
            Patches.Polygon(arr, fill=False, ec=[1, 0, 1], linewidth=2.5, zorder=3, alpha=0.9)
        )
        ax.plot(arr[:, 0], arr[:, 1], color=[1, 1, 0], marker=".", markersize=6,
                linestyle="none", zorder=4)
        n_buildings += 1

    ax.set_xlim(0, image_np.shape[1])
    ax.set_ylim(image_np.shape[0], 0)
    fig.tight_layout(pad=0)

    fig.canvas.draw()
    buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8)
    w, h = fig.canvas.get_width_height()
    out = buf.reshape(h, w, 4)[..., :3].copy()
    plt.close(fig)
    return out, n_buildings


@spaces.GPU(duration=60)
def vectorize(image_np):
    """Extract building outline polygons from an aerial RGB image tile.

    Args:
        image_np: An aerial RGB image (any size; resized to 224x224). Best
            results on 25 cm ground-sampling-distance nadir imagery.

    Returns:
        An overlay image with the predicted building polygons drawn on top of
        the input tile (magenta edges, yellow vertices).
    """
    if image_np is None:
        raise gr.Error("Please provide an input image.")

    x, arr = preprocess_image(image_np)
    x = x.to(DEVICE)

    coord_preds, perm_preds = test_generate(model, x)
    batch_polygons = coord_and_perm_to_polygons(coord_preds, perm_preds)
    polygons = batch_polygons[0]

    overlay, n_buildings = render_overlay(arr, polygons)
    return overlay


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

EXAMPLES = [
    ["example_CH.png"],
    ["image100_CH_val.png"],
    ["image0_NY_val.png"],
    ["image0_NZ_val.png"],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🏙️ Pixels, Points & Polygons — Building Vectorizer

            Predict **building outline polygons** from a single aerial RGB tile using the
            **Pix2Poly** image model from
            [*The P³ Dataset*](https://huggingface.co/papers/2505.15379) (Sulzer et al., 2025).
            A DINO ViT-S/8 encoder feeds a transformer polygon decoder that autoregressively
            emits vertices; an optimal-transport head connects them into closed polygons.

            Inputs are resized to **224×224** (the model was trained on 25 cm GSD nadir tiles).
            Weights: [`rsi/PixelsPointsPolygons`](https://huggingface.co/rsi/PixelsPointsPolygons)
            · Code: [github](https://github.com/raphaelsulzer/pixelspointspolygons)
            """
        )
        with gr.Row():
            inp = gr.Image(label="Aerial RGB tile", type="numpy", height=400)
            out = gr.Image(label="Predicted building polygons", height=400)
        run = gr.Button("Vectorize buildings", variant="primary")

        run.click(fn=vectorize, inputs=inp, outputs=out, api_name="vectorize")

        gr.Examples(
            examples=EXAMPLES,
            inputs=inp,
            outputs=out,
            fn=vectorize,
            cache_examples=True,
            cache_mode="lazy",
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)