def main():
    tf.disable_eager_execution()

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    arguments = parser.parse_args()

    full_name = arguments.experiment_name.split('/')

    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''

    codebook, dataset = factory.build_codebook_from_name(experiment_name,
                                                         experiment_group,
                                                         return_dataset=True)

    workspace_path = os.environ.get('AE_WORKSPACE_PATH')
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    ckpt_dir = u.get_checkpoint_dir(log_dir)

    train_cfg_file_path = u.get_train_config_exp_file_path(
        log_dir, experiment_name)
    train_args = configparser.ConfigParser()
    train_args.read(train_cfg_file_path)

    width = 960
    height = 720
    videoStream = WebcamVideoStream(0, width, height).start()

    gpu_options = tf.GPUOptions(allow_growth=True,
                                per_process_gpu_memory_fraction=0.9)
    config = tf.ConfigProto(gpu_options=gpu_options)
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        factory.restore_checkpoint(sess, tf.train.Saver(), ckpt_dir)

        while videoStream.isActive():
            image = videoStream.read()
            if image is None or not image.any():
                print("Failed to capture webcam image")
                exit(-1)

            # try your detector here:
            # bb_xywh = detector.detect(image)
            # image_crop = dataset.extract_square_patch(image, bb_xywh, train_args.getfloat('Dataset','PAD_FACTOR'))
            # Rs, ts = codebook.auto_pose6d(sess, image_crop, bb_xywh, K_test, 1, train_args)

            img = cv2.resize(image, (128, 128))

            R = codebook.nearest_rotation(sess, img)
            pred_view = dataset.render_rot(R, downSample=1)
            print(R)
            cv2.imshow('resized webcam input', img)
            cv2.imshow('pred view rendered', pred_view)
            cv2.waitKey(1)
예제 #2
0
    def initialize_pose_est_model(self):
        import tensorflow as tf

        rospy.loginfo("Loading AAE model")
        os.environ["CUDA_VISIBLE_DEVICES"] = self.params["pe_gpu_id"]
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
        tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
        self.workspace_path = os.environ.get('AE_WORKSPACE_PATH')
        if self.workspace_path == None:
            print 'Please define a workspace path:\n'
            print 'export AE_WORKSPACE_PATH=/path/to/workspace\n'
            exit(-1)
        # load all code books
        full_name = self.params["pe_experiment_name"].split('/')    
        experiment_name = full_name.pop()
        experiment_group = full_name.pop() if len(full_name) > 0 else ''
        self.ply_paths = glob.glob(self.params["model_dir"] + '/ply/*.ply')
        self.ply_paths.sort()
        self.ply_centered_paths = glob.glob(self.params["model_dir"] + '/ply_centered/*.ply')
        self.ply_centered_paths.sort()
        self.codebook, self.dataset = factory.build_codebook_from_name(experiment_name, experiment_group, return_dataset = True, joint=True)

        self.dims = []
        self.centroids = []
        self.cloud_objs = []
        for ply_centered in self.ply_centered_paths:
            cloud = o3d.io.read_point_cloud(ply_centered)
            self.dims.append(cloud.get_max_bound())
        for ply in self.ply_paths:
            cloud = o3d.io.read_point_cloud(ply)
            centroid = cloud.get_center()
            self.centroids.append(centroid)
            self.cloud_objs.append(cloud)

        gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction = 0.5)
        config = tf.ConfigProto(gpu_options=gpu_options)

        log_dir = u.get_log_dir(self.workspace_path, experiment_name, experiment_group)
        train_cfg_file_path = u.get_train_config_exp_file_path(log_dir, experiment_name)
        self.train_args = configparser.ConfigParser(inline_comment_prefixes="#")
        self.train_args.read(train_cfg_file_path)
        test_configpath = os.path.join(self.workspace_path, 'cfg_eval', self.params["pe_test_config"])
        test_args = configparser.ConfigParser()
        test_args.read(test_configpath)

        self.sess = tf.Session(config=config)
        saver = tf.train.Saver()
        checkpoint_file = u.get_checkpoint_basefilename(log_dir, False, latest=self.train_args.getint('Training', 'NUM_ITER'), joint=True)
        saver.restore(self.sess, checkpoint_file)
def main():
    import argparse
    import configparser
    from . import eval_utils
    parser = argparse.ArgumentParser()
    
    parser.add_argument('experiment_name')
    parser.add_argument('evaluation_name')
    parser.add_argument('--eval_cfg', default='eval.cfg', required=False)
    arguments = parser.parse_args()
    full_name = arguments.experiment_name.split('/')
    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    evaluation_name = arguments.evaluation_name
    eval_cfg = arguments.eval_cfg

    workspace_path = os.environ.get('AE_WORKSPACE_PATH')
    train_cfg_file_path = get_config_file_path(workspace_path, experiment_name, experiment_group)
    eval_cfg_file_path = get_eval_config_file_path(workspace_path, eval_cfg=eval_cfg)

    train_args = configparser.ConfigParser(inline_comment_prefixes="#")
    eval_args = configparser.ConfigParser(inline_comment_prefixes="#")
    train_args.read(train_cfg_file_path)
    eval_args.read(eval_cfg_file_path)

    dataset_name = eval_args.get('DATA','DATASET')
    scenes = eval(eval_args.get('DATA','SCENES')) if len(eval(eval_args.get('DATA','SCENES'))) > 0 else eval_utils.get_all_scenes_for_obj(eval_args)
    cam_type = eval_args.get('DATA','cam_type')
    data = dataset_name + '_' + cam_type if len(cam_type) > 0 else dataset_name

    log_dir = get_log_dir(workspace_path, experiment_name, experiment_group)
    eval_dir = get_eval_dir(log_dir, evaluation_name, data)
    
    plot_R_err_hist_vis(eval_args, eval_dir, scenes,bins=15)
    plot_t_err_hist_vis(eval_args, eval_dir, scenes,bins=15)
    plt.show()
예제 #4
0
    def __init__(self, test_configpath):

        test_args = configparser.ConfigParser()
        test_args.read(test_configpath)

        workspace_path = os.environ.get('AE_WORKSPACE_PATH')

        if workspace_path == None:
            print('Please define a workspace path:')
            print('export AE_WORKSPACE_PATH=/path/to/workspace')
            exit(-1)

        self._camPose = test_args.getboolean('CAMERA', 'camPose')
        self._camK = np.array(eval(test_args.get('CAMERA',
                                                 'K_test'))).reshape(3, 3)
        self._width = test_args.getint('CAMERA', 'width')
        self._height = test_args.getint('CAMERA', 'height')

        self._upright = test_args.getboolean('AAE', 'upright')
        self.all_experiments = eval(test_args.get('AAE', 'experiments'))

        self.class_names = eval(test_args.get('DETECTOR', 'class_names'))
        self.det_threshold = eval(test_args.get('DETECTOR', 'det_threshold'))
        self.icp = test_args.getboolean('ICP', 'icp')

        if self.icp:
            self._depth_scale = test_args.getfloat('DATA', 'depth_scale')

        self.all_codebooks = []
        self.all_train_args = []
        self.pad_factors = []
        self.patch_sizes = []

        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.allow_growth = True
        config.gpu_options.per_process_gpu_memory_fraction = test_args.getfloat(
            'MODEL', 'gpu_memory_fraction')

        self.sess = tf.Session(config=config)
        set_session(self.sess)
        self.detector = load_model(
            str(test_args.get('DETECTOR', 'detector_model_path')),
            backbone_name=test_args.get('DETECTOR', 'backbone'))
        #detector = self._load_model_with_nms(test_args)

        for i, experiment in enumerate(self.all_experiments):
            full_name = experiment.split('/')
            experiment_name = full_name.pop()
            experiment_group = full_name.pop() if len(full_name) > 0 else ''
            log_dir = utils.get_log_dir(workspace_path, experiment_name,
                                        experiment_group)
            ckpt_dir = utils.get_checkpoint_dir(log_dir)
            train_cfg_file_path = utils.get_train_config_exp_file_path(
                log_dir, experiment_name)
            print(train_cfg_file_path)
            # train_cfg_file_path = utils.get_config_file_path(workspace_path, experiment_name, experiment_group)
            train_args = configparser.ConfigParser()
            train_args.read(train_cfg_file_path)
            self.all_train_args.append(train_args)
            self.pad_factors.append(
                train_args.getfloat('Dataset', 'PAD_FACTOR'))
            self.patch_sizes.append(
                (train_args.getint('Dataset',
                                   'W'), train_args.getint('Dataset', 'H')))

            self.all_codebooks.append(
                factory.build_codebook_from_name(experiment_name,
                                                 experiment_group,
                                                 return_dataset=False))
            saver = tf.train.Saver(var_list=tf.get_collection(
                tf.GraphKeys.GLOBAL_VARIABLES, scope=experiment_name))
            factory.restore_checkpoint(self.sess, saver, ckpt_dir)

            # if self.icp:
            #     assert len(self.all_experiments) == 1, 'icp currently only works for one object'
            #     # currently works only for one object
            #     from auto_pose.icp import icp
            #     self.icp_handle = icp.ICP(train_args)
        if test_args.getboolean('ICP', 'icp'):
            from auto_pose.icp import icp
            self.icp_handle = icp.ICP(test_args, self.all_train_args)
예제 #5
0
parser = argparse.ArgumentParser()
parser.add_argument("experiment_name")
arguments = parser.parse_args()

full_name = arguments.experiment_name.split('/')

experiment_name = full_name.pop()
experiment_group = full_name.pop() if len(full_name) > 0 else ''

codebook, dataset = factory.build_codebook_from_name(experiment_name,
                                                     experiment_group,
                                                     return_dataset=True)

workspace_path = os.environ.get('AE_WORKSPACE_PATH')
log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
ckpt_dir = u.get_checkpoint_dir(log_dir)

train_cfg_file_path = u.get_train_config_exp_file_path(log_dir,
                                                       experiment_name)
train_args = configparser.ConfigParser()
train_args.read(train_cfg_file_path)

width = 960
height = 720
videoStream = WebcamVideoStream(0, width, height).start()

with tf.Session() as sess:
    factory.restore_checkpoint(sess, tf.train.Saver(), ckpt_dir)

    while videoStream.isActive():
def detection(detection_graph, category_index, score, expand):
    print("> Building Graph")
    print(category_index)
    # Session Config: allow seperate GPU/CPU adressing and limit memory allocation
    config = tf.ConfigProto(allow_soft_placement=True,
                            log_device_placement=log_device)
    config.gpu_options.allow_growth = allow_memory_growth
    config.gpu_options.per_process_gpu_memory_fraction = 0.4  ###Jetson only
    cur_frames = 0
    with detection_graph.as_default():
        #run_meta = tf.RunMetadata()
        with tf.Session(graph=detection_graph, config=config) as sess:
            # Define Input and Ouput tensors
            image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
            detection_boxes = detection_graph.get_tensor_by_name(
                'detection_boxes:0')
            detection_scores = detection_graph.get_tensor_by_name(
                'detection_scores:0')
            detection_classes = detection_graph.get_tensor_by_name(
                'detection_classes:0')
            num_detections = detection_graph.get_tensor_by_name(
                'num_detections:0')
            if split_model:
                score_out = detection_graph.get_tensor_by_name(
                    'Postprocessor/convert_scores:0')
                expand_out = detection_graph.get_tensor_by_name(
                    'Postprocessor/ExpandDims_1:0')
                score_in = detection_graph.get_tensor_by_name(
                    'Postprocessor/convert_scores_1:0')
                expand_in = detection_graph.get_tensor_by_name(
                    'Postprocessor/ExpandDims_1_1:0')
                # Threading
                gpu_worker = SessionWorker("GPU", detection_graph, config)
                cpu_worker = SessionWorker("CPU", detection_graph, config)
                gpu_opts = [score_out, expand_out]
                cpu_opts = [
                    detection_boxes, detection_scores, detection_classes,
                    num_detections
                ]
                gpu_counter = 0
                cpu_counter = 0

            for i, experiment_name in enumerate(arguments.experiment_names):

                full_name = experiment_name.split('/')
                experiment_name = full_name.pop()
                experiment_group = full_name.pop(
                ) if len(full_name) > 0 else ''

                train_cfg_file_path = utils.get_config_file_path(
                    workspace_path, experiment_name, experiment_group)
                train_args = configparser.ConfigParser()
                train_args.read(train_cfg_file_path)
                h_train, w_train, c = train_args.getint(
                    'Dataset', 'H'), train_args.getint('Dataset',
                                                       'W'), train_args.getint(
                                                           'Dataset', 'C')
                model_paths.append(train_args.get('Paths', 'MODEL_PATH'))
                all_train_args.append(train_args)

                log_dir = utils.get_log_dir(workspace_path, experiment_name,
                                            experiment_group)
                ckpt_dir = utils.get_checkpoint_dir(log_dir)

                all_codebooks.append(
                    factory.build_codebook_from_name(experiment_name,
                                                     experiment_group,
                                                     return_dataset=False))
                factory.restore_checkpoint(
                    sess,
                    tf.train.Saver(var_list=tf.get_collection(
                        tf.GraphKeys.GLOBAL_VARIABLES, scope=experiment_name)),
                    ckpt_dir)

            #opts = tf.profiler.ProfileOptionBuilder.float_operation()
            #flops = tf.profiler.profile(sess.graph, run_meta=run_meta, cmd='op', options=opts)
            #exit()
            # i_class_mapping = {v: k for k, v in class_i_mapping.iteritems()}
            renderer = meshrenderer_phong.Renderer(model_paths, 1)

            # Start Video Stream and FPS calculation
            fps = FPS2(fps_interval).start()
            video_stream = WebcamVideoStream(video_input, width,
                                             height).start()
            cur_frames = 0
            print("> Press 'q' to Exit, 'a' to start auto_pose")
            print('> Starting Detection')
            while video_stream.isActive():
                # actual Detection
                if split_model:
                    # split model in seperate gpu and cpu session threads
                    if gpu_worker.is_sess_empty():
                        # read video frame, expand dimensions and convert to rgb
                        image = video_stream.read()
                        image_expanded = np.expand_dims(cv2.cvtColor(
                            image, cv2.COLOR_BGR2RGB),
                                                        axis=0)
                        # put new queue
                        gpu_feeds = {image_tensor: image_expanded}
                        if visualize:
                            gpu_extras = image  # for visualization frame
                        else:
                            gpu_extras = None
                        gpu_worker.put_sess_queue(gpu_opts, gpu_feeds,
                                                  gpu_extras)

                    g = gpu_worker.get_result_queue()
                    if g is None:
                        # gpu thread has no output queue. ok skip, let's check cpu thread.
                        gpu_counter += 1
                    else:
                        # gpu thread has output queue.
                        gpu_counter = 0
                        score, expand, image = g["results"][0], g["results"][
                            1], g["extras"]

                        if cpu_worker.is_sess_empty():
                            # When cpu thread has no next queue, put new queue.
                            # else, drop gpu queue.
                            cpu_feeds = {score_in: score, expand_in: expand}
                            cpu_extras = image
                            cpu_worker.put_sess_queue(cpu_opts, cpu_feeds,
                                                      cpu_extras)

                    c = cpu_worker.get_result_queue()
                    if c is None:
                        # cpu thread has no output queue. ok, nothing to do. continue
                        cpu_counter += 1
                        time.sleep(0.005)
                        continue  # If CPU RESULT has not been set yet, no fps update
                    else:
                        cpu_counter = 0
                        boxes, scores, classes, num, image = c["results"][
                            0], c["results"][1], c["results"][2], c["results"][
                                3], c["extras"]
                else:
                    # default session
                    image = video_stream.read()
                    image_expanded = np.expand_dims(cv2.cvtColor(
                        image, cv2.COLOR_BGR2RGB),
                                                    axis=0)
                    boxes, scores, classes, num = sess.run(
                        [
                            detection_boxes, detection_scores,
                            detection_classes, num_detections
                        ],
                        feed_dict={image_tensor: image_expanded})

                # Visualization of the results of a detection.

                H, W = image.shape[:2]

                img_crops = []
                det_bbs = []
                det_classes = []
                det_scores = []

                det_aae_bbs = []
                det_aae_objects_k = []
                #print vis_img.shape
                boxes = np.squeeze(boxes)
                scores = np.squeeze(scores)
                classes = np.squeeze(classes).astype(np.int32)

                highest_class_score = {clas: 0.0 for clas in classes}
                for box, score, clas in zip(boxes, scores, classes):
                    if score > det_th and score > highest_class_score[clas]:

                        highest_class_score[clas] = score
                        ymin, xmin, ymax, xmax = (np.array(box) * np.array(
                            [height, width, height, width])).astype(np.int32)

                        h, w = (ymax - ymin, xmax - xmin)
                        det_bbs.append([xmin, ymin, w, h])
                        det_classes.append(clas)
                        det_scores.append(score)
                        if clas in clas_k_map:
                            det_aae_bbs.append([xmin, ymin, w, h])

                            det_aae_objects_k.append(clas_k_map[clas])

                            size = int(
                                np.maximum(h, w) *
                                train_args.getfloat('Dataset', 'PAD_FACTOR'))
                            cx = xmin + (xmax - xmin) / 2
                            cy = ymin + (ymax - ymin) / 2

                            left = np.maximum(cx - size / 2, 0)
                            top = np.maximum(cy - size / 2, 0)

                            img_crop = image[top:cy + size / 2,
                                             left:cx + size / 2]
                            img_crop = cv2.resize(img_crop, (h_train, w_train))

                            img_crop = img_crop / 255.
                            img_crops.append(img_crop)

                if len(det_aae_bbs) > 0:

                    Rs = []
                    ts = []
                    for k, bb, img_crop in zip(det_aae_objects_k, det_aae_bbs,
                                               img_crops):
                        R, t = all_codebooks[k].auto_pose6d(sess,
                                                            img_crop,
                                                            bb,
                                                            K_test,
                                                            1,
                                                            all_train_args[k],
                                                            upright=False)
                        Rs.append(R.squeeze())
                        ts.append(t.squeeze())

                    Rs = np.array(Rs)
                    ts = np.array(ts)

                    bgr_y, _, _ = renderer.render_many(
                        obj_ids=np.array(det_aae_objects_k).astype(np.int32),
                        W=width / arguments.down,
                        H=height / arguments.down,
                        K=K_down,
                        Rs=Rs,
                        ts=ts,
                        near=1.,
                        far=10000.,
                        random_light=False,
                        # calc_bbs=False,
                        # depth=False
                    )

                    bgr_y = cv2.resize(bgr_y, (width, height))

                    g_y = np.zeros_like(bgr_y)
                    g_y[:, :, 1] = bgr_y[:, :, 1]
                    im_bg = cv2.bitwise_and(image,
                                            image,
                                            mask=(g_y[:, :, 1] == 0).astype(
                                                np.uint8))
                    image = cv2.addWeighted(im_bg, 1, g_y, 1, 0)

                for bb, score, clas in zip(det_bbs, det_scores, det_classes):
                    xmin, ymin, xmax, ymax = bb[0], bb[
                        1], bb[2] + bb[0], bb[1] + bb[3]
                    cv2.putText(
                        image,
                        '%s : %1.3f' % (category_index[clas]['name'], score),
                        (xmin, ymax + 20), cv2.FONT_ITALIC, .5,
                        color_dict[clas - 1], 2)
                    cv2.rectangle(image, (xmin, ymin), (xmax, ymax),
                                  color_dict[clas - 1], 2)

                if vis_text:
                    cv2.putText(image, "fps: {}".format(fps.fps_local()),
                                (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                                (77, 255, 9), 2)
                cv2.imshow('object_detection', image)
                # Exit Option
                key = cv2.waitKey(1)
                if key == ord('q'):
                    break

                fps.update()

    # End everything
    if split_model:
        gpu_worker.stop()
        cpu_worker.stop()
    fps.stop()
    video_stream.stop()
    cv2.destroyAllWindows()
    print('> [INFO] elapsed time (total): {:.2f}'.format(fps.elapsed()))
    print('> [INFO] approx. FPS: {:.2f}'.format(fps.fps()))
def main():
    workspace_path = os.environ.get('AE_WORKSPACE_PATH')

    if workspace_path == None:
        print('Please define a workspace path:\n')
        print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
        exit(-1)

    gentle_stop = np.array((1,), dtype=np.bool)
    gentle_stop[0] = False
    def on_ctrl_c(signal, frame):
        gentle_stop[0] = True
    signal.signal(signal.SIGINT, on_ctrl_c)

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    parser.add_argument("-d", action='store_true', default=False)
    parser.add_argument("-gen", action='store_true', default=False)
    parser.add_argument("-vis_emb", action='store_true', default=False)
    parser.add_argument('--at_step', default=None,  type=int, required=False)


    arguments = parser.parse_args()

    full_name = arguments.experiment_name.split('/')
    
    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    
    debug_mode = arguments.d
    generate_data = arguments.gen
    at_step = arguments.at_step

    cfg_file_path = u.get_config_file_path(workspace_path, experiment_name, experiment_group)
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    checkpoint_file = u.get_checkpoint_basefilename(log_dir)
    ckpt_dir = u.get_checkpoint_dir(log_dir)
    train_fig_dir = u.get_train_fig_dir(log_dir)
    dataset_path = u.get_dataset_path(workspace_path)
    
    if not os.path.exists(cfg_file_path):
        print('Could not find config file:\n')
        print('{}\n'.format(cfg_file_path))
        exit(-1)
        

    args = configparser.ConfigParser()
    args.read(cfg_file_path)

    num_iter = args.getint('Training', 'NUM_ITER') if not debug_mode else np.iinfo(np.int32).max
    save_interval = args.getint('Training', 'SAVE_INTERVAL')
    num_gpus = 1
    model_type = args.get('Dataset', 'MODEL')

    with tf.variable_scope(experiment_name, reuse=tf.AUTO_REUSE):
        
        dataset = factory.build_dataset(dataset_path, args)
        multi_queue = factory.build_multi_queue(dataset, args)
        dev_splits = np.array_split(np.arange(24), num_gpus)

        iterator = multi_queue.create_iterator(dataset_path, args)
        all_object_views = tf.concat([inp[0] for inp in multi_queue.next_element],0)

        bs = multi_queue._batch_size
        encoding_splits = []
        for dev in range(num_gpus):
            with tf.device('/device:GPU:%s' % dev):   
                encoder = factory.build_encoder(all_object_views[dev_splits[dev][0]*bs:(dev_splits[dev][-1]+1)*bs], args, is_training=False)
                encoding_splits.append(tf.split(encoder.z, len(dev_splits[dev]),0))

    with tf.variable_scope(experiment_name):
        decoders = []
        for dev in range(num_gpus):     
            with tf.device('/device:GPU:%s' % dev):  
                for j,i in enumerate(dev_splits[dev]):
                    decoders.append(factory.build_decoder(multi_queue.next_element[i], encoding_splits[dev][j], args, is_training=False, idx=i))

        ae = factory.build_ae(encoder, decoders, args)
        codebook = factory.build_codebook(encoder, dataset, args)
        train_op = factory.build_train_op(ae, args)
        saver = tf.train.Saver(save_relative_paths=True)

    dataset.load_bg_images(dataset_path)
    multi_queue.create_tfrecord_training_images(dataset_path, args)

    widgets = ['Training: ', progressbar.Percentage(),
         ' ', progressbar.Bar(),
         ' ', progressbar.Counter(), ' / %s' % num_iter,
         ' ', progressbar.ETA(), ' ']
    bar = progressbar.ProgressBar(maxval=num_iter,widgets=widgets)


    gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction = 0.9)
    config = tf.ConfigProto(gpu_options=gpu_options,log_device_placement=True,allow_soft_placement=True)

    with tf.Session(config=config) as sess:

        sess.run(multi_queue.bg_img_init.initializer)
        sess.run(iterator.initializer)


        chkpt = tf.train.get_checkpoint_state(ckpt_dir)
        if chkpt and chkpt.model_checkpoint_path:
            if at_step is None:
                checkpoint_file_basename = u.get_checkpoint_basefilename(log_dir,latest=args.getint('Training', 'NUM_ITER'))
            else:
                checkpoint_file_basename = u.get_checkpoint_basefilename(log_dir,latest=at_step)
            print('loading ', checkpoint_file_basename)
            saver.restore(sess, checkpoint_file_basename)
        else:            
            if encoder._pre_trained_model != 'False':
                encoder.saver.restore(sess, encoder._pre_trained_model)
                all_vars = set([var for var in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)])
                var_list = all_vars.symmetric_difference([v[1] for v in list(encoder.fil_var_list.items())])
                sess.run(tf.variables_initializer(var_list))
                print(sess.run(tf.report_uninitialized_variables()))
            else:
                sess.run(tf.global_variables_initializer())

        if not debug_mode:
            print('Training with %s model' % args.get('Dataset','MODEL'), os.path.basename(args.get('Paths','MODEL_PATH')))
            bar.start()

        while True:

            this,_,reconstr_train,enc_z  = sess.run([multi_queue.next_element,multi_queue.next_bg_element,[decoder.x for decoder in decoders], encoder.z])

            this_x = np.concatenate([el[0] for el in this])
            this_y = np.concatenate([el[2] for el in this])
            print(this_x.shape)
            reconstr_train = np.concatenate(reconstr_train)
            print(this_x.shape)
            cv2.imshow('sample batch', np.hstack(( u.tiles(this_x, 4, 6), u.tiles(reconstr_train, 4,6),u.tiles(this_y, 4, 6))) )
            k = cv2.waitKey(0)

            idx = np.random.randint(0,24)
            this_y = np.repeat(this_y[idx:idx+1, :, :], 24, axis=0)
            reconstr_train = sess.run([decoder.x for decoder in decoders],feed_dict={encoder._input:this_y})
            reconstr_train = np.array(reconstr_train)
            print(reconstr_train.shape)
            reconstr_train = reconstr_train.squeeze()
            cv2.imshow('sample batch 2', np.hstack((u.tiles(this_y, 4, 6), u.tiles(reconstr_train, 4, 6))))
            k = cv2.waitKey(0)
            if k == 27:
                break
            if gentle_stop[0]:
                break

        if not debug_mode:
            bar.finish()
        if not gentle_stop[0] and not debug_mode:
            print('To create the embedding run:\n')
            print('ae_embed {}\n'.format(full_name))
예제 #8
0
def main():
    tf.disable_eager_execution()

    workspace_path = os.environ.get('AE_WORKSPACE_PATH')

    if workspace_path == None:
        print('Please define a workspace path:\n')
        print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
        exit(-1)

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    parser.add_argument('--at_step', default=None, required=False)
    arguments = parser.parse_args()
    full_name = arguments.experiment_name.split('/')

    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    at_step = arguments.at_step

    cfg_file_path = u.get_config_file_path(workspace_path, experiment_name,
                                           experiment_group)
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    checkpoint_file = u.get_checkpoint_basefilename(log_dir)
    ckpt_dir = u.get_checkpoint_dir(log_dir)
    dataset_path = u.get_dataset_path(workspace_path)

    print(checkpoint_file)
    print(ckpt_dir)
    print('#' * 20)

    if not os.path.exists(cfg_file_path):
        print('Could not find config file:\n')
        print('{}\n'.format(cfg_file_path))
        exit(-1)

    args = configparser.ConfigParser()
    args.read(cfg_file_path)

    with tf.variable_scope(experiment_name):
        dataset = factory.build_dataset(dataset_path, args)
        queue = factory.build_queue(dataset, args)
        encoder = factory.build_encoder(queue.x, args)
        decoder = factory.build_decoder(queue.y, encoder, args)
        ae = factory.build_ae(encoder, decoder, args)
        codebook = factory.build_codebook(encoder, dataset, args)
        saver = tf.train.Saver(save_relative_paths=True)

    batch_size = args.getint('Training', 'BATCH_SIZE')
    model = args.get('Dataset', 'MODEL')

    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
    config = tf.ConfigProto(gpu_options=gpu_options)

    with tf.Session(config=config) as sess:

        print(ckpt_dir)
        print('#' * 20)

        factory.restore_checkpoint(sess, saver, ckpt_dir, at_step=at_step)

        # chkpt = tf.train.get_checkpoint_state(ckpt_dir)
        # if chkpt and chkpt.model_checkpoint_path:
        #     print chkpt.model_checkpoint_path
        #     saver.restore(sess, chkpt.model_checkpoint_path)
        # else:
        #     print 'No checkpoint found. Expected one in:\n'
        #     print '{}\n'.format(ckpt_dir)
        #     exit(-1)

        if model == 'dsprites':
            codebook.update_embedding_dsprites(sess, args)
        else:
            codebook.update_embedding(sess, batch_size)

        print('Saving new checkoint ..')

        saver.save(sess, checkpoint_file, global_step=ae.global_step)

        print('done')
예제 #9
0
def main():
    '''
    lxc:
    use_euclidean means the similarity between test embedding and template embedding 
    are computed using Euclidean Distance
    '''
    #use_euclidean = False

    parser = argparse.ArgumentParser()

    parser.add_argument('experiment_name')
    parser.add_argument('evaluation_name')
    parser.add_argument('--eval_cfg', default='eval.cfg', required=False)
    parser.add_argument('--at_step', default=None, required=False)
    arguments = parser.parse_args()
    full_name = arguments.experiment_name.split('/')
    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    evaluation_name = arguments.evaluation_name
    eval_cfg = arguments.eval_cfg
    at_step = arguments.at_step

    workspace_path = os.environ.get('AE_WORKSPACE_PATH')
    train_cfg_file_path = u.get_config_file_path(workspace_path,
                                                 experiment_name,
                                                 experiment_group)
    eval_cfg_file_path = u.get_eval_config_file_path(workspace_path,
                                                     eval_cfg=eval_cfg)

    train_args = configparser.ConfigParser()
    eval_args = configparser.ConfigParser()
    train_args.read(train_cfg_file_path)
    eval_args.read(eval_cfg_file_path)

    #[DATA]
    # target data params
    dataset_name = eval_args.get('DATA', 'DATASET')
    obj_id = eval_args.getint('DATA', 'OBJ_ID')
    scenes = eval(eval_args.get(
        'DATA', 'SCENES')) if len(eval(eval_args.get(
            'DATA',
            'SCENES'))) > 0 else eval_utils.get_all_scenes_for_obj(eval_args)
    cam_type = eval_args.get('DATA', 'cam_type')
    model_type = 'reconst' if dataset_name == 'tless' else ''  # model_type set to reconst only for tless.

    data_params = dataset_params.get_dataset_params(dataset_name,
                                                    model_type=model_type,
                                                    train_type='',
                                                    test_type=cam_type,
                                                    cam_type=cam_type)
    target_models_info = inout.load_yaml(
        data_params['models_info_path'])  # lxc

    # source data params, lxc
    source_dataset_name = 'toyotalight'
    # source_dataset_name = train_args.get('DATA','DATASET') # TODO train args no section DATA
    # source_obj_id = train_args.getint('DATA','OBJ_ID') # TODO train args no section DATA
    source_obj_id = int(train_cfg_file_path[-6:-4])  # TODO workaround
    source_data_params = dataset_params.get_dataset_params(source_dataset_name,
                                                           model_type='',
                                                           train_type='',
                                                           test_type='',
                                                           cam_type='')
    # for tless temporarily.
    # source_data_params = dataset_params.get_dataset_params(source_dataset_name, model_type='', train_type='', test_type='kinect', cam_type='kinect')
    source_models_info = inout.load_yaml(
        source_data_params['models_info_path'])
    print("source_models_info_path:", source_data_params['models_info_path'])
    # 'diameter' is not equal to sqrt(x^2+y^2+z^2) for hinterstoisser, rutgers, tless, tejaniDB. etc.
    # for toyotalight, 'diameter' == sqrt(...).
    target_models_3Dlength = np.linalg.norm([
        target_models_info[obj_id][key]
        for key in ['size_x', 'size_y', 'size_z']
    ])
    source_models_3Dlength = np.linalg.norm([
        source_models_info[source_obj_id][key]
        for key in ['size_x', 'size_y', 'size_z']
    ])

    target_source_length_ratio = target_models_3Dlength / source_models_3Dlength
    print("target_source_length_ratio:", target_source_length_ratio)
    print("source id {:02d}, target id {:02d}".format(source_obj_id, obj_id))
    print('basepath: ', data_params['base_path'])
    #[BBOXES]
    estimate_bbs = eval_args.getboolean('BBOXES', 'ESTIMATE_BBS')
    #[METRIC]
    top_nn = eval_args.getint('METRIC', 'TOP_N')
    #[EVALUATION]
    icp = eval_args.getboolean('EVALUATION', 'ICP')

    evaluation_name = evaluation_name + '_icp' if icp else evaluation_name
    evaluation_name = evaluation_name + '_bbest' if estimate_bbs else evaluation_name

    data = dataset_name + '_' + cam_type if len(cam_type) > 0 else dataset_name

    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    ckpt_dir = u.get_checkpoint_dir(log_dir)
    eval_dir = u.get_eval_dir(log_dir, evaluation_name, data)

    # if eval_args.getboolean('EVALUATION','EVALUATE_ERRORS'):
    #     eval_loc.match_and_eval_performance_scores(eval_args, eval_dir)
    #     exit()

    if not os.path.exists(eval_dir):
        os.makedirs(eval_dir)
    shutil.copy2(eval_cfg_file_path, eval_dir)

    print "eval_args: ", eval_args

    codebook, dataset, decoder = factory.build_codebook_from_name(
        experiment_name,
        experiment_group,
        return_dataset=True,
        return_decoder=True)
    dataset.renderer
    gpu_options = tf.GPUOptions(allow_growth=True,
                                per_process_gpu_memory_fraction=0.5)
    config = tf.ConfigProto(gpu_options=gpu_options)

    sess = tf.Session(config=config)
    factory.restore_checkpoint(sess,
                               tf.train.Saver(),
                               ckpt_dir,
                               at_step=at_step)

    if estimate_bbs:
        #Object Detection, seperate from main
        # sys.path.append('/net/rmc-lx0050/home_local/sund_ma/src/SSD_Tensorflow')
        # from ssd_detector import SSD_detector
        # #TODO: set num_classes, network etc.
        # ssd = SSD_detector(sess, num_classes=31, net_shape=(300,300))
        from rmcssd.bin import detector
        ssd = detector.Detector(eval_args.get('BBOXES', 'CKPT'))

    t_errors = []
    R_errors = []
    all_test_visibs = []

    test_embeddings = []
    for scene_id in scenes:

        test_imgs = eval_utils.load_scenes(scene_id, eval_args)
        test_imgs_depth = eval_utils.load_scenes(
            scene_id, eval_args, depth=True) if icp else None

        if estimate_bbs:
            print eval_args.get('BBOXES', 'EXTERNAL')
            if eval_args.get('BBOXES', 'EXTERNAL') == 'False':
                bb_preds = {}
                for i, img in enumerate(test_imgs):
                    print img.shape
                    bb_preds[i] = ssd.detectSceneBBs(img,
                                                     min_score=.2,
                                                     nms_threshold=.45)
                # inout.save_yaml(os.path.join(scene_res_dir,'bb_preds.yml'), bb_preds)
                print bb_preds
            else:
                bb_preds = inout.load_yaml(
                    os.path.join(eval_args.get('BBOXES', 'EXTERNAL'),
                                 '{:02d}.yml'.format(scene_id)))

            test_img_crops, test_img_depth_crops, bbs, bb_scores, visibilities = eval_utils.generate_scene_crops(
                test_imgs, test_imgs_depth, bb_preds, eval_args, train_args)
        else:
            # test_img_crops: each crop contains some bbox(es) for specified object id.
            test_img_crops, test_img_depth_crops, bbs, bb_scores, visibilities = eval_utils.get_gt_scene_crops(
                scene_id, eval_args, train_args)

        if len(test_img_crops) == 0:
            print 'ERROR: object %s not in scene %s' % (obj_id, scene_id)
            exit()

        info = inout.load_info(
            data_params['scene_info_mpath'].format(scene_id))
        Ks_test = [np.array(v['cam_K']).reshape(3, 3) for v in info.values()]

        ######remove
        gts = inout.load_gt(data_params['scene_gt_mpath'].format(scene_id))
        visib_gts = inout.load_yaml(data_params['scene_gt_stats_mpath'].format(
            scene_id, 15))
        #######
        W_test, H_test = data_params['test_im_size']

        icp_renderer = icp_utils.SynRenderer(train_args) if icp else None
        noof_scene_views = eval_utils.noof_scene_views(scene_id, eval_args)

        test_embeddings.append([])

        scene_res_dir = os.path.join(
            eval_dir, '{scene_id:02d}'.format(scene_id=scene_id))
        if not os.path.exists(scene_res_dir):
            os.makedirs(scene_res_dir)

        for view in xrange(
                noof_scene_views
        ):  # for example, LINEMOD ape noof_scene_views = 1236
            try:
                # only a specified object id is selected throughout the whole scene views.
                test_crops, test_crops_depth, test_bbs, test_scores, test_visibs = eval_utils.select_img_crops(
                    test_img_crops[view][obj_id],
                    test_img_depth_crops[view][obj_id] if icp else None,
                    bbs[view][obj_id], bb_scores[view][obj_id],
                    visibilities[view][obj_id], eval_args)
            except:
                print 'no detections'
                continue

            print view
            preds = {}
            pred_views = []
            all_test_visibs.append(test_visibs[0])
            t_errors_crop = []
            R_errors_crop = []

            for i, (test_crop, test_bb, test_score) in enumerate(
                    zip(test_crops, test_bbs, test_scores)):
                # each test_crop is a ground truth patch
                if train_args.getint('Dataset', 'C') == 1:
                    test_crop = cv2.cvtColor(test_crop,
                                             cv2.COLOR_BGR2GRAY)[:, :, None]
                start = time.time()
                '''modify here to change the pose estimation algorithm. lxc'''

                Rs_est, ts_est = codebook.auto_pose6d(
                    sess,
                    test_crop,
                    test_bb,
                    Ks_test[view].copy(),
                    top_nn,
                    train_args,
                    target_source_length_ratio=target_source_length_ratio)
                ae_time = time.time() - start
                run_time = ae_time + bb_preds[view][0][
                    'det_time'] if estimate_bbs else ae_time

                if eval_args.getboolean('PLOT', 'EMBEDDING_PCA'):
                    test_embeddings[-1].append(
                        codebook.test_embedding(sess,
                                                test_crop,
                                                normalized=True))

                # icp = False if view<350 else True
                #TODO:
                Rs_est_old, ts_est_old = Rs_est.copy(), ts_est.copy()
                for p in xrange(top_nn):
                    if icp:
                        start = time.time()
                        # icp only along tz
                        R_est_refined, t_est_refined = icp_utils.icp_refinement(
                            test_crops_depth[i],
                            icp_renderer,
                            Rs_est[p],
                            ts_est[p],
                            Ks_test[view].copy(), (W_test, H_test),
                            depth_only=True,
                            max_mean_dist_factor=5.0)
                        print ts_est[p]
                        print t_est_refined
                        # x,y update,does not change tz:
                        _, ts_est_refined = codebook.auto_pose6d(
                            sess,
                            test_crop,
                            test_bb,
                            Ks_test[view].copy(),
                            top_nn,
                            train_args,
                            depth_pred=t_est_refined[2])
                        # commented by lxc
                        # _, ts_est_refined, _ = codebook.auto_pose6d(sess, test_crop, test_bb, Ks_test[view].copy(), top_nn, train_args,depth_pred=t_est_refined[2])
                        t_est_refined = ts_est_refined[p]
                        # rotation icp, only accepted if below 20 deg change
                        R_est_refined, _ = icp_utils.icp_refinement(
                            test_crops_depth[i],
                            icp_renderer,
                            R_est_refined,
                            t_est_refined,
                            Ks_test[view].copy(), (W_test, H_test),
                            no_depth=True)
                        print Rs_est[p]
                        print R_est_refined
                        icp_time = time.time() - start
                        Rs_est[p], ts_est[p] = R_est_refined, t_est_refined
                    preds.setdefault('ests', []).append({
                        'score': test_score,
                        'R': Rs_est[p],
                        't': ts_est[p]
                    })
                run_time = run_time + icp_time if icp else run_time

                min_t_err, min_R_err = eval_plots.print_trans_rot_errors(
                    gts[view], obj_id, ts_est, ts_est_old, Rs_est, Rs_est_old)
                t_errors_crop.append(min_t_err)
                R_errors_crop.append(min_R_err)

                if eval_args.getboolean('PLOT', 'RECONSTRUCTION'):
                    eval_plots.plot_reconstruction_test(
                        sess, codebook._encoder, decoder, test_crop)
                    # eval_plots.plot_reconstruction_train(sess, decoder, nearest_train_codes[0])
                if eval_args.getboolean('PLOT',
                                        'NEAREST_NEIGHBORS') and not icp:
                    for R_est, t_est in zip(Rs_est, ts_est):
                        pred_views.append(
                            dataset.render_rot(R_est, downSample=2))
                    eval_plots.show_nearest_rotation(pred_views, test_crop,
                                                     view)
                if eval_args.getboolean('PLOT', 'SCENE_WITH_ESTIMATE'):
                    eval_plots.plot_scene_with_estimate(
                        test_imgs[view].copy(),
                        icp_renderer.renderer if icp else dataset.renderer,
                        Ks_test[view].copy(), Rs_est_old[0], ts_est_old[0],
                        Rs_est[0], ts_est[0], test_bb, test_score, obj_id,
                        gts[view], bb_preds[view] if estimate_bbs else None)

                if cv2.waitKey(1) == 32:
                    cv2.waitKey(0)

            t_errors.append(t_errors_crop[np.argmin(
                np.linalg.norm(np.array(t_errors_crop), axis=1))])
            R_errors.append(R_errors_crop[np.argmin(
                np.linalg.norm(np.array(t_errors_crop), axis=1))])

            # save predictions in sixd format
            res_path = os.path.join(scene_res_dir,
                                    '%04d_%02d.yml' % (view, obj_id))
            inout.save_results_sixd17(res_path, preds, run_time=run_time)

    if not os.path.exists(os.path.join(eval_dir, 'latex')):
        os.makedirs(os.path.join(eval_dir, 'latex'))
    if not os.path.exists(os.path.join(eval_dir, 'figures')):
        os.makedirs(os.path.join(eval_dir, 'figures'))
    '''evaluation code
        dataset_renderer renders source object model for evaluation;
        If we need target object model for evaluation, go get a new renderer.
    '''

    if eval_args.getboolean('EVALUATION', 'COMPUTE_ERRORS'):
        eval_calc_errors.eval_calc_errors(eval_args,
                                          eval_dir,
                                          dataset_renderer=dataset.renderer)
    if eval_args.getboolean('EVALUATION', 'EVALUATE_ERRORS'):
        eval_loc.match_and_eval_performance_scores(eval_args, eval_dir)
    '''plot code'''
    cyclo = train_args.getint('Embedding', 'NUM_CYCLO')
    if eval_args.getboolean('PLOT', 'EMBEDDING_PCA'):
        embedding = sess.run(codebook.embedding_normalized)
        eval_plots.compute_pca_plot_embedding(eval_dir,
                                              embedding[::cyclo],
                                              np.array(test_embeddings[0]),
                                              obj_id=obj_id)
    if eval_args.getboolean('PLOT', 'VIEWSPHERE'):
        eval_plots.plot_viewsphere_for_embedding(
            dataset.viewsphere_for_embedding[::cyclo], eval_dir, obj_id=obj_id)
    if eval_args.getboolean('PLOT', 'CUM_T_ERROR_HIST'):
        eval_plots.plot_t_err_hist(np.array(t_errors), eval_dir, obj_id=obj_id)
        eval_plots.plot_t_err_hist2(np.array(t_errors),
                                    eval_dir,
                                    obj_id=obj_id)
    if eval_args.getboolean('PLOT', 'CUM_R_ERROR_HIST'):
        eval_plots.plot_R_err_hist(eval_args, eval_dir, scenes)
        eval_plots.plot_R_err_hist2(np.array(R_errors),
                                    eval_dir,
                                    obj_id=obj_id)
    if eval_args.getboolean('PLOT', 'CUM_VSD_ERROR_HIST'):
        eval_plots.plot_vsd_err_hist(eval_args, eval_dir, scenes)
    if eval_args.getboolean('PLOT', 'VSD_OCCLUSION'):
        eval_plots.plot_vsd_occlusion(eval_args, eval_dir, scenes,
                                      np.array(all_test_visibs))
    if eval_args.getboolean('PLOT', 'R_ERROR_OCCLUSION'):
        eval_plots.plot_re_rect_occlusion(eval_args, eval_dir, scenes,
                                          np.array(all_test_visibs))
    if eval_args.getboolean('PLOT', 'ANIMATE_EMBEDDING_PCA'):
        eval_plots.animate_embedding_path(test_embeddings[0])
    if eval_args.getboolean('PLOT', 'RECONSTRUCTION_TEST_BATCH'):
        eval_plots.plot_reconstruction_test_batch(sess,
                                                  codebook,
                                                  decoder,
                                                  test_img_crops,
                                                  noof_scene_views,
                                                  obj_id,
                                                  eval_dir=eval_dir)
        # plt.show()

        # calculate 6D pose errors
        # print 'exiting ...'
        # eval_calc_errors.eval_calc_errors(eval_args, eval_dir)
        # calculate 6D pose errors

    report = latex_report.Report(eval_dir, log_dir)
    report.write_configuration(train_cfg_file_path, eval_cfg_file_path)
    report.merge_all_tex_files()
    report.include_all_figures()
    report.save(open_pdf=False)
예제 #10
0
def main():

    parser = argparse.ArgumentParser()

    parser.add_argument('experiment_name')
    parser.add_argument('evaluation_name')
    parser.add_argument('--eval_cfg', default='eval.cfg', required=False)
    parser.add_argument('--at_step', default=None, type=str, required=False)
    parser.add_argument('--model_path', default=None, required=True)
    arguments = parser.parse_args()
    full_name = arguments.experiment_name.split('/')
    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    evaluation_name = arguments.evaluation_name
    eval_cfg = arguments.eval_cfg
    at_step = arguments.at_step
    model_path = arguments.model_path

    workspace_path = os.environ.get('AE_WORKSPACE_PATH')
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    train_cfg_file_path = u.get_train_config_exp_file_path(
        log_dir, experiment_name)
    eval_cfg_file_path = u.get_eval_config_file_path(workspace_path,
                                                     eval_cfg=eval_cfg)

    train_args = configparser.ConfigParser(inline_comment_prefixes="#")
    eval_args = configparser.ConfigParser(inline_comment_prefixes="#")
    train_args.read(train_cfg_file_path)
    eval_args.read(eval_cfg_file_path)

    #[DATA]
    dataset_name = eval_args.get('DATA', 'DATASET')
    obj_id = eval_args.getint('DATA', 'OBJ_ID')
    scenes = eval(eval_args.get(
        'DATA', 'SCENES')) if len(eval(eval_args.get(
            'DATA',
            'SCENES'))) > 0 else eval_utils.get_all_scenes_for_obj(eval_args)
    cam_type = eval_args.get('DATA', 'cam_type')
    data_params = dataset_params.get_dataset_params(dataset_name,
                                                    model_type='',
                                                    train_type='',
                                                    test_type=cam_type,
                                                    cam_type=cam_type)
    #[BBOXES]
    estimate_bbs = eval_args.getboolean('BBOXES', 'ESTIMATE_BBS')
    gt_masks = eval_args.getboolean('BBOXES', 'gt_masks')
    estimate_masks = eval_args.getboolean('BBOXES', 'estimate_masks')

    #[METRIC]
    top_nn = eval_args.getint('METRIC', 'TOP_N')
    #[EVALUATION]
    icp = eval_args.getboolean('EVALUATION', 'ICP')
    gt_trans = eval_args.getboolean('EVALUATION', 'gt_trans')
    iterative_code_refinement = eval_args.getboolean(
        'EVALUATION', 'iterative_code_refinement')

    H_AE = train_args.getint('Dataset', 'H')
    W_AE = train_args.getint('Dataset', 'W')

    evaluation_name = evaluation_name + '_icp' if icp else evaluation_name
    evaluation_name = evaluation_name + '_bbest' if estimate_bbs else evaluation_name
    evaluation_name = evaluation_name + '_maskest' if estimate_masks else evaluation_name
    evaluation_name = evaluation_name + '_gttrans' if gt_trans else evaluation_name
    evaluation_name = evaluation_name + '_gtmasks' if gt_masks else evaluation_name
    evaluation_name = evaluation_name + '_refined' if iterative_code_refinement else evaluation_name

    data = dataset_name + '_' + cam_type if len(cam_type) > 0 else dataset_name

    if at_step is None:
        checkpoint_file = u.get_checkpoint_basefilename(
            log_dir,
            False,
            latest=train_args.getint('Training', 'NUM_ITER'),
            joint=True)
    else:
        checkpoint_file = u.get_checkpoint_basefilename(log_dir,
                                                        False,
                                                        latest=at_step,
                                                        joint=True)
    print(checkpoint_file)
    eval_dir = u.get_eval_dir(log_dir, evaluation_name, data)

    if not os.path.exists(eval_dir):
        os.makedirs(eval_dir)
    shutil.copy2(eval_cfg_file_path, eval_dir)

    codebook, dataset = factory.build_codebook_from_name(experiment_name,
                                                         experiment_group,
                                                         return_dataset=True,
                                                         joint=True)
    dataset._kw['model_path'] = [model_path]
    dataset._kw['model'] = 'cad' if 'cad' in model_path else 'reconst'
    dataset._kw['model'] = 'reconst' if 'reconst' in model_path else 'cad'

    gpu_options = tf.GPUOptions(allow_growth=True,
                                per_process_gpu_memory_fraction=0.5)
    config = tf.ConfigProto(gpu_options=gpu_options)

    sess = tf.Session(config=config)
    saver = tf.train.Saver()
    saver.restore(sess, checkpoint_file)

    t_errors = []
    R_errors = []
    all_test_visibs = []

    external_path = eval_args.get('BBOXES', 'EXTERNAL')

    test_embeddings = []
    for scene_id in scenes:

        test_imgs = eval_utils.load_scenes(scene_id, eval_args)
        test_imgs_depth = eval_utils.load_scenes(
            scene_id, eval_args, depth=True) if icp else None

        if estimate_bbs:
            print(external_path)
            if external_path == 'False':
                bb_preds = {}
                for i, img in enumerate(test_imgs):
                    print((img.shape))
                    bb_preds[i] = ssd.detectSceneBBs(img,
                                                     min_score=.05,
                                                     nms_threshold=.45)
                print(bb_preds)
            else:
                if estimate_masks:
                    bb_preds = inout.load_yaml(
                        os.path.join(
                            external_path,
                            '{:02d}/mask_rcnn_predict.yml'.format(scene_id)))
                    print(list(bb_preds[0][0].keys()))
                    mask_paths = glob.glob(
                        os.path.join(external_path,
                                     '{:02d}/masks/*.npy'.format(scene_id)))
                    maskrcnn_scene_masks = [np.load(mp) for mp in mask_paths]
                else:
                    maskrcnn_scene_masks = None
                    bb_preds = inout.load_yaml(
                        os.path.join(external_path,
                                     '{:02d}.yml'.format(scene_id)))

            test_img_crops, test_img_depth_crops, bbs, bb_scores, visibilities = eval_utils.generate_scene_crops(
                test_imgs,
                test_imgs_depth,
                bb_preds,
                eval_args, (H_AE, W_AE),
                inst_masks=maskrcnn_scene_masks)
        else:
            test_img_crops, test_img_depth_crops, bbs, bb_scores, visibilities = eval_utils.get_gt_scene_crops(
                scene_id,
                eval_args,
                train_args,
                load_gt_masks=external_path if gt_masks else gt_masks)

        if len(test_img_crops) == 0:
            print(('ERROR: object %s not in scene %s' % (obj_id, scene_id)))
            exit()

        info = inout.load_info(
            data_params['scene_info_mpath'].format(scene_id))
        Ks_test = [
            np.array(v['cam_K']).reshape(3, 3) for v in list(info.values())
        ]

        ######remove
        gts = inout.load_gt(data_params['scene_gt_mpath'].format(scene_id))
        visib_gts = inout.load_yaml(data_params['scene_gt_stats_mpath'].format(
            scene_id, 15))
        #######
        W_test, H_test = data_params['test_im_size']

        icp_renderer = icp_utils.SynRenderer(
            train_args, dataset._kw['model_path'][0]) if icp else None
        noof_scene_views = eval_utils.noof_scene_views(scene_id, eval_args)

        test_embeddings.append([])

        scene_res_dir = os.path.join(
            eval_dir, '{scene_id:02d}'.format(scene_id=scene_id))
        if not os.path.exists(scene_res_dir):
            os.makedirs(scene_res_dir)

        for view in range(noof_scene_views):
            try:
                test_crops, test_crops_depth, test_bbs, test_scores, test_visibs = eval_utils.select_img_crops(
                    test_img_crops[view][obj_id],
                    test_img_depth_crops[view][obj_id] if icp else None,
                    bbs[view][obj_id], bb_scores[view][obj_id],
                    visibilities[view][obj_id], eval_args)
            except:
                print('no detections')
                continue

            print(view)
            preds = {}
            pred_views = []
            all_test_visibs.append(test_visibs[0])
            t_errors_crop = []
            R_errors_crop = []

            for i, (test_crop, test_bb, test_score) in enumerate(
                    zip(test_crops, test_bbs, test_scores)):

                start = time.time()
                if train_args.getint('Dataset', 'C') == 1:
                    test_crop = cv2.cvtColor(test_crop,
                                             cv2.COLOR_BGR2GRAY)[:, :, None]
                Rs_est, ts_est, _ = codebook.auto_pose6d(
                    sess,
                    test_crop,
                    test_bb,
                    Ks_test[view].copy(),
                    top_nn,
                    train_args,
                    codebook._get_codebook_name(model_path),
                    refine=iterative_code_refinement)
                Rs_est_old, ts_est_old = Rs_est.copy(), ts_est.copy()
                ae_time = time.time() - start

                if eval_args.getboolean('PLOT', 'EMBEDDING_PCA'):
                    test_embeddings[-1].append(
                        codebook.test_embedding(sess,
                                                test_crop,
                                                normalized=True))

                if eval_args.getboolean('EVALUATION', 'gt_trans'):
                    ts_est = np.empty((top_nn, 3))
                    for n in range(top_nn):
                        smallest_diff = np.inf
                        for visib_gt, gt in zip(visib_gts[view], gts[view]):
                            if gt['obj_id'] == obj_id:
                                diff = np.sum(
                                    np.abs(gt['obj_bb'] -
                                           np.array(visib_gt['bbox_obj'])))
                                if diff < smallest_diff:
                                    smallest_diff = diff
                                    gt_obj = gt.copy()
                                    print('Im there')
                        ts_est[n] = np.array(gt_obj['cam_t_m2c']).reshape(-1)

                try:
                    run_time = ae_time + bb_preds[view][0][
                        'det_time'] if estimate_bbs else ae_time
                except:
                    run_time = ae_time

                for p in range(top_nn):
                    if icp:
                        # note: In the CVPR paper a different ICP was used
                        start = time.time()
                        # depth icp
                        R_est_refined, t_est_refined = icp_utils.icp_refinement(
                            test_crops_depth[i],
                            icp_renderer,
                            Rs_est[p],
                            ts_est[p],
                            Ks_test[view].copy(), (W_test, H_test),
                            depth_only=True,
                            max_mean_dist_factor=5.0)
                        print(t_est_refined)

                        # x,y update,does not change tz:
                        _, ts_est_refined, _ = codebook.auto_pose6d(
                            sess,
                            test_crop,
                            test_bb,
                            Ks_test[view].copy(),
                            top_nn,
                            train_args,
                            codebook._get_codebook_name(model_path),
                            depth_pred=t_est_refined[2],
                            refine=iterative_code_refinement)

                        t_est_refined = ts_est_refined[p]

                        # rotation icp, only accepted if below 20 deg change
                        R_est_refined, _ = icp_utils.icp_refinement(
                            test_crops_depth[i],
                            icp_renderer,
                            R_est_refined,
                            t_est_refined,
                            Ks_test[view].copy(), (W_test, H_test),
                            no_depth=True)
                        print((Rs_est[p]))
                        print(R_est_refined)

                        icp_time = time.time() - start
                        Rs_est[p], ts_est[p] = R_est_refined, t_est_refined

                    preds.setdefault('ests', []).append({
                        'score': test_score,
                        'R': Rs_est[p],
                        't': ts_est[p]
                    })
                run_time = run_time + icp_time if icp else run_time

                min_t_err, min_R_err = eval_plots.print_trans_rot_errors(
                    gts[view], obj_id, ts_est, ts_est_old, Rs_est, Rs_est_old)
                t_errors_crop.append(min_t_err)
                R_errors_crop.append(min_R_err)

                if eval_args.getboolean('PLOT',
                                        'NEAREST_NEIGHBORS') and not icp:
                    for R_est, t_est in zip(Rs_est, ts_est):
                        pred_views.append(
                            dataset.render_rot(R_est, downSample=2))
                    eval_plots.show_nearest_rotation(pred_views, test_crop,
                                                     view)
                if eval_args.getboolean('PLOT', 'SCENE_WITH_ESTIMATE'):
                    eval_plots.plot_scene_with_estimate(
                        test_imgs[view].copy(),
                        icp_renderer.renderer if icp else dataset.renderer,
                        Ks_test[view].copy(), Rs_est_old[0], ts_est_old[0],
                        Rs_est[0], ts_est[0], test_bb, test_score, obj_id,
                        gts[view], bb_preds[view] if estimate_bbs else None)

                if cv2.waitKey(1) == 32:
                    cv2.waitKey(0)

            t_errors.append(t_errors_crop[np.argmin(
                np.linalg.norm(np.array(t_errors_crop), axis=1))])
            R_errors.append(R_errors_crop[np.argmin(
                np.linalg.norm(np.array(t_errors_crop), axis=1))])

            # save predictions in sixd format
            res_path = os.path.join(scene_res_dir,
                                    '%04d_%02d.yml' % (view, obj_id))
            inout.save_results_sixd17(res_path, preds, run_time=run_time)

    if not os.path.exists(os.path.join(eval_dir, 'latex')):
        os.makedirs(os.path.join(eval_dir, 'latex'))
    if not os.path.exists(os.path.join(eval_dir, 'figures')):
        os.makedirs(os.path.join(eval_dir, 'figures'))

    if eval_args.getboolean('EVALUATION', 'COMPUTE_ERRORS'):
        eval_calc_errors.eval_calc_errors(eval_args, eval_dir)
    if eval_args.getboolean('EVALUATION', 'EVALUATE_ERRORS'):
        eval_loc.match_and_eval_performance_scores(eval_args, eval_dir)

    cyclo = train_args.getint('Embedding', 'NUM_CYCLO')
    if eval_args.getboolean('PLOT', 'EMBEDDING_PCA'):
        embedding = sess.run(codebook.embedding_normalized)
        eval_plots.compute_pca_plot_embedding(eval_dir, embedding[::cyclo],
                                              np.array(test_embeddings[0]))
    if eval_args.getboolean('PLOT', 'VIEWSPHERE'):
        eval_plots.plot_viewsphere_for_embedding(
            dataset.viewsphere_for_embedding[::cyclo], eval_dir)
    if eval_args.getboolean('PLOT', 'CUM_T_ERROR_HIST'):
        eval_plots.plot_t_err_hist(np.array(t_errors), eval_dir)
        eval_plots.plot_t_err_hist2(np.array(t_errors), eval_dir)
    if eval_args.getboolean('PLOT', 'CUM_R_ERROR_HIST'):
        eval_plots.plot_R_err_recall(eval_args, eval_dir, scenes)
        eval_plots.plot_R_err_hist2(np.array(R_errors), eval_dir)
    if eval_args.getboolean('PLOT', 'CUM_VSD_ERROR_HIST'):
        try:
            eval_plots.plot_vsd_err_hist(eval_args, eval_dir, scenes)
        except:
            pass
    if eval_args.getboolean('PLOT', 'VSD_OCCLUSION'):
        try:
            eval_plots.plot_vsd_occlusion(eval_args, eval_dir, scenes,
                                          np.array(all_test_visibs))
        except:
            pass
    if eval_args.getboolean('PLOT', 'R_ERROR_OCCLUSION'):
        try:
            eval_plots.plot_re_rect_occlusion(eval_args, eval_dir, scenes,
                                              np.array(all_test_visibs))
        except:
            pass
    if eval_args.getboolean('PLOT', 'ANIMATE_EMBEDDING_PCA'):
        eval_plots.animate_embedding_path(test_embeddings[0])

    report = latex_report.Report(eval_dir, log_dir)
    report.write_configuration(train_cfg_file_path, eval_cfg_file_path)
    report.merge_all_tex_files()
    report.include_all_figures()
    report.save(open_pdf=True)
    def __init__(self, test_config_path):

        test_args = self.get_params(test_config_path)

        workspace_path = os.environ.get('AE_WORKSPACE_PATH')

        if workspace_path == None:
            print('Please define a workspace path:\n')
            print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
            exit(-1)

        self._process_requirements = ['color_img', 'camK', 'bboxes']
        self._full_model_name = test_args.get('mp_encoder', 'full_model_name')
        if test_args.getboolean('mp_encoder', 'camPose'):
            self._process_requirements.append('camPose')
        self._camPose = test_args.getboolean('mp_encoder', 'camPose')
        self._upright = test_args.getboolean('mp_encoder', 'upright')
        self._topk = test_args.getint('mp_encoder', 'topk')
        if self._topk > 1:
            print('ERROR: topk > 1 not implemented yet')
            exit()

        self._image_format = {
            'color_format': test_args.get('mp_encoder', 'color_format'),
            'color_data_type':
            eval(test_args.get('mp_encoder', 'color_data_type')),
            'depth_data_type':
            eval(test_args.get('mp_encoder', 'depth_data_type'))
        }

        # self.vis = test_args.getboolean('mp_encoder','pose_visualization')

        # self.all_experiments = eval(test_args.get('mp_encoder','experiments'))
        # self.class_2_codebook = eval(test_args.get('mp_encoder','class_2_codebook'))

        config = tf.ConfigProto(allow_soft_placement=True)
        config.gpu_options.allow_growth = True
        config.gpu_options.per_process_gpu_memory_fraction = test_args.getfloat(
            'mp_encoder', 'gpu_memory_fraction')

        full_name = self._full_model_name.split('/')
        experiment_name = full_name.pop()
        experiment_group = full_name.pop() if len(full_name) > 0 else ''
        log_dir = utils.get_log_dir(workspace_path, experiment_name,
                                    experiment_group)

        train_cfg_file_path = utils.get_train_config_exp_file_path(
            log_dir, experiment_name)
        print(train_cfg_file_path)
        # train_cfg_file_path = utils.get_config_file_path(workspace_path, experiment_name, experiment_group)
        self.train_args = configparser.ConfigParser(
            inline_comment_prefixes="#")
        self.train_args.read(train_cfg_file_path)

        checkpoint_file = utils.get_checkpoint_basefilename(
            log_dir,
            False,
            latest=self.train_args.getint('Training', 'NUM_ITER'),
            joint=True)

        self.codebook_multi, self.dataset = ae_factory.build_codebook_from_name(
            experiment_name, experiment_group, return_dataset=True, joint=True)
        encoder = self.codebook_multi._encoder

        try:
            base_path = test_args.get('mp_encoder', 'base_path')
            class_2_objs = eval(test_args.get('mp_encoder', 'class_2_objs'))
            self.class_2_objpath, self.class_2_codebook = {}, {}
            for class_name, obj_path in class_2_objs.items():
                self.class_2_objpath[class_name] = os.path.join(
                    base_path, obj_path)
                self.class_2_codebook[
                    class_name] = self.codebook_multi._get_codebook_name(
                        self.class_2_objpath[class_name])
        except:
            self.class_2_codebook = eval(
                test_args.get('mp_encoder', 'class_2_codebook'))
        self.sess = tf.Session(config=config)
        saver = tf.train.Saver()
        saver.restore(self.sess, checkpoint_file)

        self.pad_factor = self.train_args.getfloat('Dataset', 'PAD_FACTOR')
        self.patch_size = (self.train_args.getint('Dataset', 'W'),
                           self.train_args.getint('Dataset', 'H'))
def main():
    workspace_path = os.environ.get('AE_WORKSPACE_PATH')

    if workspace_path is None:
        print('Please define a workspace path:\n')
        print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
        exit(-1)

    gentle_stop = np.array((1, ), dtype=np.bool)
    gentle_stop[0] = False

    def on_ctrl_c(signal, frame):
        gentle_stop[0] = True

    signal.signal(signal.SIGINT, on_ctrl_c)

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    parser.add_argument("-d", action='store_true', default=False)
    parser.add_argument("-gen", action='store_true', default=False)
    parser.add_argument('--at_step', default=None, type=int, required=False)

    arguments = parser.parse_args()

    full_name = arguments.experiment_name.split('/')

    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''

    debug_mode = arguments.d
    generate_data = arguments.gen
    at_step = arguments.at_step

    cfg_file_path = u.get_config_file_path(workspace_path, experiment_name,
                                           experiment_group)
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    checkpoint_file = u.get_checkpoint_basefilename(log_dir)
    ckpt_dir = u.get_checkpoint_dir(log_dir)
    train_fig_dir = u.get_train_fig_dir(log_dir)
    dataset_path = u.get_dataset_path(workspace_path)

    if not os.path.exists(cfg_file_path):
        print('Could not find config file:\n')
        print(('{}\n'.format(cfg_file_path)))
        exit(-1)

    if not os.path.exists(ckpt_dir):
        os.makedirs(ckpt_dir)
    if not os.path.exists(train_fig_dir):
        os.makedirs(train_fig_dir)
    if not os.path.exists(dataset_path):
        os.makedirs(dataset_path)

    args = configparser.ConfigParser(inline_comment_prefixes="#")
    args.read(cfg_file_path)

    shutil.copy2(cfg_file_path, log_dir)

    num_iter = args.getint(
        'Training', 'NUM_ITER') if not debug_mode else np.iinfo(np.int32).max
    save_interval = args.getint('Training', 'SAVE_INTERVAL')
    num_gpus = args.getint('Training', 'NUM_GPUS')

    with tf.device('/device:CPU:0'):
        with tf.variable_scope(experiment_name, reuse=tf.AUTO_REUSE):

            dataset = factory.build_dataset(dataset_path, args)
            multi_queue = factory.build_multi_queue(dataset, args)
            if generate_data:
                # dataset.load_bg_images(dataset_path)
                multi_queue.create_tfrecord_training_images(dataset_path, args)
                print('finished generating training images')
                exit()

            dev_splits = np.array_split(np.arange(multi_queue._num_objects),
                                        num_gpus)

            iterator = multi_queue.create_iterator(dataset_path, args)

            all_x, all_y = list(
                zip(*[(inp[0], inp[2]) for inp in multi_queue.next_element]))
            all_x, all_y = tf.concat(all_x, axis=0), tf.concat(all_y, axis=0)
            print(all_x.shape)
            encoding_splits = []
            for dev in range(num_gpus):
                with tf.device('/device:GPU:%s' % dev):
                    sta = dev_splits[dev][0] * multi_queue._batch_size
                    end = (dev_splits[dev][-1] + 1) * multi_queue._batch_size
                    print(sta, end)
                    encoder = factory.build_encoder(all_x[sta:end],
                                                    args,
                                                    target=all_y[sta:end],
                                                    is_training=True)
                    encoding_splits.append(
                        tf.split(encoder.z, len(dev_splits[dev]), 0))

        with tf.variable_scope(experiment_name):
            decoders = []
            for dev in range(num_gpus):
                with tf.device('/device:GPU:%s' % dev):
                    for j, i in enumerate(dev_splits[dev]):
                        print(len(encoding_splits))
                        decoders.append(
                            factory.build_decoder(multi_queue.next_element[i],
                                                  encoding_splits[dev][j],
                                                  args,
                                                  is_training=True,
                                                  idx=i))

            ae = factory.build_ae(encoder, decoders, args)
            codebook = factory.build_codebook(encoder, dataset, args)
            train_op = factory.build_train_op(ae, args)
            saver = tf.train.Saver(save_relative_paths=True, max_to_keep=1)

        # dataset.get_training_images(dataset_path, args)
    # dataset.load_bg_images(dataset_path)
    multi_queue.create_tfrecord_training_images(dataset_path, args)

    if generate_data:
        print(('finished generating synthetic training data for ' +
               experiment_name))
        print('exiting...')
        exit()

    widgets = [
        'Training: ',
        progressbar.Percentage(), ' ',
        progressbar.Bar(), ' ',
        progressbar.Counter(),
        ' / %s' % num_iter, ' ',
        progressbar.ETA(), ' '
    ]
    bar = progressbar.ProgressBar(maxval=num_iter, widgets=widgets)

    gpu_options = tf.GPUOptions(allow_growth=True,
                                per_process_gpu_memory_fraction=0.9)
    config = tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True)

    with tf.Session(config=config) as sess:

        sess.run(multi_queue.bg_img_init.initializer)
        sess.run(iterator.initializer)

        u.create_summaries(multi_queue, decoders, ae)
        merged_loss_summary = tf.summary.merge_all()
        summary_writer = tf.summary.FileWriter(ckpt_dir, sess.graph)

        chkpt = tf.train.get_checkpoint_state(ckpt_dir)
        if chkpt and chkpt.model_checkpoint_path:
            if at_step is None:
                # checkpoint_file_basename = u.get_checkpoint_basefilename(log_dir,latest=args.getint('Training', 'NUM_ITER'))
                checkpoint_file_basename = chkpt.model_checkpoint_path
            else:
                checkpoint_file_basename = u.get_checkpoint_basefilename(
                    log_dir, latest=at_step)
            print(('loading ', checkpoint_file_basename))
            saver.restore(sess, checkpoint_file_basename)
            # except:
            #     print 'loading ', chkpt.model_checkpoint_path
            #     saver.restore(sess, chkpt.model_checkpoint_path)
        else:
            if encoder._pre_trained_model != 'False':
                encoder.saver.restore(sess, encoder._pre_trained_model)
                all_vars = set([
                    var
                    for var in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
                ])
                var_list = all_vars.symmetric_difference(
                    [v[1] for v in list(encoder.fil_var_list.items())])
                sess.run(tf.variables_initializer(var_list))
                print(sess.run(tf.report_uninitialized_variables()))
            else:
                sess.run(tf.global_variables_initializer())

        if not debug_mode:
            print(('Training with %s model' % args.get('Dataset', 'MODEL'),
                   os.path.basename(args.get('Paths', 'MODEL_PATH'))))
            bar.start()

        for i in range(encoder.global_step.eval(), num_iter):
            if not debug_mode:
                # print 'before optimize'
                sess.run([train_op, multi_queue.next_bg_element])
                # print 'after optimize'
                if (i + 1) % 100 == 0:
                    merged_summaries = sess.run(merged_loss_summary)
                    summary_writer.add_summary(merged_summaries, i)

                bar.update(i)

                if (i + 1) % save_interval == 0:
                    saver.save(sess,
                               checkpoint_file,
                               global_step=encoder.global_step)

                    # this_x, this_y = sess.run([queue.x, queue.y])
                    # reconstr_train = sess.run(decoder.x,feed_dict={queue.x:this_x})

                    this, reconstr_train = sess.run([
                        multi_queue.next_element,
                        [decoder.x for decoder in decoders]
                    ])
                    this_x = np.concatenate([el[0] for el in this])
                    this_y = np.concatenate([el[2] for el in this])
                    # reconstr_train = sess.run(,feed_dict={queue.x:this_x})
                    reconstr_train = np.concatenate(reconstr_train)
                    for imgs in [this_x, this_y, reconstr_train]:
                        np.random.seed(0)
                        np.random.shuffle(imgs)
                    train_imgs = np.hstack(
                        (u.tiles(this_x, 4,
                                 4), u.tiles(reconstr_train, 4,
                                             4), u.tiles(this_y, 4, 4)))
                    cv2.imwrite(
                        os.path.join(train_fig_dir,
                                     'training_images_%s.png' % i),
                        train_imgs * 255)
            else:

                this, _, reconstr_train = sess.run([
                    multi_queue.next_element, multi_queue.next_bg_element,
                    [decoder.x for decoder in decoders]
                ])

                this_x = np.concatenate([el[0] for el in this])
                this_y = np.concatenate([el[2] for el in this])
                print(this_x.shape, reconstr_train[0].shape,
                      len(reconstr_train))
                reconstr_train = np.concatenate(reconstr_train, axis=0)
                for imgs in [this_x, this_y, reconstr_train]:
                    np.random.seed(0)
                    np.random.shuffle(imgs)
                print(this_x.shape)
                cv2.imshow(
                    'sample batch',
                    np.hstack((u.tiles(this_x, 4,
                                       6), u.tiles(reconstr_train, 4,
                                                   6), u.tiles(this_y, 4, 6))))
                k = cv2.waitKey(0)
                if k == 27:
                    break

            if gentle_stop[0]:
                break

        if not debug_mode:
            bar.finish()
        if not gentle_stop[0] and not debug_mode:
            print('To create the embedding run:\n')
            print(('ae_embed {}\n'.format(full_name)))
def main():
    workspace_path = os.environ.get('AE_WORKSPACE_PATH')

    if workspace_path == None:
        print('Please define a workspace path:\n')
        print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
        exit(-1)

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    parser.add_argument('--at_step', default=None, type=int, required=False)
    parser.add_argument('--model_path', type=str, required=True)
    arguments = parser.parse_args()
    full_name = arguments.experiment_name.split('/')

    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    at_step = arguments.at_step
    model_path = arguments.model_path

    cfg_file_path = u.get_config_file_path(workspace_path, experiment_name,
                                           experiment_group)
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)

    ckpt_dir = u.get_checkpoint_dir(log_dir)
    dataset_path = u.get_dataset_path(workspace_path)

    if not os.path.exists(cfg_file_path):
        print('Could not find config file:\n')
        print('{}\n'.format(cfg_file_path))
        exit(-1)

    args = configparser.ConfigParser()
    args.read(cfg_file_path)
    iteration = args.getint('Training',
                            'NUM_ITER') if at_step is None else at_step

    checkpoint_file_basename = u.get_checkpoint_basefilename(log_dir,
                                                             latest=iteration,
                                                             joint=True)
    if not tf.train.checkpoint_exists(checkpoint_file_basename):
        checkpoint_file_basename = u.get_checkpoint_basefilename(
            log_dir, latest=iteration, joint=False)

    checkpoint_single_encoding = u.get_checkpoint_basefilename(
        log_dir, latest=iteration, model_path=model_path)
    target_checkpoint_file = u.get_checkpoint_basefilename(log_dir, joint=True)

    print(checkpoint_file_basename)
    print(target_checkpoint_file)
    print(ckpt_dir)
    print('#' * 20)

    with tf.variable_scope(experiment_name):
        dataset = factory.build_dataset(dataset_path, args)
        queue = factory.build_queue(dataset, args)
        encoder = factory.build_encoder(queue.x, args)
        # decoder = factory.build_decoder(queue.y, encoder, args)
        # ae = factory.build_ae(encoder, decoder, args)
        # before_cb = set(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES))
        codebook_multi = factory.build_codebook_multi(
            encoder, dataset, args, checkpoint_file_basename)
        restore_saver = tf.train.Saver(save_relative_paths=True,
                                       max_to_keep=100)

        codebook_multi.add_new_codebook_to_graph(model_path)
        # inters_vars = before_cb.intersection(set(tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)))
        saver = tf.train.Saver(save_relative_paths=True, max_to_keep=100)

    batch_size = args.getint('Training', 'BATCH_SIZE') * len(
        eval(args.get('Paths', 'MODEL_PATH')))
    model = args.get('Dataset', 'MODEL')

    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)
    config = tf.ConfigProto(gpu_options=gpu_options)

    with tf.Session(config=config) as sess:
        print(ckpt_dir)
        # print sess.run(encoder.global_step)
        print('#' * 20)

        # factory.restore_checkpoint(sess, saver, ckpt_dir, at_step=at_step)
        sess.run(tf.global_variables_initializer())
        restore_saver.restore(sess, checkpoint_file_basename)

        print('#' * 20)
        # chkpt = tf.train.get_checkpoint_state(ckpt_dir)
        # if chkpt and chkpt.model_checkpoint_path:
        #     print chkpt.model_checkpoint_path
        #     saver.restore(sess, chkpt.model_checkpoint_path)
        # else:
        #     print 'No checkpoint found. Expected one in:\n'
        #     print '{}\n'.format(ckpt_dir)
        #     exit(-1)

        try:
            loaded_emb = tf.train.load_variable(
                checkpoint_single_encoding,
                experiment_name + '/embedding_normalized')
            loaded_obj_bbs = tf.train.load_variable(
                checkpoint_single_encoding,
                experiment_name + '/embed_obj_bbs_var')
        except:
            loaded_emb = None
            loaded_obj_bbs = None

        if model == 'dsprites':
            codebook_multi.update_embedding_dsprites(sess, args)
        else:
            codebook_multi.update_embedding(sess,
                                            batch_size,
                                            model_path,
                                            loaded_emb=loaded_emb,
                                            loaded_obj_bbs=loaded_obj_bbs)

        print('Saving new checkoint ..')

        saver.save(sess, target_checkpoint_file, global_step=iteration)

        print('done')
예제 #14
0
def main():
    tf.disable_eager_execution()
    
    workspace_path = os.environ.get('AE_WORKSPACE_PATH')

    if workspace_path is None:
        print('Please define a workspace path:\n')
        print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
        exit(-1)

    gentle_stop = np.array((1,), dtype=np.bool)
    gentle_stop[0] = False
    def on_ctrl_c(signal, frame):
        gentle_stop[0] = True
    signal.signal(signal.SIGINT, on_ctrl_c)

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    parser.add_argument("-d", action='store_true', default=False)
    parser.add_argument("-gen", action='store_true', default=False)
    arguments = parser.parse_args()

    full_name = arguments.experiment_name.split('/')

    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''

    debug_mode = arguments.d
    generate_data = arguments.gen

    cfg_file_path = u.get_config_file_path(workspace_path, experiment_name, experiment_group)
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)
    checkpoint_file = u.get_checkpoint_basefilename(log_dir)
    ckpt_dir = u.get_checkpoint_dir(log_dir)
    train_fig_dir = u.get_train_fig_dir(log_dir)
    dataset_path = u.get_dataset_path(workspace_path)

    if not os.path.exists(cfg_file_path):
        print('Could not find config file:\n')
        print('{}\n'.format(cfg_file_path))
        exit(-1)

    if not os.path.exists(ckpt_dir):
        os.makedirs(ckpt_dir)
    if not os.path.exists(train_fig_dir):
        os.makedirs(train_fig_dir)
    if not os.path.exists(dataset_path):
        os.makedirs(dataset_path)

    args = configparser.ConfigParser()
    args.read(cfg_file_path)

    shutil.copy2(cfg_file_path, log_dir)

    with tf.variable_scope(experiment_name):
        dataset = factory.build_dataset(dataset_path, args)
        queue = factory.build_queue(dataset, args)
        encoder = factory.build_encoder(queue.x, args, is_training=True)
        decoder = factory.build_decoder(queue.y, encoder, args, is_training=True)
        ae = factory.build_ae(encoder, decoder, args)
        codebook = factory.build_codebook(encoder, dataset, args)
        train_op = factory.build_train_op(ae, args)
        saver = tf.train.Saver(save_relative_paths=True)

    num_iter = args.getint('Training', 'NUM_ITER') if not debug_mode else 100000
    save_interval = args.getint('Training', 'SAVE_INTERVAL')
    model_type = args.get('Dataset', 'MODEL')

    if model_type=='dsprites':
        dataset.get_sprite_training_images(args)
    else:
        dataset.get_training_images(dataset_path, args)
        dataset.load_bg_images(dataset_path)

    if generate_data:
        print('finished generating synthetic training data for ' + experiment_name)
        print('exiting...')
        exit()

    widgets = ['Training: ', progressbar.Percentage(),
         ' ', progressbar.Bar(),
         ' ', progressbar.Counter(), ' / %s' % num_iter,
         ' ', progressbar.ETA(), ' ']
    bar = progressbar.ProgressBar(maxval=num_iter,widgets=widgets)


    gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction=0.8)
    config = tf.ConfigProto(gpu_options=gpu_options)

    with tf.Session(config=config) as sess:

        chkpt = tf.train.get_checkpoint_state(ckpt_dir)
        if chkpt and chkpt.model_checkpoint_path:
            saver.restore(sess, chkpt.model_checkpoint_path)
        else:
            sess.run(tf.global_variables_initializer())

        merged_loss_summary = tf.summary.merge_all()
        summary_writer = tf.summary.FileWriter(ckpt_dir, sess.graph)


        if not debug_mode:
            print('Training with %s model' % args.get('Dataset','MODEL'), os.path.basename(args.get('Paths','MODEL_PATH')))
            bar.start()

        queue.start(sess)
        for i in range(ae.global_step.eval(), num_iter):
            if not debug_mode:
                sess.run(train_op)
                if i % 10 == 0:
                    loss = sess.run(merged_loss_summary)
                    summary_writer.add_summary(loss, i)

                bar.update(i)
                if (i+1) % save_interval == 0:
                    saver.save(sess, checkpoint_file, global_step=ae.global_step)

                    this_x, this_y = sess.run([queue.x, queue.y])
                    reconstr_train = sess.run(decoder.x,feed_dict={queue.x:this_x})
                    train_imgs = np.hstack(( u.tiles(this_x, 4, 4), u.tiles(reconstr_train, 4,4),u.tiles(this_y, 4, 4)))
                    cv2.imwrite(os.path.join(train_fig_dir,'training_images_%s.png' % i), train_imgs*255)
            else:

                this_x, this_y = sess.run([queue.x, queue.y])
                reconstr_train = sess.run(decoder.x,feed_dict={queue.x:this_x})
                cv2.imshow('sample batch', np.hstack(( u.tiles(this_x, 3, 3), u.tiles(reconstr_train, 3,3),u.tiles(this_y, 3, 3))) )
                k = cv2.waitKey(0)
                if k == 27:
                    break

            if gentle_stop[0]:
                break

        queue.stop(sess)
        if not debug_mode:
            bar.finish()
        if not gentle_stop[0] and not debug_mode:
            print('To create the embedding run:\n')
            print('ae_embed {}\n'.format(full_name))
예제 #15
0
def main():
    workspace_path = os.environ.get('AE_WORKSPACE_PATH')

    if workspace_path == None:
        print('Please define a workspace path:\n')
        print('export AE_WORKSPACE_PATH=/path/to/workspace\n')
        exit(-1)

    gentle_stop = np.array((1,), dtype=np.bool)
    gentle_stop[0] = False

    def on_ctrl_c(signal, frame):
        gentle_stop[0] = True

    signal.signal(signal.SIGINT, on_ctrl_c)

    parser = argparse.ArgumentParser()
    parser.add_argument("experiment_name")
    parser.add_argument('--model_path', default=False)
    parser.add_argument('--config_path', default='eval_latent_template.cfg', required=True)
    parser.add_argument("-d", action='store_true', default=False)
    parser.add_argument("-gen", action='store_true', default=False)
    parser.add_argument("-vis_emb", action='store_true', default=False)
    parser.add_argument('--at_step', default=None,  type=int, required=False)


    arguments = parser.parse_args()

    full_name = arguments.experiment_name.split('/')
    model_path = arguments.model_path
    joint = False if model_path else True
    
    experiment_name = full_name.pop()
    experiment_group = full_name.pop() if len(full_name) > 0 else ''
    
    at_step = arguments.at_step

    cfg_file_path = u.get_config_file_path(workspace_path, experiment_name, experiment_group)
    latent_cfg_file_path = u.get_eval_config_file_path(workspace_path, arguments.config_path)
    log_dir = u.get_log_dir(workspace_path, experiment_name, experiment_group)

    if not os.path.exists(cfg_file_path):
        print('Could not find config file:\n')
        print('{}\n'.format(cfg_file_path))
        exit(-1)

    args = configparser.ConfigParser(inline_comment_prefixes="#")
    args.read(cfg_file_path)

    args_latent = configparser.ConfigParser(inline_comment_prefixes="#")
    args_latent.read(latent_cfg_file_path)

    if at_step is None:
        checkpoint_file = u.get_checkpoint_basefilename(log_dir, model_path, latest=args.getint('Training', 'NUM_ITER'), joint=joint)
    else:
        checkpoint_file = u.get_checkpoint_basefilename(log_dir, model_path, latest=at_step, joint=joint)

    model_type = args.get('Dataset', 'MODEL')

    codebook, dataset = factory.build_codebook_from_name(experiment_name, experiment_group, return_dataset = True, joint=joint)
    encoder = codebook._encoder
    
    gpu_options = tf.GPUOptions(allow_growth=True, per_process_gpu_memory_fraction = 0.5)
    config = tf.ConfigProto(gpu_options=gpu_options)
    sess = tf.Session(config=config)
    saver = tf.train.Saver()
    saver.restore(sess, checkpoint_file)

    ######
    dataset_exp = args_latent.get('Data', 'dataset')
    base_path = args_latent.get('Data', 'base_path')
    split = args_latent.get('Data', 'split')
    num_obj = args_latent.getint('Data', 'num_obj')
    num_views = args_latent.getint('Data', 'num_views')
    test_class = args_latent.get('Data', 'test_class')

    # for test_class in test_classes:
    models = sorted(glob.glob(os.path.join(base_path, test_class, split, '*_normalized.off')))

    if split == 'test': or split=='train':
        if os.path.exists(os.path.dirname(models[0])):
            dataset._kw['model_path'] = models[0:num_obj]
        else:
            print((models[0], ' does not exist'))