239 lines
6.7 KiB
Go
239 lines
6.7 KiB
Go
package data
|
|
|
|
import (
|
|
"archive/zip"
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"github.com/cyrilix/robocar-tools/record"
|
|
"github.com/disintegration/imaging"
|
|
"go.uber.org/zap"
|
|
"image"
|
|
"image/jpeg"
|
|
"io/ioutil"
|
|
"os"
|
|
"path"
|
|
"regexp"
|
|
"strings"
|
|
)
|
|
|
|
var camSubDir = "cam"
|
|
|
|
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)
|
|
}
|
|
if err != nil {
|
|
return fmt.Errorf("unable to build archive: %w", err)
|
|
}
|
|
err = ioutil.WriteFile(archiveName, content, os.FileMode(0755))
|
|
if err != nil {
|
|
return fmt.Errorf("unable to write archive content to disk: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
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)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("unable to list directory in %v dir: %w", basedir, err)
|
|
}
|
|
|
|
imgCams := make([]string, 0)
|
|
records := make([]string, 0)
|
|
|
|
for _, dirItem := range dirItems {
|
|
l.Infof("process %v directory", dirItem.Name())
|
|
imgDir := path.Join(basedir, dirItem.Name(), camSubDir)
|
|
imgs, err := ioutil.ReadDir(imgDir)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("unable to list cam images in directory %v: %w", imgDir, err)
|
|
}
|
|
|
|
for _, img := range imgs {
|
|
idx, err := indexFromFile(img.Name())
|
|
if err != nil {
|
|
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.FileNameFormat, idx)))
|
|
imgCams = append(imgCams, path.Join(basedir, dirItem.Name(), camSubDir, img.Name()))
|
|
}
|
|
}
|
|
|
|
if sliceSize > 0 {
|
|
imgCams, records, err = applySlice(imgCams, records, sliceSize)
|
|
}
|
|
|
|
// Create a buffer to write our archive to.
|
|
buf := new(bytes.Buffer)
|
|
// Create a new zip archive.
|
|
w := zip.NewWriter(buf)
|
|
|
|
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, imgWidth, imgHeight, horizon, true)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("unable to build archive: %w", err)
|
|
}
|
|
}
|
|
|
|
err = w.Close()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("unable to close zip archive: %w", err)
|
|
}
|
|
content, err := ioutil.ReadAll(buf)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("unable to generate archive content: %w", err)
|
|
}
|
|
l.Info("archive built\n")
|
|
return content, 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 {
|
|
zap.S().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(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, imgWidth, imgHeight, horizon)
|
|
if err != nil {
|
|
return fmt.Errorf("unable to cam files in zip archive: %w", err)
|
|
}
|
|
|
|
return err
|
|
}
|
|
|
|
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 {
|
|
return fmt.Errorf("unable to read img: %w", err)
|
|
}
|
|
|
|
_, imgName := path.Split(im)
|
|
if flipImage || imgWidth > 0 && imgHeight > 0 || horizon > 0 {
|
|
img, _, err := image.Decode(bytes.NewReader(imgContent))
|
|
if err != nil {
|
|
zap.S().Fatalf("unable to decode jpeg image: %v", err)
|
|
}
|
|
|
|
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()
|
|
}
|
|
|
|
err = addToArchive(w, imgName, imgContent)
|
|
if err != nil {
|
|
return fmt.Errorf("unable to create new img entry in archive: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func addJsonFiles(recordFiles []string, imgCam []string, flipImage bool, 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: %w", err)
|
|
}
|
|
var rcd record.Record
|
|
err = json.Unmarshal(content, &rcd)
|
|
if err != nil {
|
|
return fmt.Errorf("unable to unmarshal record: %w", err)
|
|
}
|
|
_, camName := path.Split((imgCam)[idx])
|
|
|
|
if flipImage {
|
|
rcd.UserAngle = rcd.UserAngle * -1
|
|
rcd.CamImageArray = fmt.Sprintf("flip_%s", camName)
|
|
} else {
|
|
rcd.CamImageArray = camName
|
|
}
|
|
|
|
recordBytes, err := json.Marshal(&rcd)
|
|
if err != nil {
|
|
return fmt.Errorf("unable to marshal %v record: %w", rcd, err)
|
|
}
|
|
|
|
_, recordName := path.Split(r)
|
|
if flipImage {
|
|
recordName = strings.ReplaceAll(recordName, "record", "record_flip")
|
|
}
|
|
err = addToArchive(w, recordName, recordBytes)
|
|
if err != nil {
|
|
return fmt.Errorf("unable to create new record in archive: %w", 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: %w", name, err)
|
|
}
|
|
|
|
_, err = recordWriter.Write(content)
|
|
if err != nil {
|
|
return fmt.Errorf("unable to add content in %v zip archive: %w", name, err)
|
|
}
|
|
return nil
|
|
}
|