def METHOD_NAME(unfolded_inp, im_size):
    out_vals = np.argmax(unfolded_inp[:, :, :], axis=0)
    colors = np.zeros((im_size, im_size, 3), dtype=np.uint8)
    for ii, row in enumerate(out_vals):
        for j, val in enumerate(row):
            if val != 3:
                if val == 1:
                    val = 2
                elif val == 2:
                    val = 1
                colors[ii, j, val] = 255
    return colors