From 28779d4f7c85d1a432df5bf4ef2ac01ba27004d8 Mon Sep 17 00:00:00 2001
From: Anthony <Anthbs@gmail.com>
Date: Mon, 13 Aug 2018 12:45:10 +1200
Subject: [PATCH] Fix KITTI export missing training images

Fix for no training images due to posFrameIndex only being incremented after label generation but all images are saved before the labels are generated.
---
 src/lib/detection_algorithms/kitti/exporter.js | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/lib/detection_algorithms/kitti/exporter.js b/src/lib/detection_algorithms/kitti/exporter.js
index 76d4541ef7..ce58dfd031 100644
--- a/src/lib/detection_algorithms/kitti/exporter.js
+++ b/src/lib/detection_algorithms/kitti/exporter.js
@@ -38,14 +38,18 @@ function Exporter(exportDirPath, classes, taggedFramesCount, frameWidth, frameHe
     self.frameWidth = frameWidth;
     self.frameHeight = frameHeight;
     self.testFrameIndices = null;
-    self.posFrameIndex = null;
+    self.testFrameNames = null;
+    self.posFrameLabelIndex = null;
+    self.posFrameImageIndex = null;
     self.testSplit = testSplit || 0.2;
 
     // Prepare everything for exporting (e.g. create metadata files,
     // directories, ..)    
     // Returns: A Promise object that resolves when the operation completes
     this.init = function init() {
-        self.posFrameIndex = 0;
+        self.posFrameLabelIndex = 0;
+        self.posFrameImageIndex = 0;
+        self.testFrameNames = [];
         self.testFrameIndices = detectionUtils.generateTestIndecies(self.testSplit, taggedFramesCount);
         self.filesTouched = {};  // Keep track of files we've touched so far
         return new Promise((resolve, reject) => {
@@ -87,13 +91,17 @@ function Exporter(exportDirPath, classes, taggedFramesCount, frameWidth, frameHe
                 detectionUtils.ensureDirExists.bind(null, self.validImagesDirPath),
                 detectionUtils.ensureDirExists.bind(null, self.validLabelsDirPath),
                 function saveImage(cb) {
-                    var isTestFrame = (self.testFrameIndices.includes(self.posFrameIndex));
+                    var isTestFrame = (self.testFrameIndices.includes(self.posFrameImageIndex));
                     var outputDirPath = (isTestFrame ? self.validImagesDirPath : self.trainImagesDirPath)
                     var imageFilePath = path.join(outputDirPath, frameFileName);
                     fs.writeFile(imageFilePath, frameBuffer, cb);
+                    self.posFrameImageIndex++;
+                    if(isTestFrame) {
+                        self.testFrameNames.push(frameFileName);
+                    }
                 },
                 function saveLabel(cb) {
-                    var isTestFrame = (self.testFrameIndices.includes(self.posFrameIndex));
+                    var isTestFrame = (self.testFrameNames.includes(frameFileName));
                     var outputDirPath = (isTestFrame ? self.validLabelsDirPath : self.trainLabelsDirPath);
                     var labelFileName = path.parse(frameFileName).name + '.txt';
                     var labelFilePath = path.join(outputDirPath, labelFileName);
@@ -119,7 +127,7 @@ function Exporter(exportDirPath, classes, taggedFramesCount, frameWidth, frameHe
                                                  rotation_y.toFixed(1));
                     }
                     fs.writeFile(labelFilePath, labelData, cb);
-                    self.posFrameIndex++;
+                    self.posFrameLabelIndex++;
                 },
             ], (err) => {
                 if (err) {