fix: fix index count was not correct, add index manager

fix: fix index count was not correct, add index manager for update albums
greyovo · Feb 4, 2025 · 4d2336c · 4d2336c
1 parent 05a2451
commit 4d2336c
Show file tree

Hide file tree

Showing 15 changed files with 435 additions and 82 deletions.
diff --git a/app/src/main/java/me/grey/picquery/common/AppModules.kt b/app/src/main/java/me/grey/picquery/common/AppModules.kt
@@ -64,6 +64,7 @@ private val domainModules = module {
         AlbumManager(
             albumRepository = get(),
             photoRepository = get(),
+            embeddingRepository = get(),
             imageSearcher = get(),
             ioDispatcher = get()
         )

diff --git a/app/src/main/java/me/grey/picquery/common/Routes.kt b/app/src/main/java/me/grey/picquery/common/Routes.kt
@@ -4,5 +4,6 @@ enum class Routes {
     Home,
     Search,
     Display,
+    IndexMgr,
     Setting
 }
diff --git a/app/src/main/java/me/grey/picquery/data/dao/EmbeddingDao.kt b/app/src/main/java/me/grey/picquery/data/dao/EmbeddingDao.kt
@@ -36,8 +36,11 @@ interface EmbeddingDao {
     )
     fun getByAlbumIdList(albumIds: List<Long>, limit: Int, offset: Int): List<Embedding>
 
-//    @Insert
-//    fun insertAll(embeddings: List<Embedding>)
+    @Query(
+        "DELETE FROM $tableName WHERE album_id=(:albumId)"
+    )
+    fun removeByAlbumId(albumId: Long): Unit
+
 
     @Upsert
     fun upsertAll(embeddings: List<Embedding>)

diff --git a/app/src/main/java/me/grey/picquery/data/data_source/AlbumRepository.kt b/app/src/main/java/me/grey/picquery/data/data_source/AlbumRepository.kt
@@ -92,4 +92,8 @@ class AlbumRepository(
     fun addAllSearchableAlbum(album: List<Album>) {
         return database.albumDao().upsertAll(album)
     }
+
+    fun removeSearchableAlbum(singleAlbum: Album) {
+        database.albumDao().delete(singleAlbum)
+    }
 }
diff --git a/app/src/main/java/me/grey/picquery/data/data_source/EmbeddingRepository.kt b/app/src/main/java/me/grey/picquery/data/data_source/EmbeddingRepository.kt
@@ -81,4 +81,8 @@ class EmbeddingRepository(
     fun updateAll(list: List<Embedding>) {
         return dataSource.upsertAll(list)
     }
+
+    fun removeByAlbum(album: Album){
+        return dataSource.removeByAlbumId(album.id)
+    }
 }
diff --git a/app/src/main/java/me/grey/picquery/domain/AlbumManager.kt b/app/src/main/java/me/grey/picquery/domain/AlbumManager.kt
@@ -8,22 +8,24 @@ import androidx.compose.runtime.mutableStateOf
 import com.google.accompanist.permissions.ExperimentalPermissionsApi
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.asFlow
 import kotlinx.coroutines.flow.flatMapConcat
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
-import me.grey.picquery.PicQueryApplication
 import me.grey.picquery.PicQueryApplication.Companion.context
 import me.grey.picquery.R
 import me.grey.picquery.common.showToast
 import me.grey.picquery.data.data_source.AlbumRepository
+import me.grey.picquery.data.data_source.EmbeddingRepository
 import me.grey.picquery.data.data_source.PhotoRepository
 import me.grey.picquery.data.model.Album
 import me.grey.picquery.ui.albums.IndexingAlbumState
+import java.util.concurrent.atomic.AtomicInteger
 
 class AlbumManager(
     private val albumRepository: AlbumRepository,
     private val photoRepository: PhotoRepository,
+    private val embeddingRepository: EmbeddingRepository,
     private val imageSearcher: ImageSearcher,
     private val ioDispatcher: CoroutineDispatcher
 ) {
@@ -37,14 +39,16 @@ class AlbumManager(
         get() = indexingAlbumState.value.isBusy
 
     private val albumList = mutableStateListOf<Album>()
-    val searchableAlbumList = mutableStateListOf<Album>()
-    val unsearchableAlbumList = mutableStateListOf<Album>()
+    val searchableAlbumList = MutableStateFlow<List<Album>>(emptyList())
+    val unsearchableAlbumList = MutableStateFlow<List<Album>>(emptyList())
     val albumsToEncode = mutableStateListOf<Album>()
 
     private fun searchableAlbumFlow() = albumRepository.getSearchableAlbumFlow()
 
     private var initialized = false
 
+    fun getAlbumList() = albumList
+
     suspend fun initAllAlbumList() {
         if (initialized) return
         withContext(ioDispatcher) {
@@ -57,18 +61,18 @@ class AlbumManager(
         }
     }
 
-    private suspend fun initDataFlow() {
+    suspend fun initDataFlow() {
         searchableAlbumFlow().collect {
             // 从数据库中检索已经索引的相册
             // 有些相册可能已经索引但已被删除，因此要从全部相册中筛选，而不能直接返回数据库的结果
-            searchableAlbumList.clear()
-            searchableAlbumList.addAll(it)
-            searchableAlbumList.sortByDescending { album: Album -> album.count }
+            val res = it.toMutableList().sortedByDescending { album: Album -> album.count }
+            searchableAlbumList.emit(res)
+            Log.d(TAG, "Searchable albums: ${it.size}")
             // 从全部相册减去已经索引的ID，就是未索引的相册
             val unsearchable = albumList.filter { all -> !it.contains(all) }
-            unsearchableAlbumList.clear()
-            unsearchableAlbumList.addAll(unsearchable)
-            unsearchableAlbumList.sortByDescending { album: Album -> album.count }
+
+            unsearchableAlbumList.emit(unsearchable.toMutableList().sortedByDescending { album: Album -> album.count })
+            Log.d(TAG, "Unsearchable albums: ${unsearchable.size}")
         }
     }
 
@@ -81,9 +85,9 @@ class AlbumManager(
     }
 
     fun toggleSelectAllAlbums() {
-        if (albumsToEncode.size != unsearchableAlbumList.size) {
+        if (albumsToEncode.size != unsearchableAlbumList.value?.size) {
             albumsToEncode.clear()
-            albumsToEncode.addAll(unsearchableAlbumList)
+            albumsToEncode.addAll(unsearchableAlbumList.value)
         } else {
             albumsToEncode.clear()
         }
@@ -92,11 +96,11 @@ class AlbumManager(
     /**
      * 获取多个相册的照片流
      */
-@OptIn(ExperimentalCoroutinesApi::class)
-private fun getPhotosFlow(albums: List<Album>) = albums.asFlow()
-    .flatMapConcat { album ->
-        photoRepository.getPhotoListByAlbumIdPaginated(album.id)
-    }
+    @OptIn(ExperimentalCoroutinesApi::class)
+    private fun getPhotosFlow(albums: List<Album>) = albums.asFlow()
+        .flatMapConcat { album ->
+            photoRepository.getPhotoListByAlbumIdPaginated(album.id)
+        }
 
     /**
      * 获取相册列表的照片总数
@@ -105,62 +109,64 @@ private fun getPhotosFlow(albums: List<Album>) = albums.asFlow()
         albums.sumOf { album -> photoRepository.getImageCountInAlbum(album.id) }
     }
 
-    fun encodeAlbums(albums: List<Album>) {
-        PicQueryApplication.applicationScope.launch {
-            if (albums.isEmpty()) {
-                showToast(context.getString(R.string.no_album_selected))
-                return@launch
-            }
+    suspend fun encodeAlbums(albums: List<Album>) {
+        if (albums.isEmpty()) {
+            showToast(context.getString(R.string.no_album_selected))
+            return
+        }
 
-            indexingAlbumState.value = 
-                IndexingAlbumState(status = IndexingAlbumState.Status.Loading)
-
-            try {
-                // 先获取总照片数
-                val totalPhotos = getTotalPhotoCount(albums)
-                var processedPhotos = 0
-                var success = true
-
-                // 分批处理照片
-                getPhotosFlow(albums).collect { photoChunk ->
-                    // 对每一批照片进行编码
-                    val chunkSuccess = imageSearcher.encodePhotoListV2(photoChunk) { cur, total, cost ->
-                        processedPhotos += cur
-                        indexingAlbumState.value = indexingAlbumState.value.copy(
-                            current = processedPhotos,
-                            total = totalPhotos,
-                            cost = cost,
-                            status = IndexingAlbumState.Status.Indexing
-                        )
-                    }
-
-                    if (!chunkSuccess) {
-                        success = false
-                        Log.w(TAG, "Failed to encode photo chunk, size: ${photoChunk.size}")
-                    }
-                }
+        indexingAlbumState.value =
+            IndexingAlbumState(status = IndexingAlbumState.Status.Loading)
+
+        try {
+            val totalPhotos = getTotalPhotoCount(albums)
+            val processedPhotos = AtomicInteger(0)
+            var success = true
 
-                if (success) {
-                    Log.i(TAG, "Encoded ${albums.size} album(s) with $totalPhotos photos!")
-                    withContext(ioDispatcher) {
-                        albumRepository.addAllSearchableAlbum(albums)
-                    }
+            getPhotosFlow(albums).collect { photoChunk ->
+
+                val chunkSuccess = imageSearcher.encodePhotoListV2(photoChunk) { cur, total, cost ->
+                    Log.d(TAG, "Encoded $cur/$total photos, cost: $cost")
+                    processedPhotos.addAndGet(cur)
                     indexingAlbumState.value = indexingAlbumState.value.copy(
-                        status = IndexingAlbumState.Status.Finish
+                        current = processedPhotos.get(),
+                        total = totalPhotos,
+                        cost = cost,
+                        status = IndexingAlbumState.Status.Indexing
                     )
-                } else {
-                    Log.w(TAG, "encodePhotoList failed! Maybe too much request.")
                 }
-            } catch (e: Exception) {
-                Log.e(TAG, "Error encoding albums", e)
+
+                if (!chunkSuccess) {
+                    success = false
+                    Log.w(TAG, "Failed to encode photo chunk, size: ${photoChunk.size}")
+                }
+            }
+
+            if (success) {
+                Log.i(TAG, "Encoded ${albums.size} album(s) with $totalPhotos photos!")
+                withContext(ioDispatcher) {
+                    albumRepository.addAllSearchableAlbum(albums)
+                }
                 indexingAlbumState.value = indexingAlbumState.value.copy(
-                    status = IndexingAlbumState.Status.Error
+                    status = IndexingAlbumState.Status.Finish
                 )
+            } else {
+                Log.w(TAG, "encodePhotoList failed! Maybe too much request.")
             }
+        } catch (e: Exception) {
+            Log.e(TAG, "Error encoding albums", e)
+            indexingAlbumState.value = indexingAlbumState.value.copy(
+                status = IndexingAlbumState.Status.Error
+            )
         }
     }
 
     fun clearIndexingState() {
         indexingAlbumState.value = IndexingAlbumState()
     }
+
+    fun removeSingleAlbumIndex(album: Album) {
+        embeddingRepository.removeByAlbum(album)
+        albumRepository.removeSearchableAlbum(album)
+    }
 }
diff --git a/app/src/main/java/me/grey/picquery/domain/ImageSearcher.kt b/app/src/main/java/me/grey/picquery/domain/ImageSearcher.kt
@@ -10,6 +10,8 @@ import androidx.compose.runtime.mutableFloatStateOf
 import androidx.compose.runtime.mutableStateListOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.ui.graphics.vector.ImageVector
+import androidx.core.content.res.ResourcesCompat
+import androidx.core.graphics.drawable.toBitmap
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -66,6 +68,8 @@ class ImageSearcher(
         private const val DEFAULT_MATCH_THRESHOLD = 0.01
         private const val TOP_K = 30
         private const val SEARCH_BATCH_SIZE = 1000
+
+
     }
 
     val searchRange = mutableStateListOf<Album>()
@@ -136,16 +140,8 @@ class ImageSearcher(
                 .chunked(100)
                 .onEach { Log.d(TAG, "onEach: ${it.size}") }
                 .onCompletion {
-                    val cost = max((System.currentTimeMillis() - startTime)/1000f, 1f)
-                    Log.d(
-                        TAG,
-                        "[OK] encode ${photos.size} pics in $cost s, " +
-                                "avg speed: ${photos.size / cost} pic/s"
-                    )
-                    progressCallback?.invoke(photos.size, photos.size, 0)
                     embeddingRepository.updateCache()
                     encodingLock = false
-
                 }
                 .collect {
                     val loops = 1
@@ -165,13 +161,14 @@ class ImageSearcher(
                         }
                         deferreds.awaitAll()
                     }
+                    cur.set(it.size)
+
                     progressCallback?.invoke(
                         cur.get(),
                         photos.size,
                         cost / it.size,
                     )
                     Log.d(TAG, "cost: ${cost}")
-                    cur.addAndGet(it.size)
                 }
         }
         return true
@@ -221,7 +218,7 @@ class ImageSearcher(
         image: Bitmap,
         range: List<Album> = searchRange,
         onSuccess: suspend (List<Long>?) -> Unit,
-        ) {
+    ) {
         return withContext(dispatcher) {
             if (searchingLock) {
                 return@withContext
@@ -255,7 +252,7 @@ class ImageSearcher(
             embeddings.collect { chunk ->
                 Log.d(TAG, "Processing chunk: ${chunk.size}")
                 totalProcessed += chunk.size
-                
+
                 for (emb in chunk) {
                     val sim = calculateSimilarity(emb.data.toFloatArray(), textFeat)
                     Log.d(TAG, "similarity: ${emb.photoId} -> $sim")

diff --git a/app/src/main/java/me/grey/picquery/ui/AppNavHost.kt b/app/src/main/java/me/grey/picquery/ui/AppNavHost.kt
@@ -14,6 +14,7 @@ import me.grey.picquery.common.Animation.popInAnimation
 import me.grey.picquery.common.Routes
 import me.grey.picquery.ui.display.DisplayScreen
 import me.grey.picquery.ui.home.HomeScreen
+import me.grey.picquery.ui.indexmgr.IndexMgrScreen
 import me.grey.picquery.ui.search.SearchScreen
 import me.grey.picquery.ui.setting.SettingScreen
 
@@ -64,6 +65,11 @@ fun AppNavHost(
                 },
             )
         }
+        composable(Routes.IndexMgr.name) {
+            IndexMgrScreen(
+                onNavigateBack = { navController.popBackStack() },
+            )
+        }
         composable(
             Routes.Setting.name,
             enterTransition = { popInAnimation },
@@ -72,6 +78,9 @@ fun AppNavHost(
         ) {
             SettingScreen(
                 onNavigateBack = { navController.popBackStack() },
+                navigateToIndexMgr = {
+                    navController.navigate(Routes.IndexMgr.name)
+                }
             )
         }
     }

diff --git a/app/src/main/java/me/grey/picquery/ui/display/DisplayScreen.kt b/app/src/main/java/me/grey/picquery/ui/display/DisplayScreen.kt
@@ -66,6 +66,7 @@ fun DisplayScreen(
     )
 
     LaunchedEffect(initialPage) {
+        if (initialPage == -1) return@LaunchedEffect
         displayViewModel.loadPhotos()
         if (photoList.isNotEmpty()) {
             pagerState.scrollToPage(initialPage)
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,5 +4,6 @@ enum class Routes { @@
         Home,
         Search,
         Display,
+        IndexMgr,
         Setting
     }