def run_batch(batch_idx, val, batch_loader, tracker_cnn, criterion, optimizer, history, save_debug_image):
"""Train or validate on a single batch."""
train = not val
time_cbatch_start = time.time()
inputs, outputs_gt = batch_loader.get_batch()
if Config.GPU >= 0:
inputs = to_cuda(to_variable(inputs, volatile=val), Config.GPU)
outputs_gt_bins = to_cuda(to_variable(np.argmax(outputs_gt, axis=1), volatile=val, requires_grad=False), Config.GPU)
outputs_gt = to_cuda(to_variable(outputs_gt, volatile=val, requires_grad=False), Config.GPU)
time_cbatch_end = time.time()
time_fwbw_start = time.time()
if train:
optimizer.zero_grad()
outputs_pred = tracker_cnn(inputs)
outputs_pred_sm = F.softmax(outputs_pred)
loss = criterion(outputs_pred, outputs_gt_bins)
if train:
loss.backward()
optimizer.step()
time_fwbw_end = time.time()
loss = loss.data.cpu().numpy()[0]
outputs_pred_np = to_numpy(outputs_pred_sm)
outputs_gt_np = to_numpy(outputs_gt)
acc = np.sum(np.equal(np.argmax(outputs_pred_np, axis=1), np.argmax(outputs_gt_np, axis=1))) / BATCH_SIZE
history.add_value("loss", "train" if train else "val", batch_idx, loss, average=val)
history.add_value("acc", "train" if train else "val", batch_idx, acc, average=val)
print("[%s] Batch %05d | loss %.8f | acc %.2f | cbatch %.04fs | fwbw %.04fs" % ("T" if train else "V", batch_idx, loss, acc, time_cbatch_end - time_cbatch_start, time_fwbw_end - time_fwbw_start))
if save_debug_image:
debug_img = generate_debug_image(inputs, outputs_gt, outputs_pred_sm)
misc.imsave("debug_img_%s.jpg" % ("train" if train else "val"), debug_img)
def draw_frame_grids(self, scr, grids):
grids_meta = [(0, "street boundaries"),
(3, "crashables (except cars)"), (7, "street markings"),
(4, "current lane"), (1, "cars"), (2, "cars in mirrors")]
titles = [title for idx, title in grids_meta]
grids = to_numpy(grids[0])
grids = [grids[idx] for idx, title in grids_meta]
#self.grid_to_graph(scr, grids[0])
bgcolor = [0, 0, 0]
image = np.zeros((720, 1280, 3), dtype=np.uint8) + bgcolor
scr_main = ia.imresize_single_image(
scr, (int(720 * 0.58), int(1280 * 0.58)))
#util.draw_image(image, y=720-scr_main.shape[0], x=1080-scr_main.shape[1], other_img=scr_main, copy=False)
util.draw_image(image,
y=int((image.shape[0] - scr_main.shape[0]) / 2),
x=1280 - scr_main.shape[1] - 2,
other_img=scr_main,
copy=False)
image = util.draw_text(
image,
x=1280 - (scr_main.shape[1] // 2) - 125,
y=image.shape[0] - int(
(image.shape[0] - scr_main.shape[0]) / 2) + 10,
text="Framerate matches the one that the model sees (10fps).",
size=10,
color=[128, 128, 128])
grid_rel_size = 0.19
scr_small = ia.imresize_single_image(
scr, (int(720 * grid_rel_size), int(1280 * grid_rel_size)))
grid_hms = []
for grid, title in zip(grids, titles):
grid = (grid * 255).astype(np.uint8)[:, :, np.newaxis]
grid = ia.imresize_single_image(
grid, (int(720 * grid_rel_size), int(1280 * grid_rel_size)),
interpolation="nearest")
grid_hm = util.draw_heatmap_overlay(scr_small, grid / 255)
grid_hm = np.pad(grid_hm, ((2, 0), (2, 2), (0, 0)),
mode="constant",
constant_values=np.average(bgcolor))
#grid_hm = np.pad(grid_hm, ((0, 20), (0, 0), (0, 0)), mode="constant", constant_values=0)
#grid_hm[-20:, 2:-2, :] = [128, 128, 255]
#grid_hm = util.draw_text(grid_hm, x=4, y=grid_hm.shape[0]-16, text=title, size=10, color=[255, 255, 255])
grid_hm = np.pad(grid_hm, ((40, 0), (0, 0), (0, 0)),
mode="constant",
constant_values=0)
grid_hm = util.draw_text(grid_hm,
x=4,
y=20,
text=title,
size=12,
color=[255, 255, 255])
grid_hms.append(grid_hm)
grid_hms = ia.draw_grid(grid_hms, cols=2)
util.draw_image(image, y=70, x=0, other_img=grid_hms, copy=False)
return image
def remove_unannotated_atts_gt(outputs_atts_pred, outputs_atts_gt,
atts_annotated):
"""Zero-grad attribute outputs for which there is no annotation data for an
example."""
gt2 = np.copy(outputs_atts_gt)
pred = to_numpy(outputs_atts_pred)
for b_idx in xrange(atts_annotated.shape[0]):
if atts_annotated[b_idx, 0] == 0:
gt2[b_idx, ...] = pred[b_idx, ...]
return gt2
def remove_unannotated_grids_gt(outputs_grids_pred, outputs_grids_gt,
grids_annotated):
"""Zero-grad grid outputs for which there is no annotation data for an
example."""
gt2 = np.copy(outputs_grids_gt)
pred = to_numpy(outputs_grids_pred)
for b_idx in xrange(grids_annotated.shape[0]):
for grid_idx in xrange(grids_annotated.shape[1]):
if grids_annotated[b_idx, grid_idx] == 0:
gt2[b_idx, grid_idx, ...] = pred[b_idx, grid_idx, ...]
return gt2
def generate_debug_image(inputs, outputs_gt, outputs_pred):
"""Draw an image with current ground truth and predictions for debug purposes."""
current_image = inputs.data[0].cpu().numpy()
current_image = np.clip(current_image * 255, 0, 255).astype(np.uint8).transpose((1, 2, 0))
current_image = ia.imresize_single_image(current_image, (32*4, 64*4))
h, w = current_image.shape[0:2]
outputs_gt = to_numpy(outputs_gt)[0]
outputs_pred = to_numpy(outputs_pred)[0]
binwidth = 6
outputs_grid = np.zeros((20+2, outputs_gt.shape[0]*binwidth, 3), dtype=np.uint8)
for angle_bin_idx in xrange(outputs_gt.shape[0]):
val = outputs_pred[angle_bin_idx]
x_start = angle_bin_idx*binwidth
x_end = (angle_bin_idx+1)*binwidth
fill_start = 1
fill_end = 1 + int(20*val)
#print(angle_bin_idx, x_start, x_end, fill_start, fill_end, outputs_grid.shape, outputs_grid[fill_start:fill_end, x_start+1:x_end].shape)
if fill_start < fill_end:
outputs_grid[fill_start:fill_end, x_start+1:x_end] = [255, 255, 255]
bordercol = [128, 128, 128] if outputs_gt[angle_bin_idx] < 1 else [0, 0, 255]
outputs_grid[0:22, x_start:x_start+1] = bordercol
outputs_grid[0:22, x_end:x_end+1] = bordercol
outputs_grid[0, x_start:x_end+1] = bordercol
outputs_grid[21, x_start:x_end+1] = bordercol
outputs_grid = outputs_grid[::-1, :, :]
bin_gt = np.argmax(outputs_gt)
bin_pred = np.argmax(outputs_pred)
angles = [(binidx*ANGLE_BIN_SIZE) - 180 for binidx in [bin_gt, bin_pred]]
#print(outputs_grid.shape)
current_image = np.pad(current_image, ((0, 128), (0, 400), (0, 0)), mode="constant")
current_image[h+4:h+4+22, 4:4+outputs_grid.shape[1], :] = outputs_grid
current_image = util.draw_text(current_image, x=4, y=h+4+22+4, text="GT: %03.2fdeg\nPR: %03.2fdeg" % (angles[0], angles[1]), size=10)
return current_image
def draw_frame_attributes(self, scr, atts):
atts = atts[0]
mincolf = 0.2
#print("space_front raw", atts[33:37], F.softmax(atts[33:37]))
#print("space_left raw", atts[37:41], F.softmax(atts[37:41]))
#print("space_right raw", atts[41:45], F.softmax(atts[41:45].unsqueeze(0)).squeeze(0))
road_type = simplesoftmax(to_numpy(atts[0:10]))
intersection = simplesoftmax(to_numpy(atts[10:17]))
direction = simplesoftmax(to_numpy(atts[17:20]))
lane_count = simplesoftmax(to_numpy(atts[20:25]))
curve = simplesoftmax(to_numpy(atts[25:33]))
space_front = simplesoftmax(to_numpy(atts[33:37]))
space_left = simplesoftmax(to_numpy(atts[37:41]))
space_right = simplesoftmax(to_numpy(atts[41:45]))
offroad = simplesoftmax(to_numpy(atts[45:48]))
bgcolor = [0, 0, 0]
image = np.zeros((720, 1280, 3), dtype=np.uint8) + bgcolor
scr_main = ia.imresize_single_image(
scr, (int(720 * 0.58), int(1280 * 0.58)))
util.draw_image(image,
y=int((image.shape[0] - scr_main.shape[0]) / 2),
x=1280 - scr_main.shape[1] - 2,
other_img=scr_main,
copy=False)
image = util.draw_text(
image,
x=1280 - (scr_main.shape[1] // 2) - 125,
y=image.shape[0] - int(
(image.shape[0] - scr_main.shape[0]) / 2) + 10,
text="Framerate matches the one that the model sees (10fps).",
size=10,
color=[128, 128, 128])
# ---------------
# Curve
# ---------------
"""
street = np.zeros((65, 65, 3), dtype=np.float32)
street[:, 0:2, :] = 255
street[:, -2:, :] = 255
street[:, 32:35, :] = 255
street_left_strong = curve(street
"""
curve_left_strong = 255 - ndimage.imread(
"../images/video/curve-left-strong.png", mode="RGB")
curve_left_medium = 255 - ndimage.imread(
"../images/video/curve-left-medium.png", mode="RGB")
curve_left_slight = 255 - ndimage.imread(
"../images/video/curve-left-slight.png", mode="RGB")
curve_straight = 255 - ndimage.imread(
"../images/video/curve-straight.png", mode="RGB")
curve_right_strong = np.fliplr(curve_left_strong)
curve_right_medium = np.fliplr(curve_left_medium)
curve_right_slight = np.fliplr(curve_left_slight)
curve_straight = (curve_straight *
np.clip(curve[0], mincolf, 1.0)).astype(np.uint8)
curve_left_slight = (curve_left_slight *
np.clip(curve[1], mincolf, 1.0)).astype(np.uint8)
curve_left_medium = (curve_left_medium *
np.clip(curve[2], mincolf, 1.0)).astype(np.uint8)
curve_left_strong = (curve_left_strong *
np.clip(curve[3], mincolf, 1.0)).astype(np.uint8)
curve_right_slight = (curve_right_slight *
np.clip(curve[4], mincolf, 1.0)).astype(np.uint8)
curve_right_medium = (curve_right_medium *
np.clip(curve[5], mincolf, 1.0)).astype(np.uint8)
curve_right_strong = (curve_right_strong *
np.clip(curve[6], mincolf, 1.0)).astype(np.uint8)
def add_perc(curve_img, perc, x_correct):
col = np.clip(255 * perc, mincolf * 255, 255)
col = np.array([col, col, col], dtype=np.uint8)
curve_img_pad = np.pad(curve_img, ((0, 20), (0, 0), (0, 0)),
mode="constant",
constant_values=0)
x = int(curve_img_pad.shape[1] / 2) - 6
if (perc * 100) >= 100:
x = x - 9
elif (perc * 100) >= 10:
x = x - 6
x = x + x_correct
curve_img_pad = util.draw_text(curve_img_pad,
x=x,
y=curve_img_pad.shape[0] - 15,
text="%.0f%%" % (perc * 100, ),
color=col,
size=9)
return curve_img_pad
curve_straight = add_perc(curve_straight, curve[0], x_correct=0)
curve_left_slight = add_perc(curve_left_slight, curve[1], x_correct=3)
curve_left_medium = add_perc(curve_left_medium, curve[2], x_correct=1)
curve_left_strong = add_perc(curve_left_strong, curve[3], x_correct=-1)
curve_right_slight = add_perc(curve_right_slight,
curve[4],
x_correct=-3)
curve_right_medium = add_perc(curve_right_medium,
curve[5],
x_correct=-2)
curve_right_strong = add_perc(curve_right_strong,
curve[6],
x_correct=0)
curves = np.hstack([
curve_left_strong, curve_left_medium, curve_left_slight,
curve_straight, curve_right_slight, curve_right_medium,
curve_right_strong
])
curves = np.pad(curves, ((50, 0), (20, 0), (0, 0)),
mode="constant",
constant_values=0)
curves = util.draw_text(curves,
x=4,
y=4,
text="Curve",
color=[255, 255, 255])
util.draw_image(image, y=50, x=2, other_img=curves, copy=False)
# ---------------
# Lane count
# ---------------
pics = []
for lc_idx in range(4):
col = int(np.clip(255 * lane_count[lc_idx], 255 * mincolf, 255))
col = np.array([col, col, col], dtype=np.uint8)
lc = lc_idx + 1
marking_width = 2
street = np.zeros((64, 64, 3), dtype=np.float32)
street[:, 0:marking_width, :] = col
street[:, -marking_width:, :] = col
inner_width = street.shape[1] - 2 * marking_width
lane_width = int((inner_width - (lc - 1) * marking_width) // lc)
start = marking_width
for i in range(lc - 1):
mstart = start + lane_width
mend = mstart + marking_width
street[1::6, mstart:mend, :] = col
street[2::6, mstart:mend, :] = col
street[3::6, mstart:mend, :] = col
start = mend
x = 14 + 24
if lane_count[lc_idx] * 100 >= 10:
x = x - 8
elif lane_count[lc_idx] * 100 >= 100:
x = x - 12
street = np.pad(street, ((0, 20), (14, 14), (0, 0)),
mode="constant",
constant_values=0)
street = util.draw_text(street,
x=x,
y=street.shape[0] - 14,
text="%.0f%%" %
(lane_count[lc_idx] * 100, ),
size=9,
color=col)
pics.append(street)
pics = np.hstack(pics)
pics = np.pad(pics, ((55, 0), (20, 0), (0, 0)),
mode="constant",
constant_values=0)
pics = util.draw_text(pics,
x=4,
y=4,
text="Lane Count",
color=[255, 255, 255])
util.draw_image(image, y=250, x=2, other_img=pics, copy=False)
# ---------------
# Space
# ---------------
truck = np.zeros((100, 55, 3), dtype=np.uint8)
truck[0:2, :, :] = 255
truck[0:20, 0:2, :] = 255
truck[0:20, -2:, :] = 255
truck[20:22, :, :] = 255
truck[22:25, 25:27, :] = 255
truck[22:25, 29:31, :] = 255
truck[24:26, :, :] = 255
truck[24:, 0:2, :] = 255
truck[24:, -2:, :] = 255
truck[24:, -2:, :] = 255
truck[-2:, :, :] = 255
truck_full = np.pad(truck, ((50, 50), (100, 50), (0, 0)),
mode="constant",
constant_values=np.average(bgcolor))
#print("space_front", space_front)
#print("space_right", space_right)
#print("space_left", space_left)
fill_top = 1 * space_front[0] + 0.6 * space_front[
1] + 0.25 * space_front[2] + 0 * space_front[3]
fill_right = 1 * space_right[0] + 0.6 * space_right[
1] + 0.25 * space_right[2] + 0 * space_right[3]
fill_left = 1 * space_left[0] + 0.6 * space_left[
1] + 0.25 * space_left[2] + 0 * space_left[3]
r_outer_top = 8 + int((30 - 8) * fill_top)
r_outer_right = 8 + int((30 - 8) * fill_right)
r_outer_left = 8 + int((30 - 8) * fill_left)
def fill_to_text(fill):
col = np.array([255, 255, 255], dtype=np.uint8)
if fill > 0.75:
text = "plenty"
elif fill > 0.5:
text = "some"
elif fill > 0.25:
text = "low"
else:
text = "minimal"
return text, col
#top
truck_full = util.draw_direction_circle(truck_full,
y=33,
x=100 + 27,
r_inner=8,
r_outer=30,
angle_start=-60,
angle_end=60,
color_border=[255, 255, 255],
color_fill=[0, 0, 0])
truck_full = util.draw_direction_circle(truck_full,
y=33,
x=100 + 27,
r_inner=8,
r_outer=r_outer_top,
angle_start=-60,
angle_end=60,
color_border=[255, 255, 255],
color_fill=[255, 255, 255])
#text, col = fill_to_text(fill_top)
#truck_full = util.draw_text(truck_full, x=100+27, y=15, text=text, size=9, color=col)
# right
truck_full = util.draw_direction_circle(truck_full,
y=100,
x=170,
r_inner=8,
r_outer=30,
angle_start=30,
angle_end=180 - 30,
color_border=[255, 255, 255],
color_fill=[0, 0, 0])
truck_full = util.draw_direction_circle(truck_full,
y=100,
x=170,
r_inner=8,
r_outer=r_outer_right,
angle_start=30,
angle_end=180 - 30,
color_border=[255, 255, 255],
color_fill=[255, 255, 255])
#text, col = fill_to_text(fill_right)
#truck_full = util.draw_text(truck_full, x=170, y=100, text=text, size=9, color=col)
# left
truck_full = util.draw_direction_circle(truck_full,
y=100,
x=83,
r_inner=8,
r_outer=30,
angle_start=180 + 30,
angle_end=360 - 30,
color_border=[255, 255, 255],
color_fill=[0, 0, 0])
truck_full = util.draw_direction_circle(truck_full,
y=100,
x=83,
r_inner=8,
r_outer=r_outer_left,
angle_start=180 + 30,
angle_end=360 - 30,
color_border=[255, 255, 255],
color_fill=[255, 255, 255])
#text, col = fill_to_text(fill_left)
#truck_full = util.draw_text(truck_full, x=75, y=100, text=text, size=9, color=col)
truck_full = np.pad(truck_full, ((50, 0), (110, 0), (0, 0)),
mode="constant",
constant_values=0)
truck_full = util.draw_text(truck_full,
x=4,
y=4,
text="Space",
color=[255, 255, 255])
util.draw_image(image, y=450, x=10, other_img=truck_full, copy=False)
return image
def generate_debug_image(images, images_prev, \
outputs_ae_gt, outputs_grids_gt, outputs_atts_gt, \
outputs_multiactions_gt, outputs_flow_gt, outputs_canny_gt, \
outputs_flipped_gt, \
outputs_ae_pred, outputs_grids_pred, outputs_atts_pred, \
outputs_multiactions_pred, outputs_flow_pred, outputs_canny_pred, \
outputs_flipped_pred, \
grids_annotated, atts_annotated):
image = to_numpy(images)[0]
grids_annotated = grids_annotated[0]
atts_annotated = atts_annotated[0]
ae_gt = to_numpy(outputs_ae_gt)[0]
grids_gt = to_numpy(outputs_grids_gt)[0]
atts_gt = to_numpy(outputs_atts_gt)[0]
multiactions_gt = to_numpy(outputs_multiactions_gt)[0]
flow_gt = to_numpy(outputs_flow_gt)[0]
canny_gt = to_numpy(outputs_canny_gt)[0]
flipped_gt = to_numpy(outputs_flipped_gt)[0]
ae_pred = to_numpy(outputs_ae_pred)[0]
grids_pred = to_numpy(outputs_grids_pred)[0]
atts_pred = to_numpy(outputs_atts_pred)[0]
multiactions_pred = to_numpy(outputs_multiactions_pred)[0]
flow_pred = to_numpy(outputs_flow_pred)[0]
canny_pred = to_numpy(outputs_canny_pred)[0]
flipped_pred = to_numpy(outputs_flipped_pred)[0]
image = (np.squeeze(image).transpose(1, 2, 0) * 255).astype(np.uint8)
ae_pred = (np.squeeze(ae_pred).transpose(1, 2, 0) * 255).astype(np.uint8)
grids_gt = (np.squeeze(grids_gt).transpose(1, 2, 0) * 255).astype(np.uint8)
grids_pred = (np.squeeze(grids_pred).transpose(1, 2, 0) * 255).astype(
np.uint8)
atts_gt = np.squeeze(atts_gt)
atts_pred = np.squeeze(atts_pred)
multiactions_gt = np.squeeze(multiactions_gt)
multiactions_pred = np.squeeze(multiactions_pred)
flow_gt = (flow_gt.transpose(1, 2, 0) * 255).astype(np.uint8)
flow_pred = (flow_pred.transpose(1, 2, 0) * 255).astype(np.uint8)
canny_gt = (canny_gt.transpose(1, 2, 0) * 255).astype(np.uint8)
canny_pred = (canny_pred.transpose(1, 2, 0) * 255).astype(np.uint8)
#print(image.shape, grid_gt.shape, grid_pred.shape)
h, w = int(image.shape[0] * 0.5), int(image.shape[1] * 0.5)
#h, w = image.shape[0:2]
image_rs = ia.imresize_single_image(image, (h, w), interpolation="cubic")
grids_vis = []
for i in xrange(grids_gt.shape[2]):
grid_gt_rs = ia.imresize_single_image(grids_gt[..., i][:, :,
np.newaxis],
(h, w),
interpolation="cubic")
grid_pred_rs = ia.imresize_single_image(grids_pred[..., i][:, :,
np.newaxis],
(h, w),
interpolation="cubic")
grid_gt_hm = util.draw_heatmap_overlay(image_rs,
np.squeeze(grid_gt_rs) / 255)
grid_pred_hm = util.draw_heatmap_overlay(
image_rs,
np.squeeze(grid_pred_rs) / 255)
if grids_annotated[i] == 0:
grid_gt_hm[::4, ::4, :] = [255, 0, 0]
grids_vis.append(np.hstack((grid_gt_hm, grid_pred_hm)))
"""
lst = [image[0:3]] \
+ [image[3:6]] \
+ [ia.imresize_single_image(ae_pred[:, :, 0:3], (image.shape[0], image.shape[1]), interpolation="cubic")] \
+ [ia.imresize_single_image(ae_pred[:, :, 3:6], (image.shape[0], image.shape[1]), interpolation="cubic")] \
+ [ia.imresize_single_image(np.tile(flow_pred[:, :, 0][:, :, np.newaxis], (1, 1, 3)), (image.shape[0], image.shape[1]), interpolation="cubic")] \
+ [ia.imresize_single_image(np.tile(flow_pred[:, :, 1][:, :, np.newaxis], (1, 1, 3)), (image.shape[0], image.shape[1]), interpolation="cubic")] \
+ grids_vis
print([s.shape for s in lst])
"""
def downscale(im):
return ia.imresize_single_image(
im, (image.shape[0] // 2, image.shape[1] // 2),
interpolation="cubic")
def to_rgb(im):
if im.ndim == 2:
im = im[:, :, np.newaxis]
return np.tile(im, (1, 1, 3))
#print(canny_gt.shape, canny_gt[...,0].shape, to_rgb(canny_gt[...,0]).shape, downscale(to_rgb(canny_gt[...,0])).shape)
#print(canny_pred.shape, canny_gt[...,0].shape, to_rgb(canny_pred[...,0]).shape, downscale(to_rgb(canny_pred[...,0])).shape)
current_image = np.vstack(
#[image[:, :, 0:3]]
[image[:, :, 0:3]] + grids_vis + [
np.hstack([
downscale(ae_pred[:, :, 0:3]),
downscale(to_rgb(ae_pred[:, :, 3]))
])
] + [
np.hstack([
downscale(to_rgb(ae_pred[:, :, 4])),
# downscale(to_rgb(ae_pred[:, :, 5]))
np.zeros_like(downscale(to_rgb(ae_pred[:, :, 4])))
])
] + [
np.hstack([
downscale(to_rgb(flow_gt[..., 0])),
downscale(to_rgb(flow_pred[..., 0]))
])
] + [
np.hstack([
downscale(to_rgb(canny_gt[..., 0])),
downscale(to_rgb(canny_pred[..., 0]))
])
])
y_grids_start = image.shape[0]
grid_height = grids_vis[0].shape[0]
for i, name in enumerate(train.GRIDS_ORDER):
current_image = util.draw_text(current_image,
x=2,
y=y_grids_start +
(i + 1) * grid_height - 12,
text=name,
size=8,
color=[0, 255, 0])
current_image = np.pad(current_image, ((0, 280), (0, 280), (0, 0)),
mode="constant",
constant_values=0)
texts = []
att_idx = 0
for i, att_group in enumerate(ATTRIBUTE_GROUPS):
texts.append(att_group.name_shown)
for j, att in enumerate(att_group.attributes):
if atts_annotated[0] == 0:
texts.append(" %s | ? | %.2f" % (att.name, atts_pred[att_idx]))
else:
texts.append(" %s | %.2f | %.2f" %
(att.name, atts_gt[att_idx], atts_pred[att_idx]))
att_idx += 1
current_image = util.draw_text(current_image,
x=current_image.shape[1] - 256 + 1,
y=1,
text="\n".join(texts),
size=8,
color=[0, 255, 0])
ma_texts = ["multiactions (prev avg, next avg, curr, next)"]
counter = 0
while counter < multiactions_gt.shape[0]:
ma_texts_sub = ([], [])
for j in xrange(9):
ma_texts_sub[0].append("%.2f" % (multiactions_gt[counter], ))
ma_texts_sub[1].append("%.2f" % (multiactions_pred[counter], ))
counter += 1
ma_texts.append(" ".join(ma_texts_sub[0]))
ma_texts.append(" ".join(ma_texts_sub[1]))
ma_texts.append("")
current_image = util.draw_text(current_image,
x=current_image.shape[1] - 256 + 1,
y=650,
text="\n".join(ma_texts),
size=8,
color=[0, 255, 0])
flipped_texts = [
"flipped", " ".join(
["%.2f" % (flipped_gt[i], ) for i in xrange(flipped_gt.shape[0])]),
" ".join([
"%.2f" % (flipped_pred[i], ) for i in xrange(flipped_pred.shape[0])
])
]
current_image = util.draw_text(current_image,
x=current_image.shape[1] - 256 + 1,
y=810,
text="\n".join(flipped_texts),
size=8,
color=[0, 255, 0])
return current_image
def main():
"""Initialize/load model, dataset, optimizers, history and loss
plotter, augmentation sequence. Then start training loop."""
parser = argparse.ArgumentParser(description="Train semisupervised model")
parser.add_argument('--nocontinue',
default=False,
action="store_true",
help="Whether to NOT continue the previous experiment",
required=False)
parser.add_argument(
'--withshortcuts',
default=False,
action="store_true",
help=
"Whether to train a model with shortcuts from downscaling to upscaling layers.",
required=False)
args = parser.parse_args()
checkpoint_fp = "train_semisupervised_model%s.tar" % (
"_withshortcuts" if args.withshortcuts else "", )
if os.path.isfile(checkpoint_fp) and not args.nocontinue:
checkpoint = torch.load(checkpoint_fp)
else:
checkpoint = None
# load or initialize loss history
if checkpoint is not None:
history = plotting.History.from_string(checkpoint["history"])
else:
history = plotting.History()
history.add_group("loss-ae", ["train", "val"], increasing=False)
history.add_group("loss-grids", ["train", "val"], increasing=False)
history.add_group("loss-atts", ["train", "val"], increasing=False)
history.add_group("loss-multiactions", ["train", "val"],
increasing=False)
history.add_group("loss-flow", ["train", "val"], increasing=False)
history.add_group("loss-canny", ["train", "val"], increasing=False)
history.add_group("loss-flipped", ["train", "val"], increasing=False)
# initialize loss plotter
loss_plotter = plotting.LossPlotter(
history.get_group_names(),
history.get_groups_increasing(),
save_to_fp="train_semisupervised_plot%s.jpg" %
("_withshortcuts" if args.withshortcuts else "", ))
loss_plotter.start_batch_idx = 100
# initialize and load model
predictor = models.Predictor(
) if not args.withshortcuts else models.PredictorWithShortcuts()
if checkpoint is not None:
predictor.load_state_dict(checkpoint["predictor_state_dict"])
predictor.train()
# initialize optimizer
optimizer_predictor = optim.Adam(predictor.parameters())
# initialize losses
criterion_ae = nn.MSELoss()
criterion_grids = nn.BCELoss()
criterion_atts = nn.BCELoss()
criterion_multiactions = nn.BCELoss()
criterion_flow = nn.BCELoss()
criterion_canny = nn.BCELoss()
criterion_flipped = nn.BCELoss()
# send everything to gpu
if GPU >= 0:
predictor.cuda(GPU)
criterion_ae.cuda(GPU)
criterion_grids.cuda(GPU)
criterion_atts.cuda(GPU)
criterion_multiactions.cuda(GPU)
criterion_flow.cuda(GPU)
criterion_canny.cuda(GPU)
criterion_flipped.cuda(GPU)
# initialize image augmentation cascade
rarely = lambda aug: iaa.Sometimes(0.1, aug)
sometimes = lambda aug: iaa.Sometimes(0.2, aug)
often = lambda aug: iaa.Sometimes(0.3, aug)
# no hflips here, because that would mess up the optimal steering direction
# no grayscale here, because that doesn't play well with the grayscale
# previous images
# no coarse dropout, because then the model would have to magically guess
# things like edges or flow
augseq = iaa.Sequential(
[
often(iaa.Crop(percent=(0, 0.05))),
sometimes(iaa.GaussianBlur(
(0, 0.2))), # blur images with a sigma between 0 and 3.0
often(
iaa.AdditiveGaussianNoise(
loc=0, scale=(0.0, 0.01 * 255),
per_channel=0.5)), # add gaussian noise to images
often(iaa.Dropout((0.0, 0.05), per_channel=0.5)),
rarely(iaa.Sharpen(alpha=(0, 0.7),
lightness=(0.75, 1.5))), # sharpen images
rarely(iaa.Emboss(alpha=(0, 0.7),
strength=(0, 2.0))), # emboss images
rarely(
iaa.Sometimes(
0.5,
iaa.EdgeDetect(alpha=(0, 0.4)),
iaa.DirectedEdgeDetect(alpha=(0, 0.4),
direction=(0.0, 1.0)),
)),
often(iaa.Add(
(-20, 20), per_channel=0.5
)), # change brightness of images (by -10 to 10 of original value)
often(iaa.Multiply((0.8, 1.2), per_channel=0.25)
), # change brightness of images (50-150% of original value)
often(iaa.ContrastNormalization(
(0.8, 1.2),
per_channel=0.5)), # improve or worsen the contrast
sometimes(
iaa.Affine(scale={
"x": (0.9, 1.1),
"y": (0.9, 1.1)
},
translate_percent={
"x": (-0.07, 0.07),
"y": (-0.07, 0.07)
},
rotate=(0, 0),
shear=(0, 0),
order=[0, 1],
cval=(0, 255),
mode=ia.ALL))
],
random_order=True # do all of the above in random order
)
# load datasets
print("Loading dataset...")
if USE_COMPRESSED_ANNOTATIONS:
examples = load_dataset_annotated_compressed()
else:
examples = load_dataset_annotated()
#examples_annotated_ids = set([ex.state_idx for ex in examples])
examples_annotated_ids = set()
examples_autogen_val = load_dataset_autogen(val=True,
nb_load=NB_AUTOGEN_VAL,
not_in=examples_annotated_ids)
examples_autogen_train = load_dataset_autogen(
val=False, nb_load=NB_AUTOGEN_TRAIN, not_in=examples_annotated_ids)
random.shuffle(examples)
random.shuffle(examples_autogen_val)
random.shuffle(examples_autogen_train)
examples_val = examples[0:NB_VAL_SPLIT]
examples_train = examples[NB_VAL_SPLIT:]
# initialize background batch loaders
#memory = replay_memory.ReplayMemory.get_instance_supervised()
batch_loader_train = BatchLoader(examples_train,
examples_autogen_train,
augseq=augseq,
queue_size=15,
nb_workers=4,
threaded=False)
batch_loader_val = BatchLoader(examples_val,
examples_autogen_val,
augseq=iaa.Noop(),
queue_size=NB_VAL_BATCHES,
nb_workers=1,
threaded=False)
# training loop
print("Training...")
start_batch_idx = 0 if checkpoint is None else checkpoint["batch_idx"] + 1
for batch_idx in xrange(start_batch_idx, NB_BATCHES):
# train on batch
# load batch data
time_cbatch_start = time.time()
(inputs,
inputs_prev), (outputs_ae_gt, outputs_grids_gt_orig,
outputs_atts_gt_orig, outputs_multiactions_gt,
outputs_flow_gt, outputs_canny_gt,
outputs_flipped_gt), (
grids_annotated,
atts_annotated) = batch_loader_train.get_batch()
inputs = to_cuda(to_variable(inputs), GPU)
inputs_prev = to_cuda(to_variable(inputs_prev), GPU)
outputs_ae_gt = to_cuda(
to_variable(outputs_ae_gt, requires_grad=False), GPU)
outputs_multiactions_gt = to_cuda(
to_variable(outputs_multiactions_gt, requires_grad=False), GPU)
outputs_flow_gt = to_cuda(
to_variable(outputs_flow_gt, requires_grad=False), GPU)
outputs_canny_gt = to_cuda(
to_variable(outputs_canny_gt, requires_grad=False), GPU)
outputs_flipped_gt = to_cuda(
to_variable(outputs_flipped_gt, requires_grad=False), GPU)
time_cbatch_end = time.time()
# predict and compute losses
time_fwbw_start = time.time()
optimizer_predictor.zero_grad()
(outputs_ae_pred, outputs_grids_pred, outputs_atts_pred,
outputs_multiactions_pred, outputs_flow_pred, outputs_canny_pred,
outputs_flipped_pred, emb) = predictor(inputs, inputs_prev)
# zero-grad some outputs where annotations are not available for specific examples
outputs_grids_gt = remove_unannotated_grids_gt(outputs_grids_pred,
outputs_grids_gt_orig,
grids_annotated)
outputs_grids_gt = to_cuda(
to_variable(outputs_grids_gt, requires_grad=False), GPU)
outputs_atts_gt = remove_unannotated_atts_gt(outputs_atts_pred,
outputs_atts_gt_orig,
atts_annotated)
outputs_atts_gt = to_cuda(
to_variable(outputs_atts_gt, requires_grad=False), GPU)
loss_ae = criterion_ae(outputs_ae_pred, outputs_ae_gt)
loss_grids = criterion_grids(outputs_grids_pred, outputs_grids_gt)
loss_atts = criterion_atts(outputs_atts_pred, outputs_atts_gt)
loss_multiactions = criterion_multiactions(outputs_multiactions_pred,
outputs_multiactions_gt)
loss_flow = criterion_flow(outputs_flow_pred, outputs_flow_gt)
loss_canny = criterion_canny(outputs_canny_pred, outputs_canny_gt)
loss_flipped = criterion_flipped(outputs_flipped_pred,
outputs_flipped_gt)
losses_grad_lst = [
loss.data.new().resize_as_(loss.data).fill_(w) for loss, w in zip([
loss_ae, loss_grids, loss_atts, loss_multiactions, loss_flow,
loss_canny, loss_flipped
], [
LOSS_AE_WEIGHTING, LOSS_GRIDS_WEIGHTING,
LOSS_ATTRIBUTES_WEIGHTING, LOSS_MULTIACTIONS_WEIGHTING,
LOSS_FLOW_WEIGHTING, LOSS_CANNY_WEIGHTING,
LOSS_FLIPPED_WEIGHTING
])
]
torch.autograd.backward([
loss_ae, loss_grids, loss_atts, loss_multiactions, loss_flow,
loss_canny, loss_flipped
], losses_grad_lst)
optimizer_predictor.step()
time_fwbw_end = time.time()
# add losses to history and output a message
loss_ae_value = to_numpy(loss_ae)[0]
loss_grids_value = to_numpy(loss_grids)[0]
loss_atts_value = to_numpy(loss_atts)[0]
loss_multiactions_value = to_numpy(loss_multiactions)[0]
loss_flow_value = to_numpy(loss_flow)[0]
loss_canny_value = to_numpy(loss_canny)[0]
loss_flipped_value = to_numpy(loss_flipped)[0]
history.add_value("loss-ae", "train", batch_idx, loss_ae_value)
history.add_value("loss-grids", "train", batch_idx, loss_grids_value)
history.add_value("loss-atts", "train", batch_idx, loss_atts_value)
history.add_value("loss-multiactions", "train", batch_idx,
loss_multiactions_value)
history.add_value("loss-flow", "train", batch_idx, loss_flow_value)
history.add_value("loss-canny", "train", batch_idx, loss_canny_value)
history.add_value("loss-flipped", "train", batch_idx,
loss_flipped_value)
print(
"[T] Batch %05d L[ae=%.4f, grids=%.4f, atts=%.4f, multiactions=%.4f, flow=%.4f, canny=%.4f, flipped=%.4f] T[cbatch=%.04fs, fwbw=%.04fs]"
% (batch_idx, loss_ae_value, loss_grids_value, loss_atts_value,
loss_multiactions_value, loss_flow_value, loss_canny_value,
loss_flipped_value, time_cbatch_end - time_cbatch_start,
time_fwbw_end - time_fwbw_start))
# genrate a debug image showing batch predictions and ground truths
if (batch_idx + 1) % 20 == 0:
debug_img = generate_debug_image(
inputs, inputs_prev, outputs_ae_gt, outputs_grids_gt_orig,
outputs_atts_gt_orig, outputs_multiactions_gt, outputs_flow_gt,
outputs_canny_gt, outputs_flipped_gt, outputs_ae_pred,
outputs_grids_pred, outputs_atts_pred,
outputs_multiactions_pred, outputs_flow_pred,
outputs_canny_pred, outputs_flipped_pred, grids_annotated,
atts_annotated)
misc.imsave(
"train_semisupervised_debug_img%s.jpg" %
("_withshortcuts" if args.withshortcuts else "", ), debug_img)
# run N validation batches
# TODO merge this with training stuff above (one function for both)
if (batch_idx + 1) % VAL_EVERY == 0:
predictor.eval()
loss_ae_total = 0
loss_grids_total = 0
loss_atts_total = 0
loss_multiactions_total = 0
loss_flow_total = 0
loss_canny_total = 0
loss_flipped_total = 0
for i in xrange(NB_VAL_BATCHES):
time_cbatch_start = time.time()
(inputs, inputs_prev), (
outputs_ae_gt, outputs_grids_gt_orig, outputs_atts_gt_orig,
outputs_multiactions_gt, outputs_flow_gt, outputs_canny_gt,
outputs_flipped_gt), (
grids_annotated,
atts_annotated) = batch_loader_val.get_batch()
inputs = to_cuda(to_variable(inputs, volatile=True), GPU)
inputs_prev = to_cuda(to_variable(inputs_prev, volatile=True),
GPU)
outputs_ae_gt = to_cuda(
to_variable(outputs_ae_gt, volatile=True), GPU)
outputs_multiactions_gt = to_cuda(
to_variable(outputs_multiactions_gt, volatile=True), GPU)
outputs_flow_gt = to_cuda(
to_variable(outputs_flow_gt, volatile=True), GPU)
outputs_canny_gt = to_cuda(
to_variable(outputs_canny_gt, volatile=True), GPU)
outputs_flipped_gt = to_cuda(
to_variable(outputs_flipped_gt, volatile=True), GPU)
time_cbatch_end = time.time()
time_fwbw_start = time.time()
(outputs_ae_pred, outputs_grids_pred, outputs_atts_pred,
outputs_multiactions_pred, outputs_flow_pred,
outputs_canny_pred, outputs_flipped_pred,
emb) = predictor(inputs, inputs_prev)
outputs_grids_gt = remove_unannotated_grids_gt(
outputs_grids_pred, outputs_grids_gt_orig, grids_annotated)
outputs_grids_gt = to_cuda(
to_variable(outputs_grids_gt, volatile=True), GPU)
outputs_atts_gt = remove_unannotated_atts_gt(
outputs_atts_pred, outputs_atts_gt_orig, atts_annotated)
outputs_atts_gt = to_cuda(
to_variable(outputs_atts_gt, volatile=True), GPU)
loss_ae = criterion_ae(outputs_ae_pred, outputs_ae_gt)
loss_grids = criterion_grids(outputs_grids_pred,
outputs_grids_gt)
loss_atts = criterion_atts(outputs_atts_pred, outputs_atts_gt)
loss_multiactions = criterion_multiactions(
outputs_multiactions_pred, outputs_multiactions_gt)
loss_flow = criterion_flow(outputs_flow_pred, outputs_flow_gt)
loss_canny = criterion_canny(outputs_canny_pred,
outputs_canny_gt)
loss_flipped = criterion_flipped(outputs_flipped_pred,
outputs_flipped_gt)
time_fwbw_end = time.time()
loss_ae_value = to_numpy(loss_ae)[0]
loss_grids_value = to_numpy(loss_grids)[0]
loss_atts_value = to_numpy(loss_atts)[0]
loss_multiactions_value = to_numpy(loss_multiactions)[0]
loss_flow_value = to_numpy(loss_flow)[0]
loss_canny_value = to_numpy(loss_canny)[0]
loss_flipped_value = to_numpy(loss_flipped)[0]
loss_ae_total += loss_ae_value
loss_grids_total += loss_grids_value
loss_atts_total += loss_atts_value
loss_multiactions_total += loss_multiactions_value
loss_flow_total += loss_flow_value
loss_canny_total += loss_canny_value
loss_flipped_total += loss_flipped_value
print(
"[V] Batch %05d L[ae=%.4f, grids=%.4f, atts=%.4f, multiactions=%.4f, flow=%.4f, canny=%.4f, flipped=%.4f] T[cbatch=%.04fs, fwbw=%.04fs]"
%
(batch_idx, loss_ae_value, loss_grids_value,
loss_atts_value, loss_multiactions_value, loss_flow_value,
loss_canny_value, loss_flipped_value, time_cbatch_end -
time_cbatch_start, time_fwbw_end - time_fwbw_start))
if i == 0:
debug_img = generate_debug_image(
inputs, inputs_prev, outputs_ae_gt,
outputs_grids_gt_orig, outputs_atts_gt_orig,
outputs_multiactions_gt, outputs_flow_gt,
outputs_canny_gt, outputs_flipped_gt, outputs_ae_pred,
outputs_grids_pred, outputs_atts_pred,
outputs_multiactions_pred, outputs_flow_pred,
outputs_canny_pred, outputs_flipped_pred,
grids_annotated, atts_annotated)
misc.imsave(
"train_semisupervised_debug_img_val%s.jpg" %
("_withshortcuts" if args.withshortcuts else "", ),
debug_img)
history.add_value("loss-ae", "val", batch_idx,
loss_ae_total / NB_VAL_BATCHES)
history.add_value("loss-grids", "val", batch_idx,
loss_grids_total / NB_VAL_BATCHES)
history.add_value("loss-atts", "val", batch_idx,
loss_atts_total / NB_VAL_BATCHES)
history.add_value("loss-multiactions", "val", batch_idx,
loss_multiactions_total / NB_VAL_BATCHES)
history.add_value("loss-flow", "val", batch_idx,
loss_flow_total / NB_VAL_BATCHES)
history.add_value("loss-canny", "val", batch_idx,
loss_canny_total / NB_VAL_BATCHES)
history.add_value("loss-flipped", "val", batch_idx,
loss_flipped_total / NB_VAL_BATCHES)
predictor.train()
# generate loss plot
if (batch_idx + 1) % PLOT_EVERY == 0:
loss_plotter.plot(history)
# every N batches, save a checkpoint
if (batch_idx + 1) % SAVE_EVERY == 0:
checkpoint_fp = "train_semisupervised_model%s.tar" % (
"_withshortcuts" if args.withshortcuts else "", )
torch.save(
{
"batch_idx": batch_idx,
"history": history.to_string(),
"predictor_state_dict": predictor.state_dict(),
}, checkpoint_fp)
# refresh automatically generated examples (autoencoder, canny edge stuff etc.)
if (batch_idx + 1) % 1000 == 0:
print("Refreshing autogen dataset...")
batch_loader_train.join()
examples_autogen_train = load_dataset_autogen(
val=False,
nb_load=NB_AUTOGEN_TRAIN,
not_in=examples_annotated_ids)
batch_loader_train = BatchLoader(examples_train,
examples_autogen_train,
augseq=augseq,
queue_size=15,
nb_workers=4,
threaded=False)
def generate_overview_image(current_state, last_state, \
action_up_down_bpe, action_left_right_bpe, \
memory, memory_val, \
ticks, last_train_tick, \
plans, plan_to_rewards_direct, plan_to_reward_indirect, \
plan_to_reward, plans_ranking, current_plan, best_plan_ae_decodings,
idr_v, idr_adv,
grids, args):
h, w = current_state.screenshot_rs.shape[0:2]
scr = np.copy(current_state.screenshot_rs)
scr = ia.imresize_single_image(scr, (h // 2, w // 2))
if best_plan_ae_decodings is not None:
ae_decodings = (to_numpy(best_plan_ae_decodings) * 255).astype(
np.uint8).transpose((0, 2, 3, 1))
ae_decodings = [
ia.imresize_single_image(ae_decodings[i, ...], (h // 4, w // 4))
for i in xrange(ae_decodings.shape[0])
]
ae_decodings = ia.draw_grid(ae_decodings, cols=5)
#ae_decodings = np.vstack([
# np.hstack(ae_decodings[0:5]),
# np.hstack(ae_decodings[5:10])
#])
else:
ae_decodings = np.zeros((1, 1, 3), dtype=np.uint8)
if grids is not None:
scr_rs = ia.imresize_single_image(scr, (h // 4, w // 4))
grids = (to_numpy(grids)[0] * 255).astype(np.uint8)
grids = [
ia.imresize_single_image(grids[i, ...][:, :, np.newaxis],
(h // 4, w // 4))
for i in xrange(grids.shape[0])
]
grids = [
util.draw_heatmap_overlay(
scr_rs,
np.squeeze(grid / 255).astype(np.float32)) for grid in grids
]
grids = ia.draw_grid(grids, cols=4)
else:
grids = np.zeros((1, 1, 3), dtype=np.uint8)
plans_text = []
if idr_v is not None and idr_adv is not None:
idr_v = to_numpy(idr_v[0])
idr_adv = to_numpy(idr_adv[0])
plans_text.append("V(s): %+07.2f" % (idr_v[0], ))
adv_texts = []
curr = []
for i, ma in enumerate(actionslib.ALL_MULTIACTIONS):
curr.append("A(%s%s): %+07.2f" %
(ma[0] if ma[0] != "~WS" else "_",
ma[1] if ma[1] != "~AD" else "_", idr_adv[i]))
if (i + 1) % 3 == 0 or (i + 1) == len(actionslib.ALL_MULTIACTIONS):
adv_texts.append(" ".join(curr))
curr = []
plans_text.extend(adv_texts)
if current_plan is not None:
plans_text.append("")
plans_text.append("Current Plan:")
actions_ud_text = []
actions_lr_text = []
for multiaction in current_plan:
actions_ud_text.append(
"%s" % (multiaction[0] if multiaction[0] != "~WS" else "_", ))
actions_lr_text.append(
"%s" % (multiaction[1] if multiaction[1] != "~AD" else "_", ))
plans_text.extend(
[" ".join(actions_ud_text), " ".join(actions_lr_text)])
plans_text.append("")
plans_text.append("Best Plans:")
if plan_to_rewards_direct is not None:
for plan_idx in plans_ranking[::-1][0:5]:
plan = plans[plan_idx]
rewards_direct = plan_to_rewards_direct[plan_idx]
reward_indirect = plan_to_reward_indirect[plan_idx]
reward = plan_to_reward[plan_idx]
actions_ud_text = []
actions_lr_text = []
rewards_text = []
for multiaction in plan:
actions_ud_text.append(
"%s" %
(multiaction[0] if multiaction[0] != "~WS" else "_", ))
actions_lr_text.append(
"%s" %
(multiaction[1] if multiaction[1] != "~AD" else "_", ))
for rewards_t in rewards_direct:
rewards_text.append("%+04.1f" % (rewards_t, ))
rewards_text.append("| %+07.2f (V(s')=%+07.2f)" %
(reward, reward_indirect))
plans_text.extend([
"", " ".join(actions_ud_text), " ".join(actions_lr_text),
" ".join(rewards_text)
])
plans_text = "\n".join(plans_text)
stats_texts = [
"u/d bpe: %s" % (action_up_down_bpe.rjust(5)),
" l/r bpe: %s" % (action_left_right_bpe.rjust(5)),
"u/d ape: %s %s" %
(current_state.action_up_down.rjust(5),
"[C]" if action_up_down_bpe != current_state.action_up_down else ""),
" l/r ape: %s %s" %
(current_state.action_left_right.rjust(5), "[C]"
if action_left_right_bpe != current_state.action_left_right else ""),
"speed: %03d" % (current_state.speed, )
if current_state.speed is not None else "speed: None",
"is_reverse: yes" if current_state.is_reverse else "is_reverse: no",
"is_damage_shown: yes" if current_state.is_damage_shown else
"is_damage_shown: no", "is_offence_shown: yes"
if current_state.is_offence_shown else "is_offence_shown: no",
"steering wheel: %05.2f (%05.2f)" %
(current_state.steering_wheel_cnn,
current_state.steering_wheel_raw_cnn),
"reward for last state: %05.2f" % (last_state.reward, )
if last_state is not None else "reward for last state: None",
"p_explore: %.2f%s" %
(current_state.p_explore if args.p_explore is None else args.p_explore,
"" if args.p_explore is None else " (constant)"),
"memory size (train/val): %06d / %06d" %
(memory.size, memory_val.size),
"ticks: %06d" % (ticks, ),
"last train: %06d" % (last_train_tick, )
]
stats_text = "\n".join(stats_texts)
all_texts = plans_text + "\n\n\n" + stats_text
result = np.zeros((720, 590, 3), dtype=np.uint8)
util.draw_image(result, x=0, y=0, other_img=scr, copy=False)
util.draw_image(result,
x=0,
y=scr.shape[0] + 10,
other_img=ae_decodings,
copy=False)
util.draw_image(result,
x=0,
y=scr.shape[0] + 10 + ae_decodings.shape[0] + 10,
other_img=grids,
copy=False)
result = util.draw_text(result,
x=0,
y=scr.shape[0] + 10 + ae_decodings.shape[0] + 10 +
grids.shape[0] + 10,
size=8,
text=all_texts,
color=[255, 255, 255])
return result
def generate_training_debug_image(inputs_supervised, inputs_supervised_prev, \
outputs_dr_preds, outputs_dr_gt, \
outputs_idr_preds, outputs_idr_gt, \
outputs_successor_preds, outputs_successor_gt, \
outputs_ae_preds, outputs_ae_gt, \
outputs_dr_successors_preds, outputs_dr_successors_gt, \
outputs_idr_successors_preds, outputs_idr_successors_gt,
multiactions):
imgs_in = to_numpy(inputs_supervised)[0]
imgs_in = np.clip(imgs_in * 255, 0, 255).astype(np.uint8).transpose(
(1, 2, 0))
imgs_in_prev = to_numpy(inputs_supervised_prev)[0]
imgs_in_prev = np.clip(imgs_in_prev * 255, 0,
255).astype(np.uint8).transpose((1, 2, 0))
h, w = imgs_in.shape[0:2]
imgs_in = np.vstack([
np.hstack([
downscale(imgs_in[..., 0:3]),
downscale(to_rgb(imgs_in_prev[..., 0]))
]),
#np.hstack([downscale(to_rgb(imgs_in_prev[..., 1])), downscale(to_rgb(imgs_in_prev[..., 2]))])
np.hstack([
downscale(to_rgb(imgs_in_prev[..., 1])),
np.zeros_like(imgs_in[..., 0:3])
])
])
h_imgs = imgs_in.shape[0]
ae_gt = np.clip(to_numpy(outputs_ae_gt)[0] * 255, 0,
255).astype(np.uint8).transpose((1, 2, 0))
ae_preds = np.clip(to_numpy(outputs_ae_preds)[0] * 255, 0,
255).astype(np.uint8).transpose((1, 2, 0))
"""
imgs_ae = np.vstack([
downscale(ae_preds[..., 0:3]),
downscale(to_rgb(ae_preds[..., 3])),
downscale(to_rgb(ae_preds[..., 4])),
downscale(to_rgb(ae_preds[..., 5]))
])
"""
imgs_ae = np.hstack([downscale(ae_gt), downscale(ae_preds)])
h_ae = imgs_ae.shape[0]
outputs_successor_dr_grid = draw_successor_dr_grid(
to_numpy(F.softmax(outputs_dr_successors_preds[:, 0, :])),
to_numpy(outputs_dr_successors_gt[:, 0]),
upscale_factor=(2, 4))
outputs_dr_preds = to_numpy(F.softmax(outputs_dr_preds))[0]
outputs_dr_gt = to_numpy(outputs_dr_gt)[0]
grid_preds = output_grid_to_image(outputs_dr_preds[np.newaxis, :],
upscale_factor=(2, 4))
grid_gt = output_grid_to_image(outputs_dr_gt[np.newaxis, :],
upscale_factor=(2, 4))
imgs_dr = np.hstack([
grid_gt,
np.zeros((grid_gt.shape[0], 4, 3), dtype=np.uint8), grid_preds,
np.zeros((grid_gt.shape[0], 8, 3), dtype=np.uint8),
outputs_successor_dr_grid
])
successor_multiactions_str = " ".join([
"%s%s" %
(ma[0] if ma[0] != "~WS" else "_", ma[1] if ma[1] != "~AD" else "_")
for ma in multiactions[0]
])
imgs_dr = np.pad(imgs_dr, ((30, 0), (0, 300), (0, 0)),
mode="constant",
constant_values=0)
imgs_dr = util.draw_text(
imgs_dr,
x=0,
y=0,
text="DR curr bins gt:%s, pred:%s | successor preds\nsucc. mas: %s" %
(str(np.argmax(outputs_dr_gt)), str(
np.argmax(outputs_dr_preds)), successor_multiactions_str),
size=9)
h_dr = imgs_dr.shape[0]
outputs_idr_preds = np.squeeze(to_numpy(outputs_idr_preds)[0])
outputs_idr_gt = np.squeeze(to_numpy(outputs_idr_gt)[0])
idr_text = [
"[IndirectReward A0]",
" gt: %.2f" % (outputs_idr_gt[..., 0], ),
" pr: %.2f" % (outputs_idr_preds[..., 0], ), "[IndirectReward A1]",
" gt: %.2f" % (outputs_idr_gt[..., 1], ),
" pr: %.2f" % (outputs_idr_preds[..., 1], ), "[IndirectReward A2]",
" gt: %.2f" % (outputs_idr_gt[..., 2], ),
" pr: %.2f" % (outputs_idr_preds[..., 2], )
]
idr_text = "\n".join(idr_text)
outputs_successor_preds = np.squeeze(
to_numpy(outputs_successor_preds)[:, 0, :])
outputs_successor_gt = np.squeeze(to_numpy(outputs_successor_gt)[:, 0, :])
distances = np.average((outputs_successor_preds - outputs_successor_gt)**2,
axis=1)
successors_text = [
"[Successors]",
" Distances:",
" " + " ".join(["%02.2f" % (d, ) for d in distances]),
" T=0 gt/pred:",
" " + " ".join(
["%+02.2f" % (val, ) for val in outputs_successor_gt[0, 0:25]]),
" " + " ".join(
["%+02.2f" % (val, ) for val in outputs_successor_preds[0, 0:25]]),
" T=1 gt/pred:",
" " + " ".join(
["%+02.2f" % (val, ) for val in outputs_successor_gt[1, 0:25]]),
" " + " ".join(
["%+02.2f" % (val, ) for val in outputs_successor_preds[1, 0:25]]),
" T=2 gt/pred:",
" " + " ".join(
["%+02.2f" % (val, ) for val in outputs_successor_gt[2, 0:25]]),
" " + " ".join(
["%+02.2f" % (val, ) for val in outputs_successor_preds[2, 0:25]]),
]
successors_text = "\n".join(successors_text)
outputs_dr_successors_preds = np.squeeze(
to_numpy(outputs_dr_successors_preds)[:, 0, :])
outputs_dr_successors_gt = np.squeeze(
to_numpy(outputs_dr_successors_gt)[:, 0, :])
bins_dr_successors_preds = np.argmax(outputs_dr_successors_preds, axis=1)
bins_dr_successors_gt = np.argmax(outputs_dr_successors_gt, axis=1)
successors_dr_text = [
"[Direct rewards bins of successors]",
" gt: " + " ".join(["%d" % (b, ) for b in bins_dr_successors_gt]),
" pred: " + " ".join(["%d" % (b, ) for b in bins_dr_successors_preds])
]
successors_dr_text = "\n".join(successors_dr_text)
outputs_idr_successors_preds = np.squeeze(
to_numpy(outputs_idr_successors_preds)[:, 0, :])
outputs_idr_successors_gt = np.squeeze(
to_numpy(outputs_idr_successors_gt)[:, 0, :])
successors_idr_text = [
"[Indirect rewards of successors A0]", " gt: " + " ".join(
["%+03.2f" % (v, ) for v in outputs_idr_successors_gt[..., 0]]),
" pred: " + " ".join(
["%+03.2f" % (v, ) for v in outputs_idr_successors_preds[..., 0]]),
"[Indirect rewards of successors A1]", " gt: " + " ".join(
["%+03.2f" % (v, ) for v in outputs_idr_successors_gt[..., 1]]),
" pred: " + " ".join(
["%+03.2f" % (v, ) for v in outputs_idr_successors_preds[..., 1]]),
"[Indirect rewards of successors A2]", " gt: " + " ".join(
["%+03.2f" % (v, ) for v in outputs_idr_successors_gt[..., 2]]),
" pred: " + " ".join(
["%+03.2f" % (v, ) for v in outputs_idr_successors_preds[..., 2]])
]
successors_idr_text = "\n".join(successors_idr_text)
result = np.zeros((950, 320, 3), dtype=np.uint8)
spacing = 4
util.draw_image(result, x=0, y=0, other_img=imgs_in, copy=False)
util.draw_image(result,
x=0,
y=h_imgs + spacing,
other_img=imgs_ae,
copy=False)
util.draw_image(result,
x=0,
y=h_imgs + spacing + h_ae + spacing,
other_img=imgs_dr,
copy=False)
result = util.draw_text(
result,
x=0,
y=h_imgs + spacing + h_ae + spacing + h_dr + spacing,
text=idr_text + "\n" + successors_text + "\n" + successors_dr_text +
"\n" + successors_idr_text,
size=9)
return result
