robocar-tools/record/record.go

103 lines
2.7 KiB
Go

package record
import (
"encoding/json"
"fmt"
"github.com/cyrilix/robocar-base/service"
"github.com/cyrilix/robocar-protobuf/go/events"
mqtt "github.com/eclipse/paho.mqtt.golang"
"github.com/golang/protobuf/proto"
log "github.com/sirupsen/logrus"
"io/ioutil"
"os"
)
func New(client mqtt.Client, recordsDir, recordTopic string) (*Recorder, error) {
err := os.MkdirAll(recordsDir, os.FileMode(0755))
if err != nil {
return nil, fmt.Errorf("unable to create %v directory: %v", recordsDir, err)
}
return &Recorder{
client: client,
recordsDir: recordsDir,
recordTopic: recordTopic,
cancel: make(chan interface{}),
}, nil
}
type Recorder struct {
client mqtt.Client
recordsDir string
recordTopic string
cancel chan interface{}
}
var RecorNameFormat = "record_%s.json"
func (r *Recorder) Start() error {
err := service.RegisterCallback(r.client, r.recordTopic, r.onRecordMsg)
if err != nil {
return fmt.Errorf("unable to start recorder part: %v", err)
}
<-r.cancel
return nil
}
func (r *Recorder) Stop() {
service.StopService("record", r.client, r.recordTopic)
close(r.cancel)
}
func (r *Recorder) onRecordMsg(_ mqtt.Client, message mqtt.Message) {
var msg events.RecordMessage
err := proto.Unmarshal(message.Payload(), &msg)
if err != nil {
log.Errorf("unable to unmarshal protobuf %T: %v", msg, err)
return
}
fmt.Printf("record %s: %s\r", msg.GetRecordSet(), msg.GetFrame().GetId().GetId())
recordDir := fmt.Sprintf("%s/%s", r.recordsDir, msg.GetRecordSet())
imgDir := fmt.Sprintf("%s/cam", recordDir)
imgName := fmt.Sprintf("%s/cam-image_array_%s.jpg", imgDir, msg.GetFrame().GetId().GetId())
err = os.MkdirAll(imgDir, os.FileMode(0755))
if err != nil {
log.Errorf("unable to create %v directory: %v", imgDir, err)
return
}
err = ioutil.WriteFile(imgName, msg.GetFrame().GetFrame(), os.FileMode(0755))
if err != nil {
log.Errorf("unable to write img file %v: %v", imgName, err)
return
}
jsonDir := fmt.Sprintf("%s/", recordDir)
recordName := fmt.Sprintf("%s/%s", jsonDir, fmt.Sprintf(RecorNameFormat, msg.GetFrame().GetId().GetId()))
err = os.MkdirAll(jsonDir, os.FileMode(0755))
if err != nil {
log.Errorf("unable to create %v directory: %v", jsonDir, err)
return
}
record := Record{
UserAngle: msg.GetSteering().GetSteering(),
CamImageArray: imgName,
}
jsonBytes, err := json.Marshal(&record)
if err != nil {
log.Errorf("unable to marshal json content: %v", err)
return
}
err = ioutil.WriteFile(recordName, jsonBytes, 0755)
if err != nil {
log.Errorf("unable to write json file %v: %v", recordName, err)
}
}
type Record struct {
UserAngle float32 `json:"user/angle,"`
CamImageArray string `json:"cam/image_array,"`
}