Training using multiple datasets

lib/fast_rcnn/config.py

@@ -149,7 +149,7 @@
 __C.TEST.BBOX_REG = True
 
 # Propose boxes
-__C.TEST.HAS_RPN = False
+__C.TEST.HAS_RPN = True
 
 # Test using these proposals
 __C.TEST.PROPOSAL_METHOD = 'selective_search'

lib/fast_rcnn/train.py

@@ -95,6 +95,20 @@ def train_model(self, max_iters):
         last_snapshot_iter = -1
         timer = Timer()
         model_paths = []
+
+	# Check if the max_iter model exists 
+        infix = ('_' + cfg.TRAIN.SNAPSHOT_INFIX
+                 if cfg.TRAIN.SNAPSHOT_INFIX != '' else '')
+        filename = (self.solver_param.snapshot_prefix + infix +
+                    '_iter_{:d}'.format(max_iters) + '.caffemodel')
+        filename = os.path.join(self.output_dir, filename)
+
+	# If exists append its path and return
+	if os.path.isfile(filename):
+            print "Max iter model exists."
+	    model_paths.append(filename)
+	    return model_paths
+
         while self.solver.iter < max_iters:
             # Make one SGD update
             timer.tic()

tools/train_faster_rcnn_alt_opt.py

@@ -18,6 +18,7 @@
 from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list, get_output_dir
 from datasets.factory import get_imdb
 from rpn.generate import imdb_proposals
+import datasets.imdb
 import argparse
 import pprint
 import numpy as np
@@ -57,14 +58,33 @@ def parse_args():
     args = parser.parse_args()
     return args
 
-def get_roidb(imdb_name, rpn_file=None):
-    imdb = get_imdb(imdb_name)
-    print 'Loaded dataset `{:s}` for training'.format(imdb.name)
-    imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
-    print 'Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD)
-    if rpn_file is not None:
-        imdb.config['rpn_file'] = rpn_file
-    roidb = get_training_roidb(imdb)
+def combined_roidb(imdb_names, rpn_files=None):
+    def get_roidb(imdb_name, rpn_file=None):
+        imdb = get_imdb(imdb_name)
+        print 'Loaded dataset `{:s}` for training'.format(imdb.name)
+        imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
+        print 'Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD)
+        if rpn_file is not None:
+            imdb.config['rpn_file'] = rpn_file
+        roidb = get_training_roidb(imdb)
+        return roidb
+
+    if '+' in imdb_names:  # Multiple databases e.g. voc_2007_trainval+voc_2012_trainval
+        imdbs_list = imdb_names.split('+')
+        imdb = datasets.imdb.imdb(imdb_names)
+        if rpn_files is None:
+            rpn_files = [None] * len(imdbs_list)
+
+        assert(len(imdbs_list) == len(rpn_files))
+        for i, imdb_name in enumerate(imdbs_list):
+            if i == 0:
+                roidb = get_roidb(imdb_name, rpn_files[i])
+            else:
+                roidb.extend( get_roidb(imdb_name, rpn_files[i]) )
+    else:   # Single database case e.g. voc_2007_trainval
+        roidb = get_roidb(imdb_names, rpn_files)
+        imdb= get_imdb(imdb_names)
+
     return roidb, imdb
 
 def get_solvers(net_name):
@@ -119,7 +139,7 @@ def train_rpn(queue=None, imdb_name=None, init_model=None, solver=None,
     import caffe
     _init_caffe(cfg)
 
-    roidb, imdb = get_roidb(imdb_name)
+    roidb, imdb = combined_roidb(imdb_name)
     print 'roidb len: {}'.format(len(roidb))
     output_dir = get_output_dir(imdb)
     print 'Output will be saved to `{:s}`'.format(output_dir)
@@ -148,26 +168,47 @@ def rpn_generate(queue=None, imdb_name=None, rpn_model_path=None, cfg=None,
     import caffe
     _init_caffe(cfg)
 
-    # NOTE: the matlab implementation computes proposals on flipped images, too.
-    # We compute them on the image once and then flip the already computed
-    # proposals. This might cause a minor loss in mAP (less proposal jittering).
-    imdb = get_imdb(imdb_name)
-    print 'Loaded dataset `{:s}` for proposal generation'.format(imdb.name)
-
-    # Load RPN and configure output directory
-    rpn_net = caffe.Net(rpn_test_prototxt, rpn_model_path, caffe.TEST)
-    output_dir = get_output_dir(imdb)
-    print 'Output will be saved to `{:s}`'.format(output_dir)
-    # Generate proposals on the imdb
-    rpn_proposals = imdb_proposals(rpn_net, imdb)
-    # Write proposals to disk and send the proposal file path through the
+    if '+' in imdb_name:
+        imdbs_list = imdb_name.split('+')
+        imdb = datasets.imdb.imdb(imdb_name)
+        output_dir = get_output_dir(imdb, None)
+        print 'Output will be saved to `{:s}`'.format(output_dir)
+    else:
+        imdbs_list = [imdb_name]
+        output_dir = None   # Gets set later for single database case
+    rpn_proposals_path = [None] * len(imdbs_list)
+    for i, imdb_name in enumerate(imdbs_list): 
+        # NOTE: the matlab implementation computes proposals on flipped images, too.
+        # We compute them on the image once and then flip the already computed
+        # proposals. This might cause a minor loss in mAP (less proposal jittering).
+        imdb = get_imdb(imdb_name)
+        print 'Loaded dataset `{:s}` for proposal generation'.format(imdb.name)
+        if output_dir is None:
+            output_dir = get_output_dir(imdb, None)
+            print 'Output will be saved to `{:s}`'.format(output_dir)
+        # Load RPN and configure output directory
+        rpn_net = caffe.Net(rpn_test_prototxt, rpn_model_path, caffe.TEST)
+
+        rpn_net_name = os.path.splitext(os.path.basename(rpn_model_path))[0]
+        rpn_proposals_path[i] = os.path.join(
+            output_dir, rpn_net_name + '_' + imdb_name + '_proposals.pkl')
+
+        # Check if rpn proposals have already been computed
+        # If so, don't recompute 
+        if not os.path.isfile(rpn_proposals_path[i]):
+            # Generate proposals on the imdb
+            rpn_proposals = imdb_proposals(rpn_net, imdb)
+
+            # Write proposals to disk
+            with open(rpn_proposals_path[i], 'wb') as f:
+                cPickle.dump(rpn_proposals, f, cPickle.HIGHEST_PROTOCOL)
+            print 'Wrote RPN proposals to {}'.format(rpn_proposals_path[i])
+        else:
+            print "Proposals for " + imdb_name + " exist already."
+    if len(rpn_proposals_path) == 1:
+        rpn_proposals_path = rpn_proposals_path[0]
+    # Send the proposal file path through the
     # multiprocessing queue
-    rpn_net_name = os.path.splitext(os.path.basename(rpn_model_path))[0]
-    rpn_proposals_path = os.path.join(
-        output_dir, rpn_net_name + '_proposals.pkl')
-    with open(rpn_proposals_path, 'wb') as f:
-        cPickle.dump(rpn_proposals, f, cPickle.HIGHEST_PROTOCOL)
-    print 'Wrote RPN proposals to {}'.format(rpn_proposals_path)
     queue.put({'proposal_path': rpn_proposals_path})
 
 def train_fast_rcnn(queue=None, imdb_name=None, init_model=None, solver=None,
@@ -186,8 +227,8 @@ def train_fast_rcnn(queue=None, imdb_name=None, init_model=None, solver=None,
     import caffe
     _init_caffe(cfg)
 
-    roidb, imdb = get_roidb(imdb_name, rpn_file=rpn_file)
-    output_dir = get_output_dir(imdb)
+    roidb, imdb = combined_roidb(imdb_name, rpn_files=rpn_file)
+    output_dir = get_output_dir(imdb, None)
     print 'Output will be saved to `{:s}`'.format(output_dir)
     # Train Fast R-CNN
     model_paths = train_net(solver, roidb, output_dir,
@@ -240,7 +281,7 @@ def train_fast_rcnn(queue=None, imdb_name=None, init_model=None, solver=None,
     p.start()
     rpn_stage1_out = mp_queue.get()
     p.join()
-
+    
     print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'
     print 'Stage 1 RPN, generate proposals'
     print '~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~'