eval_modelnet40.py

import tensorflow as tf
import numpy as np
import modelnet_provider as modelnet
import importlib
import util
import sys
import os
import datetime 
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
# The GPU id to use, usually either "0" or "1"
os.environ["CUDA_VISIBLE_DEVICES"]="0"
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
def eval(args):


    num_points = 2048   # ModelNet40
    num_channels = 3

    model_name = args["model"]
    model = importlib.import_module(model_name)
    print('Using model ', model_name)

    sort_cloud = args['sort_cloud']
    sort_method= args['sort_method']
    if len(sys.argv) > 1:
        epoch = int(sys.argv[1])
    else:
        epoch = 0

    batch_size = args["batch_size"]
    d = modelnet.DataConsumer(file='data/modelnet40_ply_hdf5_2048/test_files.txt', batch_size=batch_size, num_points=num_points, num_channels=num_channels, test=True, sort_cloud=sort_cloud, sort_method=sort_method)
    num_class = 40
    SHAPE_NAMES = [line.rstrip() for line in \
    open(os.path.join(BASE_DIR, 'data/modelnet40_ply_hdf5_2048/shape_names.txt'))] 
  
#def log_string(out_str):
#    LOG_FOUT.write(out_str+'\n')
#    LOG_FOUT.flush()
#    print(out_str)
    
    print('Categories: ', num_class)
    print('Sort cloud: ', sort_cloud)
    
    use_gpu = args["use_gpu"]
    if use_gpu:
        str_device = '/gpu:1'
    else:
        str_device = '/cpu:0'
    print('Using device', str_device)
    
    with tf.Graph().as_default(), tf.device(str_device):

        if not use_gpu:
            config = tf.ConfigProto(
                device_count = {'GPU': 1},
                log_device_placement = False
            )
        else:
            config = tf.ConfigProto(
                allow_soft_placement = True,
                log_device_placement = False
            )
            
        with tf.Session(config=config) as session:
            network = model.PointConvNet(num_class)

            batch_points_placeholder = tf.placeholder(tf.float32, shape=(batch_size, num_points, 3))
            batch_input_placeholder = tf.placeholder(tf.float32, shape=(batch_size, num_points, num_channels))
            batch_label_placeholder = tf.placeholder(tf.int32, shape=(batch_size))

            with tf.variable_scope("pointcnn") as scope:
                batch_prediction_placeholder = network.model(batch_points_placeholder, batch_input_placeholder, is_training=False)

            saver = tf.train.Saver(max_to_keep=None)
            file = "./16-layers/{}_snapshot_{}.tf".format(model_name, epoch)
            print('Testing with model ', file)
            saver.restore(session, file)

            total_correct = 0
            total_seen = 0
            total_correct_class = np.zeros((num_class))
            total_seen_class = np.zeros((num_class))

            num_batches = 0

            while True:
                points, points_data, gt_label = d.get_batch_point_cloud()

                feed_dict = {
                    batch_points_placeholder : points,
                    batch_input_placeholder : points_data,
                    batch_label_placeholder : gt_label
                }
                [pred_val] = session.run([batch_prediction_placeholder], feed_dict=feed_dict)

                pred_label = np.argmax(pred_val, axis=1)

                correct = np.sum(pred_label == gt_label)
                total_correct += correct
                total_seen += batch_size

                for i in range(batch_size):
                    l = gt_label[i]
                    total_seen_class[l] += 1
                    total_correct_class[l] += (pred_label[i] == l)

                num_batches += 1
                print('Current accuracy  : %f' % (correct / float(batch_size)))
                
                if not d.has_next_batch():
                    break
                d.next_batch()
                

            mean_accuracy = (total_correct / float(total_seen))
            mean_class_accuracy = (np.mean(total_correct_class / total_seen_class))
            print('Mean accuracy      : %f' % mean_accuracy)
            print('Avg class accuracy : %f' % mean_class_accuracy)
            
            class_accuracies = np.array(total_correct_class)/np.array(total_seen_class,dtype=np.float)
            for i, name in enumerate(SHAPE_NAMES):
                 print('%10s:\t%0.3f' % (name, class_accuracies[i]))

            with open("test_file_{}.txt".format(model_name), 'a') as f:
                f.write(str(datetime.datetime.now().strftime("%c")) + ', ' + \
                        str(epoch) + ', ' + str(mean_accuracy) + ', ' + str(mean_class_accuracy) + '\n')


if __name__ == "__main__":
    args = util.parse_arguments("param.json")
    eval(args)