File size: 3,550 Bytes
36b249e | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | 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 |