Merge pull request #11 from RobotLocomotion/pf-save-depth-no-pyc

Save depth images (cleaned up)

Author: peteflorence

.gitignore

@@ -1 +1,4 @@
 data/
+*pyc
+automation/scp_config.yaml
+src/ply

modules/labelfusion/imagecapture.py

@@ -17,7 +17,7 @@
 class ImageCapture(object):
 
     def __init__(self, imageManager, fileSaveLocation,
-                 cameraName = "OPENNI_FRAME_LEFT", setupCallback=False):
+                 cameraName="OPENNI_FRAME_LEFT", setupCallback=False):
         self.imageManager = imageManager
         self.fileSaveLocation = fileSaveLocation
         self.cameraName = cameraName
@@ -31,27 +31,28 @@ def setupCallback(self):
         lcmUtils.addSubscriber("OPENNI_FRAME", lcmbotcore.images_t(),
                                self.onImageMessage)
 
-    def saveImage(self, extension="rgb.png"):
+    def saveImage(self, saveUtime=True, extension="rgb.png"):
         # construct filename where this image will be saved
         baseFilename = utils.convertImageIDToPaddedString(self.counter) + "_"
         baseFilename = os.path.join(self.fileSaveLocation, baseFilename)
-        rgbFilename = baseFilename + extension
-        utimeFilename = baseFilename + "utime.txt"
+        imageFilename = baseFilename + extension
         self.counter += 1
 
         # get the image and it's utime
         self.imageManager.updateImages()
         image = self.imageManager.getImage(self.cameraName)
-        utime = self.imageManager.getUtime(self.cameraName)
         image = filterUtils.flipImage(image)
-        print 'writing:', rgbFilename
-        ImageCapture.writeImage(image, rgbFilename)
-
-        # now save the utime
-        print 'writing:', utimeFilename
-        text_file = open(utimeFilename, "w")
-        text_file.write(str(utime))
-        text_file.close()
+        print 'writing:', imageFilename
+        ImageCapture.writeImage(image, imageFilename)
+
+        if saveUtime:
+            utime = self.imageManager.getUtime(self.cameraName)
+            utimeFilename = baseFilename + "utime.txt"
+            # now save the utime
+            print 'writing:', utimeFilename
+            text_file = open(utimeFilename, "w")
+            text_file.write(str(utime))
+            text_file.close()
 
     @staticmethod
     def writeImage(image, filename):
@@ -80,8 +81,7 @@ def onImageMessage(self, msg):
 
 
     @staticmethod
-    def readFromLogFile(lcmLogFilename, fileSaveLocation, channelName="OPENNI_FRAME",
-                        cameraName="OPENNI_FRAME_LEFT"):
+    def readFromLogFile(lcmLogFilename, fileSaveLocation, channelName="OPENNI_FRAME", cameraName="OPENNI_FRAME_LEFT", saveDepth=False):
         """
         Reads from lcmlog located at filename. Goes through each
         images_t() message on OPENNI_FRAME channel and saves it
@@ -94,26 +94,37 @@ def readFromLogFile(lcmLogFilename, fileSaveLocation, channelName="OPENNI_FRAME"
         if not os.path.isdir(fileSaveLocation):
             os.makedirs(fileSaveLocation)
 
-        # first construct imageManager object
+
+        # construct imageManager object
         imageManager = cameraview.ImageManager()
-        imageManager.queue.addCameraStream(channelName, cameraName, lcmbotcore.images_t.LEFT)
+        if saveDepth:
+            imageManager.queue.addCameraStream(channelName, cameraName, lcmbotcore.images_t.DEPTH_MM_ZIPPED)
+        else:
+            imageManager.queue.addCameraStream(channelName, cameraName, lcmbotcore.images_t.LEFT)
         imageManager.addImage(cameraName)
 
         # open the lcm log
         imageManager.queue.openLCMFile(lcmLogFilename)
 
         imageCapture = ImageCapture(imageManager, fileSaveLocation,
-                 cameraName = "OPENNI_FRAME_LEFT", setupCallback=False)
+                 cameraName=cameraName, setupCallback=False)
 
         while imageManager.queue.readNextImagesMessage():
-            imageCapture.saveImage()
+            if saveDepth:
+                imageCapture.saveImage(saveUtime=False, extension="depth.png")
+            else:
+                imageCapture.saveImage(extension="rgb.png")
 
         print "reached end of lcm log"
         return
 
-def captureImages(logFolder):
-    paths = utils.getFilenames(logFolder)
-    ImageCapture.readFromLogFile(paths['lcmlog'], paths['images'])
+def captureImages(logFolder, saveDepth):
+    corlPaths = utils.getFilenames(logFolder)
+    
+    if saveDepth:
+        ImageCapture.readFromLogFile(corlPaths['lcmlog'], corlPaths['images'], cameraName="OPENNI_FRAME_DEPTH_MM_ZIPPED", saveDepth=True)
+
+    ImageCapture.readFromLogFile(corlPaths['lcmlog'], corlPaths['images'], cameraName="OPENNI_FRAME_LEFT")
 
 
 def test():

scripts/bin/run_create_data

@@ -7,12 +7,22 @@ cd ..
 up_dir=${PWD##*/}
 cd $current_dir
 
+if [ $# -eq 0 ] || [ $1 != "-d" ]; then
+    echo "Not extracting depth images.  Use -d to extract them."
+	saveDepth="False"
+elif [ $1 = "-d" ]; then
+	echo "-d: Extracting depth images as well"
+	saveDepth="True"
+fi
+
 echo "First extracting images..."
 
-directorPython $LABELFUSION_SOURCE_DIR/scripts/extractImagesFromLog.py --bot-config $LABELFUSION_SOURCE_DIR/config/bot_frames.cfg --logFolder $up_dir/$current_dir
+directorPython $LABELFUSION_SOURCE_DIR/scripts/extractImagesFromLog.py --bot-config $LABELFUSION_SOURCE_DIR/config/bot_frames.cfg --logFolder $up_dir/$current_dir --saveDepth $saveDepth
 
 echo "Now rendering training images"
 
+exit 1
+
 directorPython $LABELFUSION_SOURCE_DIR/scripts/renderTrainingImages.py --bot-config $LABELFUSION_SOURCE_DIR/config/bot_frames.cfg --logFolder $up_dir/$current_dir
 
 echo "Done."

scripts/extractImagesFromLog.py

@@ -21,7 +21,15 @@
 
 	parser = drcargs.getGlobalArgParser().getParser()
 	parser.add_argument('--logFolder', type=str, dest='logFolder', help='location of top level folder for this log, relative to LabelFusion/data')
+	parser.add_argument('--saveDepth', type=str, dest='saveDepth', help='string (True or False) whether or not to extract depth images from log')
 	args = parser.parse_args()
 	print "logFolder = ", args.logFolder
-	labelfusion.imagecapture.captureImages(args.logFolder)
+
+	if args.saveDepth=="True":
+		saveDepth=True
+	else:
+		saveDepth=False
+	print saveDepth
+
+	labelfusion.imagecapture.captureImages(args.logFolder, saveDepth)