def _generate_item(self, index):
change = False
size_fits = False
while not change and not size_fits:
tform, text, task_data, schematic = self.generate_task()
new_schematic = tform(schematic, **task_data["tform_data"])
if max(new_schematic.shape) <= self.sidelength:
size_fits = True
new_schematic = shape_transforms.moment_at_center(
new_schematic, self.sidelength)
diff = (schematic - new_schematic).any()
if self.allow_same:
change = True
if not diff:
text = ["no", "change"]
else:
if diff:
change = True
w = self.maybe_words_to_tensor(text)
schematic = self.maybe_hash_idm(torch.LongTensor(schematic))
new_schematic = self.maybe_hash_idm(torch.LongTensor(new_schematic))
if task_data["inverse"]:
return w, new_schematic, schematic
else:
return w, schematic, new_schematic
def generate_task(self):
size_fits = False
while not size_fits:
schematic, _ = blocks_list_to_npy(get_schematic(), xyz=True)
if max(schematic.shape) < self.sidelength:
size_fits = True
schematic = shape_transforms.moment_at_center(
schematic, self.sidelength)
tform_name = self.tform_names[self.template_sampler.sample()]
tform, text, task_data = self.templates[tform_name](schematic)
return tform, text, task_data, schematic