def get_cd(imitate_net, data, one_hot_labels, program_len):
batch_size = data.shape[1]
beam_width = 10
all_beams, next_beams_prob, all_inputs = imitate_net.beam_search(
[data, one_hot_labels], beam_width, program_len)
beam_labels = beams_parser(all_beams, batch_size, beam_width=beam_width)
beam_labels_numpy = np.zeros((batch_size * beam_width, program_len),
dtype=np.int32)
for i in range(batch_size):
beam_labels_numpy[i * beam_width:(i + 1) *
beam_width, :] = beam_labels[i]
# find expression from these predicted beam labels
expressions = [""] * batch_size * beam_width
for i in range(batch_size * beam_width):
for j in range(program_len):
expressions[i] += unique_draw[beam_labels_numpy[i, j]]
for index, prog in enumerate(expressions):
expressions[index] = prog.split("$")[0]
predicted_images = image_from_expressions(parser, expressions)
target_images = data_[-1, :, 0, :, :].astype(dtype=bool)
target_images_new = np.repeat(target_images, axis=0, repeats=beam_width)
beam_CD = chamfer(target_images_new, predicted_images)
CD = np.zeros((batch_size, 1))
for r in range(batch_size):
CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width])
return np.sum(CD)
def infer_programs(imitate_net, self_training=False, ab=None):
config = read_config.Config("config_cad.yml")
# Load the terminals symbols of the grammar
with open("terminals.txt", "r") as file:
unique_draw = file.readlines()
for index, e in enumerate(unique_draw):
unique_draw[index] = e[0:-1]
config.train_size = 10000
config.test_size = 3000
imitate_net.eval()
imitate_net.epsilon = 0
parser = ParseModelOutput(unique_draw, max_len // 2 + 1, max_len,
config.canvas_shape)
generator = Generator()
test_gen = generator.test_gen(batch_size=config.batch_size,
path="data/cad/cad.h5",
if_augment=False)
pred_expressions = []
Rs = 0
CDs = 0
Target_images = []
for batch_idx in range(config.test_size // config.batch_size):
with torch.no_grad():
print(f"Inferring test cad batch: {batch_idx}")
data_ = next(test_gen)
labels = np.zeros((config.batch_size, max_len), dtype=np.int32)
one_hot_labels = prepare_input_op(labels, len(unique_draw))
one_hot_labels = torch.from_numpy(one_hot_labels).to(device)
data = torch.from_numpy(data_).to(device)
all_beams, next_beams_prob, all_inputs = imitate_net.beam_search(
[data[-1, :, 0, :, :], one_hot_labels], beam_width, max_len)
beam_labels = beams_parser(all_beams,
data_.shape[1],
beam_width=beam_width)
beam_labels_numpy = np.zeros(
(config.batch_size * beam_width, max_len), dtype=np.int32)
Target_images.append(data_[-1, :, 0, :, :])
for i in range(data_.shape[1]):
beam_labels_numpy[i * beam_width:(i + 1) *
beam_width, :] = beam_labels[i]
# find expression from these predicted beam labels
expressions = [""] * config.batch_size * beam_width
for i in range(config.batch_size * beam_width):
for j in range(max_len):
expressions[i] += unique_draw[beam_labels_numpy[i, j]]
for index, prog in enumerate(expressions):
expressions[index] = prog.split("$")[0]
pred_expressions += expressions
predicted_images = image_from_expressions(parser, expressions)
target_images = data_[-1, :, 0, :, :].astype(dtype=bool)
target_images_new = np.repeat(target_images,
axis=0,
repeats=beam_width)
beam_CD = chamfer(target_images_new, predicted_images)
CD = np.zeros((config.batch_size, 1))
for r in range(config.batch_size):
CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width])
CDs += np.mean(CD)
for j in range(0, config.batch_size):
f, a = plt.subplots(1, beam_width + 1, figsize=(30, 3))
a[0].imshow(data_[-1, j, 0, :, :], cmap="Greys_r")
a[0].axis("off")
a[0].set_title("target")
for i in range(1, beam_width + 1):
a[i].imshow(predicted_images[j * beam_width + i - 1],
cmap="Greys_r")
a[i].set_title("{}".format(i))
a[i].axis("off")
plt.savefig("best_lest/" +
"{}.png".format(batch_idx * config.batch_size + j),
transparent=0)
plt.close("all")
# with open("best_st_expressions.txt", "w") as file:
# for e in pred_expressions:
# file.write(f"{e}\n")
# break
return CDs / (config.test_size // config.batch_size)
def infer_programs(imitate_net, path, self_training=False, ab=None):
save_viz = False
config = read_config.Config("config_cad.yml")
# Load the terminals symbols of the grammar
with open("terminals.txt", "r") as file:
unique_draw = file.readlines()
for index, e in enumerate(unique_draw):
unique_draw[index] = e[0:-1]
config.train_size = 10000
config.test_size = 3000
imitate_net.eval()
imitate_net.epsilon = 0
parser = ParseModelOutput(unique_draw, max_len // 2 + 1, max_len,
config.canvas_shape)
pred_expressions = []
if ab is not None:
pred_labels = np.zeros((config.train_size * ab, max_len))
else:
pred_labels = np.zeros((config.train_size, max_len))
image_path = f"{path}/images/"
results_path = f"{path}/results/"
labels_path = f"{path}/labels/"
os.makedirs(os.path.dirname(image_path), exist_ok=True)
os.makedirs(os.path.dirname(results_path), exist_ok=True)
os.makedirs(os.path.dirname(labels_path), exist_ok=True)
os.makedirs(os.path.dirname(labels_path + "val/"), exist_ok=True)
generator = Generator()
train_gen = generator.train_gen(batch_size=config.batch_size,
path="data/cad/cad.h5",
if_augment=False)
val_gen = generator.val_gen(batch_size=config.batch_size,
path="data/cad/cad.h5",
if_augment=False)
Rs = 0
CDs = 0
Target_images = []
start = time.time()
pred_images = np.zeros((config.train_size, 64, 64))
for batch_idx in range(config.train_size // config.batch_size):
with torch.no_grad():
print(f"Inferring cad batch: {batch_idx}")
data_ = next(train_gen)
labels = np.zeros((config.batch_size, max_len), dtype=np.int32)
one_hot_labels = prepare_input_op(labels, len(unique_draw))
one_hot_labels = torch.from_numpy(one_hot_labels).to(device)
data = torch.from_numpy(data_).to(device)
all_beams, next_beams_prob, all_inputs = imitate_net.beam_search(
[data[-1, :, 0, :, :], one_hot_labels], beam_width, max_len)
beam_labels = beams_parser(all_beams,
data_.shape[1],
beam_width=beam_width)
beam_labels_numpy = np.zeros(
(config.batch_size * beam_width, max_len), dtype=np.int32)
Target_images.append(data_[-1, :, 0, :, :])
for i in range(data_.shape[1]):
beam_labels_numpy[i * beam_width:(i + 1) *
beam_width, :] = beam_labels[i]
# find expression from these predicted beam labels
expressions = [""] * config.batch_size * beam_width
for i in range(config.batch_size * beam_width):
for j in range(max_len):
expressions[i] += unique_draw[beam_labels_numpy[i, j]]
for index, prog in enumerate(expressions):
expressions[index] = prog.split("$")[0]
pred_expressions += expressions
predicted_images = image_from_expressions(parser, expressions)
target_images = data_[-1, :, 0, :, :].astype(dtype=bool)
target_images_new = np.repeat(target_images,
axis=0,
repeats=beam_width)
# beam_R = np.sum(np.logical_and(target_images_new, predicted_images),
# (1, 2)) / np.sum(np.logical_or(target_images_new, predicted_images), (1, 2))
#
# R = np.zeros((config.batch_size, 1))
# for r in range(config.batch_size):
# R[r, 0] = max(beam_R[r * beam_width:(r + 1) * beam_width])
#
# Rs += np.mean(R)
beam_CD = chamfer(target_images_new, predicted_images)
# select best expression by chamfer distance
if ab is None:
best_labels = np.zeros((config.batch_size, max_len))
for r in range(config.batch_size):
idx = np.argmin(beam_CD[r * beam_width:(r + 1) *
beam_width])
best_labels[r] = beam_labels[r][idx]
pred_labels[batch_idx *
config.batch_size:batch_idx * config.batch_size +
config.batch_size] = best_labels
else:
best_labels = np.zeros((config.batch_size * ab, max_len))
for r in range(config.batch_size):
sorted_idx = np.argsort(beam_CD[r * beam_width:(r + 1) *
beam_width])[:ab]
best_labels[r * ab:r * ab +
ab] = beam_labels[r][sorted_idx]
pred_labels[batch_idx * config.batch_size *
ab:batch_idx * config.batch_size * ab +
config.batch_size * ab] = best_labels
CD = np.zeros((config.batch_size, 1))
for r in range(config.batch_size):
CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width])
pred_images[batch_idx * config.batch_size +
r] = predicted_images[r * beam_width + np.argmin(
beam_CD[r * beam_width:(r + 1) * beam_width])]
CDs += np.mean(CD)
if save_viz:
for j in range(0, config.batch_size):
f, a = plt.subplots(1, beam_width + 1, figsize=(30, 3))
a[0].imshow(data_[-1, j, 0, :, :], cmap="Greys_r")
a[0].axis("off")
a[0].set_title("target")
for i in range(1, beam_width + 1):
a[i].imshow(predicted_images[j * beam_width + i - 1],
cmap="Greys_r")
a[i].set_title("{}".format(i))
a[i].axis("off")
plt.savefig(
image_path +
"{}.png".format(batch_idx * config.batch_size + j),
transparent=0)
plt.close("all")
save_viz = False
print("Inferring cad average chamfer distance: {}".format(
CDs / (config.train_size // config.batch_size)),
flush=True)
Rs = Rs / (config.train_size // config.batch_size)
CDs = CDs / (config.train_size // config.batch_size)
print(Rs, CDs)
results = {"iou": Rs, "chamferdistance": CDs}
with open(results_path + "results_beam_width_{}.org".format(beam_width),
'w') as outfile:
json.dump(results, outfile)
torch.save(pred_labels, labels_path + "labels.pt")
# torch.save(pred_images, labels_path + "images.pt")
if self_training:
if ab is None:
torch.save(np.concatenate(Target_images, axis=0),
labels_path + "images.pt")
else:
torch.save(
np.repeat(np.concatenate(Target_images, axis=0), ab, axis=0),
labels_path + "images.pt")
test_gen = generator.test_gen(batch_size=config.batch_size,
path="data/cad/cad.h5",
if_augment=False)
pred_expressions = []
Rs = 0
CDs = 0
Target_images = []
for batch_idx in range(config.test_size // config.batch_size):
with torch.no_grad():
print(f"Inferring test cad batch: {batch_idx}")
data_ = next(test_gen)
labels = np.zeros((config.batch_size, max_len), dtype=np.int32)
one_hot_labels = prepare_input_op(labels, len(unique_draw))
one_hot_labels = torch.from_numpy(one_hot_labels).to(device)
data = torch.from_numpy(data_).to(device)
all_beams, next_beams_prob, all_inputs = imitate_net.beam_search(
[data[-1, :, 0, :, :], one_hot_labels], beam_width, max_len)
beam_labels = beams_parser(all_beams,
data_.shape[1],
beam_width=beam_width)
beam_labels_numpy = np.zeros(
(config.batch_size * beam_width, max_len), dtype=np.int32)
Target_images.append(data_[-1, :, 0, :, :])
for i in range(data_.shape[1]):
beam_labels_numpy[i * beam_width:(i + 1) *
beam_width, :] = beam_labels[i]
# find expression from these predicted beam labels
expressions = [""] * config.batch_size * beam_width
for i in range(config.batch_size * beam_width):
for j in range(max_len):
expressions[i] += unique_draw[beam_labels_numpy[i, j]]
for index, prog in enumerate(expressions):
expressions[index] = prog.split("$")[0]
pred_expressions += expressions
predicted_images = image_from_expressions(parser, expressions)
target_images = data_[-1, :, 0, :, :].astype(dtype=bool)
target_images_new = np.repeat(target_images,
axis=0,
repeats=beam_width)
beam_CD = chamfer(target_images_new, predicted_images)
CD = np.zeros((config.batch_size, 1))
for r in range(config.batch_size):
CD[r, 0] = min(beam_CD[r * beam_width:(r + 1) * beam_width])
CDs += np.mean(CD)
print(f"TEST CD: {CDs / (config.test_size // config.batch_size)}")
end = time.time()
print(f"Inference time: {end-start}")
jitter_program=jit)
for k in dataset_sizes.keys():
test_batch_size = config.batch_size
for _ in range(dataset_sizes[k][1] // test_batch_size):
data_, labels = next(test_gen_objs[k])
one_hot_labels = prepare_input_op(labels,
len(generator.unique_draw))
one_hot_labels = Variable(torch.from_numpy(one_hot_labels)).cuda()
data = Variable(torch.from_numpy(data_), volatile=True).cuda()
labels = Variable(torch.from_numpy(labels)).cuda()
all_beams, next_beams_prob, all_inputs = imitate_net.beam_search(
[data, one_hot_labels], beam_width, maxx_len)
targ_prog = parser.labels2exps(labels, k)
beam_labels = beams_parser(all_beams,
test_batch_size,
beam_width=beam_width)
beam_labels_numpy = np.zeros(
(test_batch_size * beam_width, maxx_len), dtype=np.int32)
for i in range(test_batch_size):
beam_labels_numpy[i * beam_width:(i + 1) *
beam_width, :] = beam_labels[i]
# find expression from these predicted beam labels
expressions = [""] * test_batch_size * beam_width
for i in range(test_batch_size * beam_width):
for j in range(maxx_len):
expressions[i] += generator.unique_draw[beam_labels_numpy[
i, j]]
CDs = 0
Target_images = []
for batch_idx in range(config.test_size // config.batch_size):
with torch.no_grad():
print(batch_idx)
data_ = next(test_gen)
labels = np.zeros((config.batch_size, max_len), dtype=np.int32)
one_hot_labels = prepare_input_op(labels, len(unique_draw))
one_hot_labels = Variable(torch.from_numpy(one_hot_labels)).cuda()
data = Variable(torch.from_numpy(data_)).cuda()
all_beams, next_beams_prob, all_inputs = imitate_net.beam_search(
[data, one_hot_labels], beam_width, max_len)
beam_labels = beams_parser(all_beams,
data_.shape[1],
beam_width=beam_width)
beam_labels_numpy = np.zeros(
(config.batch_size * beam_width, max_len), dtype=np.int32)
Target_images.append(data_[-1, :, 0, :, :])
for i in range(data_.shape[1]):
beam_labels_numpy[i * beam_width:(i + 1) *
beam_width, :] = beam_labels[i]
# find expression from these predicted beam labels
expressions = [""] * config.batch_size * beam_width
for i in range(config.batch_size * beam_width):
for j in range(max_len):
expressions[i] += unique_draw[beam_labels_numpy[i, j]]
for index, prog in enumerate(expressions):