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