import json
from pathlib import Path

import torch
from PIL import Image, ImageDraw
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as TF


ROOT = Path(__file__).resolve().parent
IMAGE_SIZE = 518
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


def prepare_image(filename: str, obb: list[float] | None = None):
    image = Image.open(ROOT / "images" / filename).convert("RGB")
    mask = Image.new("L", image.size, 0 if obb else 255)

    if obb:
        points = [
            (int(x), int(y))
            for x, y in zip(obb[::2], obb[1::2])
        ]
        ImageDraw.Draw(mask).polygon(points, fill=255)

    width, height = image.size
    scale = min(IMAGE_SIZE / width, IMAGE_SIZE / height)
    new_width = min(IMAGE_SIZE, round(width * scale))
    new_height = min(IMAGE_SIZE, round(height * scale))
    size = [new_height, new_width]

    image = TF.resize(
        image,
        size,
        interpolation=InterpolationMode.BILINEAR,
        antialias=True,
    )
    mask = TF.resize(
        mask,
        size,
        interpolation=InterpolationMode.NEAREST,
    )

    pad_x = IMAGE_SIZE - new_width
    pad_y = IMAGE_SIZE - new_height
    padding = [
        pad_x // 2,
        pad_y // 2,
        pad_x - pad_x // 2,
        pad_y - pad_y // 2,
    ]
    image = TF.pad(image, padding, fill=255)
    mask = TF.pad(mask, padding, fill=0)

    image = TF.to_tensor(image)
    image = TF.normalize(
        image,
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    )
    mask = TF.to_tensor(mask)

    return image.unsqueeze(0).to(DEVICE), mask.unsqueeze(0).to(DEVICE)


with (ROOT / "data.jsonl").open("r", encoding="utf-8") as file:
    items = [json.loads(line) for line in file if line.strip()]

model = torch.jit.load(ROOT / "dino_hatching.pt", map_location=DEVICE).eval()
wall_types = list(dict.fromkeys(item["wall_type"] for item in items))
walls = {name: prepare_image(name) for name in wall_types}

print("\nScore matrix")
print(" " * 6 + "".join(f"W{i}".rjust(8) for i in range(1, len(wall_types) + 1)))

with torch.inference_mode():
    for index, item in enumerate(items, start=1):
        plan_image, plan_mask = prepare_image(
            item["plan_image"],
            item["plan_obb"],
        )
        scores = []

        for wall_type in wall_types:
            wall_image, wall_mask = walls[wall_type]
            logit = model(wall_image, wall_mask, plan_image, plan_mask)
            scores.append(f"{torch.sigmoid(logit).item():.4f}")

        print(f"P{index:<5}" + "".join(f"{score:>8}" for score in scores))

print("\nRows:")
for index, item in enumerate(items, start=1):
    print(f"  P{index}: {item['plan_image']}")

print("Columns:")
for index, wall_type in enumerate(wall_types, start=1):
    print(f"  W{index}: {wall_type}")

# Score matrix
#             W1      W2      W3      W4      W5
# P1      0.9965  0.0005  0.0006  0.0044  0.0317
# P2      0.0026  0.9932  0.9916  0.0010  0.0028
# P3      0.0001  0.9984  0.9988  0.0001  0.0036
# P4      0.0005  0.0001  0.0002  0.9979  0.0002
# P5      0.0103  0.0004  0.0006  0.0001  0.9971

# Rows:
#   P1: valid_pair_000001_plan.png
#   P2: valid_pair_000004_plan.png
#   P3: valid_pair_000010_plan.png
#   P4: valid_pair_000013_plan.png
#   P5: valid_pair_000016_plan.png
# Columns:
#   W1: valid_pair_000001_legend.png
#   W2: valid_pair_000004_legend.png
#   W3: valid_pair_000010_legend.png
#   W4: valid_pair_000013_legend.png
#   W5: valid_pair_000016_legend.png

score is the probability that both hatchings belong to the same wall type.

1.0 → same wall type
0.0 → different wall type

Metrics

valid loss: 0.0442
valid_acc: 0.9890

GreenMap/siamese-dinov2-wall-hatching-matcher

Siamese DINOv2 Wall Hatching Matcher

Model

Inference

Metrics

GreenMap/
siamese-dinov2-wall-hatching-matcher