Skip to content

Commit

Permalink
Fix issues with SSD model inference (#398)
Browse files Browse the repository at this point in the history 
* Fix class predictions decoding for SSDMobilenet

* Fix wrong channel ordering

* Add comment for SSD class predictions mapping
  • Loading branch information
@ermolenkodev
ermolenkodev authored Jun 1, 2022
1 parent 0955718 commit 952c5ef
Show file tree
Hide file tree
Showing 5 changed files with 28 additions and 18 deletions.
7 changes: 5 additions & 2 deletions api/src/main/kotlin/org/jetbrains/kotlinx/dl/dataset/handler/CocoUtils.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,13 @@ package org.jetbrains.kotlinx.dl.dataset.handler
/**
* 80 object categories in COCO dataset.
*
* Note that output class labels of SSD model do not correspond to ids in COCO annotations.
* If you want to evaluate this model on the COCO validation/test set, you need to convert class predictions using appropriate mapping.
*
* @see <a href="https://cocodataset.org/#home">
* COCO dataset</a>
*/
public val cocoCategories: Map<Int, String> = mapOf(
public val cocoCategoriesForSSD: Map<Int, String> = mapOf(
1 to "person",
2 to "bicycle",
3 to "car",
Expand Down Expand Up @@ -101,7 +104,7 @@ public val cocoCategories: Map<Int, String> = mapOf(
* @see <a href="https://cocodataset.org/#home">
* COCO dataset</a>
*/
public val cocoCategoriesForEfficientDet: Map<Int, String> = mapOf(
public val cocoCategories: Map<Int, String> = mapOf(
1 to "person",
2 to "bicycle",
3 to "car",
Expand Down
18 changes: 9 additions & 9 deletions onnx/src/main/kotlin/org/jetbrains/kotlinx/dl/api/inference/onnx/ONNXModels.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -617,7 +617,7 @@ public object ONNXModels {
ModelType<T, U> {
/**
* This model is a real-time neural network for object detection that detects 80 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForSSD]).
*
* The model have an input with the shape is (1x3x1200x1200).
*
Expand Down Expand Up @@ -660,7 +660,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 80 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForSSD]).
*
* SSD-MobilenetV1 is an object detection model that uses a Single Shot MultiBox Detector (SSD) approach
* to predict object classes for boundary boxes.
Expand Down Expand Up @@ -703,7 +703,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down Expand Up @@ -740,7 +740,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down Expand Up @@ -777,7 +777,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down Expand Up @@ -814,7 +814,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down Expand Up @@ -851,7 +851,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down Expand Up @@ -888,7 +888,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down Expand Up @@ -925,7 +925,7 @@ public object ONNXModels {

/**
* This model is a real-time neural network for object detection that detects 90 different classes
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet]).
* (labels are available in [org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories]).
*
* Internally it uses the EfficientNets as backbone networks.
*
Expand Down
4 changes: 2 additions & 2 deletions ...tbrains/kotlinx/dl/api/inference/onnx/objectdetection/EfficientDetObjectDetectionModel.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ package org.jetbrains.kotlinx.dl.api.inference.onnx.objectdetection
import org.jetbrains.kotlinx.dl.api.inference.objectdetection.DetectedObject
import org.jetbrains.kotlinx.dl.api.inference.onnx.ONNXModels
import org.jetbrains.kotlinx.dl.api.inference.onnx.OnnxInferenceModel
import org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForEfficientDet
import org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories
import org.jetbrains.kotlinx.dl.dataset.image.ColorMode
import org.jetbrains.kotlinx.dl.dataset.preprocessor.image.convert
import org.jetbrains.kotlinx.dl.dataset.preprocessor.image.resize
Expand Down Expand Up @@ -42,7 +42,7 @@ public class EfficientDetObjectDetectionModel : OnnxInferenceModel() {
val probability = items[i][5]
if (probability != 0.0f) {
val detectedObject = DetectedObject(
classLabel = cocoCategoriesForEfficientDet[items[i][6].toInt()]!!,
classLabel = cocoCategories[items[i][6].toInt()]!!,
probability = probability,
// left, bot, right, top
xMin = minOf(items[i][2] / inputShape[2], 1.0f),
Expand Down
2 changes: 1 addition & 1 deletion ...rains/kotlinx/dl/api/inference/onnx/objectdetection/SSDMobileNetV1ObjectDetectionModel.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ public class SSDMobileNetV1ObjectDetectionModel : OnnxInferenceModel() {
outputHeight = this@SSDMobileNetV1ObjectDetectionModel.inputShape[1].toInt()
outputWidth = this@SSDMobileNetV1ObjectDetectionModel.inputShape[2].toInt()
}
convert { colorMode = ColorMode.BGR }
convert { colorMode = ColorMode.RGB }
}
}

Expand Down
15 changes: 11 additions & 4 deletions ...in/org/jetbrains/kotlinx/dl/api/inference/onnx/objectdetection/SSDObjectDetectionModel.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ package org.jetbrains.kotlinx.dl.api.inference.onnx.objectdetection
import org.jetbrains.kotlinx.dl.api.inference.objectdetection.DetectedObject
import org.jetbrains.kotlinx.dl.api.inference.onnx.ONNXModels
import org.jetbrains.kotlinx.dl.api.inference.onnx.OnnxInferenceModel
import org.jetbrains.kotlinx.dl.dataset.handler.cocoCategories
import org.jetbrains.kotlinx.dl.dataset.handler.cocoCategoriesForSSD
import org.jetbrains.kotlinx.dl.dataset.image.ColorMode
import org.jetbrains.kotlinx.dl.dataset.preprocessor.image.convert
import org.jetbrains.kotlinx.dl.dataset.preprocessor.image.resize
Expand All @@ -27,6 +27,13 @@ private const val INPUT_SIZE = 1200
*
* It internally uses [ONNXModels.ObjectDetection.SSD] trained on the COCO dataset.
*
* Note that output class labels do not correspond to ids in COCO annotations.
* If you want to evaluate this model on the COCO validation/test set, you need to convert class predictions using appropriate mapping.
*
*
* @see <a href="https://github.com/lji72/inference/blob/tf_ssd_resent34_align_onnx/others/cloud/single_stage_detector/tensorflow/dataset_config/coco_labelmap.txt">
* Example mapping</a>
*
* @since 0.3
*/
public class SSDObjectDetectionModel : OnnxInferenceModel() {
Expand All @@ -39,7 +46,7 @@ public class SSDObjectDetectionModel : OnnxInferenceModel() {
* @param [topK] The number of the detected objects with the highest score to be returned.
* @return List of [DetectedObject] sorted by score.
*/
public fun detectObjects(inputData: FloatArray, topK: Int = 5): List<DetectedObject> {
public fun detectObjects(inputData: FloatArray, topK: Int = 5): List<DetectedObject> {
val rawPrediction = this.predictRaw(inputData)

val foundObjects = mutableListOf<DetectedObject>()
Expand All @@ -50,7 +57,7 @@ public class SSDObjectDetectionModel : OnnxInferenceModel() {

for (i in 0 until numberOfFoundObjects) {
val detectedObject = DetectedObject(
classLabel = cocoCategories[classIndices[i].toInt()]!!,
classLabel = cocoCategoriesForSSD[classIndices[i].toInt()]!!,
probability = probabilities[i],
// left, bot, right, top
xMin = boxes[i][0],
Expand Down Expand Up @@ -86,7 +93,7 @@ public class SSDObjectDetectionModel : OnnxInferenceModel() {
outputHeight = INPUT_SIZE
outputWidth = INPUT_SIZE
}
convert { colorMode = ColorMode.BGR }
convert { colorMode = ColorMode.RGB }
}
}

Expand Down

0 comments on commit 952c5ef

Please sign in to comment.