add model and example

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+dino_hatching.pt filter=lfs diff=lfs merge=lfs -text

dino_hatching.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55359bcfd93e4952185eaefd51afbd03ff62a8666b21b88ccf8473258ac95b1f
+size 351351779

example/inference.py

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