[train-archive] Implement command to generate train.zip file

data/data.go

@@ -0,0 +1,179 @@
+package data
+
+import (
+	"archive/zip"
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/cyrilix/robocar-tools/record"
+	log "github.com/sirupsen/logrus"
+	"io/ioutil"
+	"os"
+	"path"
+	"regexp"
+)
+
+var camSubDir = "cam"
+
+func BuildArchive(basedir string, archiveName string, sliceSize int) error {
+	dirItems, err := ioutil.ReadDir(basedir)
+	if err != nil {
+		return fmt.Errorf("unable to list directory in %v dir: %v", basedir, err)
+	}
+
+	imgCams := make([]string, 0)
+	records := make([]string, 0)
+
+	for _, dirItem := range dirItems {
+		log.Debugf("process %v directory", dirItem)
+		imgDir := path.Join(basedir, dirItem.Name(), camSubDir)
+		imgs, err := ioutil.ReadDir(imgDir)
+		if err != nil {
+			return fmt.Errorf("unable to list cam images in directory %v: %v", imgDir, err)
+		}
+
+		for _, img := range imgs {
+			idx, err := indexFromFile(img.Name())
+			if err != nil {
+				return fmt.Errorf("unable to find index in cam image name %v: %v", img.Name(), err)
+			}
+			log.Debugf("found image with index %v", idx)
+			records = append(records, path.Join(basedir, dirItem.Name(), fmt.Sprintf(record.RecorNameFormat, idx)))
+			imgCams = append(imgCams, path.Join(basedir, dirItem.Name(), camSubDir, img.Name()))
+		}
+
+	}
+
+	if sliceSize > 0{
+		imgCams, records, err = applySlice(imgCams, records, sliceSize)
+	}
+
+	content, err := buildArchiveContent(&imgCams, &records)
+	if err != nil {
+		return fmt.Errorf("unable to build archive: %v", err)
+	}
+	err = ioutil.WriteFile(archiveName, *content, os.FileMode(0755))
+	if err != nil {
+		return fmt.Errorf("unable to write archive content to disk: %v", err)
+	}
+	return nil
+}
+
+func applySlice(imgCams []string, records []string, sliceSize int) ([]string, []string, error) {
+	// Add sliceSize images shift
+	i := imgCams[:len(imgCams)-sliceSize]
+	r := records[sliceSize:]
+
+	return i, r, nil
+}
+
+var indexRegexp *regexp.Regexp
+
+func init() {
+	re, err := regexp.Compile("image_array_(?P<idx>[0-9]+)\\.jpg$")
+	if err != nil {
+		log.Fatalf("unable to compile regex: %v", err)
+	}
+	indexRegexp = re
+}
+
+func indexFromFile(fileName string) (string, error) {
+	matches := findNamedMatches(indexRegexp, fileName)
+	if matches["idx"] == "" {
+		return "", fmt.Errorf("no index in filename")
+	}
+	return matches["idx"], nil
+}
+
+func findNamedMatches(regex *regexp.Regexp, str string) map[string]string {
+	match := regex.FindStringSubmatch(str)
+
+	results := map[string]string{}
+	for i, name := range match {
+		results[regex.SubexpNames()[i]] = name
+	}
+	return results
+}
+
+func buildArchiveContent(imgFiles *[]string, recordFiles *[]string) (*[]byte, error) {
+	// Create a buffer to write our archive to.
+	buf := new(bytes.Buffer)
+
+	// Create a new zip archive.
+	w := zip.NewWriter(buf)
+
+	err := addJsonFiles(recordFiles, imgFiles, w)
+	if err != nil {
+		return nil, fmt.Errorf("unable to write json files in zip archive: %v", err)
+	}
+
+	err = addCamImages(imgFiles, w)
+	if err != nil {
+		return nil, fmt.Errorf("unable to cam files in zip archive: %v", err)
+	}
+
+	err = w.Close()
+	if err != nil {
+		return nil, fmt.Errorf("unable to build archive: %v", err)
+	}
+
+	content, err := ioutil.ReadAll(buf)
+	return &content, err
+}
+
+func addCamImages(imgFiles *[]string, w *zip.Writer) error {
+	for _, img := range *imgFiles {
+		imgContent, err := ioutil.ReadFile(img)
+		if err != nil {
+			return fmt.Errorf("unable to read img: %v", err)
+		}
+
+		_, imgName := path.Split(img)
+		err = addToArchive(w, imgName, &imgContent)
+		if err != nil {
+			return fmt.Errorf("unable to create new img entry in archive: %v", err)
+		}
+	}
+	return nil
+}
+
+func addJsonFiles(recordFiles *[]string, imgCam *[]string, w *zip.Writer) error {
+	for idx, r := range *recordFiles {
+		content, err := ioutil.ReadFile(r)
+		if err != nil {
+			return fmt.Errorf("unable to read json content: %v", err)
+		}
+		var rcd record.Record
+		err = json.Unmarshal(content, &rcd)
+		if err != nil {
+			return fmt.Errorf("unable to unmarshal record: %v", err)
+		}
+		_, camName := path.Split((*imgCam)[idx])
+		rcd.CamImageArray = camName
+
+		recordBytes, err := json.Marshal(&rcd)
+		if err != nil {
+			return fmt.Errorf("unable to marshal %v record: %v", rcd, err)
+		}
+
+		_, recordName := path.Split(r)
+		err = addToArchive(w, recordName, &recordBytes)
+		if err != nil {
+			return fmt.Errorf("unable to create new record in archive: %v", err)
+		}
+	}
+	return nil
+}
+
+func addToArchive(w *zip.Writer, name string, content *[]byte) error {
+	recordWriter, err := w.Create(name)
+	if err != nil {
+		return fmt.Errorf("unable to create new entry %v in archive: %v", name, err)
+	}
+
+	_, err = recordWriter.Write(*content)
+	if err != nil {
+		return fmt.Errorf("unable to add content in %v zip archive: %v", name, err)
+	}
+	return nil
+}

data/data_test.go

@@ -0,0 +1,118 @@
+package data
+
+import (
+	"archive/zip"
+	"encoding/json"
+	"fmt"
+	"github.com/cyrilix/robocar-tools/record"
+	log "github.com/sirupsen/logrus"
+	"io/ioutil"
+	"os"
+	"path"
+	"strings"
+	"testing"
+)
+
+func TestBuildArchive(t *testing.T) {
+	tmpDir, err := ioutil.TempDir("", "buildarchive")
+	if err != nil {
+		t.Fatalf("unable to make tmpdir: %v", err)
+	}
+	defer func() {
+		err := os.RemoveAll(tmpDir)
+		if err != nil {
+			log.Warnf("unable to remove tempdir %v: %v", tmpDir, err)
+		}
+	}()
+
+	archive := path.Join(tmpDir, "train.zip")
+
+	expectedRecordFiles, expectedImgFiles := expectedFiles()
+
+	err = BuildArchive("testdata", archive, 0)
+	if err != nil {
+		t.Errorf("unable to build archive: %v", err)
+	}
+
+	r, err := zip.OpenReader(archive)
+	if err != nil {
+		t.Errorf("unable to read archive, %v", err)
+	}
+	defer r.Close()
+
+	if len(r.File) != len(expectedImgFiles)+len(expectedRecordFiles) {
+		t.Errorf("bad number of files in archive: %v, wants %v", len(r.File), len(expectedImgFiles)+len(expectedRecordFiles))
+	}
+
+	// Iterate through the files in the archive,
+	// printing some of their contents.
+	for _, f := range r.File {
+		filename := f.Name
+		if filename[len(filename)-4:] == "json" {
+			expectedRecordFiles[filename] = true
+			expectedtImgName := strings.Replace(filename, "record", "cam-image_array", 1)
+			expectedtImgName = strings.Replace(expectedtImgName, "json", "jpg", 1)
+			checkJsonContent(t, f, expectedtImgName)
+			continue
+		}
+		if filename[len(filename)-3:] == "jpg" {
+			expectedImgFiles[filename] = true
+			continue
+		}
+		t.Errorf("unexpected file in archive: %v", filename)
+	}
+
+	checkAllFilesAreFoundInArchive(expectedRecordFiles, t, expectedImgFiles)
+}
+
+func checkAllFilesAreFoundInArchive(expectedRecordFiles map[string]bool, t *testing.T, expectedImgFiles map[string]bool) {
+	for f, found := range expectedRecordFiles {
+		if !found {
+			t.Errorf("%v not found in archive", f)
+		}
+	}
+	for f, found := range expectedImgFiles {
+		if !found {
+			t.Errorf("%v not found in archive", f)
+		}
+	}
+}
+
+func checkJsonContent(t *testing.T, f *zip.File, expectedCamImage string) {
+	rc, err := f.Open()
+	if err != nil {
+		t.Errorf("unable to read file content of %v: %v", f.Name, err)
+	}
+	defer rc.Close()
+
+	content, err := ioutil.ReadAll(rc)
+	if err != nil {
+		t.Errorf("%v has invalid json content: %v", f.Name, err)
+	}
+	var rcd record.Record
+	err = json.Unmarshal(content, &rcd)
+	if err != nil {
+		t.Errorf("unable to unmarshal json content of%v: %v", f.Name, err)
+	}
+
+	if rcd.CamImageArray != expectedCamImage {
+		t.Errorf("record %v: invalid image ref: %v, wants %v", f.Name, rcd.CamImageArray, expectedCamImage)
+	}
+	if rcd.UserAngle == 0. {
+		t.Errorf("record %v: user angle has not been initialised", f.Name)
+	}
+}
+
+func expectedFiles() (map[string]bool, map[string]bool) {
+	expectedRecordFiles := make(map[string]bool)
+	expectedImgFiles := make(map[string]bool)
+	for i := 1; i <= 8; i++ {
+		expectedRecordFiles[fmt.Sprintf("record_%07d.json", i)] = false
+		expectedImgFiles[fmt.Sprintf("cam-image_array_%07d.jpg", i)] = false
+	}
+	for i := 101; i <= 106; i++ {
+		expectedRecordFiles[fmt.Sprintf("record_%07d.json", i)] = false
+		expectedImgFiles[fmt.Sprintf("cam-image_array_%07d.jpg", i)] = false
+	}
+	return expectedRecordFiles, expectedImgFiles
+}

data/testdata/2020021819-3/record_0000001.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000001.jpg"}

data/testdata/2020021819-3/record_0000002.json

@@ -0,0 +1 @@
+{"user/angle":0.045098066,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000002.jpg"}

data/testdata/2020021819-3/record_0000003.json

@@ -0,0 +1 @@
+{"user/angle":0.045098066,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000003.jpg"}

data/testdata/2020021819-3/record_0000004.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000004.jpg"}

data/testdata/2020021819-3/record_0000005.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000005.jpg"}

data/testdata/2020021819-3/record_0000006.json

@@ -0,0 +1 @@
+{"user/angle":0.043137312,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000006.jpg"}

data/testdata/2020021819-3/record_0000007.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000007.jpg"}

data/testdata/2020021819-3/record_0000008.json

@@ -0,0 +1 @@
+{"user/angle":0.045098066,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000008.jpg"}

data/testdata/2020021819-4/record_0000101.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000101.jpg"}

data/testdata/2020021819-4/record_0000102.json

@@ -0,0 +1 @@
+{"user/angle":0.045098066,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000102.jpg"}

data/testdata/2020021819-4/record_0000103.json

@@ -0,0 +1 @@
+{"user/angle":0.045098066,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000103.jpg"}

data/testdata/2020021819-4/record_0000104.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000104.jpg"}

data/testdata/2020021819-4/record_0000105.json

@@ -0,0 +1 @@
+{"user/angle":0.04117644,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000105.jpg"}

data/testdata/2020021819-4/record_0000106.json

@@ -0,0 +1 @@
+{"user/angle":0.043137312,"cam/image_array":"/tmp/record//2020021819-3/cam/cam-image_array_0000106.jpg"}