Solving Masked Jigsaw Puzzles with Diffusion Vision Transformers

TL;DR

The paper addresses the challenge of solving image and temporal jigsaw puzzles with many pieces and missing data. It introduces JPDVT, a unified approach that models puzzles as unordered sets of piece-content embeddings paired with positional encodings and solves them via a conditional diffusion denoising process conditioned on the visible content. By leveraging a diffusion-transformer architecture with adaptive normalization and relative positional embeddings, the method achieves state-of-the-art results on both image and video puzzles, including scenarios with missing pieces and large piece counts. The work demonstrates strong empirical performance, robust reconstruction capabilities, and applicability to downstream tasks like temporal super-resolution, offering a scalable, generalizable solution for reordering and inpainting in complex visual data.

Abstract

Solving image and video jigsaw puzzles poses the challenging task of rearranging image fragments or video frames from unordered sequences to restore meaningful images and video sequences. Existing approaches often hinge on discriminative models tasked with predicting either the absolute positions of puzzle elements or the permutation actions applied to the original data. Unfortunately, these methods face limitations in effectively solving puzzles with a large number of elements. In this paper, we propose JPDVT, an innovative approach that harnesses diffusion transformers to address this challenge. Specifically, we generate positional information for image patches or video frames, conditioned on their underlying visual content. This information is then employed to accurately assemble the puzzle pieces in their correct positions, even in scenarios involving missing pieces. Our method achieves state-of-the-art performance on several datasets.

Figures (13)

• Figure 1: Top: given unordered image fragments, some of them masked, we want to reconstruct the original image. Bottom: given shuffled video frames, some of them masked, we want to reconstruct the original video.
• Figure 2: Top. Each piece has a positional encoding and an embedding of its visual content. The forward diffusion step (highlighted in the red box), gradually adds noise to the positional encodings. The reverse generation reconstructs the positional encodings, conditioned by the provided visual content. Bottom. Samples of 1D and 2D positional encodings, for video frames and image tiles, both in their original form and with added noise.
• Figure 3: Left: Architecture employed to train the proposed diffusion transformer model. Image patches, accompanied by their corresponding position encoding tokens, undergo permutation in a consistent manner. During forward diffusion, noise is introduced to all positional encoding tokens and a subset of embedded patch tokens. The transformer is trained to execute a reverse diffusion process, mitigating the introduced noise. Right: Inference process. During inference, the observed patches are padded with Gaussian noise to fill the entire image. Linearly embedded patch tokens are then concatenated with positional encoding tokens initialized with Gaussian noise. The transformer model is designed to reconstruct the visual content of the missing patches and to determine the positions for all patches, simultaneously.
• Figure 4: Qualitative results for puzzles for Imagenet-1K dataset with 1, 2, and 3 missing pieces. Left: without cropping. Right: with cropping. First column shows the input puzzle, second column shows the sorted result (denoted as JPDVT-S), third column shows the sorted and inpainted result (denoted as JPDVT-SI), and last column shows the ground truth images.
• Figure 5: Video reshuffle results are showcased on various datasets, including 20 frames on MovingMNIST, 32 frames on CLEVRER, 32 frames on UCF101, and 32 frames on QST. Due to space limitations, only a subset of frames is presented in this illustration.