rosbag_always.py
Go to the documentation of this file.
00001 #!/usr/bin/env python
00002 
00003 import rospy
00004 import time
00005 import subprocess
00006 import signal
00007 import os
00008 import sys
00009 import argparse
00010 import re
00011 import shutil
00012 try:
00013     import colorama
00014 except:
00015     print "Please install colorama by pip install colorama"
00016     sys.exit(1)
00017 from colorama import Fore, Style
00018 from jsk_topic_tools.master_util import isMasterAlive
00019 
00020 
00021 def runROSBag(topics, size, save_dir):
00022     """
00023     run rosbag and return Popen object
00024     """
00025     cmd = 'roslaunch jsk_data rosbag_always_run_rosbag.launch'
00026     formatted_topics = [t for t in topics.split(' ') if t]
00027     args = cmd.split(' ') + ["TOPICS:=" + topics + ""] + ["SIZE:=" + size] + ["OUTPUT:=" + save_dir + '/rosbag']
00028     print args
00029     return subprocess.Popen(args)
00030 
00031 def parseBagFile(bag):
00032     # bag file name is ...
00033     # 'rosbag_YYYY-MM-DD-HH-mm-SS_i.bag'
00034     regexp = 'rosbag_(\d\d\d\d)-(\d\d)-(\d\d)-(\d\d)-(\d\d)-(\d\d)_([\d]+).bag'
00035     result = re.match(regexp, bag)
00036     if not result:
00037         return None
00038     else:
00039         return [result.group(f) for f in range(1, 8)]
00040 
00041 def mkdirForBag(root_dir, bag):
00042     # mkdir like ${root_dir}/${YYYY}/${MM}/${DD}/${HH}
00043     parse = parseBagFile(bag)
00044     YYYY = parse[0]
00045     MM = parse[1]
00046     DD = parse[2]
00047     HH = parse[3]
00048     directory = os.path.join(root_dir, YYYY, MM, DD, HH)
00049     if not os.path.exists(directory):
00050         os.makedirs(directory)
00051     return directory
00052     
00053 def moveBagFiles(save_dir, bags):
00054     for bag in bags:
00055         move_dir = mkdirForBag(save_dir, bag)
00056         from_file = os.path.join(save_dir, bag)
00057         to_file = os.path.join(move_dir, bag)
00058         print 'moving file %s -> %s' % (from_file, to_file)
00059         shutil.move(from_file, to_file)
00060                     
00061 def watchFileSystem(save_dir, max_size):
00062     files = os.listdir(save_dir)
00063     target_bags = [f for f in files 
00064                    if f.startswith('rosbag') 
00065                    and f.endswith('.bag')]
00066     moveBagFiles(save_dir, target_bags)
00067     # check size of directory
00068     checkDirectorySize(save_dir, int(max_size))
00069 
00070 def getDirectorySize(start_path = '.'):
00071     "size unit is Mega bytes"
00072     total_size = 0
00073     for dirpath, dirnames, filenames in os.walk(start_path):
00074         for f in filenames:
00075             fp = os.path.join(dirpath, f)
00076             try:
00077                 total_size += os.path.getsize(fp)
00078             except:
00079                 pass
00080     return total_size / 1000.0 / 1000.0
00081 
00082 def keyFuncToSortBag(bag):
00083     parse = parseBagFile(os.path.basename(bag))
00084     parse[len(parse) - 1] = str(int(parse[len(parse) - 1])).zfill(4)
00085     concatenated_string = reduce(lambda x, y: x + y, parse)
00086     return int(concatenated_string)
00087     
00088 def listBagsSortedByDate(save_dir):
00089     bags = []
00090     for dirpath, dirnames, filenames in os.walk(save_dir):
00091         for f in filenames:
00092             if f.startswith('rosbag') and (f.endswith('.bag') or f.endswith('.active')):
00093                 bags.append(os.path.join(dirpath, f))
00094     return sorted(bags, key=keyFuncToSortBag)
00095     
00096 def removeOldFiles(save_dir, max_size, current_size):
00097     files = listBagsSortedByDate(save_dir)
00098     remove_size = current_size - max_size
00099     for f in files:
00100         the_size = os.path.getsize(f)
00101         print Fore.GREEN + 'removing %s (%d)' % (f, the_size / 1000 / 1000) + Fore.RESET
00102         os.remove(f)
00103         # Send desktop notification
00104         subprocess.check_output(['notify-send', "Removed %s (%d)" % (f, the_size / 1000 / 1000)])
00105         remove_size = remove_size - the_size / 1000.0 / 1000.0
00106         if remove_size < 0:
00107             return
00108     
00109 def checkDirectorySize(save_dir, max_size):
00110     size = getDirectorySize(save_dir)
00111     # print 'current directory size is %fM (max is %dM)' % (size, int(max_size))
00112     if size > max_size:
00113         removeOldFiles(save_dir, max_size, size)
00114     
00115 g_rosbag_process = False
00116 
00117 def restartROSBag(topics, size, save_dir):
00118     global g_rosbag_process
00119     print 'Running rosbag...'
00120     g_rosbag_process = runROSBag(topics, size, save_dir)
00121 
00122 def killChildProcesses(ppid):
00123     output = subprocess.check_output(['ps', '--ppid=' + str(ppid), '--no-headers'])
00124     for process_line in output.split('\n'):
00125         strip_process_line = process_line.strip()
00126         if strip_process_line:
00127             pid = strip_process_line.split(' ')[0]
00128             name = strip_process_line.split(' ')[-1]
00129             print 'killing %s' % (name)
00130             os.kill(int(pid), signal.SIGINT)
00131 
00132 def killROSBag():
00133     global g_rosbag_process
00134     if g_rosbag_process:
00135         print 'Killing rosbag ...'
00136         rosbag_pid = g_rosbag_process.pid
00137         try:
00138             killChildProcesses(rosbag_pid)
00139             g_rosbag_process.send_signal(subprocess.signal.SIGINT)
00140         except:
00141             pass
00142     
00143 def main(topics, size, save_dir, max_size, rate = 1):
00144     if not os.path.exists(save_dir):
00145         os.makedirs(save_dir)
00146     previous_master_state = None
00147     try:
00148         while True:
00149             master_state = isMasterAlive()
00150             if not master_state and previous_master_state:
00151                 print "kill rosbag"
00152                 killROSBag()
00153             elif master_state and not previous_master_state:
00154                 print "restart rosbag"
00155                 restartROSBag(topics, size, save_dir)
00156             watchFileSystem(save_dir, max_size)
00157             previous_master_state = master_state
00158             time.sleep(1.0 / rate)
00159     except Exception, e:
00160         time.sleep(1)
00161         watchFileSystem(save_dir, max_size)
00162     finally:
00163         killROSBag()
00164 
00165 
00166         
00167 if __name__ == "__main__":
00168     parser = argparse.ArgumentParser(description='rosbag record regardless of rosmaster status')
00169     parser.add_argument('--topics', help="topics to record", required=True)
00170     parser.add_argument('--size', help="size of each rosbag", required=True)
00171     parser.add_argument('--save-dir', help="directory to store rosbag", required=True)
00172     parser.add_argument('--max-size', help="maximum size of rosbags in save_dir", required=True, type=int)
00173     args = parser.parse_args()
00174     main(args.topics, args.size, args.save_dir, args.max_size)


jsk_data
Author(s):
autogenerated on Fri Sep 8 2017 03:39:16