MainActivity.kt

package com.example.dlod

import android.Manifest
import android.content.pm.PackageManager
import android.graphics.*
import android.os.Bundle
import android.util.Log
import androidx.activity.ComponentActivity
import androidx.activity.compose.setContent
import androidx.activity.result.contract.ActivityResultContracts
import androidx.camera.core.*
import androidx.camera.lifecycle.ProcessCameraProvider
import androidx.camera.view.PreviewView
import androidx.compose.foundation.Canvas
import androidx.compose.foundation.Image
import androidx.compose.foundation.clickable
import androidx.compose.foundation.layout.*
import androidx.compose.material.icons.Icons
import androidx.compose.material.icons.filled.PhotoCamera
import androidx.compose.material3.*
import androidx.compose.runtime.*
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.asImageBitmap
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.compose.ui.res.painterResource
import androidx.compose.ui.res.stringResource
import androidx.compose.ui.text.style.TextAlign
import androidx.compose.ui.unit.dp
import androidx.compose.ui.unit.sp
import androidx.compose.ui.viewinterop.AndroidView
import androidx.core.content.ContextCompat
import androidx.lifecycle.lifecycleScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import org.tensorflow.lite.support.image.TensorImage
import org.tensorflow.lite.task.vision.detector.Detection
import org.tensorflow.lite.task.vision.detector.ObjectDetector
import java.util.concurrent.ExecutorService
import java.util.concurrent.Executors
import android.graphics.YuvImage
import android.graphics.ImageFormat
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.geometry.Size
import androidx.compose.ui.graphics.drawscope.Stroke
import androidx.compose.foundation.layout.offset
import androidx.compose.foundation.background
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.ui.platform.LocalDensity
import androidx.compose.ui.layout.onGloballyPositioned

class MainActivity : ComponentActivity() {
    companion object {
        const val TAG = "TFLite - ODT"
        private const val MAX_FONT_SIZE = 96F
    }

    private lateinit var cameraExecutor: ExecutorService
    private var isRealTimeMode by mutableStateOf(false)
    private var selectedBitmap by mutableStateOf<Bitmap?>(null)
    private var detectionResults by mutableStateOf<List<DetectionResult>>(emptyList())

    // Camera permission request
    private val requestPermissionLauncher = registerForActivityResult(
        ActivityResultContracts.RequestPermission()
    ) { isGranted: Boolean ->
        if (isGranted) {
            isRealTimeMode = true
        } else {
            Log.e(TAG, "Camera permission denied")
        }
    }

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        cameraExecutor = Executors.newSingleThreadExecutor()

        setContent {
            MaterialTheme {
                Surface(
                    modifier = Modifier.fillMaxSize(),
                    color = MaterialTheme.colorScheme.background
                ) {
                    if (isRealTimeMode) {
                        CameraScreen(
                            onBackClick = {
                                isRealTimeMode = false
                                detectionResults = emptyList() // Clear detection results
                            },
                            detectionResults = detectionResults,
                            onDetectionResults = { results -> detectionResults = results }
                        )
                    } else {
                        MainScreen(
                            selectedBitmap = selectedBitmap,
                            onCaptureImageClick = {
                                checkCameraPermissionAndStart()
                            },
                            onSampleImageClick = { drawableRes ->
                                val bitmap = getSampleImage(drawableRes)
                                setViewAndDetect(bitmap)
                            }
                        )
                    }
                }
            }
        }
    }

    private fun checkCameraPermissionAndStart() {
        when {
            ContextCompat.checkSelfPermission(
                this,
                Manifest.permission.CAMERA
            ) == PackageManager.PERMISSION_GRANTED -> {
                isRealTimeMode = true
            }
            else -> {
                requestPermissionLauncher.launch(Manifest.permission.CAMERA)
            }
        }
    }

    override fun onDestroy() {
        super.onDestroy()
        cameraExecutor.shutdown()
    }

    private fun runObjectDetection(bitmap: Bitmap): List<DetectionResult> {
        return try {
            val image = TensorImage.fromBitmap(bitmap)
            val options = ObjectDetector.ObjectDetectorOptions.builder()
                .setMaxResults(5)
                .setScoreThreshold(0.5f)
                .build()
            val detector = ObjectDetector.createFromFileAndOptions(
                this,
                "1.tflite",
                options
            )
            val results = detector.detect(image)
            debugPrint(results)

            results.map {
                val category = it.categories.first()
                val text = "${category.label}, ${category.score.times(100).toInt()}%"
                DetectionResult(it.boundingBox, text)
            }
        } catch (e: Exception) {
            Log.e(TAG, "Object detection failed", e)
            emptyList()
        }
    }

    private fun debugPrint(results: List<Detection>) {
        for ((i, obj) in results.withIndex()) {
            val box = obj.boundingBox
            Log.d(TAG, "Detected object: $i")
            Log.d(TAG, "  boundingBox: (${box.left}, ${box.top}) - (${box.right},${box.bottom})")
            for ((j, category) in obj.categories.withIndex()) {
                Log.d(TAG, "    Label $j: ${category.label}")
                val confidence: Int = category.score.times(100).toInt()
                Log.d(TAG, "    Confidence: ${confidence}%")
            }
        }
    }

    private fun setViewAndDetect(bitmap: Bitmap) {
        selectedBitmap = bitmap
        lifecycleScope.launch(Dispatchers.Default) {
            val results = runObjectDetection(bitmap)
            val imgWithResult = drawDetectionResult(bitmap, results)
            runOnUiThread {
                selectedBitmap = imgWithResult
            }
        }
    }

    private fun getSampleImage(drawable: Int): Bitmap {
        return BitmapFactory.decodeResource(resources, drawable, BitmapFactory.Options().apply {
            inMutable = true
        })
    }

    private fun drawDetectionResult(
        bitmap: Bitmap,
        detectionResults: List<DetectionResult>
    ): Bitmap {
        val outputBitmap = bitmap.copy(Bitmap.Config.ARGB_8888, true)
        val canvas = Canvas(outputBitmap)
        val pen = Paint()
        pen.textAlign = Paint.Align.LEFT

        detectionResults.forEach {
            pen.color = Color.RED
            pen.strokeWidth = 8F
            pen.style = Paint.Style.STROKE
            val box = it.boundingBox
            canvas.drawRect(box, pen)

            val tagSize = Rect(0, 0, 0, 0)
            pen.style = Paint.Style.FILL_AND_STROKE
            pen.color = Color.YELLOW
            pen.strokeWidth = 2F
            pen.textSize = MAX_FONT_SIZE
            pen.getTextBounds(it.text, 0, it.text.length, tagSize)
            val fontSize: Float = pen.textSize * box.width() / tagSize.width()

            if (fontSize < pen.textSize) pen.textSize = fontSize

            var margin = (box.width() - tagSize.width()) / 2.0F
            if (margin < 0F) margin = 0F
            canvas.drawText(
                it.text, box.left + margin,
                box.top + tagSize.height().times(1F), pen
            )
        }
        return outputBitmap
    }

    @Composable
    fun CameraScreen(
        onBackClick: () -> Unit,
        detectionResults: List<DetectionResult>,
        onDetectionResults: (List<DetectionResult>) -> Unit
    ) {
        val context = LocalContext.current
        val lifecycleOwner = LocalLifecycleOwner.current
        val density = LocalDensity.current.density
        val previewView = remember { PreviewView(context) }
        var previewSize by remember { mutableStateOf(androidx.compose.ui.geometry.Size.Zero) }

        DisposableEffect(Unit) {
            val cameraProviderFuture = ProcessCameraProvider.getInstance(context)
            cameraProviderFuture.addListener({
                val cameraProvider = cameraProviderFuture.get()

                val preview = Preview.Builder().build().also {
                    it.setSurfaceProvider(previewView.surfaceProvider)
                }

                val imageAnalyzer = ImageAnalysis.Builder()
                    .setBackpressureStrategy(ImageAnalysis.STRATEGY_KEEP_ONLY_LATEST)
                    .build()
                    .also {
                        it.setAnalyzer(cameraExecutor, ObjectDetectionAnalyzer(previewSize) { results ->
                            onDetectionResults(results)
                        })
                    }

                val cameraSelector = CameraSelector.DEFAULT_BACK_CAMERA

                try {
                    cameraProvider.unbindAll()
                    cameraProvider.bindToLifecycle(
                        lifecycleOwner,
                        cameraSelector,
                        preview,
                        imageAnalyzer
                    )
                } catch (exc: Exception) {
                    Log.e(TAG, "Use case binding failed", exc)
                }
            }, ContextCompat.getMainExecutor(context))

            onDispose {
                cameraProviderFuture.get().unbindAll()
            }
        }

        Box(modifier = Modifier.fillMaxSize()) {
            AndroidView(
                factory = { previewView },
                modifier = Modifier
                    .fillMaxSize()
                    .onGloballyPositioned { coordinates ->
                        previewSize = androidx.compose.ui.geometry.Size(
                            coordinates.size.width.toFloat(),
                            coordinates.size.height.toFloat()
                        )
                    }
            )

            // Single Canvas for all detection results
            Canvas(modifier = Modifier.fillMaxSize()) {
                detectionResults.forEach { result ->
                    val box = result.boundingBox

                    // Draw bounding box
                    drawRect(
                        color = androidx.compose.ui.graphics.Color.Red,
                        topLeft = Offset(box.left, box.top),
                        size = Size(box.width(), box.height()),
                        style = Stroke(width = 8f)
                    )
                }
            }

            // Text overlays
            detectionResults.forEach { result ->
                val box = result.boundingBox
                if (box.top > 30f) {
                    Text(
                        text = result.text,
                        color = androidx.compose.ui.graphics.Color.Yellow,
                        fontSize = 12.sp,
                        modifier = Modifier
                            .offset(
                                x = (box.left / density).dp,
                                y = ((box.top - 30f) / density).dp
                            )
                            .background(
                                androidx.compose.ui.graphics.Color.Black.copy(alpha = 0.7f),
                                shape = RoundedCornerShape(4.dp)
                            )
                            .padding(4.dp)
                    )
                }
            }

            Button(
                onClick = {
                    onDetectionResults(emptyList())
                    onBackClick()
                },
                modifier = Modifier
                    .align(Alignment.TopStart)
                    .padding(16.dp)
            ) {
                Text("Back")
            }
        }
    }

    inner class ObjectDetectionAnalyzer(
        private val previewSize: androidx.compose.ui.geometry.Size,
        private val onResults: (List<DetectionResult>) -> Unit
    ) : ImageAnalysis.Analyzer {

        private var lastAnalyzedTimestamp = 0L
        private val analysisIntervalMs = 500L

        override fun analyze(imageProxy: ImageProxy) {
            val currentTimestamp = System.currentTimeMillis()
            if (currentTimestamp - lastAnalyzedTimestamp >= analysisIntervalMs) {
                val bitmap = imageProxyToBitmap(imageProxy)
                if (bitmap != null && previewSize.width > 0 && previewSize.height > 0) {
                    val results = runObjectDetection(bitmap)
                    val transformedResults = transformCoordinates(results, bitmap)
                    onResults(transformedResults)
                    lastAnalyzedTimestamp = currentTimestamp
                }
            }
            imageProxy.close()
        }

        private fun transformCoordinates(
            results: List<DetectionResult>,
            rotatedBitmap: Bitmap // Bitmap is already rotated to match display orientation
        ): List<DetectionResult> {
            if (previewSize.width == 0f || previewSize.height == 0f || rotatedBitmap.width == 0 || rotatedBitmap.height == 0) {
                Log.w(TAG, "Cannot transform coordinates: Invalid dimensions. Preview: $previewSize, Bitmap: ${rotatedBitmap.width}x${rotatedBitmap.height}")
                return results // Not enough info to transform
            }

            val modelInputWidth = rotatedBitmap.width.toFloat()
            val modelInputHeight = rotatedBitmap.height.toFloat()

            val viewWidth = previewSize.width
            val viewHeight = previewSize.height

            // Calculate scale and offset for PreviewView.ScaleType.FILL_CENTER
            // The image is scaled uniformly to fill the view, and centered.
            // Excess parts of the image are cropped.
            val scale: Float
            val offsetX: Float
            val offsetY: Float

            val modelAspectRatio = modelInputWidth / modelInputHeight
            val viewAspectRatio = viewWidth / viewHeight

            if (modelAspectRatio > viewAspectRatio) {
                // Model input is wider than the view (relative to aspect ratio).
                // Scale to fit view height, width will be larger and cropped.
                scale = viewHeight / modelInputHeight
                offsetX = (viewWidth - modelInputWidth * scale) / 2f
                offsetY = 0f
            } else {
                // Model input is taller than or same aspect as the view.
                // Scale to fit view width, height will be larger and cropped.
                scale = viewWidth / modelInputWidth
                offsetX = 0f
                offsetY = (viewHeight - modelInputHeight * scale) / 2f
            }
            
            // Log.d(TAG, "Transform: model(${modelInputWidth}x${modelInputHeight}), view(${viewWidth}x${viewHeight}), scale=$scale, offsetX=$offsetX, offsetY=$offsetY")

            return results.map { result ->
                val originalBox = result.boundingBox
                
                val transformedLeft = originalBox.left * scale + offsetX
                val transformedTop = originalBox.top * scale + offsetY
                val transformedRight = originalBox.right * scale + offsetX
                val transformedBottom = originalBox.bottom * scale + offsetY

                val transformedBox = RectF(
                    transformedLeft,
                    transformedTop,
                    transformedRight,
                    transformedBottom
                )
                // Log.d(TAG, "OriginalBox: $originalBox -> TransformedBox: $transformedBox")
                DetectionResult(transformedBox, result.text)
            }
        }

        private fun imageProxyToBitmap(imageProxy: ImageProxy): Bitmap? {
            return try {
                val yBuffer = imageProxy.planes[0].buffer // Y
                val vuBuffer = imageProxy.planes[2].buffer // VU

                val ySize = yBuffer.remaining()
                val vuSize = vuBuffer.remaining()

                val nv21 = ByteArray(ySize + vuSize)

                yBuffer.get(nv21, 0, ySize)
                vuBuffer.get(nv21, ySize, vuSize)

                val yuvImage = YuvImage(nv21, ImageFormat.NV21, imageProxy.width, imageProxy.height, null)
                val out = java.io.ByteArrayOutputStream()
                // Consider increasing JPEG quality if detection accuracy is impacted, e.g., 90
                yuvImage.compressToJpeg(Rect(0, 0, imageProxy.width, imageProxy.height), 80, out)
                val imageBytes = out.toByteArray()
                var bitmap = BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.size)

                // Apply rotation to the bitmap to match the display orientation
                val rotationDegrees = imageProxy.imageInfo.rotationDegrees
                if (rotationDegrees != 0) {
                    val matrix = Matrix()
                    matrix.postRotate(rotationDegrees.toFloat())
                    val rotatedBitmap = Bitmap.createBitmap(bitmap, 0, 0, bitmap.width, bitmap.height, matrix, true)
                    if (rotatedBitmap != bitmap) { // Avoid recycling if createBitmap returned the same instance
                        bitmap.recycle()
                    }
                    bitmap = rotatedBitmap
                }
                bitmap
            } catch (e: Exception) {
                Log.e(TAG, "Error converting imageProxy to bitmap: ${e.message}", e)
                null
            }
        }
    }
}

@Composable
fun MainScreen(
    selectedBitmap: Bitmap?,
    onCaptureImageClick: () -> Unit,
    onSampleImageClick: (Int) -> Unit
) {
    Column(
        modifier = Modifier
            .fillMaxSize()
            .padding(16.dp)
    ) {
        Box(
            modifier = Modifier
                .fillMaxWidth()
                .weight(1f),
            contentAlignment = Alignment.Center
        ) {
            if (selectedBitmap != null) {
                Image(
                    bitmap = selectedBitmap.asImageBitmap(),
                    contentDescription = null,
                    modifier = Modifier.fillMaxSize(),
                    contentScale = androidx.compose.ui.layout.ContentScale.Fit
                )
            } else {
                Text(
                    text = stringResource(R.string.tv_placeholder),
                    fontSize = 40.sp,
                    textAlign = TextAlign.Center,
                    modifier = Modifier.fillMaxWidth()
                )
            }
        }

        Text(
            text = stringResource(R.string.tv_description),
            modifier = Modifier
                .fillMaxWidth()
                .padding(top = 10.dp, bottom = 10.dp)
        )

        Row(
            modifier = Modifier
                .fillMaxWidth()
                .padding(bottom = 16.dp),
            horizontalArrangement = Arrangement.Center
        ) {
            Image(
                painter = painterResource(R.drawable.img_meal_one),
                contentDescription = null,
                modifier = Modifier
                    .size(80.dp)
                    .clickable { onSampleImageClick(R.drawable.img_meal_one) },
                contentScale = androidx.compose.ui.layout.ContentScale.Crop
            )

            Spacer(modifier = Modifier.width(16.dp))

            Image(
                painter = painterResource(R.drawable.img_meal_two),
                contentDescription = null,
                modifier = Modifier
                    .size(80.dp)
                    .clickable { onSampleImageClick(R.drawable.img_meal_two) },
                contentScale = androidx.compose.ui.layout.ContentScale.Crop
            )

            Spacer(modifier = Modifier.width(16.dp))

            Image(
                painter = painterResource(R.drawable.img_meal_three),
                contentDescription = null,
                modifier = Modifier
                    .size(80.dp)
                    .clickable { onSampleImageClick(R.drawable.img_meal_three) },
                contentScale = androidx.compose.ui.layout.ContentScale.Crop
            )
        }

        Button(
            onClick = onCaptureImageClick,
            modifier = Modifier.align(Alignment.CenterHorizontally)
        ) {
            Icon(
                imageVector = Icons.Filled.PhotoCamera,
                contentDescription = null,
                modifier = Modifier.padding(end = 8.dp)
            )
            Text(
                text = "Start Real-time Detection",
                style = MaterialTheme.typography.labelLarge
            )
        }
    }
}

data class DetectionResult(val boundingBox: RectF, val text: String)