-
Notifications
You must be signed in to change notification settings - Fork 110
/
rosbag_parser.go
104 lines (88 loc) · 2.44 KB
/
rosbag_parser.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
// Package ros implements functionality that bridges the gap between `rdk` and ROS
package ros
import (
"encoding/json"
"io"
"os"
"github.com/edaniels/gobag/rosbag"
"github.com/pkg/errors"
"go.viam.com/utils"
)
// ReadBag reads the contents of a rosbag into a gobag data structure.
func ReadBag(filename string) (*rosbag.RosBag, error) {
//nolint:gosec
f, err := os.Open(filename)
defer utils.UncheckedErrorFunc(f.Close)
if err != nil {
return nil, errors.Wrapf(err, "unable to open input file")
}
rb := rosbag.NewRosBag()
if err := rb.Read(f); err != nil {
return nil, errors.Wrapf(err, "unable to create ros bag, error")
}
return rb, nil
}
// WriteTopicsJSON writes data from a rosbag into JSON files, filtered and sorted by topic.
func WriteTopicsJSON(rb *rosbag.RosBag, startTime, endTime int64, topicsFilter []string) error {
var timeFilterFunc func(int64) bool
if startTime == 0 || endTime == 0 {
timeFilterFunc = func(timestamp int64) bool {
return true
}
} else {
timeFilterFunc = func(timestamp int64) bool {
return timestamp >= startTime && timestamp <= endTime
}
}
var topicFilterFunc func(string) bool
if len(topicsFilter) == 0 {
topicFilterFunc = func(string) bool {
return true
}
} else {
topicsFilterMap := make(map[string]bool)
for _, topic := range topicsFilter {
topicsFilterMap[topic] = true
}
topicFilterFunc = func(topic string) bool {
_, ok := topicsFilterMap[topic]
return ok
}
}
if err := rb.ParseTopicsToJSON("", timeFilterFunc, topicFilterFunc, false); err != nil {
return errors.Wrapf(err, "error while parsing bag to JSON")
}
return nil
}
// AllMessagesForTopic returns all messages for a specific topic in the ros bag.
func AllMessagesForTopic(rb *rosbag.RosBag, topic string) ([]map[string]interface{}, error) {
if err := rb.ParseTopicsToJSON(
"",
func(int64) bool { return true },
func(t string) bool { return t == topic },
false,
); err != nil {
return nil, errors.Wrapf(err, "error while parsing bag to JSON")
}
msgs := rb.TopicsAsJSON[topic]
if msgs == nil {
return nil, errors.Errorf("no messages for topic %s", topic)
}
all := []map[string]interface{}{}
for {
data, err := msgs.ReadBytes('\n')
if err != nil {
if errors.Is(err, io.EOF) {
break
}
return nil, err
}
message := map[string]interface{}{}
err = json.Unmarshal(data, &message)
if err != nil {
return nil, err
}
all = append(all, message)
}
return all, nil
}