[Python] rembg로 배경을 지워보기

이 포스트에서는 딥러닝 모델로 물체 주변의 배경을 지워주는 라이브러리인 rembg를 소개한다.

rembg 소개

rembg는 2020년 어워드를 받은 U²-net이라는 모델의 pretrained 모델을 사용할 수 있게 만들어진 라이브러리로, 배경을 지우는데 탁월한 성능을 보여준다.

 

 

깃허브 공식 리포에서 보여주는 배경 지우기 예시. 단색 배경을 마스킹하여 지우는 기존의 크로마키 기법보다 유연하게 사용할 수 있는 강점이 있다.

사용해 보기

rembg 공식 리포에서는 cpu 또는 gpu 사용에 따라 필요한 설치를 구분했다.

 

pip install rembg[cpu]

 

or

pip install rembg[gpu]

 

gpu 사용을 하려면 onnxruntime-gpu 지원 여부를 확인해야 한다고 한다.

 

https://onnxruntime.ai 에서 확인 및 설치를 진행할 수 있다.

 

rembg는 단순 바이너리로 읽어드린 것부터, opencv로 읽은 numpy array까지 다양한 환경에서 처리를 수행할 수 있다. 이 포스트에서는 단순한 확인을 위해 바이너리로 읽어서 처리를 수행해 보자.

 

import rembg


if __name__ == "__main__":
    with open("test.png", 'rb') as f:
        with open("test_result.png", 'wb') as f2:
            img_in = f.read()
            img_out = rembg.remove(img_in)
            f2.write(img_out)

 

 

최초 실행하면 인터넷에서 모델을 다운로드한다. cpu 모델로 사용하면 1000*1000 이하의 이미지에 대해 3초 내외의 시간으로 결과를 출력한다.

적용 결과

 

배경과 배경이 아닌 대상이 모호한 경우를 제외하면 현실 인물, 가상 인물, 물체와 같은 여러 상황에서도 pretrained 모델로 훌륭한 결과를 보여주고 있다.

 

이처럼 rembg는 배경을 지우는 기능을 사용하고자 할 때 매우 좋은 선택지가 될 수 있을 것이다.