Merge pull request #12 from skumra/release/0.2.0

Release/0.2.0

cleanup.sh

@@ -0,0 +1,3 @@
+#!/usr/bin/env bash
+
+find logs/ -maxdepth 1 -type d | grep -v ^\\.$ | xargs -n 1 du -s | while read size name ; do if [ $size -le 10485 ] ; then echo rm -rf $name ; fi done

evaluate.py

@@ -1,8 +1,11 @@
 import argparse
 import logging
+import time
 
+import numpy as np
 import torch.utils.data
 
+from hardware.device import get_device
 from inference.post_process import post_process_output
 from utils.data import get_dataset
 from utils.dataset_processing import evaluation, grasp
@@ -12,26 +15,49 @@
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description='Evaluate network')
+    parser = argparse.ArgumentParser(description='Evaluate networks')
 
     # Network
-    parser.add_argument('--network', type=str, help='Path to saved network to evaluate')
-
-    # Dataset & Data & Training
-    parser.add_argument('--dataset', type=str, help='Dataset Name ("cornell" or "jaquard")')
-    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=1, help='Use RGB image for evaluation (1/0)')
-    parser.add_argument('--augment', action='store_true', help='Whether data augmentation should be applied')
-    parser.add_argument('--split', type=float, default=0.9, help='Fraction of data for training (remainder is validation)')
+    parser.add_argument('--network', metavar='N', type=str, nargs='+',
+                        help='Path to saved networks to evaluate')
+
+    # Dataset
+    parser.add_argument('--dataset', type=str,
+                        help='Dataset Name ("cornell" or "jaquard")')
+    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=1,
+                        help='Use RGB image for evaluation (1/0)')
+    parser.add_argument('--augment', action='store_true',
+                        help='Whether data augmentation should be applied')
+    parser.add_argument('--split', type=float, default=0.9,
+                        help='Fraction of data for training (remainder is validation)')
+    parser.add_argument('--ds-shuffle', action='store_true', default=False,
+                        help='Shuffle the dataset')
     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('--num-workers', type=int, default=8, 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('--num-workers', type=int, default=8,
+                        help='Dataset workers')
+
+    # Evaluation
+    parser.add_argument('--n-grasps', type=int, default=1,
+                        help='Number of grasps to consider per image')
+    parser.add_argument('--iou-threshold', type=float, default=0.25,
+                        help='Threshold for IOU matching')
+    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')
+
+    # Misc.
+    parser.add_argument('--vis', action='store_true',
+                        help='Visualise the network output')
+    parser.add_argument('--cpu', dest='force_cpu', action='store_true', default=False,
+                        help='Force code to run in CPU mode')
+    parser.add_argument('--random-seed', type=int, default=123,
+                        help='Random seed for numpy')
 
     args = parser.parse_args()
 
@@ -46,72 +72,98 @@ def parse_args():
 if __name__ == '__main__':
     args = parse_args()
 
-    # Load Network
-    net = torch.load(args.network)
-    device = torch.device("cuda:0")
+    # Get the compute device
+    device = get_device(args.force_cpu)
 
     # Load Dataset
     logging.info('Loading {} Dataset...'.format(args.dataset.title()))
     Dataset = get_dataset(args.dataset)
-    test_dataset = Dataset(args.dataset_path, start=args.split, end=1.0, ds_rotate=args.ds_rotate,
-                           random_rotate=args.augment, random_zoom=args.augment,
-                           include_depth=args.use_depth, include_rgb=args.use_rgb)
+    test_dataset = Dataset(args.dataset_path,
+                           ds_rotate=args.ds_rotate,
+                           random_rotate=args.augment,
+                           random_zoom=args.augment,
+                           include_depth=args.use_depth,
+                           include_rgb=args.use_rgb)
+
+    indices = list(range(test_dataset.length))
+    split = int(np.floor(args.split * test_dataset.length))
+    if args.ds_shuffle:
+        np.random.seed(args.random_seed)
+        np.random.shuffle(indices)
+    val_indices = indices[split:]
+    val_sampler = torch.utils.data.sampler.SubsetRandomSampler(val_indices)
+    logging.info('Validation size: {}'.format(len(val_indices)))
+
     test_data = torch.utils.data.DataLoader(
         test_dataset,
         batch_size=1,
-        shuffle=False,
-        num_workers=args.num_workers
+        num_workers=args.num_workers,
+        sampler=val_sampler
     )
     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
-
-    with torch.no_grad():
-        for idx, (x, y, didx, rot, zoom) in enumerate(test_data):
-            logging.info('Processing {}/{}'.format(idx+1, len(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:
-                    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:
-                save_results(
-                    rgb_img=test_data.dataset.get_rgb(didx, rot, zoom, normalise=False),
-                    depth_img=test_data.dataset.get_depth(didx, rot, zoom),
-                    grasp_q_img=q_img,
-                    grasp_angle_img=ang_img,
-                    no_grasps=args.n_grasps,
-                    grasp_width_img=width_img
-                )
-
-    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))
+    for network in args.network:
+        logging.info('\nEvaluating model {}'.format(network))
+
+        # Load Network
+        net = torch.load(network)
+
+        results = {'correct': 0, 'failed': 0}
+
+        if args.jacquard_output:
+            jo_fn = network + '_jacquard_output.txt'
+            with open(jo_fn, 'w') as f:
+                pass
+
+        start_time = time.time()
+
+        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,
+                                                       threshold=args.iou_threshold
+                                                       )
+                    if s:
+                        results['correct'] += 1
+                    else:
+                        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:
+                    save_results(
+                        rgb_img=test_data.dataset.get_rgb(didx, rot, zoom, normalise=False),
+                        depth_img=test_data.dataset.get_depth(didx, rot, zoom),
+                        grasp_q_img=q_img,
+                        grasp_angle_img=ang_img,
+                        no_grasps=args.n_grasps,
+                        grasp_width_img=width_img
+                    )
+
+        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))
+
+        del net
+        torch.cuda.empty_cache()

hardware/calibrate_camera.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import time
 
@@ -21,6 +22,7 @@ def __init__(self,
         self.calib_grid_step = calib_grid_step
         self.checkerboard_offset_from_tool = checkerboard_offset_from_tool
 
+        # Cols: min max, Rows: x y z (define workspace limits in robot coordinates)
         self.workspace_limits = workspace_limits
 
         self.camera = RealSenseCamera(device_id=cam_id)
@@ -108,21 +110,21 @@ def _generate_grid(self):
     def run(self):
         # Connect to camera
         self.camera.connect()
-        print(self.camera.intrinsics)
+        logging.debug(self.camera.intrinsics)
 
-        print('Collecting data...')
+        logging.info('Collecting data...')
 
         calib_grid_pts = self._generate_grid()
 
-        print('Total grid points: ', calib_grid_pts.shape[0])
+        logging.info('Total grid points: ', calib_grid_pts.shape[0])
 
         for tool_position in calib_grid_pts:
-            print('Requesting move to tool position: ', tool_position)
+            logging.info('Requesting move to tool position: ', tool_position)
             np.save(self.tool_position, tool_position)
             np.save(self.move_completed, 0)
             while not np.load(self.move_completed):
                 time.sleep(0.1)
-            #wait for robot to be stable    
+            # Wait for robot to be stable
             time.sleep(2)
 
             # Find checkerboard center
@@ -162,57 +164,23 @@ def run(self):
                 cv2.imshow('Calibration', vis)
                 cv2.waitKey(10)
             else:
-                print('Checker board not found')
+                logging.info('Checker board not found')
 
         self.measured_pts = np.asarray(self.measured_pts)
         self.observed_pts = np.asarray(self.observed_pts)
         self.observed_pix = np.asarray(self.observed_pix)
 
         # Optimize z scale w.r.t. rigid transform error
-        print('Calibrating...')
+        logging.info('Calibrating...')
         z_scale_init = 1
         optim_result = optimize.minimize(self._get_rigid_transform_error, np.asarray(z_scale_init), method='Nelder-Mead')
         camera_depth_offset = optim_result.x
 
         # Save camera optimized offset and camera pose
-        print('Saving...')
+        logging.info('Saving...')
         np.savetxt('saved_data/camera_depth_scale.txt', camera_depth_offset, delimiter=' ')
         rmse = self._get_rigid_transform_error(camera_depth_offset)
-        print('RMSE: ', rmse)
+        logging.info('RMSE: ', rmse)
         camera_pose = np.linalg.inv(self.world2camera)
         np.savetxt('saved_data/camera_pose.txt', camera_pose, delimiter=' ')
-        print('Done.')
-
-        # ---------------------------------------------
-
-        # np.savetxt('saved_data/measured_pts.txt', self.measured_pts, delimiter=' ')
-        # np.savetxt('saved_data/observed_pts.txt', self.observed_pts, delimiter=' ')
-        # np.savetxt('saved_data/observed_pix.txt', self.observed_pix, delimiter=' ')
-        # measured_pts = np.loadtxt('saved_data/measured_pts.txt', delimiter=' ')
-        # observed_pts = np.loadtxt('saved_data/observed_pts.txt', delimiter=' ')
-        # observed_pix = np.loadtxt('saved_data/observed_pix.txt', delimiter=' ')
-
-        # fig = plt.figure()
-        # ax = fig.add_subplot(111, projection='3d')
-        # ax.scatter(measured_pts[:,0],measured_pts[:,1],measured_pts[:,2], c='blue')
-
-        # print(camera_depth_offset)
-        # R, t = self._get_rigid_transform(np.asarray(measured_pts), np.asarray(observed_pts))
-        # t.shape = (3,1)
-        # camera_pose = np.concatenate((np.concatenate((R, t), axis=1),np.array([[0, 0, 0, 1]])), axis=0)
-        # camera2robot = np.linalg.inv(camera_pose)
-        # t_observed_pts = np.transpose(np.dot(camera2robot[0:3,0:3],np.transpose(observed_pts)) + np.tile(camera2robot[0:3,3:],(1,observed_pts.shape[0])))
-
-        # ax.scatter(t_observed_pts[:,0],t_observed_pts[:,1],t_observed_pts[:,2], c='red')
-
-        # new_observed_pts = observed_pts.copy()
-        # new_observed_pts[:,2] = new_observed_pts[:,2] * camera_depth_offset[0]
-        # R, t = self._get_rigid_transform(np.asarray(measured_pts), np.asarray(new_observed_pts))
-        # t.shape = (3,1)
-        # camera_pose = np.concatenate((np.concatenate((R, t), axis=1),np.array([[0, 0, 0, 1]])), axis=0)
-        # camera2robot = np.linalg.inv(camera_pose)
-        # t_new_observed_pts = np.transpose(np.dot(camera2robot[0:3,0:3],np.transpose(new_observed_pts)) + np.tile(camera2robot[0:3,3:],(1,new_observed_pts.shape[0])))
-
-        # ax.scatter(t_new_observed_pts[:,0],t_new_observed_pts[:,1],t_new_observed_pts[:,2], c='green')
-
-        # plt.show()
+        logging.info('Done.')

hardware/camera.py

@@ -13,7 +13,6 @@ def __init__(self,
                  width=640,
                  height=480,
                  fps=6):
-
         self.device_id = device_id
         self.width = width
         self.height = height
@@ -53,7 +52,7 @@ def get_image_bundle(self):
 
         depth_image = np.expand_dims(depth_image, axis=2)
 
-        return{
+        return {
             'rgb': color_image,
             'aligned_depth': depth_image,
         }
@@ -79,4 +78,3 @@ def plot_image_bundle(self):
     cam.connect()
     while True:
         cam.plot_image_bundle()
-

hardware/device.py

@@ -0,0 +1,19 @@
+import logging
+
+import torch
+
+logging.basicConfig(level=logging.INFO)
+
+
+def get_device(force_cpu):
+    # Check if CUDA can be used
+    if torch.cuda.is_available() and not force_cpu:
+        logging.info("CUDA detected. Running with GPU acceleration.")
+        device = torch.device("cuda")
+    elif force_cpu:
+        logging.info("CUDA detected, but overriding with option '--cpu'. Running with only CPU.")
+        device = torch.device("cpu")
+    else:
+        logging.info("CUDA is *NOT* detected. Running with only CPU.")
+        device = torch.device("cpu")
+    return device