Fix ssdLightAPITest (Kotlin#440)

* Turn Coco from class to enum
ermolenkodev · Sep 8, 2022 · e4907ea · e4907ea
1 parent b21b786
commit e4907ea
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 14 deletions.
diff --git a/dataset/src/commonMain/kotlin/org/jetbrains/kotlinx/dl/dataset/CocoUtils.kt b/dataset/src/commonMain/kotlin/org/jetbrains/kotlinx/dl/dataset/CocoUtils.kt
@@ -5,10 +5,16 @@
 
 package org.jetbrains.kotlinx.dl.dataset
 
-public class Coco(public val version: CocoVersion, zeroIndexed: Boolean = false) {
-    public val labels: Map<Int, String> = when (version) {
-        CocoVersion.V2014 -> if (zeroIndexed) toZeroIndexed(cocoCategories2014) else cocoCategories2014
-        CocoVersion.V2017 -> if (zeroIndexed) toZeroIndexed(cocoCategories2017) else cocoCategories2017
+
+public enum class Coco {
+    V2014,
+    V2017;
+
+    public fun labels(zeroIndexed: Boolean = false) : Map<Int, String> {
+        return when (this) {
+            V2014 -> if (zeroIndexed) toZeroIndexed(cocoCategories2014) else cocoCategories2014
+            V2017 -> if (zeroIndexed) toZeroIndexed(cocoCategories2017) else cocoCategories2017
+        }
     }
 
     private fun toZeroIndexed(labels: Map<Int, String>) : Map<Int, String> {
@@ -20,7 +26,6 @@ public class Coco(public val version: CocoVersion, zeroIndexed: Boolean = false)
     }
 }
 
-
 /**
  * 80 object categories in COCO dataset.
  *

diff --git a/...tbrains/kotlinx/dl/api/inference/onnx/objectdetection/SSDMobileNetObjectDetectionModel.kt b/...tbrains/kotlinx/dl/api/inference/onnx/objectdetection/SSDMobileNetObjectDetectionModel.kt
@@ -2,6 +2,7 @@ package org.jetbrains.kotlinx.dl.api.inference.onnx.objectdetection
 
 import android.graphics.Bitmap
 import org.jetbrains.kotlinx.dl.api.inference.InferenceModel
+import org.jetbrains.kotlinx.dl.api.inference.objectdetection.DetectedObject
 import org.jetbrains.kotlinx.dl.api.inference.onnx.OnnxInferenceModel
 import org.jetbrains.kotlinx.dl.api.inference.onnx.executionproviders.ExecutionProvider.CPU
 import org.jetbrains.kotlinx.dl.dataset.Coco
@@ -22,18 +23,18 @@ public class SSDMobileNetObjectDetectionModel(override val internalModel: OnnxIn
     SSDObjectDetectionModelBase<Bitmap>(SSD_MOBILENET_METADATA),
     InferenceModel by internalModel {
 
-    override val classLabels: Map<Int, String> = Coco(V2017, zeroIndexed = true).labels
+    override val classLabels: Map<Int, String> = Coco.V2017.labels(zeroIndexed = true)
 
     private var targetRotation = 0f
 
+    override lateinit var preprocessing: Operation<Bitmap, Pair<FloatArray, TensorShape>>
+        private set
+
     public constructor (modelBytes: ByteArray) : this(OnnxInferenceModel(modelBytes)) {
         internalModel.initializeWith(CPU())
         preprocessing = buildPreprocessingPipeline()
     }
 
-    override lateinit var preprocessing: Operation<Bitmap, Pair<FloatArray, TensorShape>>
-        private set
-
     public fun setTargetRotation(targetRotation: Float) {
         if (this.targetRotation == targetRotation) return
 

diff --git a/...n/org/jetbrains/kotlinx/dl/api/inference/onnx/objectdetection/ObjectDetectionModelBase.kt b/...n/org/jetbrains/kotlinx/dl/api/inference/onnx/objectdetection/ObjectDetectionModelBase.kt
@@ -68,8 +68,7 @@ public abstract class EfficientDetObjectDetectionModelBase<I> : ObjectDetectionM
 /**
  * Base class for object detection model based on SSD architecture.
  */
-public abstract class SSDObjectDetectionModelBase<I>(private val metadata: SSDModelMetadata) : ObjectDetectionModelBase<I>() {
-
+public abstract class SSDObjectDetectionModelBase<I>(protected val metadata: SSDModelMetadata) : ObjectDetectionModelBase<I>() {
     override fun convert(output: Map<String, Any>): List<DetectedObject> {
         val boxes = (output[metadata.outputBoxesName] as Array<Array<FloatArray>>)[0]
         val classIndices = (output[metadata.outputClassesName] as Array<FloatArray>)[0]

diff --git a/...tbrains/kotlinx/dl/api/inference/onnx/objectdetection/EfficientDetObjectDetectionModel.kt b/...tbrains/kotlinx/dl/api/inference/onnx/objectdetection/EfficientDetObjectDetectionModel.kt
@@ -46,7 +46,7 @@ public class EfficientDetObjectDetectionModel(override val internalModel: OnnxIn
             // model is quite sensitive for this
             .convert { colorMode = ColorMode.RGB }
             .toFloatArray { }
-    override val classLabels: Map<Int, String> = Coco(V2017).labels
+    override val classLabels: Map<Int, String> = Coco.V2017.labels(zeroIndexed = false)
 
     /**
      * Constructs the object detection model from a given path.

diff --git a/...rains/kotlinx/dl/api/inference/onnx/objectdetection/SSDMobileNetV1ObjectDetectionModel.kt b/...rains/kotlinx/dl/api/inference/onnx/objectdetection/SSDMobileNetV1ObjectDetectionModel.kt
@@ -55,7 +55,7 @@ public class SSDMobileNetV1ObjectDetectionModel(override val internalModel: Onnx
             .toFloatArray { }
             .call(ONNXModels.ObjectDetection.SSDMobileNetV1.preprocessor)
 
-    override val classLabels: Map<Int, String> = Coco(V2017).labels
+    override val classLabels: Map<Int, String> = Coco.V2017.labels()
 
     /**
      * Constructs the object detection model from a given path.

diff --git a/...in/org/jetbrains/kotlinx/dl/api/inference/onnx/objectdetection/SSDObjectDetectionModel.kt b/...in/org/jetbrains/kotlinx/dl/api/inference/onnx/objectdetection/SSDObjectDetectionModel.kt
@@ -58,7 +58,7 @@ public class SSDObjectDetectionModel(override val internalModel: OnnxInferenceMo
             .toFloatArray { }
             .call(ONNXModels.ObjectDetection.SSD.preprocessor)
 
-    override val classLabels: Map<Int, String> = Coco(CocoVersion.V2014).labels
+    override val classLabels: Map<Int, String> = Coco.V2014.labels()
 
     /**
      * Constructs the object detection model from a given path.
@@ -80,6 +80,30 @@ public class SSDObjectDetectionModel(override val internalModel: OnnxInferenceMo
         return detectObjects(ImageConverter.toBufferedImage(imageFile), topK)
     }
 
+    // TODO remove code duplication due to different type of class labels array
+    override fun convert(output: Map<String, Any>): List<DetectedObject> {
+        val boxes = (output[metadata.outputBoxesName] as Array<Array<FloatArray>>)[0]
+        val classIndices = (output[metadata.outputClassesName] as Array<LongArray>)[0]
+        val probabilities = (output[metadata.outputScoresName] as Array<FloatArray>)[0]
+        val numberOfFoundObjects = boxes.size
+
+        val foundObjects = mutableListOf<DetectedObject>()
+        for (i in 0 until numberOfFoundObjects) {
+            val detectedObject = DetectedObject(
+                classLabel = if (classIndices[i].toInt() in classLabels.keys) classLabels[classIndices[i].toInt()]!! else "Unknown",
+                probability = probabilities[i],
+                // left, bot, right, top
+                xMin = boxes[i][metadata.xMinIdx],
+                yMin = boxes[i][metadata.yMinIdx],
+                xMax = boxes[i][metadata.xMinIdx + 2],
+                yMax = boxes[i][metadata.yMinIdx + 2]
+            )
+            foundObjects.add(detectedObject)
+        }
+        return foundObjects
+    }
+
+
     override fun copy(
         copiedModelName: String?,
         saveOptimizerState: Boolean,