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

This commit is contained in:
Pierre-Yves Nicolas
2025-06-19 17:52:03 +02:00
parent 0c5a219783
commit 453a8b3fb0
2 changed files with 14 additions and 7 deletions

View File

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

View File

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