ONNX 포맷의 이미지 분류 모델을 사용할 때 입력 이미지의 형식을 변환해야 하는 경우가 있습니다. 공식 문서에서는 Kotlin 예제가 제공되지만, Java를 사용하는 Android 개발자에게는 어려움이 있을 수 있습니다. 이 글에서는 Java로 이미지를 전처리하는 방법과 ONNX 모델을 실행하는 전체 과정을 소개합니다.

1. 이미지 전처리 클래스

아래 ImagePreProcessor 클래스는 Bitmap 이미지를 ONNX 모델이 요구하는 FloatBuffer 형식으로 변환합니다. 크기는 224x224, 채널은 RGB로 가정하며, ImageNet 표준 정규화 값을 사용합니다.

import android.graphics.Bitmap;
import java.nio.FloatBuffer;

public class ImagePreProcessor {
    private static final int BATCH = 1;
    private static final int CHANNELS = 3;
    private static final int WIDTH = 224;
    private static final int HEIGHT = 224;

    public static FloatBuffer process(Bitmap original) {
        int totalPixels = WIDTH * HEIGHT;
        FloatBuffer buffer = FloatBuffer.allocate(BATCH * CHANNELS * totalPixels);
        buffer.rewind();

        int[] pixels = new int[totalPixels];
        original.getPixels(pixels, 0, original.getWidth(), 0, 0, original.getWidth(), original.getHeight());

        for (int y = 0; y < HEIGHT; y++) {
            for (int x = 0; x < WIDTH; x++) {
                int idx = y * WIDTH + x;
                int color = pixels[idx];

                float r = ((color >> 16 & 0xFF) / 255.0f - 0.485f) / 0.229f;
                float g = ((color >> 8 & 0xFF) / 255.0f - 0.456f) / 0.224f;
                float b = ((color & 0xFF) / 255.0f - 0.406f) / 0.225f;

                buffer.put(idx, r);
                buffer.put(idx + totalPixels, g);
                buffer.put(idx + 2 * totalPixels, b);
            }
        }

        buffer.rewind();
        return buffer;
    }
}

2. 모델 로딩 및 추론 실행

다음 메서드는 ONNX 런타임을 초기화하고, Assets에서 모델 파일을 읽어 세션을 생성한 후 이미지를 전처리하여 추론을 수행합니다. URI 또는 Bitmap 중 하나로 입력을 받을 수 있습니다.

private float[] inferFromImage(Uri imageUri, Bitmap inputBitmap) throws IOException {
    float[][] rawOutput = new float[1][];
    OrtEnvironment env = OrtEnvironment.getEnvironment();
    AssetManager assets = getAssets();
    OrtSession.SessionOptions opts = new OrtSession.SessionOptions();

    try (InputStream modelStream = assets.open("model.onnx");
         ByteArrayOutputStream baos = new ByteArrayOutputStream()) {

        byte[] buf = new byte[4096];
        int read;
        while ((read = modelStream.read(buf)) != -1) {
            baos.write(buf, 0, read);
        }
        byte[] modelBytes = baos.toByteArray();
        OrtSession session = env.createSession(modelBytes, opts);

        // Bitmap 준비
        Bitmap sourceBitmap;
        if (inputBitmap != null) {
            sourceBitmap = inputBitmap;
        } else if (imageUri != null) {
            InputStream imageStream = getContentResolver().openInputStream(imageUri);
            sourceBitmap = BitmapFactory.decodeStream(imageStream);
        } else {
            throw new IllegalArgumentException("이미지 소스가 없습니다.");
        }

        Bitmap resized = Bitmap.createScaledBitmap(sourceBitmap, 224, 224, true);
        FloatBuffer processed = ImagePreProcessor.process(resized);
        OnnxTensor tensor = OnnxTensor.createTensor(env, processed, new long[]{1, 3, 224, 224});

        Map<String, OnnxTensor> inputMap = new HashMap<>();
        inputMap.put("input", tensor);

        OrtSession.Result result = session.run(inputMap);
        OnnxTensor outputTensor = (OnnxTensor) result.get(0);
        rawOutput = (float[][]) outputTensor.getValue();

        // 결과 출력 (디버깅용)
        for (float[] row : rawOutput) {
            for (float val : row) {
                System.out.print(val + " ");
            }
            System.out.println();
        }

    } catch (OrtException e) {
        e.printStackTrace();
    }

    return rawOutput[0];
}