Extract polygon targets from masks

scripts/train_sampoly_polygon.py

@@ -12,6 +12,7 @@ from dataclasses import asdict, dataclass
 from pathlib import Path
 
 import torch
+import numpy as np
 from PIL import Image, ImageDraw
 from torch import Tensor, nn
 import torch.nn.functional as F
@@ -140,6 +141,42 @@ def draw_targets(polygons: list[Tensor], size: int) -> tuple[Tensor, Tensor, Ten
     return mask, boundary, vertex
 
 
+def polygons_from_binary_mask(mask: Image.Image, vertices_per_polygon: int) -> list[Tensor]:
+    mask_np = np.asarray(mask.convert("L"))
+    binary = (mask_np > 0).astype(np.uint8)
+    if binary.max() == 0:
+        return []
+    try:
+        import cv2
+
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        polygons = []
+        h, w = binary.shape
+        min_area = max(4.0, 0.0005 * h * w)
+        for contour in contours:
+            if cv2.contourArea(contour) < min_area:
+                continue
+            pts = [(float(p[0][0]) / max(w - 1, 1), float(p[0][1]) / max(h - 1, 1)) for p in contour]
+            sampled = resample_polygon(pts, vertices_per_polygon)
+            polygons.append(torch.tensor(sampled, dtype=torch.float32).clamp(0, 1))
+        return polygons
+    except Exception:
+        ys, xs = np.where(binary > 0)
+        if len(xs) == 0:
+            return []
+        h, w = binary.shape
+        x1, x2 = xs.min() / max(w - 1, 1), xs.max() / max(w - 1, 1)
+        y1, y2 = ys.min() / max(h - 1, 1), ys.max() / max(h - 1, 1)
+        sampled = resample_polygon([(x1, y1), (x2, y1), (x2, y2), (x1, y2)], vertices_per_polygon)
+        return [torch.tensor(sampled, dtype=torch.float32).clamp(0, 1)]
+
+
+def boundary_from_mask(mask_tensor: Tensor) -> Tensor:
+    pooled_max = F.max_pool2d(mask_tensor.unsqueeze(0), kernel_size=3, stride=1, padding=1)
+    pooled_min = -F.max_pool2d(-mask_tensor.unsqueeze(0), kernel_size=3, stride=1, padding=1)
+    return (pooled_max - pooled_min).squeeze(0).clamp(0, 1)
+
+
 class PolygonDataset(Dataset):
     def __init__(self, root: str | Path, split: str, image_size: int, vertices_per_polygon: int) -> None:
         self.root = Path(root)
@@ -187,8 +224,12 @@ class PolygonDataset(Dataset):
         if not polygons and mask_path.exists():
             mask = Image.open(mask_path).convert("L").resize((self.image_size, self.image_size), Image.NEAREST)
             mask_tensor = (TF.to_tensor(mask) > 0.5).float()
-            boundary = torch.zeros_like(mask_tensor)
-            vertex = torch.zeros_like(mask_tensor)
+            polygons = polygons_from_binary_mask(mask, self.vertices_per_polygon)
+            if polygons:
+                _, boundary, vertex = draw_targets(polygons, self.image_size)
+            else:
+                boundary = boundary_from_mask(mask_tensor)
+                vertex = torch.zeros_like(mask_tensor)
         else:
             mask_tensor, boundary, vertex = draw_targets(polygons, self.image_size)
         return {