PolyMaX: Predizione Densa Generale con Mask Transformer

PolyMaX: General Dense Prediction with Mask Transformer

November 9, 2023

Autori: Xuan Yang, Liangzhe Yuan, Kimberly Wilber, Astuti Sharma, Xiuye Gu, Siyuan Qiao, Stephanie Debats, Huisheng Wang, Hartwig Adam, Mikhail Sirotenko, Liang-Chieh Chen

cs.AI

Abstract

I compiti di predizione densa, come la segmentazione semantica, la stima della profondità e la predizione delle normali di superficie, possono essere facilmente formulati come classificazione per pixel (output discreti) o regressione (output continui). Questo paradigma di predizione per pixel è rimasto popolare grazie alla diffusione delle reti completamente convoluzionali. Tuttavia, sulla recente frontiera dei compiti di segmentazione, la comunità ha assistito a un cambiamento di paradigma dalla predizione per pixel alla predizione a cluster con l'emergere delle architetture transformer, in particolare i mask transformer, che predicono direttamente un'etichetta per una maschera invece che per un pixel. Nonostante questo cambiamento, i metodi basati sul paradigma di predizione per pixel dominano ancora i benchmark per gli altri compiti di predizione densa che richiedono output continui, come la stima della profondità e la predizione delle normali di superficie. Ispirati dal successo di DORN e AdaBins nella stima della profondità, ottenuto discretizzando lo spazio di output continuo, proponiamo di generalizzare il metodo basato sulla predizione a cluster ai compiti generali di predizione densa. Questo ci permette di unificare i compiti di predizione densa con il framework dei mask transformer. In modo notevole, il modello risultante PolyMaX dimostra prestazioni all'avanguardia su tre benchmark del dataset NYUD-v2. Speriamo che il nostro design semplice ma efficace possa ispirare ulteriori ricerche sullo sfruttamento dei mask transformer per più compiti di predizione densa. Codice e modello saranno resi disponibili.

English

Dense prediction tasks, such as semantic segmentation, depth estimation, and surface normal prediction, can be easily formulated as per-pixel classification (discrete outputs) or regression (continuous outputs). This per-pixel prediction paradigm has remained popular due to the prevalence of fully convolutional networks. However, on the recent frontier of segmentation task, the community has been witnessing a shift of paradigm from per-pixel prediction to cluster-prediction with the emergence of transformer architectures, particularly the mask transformers, which directly predicts a label for a mask instead of a pixel. Despite this shift, methods based on the per-pixel prediction paradigm still dominate the benchmarks on the other dense prediction tasks that require continuous outputs, such as depth estimation and surface normal prediction. Motivated by the success of DORN and AdaBins in depth estimation, achieved by discretizing the continuous output space, we propose to generalize the cluster-prediction based method to general dense prediction tasks. This allows us to unify dense prediction tasks with the mask transformer framework. Remarkably, the resulting model PolyMaX demonstrates state-of-the-art performance on three benchmarks of NYUD-v2 dataset. We hope our simple yet effective design can inspire more research on exploiting mask transformers for more dense prediction tasks. Code and model will be made available.