Rust에서 RF-DETR ONNX 모델을 사용한 이미지 객체 식별

이전 프로젝트에서 YOLOX를 사용해볼 수 있는 기회가 있었습니다. 당시에 Rust로 프로젝트를 제작하고 있었기 때문에 PyTorch로 학습된 모델을 Rust에서 돌려보기위해 여러 방법들을 시도해보았습니다. 이 때 ONNX을 알게 되어 onnx로 모델을 변환하고 실행하는 시도도 해보았습니다. 최종적으로는 속도 차이로 tch-rs(파이토치 러스트 바인딩)을 사용하여 실행하는 것으로 결론을 내렸습니다. 당시에 아쉬웠던 부분은 YOLOX의 동작 속도였는데요. 우연찮게 LinkedIn에서 Roboflow의 RF-DETR 에 대한 게시물을 발견하게 되었고 나중에 기회가 되면 써봐야지 했습니다.

 

그리고 드디어 여유 시간이 생겨 실제로 만들어 볼 수 있게 되었습니다.

 

 

프로젝트에서 ONNX 모델을 사용할 예정이기 때문에 ONNX Runtime 설치가 필요합니다.

아래의 사이트에 방문하셔서 여러 환경에 따른 설치 방법을 안내 받으실 수 있습니다.

https://onnxruntime.ai/

 

 

여기에서는 macOS에서 사용하는 것을 가정하고 있으므로 아래의 명령어를 사용하여 설치하시기 바랍니다.

pip install onnxruntime

 

 

코드가 제법 길기 때문에 아래의 저장소에 있는 프로젝트를 가져오겠습니다.

https://github.com/cellaxon/rf_detr_onnx_test

git clone https://github.com/cellaxon/rf_detr_onnx_test
cd rf_detr_onnx_test

 

 

다음으로는 허깅페이스에서 모델 파일을 다운로드 받으셔야 합니다.

아래의 사이트에 방문하셔서 model.onnx 파일을 다운로드 받으시기 바랍니다.

https://huggingface.co/onnx-community/rfdetr_base-ONNX/tree/main/onnx

 

rf_detr_onnx_test 폴더 하위에 assets/models 폴더를 생성하시고
다운로드 받으신 모델 파일을 복사해주세요

# assets/models 폴더 생성 (없는 경우)
mkdir -p assets/models

# 모델 파일을 assets/models/model.onnx로 이동
mv ~/Downloads/model.onnx assets/models/

 

 

이제 실행해보겠습니다.

cargo run --release

 

 

프로그램 실행 화면

 

추론을 시작하려면 좌측 상단의 Select Image 버튼을 누른 다음 대상 이미지를 선택하시면 됩니다.

 

 

잘 되시나요?

 

처음에는 터미널에서 실행하고 파일로 이미지를 출력하는 형태로 만들었는데 아무래도 결과를 바로 볼 수 있는 편이 좋을 것 같아 GUI를 붙이고 이것저것 테스트하면서 뜯어고치다보니 여기까지 왔습니다.

 

GUI는 Dioxus를 써보려했으나 Cursor가 몇 번 시도 후에 egui로 방향을 틀어버리더라고요. 익숙해서 좋긴했지만 Dioxus는 언제 써볼 수 있을런지 모르겠습니다. AI가 스스로 포기하는 걸 보면 아직 쓰면 안될 것 같기도 하고요.

 

처음에는 출력하는 ID가 COCO class 인 줄 알았는데 아니었습니다.

고양이가 개로 나오고 개는 말로 표시되는 문제가..
아래의 파일을 참고하시면 됩니다.

https://huggingface.co/onnx-community/rfdetr_base-ONNX/blob/main/config.json

 

모델이 입력받는 이미지의 크기가 560x560으로 고정되어 있어 입력으로 넣는 이미지를 비율 그대로 축소해서 넣은 다음 추론을 실행하고 결과를 다시 원본 이미지의 위치로 변환한 뒤에 표시해주도록 만들었습니다.

 

아무래도 Rust에서 이런 작업을 하시는 분들이 없어서인지 여기까지 오는데 생각보다 애로 사항이 많았습니다만 다행히 여기까지 만들 수 있었습니다.

 

이 다음으로 뭔가 만들게 된다면 영상 파일이나 카메라 영상을 OpenCV로 읽어서 실시간 추론을 해보면 좋을 것 같네요.