[driver] Implement way to limit processing; Skip already processed

bob/db/hci_tagging/driver.py

@@ -16,6 +16,8 @@ import bob.io.base
 import bob.db.base
 from bob.db.base.driver import Interface as BaseInterface
 
+from . import utils
+
 DATABASE_LOCATION = '/idiap/resource/database/HCI_Tagging'
 
 
@@ -51,6 +53,8 @@ def create_meta(args):
   objects = db.objects()
   if args.selftest:
     objects = objects[:5]
+  if args.limit:
+    objects = objects[:args.limit]
 
   if args.grid_count:
     print len(objects)
@@ -67,8 +71,12 @@ def create_meta(args):
   basedir = pkg_resources.resource_filename(__name__, 'data')
 
   for obj in objects:
-    print "Creating meta data for `%s'..." % obj.make_path()
+    output = obj.make_path(basedir, '.hdf5')
+    if os.path.exists(output) and not args.force:
+      print "Skipping `%s' (meta file exists)" % obj.make_path()
+      continue
     try:
+      print "Creating meta data for `%s'..." % obj.make_path()
       max_faces = 0
       if args.selftest: max_faces = 2
       detections = obj.run_face_detector(args.directory, max_faces)
@@ -79,7 +87,6 @@ def create_meta(args):
         quality.append(det['quality'])
         bb = det['boundingbox']
         arr.append([k, bb.topleft[1], bb.topleft[0], bb.size[1], bb.size[0]])
-      output = obj.make_path(basedir, '.hdf5')
       outdir = os.path.dirname(output)
       if not os.path.exists(outdir): os.makedirs(outdir)
       h5 = bob.io.base.HDF5File(output, 'w')
@@ -97,6 +104,59 @@ def create_meta(args):
   return 0
 
 
+def debug(args):
+  """Debugs the face detection and heart-rate estimation"""
+
+
+  from . import Database
+  db = Database()
+
+  objects = db.objects()
+  if args.selftest:
+    objects = objects[:5]
+  if args.limit:
+    objects = objects[:args.limit]
+
+  if args.grid_count:
+    print len(objects)
+    sys.exit(0)
+
+  # if we are on a grid environment, just find what I have to process.
+  if os.environ.has_key('SGE_TASK_ID'):
+    pos = int(os.environ['SGE_TASK_ID']) - 1
+    if pos >= len(objects):
+      raise RuntimeError, "Grid request for job %d on a setup with %d jobs" % \
+          (pos, len(objects))
+    objects = [objects[pos]]
+
+  basedir = 'debug'
+
+  for obj in objects:
+    print "Creating debug data for `%s'..." % obj.make_path()
+    try:
+
+      detections = obj.load_face_detections()
+      if detections is None:
+        detections = obj.run_face_detector(args.directory)
+        detections = dict([(k, (v.topleft[1], v.topleft[0], v.size[1], v.size[0])) for k,v in detections.items()])
+      # save annotated video file
+      output = obj.make_path(args.output_directory, '.avi')
+      print "Annotating video `%s'" % output
+      utils.annotate_video(obj.load_video(args.directory), detections, output)
+
+      print "Annotating heart-rate `%s'" % output
+      output = obj.make_path(args.output_directory, '.pdf')
+      utils.explain_heartrate(obj, args.directory, output)
+
+    finally:
+      if args.selftest:
+        if os.path.exists(args.output_directory):
+          import shutil
+          shutil.rmtree(args.output_directory)
+
+  return 0
+
+
 def checkfiles(args):
   """Checks the existence of the files based on your criteria"""
 
@@ -177,5 +237,17 @@ class Interface(BaseInterface):
     meta_parser = subparsers.add_parser('mkmeta', help=create_meta.__doc__)
     meta_parser.add_argument('-d', '--directory', dest="directory", default=DATABASE_LOCATION, help="This path points to the location where the database raw files are installed (defaults to '%(default)s')")
     meta_parser.add_argument('--grid-count', dest="grid_count", default=False, action='store_true', help=SUPPRESS)
+    meta_parser.add_argument('--force', dest="force", default=False, action='store_true', help='If set, will overwrite existing meta files if they exist. Otherwise, just run on unexisting data')
+    meta_parser.add_argument('--limit', dest="limit", default=0, type=int, help="Limits the number of objects to treat (defaults to '%(default)')")
     meta_parser.add_argument('--self-test', dest="selftest", default=False, action='store_true', help=SUPPRESS)
     meta_parser.set_defaults(func=create_meta) #action
+
+    # debug
+    debug_message = debug.__doc__
+    debug_parser = subparsers.add_parser('debug', help=debug.__doc__)
+    debug_parser.add_argument('-d', '--directory', dest="directory", default=DATABASE_LOCATION, help="This path points to the location where the database raw files are installed (defaults to '%(default)s')")
+    debug_parser.add_argument('-o', '--output-directory', dest="output_directory", default='debug', help="This path points to the location where the debugging results will be stored (defaults to '%(default)s')")
+    debug_parser.add_argument('--grid-count', dest="grid_count", default=False, action='store_true', help=SUPPRESS)
+    debug_parser.add_argument('--limit', dest="limit", default=0, type=int, help="Limits the number of objects to treat (defaults to '%(default)')")
+    debug_parser.add_argument('--self-test', dest="selftest", default=False, action='store_true', help=SUPPRESS)
+    debug_parser.set_defaults(func=debug) #action