2023.06

Meta - Segment Anything

https://github.com/facebookresearch/segment-anything

Polygon 형태의 작업자의 Labeling 작업효율을 극대화 하기 위한 신규 기능으로 Meta 에서 오픈소스로 공개한 Segment Anything 모델을 사용한 추론 기능을 넣기로 했다.

이를 구현하기 위해서 중요한 것은 MLOps 측에서는 1차 추론(Image Embedding(*.npy) 파일 생성) 까지만 처리하고, 실제 작업자가 작업을 할때는 이미지에 매칭되는 embedding 파일을 기반으로 Inference Session 을 생성해서 클라이언트 단에서 직접 2차 추론을 구현해야 했다.

위의 데모를 참고해서 실제 플랫폼에서 사용하기 위해 추론 기능을 옮기는 작업을 수행했다.

실행 환경은 storybook 7 (react-vite framework) 에서 테스트 했다.

onnxruntime-web 를 활용해서 추론 세션을 생성하는데, 이 세션에서 정적 wasm 파일을 hosting 되는 위치에서 runtime 에 직접 사용하기 때문에 정적 파일 serving 을 위해 vite-plugin-static-copy 플러그인을 연동시켜 줬다. (webpack 의 경우 copy-webpack-plugin 을 사용하면 된다)

// main.ts
import { mergeConfig } from 'vite';
import { viteStaticCopy } from 'vite-plugin-static-copy';

import type { StorybookConfig } from '@storybook/react-vite';

const config: StorybookConfig = {
	...
  async viteFinal(config, {}) {
    return mergeConfig(config, {
      plugins: [
        viteStaticCopy({
          targets: [
            {
              src: 'node_modules/onnxruntime-web/dist/*.wasm',
              dest: '/',
            }
          ]
        })
      ]
    });
  },
	...
};
export default config;

데모는 기능이 쪼개져 있어서 Custom Hook 으로 모듈화 시켜서 사용하도록 구성했다.

// useSegmentAnything.ts
const useSegmentAnything = ({
  image,
  modelSrc,
  imageEmbeddingSrc,
  hoverWait = 15,
}) => {
  // ONNX model
  const [model, setModel] = useState<InferenceSession | null>(null);

  // Image embedding tensor
  const [tensor, setTensor] = useState<Tensor | null>(null);

	...
}

데모에서 구현한 기능을 재구현하기 위해 여러가지를 신경써야 했다.

1. Bounding Box 를 그리는 중에 실시간 추론시 버벅임을 최소화 하기 위한 개선 - debounce 적용

// useSegmentAnything.ts
// Box 생성 중에도 추론이 가능하면서 드로잉 성능을 위해 debounce 적용
  const setBox = useCallback(
    debounce((box: Float32Array) => _setBox(box), 50, {
      leading: true,
      trailing: false,
    }),
    [],
  );