""" 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)