Add on-device AI image enhancement feature to enhance images

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.

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)
 }

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)
 }

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)
+    }
 }

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,21 @@ fun MainNavigation(
                                 )
                             }
 
+                            is Pane.ImageEnhancement -> NavEntry(backStackKey) {
+                                if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
+                                    ImageEnhancementScreen(
+                                        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) {

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?,

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,

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,

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 -> {