Waymo 데이터셋을 직접 전처리해서 street gaussian학습/평가해보기

이전글에서는 street gaussian 저자들이 이미 모든 전처리를 다 해놓은 데이터셋을 받아 학습을 돌려봤었다.

Street-gaussian 학습환경 세팅 & 학습시켜보기

 

이번에는 raw한 웨이모 데이터셋 그 자체를 받아 직접 처리를 하여 돌려보자.

 

기본적으로 이전 글의 내용에서 추가적으로 이어진다고 생각하면 좋다.(환경, 설치 패키지 등)

---

웨이모 데이터셋 다운로드

웨이모 데이터셋은 google cloud에 업로드 되어있다.

로컬 PC로 다운받고 다시 서버로 보내는 귀찮은 작업을 하지 않기 위해 아래 명령어를 통해 툴을 설치하자.

나는 서버를 사용하기 때문에 이렇게 하지만 로컬에서 코드를 돌린다면 안해도 된다.

sudo snap install google-cloud-cli --classic

 

웨이모 데이터셋은 구글로그인으로 인증된 사람에게만 다운로드를 가능하게 해놓았다고 한다.

아래 명령어를 실행시키면 터미널에 긴 링크가 하나 나오는데, 이 링크를 복사해 브라우저에 띄워보자.

 

gcloud auth login

구글 로그인을 하면 아래와 같은 창이 뜬다.

밑의 키를 복사해서 터미널이 붙여넣자.

이제 다운받을 수 있다.

나는 가장 위에 있는 10017... 번호의 segment를 다운했다.

아래 명령어로 다운 가능하다.

이때 data/waymo 경로는 이전 글에서도 있던 경로인데, 이번 데이터를 저장하기 위해 하위에 raw폴더를 만들어놓았다.

이 명령어를 실행시키기 전에, 명령어에 넣은 경로의 폴더가 있는지 확인하자.

이 방식은 해당 경로에 직접 다운이 된다. 만약 심볼릭 링크를 사용할 예정이면 다른 경로를 입력해야한다.

gcloud storage cp "gs://waymo_open_dataset_v_1_4_1/individual_files/training/segment-10017090168044687777_6380_000_6400_000_with_camera_labels.tfrecord" data/waymo/raw/

 

다운 받은 결과이다.

 

*혹시나 그래도 권한이 없다고 뜬다면 웨이모 오픈 데이터셋 공식 홈피에 들어가서 회원가입하자.

간단히 이름과 구글계정, 소속 정도만 입력하면 된다.

https://waymo.com/open/

 

---

데이터 전처리

 

1. 파일 생성 및 수정

script/waymo/waymo_splits경로 밑에 my_test.txt라는 파일을 하나 만들고 내용은 아래처럼 하자.

# scene_id, seg_name, start_timestep, end_timestep, scene_type
0,seg10017090168044687777_6380_000_6400_000_with_camera_labels,0,-1,dynamic

또 waymo_splits/waymo_converter.py파일의 대략 548번째 라인 정도를 보면 assert문이 하나 있는데 주석처리하자.

 

 

2. converter실행

이제 아래 코드를 실행하자.

python script/waymo/waymo_converter.py \
  --root_dir data/waymo/raw \
  --save_dir data/waymo/training \
  --split_file script/waymo/waymo_splits/my_test.txt \
  --segment_file script/waymo/waymo_splits/segment_list_train.txt

 

이 코드가 하는 일은 아래와 같다.

 

1. /raw폴더에서 해당 .tfrecord파일을 찾고 공식 리스트(segment_list_train.txt)에 있는지 확인한다.

2. 이어서 tfrecord는 바이너리파일이므로 이를 프레임별로 데이터를 가공한다.

3. 5개의 카메라에서 나온 영상을 프레임으로 나눠 jpg파일로 저장한다.

4. 각 객체의 2D BB정보를 파싱한다.

5. LiDAR센서의 점 데이터를 월드 좌표계로 변환한다.

6. cameras.json을 만들어 모든 프레임의 카메라 관련 정보를 저장한다.

7. track_info.txt를 만들어 각 객체가 시간에 따라 어디로 움직이는지 정보를 저장한다.

 

실행한 터미널 결과를 보면 흐름이 나와있다.

(street-gaussian) cv@cv-Super-Server:~/project/street_gaussians$ python script/waymo/waymo_converter.py   --root_dir data/waymo/raw   --save_dir data/waymo/training   --split_file script/waymo/waymo_splits/my_test.txt   --segment_file script/waymo/waymo_splits/segment_list_train.txt
Processing sequence segment-10017090168044687777_6380_000_6400_000_with_camera_labels...
Saving to data/waymo/training/000
Processing ego pose...
198it [00:01, 192.16it/s]
Processing camera calibration...
198it [00:00, 401.80it/s]
Processing image data...
198it [02:50,  1.16it/s]
Processing image data done...
Processing LiDAR data...
198it [05:02,  1.53s/it]
Processing LiDAR data done...
Processing tracking data...
198it [00:30,  6.47it/s]
Saving dynamic mask ...
198it [00:11, 17.69it/s]
Saving dynamic mask done...

 

아래는 실행 전 후의 폴더 구조 변화이다.

 

3. LiDAR Depth생성

아래 코드를 실행시켜 pointcloud.npz을 보고 lidar_depth라는 폴더를 만들게 한다.

python script/waymo/generate_lidar_depth.py --datadir data/waymo/training/000

 

해당 000폴더 안에 lidar_depth 폴더가 생성된 것을 볼 수 있다.

 

---

GroundingDINO, SAM 설치

 

이것들은 하늘 마스크를 만드는데 필요하다.

GroundingDINO는 텍스트를 입력받아('하늘') 이미지에서 해당 부분을 박스를 친다.

SAM은 그 박스를 기반으로 더 정확하게 픽셀단위로 '하늘'을 마스킹 한다.

 

프로젝트 폴더 루트(/street_gaussian)에서 아래 명령어를 순서대로 실행한다.

git clone https://github.com/IDEA-Research/GroundingDINO.git
cd GroundingDINO
export CC=gcc
export CXX=g++
pip install --no-build-isolation -e .

 

다시 상위폴더(/street_gaussian)로 나간 후

가중치 파일 2개를 다운받는다.

각각 GroundingDINO와 SAM의 가중치이다.

두 파일의 용량을 합치면 거의 3GB가까이 된다.

cd ..
wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth

 

그리고 GroundDINO가 텍스트 입력을 요구하므로 트랜스포머를 깔아야한다.

트랜스포머가 깔려있긴 할텐데 더 높은 버전이기 때문에 버전 일치를 위해 아래 명령어를 실행해 낮춰준다.

원래 이 버전 충돌 문제는 파이토치와 트랜스포머의 버전이 다른 것 때문인데, 파이토치를 올려도 된다.

하지만 우리의 경우 파이토치를 올렸다가는 3DGS관련 CUDA와 또 골치 아픈 문제가 생길것이기 때문에 트랜스포머를 낮춘다.

pip install transformers==4.36.2

 

최종적으로 아래 명령어를 실행하면 하늘을 마스킹한 결과가 나온다.

python script/waymo/generate_sky_mask.py \
  --datadir data/waymo/training/000 \
  --sam_checkpoint sam_vit_h_4b8939.pth

 

아래는 프레임0의 0번 카메라에 대한 마스킹 결과이다.

 

---

학습 돌리기

street_gaussian/configs/default.yaml파일을 열면 ???라고 써져있는 부분들이 있다.

현재 학습시킬 데이터셋에 맞게 채워주면된다.

몇가지 채운 값들에 대해 설명하자면,

selected_frames는 주석처리를 하므로써 모든 프레임에 대해 학습할 것이라는 의미를 전달했고, 

white_background는 false라고 설정했다.3dgs의 배경색을 생각하면 된다.

lamdba_sky_scale은 하늘 데이터를 얼마나 믿고 학습에 사용할지를 얘기한다.

주석에서 측면 카메라에서 왜곡이 있는 경우도 있다고 하니 0.5로 설정했다.

task: waymo_raw
source_path: ./data/waymo/training/000
exp_name: waymo_raw_000
to_cuda: true # move training cameras to GPU 

data:
  split_test: -1
  split_train: 1
  type: Waymo
  white_background: false
  #selected_frames: ???
  cameras: [0, 1, 2, 3, 4]
  extent: 10 # the unit here is meter
  use_colmap: true
  filter_colmap: true
  box_scale: 1.5

model:
  gaussian:
    sh_degree: 1
    fourier_dim: 5
    fourier_scale: 1.
    flip_prob: 0.2
  nsg:
    include_bkgd: true
    include_obj: true
    include_sky: true
    opt_track: true
  sky:
    resolution: 1024
    white_background: false  # set the default color based on the current time of day

train:
  iterations: 30000
  test_iterations: [7000, 30000]
  save_iterations: [30000]
  checkpoint_iterations: [30000]

optim:
  prune_box_interval: 100
  densification_interval: 100
  densify_from_iter: 500
  densify_grad_threshold_bkgd: 0.0006
  densify_grad_abs_bkgd: True
  densify_grad_threshold_obj: 0.0002
  densify_grad_abs_obj: False
  densify_grad_threshold: 0.0002
  densify_until_iter: 25000
  feature_lr: 0.0025
  max_screen_size: 20
  min_opacity: 0.005
  opacity_lr: 0.05
  opacity_reset_interval: 3000
  percent_big_ws: 0.1
  percent_dense: 0.01
  position_lr_delay_mult: 0.01
  position_lr_final: 1.6e-06
  position_lr_init: 0.00016
  position_lr_max_steps: 50000
  rotation_lr: 0.001
  scaling_lr: 0.005
  semantic_lr: 0.01

  lambda_dssim: 0.2
  lambda_sky: 0.05
  lambda_sky_scale: [1.0, 1.0, 1.0, 0.5, 0.5] # the sky mask is sometimes inaccurate, you can set the loss scale for each camera
  lambda_mask: 0.1
  lambda_reg: 0.1
  lambda_depth_lidar: 0.1

  track_position_lr_delay_mult: 0.2
  track_position_lr_init: 0.0005
  track_position_lr_final: 0.0001
  track_position_max_steps: 30000

  track_rotation_lr_delay_mult: 0.5
  track_rotation_lr_init: 0.00001
  track_rotation_lr_final: 0.000005
  track_rotation_max_steps: 30000

render:
  fps: 24
  concat_cameras: [1, 0, 2]

 

이제 아래 명령어로 학습을 시작시키면 된다.

python train.py --config configs/default.yaml

GPU가 여러개인 경우 특정 GPU를 지정해 사용하고 싶다면 아래처럼 한다.

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=1 python train.py --config configs/default.yaml

 

이게 VRAM을 많이 잡아먹는데, 그래서 제한된 VRAM안에서 학습을 돌리는 몇가지 방법을 소개한다.

 

1. resolution 설정

default.yaml파일의 data 부분에 resolution을 추가한다.

data:
	resolution: 4
    ...

 

2. to_cuda설정

default.yaml파일을 보면, 4번째 줄에 아래처럼 되어있다.

이게 true로 되어있으면 모든 카메라의 이미지 데이터를 다 GPU로 보내놓느다.

빠르지만 VRAM을 많이 먹기 때문에 false로 바꿔서 필요할 때만 CPU에서 GPU로 데이터를 보내는 방식을 사용할 수 있다.

to_cuda: true

 

각자의 GPU상황에 맞게 위 방식들의 적용여부를 정하면 될 듯하다.

---

평가하기

아래 명령어를 쳐서 영상을 만들어보자.

이때 옵션으로 evaluate나 trajectory를 넣을 수 있다.

evaluate는 metric등을 사용해 지표계산을 하기 위해 사용한다.

각 프레임별, 카메라별 gt, rgb, depth, diff의 이미지를 출력한다. 영상은 나오지 않는다.

trajectory는 어딘가에 시연할 용도로 부드러운 영상을 만드는 옵션이다.

CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 python render.py --config output/waymo_raw/waymo_raw_000/configs/config_000000.yaml mode evaluate