Panoptic 분할 및 마스크 인페인팅을 위한 간단한 잠재 확산 접근법

이 저장소에는 Panoptic 분할 및 마스크 인페인팅을 위한 간단한 잠재 확산 접근 방식인 LDMSeg의 Pytorch 구현이 포함되어 있습니다. 제공된 코드에는 교육과 평가가 모두 포함되어 있습니다.

Panoptic 분할 및 마스크 인페인팅을 위한 간단한 잠재 확산 접근법
바우터 반 간스베케(Wouter Van Gansbeke)와 베르트 드 브라반데레(Bert De Brabandere)

내용물

소개
설치
- 자동 설치
- 수동 설치
훈련
- 1단계: 자동 인코더 학습
- 2단계: LDM 학습
사전 훈련된 모델
소환
특허
감사의 말씀

? 소개

이 논문은 파놉틱 분할 작업을 해결하기 위한 조건부 잠재 확산 접근법을 제시합니다. 목표는 특수 아키텍처(예: 지역 제안 네트워크 또는 객체 쿼리), 복잡한 손실 함수(예: 헝가리어 일치 또는 경계 상자 기반) 및 추가 후처리 방법(예: 클러스터링, NMS)에 대한 필요성을 생략하는 것입니다. , 또는 개체 붙여넣기). 결과적으로 우리는 작업에 구애받지 않는 프레임워크인 Stable Diffusion을 사용합니다. 제시된 접근 방식은 두 단계로 구성됩니다. (1) 얕은 자동 인코더를 사용하여 Panoptic 분할 마스크를 잠재 공간에 투영합니다. (2) RGB 이미지를 조건으로 잠재 공간에서 확산 모델을 훈련합니다.

주요 기여 : 우리의 기여는 세 가지입니다:

생성 프레임워크 : Panoptic 분할을 위한 LDM(Latent Diffusion Model)을 기반으로 한 완전 생성적 접근 방식을 제안합니다. 우리의 접근 방식은 Stable Diffusion을 기반으로 단순성을 추구하고 컴퓨팅을 용이하게 합니다. 먼저 사전 정의된 클래스에서 Panoptic 분할을 해방하기 위해 클래스에 구애받지 않는 설정을 연구합니다.
범용 설계 : 우리의 접근 방식은 대부분의 일반적인 방법에 존재하는 특수 아키텍처, 복잡한 손실 기능 및 객체 감지 모듈을 우회합니다. 여기에서 잡음 제거 목표는 객체 쿼리, 지역 제안 및 헝가리어 매칭의 필요성을 생략합니다. 이 간단하고 일반적인 접근 방식은 깊이 예측, 돌출성 추정 등과 같은 광범위한 밀집 예측 작업에 대한 향후 확장을 위한 길을 열어줍니다.
마스크 인페인팅 : 장면 중심 데이터세트에 대한 접근 방식을 성공적으로 적용하고 다양한 희소성 수준에 대한 마스크 인페인팅 기능을 시연합니다. 이 접근 방식은 전역 마스크 인페인팅에 대한 유망한 결과를 보여줍니다.

? 설치

코드는 최신 Pytorch 버전(예: 2.0)에서 실행됩니다. 또한 Anaconda를 사용하여 Python 환경을 만들 수 있습니다.

 conda create -n LDMSeg python=3.11
conda activate LDMSeg

자동 설치

자동 설치를 따르는 것이 좋습니다( tools/scripts/install_env.sh 참조). 편집 가능 모드에서 프로젝트를 설치하려면 다음 명령을 실행하십시오. 모든 종속성은 자동으로 설치됩니다. 이것이 항상 작동하지 않을 수도 있으므로(예: CUDA 또는 gcc 문제로 인해) 수동 설치 단계를 살펴보시기 바랍니다.

python -m pip install -e .
pip install git+https://github.com/facebookresearch/detectron2.git
pip install git+https://github.com/cocodataset/panopticapi.git

수동 설치

가장 중요한 패키지는 다음과 같이 pip를 사용하여 빠르게 설치할 수 있습니다.

pip install torch torchvision einops                            # Main framework
pip install diffusers transformers xformers accelerate timm     # For using pretrained models
pip install scipy opencv-python                                 # For augmentations or loss
pip install pyyaml easydict hydra-core                          # For using config files
pip install termcolor wandb                                     # For printing and logging

내 환경의 복사본은 data/environment.yml 참조하세요. 또한 detectorron2와 panopticapi의 일부 종속성에 의존합니다. 해당 문서를 따르십시오.

?️ 데이터세트

현재 COCO 데이터 세트를 지원합니다. 이미지와 해당 팬옵틱 분할 마스크를 설치하려면 문서를 따르십시오. 또한 COCO 데이터세트에 대한 몇 가지 예를 보려면 ldmseg/data/ 디렉터리를 살펴보세요. 참고로 채택된 구조는 상당히 표준적이어야 합니다.

 .
└── coco
    ├── annotations
    ├── panoptic_semseg_train2017
    ├── panoptic_semseg_val2017
    ├── panoptic_train2017 -> annotations/panoptic_train2017
    ├── panoptic_val2017 -> annotations/panoptic_val2017
    ├── test2017
    ├── train2017
    └── val2017

마지막으로 configs/env/root_paths.yml 의 경로를 데이터 세트 루트와 원하는 출력 디렉터리로 각각 변경합니다.

⏳ 훈련

제시된 접근 방식은 두 가지로 나누어집니다. 첫째, 낮은 차원 공간(예: 64x64)에서 분할 맵을 나타내도록 자동 인코더를 훈련합니다. 다음으로 사전 훈련된 LDM(Latent Diffusion Model), 특히 Stable Diffusion에서 시작하여 RGB 이미지에서 팬옵틱 마스크를 생성할 수 있는 모델을 훈련합니다. 모델은 다음 명령을 실행하여 훈련할 수 있습니다. 기본적으로 tools/configs/base/base.yaml 에 정의된 기본 구성 파일을 사용하여 COCO 데이터 세트를 학습합니다. hydra 패키지를 사용하므로 이 파일은 자동으로 로드됩니다.

1단계: Panoptic Segmentation Map에서 자동 인코더 훈련

 python - W ignore tools / main_ae . py 
    datasets = coco 
    base . train_kwargs . fp16 = True 
    base . optimizer_name = adamw 
    base . optimizer_kwargs . lr = 1e-4 
    base . optimizer_kwargs . weight_decay = 0.05

인수 전달에 대한 자세한 내용은 tools/scripts/train_ae.sh 에서 확인할 수 있습니다. 예를 들어, 총 배치 크기가 16인 23GB의 단일 GPU에서 50,000번의 반복에 대해 이 모델을 실행합니다.

2단계: RGB 이미지를 조건으로 한 Panoptic 분할을 위한 LDM 훈련

 python - W ignore tools / main_ldm . py 
    datasets = coco 
    base . train_kwargs . gradient_checkpointing = True 
    base . train_kwargs . fp16 = True 
    base . train_kwargs . weight_dtype = float16 
    base . optimizer_zero_redundancy = True 
    base . optimizer_name = adamw 
    base . optimizer_kwargs . lr = 1e-4 
    base . optimizer_kwargs . weight_decay = 0.05 
    base . scheduler_kwargs . weight = 'max_clamp_snr' 
    base . vae_model_kwargs . pretrained_path = '$AE_MODEL'

$AE_MODEL 이전 단계에서 얻은 모델의 경로를 나타냅니다. 인수 전달에 대한 자세한 내용은 tools/scripts/train_diffusion.sh 에서 확인할 수 있습니다. 예를 들어, 총 배치 크기가 256인 16GB GPU 8개에서 200,000번의 반복에 대해 이 모델을 실행했습니다.

사전 훈련된 모델

우리는 여러 가지 훈련된 모델을 출시할 계획입니다. (클래스에 구애받지 않는) PQ 메트릭은 COCO 검증 세트에 제공됩니다.

모델	#Params	데이터세트	반복	PQ	평방	RQ	다운로드 링크
AE	~200만	머리	66,000	-	-	-	다운로드(23MB)
LDM	~800M	머리	200,000	51.7	82.0	63.0	다운로드(3.3GB)

참고: 덜 강력한 AE(즉, 다운샘플링 또는 업샘플링 레이어가 적음)는 추가적인 미세 조정을 수행하지 않으므로 인페인팅에 도움이 되는 경우가 많습니다.

평가는 다음과 같아야 합니다.

 python - W ignore tools / main_ldm . py 
    datasets = coco 
    base . sampling_kwargs . num_inference_steps = 50 
    base . eval_only = True 
    base . load_path = $ PRETRAINED_MODEL_PATH

필요한 경우 매개변수를 추가할 수 있습니다. --base.eval_kwargs.count_th 700 또는 --base.eval_kwargs.mask_th 0.9 와 같은 더 높은 임계값은 수치를 더욱 높일 수 있습니다. 그러나 평가를 위해 0.5로 임계값을 설정하고 512보다 작은 영역의 세그먼트를 제거하여 표준 값을 사용합니다.

위에서 사전 학습된 모델을 평가하려면 tools/scripts/eval.sh 실행하세요.

여기서는 결과를 시각화합니다.

? 소환

이 저장소가 귀하의 연구에 유용하다고 생각되면 다음 논문을 인용하는 것을 고려해 보십시오.

 @article { vangansbeke2024ldmseg ,
  title = { a simple latent diffusion approach for panoptic segmentation and mask inpainting } ,
  author = { Van Gansbeke, Wouter and De Brabandere, Bert } ,
  journal = { arxiv preprint arxiv:2401.10227 } ,
  year = { 2024 }
}