litert-runtime

AiPackManagerFactory.getInstance(context).getPackLocation(name)?.assetsPath()

import android.content.Context
import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.CompiledModel
import com.google.ai.edge.litert.Environment
import com.google.ai.edge.litert.TensorBuffer

const val INPUT_SIZE = 224 * 224 * 3
private val environment by lazy { Environment.create() }

class Classifier(context: Context, accelerator: Accelerator) : AutoCloseable {
    private val model = CompiledModel.create(
        context.assets,
        "model.tflite",
        CompiledModel.Options(accelerator),
        environment
    )
    private val inputs = mutableListOf<TensorBuffer>()
    private val outputs = mutableListOf<TensorBuffer>()

    init {
        try {
            inputs += model.createInputBuffers()
            outputs += model.createOutputBuffers()
            infer(FloatArray(INPUT_SIZE))
        } catch (e: Exception) {
            close()
            throw e
        }
    }

    fun infer(input: FloatArray): FloatArray {
        inputs[0].writeFloat(input)
        model.run(inputs, outputs)
        return outputs[0].readFloat()
    }

    override fun close() {
        (inputs + outputs).forEach { it.close() }
        model.close()
    }
}

import android.app.Application
import androidx.lifecycle.AndroidViewModel
import com.google.ai.edge.litert.Accelerator
import com.google.ai.edge.litert.LiteRtException
import java.util.concurrent.Executors
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.asCoroutineDispatcher
import kotlinx.coroutines.cancel
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch

data class UiState(
    val ready: Boolean = false,
    val accelerator: Accelerator = Accelerator.CPU,
    val result: List<Float>? = null,
    val error: String? = null
)

class MainViewModel(app: Application) : AndroidViewModel(app) {
    private val executor = Executors.newSingleThreadExecutor()
    private val scope = CoroutineScope(
        SupervisorJob() + executor.asCoroutineDispatcher()
    )
    private var classifier: Classifier? = null
    private val _uiState = MutableStateFlow(UiState())
    val uiState: StateFlow<UiState> = _uiState.asStateFlow()

    fun load(preferred: Accelerator = Accelerator.GPU) {
        scope.launch {
            if (classifier != null) return@launch
            val appContext = getApplication<Application>()
            val selected = try {
                classifier = Classifier(appContext, preferred)
                preferred
            } catch (e: LiteRtException) {
                if (preferred != Accelerator.CPU) {
                    try {
                        classifier = Classifier(appContext, Accelerator.CPU)
                        Accelerator.CPU
                    } catch (cpuError: LiteRtException) {
                        _uiState.value = UiState(error = cpuError.message)
                        return@launch
                    }
                } else {
                    _uiState.value = UiState(error = e.message)
                    return@launch
                }
            }
            _uiState.value = UiState(ready = true, accelerator = selected)
        }
    }

    fun classify(input: FloatArray) {
        scope.launch {
            try {
                classifier?.infer(input)?.let { result ->
                    _uiState.update {
                        it.copy(result = result.toList(), error = null)
                    }
                }
            } catch (e: LiteRtException) {
                _uiState.update { it.copy(error = e.message) }
            }
        }
    }

    override fun onCleared() {
        scope.launch { classifier?.close() }.invokeOnCompletion {
            scope.cancel()
        }
        executor.shutdown()
    }
}

import android.graphics.Bitmap
import androidx.compose.foundation.layout.Column
import androidx.compose.material3.Button
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.lifecycle.compose.collectAsStateWithLifecycle
import androidx.lifecycle.viewmodel.compose.viewModel
import com.google.ai.edge.litert.Accelerator

@Composable
fun MainScreen(
    bitmap: Bitmap,
    viewModel: MainViewModel = viewModel()
) {
    val state by viewModel.uiState.collectAsStateWithLifecycle()
    LaunchedEffect(Unit) { viewModel.load(Accelerator.GPU) }
    Column {
        state.error?.let { Text(it) }
        Button(
            onClick = { viewModel.classify(preprocess(bitmap)) },
            enabled = state.ready
        ) {
            Text("Run (${state.accelerator})")
        }
        state.result?.let { scores ->
            Text("Top class: ${scores.indices.maxBy { i -> scores[i] }}")
        }
    }
}

import android.graphics.Bitmap

fun preprocess(
    bitmap: Bitmap,
    size: Int = 224,
    mean: Float = 127.5f,
    std: Float = 127.5f
): FloatArray {
    val scaled = checkNotNull(
        Bitmap.createScaledBitmap(bitmap, size, size, true)
            .copy(Bitmap.Config.ARGB_8888, false)
    )
    val pixels = IntArray(size * size).also {
        scaled.getPixels(it, 0, size, 0, 0, size, size)
    }
    val out = FloatArray(size * size * 3)
    pixels.forEachIndexed { i, p ->
        out[i * 3] = ((p shr 16 and 0xFF) - mean) / std
        out[i * 3 + 1] = ((p shr 8 and 0xFF) - mean) / std
        out[i * 3 + 2] = ((p and 0xFF) - mean) / std
    }
    return out
}