X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=blobdiff_plain;f=picoclvr.py;h=5da39436749ac7dcff688a0fea476a8d507d3371;hb=27dd45177a8552a8b8c9e3a4d9388844fa1d4d27;hp=cc937af2a10ca89cfb4328c604e790c6866252a6;hpb=bd9e5951a5741f7e3e44fc03379795eff83242d6;p=culture.git
diff --git a/picoclvr.py b/picoclvr.py
index cc937af..5da3943 100755
--- a/picoclvr.py
+++ b/picoclvr.py
@@ -196,13 +196,11 @@ def generate(
nb_colors=5,
pruner=None,
):
-
assert nb_colors >= max_nb_squares and nb_colors <= len(color_name2rgb) - 1
descr = []
for n in range(nb):
-
nb_squares = torch.randint(max_nb_squares, (1,)) + 1
square_position = torch.randperm(height * width)[:nb_squares]
@@ -241,15 +239,8 @@ def generate(
# Extracts the image after in descr as a 1x3xHxW tensor
-def descr2img(descr, n, height, width):
-
- if type(descr) == list:
- return torch.cat([descr2img(d, n, height, width) for d in descr], 0)
-
- if type(n) == list:
- return torch.cat([descr2img(descr, k, height, width) for k in n], 0).unsqueeze(
- 0
- )
+def descr2img(descr, height, width):
+ result = []
def token2color(t):
try:
@@ -257,15 +248,15 @@ def descr2img(descr, n, height, width):
except KeyError:
return [128, 128, 128]
- d = descr.split("")
- d = d[n + 1] if len(d) > n + 1 else ""
- d = d.strip().split(" ")[: height * width]
- d = d + [""] * (height * width - len(d))
- d = [token2color(t) for t in d]
- img = torch.tensor(d).permute(1, 0)
- img = img.reshape(1, 3, height, width)
+ for d in descr:
+ d = d.split("")[1]
+ d = d.strip().split(" ")[: height * width]
+ d = d + [""] * (height * width - len(d))
+ d = [token2color(t) for t in d]
+ img = torch.tensor(d).permute(1, 0).reshape(1, 3, height, width)
+ result.append(img)
- return img
+ return torch.cat(result, 0)
######################################################################
@@ -274,7 +265,6 @@ def descr2img(descr, n, height, width):
def descr2properties(descr, height, width):
-
if type(descr) == list:
return [descr2properties(d, height, width) for d in descr]
@@ -319,7 +309,6 @@ def descr2properties(descr, height, width):
def nb_properties(descr, height, width, pruner=None):
-
if type(descr) == list:
return [nb_properties(d, height, width, pruner) for d in descr]
@@ -353,7 +342,7 @@ if __name__ == "__main__":
for d in descr:
f.write(f"{d}\n\n")
- img = descr2img(descr, n=0, height=12, width=16)
+ img = descr2img(descr, height=12, width=16)
if img.size(0) == 1:
img = F.pad(img, (1, 1, 1, 1), value=64)