前言
其实大概三四个月前就想写一篇文章来介绍移动端 AI 搜图的一些进展,不过由于本人的精力有限和一些其他的原因,没有及时更新。所以也就拖更很久,好在春节有些时间可以把之前的一些知识总结,更好的展现给大家。
相信用 Android 手机的同学多少都有一些感觉,Android 手机上的相册都多了一个搜图的功能,例如小米手机或是 Oppo 手机都上线了类似的功能,输入文字可以获得相关的图片。下面展示一下小米相册里面的搜图功能:
我们可以看到相册成功理解我们输入的文字查找出一些相关的图片出来,那么相册是如何识别文字然后将文字与图像匹配上的呢?且听我下节分说。
Clip 模型,AI 搜图的核心大脑
Clip 的模型的介绍
大概三年前 OpenAi 发布了 Clip 模型, 当然后续 Apple 基于 Clip 模型开发了更适合在移动端运行的 mobile clip 模型。
CLIP (Contrastive Language-Image Pre-Training) is a neural network trained on a variety of (image, text) pairs. It can be instructed in natural language to predict the most relevant text snippet, given an image, without directly optimizing for the task, similarly to the zero-shot capabilities of GPT-2 and 3. We found CLIP matches the performance of the original ResNet50 on ImageNet "zero-shot" without using any of the original 1.28M labeled examples, overcoming several major challenges in computer vision.
下面我对上面的介绍稍微做一下解读方便大家的理解,首先 Clip 模型是对一组海量的 text,image 的一一对照训练出来,举例说明(text:一张大海的图片,image: 大海.jpg)。
这个模型具有 zero-shot 的能力,如何理解这个 zero-shot呢?首先字面意思 zero-shot 即是没开一枪的情况,可以衍生含义为:对于某个未见过的图像、文本通过语义以识别他们。这与我们之前听过的 yolo 模型有很大的不一样,yolo 模型我们需要大量的数据标注,例如我们需要给模型投喂大量的狗狗图像,模型才具有识别狗狗的能力。如果没有投喂过猫猫的数据,那么 yolo 模型便不具有识别猫猫的能力。那么他是如何实现这个 zero-shot 功能的呢?
Clip 模型的原理
借用一下 Clip 官方的示意图,给大家大概讲解一下 Clip 训练过程。首先 text 被 TextEncoder 编码成一组向量,Imange 被 ImageEncoder 编码成一组向量。然后训练过程就是让 text 与 image 的向量尽可能的接近,让一组数据尽可能的落在相邻的向量空间。 当一个模型被正确训练之后,相同语义的 image 和 text,总是能够编码成相似的向量,然后我们通过通用的方式,例如欧式距离、余弦相似度、点乘来比较向量的相似度。如果了解过线性代数的同学对这些概念肯定都会比较熟悉,不过不了解也没关系,你只要理解他们都被编码成相似的向量。
如何理解两个向量很近呢,就是他们之间的距离很近,假设有你要比较两个人是否相似,肯定看他的五官(眼耳口鼻喉),然后提取出向量,比较两个向量的相似度是否相似来判断他们是否长的很像。
什么是向量?
为了简化大家的理解,在这里大家可以将向量理解为一个有序的数组,主要用来表示处理数据的特征。 举例说明,假设有一个向量来表示人的身材,我们的向量有两个维度,身高和体重,我们的模型主要目的就是提取人的身高和体重这两个特征,将它们转化成向量,方便留给后续的处理。这个提取物体特征的过程我们称之为 Encode,至于他的详细处理过程通过卷积,损失函数等等细节在这里我们不做过度了解。
如何开发一个 AI 搜图应用
下面我通过开源项目 PicQuery 来讲解如何开发一个 AI 搜图的 Android 应用。 PicQuery 的图片搜索功能基于 Clip 技术,可以切换使用 mobile clip,主要由以下关键组件构成:
- 图像编码器(Image Encoder)
- 文本编码器(Text Encoder)
- 相似度计算模块
- 嵌入向量存储库
预处理
在搜索的第一步应该是对相册的图片做预处理,使用 ImageEncoder 提取每个 image 的特征向量, 然后保存在数据库之中,这样在输入 text 搜索的时候才能避免每次都重复计算相册里面的 image 向量。
关于图片预处理有几点需要注意,避免一次性的加载所有相册的图片降低内存占用,这样可以减少 OOM 的产生,分批次处理 image 然后保存到数据库之中。
kotlin
// 分批处理照片
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(
current = processedPhotos.get(),
total = totalPhotos,
cost = cost,
status = IndexingAlbumState.Status.Indexing
)
}
if (!chunkSuccess) {
success = false
Log.w(TAG, "Failed to encode photo chunk, size: ${photoChunk.size}")
}
}
首先分批次从数据库读取图片,然后交给 ImageEncoder 进行处理。整个过程使用了 kotlin flow 来管理生产消费过程,极大的简化了复杂场景的流程处理,降低内存消耗。对于正常图片获取处理流程,如果下游 ImageEncoder 还没有消费完,上游会 suspend,避免一次性 produce 所有数据,干爆内存。 当然由于我个人手机相册图片只有 1500 多张,并没有验证过更多的图片的场景,如果有的话,欢迎小伙伴验证一下。关于 kotlin flow 实现生产消费者模型,如果想了解更多可以见本人的Kotlin flow 与生产消费者模型。
关于 ImageEncoderONNX
项目里面目前关于图片 Encode 是使用了 ONNX runtime框架来加载 clip 模型,当然也是支持 tensorflow lite 模,不过模型导出目前自测还有些小问题。对于业余爱好者想体验 AI 功能,自己去导出一些模型来开发,个人感觉 onnx 模型上周要比 tensorlite 要更容易,不过 tensorlite 在 Android 手机上的性能是要强于 onnx 的。
kotlin
open class ImageEncoderONNX(
private val dim: Long,
modelPath: String,
context: Context,
private val preprocessor: Preprocessor,
private val dispatcher: CoroutineDispatcher
) : ImageEncoder {
private var ortSession: OrtSession? = null
init {
ortSession = ortEnv.createSession(
AssetUtil.assetFilePath(context, modelPath),
options
)
}
```
override suspend fun encodeBatch(bitmaps: List<Bitmap>): List<FloatArray> {
...
}
}
关于 TextEncoderONNX:
文本编码器同样基于 ONNX,支持将文本转换为语义向量:
kotlin
abstract class TextEncoderONNX(private val context: Context) : TextEncoder {
private var tokenizer: BPETokenizer? = null
override fun encode(input: String): FloatArray {
val token = tokenizer!!.tokenize(input)
// 将文本转换为定长向量表示
}
}
AI 搜图流程介绍
搜索过程通过 ImageSearcher 类来调度:
- 对输入的图像或者文本编码,产生向量,
- 然后计算向量相似度
- 返回相似图片
kotlin
fun search(
text: String,
range: List<Album> = searchRange,
onSuccess: suspend (List<Long>?) -> Unit,
) {
// 文本编码
val textFeature = textEncoder.encode(text)
// 向量相似度搜索
searchWithVector(range, textFeature)
}
向量相似性计算
PicQuery 采用余弦相似度作为核心相似性度量:
kotlin
fun calculateSimilarity(
vector1: FloatArray,
vector2: FloatArray
): Float {
// 计算两个向量的余弦相似度
val dotProduct = vector1.zip(vector2).map { it.first * it.second }.sum()
val magnitude1 = sqrt(vector1.map { it * it }.sum())
val magnitude2 = sqrt(vector2.map { it * it }.sum())
return dotProduct / (magnitude1 * magnitude2)
}
关键技术亮点
-
多模态搜索 支持文本描述和图像相似性搜索,极大地提升了用户检索体验。
-
高效向量存储 使用 Room 数据库存储图像嵌入向量,支持快速检索:
kotlin
kotlin
@Dao
interface EmbeddingDao {
@Query("SELECT * FROM embeddings")
fun getAllEmbeddings(): List<Embedding>
}
``
3. 后台异步处理
利用 WorkManager 和协程进行图像索引:
kotlin
```kotlin
class ImageSimilarityCalculationWorker(
context: Context,
params: WorkerParameters
) : CoroutineWorker(context, params) {
override suspend fun doWork(): Result {
// 后台计算图像相似性
val embeddings = embeddingRepository.getAllEmbeddings()
// 计算并存储相似性
}
}
性能与优化策略
- 使用 kotlin flow 来优化 image 向量生产读取过程避免一次加载过多数据到内存
- 批量编码:支持多图像并行处理
技术挑战与解决方案
模型大小与推理性能
目前项目里面优先使用的是量化过后的 clip 模型,他占用内存较小,虽然精度不如原始的模型,但胜在体积小运算速度快,在移动端是较为理想的选择。理论上 mobile clip 的模型的执行效率是要好过 clip 模型的,不过由于 mobile clip 在量化模型的导出上存在一些问题,目前使用的不是 int8 的量化模型,虽说准确率稍微高一点,但是执行速度上会慢不少,如果相册内部有非常多照片的同学建议优先选择 clip 模型,对准确度要求更高的同学建议尝试 mobile clip 模型。
结语
PicQuery 通过创新的多模态搜索技术,为移动设备上的图片检索提供了一个高效、智能的解决方案。通过深度学习、向量检索和高性能计算,为用户带来前所未有的图片搜索体验。如果你也想在Android 手机上体验一下 AI 搜图,欢迎来下载体验,相信定会让你有所得。