From 95bd3bd823b0eb690d9179837cb5713ea16e74f9 Mon Sep 17 00:00:00 2001
From: Pierre-Yves Nicolas <6371790+pynicolas@users.noreply.github.com>
Date: Tue, 27 May 2025 08:10:26 +0200
Subject: [PATCH] Fix computation of the number of objects that are detected

---
 .../main/java/org/mydomain/myscan/MainViewModel.kt    | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/app/src/main/java/org/mydomain/myscan/MainViewModel.kt b/app/src/main/java/org/mydomain/myscan/MainViewModel.kt
index 7a3a931..f283f8e 100644
--- a/app/src/main/java/org/mydomain/myscan/MainViewModel.kt
+++ b/app/src/main/java/org/mydomain/myscan/MainViewModel.kt
@@ -32,7 +32,7 @@ class MainViewModel(private val imageSegmentationService: ImageSegmentationServi
             imageSegmentationService.initialize()
             imageSegmentationService.segmentation
                 .filterNotNull()
-                .map { UiState("Found ${it.segmentation.masks.size} objects!", it.inferenceTime) }
+                .map { UiState("Found ${numberOfObjectsDetected(it.segmentation)} objects!", it.inferenceTime) }
                 .collect {
                     Log.d("MyScan", "New UIstate ${it}")
                     _uiState.value = it
@@ -40,6 +40,15 @@ class MainViewModel(private val imageSegmentationService: ImageSegmentationServi
         }
     }
 
+    fun numberOfObjectsDetected(segmentation: ImageSegmentationService.Segmentation) : Int {
+        val tensor = segmentation.masks[0];
+        val buffer = tensor.buffer
+        val uniqueValues = HashSet<Int>()
+        for (i in 0..tensor.width * tensor.height - 1) {
+            uniqueValues.add(buffer[i].toInt())
+        }
+        return uniqueValues.size - 1;
+    }
 
     fun segment(imageProxy: ImageProxy) {
         Log.d("MyScan", "MainViewModel.Calling segment")