[Android/Kotlin] 안드로이드에서 코틀린으로 TFLite 모델 돌릴 때 생길 수 있는 RGB, BGR 오류 / Pytorch와 TFLite 출력이 다른 이유

안드로이드에서 코틀린으로 TFLite 모델 돌릴 때 생길 수 있는 RGB, BGR 오류

부제: Pytorch와 TFLite 출력이 다른 이유, RGB2BGR

 

현재 PyTorch로 학습된 모델을 TFLite로 변환하여 안드로이드에서 동작시키는 코드를 짜고 있다.

Convert를 다 하고, Android Studio에서 Kotlin으로 TFlite 모델을 동작시킬 때 논리적 오류가 발생했다.

 

환경 별 모델 Output 비교

python, kotlin을 사용하여 PyTorch, TFLite 모델을 동작시켰을 때 나오는 output 비교.

 

왼쪽 위는 레이블, 오른쪽 위는 Python으로 PyTorch를 동작시켰을 때 나오는 output, 왼쪽 아래는 Python에서 TFLite를 동작시켰을 때 나오는 output, 오른쪽 아래는 Kotlin에서 TFLite를 동작시켰을 때 나오는 output이다.

 

Kotlin에서 TFLite output을 보면 레이블과 비교했을 때, 형편없는 결과를 냈다고 볼 수 있다.

Python에서 PyTorch output과 Python에서 TFLite ouput을 보면 레이블과 크게 차이가 나지 않을 뿐더러 두 개만 놓고서도 차이가 거의 없는 모습이다. 즉, PyTorch에서 TFLite로 Convert시킬 때의 오류는 아니라는 것이다.

 

이 문제를 해결하기 위해 약 일주일이나 고생을 했다.. 결국 찾은 해답은 바로 Kotlin과 Python은 이미지를 읽을 때 가져오는 컬러 채널 순서가 다르다는 것이다.

 

Python에서는 컬러 이미지를 BGR 순서로 읽어오고, Kotlin에서는 컬러 이미지를 RGB 순서로 읽어온다. PyTorch 모델을 Convert 시킬 때도 유의했던 부분인데, 또 놓쳐서 시간을 낭비했다.

 

그래서 Android Studio에서 아래 코드를 이용하여 input을 RGB 에서 BGR로 바꿔주었다.

 

Kotlin에서 컬러 이미지 RGB을 BGR으로 변환하는 코드. RGB2BGR

// asset 폴더에서 이미지 읽어와서 Bitmap으로 저장
val assetManager = resources.assets
val inputStream = assetManager.open("00004_kf94.png")
val bitmap = BitmapFactory.decodeStream(inputStream)

// rgb matrix
val rgb_mat = Mat(256, 256, CvType.CV_8UC3)
Utils.bitmapToMat(bitmap, rgb_mat) // 읽어온 Bitmap을 Mat으로. 기본은 RGB 형태

// bgr matrix
val bgr_mat = Mat(256, 256, CvType.CV_8UC3) // BGR을 저장할 Mat 선언

// Convert: RGB to BGR
Imgproc.cvtColor(rgb_mat, bgr_mat, Imgproc.COLOR_RGB2BGR)


// mat 을 bitmap으로 (모델의 input으로 넣기 위함)
var input_BitmapBuffer = Bitmap.createBitmap(256, 256, Bitmap.Config.ARGB_8888) //ARGB_8888
Utils.matToBitmap(bgr_mat, input_BitmapBuffer)

 

Result

RGB2BGR 후 모델 Output

 

확실히 성능이 더 좋아진 것을 볼 수 있다. 안타깝게도 아직 Python에서 돌린 결과와 다른 부분이 있긴 하지만 눈에 띄는 노이즈는 많이 사라진 것 같다. 다른 원인을 찾기 위해 추가적인 연구를 진행할 예정이다. 또한, PyTorch 모델 학습 시 RGB 채널로 구성된 영상을 더 넣는 방식의 Data Augmentation도 고려하고 있다.

 

* 한가지 더 주의사항

방법에 따라 이미지 전처리를 위해 Normalize 하는 부분에도 python에서 하던 순서 그래도 사용하지 말고, RGB로 바꿔서 사용해야 할 수도 있다.

    // python 순서 BGR
//    private val parsing_mean: FloatArray = floatArrayOf(0.485f * 255, 0.456f * 255, 0.406f * 255)
//    private val parsing_std: FloatArray = floatArrayOf(0.229f * 255, 0.224f * 255, 0.225f * 255)

    // kotlin 순서 RGB
    private val parsing_mean: FloatArray = floatArrayOf(0.456f * 255, 0.406f * 255, 0.485f * 255)
    private val parsing_std: FloatArray = floatArrayOf(0.224f * 255, 0.225f * 255, 0.229f * 255)