robocars/robocar-tools
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

feat: add display record command and refactor training command

Browse Source
This commit is contained in:
Cyrille Nofficial
2022-06-09 12:19:54 +02:00
parent 02db9e241e
commit 8b8d53af58
8 changed files with 255 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
pkg/data/data.go
View File

@ -19,8 +19,8 @@ import (
var camSubDir = "cam"
func WriteArchive(basedir string, archiveName string, sliceSize int, flipImages bool) error {
content, err := BuildArchive(basedir, sliceSize, flipImages)
func WriteArchive(basedir string, archiveName string, sliceSize int, imgWidth, imgHeight int, horizon int, flipImages bool) error {
content, err := BuildArchive(basedir, sliceSize, imgWidth, imgHeight, horizon, flipImages)
if err != nil {
return fmt.Errorf("unable to build archive: %w", err)
}
@ -34,7 +34,7 @@ func WriteArchive(basedir string, archiveName string, sliceSize int, flipImages
return nil
}
func BuildArchive(basedir string, sliceSize int, flipImages bool) ([]byte, error) {
func BuildArchive(basedir string, sliceSize int, imgWidth, imgHeight int, horizon int, flipImages bool) ([]byte, error) {
l := zap.S()
l.Infof("build zip archive from %s\n", basedir)
dirItems, err := ioutil.ReadDir(basedir)
@ -59,7 +59,7 @@ func BuildArchive(basedir string, sliceSize int, flipImages bool) ([]byte, error
return nil, fmt.Errorf("unable to find index in cam image name %v: %w", img.Name(), err)
}
l.Debugf("found image with index %v", idx)
records = append(records, path.Join(basedir, dirItem.Name(), fmt.Sprintf(record.RecorNameFormat, idx)))
records = append(records, path.Join(basedir, dirItem.Name(), fmt.Sprintf(record.FileNameFormat, idx)))
imgCams = append(imgCams, path.Join(basedir, dirItem.Name(), camSubDir, img.Name()))
}
}
@ -73,12 +73,12 @@ func BuildArchive(basedir string, sliceSize int, flipImages bool) ([]byte, error
// Create a new zip archive.
w := zip.NewWriter(buf)
err = buildArchiveContent(w, imgCams, records, false)
err = buildArchiveContent(w, imgCams, records, imgWidth, imgHeight, horizon, false)
if err != nil {
return nil, fmt.Errorf("unable to build archive: %w", err)
}
if flipImages {
err = buildArchiveContent(w, imgCams, records, true)
err = buildArchiveContent(w, imgCams, records, imgWidth, imgHeight, horizon, true)
if err != nil {
return nil, fmt.Errorf("unable to build archive: %w", err)
}
@ -132,13 +132,13 @@ func findNamedMatches(regex *regexp.Regexp, str string) map[string]string {
return results
}
func buildArchiveContent(w *zip.Writer, imgFiles []string, recordFiles []string, withFlipImages bool) error {
func buildArchiveContent(w *zip.Writer, imgFiles []string, recordFiles []string, imgWidth, imgHeight int, horizon int, withFlipImages bool) error {
err := addJsonFiles(recordFiles, imgFiles, withFlipImages, w)
if err != nil {
return fmt.Errorf("unable to write json files in zip archive: %w", err)
}
err = addCamImages(imgFiles, withFlipImages, w)
err = addCamImages(imgFiles, withFlipImages, w, imgWidth, imgHeight, horizon)
if err != nil {
return fmt.Errorf("unable to cam files in zip archive: %w", err)
}
@ -146,7 +146,7 @@ func buildArchiveContent(w *zip.Writer, imgFiles []string, recordFiles []string,
return err
}
func addCamImages(imgFiles []string, flipImage bool, w *zip.Writer) error {
func addCamImages(imgFiles []string, flipImage bool, w *zip.Writer, imgWidth, imgHeight int, horizon int) error {
for _, im := range imgFiles {
imgContent, err := ioutil.ReadFile(im)
if err != nil {
@ -154,17 +154,29 @@ func addCamImages(imgFiles []string, flipImage bool, w *zip.Writer) error {
}
_, imgName := path.Split(im)
if flipImage {
if flipImage || imgWidth > 0 && imgHeight > 0 || horizon > 0 {
img, _, err := image.Decode(bytes.NewReader(imgContent))
if err != nil {
zap.S().Fatalf("unable to decode peg image: %v", err)
zap.S().Fatalf("unable to decode jpeg image: %v", err)
}
imgFlip := imaging.FlipH(img)
var bytesBuff bytes.Buffer
err = jpeg.Encode(&bytesBuff, imgFlip, nil)
if imgWidth > 0 && imgHeight > 0 {
bounds := img.Bounds()
if bounds.Dx() != imgWidth || bounds.Dy() != imgWidth {
zap.S().Debugf("resize image %v from %dx%d to %dx%d", im, bounds.Dx(), bounds.Dy(), imgWidth, imgHeight)
img = imaging.Resize(img, imgWidth, imgHeight, imaging.NearestNeighbor)
}
}
if flipImage {
img = imaging.FlipH(img)
imgName = fmt.Sprintf("flip_%s", imgName)
}
if horizon > 0 {
img = imaging.Crop(img, image.Rect(0, horizon, img.Bounds().Dx(), img.Bounds().Dy()))
}
var bytesBuff bytes.Buffer
err = jpeg.Encode(&bytesBuff, img, nil)
imgContent = bytesBuff.Bytes()
imgName = fmt.Sprintf("flip_%s", imgName)
}
err = addToArchive(w, imgName, imgContent)

2
pkg/data/data_test.go
View File

@ -29,7 +29,7 @@ func TestBuildArchive(t *testing.T) {
expectedRecordFiles, expectedImgFiles := expectedFiles()
err = WriteArchive("testdata", archive, 0)
err = WriteArchive("testdata", archive, 0, 160, 120, 0, false)
if err != nil {
t.Errorf("unable to build archive: %v", err)
}

130
pkg/display/record.go Normal file
View File

@ -0,0 +1,130 @@
package display
import (
"fmt"
"github.com/cyrilix/robocar-protobuf/go/events"
mqtt "github.com/eclipse/paho.mqtt.golang"
"github.com/golang/protobuf/proto"
"go.uber.org/zap"
"gocv.io/x/gocv"
"image"
"image/color"
)
func NewRecordDisplay(client mqtt.Client, recordTopic string) *Record {
return &Record{
client: client,
recordTopic: recordTopic,
window: gocv.NewWindow("recordTopic"),
recordChan: make(chan *events.RecordMessage),
cancel: make(chan interface{}),
}
}
type Record struct {
client mqtt.Client
recordTopic string
window *gocv.Window
recordChan chan *events.RecordMessage
cancel chan interface{}
}
func (r *Record) Start() error {
if err := r.registerCallbacks(); err != nil {
return fmt.Errorf("unable to start service: %v", err)
}
var rec *events.RecordMessage
var objectsMsg events.ObjectsMessage
var roadMsg events.RoadMessage
for {
select {
case newRecord := <-r.recordChan:
rec = newRecord
case <-r.cancel:
return nil
}
go r.drawRecord(rec, &objectsMsg, &roadMsg)
}
}
func (r *Record) Stop() {
defer r.window.Close()
close(r.cancel)
StopService("record-display", r.client, r.recordTopic)
}
func (r *Record) onRecord(_ mqtt.Client, message mqtt.Message) {
var msg events.RecordMessage
err := proto.Unmarshal(message.Payload(), &msg)
if err != nil {
zap.S().Errorf("unable to unmarshal protobuf FrameMessage: %v", err)
return
}
message.Ack()
r.recordChan <- &msg
}
func (r *Record) registerCallbacks() error {
err := RegisterCallback(r.client, r.recordTopic, r.onRecord)
if err != nil {
return err
}
return nil
}
func (r *Record) drawRecord(rec *events.RecordMessage, objects *events.ObjectsMessage, road *events.RoadMessage) {
img, err := gocv.IMDecode(rec.GetFrame().GetFrame(), gocv.IMReadUnchanged)
if err != nil {
zap.S().Errorf("unable to decode image: %v", err)
return
}
defer img.Close()
steering := rec.GetSteering().GetSteering()
r.drawSteering(&img, steering)
r.window.IMShow(img)
r.window.WaitKey(1)
}
func (r *Record) drawSteering(img *gocv.Mat, steering float32) {
gocv.PutText(
img,
fmt.Sprintf("Steering: %.3f", steering),
image.Point{X: 20, Y: 20},
gocv.FontHersheyPlain,
1.,
color.RGBA{R: 0, G: 255, B: 0, A: 255},
1,
)
}
func StopService(name string, client mqtt.Client, topics ...string) {
zap.S().Infof("Stop %s service", name)
token := client.Unsubscribe(topics...)
token.Wait()
if token.Error() != nil {
zap.S().Errorf("unable to unsubscribe service: %v", token.Error())
}
client.Disconnect(50)
}
func RegisterCallback(client mqtt.Client, topic string, callback mqtt.MessageHandler) error {
zap.S().Infof("Register callback on topic %v", topic)
token := client.Subscribe(topic, 0, callback)
token.Wait()
if token.Error() != nil {
return fmt.Errorf("unable to register callback on topic %s: %v", topic, token.Error())
}
return nil
}

12
pkg/train/train.go
View File

@ -35,10 +35,10 @@ type Training struct {
outputBucket string
}
func (t *Training) TrainDir(ctx context.Context, jobName, basedir string, imgHeight, imgWidth int, sliceSize int, withFlipImage bool, outputModelFile string, enableSpotTraining bool) error {
func (t *Training) TrainDir(ctx context.Context, jobName, basedir string, imgWidth, imgHeight, sliceSize int, horizon int, withFlipImage bool, outputModelFile string, enableSpotTraining bool) error {
l := zap.S()
l.Infof("run training with data from %s", basedir)
archive, err := data.BuildArchive(basedir, sliceSize, withFlipImage)
archive, err := data.BuildArchive(basedir, sliceSize, imgWidth, imgHeight, horizon, withFlipImage)
if err != nil {
return fmt.Errorf("unable to build data archive: %w", err)
}
@ -110,7 +110,7 @@ func (t *Training) runTraining(ctx context.Context, jobName string, slideSize in
S3OutputPath: aws.String(t.outputBucket),
},
ResourceConfig: &types.ResourceConfig{
InstanceCount: 1,
InstanceCount: 1,
//InstanceType: types.TrainingInstanceTypeMlP2Xlarge,
InstanceType: types.TrainingInstanceTypeMlG4dnXlarge,
VolumeSizeInGB: 1,
@ -168,7 +168,7 @@ func (t *Training) runTraining(ctx context.Context, jobName string, slideSize in
}
switch status.TrainingJobStatus {
case types.TrainingJobStatusInProgress:
l.Infof("job in progress: %v - %v - %v", status.TrainingJobStatus, status.SecondaryStatus, *status.SecondaryStatusTransitions[len(status.SecondaryStatusTransitions) - 1].StatusMessage)
l.Infof("job in progress: %v - %v - %v", status.TrainingJobStatus, status.SecondaryStatus, *status.SecondaryStatusTransitions[len(status.SecondaryStatusTransitions)-1].StatusMessage)
continue
case types.TrainingJobStatusFailed:
return fmt.Errorf("job %s finished with status %v", jobName, status.TrainingJobStatus)
@ -198,5 +198,5 @@ func ListJob(ctx context.Context) error {
for _, job := range jobs.TrainingJobSummaries {
fmt.Printf("%s\t\t%s\n", *job.TrainingJobName, job.TrainingJobStatus)
}
return nil
}
return nil
}