latent diffusion segmentation
1.0.0
Este repositório contém a implementação Pytorch do LDMSeg: uma abordagem simples de difusão latente para segmentação panóptica e pintura de máscara. O código fornecido inclui treinamento e avaliação.
Uma abordagem simples de difusão latente para segmentação panóptica e pintura de máscara
Wouter Van Gansbeke e Bert De Brabandere
Este artigo apresenta uma abordagem de difusão latente condicional para enfrentar a tarefa de segmentação panóptica. O objetivo é omitir a necessidade de arquiteturas especializadas (por exemplo, redes de proposta de região ou consultas de objetos), funções de perda complexas (por exemplo, correspondência húngara ou baseada em caixas delimitadoras) e métodos adicionais de pós-processamento (por exemplo, clustering, NMS ou colagem de objetos). Como resultado, contamos com a Difusão Estável, que é uma estrutura independente de tarefas. A abordagem apresentada consiste em duas etapas: (1) projetar as máscaras de segmentação panópticas para um espaço latente com um autocodificador raso; (2) treinar um modelo de difusão no espaço latente, condicionado a imagens RGB.
Principais contribuições : Nossas contribuições são triplas:
O código é executado com versões recentes do Pytorch, por exemplo, 2.0. Além disso, você pode criar um ambiente python com Anaconda:
conda create -n LDMSeg python=3.11
conda activate LDMSeg
Recomendamos seguir a instalação automática (veja tools/scripts/install_env.sh ). Execute os seguintes comandos para instalar o projeto no modo editável. Observe que todas as dependências serão instaladas automaticamente. Como isso pode nem sempre funcionar (por exemplo, devido a problemas de CUDA ou gcc), dê uma olhada nas etapas de instalação manual.
python -m pip install -e .
pip install git+https://github.com/facebookresearch/detectron2.git
pip install git+https://github.com/cocodataset/panopticapi.gitOs pacotes mais importantes podem ser instalados rapidamente com pip como:
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 Consulte data/environment.yml para obter uma cópia do meu ambiente. Também contamos com algumas dependências do detectron2 e do panopticapi. Por favor, siga seus documentos.
Atualmente oferecemos suporte ao conjunto de dados COCO. Siga a documentação para instalar as imagens e suas máscaras de segmentação panóptica correspondentes. Além disso, dê uma olhada no diretório ldmseg/data/ para obter alguns exemplos no conjunto de dados COCO. Como nota lateral, a estrutura adotada deve ser bastante padronizada:
.
└── coco
├── annotations
├── panoptic_semseg_train2017
├── panoptic_semseg_val2017
├── panoptic_train2017 -> annotations/panoptic_train2017
├── panoptic_val2017 -> annotations/panoptic_val2017
├── test2017
├── train2017
└── val2017
Por último, mas não menos importante, altere os caminhos em configs/env/root_paths.yml para a raiz do conjunto de dados e o diretório de saída desejado, respectivamente.
A abordagem apresentada tem duas vertentes: primeiro, treinamos um codificador automático para representar mapas de segmentação em um espaço dimensional inferior (por exemplo, 64x64). A seguir, partimos de Modelos de Difusão Latente (LDM) pré-treinados, particularmente Difusão Estável, para treinar um modelo que pode gerar máscaras panópticas a partir de imagens RGB. Os modelos podem ser treinados executando os seguintes comandos. Por padrão, treinaremos no conjunto de dados COCO com o arquivo de configuração base definido em tools/configs/base/base.yaml . Observe que este arquivo será carregado automaticamente, pois dependemos do pacote hydra .
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 Mais detalhes sobre a passagem de argumentos podem ser encontrados em tools/scripts/train_ae.sh . Por exemplo, executo este modelo para 50 mil iterações em uma única GPU de 23 GB com um tamanho total de lote de 16.
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 denota o caminho para o modelo obtido na etapa anterior. Mais detalhes sobre a passagem de argumentos podem ser encontrados em tools/scripts/train_diffusion.sh . Por exemplo, executei este modelo para 200 mil iterações em 8 GPUs de 16 GB com um tamanho total de lote de 256.
Estamos planejando lançar vários modelos treinados. A métrica PQ (independente de classe) é fornecida no conjunto de validação COCO.
| Modelo | #Params | Conjunto de dados | Iters | QP | QUADRADO | RQ | Baixar link |
|---|---|---|---|---|---|---|---|
| EA | ~2 milhões | COCO | 66k | - | - | - | Baixar (23 MB) |
| LDM | ~800 milhões | COCO | 200 mil | 51,7 | 82,0 | 63,0 | Baixar (3,3 GB) |
Nota: Um AE menos poderoso (ou seja, menos camadas de downsampling ou upsampling) muitas vezes pode beneficiar a pintura interna, já que não realizamos ajustes finos adicionais.
A avaliação deve ser semelhante a:
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 Você pode adicionar parâmetros, se necessário. Limites mais altos, como --base.eval_kwargs.count_th 700 ou --base.eval_kwargs.mask_th 0.9 podem aumentar ainda mais os números. No entanto, usamos valores padrão estabelecendo limite de 0,5 e removendo segmentos com área menor que 512 para a avaliação.
Para avaliar um modelo pré-treinado acima, execute tools/scripts/eval.sh .
Aqui, visualizamos os resultados:
Se você achar este repositório útil para sua pesquisa, considere citar o seguinte artigo:
@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 }
}Para qualquer dúvida, entre em contato com o autor principal.
Este software é lançado sob uma licença Creative Commons que permite apenas uso pessoal e de pesquisa. Para obter uma licença comercial, entre em contato com os autores. Você pode ver um resumo da licença aqui.
Agradeço todos os repositórios públicos (veja também as referências no código), e em particular pelas bibliotecas detectron2 e difusores.
