update model

README.md

@@ -19,7 +19,7 @@ pipeline_tag: image-feature-extraction
 
 # Siamese DINOv2 Wall Hatching Matcher
 
-A TorchScript model for matching wall hatchings from architectural blueprints with legend patterns.
+A TorchScript model for matching wall hatchings from architectural blueprints with legend patterns. Example of use in the repository.
 
 ## Model
 
@@ -32,104 +32,124 @@ A TorchScript model for matching wall hatchings from architectural blueprints wi
 ## Inference
 
 ```python
-import cv2
-import numpy as np
+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,
+    )
 
-from PIL import Image
-from torchvision import transforms
-
-input_data = {
-    "legends": ["legend_correct.png", "legend_wrong.png"],
-    "plan_image": "wall.png", 
-    "plan_obb": [ 
-        20.672607421875,    # x1
-        20.71624755859375,  # y1
-        42.37445068359375,  # x2
-        20.71624755859375,  # y2
-        42.37445068359375,  # ...
-        111.15782165527344, 
-        20.672607421875, 
-        111.15782165527344
-        ] # mask bbox for wall image
-    }
-
-
-IMAGE_SIZE = 518 # don't change
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-image_tf = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-    transforms.ToTensor(),
-    transforms.Normalize(
+    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 = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-    transforms.ToTensor(),
-])
-
-
-def create_full_mask(size):
-    return Image.new("L", size, 255)
-
-
-def create_obb_mask(size, obb):
-    w, h = size
-
-    points = np.array(obb, dtype=np.int32).reshape(4, 2)
-
-    mask = np.zeros((h, w), dtype=np.uint8)
-    cv2.fillPoly(mask, [points], 255)
-
-    return Image.fromarray(mask)
-
-
-def prepare(image_path, obb=None):
-    image = Image.open(image_path).convert("RGB")
+    )
+    mask = TF.to_tensor(mask)
 
-    if obb is None:
-        mask = create_full_mask(image.size)
-    else:
-        mask = create_obb_mask(image.size, obb)
+    return image.unsqueeze(0).to(DEVICE), mask.unsqueeze(0).to(DEVICE)
 
-    image = image_tf(image).unsqueeze(0).to(device)
-    mask = mask_tf(mask).unsqueeze(0).to(device)
 
-    return image, mask
+with (ROOT / "data.jsonl").open("r", encoding="utf-8") as file:
+    items = [json.loads(line) for line in file if line.strip()]
 
-for legend in input_data["legends"]:
-    legend_image, legend_mask = prepare(legend)
+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}
 
-    plan_image, plan_mask = prepare(
-        input_data["plan_image"],
-        input_data["plan_obb"],
-    )
+print("\nScore matrix")
+print(" " * 6 + "".join(f"W{i}".rjust(8) for i in range(1, len(wall_types) + 1)))
 
-    model = torch.jit.load(
-        "dino_hatching.pt",
-        map_location=device,
-    )
-    model.eval()
-
-    with torch.no_grad():
-        logit = model(
-            legend_image,
-            legend_mask,
-            plan_image,
-            plan_mask,
+with torch.inference_mode():
+    for index, item in enumerate(items, start=1):
+        plan_image, plan_mask = prepare_image(
+            item["plan_image"],
+            item["plan_obb"],
         )
-
-    score = torch.sigmoid(logit).item()
-
-    print(f"{legend}: {score}")
-
-# legend_correct.png: 0.9882104396820068
-# legend_wrong.png: 0.0016661642584949732
+        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.

dino_hatching.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:39684e4034dc7dd8f8f1266678891e2892ed84e76b807d484498bc32b73f75d1
-size 351342499
+oid sha256:dc6b05869819984378bbcf3f4c8f693a32bf456f5b237a06900739dcdb5166c0
+size 351563175

example/data.jsonl

@@ -0,0 +1,5 @@
+{"plan_image":"valid_pair_000001_plan.png","plan_obb":[96.0,41.0,20.0,41.0,20.0,20.0,96.0,20.0],"wall_type":"valid_pair_000001_legend.png"}
+{"plan_image":"valid_pair_000004_plan.png","plan_obb":[48.0,63.0,20.0,63.0,20.0,20.0,48.0,20.0],"wall_type":"valid_pair_000004_legend.png"}
+{"plan_image":"valid_pair_000010_plan.png","plan_obb":[74.0,33.0,20.0,33.0,20.0,20.0,74.0,20.0],"wall_type":"valid_pair_000010_legend.png"}
+{"plan_image":"valid_pair_000013_plan.png","plan_obb":[160.0,30.0,20.0,30.0,20.0,20.0,160.0,20.0],"wall_type":"valid_pair_000013_legend.png"}
+{"plan_image":"valid_pair_000016_plan.png","plan_obb":[179.0,40.0,20.0,40.0,20.0,20.0,179.0,20.0],"wall_type":"valid_pair_000016_legend.png"}

example/example.py

@@ -0,0 +1,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

example/inference.py

@@ -1,98 +0,0 @@
-import cv2
-import numpy as np
-import torch
-
-from PIL import Image
-from torchvision import transforms
-
-input_data = {
-    "legends": ["legend_correct.png", "legend_wrong.png"],
-    "plan_image": "wall.png", 
-    "plan_obb": [ 
-        20.672607421875,    # x1
-        20.71624755859375,  # y1
-        42.37445068359375,  # x2
-        20.71624755859375,  # y2
-        42.37445068359375,  # ...
-        111.15782165527344, 
-        20.672607421875, 
-        111.15782165527344
-        ] # mask bbox for wall image
-    }
-
-
-IMAGE_SIZE = 518
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-image_tf = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225],
-    ),
-])
-
-mask_tf = transforms.Compose([
-    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
-    transforms.ToTensor(),
-])
-
-
-def create_full_mask(size):
-    return Image.new("L", size, 255)
-
-
-def create_obb_mask(size, obb):
-    w, h = size
-
-    points = np.array(obb, dtype=np.int32).reshape(4, 2)
-
-    mask = np.zeros((h, w), dtype=np.uint8)
-    cv2.fillPoly(mask, [points], 255)
-
-    return Image.fromarray(mask)
-
-
-def prepare(image_path, obb=None):
-    image = Image.open(image_path).convert("RGB")
-
-    if obb is None:
-        mask = create_full_mask(image.size)
-    else:
-        mask = create_obb_mask(image.size, obb)
-
-    image = image_tf(image).unsqueeze(0).to(device)
-    mask = mask_tf(mask).unsqueeze(0).to(device)
-
-    return image, mask
-
-for legend in input_data["legends"]:
-    legend_image, legend_mask = prepare(legend)
-
-    plan_image, plan_mask = prepare(
-        input_data["plan_image"],
-        input_data["plan_obb"],
-    )
-
-    model = torch.jit.load(
-        "dino_hatching.pt",
-        map_location=device,
-    )
-    model.eval()
-
-    with torch.no_grad():
-        logit = model(
-            legend_image,
-            legend_mask,
-            plan_image,
-            plan_mask,
-        )
-
-    score = torch.sigmoid(logit).item()
-
-    print(f"{legend}: {score}")
-
-# legend_correct.png: 0.9882104396820068
-# legend_wrong.png: 0.0016661642584949732