Split MainViewModel: extract CameraViewModel

app/src/main/java/org/fairscan/app/MainActivity.kt

@@ -33,6 +33,7 @@ import androidx.activity.compose.setContent
 import androidx.activity.enableEdgeToEdge
 import androidx.activity.result.contract.ActivityResultContracts
 import androidx.activity.viewModels
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.ui.platform.LocalContext
 import androidx.core.content.ContextCompat.checkSelfPermission
@@ -57,6 +58,8 @@ import org.fairscan.app.ui.screens.DocumentScreen
 import org.fairscan.app.ui.screens.ExportScreenWrapper
 import org.fairscan.app.ui.screens.HomeScreen
 import org.fairscan.app.ui.screens.LibrariesScreen
+import org.fairscan.app.ui.screens.camera.CameraEvent
+import org.fairscan.app.ui.screens.camera.CameraViewModel
 import org.opencv.android.OpenCVLoader
 
 private const val PDF_MIME_TYPE = "application/pdf"
@@ -70,11 +73,19 @@ class MainActivity : ComponentActivity() {
         lifecycleScope.launch(Dispatchers.IO) {
             viewModel.cleanUpOldPdfs(1000 * 3600)
         }
+        val cameraViewModel: CameraViewModel by viewModels { CameraViewModel.getFactory(this) }
         enableEdgeToEdge()
         setContent {
+            LaunchedEffect(Unit) {
+                cameraViewModel.events.collect { event ->
+                    when (event) {
+                        is CameraEvent.ImageCaptured -> viewModel.handleImageCaptured(event.jpegBytes)
+                    }
+                }
+            }
             val context = LocalContext.current
             val currentScreen by viewModel.currentScreen.collectAsStateWithLifecycle()
-            val liveAnalysisState by viewModel.liveAnalysisState.collectAsStateWithLifecycle()
+            val liveAnalysisState by cameraViewModel.liveAnalysisState.collectAsStateWithLifecycle()
             val document by viewModel.documentUiModel.collectAsStateWithLifecycle()
             val cameraPermission = rememberCameraPermissionState()
             val savePdf = { savePdf(viewModel.getFinalPdf(), viewModel) }
@@ -113,9 +124,10 @@ class MainActivity : ComponentActivity() {
                     is Screen.Main.Camera -> {
                         CameraScreen(
                             viewModel,
+                            cameraViewModel,
                             navigation,
                             liveAnalysisState,
-                            onImageAnalyzed = { image -> viewModel.liveAnalysis(image) },
+                            onImageAnalyzed = { image -> cameraViewModel.liveAnalysis(image) },
                             onFinalizePressed = navigation.toDocumentScreen,
                             cameraPermission = cameraPermission
                         )

app/src/main/java/org/fairscan/app/MainViewModel.kt

@@ -19,7 +19,6 @@ import android.graphics.Bitmap
 import android.graphics.BitmapFactory
 import android.os.Environment
 import android.util.Log
-import androidx.camera.core.ImageProxy
 import androidx.core.net.toUri
 import androidx.datastore.core.DataStore
 import androidx.lifecycle.ViewModel
@@ -33,7 +32,6 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
-import kotlinx.coroutines.flow.filterNotNull
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.flow.update
@@ -43,25 +41,18 @@ import org.fairscan.app.data.GeneratedPdf
 import org.fairscan.app.data.ImageRepository
 import org.fairscan.app.data.PdfFileManager
 import org.fairscan.app.data.recentDocumentsDataStore
-import org.fairscan.app.domain.ImageSegmentationService
-import org.fairscan.app.domain.detectDocumentQuad
-import org.fairscan.app.domain.extractDocument
-import org.fairscan.app.domain.scaledTo
 import org.fairscan.app.platform.AndroidPdfWriter
 import org.fairscan.app.platform.OpenCvTransformations
 import org.fairscan.app.ui.NavigationState
-import org.fairscan.app.ui.state.PdfGenerationUiState
-import org.fairscan.app.ui.state.RecentDocumentUiState
 import org.fairscan.app.ui.Screen
 import org.fairscan.app.ui.state.DocumentUiModel
-import org.fairscan.app.ui.state.LiveAnalysisState
+import org.fairscan.app.ui.state.PdfGenerationUiState
-import java.io.ByteArrayOutputStream
+import org.fairscan.app.ui.state.RecentDocumentUiState
 import java.io.File
 
 const val THUMBNAIL_SIZE_DP = 120
 
 class MainViewModel(
-    private val imageSegmentationService: ImageSegmentationService,
     private val imageRepository: ImageRepository,
     private val pdfFileManager: PdfFileManager,
     private val recentDocumentsDataStore: DataStore<RecentDocuments>,
@@ -74,7 +65,6 @@ class MainViewModel(
                 val density = context.resources.displayMetrics.density
                 val thumbnailSizePx = (THUMBNAIL_SIZE_DP * density).toInt()
                 return MainViewModel(
-                    ImageSegmentationService(context),
                     ImageRepository(context.filesDir, OpenCvTransformations(), thumbnailSizePx),
                     PdfFileManager(
                         File(context.cacheDir, "pdfs"),
@@ -87,10 +77,6 @@ class MainViewModel(
         }
     }
 
-    private var _liveAnalysisState = MutableStateFlow(LiveAnalysisState())
-    val liveAnalysisState: StateFlow<LiveAnalysisState> = _liveAnalysisState.asStateFlow()
-    private var lastSuccessfulLiveAnalysisState: LiveAnalysisState? = null
-
     private val _navigationState = MutableStateFlow(NavigationState.initial())
     val currentScreen: StateFlow<Screen> = _navigationState.map { it.current }
         .stateIn(viewModelScope, SharingStarted.Eagerly, _navigationState.value.current)
@@ -109,76 +95,6 @@ class MainViewModel(
             initialValue = DocumentUiModel(persistentListOf(), ::getBitmap, ::getThumbnail)
         )
 
-    private val _captureState = MutableStateFlow<CaptureState>(CaptureState.Idle)
-    val captureState: StateFlow<CaptureState> = _captureState
-
-    init {
-        viewModelScope.launch {
-            imageSegmentationService.initialize()
-            imageSegmentationService.segmentation
-                .filterNotNull()
-                .map {
-                    // TODO Should we really call toBinaryMask if it's used only in debug mode?
-                    val binaryMask = it.segmentation.toBinaryMask()
-                    LiveAnalysisState(
-                        inferenceTime = it.inferenceTime,
-                        binaryMask = binaryMask,
-                        documentQuad = detectDocumentQuad(it.segmentation, isLiveAnalysis = true),
-                        timestamp = System.currentTimeMillis(),
-                    )
-                }
-                .collect {
-                    _liveAnalysisState.value = it
-                    if (it.documentQuad != null) {
-                        lastSuccessfulLiveAnalysisState = it
-                    }
-                }
-        }
-    }
-
-    sealed class CaptureState {
-        open val frozenImage: Bitmap? = null
-
-        object Idle : CaptureState()
-        data class Capturing(override val frozenImage: Bitmap) : CaptureState()
-        data class CaptureError(override val frozenImage: Bitmap) : CaptureState()
-        data class CapturePreview(
-            override val frozenImage: Bitmap,
-            val processed: Bitmap
-        ) : CaptureState()
-    }
-
-
-    fun onCapturePressed(frozenImage: Bitmap) {
-        _captureState.value = CaptureState.Capturing(frozenImage)
-    }
-
-    private fun onCaptureProcessed(captured: Bitmap?) {
-        val current = _captureState.value
-        _captureState.value = when {
-            current is CaptureState.Capturing && captured != null ->
-                CaptureState.CapturePreview(current.frozenImage, captured)
-            current is CaptureState.Capturing ->
-                CaptureState.CaptureError(current.frozenImage)
-            else -> CaptureState.Idle
-        }
-    }
-
-    fun liveAnalysis(imageProxy: ImageProxy) {
-        if (_captureState.value !is CaptureState.Idle) {
-            imageProxy.close()
-            return
-        }
-
-        viewModelScope.launch {
-            imageSegmentationService.runSegmentationAndEmit(
-                imageProxy.toBitmap(),
-                imageProxy.imageInfo.rotationDegrees,
-            )
-            imageProxy.close()
-        }
-    }
-
     fun navigateTo(destination: Screen) {
         _navigationState.update { it.navigateTo(destination) }
     }
@@ -187,60 +103,6 @@ class MainViewModel(
         _navigationState.update { stack -> stack.navigateBack() }
     }
 
-    fun onImageCaptured(imageProxy: ImageProxy?) {
-        if (imageProxy != null) {
-            viewModelScope.launch {
-                val image = processCapturedImage(imageProxy)
-                imageProxy.close()
-                onCaptureProcessed(image)
-            }
-        } else {
-            onCaptureProcessed(null)
-        }
-    }
-
-    private suspend fun processCapturedImage(imageProxy: ImageProxy): Bitmap? = withContext(Dispatchers.IO) {
-        var corrected: Bitmap? = null
-        val bitmap = imageProxy.toBitmap()
-        val segmentation = imageSegmentationService.runSegmentationAndReturn(bitmap, 0)
-        if (segmentation != null) {
-            val mask = segmentation.segmentation
-            var quad = detectDocumentQuad(mask, isLiveAnalysis = false)
-            if (quad == null) {
-                val now = System.currentTimeMillis()
-                lastSuccessfulLiveAnalysisState?.timestamp?.let {
-                    val offset = now - it
-                    Log.i("Quad", "Last successful live analysis was $offset ms ago")
-                }
-                val recentLive = lastSuccessfulLiveAnalysisState?.takeIf {
-                    now - it.timestamp <= 1500
-                }
-                val rotations = (-imageProxy.imageInfo.rotationDegrees / 90) + 4
-                quad = recentLive?.documentQuad?.rotate90(rotations, mask.width, mask.height)
-                if (quad != null) {
-                    Log.i("Quad", "Using quad taken in live analysis; rotations=$rotations")
-                }
-            }
-            if (quad != null) {
-                val resizedQuad = quad.scaledTo(mask.width, mask.height, bitmap.width, bitmap.height)
-                corrected = extractDocument(bitmap, resizedQuad, imageProxy.imageInfo.rotationDegrees)
-            }
-        }
-        return@withContext corrected
-    }
-
-    fun addProcessedImage(quality: Int = 75) {
-        val current = _captureState.value
-        if (current is CaptureState.CapturePreview) {
-            val outputStream = ByteArrayOutputStream()
-            current.processed.compress(Bitmap.CompressFormat.JPEG, quality, outputStream)
-            val jpegBytes = outputStream.toByteArray()
-            imageRepository.add(jpegBytes)
-            _pageIds.value = imageRepository.imageIds()
-        }
-        _captureState.value = CaptureState.Idle
-    }
-
     fun rotateImage(id: String, clockwise: Boolean) {
         viewModelScope.launch {
             imageRepository.rotate(id, clockwise)
@@ -253,10 +115,6 @@ class MainViewModel(
         _pageIds.value = imageRepository.imageIds()
     }
 
-    fun afterCaptureError() {
-        _captureState.value = CaptureState.Idle
-    }
-
     fun deletePage(id: String) {
         imageRepository.delete(id)
         _pageIds.value = imageRepository.imageIds()
@@ -392,6 +250,11 @@ class MainViewModel(
             }
         }
     }
+
+    fun handleImageCaptured(jpegBytes: ByteArray) {
+        imageRepository.add(jpegBytes)
+        _pageIds.value = imageRepository.imageIds()
+    }
 }
 
 // TODO Move somewhere else: ViewModel should not depend on that

app/src/main/java/org/fairscan/app/ui/screens/camera/CameraScreen.kt

@@ -82,19 +82,18 @@ import androidx.compose.ui.unit.sp
 import androidx.lifecycle.compose.collectAsStateWithLifecycle
 import kotlinx.collections.immutable.toImmutableList
 import kotlinx.coroutines.delay
-import org.fairscan.app.ui.components.CameraPermissionState
-import org.fairscan.app.ui.state.LiveAnalysisState
 import org.fairscan.app.MainViewModel
-import org.fairscan.app.MainViewModel.CaptureState
-import org.fairscan.app.ui.Navigation
 import org.fairscan.app.R
+import org.fairscan.app.ui.Navigation
 import org.fairscan.app.ui.Screen
+import org.fairscan.app.ui.components.CameraPermissionState
 import org.fairscan.app.ui.components.CommonPageListState
 import org.fairscan.app.ui.components.MainActionButton
 import org.fairscan.app.ui.components.MyScaffold
 import org.fairscan.app.ui.components.pageCountText
 import org.fairscan.app.ui.dummyNavigation
 import org.fairscan.app.ui.fakeDocument
+import org.fairscan.app.ui.state.LiveAnalysisState
 import org.fairscan.app.ui.theme.FairScanTheme
 
 data class CameraUiState(
@@ -113,6 +112,7 @@ const val ANIMATION_DURATION = 200
 @Composable
 fun CameraScreen(
     viewModel: MainViewModel,
+    cameraViewModel: CameraViewModel,
     navigation: Navigation,
     liveAnalysisState: LiveAnalysisState,
     onImageAnalyzed: (ImageProxy) -> Unit,
@@ -132,11 +132,11 @@ fun CameraScreen(
         onDispose { captureController.shutdown() }
     }
 
-    val captureState by viewModel.captureState.collectAsStateWithLifecycle()
+    val captureState by cameraViewModel.captureState.collectAsStateWithLifecycle()
     if (captureState is CaptureState.CapturePreview) {
         LaunchedEffect(captureState) {
             delay(CAPTURED_IMAGE_DISPLAY_DURATION)
-            viewModel.addProcessedImage()
+            cameraViewModel.addProcessedImage()
         }
     }
 
@@ -146,7 +146,7 @@ fun CameraScreen(
             showDetectionError = true
             delay(1000)
             showDetectionError = false
-            viewModel.afterCaptureError()
+            cameraViewModel.afterCaptureError()
         }
     }
 
@@ -185,9 +185,9 @@ fun CameraScreen(
         onCapture = {
             previewView?.bitmap?.let {
                 Log.i("FairScan", "Pressed <Capture>")
-                viewModel.onCapturePressed(it)
+                cameraViewModel.onCapturePressed(it)
                 captureController.takePicture(
-                    onImageCaptured = { imageProxy -> viewModel.onImageCaptured(imageProxy) }
+                    onImageCaptured = { imageProxy -> cameraViewModel.onImageCaptured(imageProxy) }
                 )
             }
         },

app/src/main/java/org/fairscan/app/ui/screens/camera/CameraViewModel.kt

@@ -0,0 +1,194 @@
+/*
+ * Copyright 2025 Pierre-Yves Nicolas
+ *
+ * This program is free software: you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License as published by the Free
+ * Software Foundation, either version 3 of the License, or (at your option)
+ * any later version.
+ * This program is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
+ * more details.
+ * You should have received a copy of the GNU General Public License along with
+ * this program. If not, see <https://www.gnu.org/licenses/>.
+ */
+package org.fairscan.app.ui.screens.camera
+
+import android.content.Context
+import android.graphics.Bitmap
+import android.util.Log
+import androidx.camera.core.ImageProxy
+import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
+import androidx.lifecycle.viewmodel.CreationExtras
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asSharedFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.filterNotNull
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
+import org.fairscan.app.domain.ImageSegmentationService
+import org.fairscan.app.domain.detectDocumentQuad
+import org.fairscan.app.domain.extractDocument
+import org.fairscan.app.domain.scaledTo
+import org.fairscan.app.ui.state.LiveAnalysisState
+import java.io.ByteArrayOutputStream
+
+sealed interface CameraEvent {
+    data class ImageCaptured(val jpegBytes: ByteArray) : CameraEvent
+}
+
+class CameraViewModel(
+    private val imageSegmentationService: ImageSegmentationService
+): ViewModel() {
+
+    companion object {
+        fun getFactory(context: Context) = object : ViewModelProvider.Factory {
+            @Suppress("UNCHECKED_CAST")
+            override fun <T : ViewModel> create(modelClass: Class<T>, extras: CreationExtras): T {
+                return CameraViewModel(ImageSegmentationService(context)) as T
+            }
+        }
+    }
+
+    private val _events = MutableSharedFlow<CameraEvent>()
+    val events = _events.asSharedFlow()
+
+    private var _liveAnalysisState = MutableStateFlow(LiveAnalysisState())
+    val liveAnalysisState: StateFlow<LiveAnalysisState> = _liveAnalysisState.asStateFlow()
+    private var lastSuccessfulLiveAnalysisState: LiveAnalysisState? = null
+
+    private val _captureState = MutableStateFlow<CaptureState>(CaptureState.Idle)
+    val captureState: StateFlow<CaptureState> = _captureState
+
+    init {
+        viewModelScope.launch {
+            imageSegmentationService.initialize()
+            imageSegmentationService.segmentation
+                .filterNotNull()
+                .map {
+                    // TODO Should we really call toBinaryMask if it's used only in debug mode?
+                    val binaryMask = it.segmentation.toBinaryMask()
+                    LiveAnalysisState(
+                        inferenceTime = it.inferenceTime,
+                        binaryMask = binaryMask,
+                        documentQuad = detectDocumentQuad(it.segmentation, isLiveAnalysis = true),
+                        timestamp = System.currentTimeMillis(),
+                    )
+                }
+                .collect {
+                    _liveAnalysisState.value = it
+                    if (it.documentQuad != null) {
+                        lastSuccessfulLiveAnalysisState = it
+                    }
+                }
+        }
+    }
+
+    fun onCapturePressed(frozenImage: Bitmap) {
+        _captureState.value = CaptureState.Capturing(frozenImage)
+    }
+
+    private fun onCaptureProcessed(captured: Bitmap?) {
+        val current = _captureState.value
+        _captureState.value = when {
+            current is CaptureState.Capturing && captured != null ->
+                CaptureState.CapturePreview(current.frozenImage, captured)
+            current is CaptureState.Capturing ->
+                CaptureState.CaptureError(current.frozenImage)
+            else -> CaptureState.Idle
+        }
+    }
+
+    fun liveAnalysis(imageProxy: ImageProxy) {
+        if (_captureState.value !is CaptureState.Idle) {
+            imageProxy.close()
+            return
+        }
+
+        viewModelScope.launch {
+            imageSegmentationService.runSegmentationAndEmit(
+                imageProxy.toBitmap(),
+                imageProxy.imageInfo.rotationDegrees,
+            )
+            imageProxy.close()
+        }
+    }
+
+    fun onImageCaptured(imageProxy: ImageProxy?) {
+        if (imageProxy != null) {
+            viewModelScope.launch {
+                val image = processCapturedImage(imageProxy)
+                imageProxy.close()
+                onCaptureProcessed(image)
+            }
+        } else {
+            onCaptureProcessed(null)
+        }
+    }
+
+    private suspend fun processCapturedImage(imageProxy: ImageProxy): Bitmap? = withContext(Dispatchers.IO) {
+        var corrected: Bitmap? = null
+        val bitmap = imageProxy.toBitmap()
+        val segmentation = imageSegmentationService.runSegmentationAndReturn(bitmap, 0)
+        if (segmentation != null) {
+            val mask = segmentation.segmentation
+            var quad = detectDocumentQuad(mask, isLiveAnalysis = false)
+            if (quad == null) {
+                val now = System.currentTimeMillis()
+                lastSuccessfulLiveAnalysisState?.timestamp?.let {
+                    val offset = now - it
+                    Log.i("Quad", "Last successful live analysis was $offset ms ago")
+                }
+                val recentLive = lastSuccessfulLiveAnalysisState?.takeIf {
+                    now - it.timestamp <= 1500
+                }
+                val rotations = (-imageProxy.imageInfo.rotationDegrees / 90) + 4
+                quad = recentLive?.documentQuad?.rotate90(rotations, mask.width, mask.height)
+                if (quad != null) {
+                    Log.i("Quad", "Using quad taken in live analysis; rotations=$rotations")
+                }
+            }
+            if (quad != null) {
+                val resizedQuad = quad.scaledTo(mask.width, mask.height, bitmap.width, bitmap.height)
+                corrected = extractDocument(bitmap, resizedQuad, imageProxy.imageInfo.rotationDegrees)
+            }
+        }
+        return@withContext corrected
+    }
+
+    fun addProcessedImage(quality: Int = 75) {
+        val current = _captureState.value
+        if (current is CaptureState.CapturePreview) {
+            val outputStream = ByteArrayOutputStream()
+            current.processed.compress(Bitmap.CompressFormat.JPEG, quality, outputStream)
+            val jpegBytes = outputStream.toByteArray()
+            viewModelScope.launch {
+                _events.emit(CameraEvent.ImageCaptured(jpegBytes))
+            }
+        }
+        _captureState.value = CaptureState.Idle
+    }
+
+    fun afterCaptureError() {
+        _captureState.value = CaptureState.Idle
+    }
+
+}
+
+sealed class CaptureState {
+    open val frozenImage: Bitmap? = null
+
+    object Idle : CaptureState()
+    data class Capturing(override val frozenImage: Bitmap) : CaptureState()
+    data class CaptureError(override val frozenImage: Bitmap) : CaptureState()
+    data class CapturePreview(
+        override val frozenImage: Bitmap,
+        val processed: Bitmap
+    ) : CaptureState()
+}