From e9af63c57b15f1e7c0b46362d4796b7b181122e6 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sun, 27 Sep 2020 02:33:35 +0800
Subject: [PATCH 01/25] Update libs

---
 CVAuto       | 2 +-
 waicoolUtils | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/CVAuto b/CVAuto
index 883aef596e..1d26d384a6 160000
--- a/CVAuto
+++ b/CVAuto
@@ -1 +1 @@
-Subproject commit 883aef596ee79004e66d0dcce8b621d8d681e540
+Subproject commit 1d26d384a6693b700bf69fc28c7a89fc99a0aa79
diff --git a/waicoolUtils b/waicoolUtils
index 7d93114e0a..06ef2bd68c 160000
--- a/waicoolUtils
+++ b/waicoolUtils
@@ -1 +1 @@
-Subproject commit 7d93114e0ae99deb625f95ebc9fa86f744de3c46
+Subproject commit 06ef2bd68c1adeedfeb42953f4e34996e07e288b

From ad54520e4e3a105cb3061584bb90c66528f402ba Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Mon, 20 Jul 2020 16:05:14 +0800
Subject: [PATCH 02/25] DebugView: Add image dynamic resize

---
 .../waicool20/wai2k/util/ai/GFLModelLoader.kt | 22 +++++++++++++++++++
 src/main/resources/views/debug.fxml           |  4 ++--
 2 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
new file mode 100644
index 0000000000..b924f5ce4d
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
@@ -0,0 +1,22 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+
diff --git a/src/main/resources/views/debug.fxml b/src/main/resources/views/debug.fxml
index d638a54752..e91c38daa4 100644
--- a/src/main/resources/views/debug.fxml
+++ b/src/main/resources/views/debug.fxml
@@ -63,9 +63,9 @@
         </HBox>
       <Separator prefWidth="200.0" />
       <Label text="On-Screen OCR Testing:" />
-      <HBox alignment="CENTER" VBox.vgrow="ALWAYS">
+      <HBox fx:id="container" alignment="CENTER" minHeight="300.0" minWidth="600.0" VBox.vgrow="ALWAYS">
          <children>
-            <ImageView fx:id="ocrImageView" fitHeight="300.0" fitWidth="600.0" pickOnBounds="true" preserveRatio="true" HBox.hgrow="ALWAYS" />
+            <ImageView fx:id="ocrImageView" fitHeight="${container.height}" fitWidth="${container.width}" pickOnBounds="true" preserveRatio="true" HBox.hgrow="ALWAYS" />
          </children>
       </HBox>
       <HBox alignment="CENTER_LEFT" spacing="5.0">

From 55a9297ae38d53c7780e0fe74cbfd24f004aedb6 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Mon, 20 Jul 2020 16:06:55 +0800
Subject: [PATCH 03/25] AI: Implement model loading

---
 build.gradle.kts                              |   2 +
 .../waicool20/wai2k/util/ai/GFLModelLoader.kt |  33 ++++++
 .../wai2k/util/ai/YoloPreProcessor.kt         |  60 ++++++++++
 .../waicool20/wai2k/util/ai/YoloTranslator.kt | 105 ++++++++++++++++++
 4 files changed, 200 insertions(+)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt

diff --git a/build.gradle.kts b/build.gradle.kts
index 47ed62ba03..6b42217931 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -60,6 +60,8 @@ dependencies {
     implementation("org.controlsfx", "controlsfx", "8.40.14")
     implementation("org.reflections", "reflections", "0.9.12")
     implementation("com.squareup.okhttp3:okhttp:4.7.2")
+    implementation("ai.djl.pytorch:pytorch-engine:0.6.0")
+    implementation("ai.djl.pytorch:pytorch-native-auto:1.5.0")
 
     implementation("net.sourceforge.tess4j", "tess4j", "4.5.1") {
         exclude("org.ghost4j")
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
index b924f5ce4d..e8966e2fbd 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
@@ -19,4 +19,37 @@
 
 package com.waicool20.wai2k.util.ai
 
+import ai.djl.Device
+import ai.djl.Model
+import ai.djl.pytorch.engine.PtEngine
+import java.nio.file.Files
+import java.nio.file.Path
 
+object GFLModelLoader {
+    private val engine by lazy { PtEngine.getInstance() }
+
+    fun loadModel(path: Path): Model {
+        require(Files.isRegularFile(path)) { "Must be path to model file" }
+        require(Files.exists(path)) { "Model file does not exist" }
+        require("$path".endsWith(".pt")) { "Model must have .pt extension" }
+        val name = "${path.fileName}".dropLastWhile { it != '.' }.dropLast(1)
+        return engine.newModel(name, Device.cpu()).apply {
+            load(path.parent, name)
+            setProperty("InputSize", "640")
+            setProperty("Classes",
+                    listOf(
+                            "Node",
+                            "Command Post",
+                            "Heliport",
+                            "Enemy",
+                            "Sangvis Ferri",
+                            "Military",
+                            "Radar",
+                            "Paradeus",
+                            "Supply Crate",
+                            "Friendly"
+                    ).joinToString()
+            )
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt
new file mode 100644
index 0000000000..629d60d1f2
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt
@@ -0,0 +1,60 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.Model
+import ai.djl.modality.cv.util.NDImageUtils
+import ai.djl.ndarray.NDArray
+import ai.djl.ndarray.types.Shape
+import ai.djl.translate.Transform
+import kotlin.math.roundToInt
+
+class YoloPreProcessor(model: Model) : Transform {
+    private val size = model.getProperty("InputSize")?.toLong()
+            ?: error("Model property 'InputSize' must be set")
+
+    override fun transform(array: NDArray): NDArray {
+        // Array shape: height, width, layers
+        var out = array
+        val layers = out.shape[2]
+        val width = out.shape[1].toDouble()
+        val height = out.shape[0].toDouble()
+
+        when {
+            width > height -> {
+                val newHeight = (size * (height / width)).roundToInt()
+                out = NDImageUtils.resize(array, size.toInt(), newHeight)
+                val concatWith = out.manager.zeros(Shape(size - newHeight, size, layers))
+                out = out.concat(concatWith, 0)
+            }
+            height < width -> {
+                val newWidth = (size * (width / height)).roundToInt()
+                out = NDImageUtils.resize(array, newWidth, size.toInt())
+                val concatWith = out.manager.zeros(Shape(size, size - newWidth, layers))
+                out = out.concat(concatWith, 1)
+            }
+            width == height -> {
+                out = NDImageUtils.resize(array, size.toInt())
+            }
+        }
+        // Transpose to layers, height, width and normalize
+        return out.transpose(2, 0, 1).div(255)
+    }
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
new file mode 100644
index 0000000000..449f278239
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
@@ -0,0 +1,105 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.Model
+import ai.djl.modality.cv.output.DetectedObjects
+import ai.djl.modality.cv.output.Rectangle
+import ai.djl.modality.cv.translator.ObjectDetectionTranslator
+import ai.djl.ndarray.NDList
+import ai.djl.ndarray.index.NDIndex
+import ai.djl.translate.Pipeline
+import ai.djl.translate.TranslatorContext
+
+class YoloTranslator(
+        model: Model,
+        threshold: Double,
+        inputImageSize: Pair<Double, Double>,
+        val iouThreshold: Double = 0.5
+) : ObjectDetectionTranslator(
+        Builder()
+                .setPipeline(Pipeline(YoloPreProcessor(model)))
+                .optThreshold(threshold.toFloat())
+                .optSynset(model.getClasses() ?: error("Model property 'Classes' must be set"))
+                .optRescaleSize(inputImageSize.first, inputImageSize.second)
+) {
+    companion object {
+        private fun Model.getClasses(): List<String>? {
+            return getProperty("Classes")?.split(",")?.map { it.trim() }
+        }
+    }
+    private class Builder : ObjectDetectionTranslator.BaseBuilder<Builder>() {
+        override fun self() = this
+    }
+
+    override fun processOutput(ctx: TranslatorContext, list: NDList): DetectedObjects {
+        var output = list[0]
+        val inputArraySize = list[1].shape[1] * 32
+        val mask = output[NDIndex(":, 4")].gte(threshold).repeat(15).reshape(output.shape)
+        output = output.booleanMask(mask)
+        output = output.reshape(output.shape[0] / 15, 15)
+        val objects = mutableListOf<DetectedObjects.DetectedObject>()
+        for (i in 0 until output.shape[0]) {
+            // Array format is x1, y1, x2, y2, conf, cls
+            val detection = output[i].toFloatArray()
+            val (centerX, centerY, w, h, p) = detection
+            var x = (centerX - w / 2).toDouble() / inputArraySize
+            var y = (centerY - h / 2).toDouble() / inputArraySize
+            var width = w.toDouble() / inputArraySize
+            var height = h.toDouble() / inputArraySize
+
+            when {
+                imageWidth > imageHeight -> {
+                    val scale = imageWidth / imageHeight
+                    y *= scale
+                    height *= scale
+                }
+                imageWidth < imageHeight -> {
+                    val scale = imageHeight / imageWidth
+                    x *= scale
+                    width *= scale
+                }
+                imageWidth == imageHeight -> Unit // Do Nothing
+            }
+
+            val c = detection.slice(5..detection.lastIndex)
+            val cMaxIdx = c.indexOf(c.max())
+            objects.add(DetectedObjects.DetectedObject(classes[cMaxIdx], p.toDouble(), Rectangle(x, y, width, height)))
+        }
+        val keep = nms(objects)
+        return DetectedObjects(
+                keep.map { it.className },
+                keep.map { it.probability },
+                keep.map { it.boundingBox }
+        )
+    }
+
+    fun nms(boxes: List<DetectedObjects.DetectedObject>): List<DetectedObjects.DetectedObject> {
+        val input = boxes.toMutableList()
+        val output = mutableListOf<DetectedObjects.DetectedObject>()
+        while (input.isNotEmpty()) {
+            val best = input.maxBy { it.probability } ?: continue
+            input.remove(best)
+            input.removeAll { it.className == best.className && it.boundingBox.getIoU(best.boundingBox) >= iouThreshold }
+            output.add(best)
+        }
+        return output
+    }
+}

From a0a46132c1ccb83d40d2ad69f6f0335d929510f0 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Mon, 20 Jul 2020 16:07:37 +0800
Subject: [PATCH 04/25] DebugView: Implement node detect

---
 src/main/kotlin/com/waicool20/wai2k/Wai2K.kt  |  3 ++
 .../com/waicool20/wai2k/views/DebugView.kt    | 42 +++++++++++--------
 src/main/resources/views/debug.fxml           |  5 ---
 3 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/Wai2K.kt b/src/main/kotlin/com/waicool20/wai2k/Wai2K.kt
index 705006ac25..4b8f56ab82 100644
--- a/src/main/kotlin/com/waicool20/wai2k/Wai2K.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/Wai2K.kt
@@ -48,6 +48,9 @@ class Wai2K : App(Wai2KWorkspace::class) {
             } catch (t: Throwable) {
                 logger.warn("Could not set locale to C, application may crash if using tesseract 4.0+")
             }
+            // Set inference thread count, anything above 8 seems to be ignored
+            val cores = Runtime.getRuntime().availableProcessors().coerceAtMost(8)
+            CLib.setEnv("OMP_NUM_THREADS", cores.toString(), true)
         }
 
         private fun isRunningJar(): Boolean {
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
index 6409e5cc8f..4e3fbf730e 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
@@ -19,16 +19,17 @@
 
 package com.waicool20.wai2k.views
 
+import ai.djl.modality.cv.ImageFactory
 import com.waicool20.cvauto.android.ADB
 import com.waicool20.cvauto.android.AndroidDevice
 import com.waicool20.cvauto.core.Region
 import com.waicool20.cvauto.core.template.FileTemplate
-import com.waicool20.cvauto.util.asBufferedImage
 import com.waicool20.cvauto.util.asGrayF32
 import com.waicool20.wai2k.config.Wai2KContext
 import com.waicool20.wai2k.script.ScriptRunner
 import com.waicool20.wai2k.util.Ocr
-import com.waicool20.wai2k.util.extractNodes
+import com.waicool20.wai2k.util.ai.GFLModelLoader
+import com.waicool20.wai2k.util.ai.YoloTranslator
 import com.waicool20.wai2k.util.useCharFilter
 import com.waicool20.waicoolutils.javafx.CoroutineScopeView
 import com.waicool20.waicoolutils.javafx.addListener
@@ -43,6 +44,7 @@ import kotlinx.coroutines.*
 import kotlinx.coroutines.javafx.JavaFx
 import net.sourceforge.tess4j.ITesseract
 import tornadofx.*
+import java.awt.image.BufferedImage
 import java.nio.file.Files
 import java.nio.file.Paths
 import kotlin.time.ExperimentalTime
@@ -81,6 +83,15 @@ class DebugView : CoroutineScopeView() {
 
     private val logger = loggerFor<DebugView>()
 
+    private val predictor by lazy {
+        try {
+            val model = GFLModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("gfl.pt"))
+            model.newPredictor(YoloTranslator(model, 0.4, 2160.0 to 1080.0))
+        } catch (e: Exception) {
+            null
+        }
+    }
+
     init {
         title = "WAI2K - Debugging tools"
     }
@@ -137,20 +148,18 @@ class DebugView : CoroutineScopeView() {
                 }
             }
             while (isActive) {
-                var image = device.screens[0].capture().let {
-                    if (wSpinner.value > 0 && hSpinner.value > 0) {
-                        it.getSubimage(xSpinner.value, ySpinner.value, wSpinner.value, hSpinner.value)
-                    } else it
-                }
-                if (filterBlueCheckBox.isSelected || filterWhiteCheckBox.isSelected || filterYellowCheckBox.isSelected) {
-                    image = image.extractNodes(
-                        includeBlue = filterBlueCheckBox.isSelected,
-                        includeWhite = filterWhiteCheckBox.isSelected,
-                        includeYellow = filterYellowCheckBox.isSelected
-                    ).asBufferedImage()
-                }
-                withContext(Dispatchers.JavaFx) {
-                    ocrImageView.image = SwingFXUtils.toFXImage(image, null)
+                val predictor = this@DebugView.predictor
+                if (predictor == null) {
+                    withContext(Dispatchers.JavaFx) {
+                        ocrImageView.image = SwingFXUtils.toFXImage(device.screens[0].capture(), null)
+                    }
+                } else {
+                    val image = ImageFactory.getInstance().fromImage(device.screens[0].capture())
+                    val objects = predictor.predict(image)
+                    image.drawBoundingBoxes(objects)
+                    withContext(Dispatchers.JavaFx) {
+                        ocrImageView.image = SwingFXUtils.toFXImage(image.wrappedImage as BufferedImage, null)
+                    }
                 }
             }
         }
@@ -231,7 +240,6 @@ class DebugView : CoroutineScopeView() {
                 }
                 logger.info("Result: \n${getOCR().doOCR(image)}\n----------")
             }
-
         }
     }
 
diff --git a/src/main/resources/views/debug.fxml b/src/main/resources/views/debug.fxml
index e91c38daa4..28b1aa3241 100644
--- a/src/main/resources/views/debug.fxml
+++ b/src/main/resources/views/debug.fxml
@@ -72,11 +72,6 @@
          <children>
             <Button fx:id="OCRButton" mnemonicParsing="false" text="OCR" />
             <Button fx:id="resetOCRButton" mnemonicParsing="false" text="Reset" />
-            <Separator orientation="VERTICAL" />
-            <Label text="Filter Nodes:" />
-            <CheckBox fx:id="filterBlueCheckBox" mnemonicParsing="false" text="Blue" />
-            <CheckBox fx:id="filterWhiteCheckBox" mnemonicParsing="false" text="White" />
-            <CheckBox fx:id="filterYellowCheckBox" mnemonicParsing="false" text="Yellow" />
          </children>
       </HBox>
       <Separator prefWidth="200.0" />

From 0cdd22ec1082292e35f06601218c52d9fdcd779c Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Tue, 21 Jul 2020 22:19:07 +0800
Subject: [PATCH 05/25] AI: Implement annotator

---
 .../waicool20/wai2k/util/ai/GFLModelLoader.kt |  30 ++---
 .../com/waicool20/wai2k/views/DebugView.kt    |  77 ++++++++++-
 src/main/resources/views/debug.fxml           | 123 +++++++++---------
 3 files changed, 152 insertions(+), 78 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
index e8966e2fbd..78faa3dfa8 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
@@ -26,7 +26,20 @@ import java.nio.file.Files
 import java.nio.file.Path
 
 object GFLModelLoader {
-    private val engine by lazy { PtEngine.getInstance() }
+    val engine by lazy { PtEngine.getInstance() }
+
+    val classes = listOf(
+            "Node",
+            "Command Post",
+            "Heliport",
+            "Enemy",
+            "Sangvis Ferri",
+            "Military",
+            "Radar",
+            "Paradeus",
+            "Supply Crate",
+            "Friendly"
+    )
 
     fun loadModel(path: Path): Model {
         require(Files.isRegularFile(path)) { "Must be path to model file" }
@@ -36,20 +49,7 @@ object GFLModelLoader {
         return engine.newModel(name, Device.cpu()).apply {
             load(path.parent, name)
             setProperty("InputSize", "640")
-            setProperty("Classes",
-                    listOf(
-                            "Node",
-                            "Command Post",
-                            "Heliport",
-                            "Enemy",
-                            "Sangvis Ferri",
-                            "Military",
-                            "Radar",
-                            "Paradeus",
-                            "Supply Crate",
-                            "Friendly"
-                    ).joinToString()
-            )
+            setProperty("Classes", classes.joinToString())
         }
     }
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
index 4e3fbf730e..198edc9d8b 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
@@ -20,6 +20,7 @@
 package com.waicool20.wai2k.views
 
 import ai.djl.modality.cv.ImageFactory
+import ai.djl.modality.cv.output.DetectedObjects
 import com.waicool20.cvauto.android.ADB
 import com.waicool20.cvauto.android.AndroidDevice
 import com.waicool20.cvauto.core.Region
@@ -39,6 +40,7 @@ import javafx.scene.control.*
 import javafx.scene.control.SpinnerValueFactory.IntegerSpinnerValueFactory
 import javafx.scene.image.ImageView
 import javafx.scene.layout.VBox
+import javafx.stage.DirectoryChooser
 import javafx.stage.FileChooser
 import kotlinx.coroutines.*
 import kotlinx.coroutines.javafx.JavaFx
@@ -47,9 +49,16 @@ import tornadofx.*
 import java.awt.image.BufferedImage
 import java.nio.file.Files
 import java.nio.file.Paths
+import javax.xml.parsers.DocumentBuilderFactory
+import javax.xml.transform.OutputKeys
+import javax.xml.transform.TransformerFactory
+import javax.xml.transform.dom.DOMSource
+import javax.xml.transform.stream.StreamResult
+import kotlin.streams.asSequence
 import kotlin.time.ExperimentalTime
 import kotlin.time.measureTimedValue
 
+
 class DebugView : CoroutineScopeView() {
     override val root: VBox by fxml("/views/debug.fxml")
     private val openButton: Button by fxid()
@@ -64,9 +73,8 @@ class DebugView : CoroutineScopeView() {
     private val ocrImageView: ImageView by fxid()
     private val OCRButton: Button by fxid()
     private val resetOCRButton: Button by fxid()
-    private val filterBlueCheckBox: CheckBox by fxid()
-    private val filterWhiteCheckBox: CheckBox by fxid()
-    private val filterYellowCheckBox: CheckBox by fxid()
+    private val annotateSetButton: Button by fxid()
+    private val saveAnnotationsCheckBox: CheckBox by fxid()
 
     private val useLSTMCheckBox: CheckBox by fxid()
     private val filterCheckBox: CheckBox by fxid()
@@ -104,6 +112,7 @@ class DebugView : CoroutineScopeView() {
         assetOCRButton.setOnAction { doAssetOCR() }
         OCRButton.setOnAction { doOCR() }
         resetOCRButton.setOnAction { createNewRenderJob() }
+        annotateSetButton.setOnAction { annotateSet() }
     }
 
     private fun uiSetup() {
@@ -254,4 +263,66 @@ class DebugView : CoroutineScopeView() {
         }
         return ocr
     }
+
+    private fun annotateSet() {
+        launch(Dispatchers.IO) {
+            val predictor = predictor ?: return@launch
+
+            val dir = withContext(Dispatchers.JavaFx) {
+                DirectoryChooser().apply {
+                    title = "Annotate which directory?"
+                }.showDialog(null)?.toPath()
+            } ?: return@launch
+
+            val output = dir.resolve("out")
+            logger.info("Annotating images in $dir")
+            Files.createDirectories(output)
+
+            val doc = DocumentBuilderFactory.newInstance().newDocumentBuilder().newDocument()
+            val root = doc.createElement("annotations").also { doc.appendChild(it) }
+            val version = doc.createElement("version").apply {
+                appendChild(doc.createTextNode("1.1"))
+            }
+            root.appendChild(version)
+
+            Files.walk(dir).asSequence()
+                    .filterNot { it.parent.endsWith("out") }
+                    .filter { "$it".endsWith(".png", true) || "$it".endsWith(".jpg", true) }
+                    .sorted()
+                    .forEachIndexed { i, path ->
+                        val image = ImageFactory.getInstance().fromFile(path)
+                        val objects = predictor.predict(image)
+                        val imageNode = doc.createElement("image").apply {
+                            setAttribute("id", "$i")
+                            setAttribute("name", "${dir.parent.relativize(path)}")
+                            setAttribute("width", "${image.width}")
+                            setAttribute("height", "${image.height}")
+                        }
+                        objects.items<DetectedObjects.DetectedObject>().forEach { obj ->
+                            val bbox = obj.boundingBox.bounds
+                            doc.createElement("box").apply {
+                                setAttribute("label", obj.className)
+                                setAttribute("occluded", "0")
+                                setAttribute("xtl", "${bbox.x * image.width}")
+                                setAttribute("ytl", "${bbox.y * image.height}")
+                                setAttribute("xbr", "${(bbox.x + bbox.width) * image.width}")
+                                setAttribute("ybr", "${(bbox.y + bbox.height) * image.height}")
+                            }.also { imageNode.appendChild(it) }
+                        }
+                        root.appendChild(imageNode)
+                        if (saveAnnotationsCheckBox.isSelected) {
+                            image.drawBoundingBoxes(objects)
+                            image.save(Files.newOutputStream(output.resolve(path.fileName)), "png")
+                        }
+                        logger.info("Image: $path\n$objects")
+                    }
+
+            TransformerFactory.newInstance().newTransformer().apply {
+                setOutputProperty(OutputKeys.INDENT, "yes")
+                setOutputProperty("{http://xml.apache.org/xslt}indent-amount", "2");
+            }.transform(DOMSource(doc), StreamResult(Files.newOutputStream(output.resolve("annotations.xml"))))
+
+            logger.info("All annotations done")
+        }
+    }
 }
\ No newline at end of file
diff --git a/src/main/resources/views/debug.fxml b/src/main/resources/views/debug.fxml
index 28b1aa3241..89b60dc174 100644
--- a/src/main/resources/views/debug.fxml
+++ b/src/main/resources/views/debug.fxml
@@ -34,75 +34,78 @@
 
 <VBox maxHeight="-Infinity" maxWidth="-Infinity" minHeight="-Infinity" minWidth="-Infinity" spacing="10.0" xmlns="http://javafx.com/javafx/8.0.171" xmlns:fx="http://javafx.com/fxml/1">
     <children>
-      <HBox alignment="CENTER_LEFT" spacing="5.0">
-         <children>
-            <Label text="X:" />
-            <Spinner editable="true" prefWidth="100.0" fx:id="xSpinner" />
-            <Label text="Y:" />
-            <Spinner fx:id="ySpinner" editable="true" prefWidth="100.0" />
-            <Label text="W:" />
-            <Spinner fx:id="wSpinner" editable="true" prefWidth="100.0" />
-            <Label text="H:" />
-            <Spinner fx:id="hSpinner" editable="true" prefWidth="100.0" />
-         </children>
-         <VBox.margin>
-            <Insets left="10.0" />
-         </VBox.margin>
-      </HBox>
-      <Label text="Asset Testing:" />
+        <HBox alignment="CENTER_LEFT" spacing="5.0">
+            <children>
+                <Label text="X:" />
+                <Spinner editable="true" prefWidth="100.0" fx:id="xSpinner" />
+                <Label text="Y:" />
+                <Spinner fx:id="ySpinner" editable="true" prefWidth="100.0" />
+                <Label text="W:" />
+                <Spinner fx:id="wSpinner" editable="true" prefWidth="100.0" />
+                <Label text="H:" />
+                <Spinner fx:id="hSpinner" editable="true" prefWidth="100.0" />
+            </children>
+            <VBox.margin>
+                <Insets left="10.0" />
+            </VBox.margin>
+        </HBox>
+        <Label text="Asset Testing:" />
         <HBox alignment="CENTER_LEFT" spacing="10.0">
             <children>
                 <TextField fx:id="pathField" prefWidth="300.0" promptText="/path/to/asset" />
                 <Button fx:id="openButton" mnemonicParsing="false" text="Open" />
-              <Button fx:id="testButton" mnemonicParsing="false" text="Test" />
-            <Button fx:id="assetOCRButton" mnemonicParsing="false" text="OCR" />
+                <Button fx:id="testButton" mnemonicParsing="false" text="Test" />
+                <Button fx:id="assetOCRButton" mnemonicParsing="false" text="OCR" />
             </children>
             <VBox.margin>
                 <Insets left="10.0" />
             </VBox.margin>
         </HBox>
-      <Separator prefWidth="200.0" />
-      <Label text="On-Screen OCR Testing:" />
-      <HBox fx:id="container" alignment="CENTER" minHeight="300.0" minWidth="600.0" VBox.vgrow="ALWAYS">
-         <children>
-            <ImageView fx:id="ocrImageView" fitHeight="${container.height}" fitWidth="${container.width}" pickOnBounds="true" preserveRatio="true" HBox.hgrow="ALWAYS" />
-         </children>
-      </HBox>
-      <HBox alignment="CENTER_LEFT" spacing="5.0">
-         <children>
-            <Button fx:id="OCRButton" mnemonicParsing="false" text="OCR" />
-            <Button fx:id="resetOCRButton" mnemonicParsing="false" text="Reset" />
-         </children>
-      </HBox>
-      <Separator prefWidth="200.0" />
-      <Label text="OCR Options" />
-      <VBox alignment="CENTER_LEFT" spacing="15.0">
-         <children>
-            <CheckBox fx:id="useLSTMCheckBox" mnemonicParsing="false" text="Use LSTM Engine" />
-            <CheckBox fx:id="filterCheckBox" mnemonicParsing="false" text="Use filters:" />
-            <VBox fx:id="filterOptionsVBox" spacing="5.0">
-               <children>
-                  <RadioButton fx:id="digitsOnlyRadioButton" mnemonicParsing="false" selected="true" text="Digits Only [0-9]">
-                     <toggleGroup>
-                        <ToggleGroup fx:id="filterOptions" />
-                     </toggleGroup>
-                  </RadioButton>
-                  <HBox alignment="CENTER_LEFT" spacing="5.0">
-                     <children>
-                        <RadioButton fx:id="customRadioButton" mnemonicParsing="false" text="Custom:" toggleGroup="$filterOptions" />
-                        <TextField fx:id="allowedCharsTextField" promptText="Allowed Characters" />
-                     </children>
-                  </HBox>
-               </children>
-               <VBox.margin>
-                  <Insets left="10.0" />
-               </VBox.margin>
-            </VBox>
-         </children>
-         <VBox.margin>
-            <Insets left="10.0" />
-         </VBox.margin>
-      </VBox>
+        <Separator prefWidth="200.0" />
+        <Label text="Preview:" />
+        <HBox fx:id="container" alignment="CENTER" minHeight="300.0" minWidth="600.0" VBox.vgrow="ALWAYS">
+            <children>
+                <ImageView fx:id="ocrImageView" fitHeight="${container.height}" fitWidth="${container.width}" pickOnBounds="true" preserveRatio="true" HBox.hgrow="ALWAYS" />
+            </children>
+        </HBox>
+        <HBox alignment="CENTER_LEFT" spacing="5.0">
+            <children>
+                <Button fx:id="OCRButton" mnemonicParsing="false" text="OCR" />
+                <Button fx:id="resetOCRButton" mnemonicParsing="false" text="Reset" />
+                <Separator orientation="VERTICAL" />
+                <Button fx:id="annotateSetButton" mnemonicParsing="false" text="Annotate Set" />
+                <CheckBox fx:id="saveAnnotationsCheckBox" mnemonicParsing="false" text="Draw and Save Annotations" />
+            </children>
+        </HBox>
+        <Separator prefWidth="200.0" />
+        <Label text="OCR Options" />
+        <VBox alignment="CENTER_LEFT" spacing="15.0">
+            <children>
+                <CheckBox fx:id="useLSTMCheckBox" mnemonicParsing="false" text="Use LSTM Engine" />
+                <CheckBox fx:id="filterCheckBox" mnemonicParsing="false" text="Use filters:" />
+                <VBox fx:id="filterOptionsVBox" spacing="5.0">
+                    <children>
+                        <RadioButton fx:id="digitsOnlyRadioButton" mnemonicParsing="false" selected="true" text="Digits Only [0-9]">
+                            <toggleGroup>
+                                <ToggleGroup fx:id="filterOptions" />
+                            </toggleGroup>
+                        </RadioButton>
+                        <HBox alignment="CENTER_LEFT" spacing="5.0">
+                            <children>
+                                <RadioButton fx:id="customRadioButton" mnemonicParsing="false" text="Custom:" toggleGroup="$filterOptions" />
+                                <TextField fx:id="allowedCharsTextField" promptText="Allowed Characters" />
+                            </children>
+                        </HBox>
+                    </children>
+                    <VBox.margin>
+                        <Insets left="10.0" />
+                    </VBox.margin>
+                </VBox>
+            </children>
+            <VBox.margin>
+                <Insets left="10.0" />
+            </VBox.margin>
+        </VBox>
     </children>
     <padding>
         <Insets bottom="10.0" left="10.0" right="10.0" top="10.0" />

From 3a7fec246b927743e8abb908dc2afd78feb68c3b Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Wed, 22 Jul 2020 14:11:31 +0800
Subject: [PATCH 06/25] AI: Implement type hierarchy

---
 .../waicool20/wai2k/util/ai/GFLModelLoader.kt | 14 ----
 .../com/waicool20/wai2k/util/ai/GFLObject.kt  | 71 +++++++++++++++++++
 .../waicool20/wai2k/util/ai/YoloTranslator.kt | 55 +++++++-------
 .../com/waicool20/wai2k/views/DebugView.kt    | 14 ++--
 4 files changed, 105 insertions(+), 49 deletions(-)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
index 78faa3dfa8..16daea3ac0 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
@@ -28,19 +28,6 @@ import java.nio.file.Path
 object GFLModelLoader {
     val engine by lazy { PtEngine.getInstance() }
 
-    val classes = listOf(
-            "Node",
-            "Command Post",
-            "Heliport",
-            "Enemy",
-            "Sangvis Ferri",
-            "Military",
-            "Radar",
-            "Paradeus",
-            "Supply Crate",
-            "Friendly"
-    )
-
     fun loadModel(path: Path): Model {
         require(Files.isRegularFile(path)) { "Must be path to model file" }
         require(Files.exists(path)) { "Model file does not exist" }
@@ -49,7 +36,6 @@ object GFLModelLoader {
         return engine.newModel(name, Device.cpu()).apply {
             load(path.parent, name)
             setProperty("InputSize", "640")
-            setProperty("Classes", classes.joinToString())
         }
     }
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt
new file mode 100644
index 0000000000..54a46998a3
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt
@@ -0,0 +1,71 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.modality.cv.output.DetectedObjects
+import ai.djl.modality.cv.output.Rectangle
+
+sealed class GFLObject {
+    abstract val probability: Double
+    abstract val bbox: Rectangle
+
+    open class Node(override val probability: Double, override val bbox: Rectangle) : GFLObject()
+    class CommandPost(override val probability: Double, override val bbox: Rectangle) : Node(probability, bbox)
+    class Heliport(override val probability: Double, override val bbox: Rectangle) : Node(probability, bbox)
+    class SupplyCrate(override val probability: Double, override val bbox: Rectangle) : Node(probability, bbox)
+    class Radar(override val probability: Double, override val bbox: Rectangle) : Node(probability, bbox)
+
+    open class Enemy(override val probability: Double, override val bbox: Rectangle) : GFLObject()
+    class SangvisFerri(override val probability: Double, override val bbox: Rectangle) : Enemy(probability, bbox)
+    class Military(override val probability: Double, override val bbox: Rectangle) : Enemy(probability, bbox)
+    class Paradeus(override val probability: Double, override val bbox: Rectangle) : Enemy(probability, bbox)
+
+    open class Friendly(override val probability: Double, override val bbox: Rectangle) : GFLObject()
+
+    companion object {
+        /**
+         * All possible values of [GFLObject], must be declared in the order the model was created with
+         */
+        val values = listOf(
+                Node::class,
+                CommandPost::class,
+                Heliport::class,
+                Enemy::class,
+                SangvisFerri::class,
+                Military::class,
+                Radar::class,
+                Paradeus::class,
+                SupplyCrate::class,
+                Friendly::class
+        )
+    }
+
+    override fun toString(): String {
+        return this::class.simpleName?.replace(Regex("(.)([A-Z])"), "$1 $2") ?: "Unknown"
+    }
+}
+
+fun List<GFLObject>.toDetectedObjects(): DetectedObjects {
+    return DetectedObjects(
+            map { it.toString() },
+            map { it.probability },
+            map { it.bbox }
+    )
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
index 449f278239..b0ddeea92a 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
@@ -20,42 +20,34 @@
 package com.waicool20.wai2k.util.ai
 
 import ai.djl.Model
-import ai.djl.modality.cv.output.DetectedObjects
 import ai.djl.modality.cv.output.Rectangle
-import ai.djl.modality.cv.translator.ObjectDetectionTranslator
+import ai.djl.modality.cv.translator.BaseImageTranslator
 import ai.djl.ndarray.NDList
 import ai.djl.ndarray.index.NDIndex
 import ai.djl.translate.Pipeline
 import ai.djl.translate.TranslatorContext
+import kotlin.reflect.full.primaryConstructor
+import kotlin.reflect.full.superclasses
 
 class YoloTranslator(
         model: Model,
-        threshold: Double,
-        inputImageSize: Pair<Double, Double>,
+        val threshold: Double,
+        val inputImageSize: Pair<Double, Double>,
         val iouThreshold: Double = 0.5
-) : ObjectDetectionTranslator(
-        Builder()
-                .setPipeline(Pipeline(YoloPreProcessor(model)))
-                .optThreshold(threshold.toFloat())
-                .optSynset(model.getClasses() ?: error("Model property 'Classes' must be set"))
-                .optRescaleSize(inputImageSize.first, inputImageSize.second)
+) : BaseImageTranslator<List<GFLObject>>(
+        Builder().setPipeline(Pipeline(YoloPreProcessor(model)))
 ) {
-    companion object {
-        private fun Model.getClasses(): List<String>? {
-            return getProperty("Classes")?.split(",")?.map { it.trim() }
-        }
-    }
-    private class Builder : ObjectDetectionTranslator.BaseBuilder<Builder>() {
+    private class Builder : BaseImageTranslator.BaseBuilder<Builder>() {
         override fun self() = this
     }
 
-    override fun processOutput(ctx: TranslatorContext, list: NDList): DetectedObjects {
+    override fun processOutput(ctx: TranslatorContext, list: NDList): List<GFLObject> {
         var output = list[0]
         val inputArraySize = list[1].shape[1] * 32
         val mask = output[NDIndex(":, 4")].gte(threshold).repeat(15).reshape(output.shape)
         output = output.booleanMask(mask)
         output = output.reshape(output.shape[0] / 15, 15)
-        val objects = mutableListOf<DetectedObjects.DetectedObject>()
+        val objects = mutableListOf<GFLObject>()
         for (i in 0 until output.shape[0]) {
             // Array format is x1, y1, x2, y2, conf, cls
             val detection = output[i].toFloatArray()
@@ -65,6 +57,8 @@ class YoloTranslator(
             var width = w.toDouble() / inputArraySize
             var height = h.toDouble() / inputArraySize
 
+            val (imageWidth, imageHeight) = inputImageSize
+
             when {
                 imageWidth > imageHeight -> {
                     val scale = imageWidth / imageHeight
@@ -79,25 +73,30 @@ class YoloTranslator(
                 imageWidth == imageHeight -> Unit // Do Nothing
             }
 
+            x = x.coerceIn(0.0, 1.0)
+            y = y.coerceIn(0.0, 1.0)
+            width = width.coerceIn(0.0, 1.0)
+            height = height.coerceIn(0.0, 1.0)
+
             val c = detection.slice(5..detection.lastIndex)
             val cMaxIdx = c.indexOf(c.max())
-            objects.add(DetectedObjects.DetectedObject(classes[cMaxIdx], p.toDouble(), Rectangle(x, y, width, height)))
+            val obj = try {
+                GFLObject.values[cMaxIdx].primaryConstructor?.call(p.toDouble(), Rectangle(x, y, width, height))
+            } catch (e: Exception) {
+                null
+            }
+            if (obj != null) objects.add(obj)
         }
-        val keep = nms(objects)
-        return DetectedObjects(
-                keep.map { it.className },
-                keep.map { it.probability },
-                keep.map { it.boundingBox }
-        )
+        return nms(objects)
     }
 
-    fun nms(boxes: List<DetectedObjects.DetectedObject>): List<DetectedObjects.DetectedObject> {
+    fun nms(boxes: List<GFLObject>): List<GFLObject> {
         val input = boxes.toMutableList()
-        val output = mutableListOf<DetectedObjects.DetectedObject>()
+        val output = mutableListOf<GFLObject>()
         while (input.isNotEmpty()) {
             val best = input.maxBy { it.probability } ?: continue
             input.remove(best)
-            input.removeAll { it.className == best.className && it.boundingBox.getIoU(best.boundingBox) >= iouThreshold }
+            input.removeAll { it::class.superclasses == best::class.superclasses && it.bbox.getIoU(best.bbox) >= iouThreshold }
             output.add(best)
         }
         return output
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
index 198edc9d8b..e47783aec8 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
@@ -20,7 +20,6 @@
 package com.waicool20.wai2k.views
 
 import ai.djl.modality.cv.ImageFactory
-import ai.djl.modality.cv.output.DetectedObjects
 import com.waicool20.cvauto.android.ADB
 import com.waicool20.cvauto.android.AndroidDevice
 import com.waicool20.cvauto.core.Region
@@ -31,6 +30,7 @@ import com.waicool20.wai2k.script.ScriptRunner
 import com.waicool20.wai2k.util.Ocr
 import com.waicool20.wai2k.util.ai.GFLModelLoader
 import com.waicool20.wai2k.util.ai.YoloTranslator
+import com.waicool20.wai2k.util.ai.toDetectedObjects
 import com.waicool20.wai2k.util.useCharFilter
 import com.waicool20.waicoolutils.javafx.CoroutineScopeView
 import com.waicool20.waicoolutils.javafx.addListener
@@ -94,7 +94,7 @@ class DebugView : CoroutineScopeView() {
     private val predictor by lazy {
         try {
             val model = GFLModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("gfl.pt"))
-            model.newPredictor(YoloTranslator(model, 0.4, 2160.0 to 1080.0))
+            model.newPredictor(YoloTranslator(model, 0.6, 2160.0 to 1080.0))
         } catch (e: Exception) {
             null
         }
@@ -165,7 +165,7 @@ class DebugView : CoroutineScopeView() {
                 } else {
                     val image = ImageFactory.getInstance().fromImage(device.screens[0].capture())
                     val objects = predictor.predict(image)
-                    image.drawBoundingBoxes(objects)
+                    image.drawBoundingBoxes(objects.toDetectedObjects())
                     withContext(Dispatchers.JavaFx) {
                         ocrImageView.image = SwingFXUtils.toFXImage(image.wrappedImage as BufferedImage, null)
                     }
@@ -298,10 +298,10 @@ class DebugView : CoroutineScopeView() {
                             setAttribute("width", "${image.width}")
                             setAttribute("height", "${image.height}")
                         }
-                        objects.items<DetectedObjects.DetectedObject>().forEach { obj ->
-                            val bbox = obj.boundingBox.bounds
+                        objects.forEach { obj ->
+                            val bbox = obj.bbox
                             doc.createElement("box").apply {
-                                setAttribute("label", obj.className)
+                                setAttribute("label", "$obj")
                                 setAttribute("occluded", "0")
                                 setAttribute("xtl", "${bbox.x * image.width}")
                                 setAttribute("ytl", "${bbox.y * image.height}")
@@ -311,7 +311,7 @@ class DebugView : CoroutineScopeView() {
                         }
                         root.appendChild(imageNode)
                         if (saveAnnotationsCheckBox.isSelected) {
-                            image.drawBoundingBoxes(objects)
+                            image.drawBoundingBoxes(objects.toDetectedObjects())
                             image.save(Files.newOutputStream(output.resolve(path.fileName)), "png")
                         }
                         logger.info("Image: $path\n$objects")

From 4c82a1dc4c53f83cfa6dc579e0313c5cff583dc1 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Thu, 23 Jul 2020 02:24:36 +0800
Subject: [PATCH 07/25] AI: Fix same class overlap

---
 .../kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt  |  6 ++++--
 .../com/waicool20/wai2k/util/ai/YoloTranslator.kt    | 12 +++++++++---
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt
index 54a46998a3..cb7b77ad61 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLObject.kt
@@ -32,12 +32,14 @@ sealed class GFLObject {
     class SupplyCrate(override val probability: Double, override val bbox: Rectangle) : Node(probability, bbox)
     class Radar(override val probability: Double, override val bbox: Rectangle) : Node(probability, bbox)
 
-    open class Enemy(override val probability: Double, override val bbox: Rectangle) : GFLObject()
+    abstract class Unit : GFLObject()
+
+    open class Enemy(override val probability: Double, override val bbox: Rectangle) : Unit()
     class SangvisFerri(override val probability: Double, override val bbox: Rectangle) : Enemy(probability, bbox)
     class Military(override val probability: Double, override val bbox: Rectangle) : Enemy(probability, bbox)
     class Paradeus(override val probability: Double, override val bbox: Rectangle) : Enemy(probability, bbox)
 
-    open class Friendly(override val probability: Double, override val bbox: Rectangle) : GFLObject()
+    open class Friendly(override val probability: Double, override val bbox: Rectangle) : Unit()
 
     companion object {
         /**
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
index b0ddeea92a..5d61b7b0d8 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
@@ -26,14 +26,15 @@ import ai.djl.ndarray.NDList
 import ai.djl.ndarray.index.NDIndex
 import ai.djl.translate.Pipeline
 import ai.djl.translate.TranslatorContext
+import kotlin.reflect.full.isSubclassOf
+import kotlin.reflect.full.isSuperclassOf
 import kotlin.reflect.full.primaryConstructor
-import kotlin.reflect.full.superclasses
 
 class YoloTranslator(
         model: Model,
         val threshold: Double,
         val inputImageSize: Pair<Double, Double>,
-        val iouThreshold: Double = 0.5
+        val iouThreshold: Double = 0.4
 ) : BaseImageTranslator<List<GFLObject>>(
         Builder().setPipeline(Pipeline(YoloPreProcessor(model)))
 ) {
@@ -96,7 +97,12 @@ class YoloTranslator(
         while (input.isNotEmpty()) {
             val best = input.maxBy { it.probability } ?: continue
             input.remove(best)
-            input.removeAll { it::class.superclasses == best::class.superclasses && it.bbox.getIoU(best.bbox) >= iouThreshold }
+            input.removeAll {
+                (it::class == best::class ||
+                        it::class.isSubclassOf(best::class) ||
+                        it::class.isSuperclassOf(best::class)) &&
+                        it.bbox.getIoU(best.bbox) >= iouThreshold
+            }
             output.add(best)
         }
         return output

From 0e1c600447d5609622d754d34acd33af02d42b14 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 25 Jul 2020 14:35:54 +0800
Subject: [PATCH 08/25] Implement Coherent Point Drift

---
 .idea/inspectionProfiles/Project_Default.xml  |   2 +-
 .../wai2k/util/ai/NDArrayExtensions.kt        | 101 ++++++++++++++
 .../wai2k/util/ai/ShapeExtensions.kt          |  33 +++++
 .../wai2k/util/cpd/AffineCPDRegistration.kt   |  77 ++++++++++
 .../wai2k/util/cpd/CPDRegistration.kt         | 131 ++++++++++++++++++
 5 files changed, 343 insertions(+), 1 deletion(-)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/NDArrayExtensions.kt
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/ShapeExtensions.kt
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/cpd/AffineCPDRegistration.kt
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/cpd/CPDRegistration.kt

diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
index cfe33df0cb..bdf373599f 100644
--- a/.idea/inspectionProfiles/Project_Default.xml
+++ b/.idea/inspectionProfiles/Project_Default.xml
@@ -3,7 +3,7 @@
     <option name="myName" value="Project Default" />
     <inspection_tool class="DuplicatedCode" enabled="true" level="WEAK WARNING" enabled_by_default="true">
       <Languages>
-        <language minSize="47" name="Kotlin" />
+        <language minSize="64" name="Kotlin" />
       </Languages>
     </inspection_tool>
   </profile>
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/NDArrayExtensions.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/NDArrayExtensions.kt
new file mode 100644
index 0000000000..9155cd9c94
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/NDArrayExtensions.kt
@@ -0,0 +1,101 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.modality.cv.output.Point
+import ai.djl.ndarray.NDArray
+import ai.djl.ndarray.NDManager
+import ai.djl.ndarray.types.DataType
+
+operator fun NDArray.plus(array: NDArray): NDArray = add(array)
+operator fun NDArray.plus(number: Number): NDArray = add(number)
+operator fun NDArray.minus(array: NDArray): NDArray = sub(array)
+operator fun NDArray.minus(number: Number): NDArray = sub(number)
+operator fun NDArray.times(array: NDArray): NDArray = mul(array)
+operator fun NDArray.times(number: Number): NDArray = mul(number)
+
+operator fun NDArray.unaryMinus(): NDArray = neg()
+
+operator fun Number.plus(array: NDArray): NDArray = array + this
+operator fun Number.minus(array: NDArray): NDArray = array.neg() + this
+operator fun Number.times(array: NDArray): NDArray = array * this
+operator fun Number.div(array: NDArray): NDArray = array.pow(-1) * this
+
+operator fun NDArray.plusAssign(array: NDArray) {
+    addi(array)
+}
+
+operator fun NDArray.plusAssign(number: Number) {
+    addi(number)
+}
+
+operator fun NDArray.minusAssign(array: NDArray) {
+    subi(array)
+}
+
+operator fun NDArray.minusAssign(number: Number) {
+    subi(number)
+}
+
+operator fun NDArray.timesAssign(array: NDArray) {
+    muli(array)
+}
+
+operator fun NDArray.timesAssign(number: Number) {
+    muli(number)
+}
+
+operator fun NDArray.divAssign(array: NDArray) {
+    divi(array)
+}
+
+operator fun NDArray.divAssign(number: Number) {
+    divi(number)
+}
+
+fun NDArray.sum(vararg axes: Int) = sum(axes)
+
+fun NDArray.diag() = manager.eye(shape[0].toInt()) * tile(shape[0]).reshape(shape[0], shape[0])
+
+fun NDArray._trace() = (manager.eye(shape.dimension()) * this).sum()
+fun NDArray._stack(repeats: Long) = tile(repeats).reshape(repeats, shape[0])
+
+fun NDArray.toPoint2D(): List<Point> {
+    return when (dataType) {
+        DataType.FLOAT32 -> {
+            val array = toFloatArray()
+            List((size() / 2).toInt()) { i -> Point(array[i * 2].toDouble(), array[i * 2 + 1].toDouble()) }
+        }
+        DataType.FLOAT64 -> {
+            val array = toDoubleArray()
+            List((size() / 2).toInt()) { i -> Point(array[i * 2], array[i * 2 + 1]) }
+        }
+        else -> error("Unsupported data type: $dataType")
+    }
+}
+
+fun List<Point>.toNDArray(): NDArray {
+    val array = DoubleArray(size * 2)
+    for (i in indices) {
+        array[i * 2] = get(i).x
+        array[i * 2 + 1] = get(i).y
+    }
+    return NDManager.newBaseManager().create(array).reshape(size.toLong(), 2)
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/ShapeExtensions.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/ShapeExtensions.kt
new file mode 100644
index 0000000000..7ee7dd2d66
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/ShapeExtensions.kt
@@ -0,0 +1,33 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.ndarray.types.Shape
+
+operator fun Shape.component1() = get(0)
+operator fun Shape.component2() = get(1)
+operator fun Shape.component3() = get(2)
+operator fun Shape.component4() = get(3)
+operator fun Shape.component5() = get(4)
+operator fun Shape.component6() = get(5)
+operator fun Shape.component7() = get(6)
+operator fun Shape.component8() = get(7)
+operator fun Shape.component9() = get(8)
+operator fun Shape.component10() = get(9)
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/cpd/AffineCPDRegistration.kt b/src/main/kotlin/com/waicool20/wai2k/util/cpd/AffineCPDRegistration.kt
new file mode 100644
index 0000000000..8ae978c44e
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/cpd/AffineCPDRegistration.kt
@@ -0,0 +1,77 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.cpd
+
+import ai.djl.ndarray.NDArray
+import ai.djl.ndarray.types.Shape
+import com.waicool20.wai2k.util.ai.*
+import kotlin.math.abs
+import kotlin.math.ln
+
+/**
+ * Coherent Point Drift Algorithm, this estimates and registers a source point set to a target
+ * point set. This class provides affine transformation registration.
+ *
+ * Magical code ported from: https://github.com/siavashk/pycpd/blob/master/pycpd/affine_registration.py
+ */
+class AffineCPDRegistration(
+        X: NDArray,
+        Y: NDArray,
+        initB: NDArray? = null,
+        initT: NDArray? = null
+) : CPDRegistration(X, Y) {
+    private var B = initB ?: manager.eye(D.toInt())
+    private var t = initT ?: manager.zeros(Shape(1, D))
+    private lateinit var X_hat: NDArray
+    private lateinit var Y_hat: NDArray
+    private lateinit var A: NDArray
+    private lateinit var YPY: NDArray
+
+    override fun updateTransform() {
+        val muX = P.dot(target).sum(0) / Np
+        val muY = P.transpose().dot(source).sum(0) / Np
+
+        X_hat = target - muX._stack(N)
+        Y_hat = source - muY._stack(M)
+        A = X_hat.transpose().dot(P.transpose()).dot(Y_hat)
+        YPY = Y_hat.transpose().dot(P1.diag()).dot(Y_hat)
+        B = solve(YPY.transpose(), A.transpose())
+        t = muX.transpose() - (B.transpose().matMul(muY.transpose()))
+    }
+
+    override fun transformPointCloud() {
+        TY = source.dot(B) + t.tile(M).reshape(M, D)
+    }
+
+    override fun updateVariance() {
+        val qprev = error
+        // Cant use .trace(), not implemented in PtNDArray
+        val trAB = A.dot(B)._trace().getDouble()
+        val xPx = Pt1.transpose().dot((X_hat * X_hat).sum(1)).getDouble()
+        val trBYPYP = B.dot(YPY).dot(B)._trace().getDouble()
+        error = (xPx - 2 * trAB + trBYPYP) / (2 * sigma2) + D * Np / 2 * ln(sigma2)
+        diff = abs(error - qprev)
+        sigma2 = (xPx - trAB) / (Np * D)
+
+        if (sigma2 <= 0) {
+            sigma2 = tolerance / 10
+        }
+    }
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/cpd/CPDRegistration.kt b/src/main/kotlin/com/waicool20/wai2k/util/cpd/CPDRegistration.kt
new file mode 100644
index 0000000000..ab77e4ddd4
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/cpd/CPDRegistration.kt
@@ -0,0 +1,131 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.cpd
+
+import ai.djl.ndarray.NDArray
+import ai.djl.ndarray.NDManager
+import ai.djl.ndarray.types.Shape
+import com.waicool20.wai2k.util.ai.*
+import kotlin.math.PI
+import kotlin.math.pow
+
+/**
+ * Coherent Point Drift Algorithm, this estimates and registers a source point set to a target
+ * point set. This is the base class where all transformation types inherit from.
+ *
+ * Magical code ported from: https://github.com/siavashk/pycpd/blob/master/pycpd/emregistration.py
+ */
+abstract class CPDRegistration(
+        val target: NDArray,
+        val source: NDArray,
+        initSigma2: Double? = null,
+        val tolerance: Double = 0.001,
+        private val w: Double = 0.0
+) : Iterator<CPDRegistration>, Sequence<CPDRegistration> {
+    protected val manager = NDManager.newBaseManager()
+
+    init {
+        require(target.shape.dimension() == 2) { "Target point cloud must be 2D NDArray" }
+        require(source.shape.dimension() == 2) { "Source point cloud must be 2D NDArray" }
+        if (initSigma2 != null) {
+            require(initSigma2 > 0) { "Expected positive sigma2 value: ${this.sigma2}" }
+        }
+        require(tolerance > 0) { "Tolerance must be larger than 0: $tolerance" }
+        require(w in 0.0..1.0) { "w must be within 0 until 1: $w" }
+    }
+
+    protected var sigma2 = initSigma2 ?: run {
+        val (N, D) = target.shape
+        val (M, _) = source.shape
+        val diff = target.expandDims(0) - source.expandDims(1)
+        val err = diff.square()
+        (err.sum() / (D * M * N)).getDouble()
+    }
+
+    var TY = source
+        protected set
+    var error = Double.POSITIVE_INFINITY
+        protected set
+
+    protected val N = target.shape[0]
+    protected val D = target.shape[1]
+    protected val M = source.shape[0]
+
+    protected var diff = Double.POSITIVE_INFINITY
+    protected var P = manager.zeros(Shape(M, N))
+    protected var Pt1 = manager.zeros(Shape(N))
+    protected var P1 = manager.zeros(Shape(M))
+    protected var Np = 0.0
+
+    protected abstract fun updateTransform()
+    protected abstract fun transformPointCloud()
+    protected abstract fun updateVariance()
+
+    override fun iterator(): Iterator<CPDRegistration> {
+        return this
+    }
+
+    override fun hasNext() = true
+
+    override fun next(): CPDRegistration {
+        expectation()
+        maximization()
+        return this
+    }
+
+    fun getTranslatedPoints() = TY.toPoint2D()
+
+    private fun expectation() {
+        P = (target.expandDims(0) - TY.expandDims(1)).square().sum(2)
+
+        val c = run {
+            var c = (2 * PI * sigma2).pow(D / 2.0)
+            c = c * w / (1 - w)
+            c * M / N
+        }
+
+        P = (-P / (2 * sigma2)).exp()
+        val den = P.sum(0)._stack(M) + Double.MIN_VALUE + c
+
+        P = P / den
+        Pt1 = P.sum(0)
+        P1 = P.sum(1)
+        Np = P1.sum().getDouble()
+    }
+
+    private fun maximization() {
+        updateTransform()
+        transformPointCloud()
+        updateVariance()
+    }
+
+    protected fun solve(a: NDArray, b: NDArray): NDArray {
+        val A = a.toDoubleArray()
+        val detA = A[0] * A[3] - A[1] * A[2]
+        val temp = A[0]
+        A[0] = A[3]
+        A[1] = -A[1]
+        A[2] = -A[2]
+        A[3] = temp
+        a.set(A)
+        val invA = manager.create(A, a.shape) / detA
+        return invA.dot(b)
+    }
+}
\ No newline at end of file

From 5f96b10800a3da0ce91104d0c49bb5a0f4a1dad8 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 21:16:32 +0800
Subject: [PATCH 09/25] Update libs

Update libs
---
 build.gradle.kts | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/build.gradle.kts b/build.gradle.kts
index 6b42217931..e4efdde3ff 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -60,8 +60,8 @@ dependencies {
     implementation("org.controlsfx", "controlsfx", "8.40.14")
     implementation("org.reflections", "reflections", "0.9.12")
     implementation("com.squareup.okhttp3:okhttp:4.7.2")
-    implementation("ai.djl.pytorch:pytorch-engine:0.6.0")
-    implementation("ai.djl.pytorch:pytorch-native-auto:1.5.0")
+    implementation("ai.djl.pytorch:pytorch-engine:0.8.0")
+    implementation("ai.djl.pytorch:pytorch-native-auto:1.6.0")
 
     implementation("net.sourceforge.tess4j", "tess4j", "4.5.1") {
         exclude("org.ghost4j")

From d19037eceb5d3bc70e384e9bbc7971f89cbb79da Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Mon, 28 Sep 2020 01:45:17 +0800
Subject: [PATCH 10/25] LoaderView: Pre-load ai libs

---
 src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt b/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt
index cba1145564..780693406c 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt
@@ -29,6 +29,7 @@ import com.waicool20.wai2k.config.Wai2KContext
 import com.waicool20.wai2k.config.Wai2KProfile
 import com.waicool20.wai2k.script.ScriptContext
 import com.waicool20.wai2k.script.ScriptRunner
+import com.waicool20.wai2k.util.ai.GFLModelLoader
 import com.waicool20.waicoolutils.javafx.CoroutineScopeView
 import com.waicool20.waicoolutils.logging.LoggingEventBus
 import com.waicool20.waicoolutils.logging.loggerFor
@@ -86,6 +87,8 @@ class LoaderView : CoroutineScopeView() {
         loadWai2KProfile()
         loadScriptRunner()
         FileTemplate.checkPaths.add(wai2KConfig.assetsDirectory)
+        logger.info("Loading detection model...")
+        GFLModelLoader.loadModel(wai2KConfig.assetsDirectory.resolve("gfl.pt")).close()
         closeAndShowMainApp()
     }
 

From 32bd83fe444f65a5b407312cc618f3bcf421ec97 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Mon, 28 Sep 2020 01:45:43 +0800
Subject: [PATCH 11/25] Utils: Added more util extensions

---
 .../com/waicool20/wai2k/util/ImageUtils.kt    | 12 ++++---
 .../wai2k/util/ai/PredictorExtensions.kt      | 34 +++++++++++++++++++
 2 files changed, 42 insertions(+), 4 deletions(-)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/PredictorExtensions.kt

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ImageUtils.kt b/src/main/kotlin/com/waicool20/wai2k/util/ImageUtils.kt
index d085799895..b3d0195d74 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ImageUtils.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ImageUtils.kt
@@ -49,13 +49,17 @@ fun BufferedImage.extractNodes(
         val yellowNodes = hsv.clone().apply { hsvFilter(hueRange = 40..50, satRange = 12..100) }.getBand(2)
         img += yellowNodes
     }
+
+    return img.binarizeImage(0.75)
+}
+
+fun GrayF32.binarizeImage(threshold: Double = 0.4): GrayF32 {
     for (y in 0 until height) {
-        var index = img.startIndex + y * img.stride
+        var index = startIndex + y * stride
         for (x in 0 until width) {
-            img.data[index] = if (img.data[index] >= 175) 255f else 0f
+            data[index] = if (data[index] >= 255 * threshold) 255f else 0f
             index++
         }
     }
-
-    return img
+    return this
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/PredictorExtensions.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/PredictorExtensions.kt
new file mode 100644
index 0000000000..370c1aa235
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/PredictorExtensions.kt
@@ -0,0 +1,34 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.inference.Predictor
+import ai.djl.modality.cv.Image
+import ai.djl.modality.cv.ImageFactory
+import com.waicool20.cvauto.core.Region
+import java.awt.image.BufferedImage
+
+fun <T> Predictor<Image, T>.predict(image: BufferedImage): T {
+    return predict(ImageFactory.getInstance().fromImage(image))
+}
+
+fun <T> Predictor<Image, T>.predict(region: Region<*>): T {
+    return predict(ImageFactory.getInstance().fromImage(region.capture()))
+}
\ No newline at end of file

From b698e64bf47bbc30dcaf52e4cd7adbd9efb619e5 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Thu, 1 Oct 2020 02:07:55 +0800
Subject: [PATCH 12/25] YoloTranslator: Auto detect image width/height

---
 .../com/waicool20/wai2k/util/ai/YoloTranslator.kt   | 13 ++++++++++---
 .../kotlin/com/waicool20/wai2k/views/DebugView.kt   |  2 +-
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
index 5d61b7b0d8..9abc210269 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
@@ -20,6 +20,7 @@
 package com.waicool20.wai2k.util.ai
 
 import ai.djl.Model
+import ai.djl.modality.cv.Image
 import ai.djl.modality.cv.output.Rectangle
 import ai.djl.modality.cv.translator.BaseImageTranslator
 import ai.djl.ndarray.NDList
@@ -33,7 +34,6 @@ import kotlin.reflect.full.primaryConstructor
 class YoloTranslator(
         model: Model,
         val threshold: Double,
-        val inputImageSize: Pair<Double, Double>,
         val iouThreshold: Double = 0.4
 ) : BaseImageTranslator<List<GFLObject>>(
         Builder().setPipeline(Pipeline(YoloPreProcessor(model)))
@@ -42,6 +42,15 @@ class YoloTranslator(
         override fun self() = this
     }
 
+    private var imageWidth = -1
+    private var imageHeight = -1
+
+    override fun processInput(ctx: TranslatorContext, input: Image): NDList {
+        imageWidth = input.width
+        imageHeight = input.height
+        return super.processInput(ctx, input)
+    }
+
     override fun processOutput(ctx: TranslatorContext, list: NDList): List<GFLObject> {
         var output = list[0]
         val inputArraySize = list[1].shape[1] * 32
@@ -58,8 +67,6 @@ class YoloTranslator(
             var width = w.toDouble() / inputArraySize
             var height = h.toDouble() / inputArraySize
 
-            val (imageWidth, imageHeight) = inputImageSize
-
             when {
                 imageWidth > imageHeight -> {
                     val scale = imageWidth / imageHeight
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
index e47783aec8..85fb6b7527 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
@@ -94,7 +94,7 @@ class DebugView : CoroutineScopeView() {
     private val predictor by lazy {
         try {
             val model = GFLModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("gfl.pt"))
-            model.newPredictor(YoloTranslator(model, 0.6, 2160.0 to 1080.0))
+            model.newPredictor(YoloTranslator(model, 0.6))
         } catch (e: Exception) {
             null
         }

From 17b66ba300866633bbe86be2f9e41a4075852ef5 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 20:38:01 +0800
Subject: [PATCH 13/25] ModelLoader: Don't set additional properties; Renamed
 file

---
 .../wai2k/util/ai/{GFLModelLoader.kt => ModelLoader.kt}      | 3 +--
 src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt       | 5 +++--
 src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt      | 5 +++--
 3 files changed, 7 insertions(+), 6 deletions(-)
 rename src/main/kotlin/com/waicool20/wai2k/util/ai/{GFLModelLoader.kt => ModelLoader.kt} (95%)

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/ModelLoader.kt
similarity index 95%
rename from src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
rename to src/main/kotlin/com/waicool20/wai2k/util/ai/ModelLoader.kt
index 16daea3ac0..49929e3b2b 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/GFLModelLoader.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/ModelLoader.kt
@@ -25,7 +25,7 @@ import ai.djl.pytorch.engine.PtEngine
 import java.nio.file.Files
 import java.nio.file.Path
 
-object GFLModelLoader {
+object ModelLoader {
     val engine by lazy { PtEngine.getInstance() }
 
     fun loadModel(path: Path): Model {
@@ -35,7 +35,6 @@ object GFLModelLoader {
         val name = "${path.fileName}".dropLastWhile { it != '.' }.dropLast(1)
         return engine.newModel(name, Device.cpu()).apply {
             load(path.parent, name)
-            setProperty("InputSize", "640")
         }
     }
 }
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
index 85fb6b7527..e12de4e841 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
@@ -28,7 +28,7 @@ import com.waicool20.cvauto.util.asGrayF32
 import com.waicool20.wai2k.config.Wai2KContext
 import com.waicool20.wai2k.script.ScriptRunner
 import com.waicool20.wai2k.util.Ocr
-import com.waicool20.wai2k.util.ai.GFLModelLoader
+import com.waicool20.wai2k.util.ai.ModelLoader
 import com.waicool20.wai2k.util.ai.YoloTranslator
 import com.waicool20.wai2k.util.ai.toDetectedObjects
 import com.waicool20.wai2k.util.useCharFilter
@@ -93,7 +93,8 @@ class DebugView : CoroutineScopeView() {
 
     private val predictor by lazy {
         try {
-            val model = GFLModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("gfl.pt"))
+            val model = ModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("gfl.pt"))
+            model.setProperty("InputSize", "640")
             model.newPredictor(YoloTranslator(model, 0.6))
         } catch (e: Exception) {
             null
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt b/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt
index 780693406c..513533a39c 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/LoaderView.kt
@@ -19,6 +19,7 @@
 
 package com.waicool20.wai2k.views
 
+import ai.djl.Device
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.fasterxml.jackson.module.kotlin.readValue
 import com.waicool20.cvauto.android.ADB
@@ -29,7 +30,7 @@ import com.waicool20.wai2k.config.Wai2KContext
 import com.waicool20.wai2k.config.Wai2KProfile
 import com.waicool20.wai2k.script.ScriptContext
 import com.waicool20.wai2k.script.ScriptRunner
-import com.waicool20.wai2k.util.ai.GFLModelLoader
+import com.waicool20.wai2k.util.ai.ModelLoader
 import com.waicool20.waicoolutils.javafx.CoroutineScopeView
 import com.waicool20.waicoolutils.logging.LoggingEventBus
 import com.waicool20.waicoolutils.logging.loggerFor
@@ -88,7 +89,7 @@ class LoaderView : CoroutineScopeView() {
         loadScriptRunner()
         FileTemplate.checkPaths.add(wai2KConfig.assetsDirectory)
         logger.info("Loading detection model...")
-        GFLModelLoader.loadModel(wai2KConfig.assetsDirectory.resolve("gfl.pt")).close()
+        ModelLoader.engine.newModel("Loading", Device.defaultDevice()).close()
         closeAndShowMainApp()
     }
 

From 82e5af8a055f6d4a5417ba83103413cf2038c5bc Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 20:39:22 +0800
Subject: [PATCH 14/25] AI: Implemented matching model

---
 .../waicool20/wai2k/util/ai/MatchingModel.kt  |  73 ++++++++++++
 .../wai2k/util/ai/MatchingPreProcessor.kt     |  33 ++++++
 .../wai2k/util/ai/MatchingTranslator.kt       | 108 ++++++++++++++++++
 3 files changed, 214 insertions(+)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingModel.kt
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingModel.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingModel.kt
new file mode 100644
index 0000000000..533eed1c9c
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingModel.kt
@@ -0,0 +1,73 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.Device
+import ai.djl.Model
+import ai.djl.ndarray.NDList
+import ai.djl.ndarray.NDManager
+import ai.djl.ndarray.types.Shape
+import ai.djl.nn.AbstractBlock
+import ai.djl.pytorch.engine.PtEngine
+import ai.djl.training.ParameterStore
+import ai.djl.util.PairList
+import java.nio.file.Path
+
+class MatchingModel(
+    superpointModelPath: Path,
+    superglueModelPath: Path
+) : Model by engine.newModel("MatchingModel", Device.defaultDevice()) {
+    companion object {
+        private val engine by lazy { PtEngine.getInstance() }
+    }
+
+    private val superpoint = ModelLoader.loadModel(superpointModelPath)
+    private val superglue = ModelLoader.loadModel(superglueModelPath)
+
+    init {
+        block = object : AbstractBlock(2) {
+            override fun forward(parameterStore: ParameterStore, inputs: NDList, training: Boolean, params: PairList<String, Any>?): NDList {
+                val (img0, img1) = inputs
+                val pred0 = superpoint.block.forward(parameterStore, NDList(img0), training, params)
+                val pred1 = superpoint.block.forward(parameterStore, NDList(img1), training, params)
+                val data = NDList().apply {
+                    add(img0)
+                    addAll(pred0)
+                    add(img1)
+                    addAll(pred1)
+                }
+                return NDList().apply {
+                    addAll(pred0)
+                    addAll(pred1)
+                    addAll(superglue.block.forward(parameterStore, data, training, params))
+                }
+            }
+
+            override fun getOutputShapes(manager: NDManager, inputShapes: Array<out Shape>): Array<Shape> {
+                return emptyArray()
+            }
+        }
+    }
+
+    override fun close() {
+        superpoint.close()
+        superglue.close()
+    }
+}
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt
new file mode 100644
index 0000000000..6468a48159
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt
@@ -0,0 +1,33 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.modality.cv.transform.Resize
+import ai.djl.ndarray.NDArray
+import ai.djl.translate.Transform
+
+class MatchingPreProcessor(width: Int, height: Int): Transform {
+    private val resize = Resize(width, height)
+    override fun transform(array: NDArray): NDArray {
+        return resize
+            .transform(array)
+            .transpose(2, 0, 1) / 255
+    }
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt
new file mode 100644
index 0000000000..431b709208
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt
@@ -0,0 +1,108 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.modality.cv.Image
+import ai.djl.ndarray.NDList
+import ai.djl.ndarray.index.NDIndex
+import ai.djl.translate.Batchifier
+import ai.djl.translate.Pipeline
+import ai.djl.translate.Translator
+import ai.djl.translate.TranslatorContext
+import boofcv.struct.geo.AssociatedPair
+import com.waicool20.cvauto.util.wrapper.Config
+import com.waicool20.cvauto.util.wrapper.KFactoryMultiViewRobust
+import georegression.struct.homography.Homography2D_F64
+import georegression.struct.point.Point2D_F64
+
+class MatchingTranslator(
+    private val resizeWidth: Int = 640,
+    private val resizeHeight: Int = 480
+) : Translator<Pair<Image, Image>, Homography2D_F64> {
+
+    private var img0Width = -1
+    private var img0Height = -1
+    private var img1Width = -1
+    private var img1Height = -1
+
+    override fun getBatchifier() = Batchifier.STACK
+
+    override fun getPipeline() = Pipeline(MatchingPreProcessor(resizeWidth, resizeHeight))
+
+    override fun processInput(ctx: TranslatorContext, input: Pair<Image, Image>): NDList {
+        val (img0, img1) = input
+        img0Width = img0.width
+        img0Height = img0.height
+        img1Width = img1.width
+        img1Height = img1.height
+
+        val img0array = img0.toNDArray(ctx.ndManager, Image.Flag.GRAYSCALE)
+        val img1array = img1.toNDArray(ctx.ndManager, Image.Flag.GRAYSCALE)
+        return NDList().apply {
+            addAll(pipeline.transform(NDList(img0array)))
+            addAll(pipeline.transform(NDList(img1array)))
+        }
+    }
+
+    override fun processOutput(ctx: TranslatorContext, list: NDList): Homography2D_F64 {
+        var kpts0 = list[0]
+        val kpts1Temp = list[3]
+        var matches = list[6]
+        var conf = list[8]
+
+        val keepMask = matches.gt(-1)
+
+        matches = matches.booleanMask(keepMask)
+        conf = conf.booleanMask(keepMask)
+
+        kpts0 = kpts0.booleanMask(keepMask.stack(keepMask, 1), 1)
+        kpts0 = kpts0.reshape(kpts0.shape[0] / 2, 2)
+
+        var kpts1 = kpts0.zerosLike()
+
+        for ((i, j) in matches.toLongArray().withIndex()) {
+            kpts1.set(NDIndex(i.toLong()), kpts1Temp.get(j))
+        }
+
+        val scales0 = ctx.ndManager.create(floatArrayOf(img0Width.toFloat() / resizeWidth, img0Height.toFloat() / resizeHeight))
+        val scales1 = ctx.ndManager.create(floatArrayOf(img1Width.toFloat() / resizeWidth, img1Height.toFloat() / resizeHeight))
+
+        kpts0 = kpts0 * scales0
+        kpts1 = kpts1 * scales1
+
+        val points1 = kpts0.toPoint2D().map { Point2D_F64(it.x, it.y) }
+        val points2 = kpts1.toPoint2D().map { Point2D_F64(it.x, it.y) }
+        val confArr = conf.toFloatArray()
+
+        val pairs = mutableListOf<AssociatedPair>()
+        for (i in points1.indices) {
+            if (confArr[i] > 0.45) {
+                pairs.add(AssociatedPair(points1[i], points2[i], false))
+            }
+        }
+
+        val modelMatcher = KFactoryMultiViewRobust.homographyRansac(
+            Config.Ransac(60, 3.0)
+        )
+
+        check(modelMatcher.process(pairs)) { "Model matching failed!" }
+        return modelMatcher.modelParameters.copy()
+    }
+}
\ No newline at end of file

From 6f75f81b0d0d01097987379770b13bf50959d907 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 21:04:12 +0800
Subject: [PATCH 15/25] HomographyMapRunner: Use matching model

---
 .../modules/combat/HomographyMapRunner.kt     | 100 ++++++++----------
 1 file changed, 42 insertions(+), 58 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
index d19b1a9364..fa96ccad4b 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
@@ -19,22 +19,26 @@
 
 package com.waicool20.wai2k.script.modules.combat
 
-import boofcv.struct.image.GrayF32
+import ai.djl.metric.Metrics
+import ai.djl.modality.cv.Image
+import ai.djl.modality.cv.ImageFactory
 import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.fasterxml.jackson.module.kotlin.readValue
 import com.waicool20.cvauto.android.AndroidRegion
-import com.waicool20.cvauto.util.homography
 import com.waicool20.cvauto.util.transformRect
 import com.waicool20.wai2k.config.Wai2KConfig
 import com.waicool20.wai2k.config.Wai2KProfile
-import com.waicool20.wai2k.script.NodeNotFoundException
 import com.waicool20.wai2k.script.ScriptRunner
-import com.waicool20.wai2k.util.extractNodes
+import com.waicool20.wai2k.util.ai.MatchingModel
+import com.waicool20.wai2k.util.ai.MatchingTranslator
 import com.waicool20.waicoolutils.logging.loggerFor
 import georegression.struct.homography.Homography2D_F64
-import kotlinx.coroutines.*
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.async
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.runBlocking
 import java.nio.file.Files
-import javax.imageio.ImageIO
+import java.nio.file.Paths
 import kotlin.math.max
 import kotlin.math.pow
 import kotlin.random.Random
@@ -60,18 +64,24 @@ abstract class HomographyMapRunner(
         private const val maxSideDiff = 5.0
     }
 
+    private val metrics = Metrics()
+
+    private val predictor by lazy {
+        val model = MatchingModel(
+            config.assetsDirectory.resolve("SuperPoint.pt"),
+            config.assetsDirectory.resolve("SuperGlue.pt")
+        )
+        model.newPredictor(MatchingTranslator(480, 360)).also { it.setMetrics(metrics) }
+    }
+
     /**
      * Map homography cache
      */
     private var mapH: Homography2D_F64? = null
 
-    protected open val extractBlueNodes: Boolean = true
-    protected open val extractWhiteNodes: Boolean = false
-    protected open val extractYellowNodes: Boolean = true
-
     final override val nodes: List<MapNode>
 
-    val fullMap: GrayF32
+    val fullMap: Image
 
     init {
         val n = async(Dispatchers.IO) {
@@ -83,12 +93,7 @@ abstract class HomographyMapRunner(
             }
         }
         val fm = async(Dispatchers.IO) {
-            val path = config.assetsDirectory.resolve("$PREFIX/map.png")
-            if (Files.exists(path)) {
-                ImageIO.read(path.toFile()).extractNodes(extractBlueNodes, extractWhiteNodes, extractYellowNodes)
-            } else {
-                GrayF32()
-            }
+            ImageFactory.getInstance().fromFile(Paths.get("$PREFIX/map.png"))
         }
 
         nodes = runBlocking { n.await() }
@@ -97,27 +102,28 @@ abstract class HomographyMapRunner(
 
     override suspend fun MapNode.findRegion(): AndroidRegion {
         val window = mapRunnerRegions.window
-        var h: Homography2D_F64? = null
-        while (h == null) {
-            h = try {
-                mapH
-                    ?: fullMap.homography(window.capture().extractNodes(extractBlueNodes, extractWhiteNodes, extractYellowNodes))
-            } catch (e: IllegalStateException) {
-                continue
-            }
+
+        val h = mapH ?: run {
+            logger.info("Finding map transformation")
+            val prediction = predictor.predict(fullMap to ImageFactory.getInstance().fromImage(window.capture()))
+            logger.debug("Homography prediction metrics:")
+            logger.debug("Preprocess: ${metrics.latestMetric("Preprocess").value.toLong() / 1000000} ms")
+            logger.debug("Inference: ${metrics.latestMetric("Inference").value.toLong() / 1000000} ms")
+            logger.debug("Postprocess: ${metrics.latestMetric("Postprocess").value.toLong() / 1000000} ms")
+            logger.debug("Total: ${metrics.latestMetric("Total").value.toLong() / 1000000} ms")
+            mapH = prediction
+            prediction
         }
 
         suspend fun retry(): AndroidRegion {
-            if (Random.nextBoolean()) {
-                logger.info("Zoom out")
-                region.pinch(
-                    Random.nextInt(500, 700),
-                    Random.nextInt(300, 400),
-                    0.0,
-                    500
-                )
-                delay(1000)
-            }
+            logger.info("Zoom out")
+            region.pinch(
+                Random.nextInt(500, 700),
+                Random.nextInt(300, 400),
+                0.0,
+                500
+            )
+            delay(1000)
             mapH = null
             return findRegion()
         }
@@ -192,28 +198,6 @@ abstract class HomographyMapRunner(
             delay(200)
             return findRegion()
         }
-
-        while (isActive) {
-            val targets = mutableListOf<Pair<Int, Int>>()
-            val img = roi.capture().extractNodes()
-            for (y in 0 until img.height) {
-                var index = img.startIndex + y * img.stride
-                for (x in 0 until img.width) {
-                    if (img.data[index++] >= 175) targets += x to y
-                }
-            }
-            yield()
-            if (targets.isEmpty()) {
-                logger.debug("No targets found, retry")
-                if (Random.nextBoolean()) continue else return retry()
-            }
-            logger.debug("${targets.size} target candidates for node $this")
-            val target = targets.random().let { (cX, cY) ->
-                region.subRegionAs<AndroidRegion>(roi.x + cX, roi.y + cY, 5, 5)
-            }
-            logger.debug("Node target: (x=${target.x},y=${target.y})")
-            return target
-        }
-        throw NodeNotFoundException(this)
+        return roi
     }
 }
\ No newline at end of file

From 215d02c8952517410296babc810328e7b3503fb8 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 21:04:26 +0800
Subject: [PATCH 16/25] Map0-2: Add pre zoom

---
 .../wai2k/script/modules/combat/maps/Map0_2.kt     | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt
index 818a945e24..b6e0fe49e2 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt
@@ -25,7 +25,10 @@ import com.waicool20.wai2k.config.Wai2KProfile
 import com.waicool20.wai2k.script.ScriptRunner
 import com.waicool20.wai2k.script.modules.combat.HomographyMapRunner
 import com.waicool20.waicoolutils.logging.loggerFor
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.yield
+import kotlin.math.roundToLong
+import kotlin.random.Random
 
 class Map0_2(
     scriptRunner: ScriptRunner,
@@ -37,7 +40,16 @@ class Map0_2(
     override val isCorpseDraggingMap = true
 
     override suspend fun execute() {
-        nodes[2].findRegion() // Try focus boss node to get map centered
+        if (gameState.requiresMapInit) {
+            logger.info("Zoom out")
+            region.pinch(
+                Random.nextInt(900, 1000),
+                Random.nextInt(300, 400),
+                0.0,
+                1000)
+            delay((900 * gameState.delayCoefficient).roundToLong()) //Wait to settle
+            gameState.requiresMapInit = false
+        }
         val rEchelons = deployEchelons(nodes[14], nodes[13])
         mapRunnerRegions.startOperation.click(); yield()
         waitForGNKSplash()

From 03d83e80e831283d4158a97808bec49cce81ec19 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 21:06:54 +0800
Subject: [PATCH 17/25] Move models to asset sub directory

---
 .../wai2k/script/modules/combat/HomographyMapRunner.kt        | 4 ++--
 src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt        | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
index fa96ccad4b..a053bb08f7 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
@@ -68,8 +68,8 @@ abstract class HomographyMapRunner(
 
     private val predictor by lazy {
         val model = MatchingModel(
-            config.assetsDirectory.resolve("SuperPoint.pt"),
-            config.assetsDirectory.resolve("SuperGlue.pt")
+            config.assetsDirectory.resolve("models/SuperPoint.pt"),
+            config.assetsDirectory.resolve("models/SuperGlue.pt")
         )
         model.newPredictor(MatchingTranslator(480, 360)).also { it.setMetrics(metrics) }
     }
diff --git a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
index e12de4e841..1313c7fd5c 100644
--- a/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/views/DebugView.kt
@@ -93,7 +93,7 @@ class DebugView : CoroutineScopeView() {
 
     private val predictor by lazy {
         try {
-            val model = ModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("gfl.pt"))
+            val model = ModelLoader.loadModel(wai2KContext.wai2KConfig.assetsDirectory.resolve("models/gfl.pt"))
             model.setProperty("InputSize", "640")
             model.newPredictor(YoloTranslator(model, 0.6))
         } catch (e: Exception) {

From e3150beb2cf61bb03338f60f39f462524a5b9e12 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 21:33:59 +0800
Subject: [PATCH 18/25] Update libs

---
 CVAuto       | 2 +-
 waicoolUtils | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/CVAuto b/CVAuto
index 1d26d384a6..883aef596e 160000
--- a/CVAuto
+++ b/CVAuto
@@ -1 +1 @@
-Subproject commit 1d26d384a6693b700bf69fc28c7a89fc99a0aa79
+Subproject commit 883aef596ee79004e66d0dcce8b621d8d681e540
diff --git a/waicoolUtils b/waicoolUtils
index 06ef2bd68c..7d93114e0a 160000
--- a/waicoolUtils
+++ b/waicoolUtils
@@ -1 +1 @@
-Subproject commit 06ef2bd68c1adeedfeb42953f4e34996e07e288b
+Subproject commit 7d93114e0ae99deb625f95ebc9fa86f744de3c46

From 479c881b1c8b41df38dfc3adc7753a40c539f487 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Fri, 2 Oct 2020 21:55:56 +0800
Subject: [PATCH 19/25] HomographyMapRunner: Fix map file loading

---
 .../wai2k/script/modules/combat/HomographyMapRunner.kt       | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
index a053bb08f7..3b3cc2bb26 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
@@ -38,7 +38,6 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.runBlocking
 import java.nio.file.Files
-import java.nio.file.Paths
 import kotlin.math.max
 import kotlin.math.pow
 import kotlin.random.Random
@@ -58,7 +57,7 @@ abstract class HomographyMapRunner(
         private const val minScroll = 75
 
         /**
-         * Difference theresholds
+         * Difference thresholds
          */
         private const val maxMapDiff = 80.0
         private const val maxSideDiff = 5.0
@@ -93,7 +92,7 @@ abstract class HomographyMapRunner(
             }
         }
         val fm = async(Dispatchers.IO) {
-            ImageFactory.getInstance().fromFile(Paths.get("$PREFIX/map.png"))
+            ImageFactory.getInstance().fromFile(config.assetsDirectory.resolve("$PREFIX/map.png"))
         }
 
         nodes = runBlocking { n.await() }

From 9a03a61d1d137525f011399bb6df108d800a42eb Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 3 Oct 2020 00:13:48 +0800
Subject: [PATCH 20/25] YoloTranslator: Process graphics before turning into
 NDArray

---
 .../util/ai/TransposeNormalizeTransform.kt    | 29 +++++++
 .../wai2k/util/ai/YoloPreProcessor.kt         | 60 -------------
 .../waicool20/wai2k/util/ai/YoloTranslator.kt | 87 +++++++++++++++----
 3 files changed, 99 insertions(+), 77 deletions(-)
 create mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/TransposeNormalizeTransform.kt
 delete mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/TransposeNormalizeTransform.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/TransposeNormalizeTransform.kt
new file mode 100644
index 0000000000..326b5a609f
--- /dev/null
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/TransposeNormalizeTransform.kt
@@ -0,0 +1,29 @@
+/*
+ * GPLv3 License
+ *
+ *  Copyright (c) WAI2K by waicool20
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package com.waicool20.wai2k.util.ai
+
+import ai.djl.ndarray.NDArray
+import ai.djl.translate.Transform
+
+class TransposeNormalizeTransform: Transform {
+    override fun transform(array: NDArray): NDArray {
+        return array.transpose(2, 0, 1).div(255f)
+    }
+}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt
deleted file mode 100644
index 629d60d1f2..0000000000
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloPreProcessor.kt
+++ /dev/null
@@ -1,60 +0,0 @@
-/*
- * GPLv3 License
- *
- *  Copyright (c) WAI2K by waicool20
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program.  If not, see <http://www.gnu.org/licenses/>.
- */
-
-package com.waicool20.wai2k.util.ai
-
-import ai.djl.Model
-import ai.djl.modality.cv.util.NDImageUtils
-import ai.djl.ndarray.NDArray
-import ai.djl.ndarray.types.Shape
-import ai.djl.translate.Transform
-import kotlin.math.roundToInt
-
-class YoloPreProcessor(model: Model) : Transform {
-    private val size = model.getProperty("InputSize")?.toLong()
-            ?: error("Model property 'InputSize' must be set")
-
-    override fun transform(array: NDArray): NDArray {
-        // Array shape: height, width, layers
-        var out = array
-        val layers = out.shape[2]
-        val width = out.shape[1].toDouble()
-        val height = out.shape[0].toDouble()
-
-        when {
-            width > height -> {
-                val newHeight = (size * (height / width)).roundToInt()
-                out = NDImageUtils.resize(array, size.toInt(), newHeight)
-                val concatWith = out.manager.zeros(Shape(size - newHeight, size, layers))
-                out = out.concat(concatWith, 0)
-            }
-            height < width -> {
-                val newWidth = (size * (width / height)).roundToInt()
-                out = NDImageUtils.resize(array, newWidth, size.toInt())
-                val concatWith = out.manager.zeros(Shape(size, size - newWidth, layers))
-                out = out.concat(concatWith, 1)
-            }
-            width == height -> {
-                out = NDImageUtils.resize(array, size.toInt())
-            }
-        }
-        // Transpose to layers, height, width and normalize
-        return out.transpose(2, 0, 1).div(255)
-    }
-}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
index 9abc210269..e4d779840e 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/YoloTranslator.kt
@@ -21,12 +21,26 @@ package com.waicool20.wai2k.util.ai
 
 import ai.djl.Model
 import ai.djl.modality.cv.Image
+import ai.djl.modality.cv.ImageFactory
 import ai.djl.modality.cv.output.Rectangle
 import ai.djl.modality.cv.translator.BaseImageTranslator
+import ai.djl.modality.cv.util.NDImageUtils
+import ai.djl.ndarray.NDArray
 import ai.djl.ndarray.NDList
 import ai.djl.ndarray.index.NDIndex
+import ai.djl.ndarray.types.Shape
 import ai.djl.translate.Pipeline
+import ai.djl.translate.Transform
 import ai.djl.translate.TranslatorContext
+import com.waicool20.waicoolutils.createCompatibleCopy
+import com.waicool20.waicoolutils.pad
+import java.awt.Color
+import java.awt.Graphics2D
+import java.awt.RenderingHints
+import java.awt.image.BufferedImage
+import java.awt.image.ImageObserver
+import java.util.concurrent.CountDownLatch
+import kotlin.math.roundToInt
 import kotlin.reflect.full.isSubclassOf
 import kotlin.reflect.full.isSuperclassOf
 import kotlin.reflect.full.primaryConstructor
@@ -36,19 +50,51 @@ class YoloTranslator(
         val threshold: Double,
         val iouThreshold: Double = 0.4
 ) : BaseImageTranslator<List<GFLObject>>(
-        Builder().setPipeline(Pipeline(YoloPreProcessor(model)))
+        Builder().setPipeline(Pipeline(TransposeNormalizeTransform()))
 ) {
+    private val size = model.getProperty("InputSize")?.toInt()
+        ?: error("Model property 'InputSize' must be set")
+
     private class Builder : BaseImageTranslator.BaseBuilder<Builder>() {
         override fun self() = this
     }
 
-    private var imageWidth = -1
-    private var imageHeight = -1
+    private var imageWidth = -1.0
+    private var imageHeight = -1.0
 
     override fun processInput(ctx: TranslatorContext, input: Image): NDList {
-        imageWidth = input.width
-        imageHeight = input.height
-        return super.processInput(ctx, input)
+        imageWidth = input.width.toDouble()
+        imageHeight = input.height.toDouble()
+
+        val inputImage= input.wrappedImage as BufferedImage
+        val networkInput = inputImage.createCompatibleCopy(size, size)
+        val g = (networkInput.graphics as Graphics2D).apply {
+            setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR)
+            paint = Color.BLACK
+            fillRect(0, 0, size, size)
+        }
+
+        if (input.width < size && input.height < size) {
+            g.drawImage(inputImage, 0, 0, null)
+        } else {
+            val width = input.width.toDouble()
+            val height = input.height.toDouble()
+            when {
+                width > height -> {
+                    val newHeight = (size * (height / width)).roundToInt()
+                    g.drawImage(inputImage, 0, 0, size, newHeight, null)
+                }
+                height < width -> {
+                    val newWidth = (size * (width / height)).roundToInt()
+                    g.drawImage(inputImage, 0, 0, newWidth, size, null)
+                }
+                width == height -> {
+                    g.drawImage(inputImage, 0, 0, size, size, null)
+                }
+            }
+        }
+        g.dispose()
+        return super.processInput(ctx, ImageFactory.getInstance().fromImage(networkInput))
     }
 
     override fun processOutput(ctx: TranslatorContext, list: NDList): List<GFLObject> {
@@ -67,18 +113,25 @@ class YoloTranslator(
             var width = w.toDouble() / inputArraySize
             var height = h.toDouble() / inputArraySize
 
-            when {
-                imageWidth > imageHeight -> {
-                    val scale = imageWidth / imageHeight
-                    y *= scale
-                    height *= scale
-                }
-                imageWidth < imageHeight -> {
-                    val scale = imageHeight / imageWidth
-                    x *= scale
-                    width *= scale
+            if (imageWidth < size && imageHeight < size) {
+                x *= size / imageWidth
+                y *= size / imageHeight
+                width *= size / imageWidth
+                height *= size / imageHeight
+            } else {
+                when {
+                    imageWidth > imageHeight -> {
+                        val scale = imageWidth / imageHeight
+                        y *= scale
+                        height *= scale
+                    }
+                    imageWidth < imageHeight -> {
+                        val scale = imageHeight / imageWidth
+                        x *= scale
+                        width *= scale
+                    }
+                    imageWidth == imageHeight -> Unit // Do Nothing
                 }
-                imageWidth == imageHeight -> Unit // Do Nothing
             }
 
             x = x.coerceIn(0.0, 1.0)

From c0d439396f02f7c51376937902703830a0b6e9f0 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 3 Oct 2020 00:17:54 +0800
Subject: [PATCH 21/25] Removed MatchingPreProcessor

---
 .../wai2k/util/ai/MatchingPreProcessor.kt     | 33 -------------------
 .../wai2k/util/ai/MatchingTranslator.kt       |  3 +-
 2 files changed, 2 insertions(+), 34 deletions(-)
 delete mode 100644 src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt

diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt
deleted file mode 100644
index 6468a48159..0000000000
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingPreProcessor.kt
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * GPLv3 License
- *
- *  Copyright (c) WAI2K by waicool20
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program.  If not, see <http://www.gnu.org/licenses/>.
- */
-
-package com.waicool20.wai2k.util.ai
-
-import ai.djl.modality.cv.transform.Resize
-import ai.djl.ndarray.NDArray
-import ai.djl.translate.Transform
-
-class MatchingPreProcessor(width: Int, height: Int): Transform {
-    private val resize = Resize(width, height)
-    override fun transform(array: NDArray): NDArray {
-        return resize
-            .transform(array)
-            .transpose(2, 0, 1) / 255
-    }
-}
\ No newline at end of file
diff --git a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt
index 431b709208..cc58d69df8 100644
--- a/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/util/ai/MatchingTranslator.kt
@@ -20,6 +20,7 @@
 package com.waicool20.wai2k.util.ai
 
 import ai.djl.modality.cv.Image
+import ai.djl.modality.cv.transform.Resize
 import ai.djl.ndarray.NDList
 import ai.djl.ndarray.index.NDIndex
 import ai.djl.translate.Batchifier
@@ -44,7 +45,7 @@ class MatchingTranslator(
 
     override fun getBatchifier() = Batchifier.STACK
 
-    override fun getPipeline() = Pipeline(MatchingPreProcessor(resizeWidth, resizeHeight))
+    override fun getPipeline() = Pipeline(Resize(resizeWidth, resizeHeight), TransposeNormalizeTransform())
 
     override fun processInput(ctx: TranslatorContext, input: Pair<Image, Image>): NDList {
         val (img0, img1) = input

From a7dfbed02e00b0e41835825c09f984c8d2548bc9 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 3 Oct 2020 00:26:10 +0800
Subject: [PATCH 22/25] MapRunner: Add map lifecycle function begin/cleanup

---
 .../wai2k/script/modules/CombatSimModule.kt     |  2 +-
 .../wai2k/script/modules/combat/MapRunner.kt    | 17 +++++++++++++++--
 .../wai2k/script/modules/combat/maps/Map0_2.kt  |  2 +-
 .../script/modules/combat/maps/Map0_2_EX.kt     |  2 +-
 .../script/modules/combat/maps/Map10_4E.kt      |  2 +-
 .../script/modules/combat/maps/Map10_4E_Drag.kt |  2 +-
 .../wai2k/script/modules/combat/maps/Map11_5.kt |  2 +-
 .../wai2k/script/modules/combat/maps/Map1_6.kt  |  2 +-
 .../wai2k/script/modules/combat/maps/Map2_6.kt  |  2 +-
 .../wai2k/script/modules/combat/maps/Map3_4E.kt |  2 +-
 .../wai2k/script/modules/combat/maps/Map3_6.kt  |  2 +-
 .../wai2k/script/modules/combat/maps/Map4_6.kt  |  2 +-
 .../script/modules/combat/maps/Map4_6_Data.kt   |  2 +-
 .../wai2k/script/modules/combat/maps/Map5_6.kt  |  2 +-
 .../wai2k/script/modules/combat/maps/Map6_4N.kt |  2 +-
 .../wai2k/script/modules/combat/maps/Map6_6.kt  |  2 +-
 .../wai2k/script/modules/combat/maps/Map8_1N.kt |  2 +-
 17 files changed, 31 insertions(+), 18 deletions(-)

diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/CombatSimModule.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/CombatSimModule.kt
index ffed112247..bcbbe47418 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/CombatSimModule.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/CombatSimModule.kt
@@ -64,7 +64,7 @@ class CombatSimModule(
 
     private val mapRunner = object : AbsoluteMapRunner(scriptRunner, region, config, profile) {
         override val isCorpseDraggingMap = false
-        override suspend fun execute() {
+        override suspend fun begin() {
             if (profile.combatSimulation.neuralFragment == Level.OFF) return
 
             if (OffsetDateTime.now(ZoneOffset.ofHours(-8)).dayOfWeek !in dataSimDays) {
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/MapRunner.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/MapRunner.kt
index 91aee6b4ea..ee7223dd9b 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/MapRunner.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/MapRunner.kt
@@ -116,9 +116,22 @@ abstract class MapRunner(
     abstract val isCorpseDraggingMap: Boolean
 
     /**
-     * Main execution function that is executed when map is entered
+     * Main run function that goes through whole life cycle of MapRunner
      */
-    abstract suspend fun execute()
+    suspend fun execute() {
+        begin()
+        cleanup()
+    }
+
+    /**
+     * Function that is executed when map is entered
+     */
+    abstract suspend fun begin()
+
+    /**
+     * Cleanup function run after execute()
+     */
+    open suspend fun cleanup() = Unit
 
     /**
      * Executes when entering a battle
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt
index b6e0fe49e2..71c7e08fc4 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2.kt
@@ -39,7 +39,7 @@ class Map0_2(
     private val logger = loggerFor<Map0_2>()
     override val isCorpseDraggingMap = true
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             region.pinch(
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2_EX.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2_EX.kt
index 6b58b58a99..7daca659e9 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2_EX.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map0_2_EX.kt
@@ -40,7 +40,7 @@ class Map0_2_EX(
     private val logger = loggerFor<Map0_2_EX>()
     override val isCorpseDraggingMap = true
 
-    override suspend fun execute() {
+    override suspend fun begin() {
 
         if (gameState.requiresMapInit) {
             // Check to see if its already good
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E.kt
index b5c3f2f0d6..7cbdbc1f07 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E.kt
@@ -38,7 +38,7 @@ class Map10_4E(
     private val logger = loggerFor<Map10_4E>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             repeat(2) {
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E_Drag.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E_Drag.kt
index 314874cb91..c31da4f5e0 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E_Drag.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map10_4E_Drag.kt
@@ -43,7 +43,7 @@ class Map10_4E_Drag(
     private val logger = loggerFor<Map10_4E_Drag>()
     override val isCorpseDraggingMap = true
 
-    override suspend fun execute() {
+    override suspend fun begin() {
 
         // Mostly empty region to the left
         val r = region.subRegionAs<AndroidRegion>(300, 500, 150, 8)
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map11_5.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map11_5.kt
index 183c5afe70..e96d2193d3 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map11_5.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map11_5.kt
@@ -38,7 +38,7 @@ class Map11_5(
     private val logger = loggerFor<Map11_5>()
     override val isCorpseDraggingMap = true
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         // No need to zoom, delay for map lag
         delay((1000 * gameState.delayCoefficient).roundToLong())
         val rEchelons = deployEchelons(nodes[1], nodes[0])
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map1_6.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map1_6.kt
index a40207e526..ac1a26be4e 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map1_6.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map1_6.kt
@@ -40,7 +40,7 @@ class Map1_6(
     private val logger = loggerFor<Map1_6>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         logger.info("Zoom out")
         region.pinch(
             Random.nextInt(700, 800),
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map2_6.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map2_6.kt
index e044b71840..143e31e296 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map2_6.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map2_6.kt
@@ -39,7 +39,7 @@ class Map2_6(
     private val logger = loggerFor<Map2_6>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         logger.info("Zoom out")
         region.pinch(
             Random.nextInt(700, 800),
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_4E.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_4E.kt
index b5a25169a1..981c859c2d 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_4E.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_4E.kt
@@ -38,7 +38,7 @@ class Map3_4E(
     private val logger = loggerFor<Map3_4E>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         logger.info("Zoom out")
         region.pinch(
             Random.nextInt(700, 800),
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_6.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_6.kt
index fd5841f4e9..de97e10ee5 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_6.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map3_6.kt
@@ -39,7 +39,7 @@ class Map3_6(
     private val logger = loggerFor<Map3_6>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         logger.info("Zoom out")
         region.pinch(
             Random.nextInt(700, 800),
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6.kt
index e76a6bf63d..c9a31cede8 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6.kt
@@ -38,7 +38,7 @@ class Map4_6(
     private val logger = loggerFor<Map4_6>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             repeat(2) {
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6_Data.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6_Data.kt
index d102e45197..8baf782ec7 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6_Data.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map4_6_Data.kt
@@ -42,7 +42,7 @@ class Map4_6_Data(
     //Allow interruption of waiting for turn if necessary
     private var combatComplete = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             repeat(2) {
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map5_6.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map5_6.kt
index c4fba3c437..bdea64ea90 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map5_6.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map5_6.kt
@@ -38,7 +38,7 @@ class Map5_6(
     private val logger = loggerFor<Map5_6>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             repeat(2) {
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_4N.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_4N.kt
index 685ffacd0a..8966f1102a 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_4N.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_4N.kt
@@ -38,7 +38,7 @@ class Map6_4N(
     private val logger = loggerFor<Map6_4N>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             zoom()
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_6.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_6.kt
index fde6014794..a6ab9992a1 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_6.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map6_6.kt
@@ -38,7 +38,7 @@ class Map6_6(
     private val logger = loggerFor<Map6_6>()
     override val isCorpseDraggingMap = false
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             repeat(2) {
diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map8_1N.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map8_1N.kt
index 3b3ccd424d..226b1a0a27 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map8_1N.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/maps/Map8_1N.kt
@@ -38,7 +38,7 @@ class Map8_1N(
     private val logger = loggerFor<Map8_1N>()
     override val isCorpseDraggingMap = true
 
-    override suspend fun execute() {
+    override suspend fun begin() {
         if (gameState.requiresMapInit) {
             logger.info("Zoom out")
             repeat(2) {

From 3fe6d8c5a07d0712182110d76c4a321bfde49447 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 3 Oct 2020 00:27:11 +0800
Subject: [PATCH 23/25] HomographyMapRunner: Clear homography cache when map is
 done

---
 .../wai2k/script/modules/combat/HomographyMapRunner.kt        | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
index 3b3cc2bb26..3460ee3574 100644
--- a/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
+++ b/src/main/kotlin/com/waicool20/wai2k/script/modules/combat/HomographyMapRunner.kt
@@ -99,6 +99,10 @@ abstract class HomographyMapRunner(
         fullMap = runBlocking { fm.await() }
     }
 
+    override suspend fun cleanup() {
+        mapH = null
+    }
+
     override suspend fun MapNode.findRegion(): AndroidRegion {
         val window = mapRunnerRegions.window
 

From 7db6ec6262da2d28009471489a12c92e996ee4b3 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 3 Oct 2020 02:08:17 +0800
Subject: [PATCH 24/25] Ignore model files

---
 .gitignore | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.gitignore b/.gitignore
index 855557650f..fb6fe2a03c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -159,3 +159,5 @@ gradle-app.setting
 /wai2k/
 
 .idea/modules.xml
+
+assets/models/

From 601764bb43189a99bfe0cebe51bdca380a1e2c61 Mon Sep 17 00:00:00 2001
From: waicool20 <waicool20@gmail.com>
Date: Sat, 3 Oct 2020 02:09:26 +0800
Subject: [PATCH 25/25] Create deploy artifacts in build sub-folder; Generate
 md5 sum for each artifact

---
 build.gradle.kts          | 34 +++++++++++++++++++++++++++++++---
 launcher/build.gradle.kts | 14 ++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/build.gradle.kts b/build.gradle.kts
index e4efdde3ff..558fd66884 100644
--- a/build.gradle.kts
+++ b/build.gradle.kts
@@ -23,6 +23,8 @@ import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
 import java.nio.file.Files
 import java.nio.file.Paths
 import java.nio.file.StandardOpenOption
+import java.security.MessageDigest
+import javax.xml.bind.DatatypeConverter
 
 plugins {
     java
@@ -100,11 +102,16 @@ tasks {
         dependencies {
             include { it.moduleGroup.startsWith("com.waicool20") }
         }
+        doLast { md5sum(archiveFile.get()) }
     }
 }
 
+task("prepare-deploy") {
+    dependsOn("build", "deps-list", "packAssets")
+}
+
 task("deps-list") {
-    val file = Paths.get("$projectDir/build/dependencies.txt")
+    val file = Paths.get("$buildDir/deploy/dependencies.txt")
     doFirst {
         if (Files.notExists(file)) {
             Files.createDirectories(file.parent)
@@ -143,7 +150,28 @@ task("versioning") {
 
 task<Zip>("packAssets") {
     archiveFileName.set("assets.zip")
-    destinationDirectory.set(file("$buildDir"))
+    destinationDirectory.set(file("$buildDir/deploy/"))
+
+    from(projectDir)
+    include("/assets/**")
+    exclude("/assets/models/**")
+    doLast { md5sum(archiveFile.get()) }
+}
+
+task<Zip>("packModels") {
+    archiveFileName.set("models.zip")
+    destinationDirectory.set(file("$buildDir/deploy/"))
+
+    from(projectDir)
+    include("/assets/models/**")
+    doLast { md5sum(archiveFile.get()) }
+}
 
-    from("$projectDir/assets")
+fun md5sum(file: RegularFile) {
+    val path = file.asFile.toPath()
+    val md5File = Paths.get("$path.md5")
+    val md5sum = MessageDigest.getInstance("MD5")
+        .digest(Files.readAllBytes(path))
+        .let { DatatypeConverter.printHexBinary(it) }
+    Files.write(md5File, md5sum.toByteArray())
 }
\ No newline at end of file
diff --git a/launcher/build.gradle.kts b/launcher/build.gradle.kts
index cbf92dcf2a..ed796284a0 100644
--- a/launcher/build.gradle.kts
+++ b/launcher/build.gradle.kts
@@ -25,6 +25,10 @@
 import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
 import org.jetbrains.kotlin.gradle.plugin.KotlinPluginWrapper
 import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+import java.nio.file.Files
+import java.nio.file.Paths
+import java.security.MessageDigest
+import javax.xml.bind.DatatypeConverter
 
 plugins {
     java
@@ -68,5 +72,15 @@ tasks {
         archiveClassifier.value("")
         archiveVersion.value("")
         exclude("kotlin/reflect/**")
+        doLast { md5sum(archiveFile.get()) }
     }
 }
+
+fun md5sum(file: RegularFile) {
+    val path = file.asFile.toPath()
+    val md5File = Paths.get("$path.md5")
+    val md5sum = MessageDigest.getInstance("MD5")
+        .digest(Files.readAllBytes(path))
+        .let { DatatypeConverter.printHexBinary(it) }
+    Files.write(md5File, md5sum.toByteArray())
+}