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