eval_net.py

import argparse
import logging
import time

import torch.utils.data

from models.common import post_process_output
from utils.dataset_processing import evaluation, grasp
from utils.data import get_dataset
from models.lgpnet import LGPNet


logging.basicConfig(level=logging.INFO)


def parse_args():
    parser = argparse.ArgumentParser(description='Evaluate LGPNet')

    # Network
    parser.add_argument('--network', type=str, help='Path to saved network to evaluate')
    parser.add_argument('--layers', type=int, default=50, help='Layers number')

    # Dataset & Data & Training
    parser.add_argument('--dataset', type=str, help='Dataset Name ("cornell" or "jacquard")')
    parser.add_argument('--dataset-path', type=str, help='Path to dataset')
    parser.add_argument('--use-depth', type=int, default=1, help='Use Depth image for evaluation (1/0)')
    parser.add_argument('--use-rgb', type=int, default=0, help='Use RGB image for evaluation (0/1)')
    parser.add_argument('--augment', action='store_true', help='Whether data augmentation should be applied')
    parser.add_argument('--ds-rotate', type=float, default=0.0,
                        help='Shift the start point of the dataset to use a different test/train split')
    parser.add_argument('--start-split', type=float, default=0.0, help='The start of the split for train dataset')
    parser.add_argument('--end-split', type=float, default=0.2, help='The end of the split for train dataset')
    parser.add_argument('--ds-shuffle', action='store_true', help='Turning OW to IW')
    parser.add_argument('--num-workers', type=int, default=4, help='Dataset workers')

    parser.add_argument('--n-grasps', type=int, default=1, help='Number of grasps to consider per image')
    parser.add_argument('--iou-eval', action='store_true', help='Compute success based on IoU metric.')
    parser.add_argument('--jacquard-output', action='store_true', help='Jacquard-dataset style output')
    parser.add_argument('--vis', action='store_true', help='Visualise the network output')
    parser.add_argument('--gpu-idx', type=str, default='0', help='GPU index')

    args = parser.parse_args()

    if args.jacquard_output and args.dataset != 'jacquard':
        raise ValueError('--jacquard-output can only be used with the --dataset jacquard option.')
    if args.jacquard_output and args.augment:
        raise ValueError('--jacquard-output can not be used with data augmentation.')

    return args


if __name__ == '__main__':
    args = parse_args()
    device = torch.device("cuda:{}".format(args.gpu_idx))
    net = torch.load(args.network)
    net = net.to(device)

    # Load Dataset
    logging.info('Loading {} Dataset...'.format(args.dataset.title()))
    Dataset = get_dataset(args.dataset)
    test_dataset = Dataset(
        args.dataset_path,
        start_split=args.start_split,
        end_split=args.end_split,
        ds_rotate=args.ds_rotate,
        is_val=True,
        ds_shuffle=args.ds_shuffle,
        random_rotate=args.augment,
        random_zoom=args.augment,
        include_depth=args.use_depth,
        include_rgb=args.use_rgb
    )

    split_dataset_method = 'OW'
    if args.ds_shuffle:
        split_dataset_method = 'IW'
    logging.info('Split dataset with {} method'.format(split_dataset_method))

    test_data = torch.utils.data.DataLoader(
        test_dataset,
        batch_size=1,
        shuffle=False,
        num_workers=args.num_workers
    )
    logging.info('Done')
    
    results = {'correct': 0, 'failed': 0}

    if args.jacquard_output:
        jo_fn = args.network + '_jacquard_output.txt'
        with open(jo_fn, 'w') as f:
            pass
    
    logging.info('Evaluating model...')
    net.eval()
    start_time = time.time()
    failed_idx = []
    with torch.no_grad():
        for idx, (x, y, didx, rot, zoom) in enumerate(test_data):
            xc = x.to(device)
            yc = [yi.to(device) for yi in y]
            lossd = net.compute_loss(xc, yc)

            q_img, ang_img, width_img = post_process_output(lossd['pred']['pos'], lossd['pred']['cos'],
                                                        lossd['pred']['sin'], lossd['pred']['width'])

            if args.iou_eval:
                s = evaluation.calculate_iou_match(q_img, ang_img, test_data.dataset.get_gtbb(didx, rot, zoom),
                                                   no_grasps=args.n_grasps,
                                                   grasp_width=width_img,
                                                   )
                if s:
                    results['correct'] += 1
                else:
                    failed_idx.append(idx)
                    results['failed'] += 1

            if args.jacquard_output:
                grasps = grasp.detect_grasps(q_img, ang_img, width_img=width_img, no_grasps=1)
                with open(jo_fn, 'a') as f:
                    for g in grasps:
                        f.write(test_data.dataset.get_jname(didx) + '\n')
                        f.write(g.to_jacquard(scale=1024 / 300) + '\n')

            if args.vis:
                evaluation.save_output(test_data.dataset.get_rgb(didx, rot, zoom, normalise=False),
                                       test_data.dataset.get_depth(didx, rot, zoom), q_img,
                                       ang_img, no_grasps=args.n_grasps, grasp_width_img=width_img, 
                                       idx_fig=idx, dataset=args.dataset, args=args)

    avg_time = (time.time() - start_time) / len(test_data)
    logging.info('Average evaluation time per image: {}ms'.format(avg_time * 1000))

    if args.iou_eval:
        logging.info('IOU Results: %d/%d = %f' % (results['correct'],
                              results['correct'] + results['failed'],
                              results['correct'] / (results['correct'] + results['failed'])))

    if args.jacquard_output:
        logging.info('Jacquard output saved to {}'.format(jo_fn))
    
    print(failed_idx)