Reassembly tasks play a fundamental role in many fields and multiple approaches exist to solve specific reassembly problems. In this context, we posit that a general unified model can effectively address them all, irrespective of the input data type (images, 3D, etc.). We introduce DiffAssemble, a Graph Neural Network (GNN)-based architecture that learns to solve reassembly tasks using a diffusion model formulation. Our method treats the elements of a set, whether pieces of 2D patch or 3D object fragments, as nodes of a spatial graph. Training is performed by introducing noise into the position and rotation of the elements and iteratively denoising them to reconstruct the coherent initial pose. DiffAssemble achieves state-of-the-art (SOTA) results in most 2D and 3D reassembly tasks and is the first learning-based approach that solves 2D puzzles for both rotation and translation. Furthermore, we highlight its remarkable reduction in run-time, performing 11 times faster than the quickest optimization-based method for puzzle solving. Code available at https://github.com/IIT-PAVIS/DiffAssemble

翻译：重组任务在众多领域中扮演着基础性角色，目前已有多种方法用于解决特定的重组问题。在此背景下，我们认为一个通用统一模型能够有效处理所有此类任务，无论输入数据类型（图像、三维数据等）如何。我们提出了DiffAssemble，一种基于图神经网络（GNN）的架构，利用扩散模型公式学习解决重组任务。该方法将集合中的元素——无论是二维图像片段还是三维物体碎片——视为空间图中的节点。训练过程中，通过对元素的位置和旋转引入噪声，并迭代去噪以重建连贯的初始姿态。DiffAssemble在大多数二维和三维重组任务中达到了最先进（SOTA）水平，并且是首个同时解决二维拼图旋转和平移问题的基于学习方法。此外，我们突显了其在运行时间上的显著减少，比最快的基于优化的拼图求解方法快11倍。代码开源于 https://github.com/IIT-PAVIS/DiffAssemble