fun getSegmentationMask()

in amazon-chime-sdk/src/main/java/com/amazonaws/services/chime/sdk/meetings/audiovideo/video/backgroundfilter/BackgroundFilterVideoFrameProcessor.kt [146:196]
41 lines of code
9 McCabe index (conditional complexity)

    fun getSegmentationMask(
        scaledInputBitmap: Bitmap
    ): Bitmap? {
        // Load model only when previous frame dimensions does not match new frame dimensions.
        if (scaledInputBitmap.width != cachedWidth || scaledInputBitmap.height != cachedHeight) {
            cachedWidth = scaledInputBitmap.width
            cachedHeight = scaledInputBitmap.height
            // Load model.
            segmentationProcessor.initialize(
                scaledInputBitmap.width,
                scaledInputBitmap.height,
                defaultInputModelShape
            )
            val modelState = segmentationProcessor.modelState
            // Logs the error or info message only the first time model is loaded to avoid exploding logs for every frame.
            if (modelState != ModelState.LOADING && modelState != ModelState.LOADED && modelStateMsg == null) {
                modelStateMsg = "Failed to load model with model state $modelState"
                modelStateMsg?.let {
                    logger.error(TAG, it)
                }
            } else {
                if (modelStateMsg == null) {
                    modelStateMsg = "Model State $modelState"
                    modelStateMsg?.let {
                        logger.info(TAG, it)
                    }
                }
            }
        }

        // Set input for model.
        val inputBuffer: ByteBuffer = segmentationProcessor.getInputBuffer()
        scaledInputBitmap.copyPixelsToBuffer(inputBuffer)

        // Predict the result of model.
        val predictResult = segmentationProcessor.predict()
        // Logs the error message only the first time predict is called to avoid exploding logs for every frame.
        if (predictResult == PredictResult.ERROR && predictMsg == null) {
            predictMsg = "Error segmenting background from foreground"
            predictMsg?.let {
                logger.error(TAG, it)
            }
        }

        // Set output after prediction.
        val outputBuffer: ByteBuffer = segmentationProcessor.getOutputBuffer()
        val outputBitmap =
            Bitmap.createBitmap(scaledInputBitmap.width, scaledInputBitmap.height, bitmapConfig)
        outputBitmap.copyPixelsFromBuffer(outputBuffer)
        return outputBitmap
    }