1. Mediapipe custom model 적용
이전 글 (mediapipe 이해)와 같은 실험환경에서 글을 작성하니 참고 바랍니다. 이번 글에서는 custom하게 만들어놓은 tflite model로 mediapipe에 적용해보겠습니다. tflite model의 task는 portrait segmentation 이며 mac book의 camera를 통해 live로 portrait segementation을 목적으로 합니다.
이제 차근차근 한 Step씩 custom한 모델을 mediapipe에 어떻게 적용시키는 지 보시려면 아래부터 쭉 읽으시면 되시고 과정은 생략하고 빠르게 결과를 확인하고 싶으실 경우 밑에 Appendix 1.을 읽어보시면 됩니다.
1.1 portrait segmentation model 다운로드
# make workspace directory
mkdir workspace
cd workspace
# download portrait segmentation model
wget https://github.com/anilsathyan7/Portrait-Segmentation/raw/master/mediapipe_slimnet/portrait_segmentation.tflite
먼저, 위 명령어를 통해 workspace 디렉토리를 만들고 해당 workspace에 portrait segmentation model을 다운받습니다. 다운로드받은 model의 input shape은 512x512x3이고 output shape은 512x512x2임을 netron으로 확인 가능합니다.
1.2 mediapipe repository 다운로드
mediapipe를 사용해야 하므로 git clone https://github.com/google/mediapipe.git
을 통해 workspace에 mediapipe를 다운로드합니다. 그리고 다운로드한 portrait model을 mv portrait_segmentation.tflite mediapipe/mediapipe/models/
통해 mediapipe model을 담는 디렉토리로 옮깁니다.
1.3 BUILD 파일 작성
bazel build를 통해 portrait model을 동작시킬 library를 만들기위해 아래 명령어들을 통해 BUILD파일을 만들어 보자.
cd mediapipe/mediapipe/examples/desktop
mkdir portrait_segmentation
cp hair_segmentation/BUILD portrait_segmentation/BUILD
BUILD파일은 hair_segmentation의 파일을 복사해왔지만 portrait segmentation을 위해 BUILD의 내용은 다음과 같이 수정해본다.
licenses(["notice"])
package(default_visibility = ["//mediapipe/examples:__subpackages__"])
#cpu only
cc_binary(
name = "portrait_segmentation_cpu",
deps = [
"//mediapipe/examples/desktop:demo_run_graph_main",
"//mediapipe/graphs/portrait_segmentation:desktop_calculators", #이후에 만들 파일
],
)
위 파일의 주요 특징은 다음과 같습니다.
- bazel build를 통해 만들어질 binary파일 이름은 portrait_segmentation_cpu 임.
- deps는 dependency를 말하며 binary파일을 만들어내는데 필요한 BUILD 파일을 정의
- desktop:demo_run_graph_main: ML model을 desktop cpu용으로 동작하도록 하는 graph를 실행하도록 함.
- portrait_segmentation:desktop_calculators: portrait model을 동작시키는 데 필요한 calculator들을 담고 있음.
mediapipe repo에는 이미 desktop:demo_run_graph_main을 실행시킬 BUILD파일이 존재하지만 portrait_segmentation:desktop_caculators는 존재하지않으므로 BUILD파일을 작성해야합니다. 밑에서 해보죠!
1.4 Graph 및 BUILD 파일 작성
이전 글에서 말씀드렸듯이 graph는 일련의 system flow을 나타내며 이를 통해 ML model을 inference하게됩니다. Graph는 .pbtxt 확장자를 가지는 파일로 작성됩니다. 그리고 여기서 작성할 BUILD 파일을 통해 graph에서 필요한(dependency있는) calculator들을 정의하게 됩니다.
먼저 다음 명령어들을 통해 graph 하위 디렉토리에 portrait segmentation 디렉토리를 만들고 기존에 mediapipe에서 제공하는 hair_segmentation BUILD파일과 .pbtxt파일을 복사해오죠. (portrait segmentation과 task가 비슷해서 가져옵니다ㅎㅎ)
cd mediapipe/mediapipe/graphs
mkdir portrait_segmentation
cp hair_segmentation/BUILD portrait_segmentation/BUILD
cp hair_segmentation/hair_segmentation_desktop_live.pbtxt portrait_segmentation/desktop_live.pbtxt
일단 BUILD 파일은 수정할 필요가 없습니다. 그 이유는 hair segmentation에서 사용하는 calculator들을 모두 사용할 것이기 때문입니다. 그렇다면 desktop_live.pbtxt만 다음과 같이 부분 수정해봅니다.
node {
calculator: "TfLiteConverterCalculator"
input_stream: "IMAGE:mask_embedded_input_video"
output_stream: "TENSORS:image_tensor"
node_options: {
[type.googleapis.com/mediapipe.TfLiteConverterCalculatorOptions] {
zero_center: false
max_num_channels: 3 #changed
}
}
}
...
node {
calculator: "TfLiteInferenceCalculator"
input_stream: "TENSORS:image_tensor"
output_stream: "TENSORS:segmentation_tensor"
input_side_packet: "CUSTOM_OP_RESOLVER:op_resolver"
node_options: {
[type.googleapis.com/mediapipe.TfLiteInferenceCalculatorOptions] {
model_path: "mediapipe/models/portrait_segmentation.tflite" #changed
use_gpu: false
}
}
}
...
node {
calculator: "RecolorCalculator"
input_stream: "IMAGE:throttled_input_video"
input_stream: "MASK:hair_mask"
output_stream: "IMAGE:output_video"
node_options: {
[type.googleapis.com/mediapipe.RecolorCalculatorOptions] {
color { r: 0 g: 0 b: 0 } #changed
mask_channel: RED
invert_mask: true #changed
}
}
}
수정된 부분에 대한 설명입니다.
- max_num_channels: 3
- cam으로 들어오는 image가 RGB channel로 3개이므로 3으로 설정
- model_path: "mediapipe/models/portrait_segmentation.tflite"
- custom한 model인 portrait_segmentation.tflite파일의 path를 명시
- color {r: 0 g: 0 b: 0} (optional)
- 저는 segmentation을 통해 masking될 부분의 색을 검은색(r:0, g:0, b:0)으로 설정
- 원하는 숫자를 입력해서 색을 바꾸셔도 됩니다!!
- invert_mask: true (optional)
- true로 설정하여 background를 위에서 정해준 검은색으로 masking시켜줌
- false일시 portrait인 사람이 검은색으로 masking 됨
1.5 Graph 및 BUILD 파일 작성
이제 1. bazel build를 통해 binary 파일을 만들고 2. bazel-bin을 통해 binary 파일을 실행해보죠.
GPU는 사용하지 않으므로 MEDIAPIPE_DISABLE_GPU=1
으로 설정하고 1.3 에서 만든 BUILD파일안의 portrait_segmentation_cpu 이름을 가진 cc_binary를 대상으로 bazel build를 합니다.
#bazel build 하기
bazel build -c opt --define MEDIAPIPE_DISABLE_GPU=1 mediapipe/examples/desktop/portrait_segmentation:portrait_segmentation_cpu
위 명령어를 실행하면 다음과 같이 출력될것이며 bazel-bin
이라는 디렉토리가 생김을 확인가능하며 bazel-bin/mediapipe/examples/desktop/portrait_segmentation/portrait_segmentation_cpu
에 binary 파일이 build를 통해 생성됨을 확인가능합니다.
이제 binary 파일을 bazel-bin으로 실행시켜 portrait segmentation 모델을 통해 live로 inference해보자.
# bazel-bin 으로 binary파일 실행
LOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/portrait_segmentation/portrait_segmentation_cpu \
--calculator_graph_config_file=mediapipe/graphs/portrait_segmentation/desktop_live.pbtxt
binary파일을 실행하는데 필요한 calculator_graph_config_file
옵션의 값으로 1.4에서 만든 graph를 담는 pbtxt파일을 명시한다. 다음은 드디어! Custom한 model로 portrait segmentation을 성공한 모습을 볼수 있다!
Appendix
1. 결과 지향적인 분들을 위해!
위의 1.4까지의 step을 생략하고 바로 결과를 확인하실 수 있게 1.4까지의 작업을 완료해놓은 github을 만들어놓았으니 해당 github를 clone하셔서 1.5 step을 따라하시면 바로 위와 같이 live로 portrait segmentation이 가능합니다. (아래 쭉 따라 치시면 되용)
git clone https://github.com/da2so/mediapipe_portrait_segmentation.git
cd mediapipe_portrait_segmentation
bazel build -c opt --define MEDIAPIPE_DISABLE_GPU=1 mediapipe/examples/desktop/portrait_segmentation:portrait_segmentation_cpu
LOG_logtostderr=1 bazel-bin/mediapipe/examples/desktop/portrait_segmentation/portrait_segmentation_cpu \
--calculator_graph_config_file=mediapipe/graphs/portrait_segmentation/desktop_live.pbtxt
'AI Engineering > TensorFlow' 카테고리의 다른 글
TensorFlow.js (2) - WebGL 기반 hand pose detection (2) | 2022.03.23 |
---|---|
TensorFlow.js (1) - TensorFlow.js 이해 및 detection 예제 (1) | 2022.03.17 |
TFLite 뽀개기 (3) - Quantization (2) | 2022.03.09 |
TFLite 뽀개기 (2) - TFLite Inference (0) | 2022.03.09 |
TFLite 뽀개기 (1) - TFLite 기본 이해 및 사용하기 (0) | 2022.03.09 |