Release code

.gitignore

@@ -0,0 +1,54 @@
+lightning_logs/
+results/
+Puzzle-Diff/
+/dataset/
+/datasets/*
+!/datasets/data_splits
+dataset_old_delete/
+.vscode/
+wandb/
+
+singularity/*.o*
+.ipynb_checkpoints
+
+
+singularity/*.o*
+exploration_analysis
+launch_experiment.sh
+singularity/*
+
+
+singularity/build/singularity/*
+singularity/*.o*
+.ipynb_checkpoints/
+
+*.sh
+*.pickle
+
+*.png
+*.ckpt
+*.sif
+*.pyc
+
+*.ipynb
+
+Mambaforge-Linux-x86_64.sh
+*.zip
+*.tar
+
+rotation_equivariant_test.ipynb
+.ipynb_checkpoints/rotation_equivariant_test-checkpoint.ipynb
+
+
+launch_experiment.sh
+
+
+singularity/*.sh
+Mambaforge-Linux-x86_64.sh
+*.zip
+*.tar
+*.ipynb
+launch_experiment.sh
+singularity/build/singularity
+
+!/page/*

.pre-commit-config.yaml

@@ -0,0 +1,43 @@
+# See https://pre-commit.com for more information
+# See https://pre-commit.com/hooks.html for more hooks
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+
+
+  - repo: https://github.com/psf/black
+    rev: 22.6.0
+    hooks:
+      - id: black
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+        name: isort (python)
+# repos:
+#   - repo: https://github.com/pre-commit/pre-commit-hooks
+#     rev: v4.3.0
+#     hooks:
+#       - id: check-added-large-files
+#       - id: trailing-whitespace
+
+#   - repo: local
+#     hooks:
+#       - id: isort
+#         name: isort
+#         stages: [commit]
+#         #language: system
+#         entry: isort
+#         types: [python]
+
+#       - id: black
+#         name: black
+#         stages: [commit]
+#         #language: system
+#         entry: black
+#         types: [python]

README.md

@@ -0,0 +1,51 @@
+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. We will release the code upon paper acceptance.
+- [Webpage](https://iit-pavis.github.io/DiffAssemble/)
+- [Code](https://github.com/IIT-PAVIS/DiffAssemble)
+- [Arxiv](soon)
+
+![Method](https://raw.githubusercontent.com/IIT-PAVIS/DiffAssemble/release/page/DiffAssemble.png)
+
+# Datasets
+
+## Puzzles
+
+- CelebA: [Link](https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)
+- WikiArt: [Link](https://paperswithcode.com/dataset/wikiart)
+
+## 3D
+- Breaking bad: [Link](https://breaking-bad-dataset.github.io/)
+
+# Environment
+- We provide the environment definition in `singularity/build/conda_env.yaml`
+- Singularity image is also available at [WIP]
+- Requirements:
+```
+  - pytorch==1.12.1
+  - cudatoolkit<=11.3.10
+  - pyg
+  - einops
+  - black
+  - pre-commit
+  - pytorch-lightning<1.8
+  - pip
+  - matplotlib
+  - wandb
+  - transformers
+  - timm
+  - kornia
+  - trimesh
+  - scipy
+```
+
+
+# Cite us!
+```
+@misc{giuliari2023positional,
+      title={Positional Diffusion: Ordering Unordered Sets with Diffusion Probabilistic Models},
+      author={Francesco Giuliari and Gianluca Scarpellini and Stuart James and Yiming Wang and Alessio {Del Bue}},
+      year={2023},
+      eprint={2303.11120},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV}
+}
+```

_config.yaml

@@ -0,0 +1,3 @@
+remote_theme: pages-themes/cayman@v0.2.0
+plugins:
+- jekyll-remote-theme