Repository: PKU-EPIC/GraspNeRF
Branch: main
Commit: 946e12a11fa2
Files: 61
Total size: 657.5 KB
Directory structure:
gitextract_hto38wet/
├── .gitignore
├── README.md
├── data_generator/
│ ├── modify_material.py
│ ├── render_pile_STD_rand.py
│ └── run_pile_rand.sh
├── requirements.txt
├── run_simgrasp.sh
├── scripts/
│ ├── sim_grasp.py
│ └── stat_expresult.py
├── src/
│ ├── gd/
│ │ ├── __init__.py
│ │ ├── baselines.py
│ │ ├── dataset.py
│ │ ├── detection.py
│ │ ├── experiments/
│ │ │ ├── __init__.py
│ │ │ └── clutter_removal.py
│ │ ├── grasp.py
│ │ ├── io.py
│ │ ├── networks.py
│ │ ├── perception.py
│ │ ├── simulation.py
│ │ ├── utils/
│ │ │ ├── __init__.py
│ │ │ ├── btsim.py
│ │ │ ├── panda_control.py
│ │ │ ├── ros_utils.py
│ │ │ └── transform.py
│ │ └── vis.py
│ ├── nr/
│ │ ├── asset.py
│ │ ├── configs/
│ │ │ └── nrvgn_sdf.yaml
│ │ ├── dataset/
│ │ │ ├── database.py
│ │ │ ├── name2dataset.py
│ │ │ └── train_dataset.py
│ │ ├── main.py
│ │ ├── network/
│ │ │ ├── aggregate_net.py
│ │ │ ├── dist_decoder.py
│ │ │ ├── ibrnet.py
│ │ │ ├── init_net.py
│ │ │ ├── loss.py
│ │ │ ├── metrics.py
│ │ │ ├── mvsnet/
│ │ │ │ ├── modules.py
│ │ │ │ └── mvsnet.py
│ │ │ ├── neus.py
│ │ │ ├── ops.py
│ │ │ ├── render_ops.py
│ │ │ ├── renderer.py
│ │ │ └── vis_encoder.py
│ │ ├── run_training.py
│ │ ├── train/
│ │ │ ├── lr_common_manager.py
│ │ │ ├── train_tools.py
│ │ │ ├── train_valid.py
│ │ │ └── trainer.py
│ │ └── utils/
│ │ ├── base_utils.py
│ │ ├── dataset_utils.py
│ │ ├── draw_utils.py
│ │ ├── field_utils.py
│ │ ├── grasp_utils.py
│ │ ├── imgs_info.py
│ │ └── view_select.py
│ └── rd/
│ ├── modify_material.py
│ ├── render.py
│ └── render_utils.py
└── train.sh
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# pyenv
.python-version
# celery beat schedule file
celerybeat-schedule
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
.venv-py2
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
# vscode
.vscode
# data
data
__pycache__
*.log
log/
output/
ckpt/
*.pth
================================================
FILE: README.md
================================================
# GraspNeRF: Multiview-based 6-DoF Grasp Detection for Transparent and Specular Objects Using Generalizable NeRF (ICRA 2023)
This is the official repository of [**GraspNeRF: Multiview-based 6-DoF Grasp Detection for Transparent and Specular Objects Using Generalizable NeRF**](https://arxiv.org/abs/2210.06575).
For more information, please visit our [**project page**](https://pku-epic.github.io/GraspNeRF/).
## Introduction
In this work, we propose a multiview RGB-based 6-DoF grasp detection network, **GraspNeRF**,
that leverages the **generalizable neural radiance field (NeRF)** to achieve material-agnostic object grasping in clutter.
Compared to the existing NeRF-based 3-DoF grasp detection methods that rely on densely captured input images and time-consuming per-scene optimization,
our system can perform zero-shot NeRF construction with **sparse RGB inputs** and reliably detect **6-DoF grasps**, both in **real-time**.
The proposed framework jointly learns generalizable NeRF and grasp detection in an **end-to-end** manner, optimizing the scene representation construction for the grasping.
For training data, we generate a large-scale photorealistic **domain-randomized synthetic dataset** of grasping in cluttered tabletop scenes that enables direct transfer to the real world.
Experiments in synthetic and real-world environments demonstrate that our method significantly outperforms all the baselines in all the experiments.
## Overview
This repository provides:
- PyTorch code, and weights of GraspNeRF.
- Grasp Simulator based on blender and pybullet.
- Multiview 6-DoF Grasping Dataset Generator and Examples.
## Dependency
1. Please run
```
pip install -r requirements.txt
```
to install dependency.
2. (optional) Please install [blender 2.93.3--Ubuntu](https://www.blender.org/) if you need simulation.
## Data & Checkpoints
1. Please generate or download and uncompress the [example data](https://drive.google.com/file/d/1Ku-EotayUhfv5DtXAvFitGzzdMF84Ve2/view?usp=share_link) to `data/` for training, and [rendering assets](https://drive.google.com/file/d/1Udvi2QQ6AtYDLUWY0oH-PO2R6kZBxJLT/view?usp=share_link) to `data/assets` for simulation.
Specifically, download [imagenet valset](https://image-net.org/data/ILSVRC/2010/ILSVRC2010_images_val.tar) to `data/assets/imagenet/images/val` which is used as random texture in simulation.
2. We provide pretrained weights for testing. Please download the [checkpoint](https://drive.google.com/file/d/1k-Cy4NO2isCBYc3az-34HEdcNxDptDgU/view?usp=share_link) to `src/nr/ckpt/test`.
## Testing
Our grasp simulation pipeline is depend on blender and pybullet. Please verify the installation before running simulation.
After the dependency and assets are ready, please run
```
bash run_simgrasp.sh
```
## Training
After the training data is ready, please run
```
bash train.sh GPU_ID
```
e.g. `bash train.sh 0`.
## Data Generator
1. Download the scene descriptor files from [GIGA](https://github.com/UT-Austin-RPL/GIGA#pre-generated-data) and [assets](https://drive.google.com/file/d/1-59zcQ8h5esT_ogjaDjtzQ6sG70WNWzU/view?usp=share_link).
2. For example, run
```
bash run_pile_rand.sh
```
in ./data_generator for pile data generation.
## Citation
If you find our work useful in your research, please consider citing:
```
@article{Dai2023GraspNeRF,
title={GraspNeRF: Multiview-based 6-DoF Grasp Detection for Transparent and Specular Objects Using Generalizable NeRF},
author={Qiyu Dai and Yan Zhu and Yiran Geng and Ciyu Ruan and Jiazhao Zhang and He Wang},
booktitle={IEEE International Conference on Robotics and Automation (ICRA)},
year={2023}
```
## License
This work and the dataset are licensed under [CC BY-NC 4.0][cc-by-nc].
[![CC BY-NC 4.0][cc-by-nc-image]][cc-by-nc]
[cc-by-nc]: https://creativecommons.org/licenses/by-nc/4.0/
[cc-by-nc-image]: https://licensebuttons.net/l/by-nc/4.0/88x31.png
## Contact
If you have any questions, please open a github issue or contact us:
Qiyu Dai: qiyudai@pku.edu.cn, Yan Zhu: zhuyan_@stu.pku.edu.cn, He Wang: hewang@pku.edu.cn
================================================
FILE: data_generator/modify_material.py
================================================
from mathutils import Vector
import bpy
import random
import math
def modify_material(mat_links, mat_nodes, material_name, mat_randomize_mode, is_texture=False, orign_base_color=None, tex_node=None, is_transfer=True, is_arm=False):
if is_transfer:
if material_name.split("_")[0] == "metal" or material_name.split("_")[0] == "porcelain" or material_name.split("_")[0] == "plasticsp" or material_name.split("_")[0] == "paintsp":
tex_mix_prop = random.uniform(0.85, 0.98)
else:
tex_mix_prop = random.uniform(0.7, 0.95)
mix_prop = random.uniform(0.6, 0.9)
if mat_randomize_mode == "specular_texmix" or mat_randomize_mode == "mixed" \
or material_name.split("_")[0] == "metal" or material_name.split("_")[0] == "porcelain" \
or material_name.split("_")[0] == "plasticsp" or material_name.split("_")[0] == "paintsp":
transfer_rand = random.randint(0,2)
else:
transfer_rand = 1
if transfer_rand == 1:
transfer_flag = True
else:
transfer_flag = False
tex_mix_prop = 1
mix_prop = 1
if not is_arm:
bs_color_rand = random.uniform(-0.2, 0.2)
else:
bs_color_rand = 0
r_rand = bs_color_rand
g_rand = bs_color_rand
b_rand = bs_color_rand
else:
tex_mix_prop = 1
mix_prop = 1
transfer_flag = False
if not is_arm:
bs_color_rand = random.uniform(-0.2, 0.2)
else:
bs_color_rand = 0
r_rand = bs_color_rand
g_rand = bs_color_rand
b_rand = bs_color_rand
bsdfnode_list = [n for n in mat_nodes if isinstance(n, bpy.types.ShaderNodeBsdfPrincipled)]
if bsdfnode_list != []:
for bsdfnode in bsdfnode_list:
if not bsdfnode.inputs[4].links: # metallic
src_value = bsdfnode.inputs[4].default_value
if material_name.split("_")[0] == "metal":
new_value = src_value + random.uniform(-0.05, 0.05)
elif material_name.split("_")[0] == "porcelain":
new_value = src_value + random.uniform(-0.05, 0.1)
elif material_name.split("_")[0] == "plasticsp":
new_value = src_value + random.uniform(-0.05, 0.1)
else:
new_value = src_value + random.uniform(-0.05, 0.05)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[4].default_value = new_value
if not bsdfnode.inputs[5].links: # specular
src_value = bsdfnode.inputs[5].default_value
#if material_name.split("_")[0] == "metal":
new_value = src_value + random.uniform(0, 0.3)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[5].default_value = new_value
if not bsdfnode.inputs[6].links: # specularTint
src_value = bsdfnode.inputs[6].default_value
new_value = src_value + random.uniform(-1, 1)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[6].default_value = new_value
if not bsdfnode.inputs[7].links: # roughness
src_value = bsdfnode.inputs[7].default_value
if material_name.split("_")[0] == "metal" or material_name.split("_")[0] == "porcelain" or material_name.split("_")[0] == "plasticsp" or material_name.split("_")[0] == "paintsp":
new_value = src_value + random.uniform(-0.2, 0.01)
else:
new_value = src_value + random.uniform(-0.03, 0.1)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[7].default_value = new_value
if not bsdfnode.inputs[8].links: # anisotropic
src_value = bsdfnode.inputs[8].default_value
new_value = src_value + random.uniform(-0.1, 0.1)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[8].default_value = new_value
if not bsdfnode.inputs[9].links: # anisotropicRotation
src_value = bsdfnode.inputs[9].default_value
new_value = src_value + random.uniform(-0.3, 0.3)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[9].default_value = new_value
if not bsdfnode.inputs[10].links: # sheen
src_value = bsdfnode.inputs[10].default_value
new_value = src_value + random.uniform(-0.1, 0.1)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[10].default_value = new_value
if not bsdfnode.inputs[11].links: # sheenTint
src_value = bsdfnode.inputs[11].default_value
new_value = src_value + random.uniform(-0.2, 0.2)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[11].default_value = new_value
if not bsdfnode.inputs[12].links: # clearcoat
src_value = bsdfnode.inputs[12].default_value
new_value = src_value + random.uniform(-0.2, 0.2)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[12].default_value = new_value
if not bsdfnode.inputs[13].links: # clearcoatGloss
src_value = bsdfnode.inputs[13].default_value
new_value = src_value + random.uniform(-0.2, 0.2)
if new_value > 1.0: new_value = 1.0
elif new_value < 0: new_value = 0.0
bsdfnode.inputs[13].default_value = new_value
## metal
if material_name == "metal_0":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.3, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 1.0) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.3, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_1":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.9, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.5, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[7].default_value = random.uniform(0.08, 0.25) # roughness
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.04, 0.5) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.3, 0.7) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.8, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Tangent"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[22])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_10":
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.5) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.3, 0.7) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
bsdf_new.location = Vector((-800, 0))
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
mix_new.location = Vector((-800, 0))
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Image Texture"].outputs[1], mat_nodes["Principled BSDF-new"].inputs[19])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[4])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_11":
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.8) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 0.8) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[4])
mat_links.new(mat_nodes["Image Texture.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_12":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.8) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 0.8) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Reroute.006"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_13":
# mat_nodes["Principled BSDF.001"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF.001"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF.001"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF.001"].inputs[8].default_value = random.uniform(0.3, 0.7) # anisotropic
# mat_nodes["Principled BSDF.001"].inputs[9].default_value = random.uniform(0.0, 0.8) # anisotropicRotation
# mat_nodes["Principled BSDF.001"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF.001"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF.001"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Mix.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF.001"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output.001"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF.001"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_14":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.5) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 0.5) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.85
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Group"].outputs[1], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Group"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_2":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.5, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.95) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Image Texture.003"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_3":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.5, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.2) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
mat_nodes["Gamma"].inputs[1].default_value = random.uniform(3.0, 4.0)
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Gamma"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_4":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.1, 0.5) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.2) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 0.5) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.5) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_5":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.2, 0.4) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.6, 0.9) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.8, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Voronoi Texture"].outputs[1], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Tangent"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[22])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[21])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_6":
# mat_nodes["BSDF guidé"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["BSDF guidé"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["BSDF guidé"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["BSDF guidé"].inputs[8].default_value = random.uniform(0.0, 0.2) # anisotropic
# mat_nodes["BSDF guidé"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["BSDF guidé"].inputs[12].default_value = random.uniform(0.0, 0.3) # clearcoat
# mat_nodes["BSDF guidé"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
mat_nodes["Valeur"].outputs[0].default_value = random.uniform(0.1, 0.3)
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
bsdf_new.location = Vector((-800, 0))
for key, input in enumerate(mat_nodes["BSDF guidé"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
mix_new.location = Vector((-800, 0))
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Mélanger.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["BSDF guidé"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Sortie de matériau"].inputs[0])
elif material_name == "metal_7":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.7, 0.9) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 0.3) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
#bsdf_new.location = Vector((-800, 0))
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
#mix_new.location = Vector((-800, 0))
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.7
mat_links.new(mat_nodes["Reroute.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Tangent"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[22])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_8":
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
bsdf_1_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_1_new.name = 'Principled BSDF-1-new'
for key, input in enumerate(mat_nodes["Principled BSDF.001"].inputs):
bsdf_1_new.inputs[key].default_value = input.default_value
bsdf_2_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_2_new.name = 'Principled BSDF-2-new'
for key, input in enumerate(mat_nodes["Principled BSDF.002"].inputs):
bsdf_2_new.inputs[key].default_value = input.default_value
bsdf_3_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_3_new.name = 'Principled BSDF-3-new'
for key, input in enumerate(mat_nodes["Principled BSDF.003"].inputs):
bsdf_3_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
mix_1_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_1_new.name = 'Mix Shader-1-new'
mix_2_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_2_new.name = 'Mix Shader-2-new'
mix_3_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_3_new.name = 'Mix Shader-3-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-1-new"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-2-new"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-3-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.6
mat_nodes["Mix Shader-1-new"].inputs[0].default_value = 0.6
mat_nodes["Mix Shader-2-new"].inputs[0].default_value = 0.6
mat_nodes["Mix Shader-3-new"].inputs[0].default_value = 0.6
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Principled BSDF-1-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Principled BSDF-2-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Principled BSDF-3-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.5
mat_nodes["Mix Shader-1-new"].inputs[0].default_value = 0.5
mat_nodes["Mix Shader-2-new"].inputs[0].default_value = 0.5
mat_nodes["Mix Shader-3-new"].inputs[0].default_value = 0.5
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-1-new"].inputs[20])
mat_links.new(mat_nodes["Bump.001"].outputs[0], mat_nodes["Principled BSDF-2-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF.001"].outputs[0], mat_nodes["Mix Shader-1-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-1-new"].outputs[0], mat_nodes["Mix Shader-1-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-1-new"].outputs[0], mat_nodes["Mix Shader"].inputs[2])
mat_links.new(mat_nodes["Principled BSDF.002"].outputs[0], mat_nodes["Mix Shader-2-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-2-new"].outputs[0], mat_nodes["Mix Shader-2-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-2-new"].outputs[0], mat_nodes["Mix Shader.001"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF.003"].outputs[0], mat_nodes["Mix Shader-3-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-3-new"].outputs[0], mat_nodes["Mix Shader-3-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-3-new"].outputs[0], mat_nodes["Mix Shader.001"].inputs[2])
elif material_name == "metal_9":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[7].default_value = random.uniform(0.01, 0.3) # roughness
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 0.3) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
mat_nodes["Anisotropic BSDF"].inputs[1].default_value = random.uniform(0.11, 0.25)
mat_nodes["Anisotropic BSDF"].inputs[2].default_value = random.uniform(0.4, 0.6)
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
mat_links.new(tex_node.outputs[0], mat_nodes["Anisotropic BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_nodes["Anisotropic BSDF"].inputs[0].default_value = list(orign_base_color)
mat_links.new(mat_nodes["Principled BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Principled BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
## porcelain
elif material_name == "porcelain_0":
if transfer_flag == True:
# if is_texture:
# mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
# else:
# mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop#0.8
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "porcelain_1":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Mix"].inputs[1].default_value
new_bs_color_r = bs_color[0] + random.uniform(-0.3, 0.3)
new_bs_color_g = bs_color[1] + random.uniform(-0.3, 0.3)
new_bs_color_b = bs_color[2] + random.uniform(-0.3, 0.3)
if new_bs_color_r < 0:
new_bs_color_r = 0.2
if new_bs_color_g < 0:
new_bs_color_g = 0.2
if new_bs_color_b < 0:
new_bs_color_b = 0.2
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Mix"].inputs[1].default_value = list(new_bs_color)
elif material_name == "porcelain_2":
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop#0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.8
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "porcelain_3":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix.001"].inputs[1])
else:
mat_nodes["Mix.001"].inputs[1].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Mix.001"].inputs[1].default_value
new_bs_color_r = bs_color[0] + random.uniform(-0.3, 0.3)
new_bs_color_g = bs_color[1] + random.uniform(-0.3, 0.3)
new_bs_color_b = bs_color[2] + random.uniform(-0.3, 0.3)
if new_bs_color_r < 0:
new_bs_color_r = 0.2
if new_bs_color_g < 0:
new_bs_color_g = 0.2
if new_bs_color_b < 0:
new_bs_color_b = 0.2
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Mix.001"].inputs[1].default_value = list(new_bs_color)
elif material_name == "porcelain_4":
if transfer_flag == True:
# if is_texture:
# mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
# mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
# else:
# mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
# mat_nodes["Glossy BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Glossy BSDF"].inputs[1].default_value = random.uniform(0.05, 0.15)
diff_new = mat_nodes.new(type='ShaderNodeBsdfDiffuse')
diff_new.name = 'Diffuse BSDF-new'
for key, input in enumerate(mat_nodes["Diffuse BSDF"].inputs):
diff_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Diffuse BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Diffuse BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
elif material_name == "porcelain_5":
if transfer_flag == True:
# if is_texture:
# mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
# mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
# else:
# mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
# mat_nodes["Glossy BSDF"].inputs[0].default_value = list(orign_base_color)
diff_new = mat_nodes.new(type='ShaderNodeBsdfDiffuse')
diff_new.name = 'Diffuse BSDF-new'
for key, input in enumerate(mat_nodes["Diffuse BSDF"].inputs):
diff_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Diffuse BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Diffuse BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
elif material_name == "porcelain_6":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Diffuse BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + random.uniform(-0.3, 0.3)
new_bs_color_g = bs_color[1] + random.uniform(-0.3, 0.3)
new_bs_color_b = bs_color[2] + random.uniform(-0.3, 0.3)
if new_bs_color_r < 0:
new_bs_color_r = 0.2
if new_bs_color_g < 0:
new_bs_color_g = 0.2
if new_bs_color_b < 0:
new_bs_color_b = 0.2
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(new_bs_color)
## plastic
elif material_name == "plastic_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_3":
mat_nodes["值(明度)"].outputs[0].default_value = random.uniform(0.05, 0.25)
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_5":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_6":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.012"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.021"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.022"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.033"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
mat_nodes["RGB.001"].outputs[0].default_value = list(orign_base_color)
"""
mat_nodes["RGB.002"].outputs[0].default_value = list(orign_base_color)
mat_nodes["RGB.003"].outputs[0].default_value = list(orign_base_color)
"""
## rubber
elif material_name == "rubber_0":
diff_new = mat_nodes.new(type='ShaderNodeBsdfDiffuse')
diff_new.name = 'Diffuse BSDF-new'
for key, input in enumerate(mat_nodes["Diffuse BSDF"].inputs):
diff_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Diffuse BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Diffuse BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
elif material_name == "rubber_1":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["RGB Curves"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "rubber_2":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["RGB Curves"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "rubber_3":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "rubber_4":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
######################## 2022.02.04
## plastic_specular
elif material_name == "plasticsp_0":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.001"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute"].inputs[0])
else:
mat_nodes["RGB.001"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB.001"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB.001"].outputs[0].default_value = list(new_bs_color)
elif material_name == "plasticsp_1":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
## paint_specular
elif material_name == "paintsp_0":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB"].outputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_1":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[2])
mat_links.new(tex_node.outputs[0], mat_nodes["Mix.001"].inputs[1])
mat_links.new(tex_node.outputs[0], mat_nodes["Hue Saturation Value"].inputs[4])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB"].outputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_2":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Group"].inputs[0])
else:
mat_nodes["Group"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Group"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Group"].inputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_3":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
else:
#bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color = [random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1), 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_4":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Glossy BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Glossy BSDF.001"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
mat_nodes["Glossy BSDF"].inputs[0].default_value = list(new_bs_color)
mat_nodes["Glossy BSDF.001"].inputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_5":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Invert"].inputs[1])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.002"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB"].outputs[0].default_value = list(new_bs_color)
## rubber
elif material_name == "rubber_5":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Mix.005"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
## plastic
elif material_name == "plastic_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_4":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_7":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_8":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Group"].inputs[0])
else:
mat_nodes["Group"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_9":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_10":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_11":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_12":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[2])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_13":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_14":
mat_nodes["Math.005"].inputs[1].default_value = random.uniform(0.05, 0.3)
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
## paper
elif material_name == "paper_0":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Mix.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "paper_1":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.8, 0.95)
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.8, 0.9)
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "paper_2":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.9, 0.95)
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.9, 0.95)
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Bright/Contrast"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
## leather
elif material_name == "leather_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
elif material_name == "leather_3":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_4":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_5":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["lether"].inputs[0])
else:
mat_nodes["lether"].inputs[0].default_value = list(orign_base_color)
## wood (全部不作处理，保留原始材质)
elif material_name == "wood_0":
pass
elif material_name == "wood_1":
pass
elif material_name == "wood_2":
pass
elif material_name == "wood_3":
pass
elif material_name == "wood_4":
pass
elif material_name == "wood_5":
pass
elif material_name == "wood_6":
pass
elif material_name == "wood_7":
pass
elif material_name == "wood_8":
pass
elif material_name == "wood_9":
pass
## fabric
elif material_name == "fabric_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "fabric_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "fabric_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
## clay
elif material_name == "clay_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_3":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
elif material_name == "clay_4":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_5":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
## glass
elif material_name == "glass_0":
mat_nodes["Mix Shader"].inputs[0].default_value = random.uniform(0.1, 0.3)
mat_nodes["Glossy BSDF"].inputs[1].default_value = random.uniform(0.1, 0.3)
elif material_name == "glass_4":
mat_nodes["Layer Weight"].inputs[0].default_value = random.uniform(0.3, 0.7)
mat_nodes["Glossy BSDF"].inputs[1].default_value = random.uniform(0.05, 0.2)
elif material_name == "glass_5":
mat_nodes["Layer Weight"].inputs[0].default_value = random.uniform(0.2, 0.4)
mat_nodes["Glossy BSDF"].inputs[1].default_value = random.uniform(0.0, 0.1)
elif material_name == "glass_14":
mat_nodes["Glass BSDF.005"].inputs[1].default_value = random.uniform(0.0, 0.1)
mat_nodes["Glass BSDF.006"].inputs[1].default_value = random.uniform(0.0, 0.1)
mat_nodes["Glass BSDF.007"].inputs[1].default_value = random.uniform(0.0, 0.1)
mat_nodes["Glass BSDF.008"].inputs[1].default_value = random.uniform(0.0, 0.1)
mat_nodes["Layer Weight"].inputs[0].default_value = random.uniform(0.81, 0.87)
mat_nodes["Layer Weight.001"].inputs[0].default_value = random.uniform(0.65, 0.71)
mat_nodes["Layer Weight.002"].inputs[0].default_value = random.uniform(0.81, 0.87)
color_value = random.uniform(0.599459, 0.70)
mat_nodes["Transparent BSDF"].inputs[0].default_value = list([color_value, color_value, color_value, 1])
# elif material_name == "glass_15":
# mat_nodes["Glass BSDF"].inputs[1].default_value = random.uniform(0.0, 0.1)
# mat_nodes["Glass BSDF"].inputs[2].default_value = random.uniform(1.325, 1.335)
# color_value = random.uniform(0.297, 0.35)
# mat_nodes["Transparent BSDF"].inputs[0].default_value = list([color_value, color_value, color_value, 1])
########################
else:
print(material_name + " no change")
def set_modify_material(obj, material, obj_texture_img_list, mat_randomize_mode, is_transfer=True):
for mat_slot in obj.material_slots:
if mat_slot.material:
if mat_slot.material.node_tree:
srcmat = material
mat = srcmat.copy()
mat.name = mat_slot.material.name # rename
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
bsdf_node = mat_slot.material.node_tree.nodes.get("Principled BSDF", None)
if bsdf_node is not None:
tex_node = mat_slot.material.node_tree.nodes.new(type='ShaderNodeTexImage')
tex_node.name = 'objtexture_tex'
tex_node.extension = 'EXTEND'
flag = random.randint(0, len(obj_texture_img_list)-1)
tex_node.image = obj_texture_img_list[flag]
tex_node_orign = mat_slot.material.node_tree.nodes.get('objtexture_tex', None)
if tex_node_orign is not None:
#mat = mat_slot.material.copy()
# Get the bl_idname to create a new node of the same type
tex_node = mat_nodes.new(tex_node_orign.bl_idname)
texture_img = bpy.data.images[tex_node_orign.image.name]
# Assign the default values from the old node to the new node
tex_node.image = texture_img
tex_node.projection = 'SPHERE'
#tex_node.location = Vector((-800, 0))
mapping_node = mat_nodes.new(type='ShaderNodeMapping')
mapping_node.name = 'objtexture_mapping'
texcoord_node = mat_nodes.new(type='ShaderNodeTexCoord')
texcoord_node.name = 'objtexture_texcoord'
mat_links.new(mapping_node.outputs[0], tex_node.inputs[0])
mat_links.new(texcoord_node.outputs[0], mapping_node.inputs[0])
modify_material(mat_links, mat_nodes, srcmat.name, mat_randomize_mode, is_texture=True, tex_node=tex_node, is_transfer=is_transfer)
else:
orign_base_color = mat_slot.material.node_tree.nodes["Principled BSDF"].inputs[0].default_value
if orign_base_color[0] == 0.0 and orign_base_color[1] == 0.0 and orign_base_color[2] == 0.0:
orign_base_color = [0.05, 0.05, 0.05, 1]
modify_material(mat_links, mat_nodes, srcmat.name, mat_randomize_mode, is_texture=False, orign_base_color=orign_base_color, is_transfer=is_transfer)
bpy.data.materials.remove(mat_slot.material)
mat_slot.material = mat
def set_modify_raw_material(obj):
for mat_slot in obj.material_slots:
if mat_slot.material:
if mat_slot.material.node_tree:
bsdf_node = mat_slot.material.node_tree.nodes.get("Principled BSDF", None)
if bsdf_node is not None:
tex_node_orign = mat_slot.material.node_tree.nodes.get("Image Texture", None)
if tex_node_orign is None:
orign_base_color = mat_slot.material.node_tree.nodes["Principled BSDF"].inputs[0].default_value
if orign_base_color[0] == 0.0 and orign_base_color[1] == 0.0 and orign_base_color[2] == 0.0:
mat = mat_slot.material.copy()
mat.name = mat_slot.material.name # rename
mat_nodes = mat.node_tree.nodes
mat_nodes["Principled BSDF"].inputs[0].default_value = list([0.05, 0.05, 0.05, 1])
bpy.data.materials.remove(mat_slot.material)
mat_slot.material = mat
def set_modify_table_material(obj, material, selected_realtable_img):
srcmat = material
#print(srcmat.name)
mat = srcmat.copy()
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
tex_node = mat_nodes.new(type='ShaderNodeTexImage')
tex_node.name = 'realtable_tex'
tex_node.extension = 'EXTEND'
tex_node.image = selected_realtable_img
mapping_node = mat_nodes.new(type='ShaderNodeMapping')
mapping_node.name = 'realtable_mapping'
texcoord_node = mat_nodes.new(type='ShaderNodeTexCoord')
texcoord_node.name = 'realtable_texcoord'
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
mat_links.new(mapping_node.outputs[0], tex_node.inputs[0])
mat_links.new(texcoord_node.outputs[2], mapping_node.inputs[0])
obj.active_material = mat
def set_modify_floor_material(obj, material, selected_realfloor_img):
srcmat = material
mat = srcmat.copy()
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
bsdfnode_list = [n for n in mat_nodes if isinstance(n, bpy.types.ShaderNodeBsdfPrincipled)]
if bsdfnode_list == []:
obj.active_material = material
else:
for bsdfnode in bsdfnode_list:
tex_node = mat_nodes.new(type='ShaderNodeTexImage')
tex_node.name = 'realfloor_tex'
tex_node.extension = 'REPEAT'
tex_node.image = selected_realfloor_img
mapping_node = mat_nodes.new(type='ShaderNodeMapping')
mapping_node.name = 'realfloor_mapping'
texcoord_node = mat_nodes.new(type='ShaderNodeTexCoord')
texcoord_node.name = 'realfloor_texcoord'
mat_links.new(tex_node.outputs[0], bsdfnode.inputs[0])
mat_links.new(mapping_node.outputs[0], tex_node.inputs[0])
mat_links.new(texcoord_node.outputs[2], mapping_node.inputs[0])
obj.active_material = mat
def set_modify_arm_material(obj, material):
for mat_slot in obj.material_slots:
if mat_slot.material:
if mat_slot.material.node_tree:
srcmat = material
mat = srcmat.copy()
mat.name = mat_slot.material.name # rename
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
rgb = random.uniform(0.50, 1.00)
orign_base_color = [rgb, rgb, rgb, 1]
modify_material(mat_links, mat_nodes, srcmat.name, mat_randomize_mode="diffuse", is_texture=False, orign_base_color=orign_base_color, is_transfer=False, is_arm=True)
bpy.data.materials.remove(mat_slot.material)
mat_slot.material = mat
================================================
FILE: data_generator/render_pile_STD_rand.py
================================================
import os
import random
import bpy
import math
import numpy as np
from mathutils import Vector, Matrix
import copy
import sys
import time
import argparse
from scipy.spatial.transform import Rotation
from bpy_extras.object_utils import world_to_camera_view
sys.path.append(os.getcwd())
from modify_material import set_modify_material, set_modify_raw_material, set_modify_table_material, set_modify_floor_material, set_modify_arm_material
argv = sys.argv
argv = argv[argv.index("--") + 1:] # get all args after "--"
RENDERING_PATH = os.getcwd()
LIGHT_EMITTER_ENERGY = 5
LIGHT_ENV_MAP_ENERGY_IR = 0.035
LIGHT_ENV_MAP_ENERGY_RGB = 1.0
CYCLES_SAMPLE = 32
DEVICE_LIST = [6]
SCENE_NUM = int(argv[0]) if len(argv) > 0 else 0
CAMERA_TYPE = "realsense"
NUM_FRAME_PER_SCENE = 24
RENDER_START_FRAME = 0
RENDER_END_FRAME = 24
RENDER_FRAMES_LIST = list(range(RENDER_START_FRAME, RENDER_END_FRAME))
look_at_shift = np.array([0,0,0])
num_point_ver = 6
num_point_hor = 4
beta_range = (random.uniform(12, 16)*math.pi/180, random.uniform(41, 45)*math.pi/180)
r = random.uniform(0.40, 0.50)
code_root = "/data/InterNeRF/renderer/renderer_giga_GPU6-0_rand_M"
emitter_pattern_path = os.path.join(code_root, "pattern", "test_pattern.png")
env_map_path = os.path.join(code_root, "envmap_lib")
default_background_texture_path = os.path.join(code_root, "texture", "texture_0.jpg")
table_CAD_model_path = os.path.join(code_root, "table_obj", "table.obj")
table_plane_CAD_model_path = os.path.join(code_root, "table_obj", "table_plane.obj")
arm_CAD_model_path = os.path.join(code_root, "table_obj", "arm.obj")
TABLE_CAD_MODEL_HEIGHT = 0.75
real_table_image_root_path = os.path.join(code_root, "realtable")
real_floor_image_root_path = os.path.join(code_root, "realfloor")
obj_texture_image_root_path = "/data/InterNeRF/data/imagenet"
obj_texture_image_idxfile = "train_paths.txt"
output_root_path = "/data/InterNeRF/data/giga_hemisphere_train_0827/pile_full/6_M_rand/pile_%d-%d"%(830*(SCENE_NUM//830), 830*(SCENE_NUM//830)+830)
raw_urdfs_and_poses_dir_path = "/data/InterNeRF/data/GIGA/data_pile_train_raw/mesh_pose_list"
# render mode
render_mode_list = {'RGB': 1,
'IR': 0,
'NOCS': 0,
'Mask': 1,
'Normal': 1}
# material randomization mode (diffuse, transparent, specular, specular_tex, specular_texmix, mixed)
my_material_randomize_mode = 'mixed'
# set depth sensor parameter
camera_width = 640
camera_height = 360
camera_fov = 69.75 / 180 * math.pi
baseline_distance = 0.055
# set background parameter
background_size = 3.
background_position = (0., 0., 0.)
background_scale = (1., 1., 1.)
# set camera randomize paramater
# start_point_range: (range_r, range_vector), range_r: (r_min, r_max), range_vector: (x_min, x_max, y_min, y_max)
# look_at_range: (x_min, x_max, y_min, y_max, z_min, z_max)
# up_range: (x_min, x_max, y_min, y_max)
start_point_range = ((0.5, 0.95), (-0.6, 0.6, -0.6, 0.6))
up_range = (-0.18, -0.18, -0.18, 0.18)
look_at_range = (background_position[0] - 0.05, background_position[0] + 0.05,
background_position[1] - 0.05, background_position[1] + 0.05,
background_position[2] - 0.05, background_position[2] + 0.05)
g_syn_light_num_lowbound = 4
g_syn_light_num_highbound = 6
g_syn_light_dist_lowbound = 8
g_syn_light_dist_highbound = 12
g_syn_light_azimuth_degree_lowbound = 0
g_syn_light_azimuth_degree_highbound = 360
g_syn_light_elevation_degree_lowbound = 0
g_syn_light_elevation_degree_highbound = 90
g_syn_light_energy_mean = 3
g_syn_light_energy_std = 0.5
g_syn_light_environment_energy_lowbound = 0
g_syn_light_environment_energy_highbound = 1
def obj_centered_camera_pos(dist, azimuth_deg, elevation_deg):
phi = float(elevation_deg) / 180 * math.pi
theta = float(azimuth_deg) / 180 * math.pi
x = (dist * math.cos(theta) * math.cos(phi))
y = (dist * math.sin(theta) * math.cos(phi))
z = (dist * math.sin(phi))
return (x, y, z)
def quaternionFromYawPitchRoll(yaw, pitch, roll):
c1 = math.cos(yaw / 2.0)
c2 = math.cos(pitch / 2.0)
c3 = math.cos(roll / 2.0)
s1 = math.sin(yaw / 2.0)
s2 = math.sin(pitch / 2.0)
s3 = math.sin(roll / 2.0)
q1 = c1 * c2 * c3 + s1 * s2 * s3
q2 = c1 * c2 * s3 - s1 * s2 * c3
q3 = c1 * s2 * c3 + s1 * c2 * s3
q4 = s1 * c2 * c3 - c1 * s2 * s3
return (q1, q2, q3, q4)
def camPosToQuaternion(cx, cy, cz):
q1a = 0
q1b = 0
q1c = math.sqrt(2) / 2
q1d = math.sqrt(2) / 2
camDist = math.sqrt(cx * cx + cy * cy + cz * cz)
cx = cx / camDist
cy = cy / camDist
cz = cz / camDist
t = math.sqrt(cx * cx + cy * cy)
tx = cx / t
ty = cy / t
yaw = math.acos(ty)
if tx > 0:
yaw = 2 * math.pi - yaw
pitch = 0
tmp = min(max(tx*cx + ty*cy, -1),1)
#roll = math.acos(tx * cx + ty * cy)
roll = math.acos(tmp)
if cz < 0:
roll = -roll
print("%f %f %f" % (yaw, pitch, roll))
q2a, q2b, q2c, q2d = quaternionFromYawPitchRoll(yaw, pitch, roll)
q1 = q1a * q2a - q1b * q2b - q1c * q2c - q1d * q2d
q2 = q1b * q2a + q1a * q2b + q1d * q2c - q1c * q2d
q3 = q1c * q2a - q1d * q2b + q1a * q2c + q1b * q2d
q4 = q1d * q2a + q1c * q2b - q1b * q2c + q1a * q2d
return (q1, q2, q3, q4)
def camRotQuaternion(cx, cy, cz, theta):
theta = theta / 180.0 * math.pi
camDist = math.sqrt(cx * cx + cy * cy + cz * cz)
cx = -cx / camDist
cy = -cy / camDist
cz = -cz / camDist
q1 = math.cos(theta * 0.5)
q2 = -cx * math.sin(theta * 0.5)
q3 = -cy * math.sin(theta * 0.5)
q4 = -cz * math.sin(theta * 0.5)
return (q1, q2, q3, q4)
def quaternionProduct(qx, qy):
a = qx[0]
b = qx[1]
c = qx[2]
d = qx[3]
e = qy[0]
f = qy[1]
g = qy[2]
h = qy[3]
q1 = a * e - b * f - c * g - d * h
q2 = a * f + b * e + c * h - d * g
q3 = a * g - b * h + c * e + d * f
q4 = a * h + b * g - c * f + d * e
return (q1, q2, q3, q4)
def quaternionToRotation(q):
w, x, y, z = q
r00 = 1 - 2 * y ** 2 - 2 * z ** 2
r01 = 2 * x * y + 2 * w * z
r02 = 2 * x * z - 2 * w * y
r10 = 2 * x * y - 2 * w * z
r11 = 1 - 2 * x ** 2 - 2 * z ** 2
r12 = 2 * y * z + 2 * w * x
r20 = 2 * x * z + 2 * w * y
r21 = 2 * y * z - 2 * w * x
r22 = 1 - 2 * x ** 2 - 2 * y ** 2
r = [[r00, r01, r02], [r10, r11, r12], [r20, r21, r22]]
return r
def quaternionToRotation_xyzw(q):
x, y, z, w = q
r00 = 1 - 2 * y ** 2 - 2 * z ** 2
r01 = 2 * x * y + 2 * w * z
r02 = 2 * x * z - 2 * w * y
r10 = 2 * x * y - 2 * w * z
r11 = 1 - 2 * x ** 2 - 2 * z ** 2
r12 = 2 * y * z + 2 * w * x
r20 = 2 * x * z + 2 * w * y
r21 = 2 * y * z - 2 * w * x
r22 = 1 - 2 * x ** 2 - 2 * y ** 2
r = [[r00, r01, r02], [r10, r11, r12], [r20, r21, r22]]
return r
def quaternionFromRotMat(rotation_matrix):
rotation_matrix = np.reshape(rotation_matrix, (1, 9))[0]
w = math.sqrt(rotation_matrix[0]+rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
x = math.sqrt(rotation_matrix[0]-rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
y = math.sqrt(-rotation_matrix[0]+rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
z = math.sqrt(-rotation_matrix[0]-rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
a = [w,x,y,z]
m = a.index(max(a))
if m == 0:
x = (rotation_matrix[7]-rotation_matrix[5])/(4*w)
y = (rotation_matrix[2]-rotation_matrix[6])/(4*w)
z = (rotation_matrix[3]-rotation_matrix[1])/(4*w)
if m == 1:
w = (rotation_matrix[7]-rotation_matrix[5])/(4*x)
y = (rotation_matrix[1]+rotation_matrix[3])/(4*x)
z = (rotation_matrix[6]+rotation_matrix[2])/(4*x)
if m == 2:
w = (rotation_matrix[2]-rotation_matrix[6])/(4*y)
x = (rotation_matrix[1]+rotation_matrix[3])/(4*y)
z = (rotation_matrix[5]+rotation_matrix[7])/(4*y)
if m == 3:
w = (rotation_matrix[3]-rotation_matrix[1])/(4*z)
x = (rotation_matrix[6]+rotation_matrix[2])/(4*z)
y = (rotation_matrix[5]+rotation_matrix[7])/(4*z)
quaternion = (w,x,y,z)
return quaternion
def quaternionFromRotMat_xyzw(rotation_matrix):
rotation_matrix = np.reshape(rotation_matrix, (1, 9))[0]
w = math.sqrt(rotation_matrix[0]+rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
x = math.sqrt(rotation_matrix[0]-rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
y = math.sqrt(-rotation_matrix[0]+rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
z = math.sqrt(-rotation_matrix[0]-rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
a = [x,y,z,w]
m = a.index(max(a))
if m == 0:
x = (rotation_matrix[7]-rotation_matrix[5])/(4*w)
y = (rotation_matrix[2]-rotation_matrix[6])/(4*w)
z = (rotation_matrix[3]-rotation_matrix[1])/(4*w)
if m == 1:
w = (rotation_matrix[7]-rotation_matrix[5])/(4*x)
y = (rotation_matrix[1]+rotation_matrix[3])/(4*x)
z = (rotation_matrix[6]+rotation_matrix[2])/(4*x)
if m == 2:
w = (rotation_matrix[2]-rotation_matrix[6])/(4*y)
x = (rotation_matrix[1]+rotation_matrix[3])/(4*y)
z = (rotation_matrix[5]+rotation_matrix[7])/(4*y)
if m == 3:
w = (rotation_matrix[3]-rotation_matrix[1])/(4*z)
x = (rotation_matrix[6]+rotation_matrix[2])/(4*z)
y = (rotation_matrix[5]+rotation_matrix[7])/(4*z)
quaternion = (x,y,z,w)
return quaternion
def rotVector(q, vector_ori):
r = quaternionToRotation(q)
x_ori = vector_ori[0]
y_ori = vector_ori[1]
z_ori = vector_ori[2]
x_rot = r[0][0] * x_ori + r[1][0] * y_ori + r[2][0] * z_ori
y_rot = r[0][1] * x_ori + r[1][1] * y_ori + r[2][1] * z_ori
z_rot = r[0][2] * x_ori + r[1][2] * y_ori + r[2][2] * z_ori
return (x_rot, y_rot, z_rot)
def cameraLPosToCameraRPos(q_l, pos_l, baseline_dis):
vector_camera_l_y = (1, 0, 0)
vector_rot = rotVector(q_l, vector_camera_l_y)
pos_r = (pos_l[0] + vector_rot[0] * baseline_dis,
pos_l[1] + vector_rot[1] * baseline_dis,
pos_l[2] + vector_rot[2] * baseline_dis)
return pos_r
def getRTFromAToB(pointCloudA, pointCloudB):
muA = np.mean(pointCloudA, axis=0)
muB = np.mean(pointCloudB, axis=0)
zeroMeanA = pointCloudA - muA
zeroMeanB = pointCloudB - muB
covMat = np.matmul(np.transpose(zeroMeanA), zeroMeanB)
U, S, Vt = np.linalg.svd(covMat)
R = np.matmul(Vt.T, U.T)
if np.linalg.det(R) < 0:
print("Reflection detected")
Vt[2, :] *= -1
R = Vt.T * U.T
T = (-np.matmul(R, muA.T) + muB.T).reshape(3, 1)
return R, T
def cameraPositionRandomize(start_point_range, look_at_range, up_range):
r_range, vector_range = start_point_range
r_min, r_max = r_range
x_min, x_max, y_min, y_max = vector_range
r = random.uniform(r_min, r_max)
x = random.uniform(x_min, x_max)
y = random.uniform(y_min, y_max)
z = math.sqrt(1 - x**2 - y**2)
vector_camera_axis = np.array([x, y, z])
x_min, x_max, y_min, y_max = up_range
x = random.uniform(x_min, x_max)
y = random.uniform(y_min, y_max)
z = math.sqrt(1 - x**2 - y**2)
up = np.array([x, y, z])
x_min, x_max, y_min, y_max, z_min, z_max = look_at_range
look_at = np.array([random.uniform(x_min, x_max),
random.uniform(y_min, y_max),
random.uniform(z_min, z_max)])
position = look_at + r * vector_camera_axis
vectorZ = - (look_at - position)/np.linalg.norm(look_at - position)
vectorX = np.cross(up, vectorZ)/np.linalg.norm(np.cross(up, vectorZ))
vectorY = np.cross(vectorZ, vectorX)/np.linalg.norm(np.cross(vectorX, vectorZ))
# points in camera coordinates
pointSensor= np.array([[0., 0., 0.], [1., 0., 0.], [0., 2., 0.], [0., 0., 3.]])
# points in world coordinates
pointWorld = np.array([position,
position + vectorX,
position + vectorY * 2,
position + vectorZ * 3])
resR, resT = getRTFromAToB(pointSensor, pointWorld)
resQ = quaternionFromRotMat(resR)
return resQ, resT
def add_noise_to_transformation_matrix(Rot, Trans, angle_std=2, translation_std=0.01):
axis = np.random.rand(3)
axis /= np.linalg.norm(axis)
angle = np.random.uniform(0, angle_std) / 180 * np.pi
Rot = Rotation.from_rotvec(angle * axis).as_matrix() @ Rot
direction = np.random.rand(3)
direction /= np.linalg.norm(direction)
length = np.random.uniform(0, translation_std)
Trans += direction.reshape(Trans.shape) * length
return Rot, Trans
def genCameraPosition(look_at):
quat_list = []
rot_list = []
trans_list = []
position_list = []
alpha = 0
alpha_delta = (2 * math.pi) / num_point_ver
for i in range(num_point_ver):
alpha = alpha + alpha_delta
flag_x = 1
flag_y = 1
alpha1 = alpha
if alpha > math.pi/2 and alpha <= math.pi:
alpha1 = math.pi - alpha #alpha - math.pi/2
flag_x = -1
flag_y = 1
elif alpha > math.pi and alpha <= math.pi*(3/2):
alpha1 = alpha - math.pi #math.pi*(3/2) - alpha
flag_x = -1
flag_y = -1
elif alpha > math.pi*(3/2):
alpha1 = math.pi*2 - alpha #alpha - math.pi*(3/2)
flag_x = 1
flag_y = -1
beta = beta_range[0]
beta_delta = (beta_range[1]-beta_range[0])/(num_point_hor-1)
for j in range(num_point_hor):
if j != 0:
beta = beta + beta_delta
x = flag_x * (r * math.sin(beta)) * math.cos(alpha1)
y = flag_y * (r * math.sin(beta)) * math.sin(alpha1)
z = r * math.cos(beta)
position = np.array([x, y, z]) + look_at
look_at = look_at
up = np.array([0, 0, 1])
vectorZ = - (look_at - position)/np.linalg.norm(look_at - position)
vectorX = np.cross(up, vectorZ)/np.linalg.norm(np.cross(up, vectorZ))
vectorY = np.cross(vectorZ, vectorX)/np.linalg.norm(np.cross(vectorX, vectorZ))
# points in camera coordinates
pointSensor= np.array([[0., 0., 0.], [1., 0., 0.], [0., 2., 0.], [0., 0., 3.]])
# points in world coordinates
pointWorld = np.array([position,
position + vectorX,
position + vectorY * 2,
position + vectorZ * 3])
# get R and T
resR, resT = getRTFromAToB(pointSensor, pointWorld)
# add noise
resR, resT = add_noise_to_transformation_matrix(resR, resT)
# add to list
resQ = quaternionFromRotMat(resR)
quat_list.append(resQ)
rot_list.append(resR)
trans_list.append(resT)
#position_list.append(position)
position_list.append(resT.reshape(3))
return quat_list, trans_list, rot_list, position_list
def quanternion_mul(q1, q2):
s1 = q1[0]
v1 = np.array(q1[1:])
s2 = q2[0]
v2 = np.array(q2[1:])
s = s1 * s2 - np.dot(v1, v2)
v = s1 * v2 + s2 * v1 + np.cross(v1, v2)
return (s, v[0], v[1], v[2])
class BlenderRenderer(object):
def __init__(self, viewport_size_x=640, viewport_size_y=360):
'''
viewport_size_x, viewport_size_y: rendering viewport resolution
'''
# remove all objects, cameras and lights
for obj in bpy.data.meshes:
bpy.data.meshes.remove(obj)
for cam in bpy.data.cameras:
bpy.data.cameras.remove(cam)
for light in bpy.data.lights:
bpy.data.lights.remove(light)
for obj in bpy.data.objects:
bpy.data.objects.remove(obj, do_unlink=True)
# remove all materials
# for item in bpy.data.materials:
# bpy.data.materials.remove(item)
render_context = bpy.context.scene.render
# add left camera
camera_l_data = bpy.data.cameras.new(name="camera_l")
camera_l_object = bpy.data.objects.new(name="camera_l", object_data=camera_l_data)
bpy.context.collection.objects.link(camera_l_object)
# add right camera
camera_r_data = bpy.data.cameras.new(name="camera_r")
camera_r_object = bpy.data.objects.new(name="camera_r", object_data=camera_r_data)
bpy.context.collection.objects.link(camera_r_object)
camera_l = bpy.data.objects["camera_l"]
camera_r = bpy.data.objects["camera_r"]
# set the camera postion and orientation so that it is in
# the front of the object
camera_l.location = (1, 0, 0)
camera_r.location = (1, 0, 0)
# add emitter light
light_emitter_data = bpy.data.lights.new(name="light_emitter", type='SPOT')
light_emitter_object = bpy.data.objects.new(name="light_emitter", object_data=light_emitter_data)
bpy.context.collection.objects.link(light_emitter_object)
light_emitter = bpy.data.objects["light_emitter"]
light_emitter.location = (1, 0, 0)
light_emitter.data.energy = LIGHT_EMITTER_ENERGY
# render setting
render_context.resolution_percentage = 100
self.render_context = render_context
self.camera_l = camera_l
self.camera_r = camera_r
self.light_emitter = light_emitter
self.model_loaded = False
self.background_added = None
self.render_context.resolution_x = viewport_size_x
self.render_context.resolution_y = viewport_size_y
self.my_material = {}
self.render_mode = 'IR'
# output setting
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.compression = 0
self.render_context.image_settings.color_mode = 'BW'
self.render_context.image_settings.color_depth = '8'
# cycles setting
self.render_context.engine = 'CYCLES'
bpy.context.scene.cycles.progressive = 'BRANCHED_PATH'
bpy.context.scene.cycles.use_denoising = True
bpy.context.scene.cycles.denoiser = 'NLM'
bpy.context.scene.cycles.film_exposure = 0.5
# self.render_context.use_antialiasing = False
##########
bpy.context.scene.view_layers["View Layer"].use_sky = True
##########
# switch on nodes
bpy.context.scene.use_nodes = True
tree = bpy.context.scene.node_tree
links = tree.links
# clear default nodes
for n in tree.nodes:
tree.nodes.remove(n)
# create input render layer node
rl = tree.nodes.new('CompositorNodeRLayers')
# create output node
self.fileOutput = tree.nodes.new(type="CompositorNodeOutputFile")
self.fileOutput.base_path = "./new_data/0000"
self.fileOutput.format.file_format = 'OPEN_EXR'
self.fileOutput.format.color_depth= '32'
self.fileOutput.file_slots[0].path = 'depth#'
# links.new(map.outputs[0], fileOutput.inputs[0])
links.new(rl.outputs[2], self.fileOutput.inputs[0])
# links.new(gamma.outputs[0], fileOutput.inputs[0])
# depth sensor pattern
self.pattern = []
# environment map
self.env_map = []
###
self.realtable_img_list = []
self.realfloor_img_list = []
self.obj_texture_img_list = []
###
def loadImages(self, pattern_path, env_map_path, real_table_image_root_path, real_floor_image_root_path):
# load pattern image
self.pattern = bpy.data.images.load(filepath=pattern_path)
# load env map
for item in os.listdir(env_map_path):
if item.split('.')[-1] == 'hdr':
self.env_map.append(bpy.data.images.load(filepath=os.path.join(env_map_path, item)))
###
# load real table images
for item in os.listdir(real_table_image_root_path):
if item.split('.')[-1] == 'jpg':
self.realtable_img_list.append(bpy.data.images.load(filepath=os.path.join(real_table_image_root_path, item)))
# load real floor images
for item in os.listdir(real_floor_image_root_path):
if item.split('.')[-1] == 'jpg':
self.realfloor_img_list.append(bpy.data.images.load(filepath=os.path.join(real_floor_image_root_path, item)))
# load obj texture images
f_teximg_idx = open(os.path.join(obj_texture_image_root_path,obj_texture_image_idxfile),"r")
lines = f_teximg_idx.readlines()
# for item in lines:
# item = item[:-1] # 去掉"\n"
# #for item in os.listdir(obj_texture_image_root_path):
# #if item.split('.')[-1] == 'jpg':
# self.obj_texture_img_list.append(bpy.data.images.load(filepath=os.path.join(obj_texture_image_root_path, "images", item)))
start = random.randint(0,99900)
end = start+100
for item in lines[start:end]:
item = item[:-1] # 去掉"\n"
self.obj_texture_img_list.append(bpy.data.images.load(filepath=os.path.join(obj_texture_image_root_path, "images", item)))
###
def addEnvMap(self):
# Get the environment node tree of the current scene
node_tree = bpy.context.scene.world.node_tree
tree_nodes = node_tree.nodes
# Clear all nodes
tree_nodes.clear()
# Add Background node
node_background = tree_nodes.new(type='ShaderNodeBackground')
# Add Environment Texture node
node_environment = tree_nodes.new('ShaderNodeTexEnvironment')
# Load and assign the image to the node property
# node_environment.image = bpy.data.images.load("/Users/zhangjiyao/Desktop/test_addon/envmap_lib/autoshop_01_1k.hdr") # Relative path
node_environment.location = -300,0
node_tex_coord = tree_nodes.new(type='ShaderNodeTexCoord')
node_tex_coord.location = -700,0
node_mapping = tree_nodes.new(type='ShaderNodeMapping')
node_mapping.location = -500,0
# Add Output node
node_output = tree_nodes.new(type='ShaderNodeOutputWorld')
node_output.location = 200,0
# Link all nodes
links = node_tree.links
links.new(node_environment.outputs["Color"], node_background.inputs["Color"])
links.new(node_background.outputs["Background"], node_output.inputs["Surface"])
links.new(node_tex_coord.outputs["Generated"], node_mapping.inputs["Vector"])
links.new(node_mapping.outputs["Vector"], node_environment.inputs["Vector"])
#### bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = 1.0
random_energy = random.uniform(LIGHT_ENV_MAP_ENERGY_RGB * 0.8, LIGHT_ENV_MAP_ENERGY_RGB * 1.2)
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = random_energy
####
def setEnvMap(self, env_map_id, rotation_elur_z):
# Get the environment node tree of the current scene
node_tree = bpy.context.scene.world.node_tree
# Get Environment Texture node
node_environment = node_tree.nodes['Environment Texture']
# Load and assign the image to the node property
node_environment.image = self.env_map[env_map_id]
node_mapping = node_tree.nodes['Mapping']
node_mapping.inputs[2].default_value[2] = rotation_elur_z
def addMaskMaterial(self, num=20):
background_material_name_list = ["mask_background", "mask_table", "mask_tableplane", "mask_arm"]
for material_name in background_material_name_list:
material_class = (bpy.data.materials.get(material_name) or bpy.data.materials.new(material_name)) # test if material exists, if it does not exist, create it:
# enable 'Use nodes'
material_class.use_nodes = True
node_tree = material_class.node_tree
# remove default nodes
material_class.node_tree.nodes.clear()
# add new nodes
node_1 = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_2= node_tree.nodes.new('ShaderNodeBrightContrast')
# link nodes
node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
node_2.inputs[0].default_value = (1, 1, 1, 1)
self.my_material[material_name] = material_class
#print("##############################", self.my_material)
for i in range(num):
class_name = str(i + 1)
# set the material of background
material_name = "mask_" + class_name
# test if material exists
# if it does not exist, create it:
material_class = (bpy.data.materials.get(material_name) or
bpy.data.materials.new(material_name))
# enable 'Use nodes'
material_class.use_nodes = True
node_tree = material_class.node_tree
# remove default nodes
material_class.node_tree.nodes.clear()
# add new nodes
node_1 = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_2= node_tree.nodes.new('ShaderNodeBrightContrast')
# link nodes
node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
if class_name.split('_')[0] == 'background' or class_name.split('_')[0] == 'table' or class_name.split('_')[0] == 'tableplane' or class_name.split('_')[0] == 'arm':
node_2.inputs[0].default_value = (1, 1, 1, 1)
else:
node_2.inputs[0].default_value = ((i + 1)/255., 0., 0., 1)
self.my_material[material_name] = material_class
def addNOCSMaterial(self):
material_name = 'coord_color'
mat = (bpy.data.materials.get(material_name) or bpy.data.materials.new(material_name))
mat.use_nodes = True
node_tree = mat.node_tree
nodes = node_tree.nodes
nodes.clear()
links = node_tree.links
links.clear()
vcol_R = nodes.new(type="ShaderNodeVertexColor")
vcol_R.layer_name = "Col_R" # the vertex color layer name
vcol_G = nodes.new(type="ShaderNodeVertexColor")
vcol_G.layer_name = "Col_G" # the vertex color layer name
vcol_B = nodes.new(type="ShaderNodeVertexColor")
vcol_B.layer_name = "Col_B" # the vertex color layer name
node_Output = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_Emission = node_tree.nodes.new('ShaderNodeEmission')
node_LightPath = node_tree.nodes.new('ShaderNodeLightPath')
node_Mix = node_tree.nodes.new('ShaderNodeMixShader')
node_Combine = node_tree.nodes.new(type="ShaderNodeCombineRGB")
# make links
node_tree.links.new(vcol_R.outputs[1], node_Combine.inputs[0])
node_tree.links.new(vcol_G.outputs[1], node_Combine.inputs[1])
node_tree.links.new(vcol_B.outputs[1], node_Combine.inputs[2])
node_tree.links.new(node_Combine.outputs[0], node_Emission.inputs[0])
node_tree.links.new(node_LightPath.outputs[0], node_Mix.inputs[0])
node_tree.links.new(node_Emission.outputs[0], node_Mix.inputs[2])
node_tree.links.new(node_Mix.outputs[0], node_Output.inputs[0])
self.my_material[material_name] = mat
def addNormalMaterial(self):
material_name = 'normal'
mat = (bpy.data.materials.get(material_name) or bpy.data.materials.new(material_name))
mat.use_nodes = True
node_tree = mat.node_tree
nodes = node_tree.nodes
nodes.clear()
links = node_tree.links
links.clear()
# Nodes :
new_node = nodes.new(type='ShaderNodeMath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (151.59744262695312, 854.5482177734375)
new_node.name = 'Math'
new_node.operation = 'MULTIPLY'
new_node.select = False
new_node.use_clamp = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node.inputs[1].default_value = 1.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = 0.0
new_node = nodes.new(type='ShaderNodeLightPath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (602.9912719726562, 1046.660888671875)
new_node.name = 'Light Path'
new_node.select = False
new_node.width = 140.0
new_node.outputs[0].default_value = 0.0
new_node.outputs[1].default_value = 0.0
new_node.outputs[2].default_value = 0.0
new_node.outputs[3].default_value = 0.0
new_node.outputs[4].default_value = 0.0
new_node.outputs[5].default_value = 0.0
new_node.outputs[6].default_value = 0.0
new_node.outputs[7].default_value = 0.0
new_node.outputs[8].default_value = 0.0
new_node.outputs[9].default_value = 0.0
new_node.outputs[10].default_value = 0.0
new_node.outputs[11].default_value = 0.0
new_node.outputs[12].default_value = 0.0
new_node = nodes.new(type='ShaderNodeOutputMaterial')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.is_active_output = True
new_node.location = (1168.93017578125, 701.84033203125)
new_node.name = 'Material Output'
new_node.select = False
new_node.target = 'ALL'
new_node.width = 140.0
new_node.inputs[2].default_value = [0.0, 0.0, 0.0]
new_node = nodes.new(type='ShaderNodeBsdfTransparent')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (731.72900390625, 721.4832763671875)
new_node.name = 'Transparent BSDF'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = [1.0, 1.0, 1.0, 1.0]
new_node = nodes.new(type='ShaderNodeCombineXYZ')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (594.4229736328125, 602.9271240234375)
new_node.name = 'Combine XYZ'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.0
new_node.inputs[1].default_value = 0.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
new_node = nodes.new(type='ShaderNodeMixShader')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (992.7239990234375, 707.2142333984375)
new_node.name = 'Mix Shader'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node = nodes.new(type='ShaderNodeEmission')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (774.0802612304688, 608.2547607421875)
new_node.name = 'Emission'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = [1.0, 1.0, 1.0, 1.0]
new_node.inputs[1].default_value = 1.0
new_node = nodes.new(type='ShaderNodeSeparateXYZ')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (-130.12167358398438, 558.1497802734375)
new_node.name = 'Separate XYZ'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = [0.0, 0.0, 0.0]
new_node.outputs[0].default_value = 0.0
new_node.outputs[1].default_value = 0.0
new_node.outputs[2].default_value = 0.0
new_node = nodes.new(type='ShaderNodeMath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (162.43240356445312, 618.8094482421875)
new_node.name = 'Math.002'
new_node.operation = 'MULTIPLY'
new_node.select = False
new_node.use_clamp = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node.inputs[1].default_value = 1.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = 0.0
new_node = nodes.new(type='ShaderNodeMath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (126.8158187866211, 364.5539855957031)
new_node.name = 'Math.001'
new_node.operation = 'MULTIPLY'
new_node.select = False
new_node.use_clamp = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node.inputs[1].default_value = -1.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = 0.0
new_node = nodes.new(type='ShaderNodeVectorTransform')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.convert_from = 'WORLD'
new_node.convert_to = 'CAMERA'
new_node.location = (-397.0209045410156, 594.7037353515625)
new_node.name = 'Vector Transform'
new_node.select = False
new_node.vector_type = 'VECTOR'
new_node.width = 140.0
new_node.inputs[0].default_value = [0.5, 0.5, 0.5]
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
new_node = nodes.new(type='ShaderNodeNewGeometry')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (-651.8067016601562, 593.0455932617188)
new_node.name = 'Geometry'
new_node.width = 140.0
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
new_node.outputs[1].default_value = [0.0, 0.0, 0.0]
new_node.outputs[2].default_value = [0.0, 0.0, 0.0]
new_node.outputs[3].default_value = [0.0, 0.0, 0.0]
new_node.outputs[4].default_value = [0.0, 0.0, 0.0]
new_node.outputs[5].default_value = [0.0, 0.0, 0.0]
new_node.outputs[6].default_value = 0.0
new_node.outputs[7].default_value = 0.0
new_node.outputs[8].default_value = 0.0
# Links :
links.new(nodes["Light Path"].outputs[0], nodes["Mix Shader"].inputs[0])
links.new(nodes["Separate XYZ"].outputs[0], nodes["Math"].inputs[0])
links.new(nodes["Separate XYZ"].outputs[1], nodes["Math.002"].inputs[0])
links.new(nodes["Separate XYZ"].outputs[2], nodes["Math.001"].inputs[0])
links.new(nodes["Vector Transform"].outputs[0], nodes["Separate XYZ"].inputs[0])
links.new(nodes["Combine XYZ"].outputs[0], nodes["Emission"].inputs[0])
links.new(nodes["Math"].outputs[0], nodes["Combine XYZ"].inputs[0])
links.new(nodes["Math.002"].outputs[0], nodes["Combine XYZ"].inputs[1])
links.new(nodes["Math.001"].outputs[0], nodes["Combine XYZ"].inputs[2])
links.new(nodes["Transparent BSDF"].outputs[0], nodes["Mix Shader"].inputs[1])
links.new(nodes["Emission"].outputs[0], nodes["Mix Shader"].inputs[2])
links.new(nodes["Mix Shader"].outputs[0], nodes["Material Output"].inputs[0])
links.new(nodes["Geometry"].outputs[1], nodes["Vector Transform"].inputs[0])
self.my_material[material_name] = mat
def addMaterialLib(self, material_class_instance_pairs):
for mat in bpy.data.materials:
name = mat.name
name_class = str(name.split('_')[0])
if name_class != 'Dots Stroke' and name_class != 'default':
#print(name)
if name_class not in self.my_material:
self.my_material[name_class] = [mat]
else:
self.my_material[name_class].append(mat) # e.g. self.my_material['metal'] = [.....]
def setCamera(self, quaternion, translation, fov, baseline_distance):
self.camera_l.data.angle = fov
self.camera_r.data.angle = self.camera_l.data.angle
cx = translation[0]
cy = translation[1]
cz = translation[2]
self.camera_l.location[0] = cx
self.camera_l.location[1] = cy
self.camera_l.location[2] = cz
self.camera_l.rotation_mode = 'QUATERNION'
self.camera_l.rotation_quaternion[0] = quaternion[0]
self.camera_l.rotation_quaternion[1] = quaternion[1]
self.camera_l.rotation_quaternion[2] = quaternion[2]
self.camera_l.rotation_quaternion[3] = quaternion[3]
self.camera_r.rotation_mode = 'QUATERNION'
self.camera_r.rotation_quaternion[0] = quaternion[0]
self.camera_r.rotation_quaternion[1] = quaternion[1]
self.camera_r.rotation_quaternion[2] = quaternion[2]
self.camera_r.rotation_quaternion[3] = quaternion[3]
cx, cy, cz = cameraLPosToCameraRPos(quaternion, (cx, cy, cz), baseline_distance)
self.camera_r.location[0] = cx
self.camera_r.location[1] = cy
self.camera_r.location[2] = cz
def setLighting(self):
# emitter
#self.light_emitter.location = self.camera_r.location
self.light_emitter.location = self.camera_l.location + 0.51 * (self.camera_r.location - self.camera_l.location)
self.light_emitter.rotation_mode = 'QUATERNION'
self.light_emitter.rotation_quaternion = self.camera_r.rotation_quaternion
# emitter setting
bpy.context.view_layer.objects.active = None
# bpy.ops.object.select_all(action="DESELECT")
self.render_context.engine = 'CYCLES'
self.light_emitter.select_set(True)
self.light_emitter.data.use_nodes = True
self.light_emitter.data.type = "POINT"
self.light_emitter.data.shadow_soft_size = 0.001
random_energy = random.uniform(LIGHT_EMITTER_ENERGY * 0.9, LIGHT_EMITTER_ENERGY * 1.1)
self.light_emitter.data.energy = random_energy
# remove default node
light_emitter = bpy.data.objects["light_emitter"].data
light_emitter.node_tree.nodes.clear()
# light_projector.node_tree.nodes.remove(light_projector.node_tree.nodes.get('光输出')) #title of the existing node when materials.new
# light_projector.node_tree.nodes.remove(light_projector.node_tree.nodes.get('自发光(发射)')) #title of the existing node when materials.new
# add new nodes
light_output = light_emitter.node_tree.nodes.new("ShaderNodeOutputLight")
node_1 = light_emitter.node_tree.nodes.new("ShaderNodeEmission")
node_2 = light_emitter.node_tree.nodes.new("ShaderNodeTexImage")
node_3 = light_emitter.node_tree.nodes.new("ShaderNodeMapping")
node_4 = light_emitter.node_tree.nodes.new("ShaderNodeVectorMath")
node_5 = light_emitter.node_tree.nodes.new("ShaderNodeSeparateXYZ")
node_6 = light_emitter.node_tree.nodes.new("ShaderNodeTexCoord")
# link nodes
light_emitter.node_tree.links.new(light_output.inputs[0], node_1.outputs[0])
light_emitter.node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
light_emitter.node_tree.links.new(node_2.inputs[0], node_3.outputs[0])
light_emitter.node_tree.links.new(node_3.inputs[0], node_4.outputs[0])
light_emitter.node_tree.links.new(node_4.inputs[0], node_6.outputs[1])
light_emitter.node_tree.links.new(node_4.inputs[1], node_5.outputs[2])
light_emitter.node_tree.links.new(node_5.inputs[0], node_6.outputs[1])
# set parameter of nodes
node_1.inputs[1].default_value = 1.0 # scale
node_2.extension = 'CLIP'
# node_2.interpolation = 'Cubic'
node_3.inputs[1].default_value[0] = 0.5
node_3.inputs[1].default_value[1] = 0.5
node_3.inputs[1].default_value[2] = 0
node_3.inputs[2].default_value[0] = 0
node_3.inputs[2].default_value[1] = 0
node_3.inputs[2].default_value[2] = 0.05
# scale of pattern
node_3.inputs[3].default_value[0] = 0.6
node_3.inputs[3].default_value[1] = 0.85
node_3.inputs[3].default_value[2] = 0
node_4.operation = 'DIVIDE'
# pattern path
node_2.image = self.pattern
def lightModeSelect(self, light_mode):
if light_mode == "RGB":
self.light_emitter.hide_render = True
# set the environment map energy
#### random_energy = random.uniform(LIGHT_ENV_MAP_ENERGY_RGB * 0.8, LIGHT_ENV_MAP_ENERGY_RGB * 1.2)
#### bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = random_energy
elif light_mode == "IR":
self.light_emitter.hide_render = False
# set the environment map energy
random_energy = random.uniform(LIGHT_ENV_MAP_ENERGY_IR * 0.8, LIGHT_ENV_MAP_ENERGY_IR * 1.2)
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = random_energy
elif light_mode == "Mask" or light_mode == "NOCS" or light_mode == "Normal":
self.light_emitter.hide_render = True
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = 0
else:
print("Not support the mode!")
def outputModeSelect(self, output_mode):
if output_mode == "RGB":
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.compression = 0
self.render_context.image_settings.color_mode = 'RGB'
self.render_context.image_settings.color_depth = '8'
bpy.context.scene.view_settings.view_transform = 'Filmic'
bpy.context.scene.render.filter_size = 1.5
self.render_context.resolution_x = 640 ### 1280
self.render_context.resolution_y = 360 ### 720
elif output_mode == "IR":
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.compression = 0
self.render_context.image_settings.color_mode = 'BW'
self.render_context.image_settings.color_depth = '8'
bpy.context.scene.view_settings.view_transform = 'Filmic'
bpy.context.scene.render.filter_size = 1.5
self.render_context.resolution_x = 640 ### 1280
self.render_context.resolution_y = 360 ### 720
elif output_mode == "Mask":
self.render_context.image_settings.file_format = 'OPEN_EXR'
self.render_context.image_settings.color_mode = 'RGB'
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.filter_size = 0
self.render_context.resolution_x = 640
self.render_context.resolution_y = 360
elif output_mode == "NOCS":
# self.render_context.image_settings.file_format = 'OPEN_EXR'
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.color_mode = 'RGB'
self.render_context.image_settings.color_depth = '8'
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.filter_size = 0
self.render_context.resolution_x = 640
self.render_context.resolution_y = 360
elif output_mode == "Normal":
self.render_context.image_settings.file_format = 'OPEN_EXR'
self.render_context.image_settings.color_mode = 'RGB'
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.filter_size = 1.5
self.render_context.resolution_x = 640
self.render_context.resolution_y = 360
else:
print("Not support the mode!")
def renderEngineSelect(self, engine_mode):
if engine_mode == "CYCLES":
self.render_context.engine = 'CYCLES'
bpy.context.scene.cycles.progressive = 'BRANCHED_PATH'
bpy.context.scene.cycles.use_denoising = True
bpy.context.scene.cycles.denoiser = 'NLM'
bpy.context.scene.cycles.film_exposure = 1.0
bpy.context.scene.cycles.aa_samples = CYCLES_SAMPLE
## Set the device_type
bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" # or "OPENCL"
## Set the device and feature set
# bpy.context.scene.cycles.device = "CPU"
## get_devices() to let Blender detects GPU device
cuda_devices, _ = bpy.context.preferences.addons["cycles"].preferences.get_devices()
print(bpy.context.preferences.addons["cycles"].preferences.compute_device_type)
for d in bpy.context.preferences.addons["cycles"].preferences.devices:
d["use"] = 1 # Using all devices, include GPU and CPU
print(d["name"], d["use"])
'''
'''
device_list = DEVICE_LIST
activated_gpus = []
for i, device in enumerate(cuda_devices):
if (i in device_list):
device.use = True
activated_gpus.append(device.name)
else:
device.use = False
elif engine_mode == "EEVEE":
bpy.context.scene.render.engine = 'BLENDER_EEVEE'
else:
print("Not support the mode!")
def addBackground(self, size, position, scale):
# set the material of background
material_name = "default_background"
# test if material exists
# if it does not exist, create it:
material_background = (bpy.data.materials.get(material_name) or
bpy.data.materials.new(material_name))
# enable 'Use nodes'
material_background.use_nodes = True
node_tree = material_background.node_tree
# remove default nodes
material_background.node_tree.nodes.clear()
# material_background.node_tree.nodes.remove(material_background.node_tree.nodes.get('Principled BSDF')) #title of the existing node when materials.new
# material_background.node_tree.nodes.remove(material_background.node_tree.nodes.get('Material Output')) #title of the existing node when materials.new
# add new nodes
node_1 = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_2 = node_tree.nodes.new('ShaderNodeBsdfPrincipled')
node_3 = node_tree.nodes.new('ShaderNodeTexImage')
# link nodes
node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
node_tree.links.new(node_2.inputs[0], node_3.outputs[0])
# add texture image
node_3.image = bpy.data.images.load(filepath=default_background_texture_path)
self.my_material['default_background'] = material_background
# add background plane
for i in range(-2, 3, 1):
for j in range(-2, 3, 1):
position_i_j = (i * size + position[0], j * size + position[1], position[2] - TABLE_CAD_MODEL_HEIGHT)
bpy.ops.mesh.primitive_plane_add(size=size, enter_editmode=False, align='WORLD', location=position_i_j, scale=scale)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
for i in range(-2, 3, 1):
for j in [-2, 2]:
position_i_j = (i * size + position[0], j * size + position[1], position[2] - 0.25)# - TABLE_CAD_MODEL_HEIGHT)
rotation_elur = (math.pi / 2., 0., 0.)
bpy.ops.mesh.primitive_plane_add(size=size, enter_editmode=False, align='WORLD', location=position_i_j, rotation = rotation_elur)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
for j in range(-2, 3, 1):
for i in [-2, 2]:
position_i_j = (i * size + position[0], j * size + position[1], position[2] - 0.25)# - TABLE_CAD_MODEL_HEIGHT)
rotation_elur = (0, math.pi / 2, 0)
bpy.ops.mesh.primitive_plane_add(size=size, enter_editmode=False, align='WORLD', location=position_i_j, rotation = rotation_elur)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
count = 0
for obj in bpy.data.objects:
if obj.type == "MESH":
obj.name = "background_" + str(count)
obj.data.name = "background_" + str(count)
obj.active_material = material_background
count += 1
self.background_added = True
def clearModel(self):
'''
# delete all meshes
for item in bpy.data.meshes:
bpy.data.meshes.remove(item)
for item in bpy.data.materials:
bpy.data.materials.remove(item)
'''
# remove all objects except background
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
bpy.data.meshes.remove(obj.data)
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
bpy.data.objects.remove(obj, do_unlink=True)
# remove all default material
for mat in bpy.data.materials:
name = mat.name.split('.')
if name[0] == 'Material':
bpy.data.materials.remove(mat)
def loadModel(self, file_path):
self.model_loaded = True
try:
if file_path.endswith('obj'):
bpy.ops.import_scene.obj(filepath=file_path)
elif file_path.endswith('3ds'):
bpy.ops.import_scene.autodesk_3ds(filepath=file_path)
elif file_path.endswith('dae'):
# Must install OpenCollada. Please read README.md
bpy.ops.wm.collada_import(filepath=file_path)
else:
self.model_loaded = False
raise Exception("Loading failed: %s" % (file_path))
except Exception:
self.model_loaded = False
def render(self, image_name="tmp", image_path=RENDERING_PATH):
# Render the object
if not self.model_loaded:
print("Model not loaded.")
return
if self.render_mode == "IR":
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("IR")
self.outputModeSelect("IR")
self.renderEngineSelect("CYCLES")
elif self.render_mode == 'RGB':
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("RGB")
self.outputModeSelect("RGB")
self.renderEngineSelect("CYCLES")
elif self.render_mode == "Mask":
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("Mask")
self.outputModeSelect("Mask")
# self.renderEngineSelect("EEVEE")
self.renderEngineSelect("CYCLES")
bpy.context.scene.cycles.use_denoising = False
bpy.context.scene.cycles.aa_samples = 1
elif self.render_mode == "NOCS":
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("NOCS")
self.outputModeSelect("NOCS")
# self.renderEngineSelect("EEVEE")
self.renderEngineSelect("CYCLES")
bpy.context.scene.cycles.use_denoising = False
bpy.context.scene.cycles.aa_samples = 1
elif self.render_mode == "Normal":
bpy.context.scene.use_nodes = True
self.fileOutput.base_path = image_path.replace("normal","depth")
self.fileOutput.file_slots[0].path = image_name[:4]+"_#"# + 'depth_#'
# set light and render mode
self.lightModeSelect("Normal")
self.outputModeSelect("Normal")
# self.renderEngineSelect("EEVEE")
self.renderEngineSelect("CYCLES")
bpy.context.scene.cycles.use_denoising = False
bpy.context.scene.cycles.aa_samples = 32
else:
print("The render mode is not supported")
return
bpy.context.scene.render.filepath = os.path.join(image_path, image_name)
bpy.ops.render.render(write_still=True) # save straight to file
def set_material_randomize_mode(self, class_material_pairs, mat_randomize_mode, instance, material_type_in_mixed_mode):
if mat_randomize_mode in ['mixed','diffuse','transparent','specular_tex','specular_texmix']:
if material_type_in_mixed_mode == 'raw':
print(instance.name, 'material type: raw')
set_modify_raw_material(instance)
else:
print(instance.name, 'material type: ', material_type_in_mixed_mode)
material = random.sample(self.my_material[material_type_in_mixed_mode], 1)[0]
set_modify_material(instance, material, self.obj_texture_img_list, mat_randomize_mode=mat_randomize_mode)
elif mat_randomize_mode == 'specular':
material = random.sample(self.my_material[material_type_in_mixed_mode], 1)[0]
print(instance.name, 'material type: ', material_type_in_mixed_mode)
set_modify_material(instance, material, self.obj_texture_img_list, mat_randomize_mode=mat_randomize_mode, is_transfer=False)
else:
print("No such mat_randomize_mode!")
def get_instance_pose(self):
instance_pose = {}
bpy.context.view_layer.update()
cam = self.camera_l
mat_rot_x = Matrix.Rotation(math.radians(180.0), 4, 'X')
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
instance_id = obj.name.split('_')[0]
mat_rel = cam.matrix_world.inverted() @ obj.matrix_world
# location
relative_location = [mat_rel.translation[0],
- mat_rel.translation[1],
- mat_rel.translation[2]]
# rotation
# relative_rotation_euler = mat_rel.to_euler() # must be converted from radians to degrees
relative_rotation_quat = [mat_rel.to_quaternion()[0],
mat_rel.to_quaternion()[1],
mat_rel.to_quaternion()[2],
mat_rel.to_quaternion()[3]]
quat_x = [0, 1, 0, 0]
quat = quanternion_mul(quat_x, relative_rotation_quat)
quat = [quat[0], - quat[1], - quat[2], - quat[3]]
instance_pose[str(instance_id)] = [quat, relative_location]
return instance_pose
def check_visible(self, threshold=(0.1, 0.9, 0.1, 0.9)):
w_min, x_max, h_min, h_max = threshold
visible_objects_list = []
bpy.context.view_layer.update()
cs, ce = self.camera_l.data.clip_start, self.camera_l.data.clip_end
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
obj_center = obj.matrix_world.translation
co_ndc = world_to_camera_view(scene, self.camera_l, obj_center)
if (w_min < co_ndc.x < x_max and
h_min < co_ndc.y < h_max and
cs < co_ndc.z < ce):
obj.select_set(True)
visible_objects_list.append(obj)
else:
obj.select_set(False)
return visible_objects_list
def setModelPosition(instance, location, quaternion):
instance.rotation_mode = 'QUATERNION'
instance.rotation_quaternion[0] = quaternion[0]
instance.rotation_quaternion[1] = quaternion[1]
instance.rotation_quaternion[2] = quaternion[2]
instance.rotation_quaternion[3] = quaternion[3]
instance.location = location
def setRigidBody(instance):
bpy.context.view_layer.objects.active = instance
object_single = bpy.context.active_object
# add rigid body constraints to cube
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.mass = 1
bpy.context.object.rigid_body.kinematic = True
bpy.context.object.rigid_body.collision_shape = 'CONVEX_HULL'
bpy.context.object.rigid_body.restitution = 0.01
bpy.context.object.rigid_body.angular_damping = 0.8
bpy.context.object.rigid_body.linear_damping = 0.99
bpy.context.object.rigid_body.kinematic = False
object_single.keyframe_insert(data_path='rigid_body.kinematic', frame=0)
def set_visiable_objects(visible_objects_list):
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
if obj in visible_objects_list:
obj.hide_render = False
else:
obj.hide_render = True
def generate_CAD_model_list(urdf_path_list):
CAD_model_list = {}
for instance_path in urdf_path_list:
class_name = 'other'
print(instance_path)
class_list = []
class_list.append([instance_path, class_name])
if class_name == 'other' and 'other' in CAD_model_list:
CAD_model_list[class_name] = CAD_model_list[class_name] + class_list
else:
CAD_model_list[class_name] = class_list
return CAD_model_list
def generate_material_type(obj_name, class_material_pairs, instance_material_except_pairs, instance_material_include_pairs, material_class_instance_pairs, material_type):
specular_type_for_ins_list = []
transparent_type_for_ins_list = []
diffuse_type_for_ins_list = []
for key in instance_material_except_pairs:
if key in material_class_instance_pairs['specular']:
specular_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['transparent']:
transparent_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['diffuse']:
diffuse_type_for_ins_list.append(key)
for key in instance_material_include_pairs:
### if ins_idx in instance_material_include_pairs[key]:
if key in material_class_instance_pairs['specular']:
specular_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['transparent']:
transparent_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['diffuse']:
diffuse_type_for_ins_list.append(key)
if material_type == "transparent":
return random.sample(transparent_type_for_ins_list, 1)[0]
elif material_type == "diffuse":
return random.sample(diffuse_type_for_ins_list, 1)[0]
elif material_type == "specular" or material_type == "specular_tex" or material_type == "specular_texmix":
return random.sample(specular_type_for_ins_list, 1)[0]
elif material_type == "mixed":
# randomly select one material class
flag = random.randint(0, 2) # D:S:T=1:2:2
# select the raw material
if flag == 0:
# flag = random.randint(0, 7) # 1:7
# if flag == 0:
# return 'raw'
# else:
return random.sample(diffuse_type_for_ins_list, 1)[0] ### 'diffuse'
# select one from specular and transparent
elif flag == 1:
return random.sample(specular_type_for_ins_list, 1)[0] ### 'specular'
else:
return random.sample(transparent_type_for_ins_list, 1)[0] ### 'transparent'
else:
print("Material type error: ", material_type)
################################
#
# Main
#
################################
### set random seed
random.seed(1143+SCENE_NUM)
# load VGN obj path and tsdf pose
urdfs_and_poses_files_list = sorted(os.listdir(raw_urdfs_and_poses_dir_path))
urdfs_and_poses_dict = np.load(os.path.join(raw_urdfs_and_poses_dir_path, urdfs_and_poses_files_list[SCENE_NUM]), allow_pickle=True)['pc']
urdf_path_list = list(urdfs_and_poses_dict[:,0])
obj_scale_list = list(urdfs_and_poses_dict[:,1])
obj_RT_list = list(urdfs_and_poses_dict[:,2])
obj_quat_list = []
obj_trans_list = []
for RT in obj_RT_list:
R = RT[:3,:3]
T = RT[:3,3]
T = T + np.array([-0.15,-0.15,-0.0503])
quat = quaternionFromRotMat(R)
obj_quat_list.append(quat)
obj_trans_list.append(T)
g_synset_name_label_pairs = {'other': 0}
material_class_instance_pairs = {'specular': ['metal','porcelain','plasticsp','paintsp'],
'transparent': ['glass'],
'diffuse': ['plastic','rubber','paper','leather','wood','clay','fabric'],
'background': ['background']}
class_material_pairs = {'specular': ['other'],
'transparent': ['other'],
'diffuse': ['other']}
instance_material_except_pairs = {'metal': [],
'porcelain': [],
'plasticsp': [],
'paintsp':[],
'glass': [],
'plastic': [],
'rubber': [],
'leather': [],
'wood':[],
'paper':[],
'fabric':[],
'clay':[],
}
instance_material_include_pairs = {
}
material_class_id_dict = {'raw': 0,
'diffuse': 1,
'transparent': 2,
'specular': 3}
material_type_id_dict = {'raw': 0,
'metal': 1,
'porcelain': 2,
'plasticsp': 3,
'paintsp':4,
'glass': 5,
'plastic': 6,
'rubber': 7,
'leather': 8,
'wood':9,
'paper':10,
'fabric':11,
'clay':12,
}
max_instance_num = 20
if not os.path.exists(output_root_path):
os.makedirs(output_root_path)
# generate CAD model list
CAD_model_list = generate_CAD_model_list(urdf_path_list)
renderer = BlenderRenderer(viewport_size_x=camera_width, viewport_size_y=camera_height)
renderer.loadImages(emitter_pattern_path, env_map_path, real_table_image_root_path, real_floor_image_root_path)
renderer.addEnvMap()
renderer.addBackground(background_size, background_position, background_scale)
renderer.addMaterialLib(material_class_instance_pairs) ###
renderer.addMaskMaterial(max_instance_num)
renderer.addNOCSMaterial()
renderer.addNormalMaterial()
renderer.clearModel()
# set scene output path
path_scene = os.path.join(output_root_path, urdfs_and_poses_files_list[SCENE_NUM][:-4])### "scene_"+str(SCENE_NUM).zfill(4))
if os.path.exists(path_scene)==False:
os.makedirs(path_scene)
# camera pose list, environment light list and background material_listz
quaternion_list = []
translation_list = []
# environment map list
env_map_id_list = []
rotation_elur_z_list = []
# background material list
background_material_list = []
# table material list
table_material_list = []
look_at = look_at_shift
quat_list, trans_list, rot_list, position_list = genCameraPosition(look_at)
np.savetxt(os.path.join(path_scene, "cam_pos_pc.txt"), np.array(position_list))
rot_array = np.array(rot_list) # (256, 3, 3)
trans_array = np.array(trans_list) # (256, 3, 1)
cam_RT = np.concatenate([rot_array, trans_array], 2)
zero_one = np.expand_dims([[0, 0, 0, 1]],0).repeat(rot_array.shape[0],axis=0)
cam_RT = np.concatenate([cam_RT, zero_one], 1) # (256, 4, 4)
np.save(os.path.join(path_scene, "camera_pose.npy"), cam_RT)
# generate camara pose list
for i in range(NUM_FRAME_PER_SCENE):
quaternion = quat_list[i]
translation = trans_list[i]
quaternion_list.append(quaternion)
translation_list.append(translation)
# generate environment map list
env_map_id_list.append(random.randint(0, len(renderer.env_map) - 1))
rotation_elur_z_list.append(random.uniform(-math.pi, math.pi))
# generate background material list
if my_material_randomize_mode == 'raw':
background_material_list.append(renderer.my_material['default_background'])
# bpy.data.objects['background'].active_material = renderer.my_material['default_background']
else:
material_selected = random.sample(renderer.my_material['background'], 2)[0]
background_material_list.append(material_selected)
# bpy.data.objects['background'].active_material = material_selected
###
material_selected = random.sample(renderer.my_material['background'], 2)[1]
table_material_list.append(material_selected)
###
#print(background_material_list, table_material_list)
# read objects from floder
meta_output = {}
select_model_list = []
select_model_list_other = []
select_model_list_transparent = []
select_model_list_dis = []
select_number = 1
for item in CAD_model_list:
if item in ['other']:
test = CAD_model_list[item]
for model in test:
select_model_list.append(model)
else:
print("No such category!")
# load table obj
renderer.loadModel(table_CAD_model_path)
obj = bpy.data.objects['table']
table_scale = 0.001
obj.scale = (table_scale, table_scale, table_scale)
y_transform = np.array([[0,0,-1],[0,1,0],[1,0,0]])
transform = y_transform
obj_world_pose_quat = quaternionFromRotMat(transform)
obj_world_pose_T_shift = np.array([random.uniform(-0.02,0.02),random.uniform(-0.0843,-0.0243),-0.0751])
obj_world_pose_T = obj_world_pose_T_shift
setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
# load table plane
obj_world_pose_T = obj_world_pose_T_shift
obj_world_pose_T[2] = 0
bpy.ops.mesh.primitive_plane_add(size=1., enter_editmode=False, align='WORLD', location=obj_world_pose_T)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
obj = bpy.data.objects['Plane']
obj.name = 'tableplane'
obj.data.name = 'tableplane'
obj.scale = (0.898, 1.3, 1.)
# load arm
renderer.loadModel(arm_CAD_model_path)
obj = bpy.data.objects['arm']
class_scale = 0.001
obj.scale = (class_scale, class_scale, class_scale)
x_transform = np.array([[1,0,0],[0,0,-1],[0,1,0]])
transform = x_transform
obj_world_pose_quat = quaternionFromRotMat(transform)
obj_world_pose_T_shift = np.array([0,random.uniform(-0.43,-0.41),0])
obj_world_pose_T = obj_world_pose_T_shift
setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
instance_id = 1
imported_obj_name_list = []
for model in select_model_list:
instance_path = model[0]
class_name = model[1]
instance_folder = model[0].split('/')[-1][:-4]
instance_name = str(instance_id) + "_" + class_name + "_" + instance_folder
material_type_in_mixed_mode = generate_material_type(instance_name, class_material_pairs, instance_material_except_pairs, instance_material_include_pairs, material_class_instance_pairs, my_material_randomize_mode)
# download CAD model and rename
renderer.loadModel(instance_path)
import_obj_name = instance_folder
obj = bpy.data.objects[import_obj_name]
obj.name = instance_name
obj.data.name = instance_name
print(len(obj_trans_list), instance_id)
obj_world_pose_T = obj_trans_list[instance_id-1]
obj_world_pose_quat = obj_quat_list[instance_id-1]
setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
# set object as rigid body
setRigidBody(obj)
# set material
renderer.set_material_randomize_mode(class_material_pairs, my_material_randomize_mode, obj, material_type_in_mixed_mode)
# generate meta file
class_scale = obj_scale_list[instance_id-1]
obj.scale = (class_scale, class_scale, class_scale)
# material type
material_class_id = None
for key in material_class_instance_pairs:
if material_type_in_mixed_mode == 'raw':
material_class_id = material_class_id_dict[material_type_in_mixed_mode]
break
elif material_type_in_mixed_mode in material_class_instance_pairs[key]:
material_class_id = material_class_id_dict[key]
break
if material_class_id == None:
print("material_class_id error!")
meta_output[str(instance_id)] = [str(instance_folder),
str(material_class_id),
str(material_type_id_dict[material_type_in_mixed_mode])
]
instance_id += 1
# set the key frame
scene = bpy.data.scenes['Scene']
# generate meta.txt
meta_dir_path = os.path.join(path_scene, 'meta')
if os.path.exists(meta_dir_path)==False:
os.makedirs(meta_dir_path)
for i in RENDER_FRAMES_LIST: #range(RENDER_START_FRAME, RENDER_END_FRAME):
# output the meta file
path_meta = os.path.join(meta_dir_path, str(i).zfill(4) + ".txt")
if os.path.exists(path_meta):
os.remove(path_meta)
file_write_obj = open(path_meta, 'w')
for index in meta_output:
file_write_obj.write(index)
file_write_obj.write(' ')
for item in meta_output[index]:
file_write_obj.write(item)
file_write_obj.write(' ')
file_write_obj.write('\n')
file_write_obj.close()
# render IR image and RGB image
if render_mode_list['IR'] or render_mode_list['RGB']:
renderer.setEnvMap(env_map_id_list[0], rotation_elur_z_list[0])
# 随机选一张real floor image
flag = random.randint(0, len(renderer.realfloor_img_list)-1)
selected_realfloor_img = renderer.realfloor_img_list[flag]
# 随机选一张real table image
flag = random.randint(0, len(renderer.realtable_img_list)-1)
selected_realtable_img = renderer.realtable_img_list[flag]
arm_material_type_in_mixed_mode = generate_material_type(None, class_material_pairs, instance_material_except_pairs, instance_material_include_pairs, material_class_instance_pairs, material_type="diffuse")
arm_material = random.sample(renderer.my_material[arm_material_type_in_mixed_mode], 1)[0]
###
for i in RENDER_FRAMES_LIST:
renderer.setCamera(quaternion_list[i], translation_list[i], camera_fov, baseline_distance)
renderer.setLighting()
for obj in bpy.data.objects:
if obj.type == "MESH" and obj.name.split('_')[0] == 'background':
if obj.name == 'background_0':
flag = random.randint(0, 3)
if flag == 0:
set_modify_floor_material(obj, background_material_list[0], selected_realfloor_img) ### renderer.realfloor_img_list)
else:
obj.active_material = background_material_list[0]
else:
background_0_obj = bpy.data.objects['background_0']
obj.active_material = background_0_obj.material_slots[0].material
elif obj.type == "MESH" and obj.name == 'table':
flag = random.randint(0, 2)
if flag == 0:
set_modify_table_material(obj, table_material_list[0], selected_realtable_img)### renderer.realtable_img_list)
else:
obj.active_material = table_material_list[0]
elif obj.type == "MESH" and obj.name == 'tableplane':
table_obj = bpy.data.objects['table']
obj.active_material = table_obj.material_slots[0].material
elif obj.type == "MESH" and obj.name == 'arm':
set_modify_arm_material(obj, arm_material)
# render IR image
if render_mode_list['IR']:
ir_l_dir_path = os.path.join(path_scene, 'ir_l')
if os.path.exists(ir_l_dir_path)==False:
os.makedirs(ir_l_dir_path)
ir_r_dir_path = os.path.join(path_scene, 'ir_r')
if os.path.exists(ir_r_dir_path)==False:
os.makedirs(ir_r_dir_path)
renderer.render_mode = "IR"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = ir_l_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
camera = bpy.data.objects['camera_r']
scene.camera = camera
save_path = ir_r_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
# render RGB image
if render_mode_list['RGB']:
rgb_dir_path = os.path.join(path_scene, 'rgb')
if os.path.exists(rgb_dir_path)==False:
os.makedirs(rgb_dir_path)
renderer.render_mode = "RGB"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = rgb_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
# render mask map and depth map
if render_mode_list['Mask']:
# set instance mask as material
for obj in bpy.data.objects:
if obj.type == "MESH":
obj.data.materials.clear()
material_name = "mask_" + obj.name.split('_')[0]
obj.active_material = renderer.my_material[material_name]
# render mask map and depth map
for i in RENDER_FRAMES_LIST:
renderer.setCamera(quaternion_list[i], translation_list[i], camera_fov, baseline_distance)
mask_dir_path = os.path.join(path_scene, 'mask')
if os.path.exists(mask_dir_path)==False:
os.makedirs(mask_dir_path)
renderer.render_mode = "Mask"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = mask_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
# render normal map
if render_mode_list['Normal']:
# set normal as material
for obj in bpy.data.objects:
if obj.type == 'MESH':
obj.data.materials.clear()
obj.active_material = renderer.my_material["normal"]
# render normal map
for i in RENDER_FRAMES_LIST:
renderer.setCamera(quaternion_list[i], translation_list[i], camera_fov, baseline_distance)
normal_dir_path = os.path.join(path_scene, 'normal')
if os.path.exists(normal_dir_path)==False:
os.makedirs(normal_dir_path)
depth_dir_path = os.path.join(path_scene, 'depth')
if os.path.exists(depth_dir_path)==False:
os.makedirs(depth_dir_path)
renderer.render_mode = "Normal"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = normal_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
context = bpy.context
for ob in context.selected_objects:
ob.animation_data_clear()
print(bpy.data.materials)
print(len(bpy.data.materials))
src_depth_list = os.listdir(os.path.join(output_root_path,urdfs_and_poses_files_list[SCENE_NUM][:-4],"depth"))
for src_depth in src_depth_list:
source = os.path.join(output_root_path,urdfs_and_poses_files_list[SCENE_NUM][:-4],"depth",src_depth)
target = source.replace("_0.exr",".exr")
os.rename(source,target)
================================================
FILE: data_generator/run_pile_rand.sh
================================================
#!/bin/bash
cd /data/InterNeRF/renderer/renderer_giga_GPU6-0_rand_M
# 830*6
mycount=0;
while (( $mycount < 100 )); do
/home/xxx/blender-2.93.3-linux-x64/blender material_lib_v2.blend --background -noaudio --python render_pile_STD_rand.py -- $mycount;
((mycount=$mycount+1));
done;
================================================
FILE: requirements.txt
================================================
torch
tensorflow
easydict
inplace-abn
plyfile
numpy
scikit-image
pyyaml
h5py
opencv-python
tqdm
matplotlib
scipy
lpips
transforms3d
kornia
sklearn
catkin_pkg
black
jupyterlab
pandas
mpi4py
open3d
pybullet==2.7.9
pytorch-ignite
tensorboard
================================================
FILE: run_simgrasp.sh
================================================
#!/bin/bash
GPUID=0
BLENDER_BIN=blender
RENDERER_ASSET_DIR=./data/assets
BLENDER_PROJ_PATH=./data/assets/material_lib_graspnet-v2.blend
SIM_LOG_DIR="./log/`date '+%Y%m%d-%H%M%S'`"
scene="pile"
object_set="pile_subdiv"
material_type="specular_and_transparent"
render_frame_list="2,6,10,14,18,22"
check_seen_scene=0
expname=0
NUM_TRIALS=200
METHOD='graspnerf'
mycount=0
while (( $mycount < $NUM_TRIALS )); do
$BLENDER_BIN $BLENDER_PROJ_PATH --background --python scripts/sim_grasp.py \
-- $mycount $GPUID $expname $scene $object_set $check_seen_scene $material_type \
$RENDERER_ASSET_DIR $SIM_LOG_DIR 0 $render_frame_list $METHOD
python ./scripts/stat_expresult.py -- $SIM_LOG_DIR $expname
((mycount=$mycount+1));
done;
python ./scripts/stat_expresult.py -- $SIM_LOG_DIR $expname
================================================
FILE: scripts/sim_grasp.py
================================================
import argparse
import os
import sys
from pathlib import Path
def main(args, round_idx, gpuid, render_frame_list):
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpuid)
sys.path.append("src")
from nr.main import GraspNeRFPlanner
if args.method == "graspnerf":
grasp_planner = GraspNeRFPlanner(args)
else:
print("No such method!")
raise NotImplementedError
from gd.experiments import clutter_removal
clutter_removal.run(
grasp_plan_fn=grasp_planner,
logdir=args.logdir,
description=args.description,
scene=args.scene,
object_set=args.object_set,
num_objects=args.num_objects,
num_rounds=args.num_rounds,
seed=args.seed,
sim_gui=args.sim_gui,
rviz=args.rviz,
round_idx = round_idx,
renderer_root_dir = args.renderer_root_dir,
gpuid = gpuid,
args = args,
render_frame_list = render_frame_list
)
class ArgumentParserForBlender(argparse.ArgumentParser):
"""
This class is identical to its superclass, except for the parse_args
method (see docstring). It resolves the ambiguity generated when calling
Blender from the CLI with a python script, and both Blender and the script
have arguments. E.g., the following call will make Blender crash because
it will try to process the script's -a and -b flags:
>>> blender --python my_script.py -a 1 -b 2
To bypass this issue this class uses the fact that Blender will ignore all
arguments given after a double-dash ('--'). The approach is that all
arguments before '--' go to Blender, arguments after go to the script.
The following calls work fine:
>>> blender --python my_script.py -- -a 1 -b 2
>>> blender --python my_script.py --
"""
def _get_argv_after_doubledash(self):
"""
Given the sys.argv as a list of strings, this method returns the
sublist right after the '--' element (if present, otherwise returns
an empty list).
"""
try:
idx = sys.argv.index("---")
return sys.argv[idx+1:] # the list after '--'
except ValueError as e: # '--' not in the list:
return []
# overrides superclass
def parse_args(self):
"""
This method is expected to behave identically as in the superclass,
except that the sys.argv list will be pre-processed using
_get_argv_after_doubledash before. See the docstring of the class for
usage examples and details.
"""
return super().parse_args(args=self._get_argv_after_doubledash())
if __name__ == "__main__":
argv = sys.argv
argv = argv[argv.index("--") + 1:] # get all args after "--"
round_idx = int(argv[0])
gpuid = int(argv[1])
expname = str(argv[2])
scene = str(argv[3])
object_set = str(argv[4])
check_seen_scene = bool(int(argv[5]))
material_type = str(argv[6])
blender_asset_dir = str(argv[7])
log_root_dir = str(argv[8])
use_gt_tsdf = bool(int(argv[9]))
render_frame_list=[int(frame_id) for frame_id in str(argv[10]).replace(' ','').split(",")]
method = str(argv[11])
print("########## Simulation Start ##########")
print("Round %d\nmethod: %s\nmaterial_type: %s\nviews: %s "%(round_idx, method, material_type, str(render_frame_list)))
print("######################################")
parser = ArgumentParserForBlender() ### argparse.ArgumentParser()
parser.add_argument("---model", type=Path, default="")
parser.add_argument("---logdir", type=Path, default=expname)
parser.add_argument("---description", type=str, default="")
parser.add_argument("---scene", type=str, choices=["pile", "packed", "single"], default=scene)
parser.add_argument("---object-set", type=str, default=object_set)
parser.add_argument("---num-objects", type=int, default=5)
parser.add_argument("---num-rounds", type=int, default=200)
parser.add_argument("---seed", type=int, default=42)
parser.add_argument("---sim-gui", type=bool, default=False)
parser.add_argument("---rviz", action="store_true")
###
parser.add_argument("---renderer_root_dir", type=str, default=blender_asset_dir)
parser.add_argument("---log_root_dir", type=str, default=log_root_dir)
parser.add_argument("---obj_texture_image_root_path", type=str, default=blender_asset_dir+"/imagenet") #TODO
parser.add_argument("---cfg_fn", type=str, default="src/nr/configs/nrvgn_sdf.yaml")
parser.add_argument('---database_name', type=str, default='vgn_syn/train/packed/packed_170-220/032cd891d9be4a16be5ea4be9f7eca2b/w_0.8', help='//')
parser.add_argument("---gen_scene_descriptor", type=bool, default=False)
parser.add_argument("---load_scene_descriptor", type=bool, default=True)
parser.add_argument("---material_type", type=str, default=material_type)
parser.add_argument("---method", type=str, default=method)
# pybullet camera parameter
parser.add_argument("---camera_focal", type=float, default=446.31) #TODO
###
args = parser.parse_args()
main(args, round_idx, gpuid, render_frame_list)
================================================
FILE: scripts/stat_expresult.py
================================================
from pathlib import Path
import sys
import pandas as pd
import os
import numpy as np
argv = sys.argv
argv = argv[argv.index("--") + 1:] # get all args after "--"
log_root_dir = str(argv[0])
expname = str(argv[1])
class Data(object):
"""Object for loading and analyzing experimental data."""
def __init__(self, logdir):
self.logdir = logdir
self.rounds = pd.read_csv(logdir / "rounds.csv")
self.grasps = pd.read_csv(logdir / "grasps.csv")
def num_rounds(self):
return len(self.rounds.index)
def num_grasps(self):
return len(self.grasps.index)
def success_rate(self):
return self.grasps["label"].mean() * 100
def percent_cleared(self):
df = (
self.grasps[["round_id", "label"]]
.groupby("round_id")
.sum()
.rename(columns={"label": "cleared_count"})
.merge(self.rounds, on="round_id")
)
return df["cleared_count"].sum() / df["object_count"].sum() * 100
def avg_planning_time(self):
return self.grasps["planning_time"].mean()
def read_grasp(self, i):
scene_id, grasp, label = io.read_grasp(self.grasps, i)
score = self.grasps.loc[i, "score"]
scene_data = np.load(self.logdir / "scenes" / (scene_id + ".npz"))
return scene_data["points"], grasp, score, label
##############################
# Combine all trials
##############################
root_path = os.path.join(log_root_dir, "exp_results", expname)
round_dir_list = sorted(os.listdir(root_path))
if not os.path.exists(root_path + "_combine"):
os.makedirs(root_path + "_combine")
df = pd.DataFrame()
for i in range(len(round_dir_list)):
df_round = pd.read_csv(os.path.join(root_path, round_dir_list[i], "grasps.csv"))
df_round["round_id"] = i
df = pd.concat([df, df_round])
df = df.reset_index(drop=True)
df.to_csv(os.path.join(root_path + "_combine", "grasps.csv"), index=False)
df = pd.DataFrame()
for i in range(len(round_dir_list)):
df_round = pd.read_csv(os.path.join(root_path, round_dir_list[i], "rounds.csv"))
df_round["round_id"] = i
df = pd.concat([df, df_round])
df = df.reset_index(drop=True)
df.to_csv(os.path.join(root_path + "_combine", "rounds.csv"), index=False)
##############################
# Print Stat
##############################
logdir = Path(os.path.join(log_root_dir, "exp_results", expname+"_combine"))
data = Data(logdir)
# First, we compute the following metrics for the experiment:
# * **Success rate**: the ratio of successful grasp executions,
# * **Percent cleared**: the percentage of objects removed during each round,
try:
print("Path: ",str(logdir))
print("Num grasps: ", data.num_grasps())
print("Success rate: ", data.success_rate())
print("Percent cleared: ", data.percent_cleared())
except:
print("[W] Incomplete results, exit")
exit()
##############################
# Calc first-time grasping SR
##############################
sum_label = 0
firstgrasp_fail_expidx_list = []
for i in range(len(round_dir_list)):
#print(i)
df_round = pd.read_csv(os.path.join(root_path, round_dir_list[i], "grasps.csv"))
df = df_round.iloc[0:1,:]
label = df[["label"]].to_numpy(np.float32)
if label.shape[0] == 0:
firstgrasp_fail_expidx_list.append(i)
continue
sum_label += label[0,0]
if label[0,0]==0:
firstgrasp_fail_expidx_list.append(i)
print("First grasp success rate: ", sum_label / len(round_dir_list))
print("First grasp fail:", len(firstgrasp_fail_expidx_list),"/",len(round_dir_list), ", exp id: ", firstgrasp_fail_expidx_list)
================================================
FILE: src/gd/__init__.py
================================================
================================================
FILE: src/gd/baselines.py
================================================
import time
from gpd_ros.msg import GraspConfigList
import numpy as np
from sensor_msgs.msg import PointCloud2
import rospy
from gd.grasp import Grasp
from gd.utils import ros_utils
from gd.utils.transform import Rotation, Transform
class GPD(object):
def __init__(self):
self.input_topic = "/cloud_stitched"
self.output_topic = "/detect_grasps/clustered_grasps"
self.cloud_pub = rospy.Publisher(self.input_topic, PointCloud2, queue_size=1)
def __call__(self, state):
points = np.asarray(state.pc.points)
msg = ros_utils.to_cloud_msg(points, frame="task")
self.cloud_pub.publish(msg)
tic = time.time()
result = rospy.wait_for_message(self.output_topic, GraspConfigList)
toc = time.time() - tic
grasps, scores = self.to_grasp_list(result)
return grasps, scores, toc
def to_grasp_list(self, grasp_configs):
grasps, scores = [], []
for grasp_config in grasp_configs.grasps:
# orientation
x_axis = ros_utils.from_vector3_msg(grasp_config.axis)
y_axis = -ros_utils.from_vector3_msg(grasp_config.binormal)
z_axis = ros_utils.from_vector3_msg(grasp_config.approach)
orientation = Rotation.from_matrix(np.vstack([x_axis, y_axis, z_axis]).T)
# position
position = ros_utils.from_point_msg(grasp_config.position)
# width
width = grasp_config.width.data
# score
score = grasp_config.score.data
if score < 0.0:
continue # negative score is larger than positive score (https://github.com/atenpas/gpd/issues/32#issuecomment-387846534)
grasps.append(Grasp(Transform(orientation, position), width))
scores.append(score)
return grasps, scores
================================================
FILE: src/gd/dataset.py
================================================
import numpy as np
from scipy import ndimage
import torch.utils.data
from gd.io import *
from gd.perception import *
from gd.utils.transform import Rotation, Transform
class Dataset(torch.utils.data.Dataset):
def __init__(self, root, augment=False):
self.root = root
self.augment = augment
self.df = read_df(root)
def __len__(self):
return len(self.df.index)
def __getitem__(self, i):
scene_id = self.df.loc[i, "scene_id"]
ori = Rotation.from_quat(self.df.loc[i, "qx":"qw"].to_numpy(np.single))
pos = self.df.loc[i, "i":"k"].to_numpy(np.single)
width = self.df.loc[i, "width"].astype(np.single)
label = self.df.loc[i, "label"].astype(np.long)
voxel_grid = read_voxel_grid(self.root, scene_id)
if self.augment:
voxel_grid, ori, pos = apply_transform(voxel_grid, ori, pos)
index = np.round(pos).astype(np.long)
rotations = np.empty((2, 4), dtype=np.single)
R = Rotation.from_rotvec(np.pi * np.r_[0.0, 0.0, 1.0])
rotations[0] = ori.as_quat()
rotations[1] = (ori * R).as_quat()
x, y, index = voxel_grid, (label, rotations, width), index
return x, y, index
def apply_transform(voxel_grid, orientation, position):
angle = np.pi / 2.0 * np.random.choice(4)
R_augment = Rotation.from_rotvec(np.r_[0.0, 0.0, angle])
z_offset = np.random.uniform(6, 34) - position[2]
t_augment = np.r_[0.0, 0.0, z_offset]
T_augment = Transform(R_augment, t_augment)
T_center = Transform(Rotation.identity(), np.r_[20.0, 20.0, 20.0])
T = T_center * T_augment * T_center.inverse()
# transform voxel grid
T_inv = T.inverse()
matrix, offset = T_inv.rotation.as_matrix(), T_inv.translation
voxel_grid[0] = ndimage.affine_transform(voxel_grid[0], matrix, offset, order=0)
# transform grasp pose
position = T.transform_point(position)
orientation = T.rotation * orientation
return voxel_grid, orientation, position
================================================
FILE: src/gd/detection.py
================================================
import time
import numpy as np
from scipy import ndimage
import torch
from gd.grasp import *
from gd.utils.transform import Transform, Rotation
from gd.networks import load_network
class VGN(object):
def __init__(self, model_path, rviz=False):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.net = load_network(model_path, self.device)
self.rviz = rviz
def __call__(self, state):
tsdf_vol = state.tsdf.get_grid()
voxel_size = state.tsdf.voxel_size
tic = time.time()
qual_vol, rot_vol, width_vol = predict(tsdf_vol, self.net, self.device)
qual_vol, rot_vol, width_vol = process(tsdf_vol, qual_vol, rot_vol, width_vol)
grasps, scores = select(qual_vol.copy(), rot_vol, width_vol)
toc = time.time() - tic
grasps, scores = np.asarray(grasps), np.asarray(scores)
if len(grasps) > 0:
p = np.random.permutation(len(grasps))
grasps = [from_voxel_coordinates(g, voxel_size) for g in grasps[p]]
scores = scores[p]
if self.rviz:
from gd import vis
vis.draw_quality(qual_vol, state.tsdf.voxel_size, threshold=0.01)
return grasps, scores, toc
def predict(tsdf_vol, net, device):
assert tsdf_vol.shape == (1, 40, 40, 40)
# move input to the GPU
tsdf_vol = torch.from_numpy(tsdf_vol).unsqueeze(0).to(device)
# forward pass
with torch.no_grad():
qual_vol, rot_vol, width_vol = net(tsdf_vol)
# move output back to the CPU
qual_vol = qual_vol.cpu().squeeze().numpy()
rot_vol = rot_vol.cpu().squeeze().numpy()
width_vol = width_vol.cpu().squeeze().numpy()
return qual_vol, rot_vol, width_vol
def process(
tsdf_vol,
qual_vol,
rot_vol,
width_vol,
gaussian_filter_sigma=1.0,
min_width=1.33,
max_width=9.33,
):
tsdf_vol = tsdf_vol.squeeze()
# smooth quality volume with a Gaussian
qual_vol = ndimage.gaussian_filter(
qual_vol, sigma=gaussian_filter_sigma, mode="nearest"
)
# mask out voxels too far away from the surface
outside_voxels = tsdf_vol > 0.5
inside_voxels = np.logical_and(1e-3 < tsdf_vol, tsdf_vol < 0.5)
valid_voxels = ndimage.morphology.binary_dilation(
outside_voxels, iterations=2, mask=np.logical_not(inside_voxels)
)
qual_vol[valid_voxels == False] = 0.0
# reject voxels with predicted widths that are too small or too large
qual_vol[np.logical_or(width_vol < min_width, width_vol > max_width)] = 0.0
return qual_vol, rot_vol, width_vol
def select(qual_vol, rot_vol, width_vol, threshold=0.90, max_filter_size=4):
# threshold on grasp quality
qual_vol[qual_vol < threshold] = 0.0
# non maximum suppression
max_vol = ndimage.maximum_filter(qual_vol, size=max_filter_size)
qual_vol = np.where(qual_vol == max_vol, qual_vol, 0.0)
mask = np.where(qual_vol, 1.0, 0.0)
# construct grasps
grasps, scores = [], []
for index in np.argwhere(mask):
grasp, score = select_index(qual_vol, rot_vol, width_vol, index)
grasps.append(grasp)
scores.append(score)
return grasps, scores
def select_index(qual_vol, rot_vol, width_vol, index):
i, j, k = index
score = qual_vol[i, j, k]
ori = Rotation.from_quat(rot_vol[:, i, j, k])
pos = np.array([i, j, k], dtype=np.float64)
width = width_vol[i, j, k]
return Grasp(Transform(ori, pos), width), score
================================================
FILE: src/gd/experiments/__init__.py
================================================
================================================
FILE: src/gd/experiments/clutter_removal.py
================================================
import collections
import uuid
import os
import numpy as np
import pandas as pd
import sys
import shutil
from pathlib import Path
from gd import io
from gd.grasp import *
from gd.simulation import ClutterRemovalSim
sys.path.append("./")
from rd.render import blender_init_scene, blender_render, blender_update_sceneobj
MAX_CONSECUTIVE_FAILURES = 2
State = collections.namedtuple("State", ["tsdf", "pc"])
def copydirs(from_file, to_file):
if not os.path.exists(to_file):
os.makedirs(to_file)
files = os.listdir(from_file)
for f in files:
if os.path.isdir(from_file + '/' + f):
copydirs(from_file + '/' + f, to_file + '/' + f)
else:
shutil.copy(from_file + '/' + f, to_file + '/' + f)
def run(
grasp_plan_fn,
logdir,
description,
scene,
object_set,
num_objects=5,
n=6,
N=None,
num_rounds=40,
seed=1,
sim_gui=False,
rviz=False,
round_idx=0,
renderer_root_dir="",
gpuid=None,
args=None,
render_frame_list=[]
):
"""Run several rounds of simulated clutter removal experiments.
Each round, m objects are randomly placed in a tray. Then, the grasping pipeline is
run until (a) no objects remain, (b) the planner failed to find a grasp hypothesis,
or (c) maximum number of consecutive failed grasp attempts.
"""
sim = ClutterRemovalSim(scene, object_set, gui=sim_gui, seed=seed, renderer_root_dir=renderer_root_dir, args=args)
logger = Logger(args.log_root_dir, logdir, description, round_idx)
# output modality
output_modality_dict = {'RGB': 1,
'IR': 0,
'NOCS': 0,
'Mask': 0,
'Normal': 0}
for n_round in range(round_idx, round_idx+1):
urdfs_and_poses_dict = sim.reset(num_objects, round_idx)
renderer, quaternion_list, translation_list, path_scene = blender_init_scene(renderer_root_dir, args.log_root_dir, args.obj_texture_image_root_path, scene, urdfs_and_poses_dict, round_idx, logdir, False, args.material_type, gpuid, output_modality_dict)
render_finished = False
render_fail_times = 0
while not render_finished and render_fail_times < 3:
try:
blender_render(renderer, quaternion_list, translation_list, path_scene, render_frame_list, output_modality_dict, args.camera_focal, is_init=True)
render_finished = True
except:
render_fail_times += 1
if not render_finished:
raise RuntimeError("Blender render failed for 3 times.")
path_scene_backup = os.path.join(path_scene + "_backup", "%d_init"%n_round)
if os.path.exists(path_scene_backup) == False:
os.makedirs(path_scene_backup)
copydirs(path_scene, path_scene_backup)
round_id = logger.last_round_id() + 1
logger.log_round(round_id, sim.num_objects)
consecutive_failures = 1
last_label = None
n_grasp = 0
while sim.num_objects > 0 and consecutive_failures < MAX_CONSECUTIVE_FAILURES:
timings = {}
timings["integration"] = 0
gt_tsdf, gt_pc, _ = sim.acquire_tsdf(n=n, N=N)
if args.method == "graspnerf":
grasps, scores, timings["planning"] = grasp_plan_fn(render_frame_list, round_idx, n_grasp, gt_tsdf)
else:
raise NotImplementedError
if len(grasps) == 0:
print("no detections found, abort this round")
break
else:
print(f"{len(grasps)} detections found.")
# execute grasp
grasp, score = grasps[0], scores[0]
(label, _), remain_obj_inws_infos = sim.execute_grasp(grasp, allow_contact=True)
# render the modified scene after grasping
obj_name_list = [str(value[0]).split("/")[-1][:-5] for value in remain_obj_inws_infos]
obj_quat_list = [value[2][[3, 0, 1, 2]] for value in remain_obj_inws_infos]
obj_trans_list = [value[3] for value in remain_obj_inws_infos]
obj_uid_list = [value[4] for value in remain_obj_inws_infos]
# update blender scene
blender_update_sceneobj(obj_name_list, obj_trans_list, obj_quat_list, obj_uid_list)
# render updated scene
render_finished = False
render_fail_times = 0
while not render_finished and render_fail_times < 3:
try:
blender_render(renderer, quaternion_list, translation_list, path_scene, render_frame_list, output_modality_dict, args.camera_focal)
render_finished = True
except:
render_fail_times += 1
if not render_finished:
raise RuntimeError("Blender render failed for 3 times.")
path_scene_backup = os.path.join(path_scene+"_backup", "%d_%d"%(n_round,n_grasp))
if os.path.exists(path_scene_backup)==False:
os.makedirs(path_scene_backup)
copydirs(path_scene, path_scene_backup)
# log the grasp
logger.log_grasp(round_id, timings, grasp, score, label)
if last_label == Label.FAILURE and label == Label.FAILURE:
consecutive_failures += 1
else:
consecutive_failures = 1
last_label = label
n_grasp += 1
class Logger(object):
def __init__(self, log_root_dir, expname, description, round_idx):
self.logdir = Path(os.path.join(log_root_dir, "exp_results", expname , "%04d"%int(round_idx)))#description
self.scenes_dir = self.logdir / "scenes"
self.scenes_dir.mkdir(parents=True, exist_ok=True)
self.rounds_csv_path = self.logdir / "rounds.csv"
self.grasps_csv_path = self.logdir / "grasps.csv"
self._create_csv_files_if_needed()
def _create_csv_files_if_needed(self):
if not self.rounds_csv_path.exists():
io.create_csv(self.rounds_csv_path, ["round_id", "object_count"])
if not self.grasps_csv_path.exists():
columns = [
"round_id",
"scene_id",
"qx",
"qy",
"qz",
"qw",
"x",
"y",
"z",
"width",
"score",
"label",
"integration_time",
"planning_time",
]
io.create_csv(self.grasps_csv_path, columns)
def last_round_id(self):
df = pd.read_csv(self.rounds_csv_path)
return -1 if df.empty else df["round_id"].max()
def log_round(self, round_id, object_count):
io.append_csv(self.rounds_csv_path, round_id, object_count)
def log_grasp(self, round_id, timings, grasp, score, label):
# log scene
scene_id = uuid.uuid4().hex
# log grasp
qx, qy, qz, qw = grasp.pose.rotation.as_quat()
x, y, z = grasp.pose.translation
width = grasp.width
label = int(label)
io.append_csv(
self.grasps_csv_path,
round_id,
scene_id,
qx,
qy,
qz,
qw,
x,
y,
z,
width,
score,
label,
timings["integration"],
timings["planning"],
)
class Data(object):
"""Object for loading and analyzing experimental data."""
def __init__(self, logdir):
self.logdir = logdir
self.rounds = pd.read_csv(logdir / "rounds.csv")
self.grasps = pd.read_csv(logdir / "grasps.csv")
def num_rounds(self):
return len(self.rounds.index)
def num_grasps(self):
return len(self.grasps.index)
def success_rate(self):
return self.grasps["label"].mean() * 100
def percent_cleared(self):
df = (
self.grasps[["round_id", "label"]]
.groupby("round_id")
.sum()
.rename(columns={"label": "cleared_count"})
.merge(self.rounds, on="round_id")
)
return df["cleared_count"].sum() / df["object_count"].sum() * 100
def avg_planning_time(self):
return self.grasps["planning_time"].mean()
def read_grasp(self, i):
scene_id, grasp, label = io.read_grasp(self.grasps, i)
score = self.grasps.loc[i, "score"]
scene_data = np.load(self.logdir / "scenes" / (scene_id + ".npz"))
return scene_data["points"], grasp, score, label
================================================
FILE: src/gd/grasp.py
================================================
import enum
class Label(enum.IntEnum):
FAILURE = 0 # grasp execution failed due to collision or slippage
SUCCESS = 1 # object was successfully removed
class Grasp(object):
"""Grasp parameterized as pose of a 2-finger robot hand.
TODO(mbreyer): clarify definition of grasp frame
"""
def __init__(self, pose, width):
self.pose = pose
self.width = width
def to_voxel_coordinates(grasp, voxel_size):
pose = grasp.pose
pose.translation /= voxel_size
width = grasp.width / voxel_size
return Grasp(pose, width)
def from_voxel_coordinates(grasp, voxel_size):
pose = grasp.pose
pose.translation *= voxel_size
width = grasp.width * voxel_size
return Grasp(pose, width)
================================================
FILE: src/gd/io.py
================================================
import json
import uuid
import numpy as np
import pandas as pd
from gd.grasp import Grasp
from gd.perception import *
from gd.utils.transform import Rotation, Transform
def write_setup(root, size, intrinsic, max_opening_width, finger_depth):
data = {
"size": size,
"intrinsic": intrinsic.to_dict(),
"max_opening_width": max_opening_width,
"finger_depth": finger_depth,
}
write_json(data, root / "setup.json")
def read_setup(root):
data = read_json(root / "setup.json")
size = data["size"]
intrinsic = CameraIntrinsic.from_dict(data["intrinsic"])
max_opening_width = data["max_opening_width"]
finger_depth = data["finger_depth"]
return size, intrinsic, max_opening_width, finger_depth
def write_sensor_data(root, depth_imgs, extrinsics):
scene_id = uuid.uuid4().hex
path = root / "scenes" / (scene_id + ".npz")
np.savez_compressed(path, depth_imgs=depth_imgs, extrinsics=extrinsics)
return scene_id
def read_sensor_data(root, scene_id):
data = np.load(root / "scenes" / (scene_id + ".npz"))
return data["depth_imgs"], data["extrinsics"]
def write_grasp(root, scene_id, grasp, label):
# TODO concurrent writes could be an issue
csv_path = root / "grasps.csv"
if not csv_path.exists():
create_csv(
csv_path,
["scene_id", "qx", "qy", "qz", "qw", "x", "y", "z", "width", "label"],
)
qx, qy, qz, qw = grasp.pose.rotation.as_quat()
x, y, z = grasp.pose.translation
width = grasp.width
append_csv(csv_path, scene_id, qx, qy, qz, qw, x, y, z, width, label)
def read_grasp(df, i):
scene_id = df.loc[i, "scene_id"]
orientation = Rotation.from_quat(df.loc[i, "qx":"qw"].to_numpy(np.double))
position = df.loc[i, "x":"z"].to_numpy(np.double)
width = df.loc[i, "width"]
label = df.loc[i, "label"]
grasp = Grasp(Transform(orientation, position), width)
return scene_id, grasp, label
def read_df(root):
return pd.read_csv(root / "grasps.csv")
def write_df(df, root):
df.to_csv(root / "grasps.csv", index=False)
def write_voxel_grid(root, scene_id, voxel_grid):
path = root / "scenes" / (scene_id + ".npz")
np.savez_compressed(path, grid=voxel_grid)
def read_voxel_grid(root, scene_id):
path = root / "scenes" / (scene_id + ".npz")
return np.load(path)["grid"]
def read_json(path):
with path.open("r") as f:
data = json.load(f)
return data
def write_json(data, path):
with path.open("w") as f:
json.dump(data, f, indent=4)
def create_csv(path, columns):
with path.open("w") as f:
f.write(",".join(columns))
f.write("\n")
def append_csv(path, *args):
row = ",".join([str(arg) for arg in args])
with path.open("a") as f:
f.write(row)
f.write("\n")
================================================
FILE: src/gd/networks.py
================================================
from builtins import super
import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy import ndimage
def get_network(name):
models = {
"conv": ConvNet(),
}
return models[name.lower()]
def load_network(path, device):
"""Construct the neural network and load parameters from the specified file.
Args:
path: Path to the model parameters. The name must conform to `vgn_name_[_...]`.
"""
model_name = path.stem.split("_")[1]
net = get_network(model_name).to(device)
net.load_state_dict(torch.load(path, map_location=device))
return net
def conv(in_channels, out_channels, kernel_size):
return nn.Conv3d(in_channels, out_channels, kernel_size, padding=kernel_size // 2)
def conv_stride(in_channels, out_channels, kernel_size):
return nn.Conv3d(
in_channels, out_channels, kernel_size, stride=2, padding=kernel_size // 2
)
class ConvNet(nn.Module):
def __init__(self):
super().__init__()
self.encoder = Encoder(1, [16, 32, 64], [5, 3, 3])
self.decoder = Decoder(64, [64, 32, 16], [3, 3, 5])
self.conv_qual = conv(16, 1, 5)
self.conv_rot = conv(16, 4, 5)
self.conv_width = conv(16, 1, 5)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
qual_out = torch.sigmoid(self.conv_qual(x))
rot_out = F.normalize(self.conv_rot(x), dim=1)
width_out = self.conv_width(x)
return qual_out, rot_out, width_out
class Encoder(nn.Module):
def __init__(self, in_channels, filters, kernels):
super().__init__()
self.conv1 = conv_stride(in_channels, filters[0], kernels[0])
self.conv2 = conv_stride(filters[0], filters[1], kernels[1])
self.conv3 = conv_stride(filters[1], filters[2], kernels[2])
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = self.conv3(x)
x = F.relu(x)
return x
class Decoder(nn.Module):
def __init__(self, in_channels, filters, kernels):
super().__init__()
self.conv1 = conv(in_channels, filters[0], kernels[0])
self.conv2 = conv(filters[0], filters[1], kernels[1])
self.conv3 = conv(filters[1], filters[2], kernels[2])
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = F.interpolate(x, 10)
x = self.conv2(x)
x = F.relu(x)
x = F.interpolate(x, 20)
x = self.conv3(x)
x = F.relu(x)
x = F.interpolate(x, 40)
return x
def count_num_trainable_parameters(net):
return sum(p.numel() for p in net.parameters() if p.requires_grad)
================================================
FILE: src/gd/perception.py
================================================
from math import cos, sin
import time
import numpy as np
import open3d as o3d
from gd.utils.transform import Transform
class CameraIntrinsic(object):
"""Intrinsic parameters of a pinhole camera model.
Attributes:
width (int): The width in pixels of the camera.
height(int): The height in pixels of the camera.
K: The intrinsic camera matrix.
"""
def __init__(self, width, height, fx, fy, cx, cy, channel=1):
self.width = width
self.height = height
self.channel = channel
self.K = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]])
@property
def fx(self):
return self.K[0, 0]
@property
def fy(self):
return self.K[1, 1]
@property
def cx(self):
return self.K[0, 2]
@property
def cy(self):
return self.K[1, 2]
def to_dict(self):
"""Serialize intrinsic parameters to a dict object."""
data = {
"width": self.width,
"height": self.height,
"channel": self.channel,
"K": self.K.flatten().tolist(),
}
return data
@classmethod
def from_dict(cls, data):
"""Deserialize intrinisic parameters from a dict object."""
intrinsic = cls(
width=data["width"],
height=data["height"],
channel=data["channel"],
fx=data["K"][0],
fy=data["K"][4],
cx=data["K"][2],
cy=data["K"][5],
)
return intrinsic
class TSDFVolume(object):
"""Integration of multiple depth images using a TSDF."""
def __init__(self, size, resolution):
self.size = size
self.resolution = resolution
self.voxel_size = self.size / self.resolution
self.sdf_trunc = 4 * self.voxel_size
self._volume = o3d.pipelines.integration.UniformTSDFVolume(
length=self.size,
resolution=self.resolution,
sdf_trunc=self.sdf_trunc,
color_type=o3d.pipelines.integration.TSDFVolumeColorType.NoColor,
)
def integrate(self, depth_img, intrinsic, extrinsic):
"""
Args:
depth_img: The depth image.
intrinsic: The intrinsic parameters of a pinhole camera model.
extrinsics: The transform from the TSDF to camera coordinates, T_eye_task.
"""
rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
o3d.geometry.Image(np.empty_like(depth_img)),
o3d.geometry.Image(depth_img),
depth_scale=1.0,
depth_trunc=2.0,
convert_rgb_to_intensity=False,
)
intrinsic = o3d.camera.PinholeCameraIntrinsic(
width=intrinsic.width,
height=intrinsic.height,
fx=intrinsic.fx,
fy=intrinsic.fy,
cx=intrinsic.cx,
cy=intrinsic.cy,
)
self._volume.integrate(rgbd, intrinsic, extrinsic)
def get_grid(self):
cloud = self._volume.extract_voxel_point_cloud()
points = np.asarray(cloud.points)
distances = np.asarray(cloud.colors)[:, [0]]
grid = np.zeros((1, 40, 40, 40), dtype=np.float32)
for idx, point in enumerate(points):
i, j, k = np.floor(point / self.voxel_size).astype(int)
grid[0, i, j, k] = distances[idx]
return grid
def get_cloud(self):
return self._volume.extract_point_cloud()
def create_tsdf(size, resolution, depth_imgs, intrinsic, extrinsics):
tsdf = TSDFVolume(size, resolution)
for i in range(depth_imgs.shape[0]):
extrinsic = Transform.from_list(extrinsics[i])
tsdf.integrate(depth_imgs[i], intrinsic, extrinsic)
return tsdf
def camera_on_sphere(origin, radius, theta, phi):
eye = np.r_[
radius * sin(theta) * cos(phi),
radius * sin(theta) * sin(phi),
radius * cos(theta),
]
target = np.array([0.0, 0.0, 0.0])
up = np.array([0.0, 0.0, 1.0]) # this breaks when looking straight down
return Transform.look_at(eye, target, up) * origin.inverse()
================================================
FILE: src/gd/simulation.py
================================================
from pathlib import Path
import time
import os
import numpy as np
import pybullet
from gd.grasp import Label
from gd.perception import *
from gd.utils import btsim, workspace_lines
from gd.utils.transform import Rotation, Transform
class ClutterRemovalSim(object):
def __init__(self, scene, object_set, gui=True, seed=None, renderer_root_dir="", args=None):
assert scene in ["pile", "packed", "single"]
self.urdf_root = Path(renderer_root_dir + "/data/urdfs")
self.scene = scene
self.object_set = object_set
self.discover_objects()
self.global_scaling = {"blocks": 1.67}.get(object_set, 1.0)
self.gui = gui
self.rng = np.random.RandomState(seed) if seed else np.random
self.world = btsim.BtWorld(self.gui)
self.gripper = Gripper(self.world)
self.size = 6 * self.gripper.finger_depth
intrinsic = CameraIntrinsic(640, 480, 540.0, 540.0, 320.0, 240.0) # TODO: cfg
self.camera = self.world.add_camera(intrinsic, 0.1, 2.0)
##
self.args = args
self.renderer_root_dir = renderer_root_dir
if self.args.load_scene_descriptor:
if self.scene == "pile":
dir_name = "pile_pile_test_200"
elif self.scene == "packed":
dir_name = "packed_packed_test_200"
elif self.scene == "single":
dir_name = "single_single_test_200"
scene_root_dir = os.path.join(renderer_root_dir, "data/mesh_pose_list", dir_name)
self.scene_descriptor_list = [os.path.join(scene_root_dir, i) for i in sorted(os.listdir(scene_root_dir))]
@property
def num_objects(self):
return max(0, self.world.p.getNumBodies() - 1) # remove table from body count
def discover_objects(self):
root = self.urdf_root / self.object_set
self.object_urdfs = [f for f in root.iterdir() if f.suffix == ".urdf"]
def save_state(self):
self._snapshot_id = self.world.save_state()
def restore_state(self):
self.world.restore_state(self._snapshot_id)
def reset(self, object_count, n_round):
self.world.reset()
self.world.set_gravity([0.0, 0.0, -9.81])
self.draw_workspace()
if self.gui:
self.world.p.resetDebugVisualizerCamera(
cameraDistance=1.0,
cameraYaw=0.0,
cameraPitch=-45,
cameraTargetPosition=[0.15, 0.50, -0.3],
)
table_height = self.gripper.finger_depth
self.place_table(table_height)
##
if self.args.gen_scene_descriptor:
if self.scene == "pile":
urdfs_and_poses_dict = self.generate_pile_scene(object_count, table_height)
return urdfs_and_poses_dict
elif self.scene == "packed":
urdfs_and_poses_dict = self.generate_packed_scene(object_count, table_height)
return urdfs_and_poses_dict
else:
raise ValueError("Invalid scene argument")
elif self.args.load_scene_descriptor:
scene_descriptor_npz = self.scene_descriptor_list[n_round]
if self.scene == "pile":
urdfs_and_poses_dict = self.generate_pile_scene(object_count, table_height, scene_descriptor_npz)
elif self.scene == "packed":
urdfs_and_poses_dict = self.generate_packed_scene(object_count, table_height, scene_descriptor_npz)
elif self.scene == "single":
urdfs_and_poses_dict = self.generate_packedsingle_scene(object_count, table_height, scene_descriptor_npz)
else:
raise ValueError("Invalid scene argument")
return urdfs_and_poses_dict
else:
raise NotImplementedError
def draw_workspace(self):
points = workspace_lines(self.size)
color = [0.5, 0.5, 0.5]
for i in range(0, len(points), 2):
self.world.p.addUserDebugLine(
lineFromXYZ=points[i], lineToXYZ=points[i + 1], lineColorRGB=color
)
def place_table(self, height):
urdf = self.urdf_root / "setup" / "plane.urdf"
pose = Transform(Rotation.identity(), [0.15, 0.15, height])
self.world.load_urdf(urdf, pose, scale=0.6)
# define valid volume for sampling grasps
lx, ux = 0.02, self.size - 0.02
ly, uy = 0.02, self.size - 0.02
lz, uz = height + 0.005, self.size
self.lower = np.r_[lx, ly, lz]
self.upper = np.r_[ux, uy, uz]
def generate_seen_scene(self, table_height, mesh_pose_npz):
# place box
urdf = self.urdf_root / "setup" / "box.urdf"
pose = Transform(Rotation.identity(), np.r_[0.02, 0.02, table_height])
box = self.world.load_urdf(urdf, pose, scale=1.3)
# read mesh_pose_npz
print("########## scene name: ", mesh_pose_npz)
if self.args.check_seen_scene:
urdfs_and_poses_dict = np.load(mesh_pose_npz, allow_pickle=True)['pc']
urdf_path_list = list(urdfs_and_poses_dict[:,0])
obj_scale_list = list(urdfs_and_poses_dict[:,1])
obj_RT_list = list(urdfs_and_poses_dict[:,2])
urdfs_and_poses_dict = {} ##
for i in range(len(urdf_path_list)):
urdf = os.path.join(self.renderer_root_dir, urdf_path_list[i].replace("_visual.obj",".urdf"))
RT = obj_RT_list[i]
R = RT[:3,:3]
T = RT[:3,3]
rotation = Rotation.from_matrix(R)
pose = Transform(rotation ,T)
scale = obj_scale_list[i]
body = self.world.load_urdf(urdf, pose, scale=scale)
body.set_pose(pose=Transform(rotation, T))
# remove box
self.world.remove_body(box)
removed_object = True
while removed_object:
removed_object, obj_body_list = self.remove_objects_outside_workspace()
for urdf, scale, rest_pose_quat, rest_pose_trans, body_uid in obj_body_list:
urdfs_and_poses_dict[body_uid] = [scale, rest_pose_quat, rest_pose_trans, str(urdf)]
return urdfs_and_poses_dict
def generate_pile_scene(self, object_count, table_height, scene_descriptor_npz=None):
# place box
urdf = self.urdf_root / "setup" / "box.urdf"
pose = Transform(Rotation.identity(), np.r_[0.02, 0.02, table_height])
box = self.world.load_urdf(urdf, pose, scale=1.3)
urdfs_and_poses_dict = {}
if self.args.gen_scene_descriptor:
urdf_path_list = self.rng.choice(self.object_urdfs, size=object_count)
elif self.args.load_scene_descriptor:
dict = np.load(scene_descriptor_npz, allow_pickle=True).item()
obj_scale_list = [value[0] for value in dict.values()]
obj_quat_list = [value[1] for value in dict.values()]
obj_xy_list = [value[2] for value in dict.values()]
if self.scene != self.object_set:
urdf_path_list = [os.path.join(self.renderer_root_dir, value[3].replace(self.scene, self.object_set)) for value in dict.values()]
else:
urdf_path_list = [os.path.join(self.renderer_root_dir, value[3]) for value in dict.values()]
# drop objects
for i in range(len(urdf_path_list)):
if self.args.gen_scene_descriptor:
rotation = Rotation.random(random_state=self.rng)
xy = self.rng.uniform(1.0 / 3.0 * self.size, 2.0 / 3.0 * self.size, 2)
pose = Transform(rotation, np.r_[xy, table_height + 0.2])
scale = self.rng.uniform(0.8, 1.0)
# save info
urdfs_and_poses_dict[i] = [scale, pose.rotation.as_quat(), xy, str(urdf_path_list[i])] # (x, y, z, w)
elif self.args.load_scene_descriptor:
rotation = Rotation.from_quat(obj_quat_list[i])
xy = obj_xy_list[i]
pose = Transform(rotation, np.r_[xy, table_height + 0.2])
scale = obj_scale_list[i]
self.world.load_urdf(urdf_path_list[i], pose, scale=self.global_scaling * scale)
self.wait_for_objects_to_rest(timeout=1.0)
# remove box
self.world.remove_body(box)
obj_body_list = self.remove_and_wait()
if self.args.gen_scene_descriptor:
return urdfs_and_poses_dict
else:
for urdf, scale, rest_pose_quat, rest_pose_trans, body_uid in obj_body_list:
urdfs_and_poses_dict[body_uid] = [scale, rest_pose_quat, rest_pose_trans, str(urdf)]
return urdfs_and_poses_dict
def generate_packed_scene(self, object_count, table_height, scene_descriptor_npz=None):
attempts = 0
max_attempts = 12
if self.args.gen_scene_descriptor:
urdfs_and_poses_dict = {}
elif self.args.load_scene_descriptor:
dict = np.load(scene_descriptor_npz, allow_pickle=True).item()
obj_scale_list = [value[0] for value in dict.values()]
obj_angle_list = [value[1] for value in dict.values()]
obj_x_list = [value[2] for value in dict.values()]
obj_y_list = [value[3] for value in dict.values()]
if self.scene != self.object_set:
urdf_path_list = [os.path.join(self.renderer_root_dir, value[4].replace(self.scene, self.object_set)) for value in dict.values()]
else:
urdf_path_list = [os.path.join(self.renderer_root_dir, value[4]) for value in dict.values()]
while self.num_objects < object_count and attempts < max_attempts:
self.save_state()
if self.args.gen_scene_descriptor:
urdf = self.rng.choice(self.object_urdfs)
x = self.rng.uniform(0.08, 0.22)
y = self.rng.uniform(0.08, 0.22)
angle = self.rng.uniform(0.0, 2.0 * np.pi)
scale = self.rng.uniform(0.7, 0.9)
# save info
urdfs_and_poses_dict[attempts] = [scale, angle, x, y, str(urdf)] # (x, y, z, w)
elif self.args.load_scene_descriptor:
urdf = urdf_path_list[attempts]
angle = obj_angle_list[attempts]
x = obj_x_list[attempts]
y = obj_y_list[attempts]
scale = obj_scale_list[attempts]
rotation = Rotation.from_rotvec(angle * np.r_[0.0, 0.0, 1.0])
z = 1.0
pose = Transform(rotation, np.r_[x, y, z])
body = self.world.load_urdf(urdf, pose, scale=self.global_scaling * scale)
lower, upper = self.world.p.getAABB(body.uid)
z = table_height + 0.5 * (upper[2] - lower[2]) + 0.002
body.set_pose(pose=Transform(rotation, np.r_[x, y, z]))
self.world.step()
if self.world.get_contacts(body):
self.world.remove_body(body)
self.restore_state()
else:
self.remove_and_wait()
attempts += 1
if self.args.gen_scene_descriptor:
return urdfs_and_poses_dict
else:
remain_obj_inws_infos = []
for body in list(self.world.bodies.values()):
urdf = self.world.bodies_urdfs[body.uid][0]
scale = self.world.bodies_urdfs[body.uid][1]
if str(urdf).split("/")[-1] != "box.urdf" and str(urdf).split("/")[-1] != "plane.urdf":
rest_pose = body.get_pose()
rest_pose_quat = rest_pose.rotation.as_quat() # (x, y, z, w)
rest_pose_trans = rest_pose.translation
remain_obj_inws_infos.append([urdf, scale, rest_pose_quat, rest_pose_trans, str(body.uid)])
urdfs_and_poses_dict = {}
for urdf, scale, rest_pose_quat, rest_pose_trans, body_uid in remain_obj_inws_infos:
urdfs_and_poses_dict[body_uid] = [scale, rest_pose_quat, rest_pose_trans, str(urdf)]
return urdfs_and_poses_dict
def generate_packedsingle_scene(self, object_count, table_height, scene_descriptor_npz=None):
attempts = 0
if self.args.gen_scene_descriptor:
urdfs_and_poses_dict = {}
elif self.args.load_scene_descriptor:
dict = np.load(scene_descriptor_npz, allow_pickle=True).item()
obj_scale_list = [value[0] for value in dict.values()]
obj_angle_list = [value[1] for value in dict.values()]
obj_x_list = [value[2] for value in dict.values()]
obj_y_list = [value[3] for value in dict.values()]
if self.scene != self.object_set:
urdf_path_list = [os.path.join(self.renderer_root_dir, value[4].replace(self.scene, self.object_set)) for value in dict.values()]
else:
urdf_path_list = [os.path.join(self.renderer_root_dir, value[4]) for value in dict.values()]
for _ in range(1):
self.save_state()
if self.args.gen_scene_descriptor:
urdf = self.rng.choice(self.object_urdfs)
x = self.rng.uniform(0.08, 0.22)
y = self.rng.uniform(0.08, 0.22)
angle = self.rng.uniform(0.0, 2.0 * np.pi)
scale = self.rng.uniform(0.7, 0.9)
# save info
urdfs_and_poses_dict[attempts] = [scale, angle, x, y, str(urdf)] # (x, y, z, w)
elif self.args.load_scene_descriptor:
urdf = urdf_path_list[attempts]
angle = obj_angle_list[attempts]
x = obj_x_list[attempts]
y = obj_y_list[attempts]
scale = obj_scale_list[attempts]
rotation = Rotation.from_rotvec(angle * np.r_[0.0, 0.0, 1.0])
z = 1.0
pose = Transform(rotation, np.r_[x, y, z])
body = self.world.load_urdf(urdf, pose, scale=self.global_scaling * scale)
lower, upper = self.world.p.getAABB(body.uid)
z = table_height + 0.5 * (upper[2] - lower[2]) + 0.002
body.set_pose(pose=Transform(rotation, np.r_[x, y, z]))
self.world.step()
if self.world.get_contacts(body):
self.world.remove_body(body)
self.restore_state()
else:
self.remove_and_wait()
attempts += 1
if self.args.gen_scene_descriptor:
return urdfs_and_poses_dict
else:
remain_obj_inws_infos = []
for body in list(self.world.bodies.values()):
urdf = self.world.bodies_urdfs[body.uid][0]
scale = self.world.bodies_urdfs[body.uid][1]
if str(urdf).split("/")[-1] != "box.urdf" and str(urdf).split("/")[-1] != "plane.urdf":
rest_pose = body.get_pose()
rest_pose_quat = rest_pose.rotation.as_quat() # (x, y, z, w)
rest_pose_trans = rest_pose.translation
remain_obj_inws_infos.append([urdf, scale, rest_pose_quat, rest_pose_trans, str(body.uid)])
urdfs_and_poses_dict = {}
for urdf, scale, rest_pose_quat, rest_pose_trans, body_uid in remain_obj_inws_infos:
urdfs_and_poses_dict[body_uid] = [scale, rest_pose_quat, rest_pose_trans, str(urdf)]
return urdfs_and_poses_dict
def acquire_tsdf(self, n, N=None):
"""Render synthetic depth images from n viewpoints and integrate into a TSDF.
If N is None, the n viewpoints are equally distributed on circular trajectory.
If N is given, the first n viewpoints on a circular trajectory consisting of N points are rendered.
"""
tsdf = TSDFVolume(self.size, 40)
high_res_tsdf = TSDFVolume(self.size, 120)
origin = Transform(Rotation.identity(), np.r_[self.size / 2, self.size / 2, 0])
r = 2.0 * self.size
theta = np.pi / 6.0
N = N if N else n
phi_list = 2.0 * np.pi * np.arange(n) / N
extrinsics = [camera_on_sphere(origin, r, theta, phi).as_matrix() for phi in phi_list]
timing = 0.0
for extrinsic in extrinsics:
depth_img = self.camera.render(extrinsic)[1]
tic = time.time()
tsdf.integrate(depth_img, self.camera.intrinsic, extrinsic)
timing += time.time() - tic
high_res_tsdf.integrate(depth_img, self.camera.intrinsic, extrinsic)
return tsdf, high_res_tsdf.get_cloud(), timing
def execute_grasp(self, grasp, remove=True, allow_contact=False):
# --grasp is the target containing pose and width
# -- flag to control whether allow collision between pre-target and target
# -- remove whether remove the objec from the scene after succesful grasp
T_world_grasp = grasp.pose
T_grasp_pregrasp = Transform(Rotation.identity(), [0.0, 0.0, -0.05])
T_world_pregrasp = T_world_grasp * T_grasp_pregrasp
# approach along z-axis of the gripper
approach = T_world_grasp.rotation.as_matrix()[:, 2]
angle = np.arccos(np.dot(approach, np.r_[0.0, 0.0, -1.0]))
if angle > np.pi / 3.0:
# side grasp, lift the object after establishing a grasp
T_grasp_pregrasp_world = Transform(Rotation.identity(), [0.0, 0.0, 0.1])
T_world_retreat = T_grasp_pregrasp_world * T_world_grasp
else:
T_grasp_retreat = Transform(Rotation.identity(), [0.0, 0.0, -0.1])
T_world_retreat = T_world_grasp * T_grasp_retreat
# move the gripper to pregrasp pose and detect the collision
self.gripper.reset(T_world_pregrasp)
if self.gripper.detect_contact():
result = Label.FAILURE, self.gripper.max_opening_width
print("0")
else:
#move the gripper to the target pose and detect collision
self.gripper.move_tcp_xyz(T_world_grasp, abort_on_contact=True)
"""
self.set_obj_pose_again(self.mesh_pose_npz)
"""
if self.gripper.detect_contact() and not allow_contact:
result = Label.FAILURE, self.gripper.max_opening_width
print("1")
else:
self.gripper.move(0.0) # shrink the gripper
# lift the gripper up along z-axis of the world frame or z-axis of the gripper frame
self.gripper.move_tcp_xyz(T_world_retreat, abort_on_contact=False)
if self.check_success(self.gripper):
result = Label.SUCCESS, self.gripper.read()
print("2")
if remove:
contacts = self.world.get_contacts(self.gripper.body)
self.world.remove_body(contacts[0].bodyB, isRemoveObjPerGrasp=True)
else:
result = Label.FAILURE, self.gripper.max_opening_width
print("3")
self.world.remove_body(self.gripper.body)
remain_obj_inws_infos = []
if remove:
remain_obj_inws_infos = self.remove_and_wait() ### wait for blender to render updated scene
return result, remain_obj_inws_infos
def remove_and_wait(self):
# wait for objects to rest while removing bodies that fell outside the workspace
removed_object = True
while removed_object:
self.wait_for_objects_to_rest()
removed_object, remain_obj_inws_infos = self.remove_objects_outside_workspace()
return remain_obj_inws_infos
def wait_for_objects_to_rest(self, timeout=2.0, tol=0.01):
timeout = self.world.sim_time + timeout
objects_resting = False
while not objects_resting and self.world.sim_time < timeout:
# simulate a quarter of a second
for _ in range(60):
self.world.step()
# check whether all objects are resting
objects_resting = True
for _, body in self.world.bodies.items():
if np.linalg.norm(body.get_velocity()) > tol:
objects_resting = False
break
def remove_objects_outside_workspace(self):
removed_object = False
remain_obj_inws_infos = [] ##
for body in list(self.world.bodies.values()):
xyz = body.get_pose().translation
if np.any(xyz < 0.0) or np.any(xyz > self.size):
self.world.remove_body(body)
removed_object = True
else:
urdf = self.world.bodies_urdfs[body.uid][0]
scale = self.world.bodies_urdfs[body.uid][1]
if str(urdf).split("/")[-1] != "box.urdf" and str(urdf).split("/")[-1] != "plane.urdf":
rest_pose = body.get_pose()
rest_pose_quat = rest_pose.rotation.as_quat() # (x, y, z, w)
rest_pose_trans = rest_pose.translation
remain_obj_inws_infos.append([urdf, scale, rest_pose_quat, rest_pose_trans, str(body.uid)])
return removed_object, remain_obj_inws_infos ##
def check_success(self, gripper):
# check that the fingers are in contact with some object and not fully closed
contacts = self.world.get_contacts(gripper.body)
res = len(contacts) > 0 and gripper.read() > 0.1 * gripper.max_opening_width
return res
class Gripper(object):
"""Simulated Panda hand."""
def __init__(self, world):
self.world = world
self.urdf_path = Path("data/assets/data/urdfs/panda/hand.urdf") #TODO put in cfg
self.max_opening_width = 0.08
self.finger_depth = 0.05
self.T_body_tcp = Transform(Rotation.identity(), [0.0, 0.0, 0.022])
self.T_tcp_body = self.T_body_tcp.inverse()
def reset(self, T_world_tcp):
T_world_body = T_world_tcp * self.T_tcp_body
self.body = self.world.load_urdf(self.urdf_path, T_world_body)
self.body.set_pose(T_world_body) # sets the position of the COM, not URDF link
self.constraint = self.world.add_constraint(
self.body,
None,
None,
None,
pybullet.JOINT_FIXED,
[0.0, 0.0, 0.0],
Transform.identity(),
T_world_body,
)
self.update_tcp_constraint(T_world_tcp)
# constraint to keep fingers centered
self.world.add_constraint(
self.body,
self.body.links["panda_leftfinger"],
self.body,
self.body.links["panda_rightfinger"],
pybullet.JOINT_GEAR,
[1.0, 0.0, 0.0],
Transform.identity(),
Transform.identity(),
).change(gearRatio=-1, erp=0.1, maxForce=50)
self.joint1 = self.body.joints["panda_finger_joint1"]
self.joint1.set_position(0.5 * self.max_opening_width, kinematics=True)
self.joint2 = self.body.joints["panda_finger_joint2"]
self.joint2.set_position(0.5 * self.max_opening_width, kinematics=True)
def update_tcp_constraint(self, T_world_tcp):
T_world_body = T_world_tcp * self.T_tcp_body
self.constraint.change(
jointChildPivot=T_world_body.translation,
jointChildFrameOrientation=T_world_body.rotation.as_quat(),
maxForce=300,
)
def set_tcp(self, T_world_tcp):
T_word_body = T_world_tcp * self.T_tcp_body
self.body.set_pose(T_word_body)
self.update_tcp_constraint(T_world_tcp)
def move_tcp_xyz(self, target, eef_step=0.002, vel=0.10, abort_on_contact=True):
T_world_body = self.body.get_pose()
T_world_tcp = T_world_body * self.T_body_tcp
diff = target.translation - T_world_tcp.translation
n_steps = int(np.linalg.norm(diff) / eef_step)
dist_step = diff / n_steps
dur_step = np.linalg.norm(dist_step) / vel
for _ in range(n_steps):
T_world_tcp.translation += dist_step
self.update_tcp_constraint(T_world_tcp)
for _ in range(int(dur_step / self.world.dt)):
self.world.step()
if abort_on_contact and self.detect_contact():
return
def detect_contact(self, threshold=5):
if self.world.get_contacts(self.body):
return True
else:
return False
def move(self, width):
self.joint1.set_position(0.5 * width)
self.joint2.set_position(0.5 * width)
for _ in range(int(0.5 / self.world.dt)):
self.world.step()
def read(self):
width = self.joint1.get_position() + self.joint2.get_position()
return width
================================================
FILE: src/gd/utils/__init__.py
================================================
def workspace_lines(size):
return [
[0.0, 0.0, 0.0],
[size, 0.0, 0.0],
[size, 0.0, 0.0],
[size, size, 0.0],
[size, size, 0.0],
[0.0, size, 0.0],
[0.0, size, 0.0],
[0.0, 0.0, 0.0],
[0.0, 0.0, size],
[size, 0.0, size],
[size, 0.0, size],
[size, size, size],
[size, size, size],
[0.0, size, size],
[0.0, size, size],
[0.0, 0.0, size],
[0.0, 0.0, 0.0],
[0.0, 0.0, size],
[size, 0.0, 0.0],
[size, 0.0, size],
[size, size, 0.0],
[size, size, size],
[0.0, size, 0.0],
[0.0, size, size],
]
================================================
FILE: src/gd/utils/btsim.py
================================================
import time
import numpy as np
import pybullet
from pybullet_utils import bullet_client
from vgn.utils.transform import Rotation, Transform
assert pybullet.isNumpyEnabled(), "Pybullet needs to be built with NumPy"
class BtWorld(object):
"""Interface to a PyBullet physics server.
Attributes:
dt: Time step of the physics simulation.
rtf: Real time factor. If negative, the simulation is run as fast as possible.
sim_time: Virtual time elpased since the last simulation reset.
"""
def __init__(self, gui=True):
connection_mode = pybullet.GUI if gui else pybullet.DIRECT
self.p = bullet_client.BulletClient(connection_mode)
self.gui = gui
self.dt = 1.0 / 240.0
self.solver_iterations = 150
self.bodies_urdfs = {} ##
self.reset()
def set_gravity(self, gravity):
self.p.setGravity(*gravity)
def load_urdf(self, urdf_path, pose, scale=1.0):
body = Body.from_urdf(self.p, urdf_path, pose, scale)
self.bodies[body.uid] = body
self.bodies_urdfs[body.uid] = [urdf_path, scale] ##
return body
def remove_body(self, body, isRemoveObjPerGrasp=False):
self.p.removeBody(body.uid)
del self.bodies[body.uid]
if isRemoveObjPerGrasp: ##
self.bodies_urdfs.pop(body.uid)
def add_constraint(self, *argv, **kwargs):
"""See `Constraint` below."""
constraint = Constraint(self.p, *argv, **kwargs)
return constraint
def add_camera(self, intrinsic, near, far):
camera = Camera(self.p, intrinsic, near, far)
return camera
def get_contacts(self, bodyA):
points = self.p.getContactPoints(bodyA.uid)
contacts = []
for point in points:
###
urdf = self.bodies_urdfs[point[2]][0]
if str(urdf).split("/")[-1] == "plane.urdf":
continue
###
contact = Contact(
bodyA=self.bodies[point[1]],
bodyB=self.bodies[point[2]],
point=point[5],
normal=point[7],
depth=point[8],
force=point[9],
)
contacts.append(contact)
return contacts
def reset(self):
self.p.resetSimulation()
self.p.setPhysicsEngineParameter(
fixedTimeStep=self.dt, numSolverIterations=self.solver_iterations
)
self.bodies = {}
self.sim_time = 0.0
def step(self):
self.p.stepSimulation()
self.sim_time += self.dt
if self.gui:
time.sleep(self.dt)
def save_state(self):
return self.p.saveState()
def restore_state(self, state_uid):
self.p.restoreState(stateId=state_uid)
def close(self):
self.p.disconnect()
class Body(object):
"""Interface to a multibody simulated in PyBullet.
Attributes:
uid: The unique id of the body within the physics server.
name: The name of the body.
joints: A dict mapping joint names to Joint objects.
links: A dict mapping link names to Link objects.
"""
def __init__(self, physics_client, body_uid):
self.p = physics_client
self.uid = body_uid
self.name = self.p.getBodyInfo(self.uid)[1].decode("utf-8")
self.joints, self.links = {}, {}
for i in range(self.p.getNumJoints(self.uid)):
joint_info = self.p.getJointInfo(self.uid, i)
joint_name = joint_info[1].decode("utf8")
self.joints[joint_name] = Joint(self.p, self.uid, i)
link_name = joint_info[12].decode("utf8")
self.links[link_name] = Link(self.p, self.uid, i)
@classmethod
def from_urdf(cls, physics_client, urdf_path, pose, scale):
body_uid = physics_client.loadURDF(
str(urdf_path),
pose.translation,
pose.rotation.as_quat(),
globalScaling=scale,
)
return cls(physics_client, body_uid)
def get_pose(self):
pos, ori = self.p.getBasePositionAndOrientation(self.uid)
return Transform(Rotation.from_quat(ori), np.asarray(pos))
def set_pose(self, pose):
self.p.resetBasePositionAndOrientation(
self.uid, pose.translation, pose.rotation.as_quat()
)
def get_velocity(self):
linear, angular = self.p.getBaseVelocity(self.uid)
return linear, angular
class Link(object):
"""Interface to a link simulated in Pybullet.
Attributes:
link_index: The index of the joint.
"""
def __init__(self, physics_client, body_uid, link_index):
self.p = physics_client
self.body_uid = body_uid
self.link_index = link_index
def get_pose(self):
link_state = self.p.getLinkState(self.body_uid, self.link_index)
pos, ori = link_state[0], link_state[1]
return Transform(Rotation.from_quat(ori), pos)
class Joint(object):
"""Interface to a joint simulated in PyBullet.
Attributes:
joint_index: The index of the joint.
lower_limit: Lower position limit of the joint.
upper_limit: Upper position limit of the joint.
effort: The maximum joint effort.
"""
def __init__(self, physics_client, body_uid, joint_index):
self.p = physics_client
self.body_uid = body_uid
self.joint_index = joint_index
joint_info = self.p.getJointInfo(body_uid, joint_index)
self.lower_limit = joint_info[8]
self.upper_limit = joint_info[9]
self.effort = joint_info[10]
def get_position(self):
joint_state = self.p.getJointState(self.body_uid, self.joint_index)
return joint_state[0]
def set_position(self, position, kinematics=False):
if kinematics:
self.p.resetJointState(self.body_uid, self.joint_index, position)
self.p.setJointMotorControl2(
self.body_uid,
self.joint_index,
pybullet.POSITION_CONTROL,
targetPosition=position,
force=self.effort,
)
class Constraint(object):
"""Interface to a constraint in PyBullet.
Attributes:
uid: The unique id of the constraint within the physics server.
"""
def __init__(
self,
physics_client,
parent,
parent_link,
child,
child_link,
joint_type,
joint_axis,
parent_frame,
child_frame,
):
"""
Create a new constraint between links of bodies.
Args:
parent:
parent_link: None for the base.
child: None for a fixed frame in world coordinates.
"""
self.p = physics_client
parent_body_uid = parent.uid
parent_link_index = parent_link.link_index if parent_link else -1
child_body_uid = child.uid if child else -1
child_link_index = child_link.link_index if child_link else -1
self.uid = self.p.createConstraint(
parentBodyUniqueId=parent_body_uid,
parentLinkIndex=parent_link_index,
childBodyUniqueId=child_body_uid,
childLinkIndex=child_link_index,
jointType=joint_type,
jointAxis=joint_axis,
parentFramePosition=parent_frame.translation,
parentFrameOrientation=parent_frame.rotation.as_quat(),
childFramePosition=child_frame.translation,
childFrameOrientation=child_frame.rotation.as_quat(),
)
def change(self, **kwargs):
self.p.changeConstraint(self.uid, **kwargs)
class Contact(object):
"""Contact point between two multibodies.
Attributes:
point: Contact point.
normal: Normal vector from ... to ...
depth: Penetration depth
force: Contact force acting on body ...
"""
def __init__(self, bodyA, bodyB, point, normal, depth, force):
self.bodyA = bodyA
self.bodyB = bodyB
self.point = point
self.normal = normal
self.depth = depth
self.force = force
class Camera(object):
"""Virtual RGB-D camera based on the PyBullet camera interface.
Attributes:
intrinsic: The camera intrinsic parameters.
"""
def __init__(self, physics_client, intrinsic, near, far):
self.intrinsic = intrinsic
self.near = near
self.far = far
self.proj_matrix = _build_projection_matrix(intrinsic, near, far)
self.p = physics_client
def render(self, extrinsic):
"""Render synthetic RGB and depth images.
Args:
extrinsic: Extrinsic parameters, T_cam_ref.
"""
# Construct OpenGL compatible view and projection matrices.
gl_view_matrix = extrinsic
gl_view_matrix[2, :] *= -1 # flip the Z axis
gl_view_matrix = gl_view_matrix.flatten(order="F")
gl_proj_matrix = self.proj_matrix.flatten(order="F")
result = self.p.getCameraImage(
width=self.intrinsic.width,
height=self.intrinsic.height,
viewMatrix=gl_view_matrix,
projectionMatrix=gl_proj_matrix,
renderer=pybullet.ER_TINY_RENDERER,
)
rgb, z_buffer = result[2][:, :, :3], result[3]
depth = (
1.0 * self.far * self.near / (self.far - (self.far - self.near) * z_buffer)
)
return rgb, depth
def _build_projection_matrix(intrinsic, near, far):
perspective = np.array(
[
[intrinsic.fx, 0.0, -intrinsic.cx, 0.0],
[0.0, intrinsic.fy, -intrinsic.cy, 0.0],
[0.0, 0.0, near + far, near * far],
[0.0, 0.0, -1.0, 0.0],
]
)
ortho = _gl_ortho(0.0, intrinsic.width, intrinsic.height, 0.0, near, far)
return np.matmul(ortho, perspective)
def _gl_ortho(left, right, bottom, top, near, far):
ortho = np.diag(
[2.0 / (right - left), 2.0 / (top - bottom), -2.0 / (far - near), 1.0]
)
ortho[0, 3] = -(right + left) / (right - left)
ortho[1, 3] = -(top + bottom) / (top - bottom)
ortho[2, 3] = -(far + near) / (far - near)
return ortho
================================================
FILE: src/gd/utils/panda_control.py
================================================
from panda_robot import PandaArm
import rospy
import numpy as np
import quaternion
from gd.utils.transform import Transform
from scipy.spatial.transform import Rotation
class PandaCommander(object):
def __init__(self):
self.name = "panda_arm"
self.r = PandaArm()
self.r.enable_robot()
rospy.loginfo("PandaCommander ready")
self.moving = False
def reset(self):
rospy.logwarn("reset and go home!")
self.r.enable_robot()
self.home()
def clear(self):
self.r.exit_control_mode()
def home(self):
self.moving = True
self.r.move_to_neutral()
self.moving = False
rospy.loginfo("PandaCommander: Arrived home!")
def goto_joints(self, joints):
self.moving = True
self.r.move_to_joint_position(joints)
self.moving = False
def get_joints(self):
return self.r.angles()
def goto_pose(self, pose):
rospy.loginfo("PandaCommander: goto pose " + str(pose.to_list()[-3:]))
self.moving = True
x, y, z, w = pose.rotation.as_quat()
self.r.move_to_cartesian_pose(pose.translation.astype(np.float32), np.quaternion(w, x, y, z))
safe = self.r.in_safe_state()
if self.r.has_collided():
rospy.logwarn("collided!")
self.r.enable_robot()
return
error = self.r.error_in_current_state()
if error or not safe:
rospy.logwarn("error or not safe! reset and run again.")
self.reset()
self.goto_pose(pose)
return
self.moving = False
def get_pose(self):
pos, rot = self.r.ee_pose()
return Transform(Rotation.from_quat([rot.x, rot.y, rot.z, rot.w]), pos)
def grasp(self, width=0.0, force=10.0):
return self.r.exec_gripper_cmd(width, force=force)
def move_gripper(self, width):
return self.r.exec_gripper_cmd(width)
def get_gripper_width(self):
return np.sum(self.r.gripper_state()['position'])
================================================
FILE: src/gd/utils/ros_utils.py
================================================
import geometry_msgs.msg
import numpy as np
import rospy
from sensor_msgs.msg import PointCloud2, PointField
import std_msgs.msg
from gd.utils.transform import Rotation, Transform
def to_point_msg(position):
"""Convert numpy array to a Point message."""
msg = geometry_msgs.msg.Point()
msg.x = position[0]
msg.y = position[1]
msg.z = position[2]
return msg
def from_point_msg(msg):
"""Convert a Point message to a numpy array."""
return np.r_[msg.x, msg.y, msg.z]
def to_vector3_msg(vector3):
"""Convert numpy array to a Vector3 message."""
msg = geometry_msgs.msg.Vector3()
msg.x = vector3[0]
msg.y = vector3[1]
msg.z = vector3[2]
return msg
def from_vector3_msg(msg):
"""Convert a Vector3 message to a numpy array."""
return np.r_[msg.x, msg.y, msg.z]
def to_quat_msg(orientation):
"""Convert a `Rotation` object to a Quaternion message."""
quat = orientation.as_quat()
msg = geometry_msgs.msg.Quaternion()
msg.x = quat[0]
msg.y = quat[1]
msg.z = quat[2]
msg.w = quat[3]
return msg
def from_quat_msg(msg):
"""Convert a Quaternion message to a Rotation object."""
return Rotation.from_quat([msg.x, msg.y, msg.z, msg.w])
def to_pose_msg(transform):
"""Convert a `Transform` object to a Pose message."""
msg = geometry_msgs.msg.Pose()
msg.position = to_point_msg(transform.translation)
msg.orientation = to_quat_msg(transform.rotation)
return msg
def to_transform_msg(transform):
"""Convert a `Transform` object to a Transform message."""
msg = geometry_msgs.msg.Transform()
msg.translation = to_vector3_msg(transform.translation)
msg.rotation = to_quat_msg(transform.rotation)
return msg
def from_transform_msg(msg):
"""Convert a Transform message to a Transform object."""
translation = from_vector3_msg(msg.translation)
rotation = from_quat_msg(msg.rotation)
return Transform(rotation, translation)
def to_color_msg(color):
"""Convert a numpy array to a ColorRGBA message."""
msg = std_msgs.msg.ColorRGBA()
msg.r = color[0]
msg.g = color[1]
msg.b = color[2]
msg.a = color[3] if len(color) == 4 else 1.0
return msg
def to_cloud_msg(points, intensities=None, frame=None, stamp=None):
"""Convert list of unstructured points to a PointCloud2 message.
Args:
points: Point coordinates as array of shape (N,3).
colors: Colors as array of shape (N,3).
frame
stamp
"""
msg = PointCloud2()
msg.header.frame_id = frame
msg.header.stamp = stamp or rospy.Time.now()
msg.height = 1
msg.width = points.shape[0]
msg.is_bigendian = False
msg.is_dense = False
msg.fields = [
PointField("x", 0, PointField.FLOAT32, 1),
PointField("y", 4, PointField.FLOAT32, 1),
PointField("z", 8, PointField.FLOAT32, 1),
]
msg.point_step = 12
data = points
if intensities is not None:
msg.fields.append(PointField("intensity", 12, PointField.FLOAT32, 1))
msg.point_step += 4
data = np.hstack([points, intensities])
msg.row_step = msg.point_step * points.shape[0]
msg.data = data.astype(np.float32).tostring()
return msg
class TransformTree(object):
def __init__(self):
import tf2_ros
self._buffer = tf2_ros.Buffer()
self._listener = tf2_ros.TransformListener(self._buffer)
self._broadcaster = tf2_ros.TransformBroadcaster()
self._static_broadcaster = tf2_ros.StaticTransformBroadcaster()
def lookup(self, target_frame, source_frame, time, timeout=rospy.Duration(0)):
msg = self._buffer.lookup_transform(target_frame, source_frame, time, timeout)
return from_transform_msg(msg.transform)
def broadcast(self, transform, target_frame, source_frame):
msg = geometry_msgs.msg.TransformStamped()
msg.header.stamp = rospy.Time.now()
msg.header.frame_id = target_frame
msg.child_frame_id = source_frame
msg.transform = to_transform_msg(transform)
self._broadcaster.sendTransform(msg)
def broadcast_static(self, transform, target_frame, source_frame):
msg = geometry_msgs.msg.TransformStamped()
msg.header.stamp = rospy.Time.now()
msg.header.frame_id = target_frame
msg.child_frame_id = source_frame
msg.transform = to_transform_msg(transform)
self._static_broadcaster.sendTransform(msg)
================================================
FILE: src/gd/utils/transform.py
================================================
import numpy as np
import scipy.spatial.transform
class Rotation(scipy.spatial.transform.Rotation):
@classmethod
def identity(cls):
return cls.from_quat([0.0, 0.0, 0.0, 1.0])
class Transform(object):
"""Rigid spatial transform between coordinate systems in 3D space.
Attributes:
rotation (scipy.spatial.transform.Rotation)
translation (np.ndarray)
"""
def __init__(self, rotation, translation):
assert isinstance(rotation, scipy.spatial.transform.Rotation)
assert isinstance(translation, (np.ndarray, list))
self.rotation = rotation
self.translation = np.asarray(translation, np.double)
def as_matrix(self):
"""Represent as a 4x4 matrix."""
return np.vstack(
(np.c_[self.rotation.as_matrix(), self.translation], [0.0, 0.0, 0.0, 1.0])
)
def to_dict(self):
"""Serialize Transform object into a dictionary."""
return {
"rotation": self.rotation.as_quat().tolist(),
"translation": self.translation.tolist(),
}
def to_list(self):
return np.r_[self.rotation.as_quat(), self.translation]
def __mul__(self, other):
"""Compose this transform with another."""
rotation = self.rotation * other.rotation
translation = self.rotation.apply(other.translation) + self.translation
return self.__class__(rotation, translation)
def transform_point(self, point):
return self.rotation.apply(point) + self.translation
def transform_vector(self, vector):
return self.rotation.apply(vector)
def inverse(self):
"""Compute the inverse of this transform."""
rotation = self.rotation.inv()
translation = -rotation.apply(self.translation)
return self.__class__(rotation, translation)
@classmethod
def from_matrix(cls, m):
"""Initialize from a 4x4 matrix."""
rotation = Rotation.from_matrix(m[:3, :3])
translation = m[:3, 3]
return cls(rotation, translation)
@classmethod
def from_dict(cls, dictionary):
rotation = Rotation.from_quat(dictionary["rotation"])
translation = np.asarray(dictionary["translation"])
return cls(rotation, translation)
@classmethod
def from_list(cls, list):
rotation = Rotation.from_quat(list[:4])
translation = list[4:]
return cls(rotation, translation)
@classmethod
def identity(cls):
"""Initialize with the identity transformation."""
rotation = Rotation.from_quat([0.0, 0.0, 0.0, 1.0])
translation = np.array([0.0, 0.0, 0.0])
return cls(rotation, translation)
@classmethod
def look_at(cls, eye, center, up):
"""Initialize with a LookAt matrix.
Returns:
T_eye_ref, the transform from camera to the reference frame, w.r.t.
which the input arguments were defined.
"""
eye = np.asarray(eye)
center = np.asarray(center)
forward = center - eye
forward /= np.linalg.norm(forward)
right = np.cross(forward, up)
right /= np.linalg.norm(right)
up = np.asarray(up) / np.linalg.norm(up)
up = np.cross(right, forward)
m = np.eye(4, 4)
m[:3, 0] = right
m[:3, 1] = -up
m[:3, 2] = forward
m[:3, 3] = eye
return cls.from_matrix(m).inverse()
================================================
FILE: src/gd/vis.py
================================================
"""Render volumes, point clouds, and grasp detections in rviz."""
import matplotlib.colors
import numpy as np
from sensor_msgs.msg import PointCloud2
import rospy
from rospy import Publisher
from visualization_msgs.msg import Marker, MarkerArray
from gd.utils import ros_utils, workspace_lines
from gd.utils.transform import Transform, Rotation
cmap = matplotlib.colors.LinearSegmentedColormap.from_list("RedGreen", ["r", "g"])
DELETE_MARKER_MSG = Marker(action=Marker.DELETEALL)
DELETE_MARKER_ARRAY_MSG = MarkerArray(markers=[DELETE_MARKER_MSG])
def draw_workspace(size):
scale = size * 0.005
pose = Transform.identity()
scale = [scale, 0.0, 0.0]
color = [0.5, 0.5, 0.5]
msg = _create_marker_msg(Marker.LINE_LIST, "task", pose, scale, color)
msg.points = [ros_utils.to_point_msg(point) for point in workspace_lines(size)]
pubs["workspace"].publish(msg)
def draw_tsdf(vol, voxel_size, threshold=0.01):
msg = _create_vol_msg(vol, voxel_size, threshold)
pubs["tsdf"].publish(msg)
def draw_points(points):
msg = ros_utils.to_cloud_msg(points, frame="task")
pubs["points"].publish(msg)
def draw_quality(vol, voxel_size, threshold=0.01):
msg = _create_vol_msg(vol, voxel_size, threshold)
pubs["quality"].publish(msg)
def draw_volume(vol, voxel_size, threshold=0.01):
msg = _create_vol_msg(vol, voxel_size, threshold)
pubs["debug"].publish(msg)
def draw_grasp(grasp, score, finger_depth):
radius = 0.1 * finger_depth
w, d = grasp.width, finger_depth
color = cmap(float(score))
markers = []
# left finger
pose = grasp.pose * Transform(Rotation.identity(), [0.0, -w / 2, d / 2])
scale = [radius, radius, d]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 0
markers.append(msg)
# right finger
pose = grasp.pose * Transform(Rotation.identity(), [0.0, w / 2, d / 2])
scale = [radius, radius, d]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 1
markers.append(msg)
# wrist
pose = grasp.pose * Transform(Rotation.identity(), [0.0, 0.0, -d / 4])
scale = [radius, radius, d / 2]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 2
markers.append(msg)
# palm
pose = grasp.pose * Transform(
Rotation.from_rotvec(np.pi / 2 * np.r_[1.0, 0.0, 0.0]), [0.0, 0.0, 0.0]
)
scale = [radius, radius, w]
msg = _create_marker_msg(Marker.CYLINDER, "task", pose, scale, color)
msg.id = 3
markers.append(msg)
pubs["grasp"].publish(MarkerArray(markers=markers))
def draw_grasps(grasps, scores, finger_depth):
markers = []
for i, (grasp, score) in enumerate(zip(grasps, scores)):
msg = _create_grasp_marker_msg(grasp, score, finger_depth)
msg.id = i
markers.append(msg)
msg = MarkerArray(markers=markers)
pubs["grasps"].publish(msg)
def clear():
pubs["workspace"].publish(DELETE_MARKER_MSG)
pubs["tsdf"].publish(ros_utils.to_cloud_msg(np.array([]), frame="task"))
pubs["points"].publish(ros_utils.to_cloud_msg(np.array([]), frame="task"))
clear_quality()
pubs["grasp"].publish(DELETE_MARKER_ARRAY_MSG)
clear_grasps()
pubs["debug"].publish(ros_utils.to_cloud_msg(np.array([]), frame="task"))
def clear_quality():
pubs["quality"].publish(ros_utils.to_cloud_msg(np.array([]), frame="task"))
def clear_grasps():
pubs["grasps"].publish(DELETE_MARKER_ARRAY_MSG)
def _create_publishers():
pubs = dict()
pubs["workspace"] = Publisher("/workspace", Marker, queue_size=1, latch=True)
pubs["tsdf"] = Publisher("/tsdf", PointCloud2, queue_size=1, latch=True)
pubs["points"] = Publisher("/points", PointCloud2, queue_size=1, latch=True)
pubs["quality"] = Publisher("/quality", PointCloud2, queue_size=1, latch=True)
pubs["grasp"] = Publisher("/grasp", MarkerArray, queue_size=1, latch=True)
pubs["grasps"] = Publisher("/grasps", MarkerArray, queue_size=1, latch=True)
pubs["debug"] = Publisher("/debug", PointCloud2, queue_size=1, latch=True)
return pubs
def _create_marker_msg(marker_type, frame, pose, scale, color):
msg = Marker()
msg.header.frame_id = frame
msg.header.stamp = rospy.Time()
msg.type = marker_type
msg.action = Marker.ADD
msg.pose = ros_utils.to_pose_msg(pose)
msg.scale = ros_utils.to_vector3_msg(scale)
msg.color = ros_utils.to_color_msg(color)
return msg
def _create_vol_msg(vol, voxel_size, threshold):
vol = vol.squeeze()
if type(threshold) is tuple:
idx_arr = np.logical_and(vol > threshold[0], vol < threshold[1])
else:
idx_arr = vol > threshold
points = np.argwhere(idx_arr) * voxel_size
values = np.expand_dims(vol[idx_arr], 1)
return ros_utils.to_cloud_msg(points, values, frame="task")
def _create_grasp_marker_msg(grasp, score, finger_depth):
radius = 0.1 * finger_depth
w, d = grasp.width, finger_depth
scale = [radius, 0.0, 0.0]
color = list(cmap(float(score)))
if score < 0.01:
color = [0., 0., 1, 0.8]
msg = _create_marker_msg(Marker.LINE_LIST, "task", grasp.pose, scale, color)
msg.points = [ros_utils.to_point_msg(point) for point in _gripper_lines(w, d)]
return msg
def _gripper_lines(width, depth):
return [
[0.0, 0.0, -depth / 2.0],
[0.0, 0.0, 0.0],
[0.0, -width / 2.0, 0.0],
[0.0, -width / 2.0, depth],
[0.0, width / 2.0, 0.0],
[0.0, width / 2.0, depth],
[0.0, -width / 2.0, 0.0],
[0.0, width / 2.0, 0.0],
]
pubs = _create_publishers()
================================================
FILE: src/nr/asset.py
================================================
import os
import numpy as np
DATA_ROOT_DIR = '../../data/traindata_example/'
VGN_TRAIN_ROOT = DATA_ROOT_DIR + 'giga_hemisphere_train_demo'
def add_scenes(root, type, filter_list=None):
scene_names = []
splits = os.listdir(root)
for split in splits:
if filter_list is not None and split not in filter_list: continue
scenes = os.listdir(os.path.join(root, split))
scene_names += [f'vgn_syn/train/{type}/{split}/{fn}/w_0.8' for fn in scenes]
return scene_names
if os.path.exists(VGN_TRAIN_ROOT):
vgn_pile_train_scene_names = sorted(add_scenes(os.path.join(VGN_TRAIN_ROOT, 'pile_full'), 'pile'), key=lambda x: x.split('/')[4])
vgn_pack_train_scene_names = sorted(add_scenes(os.path.join(VGN_TRAIN_ROOT, 'packed_full'), 'packed'), key=lambda x: x.split('/')[4])
num_scenes_pile = len(vgn_pile_train_scene_names)
num_scenes_pack = len(vgn_pack_train_scene_names)
vgn_pack_train_scene_names = vgn_pack_train_scene_names[:num_scenes_pack]
num_val_pile = 1
num_val_pack = 1
print(f"total: {num_scenes_pile + num_scenes_pack} pile: {num_scenes_pile} pack: {num_scenes_pack}")
vgn_val_scene_names = vgn_pile_train_scene_names[-num_val_pile:] + vgn_pack_train_scene_names[-num_val_pack:]
vgn_train_scene_names = vgn_pile_train_scene_names[:-num_val_pile] + vgn_pack_train_scene_names[:-num_val_pack]
VGN_SDF_DIR = DATA_ROOT_DIR + "giga_hemisphere_train_demo/scenes_tsdf_dep-nor"
VGN_TEST_ROOT = ''
VGN_TEST_ROOT_PILE = os.path.join(VGN_TEST_ROOT,'pile')
VGN_TEST_ROOT_PACK = os.path.join(VGN_TEST_ROOT,'packed')
if os.path.exists(VGN_TEST_ROOT):
fns = os.listdir(VGN_TEST_ROOT_PILE)
vgn_pile_test_scene_names = [f'vgn_syn/test/pile//{fn}/w_0.8' for fn in fns]
fns = os.listdir(VGN_TEST_ROOT_PACK)
vgn_pack_test_scene_names = [f'vgn_syn/test/packed//{fn}/w_0.8' for fn in fns]
vgn_test_scene_names = vgn_pile_test_scene_names + vgn_pack_test_scene_names
CSV_ROOT = DATA_ROOT_DIR + 'GIGA_demo'
import pandas as pd
from pathlib import Path
import time
t0 = time.time()
VGN_PACK_TRAIN_CSV = pd.read_csv(Path(CSV_ROOT + '/data_packed_train_processed_dex_noise/grasps.csv'))
VGN_PILE_TRAIN_CSV = pd.read_csv(Path(CSV_ROOT + '/data_pile_train_processed_dex_noise/grasps.csv'))
print(f"finished loading csv in {time.time() - t0} s")
VGN_PACK_TEST_CSV = None
VGN_PILE_TEST_CSV = None
================================================
FILE: src/nr/configs/nrvgn_sdf.yaml
================================================
name: test
group_name: ""
# network
fix_seed: true
network: grasp_nerf
init_net_type: cost_volume
agg_net_type: neus
use_hierarchical_sampling: true
use_depth: true
use_depth_loss: true
depth_loss_weight: 1.0
dist_decoder_cfg:
use_vis: false
fine_dist_decoder_cfg:
use_vis: false
ray_batch_num: 4096 #2048
sample_volume: true
render_rgb: true
volume_type: [sdf]
volume_resolution: 40
depth_sample_num: 40
fine_depth_sample_num: 40
agg_net_cfg:
sample_num: 40
init_s: 0.3
fix_s: 0
fine_agg_net_cfg:
sample_num: 40
init_s: 0.3
fix_s: 0
vis_vol: false
# loss
loss: [render, depth, sdf, vgn]
val_metric: [psnr_ssim, vis_img]
key_metric_name: loss_vgn # depth_mae psnr_nr_fine
key_metric_prefer: lower
use_dr_loss: false
use_dr_fine_loss: false
use_nr_fine_loss: true
render_depth: true
depth_correct_ratio: 1.0
depth_thresh: 0.8
use_dr_prediction: false
# lr
total_step: 500000
val_interval: 5000
lr_type: exp_decay
lr_cfg:
lr_init: 1.0e-4
decay_step: 100000
decay_rate: 0.5
nr_initial_training_steps: 0
# dataset
train_dataset_type: gen
train_dataset_cfg:
resolution_type: hr
type2sample_weights: { vgn_syn: 100 }
train_database_types: ['vgn_syn']
aug_pixel_center_sample: true
aug_view_select_type: hard
ref_pad_interval: 32
use_src_imgs: true
num_input_views: 6
val_set_list:
-
name: vgn_syn
type: gen
val_scene_num: -1 # if the set, use val scene list in asset.py
cfg:
use_src_imgs: true
num_input_views: 6
================================================
FILE: src/nr/dataset/database.py
================================================
import abc
import glob
import json
import os
import re
from pathlib import Path
import sys
import open3d as o3d
from utils.draw_utils import draw_gripper_o3d
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
import cv2
import numpy as np
from skimage.io import imread, imsave
from asset import VGN_TRAIN_ROOT, VGN_TEST_ROOT, VGN_PILE_TRAIN_CSV, VGN_PACK_TRAIN_CSV, VGN_PILE_TEST_CSV,VGN_PACK_TEST_CSV, VGN_SDF_DIR
from utils.draw_utils import draw_cube, draw_axis, draw_points, draw_gripper, draw_world_points
sys.path.append("../")
from gd.utils.transform import Rotation, Transform
class BaseDatabase(abc.ABC):
def __init__(self, database_name):
self.database_name = database_name
@abc.abstractmethod
def get_image(self, img_id):
pass
@abc.abstractmethod
def get_K(self, img_id):
pass
@abc.abstractmethod
def get_pose(self, img_id):
pass
@abc.abstractmethod
def get_img_ids(self,check_depth_exist=False):
pass
@abc.abstractmethod
def get_bbox(self, img_id):
pass
@abc.abstractmethod
def get_depth(self,img_id):
pass
@abc.abstractmethod
def get_mask(self,img_id):
pass
@abc.abstractmethod
def get_depth_range(self,img_id):
pass
class GraspSynDatabase(BaseDatabase):
def __init__(self, database_name):
super().__init__(database_name)
self.debug_save_dir = Path(f'output/nrvgn/{database_name}')
tp, split, scene_type, scene_split, scene_id, background_size = database_name.split('/')
background, size = background_size.split('_')
self.split = split
self.scene_id = scene_id
self.scene_type = scene_type
self.tp = tp
self.downSample = float(size)
tp2wh = {
'vgn_syn': (640, 360)
}
src_wh = tp2wh[tp]
self.img_wh = (np.array(src_wh) * self.downSample).astype(int)
root_dir = {'test': {
'vgn_syn': VGN_TEST_ROOT,
},
'train': {
'vgn_syn': VGN_TRAIN_ROOT,
},
}
if tp == 'vgn_syn':
self.root_dir = Path(root_dir[split][tp]) / (scene_type + "_full") / scene_split / scene_id
else:
raise NotImplementedError
tp2len = {'grasp_syn': 256,
'vgn_syn':24}
self.depth_img_ids = self.img_ids = list(range(tp2len[tp]))
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.K = np.array([[
892.62,
0.0,
639.5
],
[
0.0,
892.62,
359.5
],
[
0.0,
0.0,
1.0
]])
self.K[:2] = self.K[:2] * self.downSample
if self.tp == 'vgn_syn':
self.K[:2] /= 2
self.poses_ori = np.load(self.root_dir / 'camera_pose.npy')
self.poses = [np.linalg.inv(p @ self.blender2opencv)[:3,:] for p in self.poses_ori]
self.depth_thres = {
'grasp_syn': 1.5,
'vgn_syn': 0.8,
}
self.fixed_depth_range = [0.2, 0.8]
tp2bbox3d = {'grasp_syn': [[-0.35, -0.45, 0],
[0.15, 0.05, 0.5]],
'vgn_syn': [[-0.15, -0.15, -0.05],
[0.15, 0.15, 0.25]]}
self.bbox3d = tp2bbox3d[tp]
def get_split(self):
return self.split
def get_image(self, img_id):
img_filename = os.path.join(self.root_dir,
f'rgb/{img_id:04d}.png')
img = imread(img_filename)[:,:,:3]
img = cv2.resize(img, self.img_wh)
#img[self.get_mask(img_id)] = 255
return np.asarray(img, dtype=np.float32)
def get_K(self, img_id):
return self.K.astype(np.float32).copy()
def get_pose(self, img_id):
return self.poses[img_id].astype(np.float32).copy()
def get_img_ids(self,check_depth_exist=False):
if check_depth_exist: return self.depth_img_ids
return self.img_ids
def get_bbox3d(self, vis=False):
if vis:
img_id = 0
img = self.get_image(img_id)
cRb = self.poses[img_id][:3,:3]
ctb = self.poses[img_id][:3,3]
l = self.bbox3d[1][0] - self.bbox3d[0][0]
img = draw_cube(img, cRb, ctb, self.K, length=l, bias=self.bbox3d[0])
if not self.debug_save_dir.exists():
self.debug_save_dir.mkdir(parents=True)
imsave(str(self.debug_save_dir / 'bbox3d.jpg'), img)
return self.bbox3d
def get_bbox(self, img_id, vis=False):
mask = self.get_mask(img_id,'obj')
xs,ys=np.nonzero(mask)
x_min,x_max=np.min(xs,0),np.max(xs,0)
y_min,y_max=np.min(ys,0),np.max(ys,0)
if vis:
img = self.get_image(img_id)
img = cv2.rectangle(img, (y_min, x_min), (y_max, x_max), (255,0,0), 2)
imsave(str(self.debug_save_dir / 'box.jpg'), img)
return [x_min,x_max,y_min,y_max]
def _depth_existence(self,img_id):
return True
def get_depth(self, img_id):
depth_filename = os.path.join(self.root_dir,
f'depth/{img_id:04d}.exr')
depth_h = cv2.imread(depth_filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[:,:,0]
depth_h = cv2.resize(depth_h, self.img_wh, interpolation=cv2.INTER_NEAREST)
return depth_h
def get_mask(self, img_id, tp='desk'):
if tp == 'desk':
mask = self.get_depth(img_id) < self.depth_thres[self.tp]
return (mask.astype(np.bool))
elif tp == 'obj':
mask_filename = os.path.join(self.root_dir,
f'mask/{img_id:04d}.exr')
mask = cv2.imread(mask_filename, cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)[:,:,0]
mask = cv2.resize(mask, self.img_wh, interpolation=cv2.INTER_NEAREST)
cv2.imwrite('mask.jpg', mask * 256)
return ~(mask.astype(np.bool))
else:
return np.ones((self.img_wh[1], self.img_wh[0]))
def get_depth_range(self,img_id, fixed=True):
if fixed:
return np.array(self.fixed_depth_range)
depth = self.get_depth(img_id)
nf = [max(0, np.min(depth)), min(self.depth_thres[self.tp], np.max(depth))]
return np.array(nf)
def get_sdf(self):
sdf_volume = np.load( Path(VGN_SDF_DIR) / f'{self.scene_id}.npz')['grid'][0]
return sdf_volume * 2 - 1
class VGNSynDatabase(GraspSynDatabase):
def __init__(self, database_name):
super().__init__(database_name)
split = self.get_split()
if self.scene_type == 'packed':
csv = VGN_PACK_TEST_CSV if split == 'test' else VGN_PACK_TRAIN_CSV
elif self.scene_type == 'pile':
csv = VGN_PILE_TEST_CSV if split == 'test' else VGN_PILE_TRAIN_CSV
else:
return
self.df = csv
self.df = self.df[self.df["scene_id"] == self.scene_id]
assert len(self.df) > 0, f"empty grasping info {database_name}"
def visualize_grasping(self, pos, rot, w, label=None, img_id=3,save_img=False):
voxel_size = 0.3 / 40
pts_w = pos * voxel_size
width = w * voxel_size
img = self.get_image(img_id)
t = np.array([[-0.15, -0.15, -0.05]]).repeat(pts_w.shape[0], axis=0)
pts_b = pts_w + t
cRb = self.poses[img_id][:3,:3]
ctb = self.poses[img_id][:3,3] # + np.array([-0.15, -0.15, -0.05])
for gid in range(pts_w.shape[0]):
if label is not None and label[gid] == 0:
continue
btg = pts_b[gid]
wRg = rot[gid]
bRg = wRg
bTg = np.eye(4)
bTg[:3,:3] = bRg
bTg[:3,3] = btg
cTb = self.poses[img_id]
cTg = cTb @ bTg
img = draw_gripper(img, cTg[:3,:3], cTg[:3,3], self.K, width[gid], 2)
img = draw_world_points(img, pts_b[gid], cRb, ctb, self.K)
if save_img:
save_dir = str(self.debug_save_dir / f'gripper_test-{img_id}.jpg')
print("save to", save_dir)
imsave(save_dir, img)
return img
def visualize_grasping_3d(self, pos, rot, w, label=None, voxel_size = 0.3 / 40):
pts_w = pos * voxel_size
width = w * voxel_size
geometry = o3d.geometry.TriangleMesh()
for gid in range(pts_w.shape[0]):
if label is not None and label[gid] == 0:
continue
wRg = rot[gid]
y_ccw_90 = np.array([[0, 0, -1], [0, 1,0], [1, 0, 0]])
_R = wRg @ y_ccw_90
_t = pts_w[gid]
geometry_gripper = draw_gripper_o3d(_R, _t, width[gid])
geometry += geometry_gripper
o3d.io.write_triangle_mesh(str(self.debug_save_dir / f'gripper.ply'), geometry)
def get_grasp_info(self):
pos = self.df[["i","j","k"]].to_numpy(np.single)
index = np.round(pos).astype(np.long)
l = pos.shape[0]
width = self.df[["width"]].to_numpy(np.single).reshape(l)
label = self.df[["label"]].to_numpy(np.float32).reshape(l)
rotations = np.empty((l, 2, 4), dtype=np.single)
q = self.df[["qx","qy","qz","qw"]].to_numpy(np.single)
ori = Rotation.from_quat(q)
R = Rotation.from_rotvec(np.pi * np.r_[0.0, 0.0, 1.0])
rotations[:,0] = ori.as_quat()
rotations[:,1] = (ori * R).as_quat()
# for i in range(4):
# self.visualize_grasping(pos, ori.as_matrix(), width, label, i)
# exit()
return (index, label, rotations, width)
def parse_database_name(database_name:str)->BaseDatabase:
name2database={
'vgn_syn': VGNSynDatabase,
}
database_type = database_name.split('/')[0]
if database_type in name2database:
return name2database[database_type](database_name)
else:
raise NotImplementedError
def get_database_split(database: BaseDatabase, split_type='val'):
database_name = database.database_name
if split_type.startswith('val'):
splits = split_type.split('_')
depth_valid = not(len(splits)>1 and splits[1]=='all')
if database_name.startswith('vgn'):
val_ids = database.get_img_ids()[2:24:8]# TODO
train_ids = [img_id for img_id in database.get_img_ids(check_depth_exist=depth_valid) if img_id not in val_ids]
else:
raise NotImplementedError
elif split_type.startswith('test'):
splits = split_type.split('_')
depth_valid = not(len(splits)>1 and splits[1]=='all')
if database_name.startswith('vgn'):
val_ids = database.get_img_ids()[2:24:8] + [0]# TODO
train_ids = [img_id for img_id in database.get_img_ids(check_depth_exist=depth_valid) if img_id not in val_ids]
else:
raise NotImplementedError
else:
raise NotImplementedError
return train_ids, val_ids
================================================
FILE: src/nr/dataset/name2dataset.py
================================================
from dataset.train_dataset import GeneralRendererDataset, FinetuningRendererDataset
name2dataset={
'gen': GeneralRendererDataset,
'ft': FinetuningRendererDataset,
}
================================================
FILE: src/nr/dataset/train_dataset.py
================================================
from torch.utils.data import Dataset
from asset import *
from dataset.database import parse_database_name, get_database_split
import numpy as np
from utils.base_utils import get_coords_mask
from utils.dataset_utils import set_seed
from utils.imgs_info import build_imgs_info, random_crop, random_flip, pad_imgs_info, imgs_info_slice, \
imgs_info_to_torch, grasp_info_to_torch
from utils.view_select import compute_nearest_camera_indices
def select_train_ids_for_real_estate(img_ids):
num_frames = len(img_ids)
window_size = 32
shift = np.random.randint(low=-1, high=2)
id_render = np.random.randint(low=4, high=num_frames - 4 - 1)
right_bound = min(id_render + window_size + shift, num_frames - 1)
left_bound = max(0, right_bound - 2 * window_size)
candidate_ids = np.arange(left_bound, right_bound)
# remove the query frame itself with high probability
if np.random.choice([0, 1], p=[0.01, 0.99]):
candidate_ids = candidate_ids[candidate_ids != id_render]
id_feat = np.random.choice(candidate_ids, size=min(8, len(candidate_ids)), replace=False)
img_ids = np.asarray(img_ids)
return img_ids[id_render], img_ids[id_feat]
def add_depth_offset(depth,mask,region_min,region_max,offset_min,offset_max,noise_ratio,depth_length):
coords = np.stack(np.nonzero(mask), -1)[:, (1, 0)]
length = np.max(coords, 0) - np.min(coords, 0)
center = coords[np.random.randint(0, coords.shape[0])]
lx, ly = np.random.uniform(region_min, region_max, 2) * length
diff = coords - center[None, :]
mask0 = np.abs(diff[:, 0]) < lx
mask1 = np.abs(diff[:, 1]) < ly
masked_coords = coords[mask0 & mask1]
global_offset = np.random.uniform(offset_min, offset_max) * depth_length
if np.random.random() < 0.5:
global_offset = -global_offset
local_offset = np.random.uniform(-noise_ratio, noise_ratio, masked_coords.shape[0]) * depth_length + global_offset
depth[masked_coords[:, 1], masked_coords[:, 0]] += local_offset
def build_src_imgs_info_select(database, ref_ids, ref_ids_all, cost_volume_nn_num, pad_interval=-1, self_ref=True):
if not self_ref:
# ref_ids - selected ref ids for rendering
ref_idx_exp = compute_nearest_camera_indices(database, ref_ids, ref_ids_all)
ref_idx_exp = ref_idx_exp[:, 1:1 + cost_volume_nn_num]
ref_ids_all = np.asarray(ref_ids_all)
ref_ids_exp = ref_ids_all[ref_idx_exp] # rfn,nn
ref_ids_exp_ = ref_ids_exp.flatten()
ref_ids = np.asarray(ref_ids)
ref_ids_in = np.unique(np.concatenate([ref_ids_exp_, ref_ids])) # rfn'
mask0 = ref_ids_in[None, :] == ref_ids[:, None] # rfn,rfn'
ref_idx_, ref_idx = np.nonzero(mask0)
ref_real_idx = ref_idx[np.argsort(ref_idx_)] # sort
rfn, nn = ref_ids_exp.shape
mask1 = ref_ids_in[None, :] == ref_ids_exp.flatten()[:, None] # nn*rfn,rfn'
ref_cv_idx_, ref_cv_idx = np.nonzero(mask1)
ref_cv_idx = ref_cv_idx[np.argsort(ref_cv_idx_)] # sort
ref_cv_idx = ref_cv_idx.reshape([rfn, nn])
else: # not using extra view to construct cost volume
ref_ids_in = ref_ids
ref_real_idx = np.asarray(list(range(len(ref_ids))))
ref_cv_idx = np.asarray([ref_real_idx for _ in range(len(ref_ids))])
#print("ref_ids", ref_ids, "ref_ids_in", ref_ids_in, "ref_cv_idx", ref_cv_idx, "ref_real_idx", ref_real_idx)
is_aligned = not database.database_name.startswith('space')
ref_imgs_info = build_imgs_info(database, ref_ids_in, pad_interval, is_aligned)
return ref_imgs_info, ref_cv_idx, ref_real_idx
class GeneralRendererDataset(Dataset):
default_cfg={
'train_database_types':['dtu_train','space','real_iconic','real_estate','gso'],
'type2sample_weights': {'gso':20, 'dtu_train':20, 'real_iconic':20, 'space':10, 'real_estate':10},
'val_database_name': 'nerf_synthetic/lego/black_800',
'val_database_split_type': 'val',
"total_views": 24,
"num_input_views": 3,
'min_wn': 3,
'max_wn': 4,
'ref_pad_interval': 16,
'train_ray_num': 512,
'foreground_ratio': 0.5,
'resolution_type': 'hr',
"use_consistent_depth_range": True,
'use_depth_loss_for_all': False,
"use_depth": True,
"use_src_imgs": False,
"cost_volume_nn_num": 3,
"aug_gso_shrink_range_prob": 0.5,
"aug_depth_range_prob": 0.05, #TODO
'aug_depth_range_min': 0.95,
'aug_depth_range_max': 1.05,
"aug_use_depth_offset": True,
"aug_depth_offset_prob": 0.25,
"aug_depth_offset_region_min": 0.05,
"aug_depth_offset_region_max": 0.1,
'aug_depth_offset_min': 0.5,
'aug_depth_offset_max': 1.0,
'aug_depth_offset_local': 0.1,
"aug_use_depth_small_offset": True,
"aug_use_global_noise": True,
"aug_global_noise_prob": 0.5,
"aug_depth_small_offset_prob": 0.5,
"aug_forward_crop_size": (400,600),
"aug_pixel_center_sample": False,
"aug_view_select_type": "easy",
"use_consistent_min_max": False,
"revise_depth_range": False,
'load_sdf': True,
'exclude_ref_views': False,
}
def __init__(self, cfg, is_train):
self.cfg={**self.default_cfg,**cfg}
self.is_train = is_train
if is_train:
self.num=999999
self.type2scene_names,self.database_types,self.database_weights = {}, [], []
if self.cfg['resolution_type']=='hr':
type2scene_names={'vgn_syn': vgn_train_scene_names}
elif self.cfg['resolution_type']=='lr':
type2scene_names={'vgn_syn': vgn_train_scene_names}
else:
raise NotImplementedError
for database_type in self.cfg['train_database_types']:
self.type2scene_names[database_type] = type2scene_names[database_type]
self.database_types.append(database_type)
self.database_weights.append(self.cfg['type2sample_weights'][database_type])
print(f"training scene num: {len(type2scene_names[database_type])}")
assert(len(self.database_types)>0)
# normalize weights
self.database_weights=np.asarray(self.database_weights)
self.database_weights=self.database_weights/np.sum(self.database_weights)
else:
self.database = parse_database_name(self.cfg['val_database_name'])
self.ref_ids, self.que_ids = get_database_split(self.database,self.cfg['val_database_split_type'])
self.num=len(self.que_ids)
def get_database_ref_que_ids(self, index):
if self.is_train:
database_type = np.random.choice(self.database_types,1,False,p=self.database_weights)[0]
database_scene_name = np.random.choice(self.type2scene_names[database_type])
database = parse_database_name(database_scene_name)
# if there is no depth for all views, we repeat random sample until find a scene with depth
while True:
ref_ids = database.get_img_ids(check_depth_exist=True)
if len(ref_ids)==0:
database_type = np.random.choice(self.database_types, 1, False, self.database_weights)[0]
database_scene_name = np.random.choice(self.type2scene_names[database_type])
database = parse_database_name(database_scene_name)
else: break
que_id = np.random.choice(ref_ids)
if database.database_name.startswith('real_estate'):
que_id, ref_ids = select_train_ids_for_real_estate(ref_ids)
else:
database = self.database
que_id, ref_ids = self.que_ids[index], self.ref_ids
return database, que_id, np.asarray(ref_ids)
def select_working_views_impl(self, database_name, dist_idx, ref_num):
if self.cfg['aug_view_select_type']=='default':
if database_name.startswith('space') or database_name.startswith('real_estate'):
pass
elif database_name.startswith('gso'):
pool_ratio = np.random.randint(1, 5)
dist_idx = dist_idx[:min(ref_num * pool_ratio, 32)]
elif database_name.startswith('real_iconic'):
pool_ratio = np.random.randint(1, 5)
dist_idx = dist_idx[:min(ref_num * pool_ratio, 32)]
elif database_name.startswith('dtu_train'):
pool_ratio = np.random.randint(1, 3)
dist_idx = dist_idx[:min(ref_num * pool_ratio, 12)]
else:
raise NotImplementedError
elif self.cfg['aug_view_select_type']=='easy':
if database_name.startswith('space') or database_name.startswith('real_estate'):
pass
elif database_name.startswith('gso'):
pool_ratio = 3
dist_idx = dist_idx[:min(ref_num * pool_ratio, 24)]
elif database_name.startswith('real_iconic'):
pool_ratio = np.random.randint(1, 4)
dist_idx = dist_idx[:min(ref_num * pool_ratio, 20)]
elif database_name.startswith('dtu_train'):
pool_ratio = np.random.randint(1, 3)
dist_idx = dist_idx[:min(ref_num * pool_ratio, 12)]
else:
raise NotImplementedError
elif self.cfg['aug_view_select_type']=='hard':
if database_name.startswith('grasp'):
dist_idx = dist_idx[80:]
elif database_name.startswith('vgn'):
dist_idx = dist_idx[8:]
else:
raise NotImplementedError
return dist_idx
def get_ref_que_ids(self, target_id):
N = self.cfg['total_views']
interval = list(range(0, N, N // self.cfg['num_input_views']))
res = [(target_id + i) % N for i in interval]
que_id = ( np.random.choice(res) + np.random.randint(1, N // self.cfg['num_input_views']) ) % N
return res, que_id
def select_working_views(self, database, que_id, ref_ids):
database_name = database.database_name
dist_idx = compute_nearest_camera_indices(database, [que_id], ref_ids)[0]
if self.is_train:
if np.random.random()>0.02: # 2% chance to include que image
dist_idx = dist_idx[ref_ids[dist_idx]!=que_id]
ref_num = np.random.randint(self.cfg['min_wn'], self.cfg['max_wn'])
dist_idx = self.select_working_views_impl(database_name,dist_idx,ref_num)
if database_name.startswith('grasp'):
ref_id = np.random.randint(0, 256)
ref_ids = [ref_id, (ref_id + 80) % 256, (ref_id + 160) % 256]
elif database_name.startswith('vgn'):
ref_ids, que_id = self.get_ref_que_ids(np.random.randint(0, self.cfg['total_views']))
elif not database_name.startswith('real_estate'):
# we already select working views for real estate dataset
np.random.shuffle(dist_idx)
dist_idx = dist_idx[:ref_num]
ref_ids = ref_ids[dist_idx]
else:
ref_ids = ref_ids[:ref_num]
else:
if database_name.startswith('vgn'):
ref_ids, que_id = self.get_ref_que_ids(que_id)
elif database_name.startswith('grasp'):
ref_ids = [que_id, (que_id + 80) % 256, (que_id + 160) % 256]
else:
dist_idx = dist_idx[:self.cfg['min_wn']]
ref_ids = ref_ids[dist_idx]
return ref_ids, que_id
def depth_range_aug_for_gso(self, depth_range, depth, mask):
depth_range_new = depth_range.copy()
if np.random.random() < self.cfg['aug_gso_shrink_range_prob']:
rfn, _, h, w = depth.shape
far_ratios, near_ratios = [], []
for rfi in range(rfn):
depth_val = depth[rfi][mask[rfi].astype(np.bool)]
depth_val = depth_val[depth_val > 1e-3]
depth_val = depth_val[depth_val < 1e4]
depth_max = np.max(depth_val) * 1.1
depth_min = np.min(depth_val) * 0.9
near, far = depth_range[rfi]
far_ratio = depth_max / far
near_ratio = near / depth_min
far_ratios.append(far_ratio)
near_ratios.append(near_ratio)
far_ratio = np.max(far_ratios)
near_ratio = np.max(near_ratios)
if far_ratio < 1.0: depth_range_new[:, 1] *= np.random.uniform(far_ratio, 1.0)
if near_ratio < 1.0: depth_range_new[:, 0] /= np.random.uniform(near_ratio, 1.0)
if np.random.random()<0.8:
ratio0, ratio1 = np.random.uniform(0.025, 0.1, 2)
depth_range_new[:, 0] = depth_range_new[:, 0] * (1 - ratio0)
depth_range_new[:, 1] = depth_range_new[:, 1] * (1 + ratio1)
return depth_range_new
def random_change_depth_range(self, depth_range, depth, mask, database_name):
if database_name.startswith('gso'):
depth_range_new = self.depth_range_aug_for_gso(depth_range, depth, mask)
else:
depth_range_new = depth_range.copy()
if np.random.random()0
coords = get_coords_mask(que_mask_cur, self.cfg['train_ray_num'], self.cfg['foreground_ratio']).reshape([1,-1,2])
return coords
def consistent_depth_range(self, ref_imgs_info, que_imgs_info):
depth_range_all = np.concatenate([ref_imgs_info['depth_range'], que_imgs_info['depth_range']], 0)
if self.cfg['use_consistent_min_max']:
depth_range_all[:, 0] = np.min(depth_range_all)
depth_range_all[:, 1] = np.max(depth_range_all)
else:
range_len = depth_range_all[:, 1] - depth_range_all[:, 0]
max_len = np.max(range_len)
range_margin = (max_len - range_len) / 2
ref_near = depth_range_all[:, 0] - range_margin
ref_near = np.max(np.stack([ref_near, depth_range_all[:, 0] * 0.5], -1), 1)
depth_range_all[:, 0] = ref_near
depth_range_all[:, 1] = ref_near + max_len
ref_imgs_info['depth_range'] = depth_range_all[:-1]
que_imgs_info['depth_range'] = depth_range_all[-1:]
def __getitem__(self, index):
set_seed(index, self.is_train) #TODO: currently que_id and ref is all randomly sampled
database, que_id, ref_ids_all = self.get_database_ref_que_ids(index)
ref_ids, new_que_id = self.select_working_views(database, que_id, ref_ids_all)
if self.cfg['exclude_ref_views']:
que_id = new_que_id
# print(que_id, ref_ids)
if self.cfg['use_src_imgs']:
# src_imgs_info used in construction of cost volume
ref_imgs_info, ref_cv_idx, ref_real_idx = build_src_imgs_info_select(database,ref_ids,ref_ids_all,self.cfg['cost_volume_nn_num'])
else:
ref_idx = compute_nearest_camera_indices(database, ref_ids)[:,1:4] # used in cost volume construction
is_aligned = not database.database_name.startswith('space')
ref_imgs_info = build_imgs_info(database, ref_ids, -1, is_aligned)
que_imgs_info = build_imgs_info(database, [que_id], has_depth=True)
# data augmentation
depth_range_all = np.concatenate([ref_imgs_info['depth_range'],que_imgs_info['depth_range']],0)
if database.database_name.startswith('gso'): # only used in gso currently
depth_all = np.concatenate([ref_imgs_info['depth'],que_imgs_info['depth']],0)
mask_all = np.concatenate([ref_imgs_info['masks'],que_imgs_info['masks']],0)
else:
depth_all, mask_all = None, None
depth_range_all = self.random_change_depth_range(depth_range_all, depth_all, mask_all, database.database_name)
ref_imgs_info['depth_range'] = depth_range_all[:-1]
que_imgs_info['depth_range'] = depth_range_all[-1:]
if database.database_name.startswith('gso') and self.cfg['use_depth']:
depth_aug = self.add_depth_noise(ref_imgs_info['depth'], ref_imgs_info['masks'], ref_imgs_info['depth_range'])
ref_imgs_info['true_depth'] = ref_imgs_info['depth']
ref_imgs_info['depth'] = depth_aug
if database.database_name.startswith('real_estate') \
or database.database_name.startswith('real_iconic') \
or database.database_name.startswith('space'):
# crop all datasets
ref_imgs_info, que_imgs_info = random_crop(ref_imgs_info, que_imgs_info, self.cfg['aug_forward_crop_size'])
if np.random.random()<0.5:
ref_imgs_info, que_imgs_info = random_flip(ref_imgs_info, que_imgs_info)
if self.cfg['use_depth_loss_for_all'] and self.cfg['use_depth']:
if not database.database_name.startswith('gso'):
ref_imgs_info['true_depth'] = ref_imgs_info['depth']
if database.database_name.startswith('grasp') or database.database_name.startswith('vgn'):
ref_imgs_info['true_depth'] = ref_imgs_info['depth']
que_imgs_info['true_depth'] = que_imgs_info['depth']
if self.cfg['use_consistent_depth_range']:
self.consistent_depth_range(ref_imgs_info, que_imgs_info)
# generate coords
if self.is_train:
coords = self.generate_coords_for_training(database,que_imgs_info)
else:
qn, _, hn, wn = que_imgs_info['imgs'].shape
coords = np.stack(np.meshgrid(np.arange(wn),np.arange(hn)),-1)
coords = coords.reshape([1,-1,2]).astype(np.float32)
que_imgs_info['coords'] = coords
ref_imgs_info = pad_imgs_info(ref_imgs_info,self.cfg['ref_pad_interval'])
# don't feed depth to gpu
if not self.cfg['use_depth']:
if 'depth' in ref_imgs_info: ref_imgs_info.pop('depth')
if 'depth' in que_imgs_info: que_imgs_info.pop('depth')
if 'true_depth' in ref_imgs_info: ref_imgs_info.pop('true_depth')
if self.cfg['use_src_imgs']:
src_imgs_info = ref_imgs_info.copy()
ref_imgs_info = imgs_info_slice(ref_imgs_info, ref_real_idx)
ref_imgs_info['nn_ids'] = ref_cv_idx
else:
# 'nn_ids' used in constructing cost volume (specify source image ids)
ref_imgs_info['nn_ids'] = ref_idx.astype(np.int64)
if self.cfg['load_sdf']:
ref_imgs_info['sdf_gt'] = database.get_sdf()
ref_imgs_info = imgs_info_to_torch(ref_imgs_info)
que_imgs_info = imgs_info_to_torch(que_imgs_info)
outputs = {'ref_imgs_info': ref_imgs_info, 'que_imgs_info': que_imgs_info, 'scene_name': database.database_name}
if self.cfg['use_src_imgs']: outputs['src_imgs_info'] = imgs_info_to_torch(src_imgs_info)
if database.database_name.startswith('vgn'):
outputs['grasp_info'] = grasp_info_to_torch(database.get_grasp_info())
return outputs
def __len__(self):
return self.num
class FinetuningRendererDataset(Dataset):
default_cfg={
"database_name": "nerf_synthetic/lego/black_800",
"database_split_type": "val_all"
}
def __init__(self,cfg, is_train):
self.cfg={**self.default_cfg,**cfg}
self.is_train=is_train
self.train_ids, self.val_ids = get_database_split(parse_database_name(self.cfg['database_name']),self.cfg['database_split_type'])
def __getitem__(self, index):
output={'index': index}
return output
def __len__(self):
if self.is_train:
return 99999999
else:
return len(self.val_ids)
================================================
FILE: src/nr/main.py
================================================
import sys, os
import time
sys.path.append("./src/nr")
from pathlib import Path
import numpy as np
import torch
from skimage.io import imsave, imread
from network.renderer import name2network
from utils.base_utils import load_cfg, to_cuda
from utils.imgs_info import build_render_imgs_info, imgs_info_to_torch, grasp_info_to_torch
from network.renderer import name2network
from utils.base_utils import color_map_forward
from network.loss import VGNLoss
from tqdm import tqdm
from scipy import ndimage
import cv2
from gd.utils.transform import Transform, Rotation
from gd.grasp import *
def process(
tsdf_vol,
qual_vol,
rot_vol,
width_vol,
gaussian_filter_sigma=1.0,
min_width=1.33,
max_width=9.33,
tsdf_thres_high = 0.5,
tsdf_thres_low = 1e-3,
n_grasp=0
):
tsdf_vol = tsdf_vol.squeeze()
qual_vol = qual_vol.squeeze()
rot_vol = rot_vol.squeeze()
width_vol = width_vol.squeeze()
# smooth quality volume with a Gaussian
qual_vol = ndimage.gaussian_filter(
qual_vol, sigma=gaussian_filter_sigma, mode="nearest"
)
# mask out voxels too far away from the surface
outside_voxels = tsdf_vol > tsdf_thres_high
inside_voxels = np.logical_and(tsdf_thres_low < tsdf_vol, tsdf_vol < tsdf_thres_high)
valid_voxels = ndimage.morphology.binary_dilation(
outside_voxels, iterations=2, mask=np.logical_not(inside_voxels)
)
qual_vol[valid_voxels == False] = 0.0
# reject voxels with predicted widths that are too small or too large
qual_vol[np.logical_or(width_vol < min_width, width_vol > max_width)] = 0.0
return qual_vol, rot_vol, width_vol
def select(qual_vol, rot_vol, width_vol, threshold=0.90, max_filter_size=4):
qual_vol[qual_vol < threshold] = 0.0
# non maximum suppression
max_vol = ndimage.maximum_filter(qual_vol, size=max_filter_size)
qual_vol = np.where(qual_vol == max_vol, qual_vol, 0.0)
mask = np.where(qual_vol, 1.0, 0.0)
# construct grasps
grasps, scores, indexs = [], [], []
for index in np.argwhere(mask):
indexs.append(index)
grasp, score = select_index(qual_vol, rot_vol, width_vol, index)
grasps.append(grasp)
scores.append(score)
return grasps, scores, indexs
def select_index(qual_vol, rot_vol, width_vol, index):
i, j, k = index
score = qual_vol[i, j, k]
rot = rot_vol[:, i, j, k]
ori = Rotation.from_quat(rot)
pos = np.array([i, j, k], dtype=np.float64)
width = width_vol[i, j, k]
return Grasp(Transform(ori, pos), width), score
class GraspNeRFPlanner(object):
def set_params(self, args):
self.args = args
self.voxel_size = 0.3 / 40
self.bbox3d = [[-0.15, -0.15, -0.0503],[0.15, 0.15, 0.2497]]
self.tsdf_thres_high = 0
self.tsdf_thres_low = -0.85
self.renderer_root_dir = self.args.renderer_root_dir
tp, split, scene_type, scene_split, scene_id, background_size = args.database_name.split('/')
background, size = background_size.split('_')
self.split = split
self.tp = tp
self.downSample = float(size)
tp2wh = {
'vgn_syn': (640, 360)
}
src_wh = tp2wh[tp]
self.img_wh = (np.array(src_wh) * self.downSample).astype(int)
self.blender2opencv = np.array([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])
self.K = np.array([[892.62, 0.0, 639.5],
[0.0, 892.62, 359.5],
[0.0, 0.0, 1.0]])
self.K[:2] = self.K[:2] * self.downSample
if self.tp == 'vgn_syn':
self.K[:2] /= 2
self.depth_thres = {
'vgn_syn': 0.8,
}
if args.object_set == "graspnet":
dir_name = "pile_graspnet_test"
else:
if self.args.scene == "pile":
dir_name = "pile_pile_test_200"
elif self.args.scene == "packed":
dir_name = "packed_packed_test_200"
elif self.args.scene == "single":
dir_name = "single_single_test_200"
scene_root_dir = os.path.join(self.renderer_root_dir, "data/mesh_pose_list", dir_name)
self.mesh_pose_list = [i for i in sorted(os.listdir(scene_root_dir))]
self.depth_root_dir = ""
self.depth_list = []
def __init__(self, args=None, cfg_fn=None, debug_dir=None) -> None:
default_render_cfg = {
'min_wn': 3, # working view number
'ref_pad_interval': 16, # input image size should be multiple of 16
'use_src_imgs': False, # use source images to construct cost volume or not
'cost_volume_nn_num': 3, # number of source views used in cost volume
'use_depth': True, # use colmap depth in rendering or not,
}
# load render cfg
if cfg_fn is None:
self.set_params(args)
cfg = load_cfg(args.cfg_fn)
else:
cfg = load_cfg(cfg_fn)
print(f"[I] GraspNeRFPlanner: using ckpt: {cfg['name']}")
render_cfg = cfg['train_dataset_cfg'] if 'train_dataset_cfg' in cfg else {}
render_cfg = {**default_render_cfg, **render_cfg}
cfg['render_rgb'] = False # only for training. Disable in grasping.
# load model
self.net = name2network[cfg['network']](cfg)
ckpt_filename = 'model_best'
ckpt = torch.load(Path('src/nr/ckpt') / cfg["group_name"] / cfg["name"] / f'{ckpt_filename}.pth')
self.net.load_state_dict(ckpt['network_state_dict'])
self.net.cuda()
self.net.eval()
self.step = ckpt["step"]
self.output_dir = debug_dir
if debug_dir is not None:
if not Path(debug_dir).exists():
Path(debug_dir).mkdir(parents=True)
self.loss = VGNLoss({})
self.num_input_views = render_cfg['num_input_views']
print(f"[I] GraspNeRFPlanner: load model at step {self.step} of best metric {ckpt['best_para']}")
def get_image(self, img_id, round_idx):
img_filename = os.path.join(self.args.log_root_dir, "rendered_results/" + str(self.args.logdir).split("/")[-1], "rgb/%04d.png"%img_id)
img = imread(img_filename)[:,:,:3]
img = cv2.resize(img, self.img_wh)
return np.asarray(img, dtype=np.float32)
def get_pose(self, img_id):
poses_ori = np.load(Path(self.renderer_root_dir) / 'camera_pose.npy')
poses = [np.linalg.inv(p @ self.blender2opencv)[:3,:] for p in poses_ori]
return poses[img_id].astype(np.float32).copy()
def get_K(self, img_id):
return self.K.astype(np.float32).copy()
def get_depth_range(self,img_id, round_idx, fixed=False):
if fixed:
return np.array([0.2,0.8])
depth = self.get_depth(img_id, round_idx)
nf = [max(0, np.min(depth)), min(self.depth_thres[self.tp], np.max(depth))]
return np.array(nf)
def __call__(self, test_view_id, round_idx, n_grasp, gt_tsdf):
# load data for test
images = [self.get_image(i, round_idx) for i in test_view_id]
images = color_map_forward(np.stack(images, 0)).transpose([0, 3, 1, 2])
extrinsics = np.stack([self.get_pose(i) for i in test_view_id], 0)
intrinsics = np.stack([self.get_K(i) for i in test_view_id], 0)
depth_range = np.asarray([self.get_depth_range(i, round_idx, fixed = True) for i in test_view_id], dtype=np.float32)
tsdf_vol, qual_vol_ori, rot_vol_ori, width_vol_ori, toc = self.core(images, extrinsics, intrinsics, depth_range, self.bbox3d)
qual_vol, rot_vol, width_vol = process(tsdf_vol, qual_vol_ori, rot_vol_ori, width_vol_ori, tsdf_thres_high=self.tsdf_thres_high, tsdf_thres_low=self.tsdf_thres_low, n_grasp=n_grasp)
grasps, scores, indexs = select(qual_vol.copy(), rot_vol, width_vol)
grasps, scores, indexs = np.asarray(grasps), np.asarray(scores), np.asarray(indexs)
if len(grasps) > 0:
np.random.seed(self.args.seed + round_idx + n_grasp)
p = np.random.permutation(len(grasps))
grasps = [from_voxel_coordinates(g, self.voxel_size) for g in grasps[p]]
scores = scores[p]
indexs = indexs[p]
return grasps, scores, toc
def core(self,
images: np.ndarray,
extrinsics: np.ndarray,
intrinsics: np.ndarray,
depth_range=[0.2, 0.8],
bbox3d=[[-0.15, -0.15, -0.05],[0.15, 0.15, 0.25]], gt_info=None, que_id=0):
"""
@args
images: np array of shape (3, 3, h, w), image in RGB format
extrinsics: np array of shape (3, 4, 4), the transformation matrix from world to camera
intrinsics: np array of shape (3, 3, 3)
@rets
volume, label, rot, width: np array of shape (1, 1, res, res, res)
"""
_, _, h, w = images.shape
assert h % 32 == 0 and w % 32 == 0
extrinsics = extrinsics[:, :3, :]
que_imgs_info = build_render_imgs_info(extrinsics[que_id], intrinsics[que_id], (h, w), depth_range[que_id])
src_imgs_info = {'imgs': images, 'poses': extrinsics.astype(np.float32), 'Ks': intrinsics.astype(np.float32), 'depth_range': depth_range.astype(np.float32),
'bbox3d': np.array(bbox3d)}
ref_imgs_info = src_imgs_info.copy()
num_views = images.shape[0]
ref_imgs_info['nn_ids'] = np.arange(num_views).repeat(num_views, 0)
data = {'step': self.step , 'eval': True}
if not gt_info:
data['full_vol'] = True
else:
data['grasp_info'] = to_cuda(grasp_info_to_torch(gt_info))
data['que_imgs_info'] = to_cuda(imgs_info_to_torch(que_imgs_info))
data['src_imgs_info'] = to_cuda(imgs_info_to_torch(src_imgs_info))
data['ref_imgs_info'] = to_cuda(imgs_info_to_torch(ref_imgs_info))
with torch.no_grad():
t0 = time.time()
render_info = self.net(data)
t = time.time() - t0
if gt_info:
return self.loss(render_info, data, self.step, False)
label, rot, width = render_info['vgn_pred']
return render_info['volume'].cpu().numpy(), label.cpu().numpy(), rot.cpu().numpy(), width.cpu().numpy(), t
================================================
FILE: src/nr/network/aggregate_net.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from easydict import EasyDict
import numpy as np
from network.ibrnet import IBRNetWithNeuRay, IBRNetWithNeuRayNeus
from network.neus import SingleVarianceNetwork
def get_dir_diff(prj_dir,que_dir):
rfn, qn, rn, dn, _ = prj_dir.shape
dir_diff = prj_dir - que_dir.unsqueeze(0) # rfn,qn,rn,dn,3
dir_dot = torch.sum(prj_dir * que_dir.unsqueeze(0), -1, keepdim=True)
dir_diff = torch.cat([dir_diff, dir_dot], -1) # rfn,qn,rn,dn,4
dir_diff = dir_diff.reshape(rfn, qn * rn, dn, -1).permute(1, 2, 0, 3)
return dir_diff
class BaseAggregationNet(nn.Module):
default_cfg={
'sample_num': 64,
'neuray_dim': 32,
'use_img_feats': False,
}
def __init__(self, cfg):
super().__init__()
self.cfg={**self.default_cfg, **cfg}
dim = self.cfg['neuray_dim']
self.prob_embed = nn.Sequential(
nn.Linear(2+32, dim),
nn.ReLU(),
nn.Linear(dim, dim),
)
def _get_embedding(self, prj_dict, que_dir):
"""
:param prj_dict
prj_ray_feats: rfn,qn,rn,dn,f
prj_hit_prob: rfn,qn,rn,dn,1
prj_vis: rfn,qn,rn,dn,1
prj_alpha: rfn,qn,rn,dn,1
prj_rgb: rfn,qn,rn,dn,3
prj_dir: rfn,qn,rn,dn,3
:param que_dir: qn,rn,dn,3
:return: qn,rn,dn
"""
hit_prob_val = (prj_dict['hit_prob']-0.5)*2
vis_val = (prj_dict['vis']-0.5)*2
prj_hit_prob, prj_vis, prj_rgb, prj_dir, prj_ray_feats = \
hit_prob_val, vis_val, prj_dict['rgb'], prj_dict['dir'], prj_dict['ray_feats']
rfn,qn,rn,dn,_ = hit_prob_val.shape
prob_embedding = self.prob_embed(torch.cat([prj_ray_feats, prj_hit_prob, prj_vis],-1))
if que_dir is not None:
dir_diff = get_dir_diff(prj_dir, que_dir)
else:
_,qn,rn,dn,_ = prj_hit_prob.shape
dir_diff = torch.zeros((rfn, qn * rn, dn, 4)).permute(1, 2, 0, 3).to(prj_hit_prob.device)
valid_mask = prj_dict['mask']
valid_mask = valid_mask.float() # rfn,qn,rn,dn
valid_mask = valid_mask.reshape(rfn, qn * rn, dn, -1).permute(1, 2, 0, 3)
prj_img_feats = prj_dict['img_feats']
prj_img_feats = torch.cat([prj_rgb, prj_img_feats], -1)
prj_img_feats = prj_img_feats.reshape(rfn, qn * rn, dn, -1).permute(1, 2, 0, 3)
prob_embedding = prob_embedding.reshape(rfn, qn * rn, dn, -1).permute(1, 2, 0, 3)
return prj_img_feats, prob_embedding, dir_diff, valid_mask
class DefaultAggregationNet(BaseAggregationNet):
def __init__(self,cfg):
super().__init__(cfg)
dim = self.cfg['neuray_dim']
self.agg_impl = IBRNetWithNeuRay(dim,n_samples=self.cfg['sample_num'])
def forward(self, prj_dict, que_dir, que_pts=None, que_dists=None):
qn,rn,dn,_ = que_dir.shape
prj_img_feats, prob_embedding, dir_diff, valid_mask = self._get_embedding(prj_dict, que_dir)
outs = self.agg_impl(prj_img_feats, prob_embedding, dir_diff, valid_mask)
colors = outs[...,:3] # qn*rn,dn,3
density = outs[...,3] # qn*rn,dn,0
return density.reshape(qn,rn,dn), colors.reshape(qn,rn,dn,3)
class NeusAggregationNet(BaseAggregationNet):
neus_default_cfg = {
'cos_anneal_end_iter': 0,
'init_s': 0.3,
'fix_s': False
}
def __init__(self,cfg):
cfg = {**self.neus_default_cfg, **cfg}
super().__init__(cfg)
dim = self.cfg['neuray_dim']
self.agg_impl = IBRNetWithNeuRayNeus(dim,n_samples=self.cfg['sample_num'])
self.deviation_network = SingleVarianceNetwork(self.cfg['init_s'], self.cfg['fix_s'])
self.step = 0
self.cos_anneal_ratio = 1.0
def _update_cos_anneal_ratio(self):
self.cos_anneal_ratio = np.min([1.0, self.step / self.cfg['cos_anneal_end_iter']])
def _get_alpha_from_sdf(self, sdf, grad, que_dir, que_dists):
qn,rn,dn,_ = que_dir.shape
inv_s = self.deviation_network(torch.zeros([1, 3], device=sdf.device))[:, :1].clip(1e-6, 1e6) # Single parameter
inv_s = inv_s.expand(qn*rn, dn)
true_cos = (-que_dir * grad).sum(-1, keepdim=True)
# "cos_anneal_ratio" grows from 0 to 1 in the beginning training iterations. The anneal strategy below makes
# the cos value "not dead" at the beginning training iterations, for better convergence.
iter_cos = -(F.relu(-true_cos * 0.5 + 0.5) * (1.0 - self.cos_anneal_ratio) +
F.relu(-true_cos) * self.cos_anneal_ratio)[0].squeeze(-1) # always non-positive
# Estimate signed distances at section points
estimated_next_sdf = sdf + iter_cos * que_dists[0] * 0.5
estimated_prev_sdf = sdf - iter_cos * que_dists[0] * 0.5
prev_cdf = torch.sigmoid(estimated_prev_sdf * inv_s)
next_cdf = torch.sigmoid(estimated_next_sdf * inv_s)
p = prev_cdf - next_cdf
c = prev_cdf
alpha = ((p + 1e-5) / (c + 1e-5)).reshape(qn,rn,dn).clip(0.0, 1.0)
return alpha
def forward(self, prj_dict, que_dir, que_pts, que_dists, is_train):
if self.cfg['cos_anneal_end_iter'] and is_train:
self._update_cos_anneal_ratio()
qn,rn,dn,_ = que_dir.shape
prj_img_feats, prob_embedding, dir_diff, valid_mask = self._get_embedding(prj_dict, que_dir)
outs, grad = self.agg_impl(prj_img_feats, prob_embedding, dir_diff, valid_mask, que_pts)
colors = outs[...,:3] # qn*rn,dn,3
sdf = outs[...,3] # qn*rn,dn,0
if que_dists is None:
return None, sdf.reshape(qn,rn,dn), colors.reshape(qn,rn,dn,3), None, None
if is_train:
self.step += 1
self.deviation_network.set_step(self.step)
alpha = self._get_alpha_from_sdf(sdf, grad, que_dir, que_dists)
grad_error = torch.mean((torch.linalg.norm(grad.reshape(qn,rn,dn,3), ord=2, dim=-1) - 1.0) ** 2).reshape(1,1)
return alpha.reshape(qn,rn,dn), sdf.reshape(qn,rn,dn), colors.reshape(qn,rn,dn,3), grad_error, self.deviation_network.variance.reshape(1,1)
name2agg_net={
'default': DefaultAggregationNet,
'neus': NeusAggregationNet
}
================================================
FILE: src/nr/network/dist_decoder.py
================================================
import torch.nn as nn
import torch
from network.ops import AddBias
def get_near_far_points(depth, interval, depth_range, is_ref, fixed_interval=False, fixed_interval_val=0.01):
""" is_ref | not is_ref
:param depth: [...,dn] rfn,qn,rn,dn or qn,rn,dn
:param interval: [...,dn] 1,qn,rn,dn or qn,rn,dn
:param depth_range: rfn,2 or qn,2
:param is_ref:
:param fixed_interval:
:param fixed_interval_val:
:return: near far [rfn,qn,rn,dn] or [qn,rn,dn]
"""
if is_ref:
ref_near = depth_range[:, 0]
ref_far = depth_range[:, 1]
ref_near = -1 / ref_near[:, None, None, None]
ref_far = -1 / ref_far[:, None, None, None]
depth = torch.clamp(depth, min=1e-5)
depth = -1 / depth
depth = (depth - ref_near) / (ref_far - ref_near)
else:
que_near = depth_range[:, 0] # qn
que_far = depth_range[:, 1] # qn
que_near = -1 / que_near[:, None, None]
que_far = -1 / que_far[:, None, None]
depth = torch.clamp(depth, min=1e-5)
depth = -1 / depth
depth = (depth - que_near) / (que_far - que_near)
if not fixed_interval:
if is_ref:
interval_half = interval / 2
interval_ext = torch.cat([interval_half[..., 0:1], interval_half], -1)
near = depth - interval_ext[..., :-1]
far = depth + interval_ext[..., 1:]
else:
interval_half = interval / 2
first = depth[..., 0] - interval_half[..., 0]
last = depth[..., -1] + interval_half[..., -1]
depth_ext = (depth[..., :-1] + depth[..., 1:]) / 2
depth_ext = torch.cat([first[..., None], depth_ext, last[..., None]], -1)
near = depth_ext[..., :-1]
far = depth_ext[..., 1:]
else:
near = depth - fixed_interval_val/2
far = depth + fixed_interval_val/2
return near, far
class MixtureLogisticsDistDecoder(nn.Module):
default_cfg={
'feats_dim': 32,
'bias_val': 0.05,
"use_vis": True,
}
def __init__(self,cfg):
super().__init__()
self.cfg={**self.default_cfg,**cfg}
ray_feats_dim = self.cfg["feats_dim"]
run_dim = ray_feats_dim
self.mean_decoder=nn.Sequential(
nn.Linear(ray_feats_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, 2),
nn.Softplus()
)
self.var_decoder=nn.Sequential(
nn.Linear(ray_feats_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, 2),
nn.Softplus(),
AddBias(self.cfg['bias_val']),
)
self.aw_decoder=nn.Sequential(
nn.Linear(ray_feats_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, 1),
nn.Sigmoid(),
)
if self.cfg['use_vis']:
self.vis_decoder=nn.Sequential(
nn.Linear(ray_feats_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, run_dim),
nn.ELU(),
nn.Linear(run_dim, 1),
nn.Sigmoid(),
)
def forward(self, feats):
prj_mean = self.mean_decoder(feats)
prj_var = self.var_decoder(feats)
prj_aw = self.aw_decoder(feats)
if self.cfg['use_vis']:
prj_vis = self.vis_decoder(feats)
else:
prj_vis = None
return prj_mean, prj_var, prj_vis, prj_aw
def compute_prob(self, depth, interval, mean, var, vis, aw, is_ref, depth_range):
"""
:param depth: [...,dn] rfn,qn,rn,dn or qn,rn,dn
:param interval: [...,dn] 1,qn,rn,dn or qn,rn,dn
:param mean: [...,1 or dn] rfn,qn,rn,dn,2 or qn,rn,1,2
:param var: [...,1 or dn] rfn,qn,rn,dn,2 or qn,rn,1,2
:param vis: [...,1 or dn] rfn,qn,rn,dn,1 or qn,rn,1,1
:param aw: [...,1 or dn] rfn,qn,rn,dn,1 or qn,rn,1,1
:param is_ref:
:param depth_range: rfn,2 or qn,2
:return:
"""
if interval.shape != (1,0):
near, far = get_near_far_points(depth, interval, depth_range, is_ref)
else:
near, far = get_near_far_points(depth, interval, depth_range, is_ref, fixed_interval=True, fixed_interval_val=0.01)
# near and far [rfn,qn,rn,dn] or [qn,rn,dn]
mix = torch.cat([aw, 1 - aw],-1) # [...,2]
near, far = near[...,None], far[...,None]
d0 = (near - mean) * var # [...,2]
d1 = (far - mean) * var # [...,2]
cdf0 = (0.5 + 0.5 * torch.tanh(d0)) # t(z_i)
cdf1 = (0.5 + 0.5 * torch.tanh(d1)) # t(z_{i+1})
if self.cfg['use_vis']:
cdf0, cdf1 = cdf0 * vis, cdf1 * vis
visibility = 1 - cdf0
hit_prob = cdf1 - cdf0
visibility = torch.sum(visibility*mix, -1)
hit_prob = torch.sum(hit_prob*mix, -1)
eps = 1e-5
alpha_value = torch.log(hit_prob / (visibility - hit_prob + eps) + eps)
return alpha_value, visibility, hit_prob
def decode_alpha_value(self, alpha_value):
alpha_value = torch.sigmoid(alpha_value)
return alpha_value
def predict_mean(self,prj_ray_feats):
prj_mean = self.mean_decoder(prj_ray_feats)
return prj_mean
def predict_aw(self,prj_ray_feats):
return self.aw_decoder(prj_ray_feats)
name2dist_decoder={
'mixture_logistics': MixtureLogisticsDistDecoder
}
================================================
FILE: src/nr/network/ibrnet.py
================================================
import numpy as np
import torch.nn.functional as F
import torch.nn as nn
import torch
from network.neus import *
class ScaledDotProductAttention(nn.Module):
''' Scaled Dot-Product Attention '''
def __init__(self, temperature, attn_dropout=0.1):
super().__init__()
self.temperature = temperature
# self.dropout = nn.Dropout(attn_dropout)
def forward(self, q, k, v, mask=None):
attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
if mask is not None:
attn = attn.masked_fill(mask == 0, -1e9)
# attn = attn * mask
attn = F.softmax(attn, dim=-1)
# attn = self.dropout(F.softmax(attn, dim=-1))
output = torch.matmul(attn, v)
return output, attn
class PositionwiseFeedForward(nn.Module):
''' A two-feed-forward-layer module '''
def __init__(self, d_in, d_hid, dropout=0.1):
super().__init__()
self.w_1 = nn.Linear(d_in, d_hid) # position-wise
self.w_2 = nn.Linear(d_hid, d_in) # position-wise
self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
# self.dropout = nn.Dropout(dropout)
def forward(self, x):
residual = x
x = self.w_2(F.relu(self.w_1(x)))
# x = self.dropout(x)
x += residual
x = self.layer_norm(x)
return x
class MultiHeadAttention(nn.Module):
''' Multi-Head Attention module '''
def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
super().__init__()
self.n_head = n_head
self.d_k = d_k
self.d_v = d_v
self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
self.fc = nn.Linear(n_head * d_v, d_model, bias=False)
self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)
# self.dropout = nn.Dropout(dropout)
self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
def forward(self, q, k, v, mask=None):
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
residual = q
# Pass through the pre-attention projection: b x lq x (n*dv)
# Separate different heads: b x lq x n x dv
q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
# Transpose for attention dot product: b x n x lq x dv
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
if mask is not None:
mask = mask.unsqueeze(1) # For head axis broadcasting.
q, attn = self.attention(q, k, v, mask=mask)
# Transpose to move the head dimension back: b x lq x n x dv
# Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
# q = self.dropout(self.fc(q))
q = self.fc(q)
q += residual
q = self.layer_norm(q)
return q, attn
# default tensorflow initialization of linear layers
def weights_init(m):
if isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight.data)
if m.bias is not None:
nn.init.zeros_(m.bias.data)
@torch.jit.script
def fused_mean_variance(x, weight):
mean = torch.sum(x*weight, dim=2, keepdim=True)
var = torch.sum(weight * (x - mean)**2, dim=2, keepdim=True)
return mean, var
class IBRNet(nn.Module):
def __init__(self, in_feat_ch=32, n_samples=64, **kwargs):
super(IBRNet, self).__init__()
# self.args = args
self.anti_alias_pooling = False
if self.anti_alias_pooling:
self.s = nn.Parameter(torch.tensor(0.2), requires_grad=True)
activation_func = nn.ELU(inplace=True)
self.n_samples = n_samples
self.ray_dir_fc = nn.Sequential(nn.Linear(4, 16),
activation_func,
nn.Linear(16, in_feat_ch + 3),
activation_func)
self.base_fc = nn.Sequential(nn.Linear((in_feat_ch+3)*3, 64),
activation_func,
nn.Linear(64, 32),
activation_func)
self.vis_fc = nn.Sequential(nn.Linear(32, 32),
activation_func,
nn.Linear(32, 33),
activation_func,
)
self.vis_fc2 = nn.Sequential(nn.Linear(32, 32),
activation_func,
nn.Linear(32, 1),
nn.Sigmoid()
)
self.geometry_fc = nn.Sequential(nn.Linear(32*2+1, 64),
activation_func,
nn.Linear(64, 16),
activation_func)
self.ray_attention = MultiHeadAttention(4, 16, 4, 4)
self.out_geometry_fc = nn.Sequential(nn.Linear(16, 16),
activation_func,
nn.Linear(16, 1),
nn.ReLU())
self.rgb_fc = nn.Sequential(nn.Linear(32+1+4, 16),
activation_func,
nn.Linear(16, 8),
activation_func,
nn.Linear(8, 1))
self.pos_encoding = self.posenc(d_hid=16, n_samples=self.n_samples)
self.base_fc.apply(weights_init)
self.vis_fc2.apply(weights_init)
self.vis_fc.apply(weights_init)
self.geometry_fc.apply(weights_init)
self.rgb_fc.apply(weights_init)
def posenc(self, d_hid, n_samples):
def get_position_angle_vec(position):
return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_samples)])
sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i
sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1
sinusoid_table = torch.from_numpy(sinusoid_table).to("cuda:{}".format(self.args.local_rank)).float().unsqueeze(0)
return sinusoid_table
def forward(self, rgb_feat, ray_diff, mask):
'''
:param rgb_feat: rgbs and image features [n_rays, n_samples, n_views, n_feat]
:param ray_diff: ray direction difference [n_rays, n_samples, n_views, 4], first 3 channels are directions,
last channel is inner product
:param mask: mask for whether each projection is valid or not. [n_rays, n_samples, n_views, 1]
:return: rgb and density output, [n_rays, n_samples, 4]
'''
num_views = rgb_feat.shape[2]
direction_feat = self.ray_dir_fc(ray_diff)
rgb_in = rgb_feat[..., :3]
rgb_feat = rgb_feat + direction_feat
if self.anti_alias_pooling:
_, dot_prod = torch.split(ray_diff, [3, 1], dim=-1)
exp_dot_prod = torch.exp(torch.abs(self.s) * (dot_prod - 1))
weight = (exp_dot_prod - torch.min(exp_dot_prod, dim=2, keepdim=True)[0]) * mask
weight = weight / (torch.sum(weight, dim=2, keepdim=True) + 1e-8) # means it will trust the one more with more consistent view point
else:
weight = mask / (torch.sum(mask, dim=2, keepdim=True) + 1e-8)
# compute mean and variance across different views for each point
mean, var = fused_mean_variance(rgb_feat, weight) # [n_rays, n_samples, 1, n_feat]
globalfeat = torch.cat([mean, var], dim=-1) # [n_rays, n_samples, 1, 2*n_feat]
x = torch.cat([globalfeat.expand(-1, -1, num_views, -1), rgb_feat], dim=-1) # [n_rays, n_samples, n_views, 3*n_feat]
x = self.base_fc(x)
x_vis = self.vis_fc(x * weight)
x_res, vis = torch.split(x_vis, [x_vis.shape[-1]-1, 1], dim=-1)
vis = F.sigmoid(vis) * mask
x = x + x_res
vis = self.vis_fc2(x * vis) * mask
weight = vis / (torch.sum(vis, dim=2, keepdim=True) + 1e-8)
mean, var = fused_mean_variance(x, weight)
globalfeat = torch.cat([mean.squeeze(2), var.squeeze(2), weight.mean(dim=2)], dim=-1) # [n_rays, n_samples, 32*2+1]
globalfeat = self.geometry_fc(globalfeat) # [n_rays, n_samples, 16]
num_valid_obs = torch.sum(mask, dim=2)
globalfeat = globalfeat + self.pos_encoding
globalfeat, _ = self.ray_attention(globalfeat, globalfeat, globalfeat,
mask=(num_valid_obs > 1).float()) # [n_rays, n_samples, 16]
sigma = self.out_geometry_fc(globalfeat) # [n_rays, n_samples, 1]
sigma_out = sigma.masked_fill(num_valid_obs < 1, 0.) # set the sigma of invalid point to zero
# rgb computation
x = torch.cat([x, vis, ray_diff], dim=-1)
x = self.rgb_fc(x)
x = x.masked_fill(mask == 0, -1e9)
blending_weights_valid = F.softmax(x, dim=2) # color blending
rgb_out = torch.sum(rgb_in*blending_weights_valid, dim=2)
out = torch.cat([rgb_out, sigma_out], dim=-1)
return out
class IBRNetWithNeuRay(nn.Module):
def __init__(self, neuray_in_dim=32, in_feat_ch=32, n_samples=64, **kwargs):
super().__init__()
# self.args = args
self.anti_alias_pooling = False
if self.anti_alias_pooling:
self.s = nn.Parameter(torch.tensor(0.2), requires_grad=True)
activation_func = nn.ELU(inplace=True)
self.n_samples = n_samples
self.ray_dir_fc = nn.Sequential(nn.Linear(4, 16),
activation_func,
nn.Linear(16, in_feat_ch + 3),
activation_func)
self.base_fc = nn.Sequential(nn.Linear((in_feat_ch+3)*5+neuray_in_dim, 64),
activation_func,
nn.Linear(64, 32),
activation_func)
self.vis_fc = nn.Sequential(nn.Linear(32, 32),
activation_func,
nn.Linear(32, 33),
activation_func,
)
self.vis_fc2 = nn.Sequential(nn.Linear(32, 32),
activation_func,
nn.Linear(32, 1),
nn.Sigmoid()
)
self.geometry_fc = nn.Sequential(nn.Linear(32*2+1, 64),
activation_func,
nn.Linear(64, 16),
activation_func)
self.ray_attention = MultiHeadAttention(4, 16, 4, 4)
self.out_geometry_fc = nn.Sequential(nn.Linear(16, 16),
activation_func,
nn.Linear(16, 1),
nn.ReLU())
self.rgb_fc = nn.Sequential(nn.Linear(32+1+4, 16),
activation_func,
nn.Linear(16, 8),
activation_func,
nn.Linear(8, 1))
self.neuray_fc = nn.Sequential(
nn.Linear(neuray_in_dim, 8,),
activation_func,
nn.Linear(8, 1),
)
self.pos_encoding = self.posenc(d_hid=16, n_samples=self.n_samples)
self.base_fc.apply(weights_init)
self.vis_fc2.apply(weights_init)
self.vis_fc.apply(weights_init)
self.geometry_fc.apply(weights_init)
self.rgb_fc.apply(weights_init)
self.neuray_fc.apply(weights_init)
def change_pos_encoding(self,n_samples):
self.pos_encoding = self.posenc(16, n_samples=n_samples)
def posenc(self, d_hid, n_samples):
def get_position_angle_vec(position):
return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_samples)])
sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i
sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1
sinusoid_table = torch.from_numpy(sinusoid_table).to("cuda:{}".format(0)).float().unsqueeze(0)
return sinusoid_table
def forward(self, rgb_feat, neuray_feat, ray_diff, mask):
'''
:param rgb_feat: rgbs and image features [n_rays, n_samples, n_views, n_feat]
:param ray_diff: ray direction difference [n_rays, n_samples, n_views, 4], first 3 channels are directions,
last channel is inner product
:param mask: mask for whether each projection is valid or not. [n_rays, n_samples, n_views, 1]
:return: rgb and density output, [n_rays, n_samples, 4]
'''
num_views = rgb_feat.shape[2]
direction_feat = self.ray_dir_fc(ray_diff)
rgb_in = rgb_feat[..., :3]
rgb_feat = rgb_feat + direction_feat
if self.anti_alias_pooling:
_, dot_prod = torch.split(ray_diff, [3, 1], dim=-1)
exp_dot_prod = torch.exp(torch.abs(self.s) * (dot_prod - 1))
weight = (exp_dot_prod - torch.min(exp_dot_prod, dim=2, keepdim=True)[0]) * mask
weight = weight / (torch.sum(weight, dim=2, keepdim=True) + 1e-8) # means it will trust the one more with more consistent view point
else:
weight = mask / (torch.sum(mask, dim=2, keepdim=True) + 1e-8)
# neuray layer 0
weight0 = torch.sigmoid(self.neuray_fc(neuray_feat)) * weight # [rn,dn,rfn,f]
mean0, var0 = fused_mean_variance(rgb_feat, weight0) # [n_rays, n_samples, 1, n_feat]
mean1, var1 = fused_mean_variance(rgb_feat, weight) # [n_rays, n_samples, 1, n_feat]
globalfeat = torch.cat([mean0, var0, mean1, var1], dim=-1) # [n_rays, n_samples, 1, 2*n_feat]
x = torch.cat([globalfeat.expand(-1, -1, num_views, -1), rgb_feat, neuray_feat], dim=-1) # [n_rays, n_samples, n_views, 3*n_feat]
x = self.base_fc(x)
x_vis = self.vis_fc(x * weight)
x_res, vis = torch.split(x_vis, [x_vis.shape[-1]-1, 1], dim=-1)
vis = F.sigmoid(vis) * mask
x = x + x_res
vis = self.vis_fc2(x * vis) * mask
weight = vis / (torch.sum(vis, dim=2, keepdim=True) + 1e-8)
mean, var = fused_mean_variance(x, weight)
globalfeat = torch.cat([mean.squeeze(2), var.squeeze(2), weight.mean(dim=2)], dim=-1) # [n_rays, n_samples, 32*2+1]
globalfeat = self.geometry_fc(globalfeat) # [n_rays, n_samples, 16]
num_valid_obs = torch.sum(mask, dim=2)
globalfeat = globalfeat + self.pos_encoding
globalfeat, _ = self.ray_attention(globalfeat, globalfeat, globalfeat,
mask=(num_valid_obs > 1).float()) # [n_rays, n_samples, 16]
sigma = self.out_geometry_fc(globalfeat) # [n_rays, n_samples, 1]
sigma_out = sigma.masked_fill(num_valid_obs < 1, 0.) # set the sigma of invalid point to zero
# rgb computation
x = torch.cat([x, vis, ray_diff], dim=-1)
x = self.rgb_fc(x)
x = x.masked_fill(mask == 0, -1e9)
blending_weights_valid = F.softmax(x, dim=2) # color blending
rgb_out = torch.sum(rgb_in*blending_weights_valid, dim=2)
out = torch.cat([rgb_out, sigma_out], dim=-1)
return out
class IBRNetWithNeuRayNeus(nn.Module):
def __init__(self, neuray_in_dim=32, in_feat_ch=32, n_samples=64, **kwargs):
super().__init__()
# self.args = args
self.anti_alias_pooling = False
if self.anti_alias_pooling:
self.s = nn.Parameter(torch.tensor(0.2), requires_grad=True)
activation_func = nn.ELU(inplace=True)
self.n_samples = n_samples
self.ray_dir_fc = nn.Sequential(nn.Linear(4, 16),
activation_func,
nn.Linear(16, in_feat_ch + 3),
activation_func)
self.base_fc = nn.Sequential(nn.Linear((in_feat_ch+3)*5+neuray_in_dim, 64),
activation_func,
nn.Linear(64, 32),
activation_func)
self.vis_fc = nn.Sequential(nn.Linear(32, 32),
activation_func,
nn.Linear(32, 33),
activation_func,
)
self.vis_fc2 = nn.Sequential(nn.Linear(32, 32),
activation_func,
nn.Linear(32, 1),
nn.Sigmoid()
)
self.embed_fn, input_ch = get_embedder(3, input_dims=3)
self.geometry_fc = nn.Sequential(nn.Linear(32*2+1+input_ch, 64),
activation_func,
nn.Linear(64, 16),
activation_func)
self.ray_attention = MultiHeadAttention(4, 16, 4, 4)
self.out_geometry_fc = nn.Sequential(nn.Linear(16, 16),
nn.Linear(16, 1),
)
self.rgb_fc = nn.Sequential(nn.Linear(32+1+4, 16),
activation_func,
nn.Linear(16, 8),
activation_func,
nn.Linear(8, 1))
self.neuray_fc = nn.Sequential(
nn.Linear(neuray_in_dim, 8,),
activation_func,
nn.Linear(8, 1),
)
self.pos_encoding = self.posenc(d_hid=16, n_samples=self.n_samples)
self.base_fc.apply(weights_init)
self.vis_fc2.apply(weights_init)
self.vis_fc.apply(weights_init)
self.geometry_fc.apply(weights_init)
self.rgb_fc.apply(weights_init)
self.neuray_fc.apply(weights_init)
def change_pos_encoding(self,n_samples):
self.pos_encoding = self.posenc(16, n_samples=n_samples)
def posenc(self, d_hid, n_samples):
def get_position_angle_vec(position):
return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_samples)])
sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i
sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1
sinusoid_table = torch.from_numpy(sinusoid_table).to("cuda:{}".format(0)).float().unsqueeze(0)
return sinusoid_table
def forward(self, rgb_feat, neuray_feat, ray_diff, mask, que_pts):
'''
:param rgb_feat: rgbs and image features [n_rays, n_samples, n_views, n_feat]
:param ray_diff: ray direction difference [n_rays, n_samples, n_views, 4], first 3 channels are directions,
last channel is inner product
:param mask: mask for whether each projection is valid or not. [n_rays, n_samples, n_views, 1]
:return: rgb and density output, [n_rays, n_samples, 4]
'''
num_views = rgb_feat.shape[2]
direction_feat = self.ray_dir_fc(ray_diff)
rgb_in = rgb_feat[..., :3]
rgb_feat = rgb_feat + direction_feat
if self.anti_alias_pooling:
_, dot_prod = torch.split(ray_diff, [3, 1], dim=-1)
exp_dot_prod = torch.exp(torch.abs(self.s) * (dot_prod - 1))
weight = (exp_dot_prod - torch.min(exp_dot_prod, dim=2, keepdim=True)[0]) * mask
weight = weight / (torch.sum(weight, dim=2, keepdim=True) + 1e-8) # means it will trust the one more with more consistent view point
else:
weight = mask / (torch.sum(mask, dim=2, keepdim=True) + 1e-8)
# neuray layer 0
weight0 = torch.sigmoid(self.neuray_fc(neuray_feat)) * weight # [rn,dn,rfn,f]
mean0, var0 = fused_mean_variance(rgb_feat, weight0) # [n_rays, n_samples, 1, n_feat]
mean1, var1 = fused_mean_variance(rgb_feat, weight) # [n_rays, n_samples, 1, n_feat]
globalfeat = torch.cat([mean0, var0, mean1, var1], dim=-1) # [n_rays, n_samples, 1, 2*n_feat]
x = torch.cat([globalfeat.expand(-1, -1, num_views, -1), rgb_feat, neuray_feat], dim=-1) # [n_rays, n_samples, n_views, 3*n_feat]
x = self.base_fc(x)
x_vis = self.vis_fc(x * weight)
x_res, vis = torch.split(x_vis, [x_vis.shape[-1]-1, 1], dim=-1)
vis = F.sigmoid(vis) * mask
x = x + x_res
vis = self.vis_fc2(x * vis) * mask
weight = vis / (torch.sum(vis, dim=2, keepdim=True) + 1e-8)
mean, var = fused_mean_variance(x, weight)
with torch.set_grad_enabled(True):
que_pts.requires_grad_(True)
embed_pts = self.embed_fn(que_pts)[0]
globalfeat = torch.cat([mean.squeeze(2), var.squeeze(2), weight.mean(dim=2), embed_pts], dim=-1) # [n_rays, n_samples, 32*2+1]
globalfeat = self.geometry_fc(globalfeat) # [n_rays, n_samples, 16]
num_valid_obs = torch.sum(mask, dim=2)
globalfeat = globalfeat + self.pos_encoding
globalfeat, _ = self.ray_attention(globalfeat, globalfeat, globalfeat,
mask=(num_valid_obs > 1).float()) # [n_rays, n_samples, 16]
sdf = self.out_geometry_fc(globalfeat).clip(-1.0,1.0) # [n_rays, n_samples, 1]
sdf_out = sdf.masked_fill(num_valid_obs < 1, 1.) # set the sigma of invalid point to zero
d_output = torch.ones_like(sdf_out, requires_grad=False, device=sdf_out.device)
gradients = torch.autograd.grad(
outputs=sdf_out,
inputs=que_pts,
grad_outputs=d_output,
create_graph=True,
retain_graph=True,
only_inputs=True)[0]
# rgb computation
x = torch.cat([x, vis, ray_diff], dim=-1)
x = self.rgb_fc(x)
x = x.masked_fill(mask == 0, -1e9)
blending_weights_valid = F.softmax(x, dim=2) # color blending
rgb_out = torch.sum(rgb_in*blending_weights_valid, dim=2)
out = torch.cat([rgb_out, sdf_out], dim=-1)
return out, gradients
================================================
FILE: src/nr/network/init_net.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from network.ops import interpolate_feats, masked_mean_var, ResEncoder, ResUNetLight, conv3x3, ResidualBlock, conv1x1
class CostVolumeInitNet(nn.Module):
default_cfg={
'cost_volume_sn': 64,
}
def __init__(self,cfg):
super().__init__()
self.cfg={**self.default_cfg,**cfg}
imagenet_mean = torch.from_numpy(np.asarray([0.485, 0.456, 0.406], np.float32)).cuda()[None, :, None, None]
imagenet_std = torch.from_numpy(np.asarray([0.229, 0.224, 0.225], np.float32)).cuda()[None, :, None, None]
self.register_buffer('imagenet_mean', imagenet_mean)
self.register_buffer('imagenet_std', imagenet_std)
self.res_net = ResUNetLight(out_dim=32)
norm_layer = lambda dim: nn.InstanceNorm2d(dim, track_running_stats=False, affine=True)
in_dim = 32
self.out_conv = nn.Sequential(
conv3x3(in_dim, 32),
ResidualBlock(32, 32, norm_layer=norm_layer),
conv1x1(32, 32),
)
def forward(self, ref_imgs_info, src_imgs_info, is_train):
ref_feats = self.res_net(ref_imgs_info['imgs'])
return self.out_conv(torch.cat([ref_feats], 1))
name2init_net={
'cost_volume': CostVolumeInitNet,
}
================================================
FILE: src/nr/network/loss.py
================================================
import torch
import torch.nn as nn
import numpy as np
import pyquaternion as pyq
import math
from network.ops import interpolate_feats
import torch.nn.functional as F
import torchmetrics
from utils.base_utils import calc_rot_error_from_qxyzw
class Loss:
def __init__(self, keys):
"""
keys are used in multi-gpu model, DummyLoss in train_tools.py
:param keys: the output keys of the dict
"""
self.keys=keys
def __call__(self, data_pr, data_gt, step, **kwargs):
pass
class ConsistencyLoss(Loss):
default_cfg={
'use_ray_mask': False,
'use_dr_loss': False,
'use_dr_fine_loss': False,
'use_nr_fine_loss': False,
}
def __init__(self, cfg):
self.cfg={**self.default_cfg,**cfg}
super().__init__([f'loss_prob','loss_prob_fine'])
def __call__(self, data_pr, data_gt, step, **kwargs):
if 'hit_prob_self' not in data_pr: return {}
prob0 = data_pr['hit_prob_nr'].detach() # qn,rn,dn
prob1 = data_pr['hit_prob_self'] # qn,rn,dn
if self.cfg['use_ray_mask']:
ray_mask = data_pr['ray_mask'].float() # 1,rn
else:
ray_mask = 1
ce = - prob0 * torch.log(prob1 + 1e-5) - (1 - prob0) * torch.log(1 - prob1 + 1e-5)
outputs={'loss_prob': torch.mean(torch.mean(ce,-1),1)}
if 'hit_prob_nr_fine' in data_pr:
prob0 = data_pr['hit_prob_nr_fine'].detach() # qn,rn,dn
prob1 = data_pr['hit_prob_self_fine'] # qn,rn,dn
ce = - prob0 * torch.log(prob1 + 1e-5) - (1 - prob0) * torch.log(1 - prob1 + 1e-5)
outputs['loss_prob_fine']=torch.mean(torch.mean(ce,-1),1)
return outputs
class RenderLoss(Loss):
default_cfg={
'use_ray_mask': True,
'use_dr_loss': False,
'use_dr_fine_loss': False,
'use_nr_fine_loss': False,
'disable_at_eval': True,
'render_loss_weight': 0.01
}
def __init__(self, cfg):
self.cfg={**self.default_cfg,**cfg}
super().__init__([f'loss_rgb'])
def __call__(self, data_pr, data_gt, step, is_train=True, **kwargs):
if not is_train and self.cfg['disable_at_eval']:
return {}
rgb_gt = data_pr['pixel_colors_gt'] # 1,rn,3
rgb_nr = data_pr['pixel_colors_nr'] # 1,rn,3
def compute_loss(rgb_pr,rgb_gt):
loss=torch.sum((rgb_pr-rgb_gt)**2,-1) # b,n
if self.cfg['use_ray_mask']:
ray_mask = data_pr['ray_mask'].float() # 1,rn
loss = torch.sum(loss*ray_mask,1)/(torch.sum(ray_mask,1)+1e-3)
else:
loss = torch.mean(loss, 1)
return loss * self.cfg['render_loss_weight']
results = {'loss_rgb_nr': compute_loss(rgb_nr, rgb_gt)}
if self.cfg['use_dr_loss']:
rgb_dr = data_pr['pixel_colors_dr'] # 1,rn,3
results['loss_rgb_dr'] = compute_loss(rgb_dr, rgb_gt)
if self.cfg['use_dr_fine_loss']:
results['loss_rgb_dr_fine'] = compute_loss(data_pr['pixel_colors_dr_fine'], rgb_gt)
if self.cfg['use_nr_fine_loss']:
results['loss_rgb_nr_fine'] = compute_loss(data_pr['pixel_colors_nr_fine'], rgb_gt)
return results
class DepthLoss(Loss):
default_cfg={
'depth_correct_thresh': 0.02,
'depth_loss_type': 'l2',
'depth_loss_l1_beta': 0.05,
'depth_loss_weight': 1,
'disable_at_eval': True,
}
def __init__(self, cfg):
super().__init__(['loss_depth'])
self.cfg={**self.default_cfg,**cfg}
if self.cfg['depth_loss_type']=='smooth_l1':
self.loss_op=nn.SmoothL1Loss(reduction='none',beta=self.cfg['depth_loss_l1_beta'])
def __call__(self, data_pr, data_gt, step, is_train=True, **kwargs):
if not is_train and self.cfg['disable_at_eval']:
return {}
if 'true_depth' not in data_gt['ref_imgs_info']:
print('no')
return {'loss_depth': torch.zeros([1], dtype=torch.float32, device=data_pr['pixel_colors_nr'].device)}
coords = data_pr['depth_coords'] # rfn,pn,2
depth_pr = data_pr['depth_mean'] # rfn,pn
depth_maps = data_gt['ref_imgs_info']['true_depth'] # rfn,1,h,w
rfn, _, h, w = depth_maps.shape
depth_gt = interpolate_feats(
depth_maps,coords,h,w,padding_mode='border',align_corners=True)[...,0] # rfn,pn
# transform to inverse depth coordinate
depth_range = data_gt['ref_imgs_info']['depth_range'] # rfn,2
near, far = -1/depth_range[:,0:1], -1/depth_range[:,1:2] # rfn,1
def process(depth):
depth = torch.clamp(depth, min=1e-5)
depth = -1 / depth
depth = (depth - near) / (far - near)
depth = torch.clamp(depth, min=0, max=1.0)
return depth
depth_gt = process(depth_gt)
# compute loss
def compute_loss(depth_pr):
if self.cfg['depth_loss_type']=='l2':
loss = (depth_gt - depth_pr)**2
elif self.cfg['depth_loss_type']=='smooth_l1':
loss = self.loss_op(depth_gt, depth_pr)
if data_gt['scene_name'].startswith('gso'):
depth_maps_noise = data_gt['ref_imgs_info']['depth'] # rfn,1,h,w
depth_aug = interpolate_feats(depth_maps_noise, coords, h, w, padding_mode='border', align_corners=True)[..., 0] # rfn,pn
depth_aug = process(depth_aug)
mask = (torch.abs(depth_aug-depth_gt) 0:
valid_mask = data_gt['ref_imgs_info']['sdf_gt'] != -1.0
outputs['loss_sdf'] = self.loss_fn(data_gt['ref_imgs_info']['sdf_gt'] * valid_mask, data_pr['volume'][0,0] * valid_mask)[None] * self.cfg['loss_sdf_weight']
if self.cfg['loss_eikonal_weight'] > 0:
outputs['loss_eikonal'] = (data_pr['sdf_gradient_error']).mean()[None] * self.cfg['loss_eikonal_weight']
if self.cfg['record_s']:
outputs['variance'] = data_pr['s'][None]
if self.cfg['loss_s_weight'] > 0:
outputs['loss_s'] = torch.norm(data_pr['s']).mean()[None] * self.cfg['loss_s_weight']
return outputs
class VGNLoss(Loss):
default_cfg={
'loss_vgn_weight': 1e-2,
}
def __init__(self, cfg):
super().__init__(['loss_vgn'])
self.cfg={**self.default_cfg,**cfg}
def _loss_fn(self, y_pred, y, is_train):
label_pred, rotation_pred, width_pred = y_pred
_, label, rotations, width = y
loss_qual = self._qual_loss_fn(label_pred, label)
acc = self._acc_fn(label_pred, label)
loss_rot_raw = self._rot_loss_fn(rotation_pred, rotations)
loss_rot = label * loss_rot_raw
loss_width_raw = 0.01 * self._width_loss_fn(width_pred, width)
loss_width = label * loss_width_raw
loss = loss_qual + loss_rot + loss_width
loss_item = {'loss_vgn': loss.mean()[None] * self.cfg['loss_vgn_weight'],
'vgn_total_loss':loss.mean()[None],'vgn_qual_loss': loss_qual.mean()[None],
'vgn_rot_loss': loss_rot.mean()[None], 'vgn_width_loss':loss_width.mean()[None],
'vgn_qual_acc': acc[None]}
num = torch.count_nonzero(label)
angle_torch = label * self._angle_error_fn(rotation_pred, rotations, 'torch')
loss_item['vgn_rot_err'] = (angle_torch.sum() / num)[None] if num else torch.zeros((1,),device=label.device)
return loss_item
def _qual_loss_fn(self, pred, target):
return F.binary_cross_entropy(pred, target, reduction="none")
def _acc_fn(self, pred, target):
return 100 * (torch.round(pred) == target).float().sum() / target.shape[0]
def _pr_fn(self, pred, target):
p, r = torchmetrics.functional.precision_recall(torch.round(pred).to(torch.int), target.to(torch.int), 'macro',num_classes=2)
return p[None] * 100, r[None] * 100
def _rot_loss_fn(self, pred, target):
loss0 = self._quat_loss_fn(pred, target[:, 0])
loss1 = self._quat_loss_fn(pred, target[:, 1])
return torch.min(loss0, loss1)
def _angle_error_fn(self, pred, target, method='torch'):
if method == 'np':
def _angle_error(q1, q2, ):
q1 = pyq.Quaternion(q1[[3,0,1,2]])
q1 /= q1.norm
q2 = pyq.Quaternion(q2[[3,0,1,2]])
q2 /= q2.norm
qd = q1.conjugate * q2
qdv = pyq.Quaternion(0, qd.x, qd.y, qd.z)
err = 2 * math.atan2(qdv.norm, qd.w) / math.pi * 180
return min(err, 360 - err)
q1s = pred.detach().cpu().numpy()
q2s = target.detach().cpu().numpy()
err = []
for q1,q2 in zip(q1s, q2s):
err.append(min(_angle_error(q1, q2[0]), _angle_error(q1, q2[1])))
return torch.tensor(err, device = pred.device)
elif method == 'torch':
return calc_rot_error_from_qxyzw(pred, target)
else:
raise NotImplementedError
def _quat_loss_fn(self, pred, target):
return 1.0 - torch.abs(torch.sum(pred * target, dim=1))
def _width_loss_fn(self, pred, target):
return F.mse_loss(pred, target, reduction="none")
def __call__(self, data_pr, data_gt, step, is_train=True, **kwargs):
return self._loss_fn(data_pr['vgn_pred'], data_gt['grasp_info'], is_train)
name2loss={
'render': RenderLoss,
'depth': DepthLoss,
'consist': ConsistencyLoss,
'vgn': VGNLoss,
'sdf': SDFLoss
}
================================================
FILE: src/nr/network/metrics.py
================================================
from pathlib import Path
import torch
from skimage.io import imsave
from network.loss import Loss
from utils.base_utils import color_map_backward, make_dir
from skimage.metrics import structural_similarity
import numpy as np
from utils.draw_utils import concat_images_list
def compute_psnr(img_gt, img_pr, use_vis_scores=False, vis_scores=None, vis_scores_thresh=1.5):
if use_vis_scores:
mask = vis_scores >= vis_scores_thresh
mask = mask.flatten()
img_gt = img_gt.reshape([-1, 3]).astype(np.float32)[mask]
img_pr = img_pr.reshape([-1, 3]).astype(np.float32)[mask]
mse = np.mean((img_gt - img_pr) ** 2, 0)
img_gt = img_gt.reshape([-1, 3]).astype(np.float32)
img_pr = img_pr.reshape([-1, 3]).astype(np.float32)
mse = np.mean((img_gt - img_pr) ** 2, 0)
mse = np.mean(mse)
psnr = 10 * np.log10(255 * 255 / mse)
return psnr
def compute_mae(depth_pr, depth_gt):
return np.mean(np.abs(depth_pr - depth_gt))
class PSNR_SSIM(Loss):
default_cfg = {
'eval_margin_ratio': 1.0,
}
def __init__(self, cfg):
super().__init__([])
self.cfg={**self.default_cfg,**cfg}
def __call__(self, data_pr, data_gt, step, **kwargs):
rgbs_gt = data_pr['pixel_colors_gt'] # 1,rn,3
rgbs_pr = data_pr['pixel_colors_nr'] # 1,rn,3
if 'que_imgs_info' in data_gt:
h, w = data_gt['que_imgs_info']['imgs'].shape[2:]
else:
h, w = data_pr['que_imgs_info']['imgs'].shape[2:]
rgbs_pr = rgbs_pr.reshape([h,w,3]).detach().cpu().numpy()
rgbs_pr=color_map_backward(rgbs_pr)
rgbs_gt = rgbs_gt.reshape([h,w,3]).detach().cpu().numpy()
rgbs_gt = color_map_backward(rgbs_gt)
h, w, _ = rgbs_gt.shape
h_margin = int(h * (1 - self.cfg['eval_margin_ratio'])) // 2
w_margin = int(w * (1 - self.cfg['eval_margin_ratio'])) // 2
rgbs_gt = rgbs_gt[h_margin:h - h_margin, w_margin:w - w_margin]
rgbs_pr = rgbs_pr[h_margin:h - h_margin, w_margin:w - w_margin]
psnr = compute_psnr(rgbs_gt,rgbs_pr)
outputs={
'psnr_nr': torch.tensor([psnr],dtype=torch.float32),
}
def compute_psnr_prefix(suffix):
if f'pixel_colors_{suffix}' in data_pr:
rgbs_other = data_pr[f'pixel_colors_{suffix}'] # 1,rn,3
# h, w = data_pr['shape']
rgbs_other = rgbs_other.reshape([h,w,3]).detach().cpu().numpy()
rgbs_other=color_map_backward(rgbs_other)
psnr = compute_psnr(rgbs_gt,rgbs_other)
ssim = structural_similarity(rgbs_gt,rgbs_other,win_size=11,multichannel=True,data_range=255)
outputs[f'psnr_{suffix}']=torch.tensor([psnr], dtype=torch.float32)
# compute_psnr_prefix('nr')
compute_psnr_prefix('dr')
compute_psnr_prefix('nr_fine')
compute_psnr_prefix('dr_fine')
depth_pr = data_pr['render_depth'].reshape([h,w]).detach().cpu().numpy()
depth_gt = data_gt['que_imgs_info']['true_depth'][0,0].cpu().numpy()
outputs['depth_mae'] = torch.tensor([compute_mae(depth_pr, depth_gt)],dtype=torch.float32) # higher is better
return outputs
class VisualizeImage(Loss):
def __init__(self, cfg):
super().__init__([])
def __call__(self, data_pr, data_gt, step, **kwargs):
if 'que_imgs_info' in data_gt:
h, w = data_gt['que_imgs_info']['imgs'].shape[2:]
else:
h, w = data_pr['que_imgs_info']['imgs'].shape[2:]
def get_img(key):
rgbs = data_pr[key] # 1,rn,3
rgbs = rgbs.reshape([h,w,3]).detach().cpu().numpy()
rgbs = color_map_backward(rgbs)
return rgbs
outputs={}
imgs=[get_img('pixel_colors_gt'), get_img('pixel_colors_nr')]
if 'pixel_colors_dr' in data_pr: imgs.append(get_img('pixel_colors_dr'))
if 'pixel_colors_nr_fine' in data_pr: imgs.append(get_img('pixel_colors_nr_fine'))
if 'pixel_colors_dr_fine' in data_pr: imgs.append(get_img('pixel_colors_dr_fine'))
data_index=kwargs['data_index']
model_name=kwargs['model_name']
Path(f'data/vis_val/{model_name}').mkdir(exist_ok=True, parents=True)
if h<=64 and w<=64:
imsave(f'data/vis_val/{model_name}/step-{step}-index-{data_index}.png',concat_images_list(*imgs))
else:
imsave(f'data/vis_val/{model_name}/step-{step}-index-{data_index}.jpg', concat_images_list(*imgs))
return outputs
name2metrics={
'psnr_ssim': PSNR_SSIM,
'vis_img': VisualizeImage,
}
def psnr_nr(results):
return np.mean(results['psnr_nr'])
def psnr_nr_fine(results):
return np.mean(results['psnr_nr_fine'])
def depth_mae(results):
return np.mean(results['depth_mae'])
def sdf_mae(results):
return np.mean(results['sdf_mae'])
def loss_vgn(results):
if 'loss_vgn' in results:
return np.mean(results['loss_vgn'])
else:
return 1e6
name2key_metrics={
'psnr_nr': psnr_nr,
'psnr_nr_fine': psnr_nr_fine,
'depth_mae': depth_mae,
'loss_vgn': loss_vgn,
'sdf_mae': sdf_mae
}
================================================
FILE: src/nr/network/mvsnet/modules.py
================================================
import torch
from torch import nn
import torch.nn.functional as F
from inplace_abn import InPlaceABN
from kornia.utils import create_meshgrid
class ConvBnReLU(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1, norm_act=InPlaceABN):
super(ConvBnReLU, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
self.bn = norm_act(out_channels)
def forward(self, x):
return self.bn(self.conv(x))
class ConvBnReLU3D(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1, norm_act=InPlaceABN):
super(ConvBnReLU3D, self).__init__()
self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
self.bn = norm_act(out_channels)
def forward(self, x):
return self.bn(self.conv(x))
def homo_warp(src_feat, src_proj, ref_proj_inv, depth_values):
# src_feat: (B, C, H, W)
# src_proj: (B, 4, 4)
# ref_proj_inv: (B, 4, 4)
# depth_values: (B, D)
# out: (B, C, D, H, W)
B, C, H, W = src_feat.shape
D = depth_values.shape[1]
device = src_feat.device
dtype = src_feat.dtype
transform = src_proj @ ref_proj_inv
R = transform[:, :3, :3] # (B, 3, 3)
T = transform[:, :3, 3:] # (B, 3, 1)
# create grid from the ref frame
ref_grid = create_meshgrid(H, W, normalized_coordinates=False) # (1, H, W, 2)
ref_grid = ref_grid.to(device).to(dtype)
ref_grid = ref_grid.permute(0, 3, 1, 2) # (1, 2, H, W)
ref_grid = ref_grid.reshape(1, 2, H*W) # (1, 2, H*W)
ref_grid = ref_grid.expand(B, -1, -1) # (B, 2, H*W)
ref_grid = torch.cat((ref_grid, torch.ones_like(ref_grid[:,:1])), 1) # (B, 3, H*W)
ref_grid_d = ref_grid.unsqueeze(2) * depth_values.view(B, 1, D, 1) # (B, 3, D, H*W)
ref_grid_d = ref_grid_d.view(B, 3, D*H*W)
src_grid_d = R @ ref_grid_d + T # (B, 3, D*H*W)
del ref_grid_d, ref_grid, transform, R, T # release (GPU) memory
div_val = src_grid_d[:, -1:]
div_val[div_val<1e-4] = 1e-4
src_grid = src_grid_d[:, :2] / div_val # divide by depth (B, 2, D*H*W)
del src_grid_d, div_val
src_grid[:, 0] = src_grid[:, 0]/((W - 1) / 2) - 1 # scale to -1~1
src_grid[:, 1] = src_grid[:, 1]/((H - 1) / 2) - 1 # scale to -1~1
src_grid = src_grid.permute(0, 2, 1) # (B, D*H*W, 2)
src_grid = src_grid.view(B, D, H*W, 2)
warped_src_feat = F.grid_sample(src_feat, src_grid,
mode='bilinear', padding_mode='zeros',
align_corners=True) # (B, C, D, H*W)
warped_src_feat = warped_src_feat.view(B, C, D, H, W)
return warped_src_feat
def depth_regression(p, depth_values):
# p: probability volume [B, D, H, W]
# depth_values: discrete depth values [B, D]
depth_values = depth_values.view(*depth_values.shape, 1, 1)
depth = torch.sum(p * depth_values, 1)
return depth
================================================
FILE: src/nr/network/mvsnet/mvsnet.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from network.mvsnet.modules import ConvBnReLU, ConvBnReLU3D, depth_regression, homo_warp
from inplace_abn import InPlaceABN
class FeatureNet(nn.Module):
def __init__(self, norm_act=InPlaceABN):
super(FeatureNet, self).__init__()
self.inplanes = 32
self.conv0 = ConvBnReLU(3, 8, 3, 1, 1, norm_act=norm_act)
self.conv1 = ConvBnReLU(8, 8, 3, 1, 1, norm_act=norm_act)
self.conv2 = ConvBnReLU(8, 16, 5, 2, 2, norm_act=norm_act)
self.conv3 = ConvBnReLU(16, 16, 3, 1, 1, norm_act=norm_act)
self.conv4 = ConvBnReLU(16, 16, 3, 1, 1, norm_act=norm_act)
self.conv5 = ConvBnReLU(16, 32, 5, 2, 2, norm_act=norm_act)
self.conv6 = ConvBnReLU(32, 32, 3, 1, 1, norm_act=norm_act)
self.feature = nn.Conv2d(32, 32, 3, 1, 1)
def forward(self, x):
x = self.conv1(self.conv0(x))
x = self.conv4(self.conv3(self.conv2(x)))
x = self.feature(self.conv6(self.conv5(x)))
return x
class CostRegNet(nn.Module):
def __init__(self, norm_act=InPlaceABN):
super(CostRegNet, self).__init__()
self.conv0 = ConvBnReLU3D(32, 8, norm_act=norm_act)
self.conv1 = ConvBnReLU3D(8, 16, stride=2, norm_act=norm_act)
self.conv2 = ConvBnReLU3D(16, 16, norm_act=norm_act)
self.conv3 = ConvBnReLU3D(16, 32, stride=2, norm_act=norm_act)
self.conv4 = ConvBnReLU3D(32, 32, norm_act=norm_act)
self.conv5 = ConvBnReLU3D(32, 64, stride=2, norm_act=norm_act)
self.conv6 = ConvBnReLU3D(64, 64, norm_act=norm_act)
self.conv7 = nn.Sequential(
nn.ConvTranspose3d(64, 32, kernel_size=3, padding=1, output_padding=1, stride=2, bias=False),
norm_act(32))
self.conv9 = nn.Sequential(
nn.ConvTranspose3d(32, 16, kernel_size=3, padding=1, output_padding=1, stride=2, bias=False),
norm_act(16))
self.conv11 = nn.Sequential(
nn.ConvTranspose3d(16, 8, kernel_size=3, padding=1, output_padding=1, stride=2, bias=False),
norm_act(8))
self.prob = nn.Conv3d(8, 1, 3, stride=1, padding=1)
def forward(self, x):
conv0 = self.conv0(x)
conv2 = self.conv2(self.conv1(conv0))
conv4 = self.conv4(self.conv3(conv2))
x = self.conv6(self.conv5(conv4))
x = conv4 + self.conv7(x)
del conv4
x = conv2 + self.conv9(x)
del conv2
x = conv0 + self.conv11(x)
del conv0
x = self.prob(x)
return x
class MVSNet(nn.Module):
def __init__(self, norm_act=InPlaceABN):
super(MVSNet, self).__init__()
self.feature = FeatureNet(norm_act)
self.cost_regularization = CostRegNet(norm_act)
def forward(self, imgs, proj_mats, depth_values):
# imgs: (B, V, 3, H, W)
# proj_mats: (B, V, 4, 4)
# depth_values: (B, D)
B, V, _, H, W = imgs.shape
D = depth_values.shape[1]
# step 1. feature extraction
# in: images; out: 32-channel feature maps
imgs = imgs.reshape(B*V, 3, H, W)
feats = self.feature(imgs) # (B*V, F, h, w)
del imgs
feats = feats.reshape(B, V, *feats.shape[1:]) # (B, V, F, h, w)
ref_feats, src_feats = feats[:, 0], feats[:, 1:]
ref_proj, src_projs = proj_mats[:, 0], proj_mats[:, 1:]
src_feats = src_feats.permute(1, 0, 2, 3, 4) # (V-1, B, F, h, w)
src_projs = src_projs.permute(1, 0, 2, 3) # (V-1, B, 4, 4)
# step 2. differentiable homograph, build cost volume
ref_volume = ref_feats.unsqueeze(2).repeat(1, 1, D, 1, 1) # (B, F, D, h, w)
volume_sum = ref_volume
volume_sq_sum = ref_volume ** 2
del ref_volume
ref_proj = torch.inverse(ref_proj)
for src_feat, src_proj in zip(src_feats, src_projs):
warped_volume = homo_warp(src_feat, src_proj, ref_proj, depth_values)
volume_sum = volume_sum + warped_volume
volume_sq_sum = volume_sq_sum + warped_volume ** 2
del warped_volume
# aggregate multiple feature volumes by variance
volume_variance = volume_sq_sum.div_(V).sub_(volume_sum.div_(V).pow_(2))
del volume_sq_sum, volume_sum
# step 3. cost volume regularization
cost_reg = self.cost_regularization(volume_variance).squeeze(1)
prob_volume = F.softmax(cost_reg, 1) # (B, D, h, w)
depth = depth_regression(prob_volume, depth_values)
with torch.no_grad():
# sum probability of 4 consecutive depth indices
prob_volume_sum4 = 4 * F.avg_pool3d(F.pad(prob_volume.unsqueeze(1),
pad=(0, 0, 0, 0, 1, 2)),
(4, 1, 1), stride=1).squeeze(1) # (B, D, h, w)
# find the (rounded) index that is the final prediction
depth_index = depth_regression(prob_volume,
torch.arange(D,
device=prob_volume.device,
dtype=prob_volume.dtype)
).long() # (B, h, w)
# the confidence is the 4-sum probability at this index
confidence = torch.gather(prob_volume_sum4, 1,
depth_index.unsqueeze(1)).squeeze(1) # (B, h, w)
return depth, confidence
def construct_cost_volume(self, ref_imgs, ref_nn_idx, ref_prjs, depth_values, batch_num=2):
# ref_imgs rfn,3,h,w
# ref_nn_ids: rfn,nn
# ref_prjs: rfn,4,4 note it is already scaled!!!
# depth_values: rfn,dn
# return: rfn,dn,h//4,w//4
ref_feats = self.feature(ref_imgs) # rfn,f,h,w
ref_prjs_inv = torch.inverse(ref_prjs) # rfn,4,4
dn = depth_values.shape[1]
rfn, n_num = ref_nn_idx.shape
cost_reg_all = []
for rfi in range(0,rfn,batch_num):
volume_sum, volume_sum_sq = ref_feats[rfi:rfi+batch_num].unsqueeze(2), ref_feats[rfi:rfi+batch_num].unsqueeze(2)**2 # 1,f,1,h,w
volume_sum, volume_sum_sq = volume_sum.repeat(1, 1, dn, 1, 1), volume_sum_sq.repeat(1, 1, dn, 1, 1)
for ni in range(n_num):
warp_feats = homo_warp(ref_feats[ref_nn_idx[rfi:rfi+batch_num,ni]],ref_prjs[ref_nn_idx[rfi:rfi+batch_num,ni]],
ref_prjs_inv[rfi:rfi+batch_num],depth_values[rfi:rfi+batch_num]) # 1,f,dn,h,w
volume_sum += warp_feats
volume_sum_sq += warp_feats**2
volume_variance = volume_sum_sq.div_(n_num+1).sub_(volume_sum.div_(n_num+1).pow_(2)) # 1,f,dn,h,w
del volume_sum_sq, volume_sum
# 1,dn,h,w
cost_reg_all.append(self.cost_regularization(volume_variance).squeeze(1))
cost_reg_all = torch.cat(cost_reg_all,0)
return cost_reg_all
def construct_cost_volume_with_src(self, ref_imgs, src_imgs, ref_nn_idx, ref_prjs, src_prjs, depth_values, batch_num=2):
# ref_imgs rfn,3,h,w
# src_imgs srn,3,h,w
# ref_nn_ids: rfn,nn
# ref_prjs: rfn,4,4 note it is already scaled!!!
# src_prjs: src,4,4 note it is already scaled!!!
# depth_values: rfn,dn
# return: rfn,dn,h//4,w//4
ref_feats = self.feature(ref_imgs) # rfn,f,h,w
src_feats = self.feature(src_imgs) # src,f,h,w
ref_prjs_inv = torch.inverse(ref_prjs) # rfn,4,4
dn = depth_values.shape[1]
rfn, n_num = ref_nn_idx.shape
cost_reg_all = []
for rfi in range(0,rfn,batch_num):
volume_sum, volume_sum_sq = ref_feats[rfi:rfi+batch_num].unsqueeze(2), ref_feats[rfi:rfi+batch_num].unsqueeze(2)**2 # 1,f,1,h,w
volume_sum, volume_sum_sq = volume_sum.repeat(1, 1, dn, 1, 1), volume_sum_sq.repeat(1, 1, dn, 1, 1)
for ni in range(n_num):
warp_feats = homo_warp(src_feats[ref_nn_idx[rfi:rfi+batch_num,ni]],src_prjs[ref_nn_idx[rfi:rfi+batch_num,ni]],
ref_prjs_inv[rfi:rfi+batch_num],depth_values[rfi:rfi+batch_num]) # 1,f,dn,h,w
volume_sum += warp_feats
volume_sum_sq += warp_feats**2
volume_variance = volume_sum_sq.div_(n_num+1).sub_(volume_sum.div_(n_num+1).pow_(2)) # 1,f,dn,h,w
del volume_sum_sq, volume_sum
# 1,dn,h,w
cost_reg_all.append(self.cost_regularization(volume_variance).squeeze(1))
cost_reg_all = torch.cat(cost_reg_all,0)
return cost_reg_all
def extract_model_state_dict(ckpt_path, prefixes_to_ignore=[]):
checkpoint = torch.load(ckpt_path, map_location=torch.device('cpu'))
checkpoint_ = {}
if 'state_dict' in checkpoint: # if it's a pytorch-lightning checkpoint
for k, v in checkpoint['state_dict'].items():
if not k.startswith('model.'):
continue
k = k[6:] # remove 'model.'
for prefix in prefixes_to_ignore:
if k.startswith(prefix):
print('ignore', k)
break
else:
checkpoint_[k] = v
else: # if it only has model weights
for k, v in checkpoint.items():
for prefix in prefixes_to_ignore:
if k.startswith(prefix):
print('ignore', k)
break
else:
checkpoint_[k] = v
return checkpoint_
def load_ckpt(model, ckpt_path, prefixes_to_ignore=[]):
model_dict = model.state_dict()
checkpoint_ = extract_model_state_dict(ckpt_path, prefixes_to_ignore)
model_dict.update(checkpoint_)
model.load_state_dict(model_dict)
================================================
FILE: src/nr/network/neus.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class SingleVarianceNetwork(nn.Module):
def __init__(self, init_val, fix_s=-1):
super(SingleVarianceNetwork, self).__init__()
self.register_parameter('variance', nn.Parameter(torch.tensor(init_val)))
self.variance.requires_grad = False
self.step = 0
self.fix_s = fix_s
def set_step(self, step):
self.step = step
def forward(self, x):
if self.fix_s != -1 and self.step > self.fix_s:
self.variance.requires_grad = True
return torch.ones([len(x), 1], device=x.device) * torch.exp(self.variance * 10.0)
class Embedder:
def __init__(self, **kwargs):
self.kwargs = kwargs
self.create_embedding_fn()
def create_embedding_fn(self):
embed_fns = []
d = self.kwargs['input_dims']
out_dim = 0
if self.kwargs['include_input']:
embed_fns.append(lambda x: x)
out_dim += d
max_freq = self.kwargs['max_freq_log2']
N_freqs = self.kwargs['num_freqs']
if self.kwargs['log_sampling']:
freq_bands = 2. ** torch.linspace(0., max_freq, N_freqs)
else:
freq_bands = torch.linspace(2.**0., 2.**max_freq, N_freqs)
for freq in freq_bands:
for p_fn in self.kwargs['periodic_fns']:
embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq))
out_dim += d
self.embed_fns = embed_fns
self.out_dim = out_dim
def embed(self, inputs):
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
def get_embedder(multires, input_dims=3):
embed_kwargs = {
'include_input': True,
'input_dims': input_dims,
'max_freq_log2': multires-1,
'num_freqs': multires,
'log_sampling': True,
'periodic_fns': [torch.sin, torch.cos],
}
embedder_obj = Embedder(**embed_kwargs)
def embed(x, eo=embedder_obj): return eo.embed(x)
return embed, embedder_obj.out_dim
================================================
FILE: src/nr/network/ops.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=dilation, groups=groups, bias=False, dilation=dilation, padding_mode='reflect')
def conv1x1(in_planes, out_planes, stride=1):
"""1x1 convolution"""
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False, padding_mode='reflect')
def interpolate_feats(feats, points, h=None, w=None, padding_mode='zeros', align_corners=False, inter_mode='bilinear'):
"""
:param feats: b,f,h,w
:param points: b,n,2
:param h: float
:param w: float
:param padding_mode:
:param align_corners:
:param inter_mode:
:return:
"""
b, _, ch, cw = feats.shape
if h is None and w is None:
h, w = ch, cw
x_norm = points[:, :, 0] / (w - 1) * 2 - 1
y_norm = points[:, :, 1] / (h - 1) * 2 - 1
points_norm = torch.stack([x_norm, y_norm], -1).unsqueeze(1) # [srn,1,n,2]
feats_inter = F.grid_sample(feats, points_norm, mode=inter_mode, padding_mode=padding_mode, align_corners=align_corners).squeeze(2) # srn,f,n
feats_inter = feats_inter.permute(0,2,1)
return feats_inter
def masked_mean_var(feats,mask,dim=2):
mask=mask.float() # b,1,n,1
mask_sum = torch.clamp_min(torch.sum(mask,dim,keepdim=True),min=1e-4) # b,1,1,1
feats_mean = torch.sum(feats*mask,dim,keepdim=True)/mask_sum # b,f,1,1
feats_var = torch.sum((feats-feats_mean)**2*mask,dim,keepdim=True)/mask_sum # b,f,1,1
return feats_mean, feats_var
class ResidualBlock(nn.Module):
def __init__(self, dim_in, dim_out, dim_inter=None, use_norm=True, norm_layer=nn.BatchNorm2d,bias=False):
super().__init__()
if dim_inter is None:
dim_inter=dim_out
if use_norm:
self.conv=nn.Sequential(
norm_layer(dim_in),
nn.ReLU(True),
nn.Conv2d(dim_in,dim_inter,3,1,1,bias=bias,padding_mode='reflect'),
norm_layer(dim_inter),
nn.ReLU(True),
nn.Conv2d(dim_inter,dim_out,3,1,1,bias=bias,padding_mode='reflect'),
)
else:
self.conv=nn.Sequential(
nn.ReLU(True),
nn.Conv2d(dim_in,dim_inter,3,1,1),
nn.ReLU(True),
nn.Conv2d(dim_inter,dim_out,3,1,1),
)
self.short_cut=None
if dim_in!=dim_out:
self.short_cut=nn.Conv2d(dim_in,dim_out,1,1)
def forward(self, feats):
feats_out=self.conv(feats)
if self.short_cut is not None:
feats_out=self.short_cut(feats)+feats_out
else:
feats_out=feats_out+feats
return feats_out
class AddBias(nn.Module):
def __init__(self,val):
super().__init__()
self.val=val
def forward(self,x):
return x+self.val
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
base_width=64, dilation=1, norm_layer=None):
super(BasicBlock, self).__init__()
if norm_layer is None:
norm_layer = nn.BatchNorm2d
if groups != 1 or base_width != 64:
raise ValueError('BasicBlock only supports groups=1 and base_width=64')
if dilation > 1:
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = norm_layer(planes, track_running_stats=False, affine=True)
self.relu = nn.ReLU(inplace=True)
self.conv2 = conv3x3(planes, planes)
self.bn2 = norm_layer(planes, track_running_stats=False, affine=True)
self.downsample = downsample
self.stride = stride
def forward(self, x):
identity = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
out = self.conv2(out)
out = self.bn2(out)
if self.downsample is not None:
identity = self.downsample(x)
out += identity
out = self.relu(out)
return out
class conv(nn.Module):
def __init__(self, num_in_layers, num_out_layers, kernel_size, stride):
super(conv, self).__init__()
self.kernel_size = kernel_size
self.conv = nn.Conv2d(num_in_layers,
num_out_layers,
kernel_size=kernel_size,
stride=stride,
padding=(self.kernel_size - 1) // 2,
padding_mode='reflect')
self.bn = nn.InstanceNorm2d(num_out_layers, track_running_stats=False, affine=True)
def forward(self, x):
return F.elu(self.bn(self.conv(x)), inplace=True)
class upconv(nn.Module):
def __init__(self, num_in_layers, num_out_layers, kernel_size, scale):
super(upconv, self).__init__()
self.scale = scale
self.conv = conv(num_in_layers, num_out_layers, kernel_size, 1)
def forward(self, x):
x = nn.functional.interpolate(x, scale_factor=self.scale, align_corners=True, mode='bilinear')
return self.conv(x)
class ResUNetLight(nn.Module):
def __init__(self, in_dim=3, layers=(2, 3, 6, 3), out_dim=32, inplanes=32):
super(ResUNetLight, self).__init__()
# layers = [2, 3, 6, 3]
norm_layer = nn.InstanceNorm2d
self._norm_layer = norm_layer
self.dilation = 1
block = BasicBlock
replace_stride_with_dilation = [False, False, False]
self.inplanes = inplanes
self.groups = 1 # seems useless
self.base_width = 64 # seems useless
self.conv1 = nn.Conv2d(in_dim, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False,
padding_mode='reflect')
self.bn1 = norm_layer(self.inplanes, track_running_stats=False, affine=True)
self.relu = nn.ReLU(inplace=True)
self.layer1 = self._make_layer(block, 32, layers[0], stride=2)
self.layer2 = self._make_layer(block, 64, layers[1], stride=2,
dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block, 128, layers[2], stride=2,
dilate=replace_stride_with_dilation[1])
# decoder
self.upconv3 = upconv(128, 64, 3, 2)
self.iconv3 = conv(64 + 64, 64, 3, 1)
self.upconv2 = upconv(64, 32, 3, 2)
self.iconv2 = conv(32 + 32, 32, 3, 1)
# fine-level conv
self.out_conv = nn.Conv2d(32, out_dim, 1, 1)
def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
norm_layer = self._norm_layer
downsample = None
previous_dilation = self.dilation
if dilate:
self.dilation *= stride
stride = 1
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
conv1x1(self.inplanes, planes * block.expansion, stride),
norm_layer(planes * block.expansion, track_running_stats=False, affine=True),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
self.base_width, previous_dilation, norm_layer))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes, groups=self.groups,
base_width=self.base_width, dilation=self.dilation,
norm_layer=norm_layer))
return nn.Sequential(*layers)
def skipconnect(self, x1, x2):
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2))
x = torch.cat([x2, x1], dim=1)
return x
def forward(self, x):
x = self.relu(self.bn1(self.conv1(x)))
x1 = self.layer1(x)
x2 = self.layer2(x1)
x3 = self.layer3(x2)
x = self.upconv3(x3)
x = self.skipconnect(x2, x)
x = self.iconv3(x)
x = self.upconv2(x)
x = self.skipconnect(x1, x)
x = self.iconv2(x)
x_out = self.out_conv(x)
return x_out
class ResEncoder(nn.Module):
def __init__(self):
super(ResEncoder, self).__init__()
self.inplanes = 32
filters = [32, 64, 128]
layers = [2, 2, 2, 2]
out_planes = 32
norm_layer = nn.InstanceNorm2d
self._norm_layer = norm_layer
self.dilation = 1
block = BasicBlock
replace_stride_with_dilation = [False, False, False]
self.groups = 1
self.base_width = 64
self.conv1 = nn.Conv2d(12, self.inplanes, kernel_size=8, stride=2, padding=2,
bias=False, padding_mode='reflect')
self.bn1 = norm_layer(self.inplanes, track_running_stats=False, affine=True)
self.relu = nn.ReLU(inplace=True)
self.layer1 = self._make_layer(block, filters[0], layers[0], stride=2)
self.layer2 = self._make_layer(block, filters[1], layers[1], stride=2,
dilate=replace_stride_with_dilation[0])
self.layer3 = self._make_layer(block, filters[2], layers[2], stride=2,
dilate=replace_stride_with_dilation[1])
# decoder
self.upconv3 = upconv(filters[2], filters[1], 3, 2)
self.iconv3 = conv(filters[1]*2, filters[1], 3, 1)
self.upconv2 = upconv(filters[1], filters[0], 3, 2)
self.iconv2 = conv(filters[0]*2, out_planes, 3, 1)
self.out_conv = nn.Conv2d(out_planes, out_planes, 1, 1)
def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
norm_layer = self._norm_layer
downsample = None
previous_dilation = self.dilation
if dilate:
self.dilation *= stride
stride = 1
if stride != 1 or self.inplanes != planes * block.expansion:
downsample = nn.Sequential(
conv1x1(self.inplanes, planes * block.expansion, stride),
norm_layer(planes * block.expansion, track_running_stats=False, affine=True),
)
layers = []
layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
self.base_width, 1, norm_layer))
self.inplanes = planes * block.expansion
for _ in range(1, blocks):
layers.append(block(self.inplanes, planes, groups=self.groups,
base_width=self.base_width, dilation=self.dilation,
norm_layer=norm_layer))
return nn.Sequential(*layers)
def skipconnect(self, x1, x2):
diffY = x2.size()[2] - x1.size()[2]
diffX = x2.size()[3] - x1.size()[3]
x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2))
# for padding issues, see
# https://github.com/HaiyongJiang/U-Net-Pytorch-Unstructured-Buggy/commit/0e854509c2cea854e247a9c615f175f76fbb2e3a
# https://github.com/xiaopeng-liao/Pytorch-UNet/commit/8ebac70e633bac59fc22bb5195e513d5832fb3bd
x = torch.cat([x2, x1], dim=1)
return x
def forward(self, x):
x = self.relu(self.bn1(self.conv1(x)))
x1 = self.layer1(x)
x2 = self.layer2(x1)
x3 = self.layer3(x2)
x = self.upconv3(x3)
x = self.skipconnect(x2, x)
x = self.iconv3(x)
x = self.upconv2(x)
x = self.skipconnect(x1, x)
x = self.iconv2(x)
x_out = self.out_conv(x)
return x_out
================================================
FILE: src/nr/network/render_ops.py
================================================
import torch
from network.ops import interpolate_feats
def coords2rays(coords, poses, Ks):
"""
:param coords: [rfn,rn,2]
:param poses: [rfn,3,4]
:param Ks: [rfn,3,3]
:return:
ref_rays:
centers: [rfn,rn,3]
directions: [rfn,rn,3]
"""
rot = poses[:, :, :3].unsqueeze(1).permute(0, 1, 3, 2) # rfn,1,3,3
trans = -rot @ poses[:, :, 3:].unsqueeze(1) # rfn,1,3,1
rfn, rn, _ = coords.shape
centers = trans.repeat(1, rn, 1, 1).squeeze(-1) # rfn,rn,3
coords = torch.cat([coords, torch.ones([rfn, rn, 1], dtype=torch.float32, device=coords.device)], 2) # rfn,rn,3
Ks_inv = torch.inverse(Ks).unsqueeze(1)
cam_xyz = Ks_inv @ coords.unsqueeze(3)
cam_xyz = rot @ cam_xyz + trans
directions = cam_xyz.squeeze(3) - centers
# directions = directions / torch.clamp(torch.norm(directions, dim=2, keepdim=True), min=1e-4)
return centers, directions
def depth2points(que_imgs_info, que_depth):
"""
:param que_imgs_info:
:param que_depth: qn,rn,dn
:return:
"""
cneters, directions = coords2rays(que_imgs_info['coords'],que_imgs_info['poses'],que_imgs_info['Ks']) # centers, directions qn,rn,3
qn, rn, _ = cneters.shape
que_pts = cneters.unsqueeze(2) + directions.unsqueeze(2) * que_depth.unsqueeze(3) # qn,rn,dn,3
qn, rn, dn, _ = que_pts.shape
que_dir = -directions / torch.norm(directions, dim=2, keepdim=True) # qn,rn,3
que_dir = que_dir.unsqueeze(2).repeat(1, 1, dn, 1)
return que_pts, que_dir # qn,rn,dn,3
def depth2dists(depth):
device = depth.device
dists = depth[...,1:]-depth[...,:-1]
return torch.cat([dists, torch.full([*depth.shape[:-1], 1], 1e6, dtype=torch.float32, device=device)], -1)
def depth2inv_dists(depth,depth_range):
near, far = -1 / depth_range[:, 0], -1 / depth_range[:, 1]
near, far = near[:, None, None], far[:, None, None]
depth_inv = -1 / depth # qn,rn,dn
depth_inv = (depth_inv - near) / (far - near)
dists = depth2dists(depth_inv) # qn,rn,dn
return dists
def interpolate_feature_map(ray_feats, coords, mask, h, w, border_type='border'):
"""
:param ray_feats: rfn,f,h,w
:param coords: rfn,pn,2
:param mask: rfn,pn
:param h:
:param w:
:param border_type:
:return:
"""
fh, fw = ray_feats.shape[-2:]
if fh == h and fw == w:
cur_ray_feats = interpolate_feats(ray_feats, coords, h, w, border_type, True) # rfn,pn,f
else:
cur_ray_feats = interpolate_feats(ray_feats, coords, h, w, border_type, False) # rfn,pn,f
cur_ray_feats = cur_ray_feats * mask.float().unsqueeze(-1) # rfn,pn,f
return cur_ray_feats
def alpha_values2hit_prob(alpha_values):
"""
:param alpha_values: qn,rn,dn
:return: qn,rn,dn
"""
no_hit_density = torch.cat([torch.ones((*alpha_values.shape[:-1], 1))
.to(alpha_values.device), 1. - alpha_values + 1e-10], -1) # rn,k+1
hit_prob = alpha_values * torch.cumprod(no_hit_density, -1)[..., :-1] # [n,k]
return hit_prob
def project_points_coords(pts, Rt, K):
"""
:param pts: [pn,3]
:param Rt: [rfn,3,4]
:param K: [rfn,3,3]
:return:
coords: [rfn,pn,2]
invalid_mask: [rfn,pn]
"""
pn = pts.shape[0]
hpts = torch.cat([pts,torch.ones([pn,1],device=pts.device,dtype=torch.float32)],1)
srn = Rt.shape[0]
KRt = K @ Rt # rfn,3,4
last_row = torch.zeros([srn,1,4],device=pts.device,dtype=torch.float32)
last_row[:,:,3] = 1.0
H = torch.cat([KRt,last_row],1) # rfn,4,4
pts_cam = H[:,None,:,:] @ hpts[None,:,:,None]
pts_cam = pts_cam[:,:,:3,0]
depth = pts_cam[:,:,2:]
invalid_mask = torch.abs(depth)<1e-4
depth[invalid_mask] = 1e-3
pts_2d = pts_cam[:,:,:2]/depth
return pts_2d, ~(invalid_mask[...,0]), depth
def project_points_directions(poses,points):
"""
:param poses: rfn,3,4
:param points: pn,3
:return: rfn,pn,3
"""
cam_pts = -poses[:, :, :3].permute(0, 2, 1) @ poses[:, :, 3:] # rfn,3,1
dir = points.unsqueeze(0) - cam_pts.permute(0, 2, 1) # [1,pn,3] - [rfn,1,3] -> rfn,pn,3
dir = -dir / torch.clamp_min(torch.norm(dir, dim=2, keepdim=True), min=1e-5) # rfn,pn,3
return dir
def project_points_ref_views(ref_imgs_info, que_points):
"""
:param ref_imgs_info:
:param que_points: pn,3
:return:
"""
prj_pts, prj_valid_mask, prj_depth = project_points_coords(
que_points, ref_imgs_info['poses'], ref_imgs_info['Ks']) # rfn,pn,2
h,w=ref_imgs_info['imgs'].shape[-2:]
prj_img_invalid_mask = (prj_pts[..., 0] < -0.5) | (prj_pts[..., 0] >= w - 0.5) | \
(prj_pts[..., 1] < -0.5) | (prj_pts[..., 1] >= h - 0.5)
valid_mask = prj_valid_mask & (~prj_img_invalid_mask)
prj_dir = project_points_directions(ref_imgs_info['poses'], que_points) # rfn,pn,3
return prj_dir, prj_pts, prj_depth, valid_mask
def project_points_dict(ref_imgs_info, que_pts):
# project all points
qn, rn, dn, _ = que_pts.shape
prj_dir, prj_pts, prj_depth, prj_mask = project_points_ref_views(ref_imgs_info, que_pts.reshape([qn * rn * dn, 3]))
rfn, _, h, w = ref_imgs_info['imgs'].shape
prj_ray_feats = interpolate_feature_map(ref_imgs_info['ray_feats'], prj_pts, prj_mask, h, w)
prj_rgb = interpolate_feature_map(ref_imgs_info['imgs'], prj_pts, prj_mask, h, w)
prj_dict = {'dir':prj_dir, 'pts':prj_pts, 'depth':prj_depth, 'mask': prj_mask.float(), 'ray_feats':prj_ray_feats, 'rgb':prj_rgb}
# post process
for k, v in prj_dict.items():
prj_dict[k]=v.reshape(rfn,qn,rn,dn,-1)
return prj_dict
def sample_depth(depth_range, coords, sample_num, random_sample):
"""
:param depth_range: qn,2
:param sample_num:
:param random_sample:
:return:
"""
qn, rn, _ = coords.shape
device = coords.device
near, far = depth_range[:,0], depth_range[:,1] # qn,2
dn = sample_num
assert(dn>2)
interval = (1 / far - 1 / near) / (dn - 1) # qn
val = torch.arange(1, dn - 1, dtype=torch.float32, device=near.device)[None, None, :]
if random_sample:
val = val + (torch.rand(qn, rn, dn-2, dtype=torch.float32, device=device) - 0.5) * 0.999
else:
val = val + torch.zeros(qn, rn, dn-2, dtype=torch.float32, device=device)
ticks = interval[:, None, None] * val
diff = (1 / far - 1 / near)
ticks = torch.cat([torch.zeros(qn,rn,1,dtype=torch.float32,device=device),ticks,diff[:,None,None].repeat(1,rn,1)],-1)
que_depth = 1 / (1 / near[:, None, None] + ticks) # qn, dn,
que_dists = torch.cat([que_depth[...,1:],torch.full([*que_depth.shape[:-1],1],1e6,dtype=torch.float32,device=device)],-1) - que_depth
return que_depth, que_dists # qn, rn, dn
def sample_fine_depth(depth, hit_prob, depth_range, sample_num, random_sample, inv_mode=True):
"""
:param depth: qn,rn,dn
:param hit_prob: qn,rn,dn
:param depth_range: qn,2
:param sample_num:
:param random_sample:
:param inv_mode:
:return: qn,rn,dn
"""
if inv_mode:
near, far = depth_range[0,0], depth_range[0,1]
near, far = -1/near, -1/far
depth_inv = -1 / depth # qn,rn,dn
depth_inv = (depth_inv - near) / (far - near)
depth = depth_inv
depth_center = (depth[...,1:] + depth[...,:-1])/2
depth_center = torch.cat([depth[...,0:1],depth_center,depth[...,-1:]],-1) # rfn,pn,dn+1
fdn = sample_num
# Get pdf
hit_prob = hit_prob + 1e-5 # prevent nans
pdf = hit_prob / torch.sum(hit_prob, -1, keepdim=True) # rfn,pn,dn-1
cdf = torch.cumsum(pdf, -1) # rfn,pn,dn-1
cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) # rfn,pn,dn
# Take uniform samples
if not random_sample:
interval = 1 / fdn
u = 0.5*interval+torch.arange(fdn)*interval
# u = torch.linspace(0., 1., steps=fdn)
u = u.expand(list(cdf.shape[:-1]) + [fdn]) # rfn,pn,fdn
else:
u = torch.rand(list(cdf.shape[:-1]) + [fdn])
# Invert CDF
device = pdf.device
u = u.to(device).contiguous() # rfn,pn,fdn
inds = torch.searchsorted(cdf, u, right=True) # rfn,pn,fdn
below = torch.max(torch.zeros_like(inds-1), inds-1) # rfn,pn,fdn
above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds) # rfn,pn,fdn
inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2) # rfn,pn,fdn,2
matched_shape = [*inds_g.shape[:-1], cdf.shape[-1]]
cdf_g = torch.gather(cdf.unsqueeze(-2).expand(matched_shape), -1, inds_g) # rfn,pn,fdn,2
bins_g = torch.gather(depth_center.unsqueeze(-2).expand(matched_shape), -1, inds_g) # rfn,pn,fdn,2
denom = (cdf_g[...,1]-cdf_g[...,0]) # rfn,pn,fdn
denom = torch.where(denom<1e-5, torch.ones_like(denom), denom)
t = (u-cdf_g[...,0])/denom
fine_depth = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0])
if inv_mode:
near, far = depth_range[0,0], depth_range[0,1]
near, far = -1/near, -1/far
fine_depth = fine_depth * (far - near) + near
fine_depth = -1/fine_depth
return fine_depth
================================================
FILE: src/nr/network/renderer.py
================================================
import torch
import numpy as np
import torch.nn as nn
from network.aggregate_net import name2agg_net
from network.dist_decoder import name2dist_decoder
from network.init_net import name2init_net
from network.ops import ResUNetLight
from network.vis_encoder import name2vis_encoder
from network.render_ops import *
from utils.field_utils import TSDF_SAMPLE_POINTS
class NeuralRayRenderer(nn.Module):
base_cfg={
'vis_encoder_type': 'default',
'vis_encoder_cfg': {},
'dist_decoder_type': 'mixture_logistics',
'dist_decoder_cfg': {},
'agg_net_type': 'default',
'agg_net_cfg': {},
'use_hierarchical_sampling': False,
'fine_agg_net_cfg': {},
'fine_dist_decoder_cfg': {},
'fine_depth_sample_num': 64,
'fine_depth_use_all': False,
'ray_batch_num': 2048,
'depth_sample_num': 64,
'alpha_value_ground_state': -15,
'use_dr_prediction': False,
'use_nr_color_for_dr': False,
'use_self_hit_prob': False,
'use_ray_mask': True,
'ray_mask_view_num': 2,
'ray_mask_point_num': 8,
'render_depth': False,
'disable_view_dir': False,
'render_rgb': False,
'init_net_type': 'depth',
'init_net_cfg': {},
'depth_loss_coords_num': 8192,
}
def __init__(self,cfg):
super().__init__()
self.cfg = {**self.base_cfg, **cfg}
self.vis_encoder = name2vis_encoder[self.cfg['vis_encoder_type']](self.cfg['vis_encoder_cfg'])
self.dist_decoder = name2dist_decoder[self.cfg['dist_decoder_type']](self.cfg['dist_decoder_cfg'])
self.image_encoder = ResUNetLight(3, [1,2,6,4], 32, inplanes=16)
self.init_net=name2init_net[self.cfg['init_net_type']](self.cfg['init_net_cfg'])
self.agg_net = name2agg_net[self.cfg['agg_net_type']](self.cfg['agg_net_cfg'])
if self.cfg['use_hierarchical_sampling']:
self.fine_dist_decoder = name2dist_decoder[self.cfg['dist_decoder_type']](self.cfg['fine_dist_decoder_cfg'])
self.fine_agg_net = name2agg_net[self.cfg['agg_net_type']](self.cfg['fine_agg_net_cfg'])
self.use_sdf = self.cfg['agg_net_type'] in ['neus']
def predict_proj_ray_prob(self, prj_dict, ref_imgs_info, que_dists, is_fine):
rfn, qn, rn, dn, _ = prj_dict['mask'].shape
# decode ray prob
if is_fine:
prj_mean, prj_var, prj_vis, prj_aw = self.fine_dist_decoder(prj_dict['ray_feats'])
else:
prj_mean, prj_var, prj_vis, prj_aw = self.dist_decoder(prj_dict['ray_feats'])
alpha_values, visibility, hit_prob = self.dist_decoder.compute_prob(
prj_dict['depth'].squeeze(-1),que_dists.unsqueeze(0),prj_mean,prj_var,
prj_vis, prj_aw, True, ref_imgs_info['depth_range'])
# post process
prj_dict['alpha'] = alpha_values.reshape(rfn,qn,rn,dn,1) * prj_dict['mask'] + \
(1 - prj_dict['mask']) * self.cfg['alpha_value_ground_state']
prj_dict['vis'] = visibility.reshape(rfn,qn,rn,dn,1) * prj_dict['mask']
prj_dict['hit_prob'] = hit_prob.reshape(rfn,qn,rn,dn,1) * prj_dict['mask']
return prj_dict
def get_img_feats(self,ref_imgs_info, prj_dict):
rfn, _, h, w = ref_imgs_info['imgs'].shape
rfn, qn, rn, dn, _ = prj_dict['pts'].shape
img_feats = ref_imgs_info['img_feats']
prj_img_feats = interpolate_feature_map(img_feats, prj_dict['pts'].reshape(rfn, qn * rn * dn, 2),
prj_dict['mask'].reshape(rfn, qn * rn * dn), h, w,)
prj_dict['img_feats'] = prj_img_feats.reshape(rfn, qn, rn, dn, -1)
return prj_dict
def network_rendering(self, prj_dict, que_dir, que_pts, que_depth, is_fine, is_train, is_sdf=False, sdf_only=False):
net = self.fine_agg_net if is_fine else self.agg_net
que_dists = depth2dists(que_depth) if que_depth is not None else None
rendering_outputs = net(prj_dict, que_dir, que_pts, que_dists, is_train)
outputs = {}
if is_sdf:
alpha_values, outputs['sdf_values'], colors, outputs['sdf_gradient_error'], outputs['s'] = rendering_outputs
if sdf_only:
return outputs
else:
density, colors = rendering_outputs
alpha_values = 1.0 - torch.exp(-torch.relu(density))
outputs['alpha_values'] = alpha_values
outputs['colors_nr'] = colors
outputs['hit_prob_nr'] = hit_prob = alpha_values2hit_prob(alpha_values)
outputs['pixel_colors_nr'] = torch.sum(hit_prob.unsqueeze(-1)*colors,2)
return outputs
def render_by_depth(self, que_depth, que_imgs_info, ref_imgs_info, is_train, is_fine):
ref_imgs_info = ref_imgs_info.copy()
que_imgs_info = que_imgs_info.copy()
que_dists = depth2inv_dists(que_depth,que_imgs_info['depth_range'])
# generate points with query depth
que_pts, que_dir = depth2points(que_imgs_info, que_depth)
if self.cfg['disable_view_dir']:
que_dir = None
prj_dict = project_points_dict(ref_imgs_info, que_pts)
prj_dict = self.predict_proj_ray_prob(prj_dict, ref_imgs_info, que_dists, is_fine)
prj_dict = self.get_img_feats(ref_imgs_info, prj_dict)
outputs = self.network_rendering(prj_dict, que_dir, que_pts, que_depth, is_fine, is_train, is_sdf=self.use_sdf)
if 'imgs' in que_imgs_info:
outputs['pixel_colors_gt'] = interpolate_feats(
que_imgs_info['imgs'], que_imgs_info['coords'], align_corners=True)
if self.cfg['use_ray_mask']:
outputs['ray_mask'] = torch.sum(prj_dict['mask'].int(),0)>self.cfg['ray_mask_view_num'] # qn,rn,dn,1
outputs['ray_mask'] = torch.sum(outputs['ray_mask'],2)>self.cfg['ray_mask_point_num'] # qn,rn
outputs['ray_mask'] = outputs['ray_mask'][...,0]
if self.cfg['render_depth']:
# qn,rn,dn
outputs['render_depth'] = torch.sum(outputs['hit_prob_nr'] * que_depth, -1) # qn,rn
#outputs['render_depth_dr'] = torch.sum(hit_prob_dr * que_depth, -1) # qn,rn
return outputs
def fine_render_impl(self, coarse_render_info, que_imgs_info, ref_imgs_info, is_train):
fine_depth = sample_fine_depth(coarse_render_info['depth'], coarse_render_info['hit_prob'].detach(),
que_imgs_info['depth_range'], self.cfg['fine_depth_sample_num'], is_train)
# qn, rn, fdn+dn
if self.cfg['fine_depth_use_all']:
que_depth = torch.sort(torch.cat([coarse_render_info['depth'], fine_depth], -1), -1)[0]
else:
que_depth = torch.sort(fine_depth, -1)[0]
outputs = self.render_by_depth(que_depth, que_imgs_info, ref_imgs_info, is_train, True)
return outputs
def render_impl(self, que_imgs_info, ref_imgs_info, is_train):
# [qn,rn,dn]
# sample points along test ray at different depth
que_depth, _ = sample_depth(que_imgs_info['depth_range'], que_imgs_info['coords'], self.cfg['depth_sample_num'], False)
outputs = self.render_by_depth(que_depth, que_imgs_info, ref_imgs_info, is_train, False)
if self.cfg['use_hierarchical_sampling']:
coarse_render_info= {'depth': que_depth, 'hit_prob': outputs['hit_prob_nr']}
fine_outputs = self.fine_render_impl(coarse_render_info, que_imgs_info, ref_imgs_info, is_train)
for k, v in fine_outputs.items():
outputs[k + "_fine"] = v
return outputs
def sample_volume(self, ref_imgs_info):
ref_imgs_info = ref_imgs_info.copy()
res = self.cfg['volume_resolution']
que_pts = ( torch.from_numpy(TSDF_SAMPLE_POINTS).to(ref_imgs_info['imgs'].device) +
torch.tensor(ref_imgs_info['bbox3d'][0], device=ref_imgs_info['imgs'].device)
).reshape(1, res * res, res, 3)
que_pts = torch.flip(que_pts, (2,))
prj_dict = project_points_dict(ref_imgs_info, que_pts)
prj_dict = self.get_img_feats(ref_imgs_info, prj_dict)
valid_ratio = torch.sum(prj_dict['mask'],dim=(1,2,3,4)) / np.prod(list(prj_dict['mask'].shape)[1:])
if torch.mean(valid_ratio) < 0.5:
print("!! too low ratio", valid_ratio)
prj_dict = self.predict_proj_ray_prob(prj_dict, ref_imgs_info, torch.empty(0), False)
que_dir = torch.tensor([0,0,1], device=que_pts.device).reshape(1,1,1,3).repeat(1, res * res, res, 1) if not self.cfg['disable_view_dir'] else None
feat_list = []
mode = self.cfg['volume_type']
if 'image' in mode:
image_feat = torch.cat([prj_dict['rgb'], prj_dict['img_feats']], dim=-1)
mean = torch.mean(image_feat, dim=-1)
var = torch.var(image_feat, dim=-1)
feat_list.append(torch.cat([image_feat, mean, var], dim=-1).reshape(1, res, res, res, -1).permute(1,-1))
if 'alpha' in mode:
outputs = self.network_rendering(prj_dict, que_dir, que_pts, None, False, False)
feat_list.append(outputs['alpha_values'].reshape(1, 1, res, res, res))
if 'sdf' in mode:
outputs = self.network_rendering(prj_dict, que_dir, que_pts, None, False, False, is_sdf=self.use_sdf, sdf_only=True)
feat_list.append(outputs['sdf_values'].reshape(1, 1, res, res, res))
feat = torch.cat(feat_list, dim=1)
feat = torch.flip(feat, (-1,)) # we sample from top to down, so we need to flip here
return feat
def render(self, que_imgs_info, ref_imgs_info, is_train):
render_info_all = {}
ray_batch_num = self.cfg["ray_batch_num"]
coords = que_imgs_info['coords']
ray_num = coords.shape[1]
for ray_id in range(0,ray_num,ray_batch_num):
que_imgs_info['coords']=coords[:,ray_id:ray_id+ray_batch_num]
render_info = self.render_impl(que_imgs_info,ref_imgs_info,is_train)
output_keys = [k for k in render_info.keys()]
for k in output_keys:
v = render_info[k]
if k not in render_info_all:
render_info_all[k]=[]
render_info_all[k].append(v)
for k, v in render_info_all.items():
render_info_all[k]=torch.cat(v,1)
return render_info_all
def gen_depth_loss_coords(self,h,w,device):
coords = torch.stack(torch.meshgrid(torch.arange(h), torch.arange(w)), -1).reshape(-1, 2).to(device)
num = self.cfg['depth_loss_coords_num']
idxs = torch.randperm(coords.shape[0])
idxs = idxs[:num]
coords = coords[idxs]
return coords
def predict_mean_for_depth_loss(self, ref_imgs_info):
ray_feats = ref_imgs_info['ray_feats'] # rfn,f,h',w'
ref_imgs = ref_imgs_info['imgs'] # rfn,3,h,w
rfn, _, h, w = ref_imgs.shape
coords = self.gen_depth_loss_coords(h,w,ref_imgs.device) # pn,2
coords = coords.unsqueeze(0).repeat(rfn,1,1) # rfn,pn,2
batch_num = self.cfg['depth_loss_coords_num']
pn = coords.shape[1]
coords_dist_mean, coords_dist_mean_2, coords_dist_mean_fine, coords_dist_mean_fine_2 = [], [], [], []
for ci in range(0, pn, batch_num):
coords_ = coords[:,ci:ci+batch_num]
mask_ = torch.ones(coords_.shape[:2], dtype=torch.float32, device=ref_imgs.device)
coords_ray_feats_ = interpolate_feature_map(ray_feats, coords_, mask_, h, w) # rfn,pn,f
coords_dist_mean_ = self.dist_decoder.predict_mean(coords_ray_feats_) # rfn,pn
coords_dist_mean_2.append(coords_dist_mean_[..., 1])
coords_dist_mean_ = coords_dist_mean_[..., 0]
coords_dist_mean.append(coords_dist_mean_)
if self.cfg['use_hierarchical_sampling']:
coords_dist_mean_fine_ = self.fine_dist_decoder.predict_mean(coords_ray_feats_)
coords_dist_mean_fine_2.append(coords_dist_mean_fine_[..., 1])
coords_dist_mean_fine_ = coords_dist_mean_fine_[..., 0] # use 0 for depth supervision
coords_dist_mean_fine.append(coords_dist_mean_fine_)
coords_dist_mean = torch.cat(coords_dist_mean, 1)
outputs = {'depth_mean': coords_dist_mean, 'depth_coords': coords}
if len(coords_dist_mean_2)>0:
coords_dist_mean_2 = torch.cat(coords_dist_mean_2, 1)
outputs['depth_mean_2'] = coords_dist_mean_2
if self.cfg['use_hierarchical_sampling']:
coords_dist_mean_fine = torch.cat(coords_dist_mean_fine, 1)
outputs['depth_mean_fine'] = coords_dist_mean_fine
if len(coords_dist_mean_fine_2)>0:
coords_dist_mean_fine_2 = torch.cat(coords_dist_mean_fine_2, 1)
outputs['depth_mean_fine_2'] = coords_dist_mean_fine_2
return outputs
def forward(self,data):
ref_imgs_info = data['ref_imgs_info'].copy()
que_imgs_info = data['que_imgs_info'].copy()
is_train = 'eval' not in data
src_imgs_info = data['src_imgs_info'].copy() if 'src_imgs_info' in data else None
# extract image feature
ref_imgs_info['img_feats'] = self.image_encoder(ref_imgs_info['imgs'])
# calc visibility feature map of each view from mvs
ref_imgs_info['ray_feats'] = self.init_net(ref_imgs_info, src_imgs_info, is_train)
# refine visibity feature along with image feature
ref_imgs_info['ray_feats'] = self.vis_encoder(ref_imgs_info['ray_feats'], ref_imgs_info['img_feats'])
render_outputs = {}
if self.cfg['render_rgb']:
render_outputs = self.render(que_imgs_info, ref_imgs_info, is_train)
if self.cfg['sample_volume']:
render_outputs['volume'] = self.sample_volume(ref_imgs_info)
if (self.cfg['use_depth_loss'] and 'true_depth' in ref_imgs_info) or (not is_train):
render_outputs.update(self.predict_mean_for_depth_loss(ref_imgs_info))
return render_outputs
class GraspNeRF(nn.Module):
default_cfg_vgn={
'nr_initial_training_steps': 0,
'freeze_nr_after_init': False
}
def __init__(self, cfg):
super().__init__()
self.cfg={**self.default_cfg_vgn,**cfg}
from gd.networks import get_network
self.nr_net = NeuralRayRenderer(self.cfg)
self.vgn_net = get_network("conv")
def select(self, out, index):
qual_out, rot_out, width_out = out
batch_index = torch.arange(qual_out.shape[0])
label = qual_out[batch_index, :, index[:, 0], index[:, 1], index[:, 2]].squeeze()
rot = rot_out[batch_index, :, index[:, 0], index[:, 1], index[:, 2]]
width = width_out[batch_index, :, index[:, 0], index[:, 1], index[:, 2]].squeeze()
return (label, rot, width)
def forward(self, data):
if data['step'] < self.cfg['nr_initial_training_steps']:
render_outputs = super().forward(data)
with torch.no_grad():
vgn_pred = self.vgn_net(render_outputs['volume'])
elif self.cfg['freeze_nr_after_init']:
with torch.no_grad():
render_outputs = super().forward(data)
vgn_pred = self.vgn_net(render_outputs['volume'])
else:
render_outputs = self.nr_net(data)
vgn_pred = self.vgn_net(render_outputs['volume'])
if 'full_vol' not in data:
render_outputs['vgn_pred'] = self.select(vgn_pred, data['grasp_info'][0])
else:
render_outputs['vgn_pred'] = vgn_pred
return render_outputs
name2network={
'grasp_nerf': GraspNeRF,
}
================================================
FILE: src/nr/network/vis_encoder.py
================================================
import torch.nn as nn
import torch
from network.ops import conv3x3, ResidualBlock, conv1x1
class DefaultVisEncoder(nn.Module):
default_cfg={}
def __init__(self, cfg):
super().__init__()
self.cfg={**self.default_cfg,**cfg}
norm_layer = lambda dim: nn.InstanceNorm2d(dim,track_running_stats=False,affine=True)
self.out_conv=nn.Sequential(
conv3x3(64, 32),
ResidualBlock(32, 32, norm_layer=norm_layer),
ResidualBlock(32, 32, norm_layer=norm_layer),
conv1x1(32, 32),
)
def forward(self, ray_feats, imgs_feats):
feats = self.out_conv(torch.cat([imgs_feats, ray_feats],1))
return feats
name2vis_encoder={
'default': DefaultVisEncoder,
}
================================================
FILE: src/nr/run_training.py
================================================
import argparse
from train.trainer import Trainer
from utils.base_utils import load_cfg
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, default='configs/train/gen/neuray_gen_depth_train.yaml')
flags = parser.parse_args()
trainer = Trainer(load_cfg(flags.cfg))
trainer.run()
================================================
FILE: src/nr/train/lr_common_manager.py
================================================
import abc
class LearningRateManager(abc.ABC):
@staticmethod
def set_lr_for_all(optimizer, lr):
for param_group in optimizer.param_groups:
param_group['lr'] = lr
def construct_optimizer(self, optimizer, network):
# may specify different lr for different parts
# use group to set learning rate
paras = network.parameters()
return optimizer(paras, lr=1e-3)
@abc.abstractmethod
def __call__(self, optimizer, step, *args, **kwargs):
pass
class ExpDecayLR(LearningRateManager):
def __init__(self,cfg):
self.lr_init=cfg['lr_init']
self.decay_step=cfg['decay_step']
self.decay_rate=cfg['decay_rate']
self.lr_min=1e-5
def __call__(self, optimizer, step, *args, **kwargs):
lr=max(self.lr_init*(self.decay_rate**(step//self.decay_step)),self.lr_min)
self.set_lr_for_all(optimizer,lr)
return lr
class ExpDecayLRRayFeats(ExpDecayLR):
def construct_optimizer(self, optimizer, network):
paras = network.parameters()
return optimizer([para for para in paras] + network.ray_feats, lr=1e-3)
class WarmUpExpDecayLR(LearningRateManager):
def __init__(self, cfg):
self.lr_warm=cfg['lr_warm']
self.warm_step=cfg['warm_step']
self.lr_init=cfg['lr_init']
self.decay_step=cfg['decay_step']
self.decay_rate=cfg['decay_rate']
self.lr_min=1e-5
def __call__(self, optimizer, step, *args, **kwargs):
if step 0:
msg += '{} {:.5f} '.format(k.split('/')[-1], np.mean(v))
self.writer.add_scalar(k, np.mean(v), step)
self.data[k] = []
print(msg)
with open(self.rec_fn, 'a') as f:
f.write(msg + '\n')
def rec_msg(self, msg):
print(msg)
with open(self.rec_fn, 'a') as f:
f.write(msg + '\n')
class Logger:
def __init__(self, log_dir):
self.log_dir=log_dir
self.data = OrderedDict()
self.writer = SummaryWriter(log_dir=log_dir + "/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S"))
def log(self,data, prefix='train',step=None,verbose=False):
msg=f'{prefix} '
for k, v in data.items():
msg += f'{k} {v:.5f} '
self.writer.add_scalar(f'{prefix}/{k}',v,step)
if verbose:
print(msg)
with open(os.path.join(self.log_dir,f'{prefix}.txt'), 'a') as f:
f.write(msg + '\n')
def print_shape(obj):
if type(obj) == list or type(obj) == tuple:
shapes = [item.shape for item in obj]
print(shapes)
else:
print(obj.shape)
def overwrite_configs(cfg_base: dict, cfg: dict):
keysNotinBase = []
for key in cfg.keys():
if key in cfg_base.keys():
cfg_base[key] = cfg[key]
else:
keysNotinBase.append(key)
cfg_base.update({key: cfg[key]})
if len(keysNotinBase) != 0:
print('==== WARNING: These keys are not set in DEFAULT_BASE_CONFIG... ====')
print(keysNotinBase)
return cfg_base
def to_cuda(data):
if type(data)==list:
results = []
for i, item in enumerate(data):
results.append(to_cuda(item))
return results
elif type(data)==dict:
results={}
for k,v in data.items():
results[k]=to_cuda(v)
return results
elif type(data).__name__ == "Tensor":
return data.cuda()
else:
return data
def dim_extend(data_list):
results = []
for i, tensor in enumerate(data_list):
results.append(tensor[None,...])
return results
class MultiGPUWrapper(nn.Module):
def __init__(self,network,losses):
super().__init__()
self.network=network
self.losses=losses
def forward(self, data_gt):
results={}
data_pr=self.network(data_gt)
results.update(data_pr)
for loss in self.losses:
results.update(loss(data_pr,data_gt,data_gt['step']))
return results
class DummyLoss:
def __init__(self,losses):
self.keys=[]
for loss in losses:
self.keys+=loss.keys
def __call__(self, data_pr, data_gt, step):
return {key: data_pr[key] for key in self.keys}
================================================
FILE: src/nr/train/train_valid.py
================================================
import time
import torch
import numpy as np
from tqdm import tqdm
from network.metrics import name2key_metrics
from train.train_tools import to_cuda
class ValidationEvaluator:
def __init__(self,cfg):
self.key_metric_name=cfg['key_metric_name']
self.key_metric=name2key_metrics[self.key_metric_name]
def __call__(self, model, losses, eval_dataset, step, model_name, val_set_name=None):
if val_set_name is not None: model_name=f'{model_name}-{val_set_name}'
model.eval()
eval_results={}
begin=time.time()
for data_i, data in tqdm(enumerate(eval_dataset)):
data = to_cuda(data)
data['eval']=True
data['step']=step
with torch.no_grad():
outputs=model(data)
for loss in losses:
loss_results=loss(outputs, data, step, data_index=data_i, model_name=model_name, is_train=False)
for k,v in loss_results.items():
if type(v)==torch.Tensor:
v=v.detach().cpu().numpy()
if k in eval_results:
eval_results[k].append(v)
else:
eval_results[k]=[v]
for k,v in eval_results.items():
eval_results[k]=np.concatenate(v,axis=0)
key_metric_val=self.key_metric(eval_results)
if key_metric_val != 1e6:
eval_results[self.key_metric_name + '_all'] = eval_results[self.key_metric_name]
eval_results[self.key_metric_name]=key_metric_val
print('eval cost {} s'.format(time.time()-begin))
return eval_results, key_metric_val
================================================
FILE: src/nr/train/trainer.py
================================================
import os
import random
import torch
import numpy as np
from torch.nn import DataParallel
from torch.optim import Adam, SGD
from torch.utils.data import DataLoader
from tqdm import tqdm
from dataset.name2dataset import name2dataset
from network.loss import name2loss
from network.renderer import name2network
from train.lr_common_manager import name2lr_manager
from network.metrics import name2metrics
from train.train_tools import to_cuda, Logger, reset_learning_rate, MultiGPUWrapper, DummyLoss
from train.train_valid import ValidationEvaluator
from utils.dataset_utils import simple_collate_fn, dummy_collate_fn
from asset import vgn_val_scene_names
class Trainer:
default_cfg={
"optimizer_type": 'adam',
"multi_gpus": False,
"lr_type": "exp_decay",
"lr_cfg":{
"lr_init": 1.0e-4,
"decay_step": 100000,
"decay_rate": 0.5,
},
"total_step": 300000,
"train_log_step": 20,
"val_interval": 10000,
"save_interval": 1000,
"worker_num": 8,
"fix_seed": False
}
def _init_dataset(self):
self.train_set=name2dataset[self.cfg['train_dataset_type']](self.cfg['train_dataset_cfg'], True)
self.train_set=DataLoader(self.train_set,1,True,num_workers=self.cfg['worker_num'],collate_fn=dummy_collate_fn)
print(f'train set len {len(self.train_set)}')
self.val_set_list, self.val_set_names = [], []
for val_set_cfg in self.cfg['val_set_list']:
name, val_type, val_cfg = val_set_cfg['name'], val_set_cfg['type'], val_set_cfg['cfg']
if 'val_scene_num' in val_set_cfg:
num = val_set_cfg['val_scene_num']
num = len(vgn_val_scene_names) if num == -1 else num
names, val_types = [name] * num, [val_type] * num
val_cfgs = []
for i in range(num):
val_cfgs.append({**val_cfg, **{'val_database_name': vgn_val_scene_names[i]}})
else:
names, val_types, val_cfgs = [name], [val_type], [val_cfg]
for name, val_type, val_cfg in zip(names, val_types, val_cfgs):
val_set = name2dataset[val_type](val_cfg, False)
val_set = DataLoader(val_set,1,False,num_workers=self.cfg['worker_num'],collate_fn=dummy_collate_fn)
self.val_set_list.append(val_set)
self.val_set_names.append(name)
print(f"val set num: {len(self.val_set_list)}")
def _init_network(self):
self.network=name2network[self.cfg['network']](self.cfg).cuda()
# loss
self.val_losses = []
for loss_name in self.cfg['loss']:
self.val_losses.append(name2loss[loss_name](self.cfg))
self.val_metrics = []
# metrics
for metric_name in self.cfg['val_metric']:
if metric_name in name2metrics:
self.val_metrics.append(name2metrics[metric_name](self.cfg))
else:
self.val_metrics.append(name2loss[metric_name](self.cfg))
# we do not support multi gpu training for NeuRay
if self.cfg['multi_gpus']:
raise NotImplementedError
# make multi gpu network
# self.train_network=DataParallel(MultiGPUWrapper(self.network,self.val_losses))
# self.train_losses=[DummyLoss(self.val_losses)]
else:
self.train_network=self.network
self.train_losses=self.val_losses
if self.cfg['optimizer_type']=='adam':
self.optimizer = Adam
elif self.cfg['optimizer_type']=='sgd':
self.optimizer = SGD
else:
raise NotImplementedError
self.val_evaluator=ValidationEvaluator(self.cfg)
self.lr_manager=name2lr_manager[self.cfg['lr_type']](self.cfg['lr_cfg'])
self.optimizer=self.lr_manager.construct_optimizer(self.optimizer,self.network)
def __init__(self,cfg):
self.cfg={**self.default_cfg,**cfg}
if self.cfg['fix_seed']:
seed = 0
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
torch.cuda.manual_seed_all(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
print("fix seed")
self.model_name=cfg['name']
self.model_dir=os.path.join('data/model',cfg['group_name'],cfg['name'])
if not os.path.exists(self.model_dir): os.makedirs(self.model_dir)
self.pth_fn=os.path.join(self.model_dir,'model.pth')
self.best_pth_fn=os.path.join(self.model_dir,'model_best.pth')
assert self.cfg["key_metric_prefer"] in ['higher', 'lower']
self.better = lambda x, y: x > y if self.cfg["key_metric_prefer"] == 'higher' else x < y
def run(self):
self._init_dataset()
self._init_network()
self._init_logger()
best_para,start_step=self._load_model()
if self.cfg["key_metric_prefer"] == 'lower' and start_step == 0:
best_para = 1e6
train_iter=iter(self.train_set)
pbar=tqdm(total=self.cfg['total_step'],bar_format='{r_bar}')
pbar.update(start_step)
for step in range(start_step,self.cfg['total_step']):
try:
train_data = next(train_iter)
except StopIteration:
self.train_set.dataset.reset()
train_iter = iter(self.train_set)
train_data = next(train_iter)
if not self.cfg['multi_gpus']:
train_data = to_cuda(train_data)
train_data['step']=step
self.train_network.train()
self.network.train()
lr = self.lr_manager(self.optimizer, step)
self.optimizer.zero_grad()
self.train_network.zero_grad()
log_info={}
outputs=self.train_network(train_data)
for loss in self.train_losses:
loss_results = loss(outputs,train_data,step)
for k,v in loss_results.items():
log_info[k]=v
loss=0
for k,v in log_info.items():
if k.startswith('loss'):
loss=loss+torch.mean(v)
loss.backward()
self.optimizer.step()
if ((step+1) % self.cfg['train_log_step']) == 0:
self._log_data(log_info,step+1,'train')
if step==0 or (step+1)%self.cfg['val_interval']==0 or (step+1)==self.cfg['total_step']:
torch.cuda.empty_cache()
val_results={}
val_para = 0
for vi, val_set in enumerate(self.val_set_list):
val_results_cur, val_para_cur = self.val_evaluator(
self.network, self.val_losses + self.val_metrics, val_set, step,
self.model_name, val_set_name=self.val_set_names[vi])
for k,v in val_results_cur.items():
key = f'{self.val_set_names[vi]}-{k}'
if not key in val_results:
val_results[key] = v
else:
val_results[key] += v
val_para += val_para_cur
# average all items
for k,v in val_results.items():
val_results[k] /= len(self.val_set_list)
val_para /= len(self.val_set_list)
if step and self.better(val_para, best_para): # do not save the first step
print(f'New best model {self.cfg["key_metric_name"]}: {val_para:.5f} previous {best_para:.5f}')
best_para=val_para
self._save_model(step+1,best_para,self.best_pth_fn)
self._log_data(val_results,step+1,'val')
del val_results, val_para, val_para_cur, val_results_cur
if (step+1)%self.cfg['save_interval']==0:
self._save_model(step+1,best_para)
pbar.set_postfix(loss=float(loss.detach().cpu().numpy()),lr=lr)
pbar.update(1)
del loss, log_info
pbar.close()
def _load_model(self):
best_para,start_step=0,0
if os.path.exists(self.pth_fn):
checkpoint=torch.load(self.pth_fn)
best_para = checkpoint['best_para']
start_step = checkpoint['step']
self.network.load_state_dict(checkpoint['network_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
print(f'==> resuming from the latest {self.pth_fn} of step {start_step} with best metric {best_para}')
return best_para, start_step
def _save_model(self, step, best_para, save_fn=None):
save_fn = self.pth_fn if save_fn is None else save_fn
torch.save({
'step':step,
'best_para':best_para,
'network_state_dict': self.network.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict(),
},save_fn)
def _init_logger(self):
self.logger = Logger(self.model_dir)
def _log_data(self,results,step,prefix='train',verbose=False):
log_results={}
for k, v in results.items():
if isinstance(v,float) or np.isscalar(v):
log_results[k] = v
elif type(v)==np.ndarray:
log_results[k]=np.mean(v)
else:
log_results[k]=np.mean(v.detach().cpu().numpy())
self.logger.log(log_results,prefix,step,verbose)
================================================
FILE: src/nr/utils/base_utils.py
================================================
import math
import os
import cv2
import h5py
import torch
import numpy as np
import pickle
import yaml
from plyfile import PlyData
from skimage.io import imread
#######################io#########################################
from transforms3d.axangles import mat2axangle
from transforms3d.euler import euler2mat
def read_pickle(pkl_path):
with open(pkl_path, 'rb') as f:
return pickle.load(f)
def save_pickle(data, pkl_path):
os.system('mkdir -p {}'.format(os.path.dirname(pkl_path)))
with open(pkl_path, 'wb') as f:
pickle.dump(data, f)
#####################depth and image###############################
def mask_zbuffer_to_pts(mask, zbuffer, K):
ys, xs = np.nonzero(mask)
zbuffer = zbuffer[ys, xs]
u, v, f = K[0, 2], K[1, 2], K[0, 0]
depth = zbuffer / np.sqrt((xs - u + 0.5) ** 2 + (ys - v + 0.5) ** 2 + f ** 2) * f
pts = np.asarray([xs, ys, depth], np.float32).transpose()
pts[:, :2] *= pts[:, 2:]
return np.dot(pts, np.linalg.inv(K).transpose())
def mask_depth_to_pts(mask, depth, K, rgb=None):
hs, ws = np.nonzero(mask)
depth = depth[hs, ws]
pts = np.asarray([ws, hs, depth], np.float32).transpose()
pts[:, :2] *= pts[:, 2:]
if rgb is not None:
return np.dot(pts, np.linalg.inv(K).transpose()), rgb[hs, ws]
else:
return np.dot(pts, np.linalg.inv(K).transpose())
def read_render_zbuffer(dpt_pth, max_depth, min_depth):
zbuffer = imread(dpt_pth)
mask = (zbuffer > 0) & (zbuffer < 5000)
zbuffer = zbuffer.astype(np.float64) / 2 ** 16 * (max_depth - min_depth) + min_depth
return mask, zbuffer
def zbuffer_to_depth(zbuffer, K):
u, v, f = K[0, 2], K[1, 2], K[0, 0]
x = np.arange(zbuffer.shape[1])
y = np.arange(zbuffer.shape[0])
x, y = np.meshgrid(x, y)
x = np.reshape(x, [-1, 1])
y = np.reshape(y, [-1, 1])
depth = np.reshape(zbuffer, [-1, 1])
depth = depth / np.sqrt((x - u + 0.5) ** 2 + (y - v + 0.5) ** 2 + f ** 2) * f
return np.reshape(depth, zbuffer.shape)
def project_points(pts, RT, K):
pts = np.matmul(pts, RT[:, :3].transpose()) + RT[:, 3:].transpose()
pts = np.matmul(pts, K.transpose())
dpt = pts[:, 2]
mask0 = (np.abs(dpt) < 1e-4) & (np.abs(dpt) > 0)
if np.sum(mask0) > 0: dpt[mask0] = 1e-4
mask1 = (np.abs(dpt) > -1e-4) & (np.abs(dpt) < 0)
if np.sum(mask1) > 0: dpt[mask1] = -1e-4
pts2d = pts[:, :2] / dpt[:, None]
return pts2d, dpt
#######################image processing#############################
def grey_repeats(img_raw):
if len(img_raw.shape) == 2: img_raw = np.repeat(img_raw[:, :, None], 3, axis=2)
if img_raw.shape[2] > 3: img_raw = img_raw[:, :, :3]
return img_raw
def normalize_image(img, mask=None):
if mask is not None: img[np.logical_not(mask.astype(np.bool))] = 127
img = (img.transpose([2, 0, 1]).astype(np.float32) - 127.0) / 128.0
return torch.tensor(img, dtype=torch.float32)
def tensor_to_image(tensor):
return (tensor * 128 + 127).astype(np.uint8).transpose(1, 2, 0)
def equal_hist(img):
if len(img.shape) == 3:
img0 = cv2.equalizeHist(img[:, :, 0])
img1 = cv2.equalizeHist(img[:, :, 1])
img2 = cv2.equalizeHist(img[:, :, 2])
img = np.concatenate([img0[..., None], img1[..., None], img2[..., None]], 2)
else:
img = cv2.equalizeHist(img)
return img
def resize_large_image(img, resize_max):
h, w = img.shape[:2]
max_side = max(h, w)
if max_side > resize_max:
ratio = resize_max / max_side
if ratio <= 0.5: img = cv2.GaussianBlur(img, (5, 5), 1.5)
img = cv2.resize(img, (int(round(ratio * w)), int(round(ratio * h))), interpolation=cv2.INTER_LINEAR)
return img, ratio
else:
return img, 1.0
def downsample_gaussian_blur(img, ratio):
sigma = (1 / ratio) / 3
# ksize=np.ceil(2*sigma)
ksize = int(np.ceil(((sigma - 0.8) / 0.3 + 1) * 2 + 1))
ksize = ksize + 1 if ksize % 2 == 0 else ksize
img = cv2.GaussianBlur(img, (ksize, ksize), sigma, borderType=cv2.BORDER_REFLECT101)
return img
def resize_small_image(img, resize_min):
h, w = img.shape[:2]
min_side = min(h, w)
if min_side < resize_min:
ratio = resize_min / min_side
img = cv2.resize(img, (int(round(ratio * w)), int(round(ratio * h))), interpolation=cv2.INTER_LINEAR)
return img, ratio
else:
return img, 1.0
############################geometry######################################
def round_coordinates(coord, h, w):
coord = np.round(coord).astype(np.int32)
coord[coord[:, 0] < 0, 0] = 0
coord[coord[:, 0] >= w, 0] = w - 1
coord[coord[:, 1] < 0, 1] = 0
coord[coord[:, 1] >= h, 1] = h - 1
return coord
def get_img_patch(img, pt, size):
if isinstance(size, list) or isinstance(size, tuple) or isinstance(size, np.ndarray):
size_h, size_w = size
else:
size_h, size_w = size, size
h, w = img.shape[:2]
x, y = pt.astype(np.int32)
xmin = max(0, x - size_w)
xmax = min(w - 1, x + size_w)
ymin = max(0, y - size_h)
ymax = min(h - 1, y + size_h)
patch = np.full([size_h * 2, size_w * 2, 3], 127, np.uint8)
patch[ymin - y + size_h:ymax - y + size_h, xmin - x + size_w:xmax - x + size_w] = img[ymin:ymax, xmin:xmax]
return patch
def perspective_transform(pts, H):
tpts = np.concatenate([pts, np.ones([pts.shape[0], 1])], 1) @ H.transpose()
tpts = tpts[:, :2] / np.abs(tpts[:, 2:]) # todo: why only abs? this one is correct
return tpts
def get_rot_m(angle):
return np.asarray([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]], np.float32) # rn+1,3,3
def get_rot_m_batch(angle):
return np.asarray([[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]], np.float32).transpose(
[2, 0, 1])
def compute_F(K1, K2, R, t):
"""
:param K1: [3,3]
:param K2: [3,3]
:param R: [3,3]
:param t: [3,1]
:return:
"""
A = K1 @ R.T @ t # [3,1]
C = np.asarray([[0, -A[2, 0], A[1, 0]],
[A[2, 0], 0, -A[0, 0]],
[-A[1, 0], A[0, 0], 0]])
F = (np.linalg.inv(K2)).T @ R @ K1.T @ C
return F
def compute_relative_transformation(Rt0, Rt1):
"""
x1=Rx0+t
:param Rt0: x0=R0x+t0
:param Rt1: x1=R1x+t1
:return:
R1R0.T(x0-t0)+t1
"""
R = Rt1[:, :3] @ Rt0[:, :3].T
t = Rt1[:, 3] - R @ Rt0[:, 3]
return np.concatenate([R, t[:, None]], 1)
def compute_angle(rotation_diff):
trace = np.trace(rotation_diff)
trace = trace if trace <= 3 else 3
angular_distance = np.rad2deg(np.arccos((trace - 1.) / 2.))
return angular_distance
def load_h5(filename):
dict_to_load = {}
with h5py.File(filename, 'r') as f:
keys = [key for key in f.keys()]
for key in keys:
dict_to_load[key] = f[key][()] # .value
return dict_to_load
def save_h5(dict_to_save, filename):
with h5py.File(filename, 'w') as f:
for key in dict_to_save:
f.create_dataset(key, data=dict_to_save[key])
def pts_to_hpts(pts):
return np.concatenate([pts, np.ones([pts.shape[0], 1])], 1)
def hpts_to_pts(hpts):
return hpts[:, :-1] / hpts[:, -1:]
def np_skew_symmetric(v):
M = np.asarray([
[0, -v[2], v[1], ],
[v[2], 0, -v[0], ],
[-v[1], v[0], 0, ],
])
return M
def point_line_dist(hpts, lines):
"""
:param hpts: n,3 or n,2
:param lines: n,3
:return:
"""
if hpts.shape[1] == 2:
hpts = np.concatenate([hpts, np.ones([hpts.shape[0], 1])], 1)
return np.abs(np.sum(hpts * lines, 1)) / np.linalg.norm(lines[:, :2], 2, 1)
def epipolar_distance(x0, x1, F):
"""
:param x0: [n,2]
:param x1: [n,2]
:param F: [3,3]
:return:
"""
hkps0 = np.concatenate([x0, np.ones([x0.shape[0], 1])], 1)
hkps1 = np.concatenate([x1, np.ones([x1.shape[0], 1])], 1)
lines1 = hkps0 @ F.T
lines0 = hkps1 @ F
dist10 = point_line_dist(hkps0, lines0)
dist01 = point_line_dist(hkps1, lines1)
return dist10, dist01
def epipolar_distance_mean(x0, x1, F):
return np.mean(np.stack(epipolar_distance(x0, x1, F), 1), 1)
def compute_dR_dt(R0, t0, R1, t1):
# Compute dR, dt
dR = np.dot(R1, R0.T)
dt = t1 - np.dot(dR, t0)
return dR, dt
def compute_precision_recall_np(pr, gt, eps=1e-5):
tp = np.sum(gt & pr)
fp = np.sum((~gt) & pr)
fn = np.sum(gt & (~pr))
precision = (tp + eps) / (fp + tp + eps)
recall = (tp + eps) / (tp + fn + eps)
if precision < 1e-3 or recall < 1e-3:
f1 = 0.0
else:
f1 = (2 * precision * recall + eps) / (precision + recall + eps)
return precision, recall, f1
def load_cfg(path):
with open(path, 'r') as f:
return yaml.load(f, Loader=yaml.FullLoader)
def get_stem(path, suffix_len=5):
return os.path.basename(path)[:-suffix_len]
def load_component(component_func, component_cfg_fn):
component_cfg = load_cfg(component_cfg_fn)
return component_func[component_cfg['type']](component_cfg)
def interpolate_image_points(img, pts, interpolation=cv2.INTER_LINEAR):
# img [h,w,k] pts [n,2]
if len(pts) < 32767:
pts = pts.astype(np.float32)
return cv2.remap(img, pts[:, None, 0], pts[:, None, 1], borderMode=cv2.BORDER_CONSTANT, borderValue=0,
interpolation=interpolation)[:, 0]
# pn=len(pts)
# sl=int(np.ceil(np.sqrt(pn)))
# tmp_img=np.zeros([sl*sl,2],np.float32)
# tmp_img[:pn]=pts
# tmp_img=tmp_img.reshape([sl,sl,2])
# tmp_img=cv2.remap(img,tmp_img[:,:,0],tmp_img[:,:,1],borderMode=cv2.BORDER_CONSTANT,borderValue=0,interpolation=interpolation)
# return tmp_img.flatten()[:pn]
else:
results = []
for k in range(0, len(pts), 30000):
results.append(interpolate_image_points(img, pts[k:k + 30000], interpolation))
return np.concatenate(results, 0)
def transform_points_Rt(pts, R, t):
t = t.flatten()
return pts @ R.T + t[None, :]
def transform_points_pose(pts, pose):
R, t = pose[:, :3], pose[:, 3]
return pts @ R.T + t[None, :]
def quaternion_from_matrix(matrix, isprecise=False):
'''Return quaternion from rotation matrix.
If isprecise is True, the input matrix is assumed to be a precise rotation
matrix and a faster algorithm is used.
>>> q = quaternion_from_matrix(numpy.identity(4), True)
>>> numpy.allclose(q, [1, 0, 0, 0])
True
>>> q = quaternion_from_matrix(numpy.diag([1, -1, -1, 1]))
>>> numpy.allclose(q, [0, 1, 0, 0]) or numpy.allclose(q, [0, -1, 0, 0])
True
>>> R = rotation_matrix(0.123, (1, 2, 3))
>>> q = quaternion_from_matrix(R, True)
>>> numpy.allclose(q, [0.9981095, 0.0164262, 0.0328524, 0.0492786])
True
>>> R = [[-0.545, 0.797, 0.260, 0], [0.733, 0.603, -0.313, 0],
... [-0.407, 0.021, -0.913, 0], [0, 0, 0, 1]]
>>> q = quaternion_from_matrix(R)
>>> numpy.allclose(q, [0.19069, 0.43736, 0.87485, -0.083611])
True
>>> R = [[0.395, 0.362, 0.843, 0], [-0.626, 0.796, -0.056, 0],
... [-0.677, -0.498, 0.529, 0], [0, 0, 0, 1]]
>>> q = quaternion_from_matrix(R)
>>> numpy.allclose(q, [0.82336615, -0.13610694, 0.46344705, -0.29792603])
True
>>> R = random_rotation_matrix()
>>> q = quaternion_from_matrix(R)
>>> is_same_transform(R, quaternion_matrix(q))
True
>>> R = euler_matrix(0.0, 0.0, numpy.pi/2.0)
>>> numpy.allclose(quaternion_from_matrix(R, isprecise=False),
... quaternion_from_matrix(R, isprecise=True))
True
'''
M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]
if isprecise:
q = np.empty((4,))
t = np.trace(M)
if t > M[3, 3]:
q[0] = t
q[3] = M[1, 0] - M[0, 1]
q[2] = M[0, 2] - M[2, 0]
q[1] = M[2, 1] - M[1, 2]
else:
i, j, k = 1, 2, 3
if M[1, 1] > M[0, 0]:
i, j, k = 2, 3, 1
if M[2, 2] > M[i, i]:
i, j, k = 3, 1, 2
t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3]
q[i] = t
q[j] = M[i, j] + M[j, i]
q[k] = M[k, i] + M[i, k]
q[3] = M[k, j] - M[j, k]
q *= 0.5 / math.sqrt(t * M[3, 3])
else:
m00 = M[0, 0]
m01 = M[0, 1]
m02 = M[0, 2]
m10 = M[1, 0]
m11 = M[1, 1]
m12 = M[1, 2]
m20 = M[2, 0]
m21 = M[2, 1]
m22 = M[2, 2]
# symmetric matrix K
K = np.array([[m00 - m11 - m22, 0.0, 0.0, 0.0],
[m01 + m10, m11 - m00 - m22, 0.0, 0.0],
[m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],
[m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22]])
K /= 3.0
# quaternion is eigenvector of K that corresponds to largest eigenvalue
w, V = np.linalg.eigh(K)
q = V[[3, 0, 1, 2], np.argmax(w)]
if q[0] < 0.0:
np.negative(q, q)
return q
def compute_rotation_angle_diff(R_gt, R):
eps = 1e-15
q_gt = quaternion_from_matrix(R_gt)
q = quaternion_from_matrix(R)
q = q / (np.linalg.norm(q) + eps)
q_gt = q_gt / (np.linalg.norm(q_gt) + eps)
loss_q = np.maximum(eps, (1.0 - np.sum(q * q_gt) ** 2))
err_q = np.arccos(1 - 2 * loss_q)
return np.rad2deg(np.abs(err_q))
def compute_translation_angle_diff(t_gt, t):
eps = 1e-15
t = t / (np.linalg.norm(t) + eps)
t_gt = t_gt / (np.linalg.norm(t_gt) + eps)
loss_t = np.maximum(eps, (1.0 - np.sum(t * t_gt) ** 2))
err_t = np.arccos(np.sqrt(1 - loss_t))
return np.rad2deg(np.abs(err_t))
def bbox2corners(bbox):
return np.asarray([
[bbox[0], bbox[1]],
[bbox[0] + bbox[2], bbox[1]],
[bbox[0] + bbox[2], bbox[1] + bbox[3]],
[bbox[0], bbox[1] + bbox[3]],
])
def get_identity_pose():
return np.concatenate([np.identity(3), np.zeros([3, 1])], 1).astype(np.float32)
def angular_difference(R0, R1):
return np.rad2deg(mat2axangle(R0 @ R1.T)[1])
def load_ply_model(model_path):
ply = PlyData.read(model_path)
data = ply.elements[0].data
x = data['x']
y = data['y']
z = data['z']
return np.stack([x, y, z], axis=-1)
def color_map_forward(rgb):
return rgb.astype(np.float32) / 255
def color_map_backward(rgb):
rgb = rgb * 255
rgb = np.clip(rgb, a_min=0, a_max=255).astype(np.uint8)
return rgb
def rotate_image(rot, pose, K, img, mask):
if isinstance(rot, np.ndarray):
R = rot
else:
R = np.array([[np.cos(rot), -np.sin(rot), 0.0],
[np.sin(rot), np.cos(rot), 0.0],
[0, 0, 1]], dtype=np.float32)
# adjust pose
pose_adj = np.copy(pose)
pose_adj[:, :3] = R @ pose_adj[:, :3]
pose_adj[:, 3:] = R @ pose_adj[:, 3:]
# adjust image
transform = K @ R @ np.linalg.inv(K) # transform original
h, w, _ = img.shape
ys, xs = np.nonzero(mask)
coords = np.stack([xs, ys], -1).astype(np.float32)
coords_new = cv2.perspectiveTransform(coords[:, None, :], transform)[:, 0, :]
x_min, y_min = np.floor(np.min(coords_new, 0)).astype(np.int32)
x_max, y_max = np.ceil(np.max(coords_new, 0)).astype(np.int32)
th, tw = y_max - y_min, x_max - x_min
translation = np.identity(3)
translation[0, 2] = -x_min
translation[1, 2] = -y_min
K = translation @ K
transform = translation @ transform
img = cv2.warpPerspective(img, transform, (tw, th), flags=cv2.INTER_LINEAR)
return img, pose_adj, K
def resize_img(img, ratio):
# if ratio>=1.0: return img
h, w, _ = img.shape
hn, wn = int(np.round(h * ratio)), int(np.round(w * ratio))
img_out = cv2.resize(downsample_gaussian_blur(img, ratio), (wn, hn), cv2.INTER_LINEAR)
return img_out
def pad_img(img, padding_interval=8):
h, w = img.shape[:2]
hp = (padding_interval - (h % padding_interval)) % padding_interval
wp = (padding_interval - (w % padding_interval)) % padding_interval
if hp != 0 or wp != 0:
img = np.pad(img, ((0, hp), (0, wp), (0, 0)), 'edge')
return img
def pad_img_end(img, th, tw, padding_mode='edge', constant_values=0):
h, w = img.shape[:2]
hp = th - h
wp = tw - w
if hp != 0 or wp != 0:
if padding_mode == 'constant':
img = np.pad(img, ((0, hp), (0, wp), (0, 0)), padding_mode, constant_values=constant_values)
else:
img = np.pad(img, ((0, hp), (0, wp), (0, 0)), padding_mode)
return img
def pad_img_target(img, th, tw, K=np.eye(3), background_color=0):
h, w = img.shape[:2]
hp = th - h
wp = tw - w
if hp != 0 or wp != 0:
if len(img.shape) == 3:
img = np.pad(img, ((hp // 2, hp - hp // 2), (wp // 2, wp - wp // 2), (0, 0)), 'constant',
constant_values=background_color)
elif len(img.shape) == 2:
img = np.pad(img, ((hp // 2, hp - hp // 2), (wp // 2, wp - wp // 2)), 'constant',
constant_values=background_color)
else:
print(f'image shape unknown {img.shape}')
raise NotImplementedError
translation = np.identity(3)
translation[0, 2] = wp // 2
translation[1, 2] = hp // 2
K = translation @ K
return img, K
def get_coords_mask(que_mask, train_ray_num, foreground_ratio):
min_pos_num = int(train_ray_num * foreground_ratio)
y0, x0 = np.nonzero(que_mask)
y1, x1 = np.nonzero(~que_mask)
xy0 = np.stack([x0, y0], 1).astype(np.float32)
xy1 = np.stack([x1, y1], 1).astype(np.float32)
idx = np.arange(xy0.shape[0])
np.random.shuffle(idx)
xy0 = xy0[idx]
coords0 = xy0[:min_pos_num]
# still remain pixels
if min_pos_num < train_ray_num:
xy1 = np.concatenate([xy1, xy0[min_pos_num:]], 0)
idx = np.arange(xy1.shape[0])
np.random.shuffle(idx)
coords1 = xy1[idx[:(train_ray_num - min_pos_num)]]
coords = np.concatenate([coords0, coords1], 0)
else:
coords = coords0
return coords
def get_inverse_depth(depth_range, depth_num):
near, far = depth_range
interval = (1 / far - 1 / near) / (depth_num - 1)
ticks = np.arange(1, depth_num - 1)
ticks = 1 / (1 / near + ticks * interval)
return np.concatenate([np.asarray([near]).reshape([1]), ticks, np.asarray(far).reshape([1])], 0)
def pose_inverse(pose):
R = pose[:, :3].T
t = - R @ pose[:, 3:]
return np.concatenate([R, t], -1)
def pose_compose(pose0, pose1):
"""
apply pose0 first, then pose1
:param pose0:
:param pose1:
:return:
"""
t = pose1[:, :3] @ pose0[:, 3:] + pose1[:, 3:]
R = pose1[:, :3] @ pose0[:, :3]
return np.concatenate([R, t], 1)
def make_dir(dir):
if not os.path.exists(dir):
os.system(f'mkdir -p {dir}')
def to_cuda(data):
if type(data) == list:
results = []
for i, item in enumerate(data):
results.append(to_cuda(item))
return results
elif type(data) == dict:
results = {}
for k, v in data.items():
results[k] = to_cuda(v)
return results
elif type(data).__name__ == "Tensor" or type(data).__name__=="Parameter":
return data.cuda()
else:
return data
def to_cpu_numpy(data):
if type(data) == list:
results = []
for i, item in enumerate(data):
results.append(to_cpu_numpy(item))
return results
elif type(data) == dict:
results = {}
for k, v in data.items():
results[k] = to_cpu_numpy(v)
return results
elif type(data).__name__ == "Tensor" or type(data).__name__=="Parameter":
return data.detach().cpu().numpy()
else:
return data
def sample_fps_points(points, sample_num, init_center=True, index_model=False, init_first=False, init_first_index=0,
init_point=None):
sample_num = min(points.shape[0], sample_num)
output_index = []
if init_point is None:
if init_center:
init_point = np.mean(points, 0)
else:
if init_first:
init_index = init_first_index
else:
init_index = np.random.randint(0, points.shape[0])
init_point = points[init_index]
output_index.append(init_index)
output_points = [init_point]
cur_point = init_point
distance = np.full(points.shape[0], 1e8)
for k in range(sample_num - 1):
cur_distance = np.linalg.norm(cur_point[None, :] - points, 2, 1)
distance = np.min(np.stack([cur_distance, distance], 1), 1)
cur_index = np.argmax(distance)
cur_point = points[cur_index]
output_points.append(cur_point)
output_index.append(cur_index)
if index_model:
return np.asarray(output_index)
else:
return np.asarray(output_points)
def pnp(points_3d, points_2d, camera_matrix, method=cv2.SOLVEPNP_ITERATIVE):
dist_coeffs = np.zeros(shape=[8, 1], dtype='float64')
assert points_3d.shape[0] == points_2d.shape[0], 'points 3D and points 2D must have same number of vertices'
if method == cv2.SOLVEPNP_EPNP:
points_3d = np.expand_dims(points_3d, 0)
points_2d = np.expand_dims(points_2d, 0)
points_2d = np.ascontiguousarray(points_2d.astype(np.float64))
points_3d = np.ascontiguousarray(points_3d.astype(np.float64))
camera_matrix = camera_matrix.astype(np.float64)
_, R_exp, t = cv2.solvePnP(points_3d,
points_2d,
camera_matrix,
dist_coeffs,
flags=method)
R, _ = cv2.Rodrigues(R_exp)
return np.concatenate([R, t], axis=-1)
def triangulate(kps0, kps1, pose0, pose1, K0, K1):
kps0_ = hpts_to_pts(pts_to_hpts(kps0) @ np.linalg.inv(K0).T)
kps1_ = hpts_to_pts(pts_to_hpts(kps1) @ np.linalg.inv(K1).T)
pts3d = cv2.triangulatePoints(pose0.astype(np.float64), pose1.astype(np.float64),
kps0_.T.astype(np.float64), kps1_.T.astype(np.float64)).T
pts3d = pts3d[:, :3] / pts3d[:, 3:]
return pts3d
def transformation_compose_2d(trans0, trans1):
"""
@param trans0: [2,3]
@param trans1: [2,3]
@return: apply trans0 then trans1
"""
t1 = trans1[:, 2]
t0 = trans0[:, 2]
R1 = trans1[:, :2]
R0 = trans0[:, :2]
R = R1 @ R0
t = R1 @ t0 + t1
return np.concatenate([R, t[:, None]], 1)
def transformation_apply_2d(trans, points):
return points @ trans[:, :2].T + trans[:, 2:].T
def angle_to_rotation_2d(angle):
return np.asarray([[np.cos(angle), -np.sin(angle)],
[np.sin(angle), np.cos(angle)]])
def transformation_offset_2d(x, y):
return np.concatenate([np.eye(2), np.asarray([x, y])[:, None]], 1).astype(np.float32)
def transformation_scale_2d(scale):
return np.concatenate([np.diag([scale, scale]), np.zeros([2, 1])], 1).astype(np.float32)
def transformation_rotation_2d(ang):
return np.concatenate([angle_to_rotation_2d(ang), np.zeros([2, 1])], 1).astype(np.float32)
def look_at_rotation(point):
"""
@param point: point in normalized image coordinate not in pixels
@return: R
R @ x_raw -> x_lookat
"""
x, y = point
R1 = euler2mat(-np.arctan2(x, 1), 0, 0, 'syxz')
R2 = euler2mat(np.arctan2(y, 1), 0, 0, 'sxyz')
return R2 @ R1
def save_depth(fn, depth, max_val=1000):
import png
depth = np.clip(depth, a_min=0, a_max=max_val) / max_val * 65535
depth = depth.astype(np.uint16)
with open(fn, 'wb') as f:
writer = png.Writer(width=depth.shape[1], height=depth.shape[0], bitdepth=16, greyscale=True)
zgray2list = depth.tolist()
writer.write(f, zgray2list)
def compute_geodesic_distance_from_two_matrices(m1, m2):
batch = m1.shape[0]
m = torch.bmm(m1, m2.transpose(1, 2)) # batch*3*3
cos = (m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2] - 1) / 2
cos = torch.min(cos, torch.autograd.Variable(torch.ones(batch, device=m1.device)))
cos = torch.max(cos, torch.autograd.Variable(torch.ones(batch, device=m1.device)) * -1)
theta = torch.acos(cos)
theta = torch.min(theta, 2*np.pi - theta)
theta *= 180 / np.pi
return theta
def compute_rotation_matrix_from_quaternion_xyzw(quaternion, n_flag=True):
def normalize_vector(v):
batch = v.shape[0]
v_mag = torch.sqrt(v.pow(2).sum(1)) # batch
v_mag = torch.max(v_mag, torch.autograd.Variable(torch.FloatTensor([1e-8]).to(v.device)))
v_mag = v_mag.view(batch, 1).expand(batch, v.shape[1])
v = v / v_mag
return v
batch = quaternion.shape[0]
if n_flag:
quat = normalize_vector(quaternion)
else:
quat = quaternion
qx = quat[..., 0].view(batch, 1)
qy = quat[..., 1].view(batch, 1)
qz = quat[..., 2].view(batch, 1)
qw = quat[..., 3].view(batch, 1)
# Unit quaternion rotation matrices computatation
xx = qx * qx
yy = qy * qy
zz = qz * qz
xy = qx * qy
xz = qx * qz
yz = qy * qz
xw = qx * qw
yw = qy * qw
zw = qz * qw
row0 = torch.cat((1 - 2 * yy - 2 * zz, 2 * xy - 2 * zw, 2 * xz + 2 * yw), 1) # batch*3
row1 = torch.cat((2 * xy + 2 * zw, 1 - 2 * xx - 2 * zz, 2 * yz - 2 * xw), 1) # batch*3
row2 = torch.cat((2 * xz - 2 * yw, 2 * yz + 2 * xw, 1 - 2 * xx - 2 * yy), 1) # batch*3
matrix = torch.cat((row0.view(batch, 1, 3), row1.view(batch, 1, 3), row2.view(batch, 1, 3)), 1) # batch*3*3
return matrix
def calc_rot_error_from_qxyzw(rotation_pred, rotations):
pred_rmat = compute_rotation_matrix_from_quaternion_xyzw(rotation_pred)
gt_rmat = compute_rotation_matrix_from_quaternion_xyzw(rotations[:, 0])
gt_rmat_z_180 = compute_rotation_matrix_from_quaternion_xyzw(rotations[:, 1])
gt_R_error = compute_geodesic_distance_from_two_matrices(gt_rmat, pred_rmat)
gt_R_z_180_error = compute_geodesic_distance_from_two_matrices(gt_rmat_z_180, pred_rmat)
error,_ = torch.min(torch.stack([gt_R_error, gt_R_z_180_error], dim=0), dim=0)
return error
================================================
FILE: src/nr/utils/dataset_utils.py
================================================
import numpy as np
import time
import random
import torch
def dummy_collate_fn(data_list):
return data_list[0]
def simple_collate_fn(data_list):
ks=data_list[0].keys()
outputs={k:[] for k in ks}
for k in ks:
for data in data_list:
outputs[k].append(data[k])
outputs[k]=torch.stack(outputs[k],0)
return outputs
def set_seed(index,is_train):
if is_train:
np.random.seed((index+int(time.time()))%(2**16))
random.seed((index+int(time.time()))%(2**16)+1)
torch.random.manual_seed((index+int(time.time()))%(2**16)+1)
else:
np.random.seed(index % (2 ** 16))
random.seed(index % (2 ** 16) + 1)
torch.random.manual_seed(index % (2 ** 16) + 1)
================================================
FILE: src/nr/utils/draw_utils.py
================================================
import matplotlib
matplotlib.use('Agg')
import sys
sys.path.append("./src/nr")
from utils.base_utils import compute_relative_transformation, compute_F
import numpy as np
import cv2
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
from matplotlib import cm
from sklearn.decomposition import PCA
from sklearn.manifold import TSNE
import open3d as o3d
def newline(p1, p2):
ax = plt.gca()
xmin, xmax = ax.get_xbound()
if p2[0] == p1[0]:
xmin = xmax = p1[0]
ymin, ymax = ax.get_ybound()
else:
ymax = p1[1]+(p2[1]-p1[1])/(p2[0]-p1[0])*(xmax-p1[0])
ymin = p1[1]+(p2[1]-p1[1])/(p2[0]-p1[0])*(xmin-p1[0])
l = mlines.Line2D([xmin,xmax], [ymin,ymax])
ax.add_line(l)
return l
def draw_correspondence(img0, img1, kps0, kps1, matches=None, colors=None, max_draw_line_num=None, kps_color=(0,0,255),vert=False):
if len(img0.shape)==2:
img0=np.repeat(img0[:,:,None],3,2)
if len(img1.shape)==2:
img1=np.repeat(img1[:,:,None],3,2)
h0, w0 = img0.shape[:2]
h1, w1 = img1.shape[:2]
if matches is None:
assert(kps0.shape[0]==kps1.shape[0])
matches=np.repeat(np.arange(kps0.shape[0])[:,None],2,1)
if vert:
w = max(w0, w1)
h = h0 + h1
out_img = np.zeros([h, w, 3], np.uint8)
out_img[:h0, :w0] = img0
out_img[h0:, :w1] = img1
else:
h = max(h0, h1)
w = w0 + w1
out_img = np.zeros([h, w, 3], np.uint8)
out_img[:h0, :w0] = img0
out_img[:h1, w0:] = img1
for pt in kps0:
pt = np.round(pt).astype(np.int32)
cv2.circle(out_img, tuple(pt), 1, kps_color, -1)
for pt in kps1:
pt = np.round(pt).astype(np.int32)
pt = pt.copy()
if vert:
pt[1] += h0
else:
pt[0] += w0
cv2.circle(out_img, tuple(pt), 1, kps_color, -1)
if max_draw_line_num is not None and matches.shape[0]>max_draw_line_num:
np.random.seed(6033)
idxs=np.arange(matches.shape[0])
np.random.shuffle(idxs)
idxs=idxs[:max_draw_line_num]
matches= matches[idxs]
if colors is not None and (type(colors)==list or type(colors)==np.ndarray):
colors=np.asarray(colors)
colors= colors[idxs]
for mi,m in enumerate(matches):
pt = np.round(kps0[m[0]]).astype(np.int32)
pr_pt = np.round(kps1[m[1]]).astype(np.int32)
if vert:
pr_pt[1] += h0
else:
pr_pt[0] += w0
if colors is None:
cv2.line(out_img, tuple(pt), tuple(pr_pt), (0, 255, 0), 1)
elif type(colors)==list or type(colors)==np.ndarray:
color=(int(c) for c in colors[mi])
cv2.line(out_img, tuple(pt), tuple(pr_pt), tuple(color), 1)
else:
color=(int(c) for c in colors)
cv2.line(out_img, tuple(pt), tuple(pr_pt), tuple(color), 1)
return out_img
def draw_keypoints(img,kps,colors=None,radius=2):
out_img=img.copy()
for pi, pt in enumerate(kps):
pt = np.round(pt).astype(np.int32)
if colors is not None:
color=[int(c) for c in colors[pi]]
cv2.circle(out_img, tuple(pt), radius, color, -1, cv2.FILLED)
else:
cv2.circle(out_img, tuple(pt), radius, (0,255,0), -1)
return out_img
def draw_epipolar_line(F, img0, img1, pt0, color):
h1,w1=img1.shape[:2]
hpt = np.asarray([pt0[0], pt0[1], 1], dtype=np.float32)[:, None]
l = F @ hpt
l = l[:, 0]
a, b, c = l[0], l[1], l[2]
pt1 = np.asarray([0, -c / b]).astype(np.int32)
pt2 = np.asarray([w1, (-a * w1 - c) / b]).astype(np.int32)
img0 = cv2.circle(img0, tuple(pt0.astype(np.int32)), 5, color, 2)
img1 = cv2.line(img1, tuple(pt1), tuple(pt2), color, 2)
return img0, img1
def draw_epipolar_lines(F, img0, img1,num=20):
img0,img1=img0.copy(),img1.copy()
h0, w0, _ = img0.shape
h1, w1, _ = img1.shape
for k in range(num):
color = np.random.randint(0, 255, [3], dtype=np.int32)
color = [int(c) for c in color]
pt = np.random.uniform(0, 1, 2)
pt[0] *= w0
pt[1] *= h0
pt = pt.astype(np.int32)
img0, img1 = draw_epipolar_line(F, img0, img1, pt, color)
return img0, img1
def gen_color_map(error, clip_max=12.0, clip_min=2.0, cmap_name='viridis'):
rectified_error=(error-clip_min)/(clip_max-clip_min)
rectified_error[rectified_error<0]=0
rectified_error[rectified_error>=1.0]=1.0
viridis=cm.get_cmap(cmap_name,256)
colors=[viridis(e) for e in rectified_error]
return np.asarray(np.asarray(colors)[:,:3]*255,np.uint8)
def scale_float_image(image):
max_val, min_val = np.max(image), np.min(image)
image = (image - min_val) / (max_val - min_val) * 255
return image.astype(np.uint8)
def concat_images(img0,img1,vert=False):
if not vert:
h0,h1=img0.shape[0],img1.shape[0],
if h02:
if ds_type=='pca':
pca=PCA(2)
feats=pca.fit_transform(feats)
elif ds_type=='tsne':
tsne=TSNE(2)
feats=tsne.fit_transform(feats)
elif ds_type=='pca-tsne':
if d>50:
tsne=PCA(50)
feats=tsne.fit_transform(feats)
tsne=TSNE(2,100.0)
feats=tsne.fit_transform(feats)
else:
raise NotImplementedError
colors=[np.array([c[0],c[1],c[2]],np.float64)/255.0 for c in colors]
feats_min=np.min(feats,0,keepdims=True)
feats_max=np.max(feats,0,keepdims=True)
feats=(feats-(feats_min+feats_max)/2)*10/(feats_max-feats_min)
plt.scatter(feats[:,0],feats[:,1],s=0.5,c=colors)
plt.savefig(fn)
plt.close()
return feats
def draw_points(img,points):
pts=np.round(points).astype(np.int32)
h,w,_=img.shape
pts[:,0]=np.clip(pts[:,0],a_min=0,a_max=w-1)
pts[:,1]=np.clip(pts[:,1],a_min=0,a_max=h-1)
img=img.copy()
img[pts[:,1],pts[:,0]]=255
# img[pts[:,1],pts[:,0]]+=np.asarray([127,0,0],np.uint8)[None,:]
return img
def draw_bbox(img,bbox,color=None):
if color is not None:
color=[int(c) for c in color]
else:
color=(0,255,0)
img=cv2.rectangle(img,(bbox[0],bbox[1]),(bbox[0]+bbox[2],bbox[1]+bbox[3]),color)
return img
def output_points(fn,pts,colors=None):
with open(fn, 'w') as f:
for pi, pt in enumerate(pts):
f.write(f'{pt[0]:.6f} {pt[1]:.6f} {pt[2]:.6f} ')
if colors is not None:
f.write(f'{int(colors[pi,0])} {int(colors[pi,1])} {int(colors[pi,2])}')
f.write('\n')
def compute_axis_points(pose):
R=pose[:,:3] # 3,3
t=pose[:,3:] # 3,1
pts = np.concatenate([np.identity(3),np.zeros([3,1])],1) # 3,4
pts = R.T @ (pts - t)
colors = np.asarray([[255,0,0],[0,255,0,],[0,0,255],[0,0,0]],np.uint8)
return pts.T, colors
def draw_epipolar_lines_func(img0,img1,Rt0,Rt1,K0,K1):
Rt=compute_relative_transformation(Rt0,Rt1)
F=compute_F(K0,K1,Rt[:,:3],Rt[:,3:])
return concat_images_list(*draw_epipolar_lines(F,img0,img1))
def draw_axis(img, R, t, K, length=0.1, width=3, with_text=False, dist=None):
"""
Draw a 6dof axis (XYZ -> RGB) in the given rotation and translation
:param img - rgb numpy array (RGB format, not opencv BGR)
:rotation_vec - euler rotations, numpy array of length 3,
use cv2.Rodrigues(R)[0] to convert from rotation matrix
:t - 3d translation vector, in meters (dtype must be float)
:K - intrinsic calibration matrix , 3x3
:length - factor to control the axis lengths
:dist - optional distortion coefficients, numpy array of length 4. If None distortion is ignored.
"""
rotation_vec = cv2.Rodrigues(R)[0]
img = img.astype(np.float32)
dist = np.zeros(4, dtype=float) if dist is None else dist
points = length * np.float32([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 0]]).reshape(-1, 3)
axis_points, _ = cv2.projectPoints(points, rotation_vec, t, K, dist)
axis_points = axis_points.astype(np.int)
if with_text:
for pt,txt in zip(axis_points, "XYZO"):
cv2.putText(img, txt, tuple(pt.ravel()), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
img = cv2.line(img, tuple(axis_points[3].ravel()), tuple(axis_points[0].ravel()), (0, 0, 255), width)
img = cv2.line(img, tuple(axis_points[3].ravel()), tuple(axis_points[1].ravel()), (0, 255, 0), width)
img = cv2.line(img, tuple(axis_points[3].ravel()), tuple(axis_points[2].ravel()), (255, 0, 0), width)
return img
def draw_gripper(img, R, t, K, width, thickness=2, dist=None):
rotation_vec = cv2.Rodrigues(R)[0]
img = img.astype(np.float32)
dist = np.zeros(4, dtype=float) if dist is None else dist
points = np.float32([[0, -width/2, 0], [0, -width/2, 0.05], [0, width/2, 0], [0, width/2, 0.05], [0, 0, -0.05], [0, 0, 0]]).reshape(-1, 3)
axis_points, _ = cv2.projectPoints(points, rotation_vec, t, K, dist)
axis_points = axis_points.astype(np.int)
img = cv2.line(img, tuple(axis_points[0].ravel()), tuple(axis_points[2].ravel()), (0, 255, 0), thickness)
img = cv2.line(img, tuple(axis_points[0].ravel()), tuple(axis_points[1].ravel()), (0, 255, 0), thickness)
img = cv2.line(img, tuple(axis_points[2].ravel()), tuple(axis_points[3].ravel()), (0, 255, 0), thickness)
img = cv2.line(img, tuple(axis_points[-2].ravel()), tuple(axis_points[-1].ravel()), (255, 0, 0), thickness)
return img
def draw_cube(img, R, t, K, length=0.3, bias=[-0.15,-0.15,-0.], dist=None):
rotation_vec = cv2.Rodrigues(R)[0]
img = img.astype(np.float32)
dist = np.zeros(4, dtype=float) if dist is None else dist
points = length * np.float32([[0,0,0], [0,1,0], [1,1,0], [1,0,0],
[0,0,1], [0,1,1], [1,1,1], [1,0,1]])
points += np.float32(bias)
imgpts, _ = cv2.projectPoints(points, rotation_vec, t, K, dist)
imgpts = np.int32(imgpts).reshape(-1,2)
# draw ground floor in green
img = cv2.drawContours(img, [imgpts[:4]],-1,(0,255,0),3)
# draw pillars in blue color
for i,j in zip(range(4),range(4,8)):
img = cv2.line(img, tuple(imgpts[i]), tuple(imgpts[j]),(0,0,255),3)
# draw top layer in red color
img = cv2.drawContours(img, [imgpts[4:]],-1,(255,0,0),3)
return img
def draw_world_points(img, points, Rwc, twc, K, dist=None):
rotation_vec = cv2.Rodrigues(Rwc)[0]
img = img.astype(np.float32)
dist = np.zeros(4, dtype=float) if dist is None else dist
axis_points, _ = cv2.projectPoints(points, rotation_vec, twc, K, dist)
img = draw_keypoints(img, axis_points.squeeze(1))
return img
def extract_surface_points_from_volume(vol, rg, bound=(-1,1), color=(0,0,1), scale=0.3 / 40):
assert len(vol.shape) == 3
ind = np.transpose( ( (vol > rg[0]).astype(np.bool) & (vol < rg[1]).astype(np.bool) ).nonzero())
print(ind.shape)
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(ind)
colors = np.array([color]).repeat(ind.shape[0],axis=0)
if color == None:
values = vol[ind[:,0], ind[:,1], ind[:,2]]
a, b = bound
m = (a + b) / 2
r = np.where(values <= m, values - a, - values + b)
g = np.where(values <= m, 0, 1 - r)
b = np.where(values <= m, 1 - r, 0)
colors = np.stack([r, g, b], axis=-1)
else:
colors = np.array([color]).repeat(ind.shape[0],axis=0)
pcd.scale(scale, center=(0, 0, 0))
pcd.colors = o3d.utility.Vector3dVector(colors)
return pcd
def draw_volume_surface(logdir, vol, prefix='surface', rg=(-0.2,0.2)):
pcd = extract_surface_points_from_volume(vol, rg)
o3d.io.write_point_cloud(f"{logdir}/{prefix}.ply", pcd)
def create_mesh_box(width, height, depth, dx=0, dy=0, dz=0):
''' Author: chenxi-wang
Create box instance with mesh representation.
'''
box = o3d.geometry.TriangleMesh()
vertices = np.array([[0,0,0],
[width,0,0],
[0,0,depth],
[width,0,depth],
[0,height,0],
[width,height,0],
[0,height,depth],
[width,height,depth]])
vertices[:,0] += dx
vertices[:,1] += dy
vertices[:,2] += dz
triangles = np.array([[4,7,5],[4,6,7],[0,2,4],[2,6,4],
[0,1,2],[1,3,2],[1,5,7],[1,7,3],
[2,3,7],[2,7,6],[0,4,1],[1,4,5]])
box.vertices = o3d.utility.Vector3dVector(vertices)
box.triangles = o3d.utility.Vector3iVector(triangles)
return box
def draw_gripper_o3d(R, t, width, score=1, color=None):
'''
Author: chenxi-wang
**Input:**
- center: numpy array of (3,), target point as gripper center
- R: numpy array of (3,3), rotation matrix of gripper
- width: float, gripper width
- score: float, grasp quality score
**Output:**
- open3d.geometry.TriangleMesh
'''
#x, y, z = center
center = t
#print(center)
height = 0.004
finger_width = 0.004
tail_length = 0.04
depth = 0.05
depth_base = 0 # 0.02
if color is not None:
color_r, color_g, color_b = color
else:
color_r = score # red for high score
color_g = 0
color_b = 1 - score # blue for low score
left = create_mesh_box(depth + depth_base + finger_width, finger_width, height)
right = create_mesh_box(depth + depth_base + finger_width, finger_width, height)
bottom = create_mesh_box(finger_width, width, height)
tail = create_mesh_box(tail_length, finger_width, height)
left_points = np.array(left.vertices)
left_triangles = np.array(left.triangles)
left_points[:, 0] -= depth_base + finger_width
left_points[:, 1] -= width / 2 + finger_width
left_points[:, 2] -= height / 2
right_points = np.array(right.vertices)
right_triangles = np.array(right.triangles) + 8
right_points[:, 0] -= depth_base + finger_width
right_points[:, 1] += width / 2
right_points[:, 2] -= height / 2
bottom_points = np.array(bottom.vertices)
bottom_triangles = np.array(bottom.triangles) + 16
bottom_points[:, 0] -= finger_width + depth_base
bottom_points[:, 1] -= width / 2
bottom_points[:, 2] -= height / 2
tail_points = np.array(tail.vertices)
tail_triangles = np.array(tail.triangles) + 24
tail_points[:, 0] -= tail_length + finger_width + depth_base
tail_points[:, 1] -= finger_width / 2
tail_points[:, 2] -= height / 2
vertices = np.concatenate([left_points, right_points, bottom_points, tail_points], axis=0)
vertices = np.dot(R, vertices.T).T + center
triangles = np.concatenate([left_triangles, right_triangles, bottom_triangles, tail_triangles], axis=0)
colors = np.array([[color_r, color_g, color_b] for _ in range(len(vertices))])
gripper = o3d.geometry.TriangleMesh()
gripper.vertices = o3d.utility.Vector3dVector(vertices)
gripper.triangles = o3d.utility.Vector3iVector(triangles)
gripper.vertex_colors = o3d.utility.Vector3dVector(colors)
return gripper
def transform_points(points, trans):
'''
Author: chenxi-wang
**Input:**
- points: numpy array of (N,3), point cloud
- trans: numpy array of (4,4), transformation matrix
**Output:**
- numpy array of (N,3), transformed points.
'''
ones = np.ones([points.shape[0], 1], dtype=points.dtype)
points_ = np.concatenate([points, ones], axis=-1)
points_ = np.matmul(trans, points_.T).T
return points_[:, :3]
================================================
FILE: src/nr/utils/field_utils.py
================================================
import torch
import numpy as np
def generate_grid_points_old(bound_min, bound_max, resolution):
X = torch.linspace(bound_min[0], bound_max[0], resolution)
Y = torch.linspace(bound_min[1], bound_max[1], resolution)
Z = torch.linspace(bound_max[2], bound_min[2], resolution) # from top to down to be like with training rays
XYZ = torch.stack(torch.meshgrid(X, Y, Z), dim=-1)
return XYZ
RESOLUTION = 40
VOLUME_SIZE = 0.3
VOXEL_SIZE = VOLUME_SIZE / RESOLUTION
HALF_VOXEL_SIZE = VOXEL_SIZE / 2
def generate_grid_points():
points = []
for x in range(RESOLUTION):
for y in range(RESOLUTION):
for z in range(RESOLUTION):
points.append([x * VOXEL_SIZE + HALF_VOXEL_SIZE,
y * VOXEL_SIZE + HALF_VOXEL_SIZE,
z * VOXEL_SIZE + HALF_VOXEL_SIZE])
return np.array(points).astype(np.float32)
TSDF_SAMPLE_POINTS = generate_grid_points()
if __name__ == "__main__":
GT_POINTS = np.load('points.npy')
TSDF_VOLUME_MASK = np.zeros((1, 40, 40, 40), dtype=np.bool8)
idxs = []
for point in GT_POINTS:
i, j, k = np.floor(point / VOXEL_SIZE).astype(int)
TSDF_VOLUME_MASK[0, i, j, k] = True
idxs.append(i * (RESOLUTION * RESOLUTION) + j * RESOLUTION + k)
print(TSDF_SAMPLE_POINTS[idxs], GT_POINTS)
assert np.allclose(TSDF_SAMPLE_POINTS[idxs], GT_POINTS)
================================================
FILE: src/nr/utils/grasp_utils.py
================================================
import datetime
from pathlib import Path
import numpy as np
from scipy import ndimage
import sys
sys.path.append("./src")
import time
from nr.utils.base_utils import color_map_forward
from nr.utils.draw_utils import draw_cube, extract_surface_points_from_volume
from gd.utils.transform import Transform, Rotation
from skimage.io import imsave
import cv2
class Grasp(object):
"""Grasp parameterized as pose of a 2-finger robot hand.
TODO(mbreyer): clarify definition of grasp frame
"""
def __init__(self, pose, width):
self.pose = pose
self.width = width
def to_voxel_coordinates(grasp, voxel_size):
pose = grasp.pose
pose.translation /= voxel_size
width = grasp.width / voxel_size
return Grasp(pose, width)
def from_voxel_coordinates(grasp, voxel_size):
pose = grasp.pose
pose.translation *= voxel_size
width = grasp.width * voxel_size
return Grasp(pose, width)
def process(
tsdf_vol,
qual_vol,
rot_vol,
width_vol,
gaussian_filter_sigma=1.0,
min_width=0,
max_width=12,
):
tsdf_vol = tsdf_vol.squeeze()
qual_vol = qual_vol.squeeze()
rot_vol = rot_vol.squeeze()
width_vol = width_vol.squeeze()
# smooth quality volume with a Gaussian
qual_vol = ndimage.gaussian_filter(
qual_vol, sigma=gaussian_filter_sigma, mode="nearest"
)
# mask out voxels too far away from the surface
outside_voxels = tsdf_vol > 0.1
inside_voxels = np.logical_and(-1 < tsdf_vol, tsdf_vol < -0.1)
valid_voxels = ndimage.morphology.binary_dilation(
outside_voxels, iterations=2, mask=np.logical_not(inside_voxels)
)
qual_vol[valid_voxels == False] = 0.0
# reject voxels with predicted widths that are too small or too large
qual_vol[np.logical_or(width_vol < min_width, width_vol > max_width)] = 0.0
return tsdf_vol, qual_vol, rot_vol, width_vol
def select_index(qual_vol, rot_vol, width_vol, index):
i, j, k = index
score = qual_vol[i, j, k]
ori = Rotation.from_quat(rot_vol[:, i, j, k])
pos = np.array([i, j, k], dtype=np.float64)
width = width_vol[i, j, k]
return Grasp(Transform(ori, pos), width), score
def select(qual_vol, rot_vol, width_vol, threshold=0.90, max_filter_size=4):
# threshold on grasp quality
qual_vol[qual_vol < threshold] = 0.0
# non maximum suppression
max_vol = ndimage.maximum_filter(qual_vol, size=max_filter_size)
qual_vol = np.where(qual_vol == max_vol, qual_vol, 0.0)
mask = np.where(qual_vol, 1.0, 0.0)
# construct grasps
grasps, scores = [], []
for index in np.argwhere(mask):
grasp, score = select_index(qual_vol, rot_vol, width_vol, index)
grasps.append(grasp)
scores.append(score)
return grasps, scores
def sim_grasp(database, alp_vol, qual_vol, rot_vol, width_vol, top_k=10):
from utils.grasp_utils import select, process
qual_vol, rot_vol, width_vol = process(alp_vol, qual_vol, rot_vol, width_vol)
grasps, scores = select(qual_vol.copy(), rot_vol, width_vol)
grasps, scores = np.asarray(grasps), np.asarray(scores)
img = None
if len(grasps) > 0:
p = np.argsort(scores)[::-1][:top_k]
grasps = [g for g in grasps[p]]
scores = scores[p]
pos = np.array([ g.pose.translation for g in grasps ])
rot = np.array([ g.pose.rotation.as_matrix() for g in grasps ])
width = np.array([ g.width for g in grasps ])
img = database.visualize_grasping(pos, rot, width)
database.visualize_grasping_3d(pos, rot, width, scores)
return grasps, scores, img
def run_real(run_id, model, images: list, extrinsics: list, intrinsic, save_img=True):
extrinsics = np.stack(extrinsics, 0)
intrinsics = np.repeat(np.expand_dims(intrinsic, 0), extrinsics.shape[0], axis=0)
depth_range = np.repeat(np.expand_dims(np.r_[0.2, 0.8], 0), extrinsics.shape[0], axis=0).astype(np.float32)
bbox3d = [[-0.15, -0.15, 0.00], [0.15, 0.15, 0.3]]
if save_img:
save_path = f'data/grasp_capture/{run_id}'
if not Path(save_path).exists():
Path(save_path).mkdir(parents=True)
for i, img in enumerate(images):
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
# img = draw_cube(img, extrinsics[i][:3,:3], extrinsics[i][:3,3], intrinsics[i], length=0.3, bias=bbox3d[0])
cv2.imwrite(f"{save_path}/{i}.png", img)
images = color_map_forward(np.stack(images, 0)).transpose([0, 3, 1, 2])
t0 = time.time()
tsdf_vol, qual_vol, rot_vol, width_vol = model(images, extrinsics, intrinsics, depth_range=depth_range, bbox3d=bbox3d, que_id=3)
t = time.time() - t0
tsdf_vol, qual_vol, rot_vol, width_vol = process(tsdf_vol, qual_vol, rot_vol, width_vol)
grasps, scores = select(qual_vol.copy(), rot_vol, width_vol)
grasps, scores = np.asarray(grasps), np.asarray(scores)
if len(grasps) > 0:
p = np.random.permutation(len(grasps))
grasps = [from_voxel_coordinates(g, 0.3 / 40) for g in grasps[p]]
scores = scores[p]
pc = extract_surface_points_from_volume(tsdf_vol, (-0.2, 0.2))
return grasps, scores, tsdf_vol, pc, t
================================================
FILE: src/nr/utils/imgs_info.py
================================================
import numpy as np
import torch
from utils.base_utils import color_map_forward, pad_img_end
def random_crop(ref_imgs_info, que_imgs_info, target_size):
imgs = ref_imgs_info['imgs']
n, _, h, w = imgs.shape
out_h, out_w = target_size[0], target_size[1]
if out_w >= w or out_h >= h:
return ref_imgs_info
center_h = np.random.randint(low=out_h // 2 + 1, high=h - out_h // 2 - 1)
center_w = np.random.randint(low=out_w // 2 + 1, high=w - out_w // 2 - 1)
def crop(tensor):
tensor = tensor[:, :, center_h - out_h // 2:center_h + out_h // 2,
center_w - out_w // 2:center_w + out_w // 2]
return tensor
def crop_imgs_info(imgs_info):
imgs_info['imgs'] = crop(imgs_info['imgs'])
if 'depth' in imgs_info: imgs_info['depth'] = crop(imgs_info['depth'])
if 'true_depth' in imgs_info: imgs_info['true_depth'] = crop(imgs_info['true_depth'])
if 'masks' in imgs_info: imgs_info['masks'] = crop(imgs_info['masks'])
Ks = imgs_info['Ks'] # n, 3, 3
h_init = center_h - out_h // 2
w_init = center_w - out_w // 2
Ks[:,0,2]-=w_init
Ks[:,1,2]-=h_init
imgs_info['Ks']=Ks
return imgs_info
return crop_imgs_info(ref_imgs_info), crop_imgs_info(que_imgs_info)
def random_flip(ref_imgs_info,que_imgs_info):
def flip(tensor):
tensor = np.flip(tensor.transpose([0, 2, 3, 1]), 2) # n,h,w,3
tensor = np.ascontiguousarray(tensor.transpose([0, 3, 1, 2]))
return tensor
def flip_imgs_info(imgs_info):
imgs_info['imgs'] = flip(imgs_info['imgs'])
if 'depth' in imgs_info: imgs_info['depth'] = flip(imgs_info['depth'])
if 'true_depth' in imgs_info: imgs_info['true_depth'] = flip(imgs_info['true_depth'])
if 'masks' in imgs_info: imgs_info['masks'] = flip(imgs_info['masks'])
Ks = imgs_info['Ks'] # n, 3, 3
Ks[:, 0, :] *= -1
w = imgs_info['imgs'].shape[-1]
Ks[:, 0, 2] += w - 1
imgs_info['Ks'] = Ks
return imgs_info
ref_imgs_info = flip_imgs_info(ref_imgs_info)
que_imgs_info = flip_imgs_info(que_imgs_info)
return ref_imgs_info, que_imgs_info
def pad_imgs_info(ref_imgs_info,pad_interval):
ref_imgs, ref_depths, ref_masks = ref_imgs_info['imgs'], ref_imgs_info['depth'], ref_imgs_info['masks']
ref_depth_gt = ref_imgs_info['true_depth'] if 'true_depth' in ref_imgs_info else None
rfn, _, h, w = ref_imgs.shape
ph = (pad_interval - (h % pad_interval)) % pad_interval
pw = (pad_interval - (w % pad_interval)) % pad_interval
if ph != 0 or pw != 0:
ref_imgs = np.pad(ref_imgs, ((0, 0), (0, 0), (0, ph), (0, pw)), 'reflect')
ref_depths = np.pad(ref_depths, ((0, 0), (0, 0), (0, ph), (0, pw)), 'reflect')
ref_masks = np.pad(ref_masks, ((0, 0), (0, 0), (0, ph), (0, pw)), 'reflect')
if ref_depth_gt is not None:
ref_depth_gt = np.pad(ref_depth_gt, ((0, 0), (0, 0), (0, ph), (0, pw)), 'reflect')
ref_imgs_info['imgs'], ref_imgs_info['depth'], ref_imgs_info['masks'] = ref_imgs, ref_depths, ref_masks
if ref_depth_gt is not None:
ref_imgs_info['true_depth'] = ref_depth_gt
return ref_imgs_info
def build_imgs_info(database, ref_ids, pad_interval=-1, is_aligned=True, align_depth_range=False, has_mask=True, has_depth=True, replace_none_depth = False):
if not is_aligned:
assert has_depth
rfn = len(ref_ids)
ref_imgs, ref_masks, ref_depths, shapes = [], [], [], []
for ref_id in ref_ids:
img = database.get_image(ref_id)
shapes.append([img.shape[0], img.shape[1]])
ref_imgs.append(img)
ref_masks.append(database.get_mask(ref_id))
ref_depths.append(database.get_depth(ref_id))
shapes = np.asarray(shapes)
th, tw = np.max(shapes, 0)
for rfi in range(rfn):
ref_imgs[rfi] = pad_img_end(ref_imgs[rfi], th, tw, 'reflect')
ref_masks[rfi] = pad_img_end(ref_masks[rfi][:, :, None], th, tw, 'constant', 0)[..., 0]
ref_depths[rfi] = pad_img_end(ref_depths[rfi][:, :, None], th, tw, 'constant', 0)[..., 0]
ref_imgs = color_map_forward(np.stack(ref_imgs, 0)).transpose([0, 3, 1, 2])
ref_masks = np.stack(ref_masks, 0)[:, None, :, :]
ref_depths = np.stack(ref_depths, 0)[:, None, :, :]
else:
ref_imgs = color_map_forward(np.asarray([database.get_image(ref_id) for ref_id in ref_ids])).transpose([0, 3, 1, 2])
if has_mask:
ref_masks = np.asarray([database.get_mask(ref_id) for ref_id in ref_ids], dtype=np.float32)[:, None, :, :]
else:
b, _, h, w = ref_imgs.shape
ref_masks = np.ones([b, _, h, w], dtype=np.float32)
if has_depth:
ref_depths = [database.get_depth(ref_id) for ref_id in ref_ids]
if replace_none_depth:
b, _, h, w = ref_imgs.shape
for i, depth in enumerate(ref_depths):
if depth is None: ref_depths[i] = np.zeros([h, w], dtype=np.float32)
ref_depths = np.asarray(ref_depths, dtype=np.float32)[:, None, :, :]
else: ref_depths = None
ref_poses = np.asarray([database.get_pose(ref_id) for ref_id in ref_ids], dtype=np.float32)
ref_Ks = np.asarray([database.get_K(ref_id) for ref_id in ref_ids], dtype=np.float32)
ref_depth_range = np.asarray([database.get_depth_range(ref_id) for ref_id in ref_ids], dtype=np.float32)
if align_depth_range:
ref_depth_range[:,0]=np.min(ref_depth_range[:,0])
ref_depth_range[:,1]=np.max(ref_depth_range[:,1])
ref_imgs_info = {'imgs': ref_imgs, 'poses': ref_poses, 'Ks': ref_Ks, 'depth_range': ref_depth_range, 'masks': ref_masks, 'bbox3d': database.get_bbox3d()}
if has_depth: ref_imgs_info['depth'] = ref_depths
if pad_interval!=-1:
ref_imgs_info = pad_imgs_info(ref_imgs_info, pad_interval)
return ref_imgs_info
def build_render_imgs_info(que_pose,que_K,que_shape,que_depth_range):
h, w = que_shape
h, w = int(h), int(w)
que_coords = np.stack(np.meshgrid(np.arange(w), np.arange(h)), -1)
que_coords = que_coords.reshape([1, -1, 2]).astype(np.float32)
return {'poses': que_pose.astype(np.float32)[None,:,:], # 1,3,4
'Ks': que_K.astype(np.float32)[None,:,:], # 1,3,3
'coords': que_coords,
'depth_range': np.asarray(que_depth_range, np.float32)[None, :],
'shape': (h,w)}
def build_canonical_info(bbox, resolution, que_pose, que_K):
x_min,x_max,y_min,y_max = bbox
print('bbox', bbox)
que_coords = np.stack(np.meshgrid(np.linspace(y_min, y_max, 2), np.linspace(x_min, x_max, 2)), -1)
print('que_coords', que_coords)
return {'poses': que_pose.astype(np.float32)[None,:,:], # 1,3,4
'Ks': que_K.astype(np.float32)[None,:,:], # 1,3,3
'coords': que_coords,
'depth_range': np.asarray([0.5, 0.8], np.float32)[None, :],
'shape': (resolution, resolution)}
def imgs_info_to_torch(imgs_info):
for k, v in imgs_info.items():
if isinstance(v,np.ndarray):
imgs_info[k] = torch.from_numpy(v).float()
return imgs_info
def grasp_info_to_torch(info):
torch_info = []
for item in info:
torch_info.append(torch.from_numpy(item))
return torch_info
def imgs_info_slice(imgs_info, indices):
imgs_info_out={}
imgs_info_out['bbox3d'] = imgs_info['bbox3d']
for k, v in imgs_info.items():
if k != 'bbox3d' and v is not None:
imgs_info_out[k] = v[indices]
return imgs_info_out
================================================
FILE: src/nr/utils/view_select.py
================================================
import numpy as np
from dataset.database import BaseDatabase
def compute_nearest_camera_indices(database, que_ids, ref_ids=None):
if ref_ids is None: ref_ids = que_ids
ref_poses = [database.get_pose(ref_id) for ref_id in ref_ids]
ref_cam_pts = np.asarray([-pose[:, :3].T @ pose[:, 3] for pose in ref_poses])
que_poses = [database.get_pose(que_id) for que_id in que_ids]
que_cam_pts = np.asarray([-pose[:, :3].T @ pose[:, 3] for pose in que_poses])
dists = np.linalg.norm(ref_cam_pts[None, :, :] - que_cam_pts[:, None, :], 2, 2)
dists_idx = np.argsort(dists, 1)
return dists_idx
def select_working_views(ref_poses, que_poses, work_num, exclude_self=False):
ref_cam_pts = np.asarray([-pose[:, :3].T @ pose[:, 3] for pose in ref_poses])
render_cam_pts = np.asarray([-pose[:, :3].T @ pose[:, 3] for pose in que_poses])
dists = np.linalg.norm(ref_cam_pts[None, :, :] - render_cam_pts[:, None, :], 2, 2) # qn,rfn
ids = np.argsort(dists)
if exclude_self:
ids = ids[:, 1:work_num+1]
else:
ids = ids[:, :work_num]
return ids
def select_working_views_db(database: BaseDatabase, ref_ids, que_poses, work_num, exclude_self=False):
ref_ids = database.get_img_ids() if ref_ids is None else ref_ids
ref_poses = [database.get_pose(img_id) for img_id in ref_ids]
ref_ids = np.asarray(ref_ids)
ref_poses = np.asarray(ref_poses)
indices = select_working_views(ref_poses, que_poses, work_num, exclude_self)
return ref_ids[indices] # qn,wn
================================================
FILE: src/rd/modify_material.py
================================================
from mathutils import Vector
import bpy
import random
def modify_material(mat_links, mat_nodes, material_name, mat_randomize_mode, is_texture=False, orign_base_color=None,
tex_node=None, is_transfer=True, is_arm=False):
if is_transfer:
if material_name.split("_")[0] == "metal" or material_name.split("_")[0] == "porcelain" or \
material_name.split("_")[0] == "plasticsp" or material_name.split("_")[0] == "paintsp":
tex_mix_prop = random.uniform(0.85, 0.98)
else:
tex_mix_prop = random.uniform(0.7, 0.95)
mix_prop = random.uniform(0.6, 0.9)
if mat_randomize_mode == "specular_texmix" or mat_randomize_mode == "mixed" or mat_randomize_mode == "specular_and_transparent"\
or material_name.split("_")[0] == "metal" or material_name.split("_")[0] == "porcelain" \
or material_name.split("_")[0] == "plasticsp" or material_name.split("_")[0] == "paintsp":
transfer_rand = random.randint(0, 2)
else:
transfer_rand = 1
if transfer_rand == 1:
transfer_flag = True
else:
transfer_flag = False
tex_mix_prop = 1
mix_prop = 1
if not is_arm:
bs_color_rand = random.uniform(-0.2, 0.2)
else:
bs_color_rand = 0
r_rand = bs_color_rand
g_rand = bs_color_rand
b_rand = bs_color_rand
else:
tex_mix_prop = 1
mix_prop = 1
transfer_flag = False
if not is_arm:
bs_color_rand = random.uniform(-0.2, 0.2)
else:
bs_color_rand = 0
r_rand = bs_color_rand
g_rand = bs_color_rand
b_rand = bs_color_rand
bsdfnode_list = [n for n in mat_nodes if isinstance(n, bpy.types.ShaderNodeBsdfPrincipled)]
if bsdfnode_list != []:
for bsdfnode in bsdfnode_list:
if not bsdfnode.inputs[4].links: # metallic
src_value = bsdfnode.inputs[4].default_value
if material_name.split("_")[0] == "metal":
new_value = src_value + random.uniform(-0.05, 0.05)
elif material_name.split("_")[0] == "porcelain":
new_value = src_value + random.uniform(-0.05, 0.1)
elif material_name.split("_")[0] == "plasticsp":
new_value = src_value + random.uniform(-0.05, 0.1)
else:
new_value = src_value + random.uniform(-0.05, 0.05)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[4].default_value = new_value
if not bsdfnode.inputs[5].links: # specular
src_value = bsdfnode.inputs[5].default_value
# if material_name.split("_")[0] == "metal":
new_value = src_value + random.uniform(0, 0.3)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[5].default_value = new_value
if not bsdfnode.inputs[6].links: # specularTint
src_value = bsdfnode.inputs[6].default_value
new_value = src_value + random.uniform(-1, 1)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[6].default_value = new_value
if not bsdfnode.inputs[7].links: # roughness
src_value = bsdfnode.inputs[7].default_value
if material_name.split("_")[0] == "metal" or material_name.split("_")[0] == "porcelain" or \
material_name.split("_")[0] == "plasticsp" or material_name.split("_")[0] == "paintsp":
new_value = src_value + random.uniform(-0.2, 0.01)
else:
new_value = src_value + random.uniform(-0.03, 0.1)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[7].default_value = new_value
if not bsdfnode.inputs[8].links: # anisotropic
src_value = bsdfnode.inputs[8].default_value
new_value = src_value + random.uniform(-0.1, 0.1)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[8].default_value = new_value
if not bsdfnode.inputs[9].links: # anisotropicRotation
src_value = bsdfnode.inputs[9].default_value
new_value = src_value + random.uniform(-0.3, 0.3)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[9].default_value = new_value
if not bsdfnode.inputs[10].links: # sheen
src_value = bsdfnode.inputs[10].default_value
new_value = src_value + random.uniform(-0.1, 0.1)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[10].default_value = new_value
if not bsdfnode.inputs[11].links: # sheenTint
src_value = bsdfnode.inputs[11].default_value
new_value = src_value + random.uniform(-0.2, 0.2)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[11].default_value = new_value
if not bsdfnode.inputs[12].links: # clearcoat
src_value = bsdfnode.inputs[12].default_value
new_value = src_value + random.uniform(-0.2, 0.2)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[12].default_value = new_value
if not bsdfnode.inputs[13].links: # clearcoatGloss
src_value = bsdfnode.inputs[13].default_value
new_value = src_value + random.uniform(-0.2, 0.2)
if new_value > 1.0:
new_value = 1.0
elif new_value < 0:
new_value = 0.0
bsdfnode.inputs[13].default_value = new_value
## metal
if material_name == "metal_0":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.3, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 1.0) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.3, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_1":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.9, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.5, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[7].default_value = random.uniform(0.08, 0.25) # roughness
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.04, 0.5) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.3, 0.7) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.8, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Tangent"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[22])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_10":
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.5) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.3, 0.7) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
bsdf_new.location = Vector((-800, 0))
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
mix_new.location = Vector((-800, 0))
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Image Texture"].outputs[1], mat_nodes["Principled BSDF-new"].inputs[19])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[4])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_11":
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.8) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 0.8) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[4])
mat_links.new(mat_nodes["Image Texture.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_12":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.8) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 0.8) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Reroute.006"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_13":
# mat_nodes["Principled BSDF.001"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF.001"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF.001"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF.001"].inputs[8].default_value = random.uniform(0.3, 0.7) # anisotropic
# mat_nodes["Principled BSDF.001"].inputs[9].default_value = random.uniform(0.0, 0.8) # anisotropicRotation
# mat_nodes["Principled BSDF.001"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF.001"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF.001"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Mix.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF.001"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output.001"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF.001"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_14":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.5) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 0.5) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.85
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Group"].outputs[1], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Group"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_2":
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.5, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.95) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Image Texture.003"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_3":
##
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.5, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.2) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 1.0) # clearcoatGloss
mat_nodes["Gamma"].inputs[1].default_value = random.uniform(3.0, 4.0)
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Gamma"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_4":
##
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.95, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.1, 0.5) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.0, 0.2) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 0.5) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.5) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "metal_5":
##
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.2, 0.4) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.6, 0.9) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.8, 1.0) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Voronoi Texture"].outputs[1], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Tangent"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[22])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[21])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_6":
##
# mat_nodes["BSDF guidé"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["BSDF guidé"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["BSDF guidé"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["BSDF guidé"].inputs[8].default_value = random.uniform(0.0, 0.2) # anisotropic
# mat_nodes["BSDF guidé"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["BSDF guidé"].inputs[12].default_value = random.uniform(0.0, 0.3) # clearcoat
# mat_nodes["BSDF guidé"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
mat_nodes["Valeur"].outputs[0].default_value = random.uniform(0.1, 0.3)
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
bsdf_new.location = Vector((-800, 0))
for key, input in enumerate(mat_nodes["BSDF guidé"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
mix_new.location = Vector((-800, 0))
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Mélanger.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["BSDF guidé"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Sortie de matériau"].inputs[0])
elif material_name == "metal_7":
##
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[8].default_value = random.uniform(0.7, 0.9) # anisotropic
# mat_nodes["Principled BSDF"].inputs[9].default_value = random.uniform(0.0, 1.0) # anisotropicRotation
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 0.3) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
# bsdf_new.location = Vector((-800, 0))
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
# mix_new.location = Vector((-800, 0))
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.7
mat_links.new(mat_nodes["Reroute.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Tangent"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[22])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "metal_8":
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
bsdf_1_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_1_new.name = 'Principled BSDF-1-new'
for key, input in enumerate(mat_nodes["Principled BSDF.001"].inputs):
bsdf_1_new.inputs[key].default_value = input.default_value
bsdf_2_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_2_new.name = 'Principled BSDF-2-new'
for key, input in enumerate(mat_nodes["Principled BSDF.002"].inputs):
bsdf_2_new.inputs[key].default_value = input.default_value
bsdf_3_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_3_new.name = 'Principled BSDF-3-new'
for key, input in enumerate(mat_nodes["Principled BSDF.003"].inputs):
bsdf_3_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
mix_1_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_1_new.name = 'Mix Shader-1-new'
mix_2_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_2_new.name = 'Mix Shader-2-new'
mix_3_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_3_new.name = 'Mix Shader-3-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-1-new"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-2-new"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-3-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.6
mat_nodes["Mix Shader-1-new"].inputs[0].default_value = 0.6
mat_nodes["Mix Shader-2-new"].inputs[0].default_value = 0.6
mat_nodes["Mix Shader-3-new"].inputs[0].default_value = 0.6
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Principled BSDF-1-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Principled BSDF-2-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Principled BSDF-3-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.5
mat_nodes["Mix Shader-1-new"].inputs[0].default_value = 0.5
mat_nodes["Mix Shader-2-new"].inputs[0].default_value = 0.5
mat_nodes["Mix Shader-3-new"].inputs[0].default_value = 0.5
mat_links.new(mat_nodes["ColorRamp"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-1-new"].inputs[20])
mat_links.new(mat_nodes["Bump.001"].outputs[0], mat_nodes["Principled BSDF-2-new"].inputs[20])
mat_links.new(mat_nodes["Principled BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF.001"].outputs[0], mat_nodes["Mix Shader-1-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-1-new"].outputs[0], mat_nodes["Mix Shader-1-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-1-new"].outputs[0], mat_nodes["Mix Shader"].inputs[2])
mat_links.new(mat_nodes["Principled BSDF.002"].outputs[0], mat_nodes["Mix Shader-2-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-2-new"].outputs[0], mat_nodes["Mix Shader-2-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-2-new"].outputs[0], mat_nodes["Mix Shader.001"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF.003"].outputs[0], mat_nodes["Mix Shader-3-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-3-new"].outputs[0], mat_nodes["Mix Shader-3-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-3-new"].outputs[0], mat_nodes["Mix Shader.001"].inputs[2])
elif material_name == "metal_9":
##
# mat_nodes["Principled BSDF"].inputs[4].default_value = random.uniform(0.98, 1.00) # metallic
# mat_nodes["Principled BSDF"].inputs[5].default_value = random.uniform(0.5, 1.0) # specular
# mat_nodes["Principled BSDF"].inputs[6].default_value = random.uniform(0.0, 1.0) # specularTint
mat_nodes["Principled BSDF"].inputs[7].default_value = random.uniform(0.01, 0.3) # roughness
# mat_nodes["Principled BSDF"].inputs[12].default_value = random.uniform(0.0, 0.3) # clearcoat
# mat_nodes["Principled BSDF"].inputs[13].default_value = random.uniform(0.0, 0.3) # clearcoatGloss
mat_nodes["Anisotropic BSDF"].inputs[1].default_value = random.uniform(0.11, 0.25)
mat_nodes["Anisotropic BSDF"].inputs[2].default_value = random.uniform(0.4, 0.6)
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
mat_links.new(tex_node.outputs[0], mat_nodes["Anisotropic BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_nodes["Anisotropic BSDF"].inputs[0].default_value = list(orign_base_color)
mat_links.new(mat_nodes["Principled BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Principled BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
## porcelain
elif material_name == "porcelain_0":
if transfer_flag == True:
# if is_texture:
# mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
# else:
# mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = mix_prop # 0.8
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "porcelain_1":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Mix"].inputs[1].default_value
new_bs_color_r = bs_color[0] + random.uniform(-0.3, 0.3)
new_bs_color_g = bs_color[1] + random.uniform(-0.3, 0.3)
new_bs_color_b = bs_color[2] + random.uniform(-0.3, 0.3)
if new_bs_color_r < 0:
new_bs_color_r = 0.2
if new_bs_color_g < 0:
new_bs_color_g = 0.2
if new_bs_color_b < 0:
new_bs_color_b = 0.2
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Mix"].inputs[1].default_value = list(new_bs_color)
elif material_name == "porcelain_2":
if transfer_flag == True:
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = tex_mix_prop # 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.8
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "porcelain_3":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix.001"].inputs[1])
else:
mat_nodes["Mix.001"].inputs[1].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Mix.001"].inputs[1].default_value
new_bs_color_r = bs_color[0] + random.uniform(-0.3, 0.3)
new_bs_color_g = bs_color[1] + random.uniform(-0.3, 0.3)
new_bs_color_b = bs_color[2] + random.uniform(-0.3, 0.3)
if new_bs_color_r < 0:
new_bs_color_r = 0.2
if new_bs_color_g < 0:
new_bs_color_g = 0.2
if new_bs_color_b < 0:
new_bs_color_b = 0.2
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Mix.001"].inputs[1].default_value = list(new_bs_color)
elif material_name == "porcelain_4":
if transfer_flag == True:
# if is_texture:
# mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
# mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
# else:
# mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
# mat_nodes["Glossy BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Glossy BSDF"].inputs[1].default_value = random.uniform(0.05, 0.15)
diff_new = mat_nodes.new(type='ShaderNodeBsdfDiffuse')
diff_new.name = 'Diffuse BSDF-new'
for key, input in enumerate(mat_nodes["Diffuse BSDF"].inputs):
diff_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Diffuse BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Diffuse BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
elif material_name == "porcelain_5":
if transfer_flag == True:
# if is_texture:
# mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
# mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
# else:
# mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
# mat_nodes["Glossy BSDF"].inputs[0].default_value = list(orign_base_color)
diff_new = mat_nodes.new(type='ShaderNodeBsdfDiffuse')
diff_new.name = 'Diffuse BSDF-new'
for key, input in enumerate(mat_nodes["Diffuse BSDF"].inputs):
diff_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Diffuse BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Diffuse BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
elif material_name == "porcelain_6":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Diffuse BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + random.uniform(-0.3, 0.3)
new_bs_color_g = bs_color[1] + random.uniform(-0.3, 0.3)
new_bs_color_b = bs_color[2] + random.uniform(-0.3, 0.3)
if new_bs_color_r < 0:
new_bs_color_r = 0.2
if new_bs_color_g < 0:
new_bs_color_g = 0.2
if new_bs_color_b < 0:
new_bs_color_b = 0.2
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(new_bs_color)
## plastic
elif material_name == "plastic_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_3":
mat_nodes["值(明度)"].outputs[0].default_value = random.uniform(0.05, 0.25)
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_5":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_6":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.012"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.021"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.022"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.033"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
mat_nodes["RGB.001"].outputs[0].default_value = list(orign_base_color)
"""
mat_nodes["RGB.002"].outputs[0].default_value = list(orign_base_color)
mat_nodes["RGB.003"].outputs[0].default_value = list(orign_base_color)
"""
## rubber
elif material_name == "rubber_0":
diff_new = mat_nodes.new(type='ShaderNodeBsdfDiffuse')
diff_new.name = 'Diffuse BSDF-new'
for key, input in enumerate(mat_nodes["Diffuse BSDF"].inputs):
diff_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF-new"].inputs[0])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Diffuse BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Diffuse BSDF"].outputs[0], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Diffuse BSDF-new"].outputs[0], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Mix Shader"].inputs[1])
elif material_name == "rubber_1":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["RGB Curves"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "rubber_2":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["RGB Curves"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "rubber_3":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "rubber_4":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
## plastic_specular
elif material_name == "plasticsp_0":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.001"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute"].inputs[0])
else:
mat_nodes["RGB.001"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB.001"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB.001"].outputs[0].default_value = list(new_bs_color)
elif material_name == "plasticsp_1":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
## paint_specular
elif material_name == "paintsp_0":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Diffuse BSDF"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB"].outputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_1":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[2])
mat_links.new(tex_node.outputs[0], mat_nodes["Mix.001"].inputs[1])
mat_links.new(tex_node.outputs[0], mat_nodes["Hue Saturation Value"].inputs[4])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB"].outputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_2":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Group"].inputs[0])
else:
mat_nodes["Group"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Group"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Group"].inputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_3":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
else:
# bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color = [random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1), 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_4":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF"].inputs[0])
mat_links.new(tex_node.outputs[0], mat_nodes["Glossy BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Glossy BSDF"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Glossy BSDF.001"].inputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["Principled BSDF"].inputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["Principled BSDF"].inputs[0].default_value = list(new_bs_color)
mat_nodes["Glossy BSDF"].inputs[0].default_value = list(new_bs_color)
mat_nodes["Glossy BSDF.001"].inputs[0].default_value = list(new_bs_color)
elif material_name == "paintsp_5":
if transfer_flag == True:
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Invert"].inputs[1])
mat_links.new(tex_node.outputs[0], mat_nodes["Reroute.002"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
else:
bs_color = mat_nodes["RGB"].outputs[0].default_value
new_bs_color_r = bs_color[0] + r_rand
new_bs_color_g = bs_color[1] + g_rand
new_bs_color_b = bs_color[2] + b_rand
if new_bs_color_r < 0:
new_bs_color_r = 0
if new_bs_color_g < 0:
new_bs_color_g = 0
if new_bs_color_b < 0:
new_bs_color_b = 0
if new_bs_color_r > 1:
new_bs_color_r = 1
if new_bs_color_g > 1:
new_bs_color_g = 1
if new_bs_color_b > 1:
new_bs_color_b = 1
new_bs_color = [new_bs_color_r, new_bs_color_g, new_bs_color_b, 1]
mat_nodes["RGB"].outputs[0].default_value = list(new_bs_color)
## rubber
elif material_name == "rubber_5":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 1.0
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Mix.005"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
## plastic
elif material_name == "plastic_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["Principled BSDF.001"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_4":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_7":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF.001"].inputs[0])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_8":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Group"].inputs[0])
else:
mat_nodes["Group"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_9":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_10":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_11":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_12":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[2])
else:
mat_nodes["RGB"].outputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_13":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "plastic_14":
mat_nodes["Math.005"].inputs[1].default_value = random.uniform(0.05, 0.3)
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
## paper
elif material_name == "paper_0":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = 0.9
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Mix.002"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "paper_1":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.8, 0.95)
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.8, 0.9)
mat_links.new(mat_nodes["Normal Map"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Image Texture.001"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
elif material_name == "paper_2":
bsdf_new = mat_nodes.new(type='ShaderNodeBsdfPrincipled')
bsdf_new.name = 'Principled BSDF-new'
for key, input in enumerate(mat_nodes["Principled BSDF"].inputs):
bsdf_new.inputs[key].default_value = input.default_value
mix_new = mat_nodes.new(type='ShaderNodeMixShader')
mix_new.name = 'Mix Shader-new'
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF-new"].inputs["Base Color"])
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.9, 0.95)
else:
mat_nodes["Principled BSDF-new"].inputs[0].default_value = list(orign_base_color)
mat_nodes["Mix Shader-new"].inputs[0].default_value = random.uniform(0.9, 0.95)
mat_links.new(mat_nodes["Bump"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[20])
mat_links.new(mat_nodes["Bright/Contrast"].outputs[0], mat_nodes["Principled BSDF-new"].inputs[7])
mat_links.new(mat_nodes["Principled BSDF"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[1])
mat_links.new(mat_nodes["Principled BSDF-new"].outputs["BSDF"], mat_nodes["Mix Shader-new"].inputs[2])
mat_links.new(mat_nodes["Mix Shader-new"].outputs[0], mat_nodes["Material Output"].inputs["Surface"])
## leather
elif material_name == "leather_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
elif material_name == "leather_3":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_4":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "leather_5":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["lether"].inputs[0])
else:
mat_nodes["lether"].inputs[0].default_value = list(orign_base_color)
elif material_name == "wood_0":
pass
elif material_name == "wood_1":
pass
elif material_name == "wood_2":
pass
elif material_name == "wood_3":
pass
elif material_name == "wood_4":
pass
elif material_name == "wood_5":
pass
elif material_name == "wood_6":
pass
elif material_name == "wood_7":
pass
elif material_name == "wood_8":
pass
elif material_name == "wood_9":
pass
## fabric
elif material_name == "fabric_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "fabric_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "fabric_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
## clay
elif material_name == "clay_0":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_1":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_2":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_3":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
elif material_name == "clay_4":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
else:
mat_nodes["Principled BSDF"].inputs[0].default_value = list(orign_base_color)
elif material_name == "clay_5":
if is_texture:
mat_links.new(tex_node.outputs[0], mat_nodes["Mix"].inputs[1])
else:
mat_nodes["Mix"].inputs[1].default_value = list(orign_base_color)
## glass
elif material_name == "glass_0":
pass
elif material_name == "glass_4":
pass
elif material_name == "glass_5":
pass
elif material_name == "glass_14":
pass
else:
print(material_name + " no change")
def set_modify_material(obj, material, obj_texture_img_list, mat_randomize_mode, is_transfer=True):
for mat_slot in obj.material_slots:
# print("-material_slots:",mat_slot)
if mat_slot.material:
if mat_slot.material.node_tree:
# print("--material:" + str(mat_slot.material.name))
srcmat = material
mat = srcmat.copy()
mat.name = mat_slot.material.name # rename
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
bsdf_node = mat_slot.material.node_tree.nodes.get("Principled BSDF", None)
if bsdf_node is not None:
tex_node = mat_slot.material.node_tree.nodes.new(type='ShaderNodeTexImage')
tex_node.name = 'objtexture_tex'
tex_node.extension = 'EXTEND'
# random pick real table image
flag = random.randint(0, len(obj_texture_img_list) - 1)
tex_node.image = obj_texture_img_list[flag]
# check if texture node exists
tex_node_orign = mat_slot.material.node_tree.nodes.get('objtexture_tex', None)
if tex_node_orign is not None:
# mat = mat_slot.material.copy()
# Get the bl_idname to create a new node of the same type
tex_node = mat_nodes.new(tex_node_orign.bl_idname)
texture_img = bpy.data.images[tex_node_orign.image.name]
# Assign the default values from the old node to the new node
tex_node.image = texture_img
tex_node.projection = 'SPHERE'
# tex_node.location = Vector((-800, 0))
###
mapping_node = mat_nodes.new(type='ShaderNodeMapping')
mapping_node.name = 'objtexture_mapping'
texcoord_node = mat_nodes.new(type='ShaderNodeTexCoord')
texcoord_node.name = 'objtexture_texcoord'
mat_links.new(mapping_node.outputs[0], tex_node.inputs[0])
mat_links.new(texcoord_node.outputs[0], mapping_node.inputs[0])
###
modify_material(mat_links, mat_nodes, srcmat.name, mat_randomize_mode, is_texture=True,
tex_node=tex_node, is_transfer=is_transfer)
else:
orign_base_color = mat_slot.material.node_tree.nodes["Principled BSDF"].inputs[0].default_value
if orign_base_color[0] == 0.0 and orign_base_color[1] == 0.0 and orign_base_color[2] == 0.0:
orign_base_color = [0.05, 0.05, 0.05, 1]
modify_material(mat_links, mat_nodes, srcmat.name, mat_randomize_mode, is_texture=False,
orign_base_color=orign_base_color, is_transfer=is_transfer)
# apply material
bpy.data.materials.remove(mat_slot.material)
mat_slot.material = mat
def set_modify_raw_material(obj):
for mat_slot in obj.material_slots:
if mat_slot.material:
if mat_slot.material.node_tree:
bsdf_node = mat_slot.material.node_tree.nodes.get("Principled BSDF", None)
if bsdf_node is not None:
tex_node_orign = mat_slot.material.node_tree.nodes.get("Image Texture", None)
if tex_node_orign is None:
orign_base_color = mat_slot.material.node_tree.nodes["Principled BSDF"].inputs[0].default_value
if orign_base_color[0] == 0.0 and orign_base_color[1] == 0.0 and orign_base_color[2] == 0.0:
mat = mat_slot.material.copy()
mat.name = mat_slot.material.name # rename
mat_nodes = mat.node_tree.nodes
mat_nodes["Principled BSDF"].inputs[0].default_value = list([0.05, 0.05, 0.05, 1])
bpy.data.materials.remove(mat_slot.material)
mat_slot.material = mat
def set_modify_table_material(obj, material, selected_realtable_img): ### realtable_img_list):
#print("--material:" + str(mat_slot.material.name))
srcmat = material
#print(srcmat.name)
mat = srcmat.copy()
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
tex_node = mat_nodes.new(type='ShaderNodeTexImage')
tex_node.name = 'realtable_tex'
tex_node.extension = 'EXTEND'
### flag = random.randint(0, len(realtable_img_list)-1)
tex_node.image = selected_realtable_img ### realtable_img_list[flag]
mapping_node = mat_nodes.new(type='ShaderNodeMapping')
mapping_node.name = 'realtable_mapping'
texcoord_node = mat_nodes.new(type='ShaderNodeTexCoord')
texcoord_node.name = 'realtable_texcoord'
mat_links.new(tex_node.outputs[0], mat_nodes["Principled BSDF"].inputs[0])
mat_links.new(mapping_node.outputs[0], tex_node.inputs[0])
mat_links.new(texcoord_node.outputs[2], mapping_node.inputs[0])
obj.active_material = mat
def set_modify_floor_material(obj, material, selected_realfloor_img):### realfloor_img_list):
srcmat = material
mat = srcmat.copy()
mat_links = mat.node_tree.links
mat_nodes = mat.node_tree.nodes
bsdfnode_list = [n for n in mat_nodes if isinstance(n, bpy.types.ShaderNodeBsdfPrincipled)]
if bsdfnode_list == []:
obj.active_material = material
else:
for bsdfnode in bsdfnode_list:
tex_node = mat_nodes.new(type='ShaderNodeTexImage')
tex_node.name = 'realfloor_tex'
tex_node.extension = 'REPEAT'
### flag = random.randint(0, len(realfloor_img_list)-1)
tex_node.image = selected_realfloor_img ### realfloor_img_list[flag]
mapping_node = mat_nodes.new(type='ShaderNodeMapping')
mapping_node.name = 'realfloor_mapping'
texcoord_node = mat_nodes.new(type='ShaderNodeTexCoord')
texcoord_node.name = 'realfloor_texcoord'
sc = random.uniform(1.00, 4.00)
mat_links.new(tex_node.outputs[0], bsdfnode.inputs[0]) #mat_nodes[bsdf_node_name].inputs[0])
mat_links.new(mapping_node.outputs[0], tex_node.inputs[0])
mat_links.new(texcoord_node.outputs[2], mapping_node.inputs[0])
obj.active_material = mat
================================================
FILE: src/rd/render.py
================================================
import os
import random
import bpy
import math
import numpy as np
from rd.modify_material import set_modify_table_material, set_modify_floor_material
from rd.render_utils import *
def blender_init_scene(code_root, log_root_dir, obj_texture_image_root_path, scene_type, urdfs_and_poses_dict, round_idx, logdir, check_seen_scene, material_type, gpuid, output_modality_dict):
if scene_type == "pile":
seed = 1143+830+round_idx
elif scene_type == "packed":
seed = 111143+170+round_idx
else:
seed = 43+round_idx
np.random.seed(seed)
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
DEVICE_LIST = [int(gpuid)]
obj_texture_image_idxfile = "test_paths.txt"
asset_root = code_root
env_map_path = os.path.join(asset_root, "envmap_lib_test")
real_table_image_root_path = os.path.join(asset_root, "realtable_test")
real_floor_image_root_path = os.path.join(asset_root, "realfloor_test")
emitter_pattern_path = os.path.join(asset_root, "pattern", "test_pattern.png")
default_background_texture_path = os.path.join(asset_root, "texture", "texture_0.jpg")
table_CAD_model_path = os.path.join(asset_root, "table_obj", "table.obj")
output_root_path = os.path.join(log_root_dir, "rendered_results/" + str(logdir).split("/")[-1])
obj_uid_list = [str(uid) for uid in urdfs_and_poses_dict]
obj_scale_list = [value[0] for value in urdfs_and_poses_dict.values()]
obj_quat_list = [value[1][[3, 0, 1, 2]] for value in urdfs_and_poses_dict.values()]
obj_trans_list = []
for value in urdfs_and_poses_dict.values():
T = value[2]
T = T + tsdf2blender_coord_T_shift
obj_trans_list.append(T)
urdf_path_list = [os.path.join(value[3]) for value in urdfs_and_poses_dict.values()] #"/".join(code_root.split("/")[:-2]),
max_instance_num = 20
if not os.path.exists(output_root_path):
os.makedirs(output_root_path)
# generate CAD model list
CAD_model_list = generate_CAD_model_list(scene_type, urdf_path_list, obj_uid_list)
renderer = BlenderRenderer(viewport_size_x=camera_width, viewport_size_y=camera_height, DEVICE_LIST=DEVICE_LIST)
renderer.loadImages(emitter_pattern_path, env_map_path, real_table_image_root_path, real_floor_image_root_path, obj_texture_image_root_path, obj_texture_image_idxfile, check_seen_scene)
renderer.addEnvMap()
renderer.addBackground(background_size, background_position, background_scale, default_background_texture_path)
renderer.addMaterialLib(material_class_instance_pairs) ###
renderer.addMaskMaterial(max_instance_num)
renderer.addNOCSMaterial()
renderer.addNormalMaterial()
renderer.clearModel()
# set scene output path
path_scene = output_root_path #os.path.join(output_root_path, uuid.uuid4().hex, "init") ### "scene_"+str(SCENE_NUM).zfill(4))
if os.path.exists(path_scene)==False:
os.makedirs(path_scene)
# camera pose list, environment light list and background material_listz
quaternion_list = []
translation_list = []
# environment map list
env_map_id_list = []
rotation_elur_z_list = []
# background material list
background_material_list = []
# table material list
table_material_list = []
look_at = look_at_shift
quat_list, trans_list, rot_list = genCameraPosition(look_at)
rot_array = np.array(rot_list) # (256, 3, 3)
trans_array = np.array(trans_list) # (256, 3, 1)
cam_RT = np.concatenate([rot_array, trans_array], 2)
zero_one = np.expand_dims([[0, 0, 0, 1]],0).repeat(rot_array.shape[0],axis=0)
cam_RT = np.concatenate([cam_RT, zero_one], 1) # (256, 4, 4)
# generate camara pose list
for i in range(NUM_FRAME_PER_SCENE):
quaternion = quat_list[i] #### cam_pose_list[i][0]
translation = trans_list[i] #### cam_pose_list[i][1]
quaternion_list.append(quaternion)
translation_list.append(translation)
flag_env_map = random.randint(0, len(renderer.env_map) - 1)
flag_env_map_rot = random.uniform(-math.pi, math.pi)
flag_realfloor = random.randint(0, len(renderer.realfloor_img_list) - 1)
flag_realtable = random.randint(0, len(renderer.realtable_img_list) - 1)
# generate environment map list
env_map_id_list.append(flag_env_map)
rotation_elur_z_list.append(flag_env_map_rot)
# generate background material list
if my_material_randomize_mode == 'raw':
background_material_list.append(renderer.my_material['default_background'])
else:
material_selected = random.sample(renderer.my_material['background'], 1)[0] ### renderer.my_material['background'][1]
background_material_list.append(material_selected)
material_selected = random.sample(renderer.my_material['table'], 1)[0] ### renderer.my_material['table'][0]
table_material_list.append(material_selected)
# read objects from floder
meta_output = {}
select_model_list = []
select_model_list_other = []
select_model_list_transparent = []
select_model_list_dis = []
select_number = 1
for item in CAD_model_list:
if item in ['other']:
test = CAD_model_list[item]
for model in test:
select_model_list.append(model)
else:
raise ValueError("No such category!")
# table model
renderer.loadModel(table_CAD_model_path)
obj = bpy.data.objects['table']
# resize table, unit: m
class_scale = 0.001
obj.scale = (class_scale, class_scale, class_scale)
y_transform = np.array([[0,0,-1],[0,1,0],[1,0,0]])
transform = y_transform # z_transform @ y_transform
obj_world_pose_quat = quaternionFromRotMat(transform)
obj_world_pose_T_shift = np.array([0,0,-0.0751])
obj_world_pose_T = obj_world_pose_T_shift ### np.array([cam_pose_T[0],cam_pose_T[1],0]) + obj_world_pose_T_shift
setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
obj_world_pose_T_shift = np.array([0,0,0])
obj_world_pose_T = obj_world_pose_T_shift ### np.array([cam_pose_T[0],cam_pose_T[1],0]) + obj_world_pose_T_shift
# setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
bpy.ops.mesh.primitive_plane_add(size=1., enter_editmode=False, align='WORLD', location=obj_world_pose_T)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
obj = bpy.data.objects['Plane']
obj.name = 'tableplane'
obj.data.name = 'tableplane'
obj.scale = (0.898, 1.3, 1.)
###
instance_id = 1
# set object parameters
imported_obj_name_list = []
for model in select_model_list:
instance_path = model[0]
class_name = model[1]
instance_uid = model[2]
instance_folder = model[0].split('/')[-1][:-4]
instance_name = str(instance_id) + "_" + class_name + "_" + instance_folder + "_" + instance_uid ### class_folder + "_" + instance_folder
material_type_in_mixed_mode = generate_material_type(instance_name, class_material_pairs, instance_material_except_pairs, instance_material_include_pairs, material_class_instance_pairs, material_type)
# download CAD model and rename
renderer.loadModel(instance_path)
import_obj_name = instance_folder #bpy.data.objects.keys()["instance_folder"]
obj = bpy.data.objects[import_obj_name]
obj.name = instance_name
obj.data.name = instance_name
obj_world_pose_T = obj_trans_list[instance_id-1] ### obj_pose_list[instance_id-1][:3,3]
obj_world_pose_quat = obj_quat_list[instance_id-1] ### quaternionFromRotMat(obj_world_pose_R)
setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
# set object as rigid body
setRigidBody(obj)
# set material
renderer.set_material_randomize_mode(class_material_pairs, my_material_randomize_mode, obj, material_type_in_mixed_mode)
# generate meta file
class_scale = obj_scale_list[instance_id-1] ### random.uniform(g_synset_name_scale_pairs[class_name][0], g_synset_name_scale_pairs[class_name][1])
obj.scale = (class_scale, class_scale, class_scale)
# query material type
material_class_id = None
for key in material_class_instance_pairs:
if material_type_in_mixed_mode == 'raw':
material_class_id = material_class_id_dict[material_type_in_mixed_mode]
break
elif material_type_in_mixed_mode in material_class_instance_pairs[key]:
material_class_id = material_class_id_dict[key]
break
if material_class_id == None:
raise ValueError("material_class_id error!")
meta_output[str(instance_id)] = [#str(g_synset_name_label_pairs[class_name]),
### class_folder,
str(instance_folder),
### str(class_scale),
str(material_class_id), ###str(material_name_label_pairs[material_type_in_mixed_mode])]
str(material_type_id_dict[material_type_in_mixed_mode])
]
instance_id += 1
if output_modality_dict['IR'] or output_modality_dict['RGB']:
renderer.setEnvMap(env_map_id_list[0], rotation_elur_z_list[0])
# pick real floor image
selected_realfloor_img = renderer.realfloor_img_list[flag_realfloor]
# pcik real table image
selected_realtable_img = renderer.realtable_img_list[flag_realtable]
for obj in bpy.data.objects:
if obj.type == "MESH" and obj.name.split('_')[0] == 'background':
if obj.name == 'background_0':
set_modify_floor_material(obj, background_material_list[0], selected_realfloor_img) ### renderer.realfloor_img_list)
else:
background_0_obj = bpy.data.objects['background_0']
obj.active_material = background_0_obj.material_slots[0].material
elif obj.type == "MESH" and obj.name == 'table':
set_modify_table_material(obj, table_material_list[0], selected_realtable_img)### renderer.realtable_img_list)
elif obj.type == "MESH" and obj.name == 'tableplane':
table_obj = bpy.data.objects['table']
obj.active_material = table_obj.material_slots[0].material
return renderer, quaternion_list, translation_list, path_scene#, obj_trans_list, obj_quat_list
def blender_update_sceneobj(obj_name_list, obj_trans_list, obj_quat_list, obj_uid_list):
for obj_name in bpy.data.objects.keys():
obj = bpy.data.objects[obj_name]
if obj.type == 'MESH' and obj_name[0:10] != "background" and obj_name not in ['camera_l', 'camera_r', 'light_emitter', 'table', 'tableplane']:
obj_uid = obj_name.split("_")[-1]
if obj_uid not in obj_uid_list:
print("[V]blender_update_sceneobj: obj_uid [not in] obj_uid_list: ", obj_name, obj_name_list, obj_uid, obj_uid_list)
obj.hide_render = True
else:
obj.hide_render = False
obj_world_pose_T = obj_trans_list[obj_uid_list.index(obj_uid)] + tsdf2blender_coord_T_shift
obj_world_pose_quat = obj_quat_list[obj_uid_list.index(obj_uid)]
print("[V]blender_update_sceneobj: obj_uid [in] obj_uid_list: ", obj_name, obj_name_list, obj_uid, obj_uid_list, obj_world_pose_T, obj_world_pose_quat)
setModelPosition(obj, obj_world_pose_T, obj_world_pose_quat)
def blender_render(renderer, quaternion_list, translation_list, path_scene, render_frame_list, output_modality_dict, camera_focal, is_init=False):
# set the key frame
scene = bpy.data.scenes['Scene']
camera_fov = 2 * math.atan(camera_width / (2 * camera_focal))
# render IR image and RGB image
if output_modality_dict['IR'] or output_modality_dict['RGB']:
if is_init:
renderer.src_energy_for_rgb_render = bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value
for i in render_frame_list:
renderer.setCamera(quaternion_list[i], translation_list[i], camera_fov, baseline_distance)
renderer.setLighting()
# render RGB image
if output_modality_dict['RGB']:
rgb_dir_path = os.path.join(path_scene, 'rgb')
if os.path.exists(rgb_dir_path) == False:
os.makedirs(rgb_dir_path)
renderer.render_mode = "RGB"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = rgb_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
# render IR image
if output_modality_dict['IR']:
ir_l_dir_path = os.path.join(path_scene, 'ir_l')
if os.path.exists(ir_l_dir_path)==False:
os.makedirs(ir_l_dir_path)
ir_r_dir_path = os.path.join(path_scene, 'ir_r')
if os.path.exists(ir_r_dir_path)==False:
os.makedirs(ir_r_dir_path)
renderer.render_mode = "IR"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = ir_l_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
camera = bpy.data.objects['camera_r']
scene.camera = camera
save_path = ir_r_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
# render normal map and depth map
if output_modality_dict['Normal']:
# set normal as material
for obj in bpy.data.objects:
if obj.type == 'MESH':
obj.data.materials.clear()
obj.active_material = renderer.my_material["normal"]
# render normal map
for i in render_frame_list:
renderer.setCamera(quaternion_list[i], translation_list[i], camera_fov, baseline_distance)
normal_dir_path = os.path.join(path_scene, 'normal')
if os.path.exists(normal_dir_path)==False:
os.makedirs(normal_dir_path)
depth_dir_path = os.path.join(path_scene, 'depth')
if os.path.exists(depth_dir_path)==False:
os.makedirs(depth_dir_path)
renderer.render_mode = "Normal"
camera = bpy.data.objects['camera_l']
scene.camera = camera
save_path = normal_dir_path
save_name = str(i).zfill(4)
renderer.render(save_name, save_path)
context = bpy.context
for ob in context.selected_objects:
ob.animation_data_clear()
================================================
FILE: src/rd/render_utils.py
================================================
import os
import random
import bpy
import math
import numpy as np
from mathutils import Vector, Matrix
from bpy_extras.object_utils import world_to_camera_view
from rd.modify_material import set_modify_material, set_modify_raw_material, set_modify_table_material, set_modify_floor_material
# render parameter
RENDERING_PATH = os.getcwd()
LIGHT_EMITTER_ENERGY = 5
LIGHT_ENV_MAP_ENERGY_IR = 0.035
LIGHT_ENV_MAP_ENERGY_RGB = 1.0
CYCLES_SAMPLE = 32
CAMERA_TYPE = "realsense"
NUM_FRAME_PER_SCENE = 24
# view point parameter
look_at_shift = np.array([0,0,0])
num_point_ver = 6
num_point_hor = 4
beta_range = (15*math.pi/180, 45*math.pi/180)
r = 0.5
tsdf2blender_coord_T_shift = np.array([-0.15, -0.15, -0.0503])
TABLE_CAD_MODEL_HEIGHT = 0.75
# material randomization mode (transparent, specular, mixed, raw)
my_material_randomize_mode = 'mixed'
# set depth sensor parameter
camera_width = 640
camera_height = 360
baseline_distance = 0.055
# set background parameter
background_size = 3.
background_position = (0., 0., 0.)
background_scale = (1., 1., 1.)
# set camera randomize paramater
start_point_range = ((0.5, 0.95), (-0.6, 0.6, -0.6, 0.6))
up_range = (-0.18, -0.18, -0.18, 0.18)
look_at_range = (background_position[0] - 0.05, background_position[0] + 0.05,
background_position[1] - 0.05, background_position[1] + 0.05,
background_position[2] - 0.05, background_position[2] + 0.05)
g_syn_light_num_lowbound = 4
g_syn_light_num_highbound = 6
g_syn_light_dist_lowbound = 8
g_syn_light_dist_highbound = 12
g_syn_light_azimuth_degree_lowbound = 0
g_syn_light_azimuth_degree_highbound = 360
g_syn_light_elevation_degree_lowbound = 0
g_syn_light_elevation_degree_highbound = 90
g_syn_light_energy_mean = 3
g_syn_light_energy_std = 0.5
g_syn_light_environment_energy_lowbound = 0
g_syn_light_environment_energy_highbound = 1
g_shape_synset_name_pairs_all = {'02691156': 'aeroplane',
'02747177': 'ashtray',
'02773838': 'backpack',
'02801938': 'basket',
'02808440': 'tub', # bathtub
'02818832': 'bed',
'02828884': 'bench',
'02834778': 'bicycle',
'02843684': 'mailbox', # missing in objectnet3d, birdhouse, use view distribution of mailbox
'02858304': 'boat',
'02871439': 'bookshelf',
'02876657': 'bottle',
'02880940': 'bowl', # missing in objectnet3d, bowl, use view distribution of plate
'02924116': 'bus',
'02933112': 'cabinet',
'02942699': 'camera',
'02946921': 'can',
'02954340': 'cap',
'02958343': 'car',
'02992529': 'cellphone',
'03001627': 'chair',
'03046257': 'clock',
'03085013': 'keyboard',
'03207941': 'dishwasher',
'03211117': 'tvmonitor',
'03261776': 'headphone',
'03325088': 'faucet',
'03337140': 'filing_cabinet',
'03467517': 'guitar',
'03513137': 'helmet',
'03593526': 'jar',
'03624134': 'knife',
'03636649': 'lamp',
'03642806': 'laptop',
'03691459': 'speaker',
'03710193': 'mailbox',
'03759954': 'microphone',
'03761084': 'microwave',
'03790512': 'motorbike',
'03797390': 'mug', # missing in objectnet3d, mug, use view distribution of cup
'03928116': 'piano',
'03938244': 'pillow',
'03948459': 'rifle', # missing in objectnet3d, pistol, use view distribution of rifle
'03991062': 'pot',
'04004475': 'printer',
'04074963': 'remote_control',
'04090263': 'rifle',
'04099429': 'road_pole', # missing in objectnet3d, rocket, use view distribution of road_pole
'04225987': 'skateboard',
'04256520': 'sofa',
'04330267': 'stove',
'04379243': 'diningtable', # use view distribution of dining_table
'04401088': 'telephone',
'04460130': 'road_pole', # missing in objectnet3d, tower, use view distribution of road_pole
'04468005': 'train',
'04530566': 'washing_machine',
'04554684': 'dishwasher'} # washer, use view distribution of dishwasher
g_synset_name_label_pairs = {#'aeroplane': 7,
#'bottle': 1,
#'bowl': 2,
#'camera': 3,
#'can': 4,
#'car': 5,
#'mug': 6,
'other': 0}
material_class_instance_pairs = {'specular': ['metal', 'paintsp'], # 'porcelain','plasticsp',
'transparent': ['glass'],
'diffuse': ['plastic','rubber','paper','leather','wood','clay','fabric'],
'background': ['background']}
# material list
class_material_pairs = {'specular': ['other'],
'transparent': ['other'],
'diffuse': ['other']}
instance_material_except_pairs = {'metal': [],
'porcelain': [],
'plasticsp': [],
'paintsp':[],
'glass': [],#[8,9,18,19,20,24,25,26,27,28,29,30,31,32,34,43,59,72],
'plastic': [],
'rubber': [],
'leather': [],
'wood':[],
'paper':[],
'fabric':[],
'clay':[],
}
instance_material_include_pairs = {
}
material_class_id_dict = {'raw': 0,
'diffuse': 1,
'transparent': 2,
'specular': 3}
material_type_id_dict = {'raw': 0,
'metal': 1,
'porcelain': 2,
'plasticsp': 3,
'paintsp':4,
'glass': 5,
'plastic': 6,
'rubber': 7,
'leather': 8,
'wood':9,
'paper':10,
'fabric':11,
'clay':12,
}
def obj_centered_camera_pos(dist, azimuth_deg, elevation_deg):
phi = float(elevation_deg) / 180 * math.pi
theta = float(azimuth_deg) / 180 * math.pi
x = (dist * math.cos(theta) * math.cos(phi))
y = (dist * math.sin(theta) * math.cos(phi))
z = (dist * math.sin(phi))
return (x, y, z)
def quaternionFromYawPitchRoll(yaw, pitch, roll):
c1 = math.cos(yaw / 2.0)
c2 = math.cos(pitch / 2.0)
c3 = math.cos(roll / 2.0)
s1 = math.sin(yaw / 2.0)
s2 = math.sin(pitch / 2.0)
s3 = math.sin(roll / 2.0)
q1 = c1 * c2 * c3 + s1 * s2 * s3
q2 = c1 * c2 * s3 - s1 * s2 * c3
q3 = c1 * s2 * c3 + s1 * c2 * s3
q4 = s1 * c2 * c3 - c1 * s2 * s3
return (q1, q2, q3, q4)
def camPosToQuaternion(cx, cy, cz):
q1a = 0
q1b = 0
q1c = math.sqrt(2) / 2
q1d = math.sqrt(2) / 2
camDist = math.sqrt(cx * cx + cy * cy + cz * cz)
cx = cx / camDist
cy = cy / camDist
cz = cz / camDist
t = math.sqrt(cx * cx + cy * cy)
tx = cx / t
ty = cy / t
yaw = math.acos(ty)
if tx > 0:
yaw = 2 * math.pi - yaw
pitch = 0
tmp = min(max(tx*cx + ty*cy, -1),1)
#roll = math.acos(tx * cx + ty * cy)
roll = math.acos(tmp)
if cz < 0:
roll = -roll
print("%f %f %f" % (yaw, pitch, roll))
q2a, q2b, q2c, q2d = quaternionFromYawPitchRoll(yaw, pitch, roll)
q1 = q1a * q2a - q1b * q2b - q1c * q2c - q1d * q2d
q2 = q1b * q2a + q1a * q2b + q1d * q2c - q1c * q2d
q3 = q1c * q2a - q1d * q2b + q1a * q2c + q1b * q2d
q4 = q1d * q2a + q1c * q2b - q1b * q2c + q1a * q2d
return (q1, q2, q3, q4)
def camRotQuaternion(cx, cy, cz, theta):
theta = theta / 180.0 * math.pi
camDist = math.sqrt(cx * cx + cy * cy + cz * cz)
cx = -cx / camDist
cy = -cy / camDist
cz = -cz / camDist
q1 = math.cos(theta * 0.5)
q2 = -cx * math.sin(theta * 0.5)
q3 = -cy * math.sin(theta * 0.5)
q4 = -cz * math.sin(theta * 0.5)
return (q1, q2, q3, q4)
def quaternionProduct(qx, qy):
a = qx[0]
b = qx[1]
c = qx[2]
d = qx[3]
e = qy[0]
f = qy[1]
g = qy[2]
h = qy[3]
q1 = a * e - b * f - c * g - d * h
q2 = a * f + b * e + c * h - d * g
q3 = a * g - b * h + c * e + d * f
q4 = a * h + b * g - c * f + d * e
return (q1, q2, q3, q4)
def quaternionToRotation(q):
w, x, y, z = q
r00 = 1 - 2 * y ** 2 - 2 * z ** 2
r01 = 2 * x * y + 2 * w * z
r02 = 2 * x * z - 2 * w * y
r10 = 2 * x * y - 2 * w * z
r11 = 1 - 2 * x ** 2 - 2 * z ** 2
r12 = 2 * y * z + 2 * w * x
r20 = 2 * x * z + 2 * w * y
r21 = 2 * y * z - 2 * w * x
r22 = 1 - 2 * x ** 2 - 2 * y ** 2
r = [[r00, r01, r02], [r10, r11, r12], [r20, r21, r22]]
return r
def quaternionToRotation_xyzw(q):
x, y, z, w = q
r00 = 1 - 2 * y ** 2 - 2 * z ** 2
r01 = 2 * x * y + 2 * w * z
r02 = 2 * x * z - 2 * w * y
r10 = 2 * x * y - 2 * w * z
r11 = 1 - 2 * x ** 2 - 2 * z ** 2
r12 = 2 * y * z + 2 * w * x
r20 = 2 * x * z + 2 * w * y
r21 = 2 * y * z - 2 * w * x
r22 = 1 - 2 * x ** 2 - 2 * y ** 2
r = [[r00, r01, r02], [r10, r11, r12], [r20, r21, r22]]
return r
def quaternionFromRotMat(rotation_matrix):
rotation_matrix = np.reshape(rotation_matrix, (1, 9))[0]
w = math.sqrt(rotation_matrix[0]+rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
x = math.sqrt(rotation_matrix[0]-rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
y = math.sqrt(-rotation_matrix[0]+rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
z = math.sqrt(-rotation_matrix[0]-rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
a = [w,x,y,z]
m = a.index(max(a))
if m == 0:
x = (rotation_matrix[7]-rotation_matrix[5])/(4*w)
y = (rotation_matrix[2]-rotation_matrix[6])/(4*w)
z = (rotation_matrix[3]-rotation_matrix[1])/(4*w)
if m == 1:
w = (rotation_matrix[7]-rotation_matrix[5])/(4*x)
y = (rotation_matrix[1]+rotation_matrix[3])/(4*x)
z = (rotation_matrix[6]+rotation_matrix[2])/(4*x)
if m == 2:
w = (rotation_matrix[2]-rotation_matrix[6])/(4*y)
x = (rotation_matrix[1]+rotation_matrix[3])/(4*y)
z = (rotation_matrix[5]+rotation_matrix[7])/(4*y)
if m == 3:
w = (rotation_matrix[3]-rotation_matrix[1])/(4*z)
x = (rotation_matrix[6]+rotation_matrix[2])/(4*z)
y = (rotation_matrix[5]+rotation_matrix[7])/(4*z)
quaternion = (w,x,y,z)
return quaternion
def quaternionFromRotMat_xyzw(rotation_matrix):
rotation_matrix = np.reshape(rotation_matrix, (1, 9))[0]
w = math.sqrt(rotation_matrix[0]+rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
x = math.sqrt(rotation_matrix[0]-rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
y = math.sqrt(-rotation_matrix[0]+rotation_matrix[4]-rotation_matrix[8]+1 + 1e-6)/2
z = math.sqrt(-rotation_matrix[0]-rotation_matrix[4]+rotation_matrix[8]+1 + 1e-6)/2
a = [x,y,z,w]
m = a.index(max(a))
if m == 0:
x = (rotation_matrix[7]-rotation_matrix[5])/(4*w)
y = (rotation_matrix[2]-rotation_matrix[6])/(4*w)
z = (rotation_matrix[3]-rotation_matrix[1])/(4*w)
if m == 1:
w = (rotation_matrix[7]-rotation_matrix[5])/(4*x)
y = (rotation_matrix[1]+rotation_matrix[3])/(4*x)
z = (rotation_matrix[6]+rotation_matrix[2])/(4*x)
if m == 2:
w = (rotation_matrix[2]-rotation_matrix[6])/(4*y)
x = (rotation_matrix[1]+rotation_matrix[3])/(4*y)
z = (rotation_matrix[5]+rotation_matrix[7])/(4*y)
if m == 3:
w = (rotation_matrix[3]-rotation_matrix[1])/(4*z)
x = (rotation_matrix[6]+rotation_matrix[2])/(4*z)
y = (rotation_matrix[5]+rotation_matrix[7])/(4*z)
quaternion = (x,y,z,w)
return quaternion
def rotVector(q, vector_ori):
r = quaternionToRotation(q)
x_ori = vector_ori[0]
y_ori = vector_ori[1]
z_ori = vector_ori[2]
x_rot = r[0][0] * x_ori + r[1][0] * y_ori + r[2][0] * z_ori
y_rot = r[0][1] * x_ori + r[1][1] * y_ori + r[2][1] * z_ori
z_rot = r[0][2] * x_ori + r[1][2] * y_ori + r[2][2] * z_ori
return (x_rot, y_rot, z_rot)
def cameraLPosToCameraRPos(q_l, pos_l, baseline_dis):
vector_camera_l_y = (1, 0, 0)
vector_rot = rotVector(q_l, vector_camera_l_y)
pos_r = (pos_l[0] + vector_rot[0] * baseline_dis,
pos_l[1] + vector_rot[1] * baseline_dis,
pos_l[2] + vector_rot[2] * baseline_dis)
return pos_r
def getRTFromAToB(pointCloudA, pointCloudB):
muA = np.mean(pointCloudA, axis=0)
muB = np.mean(pointCloudB, axis=0)
zeroMeanA = pointCloudA - muA
zeroMeanB = pointCloudB - muB
covMat = np.matmul(np.transpose(zeroMeanA), zeroMeanB)
U, S, Vt = np.linalg.svd(covMat)
R = np.matmul(Vt.T, U.T)
if np.linalg.det(R) < 0:
print("[V]getRTFromAToB: Reflection detected")
Vt[2, :] *= -1
R = Vt.T * U.T
T = (-np.matmul(R, muA.T) + muB.T).reshape(3, 1)
return R, T
def cameraPositionRandomize(start_point_range, look_at_range, up_range):
r_range, vector_range = start_point_range
r_min, r_max = r_range
x_min, x_max, y_min, y_max = vector_range
r = random.uniform(r_min, r_max)
x = random.uniform(x_min, x_max)
y = random.uniform(y_min, y_max)
z = math.sqrt(1 - x**2 - y**2)
vector_camera_axis = np.array([x, y, z])
x_min, x_max, y_min, y_max = up_range
x = random.uniform(x_min, x_max)
y = random.uniform(y_min, y_max)
z = math.sqrt(1 - x**2 - y**2)
up = np.array([x, y, z])
x_min, x_max, y_min, y_max, z_min, z_max = look_at_range
look_at = np.array([random.uniform(x_min, x_max),
random.uniform(y_min, y_max),
random.uniform(z_min, z_max)])
position = look_at + r * vector_camera_axis
vectorZ = - (look_at - position)/np.linalg.norm(look_at - position)
vectorX = np.cross(up, vectorZ)/np.linalg.norm(np.cross(up, vectorZ))
vectorY = np.cross(vectorZ, vectorX)/np.linalg.norm(np.cross(vectorX, vectorZ))
# points in camera coordinates
pointSensor= np.array([[0., 0., 0.], [1., 0., 0.], [0., 2., 0.], [0., 0., 3.]])
# points in world coordinates
pointWorld = np.array([position,
position + vectorX,
position + vectorY * 2,
position + vectorZ * 3])
resR, resT = getRTFromAToB(pointSensor, pointWorld)
resQ = quaternionFromRotMat(resR)
return resQ, resT
def genCameraPosition(look_at):
quat_list = []
rot_list = []
trans_list = []
position_list = []
# alpha:
alpha = 0
alpha_delta = (2 * math.pi) / num_point_ver
for i in range(num_point_ver):
alpha = alpha + alpha_delta
flag_x = 1
flag_y = 1
alpha1 = alpha
if alpha > math.pi/2 and alpha <= math.pi:
alpha1 = math.pi - alpha
flag_x = -1
flag_y = 1
elif alpha > math.pi and alpha <= math.pi*(3/2):
alpha1 = alpha - math.pi
flag_x = -1
flag_y = -1
elif alpha > math.pi*(3/2):
alpha1 = math.pi*2 - alpha
flag_x = 1
flag_y = -1
beta = beta_range[0]
beta_delta = (beta_range[1]-beta_range[0])/(num_point_hor-1)
for j in range(num_point_hor):
if j != 0:
beta = beta + beta_delta
x = flag_x * (r * math.sin(beta)) * math.cos(alpha1)
y = flag_y * (r * math.sin(beta)) * math.sin(alpha1)
z = r * math.cos(beta)
position = np.array([x, y, z]) + look_at
look_at = look_at
up = np.array([0, 0, 1])
vectorZ = - (look_at - position)/np.linalg.norm(look_at - position)
vectorX = np.cross(up, vectorZ)/np.linalg.norm(np.cross(up, vectorZ))
vectorY = np.cross(vectorZ, vectorX)/np.linalg.norm(np.cross(vectorX, vectorZ))
# points in camera coordinates
pointSensor= np.array([[0., 0., 0.], [1., 0., 0.], [0., 2., 0.], [0., 0., 3.]])
# points in world coordinates
pointWorld = np.array([position,
position + vectorX,
position + vectorY * 2,
position + vectorZ * 3])
resR, resT = getRTFromAToB(pointSensor, pointWorld)
resQ = quaternionFromRotMat(resR)
quat_list.append(resQ)
rot_list.append(resR)
trans_list.append(resT)
position_list.append(position)
return quat_list, trans_list, rot_list
def quanternion_mul(q1, q2):
s1 = q1[0]
v1 = np.array(q1[1:])
s2 = q2[0]
v2 = np.array(q2[1:])
s = s1 * s2 - np.dot(v1, v2)
v = s1 * v2 + s2 * v1 + np.cross(v1, v2)
return (s, v[0], v[1], v[2])
class BlenderRenderer(object):
def __init__(self, viewport_size_x=640, viewport_size_y=360, DEVICE_LIST=None):
'''
viewport_size_x, viewport_size_y: rendering viewport resolution
'''
self.DEVICE_LIST = DEVICE_LIST
# remove all objects, cameras and lights
for obj in bpy.data.meshes:
bpy.data.meshes.remove(obj)
for cam in bpy.data.cameras:
bpy.data.cameras.remove(cam)
for light in bpy.data.lights:
bpy.data.lights.remove(light)
for obj in bpy.data.objects:
bpy.data.objects.remove(obj, do_unlink=True)
render_context = bpy.context.scene.render
# add left camera
camera_l_data = bpy.data.cameras.new(name="camera_l")
camera_l_object = bpy.data.objects.new(name="camera_l", object_data=camera_l_data)
bpy.context.collection.objects.link(camera_l_object)
# add right camera
camera_r_data = bpy.data.cameras.new(name="camera_r")
camera_r_object = bpy.data.objects.new(name="camera_r", object_data=camera_r_data)
bpy.context.collection.objects.link(camera_r_object)
camera_l = bpy.data.objects["camera_l"]
camera_r = bpy.data.objects["camera_r"]
# set the camera postion and orientation so that it is in
# the front of the object
camera_l.location = (1, 0, 0)
camera_r.location = (1, 0, 0)
# add emitter light
light_emitter_data = bpy.data.lights.new(name="light_emitter", type='SPOT')
light_emitter_object = bpy.data.objects.new(name="light_emitter", object_data=light_emitter_data)
bpy.context.collection.objects.link(light_emitter_object)
light_emitter = bpy.data.objects["light_emitter"]
light_emitter.location = (1, 0, 0)
light_emitter.data.energy = LIGHT_EMITTER_ENERGY
# render setting
render_context.resolution_percentage = 100
self.render_context = render_context
self.camera_l = camera_l
self.camera_r = camera_r
self.light_emitter = light_emitter
self.model_loaded = False
self.background_added = None
self.render_context.resolution_x = viewport_size_x
self.render_context.resolution_y = viewport_size_y
self.my_material = {}
self.render_mode = 'IR'
# output setting
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.compression = 0
self.render_context.image_settings.color_mode = 'BW'
self.render_context.image_settings.color_depth = '8'
# cycles setting
self.render_context.engine = 'CYCLES'
bpy.context.scene.cycles.progressive = 'BRANCHED_PATH'
bpy.context.scene.cycles.use_denoising = True
bpy.context.scene.cycles.denoiser = 'NLM'
bpy.context.scene.cycles.film_exposure = 0.5
# self.render_context.use_antialiasing = False
##########
bpy.context.scene.view_layers["View Layer"].use_sky = True
##########
# switch on nodes
bpy.context.scene.use_nodes = True
tree = bpy.context.scene.node_tree
links = tree.links
# clear default nodes
for n in tree.nodes:
tree.nodes.remove(n)
# create input render layer node
rl = tree.nodes.new('CompositorNodeRLayers')
# create output node
self.fileOutput = tree.nodes.new(type="CompositorNodeOutputFile")
self.fileOutput.base_path = "./new_data/0000"
self.fileOutput.format.file_format = 'OPEN_EXR'
self.fileOutput.format.color_depth= '32'
self.fileOutput.file_slots[0].path = 'depth#'
links.new(rl.outputs[2], self.fileOutput.inputs[0])
# depth sensor pattern
self.pattern = []
# environment map
self.env_map = []
self.realtable_img_list = []
self.realfloor_img_list = []
self.obj_texture_img_list = []
self.src_energy_for_rgb_render = 0
def loadImages(self, pattern_path, env_map_path, real_table_image_root_path, real_floor_image_root_path, obj_texture_image_root_path, obj_texture_image_idxfile, check_seen_scene):
# load pattern image
self.pattern = bpy.data.images.load(filepath=pattern_path)
if check_seen_scene:
env_map_path_list = os.listdir(env_map_path)
real_table_image_root_path_list = os.listdir(real_table_image_root_path)
real_floor_image_root_path_list = os.listdir(real_floor_image_root_path)
else:
env_map_path_list = sorted(os.listdir(env_map_path))
real_table_image_root_path_list = sorted(os.listdir(real_table_image_root_path))
real_floor_image_root_path_list = sorted(os.listdir(real_floor_image_root_path))
# load env map
for item in env_map_path_list:
if item.split('.')[-1] == 'hdr':
self.env_map.append(bpy.data.images.load(filepath=os.path.join(env_map_path, item)))
# load real table images
for item in real_table_image_root_path_list:
if item.split('.')[-1] == 'jpg':
self.realtable_img_list.append(bpy.data.images.load(filepath=os.path.join(real_table_image_root_path, item)))
# load real floor images
for item in real_floor_image_root_path_list:
if item.split('.')[-1] == 'jpg':
self.realfloor_img_list.append(bpy.data.images.load(filepath=os.path.join(real_floor_image_root_path, item)))
# load obj texture images
f_teximg_idx = open(os.path.join(obj_texture_image_root_path, obj_texture_image_idxfile),"r")
lines = f_teximg_idx.readlines()
for item in lines:
item = item[:-1]
self.obj_texture_img_list.append(bpy.data.images.load(filepath=os.path.join(obj_texture_image_root_path, "images", item)))
def addEnvMap(self):
# Get the environment node tree of the current scene
node_tree = bpy.context.scene.world.node_tree
tree_nodes = node_tree.nodes
# Clear all nodes
tree_nodes.clear()
# Add Background node
node_background = tree_nodes.new(type='ShaderNodeBackground')
# Add Environment Texture node
node_environment = tree_nodes.new('ShaderNodeTexEnvironment')
# Load and assign the image to the node property
# node_environment.image = bpy.data.images.load("/Users/zhangjiyao/Desktop/test_addon/envmap_lib/autoshop_01_1k.hdr") # Relative path
node_environment.location = -300,0
node_tex_coord = tree_nodes.new(type='ShaderNodeTexCoord')
node_tex_coord.location = -700,0
node_mapping = tree_nodes.new(type='ShaderNodeMapping')
node_mapping.location = -500,0
# Add Output node
node_output = tree_nodes.new(type='ShaderNodeOutputWorld')
node_output.location = 200,0
# Link all nodes
links = node_tree.links
links.new(node_environment.outputs["Color"], node_background.inputs["Color"])
links.new(node_background.outputs["Background"], node_output.inputs["Surface"])
links.new(node_tex_coord.outputs["Generated"], node_mapping.inputs["Vector"])
links.new(node_mapping.outputs["Vector"], node_environment.inputs["Vector"])
#### bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = 1.0
random_energy = random.uniform(LIGHT_ENV_MAP_ENERGY_RGB * 0.8, LIGHT_ENV_MAP_ENERGY_RGB * 1.2)
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = random_energy
####
def setEnvMap(self, env_map_id, rotation_elur_z):
# Get the environment node tree of the current scene
node_tree = bpy.context.scene.world.node_tree
# Get Environment Texture node
node_environment = node_tree.nodes['Environment Texture']
# Load and assign the image to the node property
node_environment.image = self.env_map[env_map_id]
node_mapping = node_tree.nodes['Mapping']
node_mapping.inputs[2].default_value[2] = rotation_elur_z
def addMaskMaterial(self, num=20):
background_material_name_list = ["mask_background", "mask_table", "mask_tableplane"]
for material_name in background_material_name_list:
material_class = (bpy.data.materials.get(material_name) or bpy.data.materials.new(material_name)) # test if material exists, if it does not exist, create it:
# enable 'Use nodes'
material_class.use_nodes = True
node_tree = material_class.node_tree
# remove default nodes
material_class.node_tree.nodes.clear()
# add new nodes
node_1 = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_2= node_tree.nodes.new('ShaderNodeBrightContrast')
# link nodes
node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
node_2.inputs[0].default_value = (1, 1, 1, 1)
self.my_material[material_name] = material_class
for i in range(num):
class_name = str(i + 1)
# set the material of background
material_name = "mask_" + class_name
# test if material exists
# if it does not exist, create it:
material_class = (bpy.data.materials.get(material_name) or
bpy.data.materials.new(material_name))
# enable 'Use nodes'
material_class.use_nodes = True
node_tree = material_class.node_tree
# remove default nodes
material_class.node_tree.nodes.clear()
# add new nodes
node_1 = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_2= node_tree.nodes.new('ShaderNodeBrightContrast')
# link nodes
node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
if class_name.split('_')[0] == 'background' or class_name.split('_')[0] == 'table' or class_name.split('_')[0] == 'tableplane':
node_2.inputs[0].default_value = (1, 1, 1, 1)
else:
node_2.inputs[0].default_value = ((i + 1)/255., 0., 0., 1)
self.my_material[material_name] = material_class
def addNOCSMaterial(self):
material_name = 'coord_color'
mat = (bpy.data.materials.get(material_name) or bpy.data.materials.new(material_name))
mat.use_nodes = True
node_tree = mat.node_tree
nodes = node_tree.nodes
nodes.clear()
links = node_tree.links
links.clear()
vcol_R = nodes.new(type="ShaderNodeVertexColor")
vcol_R.layer_name = "Col_R" # the vertex color layer name
vcol_G = nodes.new(type="ShaderNodeVertexColor")
vcol_G.layer_name = "Col_G" # the vertex color layer name
vcol_B = nodes.new(type="ShaderNodeVertexColor")
vcol_B.layer_name = "Col_B" # the vertex color layer name
node_Output = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_Emission = node_tree.nodes.new('ShaderNodeEmission')
node_LightPath = node_tree.nodes.new('ShaderNodeLightPath')
node_Mix = node_tree.nodes.new('ShaderNodeMixShader')
node_Combine = node_tree.nodes.new(type="ShaderNodeCombineRGB")
# make links
node_tree.links.new(vcol_R.outputs[1], node_Combine.inputs[0])
node_tree.links.new(vcol_G.outputs[1], node_Combine.inputs[1])
node_tree.links.new(vcol_B.outputs[1], node_Combine.inputs[2])
node_tree.links.new(node_Combine.outputs[0], node_Emission.inputs[0])
node_tree.links.new(node_LightPath.outputs[0], node_Mix.inputs[0])
node_tree.links.new(node_Emission.outputs[0], node_Mix.inputs[2])
node_tree.links.new(node_Mix.outputs[0], node_Output.inputs[0])
self.my_material[material_name] = mat
def addNormalMaterial(self):
material_name = 'normal'
mat = (bpy.data.materials.get(material_name) or bpy.data.materials.new(material_name))
mat.use_nodes = True
node_tree = mat.node_tree
nodes = node_tree.nodes
nodes.clear()
links = node_tree.links
links.clear()
# Nodes :
new_node = nodes.new(type='ShaderNodeMath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (151.59744262695312, 854.5482177734375)
new_node.name = 'Math'
new_node.operation = 'MULTIPLY'
new_node.select = False
new_node.use_clamp = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node.inputs[1].default_value = 1.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = 0.0
new_node = nodes.new(type='ShaderNodeLightPath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (602.9912719726562, 1046.660888671875)
new_node.name = 'Light Path'
new_node.select = False
new_node.width = 140.0
new_node.outputs[0].default_value = 0.0
new_node.outputs[1].default_value = 0.0
new_node.outputs[2].default_value = 0.0
new_node.outputs[3].default_value = 0.0
new_node.outputs[4].default_value = 0.0
new_node.outputs[5].default_value = 0.0
new_node.outputs[6].default_value = 0.0
new_node.outputs[7].default_value = 0.0
new_node.outputs[8].default_value = 0.0
new_node.outputs[9].default_value = 0.0
new_node.outputs[10].default_value = 0.0
new_node.outputs[11].default_value = 0.0
new_node.outputs[12].default_value = 0.0
new_node = nodes.new(type='ShaderNodeOutputMaterial')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.is_active_output = True
new_node.location = (1168.93017578125, 701.84033203125)
new_node.name = 'Material Output'
new_node.select = False
new_node.target = 'ALL'
new_node.width = 140.0
new_node.inputs[2].default_value = [0.0, 0.0, 0.0]
new_node = nodes.new(type='ShaderNodeBsdfTransparent')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (731.72900390625, 721.4832763671875)
new_node.name = 'Transparent BSDF'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = [1.0, 1.0, 1.0, 1.0]
new_node = nodes.new(type='ShaderNodeCombineXYZ')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (594.4229736328125, 602.9271240234375)
new_node.name = 'Combine XYZ'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.0
new_node.inputs[1].default_value = 0.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
new_node = nodes.new(type='ShaderNodeMixShader')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (992.7239990234375, 707.2142333984375)
new_node.name = 'Mix Shader'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node = nodes.new(type='ShaderNodeEmission')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (774.0802612304688, 608.2547607421875)
new_node.name = 'Emission'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = [1.0, 1.0, 1.0, 1.0]
new_node.inputs[1].default_value = 1.0
new_node = nodes.new(type='ShaderNodeSeparateXYZ')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (-130.12167358398438, 558.1497802734375)
new_node.name = 'Separate XYZ'
new_node.select = False
new_node.width = 140.0
new_node.inputs[0].default_value = [0.0, 0.0, 0.0]
new_node.outputs[0].default_value = 0.0
new_node.outputs[1].default_value = 0.0
new_node.outputs[2].default_value = 0.0
new_node = nodes.new(type='ShaderNodeMath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (162.43240356445312, 618.8094482421875)
new_node.name = 'Math.002'
new_node.operation = 'MULTIPLY'
new_node.select = False
new_node.use_clamp = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node.inputs[1].default_value = 1.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = 0.0
new_node = nodes.new(type='ShaderNodeMath')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (126.8158187866211, 364.5539855957031)
new_node.name = 'Math.001'
new_node.operation = 'MULTIPLY'
new_node.select = False
new_node.use_clamp = False
new_node.width = 140.0
new_node.inputs[0].default_value = 0.5
new_node.inputs[1].default_value = -1.0
new_node.inputs[2].default_value = 0.0
new_node.outputs[0].default_value = 0.0
new_node = nodes.new(type='ShaderNodeVectorTransform')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.convert_from = 'WORLD'
new_node.convert_to = 'CAMERA'
new_node.location = (-397.0209045410156, 594.7037353515625)
new_node.name = 'Vector Transform'
new_node.select = False
new_node.vector_type = 'VECTOR'
new_node.width = 140.0
new_node.inputs[0].default_value = [0.5, 0.5, 0.5]
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
new_node = nodes.new(type='ShaderNodeNewGeometry')
new_node.active_preview = False
new_node.color = (0.6079999804496765, 0.6079999804496765, 0.6079999804496765)
new_node.location = (-651.8067016601562, 593.0455932617188)
new_node.name = 'Geometry'
new_node.width = 140.0
new_node.outputs[0].default_value = [0.0, 0.0, 0.0]
new_node.outputs[1].default_value = [0.0, 0.0, 0.0]
new_node.outputs[2].default_value = [0.0, 0.0, 0.0]
new_node.outputs[3].default_value = [0.0, 0.0, 0.0]
new_node.outputs[4].default_value = [0.0, 0.0, 0.0]
new_node.outputs[5].default_value = [0.0, 0.0, 0.0]
new_node.outputs[6].default_value = 0.0
new_node.outputs[7].default_value = 0.0
new_node.outputs[8].default_value = 0.0
# Links :
links.new(nodes["Light Path"].outputs[0], nodes["Mix Shader"].inputs[0])
links.new(nodes["Separate XYZ"].outputs[0], nodes["Math"].inputs[0])
links.new(nodes["Separate XYZ"].outputs[1], nodes["Math.002"].inputs[0])
links.new(nodes["Separate XYZ"].outputs[2], nodes["Math.001"].inputs[0])
links.new(nodes["Vector Transform"].outputs[0], nodes["Separate XYZ"].inputs[0])
links.new(nodes["Combine XYZ"].outputs[0], nodes["Emission"].inputs[0])
links.new(nodes["Math"].outputs[0], nodes["Combine XYZ"].inputs[0])
links.new(nodes["Math.002"].outputs[0], nodes["Combine XYZ"].inputs[1])
links.new(nodes["Math.001"].outputs[0], nodes["Combine XYZ"].inputs[2])
links.new(nodes["Transparent BSDF"].outputs[0], nodes["Mix Shader"].inputs[1])
links.new(nodes["Emission"].outputs[0], nodes["Mix Shader"].inputs[2])
links.new(nodes["Mix Shader"].outputs[0], nodes["Material Output"].inputs[0])
links.new(nodes["Geometry"].outputs[1], nodes["Vector Transform"].inputs[0])
self.my_material[material_name] = mat
def addMaterialLib(self, material_class_instance_pairs):
for mat in bpy.data.materials:
name = mat.name
name_class = str(name.split('_')[0])
if name_class != 'Dots Stroke' and name_class != 'default':
if name_class not in self.my_material:
self.my_material[name_class] = [mat]
else:
self.my_material[name_class].append(mat) # e.g. self.my_material['metal'] = [.....]
###
def setCamera(self, quaternion, translation, fov, baseline_distance):
self.camera_l.data.angle = fov
self.camera_r.data.angle = self.camera_l.data.angle
cx = translation[0]
cy = translation[1]
cz = translation[2]
self.camera_l.location[0] = cx
self.camera_l.location[1] = cy
self.camera_l.location[2] = cz
self.camera_l.rotation_mode = 'QUATERNION'
self.camera_l.rotation_quaternion[0] = quaternion[0]
self.camera_l.rotation_quaternion[1] = quaternion[1]
self.camera_l.rotation_quaternion[2] = quaternion[2]
self.camera_l.rotation_quaternion[3] = quaternion[3]
self.camera_r.rotation_mode = 'QUATERNION'
self.camera_r.rotation_quaternion[0] = quaternion[0]
self.camera_r.rotation_quaternion[1] = quaternion[1]
self.camera_r.rotation_quaternion[2] = quaternion[2]
self.camera_r.rotation_quaternion[3] = quaternion[3]
cx, cy, cz = cameraLPosToCameraRPos(quaternion, (cx, cy, cz), baseline_distance)
self.camera_r.location[0] = cx
self.camera_r.location[1] = cy
self.camera_r.location[2] = cz
def setLighting(self):
# emitter
#self.light_emitter.location = self.camera_r.location
self.light_emitter.location = self.camera_l.location + 0.51 * (self.camera_r.location - self.camera_l.location)
self.light_emitter.rotation_mode = 'QUATERNION'
self.light_emitter.rotation_quaternion = self.camera_r.rotation_quaternion
# emitter setting
bpy.context.view_layer.objects.active = None
# bpy.ops.object.select_all(action="DESELECT")
self.render_context.engine = 'CYCLES'
self.light_emitter.select_set(True)
self.light_emitter.data.use_nodes = True
self.light_emitter.data.type = "POINT"
self.light_emitter.data.shadow_soft_size = 0.001
random_energy = random.uniform(LIGHT_EMITTER_ENERGY * 0.9, LIGHT_EMITTER_ENERGY * 1.1)
self.light_emitter.data.energy = random_energy
# remove default node
light_emitter = bpy.data.objects["light_emitter"].data
light_emitter.node_tree.nodes.clear()
# add new nodes
light_output = light_emitter.node_tree.nodes.new("ShaderNodeOutputLight")
node_1 = light_emitter.node_tree.nodes.new("ShaderNodeEmission")
node_2 = light_emitter.node_tree.nodes.new("ShaderNodeTexImage")
node_3 = light_emitter.node_tree.nodes.new("ShaderNodeMapping")
node_4 = light_emitter.node_tree.nodes.new("ShaderNodeVectorMath")
node_5 = light_emitter.node_tree.nodes.new("ShaderNodeSeparateXYZ")
node_6 = light_emitter.node_tree.nodes.new("ShaderNodeTexCoord")
# link nodes
light_emitter.node_tree.links.new(light_output.inputs[0], node_1.outputs[0])
light_emitter.node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
light_emitter.node_tree.links.new(node_2.inputs[0], node_3.outputs[0])
light_emitter.node_tree.links.new(node_3.inputs[0], node_4.outputs[0])
light_emitter.node_tree.links.new(node_4.inputs[0], node_6.outputs[1])
light_emitter.node_tree.links.new(node_4.inputs[1], node_5.outputs[2])
light_emitter.node_tree.links.new(node_5.inputs[0], node_6.outputs[1])
# set parameter of nodes
node_1.inputs[1].default_value = 1.0 # scale
node_2.extension = 'CLIP'
# node_2.interpolation = 'Cubic'
node_3.inputs[1].default_value[0] = 0.5
node_3.inputs[1].default_value[1] = 0.5
node_3.inputs[1].default_value[2] = 0
node_3.inputs[2].default_value[0] = 0
node_3.inputs[2].default_value[1] = 0
node_3.inputs[2].default_value[2] = 0.05
# scale of pattern
node_3.inputs[3].default_value[0] = 0.6
node_3.inputs[3].default_value[1] = 0.85
node_3.inputs[3].default_value[2] = 0
node_4.operation = 'DIVIDE'
# pattern path
node_2.image = self.pattern
def lightModeSelect(self, light_mode):
if light_mode == "RGB":
self.light_emitter.hide_render = True
###
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = self.src_energy_for_rgb_render
elif light_mode == "IR":
self.light_emitter.hide_render = False
# set the environment map energy
random_energy = random.uniform(LIGHT_ENV_MAP_ENERGY_IR * 0.8, LIGHT_ENV_MAP_ENERGY_IR * 1.2)
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = random_energy
elif light_mode == "Mask" or light_mode == "NOCS" or light_mode == "Normal":
self.light_emitter.hide_render = True
bpy.data.worlds["World"].node_tree.nodes["Background"].inputs[1].default_value = 0
else:
raise NotImplementedError
def outputModeSelect(self, output_mode):
if output_mode == "RGB":
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.compression = 0
self.render_context.image_settings.color_mode = 'RGB'
self.render_context.image_settings.color_depth = '8'
bpy.context.scene.view_settings.view_transform = 'Filmic'
bpy.context.scene.render.filter_size = 1.5
self.render_context.resolution_x = 640 ### 1280
self.render_context.resolution_y = 360 ### 720
elif output_mode == "IR":
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.compression = 0
self.render_context.image_settings.color_mode = 'BW'
self.render_context.image_settings.color_depth = '8'
bpy.context.scene.view_settings.view_transform = 'Filmic'
bpy.context.scene.render.filter_size = 1.5
self.render_context.resolution_x = 640 ### 1280
self.render_context.resolution_y = 360 ### 720
elif output_mode == "Mask":
self.render_context.image_settings.file_format = 'OPEN_EXR'
self.render_context.image_settings.color_mode = 'RGB'
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.filter_size = 0
self.render_context.resolution_x = 640
self.render_context.resolution_y = 360
elif output_mode == "NOCS":
# self.render_context.image_settings.file_format = 'OPEN_EXR'
self.render_context.image_settings.file_format = 'PNG'
self.render_context.image_settings.color_mode = 'RGB'
self.render_context.image_settings.color_depth = '8'
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.filter_size = 0
self.render_context.resolution_x = 640
self.render_context.resolution_y = 360
elif output_mode == "Normal":
self.render_context.image_settings.file_format = 'OPEN_EXR'
self.render_context.image_settings.color_mode = 'RGB'
bpy.context.scene.view_settings.view_transform = 'Raw'
bpy.context.scene.render.filter_size = 1.5
self.render_context.resolution_x = 640
self.render_context.resolution_y = 360
else:
raise NotImplementedError
def renderEngineSelect(self, engine_mode):
if engine_mode == "CYCLES":
self.render_context.engine = 'CYCLES'
bpy.context.scene.cycles.progressive = 'BRANCHED_PATH'
bpy.context.scene.cycles.use_denoising = True
bpy.context.scene.cycles.denoiser = 'NLM'
bpy.context.scene.cycles.film_exposure = 1.0
bpy.context.scene.cycles.aa_samples = CYCLES_SAMPLE
## Set the device_type
bpy.context.preferences.addons["cycles"].preferences.compute_device_type = "CUDA" # or "OPENCL"
## get_devices() to let Blender detects GPU device
cuda_devices, _ = bpy.context.preferences.addons["cycles"].preferences.get_devices()
#print(bpy.context.preferences.addons["cycles"].preferences.compute_device_type)
for d in bpy.context.preferences.addons["cycles"].preferences.devices:
d["use"] = 1 # Using all devices, include GPU and CPU
#print(d["name"], d["use"])
device_list = self.DEVICE_LIST
activated_gpus = []
for i, device in enumerate(cuda_devices):
if (i in device_list):
device.use = True
activated_gpus.append(device.name)
else:
device.use = False
elif engine_mode == "EEVEE":
bpy.context.scene.render.engine = 'BLENDER_EEVEE'
else:
print("Not support the mode!")
def addBackground(self, size, position, scale, default_background_texture_path):
# set the material of background
material_name = "default_background"
# test if material exists
# if it does not exist, create it:
material_background = (bpy.data.materials.get(material_name) or
bpy.data.materials.new(material_name))
# enable 'Use nodes'
material_background.use_nodes = True
node_tree = material_background.node_tree
# remove default nodes
material_background.node_tree.nodes.clear()
# add new nodes
node_1 = node_tree.nodes.new('ShaderNodeOutputMaterial')
node_2 = node_tree.nodes.new('ShaderNodeBsdfPrincipled')
node_3 = node_tree.nodes.new('ShaderNodeTexImage')
# link nodes
node_tree.links.new(node_1.inputs[0], node_2.outputs[0])
node_tree.links.new(node_2.inputs[0], node_3.outputs[0])
# add texture image
node_3.image = bpy.data.images.load(filepath=default_background_texture_path)
self.my_material['default_background'] = material_background
# add background plane
for i in range(-2, 3, 1):
for j in range(-2, 3, 1):
position_i_j = (i * size + position[0], j * size + position[1], position[2] - TABLE_CAD_MODEL_HEIGHT)
bpy.ops.mesh.primitive_plane_add(size=size, enter_editmode=False, align='WORLD', location=position_i_j, scale=scale)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
for i in range(-2, 3, 1):
for j in [-2, 2]:
position_i_j = (i * size + position[0], j * size + position[1], position[2] - 0.25)# - TABLE_CAD_MODEL_HEIGHT)
rotation_elur = (math.pi / 2., 0., 0.)
bpy.ops.mesh.primitive_plane_add(size=size, enter_editmode=False, align='WORLD', location=position_i_j, rotation = rotation_elur)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
for j in range(-2, 3, 1):
for i in [-2, 2]:
position_i_j = (i * size + position[0], j * size + position[1], position[2] - 0.25)# - TABLE_CAD_MODEL_HEIGHT)
rotation_elur = (0, math.pi / 2, 0)
bpy.ops.mesh.primitive_plane_add(size=size, enter_editmode=False, align='WORLD', location=position_i_j, rotation = rotation_elur)
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.type = 'PASSIVE'
bpy.context.object.rigid_body.collision_shape = 'BOX'
count = 0
for obj in bpy.data.objects:
if obj.type == "MESH":
obj.name = "background_" + str(count)
obj.data.name = "background_" + str(count)
obj.active_material = material_background
count += 1
self.background_added = True
def clearModel(self):
'''
# delete all meshes
for item in bpy.data.meshes:
bpy.data.meshes.remove(item)
for item in bpy.data.materials:
bpy.data.materials.remove(item)
'''
# remove all objects except background
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
bpy.data.meshes.remove(obj.data)
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
bpy.data.objects.remove(obj, do_unlink=True)
# remove all default material
for mat in bpy.data.materials:
name = mat.name.split('.')
if name[0] == 'Material':
bpy.data.materials.remove(mat)
def loadModel(self, file_path):
self.model_loaded = True
try:
if file_path.endswith('obj'):
bpy.ops.import_scene.obj(filepath=file_path)
elif file_path.endswith('3ds'):
bpy.ops.import_scene.autodesk_3ds(filepath=file_path)
elif file_path.endswith('dae'):
# Must install OpenCollada. Please read README.md
bpy.ops.wm.collada_import(filepath=file_path)
else:
self.model_loaded = False
raise Exception("Loading failed: %s" % (file_path))
except Exception:
self.model_loaded = False
def render(self, image_name="tmp", image_path=RENDERING_PATH):
# Render the object
if not self.model_loaded:
print("[W]render: Model not loaded.")
return
if self.render_mode == "IR":
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("IR")
self.outputModeSelect("IR")
self.renderEngineSelect("CYCLES")
elif self.render_mode == 'RGB':
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("RGB")
self.outputModeSelect("RGB")
self.renderEngineSelect("CYCLES")
elif self.render_mode == "Mask":
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("Mask")
self.outputModeSelect("Mask")
# self.renderEngineSelect("EEVEE")
self.renderEngineSelect("CYCLES")
bpy.context.scene.cycles.use_denoising = False
bpy.context.scene.cycles.aa_samples = 1
elif self.render_mode == "NOCS":
bpy.context.scene.use_nodes = False
# set light and render mode
self.lightModeSelect("NOCS")
self.outputModeSelect("NOCS")
# self.renderEngineSelect("EEVEE")
self.renderEngineSelect("CYCLES")
bpy.context.scene.cycles.use_denoising = False
bpy.context.scene.cycles.aa_samples = 1
elif self.render_mode == "Normal":
bpy.context.scene.use_nodes = True
self.fileOutput.base_path = image_path.replace("normal","depth")
self.fileOutput.file_slots[0].path = image_name[:4]+"_#"# + 'depth_#'
# set light and render mode
self.lightModeSelect("Normal")
self.outputModeSelect("Normal")
# self.renderEngineSelect("EEVEE")
self.renderEngineSelect("CYCLES")
bpy.context.scene.cycles.use_denoising = False
bpy.context.scene.cycles.aa_samples = 32
else:
print("[W]render: The render mode is not supported")
return
bpy.context.scene.render.filepath = os.path.join(image_path, image_name)
bpy.ops.render.render(write_still=True) # save straight to file
def set_material_randomize_mode(self, class_material_pairs, mat_randomize_mode, instance, material_type_in_mixed_mode):
if mat_randomize_mode in ['mixed','diffuse','transparent','specular_tex','specular_texmix','specular_and_transparent']:
if material_type_in_mixed_mode == 'raw':
print("[V]set_material_randomize_mode", instance.name, 'material type: raw')
set_modify_raw_material(instance)
else:
material = random.sample(self.my_material[material_type_in_mixed_mode], 1)[0]
print("[V]set_material_randomize_mode", instance.name, 'material type: ', material_type_in_mixed_mode)
## graspnet
set_modify_material(instance, material, self.obj_texture_img_list, mat_randomize_mode=mat_randomize_mode)
elif mat_randomize_mode == 'specular':
material = random.sample(self.my_material[material_type_in_mixed_mode], 1)[0]
print("[V]set_material_randomize_mode", instance.name, 'material type: ', material_type_in_mixed_mode)
set_modify_material(instance, material, self.obj_texture_img_list, mat_randomize_mode=mat_randomize_mode,
is_transfer=False)
else:
raise NotImplementedError("No such mat_randomize_mode!")
def get_instance_pose(self):
instance_pose = {}
bpy.context.view_layer.update()
cam = self.camera_l
mat_rot_x = Matrix.Rotation(math.radians(180.0), 4, 'X')
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
instance_id = obj.name.split('_')[0]
mat_rel = cam.matrix_world.inverted() @ obj.matrix_world
# location
relative_location = [mat_rel.translation[0],
- mat_rel.translation[1],
- mat_rel.translation[2]]
# rotation
# relative_rotation_euler = mat_rel.to_euler() # must be converted from radians to degrees
relative_rotation_quat = [mat_rel.to_quaternion()[0],
mat_rel.to_quaternion()[1],
mat_rel.to_quaternion()[2],
mat_rel.to_quaternion()[3]]
quat_x = [0, 1, 0, 0]
quat = quanternion_mul(quat_x, relative_rotation_quat)
quat = [quat[0], - quat[1], - quat[2], - quat[3]]
instance_pose[str(instance_id)] = [quat, relative_location]
return instance_pose
def check_visible(self, threshold=(0.1, 0.9, 0.1, 0.9)):
w_min, x_max, h_min, h_max = threshold
visible_objects_list = []
bpy.context.view_layer.update()
cs, ce = self.camera_l.data.clip_start, self.camera_l.data.clip_end
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
obj_center = obj.matrix_world.translation
co_ndc = world_to_camera_view(scene, self.camera_l, obj_center)
if (w_min < co_ndc.x < x_max and
h_min < co_ndc.y < h_max and
cs < co_ndc.z < ce):
obj.select_set(True)
visible_objects_list.append(obj)
else:
obj.select_set(False)
return visible_objects_list
def setModelPosition(instance, location, quaternion):
instance.rotation_mode = 'QUATERNION'
instance.rotation_quaternion[0] = quaternion[0]
instance.rotation_quaternion[1] = quaternion[1]
instance.rotation_quaternion[2] = quaternion[2]
instance.rotation_quaternion[3] = quaternion[3]
instance.location = location
###
def setRigidBody(instance):
bpy.context.view_layer.objects.active = instance
object_single = bpy.context.active_object
# add rigid body constraints to cube
bpy.ops.rigidbody.object_add()
bpy.context.object.rigid_body.mass = 1
bpy.context.object.rigid_body.kinematic = True
bpy.context.object.rigid_body.collision_shape = 'CONVEX_HULL'
bpy.context.object.rigid_body.restitution = 0.01
bpy.context.object.rigid_body.angular_damping = 0.8
bpy.context.object.rigid_body.linear_damping = 0.99
bpy.context.object.rigid_body.kinematic = False
object_single.keyframe_insert(data_path='rigid_body.kinematic', frame=0)
def set_visiable_objects(visible_objects_list):
for obj in bpy.data.objects:
if obj.type == 'MESH' and not obj.name.split('_')[0] == 'background':
if obj in visible_objects_list:
obj.hide_render = False
else:
obj.hide_render = True
def generate_CAD_model_list(scene_type, urdf_path_list, obj_uid_list):
CAD_model_list = {}
###
for idx in range(len(urdf_path_list)):
urdf_path = urdf_path_list[idx]
obj_uid = obj_uid_list[idx]
class_name = 'other'
urdf_path = str(urdf_path).replace("\\","/").split("/")
if scene_type == "blocks":
instance_path = "/".join(urdf_path[:-1]) + "/" + urdf_path[-1][:-5]+".obj"
else:
instance_path = "/".join(urdf_path[:-1])+"/"+urdf_path[-1][:-5]+"_visual.obj"
class_list = []
class_list.append([instance_path, class_name, obj_uid])
if class_name == 'other' and 'other' in CAD_model_list:
CAD_model_list[class_name] = CAD_model_list[class_name] + class_list
else:
CAD_model_list[class_name] = class_list
return CAD_model_list
def generate_material_type(obj_name, class_material_pairs, instance_material_except_pairs, instance_material_include_pairs, material_class_instance_pairs, material_type):
###
specular_type_for_ins_list = []
transparent_type_for_ins_list = []
diffuse_type_for_ins_list = []
for key in instance_material_except_pairs:
if key in material_class_instance_pairs['specular']:
specular_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['transparent']:
transparent_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['diffuse']:
diffuse_type_for_ins_list.append(key)
for key in instance_material_include_pairs:
if key in material_class_instance_pairs['specular']:
specular_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['transparent']:
transparent_type_for_ins_list.append(key)
elif key in material_class_instance_pairs['diffuse']:
diffuse_type_for_ins_list.append(key)
if material_type == "transparent":
return random.sample(transparent_type_for_ins_list, 1)[0]
elif material_type == "diffuse":
return random.sample(diffuse_type_for_ins_list, 1)[0]
elif material_type == "specular" or material_type == "specular_tex" or material_type == "specular_texmix":
return random.sample(specular_type_for_ins_list, 1)[0]
elif material_type == "specular_and_transparent":
flag = random.randint(0, 2)
if flag == 0:
return random.sample(specular_type_for_ins_list, 1)[0] ### 'specular'
else:
return random.sample(transparent_type_for_ins_list, 1)[0] ### 'transparent'
elif material_type == "mixed":
# randomly pick material class
flag = random.randint(0, 2) # D:S:T=1:2:2
# select the raw material
if flag == 0:
return random.sample(diffuse_type_for_ins_list, 1)[0] ### 'diffuse'
# select one from specular and transparent
elif flag == 1:
return random.sample(specular_type_for_ins_list, 1)[0] ### 'specular'
else:
return random.sample(transparent_type_for_ins_list, 1)[0] ### 'transparent'
else:
raise ValueError(f"Material type error: {material_type}")
================================================
FILE: train.sh
================================================
cd src/nr
CUDA_VISIBLE_DEVICES=$1 python run_training.py --cfg configs/nrvgn_sdf.yaml
cd -