支持菜品识别

2025-10-16 19:20:59 +08:00
parent d9b725cb63
commit c23aedf464
19 changed files with 579 additions and 42 deletions
@@ -0,0 +1,146 @@
+package com.sw.dualscreen.objbox
+
+import android.content.Context
+import android.graphics.Bitmap
+import com.google.gson.Gson
+import com.google.gson.reflect.TypeToken
+import com.sw.dualscreen.utils.AssetsTool
+import io.objectbox.Box
+import io.objectbox.kotlin.boxFor
+import io.objectbox.query.Query
+import org.pytorch.IValue
+import org.pytorch.Module
+import org.pytorch.torchvision.TensorImageUtils
+import java.io.File
+import java.io.FileOutputStream
+import java.io.IOException
+import java.io.InputStream
+
+object FoodModule {
+
+    private lateinit var module_mobile: Module
+    private lateinit var box: Box<Food>
+    private lateinit var embeddingsList: List<List<Float>>
+    private lateinit var labelsList: IntArray
+    private lateinit var classInfo: FoodClassInfo
+    private val NO_MEAN_RGB = floatArrayOf(0.0f, 0.0f, 0.0f)
+    private val NO_STD_RGB = floatArrayOf(1.0f, 1.0f, 1.0f)
+
+    fun init(context: Context) {
+        module_mobile = Module.load(copyAssetToCache(context, "best_embedding_model_mobile.pt"))
+        box = ObjectBox.boxStore.boxFor(Food::class)
+        if (box.all.isNotEmpty()) {
+            box.removeAll()
+        }
+        if (box.all.isEmpty()) {
+            initFoodData(context)
+        }
+    }
+
+    fun queryFood(bitmap: Bitmap, queryCount:Int = 30): MutableList<FoodQueryResult> {
+        val inputTensor = TensorImageUtils.bitmapToFloat32Tensor(
+            bitmap,
+            NO_MEAN_RGB, // [0.485, 0.456, 0.406] TORCHVISION_NORM_MEAN_RGB
+            NO_STD_RGB    // [0.229, 0.224, 0.225] TORCHVISION_NORM_STD_RGB
+        )
+        val outputTensor = module_mobile.forward(IValue.from(inputTensor)).toTensor()
+        val floatArray = outputTensor.dataAsFloatArray
+        return queryFood(floatArray,queryCount)
+    }
+
+    fun queryFood(floatArray: FloatArray, queryCount:Int = 30): MutableList<FoodQueryResult> {
+        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 resultList = mutableListOf<FoodQueryResult>()
+        query.findIdsWithScores().forEach {
+            val food = box.get(it.id)
+            resultList.add(
+                FoodQueryResult(
+                    id = it.id,
+                    name = food.name,
+                    foodIdx = food.foodIdx,
+                    score = it.score
+                )
+            )
+        }
+        return resultList
+    }
+
+
+    private fun initFoodData(context: Context) {
+        if (box.all.isNotEmpty()) {
+            return
+        }
+        val embeddingsJson = AssetsTool.readAssetsFile(context, "data/embeddings.json")
+        val labelsJson = AssetsTool.readAssetsFile(context, "data/labels.json")
+        val classInfoJson = AssetsTool.readAssetsFile(context, "data/class_info.json")
+
+        embeddingsList =
+            Gson().fromJson(embeddingsJson, object : TypeToken<List<List<Float>>>() {}.type)
+        labelsList = Gson().fromJson(labelsJson, IntArray::class.java)
+        classInfo =
+            Gson().fromJson(classInfoJson, FoodClassInfo::class.java)
+
+        val foodMap = classInfo.idx_to_class
+        embeddingsList.forEachIndexed { index, floatList ->
+            val classIdx = labelsList[index]
+            val foodName = foodMap["$classIdx"]
+            val array = floatList.toFloatArray()
+            box.put(Food(name = foodName, foodVector = array, foodIdx = index))
+        }
+    }
+
+
+    /**
+     * ，此方法的主要目的是：从assets 拷贝到 app的cache目录
+     * @param context
+     * @param fileName
+     * @return  例如是这样：/data/user/0/com.frizzle.pluginhookandroid9/cache/plugin-debug.apk
+     *
+     * 不可能反正SD
+     */
+    fun copyAssetToCache(context: Context, fileName: String): String? {
+        //  此app的缓存目录 --> 会默认在 cache目录...，可以自己去看看哦
+        val cacheDir = context.getCacheDir()
+        if (!cacheDir.exists()) {
+            cacheDir.mkdirs() // TODO 如果没有缓存目录，就创建
+        }
+        val outPath = File(cacheDir, fileName) // TODO 创建输出的文件位置
+        if (outPath.exists()) {
+            outPath.delete() // TODO 如果该文件已经存在，就删掉
+        }
+        var `is`: InputStream? = null // 读取
+        var fos: FileOutputStream? = null // 写入
+        try {
+            // 创建文件，如果创建成功，就返回true
+            val res = outPath.createNewFile()
+            if (res) {
+                `is` = context.getAssets().open(fileName) // 拿到main/assets目录的输入流，用于读取字节
+                fos = FileOutputStream(outPath) // 读取出来的字节最终写到outPath
+                val buf = ByteArray(`is`.available()) // 缓存区
+                var byteCount: Int
+
+                // 开始循环读取
+                while ((`is`.read(buf).also { byteCount = it }) != -1) {
+                    fos.write(buf, 0, byteCount)
+                }
+                return outPath.getAbsolutePath()
+            }
+        } catch (e: IOException) {
+            e.printStackTrace()
+        } finally {
+            try {
+                // TODO  一定要记得关闭资源，为了不去性能的磨损
+                fos!!.flush()
+                `is`!!.close()
+                fos.close()
+            } catch (e: IOException) {
+                e.printStackTrace()
+            }
+        }
+        return null
+    }
+
+
+}