[AI] 3D 얼굴 생성을 위한 Gaussian Head Avatar 코드 환경 설정 및 실행

기술 단어장/AI

[AI] 3D 얼굴 생성을 위한 Gaussian Head Avatar 코드 환경 설정 및 실행

MFDO 2024. 12. 9. 19:15

조금은 고생을 했던 실행 과정을 소개한다.

Gaussian Head Avatar는

CNN이 가진 고해상도의 문제 극복의 대안중 하나인

Gaussian splatting을 이용한 얼굴 생성 기법을 적용한 코드이다.

https://github.com/YuelangX/Gaussian-Head-Avata

위의 공식 코드를 직접 실행하는 과정만을 이번에는 작성하고자 한다.

관련 고충은 이곳에서 구체적으로 서술하고 있다.

또한 한글 주석을 내 맘대로 잔뜩잔뜩 추가해

자주 설정을 까먹는 나에게 많은 도움을 주고 있다.

1) 코드 주석이 필요한 이에게

https://github.com/oMFDOo/Gaussian-Head-Avatar

2) 관련 고충들이 궁금한 이에게

https://github.com/users/oMFDOo/projects/7

환경 세팅을 30번 이상 수정한 나로는 이번에 최최종 깔끔 정리일 듯 하다.

python 생태계에 아직 어색해도 실행할 수있을 정도로 최대한 쉽게 설명하고자 한다.

CUDA 설정에 대한 주의

필자는 A100 GPU를 이용하는 김에 신나게 최신 CUDA 버전인 12.6을 이용하였다.

하지만 이 순간 아래의 짜릿한 경험을 할 수 있다.

1) pytorch 버전 불일치로 인한 수동 빌드

2) kaoiln 버전 불일치로 인한 수동 빌드 및 코드 미호환

이로 인해 비교적 안정적인 11.8버전을 이용하여 CUDA를 설치하였다.

CUDA 설치 과정

혹여나 필요한 누군가를 위해 안정적으로 설치했던 과정 코드를 나눈다.

이 부분은 본인이 사용하고 있는 GPU에 따라 다르기에 꼭 자신의 버전과 일치하도록 설치하길 바란다.

# nvidia 드라이브 삭제
sudo apt-get purge nvidia-*
sudo apt-get autoremove

# 전체 삭제 여부 확인 : 뜨는게 없어야함
dpkg -l | grep nvidia

# cuda 삭제
sudo apt-get purge cuda-*
sudo apt-get autoremove --purge

# 전체 삭제 여부 확인 : 뜨는게 없어야함
dpkg -l | grep cuda
sudo apt-get purge

# 권장 드라이버 확인 및 설치
ubuntu-drivers devices
sudo apt install nvidia-driver-565
sudo reboot
# 드라이버 설치 확인
nvidia-smi

# CUDA 설치 과정
rm -rf CUDA_11.8
mkdir CUDA_11.8
cd CUDA_11.8

# 해당 내용 진행함
# https://developer.nvidia.com/cuda-11-8-0-download-archive?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=20.04&target_type=deb_local

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-ubuntu2004.pin
sudo mv cuda-ubuntu2004.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-ubuntu2004-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo dpkg -i cuda-repo-ubuntu2004-11-8-local_11.8.0-520.61.05-1_amd64.deb
sudo cp /var/cuda-repo-ubuntu2004-11-8-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda

# CUDA 환경변수 추가
vi ~/.bashrc
export PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
source ~/.bashrc

sudo reboot
# CUDA 설치 확인
nvcc -V

Gaussian Head Avatar 환경 설치 과정

: 공식 문서의 내용을 가져오며 설치 과정을 설명하도록 하겠다.

1) conda 가상환경 설치 및 활성화

1-1) conda가 설치되어있음을 가정해놓았다. 그런데 나도 당연히 있어버려서 설치 과정은 다른 분 포스트로 넘기겠다.

1-2) 이미 코드를 가지고 있음을 가정해놓았다. 그렇기에 아래 과정이 선행된다.

git clone https://github.com/YuelangX/Gaussian-Head-Avatar.git
cd Gaussian-Head-Avatar

1-3) 우리는 제시한 코드에 이름 속성을 추가할 것이다. 나는 임시로 gha라고 지었다.

# environment.yaml에 정의된 종속성을 설치하고 이름은 gha로 하겠다.
conda env create -f environment.yaml --name gha

1-4) 해당 conda 가상환경을 이용할 것이기에 활성화도 해준다.

# gha 가상환경 활성화
conda activate gha

2) Pytorch3D 설치

일반적으로 미분이 불가한 상황에 자주 처하는 3D 모델은,

인공지능 학습의 역전파 과정에 미분이 있기에 미분가능한 렌더러가 중요하다.

이곳에서는 3D 해당 모델 렌더러로 Pytorch3D를 선정하였다.

gaussian splat의 경우의 분산화 연산이 가능한 학습 및 모델링에는 nvdiffrast가 더 많은 이점을 가졌기에

훌륭한 코드적 감각을 가진이라면 모델러를 바꾸어보기 바란다.

2-1) 기존 명령어 수정없이 복사해 pytorch3d를 설치한다.

pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py38_cu113_pyt1120/download.html

3) kaolin 설치

PyTorch 기반 라이브러리로, 3D 모델링, 점군 처리, 메쉬 생성, 렌더링, 3D 복원 등의 작업을 할 수 있도록 지원한다.

pytorch 기반이라는 즉슨 정말로 버전에 예민하다. 다른 방식의 설치 시 수많은 오류를 경험할 수 있다.

3-1) 기존 명령어 수정없이 복사해 kaolin을 설치한다.

pip install kaolin==0.13.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-1.12.0_cu113.html

4) 서브 모듈 설치

난해할 수 있고 난해하다.

차근차근 따라오자.

4-1) 해당 설명에서 제시한 gaussian-splatting의 서브모듈 경로로 들어간다.

4-2) 세 가지 중 diff-gaussian-rasterization을 먼저 설치할 것이다. ctrl + 좌클릭으로 페이지를 열어주자

4-3) 한 레포지터리가 나올텐데 그 레포지터리 코드를 clone 할 것이다.

cd .. # 단순히 gaussian-head-avatar 경로를 나가란 뜻
git clone https://github.com/graphdeco-inria/diff-gaussian-rasterization.git
cd diff-gaussian-rasterization

4-4) config.h 파일을 수정해주어야 한다.

이 수정을 거치지 않고 설치하면 이후 train_gaussianhead.py 코드 실행 시 차원 수가 맞지 않는 에러가 난다.

cd cuda_rasterizer
vi config.h
### 이곳에서 아래처럼 선언된 값을
### #define NUM_CHANNELS 3 // Default 3, RGB
### 이렇게 바꾼다.
### #define NUM_CHANNELS 32 // Default 3, RGB

- vi 수정법까지 다 적는건 속상하다 이곳의 링크를 참조하자.

4-4) diff-gaussian-rasterization를 설치한다.

cd .. # 단순 diff-gaussian-rasterization 경로로 가라는 뜻
# pip install 이전 아래가 필요할 수 있다. 설치하지 않으면 오류가 난다.
sudo apt install build-essential
sudo apt install libglm-dev
pip install .
cd .. # 설치 다했으니 나가란 뜻

4-5) 현재 진행 상황

이제 simple-knn 설치로 넘어가는 과정이다.

4-6) 놀랍게도 simple-knn은 같은 방식으로 레포지토리를 찾을 수 없다.

- 그렇기에 내가 직접 찾고 구동도 올바르게 됨을 확인한 simple-knn 코드 링크에 들어간다.

- 다행히 설치는 단순하다.

git clone https://github.com/DSaurus/simple-knn.git
cd simple-knn
pip install .

5) tets_data.npz 다운로드

- 테스트를 위한 numpy 데이터를 가져오는 작업이다. 파일 만들 때의 오타가 인간미가 있어보인다.

5-1) assets 경로 생성

cd Gaussian-Head-Avatar
mkdir assets

5-2) tets_data.npz 다운로드 링크에서 다운받아온다.

5-3) 나는 visual studio code로 ssh 연결했기에 assets 경로에 드래그해서 넣었다.

- 나는 가상환경이기에 ssh 연결했는데, 각자의 방법에 맞게 옮기도록 하자

6) Dataset 다운로드

: 대강 상업적 이용은 제한하며, 공식 연구실에 대해서만 진짜 데이터 셋을 지원한다는 의미이다.

- 테스트를 위한 임시 데이터(빨간상자)를 제공하는데 이 데이터 또한 아래 약관에 서명한 이후 사용해야한다.

- 일단 이 글을 읽는 이가 약관에 동의했다는 가정하에 진행하겠다.

- 또한 임시 제공 데이터는 가공을 마쳤으므로 다른 내용은 진행할 필요가 없다.

6-1) mini_demo_dataset 다운로드 링크로 들어가 다운로드 받는다.

6-2) 이것 또한 드래그로 옮겼다.

# 현재Gaussian-Head-Avatar 경로임 
sudo apt install unzip
unzip mini_demo_dataset.zip

이를 끝으로 환경 구성이 끝났다.

Gaussian Head Avatar 실행

: 공식 문서의 내용을 가져오며 실행 과정을 설명하도록 하겠다.

1) train_meshhead.py 실행

: mesh학습을 볼 수 있다. 명령어 수정없이 다음과 같이 실행하면, 아래와 같이 학습을 진행한다.

python train_meshhead.py --config config/train_meshhead_N031.yaml

- 결과는 results의 meshhead_NeRSemble031 폴더에 저장된다.

- 결과 이미지는 recorder/Recorder.py.에서 조정 가능하며 아래 의미를 지닌다.

            # 원본 이미지
            image = log_data['data']['images'][0, 0].permute(1, 2, 0).detach().cpu().numpy()
            image = (image * 255).astype(np.uint8)[:, :, ::-1]

            # 렌더링 이미지
            render_image = log_data['data']['render_images'][0, 0, :, :, 0:3].detach().cpu().numpy()
            render_image = (render_image * 255).astype(np.uint8)[:, :, ::-1]

            # 렌더링된 법선 이미지
            render_normal = log_data['data']['render_normals'][0, 0].detach().cpu().numpy() * 0.5 + 0.5
            render_normal = (render_normal * 255).astype(np.uint8)[:, :, ::-1]

- 여러 번 학습을 거치면 아래와 같은 결과물을 볼 수 있다.

2) train_gaussianhead.py 실행

: Gaussian 학습을 볼 수 있다. 명령어 수정없이 다음과 같이 실행하면, 아래와 같이 학습을 진행한다.

python train_gaussianhead.py --config config/train_gaussianhead_N031.yaml

- 결과 이미지는 recorder/Recorder.py.에서 조정 가능하며 아래 의미를 지닌다.

            # 원본 이미지
            image = log_data['data']['images'][0].permute(1, 2, 0).detach().cpu().numpy()
            image = (image * 255).astype(np.uint8)[:, :, ::-1]

            # 렌더링된 이미지
            render_image = log_data['data']['render_images'][0, 0:3].permute(1, 2, 0).detach().cpu().numpy()
            render_image = (render_image * 255).astype(np.uint8)[:, :, ::-1]

            # 잘라낸 이미지
            cropped_image = log_data['data']['cropped_images'][0].permute(1, 2, 0).detach().cpu().numpy()
            cropped_image = (cropped_image * 255).astype(np.uint8)[:, :, ::-1]

            # 초해상도 이미지
            supres_image = log_data['data']['supres_images'][0].permute(1, 2, 0).detach().cpu().numpy()
            supres_image = (supres_image * 255).astype(np.uint8)[:, :, ::-1]

- 여러 번 학습을 거치면 아래와 같은 결과물을 볼 수 있다.

> 모델링에 방해 요소인 광원 제거 과정이 포함되었기 때문에 눈의 반사광이 사라진 모습이 두드러진다.