import problem
+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()
+ 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
+
+ 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()
+ 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
+
+ M = Z.clone()
+ Z = Z * (torch.arange(Z.size(1) * Z.size(2)) + 1).reshape(1, Z.size(1), Z.size(2))
+
+ for _ in range(100):
+ Z = F.max_pool2d(Z, 3, 1, 1) * M
+
+ 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
+
+
class Grids(problem.Problem):
named_colors = [
("white", [255, 255, 255]),
di, dj = torch.randint(2, (2,)) * 2 - 1
nb_rec = 3
c = torch.randperm(len(self.colors) - 1)[:nb_rec] + 1
- direction = torch.randint(2, (1,))
+ direction = torch.randint(2, (1,)).item()
for X, f_X in [(A, f_A), (B, f_B)]:
while True:
r = self.rec_coo(nb_rec, prevent_overlap=True)
di, dj = torch.randint(2, (2,)) * 2 - 1
nb_rec = 3
c = torch.randperm(len(self.colors) - 1)[: 2 * nb_rec] + 1
- direction = torch.randint(4, (1,))
+ direction = torch.randint(4, (1,)).item()
for X, f_X in [(A, f_A), (B, f_B)]:
r = self.rec_coo(nb_rec, prevent_overlap=True)
for n in range(nb_rec):
return no, nq, nq_diag
def task_count(self, A, f_A, B, f_B):
- N = (torch.randint(4, (1,)) + 2).item()
+ N = torch.randint(4, (1,)).item() + 2
c = torch.randperm(len(self.colors) - 1)[:N] + 1
for X, f_X in [(A, f_A), (B, f_B)]:
for X, f_X in [(A, f_A), (B, f_B)]:
while True:
di, dj = torch.randint(7, (2,)) - 3
- i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
+ i, j = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
+ )
if (
abs(di) + abs(dj) > 0
and i + 2 * di >= 0
X[...] = 0
for _ in range((self.height * self.width) // 10):
- i, j = torch.randint(self.height, (1,)), torch.randint(
- self.width, (1,)
+ i, j = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
)
X[i, j] = c[0]
f_X[i, j] = c[0]
if abs(di) + abs(dj) == 1:
break
- i, j = torch.randint(self.height, (1,)), torch.randint(self.width, (1,))
+ i, j = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
+ )
X[i, j] = c[1]
f_X[i, j] = c[1]
def task_scale(self, A, f_A, B, f_B):
c = torch.randperm(len(self.colors) - 1)[:2] + 1
- i, j = torch.randint(self.height // 2, (1,)), torch.randint(
- self.width // 2, (1,)
+ i, j = (
+ torch.randint(self.height // 2, (1,)).item(),
+ torch.randint(self.width // 2, (1,)).item(),
)
for X, f_X in [(A, f_A), (B, f_B)]:
for _ in range(3):
while True:
- i1, j1 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
- self.width // 2 + 1, (1,)
+ i1, j1 = (
+ torch.randint(self.height // 2 + 1, (1,)).item(),
+ torch.randint(self.width // 2 + 1, (1,)).item(),
)
- i2, j2 = torch.randint(self.height // 2 + 1, (1,)), torch.randint(
- self.width // 2 + 1, (1,)
+ i2, j2 = (
+ torch.randint(self.height // 2 + 1, (1,)).item(),
+ torch.randint(self.width // 2 + 1, (1,)).item(),
)
if i1 < i2 and j1 < j2 and min(i2 - i1, j2 - j1) <= 3:
break
ai, aj = i.float().mean(), j.float().mean()
- q = torch.randint(3, (1,)) + 1
+ q = torch.randint(3, (1,)).item() + 1
X[i[0] + delta // 2 - 1, j[0] + delta // 2 - 1] = c[0]
X[i[0] + delta // 2 - 1, j[0] + delta // 2 + 1] = c[0]
di, dj = torch.randint(3, (2,)) - 1
o = torch.tensor([[0.0, 1.0], [-1.0, 0.0]])
m = torch.eye(2)
- for _ in range(torch.randint(4, (1,))):
+ for _ in range(torch.randint(4, (1,)).item()):
m = m @ o
if torch.rand(1) < 0.5:
m[0, :] = -m[0, :]
c = torch.randperm(len(self.colors) - 1)[:3] + 1
dist = torch.empty(self.height + 2, self.width + 2)
for X, f_X in [(A, f_A), (B, f_B)]:
- nb_rec = torch.randint(3, (1,)) + 1
+ nb_rec = torch.randint(3, (1,)).item() + 1
while True:
r = self.rec_coo(nb_rec, prevent_overlap=True)
X[...] = 0
X[i1:i2, j1:j2] = c[0]
f_X[i1:i2, j1:j2] = c[0]
while True:
- i0, j0 = torch.randint(self.height, (1,)), torch.randint(
- self.width, (1,)
+ i0, j0 = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
)
if X[i0, j0] == 0:
break
while True:
- i1, j1 = torch.randint(self.height, (1,)), torch.randint(
- self.width, (1,)
+ i1, j1 = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
)
if X[i1, j1] == 0:
break
k = 0
for d in range(4):
while True:
- ii, jj = torch.randint(self.height, (1,)), torch.randint(
- self.width, (1,)
+ ii, jj = (
+ torch.randint(self.height, (1,)).item(),
+ torch.randint(self.width, (1,)).item(),
)
e = 0
for i in range(S):
if f_X[i + i0, j + j0] == c[d]:
X[ii + i, jj + j] = c[d]
+ def task_islands(self, A, f_A, B, f_B):
+ c = torch.randperm(len(self.colors) - 1)[:2] + 1
+ for X, f_X in [(A, f_A), (B, f_B)]:
+ while True:
+ k = torch.randperm(self.height * self.width)
+ Z = torch.zeros(self.height + 2, self.width + 2)
+
+ i0, j0 = (
+ torch.randint(self.height, (1,)).item() + 1,
+ torch.randint(self.width, (1,)).item() + 1,
+ )
+
+ Z[i0 - 1 : i0 + 2, j0 - 1 : j0 + 2] = 1
+
+ nb = 9
+
+ for q in k:
+ i, j = q % self.height + 1, q // self.height + 1
+
+ if Z[i, j] == 0:
+ r, s, t, u, v, w, x, y = (
+ Z[i - 1, j],
+ Z[i - 1, j + 1],
+ Z[i, j + 1],
+ Z[i + 1, j + 1],
+ Z[i + 1, j],
+ Z[i + 1, j - 1],
+ Z[i, j - 1],
+ Z[i - 1, j - 1],
+ )
+
+ if (
+ (nb < 16 or r + s + t + u + v + w + x + y > 0)
+ and (s == 0 or r + t > 0)
+ and (u == 0 or t + v > 0)
+ and (w == 0 or x + v > 0)
+ and (y == 0 or x + r > 0)
+ ):
+ # if r+s+t+u+v+w+x+y==0:
+ Z[i, j] = 1
+ nb += 1
+
+ if nb == self.height * self.width // 2:
+ break
+
+ if nb == self.height * self.width // 2:
+ break
+
+ M = Z.clone()
+ Z[i0, j0] = 2
+ X[...] = (Z[1:-1, 1:-1] == 1) * c[0] + (Z[1:-1, 1:-1] == 2) * c[1]
+
+ for _ in range(self.height + self.width):
+ Z[1:-1, 1:-1] = Z[1:-1, 1:-1].maximum(
+ torch.maximum(
+ torch.maximum(Z[0:-2, 1:-1], Z[2:, 1:-1]),
+ torch.maximum(Z[1:-1, 0:-2], Z[1:-1, 2:]),
+ )
+ )
+ Z *= M
+
+ f_X[...] = (Z[1:-1, 1:-1] == 1) * c[0] + (Z[1:-1, 1:-1] == 2) * c[1]
+
######################################################################
def trivial_prompts_and_answers(self, prompts, answers):
f_A = prompt[1 * (S + 1) : 1 * (S + 1) + S].view(self.height, self.width)
B = prompt[2 * (S + 1) : 2 * (S + 1) + S].view(self.height, self.width)
f_B = answer.view(self.height, self.width)
- task = tasks[torch.randint(len(tasks), (1,))]
+ task = tasks[torch.randint(len(tasks), (1,)).item()]
task(A, f_A, B, f_B)
return prompts.flatten(1), answers.flatten(1)
# nb, nrow = 8, 2
# for t in grids.all_tasks:
- for t in [
- grids.task_replace_color,
- grids.task_frame,
- ]:
+ for t in [grids.task_count]:
print(t.__name__)
prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
grids.save_quiz_illustrations(
nb = 1000
- for t in grids.all_tasks:
- # for t in [ grids.task_replace_color ]: #grids.all_tasks:
+ # for t in grids.all_tasks:
+ for t in [grids.task_islands]:
start_time = time.perf_counter()
prompts, answers = grids.generate_prompts_and_answers_(nb, tasks=[t])
delay = time.perf_counter() - start_time