[Deep Learning] 2D Human Pose Estimation을 한 줄로! onepose

https://github.com/developer0hye/onepose

GitHub - developer0hye/onepose: Human pose estimation within one line

 

현 시점(20231005)에서 SOTA 에 준하는 성능을 기록한 모델인 ViTPose 를 아주 사용하기 쉽게 파이썬 패키지화 해보았다.

 

설치 명령어는 아래와 같다.

 

pip install git+https://github.com/developer0hye/onepose.git

 

다음은 예제 코드다.

import cv2
import onepose

if __name__ == '__main__':
    img = cv2.imread('sample.png')
    model = onepose.create_model()

    keypoints = model(img)
    onepose.visualize_keypoints(img, keypoints, model.keypoint_info, model.skeleton_info)
    cv2.imshow('img', img)
    cv2.waitKey(0)

 

한 번 티스토리 대문 이미지 갖고 테스트 해보았는데 어...? 생각보다 잘된다. 테스트한 모델은 ViT-H 모델이 백본인 모델이다.  사실 다른 Pose Estimation 모델은 안돌려봐서 잘 모르겠는데 학습된 이미지 중에서 anime image가 많이 없을 것으로 예상되는데 이런 이미지에서 마저 얼추 잘 추정한 모습을 보자니 컴퓨터 비전 알고리즘의 성능이 정말 많이 올라온 것이 체감된다.

아래는 ViT-L 모델이 백본인 모델의 결과다. 이 이미지만 갖고 얘기해보자면, 다리 부분에서 확실히 H와 L간의 성능차이가 정성적으로 보인다. 그치만 이 모델또한 일반화 성능이 뛰어나다고 봐도 될 거 같다.

 

주저리 주저리

 

기존 ViTPose 프로젝트 같은 경우에는 One Drive로 Weights을 관리하여 Pre-trained Weights을 CLI 나 라이브러리를 통해서 다운로드 받기가 까다로웠다. 또한 mmpose 를 베이스로 한 프로젝트이기에 mmpose 구동에 필요한 각종 라이브러리들의 버전을 모두 맞춰줘야하는 불편함을 감수해야했다. 아래는 관련 이슈다.

 

https://github.com/ViTAE-Transformer/ViTPose/issues/101

Feature request: easy to use pre-trained models · Issue #101 · ViTAE-Transformer/ViTPose

 

그래서, 한 번 쓰기 쉽게 만들어 보자라는 생각이 들었다. 먼저, 모델 다운로드를 자동화하고 싶었고 mmpose를 걷어내고 최대한 적은 라이브러리 만으로 추론만 가능하게 만들고자했다. 모델을 다운로드 받고 원저자분한테 이런 재배포가 가능한지 메일로 문의를 드려봤는데 아직 답장을 받지는 못했다. Weights 재배포 금지에 대한 언급은 찾지 못하여 https://github.com/developer0hye/onepose/releases/tag/1.0.0, 깃헙 릴리즈 기능을 이용해 Weights을 업로드했다. 그리고 라이센스는 ViTPose의 Apache 라이센스를 그대로 가져왔다.

Release upload vitpose weights · developer0hye/onepose

 

이때, 2기가가 넘는 파일은 업로드가 불가능하여 Weights 을 FP16으로 Quantization을 거쳐서 다시 저장 후 배포했다. 그렇기에 ViT-H를 백본으로 하는 모델의 Weights 또한 업로드 할 수 있었다.

 

모델 정의하는 방식이나 지원하는 모델을 출력하는 방식은 timm 을 따랐다.