From e9d48c8f1d33bf775df36a645dcf8023289bd08e Mon Sep 17 00:00:00 2001 From: Mayuri Khinvasara Date: Tue, 9 Jun 2026 15:42:00 +0530 Subject: [PATCH 1/2] Add on-device AI image enhancement feature to enhance images in the Chat screen Major changes: - Integrate the Google Play Services Media Effect Enhancement library to support on-device AI processing. - Implement `ImageEnhancementScreen` and `EnhancementViewModel` to handle image enhancement features, including tonemapping, deblurring, and upscaling. - Add `EnhancementUtils.kt` providing coroutine-based wrappers for enhancement session management, bitmap processing, and module installation. Minor changes: - Add an "AI Enhance" action button to image bubbles within the chat UI. - Update the `ChatMessage` model to include a unique ID and add support for updating media URIs in the `ChatRepository` and `MessageDao`. - Define navigation logic for the new enhancement pane in `Main.kt` and `SocialiteNavigation.kt`. - Update the README to include documentation for the Media Effect Enhancement integration. --- README.md | 3 + app/build.gradle.kts | 2 + .../samples/socialite/data/MessageDao.kt | 3 + .../socialite/repository/ChatRepository.kt | 4 + .../android/samples/socialite/ui/Main.kt | 14 + .../samples/socialite/ui/chat/ChatMessage.kt | 1 + .../samples/socialite/ui/chat/ChatScreen.kt | 28 +- .../socialite/ui/chat/ChatViewModel.kt | 1 + .../ui/chat/component/MessageBubble.kt | 40 ++- .../ui/mediaenhancement/EnhancementUtils.kt | 202 +++++++++++ .../mediaenhancement/EnhancementViewModel.kt | 329 ++++++++++++++++++ .../ImageEnhancementScreen.kt | 301 ++++++++++++++++ .../ui/navigation/SocialiteNavigation.kt | 4 + gradle/libs.versions.toml | 2 + 14 files changed, 925 insertions(+), 9 deletions(-) create mode 100644 app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt create mode 100644 app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt create mode 100644 app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt diff --git a/README.md b/README.md index 0e2f258c..1019ee3e 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ Jetpack APIs used: - [CameraX](https://developer.android.com/jetpack/androidx/releases/camera) Capturing photos and videos - [Room](https://developer.android.com/jetpack/androidx/releases/room) Persisting messages in a SQLite database - [Window](https://developer.android.com/jetpack/androidx/releases/window) Detecting foldable device states + - [Media Effect Enhancement](https://developers.google.com/android/reference/com/google/android/gms/media/effect/enhancement/package-summary) On-device AI image enhancement powered by Google Play Services The app also integrates the Gemini API that powers chatbot capabilities: @@ -38,6 +39,7 @@ Here are the screens that make up SociaLite: [Photo Picker](https://developer.android.com/training/data-storage/shared/photopicker)). - *Camera Screen:* Clicking the camera icon in the Chat Screen opens the in-app camera for taking photos and videos. - *Video Edit Screen:* After taking a video with the in-app camera, users can do some minor edits on this screen. + - *Image Enhancement Screen:* Clicking the AI Enhance icon on an image in the Chat Screen opens this screen to apply on-device AI enhancements like tonemapping, deblurring, and upscaling. - *Settings Screen:* A basic settings screen for tasks like resetting the chat history. ## Project Structure @@ -56,6 +58,7 @@ The project is organized into several modules and directories: - `settings/`: Code for the settings screen. See the [README](app/src/main/java/com/google/android/samples/socialite/ui/home/settings/README.md) for more details. - `timeline/`: Code for the timeline screen. See the [README](app/src/main/java/com/google/android/samples/socialite/ui/home/timeline/README.md) for more details. - `videoedit/`: Code for the video editing screen. See the [README](app/src/main/java/com/google/android/samples/socialite/ui/videoedit/README.md) for more details. + - `mediaenhancement/`: Code for the on-device AI image enhancement screen. It leverages Google Play Services to perform background processing for features like tonemapping, deblurring, and upscaling. - `res/`: Application resources (layouts, drawables, values, etc.). - `assets/`: Static assets like images and shaders. - `build.gradle.kts`: Gradle build file for the app module. diff --git a/app/build.gradle.kts b/app/build.gradle.kts index 337f1f6e..1dde3f55 100644 --- a/app/build.gradle.kts +++ b/app/build.gradle.kts @@ -194,4 +194,6 @@ dependencies { implementation(libs.navigation3.runtime) implementation(libs.navigation3.ui) implementation(libs.lifecycle.viewmodel.navigation3) + + implementation(libs.play.services.media.effect.enhancement) } diff --git a/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt b/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt index 7d242d41..ee4dedf9 100644 --- a/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt +++ b/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt @@ -42,4 +42,7 @@ interface MessageDao { @Query("DELETE FROM Message") suspend fun clearAll() + + @Query("UPDATE Message SET mediaUri = :newUri WHERE id = :messageId") + suspend fun updateMessageMediaUri(messageId: Long, newUri: String) } diff --git a/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt b/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt index 5cf366f0..4a6db4a3 100644 --- a/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt +++ b/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt @@ -412,4 +412,8 @@ class ChatRepository @Inject internal constructor( } } } + // Message update needed to save the Enhanced image bitmap so it persists + suspend fun updateMessageMediaUri(messageId: Long, newUri: String) { + messageDao.updateMessageMediaUri(messageId, newUri) + } } diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt b/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt index 3388b0d1..f316e5dd 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt @@ -66,6 +66,7 @@ import com.google.android.samples.socialite.ui.navigation.TopLevelDestination import com.google.android.samples.socialite.ui.photopicker.navigation.PhotoPickerRoute import com.google.android.samples.socialite.ui.player.VideoPlayerScreen import com.google.android.samples.socialite.ui.videoedit.VideoEditScreen +import com.google.android.samples.socialite.ui.mediaenhancement.ImageEnhancementScreen import kotlinx.coroutines.CoroutineScope import kotlinx.coroutines.launch @@ -186,6 +187,9 @@ fun MainNavigation( onInspectClicked = { uri -> backStack.add(Pane.MetadataInspector(uri)) }, + onEnhanceClicked = { messageId, uri -> + backStack.add(Pane.ImageEnhancement(backStackKey.chatId, messageId, uri)) + } ) } @@ -247,6 +251,16 @@ fun MainNavigation( ) } + is Pane.ImageEnhancement -> NavEntry(backStackKey) { + ImageEnhancementScreen( + chatId = backStackKey.chatId, + messageId = backStackKey.messageId, + uri = backStackKey.uri, + onCloseButtonClicked = { backStack.removeLastOrNull() }, + onFinishEditing = { backStack.removeLastOrNull() } + ) + } + is Pane.MetadataInspector -> NavEntry(backStackKey) { MetadataInspector(backStackKey.uri) } diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt index 558a7cc0..ed61a2dc 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt @@ -19,6 +19,7 @@ package com.google.android.samples.socialite.ui.chat import android.net.Uri data class ChatMessage( + val id: Long, val text: String, val mediaUri: String?, val mediaMimeType: String?, diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt index a33250bb..70dcb203 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt @@ -59,6 +59,7 @@ import androidx.compose.runtime.LaunchedEffect import androidx.compose.runtime.getValue import androidx.compose.runtime.remember import androidx.compose.runtime.rememberCoroutineScope +import androidx.compose.runtime.setValue import androidx.compose.ui.Alignment import androidx.compose.ui.Modifier import androidx.compose.ui.draw.clip @@ -67,6 +68,7 @@ import androidx.compose.ui.focus.focusProperties import androidx.compose.ui.focus.focusRequester import androidx.compose.ui.graphics.Color import androidx.compose.ui.input.nestedscroll.nestedScroll +import androidx.compose.ui.platform.LocalContext import androidx.compose.ui.platform.LocalLayoutDirection import androidx.compose.ui.res.stringResource import androidx.compose.ui.tooling.preview.Preview @@ -88,6 +90,7 @@ import com.google.android.samples.socialite.ui.chat.component.MessageBubble import com.google.android.samples.socialite.ui.chat.component.mediaItemDropTarget import com.google.android.samples.socialite.ui.chat.component.scrollWithKeyboards import com.google.android.samples.socialite.ui.components.tryRequestFocus +import com.google.android.samples.socialite.ui.mediaenhancement.EnhancementSupportManager import com.google.android.samples.socialite.ui.rememberIconPainter @Composable @@ -102,8 +105,12 @@ fun ChatScreen( prefilledText: String? = null, prefilledImageUri: String? = null, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, viewModel: ChatViewModel = hiltViewModel(), ) { + val context = LocalContext.current + var isEnhancementSupported by remember { androidx.compose.runtime.mutableStateOf(false) } + LaunchedEffect(chatId) { viewModel.setChatId(chatId) if (prefilledText != null) { @@ -112,6 +119,7 @@ fun ChatScreen( if (prefilledImageUri != null) { viewModel.prefillInputImage(prefilledImageUri) } + isEnhancementSupported = EnhancementSupportManager.checkSupport(context) } val chat by viewModel.chat.collectAsStateWithLifecycle() val messages by viewModel.messages.collectAsStateWithLifecycle() @@ -125,6 +133,7 @@ fun ChatScreen( textFieldState = textFieldState, attachedMedia = attachedMedia, sendEnabled = sendEnabled, + isEnhancementSupported = isEnhancementSupported, onBackPressed = onBackPressed, onSendClick = { viewModel.send() }, onCameraClick = onCameraClick, @@ -133,6 +142,7 @@ fun ChatScreen( onMediaItemAttached = viewModel::attachMedia, onRemoveAttachedMediaItem = viewModel::removeAttachedMedia, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, modifier = modifier .clip(RoundedCornerShape(5)), ) @@ -174,6 +184,7 @@ private fun ChatContent( textFieldState: TextFieldState, attachedMedia: MediaItem?, sendEnabled: Boolean, + isEnhancementSupported: Boolean = false, onBackPressed: (() -> Unit)?, onSendClick: () -> Unit, onCameraClick: () -> Unit, @@ -182,6 +193,7 @@ private fun ChatContent( onMediaItemAttached: (MediaItem) -> Unit, onRemoveAttachedMediaItem: () -> Unit, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, modifier: Modifier = Modifier, ) { val topAppBarState = rememberTopAppBarState() @@ -221,8 +233,10 @@ private fun ChatContent( modifier = Modifier .fillMaxWidth() .weight(1f), + isEnhancementSupported = isEnhancementSupported, onVideoClick = onVideoClick, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, ) InputBar( textFieldState = textFieldState, @@ -309,8 +323,10 @@ private fun MessageList( contentPadding: PaddingValues, modifier: Modifier = Modifier, state: LazyListState = rememberLazyListState(), + isEnhancementSupported: Boolean = false, onVideoClick: (uri: String) -> Unit = {}, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, ) { LazyColumn( modifier = modifier, @@ -338,10 +354,12 @@ private fun MessageList( } MessageBubble( message = message, + isEnhancementSupported = isEnhancementSupported, onVideoClick = { message.mediaUri?.let { onVideoClick(it) } }, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, ) } } @@ -355,11 +373,11 @@ private fun PreviewChatContent() { ChatContent( chat = ChatDetail(ChatWithLastMessage(0L), listOf(Contact.CONTACTS[0])), messages = listOf( - ChatMessage("Hi!", null, null, 0L, false, null), - ChatMessage("Hello", null, null, 0L, true, null), - ChatMessage("world", null, null, 0L, true, null), - ChatMessage("!", null, null, 0L, true, null), - ChatMessage("Hello, world!", null, null, 0L, true, null), + ChatMessage(1, "Hi!", null, null, 0L, false, null), + ChatMessage(2, "Hello", null, null, 0L, true, null), + ChatMessage(3, "world", null, null, 0L, true, null), + ChatMessage(4, "!", null, null, 0L, true, null), + ChatMessage(5, "Hello, world!", null, null, 0L, true, null), ), textFieldState = TextFieldState("Hello"), attachedMedia = null, diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt index e124b43b..ee165531 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt @@ -74,6 +74,7 @@ class ChatViewModel @Inject constructor( } ChatMessage( + id = message.id, text = message.text, mediaUri = message.mediaUri, mediaMimeType = message.mediaMimeType, diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt index 61e46ea7..c230872c 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt @@ -17,16 +17,20 @@ package com.google.android.samples.socialite.ui.chat.component import android.util.Log +import androidx.compose.foundation.background import androidx.compose.foundation.clickable import androidx.compose.foundation.focusable import androidx.compose.foundation.indication import androidx.compose.foundation.interaction.MutableInteractionSource +import androidx.compose.foundation.layout.Box import androidx.compose.foundation.layout.Column import androidx.compose.foundation.layout.height import androidx.compose.foundation.layout.padding import androidx.compose.foundation.layout.size +import androidx.compose.foundation.shape.CircleShape import androidx.compose.foundation.text.selection.SelectionContainer import androidx.compose.material.icons.Icons +import androidx.compose.material.icons.filled.AutoAwesome import androidx.compose.material.icons.filled.Info import androidx.compose.material3.Icon import androidx.compose.material3.IconButton @@ -38,6 +42,7 @@ import androidx.compose.runtime.Composable import androidx.compose.runtime.remember import androidx.compose.ui.Alignment import androidx.compose.ui.Modifier +import androidx.compose.ui.graphics.Color import androidx.compose.ui.platform.LocalContext import androidx.compose.ui.unit.dp import androidx.compose.ui.zIndex @@ -53,8 +58,10 @@ private const val TAG = "ChatUI" internal fun MessageBubble( message: ChatMessage, modifier: Modifier = Modifier, + isEnhancementSupported: Boolean = false, onVideoClick: () -> Unit = {}, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, ) { MessageBubbleSurface( isVideoContentAttached = message.isVideoContentAttached, @@ -71,7 +78,9 @@ internal fun MessageBubble( AttachedMedia( message = message, modifier = Modifier.draggableMediaItem(message), + isEnhancementSupported = isEnhancementSupported, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, ) } } @@ -116,17 +125,40 @@ private fun MessageBubbleSurface( private fun AttachedMedia( message: ChatMessage, modifier: Modifier = Modifier, + isEnhancementSupported: Boolean = false, onInspectClicked: (uri: String) -> Unit, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit, ) { val uri = message.mediaUri if (uri != null) { ContextMenuArea(chatMessage = message) { when { message.isImageContentAttached -> { - Photo( - uri = uri, - modifier = modifier, - ) + Box { + Photo( + uri = uri, + modifier = modifier, + ) + if (isEnhancementSupported) { + IconButton( + onClick = { onEnhanceClicked(message.id, uri) }, + modifier = Modifier + .align(Alignment.TopEnd) + .padding(8.dp) + .background( + color = MaterialTheme.colorScheme.scrim.copy(alpha = 0.4f), + shape = CircleShape + ) + .zIndex(1f) + ) { + Icon( + imageVector = Icons.Filled.AutoAwesome, + contentDescription = "AI Enhance", + tint = Color.White + ) + } + } + } } message.isVideoContentAttached -> { diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt new file mode 100644 index 00000000..6e152735 --- /dev/null +++ b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt @@ -0,0 +1,202 @@ +package com.google.android.samples.socialite.ui.mediaenhancement + +import android.content.Context +import android.graphics.Bitmap +import com.google.android.gms.common.api.Status +import com.google.android.gms.media.effect.enhancement.EnhancementCallback +import com.google.android.gms.media.effect.enhancement.EnhancementClient +import com.google.android.gms.media.effect.enhancement.EnhancementOptions +import com.google.android.gms.media.effect.enhancement.EnhancementSession +import com.google.android.gms.media.effect.enhancement.EnhancementSessionCallback +import com.google.android.gms.media.effect.enhancement.Enhancement +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.suspendCancellableCoroutine +import kotlinx.coroutines.withContext +import java.util.concurrent.Executor +import kotlin.coroutines.resume +import kotlin.coroutines.resumeWithException +import android.util.Log +import android.os.Build +import androidx.annotation.RequiresApi + + +class EnhancementFailedException(val errorCode: Int, message: String) : Exception(message) + +/** + * A modern coroutine wrapper for the enhancement process. + * + * This suspend function encapsulates the entire callback-based process of creating a session, + * processing a bitmap, and handling success or failure. + * + * @param context The application context. + * @param bitmap The input bitmap to enhance. + * @param options The enhancement options to apply. + * @param executor The executor on which to run the callbacks. + * @return The enhanced [Bitmap] on success. + * @throws [EnhancementFailedException] if any step of the process fails. + */ +/** + * Extension to create an enhancement session asynchronously. + */ +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.createSessionAsync( + options: EnhancementOptions, + executor: Executor, +): EnhancementSession = withContext(Dispatchers.Main) { + suspendCancellableCoroutine { continuation -> + val callback = object : EnhancementSessionCallback { + override fun onSessionCreated(session: EnhancementSession) { + continuation.resume(session) + } + + override fun onSessionCreationFailed(status: Status) { + continuation.resumeWithException( + Exception("Session creation failed: ${status.statusMessage} (${status.statusCode})"), + ) + } + + override fun onSessionDestroyed() { + // Log or handle if needed + } + + override fun onSessionDisconnected(status: Status) { + // Log or handle if needed + } + } + + this@createSessionAsync.createSession(options, callback) + .addOnFailureListener(executor) { e -> + if (continuation.isActive) { + continuation.resumeWithException(e) + } + } + } +} + +/** + * Extension to process a bitmap using an existing session. + */ +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementSession.processBitmapAsync( + bitmap: Bitmap, + options: EnhancementOptions, +): Bitmap = suspendCancellableCoroutine { continuation -> + val callback = object : EnhancementCallback { + override fun onBitmapProcessed(bitmap: Bitmap) { + continuation.resume(bitmap) + } + + override fun onError(statusCode: Int) { + continuation.resumeWithException( + Exception("Processing failed with status code: $statusCode"), + ) + } + + override fun onSurfaceProcessed(timestamp: Long) { /* Not used in bitmap flow */ + } + } + + this.process(bitmap, options, callback) +} + +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.installModuleAsync(onProgress: (Int) -> Unit): Boolean = + suspendCancellableCoroutine { continuation -> + val callback = object : EnhancementClient.InstallStatusCallback { + override fun onDownloadPending() { + Log.d("EnhancementUtils", "onDownloadPending") + } + + override fun onDownloadStart() { + Log.d("EnhancementUtils", "onDownloadStart") + } + + override fun onDownloadPaused() { + Log.d("EnhancementUtils", "onDownloadPaused") + } + + override fun onDownloadProgressUpdate(progress: Int) { + Log.d("EnhancementUtils", "onDownloadProgressUpdate: $progress") + onProgress(progress) + } + + override fun onDownloadComplete() { + Log.d("EnhancementUtils", "onDownloadComplete") + } + + override fun onInstalled() { + Log.d("EnhancementUtils", "onInstalled") + } + + override fun onCancelled() { + Log.d("EnhancementUtils", "onCancelled") + if (continuation.isActive) continuation.resume(false) + } + + override fun onError(description: String) { + Log.e("EnhancementUtils", "onError: $description") + if (continuation.isActive) continuation.resumeWithException(Exception(description)) + } + } + + this.installModule(callback) + .addOnSuccessListener { result -> + if (continuation.isActive) continuation.resume(result) + } + .addOnFailureListener { e -> + if (continuation.isActive) continuation.resumeWithException(e) + } + } + +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.isModuleInstalledAsync(): Boolean = + suspendCancellableCoroutine { continuation -> + this.isModuleInstalled() + .addOnSuccessListener { result -> continuation.resume(result) } + .addOnFailureListener { e -> continuation.resumeWithException(e) } + } + +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.isDeviceSupportedAsync(): Boolean = + suspendCancellableCoroutine { continuation -> + this.isDeviceSupported() + .addOnSuccessListener { result -> continuation.resume(result) } + .addOnFailureListener { e -> continuation.resumeWithException(e) } + } + +object EnhancementSupportManager { + private var isSupported: Boolean? = null + + suspend fun checkSupport(context: Context): Boolean { + if (isSupported != null) return isSupported!! + + if (android.os.Build.VERSION.SDK_INT >= android.os.Build.VERSION_CODES.R) { + try { + val client = Enhancement.getClient(context.applicationContext) + + // First check if it's already installed + val installed = client.isModuleInstalledAsync() + if (!installed) { + // If not installed, attempt to install it before querying device support + try { + client.installModuleAsync { progress -> + // Silent background install progress + } + } catch (e: Exception) { + Log.e("EnhancementSupport", "Failed to silently install module", e) + } + } + + isSupported = client.isDeviceSupportedAsync() + return isSupported!! + } catch (e: Exception) { + Log.e("EnhancementSupport", "Error checking support", e) + isSupported = false + return false + } + } else { + isSupported = false + return false + } + } +} diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt new file mode 100644 index 00000000..1f6c4f03 --- /dev/null +++ b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt @@ -0,0 +1,329 @@ +package com.google.android.samples.socialite.ui.mediaenhancement + +import android.app.Application +import android.content.Context +import android.graphics.Bitmap +import android.graphics.BitmapFactory +import android.net.Uri +import android.os.Build +import android.util.Log +import androidx.annotation.RequiresApi +import androidx.lifecycle.AndroidViewModel +import androidx.lifecycle.viewModelScope +import com.google.android.gms.media.effect.enhancement.Enhancement +import com.google.android.gms.media.effect.enhancement.EnhancementClient +import com.google.android.gms.media.effect.enhancement.EnhancementMode +import com.google.android.gms.media.effect.enhancement.EnhancementOptions +import com.google.android.gms.media.effect.enhancement.EnhancementSession +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.flow.StateFlow +import kotlinx.coroutines.flow.asStateFlow +import kotlinx.coroutines.flow.update +import kotlinx.coroutines.launch +import java.util.concurrent.Executors +import dagger.hilt.android.lifecycle.HiltViewModel +import javax.inject.Inject +import com.google.android.samples.socialite.repository.ChatRepository +import java.io.File +import java.io.FileOutputStream +import androidx.core.net.toUri +import kotlinx.coroutines.withContext + +private const val TAG = "EnhancementViewModel" + +// Data class to hold all information about an image +data class ImageInfo( + val bitmap: Bitmap? = null, val latency: Long? = null, val qualityScore: Int? = null +) + +// Defines the state of the UI +data class UiState( + val originalImage: ImageInfo? = null, + val enhancedImage: ImageInfo? = null, + val isLoading: Boolean = false, + val selectedOptions: Set = setOf("Tonemap"), + val enhancementMode: Int = EnhancementMode.BITMAP, + val isModuleInstalling: Boolean = false, + val moduleInstallProgress: Int = 0, + val moduleInstallError: String? = null, + val isModuleReady: Boolean = false, + val moduleStatus: String = "Unknown", + val isDeviceSupported: Boolean = true, + val enhancementError: String? = null +) + +@HiltViewModel +@RequiresApi(Build.VERSION_CODES.R) +class EnhancementViewModel @Inject constructor( + application: Application, + private val chatRepository: ChatRepository +) : AndroidViewModel(application) { + + private val _uiState = MutableStateFlow(UiState()) + val uiState: StateFlow = _uiState.asStateFlow() + + @RequiresApi(Build.VERSION_CODES.R) + private val enhancementClient: EnhancementClient = Enhancement.getClient(application) + private val enhancementExecutor = Executors.newSingleThreadExecutor() + private var enhancementSession: EnhancementSession? = null + + init { + checkAndInstallModule() + } + + @RequiresApi(Build.VERSION_CODES.R) + private fun checkAndInstallModule() { + viewModelScope.launch { + _uiState.update { it.copy(moduleStatus = "Checking for device support...") } + try { + if (!enhancementClient.isDeviceSupportedAsync()) { + _uiState.update { + it.copy( + isDeviceSupported = false, moduleStatus = "Device not supported" + ) + } + return@launch + } + + _uiState.update { it.copy(isDeviceSupported = true) } + + if (!enhancementClient.isModuleInstalledAsync()) { + _uiState.update { + it.copy( + isModuleInstalling = true, + moduleInstallError = null, + moduleStatus = "Not Installed" + ) + } + + + _uiState.update { it.copy(moduleStatus = "Installing...") } + val installed = enhancementClient.installModuleAsync { progress -> + _uiState.update { it.copy(moduleInstallProgress = progress) } + } + if (!installed) { + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallError = "Module installation failed or cancelled.", + moduleStatus = "Install Failed" + ) + } + return@launch + } + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallProgress = 100, + isModuleReady = true, + moduleStatus = "Installed" + ) + } + } else { + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallProgress = 100, + isModuleReady = true, + moduleStatus = "Installed" + ) + } + } + } catch (e: Exception) { + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallError = "Failed to check or install module: ${e.message}", + moduleStatus = "Error" + ) + } + } + } + } + + private suspend fun createSession(): EnhancementSession? { + if (!_uiState.value.isDeviceSupported) return null + + val originalBitmap = _uiState.value.originalImage?.bitmap ?: return null + + return try { + val options = getEnhancementOptionsFor(originalBitmap, _uiState.value.selectedOptions) + enhancementClient.createSessionAsync(options, enhancementExecutor) + } catch (e: Exception) { + _uiState.update { it.copy(enhancementError = "Failed to create session: ${e.message}") } + null + } + } + + override fun onCleared() { + enhancementSession?.release() + enhancementSession = null + enhancementExecutor.shutdown() + super.onCleared() + } + + fun onImageSelected(uri: Uri, context: Context) { + viewModelScope.launch(Dispatchers.IO) { + _uiState.update { + it.copy( + isLoading = true, + originalImage = null, + enhancedImage = null, + enhancementError = null + ) + } + // Release any previous session, since we have a new image. + enhancementSession?.release() + enhancementSession = null + + val originalBitmap = decodeBitmapFromUri(uri, context) + + if (originalBitmap == null) { + Log.e(TAG, "Failed to decode bitmap from URI.") + _uiState.update { + it.copy( + isLoading = false, + enhancementError = "Failed to load image." + ) + } + return@launch + } + + val originalImageInfo = ImageInfo(bitmap = originalBitmap) + // Show the original image and stop loading. Session will be created on demand. + _uiState.update { it.copy(originalImage = originalImageInfo, isLoading = false) } + } + } + + fun onOptionSelected(option: String) { + val currentOptions = _uiState.value.selectedOptions + val newOptions = if (option in currentOptions) { + currentOptions - option + } else { + currentOptions + option + } + _uiState.update { it.copy(selectedOptions = newOptions) } + } + + fun setEnhancementMode(mode: Int) { + _uiState.update { it.copy(enhancementMode = mode) } + } + + fun enhanceImage() { + val originalBitmap = _uiState.value.originalImage?.bitmap ?: return + + viewModelScope.launch(Dispatchers.IO) { + _uiState.update { it.copy(isLoading = true, enhancementError = null) } + try { + processImage(originalBitmap) + } finally { + _uiState.update { it.copy(isLoading = false) } + } + } + } + + private suspend fun processImage(bitmap: Bitmap) { + try { + // If the session doesn't exist, create it now. + if (enhancementSession == null) { + enhancementSession = createSession() + } + + // If session creation failed or is not possible, just return. + // The createSession function will have set the error message. + val session = enhancementSession ?: return + + val options = getEnhancementOptionsFor(bitmap, _uiState.value.selectedOptions) + val enhancementStartTime = System.currentTimeMillis() + val enhancedBitmap = session.processBitmapAsync(bitmap, options) + val enhancementLatency = System.currentTimeMillis() - enhancementStartTime + + _uiState.update { + it.copy( + enhancedImage = ImageInfo(bitmap = enhancedBitmap, latency = enhancementLatency) + ) + } + } catch (e: Exception) { + Log.e(TAG, "Enhancement failed: ${e.message}", e) + _uiState.update { + it.copy( + enhancedImage = null, enhancementError = "Enhancement failed: ${e.message}" + ) + } + } + } + + private fun getEnhancementOptionsFor( + bitmap: Bitmap, + selectedOptions: Set + ): EnhancementOptions { + return EnhancementOptions( + bitmap.width, + bitmap.height, + _uiState.value.enhancementMode, + "Tonemap" in selectedOptions, + "Deblur & DeNoise" in selectedOptions, + false, + "Upscale" in selectedOptions, + false + ) + } + + private fun decodeBitmapFromUri(uri: Uri, context: Context): Bitmap? { + return try { + context.contentResolver.openInputStream(uri)?.use { inputStream -> + BitmapFactory.decodeStream(inputStream) + } + } catch (e: Exception) { + Log.e(TAG, "Error decoding bitmap from URI", e) + null + } + } + + fun saveEnhancedImage(messageId: Long, onComplete: () -> Unit) { + val enhancedBitmap = _uiState.value.enhancedImage?.bitmap ?: return + viewModelScope.launch(Dispatchers.IO) { + try { + val context = getApplication() + val directory = File(context.filesDir, "media") + if (!directory.exists()) { + directory.mkdirs() + } + + // Prevent storage bloat by deleting previous enhanced files + directory.listFiles()?.forEach { file -> + if (file.name.startsWith("enhanced_")) { + file.delete() + } + } + + val filename = "enhanced_${System.currentTimeMillis()}.jpg" + val file = File(directory, filename) + FileOutputStream(file).use { out -> + var success = enhancedBitmap.compress(Bitmap.CompressFormat.JPEG, 100, out) + if (!success) { + Log.w(TAG, "Failed to compress hardware bitmap, trying software fallback.") + val swBitmap = enhancedBitmap.copy(Bitmap.Config.ARGB_8888, false) + if (swBitmap != null) { + success = swBitmap.compress(Bitmap.CompressFormat.JPEG, 100, out) + } + } + if (!success) { + Log.e(TAG, "Bitmap compression completely failed!") + } + } + + val newUri = file.toUri().toString() + Log.d(TAG, "Saved enhanced image to $newUri, updating message $messageId") + chatRepository.updateMessageMediaUri(messageId, newUri) + + withContext(Dispatchers.Main) { + onComplete() + } + } catch (e: Exception) { + Log.e(TAG, "Failed to save image", e) + } + } + } +} diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt new file mode 100644 index 00000000..fb9f19fa --- /dev/null +++ b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt @@ -0,0 +1,301 @@ +package com.google.android.samples.socialite.ui.mediaenhancement + +import android.graphics.Bitmap +import android.net.Uri +import android.os.Build +import androidx.annotation.RequiresApi +import androidx.compose.foundation.Image +import androidx.compose.foundation.background +import androidx.compose.foundation.layout.Arrangement +import androidx.compose.foundation.layout.Box +import androidx.compose.foundation.layout.Column +import androidx.compose.foundation.layout.Row +import androidx.compose.foundation.layout.Spacer +import androidx.compose.foundation.layout.fillMaxSize +import androidx.compose.foundation.layout.fillMaxWidth +import androidx.compose.foundation.layout.height +import androidx.compose.foundation.layout.padding +import androidx.compose.foundation.lazy.LazyRow +import androidx.compose.foundation.lazy.items +import androidx.compose.material3.Button +import androidx.compose.material3.CircularProgressIndicator +import androidx.compose.material3.ExperimentalMaterial3Api +import androidx.compose.material3.FilterChip +import androidx.compose.material3.LinearProgressIndicator +import androidx.compose.material3.MaterialTheme +import androidx.compose.material3.OutlinedButton +import androidx.compose.material3.Scaffold +import androidx.compose.material3.Text +import androidx.compose.runtime.Composable +import androidx.compose.runtime.LaunchedEffect +import androidx.compose.runtime.collectAsState +import androidx.compose.runtime.getValue +import androidx.compose.ui.Alignment +import androidx.compose.ui.Modifier +import androidx.compose.ui.draw.clip +import androidx.compose.ui.graphics.Color +import androidx.compose.ui.graphics.asImageBitmap +import androidx.compose.ui.layout.ContentScale +import androidx.compose.ui.platform.LocalContext +import androidx.compose.ui.text.font.FontWeight +import androidx.compose.ui.unit.dp +import androidx.compose.ui.unit.sp +import androidx.hilt.navigation.compose.hiltViewModel +import com.google.android.gms.media.effect.enhancement.EnhancementMode +import java.text.DecimalFormat + +@RequiresApi(Build.VERSION_CODES.R) +@OptIn(ExperimentalMaterial3Api::class) +@Composable +fun ImageEnhancementScreen( + chatId: Long, + messageId: Long, + uri: String, + onCloseButtonClicked: () -> Unit, + onFinishEditing: () -> Unit, + enhancementViewModel: EnhancementViewModel = hiltViewModel() +) { + val context = LocalContext.current + val uiState by enhancementViewModel.uiState.collectAsState() + + LaunchedEffect(uri) { + enhancementViewModel.setEnhancementMode(EnhancementMode.BITMAP) + enhancementViewModel.onImageSelected(Uri.parse(uri), context) + } + + val processingOptions = listOf("Tonemap", "Deblur & DeNoise", "Upscale") + + Scaffold { paddingValues -> + Column( + modifier = Modifier + .fillMaxSize() + .padding(paddingValues) + .padding(start = 16.dp, end = 16.dp, bottom = 16.dp), + horizontalAlignment = Alignment.CenterHorizontally + ) { + Text( + text = "Module Status: ${uiState.moduleStatus}", + style = MaterialTheme.typography.bodyMedium, + fontWeight = FontWeight.Bold + ) + + Spacer(modifier = Modifier.height(12.dp)) + + Row( + modifier = Modifier.fillMaxWidth(), + verticalAlignment = Alignment.CenterVertically + ) { + LazyRow( + modifier = Modifier.weight(1f), + horizontalArrangement = Arrangement.spacedBy(8.dp) + ) { + items(processingOptions) { option -> + FilterChip( + selected = option in uiState.selectedOptions, + onClick = { enhancementViewModel.onOptionSelected(option) }, + label = { Text(option) } + ) + } + } + Spacer(modifier = Modifier.padding(start = 8.dp)) + Button( + onClick = { enhancementViewModel.enhanceImage() }, + enabled = uiState.isModuleReady && !uiState.isLoading && uiState.originalImage?.bitmap != null + ) { + val buttonText = when { + uiState.isModuleInstalling -> "Installing..." + uiState.isLoading -> "Enhancing..." + else -> "AI Enhance" + } + Text(buttonText) + } + } + + if (uiState.isModuleInstalling) { + Spacer(modifier = Modifier.height(8.dp)) + Column(horizontalAlignment = Alignment.CenterHorizontally) { + LinearProgressIndicator( + progress = { uiState.moduleInstallProgress / 100f }, + modifier = Modifier.fillMaxWidth() + ) + Text( + text = "Downloading Enhancement module (${uiState.moduleInstallProgress}%)", + style = MaterialTheme.typography.bodySmall + ) + } + } + + if (uiState.moduleInstallError != null) { + Spacer(modifier = Modifier.height(8.dp)) + Text( + text = uiState.moduleInstallError ?: "Unknown error", + color = MaterialTheme.colorScheme.error, + style = MaterialTheme.typography.bodySmall + ) + } + + if (uiState.enhancementError != null) { + Spacer(modifier = Modifier.height(8.dp)) + Text( + text = uiState.enhancementError ?: "Unknown enhancement error", + color = MaterialTheme.colorScheme.error, + style = MaterialTheme.typography.bodySmall + ) + } + + Spacer(modifier = Modifier.height(16.dp)) + + Column( + modifier = Modifier.weight(1f), + verticalArrangement = Arrangement.spacedBy(16.dp) + ) { + EnhancementImageCard( + title = "Original", + imageInfo = uiState.originalImage, + modifier = Modifier.weight(1f) + ) + EnhancementImageCard( + title = "Enhanced", + imageInfo = uiState.enhancedImage, + isLoading = uiState.isLoading, + error = uiState.enhancementError, + modifier = Modifier.weight(1f) + ) + } + + Spacer(modifier = Modifier.height(16.dp)) + + Row( + modifier = Modifier.fillMaxWidth(), + horizontalArrangement = Arrangement.SpaceEvenly + ) { + OutlinedButton( + onClick = onCloseButtonClicked, + modifier = Modifier.weight(1f).padding(end = 8.dp) + ) { + Text("Cancel") + } + Button( + onClick = { + enhancementViewModel.saveEnhancedImage(messageId, onFinishEditing) + }, + enabled = uiState.enhancedImage?.bitmap != null, + modifier = Modifier.weight(1f).padding(start = 8.dp) + ) { + Text("Confirm") + } + } + } + } +} + +@Composable +fun EnhancementImageCard( + title: String, + imageInfo: ImageInfo?, + modifier: Modifier = Modifier, + isLoading: Boolean = false, + error: String? = null +) { + Box( + modifier = modifier + .fillMaxSize() + .background(MaterialTheme.colorScheme.surfaceVariant, MaterialTheme.shapes.medium) + .clip(MaterialTheme.shapes.medium), + contentAlignment = Alignment.Center + ) { + when { + isLoading -> { + CircularProgressIndicator() + } + error != null -> { + Text( + text = "Enhancement Failed", + color = MaterialTheme.colorScheme.error, + style = MaterialTheme.typography.bodyMedium + ) + } + imageInfo?.bitmap != null -> { + Image( + bitmap = imageInfo.bitmap.asImageBitmap(), + contentDescription = title, + modifier = Modifier.fillMaxSize(), + contentScale = ContentScale.Fit + ) + InfoTag( + title = title, + latency = imageInfo.latency, + bitmap = imageInfo.bitmap, + quality = imageInfo.qualityScore, + modifier = Modifier.align(Alignment.TopStart) + ) + } + else -> { + Text( + text = if (title == "Original") "Loading..." else "Enhanced Image", + style = MaterialTheme.typography.bodyMedium, + color = MaterialTheme.colorScheme.onSurfaceVariant + ) + } + } + } +} + +private fun formatMemorySize(bytes: Int): String { + val kb = bytes / 1024.0 + val mb = kb / 1024.0 + val decimalFormat = DecimalFormat("#.##") + return when { + mb >= 1 -> "${decimalFormat.format(mb)} MB" + else -> "${decimalFormat.format(kb)} KB" + } +} + +@Composable +private fun InfoTag( + title: String, + latency: Long?, + bitmap: Bitmap?, + quality: Int?, + modifier: Modifier = Modifier +) { + Column( + modifier = modifier + .padding(4.dp) + .background(Color.Black.copy(alpha = 0.6f), MaterialTheme.shapes.small) + .padding(horizontal = 6.dp, vertical = 2.dp), + verticalArrangement = Arrangement.spacedBy(0.dp) + ) { + Text( + text = title, + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.Bold + ) + if (bitmap != null) { + val memorySize = formatMemorySize(bitmap.byteCount) + Text( + text = "$memorySize, ${bitmap.width}x${bitmap.height}", + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.SemiBold + ) + } + if (latency != null && title != "Original") { + Text( + text = "AI Latency: ${latency}ms", + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.SemiBold + ) + } + if (quality != null) { + Text( + text = "Quality: $quality/100", + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.SemiBold + ) + } + } +} diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt b/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt index 3b98e6cd..b0242dd7 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt @@ -81,6 +81,10 @@ sealed interface Pane : Parcelable { @Parcelize @Serializable data class MetadataInspector(val uri: String) : Pane + + @Parcelize + @Serializable + data class ImageEnhancement(val chatId: Long, val messageId: Long, val uri: String) : Pane } enum class TopLevelDestination( diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index a8c0fa2c..2b7fc9ec 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -63,6 +63,7 @@ visionCommon = "17.3.0" segmentationSelfie = "16.0.0-beta6" litert = "1.3.0" datastore = "1.1.4" +playServicesMediaEffectEnhancement = "16.0.0-beta04" @@ -148,6 +149,7 @@ firebase-ai = { group = "com.google.firebase", name = "firebase-ai" } firebase-analytics = { group = "com.google.firebase", name = "firebase-analytics" } firebase-messaging = {group = "com.google.firebase", name = "firebase-messaging" } vision-common = { group = "com.google.mlkit", name = "vision-common", version.ref = "visionCommon" } +play-services-media-effect-enhancement = { module = "com.google.android.gms:play-services-media-effect-enhancement", version.ref = "playServicesMediaEffectEnhancement" } [plugins] From 86dc06a580b6a46f65ca0db1fc8887867387a3bb Mon Sep 17 00:00:00 2001 From: Mayuri Khinvasara Date: Tue, 9 Jun 2026 15:42:00 +0530 Subject: [PATCH 2/2] Add on-device AI image enhancement feature to enhance images in the Chat screen Major changes: - Integrate the Google Play Services Media Effect Enhancement library to support on-device AI processing. - Implement `ImageEnhancementScreen` and `EnhancementViewModel` to handle image enhancement features, including tonemapping, deblurring, and upscaling. - Add `EnhancementUtils.kt` providing coroutine-based wrappers for enhancement session management, bitmap processing, and module installation. Minor changes: - Add an "AI Enhance" action button to image bubbles within the chat UI. - Update the `ChatMessage` model to include a unique ID and add support for updating media URIs in the `ChatRepository` and `MessageDao`. - Define navigation logic for the new enhancement pane in `Main.kt` and `SocialiteNavigation.kt`. - Update the README to include documentation for the Media Effect Enhancement integration. --- README.md | 3 + app/build.gradle.kts | 1 + .../samples/socialite/data/MessageDao.kt | 3 + .../socialite/repository/ChatRepository.kt | 11 +- .../android/samples/socialite/ui/Main.kt | 21 +- .../samples/socialite/ui/chat/ChatMessage.kt | 1 + .../samples/socialite/ui/chat/ChatScreen.kt | 28 +- .../socialite/ui/chat/ChatViewModel.kt | 1 + .../ui/chat/component/MessageBubble.kt | 40 ++- .../ui/mediaenhancement/EnhancementUtils.kt | 187 ++++++++++ .../mediaenhancement/EnhancementViewModel.kt | 336 ++++++++++++++++++ .../ImageEnhancementScreen.kt | 298 ++++++++++++++++ .../ui/navigation/SocialiteNavigation.kt | 4 + gradle/libs.versions.toml | 2 + 14 files changed, 922 insertions(+), 14 deletions(-) create mode 100644 app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt create mode 100644 app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt create mode 100644 app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt diff --git a/README.md b/README.md index 0e2f258c..1019ee3e 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ Jetpack APIs used: - [CameraX](https://developer.android.com/jetpack/androidx/releases/camera) Capturing photos and videos - [Room](https://developer.android.com/jetpack/androidx/releases/room) Persisting messages in a SQLite database - [Window](https://developer.android.com/jetpack/androidx/releases/window) Detecting foldable device states + - [Media Effect Enhancement](https://developers.google.com/android/reference/com/google/android/gms/media/effect/enhancement/package-summary) On-device AI image enhancement powered by Google Play Services The app also integrates the Gemini API that powers chatbot capabilities: @@ -38,6 +39,7 @@ Here are the screens that make up SociaLite: [Photo Picker](https://developer.android.com/training/data-storage/shared/photopicker)). - *Camera Screen:* Clicking the camera icon in the Chat Screen opens the in-app camera for taking photos and videos. - *Video Edit Screen:* After taking a video with the in-app camera, users can do some minor edits on this screen. + - *Image Enhancement Screen:* Clicking the AI Enhance icon on an image in the Chat Screen opens this screen to apply on-device AI enhancements like tonemapping, deblurring, and upscaling. - *Settings Screen:* A basic settings screen for tasks like resetting the chat history. ## Project Structure @@ -56,6 +58,7 @@ The project is organized into several modules and directories: - `settings/`: Code for the settings screen. See the [README](app/src/main/java/com/google/android/samples/socialite/ui/home/settings/README.md) for more details. - `timeline/`: Code for the timeline screen. See the [README](app/src/main/java/com/google/android/samples/socialite/ui/home/timeline/README.md) for more details. - `videoedit/`: Code for the video editing screen. See the [README](app/src/main/java/com/google/android/samples/socialite/ui/videoedit/README.md) for more details. + - `mediaenhancement/`: Code for the on-device AI image enhancement screen. It leverages Google Play Services to perform background processing for features like tonemapping, deblurring, and upscaling. - `res/`: Application resources (layouts, drawables, values, etc.). - `assets/`: Static assets like images and shaders. - `build.gradle.kts`: Gradle build file for the app module. diff --git a/app/build.gradle.kts b/app/build.gradle.kts index 337f1f6e..35388b9c 100644 --- a/app/build.gradle.kts +++ b/app/build.gradle.kts @@ -194,4 +194,5 @@ dependencies { implementation(libs.navigation3.runtime) implementation(libs.navigation3.ui) implementation(libs.lifecycle.viewmodel.navigation3) + implementation(libs.play.services.media.effect.enhancement) } diff --git a/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt b/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt index 7d242d41..ee4dedf9 100644 --- a/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt +++ b/app/src/main/java/com/google/android/samples/socialite/data/MessageDao.kt @@ -42,4 +42,7 @@ interface MessageDao { @Query("DELETE FROM Message") suspend fun clearAll() + + @Query("UPDATE Message SET mediaUri = :newUri WHERE id = :messageId") + suspend fun updateMessageMediaUri(messageId: Long, newUri: String) } diff --git a/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt b/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt index 5cf366f0..62da8326 100644 --- a/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt +++ b/app/src/main/java/com/google/android/samples/socialite/repository/ChatRepository.kt @@ -30,7 +30,6 @@ import androidx.datastore.preferences.core.edit import androidx.datastore.preferences.preferencesDataStore import com.google.android.samples.socialite.R import com.google.android.samples.socialite.data.ChatDao -import com.google.android.samples.socialite.data.ContactDao import com.google.android.samples.socialite.data.MessageDao import com.google.android.samples.socialite.data.utils.ShortsVideoList import com.google.android.samples.socialite.di.AppCoroutineScope @@ -62,7 +61,6 @@ import kotlinx.coroutines.launch class ChatRepository @Inject internal constructor( private val chatDao: ChatDao, private val messageDao: MessageDao, - private val contactDao: ContactDao, private val notificationHelper: NotificationHelper, private val widgetModelRepository: WidgetModelRepository, @AppCoroutineScope @@ -323,7 +321,7 @@ class ChatRepository @Inject internal constructor( acc.add(message) } else { if (acc.last().isIncoming == message.isIncoming) { - val lastMessage = acc.removeLast() + val lastMessage = acc.removeAt(acc.lastIndex) val combinedMessage = Message( id = lastMessage.id, chatId = chatId, @@ -342,7 +340,7 @@ class ChatRepository @Inject internal constructor( return@fold acc } - pastMessages.removeLast() + pastMessages.removeAt(pastMessages.lastIndex) val pastContents = pastMessages.mapNotNull { message: Message -> val role = if (message.isIncoming) "model" else "user" @@ -412,4 +410,9 @@ class ChatRepository @Inject internal constructor( } } } + + // Message update needed to save the Enhanced image bitmap so it persists + suspend fun updateMessageMediaUri(messageId: Long, newUri: String) { + messageDao.updateMessageMediaUri(messageId, newUri) + } } diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt b/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt index 3388b0d1..268105db 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/Main.kt @@ -20,6 +20,7 @@ package com.google.android.samples.socialite.ui import android.app.Activity import android.content.pm.ActivityInfo +import android.os.Build import androidx.activity.compose.LocalActivity import androidx.compose.animation.ExperimentalSharedTransitionApi import androidx.compose.animation.SharedTransitionLayout @@ -59,6 +60,7 @@ import com.google.android.samples.socialite.ui.home.chatlist.ChatList import com.google.android.samples.socialite.ui.home.chatlist.ChatOpenRequest import com.google.android.samples.socialite.ui.home.settings.Settings import com.google.android.samples.socialite.ui.home.timeline.Timeline +import com.google.android.samples.socialite.ui.mediaenhancement.ImageEnhancementScreen import com.google.android.samples.socialite.ui.metadata.screens.MetadataInspector import com.google.android.samples.socialite.ui.navigation.Pane import com.google.android.samples.socialite.ui.navigation.SocialiteNavSuite @@ -186,6 +188,9 @@ fun MainNavigation( onInspectClicked = { uri -> backStack.add(Pane.MetadataInspector(uri)) }, + onEnhanceClicked = { messageId, uri -> + backStack.add(Pane.ImageEnhancement(backStackKey.chatId, messageId, uri)) + }, ) } @@ -247,8 +252,22 @@ fun MainNavigation( ) } + is Pane.ImageEnhancement -> NavEntry(backStackKey) { + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) { + ImageEnhancementScreen( + chatId = backStackKey.chatId, + messageId = backStackKey.messageId, + uri = backStackKey.uri, + onCloseButtonClicked = { backStack.removeLastOrNull() }, + onFinishEditing = { backStack.removeLastOrNull() }, + ) + } + } + is Pane.MetadataInspector -> NavEntry(backStackKey) { - MetadataInspector(backStackKey.uri) + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) { + MetadataInspector(backStackKey.uri) + } } else -> NavEntry(backStackKey) { diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt index 558a7cc0..ed61a2dc 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatMessage.kt @@ -19,6 +19,7 @@ package com.google.android.samples.socialite.ui.chat import android.net.Uri data class ChatMessage( + val id: Long, val text: String, val mediaUri: String?, val mediaMimeType: String?, diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt index a33250bb..70dcb203 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatScreen.kt @@ -59,6 +59,7 @@ import androidx.compose.runtime.LaunchedEffect import androidx.compose.runtime.getValue import androidx.compose.runtime.remember import androidx.compose.runtime.rememberCoroutineScope +import androidx.compose.runtime.setValue import androidx.compose.ui.Alignment import androidx.compose.ui.Modifier import androidx.compose.ui.draw.clip @@ -67,6 +68,7 @@ import androidx.compose.ui.focus.focusProperties import androidx.compose.ui.focus.focusRequester import androidx.compose.ui.graphics.Color import androidx.compose.ui.input.nestedscroll.nestedScroll +import androidx.compose.ui.platform.LocalContext import androidx.compose.ui.platform.LocalLayoutDirection import androidx.compose.ui.res.stringResource import androidx.compose.ui.tooling.preview.Preview @@ -88,6 +90,7 @@ import com.google.android.samples.socialite.ui.chat.component.MessageBubble import com.google.android.samples.socialite.ui.chat.component.mediaItemDropTarget import com.google.android.samples.socialite.ui.chat.component.scrollWithKeyboards import com.google.android.samples.socialite.ui.components.tryRequestFocus +import com.google.android.samples.socialite.ui.mediaenhancement.EnhancementSupportManager import com.google.android.samples.socialite.ui.rememberIconPainter @Composable @@ -102,8 +105,12 @@ fun ChatScreen( prefilledText: String? = null, prefilledImageUri: String? = null, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, viewModel: ChatViewModel = hiltViewModel(), ) { + val context = LocalContext.current + var isEnhancementSupported by remember { androidx.compose.runtime.mutableStateOf(false) } + LaunchedEffect(chatId) { viewModel.setChatId(chatId) if (prefilledText != null) { @@ -112,6 +119,7 @@ fun ChatScreen( if (prefilledImageUri != null) { viewModel.prefillInputImage(prefilledImageUri) } + isEnhancementSupported = EnhancementSupportManager.checkSupport(context) } val chat by viewModel.chat.collectAsStateWithLifecycle() val messages by viewModel.messages.collectAsStateWithLifecycle() @@ -125,6 +133,7 @@ fun ChatScreen( textFieldState = textFieldState, attachedMedia = attachedMedia, sendEnabled = sendEnabled, + isEnhancementSupported = isEnhancementSupported, onBackPressed = onBackPressed, onSendClick = { viewModel.send() }, onCameraClick = onCameraClick, @@ -133,6 +142,7 @@ fun ChatScreen( onMediaItemAttached = viewModel::attachMedia, onRemoveAttachedMediaItem = viewModel::removeAttachedMedia, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, modifier = modifier .clip(RoundedCornerShape(5)), ) @@ -174,6 +184,7 @@ private fun ChatContent( textFieldState: TextFieldState, attachedMedia: MediaItem?, sendEnabled: Boolean, + isEnhancementSupported: Boolean = false, onBackPressed: (() -> Unit)?, onSendClick: () -> Unit, onCameraClick: () -> Unit, @@ -182,6 +193,7 @@ private fun ChatContent( onMediaItemAttached: (MediaItem) -> Unit, onRemoveAttachedMediaItem: () -> Unit, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, modifier: Modifier = Modifier, ) { val topAppBarState = rememberTopAppBarState() @@ -221,8 +233,10 @@ private fun ChatContent( modifier = Modifier .fillMaxWidth() .weight(1f), + isEnhancementSupported = isEnhancementSupported, onVideoClick = onVideoClick, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, ) InputBar( textFieldState = textFieldState, @@ -309,8 +323,10 @@ private fun MessageList( contentPadding: PaddingValues, modifier: Modifier = Modifier, state: LazyListState = rememberLazyListState(), + isEnhancementSupported: Boolean = false, onVideoClick: (uri: String) -> Unit = {}, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, ) { LazyColumn( modifier = modifier, @@ -338,10 +354,12 @@ private fun MessageList( } MessageBubble( message = message, + isEnhancementSupported = isEnhancementSupported, onVideoClick = { message.mediaUri?.let { onVideoClick(it) } }, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, ) } } @@ -355,11 +373,11 @@ private fun PreviewChatContent() { ChatContent( chat = ChatDetail(ChatWithLastMessage(0L), listOf(Contact.CONTACTS[0])), messages = listOf( - ChatMessage("Hi!", null, null, 0L, false, null), - ChatMessage("Hello", null, null, 0L, true, null), - ChatMessage("world", null, null, 0L, true, null), - ChatMessage("!", null, null, 0L, true, null), - ChatMessage("Hello, world!", null, null, 0L, true, null), + ChatMessage(1, "Hi!", null, null, 0L, false, null), + ChatMessage(2, "Hello", null, null, 0L, true, null), + ChatMessage(3, "world", null, null, 0L, true, null), + ChatMessage(4, "!", null, null, 0L, true, null), + ChatMessage(5, "Hello, world!", null, null, 0L, true, null), ), textFieldState = TextFieldState("Hello"), attachedMedia = null, diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt index e124b43b..ee165531 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/ChatViewModel.kt @@ -74,6 +74,7 @@ class ChatViewModel @Inject constructor( } ChatMessage( + id = message.id, text = message.text, mediaUri = message.mediaUri, mediaMimeType = message.mediaMimeType, diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt b/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt index 61e46ea7..37eb6706 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/chat/component/MessageBubble.kt @@ -17,16 +17,20 @@ package com.google.android.samples.socialite.ui.chat.component import android.util.Log +import androidx.compose.foundation.background import androidx.compose.foundation.clickable import androidx.compose.foundation.focusable import androidx.compose.foundation.indication import androidx.compose.foundation.interaction.MutableInteractionSource +import androidx.compose.foundation.layout.Box import androidx.compose.foundation.layout.Column import androidx.compose.foundation.layout.height import androidx.compose.foundation.layout.padding import androidx.compose.foundation.layout.size +import androidx.compose.foundation.shape.CircleShape import androidx.compose.foundation.text.selection.SelectionContainer import androidx.compose.material.icons.Icons +import androidx.compose.material.icons.filled.AutoAwesome import androidx.compose.material.icons.filled.Info import androidx.compose.material3.Icon import androidx.compose.material3.IconButton @@ -38,6 +42,7 @@ import androidx.compose.runtime.Composable import androidx.compose.runtime.remember import androidx.compose.ui.Alignment import androidx.compose.ui.Modifier +import androidx.compose.ui.graphics.Color import androidx.compose.ui.platform.LocalContext import androidx.compose.ui.unit.dp import androidx.compose.ui.zIndex @@ -53,8 +58,10 @@ private const val TAG = "ChatUI" internal fun MessageBubble( message: ChatMessage, modifier: Modifier = Modifier, + isEnhancementSupported: Boolean = false, onVideoClick: () -> Unit = {}, onInspectClicked: (uri: String) -> Unit = {}, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit = { _, _ -> }, ) { MessageBubbleSurface( isVideoContentAttached = message.isVideoContentAttached, @@ -71,7 +78,9 @@ internal fun MessageBubble( AttachedMedia( message = message, modifier = Modifier.draggableMediaItem(message), + isEnhancementSupported = isEnhancementSupported, onInspectClicked = onInspectClicked, + onEnhanceClicked = onEnhanceClicked, ) } } @@ -116,17 +125,40 @@ private fun MessageBubbleSurface( private fun AttachedMedia( message: ChatMessage, modifier: Modifier = Modifier, + isEnhancementSupported: Boolean = false, onInspectClicked: (uri: String) -> Unit, + onEnhanceClicked: (messageId: Long, uri: String) -> Unit, ) { val uri = message.mediaUri if (uri != null) { ContextMenuArea(chatMessage = message) { when { message.isImageContentAttached -> { - Photo( - uri = uri, - modifier = modifier, - ) + Box { + Photo( + uri = uri, + modifier = modifier, + ) + if (isEnhancementSupported) { + IconButton( + onClick = { onEnhanceClicked(message.id, uri) }, + modifier = Modifier + .align(Alignment.TopEnd) + .padding(8.dp) + .background( + color = MaterialTheme.colorScheme.scrim.copy(alpha = 0.4f), + shape = CircleShape, + ) + .zIndex(1f), + ) { + Icon( + imageVector = Icons.Filled.AutoAwesome, + contentDescription = "AI Enhance", + tint = Color.White, + ) + } + } + } } message.isVideoContentAttached -> { diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt new file mode 100644 index 00000000..88b9c4f8 --- /dev/null +++ b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementUtils.kt @@ -0,0 +1,187 @@ +package com.google.android.samples.socialite.ui.mediaenhancement + +import android.content.Context +import android.graphics.Bitmap +import android.os.Build +import android.util.Log +import androidx.annotation.RequiresApi +import com.google.android.gms.common.api.Status +import com.google.android.gms.media.effect.enhancement.Enhancement +import com.google.android.gms.media.effect.enhancement.EnhancementCallback +import com.google.android.gms.media.effect.enhancement.EnhancementClient +import com.google.android.gms.media.effect.enhancement.EnhancementOptions +import com.google.android.gms.media.effect.enhancement.EnhancementSession +import com.google.android.gms.media.effect.enhancement.EnhancementSessionCallback +import java.util.concurrent.Executor +import kotlin.coroutines.resume +import kotlin.coroutines.resumeWithException +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.suspendCancellableCoroutine +import kotlinx.coroutines.withContext + +class EnhancementFailedException( message: String) : Exception(message) + +/** + * A modern coroutine wrapper for the enhancement process. + * + * This suspend function encapsulates the entire callback-based process of creating a session, + * processing a bitmap, and handling success or failure. + * @return The enhanced [Bitmap] on success. + * @throws [EnhancementFailedException] if any step of the process fails. + */ +/** + * Extension to create an enhancement session asynchronously. + */ +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.createSessionAsync( + options: EnhancementOptions, + executor: Executor, +): EnhancementSession = withContext(Dispatchers.Main) { + suspendCancellableCoroutine { continuation -> + val callback = object : EnhancementSessionCallback { + override fun onSessionCreated(session: EnhancementSession) { + continuation.resume(session) + } + + override fun onSessionCreationFailed(status: Status) { + continuation.resumeWithException( + Exception("Session creation failed: ${status.statusMessage} (${status.statusCode})"), + ) + } + + override fun onSessionDestroyed() { + // Log or handle if needed + } + + override fun onSessionDisconnected(status: Status) { + // Log or handle if needed + } + } + + this@createSessionAsync.createSession(options, callback) + .addOnFailureListener(executor) { e -> + if (continuation.isActive) { + continuation.resumeWithException(e) + } + } + } +} + +/** + * Extension to process a bitmap using an existing session. + */ +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementSession.processBitmapAsync( + bitmap: Bitmap, + options: EnhancementOptions, +): Bitmap = suspendCancellableCoroutine { continuation -> + val callback = object : EnhancementCallback { + override fun onBitmapProcessed(bitmap: Bitmap) { + continuation.resume(bitmap) + } + + override fun onError(statusCode: Int) { + continuation.resumeWithException( + Exception("Processing failed with status code: $statusCode"), + ) + } + + override fun onSurfaceProcessed(timestamp: Long) { + /* Not used in bitmap flow */ + } + } + + this.process(bitmap, options, callback) +} + +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.installModuleAsync(onProgress: (Int) -> Unit): Boolean = + suspendCancellableCoroutine { continuation -> + val callback = object : EnhancementClient.InstallStatusCallback { + override fun onDownloadPending() { + Log.d("EnhancementUtils", "onDownloadPending") + } + + override fun onDownloadStart() { + Log.d("EnhancementUtils", "onDownloadStart") + } + + override fun onDownloadPaused() { + Log.d("EnhancementUtils", "onDownloadPaused") + } + + override fun onDownloadProgressUpdate(progress: Int) { + Log.d("EnhancementUtils", "onDownloadProgressUpdate: $progress") + onProgress(progress) + } + + override fun onDownloadComplete() { + Log.d("EnhancementUtils", "onDownloadComplete") + } + + override fun onInstalled() { + Log.d("EnhancementUtils", "onInstalled") + if (continuation.isActive) continuation.resume(true) + } + + override fun onCancelled() { + Log.d("EnhancementUtils", "onCancelled") + if (continuation.isActive) continuation.resume(false) + } + + override fun onError(description: String) { + Log.e("EnhancementUtils", "onError: $description") + if (continuation.isActive) continuation.resumeWithException(Exception(description)) + } + } + + this.installModule(callback) + .addOnSuccessListener { result -> + if (result && continuation.isActive) { + continuation.resume(true) + } + } + .addOnFailureListener { e -> + if (continuation.isActive) continuation.resumeWithException(e) + } + } + +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.isModuleInstalledAsync(): Boolean = + suspendCancellableCoroutine { continuation -> + this.isModuleInstalled() + .addOnSuccessListener { result -> continuation.resume(result) } + .addOnFailureListener { e -> continuation.resumeWithException(e) } + } + +@RequiresApi(Build.VERSION_CODES.R) +suspend fun EnhancementClient.isDeviceSupportedAsync(): Boolean = + suspendCancellableCoroutine { continuation -> + this.isDeviceSupported() + .addOnSuccessListener { result -> continuation.resume(result) } + .addOnFailureListener { e -> continuation.resumeWithException(e) } + } + +object EnhancementSupportManager { + private var isSupported: Boolean? = null + + suspend fun checkSupport(context: Context): Boolean { + if (isSupported != null) return isSupported!! + + if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) { + try { + val client = Enhancement.getClient(context.applicationContext) + + isSupported = client.isDeviceSupportedAsync() + return isSupported!! + } catch (e: Exception) { + Log.e("EnhancementSupport", "Error checking support", e) + isSupported = false + return false + } + } else { + isSupported = false + return false + } + } +} diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt new file mode 100644 index 00000000..3e2f04c2 --- /dev/null +++ b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/EnhancementViewModel.kt @@ -0,0 +1,336 @@ +package com.google.android.samples.socialite.ui.mediaenhancement + +import android.app.Application +import android.content.Context +import android.graphics.Bitmap +import android.graphics.BitmapFactory +import android.net.Uri +import android.os.Build +import android.util.Log +import androidx.annotation.RequiresApi +import androidx.core.net.toUri +import androidx.lifecycle.AndroidViewModel +import androidx.lifecycle.viewModelScope +import com.google.android.gms.media.effect.enhancement.Enhancement +import com.google.android.gms.media.effect.enhancement.EnhancementClient +import com.google.android.gms.media.effect.enhancement.EnhancementMode +import com.google.android.gms.media.effect.enhancement.EnhancementOptions +import com.google.android.gms.media.effect.enhancement.EnhancementSession +import com.google.android.samples.socialite.repository.ChatRepository +import dagger.hilt.android.lifecycle.HiltViewModel +import java.io.File +import java.io.FileOutputStream +import java.util.concurrent.Executors +import javax.inject.Inject +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.flow.MutableStateFlow +import kotlinx.coroutines.flow.StateFlow +import kotlinx.coroutines.flow.asStateFlow +import kotlinx.coroutines.flow.update +import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext + +private const val TAG = "EnhancementViewModel" + +// Data class to hold all information about an image +data class ImageInfo( + val bitmap: Bitmap? = null, + val latency: Long? = null, + val qualityScore: Int? = null, +) + +// Defines the state of the UI +data class UiState( + val originalImage: ImageInfo? = null, + val enhancedImage: ImageInfo? = null, + val isLoading: Boolean = false, + val selectedOptions: Set = setOf("Tonemap"), + val enhancementMode: Int = EnhancementMode.BITMAP, + val isModuleInstalling: Boolean = false, + val moduleInstallProgress: Int = 0, + val moduleInstallError: String? = null, + val isModuleReady: Boolean = false, + val moduleStatus: String = "Unknown", + val isDeviceSupported: Boolean = true, + val enhancementError: String? = null, +) + +@HiltViewModel +@RequiresApi(Build.VERSION_CODES.R) +class EnhancementViewModel @Inject constructor( + application: Application, + private val chatRepository: ChatRepository, +) : AndroidViewModel(application) { + + private val _uiState = MutableStateFlow(UiState()) + val uiState: StateFlow = _uiState.asStateFlow() + + @RequiresApi(Build.VERSION_CODES.R) + private val enhancementClient: EnhancementClient = Enhancement.getClient(application) + private val enhancementExecutor = Executors.newSingleThreadExecutor() + private var enhancementSession: EnhancementSession? = null + + init { + checkAndInstallModule() + } + + @RequiresApi(Build.VERSION_CODES.R) + private fun checkAndInstallModule() { + viewModelScope.launch { + _uiState.update { it.copy(moduleStatus = "Checking for device support...") } + try { + if (!enhancementClient.isDeviceSupportedAsync()) { + _uiState.update { + it.copy( + isDeviceSupported = false, + moduleStatus = "Device not supported", + ) + } + return@launch + } + + _uiState.update { it.copy(isDeviceSupported = true) } + + if (!enhancementClient.isModuleInstalledAsync()) { + _uiState.update { + it.copy( + isModuleInstalling = true, + moduleInstallError = null, + moduleStatus = "Not Installed", + ) + } + + _uiState.update { it.copy(moduleStatus = "Installing...") } + val installed = enhancementClient.installModuleAsync { progress -> + _uiState.update { it.copy(moduleInstallProgress = progress) } + } + if (!installed) { + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallError = "Module installation failed or cancelled.", + moduleStatus = "Install Failed", + ) + } + return@launch + } + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallProgress = 100, + isModuleReady = true, + moduleStatus = "Installed", + ) + } + } else { + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallProgress = 100, + isModuleReady = true, + moduleStatus = "Installed", + ) + } + } + } catch (e: Exception) { + _uiState.update { + it.copy( + isModuleInstalling = false, + moduleInstallError = "Failed to check or install module: ${e.message}", + moduleStatus = "Error", + ) + } + } + } + } + + private suspend fun createSession(): EnhancementSession? { + if (!_uiState.value.isDeviceSupported) return null + + val originalBitmap = _uiState.value.originalImage?.bitmap ?: return null + + return try { + val options = getEnhancementOptionsFor(originalBitmap, _uiState.value.selectedOptions) + enhancementClient.createSessionAsync(options, enhancementExecutor) + } catch (e: Exception) { + _uiState.update { it.copy(enhancementError = "Failed to create session: ${e.message}") } + null + } + } + + override fun onCleared() { + enhancementSession?.release() + enhancementSession = null + enhancementExecutor.shutdown() + super.onCleared() + } + + fun onImageSelected(uri: Uri) { + viewModelScope.launch(Dispatchers.IO) { + _uiState.update { + it.copy( + isLoading = true, + originalImage = null, + enhancedImage = null, + enhancementError = null, + ) + } + // Release any previous session, since we have a new image. Confined to Main thread. + withContext(Dispatchers.Main) { + enhancementSession?.release() + enhancementSession = null + } + + val originalBitmap = decodeBitmapFromUri(uri, getApplication()) + + if (originalBitmap == null) { + Log.e(TAG, "Failed to decode bitmap from URI.") + _uiState.update { + it.copy( + isLoading = false, + enhancementError = "Failed to load image.", + ) + } + return@launch + } + + val originalImageInfo = ImageInfo(bitmap = originalBitmap) + // Show the original image and stop loading. Session will be created on demand. + _uiState.update { it.copy(originalImage = originalImageInfo, isLoading = false) } + } + } + + fun onOptionSelected(option: String) { + val currentOptions = _uiState.value.selectedOptions + val newOptions = if (option in currentOptions) { + currentOptions - option + } else { + currentOptions + option + } + _uiState.update { it.copy(selectedOptions = newOptions) } + } + + fun setEnhancementMode(mode: Int) { + _uiState.update { it.copy(enhancementMode = mode) } + } + + fun enhanceImage() { + val originalBitmap = _uiState.value.originalImage?.bitmap ?: return + + viewModelScope.launch(Dispatchers.IO) { + _uiState.update { it.copy(isLoading = true, enhancementError = null) } + try { + processImage(originalBitmap) + } finally { + _uiState.update { it.copy(isLoading = false) } + } + } + } + + private suspend fun processImage(bitmap: Bitmap) { + try { + // If the session doesn't exist, create it now. + if (enhancementSession == null) { + enhancementSession = createSession() + } + + // If session creation failed or is not possible, just return. + // The createSession function will have set the error message. + val session = enhancementSession ?: return + + val options = getEnhancementOptionsFor(bitmap, _uiState.value.selectedOptions) + val enhancementStartTime = System.currentTimeMillis() + val enhancedBitmap = session.processBitmapAsync(bitmap, options) + val enhancementLatency = System.currentTimeMillis() - enhancementStartTime + + _uiState.update { + it.copy( + enhancedImage = ImageInfo(bitmap = enhancedBitmap, latency = enhancementLatency), + ) + } + } catch (e: Exception) { + Log.e(TAG, "Enhancement failed: ${e.message}", e) + _uiState.update { + it.copy( + enhancedImage = null, + enhancementError = "Enhancement failed: ${e.message}", + ) + } + } + } + + private fun getEnhancementOptionsFor( + bitmap: Bitmap, + selectedOptions: Set, + ): EnhancementOptions { + return EnhancementOptions( + bitmap.width, + bitmap.height, + _uiState.value.enhancementMode, + "Tonemap" in selectedOptions, + "Deblur & DeNoise" in selectedOptions, + false, + "Upscale" in selectedOptions, + false, + ) + } + + private fun decodeBitmapFromUri(uri: Uri, context: Context): Bitmap? { + return try { + context.contentResolver.openInputStream(uri)?.use { inputStream -> + BitmapFactory.decodeStream(inputStream) + } + } catch (e: Throwable) { + Log.e(TAG, "Error decoding bitmap from URI", e) + null + } + } + + fun saveEnhancedImage(messageId: Long, oldUri: String, onComplete: () -> Unit) { + val enhancedBitmap = _uiState.value.enhancedImage?.bitmap ?: return + viewModelScope.launch(Dispatchers.IO) { + try { + val context = getApplication() + val directory = File(context.filesDir, "media") + if (!directory.exists()) { + directory.mkdirs() + } + + // Prevent storage consumption by deleting only the previous enhanced file for this message + val oldFileUri = oldUri.toUri() + if (oldFileUri.scheme == "file") { + val oldFile = File(oldFileUri.path ?: "") + if (oldFile.exists() && oldFile.name.startsWith("enhanced_")) { + oldFile.delete() + } + } + + val filename = "enhanced_${System.currentTimeMillis()}.jpg" + val file = File(directory, filename) + FileOutputStream(file).use { out -> + var success = enhancedBitmap.compress(Bitmap.CompressFormat.JPEG, 100, out) + if (!success) { + Log.w(TAG, "Failed to compress hardware bitmap, trying software fallback.") + val swBitmap = enhancedBitmap.copy(Bitmap.Config.ARGB_8888, false) + if (swBitmap != null) { + success = swBitmap.compress(Bitmap.CompressFormat.JPEG, 100, out) + } + } + if (!success) { + Log.e(TAG, "Bitmap compression completely failed!") + } + } + + val newUri = file.toUri().toString() + Log.d(TAG, "Saved enhanced image to $newUri, updating message $messageId") + chatRepository.updateMessageMediaUri(messageId, newUri) + + withContext(Dispatchers.Main) { + onComplete() + } + } catch (e: Exception) { + Log.e(TAG, "Failed to save image", e) + } + } + } +} diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt new file mode 100644 index 00000000..7512b0c1 --- /dev/null +++ b/app/src/main/java/com/google/android/samples/socialite/ui/mediaenhancement/ImageEnhancementScreen.kt @@ -0,0 +1,298 @@ +package com.google.android.samples.socialite.ui.mediaenhancement + +import android.graphics.Bitmap +import android.os.Build +import androidx.annotation.RequiresApi +import androidx.compose.foundation.Image +import androidx.compose.foundation.background +import androidx.compose.foundation.layout.Arrangement +import androidx.compose.foundation.layout.Box +import androidx.compose.foundation.layout.Column +import androidx.compose.foundation.layout.Row +import androidx.compose.foundation.layout.Spacer +import androidx.compose.foundation.layout.fillMaxSize +import androidx.compose.foundation.layout.fillMaxWidth +import androidx.compose.foundation.layout.height +import androidx.compose.foundation.layout.padding +import androidx.compose.foundation.lazy.LazyRow +import androidx.compose.foundation.lazy.items +import androidx.compose.material3.Button +import androidx.compose.material3.CircularProgressIndicator +import androidx.compose.material3.ExperimentalMaterial3Api +import androidx.compose.material3.FilterChip +import androidx.compose.material3.LinearProgressIndicator +import androidx.compose.material3.MaterialTheme +import androidx.compose.material3.OutlinedButton +import androidx.compose.material3.Scaffold +import androidx.compose.material3.Text +import androidx.compose.runtime.Composable +import androidx.compose.runtime.LaunchedEffect +import androidx.compose.runtime.collectAsState +import androidx.compose.runtime.getValue +import androidx.compose.ui.Alignment +import androidx.compose.ui.Modifier +import androidx.compose.ui.draw.clip +import androidx.compose.ui.graphics.Color +import androidx.compose.ui.graphics.asImageBitmap +import androidx.compose.ui.layout.ContentScale +import androidx.compose.ui.text.font.FontWeight +import androidx.compose.ui.unit.dp +import androidx.compose.ui.unit.sp +import androidx.core.net.toUri +import androidx.hilt.navigation.compose.hiltViewModel +import com.google.android.gms.media.effect.enhancement.EnhancementMode +import java.text.DecimalFormat + +@RequiresApi(Build.VERSION_CODES.R) +@OptIn(ExperimentalMaterial3Api::class) +@Composable +fun ImageEnhancementScreen( + messageId: Long, + uri: String, + onCloseButtonClicked: () -> Unit, + onFinishEditing: () -> Unit, + enhancementViewModel: EnhancementViewModel = hiltViewModel(), +) { + val uiState by enhancementViewModel.uiState.collectAsState() + + LaunchedEffect(uri) { + enhancementViewModel.setEnhancementMode(EnhancementMode.BITMAP) + enhancementViewModel.onImageSelected(uri.toUri()) + } + + val processingOptions = listOf("Tonemap", "Deblur & DeNoise", "Upscale") + + Scaffold { paddingValues -> + Column( + modifier = Modifier + .fillMaxSize() + .padding(paddingValues) + .padding(start = 16.dp, end = 16.dp, bottom = 16.dp), + horizontalAlignment = Alignment.CenterHorizontally, + ) { + Text( + text = "Module Status: ${uiState.moduleStatus}", + style = MaterialTheme.typography.bodyMedium, + fontWeight = FontWeight.Bold, + ) + + Spacer(modifier = Modifier.height(12.dp)) + + Row( + modifier = Modifier.fillMaxWidth(), + verticalAlignment = Alignment.CenterVertically, + ) { + LazyRow( + modifier = Modifier.weight(1f), + horizontalArrangement = Arrangement.spacedBy(8.dp), + ) { + items(processingOptions) { option -> + FilterChip( + selected = option in uiState.selectedOptions, + onClick = { enhancementViewModel.onOptionSelected(option) }, + label = { Text(option) }, + ) + } + } + Spacer(modifier = Modifier.padding(start = 8.dp)) + Button( + onClick = { enhancementViewModel.enhanceImage() }, + enabled = uiState.isModuleReady && !uiState.isLoading && uiState.originalImage?.bitmap != null, + ) { + val buttonText = when { + uiState.isModuleInstalling -> "Installing..." + uiState.isLoading -> "Enhancing..." + else -> "AI Enhance" + } + Text(buttonText) + } + } + + if (uiState.isModuleInstalling) { + Spacer(modifier = Modifier.height(8.dp)) + Column(horizontalAlignment = Alignment.CenterHorizontally) { + LinearProgressIndicator( + progress = { uiState.moduleInstallProgress / 100f }, + modifier = Modifier.fillMaxWidth(), + ) + Text( + text = "Downloading Enhancement module (${uiState.moduleInstallProgress}%)", + style = MaterialTheme.typography.bodySmall, + ) + } + } + + if (uiState.moduleInstallError != null) { + Spacer(modifier = Modifier.height(8.dp)) + Text( + text = uiState.moduleInstallError ?: "Unknown error", + color = MaterialTheme.colorScheme.error, + style = MaterialTheme.typography.bodySmall, + ) + } + + if (uiState.enhancementError != null) { + Spacer(modifier = Modifier.height(8.dp)) + Text( + text = uiState.enhancementError ?: "Unknown enhancement error", + color = MaterialTheme.colorScheme.error, + style = MaterialTheme.typography.bodySmall, + ) + } + + Spacer(modifier = Modifier.height(16.dp)) + + Column( + modifier = Modifier.weight(1f), + verticalArrangement = Arrangement.spacedBy(16.dp), + ) { + EnhancementImageCard( + title = "Original", + imageInfo = uiState.originalImage, + modifier = Modifier.weight(1f), + ) + EnhancementImageCard( + title = "Enhanced", + imageInfo = uiState.enhancedImage, + isLoading = uiState.isLoading, + error = uiState.enhancementError, + modifier = Modifier.weight(1f), + ) + } + + Spacer(modifier = Modifier.height(16.dp)) + + Row( + modifier = Modifier.fillMaxWidth(), + horizontalArrangement = Arrangement.SpaceEvenly, + ) { + OutlinedButton( + onClick = onCloseButtonClicked, + modifier = Modifier.weight(1f).padding(end = 8.dp), + ) { + Text("Cancel") + } + Button( + onClick = { + enhancementViewModel.saveEnhancedImage(messageId, uri, onFinishEditing) + }, + enabled = uiState.enhancedImage?.bitmap != null, + modifier = Modifier.weight(1f).padding(start = 8.dp), + ) { + Text("Confirm") + } + } + } + } +} + +@Composable +fun EnhancementImageCard( + title: String, + imageInfo: ImageInfo?, + modifier: Modifier = Modifier, + isLoading: Boolean = false, + error: String? = null, +) { + Box( + modifier = modifier + .fillMaxSize() + .background(MaterialTheme.colorScheme.surfaceVariant, MaterialTheme.shapes.medium) + .clip(MaterialTheme.shapes.medium), + contentAlignment = Alignment.Center, + ) { + when { + isLoading -> { + CircularProgressIndicator() + } + error != null -> { + Text( + text = "Enhancement Failed", + color = MaterialTheme.colorScheme.error, + style = MaterialTheme.typography.bodyMedium, + ) + } + imageInfo?.bitmap != null -> { + Image( + bitmap = imageInfo.bitmap.asImageBitmap(), + contentDescription = title, + modifier = Modifier.fillMaxSize(), + contentScale = ContentScale.Fit, + ) + InfoTag( + title = title, + latency = imageInfo.latency, + bitmap = imageInfo.bitmap, + quality = imageInfo.qualityScore, + modifier = Modifier.align(Alignment.TopStart), + ) + } + else -> { + Text( + text = if (title == "Original") "Loading..." else "Enhanced Image", + style = MaterialTheme.typography.bodyMedium, + color = MaterialTheme.colorScheme.onSurfaceVariant, + ) + } + } + } +} + +private fun formatMemorySize(bytes: Int): String { + val kb = bytes / 1024.0 + val mb = kb / 1024.0 + val decimalFormat = DecimalFormat("#.##") + return when { + mb >= 1 -> "${decimalFormat.format(mb)} MB" + else -> "${decimalFormat.format(kb)} KB" + } +} + +@Composable +private fun InfoTag( + title: String, + latency: Long?, + bitmap: Bitmap?, + quality: Int?, + modifier: Modifier = Modifier, +) { + Column( + modifier = modifier + .padding(4.dp) + .background(Color.Black.copy(alpha = 0.6f), MaterialTheme.shapes.small) + .padding(horizontal = 6.dp, vertical = 2.dp), + verticalArrangement = Arrangement.spacedBy(0.dp), + ) { + Text( + text = title, + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.Bold, + ) + if (bitmap != null) { + val memorySize = formatMemorySize(bitmap.byteCount) + Text( + text = "$memorySize, ${bitmap.width}x${bitmap.height}", + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.SemiBold, + ) + } + if (latency != null && title != "Original") { + Text( + text = "AI Latency: ${latency}ms", + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.SemiBold, + ) + } + if (quality != null) { + Text( + text = "Quality: $quality/100", + color = Color.White, + fontSize = 10.sp, + fontWeight = FontWeight.SemiBold, + ) + } + } +} diff --git a/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt b/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt index 3b98e6cd..b0242dd7 100644 --- a/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt +++ b/app/src/main/java/com/google/android/samples/socialite/ui/navigation/SocialiteNavigation.kt @@ -81,6 +81,10 @@ sealed interface Pane : Parcelable { @Parcelize @Serializable data class MetadataInspector(val uri: String) : Pane + + @Parcelize + @Serializable + data class ImageEnhancement(val chatId: Long, val messageId: Long, val uri: String) : Pane } enum class TopLevelDestination( diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index a8c0fa2c..2b7fc9ec 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -63,6 +63,7 @@ visionCommon = "17.3.0" segmentationSelfie = "16.0.0-beta6" litert = "1.3.0" datastore = "1.1.4" +playServicesMediaEffectEnhancement = "16.0.0-beta04" @@ -148,6 +149,7 @@ firebase-ai = { group = "com.google.firebase", name = "firebase-ai" } firebase-analytics = { group = "com.google.firebase", name = "firebase-analytics" } firebase-messaging = {group = "com.google.firebase", name = "firebase-messaging" } vision-common = { group = "com.google.mlkit", name = "vision-common", version.ref = "visionCommon" } +play-services-media-effect-enhancement = { module = "com.google.android.gms:play-services-media-effect-enhancement", version.ref = "playServicesMediaEffectEnhancement" } [plugins]