+def grow_islands(nb, height, width, nb_seeds, nb_iterations):
+ w = torch.empty(5, 1, 3, 3)
+
+ w[0, 0] = torch.tensor(
+ [
+ [1.0, 1.0, 1.0],
+ [1.0, 0.0, 1.0],
+ [1.0, 1.0, 1.0],
+ ]
+ )
+
+ w[1, 0] = torch.tensor(
+ [
+ [-1.0, 1.0, 0.0],
+ [1.0, 0.0, 0.0],
+ [0.0, 0.0, 0.0],
+ ]
+ )
+
+ w[2, 0] = torch.tensor(
+ [
+ [0.0, 1.0, -1.0],
+ [0.0, 0.0, 1.0],
+ [0.0, 0.0, 0.0],
+ ]
+ )
+
+ w[3, 0] = torch.tensor(
+ [
+ [0.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0],
+ [0.0, 1.0, -1.0],
+ ]
+ )
+
+ w[4, 0] = torch.tensor(
+ [
+ [0.0, 0.0, 0.0],
+ [1.0, 0.0, 0.0],
+ [-1.0, 1.0, 0.0],
+ ]
+ )
+
+ Z = torch.zeros(nb, height, width)
+ U = Z.flatten(1)
+
+ for _ in range(nb_seeds):
+ M = F.conv2d(Z[:, None, :, :], w, padding=1)
+ M = torch.cat([M[:, :1], M[:, 1:].min(dim=1, keepdim=True).values], dim=1)
+ M = ((M[:, 0] == 0) & (Z == 0)).long()
+ Q = (M.flatten(1).max(dim=1).values > 0).long()[:, None]
+ M = M * torch.rand(M.size())
+ M = M.flatten(1)
+ M = F.one_hot(M.argmax(dim=1), num_classes=M.size(1))
+ U += M * Q
+
+ for _ in range(nb_iterations):
+ M = F.conv2d(Z[:, None, :, :], w, padding=1)
+ M = torch.cat([M[:, :1], M[:, 1:].min(dim=1, keepdim=True).values], dim=1)
+ M = ((M[:, 1] >= 0) & (Z == 0)).long()
+ Q = (M.flatten(1).max(dim=1).values > 0).long()[:, None]
+ M = M * torch.rand(M.size())
+ M = M.flatten(1)
+ M = F.one_hot(M.argmax(dim=1), num_classes=M.size(1))
+ U = Z.flatten(1)
+ U += M * Q
+
+ M = Z.clone()
+ Z = Z * (torch.arange(Z.size(1) * Z.size(2)) + 1).reshape(1, Z.size(1), Z.size(2))
+
+ while True:
+ W = Z.clone()
+ Z = F.max_pool2d(Z, 3, 1, 1) * M
+ if Z.equal(W):
+ break
+
+ Z = Z.long()
+ U = Z.flatten(1)
+ V = F.one_hot(U).max(dim=1).values
+ W = V.cumsum(dim=1) - V
+ N = torch.arange(Z.size(0))[:, None, None].expand_as(Z)
+ Z = W[N, Z]
+
+ return Z
+
+