修改菜品识别逻辑

This commit is contained in:
mazengfei
2025-11-04 15:46:40 +08:00
parent 9c60049843
commit cd883af92f
4 changed files with 34 additions and 15 deletions
@@ -11,6 +11,7 @@ import io.objectbox.query.Query
import org.pytorch.IValue
import org.pytorch.Module
import org.pytorch.torchvision.TensorImageUtils
import timber.log.Timber
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
@@ -18,6 +19,7 @@ import java.io.InputStream
object FoodModule {
private const val THRESHOLD = 0.8
private lateinit var module_mobile: Module
private lateinit var box: Box<Food>
private lateinit var embeddingsList: List<List<Float>>
@@ -37,7 +39,7 @@ object FoodModule {
// }
}
fun queryFood(bitmap: Bitmap, queryCount:Int = 15): List<String> {
fun queryFood(bitmap: Bitmap, queryCount: Int = 15): List<String> {
val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
bitmap,
NO_MEAN_RGB, // [0.485, 0.456, 0.406] TORCHVISION_NORM_MEAN_RGB
@@ -45,20 +47,24 @@ object FoodModule {
)
val outputTensor = module_mobile.forward(IValue.from(inputTensor)).toTensor()
val floatArray = outputTensor.dataAsFloatArray
return queryFood(floatArray,queryCount)
return queryFood(floatArray, queryCount)
}
fun queryFood(floatArray: FloatArray, queryCount:Int = 15): List<String> {
val query: Query<Food> = box.query(Food_.foodVector.nearestNeighbors(floatArray, queryCount)).build()
fun queryFood(floatArray: FloatArray, queryCount: Int = 15): List<String> {
val query: Query<Food> =
box.query(Food_.foodVector.nearestNeighbors(floatArray, queryCount)).build()
//查询比较分数
// val tempList = query.findWithScores().sortedBy { it.score }.map { "${it.get().name}|${it.get().foodIdx}|${it.score}" }
val map = mutableMapOf<String, Int>()
query.findIdsWithScores().forEach {
val food = box.get(it.id)
//FoodQueryResult(id = it.id, name = food.name, foodIdx = food.foodIdx, score = it.score)
food.name?.let { key ->
val count = map[key] ?: 0
map.put(key, count + 1)
Timber.d("${food.name}|${food.foodIdx}|${it.score}")
if (1 - it.score >= THRESHOLD) {
//FoodQueryResult(id = it.id, name = food.name, foodIdx = food.foodIdx, score = it.score)
food.name?.let { key ->
val count = map[key] ?: 0
map.put(key, count + 1)
}
}
}
val list = map.entries.sortedByDescending { it.value }.map { it.key }