Use a lock to ensure only one segmentation is run at a time

Author: pynicolas

app/src/androidTest/java/org/mydomain/myscan/DocumentDetectionTest.kt

@@ -21,6 +21,7 @@ import android.util.Log
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import kotlinx.coroutines.runBlocking
 
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -47,7 +48,9 @@ class DocumentDetectionTest {
             val bitmap = BitmapFactory.decodeStream(inputStream)
             var outputBitmap: Bitmap? = null
 
-            val segmentationResult = segmentationService.runSegmentationAndReturn(bitmap, 0)
+            val segmentationResult = runBlocking {
+                segmentationService.runSegmentationAndReturn(bitmap, 0)
+            }
             if (segmentationResult != null) {
                 val mask = segmentationResult.segmentation.toBinaryMask()
                 val quad = detectDocumentQuad(mask)

app/src/main/java/org/mydomain/myscan/ImageSegmentation.kt

@@ -25,6 +25,8 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.isActive
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 import kotlinx.coroutines.withContext
 import org.tensorflow.lite.DataType
 import org.tensorflow.lite.Interpreter
@@ -47,6 +49,7 @@ class ImageSegmentationService(private val context: Context) {
     val segmentation: StateFlow<SegmentationResult?> = _segmentation.asStateFlow()
 
     private var interpreter: Interpreter? = null
+    private val inferenceLock = Mutex()
 
     fun initialize() {
         interpreter = try {
@@ -83,18 +86,19 @@ class ImageSegmentationService(private val context: Context) {
         return SegmentationResult(segmentResult, inferenceTime)
     }
 
-    fun runSegmentationAndReturn(bitmap: Bitmap, rotationDegrees: Int): SegmentationResult? {
-        if (interpreter != null) {
-            return runSegmentation(interpreter!!, bitmap, rotationDegrees)
+    suspend fun runSegmentationAndReturn(bitmap: Bitmap, rotationDegrees: Int): SegmentationResult? {
+        if (interpreter == null) {
+            return null
+        }
+        return inferenceLock.withLock {
+            runSegmentation(interpreter!!, bitmap, rotationDegrees)
         }
-        return null
     }
 
     suspend fun runSegmentationAndEmit(bitmap: Bitmap, rotationDegrees: Int) {
-        if (interpreter == null) return
         try {
             withContext(Dispatchers.IO) {
-                val segmentationResult = runSegmentation(interpreter!!, bitmap, rotationDegrees)
+                val segmentationResult = runSegmentationAndReturn(bitmap, rotationDegrees)
                 if (isActive) {
                     _segmentation.value = segmentationResult
                 }