rosbag转视频文件

代码改编自github repo https://github.com/OSUrobotics/bag2video。

#!/usr/bin/env python3

import os
import sys
import argparse
import traceback
import multiprocessing
from itertools import repeat

import cv2
import numpy as np

import rospy
import rosbag


try:
    from cv_bridge import CvBridge
except ImportError:
    try:
        import roslib; roslib.load_manifest("bag2video")
        from cv_bridge import CvBridge
    except:
        print ("Could not find ROS package: cv_bridge")
        print ("If ROS version is pre-Groovy, try putting this package in ROS_PACKAGE_PATH")
        sys.exit(1)

# get necessary info from the rosbag for cv2.VideoWriter
def get_info(bag, topic=None, start_time=rospy.Time(0), stop_time=rospy.Time(sys.maxsize)):
    size = (0,0)
    times = []

    is_compressed = True
    # read the first message to get the image size
    bag_data = bag.read_messages(topics=topic)
    bridge = CvBridge()
    for _, msg, _ in bag_data:
        img = np.asarray(bridge.compressed_imgmsg_to_cv2(msg, 'bgr8'))
        if len(img.shape) == 0:
            img = np.asarray(bridge.imgmsg_to_cv2(msg, 'bgr8'))
            is_compressed = False
        size = (img.shape[1],img.shape[0])
        break

    # now read the rest of the messages for the rates
    bag_data = bag.read_messages(topics=topic, start_time=start_time, end_time=stop_time)
    for _, msg, _ in bag_data:
        time = msg.header.stamp
        times.append(time.to_sec())

    # check if video is properly timestamped
    if np.all(times):
        diffs = 1/np.diff(times)
    else:
        diffs = [0.0]
    return np.median(diffs), min(diffs), max(diffs), size, times, is_compressed

def calc_n_frames(times, precision=10):
    # the smallest interval should be one frame, larger intervals more
    intervals = np.diff(times)

    # check if video is properly timestamped
    if np.all(intervals):
        return np.int64(np.round(precision*intervals/min(intervals)))
    else:
        intervals[:] = 1
        return np.int64(intervals)

def write_frames(bag, writer, is_compressed, topic=None, nframes=repeat(1), start_time=rospy.Time(0), stop_time=rospy.Time(sys.maxsize), viz=False, encoding='bgr8'):
    bridge = CvBridge()
    if viz:
        cv2.namedWindow('win')
    count = 1
    iterator = bag.read_messages(topics=topic, start_time=start_time, end_time=stop_time)
    for (topic, msg, time), reps in zip(iterator, nframes):
        if is_compressed:
            img = np.asarray(bridge.compressed_imgmsg_to_cv2(msg, 'bgr8'))
        else:
            img = np.asarray(bridge.imgmsg_to_cv2(msg, 'bgr8'))
        for rep in range(reps):
            count += 1
            writer.write(img)
        if viz:
            imshow('win', img)

def imshow(win, img):
    cv2.imshow(win, img)
    cv2.waitKey(1)

def noshow(win, img):
    pass

def bag2video(bagfile, outfile, args):
    pid = os.getpid()
    print("---进程pid: {}".format(pid))
    print("bag : {}".format(bagfile))
    print("out : {}".format(outfile))

    try:
        bag = rosbag.Bag(bagfile, 'r')
        rate, minrate, maxrate, size, times, is_compressed = get_info(bag, args.topic, start_time=args.start, stop_time=args.end)
        nframes = calc_n_frames(times, args.precision)
        if rate == minrate == maxrate == 0.0:
            writer = cv2.VideoWriter(outfile, cv2.VideoWriter_fourcc(*'DIVX'), args.fps, size)
        else:
            writer = cv2.VideoWriter(outfile, cv2.VideoWriter_fourcc(*'DIVX'), np.ceil(maxrate*args.precision), size)

        write_frames(bag, writer, is_compressed, topic=args.topic, nframes=nframes, start_time=args.start, stop_time=args.end, encoding=args.encoding)

        writer.release()
        print ('---进程{}完成'.format(pid))
    except:
        print ('---进程{}未成功, bag:{}'.format(pid, bagfile))
        traceback.print_exc()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract and encode video from bag files.')
    parser.add_argument('--precision', '-p', action='store', default=1, type=int,
                        help='Precision of variable framerate interpolation. Higher numbers\
                        match the actual framerater better, but result in larger files and slower conversion times.')
    parser.add_argument('--fps', '-fps', action='store', default=15, type=int,
                        help='Frame per second of the video. default 15.')
    parser.add_argument('--viz', '-v', action='store_true', help='Display frames in a GUI window.')
    parser.add_argument('--start', '-s', action='store', default=rospy.Time(0), type=rospy.Time,
                        help='Rostime representing where to start in the bag.')
    parser.add_argument('--end', '-e', action='store', default=rospy.Time(sys.maxsize), type=rospy.Time,
                        help='Rostime representing where to stop in the bag.')
    parser.add_argument('--encoding', choices=('rgb8', 'bgr8', 'mono8'), default='bgr8',
                        help='Encoding of the deserialized image.')

    parser.add_argument('topic')
    parser.add_argument('dirname')

    args = parser.parse_args()

    if not args.viz:
        imshow = noshow

    video_dir = os.path.join(args.dirname, "output_video")
    if not os.path.exists(video_dir):
        os.mkdir(video_dir)

    for bagname in os.listdir(args.dirname):
        if bagname == "output_video":
            continue

        if ".bag" in bagname:
            print(bagname)
            bagfile = os.path.join(args.dirname, bagname)
            outfile = video_dir + "/" + bagfile.split("/")[-1].replace(".bag", ".avi")
            print(outfile)

            sub_process = multiprocessing.Process(target=bag2video, args=(bagfile, outfile, args))
            sub_process.start()

topic是视频的话题名，dirname是存放bag文件的路径名称，程序会把目录下全部的bag文件中的topic视频全部转化成视频文件。