Priceless

PIDNet Train & Test하기 본문

ComputerVision/Semantic Segmentation

PIDNet Train & Test하기

Hyun__ 2024. 5. 20. 11:15

1. Train 준비

PIDNet을 test한다

 

위 사이트의 내용을 번역? 보기 쉽게 정리하려고 한다

Train도 직접 할 수 있으니 필요하면 참고바랍니다

 

PIDNet을 사용하기 위해

Cityscapes dataset과 CamVid dataset을 사용할 수 있다

 

원하는 디렉토리에 git clone한다

우리는 아래의 data에 데이터셋을 추가해야 한다

 

 

pytorch 등의 버전은 본인의 학습 환경에 맞춘다

아래는 RTX 4060에 적합한 환경이다

conda create -y -n PIDNet python=3.10.6
conda activate PIDNet
conda install pytorch==2.0.0 torchvision==0.15.0 torchaudio==2.0.0 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install opencv-python==4.6.0.66
pip install yacs
pip install tqdm

 

 

2. Cityscapes Datasets 준비

Cityscapes 데이터셋의 양식은 HRNet을 reference한다

https://github.com/HRNet/HRNet-Semantic-Segmentation

 

Cityscapes dataset을 사용하기 위해서는 아래와 같은 구조로 저장해야 한다

 

우선 Cityscapes datasets을 다운 받는다. cityscapes dataset 사이트에 직접 다운로드 받는 방법이 있고, github를 통해서 다운로드 받는 방법이 있다. 어차피 계정을 만들어야해서 홈페이지에서 직접 다운로드 받는 것도 좋은 방법이다. 

https://www.cityscapes-dataset.com/downloads/

첫 번째와 세 번째 데이터를 다운받는다

 

다운로드 받은 후 아래와 같이 구성될 수 있도록 파일을 옮긴다

 

3. PIDNet Train

GPU 한 대를 사용하므로 명령어와 코드를 수정한다

python tools/train.py --cfg configs/cityscapes/pidnet_small_cityscapes.yaml 

 

in cityscapes/pidnet_small_cityscapes.yaml

GPUS: (0,)

 

in datasets/basedataset.py 등 여러 파일에서

1 epoch 전에 중단되는 부분에서 수정한다

np.int -> np.int32

 

4. PIDNet Test

train을 직접 수행해도 좋지만

PIDNet은 친절하게 train 결과 모델을 제공한다

아래의 파란 글씨를 눌러 모델을 다운받을 수 있다

필자는 PIDNet-S를 사용한다

 

모델을 지정된 디렉토리에 저장한다

 

원하는 이미지를 출력하기 전 성능을 테스트한다

python models/speed/pidnet_speed.py --a 'pidnet-s' --c 19 --r 1024 2048

 

필자의 환경 상의 결과로 약 96프레임이 나온다

 

 

이후 원하는 이미지를 samples 디렉토리에 저장하고

아래 명령어를 통해 테스트한다

Test code

직접 train한 모델을 이용해 test하고 싶은 경우

모델 디렉토리를 수정한다

 

이 코드로 진행할 경우

PIDNet 디렉토리 밖에 samples 디렉토리를 생성하여 이미지를 저장해야 한다

깃허브에 제공된 samples 디렉토리에 파일을 저장하면 결과가 나오지 않는다

python tools/custom.py --a pidnet-s --p pretrained_models/cityscapes/PIDNet_S_Cityscapes_test.pt --r samples --t .png

 

Linux에서 학습을 진행하는 경우

custom.py 에서 디렉토리 코드를 수정한다

"\\" -> "/"

    with torch.no_grad():
        for img_path in images_list:
            img_name = img_path.split("/")[-1]

 

 

5. Test 결과

 

참고

PIDNet github readme

https://github.com/XuJiacong/PIDNet?tab=readme-ov-file

 

PIDNet test

https://velog.io/@ssw9999/PIDNet-test-6q743h5b