import cv2 
import time, sys, os
from ros import rosbag
import roslib
import rospy
roslib.load_manifest('sensor_msgs')
from sensor_msgs.msg import Image,Imu
from geometry_msgs.msg import Vector3
from cv_bridge import CvBridge
from numpy import asarray

# import ImageFile
from PIL import ImageFile
from PIL import Image as ImagePIL

def CompSortFileNamesNr(f):
    g = os.path.splitext(os.path.split(f)[1])[0] #get the file of the
    numbertext = ''.join(c for c in g if c.isdigit())
    return int(numbertext)

def ReadIMU(IMUFile):
    '''return IMU data and timestamp of IMU'''
    IMUfp = open(IMUFile,'r')
    IMULines = IMUfp.readlines()
    #all = IMUDatas.readlines()
    IMUDatas = {}
    for l in IMULines:
        if l[0] == "#":
            continue;
        items = l.rstrip('\n').split(' ')
        IMUDatas[items[0]] = items[1:]
    
    IMUfp.close()
    return IMUDatas 

def ReadImages(assocoations):
   assofp = open(assocoations, 'r')
   asso = assofp.readlines()
   RGBImages = {}
   depthImages = {}
   for l in asso:
       if l[0] == "#":
           continue;
       items = l.rstrip('\n').split(' ')
       RGBImages[items[0]] = items[1]
       depthImages[items[2]] = items[3]

   assofp.close()
   return RGBImages, depthImages

def CreateBag(args):#assocoations, imu, output_bag_name
    '''read assocoations.txt'''
    RGBImages,depthImages = ReadImages(args[1])

    IMUDatas = ReadIMU(args[2]) #the url of IMU data

    '''Creates a bag file with camera images'''
    if not os.path.exists(args[3]):
       os.system(r'touch %s' % args[3])
    else:
       os.system(r'rm %s' % args[3])
       os.system(r'touch %s' % args[3])

    bagName = rosbag.Bag(args[3], 'w')

    try:
        for it, iData in IMUDatas.items():
            imu = Imu()
            imuStamp = rospy.rostime.Time.from_sec(float(it))
            #angular_v = Vector3()
            linear_a = Vector3()
            #angular_v.x = float(iData[0])
            #angular_v.y = float(iData[1])
            #angular_v.z = float(iData[2])
            linear_a.x = float(iData[0])
            linear_a.y = float(iData[1])
            linear_a.z = float(iData[2])
            imu.header.stamp = imuStamp
            #imu.angular_velocity = angular_v
            imu.linear_acceleration = linear_a

            bagName.write("/imu",imu,imuStamp)

        br = CvBridge()

        for imt, img in RGBImages.items():
            #img = args[2] + img; 
            print("Adding %s" % img)

            cv_image = cv2.imread(img)

            Stamp = rospy.rostime.Time.from_sec(float(imt))

            '''set image information '''
            Img = br.cv2_to_imgmsg(cv_image)
            Img.header.stamp = Stamp
            Img.header.frame_id = "camera"

            '''for mono8'''
            Img.encoding = "rgb8"
            bagName.write('/camera/rgb/image_color', Img, Stamp)

        for dt, dimg in depthImages.items():
            #dimg = args[2] + dimg; 
            print("Adding %s" % dimg)

            cv_image = cv2.imread(dimg, cv2.IMREAD_ANYDEPTH)

            '''set image information '''
            Stamp = rospy.rostime.Time.from_sec(float(dt))

            '''set image information '''
            dImg = br.cv2_to_imgmsg(cv_image)
            dImg.header.stamp = Stamp
            dImg.header.frame_id = "camera"

            #dImg.encoding = "32FC1"

            bagName.write('/camera/depth/image', dImg, Stamp)

    finally:
        bagName.close()

if __name__ == "__main__":
    print(sys.argv)

    if len(sys.argv) < 4:
        print("Usage:\n\t python generate_bags.py /path/assocoations.txt /path/accelerometer.txt output.bag")
        exit(1)

    CreateBag(sys.argv)

#!/bin/python
import roslib
import rosbag
import rospy
import cv2
import os
from sensor_msgs.msg import Image
from geometry_msgs.msg import Pose, Quaternion, Point
from cv_bridge import CvBridge
from cv_bridge import CvBridgeError


ros_bag = 'sence5.bag'  #bag包路径
save_path = 'sence5/'   #输出数据集的路径
rgb = save_path + 'rgb/'  #rgb path
depth = save_path + 'depth/'   #depth path

bridge = CvBridge()

file_handle1 = open(save_path + 'rgb.txt', 'w')
file_handle2 = open(save_path + 'depth.txt', 'w')
file_handle3 = open(save_path + 'groundtruth.txt', 'w')

with rosbag.Bag(ros_bag, 'r') as bag:
    for topic,msg,t in bag.read_messages():

        if topic == "/camera/color/image_raw":   #rgb topic
            cv_image = bridge.imgmsg_to_cv2(msg,"bgr8")
            timestr = "%.6f" %  msg.header.stamp.to_sec()   #rgb time stamp
            image_name = timestr+ ".png"
            path = "rgb/" + image_name
            file_handle1.write(timestr + " " + path + '\n')
            cv2.imwrite(rgb + image_name, cv_image)

        if topic == "/camera/aligned_depth_to_color/image_raw":  #depth topic
            cv_image = bridge.imgmsg_to_cv2(msg)
            #cv_image = bridge.imgmsg_to_cv2(msg, '32FC1')
            #cv_image = cv_image * 255
            timestr = "%.6f" %  msg.header.stamp.to_sec()   #depth time stamp
            image_name = timestr+ ".png"
            path = "depth/" + image_name
            file_handle2.write(timestr + " " + path + '\n')
            cv2.imwrite(depth + image_name, cv_image)

        if topic == '/vrpn_client_node/RigidBody_ZLT_UAV/pose': #groundtruth topic
            p = msg.pose.position
            q = msg.pose.orientation
            timestr = "%.6f" %  msg.header.stamp.to_sec()
            file_handle3.write(timestr + " " + str(round(p.x, 4)) + " " + str(round(p.y, 4)) + " " + str(round(p.z, 4)) + " ")
            file_handle3.write(str(round(q.x, 4)) + " " + str(round(q.y, 4)) + " " + str(round(q.z, 4)) + " " + str(round(q.w, 4)) + '\n')
file_handle1.close()
file_handle2.close()
file_handle3.close()