chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
+26
@@ -0,0 +1,26 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
|
||||
# Output files
|
||||
output.glb
|
||||
autotune_cache.json
|
||||
|
||||
# Virtual environment
|
||||
venv/
|
||||
.venv/
|
||||
env/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# OS
|
||||
.DS_Store
|
||||
Thumbs.db
|
||||
|
||||
tmp
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Tencent.
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,132 @@
|
||||
Pixal3D
|
||||
Copyright (c) 2026 Tencent.
|
||||
|
||||
Tencent is pleased to support the community by making Pixal3D available.
|
||||
|
||||
The open-source software and/or models included in this distribution may have
|
||||
been modified by Tencent (“Tencent Modifications”). All Tencent Modifications
|
||||
are Copyright (c) Tencent.
|
||||
|
||||
For the avoidance of doubt, Pixal3D refers solely to code, parameters, weights,
|
||||
and associated documentation made publicly available by Tencent in accordance
|
||||
with the MIT License (see the LICENSE file).
|
||||
|
||||
Responsible Use
|
||||
This model is released to support research, creative production, and practical 3D generation. Users are responsible for ensuring that their use complies with applicable laws, rights, consent requirements, and platform policies.
|
||||
|
||||
We do not support use of this model to create non-consensual sexual content, sexualized depictions of real people, content involving minors, harassment, impersonation, or privacy-violating applications.
|
||||
|
||||
The model is provided “as is” under the MIT License. This guidance does not modify the license terms but describes the intended and responsible use of the project.
|
||||
|
||||
================================================================================
|
||||
Third-Party Components
|
||||
================================================================================
|
||||
|
||||
This product includes software developed by third parties. The third-party
|
||||
components listed below remain licensed under their respective original terms.
|
||||
Pixal3D does not impose any additional restrictions beyond those specified in
|
||||
the original licenses of these third-party components. Users are required to
|
||||
comply with all applicable terms and conditions of the original licenses and
|
||||
to ensure that the use of these third-party components conforms to all
|
||||
relevant laws and regulations.
|
||||
|
||||
In case you believe there have been errors in the attribution below, you may
|
||||
submit the concerns to us for review and correction.
|
||||
|
||||
|
||||
================================================================================
|
||||
Open Model Licensed under Apache-2.0
|
||||
================================================================================
|
||||
|
||||
1. dinov2
|
||||
Copyright (c) 2025 dinov2 original author and authors
|
||||
|
||||
|
||||
================================================================================
|
||||
Open Source Software Licensed under the MIT License
|
||||
================================================================================
|
||||
|
||||
2. TRELLIS.2
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
3. Direct3D-S2
|
||||
Copyright (c) DreamTech.
|
||||
|
||||
4. MoGe
|
||||
Copyright (c) Microsoft Corporation.
|
||||
|
||||
|
||||
================================================================================
|
||||
Terms of the MIT License
|
||||
================================================================================
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
|
||||
|
||||
================================================================================
|
||||
Terms of the Apache License, Version 2.0
|
||||
================================================================================
|
||||
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
|
||||
|
||||
Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
|
||||
|
||||
Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
|
||||
|
||||
Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
|
||||
|
||||
Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
|
||||
|
||||
Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
|
||||
|
||||
Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
|
||||
|
||||
Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
@@ -0,0 +1,298 @@
|
||||
|
||||
<div align="center">
|
||||
|
||||
# Pixal3D: Pixel-Aligned 3D Generation from Images
|
||||
|
||||
<h3>SIGGRAPH 2026</h3>
|
||||
|
||||
<small>[Dong-Yang Li](https://ldyang694.github.io/)¹ · [Wang Zhao](https://thuzhaowang.github.io/)²* · [Yuxin Chen](https://orcid.org/0000-0002-7854-1072)² · [Wenbo Hu](https://wbhu.github.io/)² · [Meng-Hao Guo](https://menghaoguo.github.io/)¹ · [Fang-Lue Zhang](https://fanglue.github.io/)³ · [Ying Shan](https://www.linkedin.com/in/YingShanProfile)² · [Shi-Min Hu](https://cg.cs.tsinghua.edu.cn/shimin.htm)¹✉</small>
|
||||
|
||||
¹Tsinghua University (BNRist) ²Tencent ARC Lab ³Victoria University of Wellington
|
||||
|
||||
*Project lead ✉Corresponding author
|
||||
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<a href="https://ldyang694.github.io/projects/pixal3d/"><img src=https://img.shields.io/badge/Project%20Page-333399.svg?logo=googlehome height=22px></a>
|
||||
<a href="https://huggingface.co/spaces/TencentARC/Pixal3D"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20Demo-276cb4.svg height=22px></a>
|
||||
<a href="https://huggingface.co/TencentARC/Pixal3D"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20Models-d96902.svg height=22px></a>
|
||||
<a href="https://arxiv.org/abs/2605.10922"><img src=https://img.shields.io/badge/Arxiv-b5212f.svg?logo=arxiv height=22px></a>
|
||||
<a href="LICENSE"><img src=https://img.shields.io/badge/License-MIT-yellow.svg height=22px></a>
|
||||
</div>
|
||||
|
||||
<div align="center">
|
||||
<img src="assets/teaser.png" alt="Teaser image of Pixal3D"/>
|
||||
</div>
|
||||
|
||||
**Pixal3D** generates high-fidelity 3D assets from a single image. Unlike previous methods that loosely inject image features via attention, Pixal3D explicitly lifts pixel features into 3D through back-projection, establishing direct pixel-to-3D correspondences. This enables near-reconstruction-level fidelity with detailed geometry and PBR textures.
|
||||
|
||||
---
|
||||
|
||||
## ✨ News
|
||||
|
||||
- **May 2026**: Release training code and data preparation toolkit. 🔧
|
||||
- **May 2026**: Release the improved version based on [Trellis.2](https://github.com/microsoft/TRELLIS.2) backbone. 💪
|
||||
- **May 2026**: Release inference code and online demo. 🤗
|
||||
- **Apr 2026**: Our paper is accepted to SIGGRAPH 2026! 🎉
|
||||
|
||||
## 📌 Branches
|
||||
|
||||
| Branch | Description |
|
||||
|--------|-------------|
|
||||
| `main` | **Latest version** — improved implementation based on [Trellis.2](https://github.com/microsoft/TRELLIS.2) backbone with better performance. |
|
||||
| `paper` | **Paper version** — original implementation based on [Direct3D-S2](https://github.com/DreamTechAI/Direct3D-S2), corresponding to results reported in our SIGGRAPH 2026 paper. |
|
||||
|
||||
> If you want to reproduce the results in our paper, please switch to the `paper` branch.
|
||||
|
||||
## 🎮 Try It Online
|
||||
|
||||
You can try Pixal3D directly in your browser without any installation via our Hugging Face Gradio demo:
|
||||
|
||||
👉 [**Launch Demo**](https://huggingface.co/spaces/TencentARC/Pixal3D)
|
||||
|
||||
## 🚀 Getting Started
|
||||
|
||||
### Installation
|
||||
|
||||
#### Step 1: Follow TRELLIS.2 Installation
|
||||
|
||||
Please first follow the installation guide of [TRELLIS.2](https://github.com/microsoft/TRELLIS.2) to set up the base environment.
|
||||
|
||||
#### Step 2: Install Additional Dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
#### Step 3: Install natten
|
||||
|
||||
```bash
|
||||
NATTEN_CUDA_ARCH="xx" NATTEN_N_WORKERS=xx pip install natten==0.21.0 --no-build-isolation
|
||||
```
|
||||
|
||||
Please replace `xx` with the CUDA architecture and the number of build workers suitable for your machine.
|
||||
|
||||
#### Step 4: Install utils3d
|
||||
|
||||
```bash
|
||||
pip install https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl
|
||||
```
|
||||
|
||||
> **Note**: `requirements-hfdemo.txt` is for the Hugging Face Spaces demo (H-series GPU architecture) and may not be compatible with other architectures.
|
||||
|
||||
### Usage
|
||||
|
||||
#### Inference
|
||||
|
||||
Generate a GLB mesh from a single image:
|
||||
|
||||
```bash
|
||||
python inference.py --image assets/images/0_img.png --output ./output.glb
|
||||
```
|
||||
|
||||
**Low-VRAM mode** (reduces peak VRAM by loading models on-demand):
|
||||
|
||||
```bash
|
||||
python inference.py --image assets/images/0_img.png --output ./output.glb --low_vram
|
||||
```
|
||||
|
||||
By default, the pipeline resolution is **1536** (standard mode) or **1024** (low-VRAM mode). You can override this with `--resolution`:
|
||||
|
||||
```bash
|
||||
# Force 1536 even in low-VRAM mode
|
||||
python inference.py --image assets/images/0_img.png --output ./output.glb --low_vram --resolution 1536
|
||||
|
||||
# Force 1024 in standard mode
|
||||
python inference.py --image assets/images/0_img.png --output ./output.glb --resolution 1024
|
||||
```
|
||||
|
||||
**Tip**: If you don't have `flash_attn` installed, you can use PyTorch's built-in SDPA backend instead:
|
||||
> ```bash
|
||||
> ATTN_BACKEND=sdpa python inference.py --image assets/images/0_img.png --output ./output.glb --low_vram
|
||||
> ```
|
||||
|
||||
### Web Demo
|
||||
|
||||
We provide a Gradio web demo for Pixal3D, which allows you to generate 3D meshes from images interactively.
|
||||
|
||||
```bash
|
||||
python app.py
|
||||
```
|
||||
|
||||
Low-VRAM mode is also available for the web demo. The frontend default resolution will automatically switch to 1024 in low-VRAM mode (1536 otherwise), but can be changed manually in the UI.
|
||||
|
||||
```bash
|
||||
python app.py --low_vram
|
||||
# or via environment variable:
|
||||
LOW_VRAM=1 python app.py
|
||||
```
|
||||
## 🔧 Training
|
||||
|
||||
We provide the full training codebase for reproducing Pixal3D from scratch.
|
||||
|
||||
### Data Preparation
|
||||
|
||||
Prepare view-aligned O-Voxel data and rendered condition images by following the data toolkit instructions:
|
||||
|
||||
> 📂 **[data_toolkit/README.md](data_toolkit/README.md)**
|
||||
|
||||
### Overview
|
||||
|
||||
Pixal3D is trained as a three-stage cascade, each progressively increasing resolution:
|
||||
|
||||
| Stage | Model | Resolutions | Config Prefix |
|
||||
|-------|-------|-------------|---------------|
|
||||
| 1 | Sparse Structure | 32 → 64 | `ss_flow_img_dit_*_proj_finetune` |
|
||||
| 2 | Shape | 256 → 512 → 1024 | `slat_flow_img2shape_*_proj_finetune` |
|
||||
| 3 | Texture | 256 → 512 → 1024 | `slat_flow_imgshape2tex_*_proj_finetune` |
|
||||
|
||||
All stages use **pixel-aligned projection conditioning** and **view-aligned latents** (2 views by default). Within each stage, start from the lowest resolution and progressively fine-tune to higher resolutions by setting `finetune_ckpt` in the config.
|
||||
|
||||
### Quick Start
|
||||
|
||||
```sh
|
||||
python train.py \
|
||||
--config <CONFIG_JSON> \
|
||||
--output_dir <OUTPUT_DIR> \
|
||||
--data_dir '<DATA_DIR_JSON>'
|
||||
```
|
||||
|
||||
`--data_dir` is a JSON string describing the dataset layout. Different stages require different keys:
|
||||
|
||||
| Stage | Required keys |
|
||||
|-------|---------------|
|
||||
| Sparse Structure | `base`, `ss_latent`, `render_cond` |
|
||||
| Shape | `base`, `shape_latent`, `render_cond` |
|
||||
| Texture | `base`, `shape_latent`, `pbr_latent`, `render_cond` |
|
||||
|
||||
### Example: Training All Three Stages
|
||||
|
||||
Below we show the full training sequence using ObjaverseXL as an example. Each higher-resolution step requires updating `finetune_ckpt` in its config JSON to point to the previous checkpoint.
|
||||
|
||||
<details>
|
||||
<summary><b>Stage 1: Sparse Structure (32 → 64)</b></summary>
|
||||
|
||||
```sh
|
||||
# Resolution 32
|
||||
python train.py \
|
||||
--config configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune.json \
|
||||
--output_dir results/ss_32 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "ss_latent": "datasets/ObjaverseXL_sketchfab/ss_latents/ss_enc_conv3d_16l8_fp16_64_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
|
||||
# Resolution 64 (set finetune_ckpt → results/ss_32 checkpoint)
|
||||
python train.py \
|
||||
--config configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune_ft64.json \
|
||||
--output_dir results/ss_ft64 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "ss_latent": "datasets/ObjaverseXL_sketchfab/ss_latents/ss_enc_conv3d_16l8_fp16_64_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
```
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Stage 2: Shape (256 → 512 → 1024)</b></summary>
|
||||
|
||||
```sh
|
||||
# Resolution 256
|
||||
python train.py \
|
||||
--config configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune.json \
|
||||
--output_dir results/shape_256 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_256_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
|
||||
# Resolution 512
|
||||
python train.py \
|
||||
--config configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune_ft512.json \
|
||||
--output_dir results/shape_ft512 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_512_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
|
||||
# Resolution 1024
|
||||
python train.py \
|
||||
--config configs/gen/slat_flow_img2shape_dit_1_3B_512_bf16_proj_finetune_ft1024.json \
|
||||
--output_dir results/shape_ft1024 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_1024_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
```
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary><b>Stage 3: Texture (256 → 512 → 1024)</b></summary>
|
||||
|
||||
```sh
|
||||
# Resolution 256
|
||||
python train.py \
|
||||
--config configs/gen/slat_flow_imgshape2tex_dit_1_3B_256_bf16_proj_finetune.json \
|
||||
--output_dir results/tex_256 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_256_view", "pbr_latent": "datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_256_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
|
||||
# Resolution 512
|
||||
python train.py \
|
||||
--config configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune.json \
|
||||
--output_dir results/tex_512 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_512_view", "pbr_latent": "datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_512_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
|
||||
# Resolution 1024
|
||||
python train.py \
|
||||
--config configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune_ft1024.json \
|
||||
--output_dir results/tex_ft1024 \
|
||||
--data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_1024_view", "pbr_latent": "datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_1024_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
|
||||
```
|
||||
</details>
|
||||
|
||||
### Additional Options
|
||||
|
||||
<details>
|
||||
<summary><b>All command-line arguments</b></summary>
|
||||
|
||||
| Argument | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `--config` | Config JSON path | *required* |
|
||||
| `--output_dir` | Output directory | *required* |
|
||||
| `--data_dir` | Dataset JSON string | `./data/` |
|
||||
| `--load_dir` | Checkpoint load directory | `output_dir` |
|
||||
| `--ckpt` | Resume from step | `latest` |
|
||||
| `--auto_retry` | Retries on failure | `3` |
|
||||
| `--tryrun` | Dry run | `false` |
|
||||
| `--profile` | Profiling | `false` |
|
||||
| `--num_nodes` | Number of nodes | `1` |
|
||||
| `--node_rank` | Current node rank | `0` |
|
||||
| `--num_gpus` | GPUs per node | all |
|
||||
| `--master_addr` | Master address | `localhost` |
|
||||
| `--master_port` | Master port | `12666` |
|
||||
| `--use_wandb` | Enable W&B logging | `false` |
|
||||
| `--wandb_project` | W&B project | `trellis2-training` |
|
||||
| `--wandb_name` | W&B run name | basename of `output_dir` |
|
||||
| `--wandb_id` | W&B run ID (resume) | — |
|
||||
|
||||
</details>
|
||||
|
||||
## 🌐 Community Projects
|
||||
|
||||
We thank the community for building extensions and deployment guides for Pixal3D!
|
||||
|
||||
- [Pixal3D-ComfyUI](https://github.com/Saganaki22/Pixal3D-ComfyUI) — ComfyUI integration with deployment guides for Windows, WSL, and more.
|
||||
|
||||
## 🤗 Acknowledgements
|
||||
|
||||
This project is heavily built upon [Trellis.2](https://github.com/microsoft/TRELLIS.2) and [Direct3D-S2](https://github.com/DreamTechAI/Direct3D-S2). We sincerely thank the authors for their outstanding work on scalable 3D generation , which serves as the foundation of our codebase and model architecture.
|
||||
|
||||
We also thank the following repos for their great contributions:
|
||||
|
||||
- [Direct3D-S2](https://github.com/DreamTechAI/Direct3D-S2)
|
||||
- [Trellis](https://github.com/microsoft/TRELLIS)
|
||||
- [Trellis.2](https://github.com/microsoft/TRELLIS.2)
|
||||
|
||||
## 📄 Citation
|
||||
|
||||
If you find this work useful, please consider citing:
|
||||
|
||||
```bibtex
|
||||
@article{li2026pixal3d,
|
||||
title={Pixal3D: Pixel-Aligned 3D Generation from Images},
|
||||
author={Li, Dong-Yang and Zhao, Wang and Chen, Yuxin and Hu, Wenbo and Guo, Meng-Hao and Zhang, Fang-Lue and Shan, Ying and Hu, Shi-Min},
|
||||
journal={arXiv preprint arXiv:2605.10922},
|
||||
year={2026}
|
||||
}
|
||||
```
|
||||
|
||||
## 📜 License
|
||||
|
||||
This project is released under the [MIT License](LICENSE). The third-party components included in this project remain licensed under their respective original terms; see [NOTICE](NOTICE) for the full list of dependencies and their licenses.
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
# WeHub 来源说明
|
||||
|
||||
- 原始项目：`TencentARC/Pixal3D`
|
||||
- 原始仓库：https://github.com/TencentARC/Pixal3D
|
||||
- 导入方式：上游默认分支的最新快照
|
||||
- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
|
||||
- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
|
||||
@@ -0,0 +1,559 @@
|
||||
import os
|
||||
import subprocess
|
||||
import argparse
|
||||
import math
|
||||
import time
|
||||
import shutil
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from datetime import datetime
|
||||
from typing import *
|
||||
from PIL import Image
|
||||
|
||||
import threading
|
||||
try:
|
||||
import nest_asyncio
|
||||
nest_asyncio.apply()
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Lock for model initialization
|
||||
init_lock = threading.Lock()
|
||||
|
||||
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
||||
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
|
||||
os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
|
||||
os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
|
||||
|
||||
import spaces
|
||||
from gradio import Server
|
||||
from gradio.data_classes import FileData
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from pixal3d.modules.sparse import SparseTensor
|
||||
from pixal3d.pipelines import Pixal3DImageTo3DPipeline
|
||||
from pixal3d.renderers import EnvMap
|
||||
from pixal3d.utils import render_utils
|
||||
import o_voxel
|
||||
|
||||
# ============================================================================
|
||||
# Constants & Defaults
|
||||
# ============================================================================
|
||||
|
||||
MAX_SEED = np.iinfo(np.int32).max
|
||||
TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
|
||||
os.makedirs(TMP_DIR, exist_ok=True)
|
||||
|
||||
MODES = [
|
||||
{"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
|
||||
{"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
|
||||
{"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
|
||||
{"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
|
||||
{"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
|
||||
{"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
|
||||
]
|
||||
STEPS = 8
|
||||
|
||||
# Cascade parameters
|
||||
CASCADE_LR_RESOLUTION = 512
|
||||
CASCADE_MAX_NUM_TOKENS = 49152
|
||||
|
||||
# MoGe defaults
|
||||
MOGE_MODEL_NAME = "Ruicheng/moge-2-vitl"
|
||||
WILD_MESH_SCALE = 1.0
|
||||
WILD_EXTEND_PIXEL = 0
|
||||
WILD_IMAGE_RESOLUTION = 512
|
||||
|
||||
# Image Cond Model configs
|
||||
IMAGE_COND_CONFIGS = {
|
||||
"ss": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 16,
|
||||
},
|
||||
"shape_512": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 32,
|
||||
"use_naf_upsample": True,
|
||||
"naf_target_size": 512,
|
||||
},
|
||||
"shape_1024": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 1024,
|
||||
"grid_resolution": 64,
|
||||
"use_naf_upsample": True,
|
||||
"naf_target_size": 512,
|
||||
},
|
||||
"tex_1024": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 1024,
|
||||
"grid_resolution": 64,
|
||||
"use_naf_upsample": True,
|
||||
"naf_target_size": 1024,
|
||||
},
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Model Loading
|
||||
# ============================================================================
|
||||
|
||||
def build_image_cond_model(config: dict):
|
||||
from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor
|
||||
model = DinoV3ProjFeatureExtractor(**config)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME):
|
||||
from moge.model.v2 import MoGeModel
|
||||
moge_model = MoGeModel.from_pretrained(model_name).to(device)
|
||||
moge_model.eval()
|
||||
return moge_model
|
||||
|
||||
# Global instances (lazy loaded or loaded at start)
|
||||
pipeline = None
|
||||
moge_model = None
|
||||
envmap = None
|
||||
LOW_VRAM = os.environ.get("LOW_VRAM", "0") == "1"
|
||||
|
||||
def init_models():
|
||||
global pipeline, moge_model, envmap
|
||||
with init_lock:
|
||||
if pipeline is not None:
|
||||
return
|
||||
|
||||
# GPU / CUDA Diagnostics (runs when GPU is allocated)
|
||||
import subprocess as _sp
|
||||
print("=" * 60)
|
||||
print("[Diagnostics] PyTorch version:", torch.__version__)
|
||||
print("[Diagnostics] CUDA available:", torch.cuda.is_available())
|
||||
if torch.cuda.is_available():
|
||||
print("[Diagnostics] CUDA version:", torch.version.cuda)
|
||||
print("[Diagnostics] cuDNN version:", torch.backends.cudnn.version())
|
||||
for i in range(torch.cuda.device_count()):
|
||||
name = torch.cuda.get_device_name(i)
|
||||
cap = torch.cuda.get_device_capability(i)
|
||||
mem = torch.cuda.get_device_properties(i).total_memory / 1024**3
|
||||
print(f"[Diagnostics] GPU {i}: {name}, sm_{cap[0]}{cap[1]}, {mem:.1f} GB")
|
||||
try:
|
||||
res = _sp.run(["nvidia-smi", "--query-gpu=name,compute_cap,memory.total", "--format=csv,noheader"], capture_output=True, text=True, timeout=10)
|
||||
print("[Diagnostics] nvidia-smi:", res.stdout.strip())
|
||||
except Exception as e:
|
||||
print(f"[Diagnostics] nvidia-smi failed: {e}")
|
||||
print("=" * 60)
|
||||
|
||||
model_path = "TencentARC/Pixal3D"
|
||||
print(f"[Pipeline] Loading from {model_path}...")
|
||||
pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
|
||||
|
||||
print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
|
||||
pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
|
||||
pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
|
||||
pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
|
||||
pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
|
||||
|
||||
if LOW_VRAM:
|
||||
# Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage.
|
||||
print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...")
|
||||
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
||||
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
||||
m = getattr(pipeline, attr, None)
|
||||
if m is not None and getattr(m, 'use_naf_upsample', False):
|
||||
m._load_naf()
|
||||
pipeline._device = torch.device("cuda")
|
||||
pipeline.low_vram = True
|
||||
print("[Pipeline] Low-VRAM mode enabled.")
|
||||
else:
|
||||
# Standard mode: all models loaded to GPU at once.
|
||||
pipeline.low_vram = False
|
||||
pipeline.cuda()
|
||||
pipeline.image_cond_model_ss.cuda()
|
||||
pipeline.image_cond_model_shape_512.cuda()
|
||||
pipeline.image_cond_model_shape_1024.cuda()
|
||||
pipeline.image_cond_model_tex_1024.cuda()
|
||||
print("[NAF] Pre-loading NAF upsampler model...")
|
||||
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
||||
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
||||
m = getattr(pipeline, attr, None)
|
||||
if m is not None and getattr(m, 'use_naf_upsample', False):
|
||||
m._load_naf()
|
||||
|
||||
print("[MoGe-2] Loading model for camera estimation...")
|
||||
if LOW_VRAM:
|
||||
# Low-VRAM: load MoGe to CPU, move to GPU on-demand per request.
|
||||
moge_model = load_moge_model(device="cpu")
|
||||
print("[MoGe-2] Low-VRAM mode: MoGe stays on CPU, loaded to GPU on-demand.")
|
||||
else:
|
||||
moge_model = load_moge_model(device="cuda")
|
||||
|
||||
print("[EnvMap] Loading environment maps...")
|
||||
_base = os.path.dirname(os.path.abspath(__file__))
|
||||
_envmap_device = 'cpu' if LOW_VRAM else 'cuda'
|
||||
envmap = {
|
||||
'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/forest.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
|
||||
'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/sunset.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
|
||||
'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/courtyard.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Utilities
|
||||
# ============================================================================
|
||||
|
||||
def compute_f_pixels(camera_angle_x: float, resolution: int) -> float:
|
||||
focal_length = 16.0 / torch.tan(torch.tensor(camera_angle_x / 2.0))
|
||||
f_pixels = focal_length * resolution / 32.0
|
||||
return float(f_pixels.item())
|
||||
|
||||
def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, image_resolution):
|
||||
rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
|
||||
gp = grid_point.to(torch.float32) @ rotation_matrix.T
|
||||
gp = gp / mesh_scale / 2
|
||||
xw, yw, zw = gp[0].item(), gp[1].item(), gp[2].item()
|
||||
xt, yt = float(target_point[0].item()), float(target_point[1].item())
|
||||
f_pixels = compute_f_pixels(camera_angle_x, image_resolution)
|
||||
x_ndc = xt - image_resolution / 2.0
|
||||
y_ndc = -(yt - image_resolution / 2.0)
|
||||
distance_x = f_pixels * xw / x_ndc - yw
|
||||
return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)}
|
||||
|
||||
def get_camera_params_wild_moge(image_path, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512):
|
||||
pil_image = Image.open(image_path).convert("RGB")
|
||||
width, height = pil_image.size
|
||||
image_np = np.array(pil_image).astype(np.float32) / 255.0
|
||||
image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).to(device)
|
||||
if LOW_VRAM:
|
||||
moge_model.to(device)
|
||||
with torch.no_grad():
|
||||
output = moge_model.infer(image_tensor)
|
||||
if LOW_VRAM:
|
||||
moge_model.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
intrinsics = output["intrinsics"].squeeze().cpu().numpy()
|
||||
fx_normalized = intrinsics[0, 0]
|
||||
fx = fx_normalized * width
|
||||
camera_angle_x = 2 * math.atan(width / (2 * fx))
|
||||
|
||||
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
||||
distance = distance_from_fov(
|
||||
camera_angle_x, grid_point,
|
||||
torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
|
||||
mesh_scale, image_resolution
|
||||
)["distance_from_x"]
|
||||
return {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
|
||||
|
||||
def pack_state(shape_slat, tex_slat, res):
|
||||
state_data = {
|
||||
'shape_slat_feats': shape_slat.feats.cpu().numpy(),
|
||||
'tex_slat_feats': tex_slat.feats.cpu().numpy(),
|
||||
'coords': shape_slat.coords.cpu().numpy(),
|
||||
'res': res,
|
||||
}
|
||||
import random
|
||||
state_path = os.path.join(TMP_DIR, f"state_{int(time.time()*1000)}_{random.randint(0,9999):04d}.npz")
|
||||
np.savez_compressed(state_path, **state_data)
|
||||
return state_path
|
||||
|
||||
def unpack_state(state_path):
|
||||
data = np.load(state_path)
|
||||
shape_slat = SparseTensor(
|
||||
feats=torch.from_numpy(data['shape_slat_feats']).cuda(),
|
||||
coords=torch.from_numpy(data['coords']).cuda(),
|
||||
)
|
||||
tex_slat = shape_slat.replace(torch.from_numpy(data['tex_slat_feats']).cuda())
|
||||
return shape_slat, tex_slat, int(data['res'])
|
||||
|
||||
# ============================================================================
|
||||
# Progress Tracking (file-based, cross-process safe for @spaces.GPU)
|
||||
# ============================================================================
|
||||
|
||||
import asyncio
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi import Request
|
||||
|
||||
PROGRESS_DIR = os.path.join(TMP_DIR, '_progress')
|
||||
os.makedirs(PROGRESS_DIR, exist_ok=True)
|
||||
|
||||
_thread_local = threading.local()
|
||||
|
||||
def _progress_file(session_id: str) -> str:
|
||||
"""Return path to a session's progress JSON file."""
|
||||
return os.path.join(PROGRESS_DIR, f"{session_id}.json")
|
||||
|
||||
def _reset_progress(session_id: str):
|
||||
_thread_local.active_session = session_id
|
||||
_write_progress_file(session_id, {"stage": "Initializing...", "step": 0, "total": 0, "done": False})
|
||||
|
||||
def _update_progress(stage: str, step: int, total: int):
|
||||
session_id = getattr(_thread_local, 'active_session', '')
|
||||
if session_id:
|
||||
_write_progress_file(session_id, {"stage": stage, "step": step, "total": total, "done": False})
|
||||
|
||||
def _finish_progress():
|
||||
session_id = getattr(_thread_local, 'active_session', '')
|
||||
if session_id:
|
||||
_write_progress_file(session_id, {"done": True})
|
||||
|
||||
def _write_progress_file(session_id: str, data: dict):
|
||||
"""Atomically write progress JSON to a file (cross-process safe)."""
|
||||
path = _progress_file(session_id)
|
||||
tmp_path = path + ".tmp"
|
||||
try:
|
||||
with open(tmp_path, 'w') as f:
|
||||
json.dump(data, f)
|
||||
os.replace(tmp_path, path) # atomic on POSIX
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Monkey-patch tqdm to intercept progress
|
||||
import tqdm as _tqdm_module
|
||||
|
||||
_original_tqdm = _tqdm_module.tqdm
|
||||
|
||||
class _TqdmProgressInterceptor(_original_tqdm):
|
||||
"""Wraps tqdm to push progress updates to SSE."""
|
||||
def __init__(self, *args, **kwargs):
|
||||
self._stage_desc = kwargs.get('desc', 'Processing')
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def set_description(self, desc=None, refresh=True):
|
||||
self._stage_desc = desc or 'Processing'
|
||||
super().set_description(desc, refresh)
|
||||
|
||||
def update(self, n=1):
|
||||
super().update(n)
|
||||
_update_progress(self._stage_desc, self.n, self.total or 0)
|
||||
|
||||
# Patch tqdm globally
|
||||
_tqdm_module.tqdm = _TqdmProgressInterceptor
|
||||
# Also patch the direct import in the sampler module and render_utils
|
||||
import pixal3d.pipelines.samplers.flow_euler as _fe_module
|
||||
_fe_module.tqdm = _TqdmProgressInterceptor
|
||||
import pixal3d.utils.render_utils as _ru_module
|
||||
_ru_module.tqdm = _TqdmProgressInterceptor
|
||||
import o_voxel.postprocess as _ovp_module
|
||||
_ovp_module.tqdm = _TqdmProgressInterceptor
|
||||
|
||||
# ============================================================================
|
||||
# API Implementation
|
||||
# ============================================================================
|
||||
|
||||
app = Server()
|
||||
|
||||
@app.get("/")
|
||||
async def homepage():
|
||||
html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
|
||||
with open(html_path, "r", encoding="utf-8") as f:
|
||||
return HTMLResponse(content=f.read())
|
||||
|
||||
@app.get("/app_config")
|
||||
async def get_config():
|
||||
"""Return server configuration for frontend (e.g. LOW_VRAM mode)."""
|
||||
return JSONResponse({"low_vram": LOW_VRAM})
|
||||
|
||||
@app.get("/progress")
|
||||
async def progress_poll(request: Request):
|
||||
"""Polling endpoint for real-time progress updates during generation."""
|
||||
session_id = request.query_params.get("session_id", "")
|
||||
path = _progress_file(session_id)
|
||||
try:
|
||||
with open(path, 'r') as f:
|
||||
data = json.load(f)
|
||||
return JSONResponse(data)
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
return JSONResponse({"stage": "Waiting...", "step": 0, "total": 0, "done": False})
|
||||
|
||||
@app.api()
|
||||
@spaces.GPU(duration=30)
|
||||
def preprocess(image: FileData) -> FileData:
|
||||
init_models()
|
||||
img = Image.open(image["path"])
|
||||
processed = pipeline.preprocess_image(img)
|
||||
out_path = os.path.join(TMP_DIR, f"preprocessed_{int(time.time()*1000)}.png")
|
||||
processed.save(out_path)
|
||||
return FileData(path=out_path)
|
||||
|
||||
@app.api()
|
||||
@spaces.GPU(duration=120)
|
||||
def generate_3d(
|
||||
image: FileData,
|
||||
seed: int,
|
||||
resolution: int,
|
||||
ss_guidance_strength: float = 7.5,
|
||||
ss_guidance_rescale: float = 0.7,
|
||||
ss_sampling_steps: int = 12,
|
||||
ss_rescale_t: float = 5.0,
|
||||
shape_slat_guidance_strength: float = 7.5,
|
||||
shape_slat_guidance_rescale: float = 0.5,
|
||||
shape_slat_sampling_steps: int = 12,
|
||||
shape_slat_rescale_t: float = 3.0,
|
||||
tex_slat_guidance_strength: float = 1.0,
|
||||
tex_slat_guidance_rescale: float = 0.0,
|
||||
tex_slat_sampling_steps: int = 12,
|
||||
tex_slat_rescale_t: float = 3.0,
|
||||
manual_fov: float = -1.0,
|
||||
fov_unit: str = "deg",
|
||||
session_id: str = "",
|
||||
) -> Dict:
|
||||
init_models()
|
||||
_reset_progress(session_id)
|
||||
_update_progress("Preprocessing & Camera Estimation", 0, 1)
|
||||
|
||||
torch.manual_seed(seed)
|
||||
hr_resolution = int(resolution)
|
||||
|
||||
img = Image.open(image["path"])
|
||||
# Image is already preprocessed by /preprocess endpoint, use directly
|
||||
image_preprocessed = img
|
||||
temp_processed_path = os.path.join(TMP_DIR, f"temp_proc_{session_id[:8]}_{int(time.time()*1000)}.png")
|
||||
image_preprocessed.save(temp_processed_path)
|
||||
|
||||
if manual_fov > 0:
|
||||
# Convert to radians based on unit
|
||||
if fov_unit == "rad":
|
||||
camera_angle_x = float(manual_fov)
|
||||
fov_deg = math.degrees(manual_fov)
|
||||
else:
|
||||
camera_angle_x = math.radians(manual_fov)
|
||||
fov_deg = float(manual_fov)
|
||||
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
||||
distance = distance_from_fov(
|
||||
camera_angle_x, grid_point,
|
||||
torch.tensor([0 - WILD_EXTEND_PIXEL, WILD_IMAGE_RESOLUTION - 1 + WILD_EXTEND_PIXEL]),
|
||||
WILD_MESH_SCALE, WILD_IMAGE_RESOLUTION
|
||||
)["distance_from_x"]
|
||||
camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': WILD_MESH_SCALE}
|
||||
print(f"[Camera] Using manual FOV: {fov_deg:.2f}° ({camera_angle_x:.4f} rad), distance: {distance:.4f}")
|
||||
else:
|
||||
camera_params = get_camera_params_wild_moge(
|
||||
temp_processed_path, device="cuda",
|
||||
mesh_scale=WILD_MESH_SCALE, extend_pixel=WILD_EXTEND_PIXEL,
|
||||
image_resolution=WILD_IMAGE_RESOLUTION,
|
||||
)
|
||||
_update_progress("Preprocessing & Camera Estimation", 1, 1)
|
||||
|
||||
ss_sampler_override = {"steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength,
|
||||
"guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t}
|
||||
shape_sampler_override = {"steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength,
|
||||
"guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t}
|
||||
tex_sampler_override = {"steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength,
|
||||
"guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t}
|
||||
|
||||
pipeline_type = f"{hr_resolution}_cascade"
|
||||
mesh_list, (shape_slat, tex_slat, res) = pipeline.run(
|
||||
image_preprocessed,
|
||||
camera_params=camera_params,
|
||||
seed=seed,
|
||||
sparse_structure_sampler_params=ss_sampler_override,
|
||||
shape_slat_sampler_params=shape_sampler_override,
|
||||
tex_slat_sampler_params=tex_sampler_override,
|
||||
preprocess_image=False,
|
||||
return_latent=True,
|
||||
pipeline_type=pipeline_type,
|
||||
max_num_tokens=CASCADE_MAX_NUM_TOKENS,
|
||||
)
|
||||
|
||||
mesh = mesh_list[0]
|
||||
state_path = pack_state(shape_slat, tex_slat, res)
|
||||
|
||||
_update_progress("Rendering views", 0, 1)
|
||||
mesh.simplify(16777216)
|
||||
cam_dist = camera_params['distance']
|
||||
near = max(0.01, cam_dist - 2.0)
|
||||
far = cam_dist + 10.0
|
||||
if LOW_VRAM:
|
||||
for v in envmap.values():
|
||||
v.image = v.image.cuda()
|
||||
if hasattr(v, '_nvdiffrec_envlight'):
|
||||
del v._nvdiffrec_envlight
|
||||
renders = render_utils.render_proj_aligned_video(
|
||||
mesh, camera_angle_x=camera_params['camera_angle_x'],
|
||||
distance=cam_dist, resolution=1024,
|
||||
num_frames=STEPS, envmap=envmap,
|
||||
near=near, far=far,
|
||||
)
|
||||
if LOW_VRAM:
|
||||
for v in envmap.values():
|
||||
if hasattr(v, '_nvdiffrec_envlight'):
|
||||
del v._nvdiffrec_envlight
|
||||
v.image = v.image.cpu()
|
||||
torch.cuda.empty_cache()
|
||||
_update_progress("Rendering views", 1, 1)
|
||||
|
||||
# Save renders and return paths
|
||||
render_files = {}
|
||||
for mode_key, frames in renders.items():
|
||||
mode_files = []
|
||||
for i, frame in enumerate(frames):
|
||||
p = os.path.abspath(os.path.join(TMP_DIR, f"render_{mode_key}_{i}_{int(time.time()*1000)}.jpg"))
|
||||
Image.fromarray(frame).save(p, quality=85)
|
||||
mode_files.append(FileData(path=p))
|
||||
render_files[mode_key] = mode_files
|
||||
|
||||
_finish_progress()
|
||||
return {
|
||||
"render_paths": render_files,
|
||||
"state_path": os.path.abspath(state_path),
|
||||
"camera_angle_x": camera_params['camera_angle_x'],
|
||||
"distance": camera_params['distance'],
|
||||
}
|
||||
|
||||
@app.api()
|
||||
@spaces.GPU(duration=240)
|
||||
def extract_glb_api(state_path: str, decimation_target: int, texture_size: int, session_id: str = "") -> FileData:
|
||||
init_models()
|
||||
_reset_progress(session_id)
|
||||
_update_progress("Decoding latent", 0, 1)
|
||||
|
||||
shape_slat, tex_slat, res = unpack_state(state_path)
|
||||
mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
|
||||
_update_progress("Decoding latent", 1, 1)
|
||||
|
||||
glb = o_voxel.postprocess.to_glb(
|
||||
vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs,
|
||||
coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout,
|
||||
grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
decimation_target=decimation_target, texture_size=texture_size,
|
||||
remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
|
||||
)
|
||||
rot = np.array([
|
||||
[-1, 0, 0, 0],
|
||||
[ 0, 0, -1, 0],
|
||||
[ 0, -1, 0, 0],
|
||||
[ 0, 0, 0, 1],
|
||||
], dtype=np.float64)
|
||||
glb.apply_transform(rot)
|
||||
|
||||
out_glb = os.path.join(TMP_DIR, f"result_{int(time.time()*1000)}.glb")
|
||||
glb.export(out_glb, extension_webp=True)
|
||||
_finish_progress()
|
||||
return FileData(path=out_glb)
|
||||
|
||||
# Mount assets and tmp for direct access
|
||||
app.mount("/assets", StaticFiles(directory="assets"), name="assets")
|
||||
app.mount("/tmp", StaticFiles(directory=TMP_DIR), name="tmp")
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
parser = argparse.ArgumentParser(description="Pixal3D Demo Server")
|
||||
parser.add_argument("--low_vram", action="store_true",
|
||||
help="Enable low-VRAM mode: models lazy-load to GPU per stage.")
|
||||
args, remaining = parser.parse_known_args()
|
||||
if args.low_vram:
|
||||
LOW_VRAM = True
|
||||
|
||||
# Re-install utils3d as in original app.py
|
||||
subprocess.run([
|
||||
sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
|
||||
"https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl"
|
||||
], check=True)
|
||||
|
||||
# Pre-initialize models before launching the server
|
||||
init_models()
|
||||
|
||||
app.launch(show_error=True, share=True)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 4.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 4.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 9.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 10 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 9.1 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 6.8 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 6.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.2 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.9 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.8 KiB |
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,15 @@
|
||||
All HDRIs are licensed as CC0.
|
||||
|
||||
These were created by Greg Zaal (Poly Haven https://polyhaven.com).
|
||||
Originals used for each HDRI:
|
||||
- City: https://polyhaven.com/a/portland_landing_pad
|
||||
- Courtyard: https://polyhaven.com/a/courtyard
|
||||
- Forest: https://polyhaven.com/a/ninomaru_teien
|
||||
- Interior: https://polyhaven.com/a/hotel_room
|
||||
- Night: Probably https://polyhaven.com/a/moonless_golf
|
||||
- Studio: Probably https://polyhaven.com/a/studio_small_01
|
||||
- Sunrise: https://polyhaven.com/a/spruit_sunrise
|
||||
- Sunset: https://polyhaven.com/a/venice_sunset
|
||||
|
||||
1K resolution of each was taken, and compressed with oiiotool:
|
||||
oiiotool input.exr --ch R,G,B -d float --compression dwab:300 --clamp:min=0.0:max=32000.0 -o output.exr
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 13 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 115 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 8.0 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 7.6 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 221 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 708 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 202 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.4 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 98 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 9.6 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 11 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 580 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 178 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.2 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 2.3 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 751 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 864 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 11 MiB |
@@ -0,0 +1,113 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 SLat Shape Flow Model (256) with View-Aligned Projection - Finetune",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "ElasticSLatFlowModel",
|
||||
"args": {
|
||||
"resolution": 16,
|
||||
"in_channels": 32,
|
||||
"out_channels": 32,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj",
|
||||
"proj_in_channels": 2048
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSLatShapeView",
|
||||
"args": {
|
||||
"resolution": 256,
|
||||
"image_size": 512,
|
||||
"min_aesthetic_score": 4.5,
|
||||
"max_tokens": 8192,
|
||||
"num_views": 2,
|
||||
"normalization": {
|
||||
"mean": [
|
||||
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
|
||||
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
|
||||
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
|
||||
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
|
||||
],
|
||||
"std": [
|
||||
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
|
||||
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
|
||||
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
|
||||
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
|
||||
]
|
||||
},
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 16,
|
||||
"batch_split": 2,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 8,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"elastic": {
|
||||
"name": "LinearMemoryController",
|
||||
"args": {
|
||||
"target_ratio": 0.75,
|
||||
"max_mem_ratio_start": 0.5
|
||||
}
|
||||
},
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 2000,
|
||||
"i_save": 5000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "uniform",
|
||||
"args": {}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 16,
|
||||
"use_naf_upsample": true,
|
||||
"naf_target_size": 256
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 SLat Shape Flow Model (512) with View-Aligned Projection + NAF Upsample - Finetune from 256 Proj Checkpoint. NAF upsamples DINOv3 32x32 features using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "ElasticSLatFlowModel",
|
||||
"args": {
|
||||
"resolution": 32,
|
||||
"in_channels": 32,
|
||||
"out_channels": 32,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj",
|
||||
"proj_in_channels": 2048
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSLatShapeView",
|
||||
"args": {
|
||||
"resolution": 512,
|
||||
"image_size": 512,
|
||||
"min_aesthetic_score": 4.5,
|
||||
"max_tokens": 8192,
|
||||
"num_views": 2,
|
||||
"normalization": {
|
||||
"mean": [
|
||||
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
|
||||
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
|
||||
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
|
||||
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
|
||||
],
|
||||
"std": [
|
||||
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
|
||||
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
|
||||
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
|
||||
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
|
||||
]
|
||||
},
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 8,
|
||||
"batch_split": 2,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 14,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"elastic": {
|
||||
"name": "LinearMemoryController",
|
||||
"args": {
|
||||
"target_ratio": 0.75,
|
||||
"max_mem_ratio_start": 0.5
|
||||
}
|
||||
},
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"finetune_ckpt": {
|
||||
"denoiser": "<path_to_slat_img2shape_256_checkpoint>"
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 1000,
|
||||
"i_save": 5000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "uniform",
|
||||
"args": {}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 32,
|
||||
"use_naf_upsample": true,
|
||||
"naf_target_size": 512
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 SLat Shape Flow Model (1024) with View-Aligned Projection + NAF Upsample - Finetune from 512 Proj Checkpoint. NAF upsamples DINOv3 32x32 features to higher resolution using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "ElasticSLatFlowModel",
|
||||
"args": {
|
||||
"resolution": 64,
|
||||
"in_channels": 32,
|
||||
"out_channels": 32,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj",
|
||||
"proj_in_channels": 2048
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSLatShapeView",
|
||||
"args": {
|
||||
"resolution": 1024,
|
||||
"image_size": 1024,
|
||||
"min_aesthetic_score": 4.5,
|
||||
"max_tokens": 32768,
|
||||
"num_views": 2,
|
||||
"normalization": {
|
||||
"mean": [
|
||||
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
|
||||
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
|
||||
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
|
||||
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
|
||||
],
|
||||
"std": [
|
||||
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
|
||||
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
|
||||
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
|
||||
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
|
||||
]
|
||||
},
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 2,
|
||||
"batch_split": 2,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 14,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 2e-5,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"elastic": {
|
||||
"name": "LinearMemoryController",
|
||||
"args": {
|
||||
"target_ratio": 0.75,
|
||||
"max_mem_ratio_start": 0.25
|
||||
}
|
||||
},
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"finetune_ckpt": {
|
||||
"denoiser": "<path_to_slat_img2shape_512_checkpoint>"
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 250,
|
||||
"i_save": 1000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "uniform",
|
||||
"args": {}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 1024,
|
||||
"grid_resolution": 64,
|
||||
"use_naf_upsample": true,
|
||||
"naf_target_size": 512
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 SLat PBR/Texture Flow Model (256) with View-Aligned Projection + NAF Upsample - Finetune. NAF upsamples DINOv3 32x32 features to 128x128 using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "ElasticSLatFlowModel",
|
||||
"args": {
|
||||
"resolution": 16,
|
||||
"in_channels": 64,
|
||||
"out_channels": 32,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj",
|
||||
"proj_in_channels": 2048
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSLatPbrView",
|
||||
"args": {
|
||||
"resolution": 256,
|
||||
"image_size": 512,
|
||||
"min_aesthetic_score": 4.5,
|
||||
"max_tokens": 8192,
|
||||
"num_views": 2,
|
||||
"full_pbr": false,
|
||||
"pbr_slat_normalization": {
|
||||
"mean": [
|
||||
3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
|
||||
0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
|
||||
-1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
|
||||
1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
|
||||
],
|
||||
"std": [
|
||||
2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
|
||||
2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
|
||||
2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
|
||||
2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
|
||||
]
|
||||
},
|
||||
"shape_slat_normalization": {
|
||||
"mean": [
|
||||
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
|
||||
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
|
||||
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
|
||||
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
|
||||
],
|
||||
"std": [
|
||||
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
|
||||
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
|
||||
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
|
||||
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
|
||||
]
|
||||
},
|
||||
"attrs": [
|
||||
"base_color",
|
||||
"metallic",
|
||||
"roughness",
|
||||
"alpha"
|
||||
],
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_pbr_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
|
||||
"pretrained_shape_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 16,
|
||||
"batch_split": 1,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 16,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"elastic": {
|
||||
"name": "LinearMemoryController",
|
||||
"args": {
|
||||
"target_ratio": 0.75,
|
||||
"max_mem_ratio_start": 0.5
|
||||
}
|
||||
},
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 1000,
|
||||
"i_save": 5000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "uniform",
|
||||
"args": {}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 16,
|
||||
"use_naf_upsample": true,
|
||||
"naf_target_size": 256
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 SLat PBR/Texture Flow Model (512) with View-Aligned Projection + NAF Upsample - Finetune from 256 Proj Checkpoint. NAF upsamples DINOv3 32x32 features to 256x256 using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "ElasticSLatFlowModel",
|
||||
"args": {
|
||||
"resolution": 32,
|
||||
"in_channels": 64,
|
||||
"out_channels": 32,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj",
|
||||
"proj_in_channels": 2048
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSLatPbrView",
|
||||
"args": {
|
||||
"resolution": 512,
|
||||
"image_size": 512,
|
||||
"min_aesthetic_score": 4.5,
|
||||
"max_tokens": 8192,
|
||||
"num_views": 2,
|
||||
"full_pbr": false,
|
||||
"pbr_slat_normalization": {
|
||||
"mean": [
|
||||
3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
|
||||
0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
|
||||
-1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
|
||||
1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
|
||||
],
|
||||
"std": [
|
||||
2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
|
||||
2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
|
||||
2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
|
||||
2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
|
||||
]
|
||||
},
|
||||
"shape_slat_normalization": {
|
||||
"mean": [
|
||||
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
|
||||
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
|
||||
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
|
||||
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
|
||||
],
|
||||
"std": [
|
||||
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
|
||||
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
|
||||
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
|
||||
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
|
||||
]
|
||||
},
|
||||
"attrs": [
|
||||
"base_color",
|
||||
"metallic",
|
||||
"roughness",
|
||||
"alpha"
|
||||
],
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_pbr_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
|
||||
"pretrained_shape_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 8,
|
||||
"batch_split": 2,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 14,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"elastic": {
|
||||
"name": "LinearMemoryController",
|
||||
"args": {
|
||||
"target_ratio": 0.75,
|
||||
"max_mem_ratio_start": 0.5
|
||||
}
|
||||
},
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"finetune_ckpt": {
|
||||
"denoiser": "<path_to_slat_imgshape2tex_256_checkpoint>"
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 1000,
|
||||
"i_save": 5000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "uniform",
|
||||
"args": {}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 32,
|
||||
"use_naf_upsample": true,
|
||||
"naf_target_size": 256
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 SLat PBR/Texture Flow Model (1024) with View-Aligned Projection + NAF Upsample - Finetune from 512 Proj Checkpoint. NAF upsamples DINOv3 32x32 features to 512x512 using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "ElasticSLatFlowModel",
|
||||
"args": {
|
||||
"resolution": 64,
|
||||
"in_channels": 64,
|
||||
"out_channels": 32,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj",
|
||||
"proj_in_channels": 2048
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSLatPbrView",
|
||||
"args": {
|
||||
"resolution": 1024,
|
||||
"image_size": 1024,
|
||||
"min_aesthetic_score": 4.5,
|
||||
"max_tokens": 32768,
|
||||
"num_views": 2,
|
||||
"full_pbr": true,
|
||||
"pbr_slat_normalization": {
|
||||
"mean": [
|
||||
3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
|
||||
0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
|
||||
-1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
|
||||
1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
|
||||
],
|
||||
"std": [
|
||||
2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
|
||||
2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
|
||||
2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
|
||||
2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
|
||||
]
|
||||
},
|
||||
"shape_slat_normalization": {
|
||||
"mean": [
|
||||
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
|
||||
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
|
||||
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
|
||||
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
|
||||
],
|
||||
"std": [
|
||||
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
|
||||
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
|
||||
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
|
||||
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
|
||||
]
|
||||
},
|
||||
"attrs": [
|
||||
"base_color",
|
||||
"metallic",
|
||||
"roughness",
|
||||
"alpha"
|
||||
],
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_pbr_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
|
||||
"pretrained_shape_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 2,
|
||||
"batch_split": 2,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 14,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 2e-5,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"elastic": {
|
||||
"name": "LinearMemoryController",
|
||||
"args": {
|
||||
"target_ratio": 0.75,
|
||||
"max_mem_ratio_start": 0.25
|
||||
}
|
||||
},
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"finetune_ckpt": {
|
||||
"denoiser": "<path_to_slat_imgshape2tex_512_checkpoint>"
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": -1,
|
||||
"i_save": 250,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "uniform",
|
||||
"args": {}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 1024,
|
||||
"grid_resolution": 64,
|
||||
"use_naf_upsample": true,
|
||||
"naf_target_size": 1024
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 Sparse Structure Flow Model with View-Aligned Projection Conditioning (32 resolution, view-based)",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "SparseStructureFlowModel",
|
||||
"args": {
|
||||
"resolution": 8,
|
||||
"in_channels": 8,
|
||||
"out_channels": 8,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSparseStructureLatentView",
|
||||
"args": {
|
||||
"min_aesthetic_score": 4.5,
|
||||
"image_size": 512,
|
||||
"num_views": 2,
|
||||
"load_camera_info": true,
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"debug": false,
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 16,
|
||||
"batch_split": 1,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 14,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 500,
|
||||
"i_save": 5000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "logitNormal",
|
||||
"args": {
|
||||
"mean": 1.0,
|
||||
"std": 1.0
|
||||
}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 8
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
{
|
||||
"_comment": "TRELLIS.2 Sparse Structure Flow Model (64) with View-Aligned Projection - Finetune from 32 Proj Checkpoint",
|
||||
"models": {
|
||||
"denoiser": {
|
||||
"name": "SparseStructureFlowModel",
|
||||
"args": {
|
||||
"resolution": 16,
|
||||
"in_channels": 8,
|
||||
"out_channels": 8,
|
||||
"model_channels": 1536,
|
||||
"cond_channels": 1024,
|
||||
"num_blocks": 30,
|
||||
"num_heads": 12,
|
||||
"mlp_ratio": 5.3334,
|
||||
"pe_mode": "rope",
|
||||
"share_mod": true,
|
||||
"initialization": "scaled",
|
||||
"qk_rms_norm": true,
|
||||
"qk_rms_norm_cross": true,
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
},
|
||||
"dataset": {
|
||||
"name": "ViewImageConditionedSparseStructureLatentView",
|
||||
"args": {
|
||||
"min_aesthetic_score": 4.5,
|
||||
"image_size": 512,
|
||||
"num_views": 2,
|
||||
"load_camera_info": true,
|
||||
"skip_aesthetic_score_datasets": ["texverse"],
|
||||
"pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
|
||||
}
|
||||
},
|
||||
"trainer": {
|
||||
"name": "ImageConditionedProjFlowMatchingCFGTrainer",
|
||||
"args": {
|
||||
"debug": false,
|
||||
"max_steps": 1000000,
|
||||
"batch_size_per_gpu": 8,
|
||||
"batch_split": 2,
|
||||
"snapshot_batch_size": 1,
|
||||
"snapshot_num_samples": 1,
|
||||
"num_workers": 14,
|
||||
"optimizer": {
|
||||
"name": "AdamW",
|
||||
"args": {
|
||||
"lr": 1e-4,
|
||||
"weight_decay": 0.01,
|
||||
"betas": [0.9, 0.95],
|
||||
"eps": 1e-8
|
||||
}
|
||||
},
|
||||
"ema_rate": [
|
||||
0.9999
|
||||
],
|
||||
"mix_precision_mode": "amp",
|
||||
"mix_precision_dtype": "bfloat16",
|
||||
"grad_clip": {
|
||||
"name": "AdaptiveGradClipper",
|
||||
"args": {
|
||||
"max_norm": 1.0,
|
||||
"clip_percentile": 95
|
||||
}
|
||||
},
|
||||
"finetune_ckpt": {
|
||||
"denoiser": "<path_to_ss_flow_32_checkpoint>"
|
||||
},
|
||||
"i_print": 500,
|
||||
"i_log": 5,
|
||||
"i_sample": 2000,
|
||||
"i_save": 5000,
|
||||
"p_uncond": 0.1,
|
||||
"t_schedule": {
|
||||
"name": "logitNormal",
|
||||
"args": {
|
||||
"mean": 1.0,
|
||||
"std": 1.0
|
||||
}
|
||||
},
|
||||
"sigma_min": 1e-5,
|
||||
"run_projection_test": false,
|
||||
"image_cond_model": {
|
||||
"name": "DinoV3ProjFeatureExtractor",
|
||||
"args": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 16
|
||||
},
|
||||
"image_attn_mode": "proj"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,212 @@
|
||||
# Dataset Preparation Toolkit
|
||||
|
||||
This toolkit provides a comprehensive pipeline for preparing 3D datasets, including downloading, processing, voxelizing, and latent encoding for SC-VAE and Flow Model training.
|
||||
|
||||
This toolkit is built upon and extended from the data processing scripts of [TRELLIS.2](https://github.com/microsoft/TRELLIS2). We gratefully acknowledge the TRELLIS.2 team for open-sourcing their data preparation pipeline, which served as the foundation for this work. Our extensions include view-aligned voxelization and latent encoding for Pixal3D.
|
||||
|
||||
### Step 1: Install Dependencies
|
||||
|
||||
Initialize the environment and install necessary dependencies:
|
||||
|
||||
```bash
|
||||
. ./data_toolkit/setup.sh
|
||||
```
|
||||
|
||||
### Step 2: Initialize Metadata
|
||||
|
||||
Before processing, load the dataset metadata.
|
||||
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py <SUBSET> --root <ROOT> [--source <SOURCE>]
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
- `SUBSET`: Target dataset subset. Options: `ObjaverseXL`, `ABO`, `HSSD`, `TexVerse` (Training sets); `SketchfabPicked`, `Toys4k` (Test sets).
|
||||
- `ROOT`: Root directory to save the data.
|
||||
- `SOURCE`: Data source (Required if `SUBSET` is `ObjaverseXL`). Options: `sketchfab`, `github`.
|
||||
|
||||
**Example:**
|
||||
Load metadata for `ObjaverseXL` (sketchfab) and save to `datasets/ObjaverseXL_sketchfab`:
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --source sketchfab --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
### Step 3: Download Data
|
||||
|
||||
Download the 3D assets to the local storage.
|
||||
|
||||
```bash
|
||||
python data_toolkit/download.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
- `RANK` / `WORLD_SIZE`: Parameters for multi-node distributed downloading.
|
||||
|
||||
**Example:**
|
||||
To download the `ObjaverseXL` subset:
|
||||
|
||||
> **Note:** The example below sets a large `WORLD_SIZE` (160,000) for demonstration purposes, meaning only a tiny fraction of the dataset will be downloaded by this single process.
|
||||
|
||||
```bash
|
||||
python data_toolkit/download.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --world_size 160000
|
||||
```
|
||||
|
||||
*Attention: Some datasets may require an interactive Hugging Face login or manual steps. Please follow any on-screen instructions.*
|
||||
|
||||
**Update Metadata:**
|
||||
After downloading, update the metadata registry:
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
If download records are missing but files already exist locally, use `--from_file` to scan and rebuild:
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --from_file
|
||||
```
|
||||
|
||||
### Step 4: Process Mesh and PBR Textures
|
||||
|
||||
Standardize 3D assets by dumping mesh and PBR textures.
|
||||
*Note: This process utilizes the CPU.*
|
||||
|
||||
```bash
|
||||
# Dump Meshes
|
||||
python data_toolkit/dump_mesh.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
|
||||
|
||||
# Dump PBR Textures
|
||||
python data_toolkit/dump_pbr.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
|
||||
|
||||
# Get statistics of the asset
|
||||
python data_toolkit/asset_stats.py --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
|
||||
```
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
python data_toolkit/dump_mesh.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
python data_toolkit/dump_pbr.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
python data_toolkit/asset_stats.py --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
**Update Metadata:**
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
### Step 5: Render Image Conditions
|
||||
|
||||
Render multi-view images for each asset. These are used both as image conditions for the generator and as camera transforms for view-aligned processing in subsequent steps.
|
||||
*Note: Blender and Pillow will be automatically installed on first run.*
|
||||
|
||||
```bash
|
||||
python data_toolkit/render_cond.py <SUBSET> --root <ROOT> [--num_views <NUM_VIEWS>] [--rank <RANK> --world_size <WORLD_SIZE>]
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
- `NUM_VIEWS`: Number of views to render per asset. Default is `2`.
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
python data_toolkit/render_cond.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
**Update Metadata:**
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
### Step 6: Convert to View-Aligned O-Voxels
|
||||
|
||||
Convert the processed meshes and textures into view-aligned O-Voxels format. Each asset is transformed according to camera views from Step 5, producing per-view voxel representations.
|
||||
*Note: This process utilizes the CPU.*
|
||||
|
||||
```bash
|
||||
python data_toolkit/dual_grid_view.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
|
||||
|
||||
python data_toolkit/voxelize_pbr_view.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
- `RESOLUTION`: Target resolutions for O-Voxels, comma-separated (e.g., `256,512,1024`). Default is `256`.
|
||||
- `VIEW_INDICES`: Specific view indices to process (e.g., `0,1,2` or `0-5`). Default processes all available views.
|
||||
|
||||
**Example:**
|
||||
Convert `ObjaverseXL` to resolution 256 for views 0-1:
|
||||
```bash
|
||||
python data_toolkit/dual_grid_view.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --resolution 256 --view_indices 0-1
|
||||
python data_toolkit/voxelize_pbr_view.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --resolution 256 --view_indices 0-1
|
||||
```
|
||||
|
||||
**Update Metadata:**
|
||||
```bash
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
### At this point, the dataset is ready for SC-VAE Training
|
||||
|
||||
### Step 7: Encode View-Aligned Latents
|
||||
|
||||
Encode view-aligned sparse structures into latents to train the first-stage generator. Each step produces per-view latent files.
|
||||
|
||||
```bash
|
||||
# 1. Encode Shape Latents (multi-view)
|
||||
python data_toolkit/encode_shape_latent_view.py --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
|
||||
|
||||
# 2. Encode PBR Latents (view-aligned)
|
||||
python data_toolkit/encode_pbr_latent_view.py --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
|
||||
|
||||
# 3. Update Metadata (Required before next step)
|
||||
python data_toolkit/build_metadata.py <SUBSET> --root <ROOT>
|
||||
|
||||
# 4. Encode Sparse Structure (SS) Latents (multi-view)
|
||||
python data_toolkit/encode_ss_latent_view.py --root <ROOT> --shape_latent_name <SHAPE_LATENT_NAME> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <SS_RESOLUTION>] [--view_indices <VIEW_INDICES>]
|
||||
```
|
||||
|
||||
**Arguments:**
|
||||
- `RESOLUTION`: Input O-Voxel resolution. Default is `1024`.
|
||||
- `SS_RESOLUTION`: Resolution for sparse structures. Default is `64`.
|
||||
- `SHAPE_LATENT_NAME`: The specific version name of the shape latent (use the `_view` variant name).
|
||||
- `VIEW_INDICES`: Specific view indices to process (e.g., `0,1,2` or `0-5`).
|
||||
|
||||
**Example:**
|
||||
```bash
|
||||
python data_toolkit/encode_shape_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 512 --view_indices 0-1
|
||||
python data_toolkit/encode_pbr_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 512 --view_indices 0-1
|
||||
python data_toolkit/encode_shape_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 1024 --view_indices 0-1
|
||||
python data_toolkit/encode_pbr_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 1024 --view_indices 0-1
|
||||
|
||||
# Update metadata
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
|
||||
# Encode SS Latents (view-aligned)
|
||||
python data_toolkit/encode_ss_latent_view.py --root datasets/ObjaverseXL_sketchfab --shape_latent_name shape_enc_next_dc_f16c32_fp16_1024_view --resolution 64 --view_indices 0-1
|
||||
|
||||
# Final Metadata Update
|
||||
python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
|
||||
```
|
||||
|
||||
### Step 8: Visualize Decoded Latents (Optional)
|
||||
|
||||
Decode latent files back to meshes, export GLB, and render a front-view image for visual inspection.
|
||||
|
||||
**Shape Latent Visualization:**
|
||||
```bash
|
||||
python data_toolkit/visualize_shape_latent.py \
|
||||
--root datasets/ObjaverseXL_sketchfab \
|
||||
--sha256 <SHA256_HASH> \
|
||||
--resolution 1024 \
|
||||
--view_idx 0
|
||||
```
|
||||
|
||||
**PBR Latent Visualization (shape + texture):**
|
||||
```bash
|
||||
python data_toolkit/visualize_pbr_latent.py \
|
||||
--root datasets/ObjaverseXL_sketchfab \
|
||||
--sha256 <SHA256_HASH> \
|
||||
--resolution 1024 \
|
||||
--view_idx 0
|
||||
```
|
||||
|
||||
Outputs are saved to `<ROOT>/vis/<SHA256>/` (shape) or `<ROOT>/vis_pbr/<SHA256>/` (PBR), including:
|
||||
- Decoded GLB mesh (with PBR textures for PBR variant)
|
||||
- Front-view rendered images (normal/depth for shape; shaded/base_color/normal etc. for PBR)
|
||||
- Copied condition renders and camera transforms from Step 5
|
||||
@@ -0,0 +1,132 @@
|
||||
import os
|
||||
import argparse
|
||||
import pickle
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--mesh_dump_root', type=str, default=None,
|
||||
help='Directory to save the mesh dumps')
|
||||
parser.add_argument('--pbr_dump_root', type=str, default=None,
|
||||
help='Directory to save the pbr dumps')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.mesh_dump_root = opt.mesh_dump_root or opt.root
|
||||
opt.pbr_dump_root = opt.pbr_dump_root or opt.root
|
||||
|
||||
os.makedirs(os.path.join(opt.root, 'asset_stats', 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'asset_stats', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'asset_stats','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if 'num_faces' in metadata.columns:
|
||||
metadata = metadata[metadata['num_faces'].isnull()]
|
||||
metadata = metadata[(metadata['mesh_dumped'] == True) | (metadata['pbr_dumped'] == True)]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
records = []
|
||||
with ThreadPoolExecutor(max_workers=opt.max_workers or os.cpu_count()) as executor, \
|
||||
tqdm(total=len(metadata), desc='Processing objects') as pbar:
|
||||
def worker(metadatum):
|
||||
try:
|
||||
sha256 = metadatum['sha256']
|
||||
if metadatum['pbr_dumped'] == True:
|
||||
with open(os.path.join(opt.pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle'), 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
|
||||
num_faces = 0
|
||||
num_vertices = 0
|
||||
for obj in dump['objects']:
|
||||
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
||||
continue
|
||||
num_faces += obj['faces'].shape[0]
|
||||
num_vertices += obj['vertices'].shape[0]
|
||||
|
||||
num_basecolor_tex = 0
|
||||
num_metallic_tex = 0
|
||||
num_roughness_tex = 0
|
||||
num_alpha_tex = 0
|
||||
for mat in dump['materials']:
|
||||
if mat['baseColorTexture'] is not None:
|
||||
num_basecolor_tex += 1
|
||||
if mat['metallicTexture'] is not None:
|
||||
num_metallic_tex += 1
|
||||
if mat['roughnessTexture'] is not None:
|
||||
num_roughness_tex += 1
|
||||
if mat['alphaTexture'] is not None:
|
||||
num_alpha_tex += 1
|
||||
|
||||
record = {
|
||||
'sha256': sha256,
|
||||
'num_faces': num_faces,
|
||||
'num_vertices': num_vertices,
|
||||
'num_basecolor_tex': num_basecolor_tex,
|
||||
'num_metallic_tex': num_metallic_tex,
|
||||
'num_roughness_tex': num_roughness_tex,
|
||||
'num_alpha_tex': num_alpha_tex,
|
||||
}
|
||||
records.append(record)
|
||||
else:
|
||||
with open(os.path.join(opt.mesh_dump_root,'mesh_dumps', f'{sha256}.pickle'), 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
|
||||
num_faces = 0
|
||||
num_vertices = 0
|
||||
for obj in dump['objects']:
|
||||
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
||||
continue
|
||||
num_faces += obj['faces'].shape[0]
|
||||
num_vertices += obj['vertices'].shape[0]
|
||||
|
||||
record = {
|
||||
'sha256': sha256,
|
||||
'num_faces': num_faces,
|
||||
'num_vertices': num_vertices,
|
||||
}
|
||||
records.append(record)
|
||||
pbar.update()
|
||||
except Exception as e:
|
||||
print(f'Error processing {sha256}: {e}')
|
||||
pbar.update()
|
||||
|
||||
for metadatum in metadata.to_dict('records'):
|
||||
executor.submit(worker, metadatum)
|
||||
|
||||
executor.shutdown(wait=True)
|
||||
|
||||
# save records
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.root, 'asset_stats', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,242 @@
|
||||
import argparse, sys, os, math, io
|
||||
from typing import *
|
||||
import bpy
|
||||
import bmesh
|
||||
from mathutils import Vector, Matrix
|
||||
import numpy as np
|
||||
import pickle
|
||||
|
||||
|
||||
"""=============== BLENDER ==============="""
|
||||
|
||||
IMPORT_FUNCTIONS: Dict[str, Callable] = {
|
||||
"obj": bpy.ops.import_scene.obj if bpy.app.version[0] < 4 else bpy.ops.wm.obj_import,
|
||||
"glb": bpy.ops.import_scene.gltf,
|
||||
"gltf": bpy.ops.import_scene.gltf,
|
||||
"usd": bpy.ops.import_scene.usd,
|
||||
"fbx": bpy.ops.import_scene.fbx,
|
||||
"stl": bpy.ops.import_mesh.stl if bpy.app.version[0] < 4 else bpy.ops.wm.stl_import,
|
||||
"usda": bpy.ops.import_scene.usda,
|
||||
"dae": bpy.ops.wm.collada_import,
|
||||
"ply": bpy.ops.import_mesh.ply if bpy.app.version[0] < 4 else bpy.ops.wm.ply_import,
|
||||
"abc": bpy.ops.wm.alembic_import,
|
||||
"blend": bpy.ops.wm.append,
|
||||
}
|
||||
|
||||
|
||||
def init_scene() -> None:
|
||||
"""Resets the scene to a clean state.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# delete everything
|
||||
for obj in bpy.data.objects:
|
||||
bpy.data.objects.remove(obj, do_unlink=True)
|
||||
|
||||
# delete all the materials
|
||||
for material in bpy.data.materials:
|
||||
bpy.data.materials.remove(material, do_unlink=True)
|
||||
|
||||
# delete all the textures
|
||||
for texture in bpy.data.textures:
|
||||
bpy.data.textures.remove(texture, do_unlink=True)
|
||||
|
||||
# delete all the images
|
||||
for image in bpy.data.images:
|
||||
bpy.data.images.remove(image, do_unlink=True)
|
||||
|
||||
|
||||
def load_object(object_path: str) -> None:
|
||||
"""Loads a model with a supported file extension into the scene.
|
||||
|
||||
Args:
|
||||
object_path (str): Path to the model file.
|
||||
|
||||
Raises:
|
||||
ValueError: If the file extension is not supported.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
file_extension = object_path.split(".")[-1].lower()
|
||||
if file_extension is None:
|
||||
raise ValueError(f"Unsupported file type: {object_path}")
|
||||
|
||||
if file_extension == "usdz":
|
||||
# install usdz io package
|
||||
dirname = os.path.dirname(os.path.realpath(__file__))
|
||||
usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
|
||||
bpy.ops.preferences.addon_install(filepath=usdz_package)
|
||||
# enable it
|
||||
addon_name = "io_scene_usdz"
|
||||
bpy.ops.preferences.addon_enable(module=addon_name)
|
||||
# import the usdz
|
||||
from io_scene_usdz.import_usdz import import_usdz
|
||||
|
||||
import_usdz(context, filepath=object_path, materials=True, animations=True)
|
||||
return None
|
||||
|
||||
# load from existing import functions
|
||||
import_function = IMPORT_FUNCTIONS[file_extension]
|
||||
|
||||
print(f"Loading object from {object_path}")
|
||||
if file_extension == "blend":
|
||||
import_function(directory=object_path, link=False)
|
||||
elif file_extension in {"glb", "gltf"}:
|
||||
import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS', bone_heuristic='TEMPERANCE')
|
||||
else:
|
||||
import_function(filepath=object_path)
|
||||
|
||||
|
||||
def delete_invisible_objects() -> None:
|
||||
"""Deletes all invisible objects in the scene.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
to_remove = []
|
||||
for obj in bpy.context.scene.objects:
|
||||
if obj.hide_viewport or obj.hide_render:
|
||||
obj.hide_viewport = False
|
||||
obj.hide_render = False
|
||||
obj.hide_select = False
|
||||
to_remove.append(obj)
|
||||
for obj in to_remove:
|
||||
bpy.data.objects.remove(obj, do_unlink=True)
|
||||
|
||||
# Delete invisible collections
|
||||
invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
|
||||
for col in invisible_collections:
|
||||
bpy.data.collections.remove(col)
|
||||
|
||||
|
||||
def scene_bbox() -> Tuple[Vector, Vector]:
|
||||
"""Returns the bounding box of the scene.
|
||||
|
||||
Taken from Shap-E rendering script
|
||||
(https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
|
||||
|
||||
Returns:
|
||||
Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
|
||||
"""
|
||||
bbox_min = (math.inf,) * 3
|
||||
bbox_max = (-math.inf,) * 3
|
||||
found = False
|
||||
scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
|
||||
for obj in scene_meshes:
|
||||
found = True
|
||||
for coord in obj.bound_box:
|
||||
coord = Vector(coord)
|
||||
coord = obj.matrix_world @ coord
|
||||
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
|
||||
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
|
||||
if not found:
|
||||
raise RuntimeError("no objects in scene to compute bounding box for")
|
||||
return Vector(bbox_min), Vector(bbox_max)
|
||||
|
||||
|
||||
def normalize_scene() -> Tuple[float, Vector]:
|
||||
"""Normalizes the scene by scaling and translating it to fit in a unit cube centered
|
||||
at the origin.
|
||||
|
||||
Mostly taken from the Point-E / Shap-E rendering script
|
||||
(https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
|
||||
but fix for multiple root objects: (see bug report here:
|
||||
https://github.com/openai/shap-e/pull/60).
|
||||
|
||||
Returns:
|
||||
Tuple[float, Vector]: The scale factor and the offset applied to the scene.
|
||||
"""
|
||||
scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
|
||||
if len(scene_root_objects) > 1:
|
||||
# create an empty object to be used as a parent for all root objects
|
||||
scene = bpy.data.objects.new("ParentEmpty", None)
|
||||
bpy.context.scene.collection.objects.link(scene)
|
||||
|
||||
# parent all root objects to the empty object
|
||||
for obj in scene_root_objects:
|
||||
obj.parent = scene
|
||||
else:
|
||||
scene = scene_root_objects[0]
|
||||
|
||||
bbox_min, bbox_max = scene_bbox()
|
||||
scale = 1 / max(bbox_max - bbox_min)
|
||||
scene.scale = scene.scale * scale
|
||||
|
||||
# Apply scale to matrix_world.
|
||||
bpy.context.view_layer.update()
|
||||
bbox_min, bbox_max = scene_bbox()
|
||||
offset = -(bbox_min + bbox_max) / 2
|
||||
scene.matrix_world.translation += offset
|
||||
|
||||
return scale, offset
|
||||
|
||||
|
||||
def main(arg):
|
||||
# Initialize context
|
||||
if arg.object.endswith(".blend"):
|
||||
delete_invisible_objects()
|
||||
else:
|
||||
init_scene()
|
||||
load_object(arg.object)
|
||||
print('[INFO] Scene initialized.')
|
||||
|
||||
# Normalize scene
|
||||
scale, offset = normalize_scene()
|
||||
print('[INFO] Scene normalized.')
|
||||
|
||||
# Start dumping
|
||||
depsgraph = bpy.context.evaluated_depsgraph_get()
|
||||
scene = bpy.context.scene
|
||||
output = {
|
||||
'objects': [],
|
||||
}
|
||||
|
||||
# Dumping meshes
|
||||
for obj in scene.objects:
|
||||
if obj.type != 'MESH':
|
||||
continue
|
||||
|
||||
pack = {
|
||||
"vertices": None,
|
||||
"faces": None,
|
||||
}
|
||||
|
||||
eval_obj = obj.evaluated_get(depsgraph)
|
||||
eval_mesh = eval_obj.to_mesh()
|
||||
|
||||
bm = bmesh.new()
|
||||
bm.from_mesh(eval_mesh)
|
||||
bm.transform(obj.matrix_world)
|
||||
bmesh.ops.triangulate(bm, faces=bm.faces)
|
||||
bm.to_mesh(eval_mesh)
|
||||
bm.free()
|
||||
|
||||
pack["vertices"] = np.array([
|
||||
v.co[:] for v in eval_mesh.vertices
|
||||
], dtype=np.float32) # (N, 3)
|
||||
|
||||
pack["faces"] = np.array([
|
||||
[eval_mesh.loops[i].vertex_index for i in poly.loop_indices]
|
||||
for poly in eval_mesh.polygons
|
||||
], dtype=np.int32) # (F, 3)
|
||||
|
||||
output['objects'].append(pack)
|
||||
|
||||
# Save output
|
||||
os.makedirs(os.path.dirname(arg.output_path), exist_ok=True)
|
||||
with open(arg.output_path, 'wb') as f:
|
||||
pickle.dump(output, f)
|
||||
print('[INFO] Output saved to {}.'.format(arg.output_path))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
|
||||
parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
|
||||
parser.add_argument('--output_path', type=str, default='/tmp', help='The path the output will be dumped to.')
|
||||
argv = sys.argv[sys.argv.index("--") + 1:]
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
main(args)
|
||||
|
||||
@@ -0,0 +1,485 @@
|
||||
import argparse, sys, os, math, io
|
||||
from typing import *
|
||||
import bpy
|
||||
import bmesh
|
||||
from mathutils import Vector, Matrix
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import pickle
|
||||
|
||||
|
||||
"""=============== BLENDER ==============="""
|
||||
|
||||
IMPORT_FUNCTIONS: Dict[str, Callable] = {
|
||||
"obj": bpy.ops.import_scene.obj if bpy.app.version[0] < 4 else bpy.ops.wm.obj_import,
|
||||
"glb": bpy.ops.import_scene.gltf,
|
||||
"gltf": bpy.ops.import_scene.gltf,
|
||||
"usd": bpy.ops.import_scene.usd,
|
||||
"fbx": bpy.ops.import_scene.fbx,
|
||||
"stl": bpy.ops.import_mesh.stl if bpy.app.version[0] < 4 else bpy.ops.wm.stl_import,
|
||||
"usda": bpy.ops.import_scene.usda,
|
||||
"dae": bpy.ops.wm.collada_import,
|
||||
"ply": bpy.ops.import_mesh.ply if bpy.app.version[0] < 4 else bpy.ops.wm.ply_import,
|
||||
"abc": bpy.ops.wm.alembic_import,
|
||||
"blend": bpy.ops.wm.append,
|
||||
}
|
||||
|
||||
|
||||
def init_scene() -> None:
|
||||
"""Resets the scene to a clean state.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# delete everything
|
||||
for obj in bpy.data.objects:
|
||||
bpy.data.objects.remove(obj, do_unlink=True)
|
||||
|
||||
# delete all the materials
|
||||
for material in bpy.data.materials:
|
||||
bpy.data.materials.remove(material, do_unlink=True)
|
||||
|
||||
# delete all the textures
|
||||
for texture in bpy.data.textures:
|
||||
bpy.data.textures.remove(texture, do_unlink=True)
|
||||
|
||||
# delete all the images
|
||||
for image in bpy.data.images:
|
||||
bpy.data.images.remove(image, do_unlink=True)
|
||||
|
||||
|
||||
def load_object(object_path: str) -> None:
|
||||
"""Loads a model with a supported file extension into the scene.
|
||||
|
||||
Args:
|
||||
object_path (str): Path to the model file.
|
||||
|
||||
Raises:
|
||||
ValueError: If the file extension is not supported.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
file_extension = object_path.split(".")[-1].lower()
|
||||
if file_extension is None:
|
||||
raise ValueError(f"Unsupported file type: {object_path}")
|
||||
|
||||
if file_extension == "usdz":
|
||||
# install usdz io package
|
||||
dirname = os.path.dirname(os.path.realpath(__file__))
|
||||
usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
|
||||
bpy.ops.preferences.addon_install(filepath=usdz_package)
|
||||
# enable it
|
||||
addon_name = "io_scene_usdz"
|
||||
bpy.ops.preferences.addon_enable(module=addon_name)
|
||||
# import the usdz
|
||||
from io_scene_usdz.import_usdz import import_usdz
|
||||
|
||||
import_usdz(context, filepath=object_path, materials=True, animations=True)
|
||||
return None
|
||||
|
||||
# load from existing import functions
|
||||
import_function = IMPORT_FUNCTIONS[file_extension]
|
||||
|
||||
print(f"Loading object from {object_path}")
|
||||
if file_extension == "blend":
|
||||
import_function(directory=object_path, link=False)
|
||||
elif file_extension in {"glb", "gltf"}:
|
||||
import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS', bone_heuristic='TEMPERANCE')
|
||||
else:
|
||||
import_function(filepath=object_path)
|
||||
|
||||
|
||||
def delete_invisible_objects() -> None:
|
||||
"""Deletes all invisible objects in the scene.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
to_remove = []
|
||||
for obj in bpy.context.scene.objects:
|
||||
if obj.hide_viewport or obj.hide_render:
|
||||
obj.hide_viewport = False
|
||||
obj.hide_render = False
|
||||
obj.hide_select = False
|
||||
to_remove.append(obj)
|
||||
for obj in to_remove:
|
||||
bpy.data.objects.remove(obj, do_unlink=True)
|
||||
|
||||
# Delete invisible collections
|
||||
invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
|
||||
for col in invisible_collections:
|
||||
bpy.data.collections.remove(col)
|
||||
|
||||
|
||||
def scene_bbox() -> Tuple[Vector, Vector]:
|
||||
"""Returns the bounding box of the scene.
|
||||
|
||||
Taken from Shap-E rendering script
|
||||
(https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
|
||||
|
||||
Returns:
|
||||
Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
|
||||
"""
|
||||
bbox_min = (math.inf,) * 3
|
||||
bbox_max = (-math.inf,) * 3
|
||||
found = False
|
||||
scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
|
||||
for obj in scene_meshes:
|
||||
found = True
|
||||
for coord in obj.bound_box:
|
||||
coord = Vector(coord)
|
||||
coord = obj.matrix_world @ coord
|
||||
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
|
||||
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
|
||||
if not found:
|
||||
raise RuntimeError("no objects in scene to compute bounding box for")
|
||||
return Vector(bbox_min), Vector(bbox_max)
|
||||
|
||||
|
||||
def normalize_scene() -> Tuple[float, Vector]:
|
||||
"""Normalizes the scene by scaling and translating it to fit in a unit cube centered
|
||||
at the origin.
|
||||
|
||||
Mostly taken from the Point-E / Shap-E rendering script
|
||||
(https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
|
||||
but fix for multiple root objects: (see bug report here:
|
||||
https://github.com/openai/shap-e/pull/60).
|
||||
|
||||
Returns:
|
||||
Tuple[float, Vector]: The scale factor and the offset applied to the scene.
|
||||
"""
|
||||
scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
|
||||
if len(scene_root_objects) > 1:
|
||||
# create an empty object to be used as a parent for all root objects
|
||||
scene = bpy.data.objects.new("ParentEmpty", None)
|
||||
bpy.context.scene.collection.objects.link(scene)
|
||||
|
||||
# parent all root objects to the empty object
|
||||
for obj in scene_root_objects:
|
||||
obj.parent = scene
|
||||
else:
|
||||
scene = scene_root_objects[0]
|
||||
|
||||
bbox_min, bbox_max = scene_bbox()
|
||||
scale = 1 / max(bbox_max - bbox_min)
|
||||
scene.scale = scene.scale * scale
|
||||
|
||||
# Apply scale to matrix_world.
|
||||
bpy.context.view_layer.update()
|
||||
bbox_min, bbox_max = scene_bbox()
|
||||
offset = -(bbox_min + bbox_max) / 2
|
||||
scene.matrix_world.translation += offset
|
||||
|
||||
return scale, offset
|
||||
|
||||
|
||||
# =============== NODE TREE PARSING ===============
|
||||
|
||||
def extract_image(tex_node, channels):
|
||||
image = tex_node.image
|
||||
pixels = np.array(image.pixels[:])
|
||||
data = pixels.reshape(image.size[1], image.size[0], -1)
|
||||
data = data[..., channels]
|
||||
|
||||
if data.dtype != np.uint8:
|
||||
data = np.clip(data, 0.0, 1.0)
|
||||
data = (data * 255).astype(np.uint8)
|
||||
|
||||
if len(data.shape) == 2: # Single channel
|
||||
pil_image = Image.fromarray(data, mode='L')
|
||||
elif data.shape[2] == 3:
|
||||
pil_image = Image.fromarray(data, mode='RGB')
|
||||
elif data.shape[2] == 4:
|
||||
pil_image = Image.fromarray(data, mode='RGBA')
|
||||
else:
|
||||
raise ValueError("Unsupported channel shape for image")
|
||||
|
||||
buffer = io.BytesIO()
|
||||
pil_image.save(buffer, format='PNG')
|
||||
png_bytes = buffer.getvalue()
|
||||
|
||||
return {
|
||||
'image': png_bytes,
|
||||
'interpolation': tex_node.interpolation,
|
||||
'extension': tex_node.extension,
|
||||
}
|
||||
|
||||
|
||||
def try_extract_image(link, expected_channel='RGB'):
|
||||
"""
|
||||
Tries to extract an image from a texture node link.
|
||||
Supported sub tree modes:
|
||||
- RGB:
|
||||
TEX_IMAGE ->
|
||||
- R, G, B:
|
||||
TEX_IMAGE -> SEPARATE_COLOR ->
|
||||
- A:
|
||||
TEX_IMAGE ->
|
||||
"""
|
||||
assert expected_channel in ['RGB', 'R', 'G', 'B', 'A'], "Unsupported channel"
|
||||
|
||||
if expected_channel == 'RGB':
|
||||
assert link.from_node.type == 'TEX_IMAGE', "Material is not supported"
|
||||
assert link.from_socket.name == 'Color', "Material is not supported"
|
||||
tex_node = link.from_node
|
||||
return extract_image(tex_node, [0, 1, 2])
|
||||
|
||||
if expected_channel in ['R', 'G', 'B']:
|
||||
socket_name = {
|
||||
'R': 'Red',
|
||||
'G': 'Green',
|
||||
'B': 'Blue',
|
||||
}[expected_channel]
|
||||
assert link.from_node.type == 'SEPARATE_COLOR' and link.from_node.mode == 'RGB', \
|
||||
f"Material is not supported, {link.from_node.type}, {link.from_node.mode}"
|
||||
assert link.from_socket.name == socket_name, "Material is not supported"
|
||||
sep_node = link.from_node
|
||||
assert sep_node.inputs[0].is_linked and sep_node.inputs[0].links[0].from_node.type == 'TEX_IMAGE', \
|
||||
"Material is not supported"
|
||||
assert sep_node.inputs[0].links[0].from_socket.name == 'Color', "Material is not supported"
|
||||
tex_node = sep_node.inputs[0].links[0].from_node
|
||||
channel_index = {
|
||||
'R': 0,
|
||||
'G': 1,
|
||||
'B': 2,
|
||||
}[expected_channel]
|
||||
return extract_image(tex_node, channel_index)
|
||||
|
||||
if expected_channel == 'A':
|
||||
assert link.from_node.type == 'TEX_IMAGE', "Material is not supported"
|
||||
assert link.from_socket.name == 'Alpha', "Material is not supported"
|
||||
tex_node = link.from_node
|
||||
return extract_image(tex_node, 3)
|
||||
|
||||
|
||||
def try_extract_factor(link, mode='color'):
|
||||
"""
|
||||
Tries to extract a factor from a math node link.
|
||||
Supported sub tree modes:
|
||||
- color:
|
||||
ANY -> MIX(MULTIPLY) ->
|
||||
- scalar:
|
||||
ANY -> MATH(MULTIPLY) ->
|
||||
"""
|
||||
assert mode in ['color','scalar'], "Unsupported mode"
|
||||
|
||||
if mode == 'color':
|
||||
if link.from_node.type == 'MIX':
|
||||
mix_node = link.from_node
|
||||
assert mix_node.data_type == 'RGBA' and mix_node.blend_type == 'MULTIPLY', f"Material is not supported, {mix_node.data_type}, {mix_node.blend_type}"
|
||||
assert not mix_node.inputs['Factor'].is_linked and mix_node.inputs['Factor'].default_value == 1.0, \
|
||||
"Material is not supported"
|
||||
if mix_node.inputs['A'].is_linked:
|
||||
assert not mix_node.inputs['B'].is_linked, "Material is not supported"
|
||||
return (list(mix_node.inputs['B'].default_value)[:3], mix_node.inputs['A'].links[0])
|
||||
else:
|
||||
assert not mix_node.inputs['A'].is_linked, "Material is not supported"
|
||||
assert mix_node.inputs['B'].is_linked, "Material is not supported"
|
||||
return (list(mix_node.inputs['A'].default_value)[:3], mix_node.inputs['B'].links[0])
|
||||
return ([1.0, 1.0, 1.0], link)
|
||||
|
||||
if mode =='scalar':
|
||||
if link.from_node.type == 'MATH':
|
||||
math_node = link.from_node
|
||||
assert math_node.operation == 'MULTIPLY', "Material is not supported"
|
||||
assert math_node.inputs[0].is_linked, "Material is not supported"
|
||||
assert not math_node.inputs[1].is_linked, "Material is not supported"
|
||||
return (math_node.inputs[1].default_value, math_node.inputs[0].links[0])
|
||||
return (1.0, link)
|
||||
|
||||
|
||||
def try_extract_image_with_factor(link, expected_channel='RGB'):
|
||||
"""
|
||||
Tries to extract an image and a factor from a texture node link.
|
||||
"""
|
||||
factor, link = try_extract_factor(link, 'color' if expected_channel in ['RGB'] else 'scalar')
|
||||
image = try_extract_image(link, expected_channel)
|
||||
return (factor, image)
|
||||
|
||||
|
||||
def main(arg):
|
||||
# Initialize context
|
||||
if arg.object.endswith(".blend"):
|
||||
delete_invisible_objects()
|
||||
else:
|
||||
init_scene()
|
||||
load_object(arg.object)
|
||||
print('[INFO] Scene initialized.')
|
||||
|
||||
# Normalize scene
|
||||
scale, offset = normalize_scene()
|
||||
print('[INFO] Scene normalized.')
|
||||
|
||||
# Start dumping
|
||||
depsgraph = bpy.context.evaluated_depsgraph_get()
|
||||
scene = bpy.context.scene
|
||||
output = {
|
||||
'materials': [],
|
||||
'objects': [],
|
||||
}
|
||||
|
||||
# Dumping materials
|
||||
for mat in bpy.data.materials:
|
||||
assert mat.use_nodes == True, "Material is not supported"
|
||||
|
||||
pack = {
|
||||
"baseColorFactor": [1.0, 1.0, 1.0],
|
||||
"alphaFactor": 1.0,
|
||||
"metallicFactor": 1.0,
|
||||
"roughnessFactor": 1.0,
|
||||
"alphaMode": "OPAQUE",
|
||||
"alphaCutoff": 0.5,
|
||||
"baseColorTexture": None,
|
||||
"alphaTexture": None,
|
||||
"metallicTexture": None,
|
||||
"roughnessTexture": None,
|
||||
}
|
||||
|
||||
try:
|
||||
principled_node = mat.node_tree.nodes.get('Principled BSDF')
|
||||
assert principled_node is not None, "Material is not supported"
|
||||
|
||||
# Handle base color
|
||||
if not principled_node.inputs['Base Color'].is_linked:
|
||||
pack["baseColorFactor"] = list(principled_node.inputs['Base Color'].default_value)
|
||||
else:
|
||||
link = principled_node.inputs['Base Color'].links[0]
|
||||
if link.from_node.type == 'RGB':
|
||||
pack["baseColorFactor"] = list(link.from_node.outputs[0].default_value)
|
||||
else:
|
||||
factor, image = try_extract_image_with_factor(link, 'RGB')
|
||||
pack["baseColorFactor"] = factor
|
||||
pack["baseColorTexture"] = image
|
||||
|
||||
# Handle alpha
|
||||
if not principled_node.inputs['Alpha'].is_linked:
|
||||
pack["alphaFactor"] = principled_node.inputs['Alpha'].default_value
|
||||
if pack["alphaFactor"] < 1.0:
|
||||
pack["alphaMode"] = "BLEND"
|
||||
else:
|
||||
link = principled_node.inputs['Alpha'].links[0]
|
||||
node = link.from_node
|
||||
if node.type == 'VALUE':
|
||||
pack["alphaFactor"] = node.outputs[0].default_value
|
||||
if pack["alphaFactor"] < 1.0:
|
||||
pack["alphaMode"] = "BLEND"
|
||||
else:
|
||||
pack["alphaMode"] = "BLEND"
|
||||
if node.type == 'MATH':
|
||||
if node.operation == 'ROUND':
|
||||
assert node.inputs[0].is_linked, "Material is not supported"
|
||||
pack["alphaMode"] = "MASK"
|
||||
link = node.inputs[0].links[0]
|
||||
elif node.operation == 'SUBTRACT':
|
||||
assert node.inputs[0].default_value == 1.0 and \
|
||||
node.inputs[1].is_linked and \
|
||||
node.inputs[1].links[0].from_node.type == 'MATH' and \
|
||||
node.inputs[1].links[0].from_node.operation == 'LESS_THAN', \
|
||||
"Material is not supported"
|
||||
assert node.inputs[1].links[0].from_node.inputs[0].is_linked, "Material is not supported"
|
||||
pack["alphaMode"] = "MASK"
|
||||
pack["alphaCutoff"] = node.inputs[1].links[0].from_node.inputs[1].default_value
|
||||
link = node.inputs[1].links[0].from_node.inputs[0].links[0]
|
||||
factor, image = try_extract_image_with_factor(link, 'A')
|
||||
pack["alphaFactor"] = factor
|
||||
pack["alphaTexture"] = image
|
||||
|
||||
# Handle metallic
|
||||
if not principled_node.inputs['Metallic'].is_linked:
|
||||
pack["metallicFactor"] = principled_node.inputs['Metallic'].default_value
|
||||
else:
|
||||
link = principled_node.inputs['Metallic'].links[0]
|
||||
node = link.from_node
|
||||
if node.type == 'VALUE':
|
||||
pack["metallicFactor"] = node.outputs[0].default_value
|
||||
else:
|
||||
factor, image = try_extract_image_with_factor(link, 'B')
|
||||
pack["metallicFactor"] = factor
|
||||
pack["metallicTexture"] = image
|
||||
|
||||
# Handle roughness
|
||||
if not principled_node.inputs['Roughness'].is_linked:
|
||||
pack["roughnessFactor"] = principled_node.inputs['Roughness'].default_value
|
||||
else:
|
||||
link = principled_node.inputs['Roughness'].links[0]
|
||||
node = link.from_node
|
||||
if node.type == 'VALUE':
|
||||
pack["roughnessFactor"] = node.outputs[0].default_value
|
||||
else:
|
||||
factor, image = try_extract_image_with_factor(link, 'G')
|
||||
pack["roughnessFactor"] = factor
|
||||
pack["roughnessTexture"] = image
|
||||
|
||||
output['materials'].append(pack)
|
||||
except:
|
||||
with open(arg.output_path + '_error.txt', 'w') as f:
|
||||
f.write(str([[n.name] for n in mat.node_tree.nodes]))
|
||||
raise RuntimeError("Material is not supported")
|
||||
|
||||
# Dumping meshes
|
||||
for obj in scene.objects:
|
||||
if obj.type != 'MESH':
|
||||
continue
|
||||
|
||||
pack = {
|
||||
"vertices": None,
|
||||
"faces": None,
|
||||
"uvs": None,
|
||||
"matIDs": None,
|
||||
}
|
||||
|
||||
eval_obj = obj.evaluated_get(depsgraph)
|
||||
eval_mesh = eval_obj.to_mesh()
|
||||
|
||||
bm = bmesh.new()
|
||||
bm.from_mesh(eval_mesh)
|
||||
bm.transform(obj.matrix_world)
|
||||
bmesh.ops.triangulate(bm, faces=bm.faces)
|
||||
bm.to_mesh(eval_mesh)
|
||||
bm.free()
|
||||
|
||||
pack["vertices"] = np.array([
|
||||
v.co[:] for v in eval_mesh.vertices
|
||||
], dtype=np.float32) # (N, 3)
|
||||
|
||||
pack["faces"] = np.array([
|
||||
[eval_mesh.loops[i].vertex_index for i in poly.loop_indices]
|
||||
for poly in eval_mesh.polygons
|
||||
], dtype=np.int32) # (F, 3)
|
||||
|
||||
pack["normals"] = np.array([
|
||||
[eval_mesh.loops[i].normal for i in poly.loop_indices]
|
||||
for poly in eval_mesh.polygons
|
||||
], dtype=np.float32) # (F, 3, 3)
|
||||
|
||||
if eval_mesh.uv_layers.active is not None:
|
||||
pack["uvs"] = np.array([
|
||||
[eval_mesh.uv_layers.active.data[i].uv for i in poly.loop_indices]
|
||||
for poly in eval_mesh.polygons
|
||||
], dtype=np.float32) # (F, 3, 2)
|
||||
|
||||
pack["mat_ids"] = np.array([
|
||||
bpy.data.materials.find(obj.material_slots[poly.material_index].name)
|
||||
if len(obj.material_slots) > 0 and obj.material_slots[poly.material_index].material is not None else -1
|
||||
for poly in eval_mesh.polygons
|
||||
], dtype=np.int32)
|
||||
|
||||
output['objects'].append(pack)
|
||||
|
||||
# Save output
|
||||
os.makedirs(os.path.dirname(arg.output_path), exist_ok=True)
|
||||
with open(arg.output_path, 'wb') as f:
|
||||
pickle.dump(output, f)
|
||||
print('[INFO] Output saved to {}.'.format(arg.output_path))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
|
||||
parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
|
||||
parser.add_argument('--output_path', type=str, default='/tmp', help='The path the output will be dumped to.')
|
||||
argv = sys.argv[sys.argv.index("--") + 1:]
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
main(args)
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
import subprocess
|
||||
import sys
|
||||
import ensurepip
|
||||
|
||||
ensurepip.bootstrap()
|
||||
subprocess.check_call([sys.executable, "-m", "pip", "install", "Pillow"])
|
||||
Binary file not shown.
@@ -0,0 +1,529 @@
|
||||
import argparse, sys, os, math, re, glob
|
||||
from typing import *
|
||||
import bpy
|
||||
from mathutils import Vector, Matrix
|
||||
import numpy as np
|
||||
import json
|
||||
import glob
|
||||
from PIL import Image
|
||||
|
||||
|
||||
"""=============== BLENDER ==============="""
|
||||
|
||||
IMPORT_FUNCTIONS: Dict[str, Callable] = {
|
||||
"obj": bpy.ops.import_scene.obj,
|
||||
"glb": bpy.ops.import_scene.gltf,
|
||||
"gltf": bpy.ops.import_scene.gltf,
|
||||
"usd": bpy.ops.import_scene.usd,
|
||||
"fbx": bpy.ops.import_scene.fbx,
|
||||
"stl": bpy.ops.import_mesh.stl,
|
||||
"usda": bpy.ops.import_scene.usda,
|
||||
"dae": bpy.ops.wm.collada_import,
|
||||
"ply": bpy.ops.import_mesh.ply,
|
||||
"abc": bpy.ops.wm.alembic_import,
|
||||
"blend": bpy.ops.wm.append,
|
||||
}
|
||||
|
||||
EXT = {
|
||||
'PNG': 'png',
|
||||
'JPEG': 'jpg',
|
||||
'OPEN_EXR': 'exr',
|
||||
'TIFF': 'tiff',
|
||||
'BMP': 'bmp',
|
||||
'HDR': 'hdr',
|
||||
'TARGA': 'tga'
|
||||
}
|
||||
|
||||
|
||||
def init_render(engine='CYCLES', resolution=512):
|
||||
bpy.context.scene.render.engine = engine
|
||||
bpy.context.scene.render.resolution_x = resolution
|
||||
bpy.context.scene.render.resolution_y = resolution
|
||||
bpy.context.scene.render.resolution_percentage = 100
|
||||
bpy.context.scene.render.image_settings.file_format = 'PNG'
|
||||
bpy.context.scene.render.image_settings.color_mode = 'RGBA'
|
||||
bpy.context.scene.render.film_transparent = True
|
||||
|
||||
bpy.context.scene.cycles.device = 'GPU'
|
||||
bpy.context.scene.cycles.samples = 32
|
||||
bpy.context.scene.cycles.filter_type = 'BOX'
|
||||
bpy.context.scene.cycles.filter_width = 1
|
||||
bpy.context.scene.cycles.diffuse_bounces = 1
|
||||
bpy.context.scene.cycles.glossy_bounces = 1
|
||||
bpy.context.scene.cycles.transparent_max_bounces = 3
|
||||
bpy.context.scene.cycles.transmission_bounces = 3
|
||||
bpy.context.scene.cycles.use_denoising = True
|
||||
|
||||
bpy.context.preferences.addons['cycles'].preferences.get_devices()
|
||||
bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
|
||||
|
||||
|
||||
def init_scene() -> None:
|
||||
"""Resets the scene to a clean state.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# delete everything
|
||||
for obj in bpy.data.objects:
|
||||
bpy.data.objects.remove(obj, do_unlink=True)
|
||||
|
||||
# delete all the materials
|
||||
for material in bpy.data.materials:
|
||||
bpy.data.materials.remove(material, do_unlink=True)
|
||||
|
||||
# delete all the textures
|
||||
for texture in bpy.data.textures:
|
||||
bpy.data.textures.remove(texture, do_unlink=True)
|
||||
|
||||
# delete all the images
|
||||
for image in bpy.data.images:
|
||||
bpy.data.images.remove(image, do_unlink=True)
|
||||
|
||||
|
||||
def init_camera():
|
||||
cam = bpy.data.objects.new('Camera', bpy.data.cameras.new('Camera'))
|
||||
bpy.context.collection.objects.link(cam)
|
||||
bpy.context.scene.camera = cam
|
||||
cam.data.sensor_height = cam.data.sensor_width = 32
|
||||
cam_constraint = cam.constraints.new(type='TRACK_TO')
|
||||
cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'
|
||||
cam_constraint.up_axis = 'UP_Y'
|
||||
cam_empty = bpy.data.objects.new("Empty", None)
|
||||
cam_empty.location = (0, 0, 0)
|
||||
bpy.context.scene.collection.objects.link(cam_empty)
|
||||
cam_constraint.target = cam_empty
|
||||
return cam
|
||||
|
||||
|
||||
def init_uniform_lighting():
|
||||
# Clear existing lights
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
bpy.ops.object.select_by_type(type="LIGHT")
|
||||
bpy.ops.object.delete()
|
||||
|
||||
# Create environment light
|
||||
if bpy.context.scene.world is None:
|
||||
world = bpy.data.worlds.new("World")
|
||||
bpy.context.scene.world = world
|
||||
else:
|
||||
world = bpy.context.scene.world
|
||||
|
||||
# Enabling nodes
|
||||
world.use_nodes = True
|
||||
node_tree = world.node_tree
|
||||
nodes = node_tree.nodes
|
||||
links = node_tree.links
|
||||
|
||||
# Remove default nodes
|
||||
for node in nodes:
|
||||
nodes.remove(node)
|
||||
|
||||
# Create background node
|
||||
bg_node = nodes.new(type="ShaderNodeBackground")
|
||||
bg_node.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
|
||||
bg_node.inputs["Strength"].default_value = 1.0
|
||||
output_node = nodes.new(type="ShaderNodeOutputWorld")
|
||||
links.new(bg_node.outputs["Background"], output_node.inputs["Surface"])
|
||||
|
||||
|
||||
def init_random_lighting(camera_dir: np.ndarray) -> None:
|
||||
# Clear existing lights
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
bpy.ops.object.select_by_type(type="LIGHT")
|
||||
bpy.ops.object.delete()
|
||||
|
||||
# Create environment light
|
||||
if bpy.context.scene.world is None:
|
||||
world = bpy.data.worlds.new("World")
|
||||
bpy.context.scene.world = world
|
||||
else:
|
||||
world = bpy.context.scene.world
|
||||
|
||||
# Enabling nodes
|
||||
world.use_nodes = True
|
||||
node_tree = world.node_tree
|
||||
nodes = node_tree.nodes
|
||||
links = node_tree.links
|
||||
|
||||
# Remove default nodes
|
||||
for node in nodes:
|
||||
nodes.remove(node)
|
||||
|
||||
# Random place lights
|
||||
num_lights = np.random.randint(1, 4)
|
||||
total_strength = 1.5
|
||||
for i in range(num_lights):
|
||||
new_light = bpy.data.objects.new(f"Light_{i}", bpy.data.lights.new(f"Light_{i}", type="POINT"))
|
||||
bpy.context.collection.objects.link(new_light)
|
||||
|
||||
new_light_distance = 1 / np.random.uniform(1/100, 1/10)
|
||||
new_light_dir = np.random.randn(3)
|
||||
new_light_dir[2] += 0.6
|
||||
new_light_dir = new_light_dir / np.linalg.norm(new_light_dir)
|
||||
new_light_location = new_light_dir * new_light_distance
|
||||
new_light_camera_strength_ratio = max(np.sum(camera_dir * new_light_dir) * 0.5 + 0.5, 0)
|
||||
new_light_max_energy = total_strength / (np.sum(camera_dir * new_light_dir) * 0.45 + 0.55)
|
||||
new_light_strength = np.sqrt(np.random.uniform(0.01, 1)) * new_light_max_energy
|
||||
new_light_camera_strength = new_light_camera_strength_ratio * new_light_strength
|
||||
total_strength -= new_light_camera_strength
|
||||
|
||||
new_light.location = (new_light_location[0], new_light_location[1], new_light_location[2])
|
||||
new_light.data.color = (1.0, 1.0, 1.0)
|
||||
new_light.data.energy = new_light_strength * new_light_distance**2 * 31.4
|
||||
new_light.data.shadow_soft_size = np.random.uniform(0.1, 0.1 * new_light_distance)
|
||||
|
||||
# Create background node
|
||||
bg_node = nodes.new(type="ShaderNodeBackground")
|
||||
bg_node.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
|
||||
bg_node.inputs["Strength"].default_value = total_strength
|
||||
output_node = nodes.new(type="ShaderNodeOutputWorld")
|
||||
links.new(bg_node.outputs["Background"], output_node.inputs["Surface"])
|
||||
|
||||
|
||||
def load_object(object_path: str) -> None:
|
||||
"""Loads a model with a supported file extension into the scene.
|
||||
|
||||
Args:
|
||||
object_path (str): Path to the model file.
|
||||
|
||||
Raises:
|
||||
ValueError: If the file extension is not supported.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
file_extension = object_path.split(".")[-1].lower()
|
||||
if file_extension is None:
|
||||
raise ValueError(f"Unsupported file type: {object_path}")
|
||||
|
||||
if file_extension == "usdz":
|
||||
# install usdz io package
|
||||
dirname = os.path.dirname(os.path.realpath(__file__))
|
||||
usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
|
||||
bpy.ops.preferences.addon_install(filepath=usdz_package)
|
||||
# enable it
|
||||
addon_name = "io_scene_usdz"
|
||||
bpy.ops.preferences.addon_enable(module=addon_name)
|
||||
# import the usdz
|
||||
from io_scene_usdz.import_usdz import import_usdz
|
||||
|
||||
import_usdz(context, filepath=object_path, materials=True, animations=True)
|
||||
return None
|
||||
|
||||
# load from existing import functions
|
||||
import_function = IMPORT_FUNCTIONS[file_extension]
|
||||
|
||||
print(f"Loading object from {object_path}")
|
||||
if file_extension == "blend":
|
||||
import_function(directory=object_path, link=False)
|
||||
elif file_extension in {"glb", "gltf"}:
|
||||
import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS')
|
||||
else:
|
||||
import_function(filepath=object_path)
|
||||
|
||||
|
||||
def delete_invisible_objects() -> None:
|
||||
"""Deletes all invisible objects in the scene.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# bpy.ops.object.mode_set(mode="OBJECT")
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
for obj in bpy.context.scene.objects:
|
||||
if obj.hide_viewport or obj.hide_render:
|
||||
obj.hide_viewport = False
|
||||
obj.hide_render = False
|
||||
obj.hide_select = False
|
||||
obj.select_set(True)
|
||||
bpy.ops.object.delete()
|
||||
|
||||
# Delete invisible collections
|
||||
invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
|
||||
for col in invisible_collections:
|
||||
bpy.data.collections.remove(col)
|
||||
|
||||
|
||||
def unhide_all_objects() -> None:
|
||||
"""Unhides all objects in the scene.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for obj in bpy.context.scene.objects:
|
||||
obj.hide_set(False)
|
||||
|
||||
|
||||
def convert_to_meshes() -> None:
|
||||
"""Converts all objects in the scene to meshes.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
bpy.context.view_layer.objects.active = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"][0]
|
||||
for obj in bpy.context.scene.objects:
|
||||
obj.select_set(True)
|
||||
bpy.ops.object.convert(target="MESH")
|
||||
|
||||
|
||||
def triangulate_meshes() -> None:
|
||||
"""Triangulates all meshes in the scene.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
|
||||
bpy.context.view_layer.objects.active = objs[0]
|
||||
for obj in objs:
|
||||
obj.select_set(True)
|
||||
bpy.ops.object.mode_set(mode="EDIT")
|
||||
bpy.ops.mesh.reveal()
|
||||
bpy.ops.mesh.select_all(action="SELECT")
|
||||
bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
|
||||
bpy.ops.object.mode_set(mode="OBJECT")
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
|
||||
|
||||
def scene_bbox() -> Tuple[Vector, Vector]:
|
||||
"""Returns the bounding box of the scene.
|
||||
|
||||
Taken from Shap-E rendering script
|
||||
(https://github.com/openai/shap-e/blob/main/shap_e/rendering/blender/blender_script.py#L68-L82)
|
||||
|
||||
Returns:
|
||||
Tuple[Vector, Vector]: The minimum and maximum coordinates of the bounding box.
|
||||
"""
|
||||
bbox_min = (math.inf,) * 3
|
||||
bbox_max = (-math.inf,) * 3
|
||||
found = False
|
||||
scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
|
||||
for obj in scene_meshes:
|
||||
found = True
|
||||
for coord in obj.bound_box:
|
||||
coord = Vector(coord)
|
||||
coord = obj.matrix_world @ coord
|
||||
bbox_min = tuple(min(x, y) for x, y in zip(bbox_min, coord))
|
||||
bbox_max = tuple(max(x, y) for x, y in zip(bbox_max, coord))
|
||||
if not found:
|
||||
raise RuntimeError("no objects in scene to compute bounding box for")
|
||||
return Vector(bbox_min), Vector(bbox_max)
|
||||
|
||||
|
||||
def normalize_scene() -> Tuple[float, Vector]:
|
||||
"""Normalizes the scene by scaling and translating it to fit in a unit cube centered
|
||||
at the origin.
|
||||
|
||||
Mostly taken from the Point-E / Shap-E rendering script
|
||||
(https://github.com/openai/point-e/blob/main/point_e/evals/scripts/blender_script.py#L97-L112),
|
||||
but fix for multiple root objects: (see bug report here:
|
||||
https://github.com/openai/shap-e/pull/60).
|
||||
|
||||
Returns:
|
||||
Tuple[float, Vector]: The scale factor and the offset applied to the scene.
|
||||
"""
|
||||
scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
|
||||
if len(scene_root_objects) > 1:
|
||||
# create an empty object to be used as a parent for all root objects
|
||||
scene = bpy.data.objects.new("ParentEmpty", None)
|
||||
bpy.context.scene.collection.objects.link(scene)
|
||||
|
||||
# parent all root objects to the empty object
|
||||
for obj in scene_root_objects:
|
||||
obj.parent = scene
|
||||
else:
|
||||
scene = scene_root_objects[0]
|
||||
|
||||
bbox_min, bbox_max = scene_bbox()
|
||||
scale = 1 / max(bbox_max - bbox_min)
|
||||
scene.scale = scene.scale * scale
|
||||
|
||||
# Apply scale to matrix_world.
|
||||
bpy.context.view_layer.update()
|
||||
bbox_min, bbox_max = scene_bbox()
|
||||
offset = -(bbox_min + bbox_max) / 2
|
||||
scene.matrix_world.translation += offset
|
||||
bpy.ops.object.select_all(action="DESELECT")
|
||||
|
||||
return scale, offset
|
||||
|
||||
|
||||
def get_transform_matrix(obj: bpy.types.Object) -> list:
|
||||
pos, rt, _ = obj.matrix_world.decompose()
|
||||
rt = rt.to_matrix()
|
||||
matrix = []
|
||||
for ii in range(3):
|
||||
a = []
|
||||
for jj in range(3):
|
||||
a.append(rt[ii][jj])
|
||||
a.append(pos[ii])
|
||||
matrix.append(a)
|
||||
matrix.append([0, 0, 0, 1])
|
||||
return matrix
|
||||
|
||||
|
||||
def check_mask_boundary_distance(image_path: str, threshold: int = 0) -> Tuple[bool, bool, int]:
|
||||
"""Check the rendered object's mask distance to image boundary.
|
||||
|
||||
Args:
|
||||
image_path: Path to the rendered PNG image with alpha channel.
|
||||
threshold: Alpha value threshold to consider as valid mask (0-255).
|
||||
|
||||
Returns:
|
||||
Tuple of (touches_boundary, too_far_from_boundary, min_distance):
|
||||
- touches_boundary: True if the mask touches any boundary
|
||||
- too_far_from_boundary: True if the mask is too far from all boundaries (>80 pixels)
|
||||
- min_distance: Minimum distance from mask to any boundary
|
||||
"""
|
||||
img = Image.open(image_path)
|
||||
if img.mode != 'RGBA':
|
||||
return False, False, 0
|
||||
|
||||
# Get alpha channel
|
||||
alpha = np.array(img)[:, :, 3]
|
||||
h, w = alpha.shape
|
||||
|
||||
# Find all pixels with alpha > threshold (mask pixels)
|
||||
mask_pixels = np.where(alpha > threshold)
|
||||
|
||||
if len(mask_pixels[0]) == 0:
|
||||
# No valid mask pixels
|
||||
return False, True, max(h, w)
|
||||
|
||||
# Calculate distances to each boundary
|
||||
min_row = np.min(mask_pixels[0]) # Distance to top edge
|
||||
max_row = np.max(mask_pixels[0]) # Distance to bottom edge
|
||||
min_col = np.min(mask_pixels[1]) # Distance to left edge
|
||||
max_col = np.max(mask_pixels[1]) # Distance to right edge
|
||||
|
||||
dist_top = min_row
|
||||
dist_bottom = (h - 1) - max_row
|
||||
dist_left = min_col
|
||||
dist_right = (w - 1) - max_col
|
||||
|
||||
min_distance = min(dist_top, dist_bottom, dist_left, dist_right)
|
||||
|
||||
# Check if touches boundary (distance <= 0)
|
||||
touches_boundary = min_distance <= 0
|
||||
|
||||
# Check if too far from boundary (distance > 130 pixels)
|
||||
too_far = min_distance > 130
|
||||
|
||||
return touches_boundary, too_far, min_distance
|
||||
|
||||
|
||||
def main(arg):
|
||||
if arg.object.endswith(".blend"):
|
||||
delete_invisible_objects()
|
||||
else:
|
||||
init_scene()
|
||||
load_object(arg.object)
|
||||
print('[INFO] Scene initialized.')
|
||||
|
||||
# normalize scene
|
||||
scale, offset = normalize_scene()
|
||||
print('[INFO] Scene normalized.')
|
||||
|
||||
# Initialize camera and lighting
|
||||
cam = init_camera()
|
||||
init_uniform_lighting()
|
||||
print('[INFO] Camera and lighting initialized.')
|
||||
|
||||
# ============= Render conditional views =============
|
||||
init_render(engine=arg.engine, resolution=arg.cond_resolution)
|
||||
# Create a list of views
|
||||
to_export = {
|
||||
"aabb": [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
"scale": scale,
|
||||
"offset": [offset.x, offset.y, offset.z],
|
||||
"frames": []
|
||||
}
|
||||
views = json.loads(arg.cond_views)
|
||||
|
||||
# Parameters for boundary check and radius adjustment
|
||||
max_retry = 10 # Maximum number of retries per view
|
||||
radius_increase_factor = 1.1 # Increase radius by 10% when too close to boundary
|
||||
radius_decrease_factor = 0.9 # Decrease radius by 10% when too far from boundary
|
||||
min_boundary_distance = 130 # Minimum distance to boundary in pixels
|
||||
|
||||
for i, view in enumerate(views):
|
||||
current_radius = view['radius']
|
||||
retry_count = 0
|
||||
|
||||
while retry_count < max_retry:
|
||||
cam_dir = np.array([
|
||||
np.cos(view['yaw']) * np.cos(view['pitch']),
|
||||
np.sin(view['yaw']) * np.cos(view['pitch']),
|
||||
np.sin(view['pitch'])
|
||||
])
|
||||
init_random_lighting(cam_dir)
|
||||
cam.location = (
|
||||
current_radius * cam_dir[0],
|
||||
current_radius * cam_dir[1],
|
||||
current_radius * cam_dir[2]
|
||||
)
|
||||
cam.data.lens = 16 / np.tan(view['fov'] / 2)
|
||||
|
||||
output_path = os.path.join(arg.cond_output_folder, f'{i:03d}.png')
|
||||
bpy.context.scene.render.filepath = output_path
|
||||
|
||||
# Render the scene
|
||||
bpy.ops.render.render(write_still=True)
|
||||
bpy.context.view_layer.update()
|
||||
|
||||
# Check mask boundary distance
|
||||
touches_boundary, too_far, min_dist = check_mask_boundary_distance(output_path)
|
||||
|
||||
if touches_boundary:
|
||||
# Object is too close to boundary, increase radius
|
||||
retry_count += 1
|
||||
old_radius = current_radius
|
||||
current_radius *= radius_increase_factor
|
||||
print(f'[WARNING] View {i}: Mask touches boundary (dist={min_dist}px). Increasing radius from {old_radius:.4f} to {current_radius:.4f} (retry {retry_count}/{max_retry})')
|
||||
elif too_far:
|
||||
# Object is too far from boundary (>80px), decrease radius
|
||||
retry_count += 1
|
||||
old_radius = current_radius
|
||||
current_radius *= radius_decrease_factor
|
||||
print(f'[INFO] View {i}: Mask too far from boundary (dist={min_dist}px > {min_boundary_distance}px). Decreasing radius from {old_radius:.4f} to {current_radius:.4f} (retry {retry_count}/{max_retry})')
|
||||
else:
|
||||
# Good distance, stop retrying
|
||||
if retry_count > 0:
|
||||
print(f'[INFO] View {i}: Mask boundary distance OK (dist={min_dist}px) after {retry_count} retries. Final radius: {current_radius:.4f}')
|
||||
else:
|
||||
print(f'[INFO] View {i}: Mask boundary distance OK (dist={min_dist}px). Radius: {current_radius:.4f}')
|
||||
break
|
||||
|
||||
if retry_count >= max_retry:
|
||||
print(f'[WARNING] View {i}: Max retries reached. Using final radius: {current_radius:.4f} (dist={min_dist}px)')
|
||||
|
||||
# Save camera parameters (with potentially updated radius)
|
||||
metadata = {
|
||||
"file_path": f'{i:03d}.png',
|
||||
"camera_angle_x": view['fov'],
|
||||
"transform_matrix": get_transform_matrix(cam),
|
||||
"radius": current_radius, # Save the actual radius used
|
||||
"original_radius": view['radius'], # Save original for reference
|
||||
"retries": retry_count
|
||||
}
|
||||
to_export["frames"].append(metadata)
|
||||
|
||||
# Save the camera parameters
|
||||
with open(os.path.join(arg.cond_output_folder, 'transforms.json'), 'w') as f:
|
||||
json.dump(to_export, f, indent=4)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
|
||||
parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
|
||||
parser.add_argument('--cond_views', type=str, help='JSON string of views. Contains a list of {yaw, pitch, radius, fov} object.')
|
||||
parser.add_argument('--cond_output_folder', type=str, default='/tmp', help='The path the output will be dumped to.')
|
||||
parser.add_argument('--cond_resolution', type=int, default=1024, help='Resolution of the conditional images.')
|
||||
parser.add_argument('--engine', type=str, default='CYCLES', help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
|
||||
argv = sys.argv[sys.argv.index("--") + 1:]
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
main(args)
|
||||
|
||||
@@ -0,0 +1,410 @@
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
import glob
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
from easydict import EasyDict as edict
|
||||
|
||||
|
||||
def update_metadata(path, opt):
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
timestamp = str(int(time.time()))
|
||||
os.makedirs(os.path.join(path, 'merged_records'), exist_ok=True)
|
||||
os.makedirs(os.path.join(path, 'new_records'), exist_ok=True)
|
||||
if opt.from_merged_records:
|
||||
df_files = [f for f in os.listdir(os.path.join(path, 'merged_records')) if f.endswith('.csv')]
|
||||
df_files = [f for f in df_files if int(f.split('_')[0]) >= opt.record_start]
|
||||
else:
|
||||
df_files = [f for f in os.listdir(os.path.join(path, 'new_records')) if f.startswith('part_') and f.endswith('.csv')]
|
||||
df_parts = []
|
||||
for f in df_files:
|
||||
try:
|
||||
df_parts.append(pd.read_csv(os.path.join(path, 'new_records', f)))
|
||||
except Exception as e:
|
||||
print(f"Failed to read {f}: {e}")
|
||||
if len(df_parts) > 0:
|
||||
if os.path.exists(os.path.join(path, 'metadata.csv')):
|
||||
metadata = pd.read_csv(os.path.join(path, 'metadata.csv'))
|
||||
else:
|
||||
columns = df_parts[0].columns
|
||||
metadata = pd.DataFrame(columns=columns)
|
||||
metadata.set_index('sha256', inplace=True)
|
||||
if metadata.index.duplicated().any():
|
||||
metadata = metadata.groupby(level=0).first()
|
||||
for df_part in df_parts:
|
||||
if 'sha256' in df_part.columns:
|
||||
df_part.set_index('sha256', inplace=True)
|
||||
if df_part.index.duplicated().any():
|
||||
df_part = df_part.groupby(level=0).first()
|
||||
metadata = df_part.combine_first(metadata)
|
||||
metadata.to_csv(os.path.join(path, 'metadata.csv'))
|
||||
for f in df_files:
|
||||
shutil.move(os.path.join(path, 'new_records', f), os.path.join(path, 'merged_records', f'{timestamp}_{f}'))
|
||||
return metadata
|
||||
else:
|
||||
if os.path.exists(os.path.join(path, 'metadata.csv')):
|
||||
return pd.read_csv(os.path.join(path, 'metadata.csv'))
|
||||
return None
|
||||
|
||||
|
||||
def build_downloaded_metadata_from_files(raw_root, global_metadata):
|
||||
"""Scan local files under raw_root to build download metadata.
|
||||
|
||||
Walks through raw_root to find downloaded 3D files (.glb, .obj, .fbx, .usdz, .gltf, .zip),
|
||||
matches them against global_metadata via file_identifier (uid extracted from URL) to recover
|
||||
the sha256 -> local_path mapping.
|
||||
"""
|
||||
extensions = ('.glb', '.obj', '.fbx', '.usdz', '.gltf', '.zip')
|
||||
|
||||
# Build uid -> sha256 mapping from global metadata
|
||||
uid_to_sha256 = {}
|
||||
if 'file_identifier' in global_metadata.columns:
|
||||
for _, row in global_metadata.iterrows():
|
||||
uid = str(row['file_identifier']).split('/')[-1]
|
||||
uid_to_sha256[uid] = row['sha256']
|
||||
|
||||
# Scan files
|
||||
records = []
|
||||
for dirpath, dirnames, filenames in os.walk(raw_root):
|
||||
for fname in filenames:
|
||||
if not fname.lower().endswith(extensions):
|
||||
continue
|
||||
uid = os.path.splitext(fname)[0]
|
||||
sha256 = uid_to_sha256.get(uid)
|
||||
if sha256 is not None:
|
||||
full_path = os.path.join(dirpath, fname)
|
||||
# Store path relative to parent of raw_root (i.e. download_root)
|
||||
rel_path = os.path.relpath(full_path, os.path.dirname(raw_root))
|
||||
records.append({'sha256': sha256, 'local_path': rel_path})
|
||||
|
||||
if len(records) == 0:
|
||||
return None
|
||||
|
||||
df = pd.DataFrame(records).set_index('sha256')
|
||||
print(f' [from_file] Found {len(df)} downloaded files under {raw_root}')
|
||||
|
||||
# Save as metadata.csv under raw_root
|
||||
os.makedirs(raw_root, exist_ok=True)
|
||||
df.to_csv(os.path.join(raw_root, 'metadata.csv'))
|
||||
return df
|
||||
|
||||
|
||||
# Check if directory is a multi-view directory (ending with _view or _view_fix)
|
||||
def _is_view_dir(dirname):
|
||||
return dirname.endswith('_view') or dirname.endswith('_view_fix')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--download_root', type=str, default=None,
|
||||
help='Directory to save the downloaded files')
|
||||
parser.add_argument('--thumbnail_root', type=str, default=None,
|
||||
help='Directory to save the thumbnail files')
|
||||
parser.add_argument('--render_cond_root', type=str, default=None,
|
||||
help='Directory to save the render condition files')
|
||||
parser.add_argument('--mesh_dump_root', type=str, default=None,
|
||||
help='Directory to save the mesh files')
|
||||
parser.add_argument('--pbr_dump_root', type=str, default=None,
|
||||
help='Directory to save the pbr files')
|
||||
parser.add_argument('--dual_grid_root', type=str, default=None,
|
||||
help='Directory to save the dual grid files')
|
||||
parser.add_argument('--pbr_voxel_root', type=str, default=None,
|
||||
help='Directory to save the pbr voxel files')
|
||||
parser.add_argument('--ss_latent_root', type=str, default=None,
|
||||
help='Directory to save the sparse structure latent files')
|
||||
parser.add_argument('--shape_latent_root', type=str, default=None,
|
||||
help='Directory to save the shape latent files')
|
||||
parser.add_argument('--pbr_latent_root', type=str, default=None,
|
||||
help='Directory to save the pbr latent files')
|
||||
parser.add_argument('--field', type=str, default='all',
|
||||
help='Fields to process, separated by commas')
|
||||
parser.add_argument('--from_file', action='store_true',
|
||||
help='Build metadata from file instead of from records of processings.' +
|
||||
'Useful when some processing fail to generate records but file already exists.')
|
||||
parser.add_argument('--from_merged_records', action='store_true',
|
||||
help='Build metadata from merged records')
|
||||
parser.add_argument('--record_start', type=int)
|
||||
parser.add_argument('--rebuild', action='store_true',
|
||||
help='Rebuild metadata from scratch, ignore existing metadata.')
|
||||
dataset_utils.add_args(parser)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.download_root = opt.download_root or opt.root
|
||||
opt.thumbnail_root = opt.thumbnail_root or opt.root
|
||||
opt.render_cond_root = opt.render_cond_root or opt.root
|
||||
opt.mesh_dump_root = opt.mesh_dump_root or opt.root
|
||||
opt.pbr_dump_root = opt.pbr_dump_root or opt.root
|
||||
opt.dual_grid_root = opt.dual_grid_root or opt.root
|
||||
opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
|
||||
opt.ss_latent_root = opt.ss_latent_root or opt.root
|
||||
opt.shape_latent_root = opt.shape_latent_root or opt.root
|
||||
opt.pbr_latent_root = opt.pbr_latent_root or opt.root
|
||||
|
||||
os.makedirs(opt.root, exist_ok=True)
|
||||
|
||||
opt.field = opt.field.split(',')
|
||||
|
||||
# get file list
|
||||
if os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
print('Loading previous metadata...')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv'))
|
||||
else:
|
||||
metadata = dataset_utils.get_metadata(**opt)
|
||||
metadata.to_csv(os.path.join(opt.root, 'metadata.csv'), index=False)
|
||||
|
||||
# merge downloaded
|
||||
if opt.from_file:
|
||||
downloaded_metadata = build_downloaded_metadata_from_files(
|
||||
os.path.join(opt.download_root, 'raw'), metadata)
|
||||
else:
|
||||
downloaded_metadata = update_metadata(os.path.join(opt.download_root, 'raw'), opt)
|
||||
|
||||
# merge thumbnails
|
||||
thumbnail_metadata = update_metadata(os.path.join(opt.thumbnail_root, 'thumbnails'), opt)
|
||||
|
||||
# merge aesthetic scores
|
||||
aesthetic_score_metadata = update_metadata(os.path.join(opt.root, 'aesthetic_scores'), opt)
|
||||
|
||||
# merge render conditions
|
||||
render_cond_metadata = update_metadata(os.path.join(opt.render_cond_root, 'renders_cond'), opt)
|
||||
|
||||
# merge mesh dumped
|
||||
mesh_dumped_metadata = update_metadata(os.path.join(opt.mesh_dump_root, 'mesh_dumps'), opt)
|
||||
|
||||
# merge pbr dumped
|
||||
pbr_dumped_metadata = update_metadata(os.path.join(opt.pbr_dump_root, 'pbr_dumps'), opt)
|
||||
|
||||
# merge asset stats
|
||||
asset_stats_metadata = update_metadata(os.path.join(opt.root, 'asset_stats'), opt)
|
||||
|
||||
# merge dual grid (original, no view transform)
|
||||
dual_grid_resolutions = []
|
||||
for dir in os.listdir(opt.dual_grid_root):
|
||||
if os.path.isdir(os.path.join(opt.dual_grid_root, dir)) and dir.startswith('dual_grid_') and not dir.startswith('dual_grid_view_'):
|
||||
dual_grid_resolutions.append(int(dir.split('_')[-1]))
|
||||
dual_grid_metadata = {}
|
||||
for res in dual_grid_resolutions:
|
||||
dual_grid_metadata[res] = update_metadata(os.path.join(opt.dual_grid_root, f'dual_grid_{res}'), opt)
|
||||
|
||||
# merge dual grid view (multi-view)
|
||||
dual_grid_view_resolutions = []
|
||||
for dir in os.listdir(opt.dual_grid_root):
|
||||
if os.path.isdir(os.path.join(opt.dual_grid_root, dir)) and dir.startswith('dual_grid_view_'):
|
||||
dual_grid_view_resolutions.append(int(dir.split('_')[-1]))
|
||||
dual_grid_view_metadata = {}
|
||||
for res in dual_grid_view_resolutions:
|
||||
dual_grid_view_metadata[res] = update_metadata(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}'), opt)
|
||||
|
||||
# merge pbr voxelized (single view)
|
||||
pbr_voxel_resolutions = []
|
||||
for dir in os.listdir(opt.pbr_voxel_root):
|
||||
if os.path.isdir(os.path.join(opt.pbr_voxel_root, dir)) and dir.startswith('pbr_voxels_') and not dir.startswith('pbr_voxels_view_'):
|
||||
pbr_voxel_resolutions.append(int(dir.split('_')[-1]))
|
||||
pbr_voxel_metadata = {}
|
||||
for res in pbr_voxel_resolutions:
|
||||
pbr_voxel_metadata[res] = update_metadata(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}'), opt)
|
||||
|
||||
# merge pbr voxelized view (multi-view)
|
||||
# Supports both pbr_voxels_view_{res} and pbr_voxels_view_fix_{res} directory names
|
||||
pbr_voxel_view_dirs = {} # res -> dir_name
|
||||
for dir in os.listdir(opt.pbr_voxel_root):
|
||||
if os.path.isdir(os.path.join(opt.pbr_voxel_root, dir)) and dir.startswith('pbr_voxels_view_') and not dir.startswith('pbr_voxels_view_fix_'):
|
||||
res = int(dir.split('_')[-1])
|
||||
pbr_voxel_view_dirs[res] = dir
|
||||
elif os.path.isdir(os.path.join(opt.pbr_voxel_root, dir)) and dir.startswith('pbr_voxels_view_fix_'):
|
||||
res = int(dir.split('_')[-1])
|
||||
pbr_voxel_view_dirs[res] = dir
|
||||
pbr_voxel_view_resolutions = sorted(pbr_voxel_view_dirs.keys())
|
||||
pbr_voxel_view_metadata = {}
|
||||
for res in pbr_voxel_view_resolutions:
|
||||
pbr_voxel_view_metadata[res] = update_metadata(os.path.join(opt.pbr_voxel_root, pbr_voxel_view_dirs[res]), opt)
|
||||
|
||||
# merge ss latents
|
||||
ss_latent_models = []
|
||||
if os.path.exists(os.path.join(opt.ss_latent_root, 'ss_latents')):
|
||||
ss_latent_models = os.listdir(os.path.join(opt.ss_latent_root, 'ss_latents'))
|
||||
ss_latent_metadata = {}
|
||||
for model in ss_latent_models:
|
||||
ss_latent_metadata[model] = update_metadata(os.path.join(opt.ss_latent_root, f'ss_latents/{model}'), opt)
|
||||
|
||||
# merge shape latents (original, no view transform)
|
||||
shape_latent_models = []
|
||||
if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents')):
|
||||
for dir in os.listdir(os.path.join(opt.shape_latent_root, 'shape_latents')):
|
||||
if os.path.isdir(os.path.join(opt.shape_latent_root, 'shape_latents', dir)) and not _is_view_dir(dir):
|
||||
shape_latent_models.append(dir)
|
||||
shape_latent_metadata = {}
|
||||
for model in shape_latent_models:
|
||||
shape_latent_metadata[model] = update_metadata(os.path.join(opt.shape_latent_root, f'shape_latents/{model}'), opt)
|
||||
|
||||
# merge shape latents view (multi-view, including _view and _view_fix)
|
||||
shape_latent_view_models = []
|
||||
if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents')):
|
||||
for dir in os.listdir(os.path.join(opt.shape_latent_root, 'shape_latents')):
|
||||
if os.path.isdir(os.path.join(opt.shape_latent_root, 'shape_latents', dir)) and _is_view_dir(dir):
|
||||
shape_latent_view_models.append(dir)
|
||||
shape_latent_view_metadata = {}
|
||||
for model in shape_latent_view_models:
|
||||
shape_latent_view_metadata[model] = update_metadata(os.path.join(opt.shape_latent_root, f'shape_latents/{model}'), opt)
|
||||
|
||||
# merge pbr latents (single view)
|
||||
pbr_latent_models = []
|
||||
if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
|
||||
for dir in os.listdir(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
|
||||
if os.path.isdir(os.path.join(opt.pbr_latent_root, 'pbr_latents', dir)) and not _is_view_dir(dir):
|
||||
pbr_latent_models.append(dir)
|
||||
pbr_latent_metadata = {}
|
||||
for model in pbr_latent_models:
|
||||
pbr_latent_metadata[model] = update_metadata(os.path.join(opt.pbr_latent_root, f'pbr_latents/{model}'), opt)
|
||||
|
||||
# merge pbr latents view (multi-view, including _view and _view_fix)
|
||||
pbr_latent_view_models = []
|
||||
if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
|
||||
for dir in os.listdir(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
|
||||
if os.path.isdir(os.path.join(opt.pbr_latent_root, 'pbr_latents', dir)) and _is_view_dir(dir):
|
||||
pbr_latent_view_models.append(dir)
|
||||
pbr_latent_view_metadata = {}
|
||||
for model in pbr_latent_view_models:
|
||||
pbr_latent_view_metadata[model] = update_metadata(os.path.join(opt.pbr_latent_root, f'pbr_latents/{model}'), opt)
|
||||
|
||||
# Merge all sub-metadata back into main metadata and save
|
||||
metadata = metadata.set_index('sha256')
|
||||
sub_metadata_list = [
|
||||
downloaded_metadata,
|
||||
thumbnail_metadata,
|
||||
aesthetic_score_metadata,
|
||||
render_cond_metadata,
|
||||
mesh_dumped_metadata,
|
||||
pbr_dumped_metadata,
|
||||
asset_stats_metadata,
|
||||
]
|
||||
for res in dual_grid_resolutions:
|
||||
sub_metadata_list.append(dual_grid_metadata.get(res))
|
||||
for res in dual_grid_view_resolutions:
|
||||
sub_metadata_list.append(dual_grid_view_metadata.get(res))
|
||||
for res in pbr_voxel_resolutions:
|
||||
sub_metadata_list.append(pbr_voxel_metadata.get(res))
|
||||
for res in pbr_voxel_view_resolutions:
|
||||
sub_metadata_list.append(pbr_voxel_view_metadata.get(res))
|
||||
for model in ss_latent_models:
|
||||
sub_metadata_list.append(ss_latent_metadata.get(model))
|
||||
for model in shape_latent_models:
|
||||
sub_metadata_list.append(shape_latent_metadata.get(model))
|
||||
for model in shape_latent_view_models:
|
||||
sub_metadata_list.append(shape_latent_view_metadata.get(model))
|
||||
for model in pbr_latent_models:
|
||||
sub_metadata_list.append(pbr_latent_metadata.get(model))
|
||||
for model in pbr_latent_view_models:
|
||||
sub_metadata_list.append(pbr_latent_view_metadata.get(model))
|
||||
if metadata.index.duplicated().any():
|
||||
metadata = metadata.groupby(level=0).first()
|
||||
for sub in sub_metadata_list:
|
||||
if sub is not None:
|
||||
if 'sha256' in sub.columns:
|
||||
sub = sub.set_index('sha256')
|
||||
if sub.index.duplicated().any():
|
||||
sub = sub.groupby(level=0).first()
|
||||
metadata = metadata.combine_first(sub)
|
||||
metadata = metadata.reset_index()
|
||||
metadata.to_csv(os.path.join(opt.root, 'metadata.csv'), index=False)
|
||||
print(f'Saved merged metadata with {len(metadata)} entries and columns: {list(metadata.columns)}')
|
||||
|
||||
# statistics
|
||||
num_downloaded = downloaded_metadata['local_path'].count() if downloaded_metadata is not None else 0
|
||||
with open(os.path.join(opt.root, 'statistics.txt'), 'w') as f:
|
||||
f.write('Statistics:\n')
|
||||
f.write(f' - Number of assets: {len(metadata)}\n')
|
||||
f.write(f' - Number of assets downloaded: {num_downloaded}\n')
|
||||
if thumbnail_metadata is not None:
|
||||
f.write(f' - Number of assets with thumbnails: {thumbnail_metadata["thumbnailed"].sum()}\n')
|
||||
if aesthetic_score_metadata is not None:
|
||||
f.write(f' - Number of assets with aesthetic scores: {aesthetic_score_metadata["aesthetic_score"].count()}\n')
|
||||
if render_cond_metadata is not None:
|
||||
f.write(f' - Number of assets with render conditions: {render_cond_metadata["cond_rendered"].count()}\n')
|
||||
if mesh_dumped_metadata is not None:
|
||||
f.write(f' - Number of assets with mesh dumped: {mesh_dumped_metadata["mesh_dumped"].sum()}\n')
|
||||
if pbr_dumped_metadata is not None:
|
||||
f.write(f' - Number of assets with PBR dumped: {pbr_dumped_metadata["pbr_dumped"].sum()}\n')
|
||||
if asset_stats_metadata is not None:
|
||||
f.write(f' - Number of assets with asset stats: {len(asset_stats_metadata)}\n')
|
||||
if len(dual_grid_resolutions) != 0:
|
||||
f.write(f' - Number of assets with dual grid:\n')
|
||||
for res in dual_grid_resolutions:
|
||||
if dual_grid_metadata[res] is not None:
|
||||
f.write(f' - {res}: {dual_grid_metadata[res]["dual_grid_converted"].sum()}\n')
|
||||
if len(dual_grid_view_resolutions) != 0:
|
||||
f.write(f' - Number of assets with dual grid view:\n')
|
||||
for res in dual_grid_view_resolutions:
|
||||
if dual_grid_view_metadata[res] is not None:
|
||||
col_name = f'dual_grid_view00_converted_{res}'
|
||||
if col_name in dual_grid_view_metadata[res].columns:
|
||||
f.write(f' - {res}: {dual_grid_view_metadata[res][col_name].sum()}\n')
|
||||
else:
|
||||
f.write(f' - {res}: {len(dual_grid_view_metadata[res])}\n')
|
||||
if len(pbr_voxel_resolutions) != 0:
|
||||
f.write(f' - Number of assets with PBR voxelization:\n')
|
||||
for res in sorted(pbr_voxel_resolutions):
|
||||
if pbr_voxel_metadata[res] is not None:
|
||||
f.write(f' - {res}: {pbr_voxel_metadata[res]["pbr_voxelized"].sum()}\n')
|
||||
if len(pbr_voxel_view_resolutions) != 0:
|
||||
f.write(f' - Number of assets with PBR voxelization view:\n')
|
||||
for res in sorted(pbr_voxel_view_resolutions):
|
||||
if pbr_voxel_view_metadata[res] is not None:
|
||||
dir_name = pbr_voxel_view_dirs[res]
|
||||
col_name_old = 'pbr_voxelized_view00'
|
||||
col_name_new = f'pbr_voxelized_view_fix00_{res}'
|
||||
if col_name_old in pbr_voxel_view_metadata[res].columns:
|
||||
f.write(f' - {dir_name}: {pbr_voxel_view_metadata[res][col_name_old].sum()}\n')
|
||||
elif col_name_new in pbr_voxel_view_metadata[res].columns:
|
||||
f.write(f' - {dir_name}: {pbr_voxel_view_metadata[res][col_name_new].sum()}\n')
|
||||
else:
|
||||
f.write(f' - {dir_name}: {len(pbr_voxel_view_metadata[res])}\n')
|
||||
if len(ss_latent_models) != 0:
|
||||
f.write(f' - Number of assets with sparse structure latents:\n')
|
||||
for model in ss_latent_models:
|
||||
if ss_latent_metadata[model] is not None:
|
||||
if 'ss_latent_encoded' in ss_latent_metadata[model].columns:
|
||||
f.write(f' - {model}: {ss_latent_metadata[model]["ss_latent_encoded"].sum()}\n')
|
||||
elif 'ss_latent_view00_encoded' in ss_latent_metadata[model].columns:
|
||||
f.write(f' - {model}: {ss_latent_metadata[model]["ss_latent_view00_encoded"].sum()}\n')
|
||||
else:
|
||||
f.write(f' - {model}: {len(ss_latent_metadata[model])}\n')
|
||||
if len(shape_latent_models) != 0:
|
||||
f.write(f' - Number of assets with shape latents:\n')
|
||||
for model in shape_latent_models:
|
||||
if shape_latent_metadata[model] is not None:
|
||||
f.write(f' - {model}: {shape_latent_metadata[model]["shape_latent_encoded"].sum()}\n')
|
||||
if len(shape_latent_view_models) != 0:
|
||||
f.write(f' - Number of assets with shape latents view:\n')
|
||||
for model in shape_latent_view_models:
|
||||
if shape_latent_view_metadata[model] is not None:
|
||||
col_name = 'shape_latent_view00_encoded'
|
||||
if col_name in shape_latent_view_metadata[model].columns:
|
||||
f.write(f' - {model}: {shape_latent_view_metadata[model][col_name].sum()}\n')
|
||||
else:
|
||||
f.write(f' - {model}: {len(shape_latent_view_metadata[model])}\n')
|
||||
if len(pbr_latent_models) != 0:
|
||||
f.write(f' - Number of assets with PBR latents:\n')
|
||||
for model in pbr_latent_models:
|
||||
if pbr_latent_metadata[model] is not None:
|
||||
f.write(f' - {model}: {pbr_latent_metadata[model]["pbr_latent_encoded"].sum()}\n')
|
||||
if len(pbr_latent_view_models) != 0:
|
||||
f.write(f' - Number of assets with PBR latents view:\n')
|
||||
for model in pbr_latent_view_models:
|
||||
if pbr_latent_view_metadata[model] is not None:
|
||||
col_name = 'pbr_latent_view00_encoded'
|
||||
if col_name in pbr_latent_view_metadata[model].columns:
|
||||
f.write(f' - {model}: {pbr_latent_view_metadata[model][col_name].sum()}\n')
|
||||
else:
|
||||
f.write(f' - {model}: {len(pbr_latent_view_metadata[model])}\n')
|
||||
|
||||
with open(os.path.join(opt.root, 'statistics.txt'), 'r') as f:
|
||||
print(f.read())
|
||||
@@ -0,0 +1,102 @@
|
||||
"""
|
||||
Build metadata.csv for TexVerse dataset from metadata.json.
|
||||
|
||||
Output format matches ABO metadata.csv:
|
||||
sha256, file_identifier, local_path
|
||||
|
||||
Usage:
|
||||
python data_toolkit/build_texverse_metadata.py \
|
||||
--metadata_json /path/to/TexVerse/metadata.json \
|
||||
--data_root /path/to/TexVerse \
|
||||
--output /path/to/TexVerse/metadata.csv \
|
||||
--max_workers 32
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
import argparse
|
||||
import hashlib
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def get_file_hash(file: str) -> str:
|
||||
sha256 = hashlib.sha256()
|
||||
with open(file, "rb") as f:
|
||||
for byte_block in iter(lambda: f.read(4096), b""):
|
||||
sha256.update(byte_block)
|
||||
return sha256.hexdigest()
|
||||
|
||||
|
||||
def process_one(uid, glb_paths, data_root):
|
||||
"""Find the first existing GLB file for this uid and compute its sha256."""
|
||||
for rel_path in glb_paths:
|
||||
full_path = os.path.join(data_root, rel_path)
|
||||
if os.path.exists(full_path):
|
||||
sha256 = get_file_hash(full_path)
|
||||
return {
|
||||
'sha256': sha256,
|
||||
'file_identifier': uid,
|
||||
'local_path': rel_path,
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Build metadata.csv for TexVerse')
|
||||
parser.add_argument('--metadata_json', type=str, required=True,
|
||||
help='Path to TexVerse metadata.json')
|
||||
parser.add_argument('--data_root', type=str, required=True,
|
||||
help='Root directory of TexVerse dataset (where glbs/ is)')
|
||||
parser.add_argument('--output', type=str, required=True,
|
||||
help='Output path for metadata.csv')
|
||||
parser.add_argument('--max_workers', type=int, default=16,
|
||||
help='Number of parallel workers')
|
||||
args = parser.parse_args()
|
||||
|
||||
with open(args.metadata_json, 'r') as f:
|
||||
metadata = json.load(f)
|
||||
|
||||
print(f'Total entries in metadata.json: {len(metadata)}')
|
||||
|
||||
# Load existing metadata.csv and skip already processed entries
|
||||
existing_uids = set()
|
||||
existing_records = []
|
||||
if os.path.exists(args.output):
|
||||
existing_df = pd.read_csv(args.output)
|
||||
existing_uids = set(existing_df['file_identifier'].values)
|
||||
existing_records = existing_df.to_dict('records')
|
||||
print(f'Found existing metadata.csv with {len(existing_uids)} entries, skipping them')
|
||||
|
||||
# Filter out already processed uids
|
||||
to_process = {uid: info for uid, info in metadata.items() if uid not in existing_uids}
|
||||
print(f'New entries to process: {len(to_process)}')
|
||||
|
||||
if len(to_process) == 0:
|
||||
print('Nothing to do, all entries already exist.')
|
||||
return
|
||||
|
||||
new_records = []
|
||||
with ProcessPoolExecutor(max_workers=args.max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(process_one, uid, info['glb_paths'], args.data_root): uid
|
||||
for uid, info in to_process.items()
|
||||
}
|
||||
for future in tqdm(as_completed(futures), total=len(futures), desc='Building metadata'):
|
||||
try:
|
||||
result = future.result()
|
||||
if result is not None:
|
||||
new_records.append(result)
|
||||
except Exception as e:
|
||||
uid = futures[future]
|
||||
print(f'Error processing {uid}: {e}')
|
||||
|
||||
all_records = existing_records + new_records
|
||||
df = pd.DataFrame.from_records(all_records, columns=['sha256', 'file_identifier', 'local_path'])
|
||||
df.to_csv(args.output, index=False)
|
||||
print(f'Added {len(new_records)} new entries, total {len(df)} entries in {args.output}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,120 @@
|
||||
import os
|
||||
import re
|
||||
import argparse
|
||||
import tarfile
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
from utils import get_file_hash
|
||||
|
||||
|
||||
def add_args(parser: argparse.ArgumentParser):
|
||||
pass
|
||||
|
||||
|
||||
def get_metadata(**kwargs):
|
||||
metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ABO.csv")
|
||||
return metadata
|
||||
|
||||
|
||||
def download(metadata, root, **kwargs):
|
||||
output_dir = root
|
||||
os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
|
||||
|
||||
if not os.path.exists(os.path.join(output_dir, 'raw', 'abo-3dmodels.tar')):
|
||||
try:
|
||||
os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
|
||||
os.system(f"wget -O {output_dir}/raw/abo-3dmodels.tar https://amazon-berkeley-objects.s3.amazonaws.com/archives/abo-3dmodels.tar")
|
||||
except:
|
||||
print("\033[93m")
|
||||
print("Error downloading ABO dataset. Please check your internet connection and try again.")
|
||||
print("Or, you can manually download the abo-3dmodels.tar file and place it in the {output_dir}/raw directory")
|
||||
print("Visit https://amazon-berkeley-objects.s3.amazonaws.com/index.html for more information")
|
||||
print("\033[0m")
|
||||
raise FileNotFoundError("Error downloading ABO dataset")
|
||||
|
||||
downloaded = {}
|
||||
metadata = metadata.set_index("file_identifier")
|
||||
with tarfile.open(os.path.join(output_dir, 'raw', 'abo-3dmodels.tar')) as tar:
|
||||
with ThreadPoolExecutor(max_workers=1) as executor, \
|
||||
tqdm(total=len(metadata), desc="Extracting") as pbar:
|
||||
def worker(instance: str) -> str:
|
||||
try:
|
||||
tar.extract(f"3dmodels/original/{instance}", path=os.path.join(output_dir, 'raw'))
|
||||
sha256 = get_file_hash(os.path.join(output_dir, 'raw/3dmodels/original', instance))
|
||||
pbar.update()
|
||||
return sha256
|
||||
except Exception as e:
|
||||
pbar.update()
|
||||
print(f"Error extracting for {instance}: {e}")
|
||||
return None
|
||||
|
||||
sha256s = executor.map(worker, metadata.index)
|
||||
executor.shutdown(wait=True)
|
||||
|
||||
for k, sha256 in zip(metadata.index, sha256s):
|
||||
if sha256 is not None:
|
||||
if sha256 == metadata.loc[k, "sha256"]:
|
||||
downloaded[sha256] = os.path.join('raw/3dmodels/original', k)
|
||||
else:
|
||||
print(f"Error downloading {k}: sha256s do not match")
|
||||
|
||||
return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
|
||||
|
||||
|
||||
def _process_instance(args):
|
||||
"""Worker function for ProcessPoolExecutor (must be top-level for pickling)"""
|
||||
import os
|
||||
metadatum, output_dir, func = args
|
||||
try:
|
||||
local_path = metadatum['local_path']
|
||||
sha256 = metadatum['sha256']
|
||||
file = os.path.join(output_dir, local_path)
|
||||
record = func(file, sha256)
|
||||
return record
|
||||
except Exception as e:
|
||||
print(f"Error processing object {metadatum.get('sha256', '?')}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
|
||||
import os
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from tqdm import tqdm
|
||||
|
||||
# load metadata
|
||||
metadata = metadata.to_dict('records')
|
||||
|
||||
max_workers = max_workers or os.cpu_count()
|
||||
records = []
|
||||
|
||||
# Track processed/skipped counts
|
||||
total_processed = 0
|
||||
total_skipped = 0
|
||||
|
||||
try:
|
||||
with ProcessPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(_process_instance, (m, output_dir, func)): m['sha256']
|
||||
for m in metadata
|
||||
}
|
||||
pbar = tqdm(as_completed(futures), total=len(futures), desc=desc)
|
||||
for future in pbar:
|
||||
try:
|
||||
r = future.result()
|
||||
if r is not None:
|
||||
records.append(r)
|
||||
# Update stats
|
||||
if '_processed_count' in r:
|
||||
total_processed += r['_processed_count']
|
||||
if '_skipped_count' in r:
|
||||
total_skipped += r['_skipped_count']
|
||||
# Update progress bar display
|
||||
pbar.set_postfix(processed=total_processed, skipped=total_skipped, refresh=False)
|
||||
except Exception as e:
|
||||
sha256 = futures[future]
|
||||
print(f"Error processing object {sha256}: {e}")
|
||||
except Exception as e:
|
||||
print(f"Error happened during processing: {e}")
|
||||
|
||||
return pd.DataFrame.from_records(records)
|
||||
@@ -0,0 +1,133 @@
|
||||
import os
|
||||
import argparse
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import objaverse.xl as oxl
|
||||
from utils import get_file_hash
|
||||
|
||||
|
||||
def add_args(parser: argparse.ArgumentParser):
|
||||
parser.add_argument('--source', type=str, default='sketchfab',
|
||||
help='Data source to download annotations from (github, sketchfab)')
|
||||
|
||||
|
||||
def get_metadata(source, **kwargs):
|
||||
if source == 'sketchfab':
|
||||
metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_sketchfab.csv")
|
||||
elif source == 'github':
|
||||
metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_github.csv")
|
||||
else:
|
||||
raise ValueError(f"Invalid source: {source}")
|
||||
return metadata
|
||||
|
||||
|
||||
def download(metadata, output_dir, **kwargs):
|
||||
os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
|
||||
|
||||
# download annotations
|
||||
annotations = oxl.get_annotations()
|
||||
annotations = annotations[annotations['sha256'].isin(metadata['sha256'].values)]
|
||||
|
||||
# download and render objects
|
||||
file_paths = oxl.download_objects(
|
||||
annotations,
|
||||
download_dir=os.path.join(output_dir, "raw"),
|
||||
save_repo_format="zip",
|
||||
)
|
||||
|
||||
downloaded = {}
|
||||
metadata = metadata.set_index("file_identifier")
|
||||
for k, v in file_paths.items():
|
||||
sha256 = metadata.loc[k, "sha256"]
|
||||
downloaded[sha256] = os.path.relpath(v, output_dir)
|
||||
|
||||
return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
|
||||
|
||||
|
||||
def _process_instance(args):
|
||||
"""Worker function for ProcessPoolExecutor (must be top-level for pickling)"""
|
||||
import os, tempfile, zipfile
|
||||
metadatum, output_dir, func = args
|
||||
try:
|
||||
local_path = metadatum['local_path']
|
||||
sha256 = metadatum['sha256']
|
||||
|
||||
direct_file_path = os.path.join(output_dir, local_path)
|
||||
if os.path.exists(direct_file_path):
|
||||
file = direct_file_path
|
||||
record = func(file, sha256)
|
||||
elif local_path.startswith('raw/github/repos/'):
|
||||
path_parts = local_path.split('/')
|
||||
file_name = os.path.join(*path_parts[5:])
|
||||
zip_file = os.path.join(output_dir, *path_parts[:5])
|
||||
if os.path.exists(zip_file):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
with zipfile.ZipFile(zip_file, 'r') as zip_ref:
|
||||
zip_ref.extractall(tmp_dir)
|
||||
file = os.path.join(tmp_dir, file_name)
|
||||
record = func(file, sha256)
|
||||
else:
|
||||
# zip file not found, pass local_path directly (for tasks like dual_grid_view that don't need the original file)
|
||||
file = local_path
|
||||
record = func(file, sha256)
|
||||
else:
|
||||
file = os.path.join(output_dir, local_path)
|
||||
record = func(file, sha256)
|
||||
return record
|
||||
except Exception as e:
|
||||
print(f"Error processing object {metadatum.get('sha256', '?')}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects', log_interval=500, timeout=None) -> pd.DataFrame:
|
||||
print("================")
|
||||
import os
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed, TimeoutError
|
||||
from tqdm import tqdm
|
||||
|
||||
# load metadata
|
||||
metadata = metadata.to_dict('records')
|
||||
|
||||
max_workers = max_workers or os.cpu_count()
|
||||
records = []
|
||||
|
||||
# Track processed/skipped counts
|
||||
total_processed = 0
|
||||
total_skipped = 0
|
||||
timeout_count = 0
|
||||
|
||||
try:
|
||||
with ProcessPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(_process_instance, (m, output_dir, func)): m['sha256']
|
||||
for m in metadata
|
||||
}
|
||||
pbar = tqdm(as_completed(futures), total=len(futures), desc=desc)
|
||||
for future in pbar:
|
||||
sha256 = futures[future]
|
||||
try:
|
||||
r = future.result(timeout=timeout)
|
||||
if r is not None:
|
||||
records.append(r)
|
||||
# Update stats
|
||||
if '_processed_count' in r:
|
||||
total_processed += r['_processed_count']
|
||||
if '_skipped_count' in r:
|
||||
total_skipped += r['_skipped_count']
|
||||
# Update progress bar display
|
||||
pbar.set_postfix(processed=total_processed, skipped=total_skipped, timeout=timeout_count, refresh=False)
|
||||
except TimeoutError:
|
||||
timeout_count += 1
|
||||
print(f"Timeout processing object {sha256} (>{timeout}s)")
|
||||
records.append({'sha256': sha256, 'error': f'Timeout (>{timeout}s)'})
|
||||
pbar.set_postfix(processed=total_processed, skipped=total_skipped, timeout=timeout_count, refresh=False)
|
||||
except Exception as e:
|
||||
print(f"Error processing object {sha256}: {e}")
|
||||
except Exception as e:
|
||||
print(f"Error happened during processing: {e}")
|
||||
|
||||
if timeout_count > 0:
|
||||
print(f"Total timeout: {timeout_count} objects")
|
||||
|
||||
return pd.DataFrame.from_records(records)
|
||||
@@ -0,0 +1,69 @@
|
||||
import os
|
||||
import argparse
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def add_args(parser: argparse.ArgumentParser):
|
||||
pass
|
||||
|
||||
|
||||
def get_metadata(**kwargs):
|
||||
raise NotImplementedError("TexVerse metadata should be built from metadata.json using build_texverse_metadata.py")
|
||||
|
||||
|
||||
def download(metadata, root, **kwargs):
|
||||
raise NotImplementedError("TexVerse GLB files are already available locally. No download needed.")
|
||||
|
||||
|
||||
def _process_instance(args):
|
||||
"""Worker function for ProcessPoolExecutor (must be top-level for pickling)"""
|
||||
metadatum, output_dir, func = args
|
||||
try:
|
||||
local_path = metadatum['local_path']
|
||||
sha256 = metadatum['sha256']
|
||||
file = os.path.join(output_dir, local_path)
|
||||
record = func(file, sha256)
|
||||
return record
|
||||
except Exception as e:
|
||||
print(f"Error processing object {metadatum.get('sha256', '?')}: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects', timeout=None) -> pd.DataFrame:
|
||||
import os
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed, TimeoutError
|
||||
from tqdm import tqdm
|
||||
|
||||
metadata = metadata.to_dict('records')
|
||||
|
||||
max_workers = max_workers or os.cpu_count()
|
||||
records = []
|
||||
timeout_count = 0
|
||||
|
||||
try:
|
||||
with ProcessPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = {
|
||||
executor.submit(_process_instance, (m, output_dir, func)): m['sha256']
|
||||
for m in metadata
|
||||
}
|
||||
for future in tqdm(as_completed(futures), total=len(futures), desc=desc):
|
||||
sha256 = futures[future]
|
||||
try:
|
||||
r = future.result(timeout=timeout)
|
||||
if r is not None:
|
||||
records.append(r)
|
||||
except TimeoutError:
|
||||
timeout_count += 1
|
||||
print(f"Timeout processing object {sha256} (>{timeout}s)")
|
||||
records.append({'sha256': sha256, 'error': f'Timeout (>{timeout}s)'})
|
||||
except Exception as e:
|
||||
print(f"Error processing object {sha256}: {e}")
|
||||
except Exception as e:
|
||||
print(f"Error happened during processing: {e}")
|
||||
|
||||
if timeout_count > 0:
|
||||
print(f"Total timeout: {timeout_count} objects")
|
||||
|
||||
return pd.DataFrame.from_records(records)
|
||||
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
import copy
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
from easydict import EasyDict as edict
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--download_root', type=str, default=None,
|
||||
help='Directory to download the objects')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--check_only', action='store_true',
|
||||
help='Only check if the objects are already downloaded')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.download_root = opt.download_root or opt.root
|
||||
|
||||
os.makedirs(opt.root, exist_ok=True)
|
||||
os.makedirs(opt.download_root, exist_ok=True)
|
||||
os.makedirs(os.path.join(opt.download_root, 'raw', 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
if 'local_path' in metadata.columns:
|
||||
metadata = metadata[metadata['local_path'].isna()]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
downloaded = dataset_utils.download(metadata, output_dir=opt.download_root, **opt)
|
||||
downloaded.to_csv(os.path.join(opt.download_root, 'raw', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,170 @@
|
||||
import os
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
import torch
|
||||
import pickle
|
||||
import o_voxel
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
|
||||
|
||||
def _dual_grid_mesh(file, metadatum, mesh_dump_root, root):
|
||||
sha256 = metadatum['sha256']
|
||||
try:
|
||||
pack = {'sha256': sha256}
|
||||
data = None
|
||||
for res in opt.resolution:
|
||||
need_process = False
|
||||
|
||||
# check if already processed
|
||||
if os.path.exists(os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz')):
|
||||
try:
|
||||
info = o_voxel.io.read_vxz_info(os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz'))
|
||||
pack[f'dual_grid_converted_{res}'] = True
|
||||
pack[f'dual_grid_size_{res}'] = info['num_voxel']
|
||||
except Exception as e:
|
||||
print(f'Error reading {sha256}.vxz: {e}')
|
||||
need_process = True
|
||||
else:
|
||||
need_process = True
|
||||
|
||||
# process mesh
|
||||
if need_process:
|
||||
if data is None:
|
||||
with open(os.path.join(mesh_dump_root, 'mesh_dumps', f'{sha256}.pickle'), 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
start = 0
|
||||
vertices = []
|
||||
faces = []
|
||||
for obj in dump['objects']:
|
||||
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
||||
continue
|
||||
vertices.append(obj['vertices'])
|
||||
faces.append(obj['faces'] + start)
|
||||
start += len(obj['vertices'])
|
||||
vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
|
||||
faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
|
||||
vertices_min = vertices.min(dim=0)[0]
|
||||
vertices_max = vertices.max(dim=0)[0]
|
||||
center = (vertices_min + vertices_max) / 2
|
||||
scale = 0.99999 / (vertices_max - vertices_min).max()
|
||||
vertices = (vertices - center) * scale
|
||||
assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
|
||||
data = {'vertices': vertices, 'faces': faces}
|
||||
|
||||
voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
|
||||
**data,
|
||||
grid_size=res,
|
||||
aabb=[[-0.5,-0.5,-0.5],[0.5,0.5,0.5]],
|
||||
face_weight=1.0,
|
||||
boundary_weight=0.2,
|
||||
regularization_weight=1e-2,
|
||||
timing=False,
|
||||
)
|
||||
dual_vertices = dual_vertices * res - voxel_indices
|
||||
assert torch.all(dual_vertices >= -1e-3) and torch.all(dual_vertices <= 1+1e-3), 'dual_vertices out of range'
|
||||
dual_vertices = torch.clamp(dual_vertices, 0, 1)
|
||||
dual_vertices = (dual_vertices * 255).type(torch.uint8)
|
||||
intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8)
|
||||
|
||||
o_voxel.io.write_vxz(
|
||||
os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz'),
|
||||
voxel_indices,
|
||||
{'vertices': dual_vertices, 'intersected': intersected},
|
||||
)
|
||||
|
||||
pack[f'dual_grid_converted_{res}'] = True
|
||||
pack[f'dual_grid_size_{res}'] = len(dual_vertices)
|
||||
|
||||
return pack
|
||||
except Exception as e:
|
||||
print(f'Error processing {sha256}: {e}')
|
||||
return {'sha256': sha256, 'error': str(e)}
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--mesh_dump_root', type=str, default=None,
|
||||
help='Directory to load mesh dumps')
|
||||
parser.add_argument('--dual_grid_root', type=str, default=None,
|
||||
help='Directory to save dual grids')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--resolution', type=str, default=256)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.resolution = [int(x) for x in opt.resolution.split(',')]
|
||||
opt.mesh_dump_root = opt.mesh_dump_root or opt.root
|
||||
opt.dual_grid_root = opt.dual_grid_root or opt.root
|
||||
|
||||
for res in opt.resolution:
|
||||
os.makedirs(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
for res in opt.resolution:
|
||||
if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'metadata.csv')):
|
||||
dual_grid_metadata = pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'metadata.csv')).set_index('sha256')
|
||||
dual_grid_metadata = dual_grid_metadata.rename(columns={'dual_grid_converted': f'dual_grid_converted_{res}', 'dual_grid_size': f'dual_grid_size_{res}'})
|
||||
metadata = metadata.combine_first(dual_grid_metadata)
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['mesh_dumped'] == True]
|
||||
mask = np.zeros(len(metadata), dtype=bool)
|
||||
for res in opt.resolution:
|
||||
if f'dual_grid_converted_{res}' in metadata.columns:
|
||||
mask |= metadata[f'dual_grid_converted_{res}'] != True
|
||||
else:
|
||||
mask[:] = True
|
||||
break
|
||||
metadata = metadata[mask]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
func = partial(_dual_grid_mesh, root=opt.dual_grid_root, mesh_dump_root=opt.mesh_dump_root)
|
||||
dual_grids = dataset_utils.foreach_instance(metadata, None, func, max_workers=opt.max_workers, no_file=True, desc='Dual griding')
|
||||
if 'error' in dual_grids.columns:
|
||||
errors = dual_grids[dual_grids[f'error'].notna()]
|
||||
with open('errors.txt', 'w') as f:
|
||||
f.write('\n'.join(errors['sha256'].tolist()))
|
||||
for res in opt.resolution:
|
||||
if f'dual_grid_converted_{res}' in dual_grids.columns:
|
||||
dual_grid_metadata = dual_grids[dual_grids[f'dual_grid_converted_{res}'] == True]
|
||||
if len(dual_grid_metadata) > 0:
|
||||
dual_grid_metadata = dual_grid_metadata[['sha256', f'dual_grid_converted_{res}', f'dual_grid_size_{res}']]
|
||||
dual_grid_metadata = dual_grid_metadata.rename(columns={f'dual_grid_converted_{res}': 'dual_grid_converted', f'dual_grid_size_{res}': 'dual_grid_size'})
|
||||
dual_grid_metadata.to_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
|
||||
@@ -0,0 +1,350 @@
|
||||
"""
|
||||
dual_grid_view.py - Multi-view transform dual grid processing
|
||||
Extends dual_grid.py with scale and mesh rotation logic
|
||||
Based on test_ovoxel_transform.py implementation
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import json
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
import torch
|
||||
import pickle
|
||||
import o_voxel
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
from utils import get_new_camera_matrix, transform_mesh, sphere_normalize_torch
|
||||
|
||||
|
||||
def _dual_grid_mesh_view(file, sha256, mesh_dump_root, transform_root, root, view_indices=None):
|
||||
"""
|
||||
Process multi-view dual grid conversion for a single sha256.
|
||||
|
||||
Args:
|
||||
file: local_path from metadata
|
||||
sha256: sha256 string
|
||||
mesh_dump_root: directory containing mesh dump files
|
||||
transform_root: directory containing transform json files
|
||||
root: output directory for dual grids
|
||||
view_indices: list of view indices to process, None for all views
|
||||
"""
|
||||
try:
|
||||
pack = {'sha256': sha256}
|
||||
vertices_sphere = None
|
||||
sphere_radius = None
|
||||
faces = None
|
||||
|
||||
# Load transforms
|
||||
transform_path = os.path.join(transform_root, sha256, 'transforms.json')
|
||||
if not os.path.exists(transform_path):
|
||||
print(f'Transform file not found for {sha256}, skipping')
|
||||
return {'sha256': sha256, 'error': 'Transform file not found'}
|
||||
|
||||
with open(transform_path, 'r') as f:
|
||||
transforms_json = json.load(f)
|
||||
transform_mats = transforms_json['frames']
|
||||
|
||||
# Determine views to process
|
||||
if view_indices is None:
|
||||
view_indices = list(range(len(transform_mats)))
|
||||
else:
|
||||
view_indices = [i for i in view_indices if i < len(transform_mats)]
|
||||
|
||||
# Track processed and skipped counts
|
||||
processed_count = 0
|
||||
skipped_count = 0
|
||||
|
||||
for view_idx in view_indices:
|
||||
for res in opt.resolution:
|
||||
need_process = False
|
||||
|
||||
# Check if already processed
|
||||
# Path structure: dual_grid_view_{res}/{sha256}/view{idx:02d}.vxz
|
||||
sha256_dir = os.path.join(root, f'dual_grid_view_{res}', sha256)
|
||||
vxz_path = os.path.join(sha256_dir, f'view{view_idx:02d}.vxz')
|
||||
if os.path.exists(vxz_path):
|
||||
try:
|
||||
info = o_voxel.io.read_vxz_info(vxz_path)
|
||||
pack[f'dual_grid_view{view_idx:02d}_converted_{res}'] = True
|
||||
pack[f'dual_grid_view{view_idx:02d}_size_{res}'] = info['num_voxel']
|
||||
skipped_count += 1
|
||||
except Exception as e:
|
||||
print(f'Error reading {sha256}/view{view_idx:02d}.vxz: {e}')
|
||||
need_process = True
|
||||
else:
|
||||
need_process = True
|
||||
|
||||
# Process mesh
|
||||
if need_process:
|
||||
# Lazy load mesh
|
||||
if vertices_sphere is None:
|
||||
mesh_file = os.path.join(mesh_dump_root, 'mesh_dumps', f'{sha256}.pickle')
|
||||
if not os.path.exists(mesh_file):
|
||||
print(f'Mesh dump not found for {sha256}, skipping')
|
||||
return {'sha256': sha256, 'error': 'Mesh dump not found'}
|
||||
|
||||
with open(mesh_file, 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
|
||||
start = 0
|
||||
vertices_list = []
|
||||
faces_list = []
|
||||
for obj in dump['objects']:
|
||||
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
||||
continue
|
||||
vertices_list.append(obj['vertices'])
|
||||
faces_list.append(obj['faces'] + start)
|
||||
start += len(obj['vertices'])
|
||||
|
||||
if len(vertices_list) == 0:
|
||||
print(f'No valid mesh data for {sha256}, skipping')
|
||||
return {'sha256': sha256, 'error': 'No valid mesh data'}
|
||||
|
||||
vertices = torch.from_numpy(np.concatenate(vertices_list, axis=0)).float().contiguous()
|
||||
faces = torch.from_numpy(np.concatenate(faces_list, axis=0)).long().contiguous()
|
||||
|
||||
# Sphere normalization (for multi-view transform) - CPU only
|
||||
vertices_sphere, sphere_center, sphere_radius = sphere_normalize_torch(vertices)
|
||||
|
||||
# Get transform for current view
|
||||
transform = transform_mats[view_idx]
|
||||
|
||||
# Multi-view transform - CPU only
|
||||
transformed_vertices = transform_mesh(vertices_sphere, transform)
|
||||
|
||||
# Post-transform normalization: scale by abs max to [-0.5, 0.5]^3
|
||||
# Only scale, no center shift, to preserve relative model position
|
||||
abs_max = transformed_vertices.abs().max().item()
|
||||
box_scale = 0.49999 / abs_max # Normalize to [-0.5, 0.5] range
|
||||
transformed_normalized = transformed_vertices * box_scale
|
||||
transformed_normalized_cpu = transformed_normalized.contiguous()
|
||||
|
||||
# Compute total scale (from original mesh to final normalized mesh)
|
||||
total_scale = box_scale / sphere_radius.item()
|
||||
|
||||
# Validate range
|
||||
assert torch.all(transformed_normalized_cpu >= -0.5) and torch.all(transformed_normalized_cpu <= 0.5), \
|
||||
f'vertices out of range for {sha256} view {view_idx}'
|
||||
|
||||
# Ensure vertices and faces are on CPU with correct types and contiguous memory
|
||||
# CPU only, consistent with process_dual_grid in test_ovoxel_transform.py
|
||||
vertices_for_grid = transformed_normalized_cpu.float().contiguous()
|
||||
faces_for_grid = faces.long().contiguous()
|
||||
data_for_grid = {'vertices': vertices_for_grid, 'faces': faces_for_grid}
|
||||
|
||||
# Dual grid encoding
|
||||
voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
|
||||
**data_for_grid,
|
||||
grid_size=res,
|
||||
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
face_weight=1.0,
|
||||
boundary_weight=0.2,
|
||||
regularization_weight=1e-2,
|
||||
timing=False,
|
||||
)
|
||||
|
||||
# Convert to intra-voxel offsets and quantize
|
||||
dual_vertices = dual_vertices.float()
|
||||
voxel_indices_float = voxel_indices.float()
|
||||
dual_vertices = dual_vertices * res - voxel_indices_float
|
||||
assert torch.all(dual_vertices >= -1e-3) and torch.all(dual_vertices <= 1+1e-3), \
|
||||
f'dual_vertices out of range for {sha256} view {view_idx}'
|
||||
dual_vertices = torch.clamp(dual_vertices, 0, 1)
|
||||
dual_vertices = (dual_vertices * 255).type(torch.uint8)
|
||||
intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8)
|
||||
|
||||
# Save .vxz file
|
||||
os.makedirs(sha256_dir, exist_ok=True)
|
||||
o_voxel.io.write_vxz(
|
||||
vxz_path,
|
||||
voxel_indices,
|
||||
{'vertices': dual_vertices, 'intersected': intersected},
|
||||
)
|
||||
|
||||
# Save scale info
|
||||
scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
|
||||
scale_info = {
|
||||
'sha256': sha256,
|
||||
'view_idx': view_idx,
|
||||
'total_scale': total_scale,
|
||||
'sphere_radius': sphere_radius.item(),
|
||||
'box_scale': box_scale,
|
||||
}
|
||||
with open(scale_path, 'w') as f:
|
||||
json.dump(scale_info, f, indent=2)
|
||||
|
||||
pack[f'dual_grid_view{view_idx:02d}_converted_{res}'] = True
|
||||
pack[f'dual_grid_view{view_idx:02d}_size_{res}'] = len(voxel_indices)
|
||||
pack[f'dual_grid_view{view_idx:02d}_scale_{res}'] = total_scale
|
||||
processed_count += 1
|
||||
|
||||
# Record processing stats
|
||||
pack['_processed_count'] = processed_count
|
||||
pack['_skipped_count'] = skipped_count
|
||||
|
||||
return pack
|
||||
|
||||
except Exception as e:
|
||||
print(f'Error processing {sha256}: {e}')
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return {'sha256': sha256, 'error': str(e)}
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--mesh_dump_root', type=str, default=None,
|
||||
help='Directory to load mesh dumps')
|
||||
parser.add_argument('--transform_root', type=str, default=None,
|
||||
help='Directory to load transform json files (renders_cond)')
|
||||
parser.add_argument('--dual_grid_root', type=str, default=None,
|
||||
help='Directory to save dual grids')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--view_indices', type=str, default=None,
|
||||
help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--resolution', type=str, default='256')
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.resolution = [int(x) for x in opt.resolution.split(',')]
|
||||
opt.mesh_dump_root = opt.mesh_dump_root or opt.root
|
||||
opt.transform_root = opt.transform_root or os.path.join(opt.root, 'renders_cond')
|
||||
opt.dual_grid_root = opt.dual_grid_root or opt.root
|
||||
|
||||
# Parse view_indices
|
||||
view_indices = None
|
||||
if opt.view_indices is not None:
|
||||
view_indices = []
|
||||
for part in opt.view_indices.split(','):
|
||||
if '-' in part:
|
||||
start, end = map(int, part.split('-'))
|
||||
view_indices.extend(range(start, end + 1))
|
||||
else:
|
||||
view_indices.append(int(part))
|
||||
view_indices = list(set(view_indices)) # Deduplicate
|
||||
view_indices.sort()
|
||||
|
||||
for res in opt.resolution:
|
||||
os.makedirs(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'new_records'), exist_ok=True)
|
||||
|
||||
# Get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
|
||||
# Check already processed dual_grid_view
|
||||
for res in opt.resolution:
|
||||
if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'metadata.csv')):
|
||||
dual_grid_metadata = pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'metadata.csv')).set_index('sha256')
|
||||
metadata = metadata.combine_first(dual_grid_metadata)
|
||||
|
||||
metadata = metadata.reset_index()
|
||||
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['mesh_dumped'] == True]
|
||||
|
||||
# Filter out objects with all views already processed
|
||||
if view_indices is not None:
|
||||
for res in opt.resolution:
|
||||
# Check if each specified view is already processed
|
||||
all_views_done_col = f'_all_views_done_{res}'
|
||||
metadata[all_views_done_col] = True
|
||||
for view_idx in view_indices:
|
||||
col_name = f'dual_grid_view{view_idx:02d}_converted_{res}'
|
||||
if col_name in metadata.columns:
|
||||
metadata[all_views_done_col] = metadata[all_views_done_col] & (metadata[col_name] == True)
|
||||
else:
|
||||
metadata[all_views_done_col] = False
|
||||
break
|
||||
|
||||
# Keep objects with at least one incomplete resolution
|
||||
any_incomplete = None
|
||||
for res in opt.resolution:
|
||||
all_views_done_col = f'_all_views_done_{res}'
|
||||
if all_views_done_col in metadata.columns:
|
||||
if any_incomplete is None:
|
||||
any_incomplete = ~metadata[all_views_done_col]
|
||||
else:
|
||||
any_incomplete = any_incomplete | ~metadata[all_views_done_col]
|
||||
|
||||
if any_incomplete is not None:
|
||||
before_filter = len(metadata)
|
||||
metadata = metadata[any_incomplete]
|
||||
print(f'Filtered out {before_filter - len(metadata)} already completed objects')
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
if view_indices:
|
||||
print(f'View indices to process: {view_indices}')
|
||||
else:
|
||||
print('Processing all available views')
|
||||
|
||||
# Process objects
|
||||
func = partial(_dual_grid_mesh_view,
|
||||
root=opt.dual_grid_root,
|
||||
mesh_dump_root=opt.mesh_dump_root,
|
||||
transform_root=opt.transform_root,
|
||||
view_indices=view_indices)
|
||||
dual_grids = dataset_utils.foreach_instance(metadata, opt.root, func, max_workers=opt.max_workers, desc='Dual griding views', timeout=300)
|
||||
|
||||
# Processing summary
|
||||
total_processed = dual_grids['_processed_count'].sum() if '_processed_count' in dual_grids.columns else 0
|
||||
total_skipped = dual_grids['_skipped_count'].sum() if '_skipped_count' in dual_grids.columns else 0
|
||||
print(f'\n========== Processing Summary ==========')
|
||||
print(f'Total processed (new): {int(total_processed)}')
|
||||
print(f'Total skipped (existing): {int(total_skipped)}')
|
||||
print(f'Total items: {int(total_processed + total_skipped)}')
|
||||
print(f'=========================================\n')
|
||||
|
||||
if 'error' in dual_grids.columns:
|
||||
errors = dual_grids[dual_grids['error'].notna()]
|
||||
if len(errors) > 0:
|
||||
with open('errors_view.txt', 'w') as f:
|
||||
f.write('\n'.join(errors['sha256'].tolist()))
|
||||
print(f'Errors written to errors_view.txt ({len(errors)} errors)')
|
||||
|
||||
# Save metadata
|
||||
for res in opt.resolution:
|
||||
# Collect all view-related columns
|
||||
view_cols = [col for col in dual_grids.columns if f'dual_grid_view' in col and f'_{res}' in col and 'converted' in col]
|
||||
if view_cols:
|
||||
# Save metadata for each view
|
||||
dual_grid_metadata = dual_grids[dual_grids[view_cols].any(axis=1)]
|
||||
if len(dual_grid_metadata) > 0:
|
||||
# Save simplified metadata
|
||||
cols_to_save = ['sha256'] + [col for col in dual_grids.columns if f'_{res}' in col]
|
||||
cols_to_save = [col for col in cols_to_save if col in dual_grids.columns]
|
||||
dual_grid_metadata[cols_to_save].to_csv(
|
||||
os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'new_records', f'part_{opt.rank}.csv'),
|
||||
index=False
|
||||
)
|
||||
|
||||
print('Done!')
|
||||
@@ -0,0 +1,126 @@
|
||||
import os
|
||||
import shutil
|
||||
import copy
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
from subprocess import DEVNULL, call
|
||||
import numpy as np
|
||||
import tempfile
|
||||
|
||||
|
||||
BLENDER_LINK = 'https://ftp.halifax.rwth-aachen.de/blender/release/Blender4.5/blender-4.5.1-linux-x64.tar.xz'
|
||||
BLENDER_INSTALLATION_PATH = '/tmp'
|
||||
BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64/blender'
|
||||
|
||||
def _install_blender():
|
||||
if not os.path.exists(BLENDER_PATH):
|
||||
os.system('sudo apt-get update')
|
||||
os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6 libxfixes3 libgl1')
|
||||
os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
|
||||
os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
|
||||
|
||||
|
||||
def _dump_mesh(file_path, sha256, root):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
temp_path = os.path.join(tmp_dir, f'{sha256}.pickle')
|
||||
output_path = os.path.join(root, 'mesh_dumps', f'{sha256}.pickle')
|
||||
args = [
|
||||
BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'dump_mesh.py'),
|
||||
'--',
|
||||
'--object', os.path.expanduser(file_path),
|
||||
'--output_path', os.path.expanduser(temp_path)
|
||||
]
|
||||
if file_path.endswith('.blend'):
|
||||
args.insert(1, file_path)
|
||||
|
||||
call(args, stdout=DEVNULL, stderr=DEVNULL)
|
||||
|
||||
if os.path.exists(temp_path):
|
||||
shutil.move(temp_path, output_path)
|
||||
return {'sha256': sha256, 'mesh_dumped': True}
|
||||
else:
|
||||
if os.path.exists(temp_path + '_error.txt'):
|
||||
with open(temp_path + '_error.txt', 'r') as f:
|
||||
error_msg = f.read()
|
||||
raise ValueError(f'Failed to dump mesh. File {file_path}. Error message: {error_msg}')
|
||||
else:
|
||||
raise ValueError(f'Failed to dump mesh. File {file_path}.')
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--download_root', type=str, default=None,
|
||||
help='Directory to save the downloaded files')
|
||||
parser.add_argument('--mesh_dump_root', type=str, default=None,
|
||||
help='Directory to save the mesh dumps')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.download_root = opt.download_root or opt.root
|
||||
opt.mesh_dump_root = opt.mesh_dump_root or opt.root
|
||||
|
||||
os.makedirs(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'new_records'), exist_ok=True)
|
||||
|
||||
# install blender
|
||||
print('Checking blender...', flush=True)
|
||||
_install_blender()
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
metadata = metadata[metadata['local_path'].notna()]
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
if 'mesh_dumped' in metadata.columns:
|
||||
metadata = metadata[metadata['mesh_dumped'] != True]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# filter out objects that are already processed
|
||||
sha256_list = os.listdir(os.path.join(opt.mesh_dump_root, 'mesh_dumps'))
|
||||
sha256_list = [os.path.splitext(f)[0] for f in sha256_list if f.endswith('.pickle')]
|
||||
for sha256 in sha256_list:
|
||||
records.append({'sha256': sha256, 'mesh_dumped': True})
|
||||
print(f'Found {len(sha256_list)} dumped mesh')
|
||||
metadata = metadata[~metadata['sha256'].isin(sha256_list)]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
func = partial(_dump_mesh, root=opt.mesh_dump_root)
|
||||
mesh_dumped = dataset_utils.foreach_instance(metadata, opt.download_root, func, max_workers=opt.max_workers, desc='Dumping mesh')
|
||||
mesh_dumped = pd.concat([mesh_dumped, pd.DataFrame.from_records(records)])
|
||||
mesh_dumped.to_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,127 @@
|
||||
import os
|
||||
import shutil
|
||||
import copy
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
from subprocess import DEVNULL, call
|
||||
import numpy as np
|
||||
import tempfile
|
||||
|
||||
|
||||
BLENDER_LINK = 'https://ftp.halifax.rwth-aachen.de/blender/release/Blender4.5/blender-4.5.1-linux-x64.tar.xz'
|
||||
BLENDER_INSTALLATION_PATH = '/tmp'
|
||||
BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64/blender'
|
||||
|
||||
def _install_blender():
|
||||
if not os.path.exists(BLENDER_PATH):
|
||||
os.system('sudo apt-get update')
|
||||
os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6 libxfixes3 libgl1')
|
||||
os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
|
||||
os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
|
||||
os.system(f'{BLENDER_PATH} -b --python {os.path.join(os.path.dirname(__file__), "blender_script", "install_pillow.py")}')
|
||||
|
||||
|
||||
def _dump_pbr(file_path, sha256, root):
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
temp_path = os.path.join(tmp_dir, f'{sha256}.pickle')
|
||||
output_path = os.path.join(root, 'pbr_dumps', f'{sha256}.pickle')
|
||||
args = [
|
||||
BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'dump_pbr.py'),
|
||||
'--',
|
||||
'--object', os.path.expanduser(file_path),
|
||||
'--output_path', os.path.expanduser(temp_path)
|
||||
]
|
||||
if file_path.endswith('.blend'):
|
||||
args.insert(1, file_path)
|
||||
|
||||
call(args, stdout=DEVNULL, stderr=DEVNULL)
|
||||
|
||||
if os.path.exists(temp_path):
|
||||
shutil.move(temp_path, output_path)
|
||||
return {'sha256': sha256, 'pbr_dumped': True}
|
||||
else:
|
||||
if os.path.exists(temp_path + '_error.txt'):
|
||||
with open(temp_path + '_error.txt', 'r') as f:
|
||||
error_msg = f.read()
|
||||
raise ValueError(f'Failed to dump PBR. File {file_path}. Error message: {error_msg}')
|
||||
else:
|
||||
raise ValueError(f'Failed to dump PBR. File {file_path}.')
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--download_root', type=str, default=None,
|
||||
help='Directory to save the downloaded files')
|
||||
parser.add_argument('--pbr_dump_root', type=str, default=None,
|
||||
help='Directory to save the mesh dumps')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.download_root = opt.download_root or opt.root
|
||||
opt.pbr_dump_root = opt.pbr_dump_root or opt.root
|
||||
|
||||
os.makedirs(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'new_records'), exist_ok=True)
|
||||
|
||||
# install blender
|
||||
print('Checking blender...', flush=True)
|
||||
_install_blender()
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
metadata = metadata[metadata['local_path'].notna()]
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
if 'pbr_dumped' in metadata.columns:
|
||||
metadata = metadata[metadata['pbr_dumped'] != True]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# filter out objects that are already processed
|
||||
sha256_list = os.listdir(os.path.join(opt.pbr_dump_root, 'pbr_dumps'))
|
||||
sha256_list = [os.path.splitext(f)[0] for f in sha256_list if f.endswith('.pickle')]
|
||||
for sha256 in sha256_list:
|
||||
records.append({'sha256': sha256, 'pbr_dumped': True})
|
||||
print(f'Found {len(sha256_list)} dumped PBRs')
|
||||
metadata = metadata[~metadata['sha256'].isin(sha256_list)]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
func = partial(_dump_pbr, root=opt.pbr_dump_root)
|
||||
pbr_dumped = dataset_utils.foreach_instance(metadata, opt.download_root, func, max_workers=opt.max_workers, desc='Dumping PBR')
|
||||
pbr_dumped = pd.concat([pbr_dumped, pd.DataFrame.from_records(records)])
|
||||
pbr_dumped.to_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,181 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||||
import json
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import o_voxel
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import Queue
|
||||
|
||||
import pixal3d.models as models
|
||||
import pixal3d.modules.sparse as sp
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
def is_valid_sparse_tensor(tensor):
|
||||
return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
|
||||
|
||||
def clear_cuda_error():
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--pbr_voxel_root', type=str, default=None,
|
||||
help='Directory to save the pbr voxel files')
|
||||
parser.add_argument('--pbr_latent_root', type=str, default=None,
|
||||
help='Directory to save the pbr latent files')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--resolution', type=int, default=1024,
|
||||
help='Sparse voxel resolution')
|
||||
parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16',
|
||||
help='Pretrained encoder model')
|
||||
parser.add_argument('--model_root', type=str,
|
||||
help='Root directory of models')
|
||||
parser.add_argument('--enc_model', type=str,
|
||||
help='Encoder model. if specified, use this model instead of pretrained model')
|
||||
parser.add_argument('--ckpt', type=str,
|
||||
help='Checkpoint to load')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
|
||||
opt.pbr_latent_root = opt.pbr_latent_root or opt.root
|
||||
|
||||
if opt.enc_model is None:
|
||||
latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
|
||||
encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
|
||||
else:
|
||||
latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
|
||||
cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
|
||||
encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
|
||||
ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
|
||||
encoder.load_state_dict(torch.load(ckpt_path), strict=False)
|
||||
encoder.eval()
|
||||
print(f'Loaded model from {ckpt_path}')
|
||||
|
||||
os.makedirs(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{opt.resolution}', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{opt.resolution}','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name,'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['pbr_voxelized'] == True]
|
||||
if 'pbr_latent_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata['pbr_latent_encoded'] != True]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# filter out objects that are already processed
|
||||
with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
|
||||
tqdm(total=len(metadata), desc="Filtering existing objects") as pbar:
|
||||
def check_sha256(sha256):
|
||||
if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, f'{sha256}.npz')):
|
||||
coords = np.load(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, f'{sha256}.npz'))['coords']
|
||||
records.append({'sha256': sha256, 'pbr_latent_encoded': True, 'pbr_latent_tokens': coords.shape[0]})
|
||||
pbar.update()
|
||||
executor.map(check_sha256, metadata['sha256'].values)
|
||||
executor.shutdown(wait=True)
|
||||
existing_sha256 = set(r['sha256'] for r in records)
|
||||
print(f'Found {len(existing_sha256)} processed objects')
|
||||
metadata = metadata[~metadata['sha256'].isin(existing_sha256)]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
sha256s = list(metadata['sha256'].values)
|
||||
load_queue = Queue(maxsize=32)
|
||||
with ThreadPoolExecutor(max_workers=32) as loader_executor, \
|
||||
ThreadPoolExecutor(max_workers=32) as saver_executor:
|
||||
|
||||
def loader(sha256):
|
||||
try:
|
||||
attrs = ['base_color', 'metallic', 'roughness', 'alpha']
|
||||
coords, attr = o_voxel.io.read_vxz(
|
||||
os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{opt.resolution}', f'{sha256}.vxz'),
|
||||
num_threads=4
|
||||
)
|
||||
feats = torch.concat([attr[k] for k in attrs], dim=-1) / 255.0 * 2 - 1
|
||||
x = sp.SparseTensor(
|
||||
feats.float(),
|
||||
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
||||
)
|
||||
load_queue.put((sha256, x))
|
||||
except Exception as e:
|
||||
print(f"[Loader Error] {sha256}: {e}")
|
||||
load_queue.put((sha256, None))
|
||||
|
||||
loader_executor.map(loader, sha256s)
|
||||
|
||||
def saver(sha256, pack):
|
||||
save_path = os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, f'{sha256}.npz')
|
||||
np.savez_compressed(save_path, **pack)
|
||||
records.append({'sha256': sha256, 'pbr_latent_encoded': True, 'pbr_latent_tokens': pack['coords'].shape[0]})
|
||||
|
||||
for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
|
||||
try:
|
||||
sha256, voxels = load_queue.get()
|
||||
if voxels is None:
|
||||
print(f"[Skip] {sha256}: Failed to load input")
|
||||
continue
|
||||
|
||||
num_voxels = voxels.feats.shape[0]
|
||||
|
||||
# NaN/Inf
|
||||
if not (is_valid_sparse_tensor(voxels)):
|
||||
print(f"[Skip] {sha256}: NaN/Inf in input")
|
||||
continue
|
||||
|
||||
z = encoder(voxels.cuda())
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not torch.isfinite(z.feats).all():
|
||||
print(f"[Skip] {sha256}: Non-finite latent in z.feats")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
pack = {
|
||||
'feats': z.feats.cpu().numpy().astype(np.float32),
|
||||
'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
|
||||
}
|
||||
saver_executor.submit(saver, sha256, pack)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Error] {sha256} ({num_voxels} voxels): {e}")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
saver_executor.shutdown(wait=True)
|
||||
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,271 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||||
import json
|
||||
import argparse
|
||||
import shutil
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import o_voxel
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import Queue
|
||||
from utils import parse_view_indices
|
||||
|
||||
import pixal3d.models as models
|
||||
import pixal3d.modules.sparse as sp
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
|
||||
def is_valid_sparse_tensor(tensor):
|
||||
return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
|
||||
|
||||
def clear_cuda_error():
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--pbr_voxel_root', type=str, default=None,
|
||||
help='Directory containing the pbr voxels')
|
||||
parser.add_argument('--pbr_latent_root', type=str, default=None,
|
||||
help='Directory to save the pbr latent files')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--resolution', type=int, default=1024,
|
||||
help='Sparse voxel resolution')
|
||||
parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16',
|
||||
help='Pretrained encoder model')
|
||||
parser.add_argument('--model_root', type=str,
|
||||
help='Root directory of models')
|
||||
parser.add_argument('--enc_model', type=str,
|
||||
help='Encoder model. if specified, use this model instead of pretrained model')
|
||||
parser.add_argument('--ckpt', type=str,
|
||||
help='Checkpoint to load')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--view_indices', type=str, default=None,
|
||||
help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
|
||||
parser.add_argument('--num_views', type=int, default=24,
|
||||
help='Total number of views (used when view_indices is None)')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
|
||||
opt.pbr_latent_root = opt.pbr_latent_root or opt.root
|
||||
|
||||
# Parse view_indices
|
||||
view_indices = parse_view_indices(opt.view_indices)
|
||||
if view_indices is None:
|
||||
view_indices = list(range(opt.num_views))
|
||||
|
||||
print(f'View indices to process: {view_indices}')
|
||||
|
||||
if opt.enc_model is None:
|
||||
latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
|
||||
encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
|
||||
else:
|
||||
latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
|
||||
cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
|
||||
encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
|
||||
ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
|
||||
encoder.load_state_dict(torch.load(ckpt_path), strict=False)
|
||||
encoder.eval()
|
||||
print(f'Loaded model from {ckpt_path}')
|
||||
|
||||
# Multi-view latent output directory
|
||||
latent_view_name = f'{latent_name}_view_fix'
|
||||
os.makedirs(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, 'new_records'), exist_ok=True)
|
||||
|
||||
# Get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')
|
||||
metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic')
|
||||
|
||||
# Check pbr_voxels_view_fix metadata
|
||||
pbr_voxel_view_path = os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{opt.resolution}', 'metadata.csv')
|
||||
if os.path.exists(pbr_voxel_view_path):
|
||||
pbr_voxel_metadata = pd.read_csv(pbr_voxel_view_path).set_index('sha256')
|
||||
metadata = metadata.join(pbr_voxel_metadata, how='left', rsuffix='_pbr_voxel')
|
||||
|
||||
# Check pbr_latent_view metadata (used to skip already completed tasks)
|
||||
pbr_latent_view_metadata_path = os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, 'metadata.csv')
|
||||
if os.path.exists(pbr_latent_view_metadata_path):
|
||||
pbr_latent_view_metadata = pd.read_csv(pbr_latent_view_metadata_path).set_index('sha256')
|
||||
metadata = metadata.join(pbr_latent_view_metadata, how='left', rsuffix='_pbr_latent_view')
|
||||
print(f'Loaded pbr_latent_view metadata with {len(pbr_latent_view_metadata)} records')
|
||||
else:
|
||||
print(f'Warning: pbr_latent_view metadata not found at {pbr_latent_view_metadata_path}')
|
||||
|
||||
metadata = metadata.reset_index()
|
||||
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
|
||||
# Filter to objects that have pbr_voxels_view_fix data
|
||||
# Use first view as indicator
|
||||
first_view_col = f'pbr_voxelized_view_fix{view_indices[0]:02d}_{opt.resolution}'
|
||||
if first_view_col in metadata.columns:
|
||||
metadata = metadata[metadata[first_view_col] == True]
|
||||
else:
|
||||
print(f'Warning: Column {first_view_col} not found in metadata, will check files directly')
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
records = []
|
||||
|
||||
# Build task list: (sha256, view_idx), filter already completed tasks via metadata
|
||||
all_tasks = []
|
||||
skipped_count = 0
|
||||
|
||||
# Pre-fetch completion status columns for each view
|
||||
encoded_cols = {view_idx: f'pbr_latent_view{view_idx:02d}_encoded' for view_idx in view_indices}
|
||||
|
||||
for _, row in metadata.iterrows():
|
||||
sha256 = row['sha256']
|
||||
for view_idx in view_indices:
|
||||
encoded_col = encoded_cols[view_idx]
|
||||
# Check if already marked as completed in metadata
|
||||
if encoded_col in metadata.columns and row.get(encoded_col, False) == True:
|
||||
skipped_count += 1
|
||||
continue
|
||||
all_tasks.append((sha256, view_idx))
|
||||
|
||||
# Split tasks by rank after filtering completed ones
|
||||
start = len(all_tasks) * opt.rank // opt.world_size
|
||||
end = len(all_tasks) * (opt.rank + 1) // opt.world_size
|
||||
tasks = all_tasks[start:end]
|
||||
|
||||
print(f'Total tasks: {len(all_tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process (all ranks): {len(all_tasks)}, This rank: {len(tasks)}')
|
||||
|
||||
load_queue = Queue(maxsize=32)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=32) as loader_executor, \
|
||||
ThreadPoolExecutor(max_workers=32) as saver_executor:
|
||||
|
||||
def loader(task):
|
||||
sha256, view_idx = task
|
||||
try:
|
||||
# Check if output file already exists, skip if so (but still record)
|
||||
output_path = os.path.join(
|
||||
opt.pbr_latent_root,
|
||||
'pbr_latents',
|
||||
latent_view_name,
|
||||
sha256,
|
||||
f'view{view_idx:02d}.npz'
|
||||
)
|
||||
if os.path.exists(output_path):
|
||||
try:
|
||||
data = np.load(output_path)
|
||||
num_tokens = data['coords'].shape[0]
|
||||
except Exception:
|
||||
num_tokens = -1
|
||||
records.append({
|
||||
'sha256': sha256,
|
||||
f'pbr_latent_view{view_idx:02d}_encoded': True,
|
||||
f'pbr_latent_view{view_idx:02d}_tokens': num_tokens,
|
||||
})
|
||||
load_queue.put((sha256, view_idx, None))
|
||||
return
|
||||
|
||||
# pbr_voxels_view_fix path: pbr_voxels_view_fix_{res}/{sha256}/view{idx:02d}.vxz
|
||||
vxz_path = os.path.join(
|
||||
opt.pbr_voxel_root,
|
||||
f'pbr_voxels_view_fix_{opt.resolution}',
|
||||
sha256,
|
||||
f'view{view_idx:02d}.vxz'
|
||||
)
|
||||
|
||||
if not os.path.exists(vxz_path):
|
||||
print(f"[Loader Skip] {sha256}/view{view_idx:02d}: vxz file not found")
|
||||
load_queue.put((sha256, view_idx, None))
|
||||
return
|
||||
|
||||
attrs = ['base_color', 'metallic', 'roughness', 'alpha']
|
||||
coords, attr = o_voxel.io.read_vxz(vxz_path, num_threads=4)
|
||||
feats = torch.concat([attr[k] for k in attrs], dim=-1) / 255.0 * 2 - 1
|
||||
x = sp.SparseTensor(
|
||||
feats.float(),
|
||||
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
||||
)
|
||||
load_queue.put((sha256, view_idx, x))
|
||||
except Exception as e:
|
||||
print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}")
|
||||
load_queue.put((sha256, view_idx, None))
|
||||
|
||||
loader_executor.map(loader, tasks)
|
||||
|
||||
def saver(sha256, view_idx, pack):
|
||||
sha256_dir = os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, sha256)
|
||||
os.makedirs(sha256_dir, exist_ok=True)
|
||||
save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz')
|
||||
np.savez_compressed(save_path, **pack)
|
||||
|
||||
# Copy scale json from pbr_voxels_view_fix
|
||||
src_scale_path = os.path.join(
|
||||
opt.pbr_voxel_root,
|
||||
f'pbr_voxels_view_fix_{opt.resolution}',
|
||||
sha256,
|
||||
f'view{view_idx:02d}_scale.json'
|
||||
)
|
||||
dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
|
||||
if os.path.exists(src_scale_path):
|
||||
shutil.copy2(src_scale_path, dst_scale_path)
|
||||
|
||||
records.append({
|
||||
'sha256': sha256,
|
||||
f'pbr_latent_view{view_idx:02d}_encoded': True,
|
||||
f'pbr_latent_view{view_idx:02d}_tokens': pack['coords'].shape[0]
|
||||
})
|
||||
|
||||
for _ in tqdm(range(len(tasks)), desc=f"Extracting {os.path.basename(opt.root)} PBR view latents (res={opt.resolution})"):
|
||||
try:
|
||||
sha256, view_idx, voxels = load_queue.get()
|
||||
if voxels is None:
|
||||
continue
|
||||
|
||||
num_voxels = voxels.feats.shape[0]
|
||||
|
||||
# NaN/Inf check
|
||||
if not is_valid_sparse_tensor(voxels):
|
||||
print(f"[Skip] {sha256}/view{view_idx:02d}: NaN/Inf in input")
|
||||
continue
|
||||
|
||||
z = encoder(voxels.cuda())
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not torch.isfinite(z.feats).all():
|
||||
print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent in z.feats")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
pack = {
|
||||
'feats': z.feats.cpu().numpy().astype(np.float32),
|
||||
'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
|
||||
}
|
||||
saver_executor.submit(saver, sha256, view_idx, pack)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Error] {sha256}/view{view_idx:02d} ({num_voxels} voxels): {e}")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
saver_executor.shutdown(wait=True)
|
||||
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,184 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||||
import json
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import o_voxel
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import Queue
|
||||
|
||||
import pixal3d.models as models
|
||||
import pixal3d.modules.sparse as sp
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
def is_valid_sparse_tensor(tensor):
|
||||
return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
|
||||
|
||||
def clear_cuda_error():
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--dual_grid_root', type=str, default=None,
|
||||
help='Directory to save the dual grids')
|
||||
parser.add_argument('--shape_latent_root', type=str, default=None,
|
||||
help='Directory to save the shape latent files')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--resolution', type=int, default=1024,
|
||||
help='Sparse voxel resolution')
|
||||
parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16',
|
||||
help='Pretrained encoder model')
|
||||
parser.add_argument('--model_root', type=str,
|
||||
help='Root directory of models')
|
||||
parser.add_argument('--enc_model', type=str,
|
||||
help='Encoder model. if specified, use this model instead of pretrained model')
|
||||
parser.add_argument('--ckpt', type=str,
|
||||
help='Checkpoint to load')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.dual_grid_root = opt.dual_grid_root or opt.root
|
||||
opt.shape_latent_root = opt.shape_latent_root or opt.root
|
||||
|
||||
if opt.enc_model is None:
|
||||
latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
|
||||
encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
|
||||
else:
|
||||
latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
|
||||
cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
|
||||
encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
|
||||
ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
|
||||
encoder.load_state_dict(torch.load(ckpt_path), strict=False)
|
||||
encoder.eval()
|
||||
print(f'Loaded model from {ckpt_path}')
|
||||
|
||||
os.makedirs(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_{opt.resolution}', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{opt.resolution}','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name,'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['dual_grid_converted'] == True]
|
||||
if 'shape_latent_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata['shape_latent_encoded'] != True]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# filter out objects that are already processed
|
||||
with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
|
||||
tqdm(total=len(metadata), desc="Filtering existing objects") as pbar:
|
||||
def check_sha256(sha256):
|
||||
if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, f'{sha256}.npz')):
|
||||
coords = np.load(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, f'{sha256}.npz'))['coords']
|
||||
records.append({'sha256': sha256, 'shape_latent_encoded': True, 'shape_latent_tokens': coords.shape[0]})
|
||||
pbar.update()
|
||||
executor.map(check_sha256, metadata['sha256'].values)
|
||||
executor.shutdown(wait=True)
|
||||
existing_sha256 = set(r['sha256'] for r in records)
|
||||
print(f'Found {len(existing_sha256)} processed objects')
|
||||
metadata = metadata[~metadata['sha256'].isin(existing_sha256)]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
sha256s = list(metadata['sha256'].values)
|
||||
load_queue = Queue(maxsize=32)
|
||||
with ThreadPoolExecutor(max_workers=32) as loader_executor, \
|
||||
ThreadPoolExecutor(max_workers=32) as saver_executor:
|
||||
|
||||
def loader(sha256):
|
||||
try:
|
||||
coords, attr = o_voxel.io.read_vxz(
|
||||
os.path.join(opt.dual_grid_root, f'dual_grid_{opt.resolution}', f'{sha256}.vxz'),
|
||||
num_threads=4
|
||||
)
|
||||
vertices = sp.SparseTensor(
|
||||
(attr['vertices'] / 255.0).float(),
|
||||
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
||||
)
|
||||
intersected = vertices.replace(torch.cat([
|
||||
attr['intersected'] % 2,
|
||||
attr['intersected'] // 2 % 2,
|
||||
attr['intersected'] // 4 % 2,
|
||||
], dim=-1).bool())
|
||||
load_queue.put((sha256, vertices, intersected))
|
||||
except Exception as e:
|
||||
print(f"[Loader Error] {sha256}: {e}")
|
||||
load_queue.put((sha256, None, None))
|
||||
|
||||
loader_executor.map(loader, sha256s)
|
||||
|
||||
def saver(sha256, pack):
|
||||
save_path = os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, f'{sha256}.npz')
|
||||
np.savez_compressed(save_path, **pack)
|
||||
records.append({'sha256': sha256, 'shape_latent_encoded': True, 'shape_latent_tokens': pack['coords'].shape[0]})
|
||||
|
||||
for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
|
||||
try:
|
||||
sha256, vertices, intersected = load_queue.get()
|
||||
if vertices is None or intersected is None:
|
||||
print(f"[Skip] {sha256}: Failed to load input")
|
||||
continue
|
||||
|
||||
num_voxels = vertices.feats.shape[0]
|
||||
|
||||
# NaN/Inf
|
||||
if not (is_valid_sparse_tensor(vertices) and is_valid_sparse_tensor(intersected)):
|
||||
print(f"[Skip] {sha256}: NaN/Inf in input")
|
||||
continue
|
||||
|
||||
z = encoder(vertices.cuda(), intersected.cuda())
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not torch.isfinite(z.feats).all():
|
||||
print(f"[Skip] {sha256}: Non-finite latent in z.feats")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
pack = {
|
||||
'feats': z.feats.cpu().numpy().astype(np.float32),
|
||||
'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
|
||||
}
|
||||
saver_executor.submit(saver, sha256, pack)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Error] {sha256} ({num_voxels} voxels): {e}")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
saver_executor.shutdown(wait=True)
|
||||
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,262 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||||
import json
|
||||
import argparse
|
||||
import shutil
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import o_voxel
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import Queue
|
||||
from utils import parse_view_indices
|
||||
|
||||
import pixal3d.models as models
|
||||
import pixal3d.modules.sparse as sp
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
|
||||
def is_valid_sparse_tensor(tensor):
|
||||
return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
|
||||
|
||||
def clear_cuda_error():
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--dual_grid_root', type=str, default=None,
|
||||
help='Directory containing the dual grids')
|
||||
parser.add_argument('--shape_latent_root', type=str, default=None,
|
||||
help='Directory to save the shape latent files')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--resolution', type=int, default=1024,
|
||||
help='Sparse voxel resolution')
|
||||
parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16',
|
||||
help='Pretrained encoder model')
|
||||
parser.add_argument('--model_root', type=str,
|
||||
help='Root directory of models')
|
||||
parser.add_argument('--enc_model', type=str,
|
||||
help='Encoder model. if specified, use this model instead of pretrained model')
|
||||
parser.add_argument('--ckpt', type=str,
|
||||
help='Checkpoint to load')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--view_indices', type=str, default=None,
|
||||
help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
|
||||
parser.add_argument('--num_views', type=int, default=24,
|
||||
help='Total number of views (used when view_indices is None)')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.dual_grid_root = opt.dual_grid_root or opt.root
|
||||
opt.shape_latent_root = opt.shape_latent_root or opt.root
|
||||
|
||||
# Parse view_indices
|
||||
view_indices = parse_view_indices(opt.view_indices)
|
||||
if view_indices is None:
|
||||
view_indices = list(range(opt.num_views))
|
||||
|
||||
print(f'View indices to process: {view_indices}')
|
||||
|
||||
if opt.enc_model is None:
|
||||
latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
|
||||
encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
|
||||
else:
|
||||
latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
|
||||
cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
|
||||
encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
|
||||
ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
|
||||
encoder.load_state_dict(torch.load(ckpt_path), strict=False)
|
||||
encoder.eval()
|
||||
print(f'Loaded model from {ckpt_path}')
|
||||
|
||||
# Multi-view latent output directory
|
||||
latent_view_name = f'{latent_name}_view'
|
||||
os.makedirs(os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, 'new_records'), exist_ok=True)
|
||||
|
||||
# Get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')
|
||||
metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic')
|
||||
|
||||
# Check dual_grid_view metadata
|
||||
dual_grid_view_path = os.path.join(opt.dual_grid_root, f'dual_grid_view_{opt.resolution}', 'metadata.csv')
|
||||
if os.path.exists(dual_grid_view_path):
|
||||
dual_grid_metadata = pd.read_csv(dual_grid_view_path).set_index('sha256')
|
||||
metadata = metadata.join(dual_grid_metadata, how='left', rsuffix='_dual_grid')
|
||||
|
||||
# Check shape_latent_view metadata (used to skip already completed tasks)
|
||||
shape_latent_view_metadata_path = os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, 'metadata.csv')
|
||||
if os.path.exists(shape_latent_view_metadata_path):
|
||||
shape_latent_view_metadata = pd.read_csv(shape_latent_view_metadata_path).set_index('sha256')
|
||||
metadata = metadata.join(shape_latent_view_metadata, how='left', rsuffix='_shape_latent_view')
|
||||
print(f'Loaded shape_latent_view metadata with {len(shape_latent_view_metadata)} records')
|
||||
else:
|
||||
print(f'Warning: shape_latent_view metadata not found at {shape_latent_view_metadata_path}')
|
||||
|
||||
metadata = metadata.reset_index()
|
||||
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
|
||||
# Filter to objects that have dual_grid_view data
|
||||
# Use first view as indicator
|
||||
first_view_col = f'dual_grid_view{view_indices[0]:02d}_converted_{opt.resolution}'
|
||||
if first_view_col in metadata.columns:
|
||||
metadata = metadata[metadata[first_view_col] == True]
|
||||
else:
|
||||
print(f'Warning: Column {first_view_col} not found in metadata, will check files directly')
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# Build task list: (sha256, view_idx), filter already completed tasks via metadata
|
||||
tasks = []
|
||||
skipped_count = 0
|
||||
|
||||
# Pre-fetch completion status columns for each view
|
||||
encoded_cols = {view_idx: f'shape_latent_view{view_idx:02d}_encoded' for view_idx in view_indices}
|
||||
|
||||
for _, row in metadata.iterrows():
|
||||
sha256 = row['sha256']
|
||||
for view_idx in view_indices:
|
||||
encoded_col = encoded_cols[view_idx]
|
||||
# Check if already marked as completed in metadata
|
||||
if encoded_col in metadata.columns and row.get(encoded_col, False) == True:
|
||||
skipped_count += 1
|
||||
continue
|
||||
tasks.append((sha256, view_idx))
|
||||
|
||||
print(f'Total tasks: {len(tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process: {len(tasks)}')
|
||||
|
||||
load_queue = Queue(maxsize=32)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=32) as loader_executor, \
|
||||
ThreadPoolExecutor(max_workers=32) as saver_executor:
|
||||
|
||||
def loader(task):
|
||||
sha256, view_idx = task
|
||||
try:
|
||||
# Check if output file already exists, skip if so
|
||||
output_path = os.path.join(
|
||||
opt.shape_latent_root,
|
||||
'shape_latents',
|
||||
latent_view_name,
|
||||
sha256,
|
||||
f'view{view_idx:02d}.npz'
|
||||
)
|
||||
if os.path.exists(output_path):
|
||||
load_queue.put((sha256, view_idx, None, None))
|
||||
return
|
||||
|
||||
# dual_grid_view path: dual_grid_view_{res}/{sha256}/view{idx:02d}.vxz
|
||||
vxz_path = os.path.join(
|
||||
opt.dual_grid_root,
|
||||
f'dual_grid_view_{opt.resolution}',
|
||||
sha256,
|
||||
f'view{view_idx:02d}.vxz'
|
||||
)
|
||||
|
||||
if not os.path.exists(vxz_path):
|
||||
print(f"[Loader Skip] {sha256}/view{view_idx:02d}: vxz file not found")
|
||||
load_queue.put((sha256, view_idx, None, None))
|
||||
return
|
||||
|
||||
coords, attr = o_voxel.io.read_vxz(vxz_path, num_threads=4)
|
||||
vertices = sp.SparseTensor(
|
||||
(attr['vertices'] / 255.0).float(),
|
||||
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
||||
)
|
||||
intersected = vertices.replace(torch.cat([
|
||||
attr['intersected'] % 2,
|
||||
attr['intersected'] // 2 % 2,
|
||||
attr['intersected'] // 4 % 2,
|
||||
], dim=-1).bool())
|
||||
load_queue.put((sha256, view_idx, vertices, intersected))
|
||||
except Exception as e:
|
||||
print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}")
|
||||
load_queue.put((sha256, view_idx, None, None))
|
||||
|
||||
loader_executor.map(loader, tasks)
|
||||
|
||||
def saver(sha256, view_idx, pack):
|
||||
sha256_dir = os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, sha256)
|
||||
os.makedirs(sha256_dir, exist_ok=True)
|
||||
save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz')
|
||||
np.savez_compressed(save_path, **pack)
|
||||
|
||||
# Copy scale json from dual_grid_view
|
||||
src_scale_path = os.path.join(
|
||||
opt.dual_grid_root,
|
||||
f'dual_grid_view_{opt.resolution}',
|
||||
sha256,
|
||||
f'view{view_idx:02d}_scale.json'
|
||||
)
|
||||
dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
|
||||
if os.path.exists(src_scale_path):
|
||||
shutil.copy2(src_scale_path, dst_scale_path)
|
||||
|
||||
records.append({
|
||||
'sha256': sha256,
|
||||
f'shape_latent_view{view_idx:02d}_encoded': True,
|
||||
f'shape_latent_view{view_idx:02d}_tokens': pack['coords'].shape[0]
|
||||
})
|
||||
|
||||
for _ in tqdm(range(len(tasks)), desc="Extracting view latents"):
|
||||
try:
|
||||
sha256, view_idx, vertices, intersected = load_queue.get()
|
||||
if vertices is None or intersected is None:
|
||||
continue
|
||||
|
||||
num_voxels = vertices.feats.shape[0]
|
||||
|
||||
# NaN/Inf check
|
||||
if not (is_valid_sparse_tensor(vertices) and is_valid_sparse_tensor(intersected)):
|
||||
print(f"[Skip] {sha256}/view{view_idx:02d}: NaN/Inf in input")
|
||||
continue
|
||||
|
||||
z = encoder(vertices.cuda(), intersected.cuda())
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not torch.isfinite(z.feats).all():
|
||||
print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent in z.feats")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
pack = {
|
||||
'feats': z.feats.cpu().numpy().astype(np.float32),
|
||||
'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
|
||||
}
|
||||
saver_executor.submit(saver, sha256, view_idx, pack)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Error] {sha256}/view{view_idx:02d} ({num_voxels} voxels): {e}")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
saver_executor.shutdown(wait=True)
|
||||
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,163 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||||
import json
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import Queue
|
||||
|
||||
import pixal3d.models as models
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
def is_valid_sparse_tensor(tensor):
|
||||
return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
|
||||
|
||||
def clear_cuda_error():
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--shape_latent_root', type=str, default=None,
|
||||
help='Directory to save the shape latent files')
|
||||
parser.add_argument('--ss_latent_root', type=str, default=None,
|
||||
help='Directory to save the shape latent files')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--resolution', type=int, default=64,
|
||||
help='Sparse voxel resolution')
|
||||
parser.add_argument('--shape_latent_name', type=str, default=None,
|
||||
help='Name of the shape latent files')
|
||||
parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16',
|
||||
help='Pretrained encoder model')
|
||||
parser.add_argument('--model_root', type=str,
|
||||
help='Root directory of models')
|
||||
parser.add_argument('--enc_model', type=str,
|
||||
help='Encoder model. if specified, use this model instead of pretrained model')
|
||||
parser.add_argument('--ckpt', type=str,
|
||||
help='Checkpoint to load')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.shape_latent_root = opt.shape_latent_root or opt.root
|
||||
opt.ss_latent_root = opt.ss_latent_root or opt.root
|
||||
|
||||
if opt.enc_model is None:
|
||||
latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
|
||||
encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
|
||||
else:
|
||||
latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
|
||||
cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
|
||||
encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
|
||||
ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
|
||||
encoder.load_state_dict(torch.load(ckpt_path), strict=False)
|
||||
encoder.eval()
|
||||
print(f'Loaded model from {ckpt_path}')
|
||||
|
||||
os.makedirs(os.path.join(opt.ss_latent_root, 'ss_latents', latent_name, 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents', opt.shape_latent_name, 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.shape_latent_root, 'shape_latents', opt.shape_latent_name,'metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.ss_latent_root,'ss_latents', latent_name, 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.ss_latent_root,'ss_latents', latent_name,'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['shape_latent_encoded'] == True]
|
||||
if 'ss_latent_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata['ss_latent_encoded'] != True]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# filter out objects that are already processed
|
||||
sha256_list = os.listdir(os.path.join(opt.ss_latent_root, 'ss_latents'))
|
||||
sha256_list = [os.path.splitext(f)[0] for f in sha256_list if f.endswith('.npz')]
|
||||
for sha256 in sha256_list:
|
||||
records.append({'sha256': sha256, 'ss_latent_encoded': True})
|
||||
print(f'Found {len(sha256_list)} processed objects')
|
||||
metadata = metadata[~metadata['sha256'].isin(sha256_list)]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
sha256s = list(metadata['sha256'].values)
|
||||
load_queue = Queue(maxsize=32)
|
||||
with ThreadPoolExecutor(max_workers=32) as loader_executor, \
|
||||
ThreadPoolExecutor(max_workers=32) as saver_executor:
|
||||
|
||||
def loader(sha256):
|
||||
try:
|
||||
coords = np.load(os.path.join(opt.shape_latent_root, 'shape_latents', opt.shape_latent_name, f'{sha256}.npz'))['coords']
|
||||
assert np.all(coords < opt.resolution), f"{sha256}: Invalid coords"
|
||||
coords = torch.from_numpy(coords).long()
|
||||
ss = torch.zeros(1, opt.resolution, opt.resolution, opt.resolution, dtype=torch.long)
|
||||
ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
|
||||
load_queue.put((sha256, ss))
|
||||
except Exception as e:
|
||||
print(f"[Loader Error] {sha256}: {e}")
|
||||
load_queue.put((sha256, None))
|
||||
|
||||
loader_executor.map(loader, sha256s)
|
||||
|
||||
def saver(sha256, pack):
|
||||
save_path = os.path.join(opt.ss_latent_root, 'ss_latents', latent_name, f'{sha256}.npz')
|
||||
np.savez_compressed(save_path, **pack)
|
||||
records.append({'sha256': sha256, 'ss_latent_encoded': True})
|
||||
|
||||
for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
|
||||
try:
|
||||
sha256, ss = load_queue.get()
|
||||
if ss is None:
|
||||
print(f"[Skip] {sha256}: Failed to load input")
|
||||
continue
|
||||
|
||||
ss = ss.cuda()[None].float()
|
||||
z = encoder(ss, sample_posterior=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not torch.isfinite(z).all():
|
||||
print(f"[Skip] {sha256}: Non-finite latent")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
pack = {
|
||||
'z': z[0].cpu().numpy(),
|
||||
}
|
||||
saver_executor.submit(saver, sha256, pack)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Error] {sha256}: {e}")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
saver_executor.shutdown(wait=True)
|
||||
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.ss_latent_root, 'ss_latents', latent_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,257 @@
|
||||
import os
|
||||
import sys
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
||||
import json
|
||||
import shutil
|
||||
import argparse
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
from easydict import EasyDict as edict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from queue import Queue
|
||||
from utils import parse_view_indices
|
||||
|
||||
import pixal3d.models as models
|
||||
|
||||
torch.set_grad_enabled(False)
|
||||
|
||||
def clear_cuda_error():
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--shape_latent_root', type=str, default=None,
|
||||
help='Directory containing the shape latent files')
|
||||
parser.add_argument('--ss_latent_root', type=str, default=None,
|
||||
help='Directory to save the ss latent files')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--resolution', type=int, default=32,
|
||||
help='SS latent resolution')
|
||||
parser.add_argument('--shape_latent_name', type=str, required=True,
|
||||
help='Name of the shape latent files (e.g., shape_enc_next_dc_f16c32_fp16_512)')
|
||||
parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16',
|
||||
help='Pretrained encoder model')
|
||||
parser.add_argument('--model_root', type=str,
|
||||
help='Root directory of models')
|
||||
parser.add_argument('--enc_model', type=str,
|
||||
help='Encoder model. if specified, use this model instead of pretrained model')
|
||||
parser.add_argument('--ckpt', type=str,
|
||||
help='Checkpoint to load')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--view_indices', type=str, default=None,
|
||||
help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
|
||||
parser.add_argument('--num_views', type=int, default=24,
|
||||
help='Total number of views (used when view_indices is None)')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
opt = parser.parse_args()
|
||||
opt = edict(vars(opt))
|
||||
opt.shape_latent_root = opt.shape_latent_root or opt.root
|
||||
opt.ss_latent_root = opt.ss_latent_root or opt.root
|
||||
|
||||
# Parse view_indices
|
||||
view_indices = parse_view_indices(opt.view_indices)
|
||||
if view_indices is None:
|
||||
view_indices = list(range(opt.num_views))
|
||||
|
||||
print(f'View indices to process: {view_indices}')
|
||||
|
||||
if opt.enc_model is None:
|
||||
latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
|
||||
encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
|
||||
else:
|
||||
latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
|
||||
cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
|
||||
encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
|
||||
ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
|
||||
encoder.load_state_dict(torch.load(ckpt_path), strict=False)
|
||||
encoder.eval()
|
||||
print(f'Loaded model from {ckpt_path}')
|
||||
|
||||
# Multi-view shape_latent and ss_latent directory names
|
||||
shape_latent_view_name = f'{opt.shape_latent_name}_view'
|
||||
ss_latent_view_name = f'{latent_name}_view'
|
||||
|
||||
os.makedirs(os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'new_records'), exist_ok=True)
|
||||
|
||||
# Get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')
|
||||
metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic')
|
||||
|
||||
# Check shape_latent_view metadata
|
||||
shape_latent_view_metadata_path = os.path.join(opt.shape_latent_root, 'shape_latents', shape_latent_view_name, 'metadata.csv')
|
||||
if os.path.exists(shape_latent_view_metadata_path):
|
||||
shape_latent_view_metadata = pd.read_csv(shape_latent_view_metadata_path).set_index('sha256')
|
||||
metadata = metadata.join(shape_latent_view_metadata, how='left', rsuffix='_shape_latent_view')
|
||||
print(f'Loaded shape_latent_view metadata with {len(shape_latent_view_metadata)} records')
|
||||
else:
|
||||
print(f'Warning: shape_latent_view metadata not found at {shape_latent_view_metadata_path}')
|
||||
|
||||
# Check ss_latent_view metadata (used to skip already completed tasks)
|
||||
ss_latent_view_metadata_path = os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'metadata.csv')
|
||||
if os.path.exists(ss_latent_view_metadata_path):
|
||||
ss_latent_view_metadata = pd.read_csv(ss_latent_view_metadata_path).set_index('sha256')
|
||||
metadata = metadata.join(ss_latent_view_metadata, how='left', rsuffix='_ss_latent_view')
|
||||
print(f'Loaded ss_latent_view metadata with {len(ss_latent_view_metadata)} records')
|
||||
else:
|
||||
print(f'Warning: ss_latent_view metadata not found at {ss_latent_view_metadata_path}')
|
||||
|
||||
metadata = metadata.reset_index()
|
||||
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
|
||||
# Filter to objects that have shape_latent_view data
|
||||
# Use first view as indicator
|
||||
first_view_col = f'shape_latent_view{view_indices[0]:02d}_encoded'
|
||||
if first_view_col in metadata.columns:
|
||||
metadata = metadata[metadata[first_view_col] == True]
|
||||
else:
|
||||
print(f'Warning: Column {first_view_col} not found in metadata, will check files directly')
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# Build task list: (sha256, view_idx), filter already completed tasks via metadata
|
||||
tasks = []
|
||||
skipped_count = 0
|
||||
|
||||
# Pre-fetch completion status columns for each view
|
||||
encoded_cols = {view_idx: f'ss_latent_view{view_idx:02d}_encoded' for view_idx in view_indices}
|
||||
|
||||
for _, row in metadata.iterrows():
|
||||
sha256 = row['sha256']
|
||||
for view_idx in view_indices:
|
||||
encoded_col = encoded_cols[view_idx]
|
||||
# Check if already marked as completed in metadata
|
||||
if encoded_col in metadata.columns and row.get(encoded_col, False) == True:
|
||||
skipped_count += 1
|
||||
continue
|
||||
tasks.append((sha256, view_idx))
|
||||
|
||||
print(f'Total tasks: {len(tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process: {len(tasks)}')
|
||||
|
||||
load_queue = Queue(maxsize=32)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=32) as loader_executor, \
|
||||
ThreadPoolExecutor(max_workers=32) as saver_executor:
|
||||
|
||||
def loader(task):
|
||||
sha256, view_idx = task
|
||||
try:
|
||||
# Check if output file already exists, skip if so
|
||||
output_path = os.path.join(
|
||||
opt.ss_latent_root,
|
||||
'ss_latents',
|
||||
ss_latent_view_name,
|
||||
sha256,
|
||||
f'view{view_idx:02d}.npz'
|
||||
)
|
||||
if os.path.exists(output_path):
|
||||
load_queue.put((sha256, view_idx, None))
|
||||
return
|
||||
|
||||
# shape_latent_view path: shape_latents/{shape_latent_view_name}/{sha256}/view{idx:02d}.npz
|
||||
npz_path = os.path.join(
|
||||
opt.shape_latent_root,
|
||||
'shape_latents',
|
||||
shape_latent_view_name,
|
||||
sha256,
|
||||
f'view{view_idx:02d}.npz'
|
||||
)
|
||||
|
||||
if not os.path.exists(npz_path):
|
||||
print(f"[Loader Skip] {sha256}/view{view_idx:02d}: npz file not found at {npz_path}")
|
||||
load_queue.put((sha256, view_idx, None))
|
||||
return
|
||||
|
||||
data = np.load(npz_path)
|
||||
coords = data['coords']
|
||||
|
||||
# Validate coords are within resolution range
|
||||
assert np.all(coords < opt.resolution), f"{sha256}/view{view_idx:02d}: Invalid coords (max={coords.max()}, resolution={opt.resolution})"
|
||||
|
||||
coords = torch.from_numpy(coords).long()
|
||||
ss = torch.zeros(1, opt.resolution, opt.resolution, opt.resolution, dtype=torch.long)
|
||||
ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
|
||||
|
||||
load_queue.put((sha256, view_idx, ss))
|
||||
except Exception as e:
|
||||
print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}")
|
||||
load_queue.put((sha256, view_idx, None))
|
||||
|
||||
loader_executor.map(loader, tasks)
|
||||
|
||||
def saver(sha256, view_idx, pack):
|
||||
sha256_dir = os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, sha256)
|
||||
os.makedirs(sha256_dir, exist_ok=True)
|
||||
save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz')
|
||||
np.savez_compressed(save_path, **pack)
|
||||
|
||||
# Copy scale.json from shape_latent_view directory
|
||||
src_scale_path = os.path.join(
|
||||
opt.shape_latent_root,
|
||||
'shape_latents',
|
||||
shape_latent_view_name,
|
||||
sha256,
|
||||
f'view{view_idx:02d}_scale.json'
|
||||
)
|
||||
dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
|
||||
if os.path.exists(src_scale_path) and not os.path.exists(dst_scale_path):
|
||||
shutil.copy2(src_scale_path, dst_scale_path)
|
||||
|
||||
records.append({
|
||||
'sha256': sha256,
|
||||
f'ss_latent_view{view_idx:02d}_encoded': True,
|
||||
})
|
||||
|
||||
for _ in tqdm(range(len(tasks)), desc="Extracting SS view latents"):
|
||||
try:
|
||||
sha256, view_idx, ss = load_queue.get()
|
||||
if ss is None:
|
||||
continue
|
||||
|
||||
ss = ss.cuda()[None].float()
|
||||
z = encoder(ss, sample_posterior=False)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
if not torch.isfinite(z).all():
|
||||
print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
pack = {
|
||||
'z': z[0].cpu().numpy(),
|
||||
}
|
||||
saver_executor.submit(saver, sha256, view_idx, pack)
|
||||
|
||||
except Exception as e:
|
||||
print(f"[Error] {sha256}/view{view_idx:02d}: {e}")
|
||||
clear_cuda_error()
|
||||
continue
|
||||
|
||||
saver_executor.shutdown(wait=True)
|
||||
|
||||
records = pd.DataFrame.from_records(records)
|
||||
records.to_csv(os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,154 @@
|
||||
import os
|
||||
import json
|
||||
import copy
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
from subprocess import DEVNULL, call
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from tqdm import tqdm
|
||||
import numpy as np
|
||||
from utils import sphere_hammersley_sequence
|
||||
|
||||
|
||||
BLENDER_LINK = 'https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz'
|
||||
BLENDER_INSTALLATION_PATH = '/tmp'
|
||||
BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender'
|
||||
|
||||
def _install_blender():
|
||||
if not os.path.exists(BLENDER_PATH):
|
||||
os.system('sudo apt-get update')
|
||||
os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6 libxfixes3 libgl1')
|
||||
os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
|
||||
os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
|
||||
# Install Pillow into Blender's bundled Python if not already present
|
||||
blender_python = os.path.join(BLENDER_INSTALLATION_PATH, 'blender-3.0.1-linux-x64', '3.0', 'python', 'bin', 'python3.9')
|
||||
try:
|
||||
import subprocess
|
||||
subprocess.check_call([blender_python, '-c', 'from PIL import Image'], stdout=DEVNULL, stderr=DEVNULL)
|
||||
except subprocess.CalledProcessError:
|
||||
print('Installing Pillow into Blender Python...')
|
||||
subprocess.check_call([blender_python, '-m', 'ensurepip'], stdout=DEVNULL, stderr=DEVNULL)
|
||||
subprocess.check_call([blender_python, '-m', 'pip', 'install', 'Pillow'], stdout=DEVNULL, stderr=DEVNULL)
|
||||
|
||||
|
||||
def _render_cond(file_path, sha256, root, num_cond_views):
|
||||
# Build conditional view camera
|
||||
yaws = []
|
||||
pitchs = []
|
||||
offset = (np.random.rand(), np.random.rand())
|
||||
for i in range(num_cond_views):
|
||||
y, p = sphere_hammersley_sequence(i, num_cond_views, offset)
|
||||
yaws.append(y)
|
||||
pitchs.append(p)
|
||||
fov_min, fov_max = 10, 70
|
||||
radius_min = np.sqrt(3) / 2 / np.sin(fov_max / 360 * np.pi)
|
||||
radius_max = np.sqrt(3) / 2 / np.sin(fov_min / 360 * np.pi)
|
||||
k_min = 1 / radius_max**2
|
||||
k_max = 1 / radius_min**2
|
||||
ks = np.random.uniform(k_min, k_max, (1000000,))
|
||||
radius = [1 / np.sqrt(k) for k in ks]
|
||||
fov = [2 * np.arcsin(np.sqrt(3) / 2 / r) for r in radius]
|
||||
cond_views = [{'yaw': y, 'pitch': p, 'radius': r, 'fov': f} for y, p, r, f in zip(yaws, pitchs, radius, fov)]
|
||||
|
||||
args = [
|
||||
BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'render_cond.py'),
|
||||
'--',
|
||||
'--object', os.path.expanduser(file_path),
|
||||
'--cond_views', json.dumps(cond_views),
|
||||
'--cond_resolution', '1024',
|
||||
'--cond_output_folder', os.path.join(root, 'renders_cond', sha256),
|
||||
'--engine', 'CYCLES',
|
||||
]
|
||||
if file_path.endswith('.blend'):
|
||||
args.insert(1, file_path)
|
||||
|
||||
call(args, stdout=DEVNULL, stderr=DEVNULL)
|
||||
|
||||
if os.path.exists(os.path.join(root, 'renders_cond', sha256, 'transforms.json')):
|
||||
return {'sha256': sha256, 'cond_rendered': True}
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--download_root', type=str, default=None,
|
||||
help='Directory to save the downloaded files')
|
||||
parser.add_argument('--render_cond_root', type=str, default=None,
|
||||
help='Directory to save the mesh dumps')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--num_cond_views', type=int, default=2,
|
||||
help='Number of conditional views to render')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=8)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.download_root = opt.download_root or opt.root
|
||||
opt.render_cond_root = opt.render_cond_root or opt.root
|
||||
|
||||
os.makedirs(os.path.join(opt.render_cond_root, 'renders_cond', 'new_records'), exist_ok=True)
|
||||
|
||||
# install blender
|
||||
print('Checking blender...', flush=True)
|
||||
_install_blender()
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.render_cond_root, 'renders_cond', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.render_cond_root, 'renders_cond', 'metadata.csv')).set_index('sha256'))
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
metadata = metadata[metadata['local_path'].notna()]
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
if 'cond_rendered' in metadata.columns:
|
||||
metadata = metadata[(metadata['cond_rendered'] != True)]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
records = []
|
||||
|
||||
# filter out objects that are already processed
|
||||
with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
|
||||
tqdm(total=len(metadata), desc="Filtering existing objects") as pbar:
|
||||
def check_sha256(sha256):
|
||||
if os.path.exists(os.path.join(opt.render_cond_root, 'renders_cond', sha256, 'transforms.json')):
|
||||
records.append({'sha256': sha256, 'cond_rendered': True})
|
||||
pbar.update()
|
||||
executor.map(check_sha256, metadata['sha256'].values)
|
||||
executor.shutdown(wait=True)
|
||||
existing_sha256 = set(r['sha256'] for r in records)
|
||||
metadata = metadata[~metadata['sha256'].isin(existing_sha256)]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
func = partial(_render_cond, root=opt.render_cond_root, num_cond_views=opt.num_cond_views)
|
||||
cond_rendered = dataset_utils.foreach_instance(metadata, opt.render_cond_root, func, max_workers=opt.max_workers, desc='Rendering objects')
|
||||
cond_rendered = pd.concat([cond_rendered, pd.DataFrame.from_records(records)])
|
||||
cond_rendered.to_csv(os.path.join(opt.render_cond_root, 'renders_cond', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1 @@
|
||||
pip install pillow imageio imageio-ffmpeg tqdm easydict opencv-python-headless pandas open3d objaverse huggingface_hub[cli] open_clip_torch
|
||||
@@ -0,0 +1,564 @@
|
||||
from typing import *
|
||||
import hashlib
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
def get_file_hash(file: str) -> str:
|
||||
sha256 = hashlib.sha256()
|
||||
# Read the file from the path
|
||||
with open(file, "rb") as f:
|
||||
# Update the hash with the file content
|
||||
for byte_block in iter(lambda: f.read(4096), b""):
|
||||
sha256.update(byte_block)
|
||||
return sha256.hexdigest()
|
||||
|
||||
# ===============LOW DISCREPANCY SEQUENCES================
|
||||
|
||||
PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
|
||||
|
||||
def radical_inverse(base, n):
|
||||
val = 0
|
||||
inv_base = 1.0 / base
|
||||
inv_base_n = inv_base
|
||||
while n > 0:
|
||||
digit = n % base
|
||||
val += digit * inv_base_n
|
||||
n //= base
|
||||
inv_base_n *= inv_base
|
||||
return val
|
||||
|
||||
def halton_sequence(dim, n):
|
||||
return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
|
||||
|
||||
def hammersley_sequence(dim, n, num_samples):
|
||||
return [n / num_samples] + halton_sequence(dim - 1, n)
|
||||
|
||||
def sphere_hammersley_sequence(n, num_samples, offset=(0, 0)):
|
||||
u, v = hammersley_sequence(2, n, num_samples)
|
||||
u += offset[0] / num_samples
|
||||
v += offset[1]
|
||||
u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
|
||||
theta = np.arccos(1 - 2 * u) - np.pi / 2
|
||||
phi = v * 2 * np.pi
|
||||
return [phi, theta]
|
||||
|
||||
# ==============PLY IO===============
|
||||
import struct
|
||||
import re
|
||||
import torch
|
||||
|
||||
def read_ply(filename):
|
||||
"""
|
||||
Read a PLY file and return vertices, triangle faces, and quad faces.
|
||||
|
||||
Args:
|
||||
filename (str): The file path to read from.
|
||||
|
||||
Returns:
|
||||
vertices (torch.Tensor): Tensor of shape [N, 3] containing vertex positions.
|
||||
tris (torch.Tensor): Tensor of shape [M, 3] containing triangle face indices (empty if none).
|
||||
quads (torch.Tensor): Tensor of shape [K, 4] containing quad face indices (empty if none).
|
||||
"""
|
||||
with open(filename, 'rb') as f:
|
||||
# Read the header until 'end_header' is encountered
|
||||
header_bytes = b""
|
||||
while True:
|
||||
line = f.readline()
|
||||
if not line:
|
||||
raise ValueError("PLY header not found")
|
||||
header_bytes += line
|
||||
if b"end_header" in line:
|
||||
break
|
||||
header = header_bytes.decode('utf-8')
|
||||
|
||||
# Determine if the file is in ASCII or binary format
|
||||
is_ascii = "ascii" in header
|
||||
|
||||
# Extract the number of vertices and faces from the header using regex
|
||||
vertex_match = re.search(r'element vertex (\d+)', header)
|
||||
if vertex_match:
|
||||
num_vertices = int(vertex_match.group(1))
|
||||
else:
|
||||
raise ValueError("Vertex count not found in header")
|
||||
|
||||
face_match = re.search(r'element face (\d+)', header)
|
||||
if face_match:
|
||||
num_faces = int(face_match.group(1))
|
||||
else:
|
||||
raise ValueError("Face count not found in header")
|
||||
|
||||
vertices = []
|
||||
tris = []
|
||||
quads = []
|
||||
|
||||
if is_ascii:
|
||||
# For ASCII format, read each line of vertex data (each line contains 3 floats)
|
||||
for _ in range(num_vertices):
|
||||
line = f.readline().decode('utf-8').strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split()
|
||||
vertices.append([float(parts[0]), float(parts[1]), float(parts[2])])
|
||||
|
||||
# Read face data, where the first number indicates the number of vertices for the face
|
||||
for _ in range(num_faces):
|
||||
line = f.readline().decode('utf-8').strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split()
|
||||
count = int(parts[0])
|
||||
indices = list(map(int, parts[1:]))
|
||||
if count == 3:
|
||||
tris.append(indices)
|
||||
elif count == 4:
|
||||
quads.append(indices)
|
||||
else:
|
||||
# Skip faces with other numbers of vertices (can be extended as needed)
|
||||
pass
|
||||
else:
|
||||
# For binary format: read directly from the binary stream
|
||||
# Each vertex consists of 3 floats (12 bytes per vertex)
|
||||
for _ in range(num_vertices):
|
||||
data = f.read(12)
|
||||
if len(data) < 12:
|
||||
raise ValueError("Insufficient vertex data")
|
||||
v = struct.unpack('<fff', data)
|
||||
vertices.append(v)
|
||||
|
||||
# Read face data from the binary stream
|
||||
for _ in range(num_faces):
|
||||
# First, read 1 byte indicating the number of vertices in the face
|
||||
count_data = f.read(1)
|
||||
if len(count_data) < 1:
|
||||
raise ValueError("Failed to read face vertex count")
|
||||
count = struct.unpack('<B', count_data)[0]
|
||||
if count == 3:
|
||||
data = f.read(12) # 3 * 4 bytes
|
||||
if len(data) < 12:
|
||||
raise ValueError("Insufficient data for triangle face")
|
||||
indices = struct.unpack('<3i', data)
|
||||
tris.append(indices)
|
||||
elif count == 4:
|
||||
data = f.read(16) # 4 * 4 bytes
|
||||
if len(data) < 16:
|
||||
raise ValueError("Insufficient data for quad face")
|
||||
indices = struct.unpack('<4i', data)
|
||||
quads.append(indices)
|
||||
else:
|
||||
# For faces with a different number of vertices, read count*4 bytes
|
||||
data = f.read(count * 4)
|
||||
# Skip or extend processing as needed
|
||||
raise ValueError(f"Unsupported face with {count} vertices")
|
||||
|
||||
# Convert lists to torch.Tensor
|
||||
vertices = torch.tensor(vertices, dtype=torch.float32)
|
||||
tris = torch.tensor(tris, dtype=torch.int32) if len(tris) > 0 else torch.empty((0, 3), dtype=torch.int32)
|
||||
quads = torch.tensor(quads, dtype=torch.int32) if len(quads) > 0 else torch.empty((0, 4), dtype=torch.int32)
|
||||
|
||||
return vertices, tris, quads
|
||||
|
||||
|
||||
def write_ply(filename, vertices, tris, quads, ascii=False):
|
||||
"""
|
||||
Write a mesh to a PLY file, with the option to save in ASCII or binary format.
|
||||
|
||||
Args:
|
||||
filename (str): The filename to write to.
|
||||
vertices (torch.Tensor): [N, 3] The vertex positions.
|
||||
tris (torch.Tensor): [M, 3] The triangle indices.
|
||||
quads (torch.Tensor): [K, 4] The quad indices.
|
||||
ascii (bool): If True, write in ASCII format. If False, write in binary format.
|
||||
"""
|
||||
# Convert torch tensors to numpy arrays
|
||||
vertices = vertices.numpy()
|
||||
tris = tris.numpy()
|
||||
quads = quads.numpy()
|
||||
|
||||
# Prepare the header
|
||||
num_vertices = len(vertices)
|
||||
num_faces = len(tris) + len(quads)
|
||||
|
||||
# Vertex properties
|
||||
vertex_header = "property float x\nproperty float y\nproperty float z"
|
||||
|
||||
# Face properties (the number of vertices per face is variable)
|
||||
face_header = "property list uchar int vertex_index"
|
||||
|
||||
# Start writing the PLY header
|
||||
header = f"ply\n"
|
||||
header += f"format {'ascii 1.0' if ascii else 'binary_little_endian 1.0'}\n"
|
||||
header += f"element vertex {num_vertices}\n"
|
||||
header += vertex_header + "\n"
|
||||
header += f"element face {num_faces}\n"
|
||||
header += face_header + "\n"
|
||||
header += "end_header\n"
|
||||
|
||||
# Open the file for writing
|
||||
with open(filename, 'wb' if not ascii else 'w') as f:
|
||||
# Write the header
|
||||
f.write(header if ascii else header.encode('utf-8'))
|
||||
|
||||
# Write the vertex data
|
||||
if ascii:
|
||||
for v in vertices:
|
||||
f.write(f"{v[0]} {v[1]} {v[2]}\n")
|
||||
else:
|
||||
for v in vertices:
|
||||
f.write(struct.pack('<fff', *v))
|
||||
|
||||
# Write the face data
|
||||
if ascii:
|
||||
for tri in tris:
|
||||
f.write(f"3 {tri[0]} {tri[1]} {tri[2]}\n")
|
||||
for quad in quads:
|
||||
f.write(f"4 {quad[0]} {quad[1]} {quad[2]} {quad[3]}\n")
|
||||
else:
|
||||
for tri in tris:
|
||||
f.write(struct.pack('<B3i', 3, *tri)) # 3 indices for triangle
|
||||
for quad in quads:
|
||||
f.write(struct.pack('<B4i', 4, *quad)) # 4 indices for quad
|
||||
|
||||
|
||||
# ==============IMAGE UTILS===============
|
||||
|
||||
def make_grid(images, nrow=None, ncol=None, aspect_ratio=None):
|
||||
num_images = len(images)
|
||||
if nrow is None and ncol is None:
|
||||
if aspect_ratio is not None:
|
||||
nrow = int(np.round(np.sqrt(num_images / aspect_ratio)))
|
||||
else:
|
||||
nrow = int(np.sqrt(num_images))
|
||||
ncol = (num_images + nrow - 1) // nrow
|
||||
elif nrow is None and ncol is not None:
|
||||
nrow = (num_images + ncol - 1) // ncol
|
||||
elif nrow is not None and ncol is None:
|
||||
ncol = (num_images + nrow - 1) // nrow
|
||||
else:
|
||||
assert nrow * ncol >= num_images, 'nrow * ncol must be greater than or equal to the number of images'
|
||||
|
||||
if images[0].ndim == 2:
|
||||
grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1]), dtype=images[0].dtype)
|
||||
else:
|
||||
grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1], images[0].shape[2]), dtype=images[0].dtype)
|
||||
for i, img in enumerate(images):
|
||||
row = i // ncol
|
||||
col = i % ncol
|
||||
grid[row * img.shape[0]:(row + 1) * img.shape[0], col * img.shape[1]:(col + 1) * img.shape[1]] = img
|
||||
return grid
|
||||
|
||||
|
||||
def notes_on_image(img, notes=None):
|
||||
img = np.pad(img, ((0, 32), (0, 0), (0, 0)), 'constant', constant_values=0)
|
||||
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||||
if notes is not None:
|
||||
img = cv2.putText(img, notes, (0, img.shape[0] - 4), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1)
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
return img
|
||||
|
||||
|
||||
|
||||
def text_image(text, resolution=(512, 512), max_size=0.5, h_align="left", v_align="center"):
|
||||
"""
|
||||
Draw text on an image of the given resolution. The text is automatically wrapped
|
||||
and scaled so that it fits completely within the image while preserving any explicit
|
||||
line breaks and original spacing. Horizontal and vertical alignment can be controlled
|
||||
via flags.
|
||||
|
||||
Parameters:
|
||||
text (str): The input text. Newline characters and spacing are preserved.
|
||||
resolution (tuple): The image resolution as (width, height).
|
||||
max_size (float): The maximum font size.
|
||||
h_align (str): Horizontal alignment. Options: "left", "center", "right".
|
||||
v_align (str): Vertical alignment. Options: "top", "center", "bottom".
|
||||
|
||||
Returns:
|
||||
numpy.ndarray: The resulting image (BGR format) with the text drawn.
|
||||
"""
|
||||
width, height = resolution
|
||||
# Create a white background image
|
||||
img = np.full((height, width, 3), 255, dtype=np.uint8)
|
||||
|
||||
# Set margins and compute available drawing area
|
||||
margin = 10
|
||||
avail_width = width - 2 * margin
|
||||
avail_height = height - 2 * margin
|
||||
|
||||
# Choose OpenCV font and text thickness
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
thickness = 1
|
||||
# Ratio for additional spacing between lines (relative to the height of "A")
|
||||
line_spacing_ratio = 0.5
|
||||
|
||||
def wrap_line(line, max_width, font, thickness, scale):
|
||||
"""
|
||||
Wrap a single line of text into multiple lines such that each line's
|
||||
width (measured at the given scale) does not exceed max_width.
|
||||
This function preserves the original spacing by splitting the line into tokens
|
||||
(words and whitespace) using a regular expression.
|
||||
|
||||
Parameters:
|
||||
line (str): The input text line.
|
||||
max_width (int): Maximum allowed width in pixels.
|
||||
font (int): OpenCV font identifier.
|
||||
thickness (int): Text thickness.
|
||||
scale (float): The current font scale.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of wrapped lines.
|
||||
"""
|
||||
# Split the line into tokens (words and whitespace), preserving spacing
|
||||
tokens = re.split(r'(\s+)', line)
|
||||
if not tokens:
|
||||
return ['']
|
||||
|
||||
wrapped_lines = []
|
||||
current_line = ""
|
||||
for token in tokens:
|
||||
candidate = current_line + token
|
||||
candidate_width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
|
||||
if candidate_width <= max_width:
|
||||
current_line = candidate
|
||||
else:
|
||||
# If current_line is empty, the token itself is too wide;
|
||||
# break the token character by character.
|
||||
if current_line == "":
|
||||
sub_token = ""
|
||||
for char in token:
|
||||
candidate_char = sub_token + char
|
||||
if cv2.getTextSize(candidate_char, font, scale, thickness)[0][0] <= max_width:
|
||||
sub_token = candidate_char
|
||||
else:
|
||||
if sub_token:
|
||||
wrapped_lines.append(sub_token)
|
||||
sub_token = char
|
||||
current_line = sub_token
|
||||
else:
|
||||
wrapped_lines.append(current_line)
|
||||
current_line = token
|
||||
if current_line:
|
||||
wrapped_lines.append(current_line)
|
||||
return wrapped_lines
|
||||
|
||||
def compute_text_block(scale):
|
||||
"""
|
||||
Wrap the entire text (splitting at explicit newline characters) using the
|
||||
provided scale, and then compute the overall width and height of the text block.
|
||||
|
||||
Returns:
|
||||
wrapped_lines (List[str]): The list of wrapped lines.
|
||||
block_width (int): Maximum width among the wrapped lines.
|
||||
block_height (int): Total height of the text block including spacing.
|
||||
sizes (List[tuple]): A list of (width, height) for each wrapped line.
|
||||
spacing (int): The spacing between lines (computed from the scaled "A" height).
|
||||
"""
|
||||
# Split text by explicit newlines
|
||||
input_lines = text.splitlines() if text else ['']
|
||||
wrapped_lines = []
|
||||
for line in input_lines:
|
||||
wrapped = wrap_line(line, avail_width, font, thickness, scale)
|
||||
wrapped_lines.extend(wrapped)
|
||||
|
||||
sizes = []
|
||||
for line in wrapped_lines:
|
||||
(text_size, _) = cv2.getTextSize(line, font, scale, thickness)
|
||||
sizes.append(text_size) # (width, height)
|
||||
|
||||
block_width = max((w for w, h in sizes), default=0)
|
||||
# Use the height of "A" (at the current scale) to compute line spacing
|
||||
base_height = cv2.getTextSize("A", font, scale, thickness)[0][1]
|
||||
spacing = int(line_spacing_ratio * base_height)
|
||||
block_height = sum(h for w, h in sizes) + spacing * (len(sizes) - 1) if sizes else 0
|
||||
|
||||
return wrapped_lines, block_width, block_height, sizes, spacing
|
||||
|
||||
# Use binary search to find the maximum scale that allows the text block to fit
|
||||
lo = 0.001
|
||||
hi = max_size
|
||||
eps = 0.001 # convergence threshold
|
||||
best_scale = lo
|
||||
best_result = None
|
||||
|
||||
while hi - lo > eps:
|
||||
mid = (lo + hi) / 2
|
||||
wrapped_lines, block_width, block_height, sizes, spacing = compute_text_block(mid)
|
||||
# Ensure that both width and height constraints are met
|
||||
if block_width <= avail_width and block_height <= avail_height:
|
||||
best_scale = mid
|
||||
best_result = (wrapped_lines, block_width, block_height, sizes, spacing)
|
||||
lo = mid # try a larger scale
|
||||
else:
|
||||
hi = mid # reduce the scale
|
||||
|
||||
if best_result is None:
|
||||
best_scale = 0.5
|
||||
best_result = compute_text_block(best_scale)
|
||||
|
||||
wrapped_lines, block_width, block_height, sizes, spacing = best_result
|
||||
|
||||
# Compute starting y-coordinate based on vertical alignment flag
|
||||
if v_align == "top":
|
||||
y_top = margin
|
||||
elif v_align == "center":
|
||||
y_top = margin + (avail_height - block_height) // 2
|
||||
elif v_align == "bottom":
|
||||
y_top = margin + (avail_height - block_height)
|
||||
else:
|
||||
y_top = margin + (avail_height - block_height) // 2 # default to center if invalid flag
|
||||
|
||||
# For cv2.putText, the y coordinate represents the text baseline;
|
||||
# so for the first line add its height.
|
||||
y = y_top + (sizes[0][1] if sizes else 0)
|
||||
|
||||
# Draw each line with horizontal alignment based on the flag
|
||||
for i, line in enumerate(wrapped_lines):
|
||||
line_width, line_height = sizes[i]
|
||||
if h_align == "left":
|
||||
x = margin
|
||||
elif h_align == "center":
|
||||
x = margin + (avail_width - line_width) // 2
|
||||
elif h_align == "right":
|
||||
x = margin + (avail_width - line_width)
|
||||
else:
|
||||
x = margin # default to left if invalid flag
|
||||
|
||||
cv2.putText(img, line, (x, y), font, best_scale, (0, 0, 0), thickness, cv2.LINE_AA)
|
||||
y += line_height + spacing
|
||||
|
||||
return img
|
||||
|
||||
|
||||
# ==================== View index parsing ====================
|
||||
|
||||
def parse_view_indices(view_indices_str):
|
||||
"""Parse view_indices string into a sorted deduplicated list of integers."""
|
||||
if view_indices_str is None:
|
||||
return None
|
||||
view_indices = []
|
||||
for part in view_indices_str.split(','):
|
||||
if '-' in part:
|
||||
start, end = map(int, part.split('-'))
|
||||
view_indices.extend(range(start, end + 1))
|
||||
else:
|
||||
view_indices.append(int(part))
|
||||
view_indices = list(set(view_indices))
|
||||
view_indices.sort()
|
||||
return view_indices
|
||||
|
||||
|
||||
# ==================== Multi-view transform functions ====================
|
||||
|
||||
import math
|
||||
|
||||
def get_new_camera_matrix(radius: float, yaw: float, pitch: float, dtype=torch.float32, device='cpu'):
|
||||
"""
|
||||
Compute camera-to-world 4x4 transform matrix in spherical coordinates,
|
||||
looking at origin with up=(0,1,0).
|
||||
Uses standard LookAt formula where camera local -Z points at target and local Y is up.
|
||||
An additional +90 degree rotation around local Z axis is applied to match
|
||||
Blender Track-To local axis convention.
|
||||
yaw, pitch are in radians.
|
||||
"""
|
||||
x = radius * math.cos(yaw) * math.cos(pitch)
|
||||
y = radius * math.sin(yaw) * math.cos(pitch)
|
||||
z = radius * math.sin(pitch)
|
||||
eye = torch.tensor([x, y, z], dtype=dtype, device=device)
|
||||
target = torch.zeros(3, dtype=dtype, device=device)
|
||||
up_global = torch.tensor([0.0, 1.0, 0.0], dtype=dtype, device=device)
|
||||
|
||||
f = (target - eye)
|
||||
f = f / torch.norm(f)
|
||||
r = torch.cross(f, up_global)
|
||||
r = r / torch.norm(r)
|
||||
u = torch.cross(r, f)
|
||||
|
||||
z_cam = -f
|
||||
x_cam = r
|
||||
y_cam = u
|
||||
|
||||
T = torch.eye(4, dtype=dtype, device=device)
|
||||
T[:3, 0] = x_cam
|
||||
T[:3, 1] = y_cam
|
||||
T[:3, 2] = z_cam
|
||||
T[:3, 3] = eye
|
||||
|
||||
# +90 degree rotation matrix around local Z axis
|
||||
Rz90 = torch.tensor([
|
||||
[0.0, -1.0, 0.0, 0.0],
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=dtype, device=device)
|
||||
|
||||
return T @ Rz90
|
||||
|
||||
|
||||
def transform_mesh(mesh_v, frame):
|
||||
"""
|
||||
Apply multi-view transform to mesh vertices based on camera transform matrix.
|
||||
"""
|
||||
device = mesh_v.device
|
||||
c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
|
||||
|
||||
# Old and new camera matrices
|
||||
radius = c2w_orig[:3, 3].norm().item()
|
||||
c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
|
||||
dtype=torch.float32, device=device)
|
||||
w2c_orig = torch.inverse(c2w_orig)
|
||||
|
||||
# Initial and final axis alignment matrices
|
||||
R_init = torch.tensor([
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, -1.0, 0.0],
|
||||
[0.0, 1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
R_back = torch.tensor([
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0, 0.0],
|
||||
[0.0, -1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
R_ply = torch.tensor([
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0, 0.0],
|
||||
[0.0, -1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
T_cam = c2w_new @ w2c_orig @ R_ply
|
||||
T_final = R_back @ T_cam @ R_init
|
||||
|
||||
# Apply transform
|
||||
mesh_v = mesh_v.reshape(-1, 3)
|
||||
verts_h = torch.cat([mesh_v, torch.ones((mesh_v.shape[0], 1), dtype=torch.float32, device=device)], dim=1)
|
||||
verts_trans = (T_final @ verts_h.T).T[:, :3]
|
||||
|
||||
return verts_trans
|
||||
|
||||
|
||||
def sphere_normalize_torch(vertices):
|
||||
"""
|
||||
Sphere normalization: normalize vertices based on sphere radius.
|
||||
"""
|
||||
bmin = torch.min(vertices, dim=0)[0]
|
||||
bmax = torch.max(vertices, dim=0)[0]
|
||||
bcenter = (bmax + bmin) / 2
|
||||
assert bcenter.abs().max() < 0.25, f"bcenter is not close to origin: {bcenter}"
|
||||
radius = torch.norm(vertices - bcenter, dim=-1).max()
|
||||
vertices_normalized = vertices / radius
|
||||
return vertices_normalized, bcenter, radius
|
||||
|
||||
|
||||
def save_image_with_notes(img, path, notes=None):
|
||||
"""
|
||||
Save an image with notes.
|
||||
"""
|
||||
if isinstance(img, torch.Tensor):
|
||||
img = img.cpu().numpy().transpose(1, 2, 0)
|
||||
if img.dtype == np.float32 or img.dtype == np.float64:
|
||||
img = np.clip(img * 255, 0, 255).astype(np.uint8)
|
||||
img = notes_on_image(img, notes)
|
||||
cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
|
||||
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
Decode shape + PBR latent (.npz) to textured GLB mesh and render the PBR front view.
|
||||
|
||||
Usage:
|
||||
python data_toolkit/visualize_pbr_latent.py \
|
||||
--root datasets/ObjaverseXL_sketchfab \
|
||||
--sha256 <SHA256_HASH> \
|
||||
--resolution 1024 \
|
||||
--view_idx 0
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
from PIL import Image
|
||||
|
||||
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
|
||||
|
||||
import pixal3d.models as models
|
||||
import pixal3d.modules.sparse as sp
|
||||
from pixal3d.representations import MeshWithVoxel
|
||||
from pixal3d.renderers import EnvMap
|
||||
from pixal3d.utils import render_utils
|
||||
import o_voxel
|
||||
|
||||
PBR_ATTR_LAYOUT = {
|
||||
'base_color': slice(0, 3),
|
||||
'metallic': slice(3, 4),
|
||||
'roughness': slice(4, 5),
|
||||
'alpha': slice(5, 6),
|
||||
}
|
||||
|
||||
|
||||
def load_latent(latent_file):
|
||||
"""Load a latent .npz file and return a SparseTensor on GPU."""
|
||||
data = np.load(latent_file)
|
||||
coords = torch.tensor(data['coords']).int()
|
||||
feats = torch.tensor(data['feats']).float()
|
||||
coords = torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=1)
|
||||
return sp.SparseTensor(feats.cuda(), coords.cuda())
|
||||
|
||||
|
||||
def load_envmaps(device='cuda'):
|
||||
"""Load HDRI environment maps from assets/."""
|
||||
base = os.path.join(os.path.dirname(__file__), '..', 'assets', 'hdri')
|
||||
envmaps = {}
|
||||
for name in ['forest', 'sunset', 'courtyard']:
|
||||
path = os.path.join(base, f'{name}.exr')
|
||||
if os.path.exists(path):
|
||||
img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
envmaps[name] = EnvMap(torch.tensor(img, dtype=torch.float32, device=device))
|
||||
return envmaps
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Decode shape + PBR latent to textured GLB and render")
|
||||
parser.add_argument("--root", type=str, required=True, help="Dataset root")
|
||||
parser.add_argument("--sha256", type=str, required=True, help="SHA256 of the asset")
|
||||
parser.add_argument("--resolution", type=int, default=1024, help="Decoder resolution")
|
||||
parser.add_argument("--view_idx", type=int, default=0, help="View index to decode")
|
||||
parser.add_argument("--shape_latent_name", type=str, default="shape_enc_next_dc_f16c32_fp16_1024_view",
|
||||
help="Shape latent directory name under shape_latents/")
|
||||
parser.add_argument("--pbr_latent_name", type=str, default="tex_enc_next_dc_f16c32_fp16_1024_view_fix",
|
||||
help="PBR latent directory name under pbr_latents/")
|
||||
parser.add_argument("--shape_decoder", type=str, default="microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16",
|
||||
help="Pretrained shape decoder")
|
||||
parser.add_argument("--pbr_decoder", type=str, default="microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
|
||||
help="Pretrained PBR/texture decoder")
|
||||
parser.add_argument("--texture_size", type=int, default=4096, help="GLB texture resolution")
|
||||
parser.add_argument("--decimation_target", type=int, default=1000000, help="GLB mesh decimation target")
|
||||
parser.add_argument("--output_dir", type=str, default=None, help="Output directory (default: <root>/vis_pbr/<sha256>)")
|
||||
args = parser.parse_args()
|
||||
|
||||
sha256 = args.sha256
|
||||
root = args.root
|
||||
view_idx = args.view_idx
|
||||
|
||||
# Paths
|
||||
shape_latent_dir = os.path.join(root, "shape_latents", args.shape_latent_name, sha256)
|
||||
pbr_latent_dir = os.path.join(root, "pbr_latents", args.pbr_latent_name, sha256)
|
||||
shape_file = os.path.join(shape_latent_dir, f"view{view_idx:02d}.npz")
|
||||
pbr_file = os.path.join(pbr_latent_dir, f"view{view_idx:02d}.npz")
|
||||
renders_dir = os.path.join(root, "renders_cond", sha256)
|
||||
output_dir = args.output_dir or os.path.join(root, "vis_pbr", sha256)
|
||||
|
||||
# Validate
|
||||
assert os.path.exists(shape_file), f"Shape latent not found: {shape_file}"
|
||||
assert os.path.exists(pbr_file), f"PBR latent not found: {pbr_file}"
|
||||
print(f"[Input] Shape latent: {shape_file}")
|
||||
print(f"[Input] PBR latent: {pbr_file}")
|
||||
if os.path.exists(renders_dir):
|
||||
print(f"[Input] Renders: {renders_dir}")
|
||||
|
||||
# 1. Load latents
|
||||
print("[Step 1] Loading latents...")
|
||||
shape_slat = load_latent(shape_file)
|
||||
pbr_slat = load_latent(pbr_file)
|
||||
print(f" Shape: coords {shape_slat.coords.shape}, feats {shape_slat.feats.shape}")
|
||||
print(f" PBR: coords {pbr_slat.coords.shape}, feats {pbr_slat.feats.shape}")
|
||||
|
||||
# 2. Load decoders
|
||||
print(f"[Step 2] Loading decoders...")
|
||||
shape_dec = models.from_pretrained(args.shape_decoder)
|
||||
shape_dec.set_resolution(args.resolution)
|
||||
shape_dec = shape_dec.cuda().eval()
|
||||
|
||||
pbr_dec = models.from_pretrained(args.pbr_decoder)
|
||||
pbr_dec = pbr_dec.cuda().eval()
|
||||
|
||||
# 3. Decode shape → mesh + subs, then PBR → voxel
|
||||
print("[Step 3] Decoding shape + PBR latents...")
|
||||
with torch.no_grad():
|
||||
meshes, subs = shape_dec(shape_slat, return_subs=True)
|
||||
vox = pbr_dec(pbr_slat, guide_subs=subs) * 0.5 + 0.5
|
||||
|
||||
mesh = meshes[0]
|
||||
mesh.fill_holes()
|
||||
mesh_with_voxel = MeshWithVoxel(
|
||||
mesh.vertices, mesh.faces,
|
||||
origin=[-0.5, -0.5, -0.5],
|
||||
voxel_size=1 / args.resolution,
|
||||
coords=vox[0].coords[:, 1:],
|
||||
attrs=vox[0].feats,
|
||||
voxel_shape=torch.Size([*vox[0].shape, *vox[0].spatial_shape]),
|
||||
layout=PBR_ATTR_LAYOUT,
|
||||
)
|
||||
print(f" Mesh: vertices {mesh.vertices.shape}, faces {mesh.faces.shape}")
|
||||
print(f" Voxel: coords {vox[0].coords.shape}, feats {vox[0].feats.shape}")
|
||||
|
||||
# 4. Export GLB with PBR textures
|
||||
print("[Step 4] Extracting textured GLB...")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
glb = o_voxel.postprocess.to_glb(
|
||||
vertices=mesh_with_voxel.vertices,
|
||||
faces=mesh_with_voxel.faces,
|
||||
attr_volume=mesh_with_voxel.attrs,
|
||||
coords=mesh_with_voxel.coords,
|
||||
attr_layout=PBR_ATTR_LAYOUT,
|
||||
grid_size=args.resolution,
|
||||
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
decimation_target=args.decimation_target,
|
||||
texture_size=args.texture_size,
|
||||
remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
|
||||
)
|
||||
# Apply rotation (same as inference.py)
|
||||
rot = np.array([
|
||||
[-1, 0, 0, 0],
|
||||
[ 0, 0, -1, 0],
|
||||
[ 0, -1, 0, 0],
|
||||
[ 0, 0, 0, 1],
|
||||
], dtype=np.float64)
|
||||
glb.apply_transform(rot)
|
||||
glb_path = os.path.join(output_dir, f"pbr_view{view_idx:02d}.glb")
|
||||
glb.export(glb_path, extension_webp=True)
|
||||
print(f" GLB saved: {glb_path}")
|
||||
|
||||
# 5. Render PBR front view (proj-aligned, same as app.py)
|
||||
print("[Step 5] Rendering PBR front view (proj-aligned)...")
|
||||
transforms_file = os.path.join(renders_dir, "transforms.json")
|
||||
shape_scale_file = os.path.join(shape_latent_dir, f"view{view_idx:02d}_scale.json")
|
||||
envmaps = load_envmaps(device='cuda')
|
||||
if os.path.exists(transforms_file) and os.path.exists(shape_scale_file) and envmaps:
|
||||
with open(transforms_file) as f:
|
||||
transforms = json.load(f)
|
||||
with open(shape_scale_file) as f:
|
||||
scale_info = json.load(f)
|
||||
total_scale = scale_info['total_scale']
|
||||
frame_info = transforms['frames'][view_idx]
|
||||
camera_angle_x = frame_info['camera_angle_x']
|
||||
distance = frame_info['radius']
|
||||
near = max(0.01, distance - 2.0)
|
||||
far = distance + 10.0
|
||||
# Scale mesh by 1/total_scale to match blender normalized space
|
||||
scaled_mesh = MeshWithVoxel(
|
||||
mesh_with_voxel.vertices / total_scale,
|
||||
mesh_with_voxel.faces,
|
||||
origin=[x / total_scale for x in mesh_with_voxel.origin],
|
||||
voxel_size=mesh_with_voxel.voxel_size / total_scale,
|
||||
coords=mesh_with_voxel.coords,
|
||||
attrs=mesh_with_voxel.attrs,
|
||||
voxel_shape=mesh_with_voxel.voxel_shape,
|
||||
layout=PBR_ATTR_LAYOUT,
|
||||
)
|
||||
print(f" total_scale={total_scale:.4f}, distance={distance:.4f}, fov={camera_angle_x:.4f}")
|
||||
renders = render_utils.render_proj_aligned_video(
|
||||
scaled_mesh, camera_angle_x=camera_angle_x, distance=distance,
|
||||
resolution=1024, num_frames=1, envmap=envmaps, near=near, far=far,
|
||||
)
|
||||
for key, frames in renders.items():
|
||||
for i, frame in enumerate(frames):
|
||||
img = Image.fromarray(frame)
|
||||
img_path = os.path.join(output_dir, f"decoded_{key}_view{view_idx:02d}_{i:03d}.png")
|
||||
img.save(img_path)
|
||||
print(f" Saved {len(frames)} {key} images")
|
||||
else:
|
||||
if not os.path.exists(transforms_file):
|
||||
print(" No transforms.json found, skipping rendering.")
|
||||
if not os.path.exists(shape_scale_file):
|
||||
print(" No scale file found, skipping rendering.")
|
||||
if not envmaps:
|
||||
print(" No HDRI envmaps found, skipping PBR rendering.")
|
||||
|
||||
# Free GPU
|
||||
del shape_dec, pbr_dec, shape_slat, pbr_slat, meshes, subs, vox
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# 6. Copy condition renders
|
||||
if os.path.exists(renders_dir):
|
||||
print("[Step 6] Copying condition renders...")
|
||||
for fname in sorted(os.listdir(renders_dir)):
|
||||
src = os.path.join(renders_dir, fname)
|
||||
dst = os.path.join(output_dir, fname)
|
||||
shutil.copy2(src, dst)
|
||||
print(f" {fname}")
|
||||
else:
|
||||
print("[Step 6] No condition renders found, skipping.")
|
||||
|
||||
# 7. Copy scale info
|
||||
for src_dir, prefix in [(shape_latent_dir, "shape"), (pbr_latent_dir, "pbr")]:
|
||||
scale_file = os.path.join(src_dir, f"view{view_idx:02d}_scale.json")
|
||||
if os.path.exists(scale_file):
|
||||
shutil.copy2(scale_file, os.path.join(output_dir, f"{prefix}_view{view_idx:02d}_scale.json"))
|
||||
|
||||
print(f"\n[Done] All outputs in: {output_dir}")
|
||||
print(f" Files: {sorted(os.listdir(output_dir))}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,160 @@
|
||||
"""
|
||||
Decode a view-aligned shape latent (.npz) to GLB mesh and render the front view.
|
||||
|
||||
Usage:
|
||||
python data_toolkit/visualize_shape_latent.py \
|
||||
--root datasets/ObjaverseXL_sketchfab \
|
||||
--sha256 <SHA256_HASH> \
|
||||
--resolution 1024 \
|
||||
--view_idx 0
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
import shutil
|
||||
import argparse
|
||||
import numpy as np
|
||||
import torch
|
||||
import trimesh
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
|
||||
|
||||
import pixal3d.models as models
|
||||
import pixal3d.modules.sparse as sp
|
||||
from pixal3d.utils import render_utils
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Decode shape latent to GLB and collect renders")
|
||||
parser.add_argument("--root", type=str, required=True, help="Dataset root, e.g. /local-ssd/datasets/ObjaverseXL_sketchfab")
|
||||
parser.add_argument("--sha256", type=str, required=True, help="SHA256 of the asset")
|
||||
parser.add_argument("--resolution", type=int, default=1024, help="Decoder resolution (must match latent resolution)")
|
||||
parser.add_argument("--view_idx", type=int, default=0, help="View index to decode")
|
||||
parser.add_argument("--latent_name", type=str, default="shape_enc_next_dc_f16c32_fp16_1024_view",
|
||||
help="Latent directory name under shape_latents/")
|
||||
parser.add_argument("--decoder", type=str, default="microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16",
|
||||
help="Pretrained shape decoder path (HuggingFace or local)")
|
||||
parser.add_argument("--output_dir", type=str, default=None, help="Output directory (default: <root>/vis/<sha256>)")
|
||||
args = parser.parse_args()
|
||||
|
||||
sha256 = args.sha256
|
||||
root = args.root
|
||||
view_idx = args.view_idx
|
||||
|
||||
# Paths
|
||||
latent_dir = os.path.join(root, "shape_latents", args.latent_name, sha256)
|
||||
latent_file = os.path.join(latent_dir, f"view{view_idx:02d}.npz")
|
||||
scale_file = os.path.join(latent_dir, f"view{view_idx:02d}_scale.json")
|
||||
renders_dir = os.path.join(root, "renders_cond", sha256)
|
||||
output_dir = args.output_dir or os.path.join(root, "vis", sha256)
|
||||
|
||||
# Validate
|
||||
assert os.path.exists(latent_file), f"Latent file not found: {latent_file}"
|
||||
print(f"[Input] Latent: {latent_file}")
|
||||
if os.path.exists(scale_file):
|
||||
print(f"[Input] Scale: {scale_file}")
|
||||
if os.path.exists(renders_dir):
|
||||
print(f"[Input] Renders: {renders_dir}")
|
||||
|
||||
# 1. Load latent
|
||||
print("[Step 1] Loading shape latent...")
|
||||
data = np.load(latent_file)
|
||||
coords = torch.tensor(data['coords']).int()
|
||||
feats = torch.tensor(data['feats']).float()
|
||||
# Prepend batch dim (0) to coords
|
||||
coords = torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=1)
|
||||
slat = sp.SparseTensor(feats.cuda(), coords.cuda())
|
||||
print(f" coords: {coords.shape}, feats: {feats.shape}")
|
||||
|
||||
# 2. Load decoder
|
||||
print(f"[Step 2] Loading shape decoder: {args.decoder}")
|
||||
decoder = models.from_pretrained(args.decoder)
|
||||
decoder.set_resolution(args.resolution)
|
||||
decoder = decoder.cuda().eval()
|
||||
|
||||
# 3. Decode
|
||||
print("[Step 3] Decoding shape latent → mesh...")
|
||||
with torch.no_grad():
|
||||
meshes, subs = decoder(slat, return_subs=True)
|
||||
mesh = meshes[0]
|
||||
print(f" vertices: {mesh.vertices.shape}, faces: {mesh.faces.shape}")
|
||||
|
||||
# 4. Convert to trimesh and export GLB
|
||||
print("[Step 4] Exporting GLB...")
|
||||
vertices = mesh.vertices.cpu().numpy()
|
||||
faces = mesh.faces.cpu().numpy()
|
||||
|
||||
# Apply coordinate rotation (same as inference.py)
|
||||
# Swap axes: x→-x, y→-z, z→-y
|
||||
rot = np.array([
|
||||
[-1, 0, 0],
|
||||
[ 0, 0, -1],
|
||||
[ 0, -1, 0],
|
||||
], dtype=np.float64)
|
||||
vertices = vertices @ rot.T
|
||||
|
||||
tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
|
||||
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
glb_path = os.path.join(output_dir, f"shape_view{view_idx:02d}.glb")
|
||||
tri_mesh.export(glb_path)
|
||||
print(f" GLB saved: {glb_path}")
|
||||
|
||||
# 5. Render front view (proj-aligned, same as app.py)
|
||||
print("[Step 5] Rendering decoded mesh (proj-aligned front view)...")
|
||||
transforms_file = os.path.join(renders_dir, "transforms.json")
|
||||
if os.path.exists(transforms_file) and os.path.exists(scale_file):
|
||||
with open(transforms_file) as f:
|
||||
transforms = json.load(f)
|
||||
with open(scale_file) as f:
|
||||
scale_info = json.load(f)
|
||||
total_scale = scale_info['total_scale']
|
||||
frame_info = transforms['frames'][view_idx]
|
||||
camera_angle_x = frame_info['camera_angle_x']
|
||||
distance = frame_info['radius']
|
||||
near = max(0.01, distance - 2.0)
|
||||
far = distance + 10.0
|
||||
# Scale mesh vertices by 1/total_scale to match blender normalized space
|
||||
from pixal3d.representations import Mesh
|
||||
scaled_mesh = Mesh(mesh.vertices / total_scale, mesh.faces)
|
||||
print(f" total_scale={total_scale:.4f}, distance={distance:.4f}, fov={camera_angle_x:.4f}")
|
||||
renders = render_utils.render_proj_aligned_video(
|
||||
scaled_mesh, camera_angle_x=camera_angle_x, distance=distance,
|
||||
resolution=1024, num_frames=1, near=near, far=far,
|
||||
)
|
||||
for key, frames in renders.items():
|
||||
for i, frame in enumerate(frames):
|
||||
img = Image.fromarray(frame)
|
||||
img_path = os.path.join(output_dir, f"decoded_{key}_view{view_idx:02d}_{i:03d}.png")
|
||||
img.save(img_path)
|
||||
print(f" Saved {len(frames)} {key} images")
|
||||
else:
|
||||
print(" No transforms.json or scale file found, skipping rendering.")
|
||||
|
||||
# Free decoder GPU memory
|
||||
del decoder, slat, meshes, subs
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# 6. Copy renders
|
||||
if os.path.exists(renders_dir):
|
||||
print("[Step 6] Copying condition renders...")
|
||||
for fname in sorted(os.listdir(renders_dir)):
|
||||
src = os.path.join(renders_dir, fname)
|
||||
dst = os.path.join(output_dir, fname)
|
||||
shutil.copy2(src, dst)
|
||||
print(f" {fname}")
|
||||
else:
|
||||
print("[Step 6] No condition renders found, skipping.")
|
||||
|
||||
# 7. Copy scale info
|
||||
if os.path.exists(scale_file):
|
||||
shutil.copy2(scale_file, os.path.join(output_dir, f"view{view_idx:02d}_scale.json"))
|
||||
|
||||
print(f"\n[Done] All outputs in: {output_dir}")
|
||||
print(f" Files: {sorted(os.listdir(output_dir))}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,167 @@
|
||||
import os
|
||||
import copy
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import pandas as pd
|
||||
import pickle
|
||||
import numpy as np
|
||||
import torch
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
import o_voxel
|
||||
|
||||
|
||||
def _pbr_voxelize(file, metadatum, pbr_dump_root, root):
|
||||
sha256 = metadatum['sha256']
|
||||
try:
|
||||
pack = {'sha256': sha256}
|
||||
dump = None
|
||||
for res in opt.resolution:
|
||||
need_process = False
|
||||
|
||||
# check if already processed
|
||||
if os.path.exists(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz')):
|
||||
try:
|
||||
info = o_voxel.io.read_vxz_info(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz'))
|
||||
pack[f'pbr_voxelized_{res}'] = True
|
||||
pack[f'num_pbr_voxels_{res}'] = info['num_voxel']
|
||||
except Exception as e:
|
||||
print(f'Error reading {sha256}.vxz: {e}')
|
||||
need_process = True
|
||||
else:
|
||||
need_process = True
|
||||
|
||||
# process if necessary
|
||||
if need_process:
|
||||
if dump == None:
|
||||
with open(os.path.join(pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle'), 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
# Fix dump alpha map
|
||||
for mat in dump['materials']:
|
||||
if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
|
||||
mat['alphaMode'] = 'BLEND'
|
||||
dump['materials'].append({
|
||||
"baseColorFactor": [0.8, 0.8, 0.8],
|
||||
"alphaFactor": 1.0,
|
||||
"metallicFactor": 0.0,
|
||||
"roughnessFactor": 0.5,
|
||||
"alphaMode": "OPAQUE",
|
||||
"alphaCutoff": 0.5,
|
||||
"baseColorTexture": None,
|
||||
"alphaTexture": None,
|
||||
"metallicTexture": None,
|
||||
"roughnessTexture": None,
|
||||
}) # append default material
|
||||
dump['objects'] = [
|
||||
obj for obj in dump['objects']
|
||||
if obj['vertices'].size != 0 and obj['faces'].size != 0
|
||||
]
|
||||
vertices = torch.from_numpy(np.concatenate([obj['vertices'] for obj in dump['objects']], axis=0)).float()
|
||||
vertices_min = vertices.min(dim=0)[0]
|
||||
vertices_max = vertices.max(dim=0)[0]
|
||||
center = (vertices_min + vertices_max) / 2
|
||||
scale = 0.99999 / (vertices_max - vertices_min).max()
|
||||
for obj in dump['objects']:
|
||||
obj['vertices'] = (torch.from_numpy(obj['vertices']).float() - center) * scale
|
||||
obj['vertices'] = obj['vertices'].numpy()
|
||||
obj['mat_ids'][obj['mat_ids'] == -1] = len(dump['materials']) - 1
|
||||
assert np.all(obj['mat_ids'] >= 0), 'invalid mat_ids'
|
||||
assert np.all(obj['vertices'] >= -0.5) and np.all(obj['vertices'] <= 0.5), 'vertices out of range'
|
||||
|
||||
coord, attr = o_voxel.convert.blender_dump_to_volumetric_attr(dump, grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
mip_level_offset=0, verbose=False, timing=False)
|
||||
del attr['normal']
|
||||
del attr['emissive']
|
||||
o_voxel.io.write_vxz(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz'), coord, attr)
|
||||
pack[f'pbr_voxelized_{res}'] = True
|
||||
pack[f'num_pbr_voxels_{res}'] = len(coord)
|
||||
|
||||
return pack
|
||||
except Exception as e:
|
||||
print(f'Error voxelizing {sha256}: {e}')
|
||||
return {'sha256': sha256, 'error': str(e)}
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--pbr_dump_root', type=str, default=None,
|
||||
help='Directory to load mesh dumps')
|
||||
parser.add_argument('--pbr_voxel_root', type=str, default=None,
|
||||
help='Directory to save voxelized pbr attributes')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--resolution', type=str, default=1024)
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.resolution = sorted([int(x) for x in opt.resolution.split(',')], reverse=True)
|
||||
opt.pbr_dump_root = opt.pbr_dump_root or opt.root
|
||||
opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
|
||||
|
||||
for res in opt.resolution:
|
||||
os.makedirs(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'new_records'), exist_ok=True)
|
||||
|
||||
# get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
for res in opt.resolution:
|
||||
if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'metadata.csv')):
|
||||
pbr_voxel_metadata = pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}','metadata.csv')).set_index('sha256')
|
||||
pbr_voxel_metadata = pbr_voxel_metadata.rename(columns={'pbr_voxelized': f'pbr_voxelized_{res}', 'num_pbr_voxels': f'num_pbr_voxels_{res}'})
|
||||
metadata = metadata.combine_first(pbr_voxel_metadata)
|
||||
metadata = metadata.reset_index()
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['pbr_dumped'] == True]
|
||||
mask = np.zeros(len(metadata), dtype=bool)
|
||||
for res in opt.resolution:
|
||||
if f'pbr_voxelized_{res}' in metadata.columns:
|
||||
mask |= metadata[f'pbr_voxelized_{res}'] != True
|
||||
else:
|
||||
mask[:] = True
|
||||
break
|
||||
metadata = metadata[mask]
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
|
||||
# process objects
|
||||
func = partial(_pbr_voxelize, pbr_dump_root=opt.pbr_dump_root, root=opt.pbr_voxel_root)
|
||||
pbr_voxelized = dataset_utils.foreach_instance(metadata, None, func, max_workers=opt.max_workers, no_file=True, desc='Voxelizing')
|
||||
if 'error' in pbr_voxelized.columns:
|
||||
errors = pbr_voxelized[pbr_voxelized['error'].notna()]
|
||||
with open('errors.txt', 'w') as f:
|
||||
f.write('\n'.join(errors['sha256'].tolist()))
|
||||
for res in opt.resolution:
|
||||
if f'pbr_voxelized_{res}' in pbr_voxelized.columns:
|
||||
pbr_voxel_metadata = pbr_voxelized[pbr_voxelized[f'pbr_voxelized_{res}'] == True]
|
||||
if len(pbr_voxel_metadata) > 0:
|
||||
pbr_voxel_metadata = pbr_voxel_metadata[['sha256', f'pbr_voxelized_{res}', f'num_pbr_voxels_{res}']]
|
||||
pbr_voxel_metadata = pbr_voxel_metadata.rename(columns={f'pbr_voxelized_{res}': 'pbr_voxelized', f'num_pbr_voxels_{res}': 'num_pbr_voxels'})
|
||||
pbr_voxel_metadata.to_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False)
|
||||
@@ -0,0 +1,596 @@
|
||||
"""
|
||||
voxelize_pbr_view.py - Multi-view transform PBR voxelization
|
||||
Extends voxelize_pbr.py with scale and mesh rotation logic
|
||||
Based on dual_grid_view.py and test_ovoxel_pbr_transform.py implementation
|
||||
"""
|
||||
import os
|
||||
import copy
|
||||
import sys
|
||||
import importlib
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import pandas as pd
|
||||
import pickle
|
||||
import numpy as np
|
||||
import torch
|
||||
from easydict import EasyDict as edict
|
||||
from functools import partial
|
||||
import o_voxel
|
||||
from utils import get_new_camera_matrix, sphere_normalize_torch
|
||||
|
||||
|
||||
# ==================== PBR-specific transform functions ====================
|
||||
|
||||
def transform_vertices(vertices, frame):
|
||||
"""
|
||||
Apply multi-view transform to vertices based on camera transform matrix.
|
||||
|
||||
Args:
|
||||
vertices: torch.Tensor, shape [N, 3], vertex coordinates
|
||||
frame: dict containing transform_matrix
|
||||
|
||||
Returns:
|
||||
transformed_vertices: torch.Tensor, shape [N, 3]
|
||||
"""
|
||||
device = vertices.device
|
||||
c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
|
||||
|
||||
# Old and new camera matrices
|
||||
radius = c2w_orig[:3, 3].norm().item()
|
||||
c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
|
||||
dtype=torch.float32, device=device)
|
||||
w2c_orig = torch.inverse(c2w_orig)
|
||||
|
||||
# Initial and final axis alignment matrices
|
||||
R_init = torch.tensor([
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, -1.0, 0.0],
|
||||
[0.0, 1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
R_back = torch.tensor([
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0, 0.0],
|
||||
[0.0, -1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
R_ply = torch.tensor([
|
||||
[1.0, 0.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0, 0.0],
|
||||
[0.0, -1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 0.0, 1.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
T_cam = c2w_new @ w2c_orig @ R_ply
|
||||
T_final = R_back @ T_cam @ R_init
|
||||
|
||||
# Apply transform
|
||||
vertices = vertices.reshape(-1, 3)
|
||||
verts_h = torch.cat([vertices, torch.ones((vertices.shape[0], 1), dtype=torch.float32, device=device)], dim=1)
|
||||
verts_trans = (T_final @ verts_h.T).T[:, :3]
|
||||
|
||||
return verts_trans
|
||||
|
||||
|
||||
def transform_normals(normals, frame):
|
||||
"""
|
||||
Apply multi-view transform to normals (rotation only).
|
||||
Consistent with test_ovoxel_pbr_transform.py implementation.
|
||||
|
||||
Args:
|
||||
normals: torch.Tensor or np.ndarray, shape [N, 3] or [N, 3, 3]
|
||||
frame: dict containing transform_matrix
|
||||
|
||||
Returns:
|
||||
transformed_normals: np.ndarray (always returns numpy for dump compatibility)
|
||||
"""
|
||||
is_numpy = isinstance(normals, np.ndarray)
|
||||
if is_numpy:
|
||||
normals = torch.from_numpy(normals).float()
|
||||
|
||||
device = normals.device
|
||||
original_shape = normals.shape
|
||||
|
||||
# Flatten to [N, 3] for processing
|
||||
if len(original_shape) == 3:
|
||||
normals_flat = normals.reshape(-1, 3)
|
||||
else:
|
||||
normals_flat = normals
|
||||
|
||||
c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
|
||||
|
||||
# Old and new camera matrices
|
||||
radius = c2w_orig[:3, 3].norm().item()
|
||||
c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
|
||||
dtype=torch.float32, device=device)
|
||||
w2c_orig = torch.inverse(c2w_orig)
|
||||
|
||||
# Axis alignment matrices (rotation part only, 3x3)
|
||||
R_init = torch.tensor([
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, -1.0],
|
||||
[0.0, 1.0, 0.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
R_back = torch.tensor([
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
[0.0, -1.0, 0.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
R_ply = torch.tensor([
|
||||
[1.0, 0.0, 0.0],
|
||||
[0.0, 0.0, 1.0],
|
||||
[0.0, -1.0, 0.0]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
# Use rotation part only
|
||||
T_cam_rot = c2w_new[:3, :3] @ w2c_orig[:3, :3] @ R_ply
|
||||
T_final_rot = R_back @ T_cam_rot @ R_init
|
||||
|
||||
# Apply rotation transform
|
||||
normals_trans = torch.matmul(normals_flat, T_final_rot.T)
|
||||
|
||||
# Re-normalize
|
||||
normals_trans = torch.nn.functional.normalize(normals_trans, dim=-1)
|
||||
|
||||
# Restore original shape
|
||||
if len(original_shape) == 3:
|
||||
normals_trans = normals_trans.reshape(original_shape)
|
||||
|
||||
# Always return numpy array for dump compatibility
|
||||
return normals_trans.numpy()
|
||||
|
||||
|
||||
def prepare_pbr_dump(dump):
|
||||
"""
|
||||
Prepare PBR dump data for processing.
|
||||
Consistent with voxelize_pbr.py preprocessing.
|
||||
|
||||
Args:
|
||||
dump: raw PBR dump data
|
||||
|
||||
Returns:
|
||||
processed dump data
|
||||
"""
|
||||
dump = copy.deepcopy(dump)
|
||||
|
||||
# Fix dump alpha map
|
||||
for mat in dump['materials']:
|
||||
if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
|
||||
mat['alphaMode'] = 'BLEND'
|
||||
|
||||
# Append default material
|
||||
dump['materials'].append({
|
||||
"baseColorFactor": [0.8, 0.8, 0.8],
|
||||
"alphaFactor": 1.0,
|
||||
"metallicFactor": 0.0,
|
||||
"roughnessFactor": 0.5,
|
||||
"alphaMode": "OPAQUE",
|
||||
"alphaCutoff": 0.5,
|
||||
"baseColorTexture": None,
|
||||
"alphaTexture": None,
|
||||
"metallicTexture": None,
|
||||
"roughnessTexture": None,
|
||||
})
|
||||
|
||||
# Filter out empty objects
|
||||
dump['objects'] = [
|
||||
obj for obj in dump['objects']
|
||||
if obj['vertices'].size != 0 and obj['faces'].size != 0
|
||||
]
|
||||
|
||||
return dump
|
||||
|
||||
|
||||
def transform_pbr_dump(dump, frame):
|
||||
"""
|
||||
Apply multi-view transform to entire PBR dump data.
|
||||
|
||||
Processing flow (based on test_ovoxel_pbr_transform.py):
|
||||
1. Box normalize all vertices (scale only, no center shift)
|
||||
2. Sphere normalize
|
||||
3. Apply multi-view transform
|
||||
4. Normalize back to [-0.5, 0.5]^3
|
||||
|
||||
Note: All object vertices are processed together (not per-object) for consistency.
|
||||
|
||||
Args:
|
||||
dump: PBR dump data (already preprocessed via prepare_pbr_dump)
|
||||
frame: camera frame info
|
||||
|
||||
Returns:
|
||||
transformed_dump: transformed dump data
|
||||
total_scale: total scale from original mesh to final mesh
|
||||
"""
|
||||
transformed_dump = copy.deepcopy(dump)
|
||||
|
||||
# 1. Collect all vertices
|
||||
all_vertices_list = []
|
||||
vertex_counts = []
|
||||
for obj in transformed_dump['objects']:
|
||||
all_vertices_list.append(obj['vertices'])
|
||||
vertex_counts.append(len(obj['vertices']))
|
||||
|
||||
if len(all_vertices_list) == 0:
|
||||
return transformed_dump, 1.0
|
||||
|
||||
all_vertices = np.concatenate(all_vertices_list, axis=0)
|
||||
|
||||
# 2. Box normalize (scale only, no center shift, consistent with original rendering)
|
||||
vertices_min = all_vertices.min(axis=0)
|
||||
vertices_max = all_vertices.max(axis=0)
|
||||
box_scale_init = 0.99999 / (vertices_max - vertices_min).max()
|
||||
all_vertices_box_normalized = all_vertices * box_scale_init
|
||||
|
||||
all_vertices_tensor = torch.from_numpy(all_vertices_box_normalized).float()
|
||||
|
||||
# 3. Sphere normalize all vertices together
|
||||
all_vertices_sphere, sphere_center, sphere_radius = sphere_normalize_torch(all_vertices_tensor)
|
||||
|
||||
# 4. Multi-view transform
|
||||
all_transformed = transform_vertices(all_vertices_sphere, frame)
|
||||
|
||||
# 5. Normalize back to [-0.5, 0.5]^3 (all vertices together)
|
||||
abs_max = all_transformed.abs().max().item()
|
||||
box_scale_final = 0.49999 / abs_max
|
||||
all_transformed_normalized = all_transformed * box_scale_final
|
||||
|
||||
# Compute total scale (from original mesh to final normalized mesh)
|
||||
total_scale = box_scale_init * box_scale_final / sphere_radius.item()
|
||||
|
||||
# 6. Split back to individual objects
|
||||
start_idx = 0
|
||||
for i, obj in enumerate(transformed_dump['objects']):
|
||||
end_idx = start_idx + vertex_counts[i]
|
||||
obj['vertices'] = all_transformed_normalized[start_idx:end_idx].numpy()
|
||||
start_idx = end_idx
|
||||
|
||||
# Transform normals
|
||||
if obj['normals'] is not None and obj['normals'].size > 0:
|
||||
obj['normals'] = transform_normals(obj['normals'], frame)
|
||||
|
||||
# Fix mat_ids (replace -1 with default material index)
|
||||
obj['mat_ids'][obj['mat_ids'] == -1] = len(transformed_dump['materials']) - 1
|
||||
|
||||
# Validate range
|
||||
assert np.all(obj['mat_ids'] >= 0), 'invalid mat_ids'
|
||||
assert np.all(obj['vertices'] >= -0.5) and np.all(obj['vertices'] <= 0.5), 'vertices out of range'
|
||||
|
||||
return transformed_dump, total_scale
|
||||
|
||||
|
||||
def _pbr_voxelize_view(file, sha256, pbr_dump_root, transform_root, root, view_indices=None):
|
||||
"""
|
||||
Process multi-view PBR voxelization for a single sha256.
|
||||
|
||||
Args:
|
||||
file: local_path from metadata
|
||||
sha256: sha256 string
|
||||
pbr_dump_root: directory containing PBR dump files
|
||||
transform_root: directory containing transform json files
|
||||
root: output directory for PBR voxels
|
||||
view_indices: list of view indices to process, None for all views
|
||||
"""
|
||||
try:
|
||||
pack = {'sha256': sha256}
|
||||
dump = None
|
||||
|
||||
# Load transforms
|
||||
transform_path = os.path.join(transform_root, sha256, 'transforms.json')
|
||||
if not os.path.exists(transform_path):
|
||||
print(f'Transform file not found for {sha256}, skipping')
|
||||
return {'sha256': sha256, 'error': 'Transform file not found'}
|
||||
|
||||
with open(transform_path, 'r') as f:
|
||||
transforms_json = json.load(f)
|
||||
transform_mats = transforms_json['frames']
|
||||
|
||||
# Determine views to process
|
||||
if view_indices is None:
|
||||
view_indices = list(range(len(transform_mats)))
|
||||
else:
|
||||
view_indices = [i for i in view_indices if i < len(transform_mats)]
|
||||
|
||||
# Track processed and skipped counts
|
||||
processed_count = 0
|
||||
skipped_count = 0
|
||||
|
||||
for view_idx in view_indices:
|
||||
for res in opt.resolution:
|
||||
need_process = False
|
||||
|
||||
# Check if already processed
|
||||
# Path structure: pbr_voxels_view_fix_{res}/{sha256}/view{idx:02d}.vxz
|
||||
sha256_dir = os.path.join(root, f'pbr_voxels_view_fix_{res}', sha256)
|
||||
vxz_path = os.path.join(sha256_dir, f'view{view_idx:02d}.vxz')
|
||||
if os.path.exists(vxz_path):
|
||||
try:
|
||||
info = o_voxel.io.read_vxz_info(vxz_path)
|
||||
pack[f'pbr_voxelized_view_fix{view_idx:02d}_{res}'] = True
|
||||
pack[f'num_pbr_voxels_view_fix{view_idx:02d}_{res}'] = info['num_voxel']
|
||||
skipped_count += 1
|
||||
except Exception as e:
|
||||
print(f'Error reading {sha256}/view{view_idx:02d}.vxz: {e}, will reprocess')
|
||||
need_process = True
|
||||
else:
|
||||
need_process = True
|
||||
|
||||
# Process PBR dump
|
||||
if need_process:
|
||||
# Lazy load dump
|
||||
if dump is None:
|
||||
pbr_dump_file = os.path.join(pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle')
|
||||
if not os.path.exists(pbr_dump_file):
|
||||
print(f'PBR dump not found for {sha256}, skipping')
|
||||
return {'sha256': sha256, 'error': 'PBR dump not found'}
|
||||
|
||||
with open(pbr_dump_file, 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
|
||||
# Prepare dump data
|
||||
dump = prepare_pbr_dump(dump)
|
||||
|
||||
if len(dump['objects']) == 0:
|
||||
print(f'No valid objects in PBR dump for {sha256}, skipping')
|
||||
return {'sha256': sha256, 'error': 'No valid objects in PBR dump'}
|
||||
|
||||
# Get transform for current view
|
||||
frame = transform_mats[view_idx]
|
||||
|
||||
# Multi-view transform (deep copy from original dump each time)
|
||||
transformed_dump, total_scale = transform_pbr_dump(dump, frame)
|
||||
|
||||
# PBR voxelization
|
||||
coord, attr = o_voxel.convert.blender_dump_to_volumetric_attr(
|
||||
transformed_dump,
|
||||
grid_size=res,
|
||||
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
mip_level_offset=0,
|
||||
verbose=False,
|
||||
timing=False
|
||||
)
|
||||
|
||||
# Remove normal and emissive (consistent with voxelize_pbr.py)
|
||||
del attr['normal']
|
||||
del attr['emissive']
|
||||
|
||||
# Save .vxz file
|
||||
os.makedirs(sha256_dir, exist_ok=True)
|
||||
o_voxel.io.write_vxz(vxz_path, coord, attr)
|
||||
|
||||
# Save scale info
|
||||
scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
|
||||
scale_info = {
|
||||
'sha256': sha256,
|
||||
'view_idx': view_idx,
|
||||
'total_scale': float(total_scale),
|
||||
}
|
||||
with open(scale_path, 'w') as f:
|
||||
json.dump(scale_info, f, indent=2)
|
||||
|
||||
pack[f'pbr_voxelized_view_fix{view_idx:02d}_{res}'] = True
|
||||
pack[f'num_pbr_voxels_view_fix{view_idx:02d}_{res}'] = len(coord)
|
||||
pack[f'pbr_voxel_scale_view_fix{view_idx:02d}_{res}'] = float(total_scale)
|
||||
processed_count += 1
|
||||
|
||||
# Record processing stats
|
||||
pack['_processed_count'] = processed_count
|
||||
pack['_skipped_count'] = skipped_count
|
||||
|
||||
return pack
|
||||
|
||||
except Exception as e:
|
||||
print(f'Error processing {sha256}: {e}')
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return {'sha256': sha256, 'error': str(e)}
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--root', type=str, required=True,
|
||||
help='Directory to save the metadata')
|
||||
parser.add_argument('--pbr_dump_root', type=str, default=None,
|
||||
help='Directory to load PBR dumps')
|
||||
parser.add_argument('--transform_root', type=str, default=None,
|
||||
help='Directory to load transform json files (renders_cond)')
|
||||
parser.add_argument('--pbr_voxel_root', type=str, default=None,
|
||||
help='Directory to save voxelized PBR attributes')
|
||||
parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
|
||||
help='Filter objects with aesthetic score lower than this value')
|
||||
parser.add_argument('--instances', type=str, default=None,
|
||||
help='Instances to process')
|
||||
parser.add_argument('--view_indices', type=str, default=None,
|
||||
help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
|
||||
parser.add_argument('--skip_list', type=str, default=None,
|
||||
help='Path to a file containing sha256 hashes to skip (one per line). '
|
||||
'Supports format: "sha256" or "dataset/sha256"')
|
||||
parser.add_argument('--clean_pbr_dir', type=str, default=None,
|
||||
help='Path to clean_pbr directory. Will auto-load {dataset}_clean_output.txt as ok-list, '
|
||||
'only sha256 in ok-list will be processed')
|
||||
parser.add_argument('--clean_pbr_name', type=str, default=None,
|
||||
help='Dataset name prefix for clean_pbr file (e.g., ObjaverseXL_github). '
|
||||
'Defaults to sys.argv[1] if not specified')
|
||||
dataset_utils.add_args(parser)
|
||||
parser.add_argument('--resolution', type=str, default='1024')
|
||||
parser.add_argument('--rank', type=int, default=0)
|
||||
parser.add_argument('--world_size', type=int, default=1)
|
||||
parser.add_argument('--max_workers', type=int, default=0)
|
||||
opt = parser.parse_args(sys.argv[2:])
|
||||
opt = edict(vars(opt))
|
||||
opt.resolution = sorted([int(x) for x in opt.resolution.split(',')], reverse=True)
|
||||
opt.pbr_dump_root = opt.pbr_dump_root or opt.root
|
||||
opt.transform_root = opt.transform_root or os.path.join(opt.root, 'renders_cond')
|
||||
opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
|
||||
|
||||
# Parse view_indices
|
||||
view_indices = None
|
||||
if opt.view_indices is not None:
|
||||
view_indices = []
|
||||
for part in opt.view_indices.split(','):
|
||||
if '-' in part:
|
||||
start, end = map(int, part.split('-'))
|
||||
view_indices.extend(range(start, end + 1))
|
||||
else:
|
||||
view_indices.append(int(part))
|
||||
view_indices = list(set(view_indices)) # Deduplicate
|
||||
view_indices.sort()
|
||||
|
||||
# Load skip list (sha256 hashes to skip)
|
||||
skip_set = set()
|
||||
if opt.skip_list is not None and os.path.exists(opt.skip_list):
|
||||
with open(opt.skip_list, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith('#'):
|
||||
# Support "dataset/sha256" and plain "sha256" format, extract pure sha256
|
||||
skip_set.add(line.split('/')[-1])
|
||||
print(f'Loaded {len(skip_set)} items from skip_list: {opt.skip_list}')
|
||||
|
||||
# Load clean_pbr ok-list (only process approved sha256)
|
||||
ok_set = None
|
||||
if opt.clean_pbr_dir is not None:
|
||||
dataset_name = opt.clean_pbr_name or sys.argv[1]
|
||||
clean_file = os.path.join(opt.clean_pbr_dir, f'{dataset_name}_clean_output.txt')
|
||||
if os.path.exists(clean_file):
|
||||
ok_set = set()
|
||||
with open(clean_file, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith('#'):
|
||||
ok_set.add(line.split('/')[-1])
|
||||
print(f'Loaded {len(ok_set)} ok items from clean_pbr: {clean_file}')
|
||||
else:
|
||||
print(f'Warning: clean_pbr file not found: {clean_file}, proceeding without ok-list filter')
|
||||
|
||||
for res in opt.resolution:
|
||||
os.makedirs(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'new_records'), exist_ok=True)
|
||||
|
||||
# Get file list
|
||||
if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
|
||||
raise ValueError('metadata.csv not found')
|
||||
metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
|
||||
if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
|
||||
if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
|
||||
metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
|
||||
|
||||
# Check already processed pbr_voxels_view_fix
|
||||
for res in opt.resolution:
|
||||
if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'metadata.csv')):
|
||||
pbr_voxel_metadata = pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'metadata.csv')).set_index('sha256')
|
||||
metadata = metadata.combine_first(pbr_voxel_metadata)
|
||||
|
||||
metadata = metadata.reset_index()
|
||||
|
||||
if opt.instances is None:
|
||||
if opt.filter_low_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
|
||||
metadata = metadata[metadata['pbr_dumped'] == True]
|
||||
|
||||
# Filter out objects with all views already processed
|
||||
if view_indices is not None:
|
||||
for res in opt.resolution:
|
||||
# Check if each specified view is already processed
|
||||
all_views_done_col = f'_all_views_done_{res}'
|
||||
metadata[all_views_done_col] = True
|
||||
for view_idx in view_indices:
|
||||
col_name = f'pbr_voxelized_view_fix{view_idx:02d}_{res}'
|
||||
if col_name in metadata.columns:
|
||||
metadata[all_views_done_col] = metadata[all_views_done_col] & (metadata[col_name] == True)
|
||||
else:
|
||||
metadata[all_views_done_col] = False
|
||||
break
|
||||
|
||||
# Keep objects with at least one incomplete resolution
|
||||
any_incomplete = None
|
||||
for res in opt.resolution:
|
||||
all_views_done_col = f'_all_views_done_{res}'
|
||||
if all_views_done_col in metadata.columns:
|
||||
if any_incomplete is None:
|
||||
any_incomplete = ~metadata[all_views_done_col]
|
||||
else:
|
||||
any_incomplete = any_incomplete | ~metadata[all_views_done_col]
|
||||
|
||||
if any_incomplete is not None:
|
||||
before_filter = len(metadata)
|
||||
metadata = metadata[any_incomplete]
|
||||
print(f'Filtered out {before_filter - len(metadata)} already completed objects')
|
||||
else:
|
||||
if os.path.exists(opt.instances):
|
||||
with open(opt.instances, 'r') as f:
|
||||
instances = f.read().splitlines()
|
||||
else:
|
||||
instances = opt.instances.split(',')
|
||||
metadata = metadata[metadata['sha256'].isin(instances)]
|
||||
|
||||
# Apply skip_list filter (exclude specified sha256)
|
||||
if skip_set:
|
||||
before_skip = len(metadata)
|
||||
metadata = metadata[~metadata['sha256'].isin(skip_set)]
|
||||
print(f'Skip list: filtered out {before_skip - len(metadata)} objects, {len(metadata)} remaining')
|
||||
|
||||
# Apply clean_pbr ok-list filter (only keep approved sha256)
|
||||
if ok_set is not None:
|
||||
before_ok = len(metadata)
|
||||
metadata = metadata[metadata['sha256'].isin(ok_set)]
|
||||
print(f'Ok list: kept {len(metadata)} objects out of {before_ok} (filtered {before_ok - len(metadata)})')
|
||||
|
||||
metadata = metadata.sample(frac=1, random_state=444).reset_index(drop=True)
|
||||
start = len(metadata) * opt.rank // opt.world_size
|
||||
end = len(metadata) * (opt.rank + 1) // opt.world_size
|
||||
metadata = metadata[start:end]
|
||||
|
||||
print(f'Processing {len(metadata)} objects...')
|
||||
if view_indices:
|
||||
print(f'View indices to process: {view_indices}')
|
||||
else:
|
||||
print('Processing all available views')
|
||||
|
||||
# Process objects
|
||||
func = partial(_pbr_voxelize_view,
|
||||
pbr_dump_root=opt.pbr_dump_root,
|
||||
transform_root=opt.transform_root,
|
||||
root=opt.pbr_voxel_root,
|
||||
view_indices=view_indices)
|
||||
pbr_voxelized = dataset_utils.foreach_instance(metadata, opt.root, func, max_workers=opt.max_workers, desc='Voxelizing PBR views')
|
||||
|
||||
# Processing summary
|
||||
total_processed = pbr_voxelized['_processed_count'].sum() if '_processed_count' in pbr_voxelized.columns else 0
|
||||
total_skipped = pbr_voxelized['_skipped_count'].sum() if '_skipped_count' in pbr_voxelized.columns else 0
|
||||
print(f'\n========== Processing Summary ==========')
|
||||
print(f'Total processed (new): {int(total_processed)}')
|
||||
print(f'Total skipped (existing): {int(total_skipped)}')
|
||||
print(f'Total items: {int(total_processed + total_skipped)}')
|
||||
print(f'=========================================\n')
|
||||
|
||||
if 'error' in pbr_voxelized.columns:
|
||||
errors = pbr_voxelized[pbr_voxelized['error'].notna()]
|
||||
if len(errors) > 0:
|
||||
with open('errors_pbr_view.txt', 'w') as f:
|
||||
f.write('\n'.join(errors['sha256'].tolist()))
|
||||
print(f'Errors written to errors_pbr_view.txt ({len(errors)} errors)')
|
||||
|
||||
# Save metadata
|
||||
for res in opt.resolution:
|
||||
# Collect all view-related columns
|
||||
view_cols = [col for col in pbr_voxelized.columns if f'pbr_voxelized_view_fix' in col and f'_{res}' in col]
|
||||
if view_cols:
|
||||
# Save metadata for each view
|
||||
pbr_voxel_metadata = pbr_voxelized[pbr_voxelized[view_cols].any(axis=1)]
|
||||
if len(pbr_voxel_metadata) > 0:
|
||||
# Save simplified metadata
|
||||
cols_to_save = ['sha256'] + [col for col in pbr_voxelized.columns if f'_{res}' in col]
|
||||
cols_to_save = [col for col in cols_to_save if col in pbr_voxelized.columns]
|
||||
pbr_voxel_metadata[cols_to_save].to_csv(
|
||||
os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'new_records', f'part_{opt.rank}.csv'),
|
||||
index=False
|
||||
)
|
||||
|
||||
print('Done!')
|
||||
+1493
File diff suppressed because it is too large
Load Diff
+312
@@ -0,0 +1,312 @@
|
||||
import os
|
||||
import argparse
|
||||
import math
|
||||
import time
|
||||
import torch
|
||||
import numpy as np
|
||||
import cv2
|
||||
from PIL import Image
|
||||
|
||||
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
||||
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
|
||||
os.environ.setdefault("ATTN_BACKEND", "flash_attn")
|
||||
os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
|
||||
os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
|
||||
|
||||
from pixal3d.pipelines import Pixal3DImageTo3DPipeline
|
||||
import o_voxel
|
||||
|
||||
# ============================================================================
|
||||
# Constants & Defaults
|
||||
# ============================================================================
|
||||
|
||||
MOGE_MODEL_NAME = "Ruicheng/moge-2-vitl"
|
||||
MODEL_PATH = "TencentARC/Pixal3D"
|
||||
|
||||
IMAGE_COND_CONFIGS = {
|
||||
"ss": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 16,
|
||||
},
|
||||
"shape_512": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 512,
|
||||
"grid_resolution": 32,
|
||||
"use_naf_upsample": True,
|
||||
"naf_target_size": 512,
|
||||
},
|
||||
"shape_1024": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 1024,
|
||||
"grid_resolution": 64,
|
||||
"use_naf_upsample": True,
|
||||
"naf_target_size": 512,
|
||||
},
|
||||
"tex_1024": {
|
||||
"model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
|
||||
"image_size": 1024,
|
||||
"grid_resolution": 64,
|
||||
"use_naf_upsample": True,
|
||||
"naf_target_size": 1024,
|
||||
},
|
||||
}
|
||||
|
||||
# ============================================================================
|
||||
# Model Loading
|
||||
# ============================================================================
|
||||
|
||||
def build_image_cond_model(config: dict):
|
||||
from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor
|
||||
model = DinoV3ProjFeatureExtractor(**config)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
|
||||
def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME):
|
||||
from moge.model.v2 import MoGeModel
|
||||
moge_model = MoGeModel.from_pretrained(model_name)
|
||||
moge_model = moge_model.to(device)
|
||||
moge_model.eval()
|
||||
return moge_model
|
||||
|
||||
|
||||
def init_pipeline(model_path=MODEL_PATH, device="cuda", low_vram=False):
|
||||
print(f"[Pipeline] Loading from {model_path}...")
|
||||
pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
|
||||
|
||||
print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
|
||||
pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
|
||||
pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
|
||||
pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
|
||||
pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
|
||||
|
||||
if low_vram:
|
||||
# Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage.
|
||||
# Peak VRAM = one flow model + one DinoV3, not all ~18 GB at once.
|
||||
print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...")
|
||||
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
||||
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
||||
m = getattr(pipeline, attr, None)
|
||||
if m is not None and getattr(m, 'use_naf_upsample', False):
|
||||
m._load_naf()
|
||||
pipeline._device = torch.device(device)
|
||||
pipeline.low_vram = True
|
||||
print("[Pipeline] Low-VRAM mode enabled.")
|
||||
else:
|
||||
# Standard mode: all models loaded to GPU at once (faster, needs more VRAM).
|
||||
pipeline.low_vram = False
|
||||
pipeline.cuda()
|
||||
pipeline.image_cond_model_ss.cuda()
|
||||
pipeline.image_cond_model_shape_512.cuda()
|
||||
pipeline.image_cond_model_shape_1024.cuda()
|
||||
pipeline.image_cond_model_tex_1024.cuda()
|
||||
print("[NAF] Pre-loading NAF upsampler model...")
|
||||
for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
|
||||
'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
|
||||
m = getattr(pipeline, attr, None)
|
||||
if m is not None and getattr(m, 'use_naf_upsample', False):
|
||||
m._load_naf()
|
||||
print("[Pipeline] Standard mode (all models on GPU).")
|
||||
|
||||
return pipeline
|
||||
|
||||
# ============================================================================
|
||||
# Camera Estimation
|
||||
# ============================================================================
|
||||
|
||||
def compute_f_pixels(camera_angle_x: float, resolution: int) -> float:
|
||||
focal_length = 16.0 / torch.tan(torch.tensor(camera_angle_x / 2.0))
|
||||
f_pixels = focal_length * resolution / 32.0
|
||||
return float(f_pixels.item())
|
||||
|
||||
|
||||
def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, image_resolution):
|
||||
rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
|
||||
gp = grid_point.to(torch.float32) @ rotation_matrix.T
|
||||
gp = gp / mesh_scale / 2
|
||||
xw, yw, zw = gp[0].item(), gp[1].item(), gp[2].item()
|
||||
xt, yt = float(target_point[0].item()), float(target_point[1].item())
|
||||
f_pixels = compute_f_pixels(camera_angle_x, image_resolution)
|
||||
x_ndc = xt - image_resolution / 2.0
|
||||
y_ndc = -(yt - image_resolution / 2.0)
|
||||
distance_x = f_pixels * xw / x_ndc - yw
|
||||
return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)}
|
||||
|
||||
|
||||
def get_camera_params_wild_moge(image_path, moge_model, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512):
|
||||
pil_image = Image.open(image_path).convert("RGB")
|
||||
width, height = pil_image.size
|
||||
image_np = np.array(pil_image).astype(np.float32) / 255.0
|
||||
image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).to(device)
|
||||
with torch.no_grad():
|
||||
output = moge_model.infer(image_tensor)
|
||||
intrinsics = output["intrinsics"].squeeze().cpu().numpy()
|
||||
fx_normalized = intrinsics[0, 0]
|
||||
fx = fx_normalized * width
|
||||
camera_angle_x = 2 * math.atan(width / (2 * fx))
|
||||
|
||||
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
||||
distance = distance_from_fov(
|
||||
camera_angle_x, grid_point,
|
||||
torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
|
||||
mesh_scale, image_resolution
|
||||
)["distance_from_x"]
|
||||
return {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
|
||||
|
||||
# ============================================================================
|
||||
# Main Inference
|
||||
# ============================================================================
|
||||
|
||||
def run_inference(
|
||||
image_path: str,
|
||||
output_path: str,
|
||||
seed: int = 42,
|
||||
ss_guidance_strength: float = 7.5,
|
||||
ss_guidance_rescale: float = 0.7,
|
||||
ss_sampling_steps: int = 12,
|
||||
ss_rescale_t: float = 5.0,
|
||||
shape_slat_guidance_strength: float = 7.5,
|
||||
shape_slat_guidance_rescale: float = 0.5,
|
||||
shape_slat_sampling_steps: int = 12,
|
||||
shape_slat_rescale_t: float = 3.0,
|
||||
tex_slat_guidance_strength: float = 1.0,
|
||||
tex_slat_guidance_rescale: float = 0.0,
|
||||
tex_slat_sampling_steps: int = 12,
|
||||
tex_slat_rescale_t: float = 3.0,
|
||||
mesh_scale: float = 1.0,
|
||||
extend_pixel: int = 0,
|
||||
image_resolution: int = 512,
|
||||
max_num_tokens: int = 49152,
|
||||
model_path: str = MODEL_PATH,
|
||||
manual_fov: float = -1.0,
|
||||
low_vram: bool = False,
|
||||
resolution: int = -1,
|
||||
):
|
||||
# Load models
|
||||
pipeline = init_pipeline(model_path, low_vram=low_vram)
|
||||
|
||||
# Preprocess image first — rembg loads to GPU for this call, then offloads.
|
||||
# MoGe is loaded afterwards so both never occupy VRAM at the same time.
|
||||
print(f"[Inference] Processing image: {image_path}")
|
||||
img = Image.open(image_path)
|
||||
image_preprocessed = pipeline.preprocess_image(img)
|
||||
|
||||
# Save preprocessed image for MoGe
|
||||
tmp_path = os.path.join(os.path.dirname(os.path.abspath(output_path)), f"_tmp_preprocessed_{int(time.time()*1000)}.png")
|
||||
image_preprocessed.save(tmp_path)
|
||||
|
||||
# Camera estimation
|
||||
if manual_fov > 0:
|
||||
# Use manually specified FOV (in radians)
|
||||
camera_angle_x = float(manual_fov)
|
||||
grid_point = torch.tensor([-1.0, 0.0, 0.0])
|
||||
distance = distance_from_fov(
|
||||
camera_angle_x, grid_point,
|
||||
torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
|
||||
mesh_scale, image_resolution
|
||||
)["distance_from_x"]
|
||||
camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
|
||||
print(f"[Inference] Using manual FOV: {math.degrees(manual_fov):.2f}° ({manual_fov:.4f} rad), distance={distance:.4f}")
|
||||
else:
|
||||
print("[MoGe-2] Loading model for camera estimation...")
|
||||
moge_model = load_moge_model(device="cuda")
|
||||
print("[Inference] Estimating camera parameters...")
|
||||
camera_params = get_camera_params_wild_moge(
|
||||
tmp_path, moge_model, device="cuda",
|
||||
mesh_scale=mesh_scale, extend_pixel=extend_pixel,
|
||||
image_resolution=image_resolution,
|
||||
)
|
||||
print(f" camera_angle_x={camera_params['camera_angle_x']:.4f}, distance={camera_params['distance']:.4f}")
|
||||
# MoGe is only needed for camera estimation; free its VRAM for inference.
|
||||
moge_model.cpu()
|
||||
del moge_model
|
||||
torch.cuda.empty_cache()
|
||||
os.remove(tmp_path)
|
||||
|
||||
# Run pipeline
|
||||
print("[Inference] Running 3D generation pipeline...")
|
||||
torch.manual_seed(seed)
|
||||
|
||||
ss_sampler_override = {
|
||||
"steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength,
|
||||
"guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t,
|
||||
}
|
||||
shape_sampler_override = {
|
||||
"steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength,
|
||||
"guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t,
|
||||
}
|
||||
tex_sampler_override = {
|
||||
"steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength,
|
||||
"guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t,
|
||||
}
|
||||
|
||||
pipeline_type = f"{resolution if resolution > 0 else (1024 if low_vram else 1536)}_cascade"
|
||||
print(f"[Inference] Using pipeline_type={pipeline_type}")
|
||||
mesh_list, (shape_slat, tex_slat, res) = pipeline.run(
|
||||
image_preprocessed,
|
||||
camera_params=camera_params,
|
||||
seed=seed,
|
||||
sparse_structure_sampler_params=ss_sampler_override,
|
||||
shape_slat_sampler_params=shape_sampler_override,
|
||||
tex_slat_sampler_params=tex_sampler_override,
|
||||
preprocess_image=False,
|
||||
return_latent=True,
|
||||
pipeline_type=pipeline_type,
|
||||
max_num_tokens=max_num_tokens,
|
||||
)
|
||||
|
||||
mesh = mesh_list[0]
|
||||
|
||||
# Extract GLB
|
||||
print("[Inference] Extracting GLB...")
|
||||
glb = o_voxel.postprocess.to_glb(
|
||||
vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs,
|
||||
coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout,
|
||||
grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
decimation_target=1000000, texture_size=4096,
|
||||
remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
|
||||
)
|
||||
|
||||
# Apply rotation
|
||||
rot = np.array([
|
||||
[-1, 0, 0, 0],
|
||||
[ 0, 0, -1, 0],
|
||||
[ 0, -1, 0, 0],
|
||||
[ 0, 0, 0, 1],
|
||||
], dtype=np.float64)
|
||||
glb.apply_transform(rot)
|
||||
|
||||
# Export
|
||||
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
|
||||
glb.export(output_path, extension_webp=True)
|
||||
print(f"[Done] GLB saved to: {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Pixal3D Inference: Image to GLB")
|
||||
parser.add_argument("--image", type=str, required=True, help="Path to input image")
|
||||
parser.add_argument("--output", type=str, default="./output.glb", help="Output GLB file path")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
parser.add_argument("--fov", type=float, default=-1.0,
|
||||
help="Manual camera FOV in radians (e.g. 0.2). "
|
||||
"If not set, FOV is auto-estimated via MoGe-2. "
|
||||
"Try 0.2 rad if you notice distortion.")
|
||||
parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Model path or HuggingFace repo")
|
||||
parser.add_argument("--low_vram", action="store_true",
|
||||
help="Enable low-VRAM mode: models stay on CPU and are loaded to GPU on-demand per stage. "
|
||||
"Reduces peak VRAM from ~18GB to ~10-12GB at the cost of slower inference.")
|
||||
parser.add_argument("--resolution", type=int, default=-1,
|
||||
help="Pipeline resolution (1024 or 1536). Default: 1024 if --low_vram, else 1536.")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
run_inference(
|
||||
image_path=args.image,
|
||||
output_path=args.output,
|
||||
seed=args.seed,
|
||||
manual_fov=args.fov,
|
||||
model_path=args.model_path,
|
||||
low_vram=args.low_vram,
|
||||
resolution=args.resolution,
|
||||
)
|
||||
@@ -0,0 +1,6 @@
|
||||
from . import models
|
||||
from . import modules
|
||||
from . import pipelines
|
||||
from . import renderers
|
||||
from . import representations
|
||||
from . import utils
|
||||
@@ -0,0 +1,52 @@
|
||||
import importlib
|
||||
|
||||
__attributes = {
|
||||
'FlexiDualGridDataset': 'flexi_dual_grid',
|
||||
'SparseVoxelPbrDataset':'sparse_voxel_pbr',
|
||||
|
||||
'SparseStructureLatent': 'sparse_structure_latent',
|
||||
'TextConditionedSparseStructureLatent': 'sparse_structure_latent',
|
||||
'ImageConditionedSparseStructureLatent': 'sparse_structure_latent',
|
||||
'SparseStructureLatentView': 'sparse_structure_latent',
|
||||
'ViewImageConditionedSparseStructureLatentView': 'sparse_structure_latent',
|
||||
|
||||
'SLat': 'structured_latent',
|
||||
'ImageConditionedSLat': 'structured_latent',
|
||||
'SLatShape': 'structured_latent_shape',
|
||||
'ImageConditionedSLatShape': 'structured_latent_shape',
|
||||
'SLatShapeView': 'structured_latent_shape',
|
||||
'ViewImageConditionedSLatShapeView': 'structured_latent_shape',
|
||||
'SLatPbr': 'structured_latent_svpbr',
|
||||
'ImageConditionedSLatPbr': 'structured_latent_svpbr',
|
||||
'SLatPbrView': 'structured_latent_svpbr',
|
||||
'ViewImageConditionedSLatPbrView': 'structured_latent_svpbr',
|
||||
}
|
||||
|
||||
__submodules = []
|
||||
|
||||
__all__ = list(__attributes.keys()) + __submodules
|
||||
|
||||
def __getattr__(name):
|
||||
if name not in globals():
|
||||
if name in __attributes:
|
||||
module_name = __attributes[name]
|
||||
module = importlib.import_module(f".{module_name}", __name__)
|
||||
globals()[name] = getattr(module, name)
|
||||
elif name in __submodules:
|
||||
module = importlib.import_module(f".{name}", __name__)
|
||||
globals()[name] = module
|
||||
else:
|
||||
raise AttributeError(f"module {__name__} has no attribute {name}")
|
||||
return globals()[name]
|
||||
|
||||
|
||||
# For Pylance
|
||||
if __name__ == '__main__':
|
||||
from .flexi_dual_grid import FlexiDualGridDataset
|
||||
from .sparse_voxel_pbr import SparseVoxelPbrDataset
|
||||
|
||||
from .sparse_structure_latent import SparseStructureLatent, ImageConditionedSparseStructureLatent
|
||||
from .structured_latent import SLat, ImageConditionedSLat
|
||||
from .structured_latent_shape import SLatShape, ImageConditionedSLatShape
|
||||
from .structured_latent_svpbr import SLatPbr, ImageConditionedSLatPbr
|
||||
|
||||
@@ -0,0 +1,349 @@
|
||||
from typing import *
|
||||
import json
|
||||
from abc import abstractmethod
|
||||
import os
|
||||
import json
|
||||
import torch
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class StandardDatasetBase(Dataset):
|
||||
"""
|
||||
Base class for standard datasets.
|
||||
|
||||
Args:
|
||||
roots (str): paths to the dataset
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip (one per line)
|
||||
Format: "dataset/sha256" (e.g., "ABO/6a79dbb5...")
|
||||
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
||||
(e.g., ["texverse"] for datasets without aesthetic_score)
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[List[str]] = None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# Datasets to skip aesthetic score check
|
||||
self.skip_aesthetic_score_datasets = set(skip_aesthetic_score_datasets or [])
|
||||
|
||||
# Load skip list if provided
|
||||
self.skip_set = set()
|
||||
if skip_list is not None and os.path.exists(skip_list):
|
||||
with open(skip_list, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line and not line.startswith('#'):
|
||||
self.skip_set.add(line)
|
||||
print(f'Loaded {len(self.skip_set)} items from skip_list: {skip_list}')
|
||||
|
||||
try:
|
||||
self.roots = json.loads(roots)
|
||||
root_type = 'obj'
|
||||
except:
|
||||
self.roots = roots.split(',')
|
||||
root_type = 'list'
|
||||
self.instances = []
|
||||
self.metadata = pd.DataFrame()
|
||||
|
||||
self._stats = {}
|
||||
if root_type == 'obj':
|
||||
for key, root in self.roots.items():
|
||||
self._stats[key] = {}
|
||||
metadata = pd.DataFrame(columns=['sha256']).set_index('sha256')
|
||||
|
||||
# Only merge key fields from ss_latent and render_cond
|
||||
# Exclude base, because cond_rendered=False in base/metadata.csv would incorrectly overwrite real values
|
||||
for sub_key, r in root.items():
|
||||
if sub_key == 'base':
|
||||
continue # Skip base directory
|
||||
metadata_file = os.path.join(r, 'metadata.csv')
|
||||
if os.path.exists(metadata_file):
|
||||
metadata = metadata.combine_first(pd.read_csv(metadata_file).set_index('sha256'))
|
||||
|
||||
# Read aesthetic_score separately from base (avoid reading other potentially conflicting columns)
|
||||
if 'base' in root:
|
||||
base_metadata_file = os.path.join(root['base'], 'metadata.csv')
|
||||
if os.path.exists(base_metadata_file):
|
||||
base_df = pd.read_csv(base_metadata_file).set_index('sha256')
|
||||
if 'aesthetic_score' in base_df.columns and 'aesthetic_score' not in metadata.columns:
|
||||
metadata['aesthetic_score'] = base_df['aesthetic_score']
|
||||
|
||||
self._stats[key]['Total'] = len(metadata)
|
||||
metadata, stats = self.filter_metadata(metadata, dataset_name=key)
|
||||
self._stats[key].update(stats)
|
||||
|
||||
# Filter out items in skip_list
|
||||
skipped_count = 0
|
||||
for sha256 in metadata.index.values:
|
||||
skip_key = f'{key}/{sha256}'
|
||||
if skip_key in self.skip_set:
|
||||
skipped_count += 1
|
||||
else:
|
||||
self.instances.append((root, sha256, key))
|
||||
if skipped_count > 0:
|
||||
self._stats[key]['Skipped (skip_list)'] = skipped_count
|
||||
self._stats[key]['After skip_list'] = len(metadata) - skipped_count
|
||||
|
||||
self.metadata = pd.concat([self.metadata, metadata])
|
||||
else:
|
||||
for root in self.roots:
|
||||
key = os.path.basename(root)
|
||||
self._stats[key] = {}
|
||||
metadata = pd.read_csv(os.path.join(root, 'metadata.csv'))
|
||||
self._stats[key]['Total'] = len(metadata)
|
||||
metadata, stats = self.filter_metadata(metadata, dataset_name=key)
|
||||
self._stats[key].update(stats)
|
||||
|
||||
# Filter out items in skip_list
|
||||
skipped_count = 0
|
||||
for sha256 in metadata['sha256'].values:
|
||||
skip_key = f'{key}/{sha256}'
|
||||
if skip_key in self.skip_set:
|
||||
skipped_count += 1
|
||||
else:
|
||||
self.instances.append((root, sha256, key))
|
||||
if skipped_count > 0:
|
||||
self._stats[key]['Skipped (skip_list)'] = skipped_count
|
||||
self._stats[key]['After skip_list'] = len(metadata) - skipped_count
|
||||
metadata.set_index('sha256', inplace=True)
|
||||
self.metadata = pd.concat([self.metadata, metadata])
|
||||
|
||||
@abstractmethod
|
||||
def filter_metadata(self, metadata: pd.DataFrame, dataset_name: str = None) -> Tuple[pd.DataFrame, Dict[str, int]]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_instance(self, root, instance: str) -> Dict[str, Any]:
|
||||
pass
|
||||
|
||||
def __len__(self):
|
||||
return len(self.instances)
|
||||
|
||||
def __getitem__(self, index) -> Dict[str, Any]:
|
||||
try:
|
||||
root, instance, dataset_name = self.instances[index]
|
||||
pack = self.get_instance(root, instance)
|
||||
pack['_dataset_name'] = dataset_name
|
||||
pack['_sha256'] = instance
|
||||
return pack
|
||||
except Exception as e:
|
||||
print(f'Error loading {self.instances[index][1]}: {e}')
|
||||
return self.__getitem__(np.random.randint(0, len(self)))
|
||||
|
||||
def __str__(self):
|
||||
lines = []
|
||||
lines.append(self.__class__.__name__)
|
||||
lines.append(f' - Total instances: {len(self)}')
|
||||
lines.append(f' - Sources:')
|
||||
for key, stats in self._stats.items():
|
||||
lines.append(f' - {key}:')
|
||||
for k, v in stats.items():
|
||||
lines.append(f' - {k}: {v}')
|
||||
return '\n'.join(lines)
|
||||
|
||||
|
||||
class ImageConditionedMixin:
|
||||
def __init__(self, roots, *, image_size=518, **kwargs):
|
||||
self.image_size = image_size
|
||||
super().__init__(roots, **kwargs)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
|
||||
metadata = metadata[metadata['cond_rendered'].notna()]
|
||||
stats['Cond rendered'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
pack = super().get_instance(root, instance)
|
||||
|
||||
image_root = os.path.join(root['render_cond'], instance)
|
||||
with open(os.path.join(image_root, 'transforms.json')) as f:
|
||||
metadata = json.load(f)
|
||||
n_views = len(metadata['frames'])
|
||||
view = np.random.randint(n_views)
|
||||
metadata = metadata['frames'][view]
|
||||
|
||||
image_path = os.path.join(image_root, metadata['file_path'])
|
||||
image = Image.open(image_path)
|
||||
|
||||
alpha = np.array(image.getchannel(3))
|
||||
bbox = np.array(alpha).nonzero()
|
||||
bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
|
||||
center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
|
||||
hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
|
||||
aug_hsize = hsize
|
||||
aug_center_offset = [0, 0]
|
||||
aug_center = [center[0] + aug_center_offset[0], center[1] + aug_center_offset[1]]
|
||||
aug_bbox = [int(aug_center[0] - aug_hsize), int(aug_center[1] - aug_hsize), int(aug_center[0] + aug_hsize), int(aug_center[1] + aug_hsize)]
|
||||
image = image.crop(aug_bbox)
|
||||
|
||||
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
||||
alpha = image.getchannel(3)
|
||||
image = image.convert('RGB')
|
||||
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
|
||||
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
||||
image = image * alpha.unsqueeze(0)
|
||||
pack['cond'] = image
|
||||
|
||||
return pack
|
||||
|
||||
|
||||
class ViewImageConditionedMixin:
|
||||
"""
|
||||
Mixin for view-based image-conditioned datasets.
|
||||
|
||||
This mixin is designed for datasets where ss_latent is stored per-view (view{XX}.npz),
|
||||
and needs to load the corresponding view image and scale from view{XX}_scale.json.
|
||||
|
||||
Args:
|
||||
image_size: Target image size
|
||||
load_camera_info: Whether to load camera information for view-aligned conditioning
|
||||
"""
|
||||
def __init__(self, roots, *, image_size=518, load_camera_info=False, **kwargs):
|
||||
self.image_size = image_size
|
||||
# self.load_camera_info = load_camera_info
|
||||
super().__init__(roots, **kwargs)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
|
||||
metadata = metadata[metadata['cond_rendered'].notna()]
|
||||
stats['Cond rendered'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
"""
|
||||
Get instance with view-aligned image and camera info.
|
||||
|
||||
Expects parent class to set:
|
||||
- pack['x_0']: the latent tensor
|
||||
- self._current_view_idx: the selected view index
|
||||
- self._current_latent_dir: the latent directory path
|
||||
"""
|
||||
pack = super().get_instance(root, instance)
|
||||
|
||||
# Get view_idx from parent class (set by SparseStructureLatentView)
|
||||
if not hasattr(self, '_current_view_idx'):
|
||||
raise RuntimeError("Parent class must set '_current_view_idx' before calling ViewImageConditionedMixin.get_instance")
|
||||
if not hasattr(self, '_current_latent_dir'):
|
||||
raise RuntimeError("Parent class must set '_current_latent_dir' before calling ViewImageConditionedMixin.get_instance")
|
||||
view_idx = self._current_view_idx
|
||||
latent_dir = self._current_latent_dir
|
||||
|
||||
# Load image metadata
|
||||
image_root = os.path.join(root['render_cond'], instance)
|
||||
with open(os.path.join(image_root, 'transforms.json')) as f:
|
||||
metadata = json.load(f)
|
||||
|
||||
# Load corresponding image for this view
|
||||
frame_metadata = metadata['frames'][view_idx]
|
||||
image_path = os.path.join(image_root, frame_metadata['file_path'])
|
||||
image = Image.open(image_path)
|
||||
|
||||
image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
||||
alpha = image.getchannel(3)
|
||||
image = image.convert('RGB')
|
||||
image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
|
||||
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
||||
image = image * alpha.unsqueeze(0)
|
||||
pack['cond'] = image
|
||||
|
||||
# Load camera info if requested
|
||||
|
||||
# camera_angle_x: check frame first, then root metadata
|
||||
if 'camera_angle_x' in frame_metadata:
|
||||
camera_angle_x = float(frame_metadata['camera_angle_x'])
|
||||
elif 'camera_angle_x' in metadata:
|
||||
camera_angle_x = float(metadata['camera_angle_x'])
|
||||
else:
|
||||
raise KeyError(f"'camera_angle_x' not found in transforms.json for {instance}")
|
||||
pack['camera_angle_x'] = torch.tensor(camera_angle_x, dtype=torch.float32)
|
||||
|
||||
# transform_matrix
|
||||
if 'transform_matrix' not in frame_metadata:
|
||||
raise KeyError(f"'transform_matrix' not found in frame {view_idx} for {instance}")
|
||||
transform_matrix = torch.tensor(frame_metadata['transform_matrix'], dtype=torch.float32)
|
||||
distance = torch.norm(transform_matrix[:3, 3]).item()
|
||||
|
||||
pack['camera_distance'] = torch.tensor(distance, dtype=torch.float32)
|
||||
# NOTE: Do NOT pass transform_matrix to ProjGrid.
|
||||
# shape_latent space objects are already rotated to front-view by transform_mesh,
|
||||
# so ProjGrid should use the default front_view_transform_matrix + distance.
|
||||
# pack['transform_matrix'] = transform_matrix
|
||||
|
||||
# Load mesh_scale from ss_latent directory's view{XX}_scale.json
|
||||
scale_json_path = os.path.join(latent_dir, f'view{view_idx:02d}_scale.json')
|
||||
if not os.path.exists(scale_json_path):
|
||||
raise FileNotFoundError(f"Scale file not found: {scale_json_path}")
|
||||
with open(scale_json_path) as f:
|
||||
scale_data = json.load(f)
|
||||
if 'total_scale' not in scale_data:
|
||||
raise KeyError(f"'total_scale' not found in {scale_json_path}")
|
||||
pack['mesh_scale'] = torch.tensor(float(scale_data['total_scale']), dtype=torch.float32)
|
||||
|
||||
return pack
|
||||
|
||||
|
||||
class MultiImageConditionedMixin:
|
||||
def __init__(self, roots, *, image_size=518, max_image_cond_view = 4, **kwargs):
|
||||
self.image_size = image_size
|
||||
self.max_image_cond_view = max_image_cond_view
|
||||
super().__init__(roots, **kwargs)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
|
||||
metadata = metadata[metadata['cond_rendered'].notna()]
|
||||
stats['Cond rendered'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
pack = super().get_instance(root, instance)
|
||||
|
||||
image_root = os.path.join(root['render_cond'], instance)
|
||||
with open(os.path.join(image_root, 'transforms.json')) as f:
|
||||
metadata = json.load(f)
|
||||
|
||||
n_views = len(metadata['frames'])
|
||||
n_sample_views = np.random.randint(1, self.max_image_cond_view+1)
|
||||
|
||||
assert n_views >= n_sample_views, f'Not enough views to sample {n_sample_views} unique images.'
|
||||
|
||||
sampled_views = np.random.choice(n_views, size=n_sample_views, replace=False)
|
||||
|
||||
cond_images = []
|
||||
for v in sampled_views:
|
||||
frame_info = metadata['frames'][v]
|
||||
image_path = os.path.join(image_root, frame_info['file_path'])
|
||||
image = Image.open(image_path)
|
||||
|
||||
alpha = np.array(image.getchannel(3))
|
||||
bbox = np.array(alpha).nonzero()
|
||||
bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
|
||||
center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
|
||||
hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
|
||||
aug_hsize = hsize
|
||||
aug_center = center
|
||||
aug_bbox = [
|
||||
int(aug_center[0] - aug_hsize),
|
||||
int(aug_center[1] - aug_hsize),
|
||||
int(aug_center[0] + aug_hsize),
|
||||
int(aug_center[1] + aug_hsize),
|
||||
]
|
||||
|
||||
img = image.crop(aug_bbox)
|
||||
img = img.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
|
||||
alpha = img.getchannel(3)
|
||||
img = img.convert('RGB')
|
||||
img = torch.tensor(np.array(img)).permute(2, 0, 1).float() / 255.0
|
||||
alpha = torch.tensor(np.array(alpha)).float() / 255.0
|
||||
img = img * alpha.unsqueeze(0)
|
||||
|
||||
cond_images.append(img)
|
||||
|
||||
pack['cond'] = [torch.stack(cond_images, dim=0)] # (V,3,H,W)
|
||||
return pack
|
||||
@@ -0,0 +1,173 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import pickle
|
||||
import torch
|
||||
import utils3d
|
||||
from .components import StandardDatasetBase
|
||||
from ..modules import sparse as sp
|
||||
from ..renderers import MeshRenderer
|
||||
from ..representations import Mesh
|
||||
from ..utils.data_utils import load_balanced_group_indices
|
||||
import o_voxel
|
||||
|
||||
|
||||
class FlexiDualGridVisMixin:
|
||||
@torch.no_grad()
|
||||
def visualize_sample(self, x: dict):
|
||||
mesh = x['mesh']
|
||||
|
||||
renderer = MeshRenderer({'near': 1, 'far': 3})
|
||||
renderer.rendering_options.resolution = 512
|
||||
renderer.rendering_options.ssaa = 4
|
||||
|
||||
# Build camera
|
||||
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
|
||||
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
|
||||
yaws = [y + yaws_offset for y in yaws]
|
||||
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
|
||||
|
||||
exts = []
|
||||
ints = []
|
||||
for yaw, pitch in zip(yaws, pitch):
|
||||
orig = torch.tensor([
|
||||
np.sin(yaw) * np.cos(pitch),
|
||||
np.cos(yaw) * np.cos(pitch),
|
||||
np.sin(pitch),
|
||||
]).float().cuda() * 2
|
||||
fov = torch.deg2rad(torch.tensor(30)).cuda()
|
||||
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
|
||||
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
|
||||
exts.append(extrinsics)
|
||||
ints.append(intrinsics)
|
||||
|
||||
# Build each representation
|
||||
images = []
|
||||
for m in mesh:
|
||||
image = torch.zeros(3, 1024, 1024).cuda()
|
||||
tile = [2, 2]
|
||||
for j, (ext, intr) in enumerate(zip(exts, ints)):
|
||||
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = \
|
||||
renderer.render(m.cuda(), ext, intr)['normal']
|
||||
images.append(image)
|
||||
images = torch.stack(images)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
class FlexiDualGridDataset(FlexiDualGridVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
Flexible Dual Grid Dataset
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
resolution (int): resolution of the voxel grid
|
||||
min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
roots,
|
||||
resolution: int = 1024,
|
||||
max_active_voxels: int = 1000000,
|
||||
max_num_faces: int = None,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
):
|
||||
self.resolution = resolution
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.max_active_voxels = max_active_voxels
|
||||
self.max_num_faces = max_num_faces
|
||||
self.value_range = (0, 1)
|
||||
|
||||
super().__init__(roots)
|
||||
|
||||
self.loads = [self.metadata.loc[sha256, f'dual_grid_size'] for _, sha256, _ in self.instances]
|
||||
|
||||
def __str__(self):
|
||||
lines = [
|
||||
super().__str__(),
|
||||
f' - Resolution: {self.resolution}',
|
||||
]
|
||||
return '\n'.join(lines)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
metadata = metadata[metadata[f'dual_grid_converted'] == True]
|
||||
stats['Dual Grid Converted'] = len(metadata)
|
||||
if self.min_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
metadata = metadata[metadata[f'dual_grid_size'] <= self.max_active_voxels]
|
||||
stats[f'Active Voxels <= {self.max_active_voxels}'] = len(metadata)
|
||||
if self.max_num_faces is not None:
|
||||
metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
|
||||
stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def read_mesh(self, root, instance):
|
||||
with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
start = 0
|
||||
vertices = []
|
||||
faces = []
|
||||
for obj in dump['objects']:
|
||||
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
||||
continue
|
||||
vertices.append(obj['vertices'])
|
||||
faces.append(obj['faces'] + start)
|
||||
start += len(obj['vertices'])
|
||||
vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
|
||||
faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
|
||||
vertices_min = vertices.min(dim=0)[0]
|
||||
vertices_max = vertices.max(dim=0)[0]
|
||||
center = (vertices_min + vertices_max) / 2
|
||||
scale = 0.99999 / (vertices_max - vertices_min).max()
|
||||
vertices = (vertices - center) * scale
|
||||
assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
|
||||
return {'mesh': [Mesh(vertices=vertices, faces=faces)]}
|
||||
|
||||
def read_dual_grid(self, root, instance):
|
||||
coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
|
||||
vertices = sp.SparseTensor(
|
||||
(attr['vertices'] / 255.0).float(),
|
||||
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
||||
)
|
||||
intersected = vertices.replace(torch.cat([
|
||||
attr['intersected'] % 2,
|
||||
attr['intersected'] // 2 % 2,
|
||||
attr['intersected'] // 4 % 2,
|
||||
], dim=-1).bool())
|
||||
return {'vertices': vertices, 'intersected': intersected}
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
mesh = self.read_mesh(root['mesh_dump'], instance)
|
||||
dual_grid = self.read_dual_grid(root['dual_grid'], instance)
|
||||
return {**mesh, **dual_grid}
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch, split_size=None):
|
||||
if split_size is None:
|
||||
group_idx = [list(range(len(batch)))]
|
||||
else:
|
||||
group_idx = load_balanced_group_indices([b['vertices'].feats.shape[0] for b in batch], split_size)
|
||||
packs = []
|
||||
for group in group_idx:
|
||||
sub_batch = [batch[i] for i in group]
|
||||
pack = {}
|
||||
|
||||
keys = [k for k in sub_batch[0].keys()]
|
||||
for k in keys:
|
||||
if isinstance(sub_batch[0][k], torch.Tensor):
|
||||
pack[k] = torch.stack([b[k] for b in sub_batch])
|
||||
elif isinstance(sub_batch[0][k], sp.SparseTensor):
|
||||
pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
|
||||
elif isinstance(sub_batch[0][k], list):
|
||||
pack[k] = sum([b[k] for b in sub_batch], [])
|
||||
else:
|
||||
pack[k] = [b[k] for b in sub_batch]
|
||||
|
||||
packs.append(pack)
|
||||
|
||||
if split_size is None:
|
||||
return packs[0]
|
||||
return packs
|
||||
|
||||
@@ -0,0 +1,408 @@
|
||||
import os
|
||||
import json
|
||||
from typing import *
|
||||
import numpy as np
|
||||
import torch
|
||||
import utils3d
|
||||
from PIL import Image
|
||||
from ..representations import Voxel
|
||||
from ..renderers import VoxelRenderer
|
||||
from .components import StandardDatasetBase, ImageConditionedMixin, ViewImageConditionedMixin
|
||||
from .. import models
|
||||
from ..utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
|
||||
|
||||
|
||||
class SparseStructureLatentVisMixin:
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16.json',
|
||||
ss_dec_path: Optional[str] = None,
|
||||
ss_dec_ckpt: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.ss_dec = None
|
||||
self.pretrained_ss_dec = pretrained_ss_dec
|
||||
self.ss_dec_path = ss_dec_path
|
||||
self.ss_dec_ckpt = ss_dec_ckpt
|
||||
|
||||
def _loading_ss_dec(self):
|
||||
if self.ss_dec is not None:
|
||||
return
|
||||
if self.ss_dec_path is not None:
|
||||
cfg = json.load(open(os.path.join(self.ss_dec_path, 'config.json'), 'r'))
|
||||
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
|
||||
ckpt_path = os.path.join(self.ss_dec_path, 'ckpts', f'decoder_{self.ss_dec_ckpt}.pt')
|
||||
decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
|
||||
else:
|
||||
decoder = models.from_pretrained(self.pretrained_ss_dec)
|
||||
self.ss_dec = decoder.cuda().eval()
|
||||
|
||||
def _delete_ss_dec(self):
|
||||
del self.ss_dec
|
||||
self.ss_dec = None
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_latent(self, z, batch_size=4):
|
||||
self._loading_ss_dec()
|
||||
ss = []
|
||||
if self.normalization:
|
||||
z = z * self.std.to(z.device) + self.mean.to(z.device)
|
||||
for i in range(0, z.shape[0], batch_size):
|
||||
ss.append(self.ss_dec(z[i:i+batch_size]))
|
||||
ss = torch.cat(ss, dim=0)
|
||||
self._delete_ss_dec()
|
||||
return ss
|
||||
|
||||
@torch.no_grad()
|
||||
def visualize_sample(
|
||||
self,
|
||||
x_0: Union[torch.Tensor, dict],
|
||||
camera_angle_x: Optional[torch.Tensor] = None,
|
||||
camera_distance: Optional[torch.Tensor] = None,
|
||||
mesh_scale: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
Visualize sparse structure samples.
|
||||
|
||||
Args:
|
||||
x_0: Latent tensor [B, C, D, H, W] or dict containing 'x_0'
|
||||
camera_angle_x: Optional [B] camera FOV angle in radians
|
||||
camera_distance: Optional [B] camera distance for GT view rendering
|
||||
mesh_scale: Optional [B] mesh scale factor for coordinate alignment
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
'multiview': [B, 3, 1024, 1024] - 4 fixed views rendered in 2x2 grid
|
||||
'gt_view': [B, 3, 512, 512] - GT camera view (if camera params provided)
|
||||
"""
|
||||
x_0 = x_0 if isinstance(x_0, torch.Tensor) else x_0['x_0']
|
||||
x_0 = self.decode_latent(x_0.cuda())
|
||||
|
||||
renderer = VoxelRenderer()
|
||||
renderer.rendering_options.resolution = 512
|
||||
renderer.rendering_options.ssaa = 4
|
||||
|
||||
# Build fixed camera views (4 views: 0°, 90°, 180°, 270°)
|
||||
yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
|
||||
yaw_offset = -16 / 180 * np.pi
|
||||
yaw = [y + yaw_offset for y in yaw]
|
||||
pitch = [20 / 180 * np.pi for _ in range(4)]
|
||||
fixed_exts, fixed_ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
|
||||
|
||||
# Check if we have GT camera parameters for front view rendering
|
||||
# GT view uses the fixed front_view_transform_matrix from image_conditioned_proj.py
|
||||
has_gt_camera = (
|
||||
camera_angle_x is not None and
|
||||
camera_distance is not None and
|
||||
mesh_scale is not None
|
||||
)
|
||||
|
||||
multiview_images = []
|
||||
gt_view_images = []
|
||||
|
||||
# Build each representation
|
||||
x_0 = x_0.cuda()
|
||||
for i in range(x_0.shape[0]):
|
||||
coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False)
|
||||
resolution = x_0.shape[-1]
|
||||
color = coords / resolution
|
||||
|
||||
# Standard voxel for fixed multiview rendering (origin at [-0.5, -0.5, -0.5])
|
||||
rep = Voxel(
|
||||
origin=[-0.5, -0.5, -0.5],
|
||||
voxel_size=1/resolution,
|
||||
coords=coords,
|
||||
attrs=color,
|
||||
layout={
|
||||
'color': slice(0, 3),
|
||||
}
|
||||
)
|
||||
|
||||
# Render 4 fixed views (2x2 grid)
|
||||
image = torch.zeros(3, 1024, 1024).cuda()
|
||||
tile = [2, 2]
|
||||
for j, (ext, intr) in enumerate(zip(fixed_exts, fixed_ints)):
|
||||
res = renderer.render(rep, ext, intr, colors_overwrite=color)
|
||||
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
|
||||
multiview_images.append(image)
|
||||
|
||||
# Render GT camera view using the fixed front view from image_conditioned_proj.py
|
||||
if has_gt_camera:
|
||||
# The GT view should match exactly how ProjGrid projects 3D points to 2D.
|
||||
#
|
||||
# In image_conditioned_proj.py (ProjGrid.forward):
|
||||
# 1. grid_points are in [-1, 1]^3 (from torch.linspace(-1, 1, res))
|
||||
# 2. grid_points are rotated by rotation_matrix (Y-Z swap): x'=x, y'=-z, z'=y
|
||||
# 3. grid_points are scaled: grid_points / mesh_scale / 2
|
||||
# 4. Points are projected using front_view_transform_matrix with distance
|
||||
#
|
||||
# front_view_transform_matrix (camera-to-world):
|
||||
# [[1, 0, 0, 0],
|
||||
# [0, 0, -1, -distance],
|
||||
# [0, 1, 0, 0],
|
||||
# [0, 0, 0, 1]]
|
||||
#
|
||||
# Camera is at (0, -distance, 0) in Blender coords (Z-up), looking at origin.
|
||||
#
|
||||
# To match this in VoxelRenderer:
|
||||
# 1. Voxel coords [0, res-1] map to positions via: pos = (coords + 0.5) * voxel_size + origin
|
||||
# 2. We need these positions to match ProjGrid's transformed grid_points
|
||||
# 3. Apply rotation by swapping/flipping coords, then scale voxel_size and origin
|
||||
|
||||
scale = mesh_scale[i].item()
|
||||
distance = camera_distance[i].item()
|
||||
fov = camera_angle_x[i].item()
|
||||
|
||||
# Coordinate transformation to match ProjGrid's rotation (x'=x, y'=-z, z'=y)
|
||||
# new_coords maps to rotated positions in the same grid structure
|
||||
new_coords = torch.zeros_like(coords)
|
||||
new_coords[:, 0] = coords[:, 0] # x stays
|
||||
new_coords[:, 1] = (resolution - 1) - coords[:, 2] # y' = -z (flip for negation)
|
||||
new_coords[:, 2] = coords[:, 1] # z' = y
|
||||
|
||||
# Voxel position calculation:
|
||||
# Original: pos = (coords + 0.5) / res - 0.5 -> range [-0.5, 0.5]
|
||||
# We need: pos = (coords + 0.5) * 2 / res - 1 -> range [-1, 1] (like ProjGrid)
|
||||
# Then: pos_final = pos / scale / 2 -> range [-0.5/scale, 0.5/scale]
|
||||
#
|
||||
# Combined: pos_final = ((coords + 0.5) * 2 / res - 1) / scale / 2
|
||||
# = (coords + 0.5) / res / scale - 0.5 / scale
|
||||
# = (coords + 0.5) * voxel_size + origin
|
||||
# where: voxel_size = 1 / res / scale
|
||||
# origin = -0.5 / scale
|
||||
|
||||
scaled_voxel_size = 1.0 / resolution / scale
|
||||
scaled_origin = [-0.5 / scale, -0.5 / scale, -0.5 / scale]
|
||||
|
||||
rep_scaled = Voxel(
|
||||
origin=scaled_origin,
|
||||
voxel_size=scaled_voxel_size,
|
||||
coords=new_coords,
|
||||
attrs=color,
|
||||
layout={
|
||||
'color': slice(0, 3),
|
||||
}
|
||||
)
|
||||
|
||||
# Build the fixed front view camera (same as front_view_transform_matrix)
|
||||
# Camera at (0, -distance, 0), looking at origin, up is Z
|
||||
cam_pos = torch.tensor([0.0, -distance, 0.0], device=coords.device)
|
||||
look_at = torch.tensor([0.0, 0.0, 0.0], device=coords.device)
|
||||
cam_up = torch.tensor([0.0, 0.0, 1.0], device=coords.device)
|
||||
|
||||
gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up)
|
||||
gt_int = utils3d.torch.intrinsics_from_fov_xy(
|
||||
torch.tensor(fov, device=coords.device),
|
||||
torch.tensor(fov, device=coords.device)
|
||||
)
|
||||
|
||||
# Ensure tensors are on the correct device (utils3d may not preserve device)
|
||||
gt_ext = gt_ext.to(coords.device)
|
||||
gt_int = gt_int.to(coords.device)
|
||||
|
||||
gt_res = renderer.render(rep_scaled, gt_ext, gt_int, colors_overwrite=color)
|
||||
gt_view_images.append(gt_res['color'])
|
||||
|
||||
result = {
|
||||
'multiview': torch.stack(multiview_images),
|
||||
}
|
||||
|
||||
if has_gt_camera and len(gt_view_images) > 0:
|
||||
result['gt_view'] = torch.stack(gt_view_images)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class SparseStructureLatent(SparseStructureLatentVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
Sparse structure latent dataset
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
min_aesthetic_score (float): minimum aesthetic score
|
||||
normalization (dict): normalization stats
|
||||
pretrained_ss_dec (str): name of the pretrained sparse structure decoder
|
||||
ss_dec_path (str): path to the sparse structure decoder, if given, will override the pretrained_ss_dec
|
||||
ss_dec_ckpt (str): name of the sparse structure decoder checkpoint
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip
|
||||
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
normalization: Optional[dict] = None,
|
||||
pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
|
||||
ss_dec_path: Optional[str] = None,
|
||||
ss_dec_ckpt: Optional[str] = None,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[list] = None,
|
||||
):
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.normalization = normalization
|
||||
self.value_range = (0, 1)
|
||||
|
||||
super().__init__(
|
||||
roots,
|
||||
pretrained_ss_dec=pretrained_ss_dec,
|
||||
ss_dec_path=ss_dec_path,
|
||||
ss_dec_ckpt=ss_dec_ckpt,
|
||||
skip_list=skip_list,
|
||||
skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
|
||||
)
|
||||
|
||||
if self.normalization is not None:
|
||||
self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1)
|
||||
self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
metadata = metadata[metadata['ss_latent_encoded'] == True]
|
||||
stats['With latent'] = len(metadata)
|
||||
# Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
|
||||
skip_aesthetic = (
|
||||
(dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
|
||||
('aesthetic_score' not in metadata.columns)
|
||||
)
|
||||
if skip_aesthetic:
|
||||
stats[f'Aesthetic score check skipped'] = len(metadata)
|
||||
else:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
latent = np.load(os.path.join(root['ss_latent'], f'{instance}.npz'))
|
||||
z = torch.tensor(latent['z']).float()
|
||||
if self.normalization is not None:
|
||||
z = (z - self.mean) / self.std
|
||||
|
||||
pack = {
|
||||
'x_0': z,
|
||||
}
|
||||
return pack
|
||||
|
||||
|
||||
class ImageConditionedSparseStructureLatent(ImageConditionedMixin, SparseStructureLatent):
|
||||
"""
|
||||
Image-conditioned sparse structure dataset
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class SparseStructureLatentView(SparseStructureLatentVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
View-based sparse structure latent dataset.
|
||||
|
||||
Data format: {sha256}/view{XX}.npz where each npz contains 'z' key.
|
||||
|
||||
Args:
|
||||
num_views (int): Number of views to use (0 to num_views-1). Default is 2.
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip
|
||||
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
normalization: Optional[dict] = None,
|
||||
num_views: int = 2,
|
||||
pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
|
||||
ss_dec_path: Optional[str] = None,
|
||||
ss_dec_ckpt: Optional[str] = None,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[list] = None,
|
||||
):
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.normalization = normalization
|
||||
self.num_views = num_views
|
||||
self.value_range = (0, 1)
|
||||
|
||||
super().__init__(
|
||||
roots,
|
||||
pretrained_ss_dec=pretrained_ss_dec,
|
||||
ss_dec_path=ss_dec_path,
|
||||
ss_dec_ckpt=ss_dec_ckpt,
|
||||
skip_list=skip_list,
|
||||
skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
|
||||
)
|
||||
|
||||
if self.normalization is not None:
|
||||
self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1)
|
||||
self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
# View-based ss_latent uses columns like:
|
||||
# ss_latent_view00_encoded, ss_latent_view01_encoded, ... (view format)
|
||||
# ss_latent_view_scale00_encoded, ss_latent_view_scale01_encoded, ... (view_scale format)
|
||||
# Check both formats and use whichever exists (prefer view_scale over view)
|
||||
required_view_cols = [f'ss_latent_view_scale{i:02d}_encoded' for i in range(self.num_views)]
|
||||
existing_view_cols = [col for col in required_view_cols if col in metadata.columns]
|
||||
|
||||
if not existing_view_cols:
|
||||
# Fallback to view format
|
||||
required_view_cols = [f'ss_latent_view{i:02d}_encoded' for i in range(self.num_views)]
|
||||
existing_view_cols = [col for col in required_view_cols if col in metadata.columns]
|
||||
|
||||
if existing_view_cols:
|
||||
# Filter rows where all required views are encoded
|
||||
# Note: NaN should be treated as False, so use == True for explicit comparison
|
||||
has_all_views = (metadata[existing_view_cols] == True).all(axis=1)
|
||||
metadata = metadata[has_all_views]
|
||||
stats[f'With {self.num_views} view latents'] = len(metadata)
|
||||
else:
|
||||
# Fallback: check ss_latent_encoded column
|
||||
if 'ss_latent_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata['ss_latent_encoded'] == True]
|
||||
stats['With latent'] = len(metadata)
|
||||
else:
|
||||
raise ValueError(f'No view columns found in metadata: {metadata.columns.tolist()}')
|
||||
# Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
|
||||
skip_aesthetic = (
|
||||
(dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
|
||||
('aesthetic_score' not in metadata.columns)
|
||||
)
|
||||
if skip_aesthetic:
|
||||
stats[f'Aesthetic score check skipped'] = len(metadata)
|
||||
else:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
# View-based format: directory with view{XX}.npz files
|
||||
latent_dir = os.path.join(root['ss_latent'], instance)
|
||||
|
||||
# Randomly select a view from the configured range
|
||||
view_idx = np.random.randint(0, self.num_views)
|
||||
view_file = f'view{view_idx:02d}.npz'
|
||||
|
||||
# Store view info for ViewImageConditionedMixin
|
||||
self._current_view_idx = view_idx
|
||||
self._current_latent_dir = latent_dir
|
||||
|
||||
latent = np.load(os.path.join(latent_dir, view_file))
|
||||
z = torch.tensor(latent['z']).float()
|
||||
if self.normalization is not None:
|
||||
z = (z - self.mean) / self.std
|
||||
|
||||
pack = {
|
||||
'x_0': z,
|
||||
'view_idx': view_idx,
|
||||
}
|
||||
return pack
|
||||
|
||||
|
||||
class ViewImageConditionedSparseStructureLatentView(ViewImageConditionedMixin, SparseStructureLatentView):
|
||||
"""
|
||||
Image-conditioned view-based sparse structure dataset.
|
||||
|
||||
Loads ss_latent from {sha256}/view{XX}.npz format and pairs with
|
||||
corresponding view from render_cond.
|
||||
|
||||
Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,298 @@
|
||||
import os
|
||||
import io
|
||||
from typing import Union
|
||||
import numpy as np
|
||||
import pickle
|
||||
import torch
|
||||
from PIL import Image
|
||||
import o_voxel
|
||||
import utils3d
|
||||
from .components import StandardDatasetBase
|
||||
from ..modules import sparse as sp
|
||||
from ..renderers import VoxelRenderer
|
||||
from ..representations import Voxel
|
||||
from ..representations.mesh import MeshWithPbrMaterial, TextureFilterMode, TextureWrapMode, AlphaMode, PbrMaterial, Texture
|
||||
|
||||
from ..utils.data_utils import load_balanced_group_indices
|
||||
|
||||
|
||||
def is_power_of_two(n: int) -> bool:
|
||||
return n > 0 and (n & (n - 1)) == 0
|
||||
|
||||
|
||||
def nearest_power_of_two(n: int) -> int:
|
||||
if n < 1:
|
||||
raise ValueError("n must be >= 1")
|
||||
if is_power_of_two(n):
|
||||
return n
|
||||
lower = 2 ** (n.bit_length() - 1)
|
||||
upper = 2 ** n.bit_length()
|
||||
if n - lower < upper - n:
|
||||
return lower
|
||||
else:
|
||||
return upper
|
||||
|
||||
|
||||
class SparseVoxelPbrVisMixin:
|
||||
@torch.no_grad()
|
||||
def visualize_sample(self, x: Union[sp.SparseTensor, dict]):
|
||||
x = x if isinstance(x, sp.SparseTensor) else x['x']
|
||||
|
||||
renderer = VoxelRenderer()
|
||||
renderer.rendering_options.resolution = 512
|
||||
renderer.rendering_options.ssaa = 4
|
||||
|
||||
# Build camera
|
||||
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
|
||||
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
|
||||
yaws = [y + yaws_offset for y in yaws]
|
||||
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
|
||||
|
||||
exts = []
|
||||
ints = []
|
||||
for yaw, pitch in zip(yaws, pitch):
|
||||
orig = torch.tensor([
|
||||
np.sin(yaw) * np.cos(pitch),
|
||||
np.cos(yaw) * np.cos(pitch),
|
||||
np.sin(pitch),
|
||||
]).float().cuda() * 2
|
||||
fov = torch.deg2rad(torch.tensor(30)).cuda()
|
||||
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
|
||||
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
|
||||
exts.append(extrinsics)
|
||||
ints.append(intrinsics)
|
||||
|
||||
images = {k: [] for k in self.layout}
|
||||
|
||||
# Build each representation
|
||||
x = x.cuda()
|
||||
for i in range(x.shape[0]):
|
||||
rep = Voxel(
|
||||
origin=[-0.5, -0.5, -0.5],
|
||||
voxel_size=1/self.resolution,
|
||||
coords=x[i].coords[:, 1:].contiguous(),
|
||||
attrs=None,
|
||||
layout={
|
||||
'color': slice(0, 3),
|
||||
}
|
||||
)
|
||||
for k in self.layout:
|
||||
image = torch.zeros(3, 1024, 1024).cuda()
|
||||
tile = [2, 2]
|
||||
for j, (ext, intr) in enumerate(zip(exts, ints)):
|
||||
attr = x[i].feats[:, self.layout[k]].expand(-1, 3)
|
||||
res = renderer.render(rep, ext, intr, colors_overwrite=attr)
|
||||
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
|
||||
images[k].append(image)
|
||||
|
||||
for k in self.layout:
|
||||
images[k] = torch.stack(images[k])
|
||||
|
||||
return images
|
||||
|
||||
|
||||
class SparseVoxelPbrDataset(SparseVoxelPbrVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
Sparse Voxel PBR dataset.
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
resolution (int): resolution of the voxel grid
|
||||
min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
roots,
|
||||
resolution: int = 1024,
|
||||
max_active_voxels: int = 1000000,
|
||||
max_num_faces: int = None,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
|
||||
with_mesh: bool = True,
|
||||
):
|
||||
self.resolution = resolution
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.max_active_voxels = max_active_voxels
|
||||
self.max_num_faces = max_num_faces
|
||||
self.with_mesh = with_mesh
|
||||
self.value_range = (-1, 1)
|
||||
self.channels = {
|
||||
'base_color': 3,
|
||||
'metallic': 1,
|
||||
'roughness': 1,
|
||||
'emissive': 3,
|
||||
'alpha': 1,
|
||||
}
|
||||
self.layout = {}
|
||||
start = 0
|
||||
for attr in attrs:
|
||||
self.layout[attr] = slice(start, start + self.channels[attr])
|
||||
start += self.channels[attr]
|
||||
|
||||
super().__init__(roots)
|
||||
|
||||
self.loads = [self.metadata.loc[sha256, f'num_pbr_voxels'] for _, sha256, _ in self.instances]
|
||||
|
||||
def __str__(self):
|
||||
lines = [
|
||||
super().__str__(),
|
||||
f' - Resolution: {self.resolution}',
|
||||
f' - Attributes: {list(self.layout.keys())}',
|
||||
]
|
||||
return '\n'.join(lines)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
metadata = metadata[metadata['pbr_voxelized'] == True]
|
||||
stats['PBR Voxelized'] = len(metadata)
|
||||
if self.min_aesthetic_score is not None:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
metadata = metadata[metadata['num_pbr_voxels'] <= self.max_active_voxels]
|
||||
stats[f'Active voxels <= {self.max_active_voxels}'] = len(metadata)
|
||||
if self.max_num_faces is not None:
|
||||
metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
|
||||
stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
@staticmethod
|
||||
def _texture_from_dump(pack) -> Texture:
|
||||
png_bytes = pack['image']
|
||||
image = Image.open(io.BytesIO(png_bytes))
|
||||
if image.width != image.height or not is_power_of_two(image.width):
|
||||
size = nearest_power_of_two(max(image.width, image.height))
|
||||
image = image.resize((size, size), Image.LANCZOS)
|
||||
texture = torch.tensor(np.array(image) / 255.0, dtype=torch.float32).reshape(image.height, image.width, -1)
|
||||
filter_mode = {
|
||||
'Linear': TextureFilterMode.LINEAR,
|
||||
'Closest': TextureFilterMode.CLOSEST,
|
||||
'Cubic': TextureFilterMode.LINEAR,
|
||||
'Smart': TextureFilterMode.LINEAR,
|
||||
}[pack['interpolation']]
|
||||
wrap_mode = {
|
||||
'REPEAT': TextureWrapMode.REPEAT,
|
||||
'EXTEND': TextureWrapMode.CLAMP_TO_EDGE,
|
||||
'CLIP': TextureWrapMode.CLAMP_TO_EDGE,
|
||||
'MIRROR': TextureWrapMode.MIRRORED_REPEAT,
|
||||
}[pack['extension']]
|
||||
return Texture(texture, filter_mode=filter_mode, wrap_mode=wrap_mode)
|
||||
|
||||
def read_mesh_with_texture(self, root, instance):
|
||||
with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
|
||||
dump = pickle.load(f)
|
||||
|
||||
# Fix dump alpha map
|
||||
for mat in dump['materials']:
|
||||
if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
|
||||
mat['alphaMode'] = 'BLEND'
|
||||
|
||||
# process material
|
||||
materials = []
|
||||
for mat in dump['materials']:
|
||||
materials.append(PbrMaterial(
|
||||
base_color_texture=self._texture_from_dump(mat['baseColorTexture']) if mat['baseColorTexture'] is not None else None,
|
||||
base_color_factor=mat['baseColorFactor'],
|
||||
metallic_texture=self._texture_from_dump(mat['metallicTexture']) if mat['metallicTexture'] is not None else None,
|
||||
metallic_factor=mat['metallicFactor'],
|
||||
roughness_texture=self._texture_from_dump(mat['roughnessTexture']) if mat['roughnessTexture'] is not None else None,
|
||||
roughness_factor=mat['roughnessFactor'],
|
||||
alpha_texture=self._texture_from_dump(mat['alphaTexture']) if mat['alphaTexture'] is not None else None,
|
||||
alpha_factor=mat['alphaFactor'],
|
||||
alpha_mode={
|
||||
'OPAQUE': AlphaMode.OPAQUE,
|
||||
'MASK': AlphaMode.MASK,
|
||||
'BLEND': AlphaMode.BLEND,
|
||||
}[mat['alphaMode']],
|
||||
alpha_cutoff=mat['alphaCutoff'],
|
||||
))
|
||||
materials.append(PbrMaterial(
|
||||
base_color_factor=[0.8, 0.8, 0.8],
|
||||
alpha_factor=1.0,
|
||||
metallic_factor=0.0,
|
||||
roughness_factor=0.5,
|
||||
alpha_mode=AlphaMode.OPAQUE,
|
||||
alpha_cutoff=0.5,
|
||||
)) # append default material
|
||||
|
||||
# process mesh
|
||||
start = 0
|
||||
vertices = []
|
||||
faces = []
|
||||
material_ids = []
|
||||
uv_coords = []
|
||||
for obj in dump['objects']:
|
||||
if obj['vertices'].size == 0 or obj['faces'].size == 0:
|
||||
continue
|
||||
vertices.append(obj['vertices'])
|
||||
faces.append(obj['faces'] + start)
|
||||
obj['mat_ids'][obj['mat_ids'] == -1] = len(materials) - 1
|
||||
material_ids.append(obj['mat_ids'])
|
||||
uv_coords.append(obj['uvs'] if obj['uvs'] is not None else np.zeros((obj['faces'].shape[0], 3, 2), dtype=np.float32))
|
||||
start += len(obj['vertices'])
|
||||
|
||||
vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
|
||||
faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
|
||||
material_ids = torch.from_numpy(np.concatenate(material_ids, axis=0)).long()
|
||||
uv_coords = torch.from_numpy(np.concatenate(uv_coords, axis=0)).float()
|
||||
|
||||
# Normalize vertices
|
||||
vertices_min = vertices.min(dim=0)[0]
|
||||
vertices_max = vertices.max(dim=0)[0]
|
||||
center = (vertices_min + vertices_max) / 2
|
||||
scale = 0.99999 / (vertices_max - vertices_min).max()
|
||||
vertices = (vertices - center) * scale
|
||||
assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
|
||||
|
||||
return {'mesh': [MeshWithPbrMaterial(
|
||||
vertices=vertices,
|
||||
faces=faces,
|
||||
material_ids=material_ids,
|
||||
uv_coords=uv_coords,
|
||||
materials=materials,
|
||||
)]}
|
||||
|
||||
def read_pbr_voxel(self, root, instance):
|
||||
coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
|
||||
feats = torch.concat([attr[k] for k in self.layout], dim=-1) / 255.0 * 2 - 1
|
||||
x = sp.SparseTensor(
|
||||
feats.float(),
|
||||
torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
|
||||
)
|
||||
return {'x': x}
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
if self.with_mesh:
|
||||
mesh = self.read_mesh_with_texture(root['pbr_dump'], instance)
|
||||
pbr_voxel = self.read_pbr_voxel(root['pbr_voxel'], instance)
|
||||
return {**mesh, **pbr_voxel}
|
||||
else:
|
||||
return self.read_pbr_voxel(root['pbr_voxel'], instance)
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch, split_size=None):
|
||||
if split_size is None:
|
||||
group_idx = [list(range(len(batch)))]
|
||||
else:
|
||||
group_idx = load_balanced_group_indices([b['x'].feats.shape[0] for b in batch], split_size)
|
||||
packs = []
|
||||
for group in group_idx:
|
||||
sub_batch = [batch[i] for i in group]
|
||||
pack = {}
|
||||
|
||||
keys = [k for k in sub_batch[0].keys()]
|
||||
for k in keys:
|
||||
if isinstance(sub_batch[0][k], torch.Tensor):
|
||||
pack[k] = torch.stack([b[k] for b in sub_batch])
|
||||
elif isinstance(sub_batch[0][k], sp.SparseTensor):
|
||||
pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
|
||||
elif isinstance(sub_batch[0][k], list):
|
||||
pack[k] = sum([b[k] for b in sub_batch], [])
|
||||
else:
|
||||
pack[k] = [b[k] for b in sub_batch]
|
||||
|
||||
packs.append(pack)
|
||||
|
||||
if split_size is None:
|
||||
return packs[0]
|
||||
return packs
|
||||
@@ -0,0 +1,224 @@
|
||||
import json
|
||||
import os
|
||||
from typing import *
|
||||
import numpy as np
|
||||
import torch
|
||||
import utils3d.torch
|
||||
from .components import StandardDatasetBase, ImageConditionedMixin
|
||||
from ..modules.sparse.basic import SparseTensor
|
||||
from .. import models
|
||||
from ..utils.render_utils import get_renderer
|
||||
from ..utils.data_utils import load_balanced_group_indices
|
||||
|
||||
|
||||
class SLatVisMixin:
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
|
||||
slat_dec_path: Optional[str] = None,
|
||||
slat_dec_ckpt: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.slat_dec = None
|
||||
self.pretrained_slat_dec = pretrained_slat_dec
|
||||
self.slat_dec_path = slat_dec_path
|
||||
self.slat_dec_ckpt = slat_dec_ckpt
|
||||
|
||||
def _loading_slat_dec(self):
|
||||
if self.slat_dec is not None:
|
||||
return
|
||||
if self.slat_dec_path is not None:
|
||||
cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
|
||||
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
|
||||
ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
|
||||
decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
|
||||
else:
|
||||
decoder = models.from_pretrained(self.pretrained_slat_dec)
|
||||
self.slat_dec = decoder.cuda().eval()
|
||||
|
||||
def _delete_slat_dec(self):
|
||||
del self.slat_dec
|
||||
self.slat_dec = None
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_latent(self, z, batch_size=4):
|
||||
self._loading_slat_dec()
|
||||
reps = []
|
||||
if self.normalization is not None:
|
||||
z = z * self.std.to(z.device) + self.mean.to(z.device)
|
||||
for i in range(0, z.shape[0], batch_size):
|
||||
reps.append(self.slat_dec(z[i:i+batch_size]))
|
||||
reps = sum(reps, [])
|
||||
self._delete_slat_dec()
|
||||
return reps
|
||||
|
||||
@torch.no_grad()
|
||||
def visualize_sample(self, x_0: Union[SparseTensor, dict]):
|
||||
x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
|
||||
reps = self.decode_latent(x_0.cuda())
|
||||
|
||||
# Build camera
|
||||
yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
|
||||
yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
|
||||
yaws = [y + yaws_offset for y in yaws]
|
||||
pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
|
||||
|
||||
exts = []
|
||||
ints = []
|
||||
for yaw, pitch in zip(yaws, pitch):
|
||||
orig = torch.tensor([
|
||||
np.sin(yaw) * np.cos(pitch),
|
||||
np.cos(yaw) * np.cos(pitch),
|
||||
np.sin(pitch),
|
||||
]).float().cuda() * 2
|
||||
fov = torch.deg2rad(torch.tensor(40)).cuda()
|
||||
extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
|
||||
intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
|
||||
exts.append(extrinsics)
|
||||
ints.append(intrinsics)
|
||||
|
||||
renderer = get_renderer(reps[0])
|
||||
images = []
|
||||
for representation in reps:
|
||||
image = torch.zeros(3, 1024, 1024).cuda()
|
||||
tile = [2, 2]
|
||||
for j, (ext, intr) in enumerate(zip(exts, ints)):
|
||||
res = renderer.render(representation, ext, intr)
|
||||
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
|
||||
images.append(image)
|
||||
images = torch.stack(images)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
class SLat(SLatVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
structured latent V2 dataset
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
min_aesthetic_score (float): minimum aesthetic score
|
||||
max_tokens (int): maximum number of tokens
|
||||
latent_key (str): key of the latent to be used
|
||||
normalization (dict): normalization stats
|
||||
pretrained_slat_dec (str): name of the pretrained slat decoder
|
||||
slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
|
||||
slat_dec_ckpt (str): name of the slat decoder checkpoint
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip
|
||||
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
max_tokens: int = 32768,
|
||||
latent_key: str = 'shape_latent',
|
||||
normalization: Optional[dict] = None,
|
||||
pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
|
||||
slat_dec_path: Optional[str] = None,
|
||||
slat_dec_ckpt: Optional[str] = None,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[list] = None,
|
||||
):
|
||||
self.normalization = normalization
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.max_tokens = max_tokens
|
||||
self.latent_key = latent_key
|
||||
self.value_range = (0, 1)
|
||||
|
||||
super().__init__(
|
||||
roots,
|
||||
pretrained_slat_dec=pretrained_slat_dec,
|
||||
slat_dec_path=slat_dec_path,
|
||||
slat_dec_ckpt=slat_dec_ckpt,
|
||||
skip_list=skip_list,
|
||||
skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
|
||||
)
|
||||
|
||||
self.loads = [self.metadata.loc[sha256, f'{latent_key}_tokens'] for _, sha256, _ in self.instances]
|
||||
|
||||
if self.normalization is not None:
|
||||
self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1)
|
||||
self.std = torch.tensor(self.normalization['std']).reshape(1, -1)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
metadata = metadata[metadata[f'{self.latent_key}_encoded'] == True]
|
||||
stats['With latent'] = len(metadata)
|
||||
# Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
|
||||
skip_aesthetic = (
|
||||
(dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
|
||||
('aesthetic_score' not in metadata.columns)
|
||||
)
|
||||
if skip_aesthetic:
|
||||
stats[f'Aesthetic score check skipped'] = len(metadata)
|
||||
else:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
metadata = metadata[metadata[f'{self.latent_key}_tokens'] <= self.max_tokens]
|
||||
stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
data = np.load(os.path.join(root[self.latent_key], f'{instance}.npz'))
|
||||
coords = torch.tensor(data['coords']).int()
|
||||
feats = torch.tensor(data['feats']).float()
|
||||
if self.normalization is not None:
|
||||
feats = (feats - self.mean) / self.std
|
||||
return {
|
||||
'coords': coords,
|
||||
'feats': feats,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch, split_size=None):
|
||||
if split_size is None:
|
||||
group_idx = [list(range(len(batch)))]
|
||||
else:
|
||||
group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
|
||||
packs = []
|
||||
for group in group_idx:
|
||||
sub_batch = [batch[i] for i in group]
|
||||
pack = {}
|
||||
coords = []
|
||||
feats = []
|
||||
layout = []
|
||||
start = 0
|
||||
for i, b in enumerate(sub_batch):
|
||||
coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
|
||||
feats.append(b['feats'])
|
||||
layout.append(slice(start, start + b['coords'].shape[0]))
|
||||
start += b['coords'].shape[0]
|
||||
coords = torch.cat(coords)
|
||||
feats = torch.cat(feats)
|
||||
pack['x_0'] = SparseTensor(
|
||||
coords=coords,
|
||||
feats=feats,
|
||||
)
|
||||
pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]])
|
||||
pack['x_0'].register_spatial_cache('layout', layout)
|
||||
|
||||
# collate other data
|
||||
keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']]
|
||||
for k in keys:
|
||||
if isinstance(sub_batch[0][k], torch.Tensor):
|
||||
pack[k] = torch.stack([b[k] for b in sub_batch])
|
||||
elif isinstance(sub_batch[0][k], list):
|
||||
pack[k] = sum([b[k] for b in sub_batch], [])
|
||||
else:
|
||||
pack[k] = [b[k] for b in sub_batch]
|
||||
|
||||
packs.append(pack)
|
||||
|
||||
if split_size is None:
|
||||
return packs[0]
|
||||
return packs
|
||||
|
||||
|
||||
class ImageConditionedSLat(ImageConditionedMixin, SLat):
|
||||
"""
|
||||
Image conditioned structured latent dataset
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,402 @@
|
||||
import os
|
||||
import json
|
||||
from typing import *
|
||||
import numpy as np
|
||||
import torch
|
||||
import utils3d
|
||||
from .. import models
|
||||
from .components import ImageConditionedMixin, ViewImageConditionedMixin
|
||||
from ..modules.sparse import SparseTensor
|
||||
from .structured_latent import SLatVisMixin, SLat
|
||||
from ..utils.render_utils import get_renderer, yaw_pitch_r_fov_to_extrinsics_intrinsics
|
||||
from ..utils.data_utils import load_balanced_group_indices
|
||||
|
||||
|
||||
class SLatShapeVisMixin(SLatVisMixin):
|
||||
def _loading_slat_dec(self):
|
||||
if self.slat_dec is not None:
|
||||
return
|
||||
if self.slat_dec_path is not None:
|
||||
cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
|
||||
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
|
||||
ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
|
||||
decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
|
||||
else:
|
||||
decoder = models.from_pretrained(self.pretrained_slat_dec)
|
||||
decoder.set_resolution(self.resolution)
|
||||
self.slat_dec = decoder.cuda().eval()
|
||||
|
||||
@torch.no_grad()
|
||||
def visualize_sample(
|
||||
self,
|
||||
x_0: Union[SparseTensor, dict],
|
||||
camera_angle_x: Optional[torch.Tensor] = None,
|
||||
camera_distance: Optional[torch.Tensor] = None,
|
||||
mesh_scale: Optional[torch.Tensor] = None,
|
||||
):
|
||||
"""
|
||||
Visualize shape samples.
|
||||
|
||||
Args:
|
||||
x_0: SparseTensor or dict containing 'x_0'
|
||||
camera_angle_x: Optional [B] camera FOV angle in radians
|
||||
camera_distance: Optional [B] camera distance for GT view rendering
|
||||
mesh_scale: Optional [B] mesh scale factor for coordinate alignment
|
||||
|
||||
Returns:
|
||||
dict with:
|
||||
'multiview': [B, 3, 1024, 1024] - 4 fixed views rendered in 2x2 grid (normal)
|
||||
'gt_view': [B, 3, 512, 512] - GT camera view (if camera params provided)
|
||||
"""
|
||||
x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
|
||||
reps = self.decode_latent(x_0.cuda())
|
||||
|
||||
# build fixed camera views (4 views: 0°, 90°, 180°, 270°)
|
||||
yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
|
||||
yaw_offset = -16 / 180 * np.pi
|
||||
yaw = [y + yaw_offset for y in yaw]
|
||||
pitch = [20 / 180 * np.pi for _ in range(4)]
|
||||
fixed_exts, fixed_ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
|
||||
|
||||
# Check if we have GT camera parameters for GT view rendering
|
||||
has_gt_camera = (
|
||||
camera_angle_x is not None and
|
||||
camera_distance is not None and
|
||||
mesh_scale is not None
|
||||
)
|
||||
|
||||
# render
|
||||
renderer = get_renderer(reps[0])
|
||||
multiview_images = []
|
||||
gt_view_images = []
|
||||
|
||||
for i, representation in enumerate(reps):
|
||||
# Render 4 fixed views (2x2 grid)
|
||||
image = torch.zeros(3, 1024, 1024).cuda()
|
||||
tile = [2, 2]
|
||||
|
||||
# Validate mesh data before rasterization
|
||||
verts = representation.vertices
|
||||
faces = representation.faces
|
||||
if verts.shape[0] == 0 or faces.shape[0] == 0:
|
||||
print(f"[visualize_sample] Warning: sample {i} has empty mesh, skipping")
|
||||
multiview_images.append(image)
|
||||
continue
|
||||
if faces.max() >= verts.shape[0]:
|
||||
print(f"[visualize_sample] Warning: sample {i} has out-of-bound face indices "
|
||||
f"(max face idx={faces.max().item()}, num verts={verts.shape[0]}), skipping")
|
||||
multiview_images.append(image)
|
||||
continue
|
||||
if torch.isnan(verts).any() or torch.isinf(verts).any():
|
||||
print(f"[visualize_sample] Warning: sample {i} has NaN/Inf vertices, skipping")
|
||||
multiview_images.append(image)
|
||||
continue
|
||||
|
||||
try:
|
||||
for j, (ext, intr) in enumerate(zip(fixed_exts, fixed_ints)):
|
||||
res = renderer.render(representation, ext, intr)
|
||||
image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['normal']
|
||||
except RuntimeError as e:
|
||||
print(f"[visualize_sample] Warning: render failed for sample {i}: {e}")
|
||||
image = torch.zeros(3, 1024, 1024).cuda()
|
||||
multiview_images.append(image)
|
||||
|
||||
# Render GT camera view using the fixed front view (same as sparse_structure_latent.py)
|
||||
if has_gt_camera:
|
||||
# The GT view should match exactly how ProjGrid projects 3D points to 2D.
|
||||
#
|
||||
# In image_conditioned_proj.py (ProjGrid.forward):
|
||||
# 1. grid_points are in [-1, 1]^3 (from torch.linspace(-1, 1, res))
|
||||
# 2. grid_points are rotated by rotation_matrix (Y-Z swap): x'=x, y'=-z, z'=y
|
||||
# 3. grid_points are scaled: grid_points / mesh_scale / 2
|
||||
# 4. Points are projected using front_view_transform_matrix with distance
|
||||
#
|
||||
# Mesh vertices are in [-0.5, 0.5]^3. To match ProjGrid's coordinate space,
|
||||
# we need to scale them: vertices / mesh_scale -> [-0.5/s, 0.5/s]^3
|
||||
# This is equivalent to ProjGrid's: [-1,1]^3 / scale / 2 -> [-0.5/s, 0.5/s]^3
|
||||
#
|
||||
# Camera position: ProjGrid camera is at (0, -distance, 0) in Blender coords (Z-up).
|
||||
# After inverse rotation to mesh space, camera is at (0, 0, distance).
|
||||
|
||||
scale = mesh_scale[i].item()
|
||||
distance = camera_distance[i].item()
|
||||
fov = camera_angle_x[i].item()
|
||||
device = representation.vertices.device
|
||||
|
||||
# Scale mesh vertices to match ProjGrid's projection space
|
||||
from ..representations import Mesh
|
||||
scaled_rep = Mesh(
|
||||
vertices=representation.vertices / scale,
|
||||
faces=representation.faces,
|
||||
)
|
||||
|
||||
cam_pos = torch.tensor([0.0, 0.0, distance], device=device)
|
||||
look_at = torch.tensor([0.0, 0.0, 0.0], device=device)
|
||||
cam_up = torch.tensor([0.0, 1.0, 0.0], device=device)
|
||||
|
||||
gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up)
|
||||
gt_int = utils3d.torch.intrinsics_from_fov_xy(
|
||||
torch.tensor(fov, device=device),
|
||||
torch.tensor(fov, device=device)
|
||||
)
|
||||
|
||||
gt_ext = gt_ext.to(device)
|
||||
gt_int = gt_int.to(device)
|
||||
|
||||
# Use scaled mesh renderer with appropriate near/far for smaller mesh
|
||||
mesh_half_size = 0.5 / scale
|
||||
renderer.rendering_options.near = max(0.01, distance - mesh_half_size - 0.5)
|
||||
renderer.rendering_options.far = distance + mesh_half_size + 0.5
|
||||
|
||||
try:
|
||||
gt_res = renderer.render(scaled_rep, gt_ext, gt_int)
|
||||
gt_view_images.append(gt_res['normal'])
|
||||
except RuntimeError as e:
|
||||
print(f"[visualize_sample] Warning: GT view render failed for sample {i}: {e}")
|
||||
gt_view_images.append(torch.full((3, 512, 512), 0.5, device=device))
|
||||
|
||||
result = {
|
||||
'multiview': torch.stack(multiview_images),
|
||||
}
|
||||
|
||||
if has_gt_camera and len(gt_view_images) > 0:
|
||||
result['gt_view'] = torch.stack(gt_view_images)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class SLatShape(SLatShapeVisMixin, SLat):
|
||||
"""
|
||||
structured latent for shape generation
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
resolution (int): resolution of the shape
|
||||
min_aesthetic_score (float): minimum aesthetic score
|
||||
max_tokens (int): maximum number of tokens
|
||||
latent_key (str): key of the latent to be used
|
||||
normalization (dict): normalization stats
|
||||
pretrained_slat_dec (str): name of the pretrained slat decoder
|
||||
slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
|
||||
slat_dec_ckpt (str): name of the slat decoder checkpoint
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip
|
||||
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
resolution: int,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
max_tokens: int = 32768,
|
||||
normalization: Optional[dict] = None,
|
||||
pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
|
||||
slat_dec_path: Optional[str] = None,
|
||||
slat_dec_ckpt: Optional[str] = None,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[list] = None,
|
||||
):
|
||||
super().__init__(
|
||||
roots,
|
||||
min_aesthetic_score=min_aesthetic_score,
|
||||
max_tokens=max_tokens,
|
||||
latent_key='shape_latent',
|
||||
normalization=normalization,
|
||||
pretrained_slat_dec=pretrained_slat_dec,
|
||||
slat_dec_path=slat_dec_path,
|
||||
slat_dec_ckpt=slat_dec_ckpt,
|
||||
skip_list=skip_list,
|
||||
skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
|
||||
)
|
||||
self.resolution = resolution
|
||||
|
||||
|
||||
class ImageConditionedSLatShape(ImageConditionedMixin, SLatShape):
|
||||
"""
|
||||
Image conditioned structured latent for shape generation
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class SLatShapeView(SLatShapeVisMixin, SLat):
|
||||
"""
|
||||
View-based structured latent for shape generation.
|
||||
|
||||
Data format: {sha256}/view{XX}.npz where each npz contains 'coords' and 'feats' keys.
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
resolution (int): resolution of the shape
|
||||
min_aesthetic_score (float): minimum aesthetic score
|
||||
max_tokens (int): maximum number of tokens
|
||||
num_views (int): Number of views to use (0 to num_views-1). Default is 2.
|
||||
normalization (dict): normalization stats
|
||||
pretrained_slat_dec (str): name of the pretrained slat decoder
|
||||
slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
|
||||
slat_dec_ckpt (str): name of the slat decoder checkpoint
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip
|
||||
skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
resolution: int,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
max_tokens: int = 32768,
|
||||
num_views: int = 2,
|
||||
normalization: Optional[dict] = None,
|
||||
pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
|
||||
slat_dec_path: Optional[str] = None,
|
||||
slat_dec_ckpt: Optional[str] = None,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[list] = None,
|
||||
):
|
||||
self.normalization = normalization
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.max_tokens = max_tokens
|
||||
self.num_views = num_views
|
||||
self.latent_key = 'shape_latent'
|
||||
self.value_range = (0, 1)
|
||||
|
||||
# Initialize parent with SLatVisMixin parameters
|
||||
from .components import StandardDatasetBase
|
||||
SLatVisMixin.__init__(
|
||||
self,
|
||||
roots,
|
||||
pretrained_slat_dec=pretrained_slat_dec,
|
||||
slat_dec_path=slat_dec_path,
|
||||
slat_dec_ckpt=slat_dec_ckpt,
|
||||
)
|
||||
StandardDatasetBase.__init__(self, roots, skip_list=skip_list, skip_aesthetic_score_datasets=skip_aesthetic_score_datasets)
|
||||
|
||||
self.resolution = resolution
|
||||
|
||||
# Calculate loads for load balancing
|
||||
self.loads = []
|
||||
for _, sha256, _ in self.instances:
|
||||
if f'{self.latent_key}_tokens' in self.metadata.columns:
|
||||
try:
|
||||
self.loads.append(self.metadata.loc[sha256, f'{self.latent_key}_tokens'])
|
||||
except:
|
||||
self.loads.append(self.max_tokens)
|
||||
else:
|
||||
self.loads.append(self.max_tokens)
|
||||
|
||||
if self.normalization is not None:
|
||||
self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1)
|
||||
self.std = torch.tensor(self.normalization['std']).reshape(1, -1)
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
# View-based shape_latent uses columns like shape_latent_view00_encoded, shape_latent_view01_encoded, etc.
|
||||
required_view_cols = [f'shape_latent_view{i:02d}_encoded' for i in range(self.num_views)]
|
||||
existing_view_cols = [col for col in required_view_cols if col in metadata.columns]
|
||||
|
||||
if existing_view_cols:
|
||||
# Filter rows where all required views are encoded
|
||||
# Note: NaN should be treated as False, so use == True for explicit comparison
|
||||
has_all_views = (metadata[existing_view_cols] == True).all(axis=1)
|
||||
metadata = metadata[has_all_views]
|
||||
stats[f'With {self.num_views} view latents'] = len(metadata)
|
||||
else:
|
||||
# Fallback: check shape_latent_encoded column
|
||||
if f'{self.latent_key}_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata[f'{self.latent_key}_encoded'] == True]
|
||||
stats['With latent'] = len(metadata)
|
||||
|
||||
# Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
|
||||
skip_aesthetic = (
|
||||
(dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
|
||||
('aesthetic_score' not in metadata.columns)
|
||||
)
|
||||
if skip_aesthetic:
|
||||
stats[f'Aesthetic score check skipped'] = len(metadata)
|
||||
else:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
|
||||
# Filter by max_tokens if column exists
|
||||
tokens_col = f'{self.latent_key}_tokens'
|
||||
if tokens_col in metadata.columns:
|
||||
metadata = metadata[metadata[tokens_col] <= self.max_tokens]
|
||||
stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
|
||||
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
# View-based format: directory with view{XX}.npz files
|
||||
latent_dir = os.path.join(root[self.latent_key], instance)
|
||||
|
||||
# Randomly select a view from the configured range
|
||||
view_idx = np.random.randint(0, self.num_views)
|
||||
view_file = f'view{view_idx:02d}.npz'
|
||||
|
||||
# Store view info for ViewImageConditionedMixin
|
||||
self._current_view_idx = view_idx
|
||||
self._current_latent_dir = latent_dir
|
||||
|
||||
data = np.load(os.path.join(latent_dir, view_file))
|
||||
coords = torch.tensor(data['coords']).int()
|
||||
feats = torch.tensor(data['feats']).float()
|
||||
if self.normalization is not None:
|
||||
feats = (feats - self.mean) / self.std
|
||||
return {
|
||||
'coords': coords,
|
||||
'feats': feats,
|
||||
'view_idx': view_idx,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch, split_size=None):
|
||||
if split_size is None:
|
||||
group_idx = [list(range(len(batch)))]
|
||||
else:
|
||||
group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
|
||||
packs = []
|
||||
for group in group_idx:
|
||||
sub_batch = [batch[i] for i in group]
|
||||
pack = {}
|
||||
coords = []
|
||||
feats = []
|
||||
layout = []
|
||||
start = 0
|
||||
for i, b in enumerate(sub_batch):
|
||||
coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
|
||||
feats.append(b['feats'])
|
||||
layout.append(slice(start, start + b['coords'].shape[0]))
|
||||
start += b['coords'].shape[0]
|
||||
coords = torch.cat(coords)
|
||||
feats = torch.cat(feats)
|
||||
pack['x_0'] = SparseTensor(
|
||||
coords=coords,
|
||||
feats=feats,
|
||||
)
|
||||
pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]])
|
||||
pack['x_0'].register_spatial_cache('layout', layout)
|
||||
|
||||
# collate other data
|
||||
keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']]
|
||||
for k in keys:
|
||||
if isinstance(sub_batch[0][k], torch.Tensor):
|
||||
pack[k] = torch.stack([b[k] for b in sub_batch])
|
||||
elif isinstance(sub_batch[0][k], list):
|
||||
pack[k] = sum([b[k] for b in sub_batch], [])
|
||||
else:
|
||||
pack[k] = [b[k] for b in sub_batch]
|
||||
|
||||
packs.append(pack)
|
||||
|
||||
if split_size is None:
|
||||
return packs[0]
|
||||
return packs
|
||||
|
||||
|
||||
class ViewImageConditionedSLatShapeView(ViewImageConditionedMixin, SLatShapeView):
|
||||
"""
|
||||
Image-conditioned view-based structured latent for shape generation.
|
||||
|
||||
Loads shape_latent from {sha256}/view{XX}.npz format and pairs with
|
||||
corresponding view from render_cond.
|
||||
|
||||
Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,666 @@
|
||||
import os
|
||||
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
|
||||
import json
|
||||
from typing import *
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
import utils3d
|
||||
from .. import models
|
||||
from .components import StandardDatasetBase, ImageConditionedMixin, ViewImageConditionedMixin
|
||||
from ..modules.sparse import SparseTensor, sparse_cat
|
||||
from ..representations import MeshWithVoxel
|
||||
from ..renderers import PbrMeshRenderer, EnvMap
|
||||
from ..utils.data_utils import load_balanced_group_indices
|
||||
from ..utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
|
||||
|
||||
|
||||
class SLatPbrVisMixin:
|
||||
def __init__(
|
||||
self,
|
||||
*args,
|
||||
pretrained_pbr_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
|
||||
pbr_slat_dec_path: Optional[str] = None,
|
||||
pbr_slat_dec_ckpt: Optional[str] = None,
|
||||
pretrained_shape_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
|
||||
shape_slat_dec_path: Optional[str] = None,
|
||||
shape_slat_dec_ckpt: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.pbr_slat_dec = None
|
||||
self.pretrained_pbr_slat_dec = pretrained_pbr_slat_dec
|
||||
self.pbr_slat_dec_path = pbr_slat_dec_path
|
||||
self.pbr_slat_dec_ckpt = pbr_slat_dec_ckpt
|
||||
self.shape_slat_dec = None
|
||||
self.pretrained_shape_slat_dec = pretrained_shape_slat_dec
|
||||
self.shape_slat_dec_path = shape_slat_dec_path
|
||||
self.shape_slat_dec_ckpt = shape_slat_dec_ckpt
|
||||
|
||||
def _loading_slat_dec(self):
|
||||
if self.pbr_slat_dec is not None and self.shape_slat_dec is not None:
|
||||
return
|
||||
if self.pbr_slat_dec_path is not None:
|
||||
cfg = json.load(open(os.path.join(self.pbr_slat_dec_path, 'config.json'), 'r'))
|
||||
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
|
||||
ckpt_path = os.path.join(self.pbr_slat_dec_path, 'ckpts', f'decoder_{self.pbr_slat_dec_ckpt}.pt')
|
||||
decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
|
||||
else:
|
||||
decoder = models.from_pretrained(self.pretrained_pbr_slat_dec)
|
||||
self.pbr_slat_dec = decoder.cuda().eval()
|
||||
|
||||
if self.shape_slat_dec_path is not None:
|
||||
cfg = json.load(open(os.path.join(self.shape_slat_dec_path, 'config.json'), 'r'))
|
||||
decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
|
||||
ckpt_path = os.path.join(self.shape_slat_dec_path, 'ckpts', f'decoder_{self.shape_slat_dec_ckpt}.pt')
|
||||
decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
|
||||
else:
|
||||
decoder = models.from_pretrained(self.pretrained_shape_slat_dec)
|
||||
decoder.set_resolution(self.resolution)
|
||||
self.shape_slat_dec = decoder.cuda().eval()
|
||||
|
||||
def _delete_slat_dec(self):
|
||||
del self.pbr_slat_dec
|
||||
self.pbr_slat_dec = None
|
||||
del self.shape_slat_dec
|
||||
self.shape_slat_dec = None
|
||||
|
||||
@torch.no_grad()
|
||||
def decode_latent(self, z, shape_z, batch_size=4):
|
||||
self._loading_slat_dec()
|
||||
reps = []
|
||||
if self.shape_slat_normalization is not None:
|
||||
shape_z = shape_z * self.shape_slat_std.to(z.device) + self.shape_slat_mean.to(z.device)
|
||||
if self.pbr_slat_normalization is not None:
|
||||
z = z * self.pbr_slat_std.to(z.device) + self.pbr_slat_mean.to(z.device)
|
||||
for i in range(0, z.shape[0], batch_size):
|
||||
mesh, subs = self.shape_slat_dec(shape_z[i:i+batch_size], return_subs=True)
|
||||
vox = self.pbr_slat_dec(z[i:i+batch_size], guide_subs=subs) * 0.5 + 0.5
|
||||
reps.extend([
|
||||
MeshWithVoxel(
|
||||
m.vertices, m.faces,
|
||||
origin = [-0.5, -0.5, -0.5],
|
||||
voxel_size = 1 / self.resolution,
|
||||
coords = v.coords[:, 1:],
|
||||
attrs = v.feats,
|
||||
voxel_shape = torch.Size([*v.shape, *v.spatial_shape]),
|
||||
layout = self.layout,
|
||||
)
|
||||
for m, v in zip(mesh, vox)
|
||||
])
|
||||
self._delete_slat_dec()
|
||||
return reps
|
||||
|
||||
@torch.no_grad()
|
||||
def visualize_sample(self, sample: dict):
|
||||
shape_z = sample['concat_cond'].cuda()
|
||||
z = sample['x_0'].cuda()
|
||||
reps = self.decode_latent(z, shape_z)
|
||||
|
||||
# Extract camera parameters for GT view rendering (if available)
|
||||
camera_angle_x = sample.get('camera_angle_x')
|
||||
camera_distance = sample.get('camera_distance')
|
||||
mesh_scale = sample.get('mesh_scale')
|
||||
has_gt_camera = (
|
||||
camera_angle_x is not None and
|
||||
camera_distance is not None and
|
||||
mesh_scale is not None
|
||||
)
|
||||
|
||||
# build camera
|
||||
yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
|
||||
yaw_offset = -16 / 180 * np.pi
|
||||
yaw = [y + yaw_offset for y in yaw]
|
||||
pitch = [20 / 180 * np.pi for _ in range(4)]
|
||||
exts, ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
|
||||
|
||||
# render
|
||||
renderer = PbrMeshRenderer()
|
||||
renderer.rendering_options.resolution = 512
|
||||
renderer.rendering_options.near = 1
|
||||
renderer.rendering_options.far = 100
|
||||
renderer.rendering_options.ssaa = 2
|
||||
renderer.rendering_options.peel_layers = 8
|
||||
envmap = EnvMap(torch.tensor(
|
||||
cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
|
||||
dtype=torch.float32, device='cuda'
|
||||
))
|
||||
|
||||
images = {}
|
||||
gt_view_images = {}
|
||||
for i, representation in enumerate(reps):
|
||||
# Validate mesh data before rasterization (same as shape training)
|
||||
verts = representation.vertices
|
||||
faces = representation.faces
|
||||
if verts.shape[0] == 0 or faces.shape[0] == 0:
|
||||
print(f"[visualize_sample] Warning: sample {i} has empty mesh, skipping")
|
||||
continue
|
||||
if faces.max() >= verts.shape[0]:
|
||||
print(f"[visualize_sample] Warning: sample {i} has out-of-bound face indices "
|
||||
f"(max face idx={faces.max().item()}, num verts={verts.shape[0]}), skipping")
|
||||
continue
|
||||
if torch.isnan(verts).any() or torch.isinf(verts).any():
|
||||
print(f"[visualize_sample] Warning: sample {i} has NaN/Inf vertices, skipping")
|
||||
continue
|
||||
|
||||
image = {}
|
||||
tile = [2, 2]
|
||||
try:
|
||||
for j, (ext, intr) in enumerate(zip(exts, ints)):
|
||||
res = renderer.render(representation, ext, intr, envmap=envmap)
|
||||
for k, v in res.items():
|
||||
if k not in images:
|
||||
images[k] = []
|
||||
if k not in image:
|
||||
image[k] = torch.zeros(3, 1024, 1024).cuda()
|
||||
image[k][:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = v
|
||||
for k in images.keys():
|
||||
images[k].append(image[k])
|
||||
except RuntimeError as e:
|
||||
print(f"[visualize_sample] Warning: render failed for sample {i}: {e}")
|
||||
try:
|
||||
torch.cuda.synchronize()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
|
||||
# Render GT camera view
|
||||
# Must scale mesh vertices by / mesh_scale to match ProjGrid's projection space.
|
||||
# ProjGrid maps [-1,1]^3 -> / scale / 2 -> [-0.5/s, 0.5/s]^3
|
||||
# Mesh vertices in [-0.5, 0.5]^3 -> / scale -> [-0.5/s, 0.5/s]^3 (equivalent)
|
||||
if has_gt_camera:
|
||||
try:
|
||||
scale = mesh_scale[i].item()
|
||||
distance = camera_distance[i].item()
|
||||
fov = camera_angle_x[i].item()
|
||||
device = representation.vertices.device
|
||||
|
||||
# Scale mesh and voxel to match ProjGrid's projection space
|
||||
scaled_rep = MeshWithVoxel(
|
||||
vertices=representation.vertices / scale,
|
||||
faces=representation.faces,
|
||||
origin=(representation.origin / scale).tolist(),
|
||||
voxel_size=representation.voxel_size / scale,
|
||||
coords=representation.coords,
|
||||
attrs=representation.attrs,
|
||||
voxel_shape=representation.voxel_shape,
|
||||
layout=representation.layout,
|
||||
)
|
||||
|
||||
cam_pos = torch.tensor([0.0, 0.0, distance], device=device)
|
||||
look_at = torch.tensor([0.0, 0.0, 0.0], device=device)
|
||||
cam_up = torch.tensor([0.0, 1.0, 0.0], device=device)
|
||||
|
||||
gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up)
|
||||
gt_int = utils3d.torch.intrinsics_from_fov_xy(
|
||||
torch.tensor(fov, device=device),
|
||||
torch.tensor(fov, device=device)
|
||||
)
|
||||
gt_ext = gt_ext.to(device)
|
||||
gt_int = gt_int.to(device)
|
||||
|
||||
# Update near/far for the smaller scaled mesh
|
||||
mesh_half_size = 0.5 / scale
|
||||
renderer.rendering_options.near = max(0.01, distance - mesh_half_size - 0.5)
|
||||
renderer.rendering_options.far = distance + mesh_half_size + 0.5
|
||||
|
||||
gt_res = renderer.render(scaled_rep, gt_ext, gt_int, envmap=envmap)
|
||||
for k, v in gt_res.items():
|
||||
gt_key = f'gt_view_{k}'
|
||||
if gt_key not in gt_view_images:
|
||||
gt_view_images[gt_key] = []
|
||||
gt_view_images[gt_key].append(v)
|
||||
except RuntimeError as e:
|
||||
print(f"[visualize_sample] Warning: GT view render failed for sample {i}: {e}")
|
||||
try:
|
||||
torch.cuda.synchronize()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for k in images.keys():
|
||||
images[k] = torch.stack(images[k], dim=0)
|
||||
|
||||
for k, v in gt_view_images.items():
|
||||
images[k] = torch.stack(v)
|
||||
|
||||
return images
|
||||
|
||||
|
||||
class SLatPbr(SLatPbrVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
structured latent for sparse voxel pbr dataset
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
latent_key (str): key of the latent to be used
|
||||
min_aesthetic_score (float): minimum aesthetic score
|
||||
normalization (dict): normalization stats
|
||||
resolution (int): resolution of decoded sparse voxel
|
||||
attrs (list): attributes to be decoded
|
||||
pretained_slat_dec (str): name of the pretrained slat decoder
|
||||
slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
|
||||
slat_dec_ckpt (str): name of the slat decoder checkpoint
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
resolution: int,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
max_tokens: int = 32768,
|
||||
full_pbr: bool = False,
|
||||
pbr_slat_normalization: Optional[dict] = None,
|
||||
shape_slat_normalization: Optional[dict] = None,
|
||||
attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
|
||||
pretrained_pbr_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
|
||||
pbr_slat_dec_path: Optional[str] = None,
|
||||
pbr_slat_dec_ckpt: Optional[str] = None,
|
||||
pretrained_shape_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
|
||||
shape_slat_dec_path: Optional[str] = None,
|
||||
shape_slat_dec_ckpt: Optional[str] = None,
|
||||
**kwargs
|
||||
):
|
||||
self.resolution = resolution
|
||||
self.pbr_slat_normalization = pbr_slat_normalization
|
||||
self.shape_slat_normalization = shape_slat_normalization
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.max_tokens = max_tokens
|
||||
self.full_pbr = full_pbr
|
||||
self.value_range = (0, 1)
|
||||
|
||||
super().__init__(
|
||||
roots,
|
||||
pretrained_pbr_slat_dec=pretrained_pbr_slat_dec,
|
||||
pbr_slat_dec_path=pbr_slat_dec_path,
|
||||
pbr_slat_dec_ckpt=pbr_slat_dec_ckpt,
|
||||
pretrained_shape_slat_dec=pretrained_shape_slat_dec,
|
||||
shape_slat_dec_path=shape_slat_dec_path,
|
||||
shape_slat_dec_ckpt=shape_slat_dec_ckpt,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
self.loads = [self.metadata.loc[sha256, 'pbr_latent_tokens'] for _, sha256, _ in self.instances]
|
||||
|
||||
if self.pbr_slat_normalization is not None:
|
||||
self.pbr_slat_mean = torch.tensor(self.pbr_slat_normalization['mean']).reshape(1, -1)
|
||||
self.pbr_slat_std = torch.tensor(self.pbr_slat_normalization['std']).reshape(1, -1)
|
||||
|
||||
if self.shape_slat_normalization is not None:
|
||||
self.shape_slat_mean = torch.tensor(self.shape_slat_normalization['mean']).reshape(1, -1)
|
||||
self.shape_slat_std = torch.tensor(self.shape_slat_normalization['std']).reshape(1, -1)
|
||||
|
||||
self.attrs = attrs
|
||||
self.channels = {
|
||||
'base_color': 3,
|
||||
'metallic': 1,
|
||||
'roughness': 1,
|
||||
'emissive': 3,
|
||||
'alpha': 1,
|
||||
}
|
||||
self.layout = {}
|
||||
start = 0
|
||||
for attr in attrs:
|
||||
self.layout[attr] = slice(start, start + self.channels[attr])
|
||||
start += self.channels[attr]
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
metadata = metadata[metadata['pbr_latent_encoded'] == True]
|
||||
stats['With PBR latent'] = len(metadata)
|
||||
metadata = metadata[metadata['shape_latent_encoded'] == True]
|
||||
stats['With shape latent'] = len(metadata)
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
metadata = metadata[metadata['pbr_latent_tokens'] <= self.max_tokens]
|
||||
stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
|
||||
if self.full_pbr:
|
||||
metadata = metadata[metadata['num_basecolor_tex'] > 0]
|
||||
metadata = metadata[metadata['num_metallic_tex'] > 0]
|
||||
metadata = metadata[metadata['num_roughness_tex'] > 0]
|
||||
stats['Full PBR'] = len(metadata)
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
# PBR latent
|
||||
data = np.load(os.path.join(root['pbr_latent'], f'{instance}.npz'))
|
||||
coords = torch.tensor(data['coords']).int()
|
||||
coords = torch.cat([torch.zeros_like(coords)[:, :1], coords], dim=1)
|
||||
feats = torch.tensor(data['feats']).float()
|
||||
if self.pbr_slat_normalization is not None:
|
||||
feats = (feats - self.pbr_slat_mean) / self.pbr_slat_std
|
||||
pbr_z = SparseTensor(feats, coords)
|
||||
|
||||
# Shape latent
|
||||
data = np.load(os.path.join(root['shape_latent'], f'{instance}.npz'))
|
||||
coords = torch.tensor(data['coords']).int()
|
||||
coords = torch.cat([torch.zeros_like(coords)[:, :1], coords], dim=1)
|
||||
feats = torch.tensor(data['feats']).float()
|
||||
if self.shape_slat_normalization is not None:
|
||||
feats = (feats - self.shape_slat_mean) / self.shape_slat_std
|
||||
shape_z = SparseTensor(feats, coords)
|
||||
|
||||
assert torch.equal(shape_z.coords, pbr_z.coords), \
|
||||
f"Shape latent and PBR latent have different coordinates: {shape_z.coords.shape} vs {pbr_z.coords.shape}"
|
||||
|
||||
return {
|
||||
'x_0': pbr_z,
|
||||
'concat_cond': shape_z,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch, split_size=None):
|
||||
if split_size is None:
|
||||
group_idx = [list(range(len(batch)))]
|
||||
else:
|
||||
group_idx = load_balanced_group_indices([b['x_0'].feats.shape[0] for b in batch], split_size)
|
||||
packs = []
|
||||
for group in group_idx:
|
||||
sub_batch = [batch[i] for i in group]
|
||||
pack = {}
|
||||
|
||||
keys = [k for k in sub_batch[0].keys()]
|
||||
for k in keys:
|
||||
if isinstance(sub_batch[0][k], torch.Tensor):
|
||||
pack[k] = torch.stack([b[k] for b in sub_batch])
|
||||
elif isinstance(sub_batch[0][k], SparseTensor):
|
||||
pack[k] = sparse_cat([b[k] for b in sub_batch], dim=0)
|
||||
elif isinstance(sub_batch[0][k], list):
|
||||
pack[k] = sum([b[k] for b in sub_batch], [])
|
||||
else:
|
||||
pack[k] = [b[k] for b in sub_batch]
|
||||
|
||||
packs.append(pack)
|
||||
|
||||
if split_size is None:
|
||||
return packs[0]
|
||||
return packs
|
||||
|
||||
|
||||
class ImageConditionedSLatPbr(ImageConditionedMixin, SLatPbr):
|
||||
"""
|
||||
Image conditioned structured latent dataset
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class SLatPbrView(SLatPbrVisMixin, StandardDatasetBase):
|
||||
"""
|
||||
View-based structured latent for PBR/texture generation with view-aligned projection.
|
||||
|
||||
Data format:
|
||||
PBR latent: {sha256}/view{XX}.npz (coords + feats)
|
||||
Shape latent: {sha256}/view{XX}.npz (coords + feats, from shape_latent_view dir)
|
||||
|
||||
Each view's PBR latent and Shape latent share the same sparse coordinates.
|
||||
|
||||
Args:
|
||||
roots (str): path to the dataset
|
||||
resolution (int): resolution of decoded sparse voxel
|
||||
min_aesthetic_score (float): minimum aesthetic score
|
||||
max_tokens (int): maximum number of tokens
|
||||
num_views (int): Number of views to use (0 to num_views-1). Default is 2.
|
||||
full_pbr (bool): Whether to require full PBR textures
|
||||
pbr_slat_normalization (dict): normalization stats for PBR latent
|
||||
shape_slat_normalization (dict): normalization stats for shape latent
|
||||
attrs (list): PBR attributes to decode
|
||||
pretrained_pbr_slat_dec (str): pretrained PBR decoder name
|
||||
pretrained_shape_slat_dec (str): pretrained shape decoder name
|
||||
skip_list (str, optional): path to a file containing sha256 hashes to skip
|
||||
skip_aesthetic_score_datasets (list, optional): datasets to skip aesthetic score check
|
||||
"""
|
||||
def __init__(self,
|
||||
roots: str,
|
||||
*,
|
||||
resolution: int,
|
||||
min_aesthetic_score: float = 5.0,
|
||||
max_tokens: int = 32768,
|
||||
num_views: int = 2,
|
||||
full_pbr: bool = False,
|
||||
pbr_slat_normalization: Optional[dict] = None,
|
||||
shape_slat_normalization: Optional[dict] = None,
|
||||
attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
|
||||
pretrained_pbr_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
|
||||
pbr_slat_dec_path: Optional[str] = None,
|
||||
pbr_slat_dec_ckpt: Optional[str] = None,
|
||||
pretrained_shape_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
|
||||
shape_slat_dec_path: Optional[str] = None,
|
||||
shape_slat_dec_ckpt: Optional[str] = None,
|
||||
skip_list: Optional[str] = None,
|
||||
skip_aesthetic_score_datasets: Optional[list] = None,
|
||||
):
|
||||
self.resolution = resolution
|
||||
self.pbr_slat_normalization = pbr_slat_normalization
|
||||
self.shape_slat_normalization = shape_slat_normalization
|
||||
self.min_aesthetic_score = min_aesthetic_score
|
||||
self.max_tokens = max_tokens
|
||||
self.num_views = num_views
|
||||
self.full_pbr = full_pbr
|
||||
self.value_range = (0, 1)
|
||||
self.skip_aesthetic_score_datasets = set(skip_aesthetic_score_datasets or [])
|
||||
|
||||
# Initialize visualization mixin
|
||||
SLatPbrVisMixin.__init__(
|
||||
self,
|
||||
roots,
|
||||
pretrained_pbr_slat_dec=pretrained_pbr_slat_dec,
|
||||
pbr_slat_dec_path=pbr_slat_dec_path,
|
||||
pbr_slat_dec_ckpt=pbr_slat_dec_ckpt,
|
||||
pretrained_shape_slat_dec=pretrained_shape_slat_dec,
|
||||
shape_slat_dec_path=shape_slat_dec_path,
|
||||
shape_slat_dec_ckpt=shape_slat_dec_ckpt,
|
||||
)
|
||||
StandardDatasetBase.__init__(
|
||||
self, roots,
|
||||
skip_list=skip_list,
|
||||
skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
|
||||
)
|
||||
|
||||
# Calculate loads for load balancing
|
||||
self.loads = []
|
||||
for _, sha256, _ in self.instances:
|
||||
if 'pbr_latent_tokens' in self.metadata.columns:
|
||||
try:
|
||||
self.loads.append(self.metadata.loc[sha256, 'pbr_latent_tokens'])
|
||||
except:
|
||||
self.loads.append(self.max_tokens)
|
||||
else:
|
||||
self.loads.append(self.max_tokens)
|
||||
|
||||
if self.pbr_slat_normalization is not None:
|
||||
self.pbr_slat_mean = torch.tensor(self.pbr_slat_normalization['mean']).reshape(1, -1)
|
||||
self.pbr_slat_std = torch.tensor(self.pbr_slat_normalization['std']).reshape(1, -1)
|
||||
|
||||
if self.shape_slat_normalization is not None:
|
||||
self.shape_slat_mean = torch.tensor(self.shape_slat_normalization['mean']).reshape(1, -1)
|
||||
self.shape_slat_std = torch.tensor(self.shape_slat_normalization['std']).reshape(1, -1)
|
||||
|
||||
self.attrs = attrs
|
||||
self.channels = {
|
||||
'base_color': 3,
|
||||
'metallic': 1,
|
||||
'roughness': 1,
|
||||
'emissive': 3,
|
||||
'alpha': 1,
|
||||
}
|
||||
self.layout = {}
|
||||
start = 0
|
||||
for attr in attrs:
|
||||
self.layout[attr] = slice(start, start + self.channels[attr])
|
||||
start += self.channels[attr]
|
||||
|
||||
def filter_metadata(self, metadata, dataset_name=None):
|
||||
stats = {}
|
||||
# View-based PBR latent uses columns like pbr_latent_view00_encoded, etc.
|
||||
required_pbr_view_cols = [f'pbr_latent_view{i:02d}_encoded' for i in range(self.num_views)]
|
||||
existing_pbr_view_cols = [col for col in required_pbr_view_cols if col in metadata.columns]
|
||||
|
||||
if existing_pbr_view_cols:
|
||||
has_all_pbr_views = (metadata[existing_pbr_view_cols] == True).all(axis=1)
|
||||
metadata = metadata[has_all_pbr_views]
|
||||
stats[f'With {self.num_views} PBR view latents'] = len(metadata)
|
||||
else:
|
||||
# Fallback: check pbr_latent_encoded
|
||||
if 'pbr_latent_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata['pbr_latent_encoded'] == True]
|
||||
stats['With PBR latent'] = len(metadata)
|
||||
|
||||
# Also require shape latent views
|
||||
required_shape_view_cols = [f'shape_latent_view{i:02d}_encoded' for i in range(self.num_views)]
|
||||
existing_shape_view_cols = [col for col in required_shape_view_cols if col in metadata.columns]
|
||||
|
||||
if existing_shape_view_cols:
|
||||
has_all_shape_views = (metadata[existing_shape_view_cols] == True).all(axis=1)
|
||||
metadata = metadata[has_all_shape_views]
|
||||
stats[f'With {self.num_views} shape view latents'] = len(metadata)
|
||||
else:
|
||||
if 'shape_latent_encoded' in metadata.columns:
|
||||
metadata = metadata[metadata['shape_latent_encoded'] == True]
|
||||
stats['With shape latent'] = len(metadata)
|
||||
|
||||
# Skip aesthetic score check for specified datasets
|
||||
skip_aesthetic = (
|
||||
(dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
|
||||
('aesthetic_score' not in metadata.columns)
|
||||
)
|
||||
if skip_aesthetic:
|
||||
stats[f'Aesthetic score check skipped'] = len(metadata)
|
||||
else:
|
||||
metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
|
||||
stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
|
||||
|
||||
# Filter by max_tokens if column exists
|
||||
if 'pbr_latent_tokens' in metadata.columns:
|
||||
metadata = metadata[metadata['pbr_latent_tokens'] <= self.max_tokens]
|
||||
stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
|
||||
|
||||
if self.full_pbr:
|
||||
if 'num_basecolor_tex' in metadata.columns:
|
||||
metadata = metadata[metadata['num_basecolor_tex'] > 0]
|
||||
if 'num_metallic_tex' in metadata.columns:
|
||||
metadata = metadata[metadata['num_metallic_tex'] > 0]
|
||||
if 'num_roughness_tex' in metadata.columns:
|
||||
metadata = metadata[metadata['num_roughness_tex'] > 0]
|
||||
stats['Full PBR'] = len(metadata)
|
||||
|
||||
return metadata, stats
|
||||
|
||||
def get_instance(self, root, instance):
|
||||
# Randomly select a view from the configured range
|
||||
view_idx = np.random.randint(0, self.num_views)
|
||||
view_file = f'view{view_idx:02d}.npz'
|
||||
|
||||
# Store view info for ViewImageConditionedMixin
|
||||
self._current_view_idx = view_idx
|
||||
|
||||
# Load PBR latent for this view
|
||||
pbr_latent_dir = os.path.join(root['pbr_latent'], instance)
|
||||
self._current_latent_dir = pbr_latent_dir
|
||||
|
||||
data = np.load(os.path.join(pbr_latent_dir, view_file))
|
||||
pbr_coords = torch.tensor(data['coords']).int()
|
||||
pbr_feats = torch.tensor(data['feats']).float()
|
||||
if self.pbr_slat_normalization is not None:
|
||||
pbr_feats = (pbr_feats - self.pbr_slat_mean) / self.pbr_slat_std
|
||||
|
||||
# Load Shape latent for this view (as concat_cond)
|
||||
shape_latent_dir = os.path.join(root['shape_latent'], instance)
|
||||
data = np.load(os.path.join(shape_latent_dir, view_file))
|
||||
shape_coords = torch.tensor(data['coords']).int()
|
||||
shape_feats = torch.tensor(data['feats']).float()
|
||||
if self.shape_slat_normalization is not None:
|
||||
shape_feats = (shape_feats - self.shape_slat_mean) / self.shape_slat_std
|
||||
|
||||
# Verify coordinates match
|
||||
assert torch.equal(pbr_coords, shape_coords), \
|
||||
f"PBR and shape latent coordinates mismatch for {instance}/view{view_idx:02d}"
|
||||
|
||||
return {
|
||||
'coords': pbr_coords,
|
||||
'pbr_feats': pbr_feats,
|
||||
'shape_feats': shape_feats,
|
||||
'view_idx': view_idx,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def collate_fn(batch, split_size=None):
|
||||
if split_size is None:
|
||||
group_idx = [list(range(len(batch)))]
|
||||
else:
|
||||
group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
|
||||
packs = []
|
||||
for group in group_idx:
|
||||
sub_batch = [batch[i] for i in group]
|
||||
pack = {}
|
||||
|
||||
# Build x_0 (PBR latent) and concat_cond (shape latent) as SparseTensors
|
||||
coords_list = []
|
||||
pbr_feats_list = []
|
||||
shape_feats_list = []
|
||||
layout = []
|
||||
start = 0
|
||||
for i, b in enumerate(sub_batch):
|
||||
batch_coords = torch.cat([
|
||||
torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32),
|
||||
b['coords']
|
||||
], dim=-1)
|
||||
coords_list.append(batch_coords)
|
||||
pbr_feats_list.append(b['pbr_feats'])
|
||||
shape_feats_list.append(b['shape_feats'])
|
||||
layout.append(slice(start, start + b['coords'].shape[0]))
|
||||
start += b['coords'].shape[0]
|
||||
|
||||
all_coords = torch.cat(coords_list)
|
||||
|
||||
# x_0: PBR latent
|
||||
pack['x_0'] = SparseTensor(
|
||||
coords=all_coords,
|
||||
feats=torch.cat(pbr_feats_list),
|
||||
)
|
||||
pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['pbr_feats'].shape[1:]])
|
||||
pack['x_0'].register_spatial_cache('layout', layout)
|
||||
|
||||
# concat_cond: Shape latent (same coordinates)
|
||||
pack['concat_cond'] = SparseTensor(
|
||||
coords=all_coords.clone(),
|
||||
feats=torch.cat(shape_feats_list),
|
||||
)
|
||||
pack['concat_cond']._shape = torch.Size([len(group), *sub_batch[0]['shape_feats'].shape[1:]])
|
||||
pack['concat_cond'].register_spatial_cache('layout', layout)
|
||||
|
||||
# collate other data (excluding already handled fields)
|
||||
skip_keys = {'coords', 'pbr_feats', 'shape_feats'}
|
||||
keys = [k for k in sub_batch[0].keys() if k not in skip_keys]
|
||||
for k in keys:
|
||||
if isinstance(sub_batch[0][k], torch.Tensor):
|
||||
pack[k] = torch.stack([b[k] for b in sub_batch])
|
||||
elif isinstance(sub_batch[0][k], list):
|
||||
pack[k] = sum([b[k] for b in sub_batch], [])
|
||||
else:
|
||||
pack[k] = [b[k] for b in sub_batch]
|
||||
|
||||
packs.append(pack)
|
||||
|
||||
if split_size is None:
|
||||
return packs[0]
|
||||
return packs
|
||||
|
||||
|
||||
class ViewImageConditionedSLatPbrView(ViewImageConditionedMixin, SLatPbrView):
|
||||
"""
|
||||
Image-conditioned view-based structured latent for PBR/texture generation
|
||||
with view-aligned projection.
|
||||
|
||||
Loads PBR latent and shape latent from {sha256}/view{XX}.npz format and pairs
|
||||
with corresponding view from render_cond.
|
||||
|
||||
Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json
|
||||
and provides camera parameters for 3D-to-2D projection.
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,78 @@
|
||||
import importlib
|
||||
|
||||
__attributes = {
|
||||
# Sparse Structure
|
||||
'SparseStructureEncoder': 'sparse_structure_vae',
|
||||
'SparseStructureDecoder': 'sparse_structure_vae',
|
||||
'SparseStructureFlowModel': 'sparse_structure_flow',
|
||||
|
||||
# SLat Generation
|
||||
'SLatFlowModel': 'structured_latent_flow',
|
||||
'ElasticSLatFlowModel': 'structured_latent_flow',
|
||||
|
||||
# SC-VAEs
|
||||
'SparseUnetVaeEncoder': 'sc_vaes.sparse_unet_vae',
|
||||
'SparseUnetVaeDecoder': 'sc_vaes.sparse_unet_vae',
|
||||
'FlexiDualGridVaeEncoder': 'sc_vaes.fdg_vae',
|
||||
'FlexiDualGridVaeDecoder': 'sc_vaes.fdg_vae'
|
||||
}
|
||||
|
||||
__submodules = []
|
||||
|
||||
__all__ = list(__attributes.keys()) + __submodules
|
||||
|
||||
def __getattr__(name):
|
||||
if name not in globals():
|
||||
if name in __attributes:
|
||||
module_name = __attributes[name]
|
||||
module = importlib.import_module(f".{module_name}", __name__)
|
||||
globals()[name] = getattr(module, name)
|
||||
elif name in __submodules:
|
||||
module = importlib.import_module(f".{name}", __name__)
|
||||
globals()[name] = module
|
||||
else:
|
||||
raise AttributeError(f"module {__name__} has no attribute {name}")
|
||||
return globals()[name]
|
||||
|
||||
|
||||
def from_pretrained(path: str, **kwargs):
|
||||
"""
|
||||
Load a model from a pretrained checkpoint.
|
||||
|
||||
Args:
|
||||
path: The path to the checkpoint. Can be either local path or a Hugging Face model name.
|
||||
NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively.
|
||||
**kwargs: Additional arguments for the model constructor.
|
||||
"""
|
||||
import os
|
||||
import json
|
||||
from safetensors.torch import load_file
|
||||
is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors")
|
||||
|
||||
if is_local:
|
||||
config_file = f"{path}.json"
|
||||
model_file = f"{path}.safetensors"
|
||||
else:
|
||||
from huggingface_hub import hf_hub_download
|
||||
path_parts = path.split('/')
|
||||
repo_id = f'{path_parts[0]}/{path_parts[1]}'
|
||||
model_name = '/'.join(path_parts[2:])
|
||||
config_file = hf_hub_download(repo_id, f"{model_name}.json")
|
||||
model_file = hf_hub_download(repo_id, f"{model_name}.safetensors")
|
||||
|
||||
with open(config_file, 'r') as f:
|
||||
config = json.load(f)
|
||||
model = __getattr__(config['name'])(**config['args'], **kwargs)
|
||||
model.load_state_dict(load_file(model_file), strict=False)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
# For Pylance
|
||||
if __name__ == '__main__':
|
||||
from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
|
||||
from .sparse_structure_flow import SparseStructureFlowModel
|
||||
from .structured_latent_flow import SLatFlowModel, ElasticSLatFlowModel
|
||||
|
||||
from .sc_vaes.sparse_unet_vae import SparseUnetVaeEncoder, SparseUnetVaeDecoder
|
||||
from .sc_vaes.fdg_vae import FlexiDualGridVaeEncoder, FlexiDualGridVaeDecoder
|
||||
@@ -0,0 +1,110 @@
|
||||
from typing import *
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from ...modules import sparse as sp
|
||||
from .sparse_unet_vae import (
|
||||
SparseResBlock3d,
|
||||
SparseConvNeXtBlock3d,
|
||||
|
||||
SparseResBlockDownsample3d,
|
||||
SparseResBlockUpsample3d,
|
||||
SparseResBlockS2C3d,
|
||||
SparseResBlockC2S3d,
|
||||
)
|
||||
from .sparse_unet_vae import (
|
||||
SparseUnetVaeEncoder,
|
||||
SparseUnetVaeDecoder,
|
||||
)
|
||||
from ...representations import Mesh
|
||||
from o_voxel.convert import flexible_dual_grid_to_mesh
|
||||
|
||||
|
||||
class FlexiDualGridVaeEncoder(SparseUnetVaeEncoder):
|
||||
def __init__(
|
||||
self,
|
||||
model_channels: List[int],
|
||||
latent_channels: int,
|
||||
num_blocks: List[int],
|
||||
block_type: List[str],
|
||||
down_block_type: List[str],
|
||||
block_args: List[Dict[str, Any]],
|
||||
use_fp16: bool = False,
|
||||
):
|
||||
super().__init__(
|
||||
6,
|
||||
model_channels,
|
||||
latent_channels,
|
||||
num_blocks,
|
||||
block_type,
|
||||
down_block_type,
|
||||
block_args,
|
||||
use_fp16,
|
||||
)
|
||||
|
||||
def forward(self, vertices: sp.SparseTensor, intersected: sp.SparseTensor, sample_posterior=False, return_raw=False):
|
||||
x = vertices.replace(torch.cat([
|
||||
vertices.feats - 0.5,
|
||||
intersected.feats.float() - 0.5,
|
||||
], dim=1))
|
||||
return super().forward(x, sample_posterior, return_raw)
|
||||
|
||||
|
||||
class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
model_channels: List[int],
|
||||
latent_channels: int,
|
||||
num_blocks: List[int],
|
||||
block_type: List[str],
|
||||
up_block_type: List[str],
|
||||
block_args: List[Dict[str, Any]],
|
||||
voxel_margin: float = 0.5,
|
||||
use_fp16: bool = False,
|
||||
):
|
||||
self.resolution = resolution
|
||||
self.voxel_margin = voxel_margin
|
||||
|
||||
super().__init__(
|
||||
7,
|
||||
model_channels,
|
||||
latent_channels,
|
||||
num_blocks,
|
||||
block_type,
|
||||
up_block_type,
|
||||
block_args,
|
||||
use_fp16,
|
||||
)
|
||||
|
||||
def set_resolution(self, resolution: int) -> None:
|
||||
self.resolution = resolution
|
||||
|
||||
def forward(self, x: sp.SparseTensor, gt_intersected: sp.SparseTensor = None, **kwargs):
|
||||
decoded = super().forward(x, **kwargs)
|
||||
if self.training:
|
||||
h, subs_gt, subs = decoded
|
||||
vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
|
||||
intersected_logits = h.replace(h.feats[..., 3:6])
|
||||
quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
|
||||
mesh = [Mesh(*flexible_dual_grid_to_mesh(
|
||||
v.coords[:, 1:], v.feats, i.feats, q.feats,
|
||||
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
grid_size=self.resolution,
|
||||
train=True
|
||||
)) for v, i, q in zip(vertices, gt_intersected, quad_lerp)]
|
||||
return mesh, vertices, intersected_logits, subs_gt, subs
|
||||
else:
|
||||
out_list = list(decoded) if isinstance(decoded, tuple) else [decoded]
|
||||
h = out_list[0]
|
||||
vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
|
||||
intersected = h.replace(h.feats[..., 3:6] > 0)
|
||||
quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
|
||||
mesh = [Mesh(*flexible_dual_grid_to_mesh(
|
||||
v.coords[:, 1:], v.feats, i.feats, q.feats,
|
||||
aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
||||
grid_size=self.resolution,
|
||||
train=False
|
||||
)) for v, i, q in zip(vertices, intersected, quad_lerp)]
|
||||
out_list[0] = mesh
|
||||
return out_list[0] if len(out_list) == 1 else tuple(out_list)
|
||||
@@ -0,0 +1,522 @@
|
||||
from typing import *
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint
|
||||
from ...modules.utils import convert_module_to_f16, convert_module_to_f32, zero_module
|
||||
from ...modules import sparse as sp
|
||||
from ...modules.norm import LayerNorm32
|
||||
|
||||
|
||||
class SparseResBlock3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
out_channels: Optional[int] = None,
|
||||
downsample: bool = False,
|
||||
upsample: bool = False,
|
||||
resample_mode: Literal['nearest', 'spatial2channel'] = 'nearest',
|
||||
use_checkpoint: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.downsample = downsample
|
||||
self.upsample = upsample
|
||||
self.resample_mode = resample_mode
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
assert not (downsample and upsample), "Cannot downsample and upsample at the same time"
|
||||
|
||||
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
||||
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
||||
if resample_mode == 'nearest':
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
|
||||
elif resample_mode =='spatial2channel' and not self.downsample:
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels * 8, 3)
|
||||
elif resample_mode =='spatial2channel' and self.downsample:
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels // 8, 3)
|
||||
self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
|
||||
if resample_mode == 'nearest':
|
||||
self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
|
||||
elif resample_mode =='spatial2channel' and self.downsample:
|
||||
self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
|
||||
elif resample_mode =='spatial2channel' and not self.downsample:
|
||||
self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
|
||||
self.updown = None
|
||||
if self.downsample:
|
||||
if resample_mode == 'nearest':
|
||||
self.updown = sp.SparseDownsample(2)
|
||||
elif resample_mode =='spatial2channel':
|
||||
self.updown = sp.SparseSpatial2Channel(2)
|
||||
elif self.upsample:
|
||||
self.to_subdiv = sp.SparseLinear(channels, 8)
|
||||
if resample_mode == 'nearest':
|
||||
self.updown = sp.SparseUpsample(2)
|
||||
elif resample_mode =='spatial2channel':
|
||||
self.updown = sp.SparseChannel2Spatial(2)
|
||||
|
||||
def _updown(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
|
||||
if self.downsample:
|
||||
x = self.updown(x)
|
||||
elif self.upsample:
|
||||
x = self.updown(x, subdiv.replace(subdiv.feats > 0))
|
||||
return x
|
||||
|
||||
def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
subdiv = None
|
||||
if self.upsample:
|
||||
subdiv = self.to_subdiv(x)
|
||||
h = x.replace(self.norm1(x.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
if self.resample_mode == 'spatial2channel':
|
||||
h = self.conv1(h)
|
||||
h = self._updown(h, subdiv)
|
||||
x = self._updown(x, subdiv)
|
||||
if self.resample_mode == 'nearest':
|
||||
h = self.conv1(h)
|
||||
h = h.replace(self.norm2(h.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv2(h)
|
||||
h = h + self.skip_connection(x)
|
||||
if self.upsample:
|
||||
return h, subdiv
|
||||
return h
|
||||
|
||||
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
if self.use_checkpoint:
|
||||
return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
|
||||
else:
|
||||
return self._forward(x)
|
||||
|
||||
|
||||
class SparseResBlockDownsample3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
out_channels: Optional[int] = None,
|
||||
use_checkpoint: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
||||
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
|
||||
self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
|
||||
self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
|
||||
self.updown = sp.SparseDownsample(2)
|
||||
|
||||
def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
h = x.replace(self.norm1(x.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.updown(h)
|
||||
x = self.updown(x)
|
||||
h = self.conv1(h)
|
||||
h = h.replace(self.norm2(h.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv2(h)
|
||||
h = h + self.skip_connection(x)
|
||||
return h
|
||||
|
||||
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
if self.use_checkpoint:
|
||||
return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
|
||||
else:
|
||||
return self._forward(x)
|
||||
|
||||
|
||||
class SparseResBlockUpsample3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
out_channels: Optional[int] = None,
|
||||
use_checkpoint: bool = False,
|
||||
pred_subdiv: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_checkpoint = use_checkpoint
|
||||
self.pred_subdiv = pred_subdiv
|
||||
|
||||
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
||||
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
|
||||
self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
|
||||
self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
|
||||
if self.pred_subdiv:
|
||||
self.to_subdiv = sp.SparseLinear(channels, 8)
|
||||
self.updown = sp.SparseUpsample(2)
|
||||
|
||||
def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
|
||||
if self.pred_subdiv:
|
||||
subdiv = self.to_subdiv(x)
|
||||
h = x.replace(self.norm1(x.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
|
||||
h = self.updown(h, subdiv_binarized)
|
||||
x = self.updown(x, subdiv_binarized)
|
||||
h = self.conv1(h)
|
||||
h = h.replace(self.norm2(h.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv2(h)
|
||||
h = h + self.skip_connection(x)
|
||||
if self.pred_subdiv:
|
||||
return h, subdiv
|
||||
else:
|
||||
return h
|
||||
|
||||
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
if self.use_checkpoint:
|
||||
return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
|
||||
else:
|
||||
return self._forward(x)
|
||||
|
||||
|
||||
class SparseResBlockS2C3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
out_channels: Optional[int] = None,
|
||||
use_checkpoint: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
||||
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels // 8, 3)
|
||||
self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
|
||||
self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
|
||||
self.updown = sp.SparseSpatial2Channel(2)
|
||||
|
||||
def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
h = x.replace(self.norm1(x.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv1(h)
|
||||
h = self.updown(h)
|
||||
x = self.updown(x)
|
||||
h = h.replace(self.norm2(h.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv2(h)
|
||||
h = h + self.skip_connection(x)
|
||||
return h
|
||||
|
||||
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
if self.use_checkpoint:
|
||||
return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
|
||||
else:
|
||||
return self._forward(x)
|
||||
|
||||
|
||||
class SparseResBlockC2S3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
out_channels: Optional[int] = None,
|
||||
use_checkpoint: bool = False,
|
||||
pred_subdiv: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_checkpoint = use_checkpoint
|
||||
self.pred_subdiv = pred_subdiv
|
||||
|
||||
self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
||||
self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
|
||||
self.conv1 = sp.SparseConv3d(channels, self.out_channels * 8, 3)
|
||||
self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
|
||||
self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
|
||||
if pred_subdiv:
|
||||
self.to_subdiv = sp.SparseLinear(channels, 8)
|
||||
self.updown = sp.SparseChannel2Spatial(2)
|
||||
|
||||
def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
|
||||
if self.pred_subdiv:
|
||||
subdiv = self.to_subdiv(x)
|
||||
h = x.replace(self.norm1(x.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv1(h)
|
||||
subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
|
||||
h = self.updown(h, subdiv_binarized)
|
||||
x = self.updown(x, subdiv_binarized)
|
||||
h = h.replace(self.norm2(h.feats))
|
||||
h = h.replace(F.silu(h.feats))
|
||||
h = self.conv2(h)
|
||||
h = h + self.skip_connection(x)
|
||||
if self.pred_subdiv:
|
||||
return h, subdiv
|
||||
else:
|
||||
return h
|
||||
|
||||
def forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
|
||||
if self.use_checkpoint:
|
||||
return torch.utils.checkpoint.checkpoint(self._forward, x, subdiv, use_reentrant=False)
|
||||
else:
|
||||
return self._forward(x, subdiv)
|
||||
|
||||
|
||||
class SparseConvNeXtBlock3d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int,
|
||||
mlp_ratio: float = 4.0,
|
||||
use_checkpoint: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.use_checkpoint = use_checkpoint
|
||||
|
||||
self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
|
||||
self.conv = sp.SparseConv3d(channels, channels, 3)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(channels, int(channels * mlp_ratio)),
|
||||
nn.SiLU(),
|
||||
zero_module(nn.Linear(int(channels * mlp_ratio), channels)),
|
||||
)
|
||||
|
||||
def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
h = self.conv(x)
|
||||
h = h.replace(self.norm(h.feats))
|
||||
h = h.replace(self.mlp(h.feats))
|
||||
return h + x
|
||||
|
||||
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
||||
if self.use_checkpoint:
|
||||
return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
|
||||
else:
|
||||
return self._forward(x)
|
||||
|
||||
|
||||
class SparseUnetVaeEncoder(nn.Module):
|
||||
"""
|
||||
Sparse Swin Transformer Unet VAE model.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
model_channels: List[int],
|
||||
latent_channels: int,
|
||||
num_blocks: List[int],
|
||||
block_type: List[str],
|
||||
down_block_type: List[str],
|
||||
block_args: List[Dict[str, Any]],
|
||||
use_fp16: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.num_blocks = num_blocks
|
||||
self.dtype = torch.float16 if use_fp16 else torch.float32
|
||||
self.dtype = torch.float16 if use_fp16 else torch.float32
|
||||
|
||||
self.input_layer = sp.SparseLinear(in_channels, model_channels[0])
|
||||
self.to_latent = sp.SparseLinear(model_channels[-1], 2 * latent_channels)
|
||||
|
||||
self.blocks = nn.ModuleList([])
|
||||
for i in range(len(num_blocks)):
|
||||
self.blocks.append(nn.ModuleList([]))
|
||||
for j in range(num_blocks[i]):
|
||||
self.blocks[-1].append(
|
||||
globals()[block_type[i]](
|
||||
model_channels[i],
|
||||
**block_args[i],
|
||||
)
|
||||
)
|
||||
if i < len(num_blocks) - 1:
|
||||
self.blocks[-1].append(
|
||||
globals()[down_block_type[i]](
|
||||
model_channels[i],
|
||||
model_channels[i+1],
|
||||
**block_args[i],
|
||||
)
|
||||
)
|
||||
|
||||
self.initialize_weights()
|
||||
if use_fp16:
|
||||
self.convert_to_fp16()
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
"""
|
||||
Return the device of the model.
|
||||
"""
|
||||
return next(self.parameters()).device
|
||||
|
||||
def convert_to_fp16(self) -> None:
|
||||
"""
|
||||
Convert the torso of the model to float16.
|
||||
"""
|
||||
self.blocks.apply(convert_module_to_f16)
|
||||
|
||||
def convert_to_fp32(self) -> None:
|
||||
"""
|
||||
Convert the torso of the model to float32.
|
||||
"""
|
||||
self.blocks.apply(convert_module_to_f32)
|
||||
|
||||
def initialize_weights(self) -> None:
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
self.apply(_basic_init)
|
||||
|
||||
def forward(self, x: sp.SparseTensor, sample_posterior=False, return_raw=False):
|
||||
h = self.input_layer(x)
|
||||
h = h.type(self.dtype)
|
||||
for i, res in enumerate(self.blocks):
|
||||
for j, block in enumerate(res):
|
||||
h = block(h)
|
||||
h = h.type(x.dtype)
|
||||
h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
|
||||
h = self.to_latent(h)
|
||||
|
||||
# Sample from the posterior distribution
|
||||
mean, logvar = h.feats.chunk(2, dim=-1)
|
||||
if sample_posterior:
|
||||
std = torch.exp(0.5 * logvar)
|
||||
z = mean + std * torch.randn_like(std)
|
||||
else:
|
||||
z = mean
|
||||
z = h.replace(z)
|
||||
|
||||
if return_raw:
|
||||
return z, mean, logvar
|
||||
else:
|
||||
return z
|
||||
|
||||
|
||||
class SparseUnetVaeDecoder(nn.Module):
|
||||
"""
|
||||
Sparse Swin Transformer Unet VAE model.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
out_channels: int,
|
||||
model_channels: List[int],
|
||||
latent_channels: int,
|
||||
num_blocks: List[int],
|
||||
block_type: List[str],
|
||||
up_block_type: List[str],
|
||||
block_args: List[Dict[str, Any]],
|
||||
use_fp16: bool = False,
|
||||
pred_subdiv: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.out_channels = out_channels
|
||||
self.model_channels = model_channels
|
||||
self.num_blocks = num_blocks
|
||||
self.use_fp16 = use_fp16
|
||||
self.pred_subdiv = pred_subdiv
|
||||
self.dtype = torch.float16 if use_fp16 else torch.float32
|
||||
self.low_vram = False
|
||||
|
||||
self.output_layer = sp.SparseLinear(model_channels[-1], out_channels)
|
||||
self.from_latent = sp.SparseLinear(latent_channels, model_channels[0])
|
||||
|
||||
self.blocks = nn.ModuleList([])
|
||||
for i in range(len(num_blocks)):
|
||||
self.blocks.append(nn.ModuleList([]))
|
||||
for j in range(num_blocks[i]):
|
||||
self.blocks[-1].append(
|
||||
globals()[block_type[i]](
|
||||
model_channels[i],
|
||||
**block_args[i],
|
||||
)
|
||||
)
|
||||
if i < len(num_blocks) - 1:
|
||||
self.blocks[-1].append(
|
||||
globals()[up_block_type[i]](
|
||||
model_channels[i],
|
||||
model_channels[i+1],
|
||||
pred_subdiv=pred_subdiv,
|
||||
**block_args[i],
|
||||
)
|
||||
)
|
||||
|
||||
self.initialize_weights()
|
||||
if use_fp16:
|
||||
self.convert_to_fp16()
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
"""
|
||||
Return the device of the model.
|
||||
"""
|
||||
return next(self.parameters()).device
|
||||
|
||||
def convert_to_fp16(self) -> None:
|
||||
"""
|
||||
Convert the torso of the model to float16.
|
||||
"""
|
||||
self.blocks.apply(convert_module_to_f16)
|
||||
|
||||
def convert_to_fp32(self) -> None:
|
||||
"""
|
||||
Convert the torso of the model to float32.
|
||||
"""
|
||||
self.blocks.apply(convert_module_to_f32)
|
||||
|
||||
def initialize_weights(self) -> None:
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
self.apply(_basic_init)
|
||||
|
||||
def forward(self, x: sp.SparseTensor, guide_subs: Optional[List[sp.SparseTensor]] = None, return_subs: bool = False) -> sp.SparseTensor:
|
||||
assert guide_subs is None or self.pred_subdiv == False, "Only decoders with pred_subdiv=False can be used with guide_subs"
|
||||
assert return_subs == False or self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with return_subs"
|
||||
|
||||
h = self.from_latent(x)
|
||||
h = h.type(self.dtype)
|
||||
subs_gt = []
|
||||
subs = []
|
||||
for i, res in enumerate(self.blocks):
|
||||
for j, block in enumerate(res):
|
||||
if i < len(self.blocks) - 1 and j == len(res) - 1:
|
||||
if self.pred_subdiv:
|
||||
if self.training:
|
||||
subs_gt.append(h.get_spatial_cache('subdivision'))
|
||||
h, sub = block(h)
|
||||
subs.append(sub)
|
||||
else:
|
||||
h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
|
||||
else:
|
||||
h = block(h)
|
||||
h = h.type(x.dtype)
|
||||
h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
|
||||
h = self.output_layer(h)
|
||||
if self.training and self.pred_subdiv:
|
||||
return h, subs_gt, subs
|
||||
else:
|
||||
if return_subs:
|
||||
return h, subs
|
||||
else:
|
||||
return h
|
||||
|
||||
def upsample(self, x: sp.SparseTensor, upsample_times: int) -> torch.Tensor:
|
||||
assert self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with upsampling"
|
||||
|
||||
h = self.from_latent(x)
|
||||
h = h.type(self.dtype)
|
||||
for i, res in enumerate(self.blocks):
|
||||
if i == upsample_times:
|
||||
return h.coords
|
||||
for j, block in enumerate(res):
|
||||
if i < len(self.blocks) - 1 and j == len(res) - 1:
|
||||
h, sub = block(h)
|
||||
else:
|
||||
h = block(h)
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
from contextlib import contextmanager
|
||||
from typing import *
|
||||
import math
|
||||
from ..modules import sparse as sp
|
||||
from ..utils.elastic_utils import ElasticModuleMixin
|
||||
|
||||
|
||||
class SparseTransformerElasticMixin(ElasticModuleMixin):
|
||||
def _get_input_size(self, x: sp.SparseTensor, *args, **kwargs):
|
||||
return x.feats.shape[0]
|
||||
|
||||
@contextmanager
|
||||
def with_mem_ratio(self, mem_ratio=1.0):
|
||||
if mem_ratio == 1.0:
|
||||
yield 1.0
|
||||
return
|
||||
num_blocks = len(self.blocks)
|
||||
num_checkpoint_blocks = min(math.ceil((1 - mem_ratio) * num_blocks) + 1, num_blocks)
|
||||
exact_mem_ratio = 1 - (num_checkpoint_blocks - 1) / num_blocks
|
||||
for i in range(num_blocks):
|
||||
self.blocks[i].use_checkpoint = i < num_checkpoint_blocks
|
||||
yield exact_mem_ratio
|
||||
for i in range(num_blocks):
|
||||
self.blocks[i].use_checkpoint = False
|
||||
@@ -0,0 +1,298 @@
|
||||
from typing import *
|
||||
from functools import partial
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
|
||||
from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
|
||||
from ..modules.attention import RotaryPositionEmbedder
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t, dim, max_period=10000):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
|
||||
Args:
|
||||
t: a 1-D Tensor of N indices, one per batch element.
|
||||
These may be fractional.
|
||||
dim: the dimension of the output.
|
||||
max_period: controls the minimum frequency of the embeddings.
|
||||
|
||||
Returns:
|
||||
an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
|
||||
).to(device=t.device)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
||||
t_emb = self.mlp(t_freq)
|
||||
return t_emb
|
||||
|
||||
|
||||
class SparseStructureFlowModel(nn.Module):
|
||||
"""
|
||||
Sparse Structure Flow Model for 3D generation.
|
||||
|
||||
Supports two conditioning modes:
|
||||
- "cross": Standard cross-attention with image features
|
||||
- "proj": View-aligned projection attention with camera-aware features
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
resolution: int,
|
||||
in_channels: int,
|
||||
model_channels: int,
|
||||
cond_channels: int,
|
||||
out_channels: int,
|
||||
num_blocks: int,
|
||||
num_heads: Optional[int] = None,
|
||||
num_head_channels: Optional[int] = 64,
|
||||
mlp_ratio: float = 4,
|
||||
pe_mode: Literal["ape", "rope"] = "ape",
|
||||
rope_freq: Tuple[float, float] = (1.0, 10000.0),
|
||||
dtype: str = 'float32',
|
||||
use_checkpoint: bool = False,
|
||||
share_mod: bool = False,
|
||||
initialization: str = 'vanilla',
|
||||
qk_rms_norm: bool = False,
|
||||
qk_rms_norm_cross: bool = False,
|
||||
image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross",
|
||||
proj_in_channels: Optional[int] = None,
|
||||
vae_in_channels: Optional[int] = None,
|
||||
**kwargs
|
||||
):
|
||||
super().__init__()
|
||||
self.resolution = resolution
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.cond_channels = cond_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_blocks = num_blocks
|
||||
self.num_heads = num_heads or model_channels // num_head_channels
|
||||
self.mlp_ratio = mlp_ratio
|
||||
self.pe_mode = pe_mode
|
||||
self.use_checkpoint = use_checkpoint
|
||||
self.share_mod = share_mod
|
||||
self.initialization = initialization
|
||||
self.qk_rms_norm = qk_rms_norm
|
||||
self.qk_rms_norm_cross = qk_rms_norm_cross
|
||||
self.image_attn_mode = image_attn_mode
|
||||
self.proj_in_channels = proj_in_channels
|
||||
self.vae_in_channels = vae_in_channels
|
||||
self.dtype = str_to_dtype(dtype)
|
||||
|
||||
self.t_embedder = TimestepEmbedder(model_channels)
|
||||
if share_mod:
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(model_channels, 6 * model_channels, bias=True)
|
||||
)
|
||||
|
||||
if pe_mode == "ape":
|
||||
pos_embedder = AbsolutePositionEmbedder(model_channels, 3)
|
||||
coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution] * 3], indexing='ij')
|
||||
coords = torch.stack(coords, dim=-1).reshape(-1, 3)
|
||||
pos_emb = pos_embedder(coords)
|
||||
self.register_buffer("pos_emb", pos_emb)
|
||||
elif pe_mode == "rope":
|
||||
pos_embedder = RotaryPositionEmbedder(self.model_channels // self.num_heads, 3)
|
||||
coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution] * 3], indexing='ij')
|
||||
coords = torch.stack(coords, dim=-1).reshape(-1, 3)
|
||||
rope_phases = pos_embedder(coords)
|
||||
self.register_buffer("rope_phases", rope_phases)
|
||||
|
||||
if pe_mode != "rope":
|
||||
self.rope_phases = None
|
||||
|
||||
self.input_layer = nn.Linear(in_channels, model_channels)
|
||||
|
||||
self.blocks = nn.ModuleList([
|
||||
ModulatedTransformerCrossBlock(
|
||||
model_channels,
|
||||
cond_channels,
|
||||
num_heads=self.num_heads,
|
||||
mlp_ratio=self.mlp_ratio,
|
||||
attn_mode='full',
|
||||
use_checkpoint=self.use_checkpoint,
|
||||
use_rope=(pe_mode == "rope"),
|
||||
rope_freq=rope_freq,
|
||||
share_mod=share_mod,
|
||||
qk_rms_norm=self.qk_rms_norm,
|
||||
qk_rms_norm_cross=self.qk_rms_norm_cross,
|
||||
image_attn_mode=image_attn_mode,
|
||||
proj_in_channels=proj_in_channels,
|
||||
vae_in_channels=vae_in_channels,
|
||||
)
|
||||
for _ in range(num_blocks)
|
||||
])
|
||||
|
||||
self.out_layer = nn.Linear(model_channels, out_channels)
|
||||
|
||||
self.initialize_weights()
|
||||
self.convert_to(self.dtype)
|
||||
|
||||
@property
|
||||
def device(self) -> torch.device:
|
||||
"""
|
||||
Return the device of the model.
|
||||
"""
|
||||
return next(self.parameters()).device
|
||||
|
||||
def convert_to(self, dtype: torch.dtype) -> None:
|
||||
"""
|
||||
Convert the torso of the model to the specified dtype.
|
||||
"""
|
||||
self.dtype = dtype
|
||||
self.blocks.apply(partial(convert_module_to, dtype=dtype))
|
||||
|
||||
def initialize_weights(self) -> None:
|
||||
if self.initialization == 'vanilla':
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
self.apply(_basic_init)
|
||||
|
||||
# Initialize timestep embedding MLP:
|
||||
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
||||
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
|
||||
|
||||
# Zero-out adaLN modulation layers in DiT blocks:
|
||||
if self.share_mod:
|
||||
nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
|
||||
else:
|
||||
for block in self.blocks:
|
||||
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
# Zero-out output layers:
|
||||
nn.init.constant_(self.out_layer.weight, 0)
|
||||
nn.init.constant_(self.out_layer.bias, 0)
|
||||
|
||||
elif self.initialization == 'scaled':
|
||||
# Initialize transformer layers:
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels)))
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
self.apply(_basic_init)
|
||||
|
||||
# Scaled init for to_out and ffn2
|
||||
def _scaled_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels))
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
for block in self.blocks:
|
||||
block.self_attn.to_out.apply(_scaled_init)
|
||||
# Handle cross, proj, and gated_proj modes
|
||||
if self.image_attn_mode in ("proj", "gated_proj"):
|
||||
block.cross_attn.cross_attn_block.to_out.apply(_scaled_init)
|
||||
else:
|
||||
block.cross_attn.to_out.apply(_scaled_init)
|
||||
block.mlp.mlp[2].apply(_scaled_init)
|
||||
|
||||
# Initialize input layer to make the initial representation have variance 1
|
||||
nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels))
|
||||
nn.init.zeros_(self.input_layer.bias)
|
||||
|
||||
# Initialize timestep embedding MLP:
|
||||
nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
|
||||
nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
|
||||
|
||||
# Zero-out adaLN modulation layers in DiT blocks:
|
||||
if self.share_mod:
|
||||
nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
|
||||
else:
|
||||
for block in self.blocks:
|
||||
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
# Zero-out output layers:
|
||||
nn.init.constant_(self.out_layer.weight, 0)
|
||||
nn.init.constant_(self.out_layer.bias, 0)
|
||||
|
||||
def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Forward pass.
|
||||
|
||||
Args:
|
||||
x: Input tensor [B, C, D, H, W]
|
||||
t: Timestep tensor [B]
|
||||
cond: Conditioning tensor. For "cross" mode: [B, N, D].
|
||||
For "proj" mode: dict {'global': global_cond, 'proj': proj_cond}
|
||||
or tuple of (global_cond, proj_cond)
|
||||
|
||||
Returns:
|
||||
Output tensor [B, C, D, H, W]
|
||||
"""
|
||||
assert [*x.shape] == [x.shape[0], self.in_channels, *[self.resolution] * 3], \
|
||||
f"Input shape mismatch, got {x.shape}, expected {[x.shape[0], self.in_channels, *[self.resolution] * 3]}"
|
||||
|
||||
h = x.view(*x.shape[:2], -1).permute(0, 2, 1).contiguous()
|
||||
|
||||
h = self.input_layer(h)
|
||||
if self.pe_mode == "ape":
|
||||
h = h + self.pos_emb[None]
|
||||
t_emb = self.t_embedder(t)
|
||||
if self.share_mod:
|
||||
t_emb = self.adaLN_modulation(t_emb)
|
||||
t_emb = manual_cast(t_emb, self.dtype)
|
||||
h = manual_cast(h, self.dtype)
|
||||
|
||||
# Handle different conditioning modes
|
||||
if self.image_attn_mode == 'proj':
|
||||
if isinstance(cond, dict):
|
||||
global_cond = cond['global']
|
||||
proj_cond = cond['proj']
|
||||
else:
|
||||
global_cond, proj_cond = cond
|
||||
global_cond = manual_cast(global_cond, self.dtype)
|
||||
proj_cond = manual_cast(proj_cond, self.dtype)
|
||||
cond = (global_cond, proj_cond)
|
||||
elif self.image_attn_mode == 'gated_proj':
|
||||
global_cond = manual_cast(cond['global'], self.dtype)
|
||||
proj_semantic = manual_cast(cond['proj_semantic'], self.dtype)
|
||||
proj_color = manual_cast(cond['proj_color'], self.dtype)
|
||||
cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color}
|
||||
else:
|
||||
cond = manual_cast(cond, self.dtype)
|
||||
|
||||
for block in self.blocks:
|
||||
h = block(h, t_emb, cond, self.rope_phases)
|
||||
h = manual_cast(h, x.dtype)
|
||||
h = F.layer_norm(h, h.shape[-1:])
|
||||
h = self.out_layer(h)
|
||||
|
||||
h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution] * 3).contiguous()
|
||||
|
||||
return h
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user