├── .gitignore ├── PointFlow ├── .gitignore ├── LICENSE ├── README.md ├── args.py ├── datasets.py ├── demo.py ├── install.sh ├── metrics │ ├── .gitignore │ ├── __init__.py │ ├── evaluation_metrics.py │ └── pytorch_structural_losses │ │ ├── .gitignore │ │ ├── Makefile │ │ ├── Makefile_autobot │ │ ├── __init__.py │ │ ├── match_cost.py │ │ ├── nn_distance.py │ │ ├── pybind │ │ ├── bind.cpp │ │ └── extern.hpp │ │ ├── setup.py │ │ └── src │ │ ├── approxmatch.cu │ │ ├── approxmatch.cuh │ │ ├── nndistance.cu │ │ ├── nndistance.cuh │ │ ├── structural_loss.cpp │ │ └── utils.hpp ├── models │ ├── __init__.py │ ├── cnf.py │ ├── diffeq_layers.py │ ├── flow.py │ ├── networks.py │ ├── normalization.py │ └── odefunc.py ├── pointcloud_dataset.py ├── prepare.sh ├── scripts │ ├── cutrearrange_gen_dist.sh │ ├── cutrearrangespread_gen_dist.sh │ ├── gathermove_gen_dist.sh │ ├── liftspread_gen_dist.sh │ ├── set_cutrearrange_gen_dist.sh │ └── set_cutrearrangespread_gen_dist.sh ├── test.py ├── train.py ├── utils.py └── visualization │ └── reconstruct.ipynb ├── README.md ├── assets ├── cuttingboard │ ├── Cutting_Board.mtl │ └── Cutting_Board.obj └── table │ ├── Table_Coffee_RiceChest.mtl │ ├── Table_Coffee_RiceChest.obj │ └── Table_Coffee_RiceChest │ └── _Wood_Cherry_.jpg ├── core ├── diffskill │ ├── agent.py │ ├── args.py │ ├── buffer.py │ ├── imitation_buffer.py │ ├── plan │ │ ├── compose_skills.py │ │ ├── optimizer.py │ │ └── project_func.py │ ├── train_diffskill_rgb.py │ └── train_flat3d.py ├── eval │ ├── eval_helper.py │ ├── hardcoded_eval_trajs.py │ └── sampler.py ├── models │ ├── cnn_actor.py │ ├── cnn_encoder.py │ ├── cnn_vae.py │ ├── mlp_abstraction.py │ ├── mlp_actor.py │ ├── pointflow_vae.py │ ├── pointnet_actor.py │ ├── pointnet_encoder.py │ ├── set_cost_predictor.py │ └── set_fea_predictor.py ├── pasta │ ├── agent.py │ ├── args.py │ ├── generate_dbscan_label.py │ ├── plan │ │ ├── compose_skills.py │ │ ├── compose_skills_rhp.py │ │ ├── optimizer.py │ │ └── set_trajs.py │ └── train_pasta.py ├── traj_opt │ ├── env_spec.py │ ├── gen_data.py │ ├── gen_init_target │ │ ├── crs_generator.py │ │ ├── cutrearrange_generator.py │ │ ├── gathermove_generator.py │ │ ├── liftspread_generator.py │ │ └── state_generator.py │ ├── generate_reset_motion.py │ └── solver.py └── utils │ ├── camera_info.json │ ├── core_utils.py │ ├── diffskill_utils.py │ ├── logger.py │ ├── open3d_utils.py │ ├── pasta_utils.py │ ├── pc_utils.py │ ├── plb_utils.py │ ├── simple_logger.py │ ├── torch_chamfer.py │ └── visualization_utils.py ├── environment.yml ├── gui.py ├── media ├── CutRearrange-v1_PASTA.gif ├── CutRearrangeSpread-v1_PASTA.gif ├── LiftSpread-v1_PASTA.gif └── pasta.jpeg ├── plb ├── __init__.py ├── config │ ├── __init__.py │ ├── default_config.py │ └── utils.py ├── engine │ ├── __init__.py │ ├── function.py │ ├── mpm_simulator.py │ ├── primitive │ │ ├── __init__.py │ │ ├── primitives.py │ │ ├── primive_base.py │ │ └── utils.py │ ├── renderer │ │ ├── __init__.py │ │ ├── renderer.py │ │ ├── renderer_tester.py │ │ ├── renderer_utils.py │ │ ├── test.py │ │ └── tina_renderer.py │ ├── shapes │ │ ├── __init__.py │ │ └── shape_maker.py │ └── taichi_env.py └── envs │ ├── __init__.py │ ├── env.py │ ├── env_ymls │ ├── assembly.yml │ ├── chopsticks.yml │ ├── cut_rearrange.yml │ ├── cut_rearrange_spread.yml │ ├── gather_move.yml │ ├── lift_spread.yml │ ├── move.yml │ ├── pinch.yml │ ├── rollingpin.yml │ ├── rope.yml │ ├── single_stage │ │ ├── lift.yml │ │ └── spread.yml │ ├── spread.yml │ ├── table.yml │ ├── torus.yml │ ├── triplemove.yml │ └── writer.yml │ ├── mp_wrapper.py │ ├── multitask_env.py │ └── utils.py ├── prepare.sh ├── run_scripts ├── diffskill_train_all.sh ├── difskill_plan.sh ├── download.py ├── flat3d_train_all.sh ├── generate_init_target.sh ├── pasta_plan.sh ├── pasta_train_abstraction.sh ├── pasta_train_policy.sh ├── run_dbscan.sh └── run_gbto.sh └── taichi_three ├── LICENSE ├── README.md ├── Tina-Dev.py ├── assets ├── .gitignore ├── atms.png ├── caustics.gltf ├── cloth.jpg ├── cornell.gltf ├── cornell.mtl ├── cornell.obj ├── cube.obj ├── cylinder.obj ├── logo.obj ├── lut.jpg ├── monkey.obj ├── multimtl.mtl ├── multimtl.obj ├── normal.png ├── orient.obj ├── pattern.jpg ├── plane.obj ├── shadow.mtl ├── shadow.obj ├── skybox.jpg ├── smallptwall.png ├── sphere.gltf ├── sphere.obj ├── torus.obj └── uvsphere.obj ├── bench.py ├── conv.py ├── docs ├── connect.py ├── gltf.py ├── lighting.py ├── monkey.py ├── options.py ├── particles.py ├── pathtrace.py ├── primitives.py ├── smooth.py ├── transform.py ├── triangle.py ├── volume.py └── wireframe.py ├── examples ├── cornell_box.py ├── ibl_matball.py ├── matball.py ├── mciso_mpm3d.py ├── meshgrid_cloth.py ├── meshgrid_wave.py ├── pars_mpm3d.py └── rtx_matball.py ├── lbm.py ├── melt ├── dump.py ├── lbm.py ├── lbvh.py ├── mpm.py ├── mrt.py ├── mrtlbm.py ├── rbd.py ├── vbrbd.py └── voxelizer.py ├── mesh_renderer.ipynb ├── setup.py ├── tests ├── bdpt.py ├── blooming.py ├── brdf.py ├── cookibl.py ├── emission.py ├── fpe.py ├── fxaa.py ├── glass.py ├── inter.py ├── knn.py ├── mis.py ├── mtlid.py ├── nlm.py ├── noise.py ├── path.py ├── pick.py ├── probe.py ├── ptlight.py ├── raytrace.py ├── rtao.py ├── rtx.py ├── skybox.py ├── ssao.py ├── ssr.py ├── volume.py ├── voxl.py ├── wave.py └── wire.py ├── tina ├── __init__.py ├── __main__.py ├── advans.py ├── assimp │ ├── __init__.py │ ├── gltf.py │ ├── obj.py │ ├── pfm.py │ └── tet.py ├── cli │ ├── __init__.py │ ├── mesh.py │ ├── particles.py │ └── volume.py ├── common.py ├── core │ ├── __init__.py │ ├── engine.py │ ├── lighting.py │ ├── material.py │ ├── particle.py │ ├── particle_sdf.py │ ├── renderer.py │ ├── shader.py │ ├── triangle.py │ ├── volume.py │ └── wireframe.py ├── hacker.py ├── inject.py ├── lazimp.py ├── matr │ ├── __init__.py │ ├── material.py │ ├── nodes.py │ └── wavelen.py ├── memory.py ├── mesh │ ├── __init__.py │ ├── base.py │ ├── conn.py │ ├── cull.py │ ├── export.py │ ├── grid.py │ ├── model.py │ ├── norm.py │ ├── prim.py │ ├── simple.py │ ├── trans.py │ └── wire.py ├── pars │ ├── __init__.py │ ├── base.py │ ├── export.py │ ├── simple.py │ └── trans.py ├── path │ ├── __init__.py │ ├── bidir.py │ ├── engine.py │ ├── geometry.py │ ├── particle.py │ ├── tree.py │ ├── triangle.py │ └── volume.py ├── postp │ ├── __init__.py │ ├── blooming.py │ ├── denoise.py │ ├── fxaa.py │ ├── ssao.py │ ├── ssr.py │ └── tonemap.py ├── probe.py ├── random.py ├── scene │ ├── __init__.py │ ├── raster.py │ └── tracer.py ├── shield.py ├── skybox.py ├── util │ ├── __init__.py │ ├── _mciso_data.py │ ├── accumator.py │ ├── control.py │ ├── matrix.py │ ├── mciso.py │ └── stack.py └── voxl │ ├── __init__.py │ ├── base.py │ ├── scale.py │ ├── simple.py │ └── trans.py └── view.py /.gitignore: -------------------------------------------------------------------------------- 1 | /test_match 2 | /plb/baselines 3 | /scratch/ 4 | chester 5 | *.ipynb_checkpoints* 6 | /datasets 7 | !PointFlow/datasets.py 8 | /runs/ 9 | .idea 10 | .vscode 11 | out 12 | data 13 | /datasets/ 14 | wandb 15 | carl_chester 16 | core/carl_launchers 17 | tests/carl_tests 18 | debug 19 | prepare_local.sh 20 | prepare_1.0_carl.sh 21 | # Byte-compiled / optimized / DLL files 22 | __pycache__/ 23 | taichi_three/tina/core/__pycache__/ 24 | *.py[cod] 25 | *.gz 26 | *.pyc 27 | *.m 28 | *.zip 29 | *.DS_Store 30 | local_data/ 31 | 32 | PointFlow/.ipynb_checkpoints 33 | PointFlow/checkpoints 34 | PointFlow/runs 35 | PointFlow/metrics/structural_losses/*.so 36 | PointFlow/metrics/structural_losses/*.cu.o 37 | PointFlow/metrics/structural_losses/makefile 38 | PointFlow/PyMesh 39 | PointFlow/checkpoint 40 | PointFlow/torchdiffeq/ 41 | PointFlow/demo/ 42 | 43 | # C extensions 44 | *.so 45 | 46 | # Distribution / packaging 47 | bin/ 48 | build/ 49 | develop-eggs/ 50 | dist/ 51 | eggs/ 52 | lib/ 53 | lib64/ 54 | parts/ 55 | sdist/ 56 | var/ 57 | *.egg-info/ 58 | .installed.cfg 59 | *.egg 60 | 61 | # Installer logs 62 | pip-log.txt 63 | pip-delete-this-directory.txt 64 | 65 | # Unit test / coverage reports 66 | .tox/ 67 | .coverage 68 | .cache 69 | nosetests.xml 70 | coverage.xml 71 | 72 | # Translations 73 | *.mo 74 | 75 | # Mr Developer 76 | .mr.developer.cfg 77 | .project 78 | .pydevproject 79 | 80 | # Rope 81 | .ropeproject 82 | 83 | # Django stuff: 84 | *.log 85 | *.pot 86 | 87 | # Sphinx documentation 88 | docs/_build/ 89 | -------------------------------------------------------------------------------- /PointFlow/.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__/ 2 | *__pycache__ 3 | .idea/ 4 | *.pyc 5 | data 6 | 7 | scratch.py 8 | 9 | *.m 10 | 11 | .ipynb_checkpoints 12 | 13 | checkpoints 14 | runs 15 | 16 | metrics/structural_losses/*.so 17 | metrics/structural_losses/*.cu.o 18 | metrics/structural_losses/makefile 19 | PyMesh 20 | checkpoint 21 | 22 | torchdiffeq/ 23 | demo/ 24 | -------------------------------------------------------------------------------- /PointFlow/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 Guandao Yang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /PointFlow/demo.py: -------------------------------------------------------------------------------- 1 | import open3d as o3d 2 | from datasets import get_datasets 3 | from args import get_args 4 | from models.networks import PointFlow 5 | import os 6 | import torch 7 | import numpy as np 8 | import torch.nn as nn 9 | 10 | 11 | def main(args): 12 | model = PointFlow(args) 13 | 14 | def _transform_(m): 15 | return nn.DataParallel(m) 16 | 17 | model = model.cuda() 18 | model.multi_gpu_wrapper(_transform_) 19 | 20 | print("Resume Path:%s" % args.resume_checkpoint) 21 | checkpoint = torch.load(args.resume_checkpoint) 22 | model.load_state_dict(checkpoint) 23 | model.eval() 24 | 25 | _, te_dataset = get_datasets(args) 26 | if args.resume_dataset_mean is not None and args.resume_dataset_std is not None: 27 | mean = np.load(args.resume_dataset_mean) 28 | std = np.load(args.resume_dataset_std) 29 | te_dataset.renormalize(mean, std) 30 | ds_mean = torch.from_numpy(te_dataset.all_points_mean).cuda() 31 | ds_std = torch.from_numpy(te_dataset.all_points_std).cuda() 32 | 33 | all_sample = [] 34 | with torch.no_grad(): 35 | for i in range(0, args.num_sample_shapes, args.batch_size): 36 | B = len(range(i, min(i + args.batch_size, args.num_sample_shapes))) 37 | N = args.num_sample_points 38 | _, out_pc = model.sample(B, N) 39 | out_pc = out_pc * ds_std + ds_mean 40 | all_sample.append(out_pc) 41 | 42 | sample_pcs = torch.cat(all_sample, dim=0).cpu().detach().numpy() 43 | print("Generation sample size:(%s, %s, %s)" % sample_pcs.shape) 44 | 45 | # Save the generative output 46 | os.makedirs("demo", exist_ok=True) 47 | np.save(os.path.join("demo", "model_out_smp.npy"), sample_pcs) 48 | 49 | # Visualize the demo 50 | pcl = o3d.geometry.PointCloud() 51 | for i in range(int(sample_pcs.shape[0])): 52 | print("Visualizing: %03d/%03d" % (i, sample_pcs.shape[0])) 53 | pts = sample_pcs[i].reshape(-1, 3) 54 | pcl.points = o3d.utility.Vector3dVector(pts) 55 | o3d.visualization.draw_geometries([pcl]) 56 | 57 | 58 | if __name__ == '__main__': 59 | args = get_args() 60 | main(args) 61 | -------------------------------------------------------------------------------- /PointFlow/install.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | 3 | root=`pwd` 4 | 5 | # Install dependecies 6 | conda install numpy matplotlib pillow scipy tqdm scikit-learn -y 7 | pip install tensorflow-gpu==1.13.1 8 | pip install tensorboardX==1.7 9 | 10 | # Compile CUDA kernel for CD/EMD loss 11 | cd metrics/pytorch_structural_losses/ 12 | make clean 13 | make 14 | cd $root 15 | 16 | # install torchdiffeq 17 | git clone https://github.com/rtqichen/torchdiffeq.git 18 | cd torchdiffeq 19 | pip install -e . 20 | -------------------------------------------------------------------------------- /PointFlow/metrics/.gitignore: -------------------------------------------------------------------------------- 1 | StructuralLosses 2 | -------------------------------------------------------------------------------- /PointFlow/metrics/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/PointFlow/metrics/__init__.py -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/.gitignore: -------------------------------------------------------------------------------- 1 | PyTorchStructuralLosses.egg-info/ 2 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/__init__.py: -------------------------------------------------------------------------------- 1 | #import torch 2 | 3 | #from MakePytorchBackend import AddGPU, Foo, ApproxMatch 4 | 5 | #from Add import add_gpu, approx_match 6 | 7 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/match_cost.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | from metrics.StructuralLosses.StructuralLossesBackend import ApproxMatch, MatchCost, MatchCostGrad 4 | 5 | # Inherit from Function 6 | class MatchCostFunction(Function): 7 | # Note that both forward and backward are @staticmethods 8 | @staticmethod 9 | # bias is an optional argument 10 | def forward(ctx, seta, setb): 11 | #print("Match Cost Forward") 12 | ctx.save_for_backward(seta, setb) 13 | ''' 14 | input: 15 | set1 : batch_size * #dataset_points * 3 16 | set2 : batch_size * #query_points * 3 17 | returns: 18 | match : batch_size * #query_points * #dataset_points 19 | ''' 20 | match, temp = ApproxMatch(seta, setb) 21 | ctx.match = match 22 | cost = MatchCost(seta, setb, match) 23 | return cost 24 | 25 | """ 26 | grad_1,grad_2=approxmatch_module.match_cost_grad(xyz1,xyz2,match) 27 | return [grad_1*tf.expand_dims(tf.expand_dims(grad_cost,1),2),grad_2*tf.expand_dims(tf.expand_dims(grad_cost,1),2),None] 28 | """ 29 | # This function has only a single output, so it gets only one gradient 30 | @staticmethod 31 | def backward(ctx, grad_output): 32 | #print("Match Cost Backward") 33 | # This is a pattern that is very convenient - at the top of backward 34 | # unpack saved_tensors and initialize all gradients w.r.t. inputs to 35 | # None. Thanks to the fact that additional trailing Nones are 36 | # ignored, the return statement is simple even when the function has 37 | # optional inputs. 38 | seta, setb = ctx.saved_tensors 39 | #grad_input = grad_weight = grad_bias = None 40 | grada, gradb = MatchCostGrad(seta, setb, ctx.match) 41 | grad_output_expand = grad_output.unsqueeze(1).unsqueeze(2) 42 | return grada*grad_output_expand, gradb*grad_output_expand 43 | 44 | match_cost = MatchCostFunction.apply 45 | 46 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/nn_distance.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from torch.autograd import Function 3 | # from extensions.StructuralLosses.StructuralLossesBackend import NNDistance, NNDistanceGrad 4 | from metrics.StructuralLosses.StructuralLossesBackend import NNDistance, NNDistanceGrad 5 | 6 | # Inherit from Function 7 | class NNDistanceFunction(Function): 8 | # Note that both forward and backward are @staticmethods 9 | @staticmethod 10 | # bias is an optional argument 11 | def forward(ctx, seta, setb): 12 | #print("Match Cost Forward") 13 | ctx.save_for_backward(seta, setb) 14 | ''' 15 | input: 16 | set1 : batch_size * #dataset_points * 3 17 | set2 : batch_size * #query_points * 3 18 | returns: 19 | dist1, idx1, dist2, idx2 20 | ''' 21 | dist1, idx1, dist2, idx2 = NNDistance(seta, setb) 22 | ctx.idx1 = idx1 23 | ctx.idx2 = idx2 24 | return dist1, dist2 25 | 26 | # This function has only a single output, so it gets only one gradient 27 | @staticmethod 28 | def backward(ctx, grad_dist1, grad_dist2): 29 | #print("Match Cost Backward") 30 | # This is a pattern that is very convenient - at the top of backward 31 | # unpack saved_tensors and initialize all gradients w.r.t. inputs to 32 | # None. Thanks to the fact that additional trailing Nones are 33 | # ignored, the return statement is simple even when the function has 34 | # optional inputs. 35 | seta, setb = ctx.saved_tensors 36 | idx1 = ctx.idx1 37 | idx2 = ctx.idx2 38 | grada, gradb = NNDistanceGrad(seta, setb, idx1, idx2, grad_dist1, grad_dist2) 39 | return grada, gradb 40 | 41 | nn_distance = NNDistanceFunction.apply 42 | 43 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/pybind/bind.cpp: -------------------------------------------------------------------------------- 1 | #include 2 | 3 | #include 4 | 5 | #include "pybind/extern.hpp" 6 | 7 | namespace py = pybind11; 8 | 9 | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){ 10 | m.def("ApproxMatch", &ApproxMatch); 11 | m.def("MatchCost", &MatchCost); 12 | m.def("MatchCostGrad", &MatchCostGrad); 13 | m.def("NNDistance", &NNDistance); 14 | m.def("NNDistanceGrad", &NNDistanceGrad); 15 | } 16 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/pybind/extern.hpp: -------------------------------------------------------------------------------- 1 | std::vector ApproxMatch(at::Tensor in_a, at::Tensor in_b); 2 | at::Tensor MatchCost(at::Tensor set_d, at::Tensor set_q, at::Tensor match); 3 | std::vector MatchCostGrad(at::Tensor set_d, at::Tensor set_q, at::Tensor match); 4 | 5 | std::vector NNDistance(at::Tensor set_d, at::Tensor set_q); 6 | std::vector NNDistanceGrad(at::Tensor set_d, at::Tensor set_q, at::Tensor idx1, at::Tensor idx2, at::Tensor grad_dist1, at::Tensor grad_dist2); 7 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | from torch.utils.cpp_extension import CUDAExtension, BuildExtension 3 | 4 | # Python interface 5 | setup( 6 | name='PyTorchStructuralLosses', 7 | version='0.1.0', 8 | install_requires=['torch'], 9 | packages=['StructuralLosses'], 10 | package_dir={'StructuralLosses': './'}, 11 | ext_modules=[ 12 | CUDAExtension( 13 | name='StructuralLossesBackend', 14 | include_dirs=['./', '/home/jianrenw/carl/research/dev/dynamic_abstraction/PointFlow/metrics/pytorch_structural_losses', '/home/xingyu/Projects/PointFlow/metrics/pytorch_structural_losses'], 15 | sources=[ 16 | 'pybind/bind.cpp', 17 | ], 18 | libraries=['make_pytorch'], 19 | library_dirs=['objs'], 20 | # extra_compile_args=['-g'] 21 | ) 22 | ], 23 | cmdclass={'build_ext': BuildExtension}, 24 | author='Christopher B. Choy', 25 | author_email='chrischoy@ai.stanford.edu', 26 | description='Tutorial for Pytorch C++ Extension with a Makefile', 27 | keywords='Pytorch C++ Extension', 28 | url='https://github.com/chrischoy/MakePytorchPlusPlus', 29 | zip_safe=False, 30 | ) 31 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/src/approxmatch.cuh: -------------------------------------------------------------------------------- 1 | /* 2 | template 3 | void AddGPUKernel(Dtype *in_a, Dtype *in_b, Dtype *out_c, int N, 4 | cudaStream_t stream); 5 | */ 6 | void approxmatch(int b,int n,int m,const float * xyz1,const float * xyz2,float * match,float * temp, cudaStream_t stream); 7 | void matchcost(int b,int n,int m,const float * xyz1,const float * xyz2,float * match, float * out, cudaStream_t stream); 8 | void matchcostgrad(int b,int n,int m,const float * xyz1,const float * xyz2,const float * match,float * grad1,float * grad2, cudaStream_t stream); 9 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/src/nndistance.cuh: -------------------------------------------------------------------------------- 1 | void nndistance(int b,int n,const float * xyz,int m,const float * xyz2,float * result,int * result_i,float * result2,int * result2_i, cudaStream_t stream); 2 | void nndistancegrad(int b,int n,const float * xyz1,int m,const float * xyz2,const float * grad_dist1,const int * idx1,const float * grad_dist2,const int * idx2,float * grad_xyz1,float * grad_xyz2, cudaStream_t stream); 3 | -------------------------------------------------------------------------------- /PointFlow/metrics/pytorch_structural_losses/src/utils.hpp: -------------------------------------------------------------------------------- 1 | #include 2 | #include 3 | #include 4 | 5 | class Formatter { 6 | public: 7 | Formatter() {} 8 | ~Formatter() {} 9 | 10 | template Formatter &operator<<(const Type &value) { 11 | stream_ << value; 12 | return *this; 13 | } 14 | 15 | std::string str() const { return stream_.str(); } 16 | operator std::string() const { return stream_.str(); } 17 | 18 | enum ConvertToString { to_str }; 19 | 20 | std::string operator>>(ConvertToString) { return stream_.str(); } 21 | 22 | private: 23 | std::stringstream stream_; 24 | Formatter(const Formatter &); 25 | Formatter &operator=(Formatter &); 26 | }; 27 | -------------------------------------------------------------------------------- /PointFlow/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/PointFlow/models/__init__.py -------------------------------------------------------------------------------- /PointFlow/models/flow.py: -------------------------------------------------------------------------------- 1 | from .odefunc import ODEfunc, ODEnet 2 | from .normalization import MovingBatchNorm1d 3 | from .cnf import CNF, SequentialFlow 4 | 5 | 6 | def count_nfe(model): 7 | class AccNumEvals(object): 8 | 9 | def __init__(self): 10 | self.num_evals = 0 11 | 12 | def __call__(self, module): 13 | if isinstance(module, CNF): 14 | self.num_evals += module.num_evals() 15 | 16 | accumulator = AccNumEvals() 17 | model.apply(accumulator) 18 | return accumulator.num_evals 19 | 20 | 21 | def count_parameters(model): 22 | return sum(p.numel() for p in model.parameters() if p.requires_grad) 23 | 24 | 25 | def count_total_time(model): 26 | class Accumulator(object): 27 | 28 | def __init__(self): 29 | self.total_time = 0 30 | 31 | def __call__(self, module): 32 | if isinstance(module, CNF): 33 | self.total_time = self.total_time + module.sqrt_end_time * module.sqrt_end_time 34 | 35 | accumulator = Accumulator() 36 | model.apply(accumulator) 37 | return accumulator.total_time 38 | 39 | 40 | def build_model(args, input_dim, hidden_dims, context_dim, num_blocks, conditional): 41 | def build_cnf(): 42 | diffeq = ODEnet( 43 | hidden_dims=hidden_dims, 44 | input_shape=(input_dim,), 45 | context_dim=context_dim, 46 | layer_type=args.layer_type, 47 | nonlinearity=args.nonlinearity, 48 | ) 49 | odefunc = ODEfunc( 50 | diffeq=diffeq, 51 | ) 52 | cnf = CNF( 53 | odefunc=odefunc, 54 | T=args.time_length, 55 | train_T=args.train_T, 56 | conditional=conditional, 57 | solver=args.solver, 58 | use_adjoint=args.use_adjoint, 59 | atol=args.atol, 60 | rtol=args.rtol, 61 | ) 62 | return cnf 63 | 64 | chain = [build_cnf() for _ in range(num_blocks)] 65 | if args.batch_norm: 66 | bn_layers = [MovingBatchNorm1d(input_dim, bn_lag=args.bn_lag, sync=args.sync_bn) 67 | for _ in range(num_blocks)] 68 | bn_chain = [MovingBatchNorm1d(input_dim, bn_lag=args.bn_lag, sync=args.sync_bn)] 69 | for a, b in zip(chain, bn_layers): 70 | bn_chain.append(a) 71 | bn_chain.append(b) 72 | chain = bn_chain 73 | model = SequentialFlow(chain) 74 | 75 | return model 76 | 77 | 78 | def get_point_cnf(args): 79 | dims = tuple(map(int, args.dims.split("-"))) 80 | model = build_model(args, args.input_dim, dims, args.zdim, args.num_blocks, True).cuda(args.gpu) 81 | print("Number of trainable parameters of Point CNF: {}".format(count_parameters(model))) 82 | return model 83 | 84 | 85 | def get_latent_cnf(args): 86 | dims = tuple(map(int, args.latent_dims.split("-"))) 87 | model = build_model(args, args.zdim, dims, 0, args.latent_num_blocks, False).cuda(args.gpu) 88 | print("Number of trainable parameters of Latent CNF: {}".format(count_parameters(model))) 89 | return model 90 | -------------------------------------------------------------------------------- /PointFlow/prepare.sh: -------------------------------------------------------------------------------- 1 | . activate plb_v2 2 | export PYTHONPATH=$PWD:$PWD/metrics/pytorch_structural_losses/build/lib.linux-x86_64-3.8/:$PWD/../ 3 | -------------------------------------------------------------------------------- /PointFlow/scripts/cutrearrange_gen_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Number of GPUs used for parallelized training in the paper: 4 3 | 4 | dims="512-512-512" 5 | latent_dims="256-256" 6 | num_blocks=1 7 | latent_num_blocks=1 8 | zdim=2 9 | batch_size=128 10 | tr_max_sample_points=1000 11 | te_max_sample_points=1000 12 | lr=2e-3 13 | epochs=12 14 | ds=dynabs 15 | env_name='CutRearrange-v1' 16 | data_dirs='' 17 | log_name="0001_cutrearrange_pointflow" 18 | 19 | python PointFlow/train.py \ 20 | --log_name ${log_name} \ 21 | --env_name ${env_name} \ 22 | --lr ${lr} \ 23 | --train_T False \ 24 | --dataset_type ${ds} \ 25 | --filter_buffer True \ 26 | --cached_state_path '' \ 27 | --data_dirs ${data_dirs} \ 28 | --tr_max_sample_points=${tr_max_sample_points}\ 29 | --te_max_sample_points=${te_max_sample_points}\ 30 | --dims ${dims} \ 31 | --latent_dims ${latent_dims} \ 32 | --num_blocks ${num_blocks} \ 33 | --latent_num_blocks ${latent_num_blocks} \ 34 | --batch_size ${batch_size} \ 35 | --zdim ${zdim} \ 36 | --epochs ${epochs} \ 37 | --save_freq 1 \ 38 | --viz_freq 1 \ 39 | --log_freq 1 \ 40 | --val_freq 1 \ 41 | --distributed \ 42 | --use_latent_flow 43 | 44 | echo "Done" 45 | exit 0 46 | -------------------------------------------------------------------------------- /PointFlow/scripts/cutrearrangespread_gen_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Number of GPUs used for parallelized training in the paper: 8 3 | 4 | dims="512-512-512" 5 | latent_dims="256-256" 6 | num_blocks=1 7 | latent_num_blocks=1 8 | zdim=2 9 | batch_size=1024 10 | tr_max_sample_points=1000 11 | te_max_sample_points=1000 12 | lr=2e-3 13 | epochs=100 14 | ds=dynabs 15 | env_name='CutRearrangeSpread-v1' 16 | data_dirs='' 17 | log_name="0001_cutrearrangespread_pointflow" 18 | 19 | python PointFlow/train.py \ 20 | --log_name ${log_name} \ 21 | --env_name ${env_name} \ 22 | --lr ${lr} \ 23 | --train_T False \ 24 | --dataset_type ${ds} \ 25 | --cached_state_path '' \ 26 | --data_dirs ${data_dirs} \ 27 | --tr_max_sample_points=${tr_max_sample_points}\ 28 | --te_max_sample_points=${te_max_sample_points}\ 29 | --dims ${dims} \ 30 | --latent_dims ${latent_dims} \ 31 | --num_blocks ${num_blocks} \ 32 | --latent_num_blocks ${latent_num_blocks} \ 33 | --batch_size ${batch_size} \ 34 | --zdim ${zdim} \ 35 | --epochs ${epochs} \ 36 | --save_freq 1 \ 37 | --viz_freq 1 \ 38 | --log_freq 1 \ 39 | --val_freq 1 \ 40 | --distributed \ 41 | --use_latent_flow \ 42 | 43 | echo "Done" 44 | exit 0 45 | -------------------------------------------------------------------------------- /PointFlow/scripts/gathermove_gen_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Number of GPUs used for parallelized training in the paper: 4 3 | 4 | dims="512-512-512" 5 | latent_dims="256-256" 6 | num_blocks=1 7 | latent_num_blocks=1 8 | zdim=4 9 | batch_size=64 10 | tr_max_sample_points=1000 11 | te_max_sample_points=1000 12 | lr=2e-3 13 | epochs=4 14 | ds=dynabs 15 | env_name='GatherMove-v1' 16 | data_dirs='' 17 | log_name="0001_gathermove_pointflow" 18 | 19 | python PointFlow/train.py \ 20 | --log_name ${log_name} \ 21 | --env_name ${env_name} \ 22 | --lr ${lr} \ 23 | --train_T False \ 24 | --dataset_type ${ds} \ 25 | --data_dirs ${data_dirs} \ 26 | --tr_max_sample_points=${tr_max_sample_points}\ 27 | --te_max_sample_points=${te_max_sample_points}\ 28 | --dims ${dims} \ 29 | --latent_dims ${latent_dims} \ 30 | --num_blocks ${num_blocks} \ 31 | --latent_num_blocks ${latent_num_blocks} \ 32 | --batch_size ${batch_size} \ 33 | --zdim ${zdim} \ 34 | --epochs ${epochs} \ 35 | --save_freq 1 \ 36 | --viz_freq 1 \ 37 | --log_freq 1 \ 38 | --val_freq 1 \ 39 | --distributed \ 40 | --use_latent_flow 41 | 42 | echo "Done" 43 | exit 0 44 | -------------------------------------------------------------------------------- /PointFlow/scripts/liftspread_gen_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Number of GPUs used for parallelized training in the paper: 4 3 | 4 | dims="512-512-512" 5 | latent_dims="256-256" 6 | num_blocks=1 7 | latent_num_blocks=1 8 | zdim=8 9 | batch_size=48 10 | tr_max_sample_points=1000 11 | te_max_sample_points=1000 12 | lr=2e-3 13 | epochs=8 14 | ds=dynabs 15 | env_name='LiftSpread-v1' 16 | data_dirs='' 17 | log_name="0001_liftspread_pointflow" 18 | 19 | python PointFlow/train.py \ 20 | --log_name ${log_name} \ 21 | --env_name ${env_name} \ 22 | --lr ${lr} \ 23 | --train_T False \ 24 | --dataset_type ${ds} \ 25 | --data_dirs ${data_dirs} \ 26 | --tr_max_sample_points=${tr_max_sample_points}\ 27 | --te_max_sample_points=${te_max_sample_points}\ 28 | --dims ${dims} \ 29 | --latent_dims ${latent_dims} \ 30 | --num_blocks ${num_blocks} \ 31 | --latent_num_blocks ${latent_num_blocks} \ 32 | --batch_size ${batch_size} \ 33 | --zdim ${zdim} \ 34 | --epochs ${epochs} \ 35 | --save_freq 1 \ 36 | --viz_freq 1 \ 37 | --log_freq 1 \ 38 | --val_freq 1 \ 39 | --distributed \ 40 | --use_latent_flow 41 | 42 | echo "Done" 43 | exit 0 44 | -------------------------------------------------------------------------------- /PointFlow/scripts/set_cutrearrange_gen_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Number of GPUs used for parallelized training in the paper: 8 3 | 4 | dims="512-512-512" 5 | latent_dims="256-256" 6 | num_blocks=1 7 | latent_num_blocks=1 8 | zdim=2 9 | batch_size=2048 10 | tr_max_sample_points=1000 11 | te_max_sample_points=1000 12 | lr=2e-3 13 | epochs=100 14 | ds=dynabs 15 | env_name='CutRearrange-v1' 16 | data_dirs='' 17 | log_name="0001_set_cutrearrange_pointflow" 18 | 19 | python PointFlow/train.py \ 20 | --log_name ${log_name} \ 21 | --env_name ${env_name} \ 22 | --lr ${lr} \ 23 | --train_T False \ 24 | --dataset_type ${ds} \ 25 | --filter_buffer False \ 26 | --cached_state_path '' \ 27 | --data_dirs ${data_dirs} \ 28 | --tr_max_sample_points=${tr_max_sample_points}\ 29 | --te_max_sample_points=${te_max_sample_points}\ 30 | --dims ${dims} \ 31 | --latent_dims ${latent_dims} \ 32 | --num_blocks ${num_blocks} \ 33 | --latent_num_blocks ${latent_num_blocks} \ 34 | --batch_size ${batch_size} \ 35 | --zdim ${zdim} \ 36 | --epochs ${epochs} \ 37 | --save_freq 1 \ 38 | --viz_freq 1 \ 39 | --log_freq 1 \ 40 | --val_freq 1 \ 41 | --distributed \ 42 | --use_latent_flow \ 43 | --load_set_from_buffer True 44 | 45 | echo "Done" 46 | exit 0 47 | -------------------------------------------------------------------------------- /PointFlow/scripts/set_cutrearrangespread_gen_dist.sh: -------------------------------------------------------------------------------- 1 | #! /bin/bash 2 | # Number of GPUs used for parallelized training in the paper: 10 3 | 4 | dims="512-512-512" 5 | latent_dims="256-256" 6 | num_blocks=1 7 | latent_num_blocks=1 8 | zdim=2 9 | batch_size=1024 10 | tr_max_sample_points=1000 11 | te_max_sample_points=1000 12 | lr=2e-3 13 | epochs=11 14 | ds=dynabs 15 | env_name='CutRearrangeSpread-v1' 16 | data_dirs='' 17 | log_name="0001_set_cutrearrangespread_pointflow" 18 | 19 | python PointFlow/train.py \ 20 | --log_name ${log_name} \ 21 | --env_name ${env_name} \ 22 | --lr ${lr} \ 23 | --train_T False \ 24 | --dataset_type ${ds} \ 25 | --cached_state_path '' \ 26 | --data_dirs ${data_dirs} \ 27 | --tr_max_sample_points=${tr_max_sample_points}\ 28 | --te_max_sample_points=${te_max_sample_points}\ 29 | --dims ${dims} \ 30 | --latent_dims ${latent_dims} \ 31 | --num_blocks ${num_blocks} \ 32 | --latent_num_blocks ${latent_num_blocks} \ 33 | --batch_size ${batch_size} \ 34 | --zdim ${zdim} \ 35 | --epochs ${epochs} \ 36 | --save_freq 1 \ 37 | --viz_freq 1 \ 38 | --log_freq 1 \ 39 | --val_freq 1 \ 40 | --distributed \ 41 | --use_latent_flow \ 42 | --load_set_from_buffer True 43 | 44 | echo "Done" 45 | exit 0 46 | -------------------------------------------------------------------------------- /assets/cuttingboard/Cutting_Board.mtl: -------------------------------------------------------------------------------- 1 | # 3ds Max Wavefront OBJ Exporter v0.97b - (c)2007 guruware 2 | # File Created: 24.03.2021 12:23:29 3 | 4 | newmtl M_Cutting_Board 5 | Ns 100.0000 6 | Ni 1.6000 7 | d 1.0000 8 | Tr 0.0000 9 | Tf 1.0000 1.0000 1.0000 10 | illum 2 11 | Ka 0.5880 0.5880 0.5880 12 | Kd 0.5000 0.5000 0.5000 13 | Ks 0.0000 0.0000 0.0000 14 | Ke 0.0000 0.0000 0.0000 15 | map_Ka textures/Cutting_Board_Diffuse.png 16 | map_Kd textures/Cutting_Board_Diffuse.png 17 | map_Ks textures/Cutting_Board_Specular.png 18 | map_bump textures/Cutting_Board_Normal.png 19 | bump textures/Cutting_Board_Normal.png 20 | -------------------------------------------------------------------------------- /assets/table/Table_Coffee_RiceChest.mtl: -------------------------------------------------------------------------------- 1 | # 2 | ## Alias OBJ Material File 3 | # Exported from SketchUp, (c) 2000-2012 Trimble Navigation Limited 4 | 5 | newmtl _Wood_Cherry_ 6 | Ka 0.000000 0.000000 0.000000 7 | Kd 0.588235 0.329412 0.156863 8 | Ks 0.330000 0.330000 0.330000 9 | map_Kd Table_Coffee_RiceChest/_Wood_Cherry_.jpg 10 | 11 | -------------------------------------------------------------------------------- /assets/table/Table_Coffee_RiceChest/_Wood_Cherry_.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/assets/table/Table_Coffee_RiceChest/_Wood_Cherry_.jpg -------------------------------------------------------------------------------- /core/diffskill/plan/project_func.py: -------------------------------------------------------------------------------- 1 | import torch 2 | from core.utils.diffskill_utils import traj_gaussian_logl, gaussian_logl, batch_pred 3 | 4 | 5 | def generate_batch_loss_func(agent, tids, tiled_u_obs, tiled_u_goal): 6 | """ 7 | :param tids: Tool to use at each step 8 | :param tiled_u_obs: tiled initial latent 9 | :param tiled_u_goal: tiled goal latent 10 | """ 11 | 12 | def batch_loss_func(u): 13 | u = u.view(tiled_u_obs.shape[0], len(tids), -1) 14 | traj_succ = [] 15 | for step, tid in enumerate(tids): 16 | if step == 0: 17 | curr_u = tiled_u_obs 18 | else: 19 | curr_u = u[:, step - 1] 20 | curr_goal = u[:, step] 21 | succ = batch_pred(agent.feas[tids[step]], {'obs': curr_u, 'goal': curr_goal, 'eval': True}) 22 | traj_succ.append(succ[None, :]) 23 | traj_succ = torch.cat(traj_succ, dim=0) # [plan_step, num_plans] 24 | pred_r = batch_pred(agent.reward_predictor, {'obs': u[:, len(tids) - 1], 'goal': tiled_u_goal, 'eval': True}) 25 | pred_r = torch.min(pred_r, torch.zeros_like(pred_r)) # score is negative emd, which should be less than zero 26 | traj_score = torch.max(traj_succ[0], torch.ones_like(traj_succ[0]) * 1e-2) 27 | for i in range(1, len(traj_succ)): 28 | traj_score = traj_score * torch.max(traj_succ[i], torch.ones_like(traj_succ[i]) * 1e-2) 29 | traj_score = traj_score * torch.exp(pred_r * 10.) 30 | total_loss = -traj_score 31 | 32 | # Add auxiliary losses for other steps as well. This will affect the CutRearrange 33 | for i in range(len(traj_succ)): 34 | total_loss += -0.01 * traj_succ[i] 35 | 36 | traj_zlogl = traj_gaussian_logl(u) 37 | return total_loss, {'traj_succ': traj_succ, 38 | 'pred_r': pred_r, 39 | 'traj_score': traj_score, 40 | 'traj_zlogl': traj_zlogl} 41 | 42 | return batch_loss_func 43 | 44 | 45 | def generate_project_func(args, plan_step, t_min=None, t_max=None): 46 | if args.input_mode == 'rgbd': 47 | def project(z): 48 | copy_z = z.view(-1, plan_step, args.dimz) 49 | with torch.no_grad(): 50 | z_logl = gaussian_logl(copy_z) 51 | projected_z = copy_z / torch.max(torch.sqrt((z_logl[:, :, None] / args.min_zlogl)), torch.ones(1, device=z.device)) 52 | return projected_z.view(*z.shape).detach() 53 | 54 | return project if 'min_zlogl' in args.__dict__ else None 55 | 56 | else: 57 | def project(u): 58 | projected_u = u.view(-1, plan_step, args.dimu) 59 | with torch.no_grad(): 60 | # projected_u[:, :, -3:] = torch.clip(projected_u[:, :, -3:], 0., 1.) # Clip to zero and one. #TODO Use the actual boundary 61 | trans = projected_u[:, :, -3:] 62 | trans = torch.max(t_min.reshape(1, 1, 3), trans) 63 | trans = torch.min(t_max.reshape(1, 1, 3), trans) 64 | projected_u[:, :, -3:] = trans 65 | return projected_u.view(u.shape).detach() 66 | 67 | return project 68 | -------------------------------------------------------------------------------- /core/models/cnn_actor.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from core.utils.diffskill_utils import weight_init 4 | from core.models.cnn_encoder import CNNEncoder 5 | 6 | class CNNActor(nn.Module): 7 | """ Image based actor""" 8 | 9 | def __init__(self, args, obs_shape, action_dim, hidden_dim=256): 10 | super().__init__() 11 | self.args = args 12 | self.encoder = CNNEncoder(obs_shape, args.actor_feature_dim) 13 | latent_dim = args.actor_feature_dim 14 | self.action_mlp = nn.Sequential(nn.Linear(latent_dim, hidden_dim), 15 | nn.ReLU(), 16 | nn.Linear(hidden_dim, hidden_dim), 17 | nn.ReLU(), 18 | nn.Linear(hidden_dim, hidden_dim), 19 | nn.ReLU(), 20 | nn.Linear(hidden_dim, action_dim)) 21 | self.done_mlp = nn.Sequential(nn.Linear(latent_dim, hidden_dim), 22 | nn.ReLU(), 23 | nn.Linear(hidden_dim, hidden_dim), 24 | nn.ReLU(), 25 | nn.Linear(hidden_dim, 1)) 26 | self.apply(weight_init) 27 | self.loss = nn.MSELoss(reduction='none') 28 | 29 | def forward(self, obs, detach_encoder=False): 30 | obs = self.encoder(obs, detach=detach_encoder) 31 | return self.action_mlp(obs), self.done_mlp(obs) 32 | -------------------------------------------------------------------------------- /core/models/cnn_encoder.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | 4 | 5 | class CNNEncoder(nn.Module): 6 | """Convolutional encoder for image-based observations.""" 7 | 8 | def __init__(self, obs_shape, feature_dim): 9 | super().__init__() 10 | 11 | assert len(obs_shape) == 3 12 | self.num_layers = 4 13 | self.num_filters = 32 14 | self.output_logits = False 15 | self.feature_dim = feature_dim 16 | 17 | self.convs = nn.ModuleList([ 18 | nn.Conv2d(obs_shape[-1], self.num_filters, 3, stride=2), 19 | nn.Conv2d(self.num_filters, self.num_filters, 3, stride=1), 20 | nn.Conv2d(self.num_filters, self.num_filters, 3, stride=1), 21 | nn.Conv2d(self.num_filters, self.num_filters, 3, stride=1) 22 | ]) 23 | 24 | self.head = nn.Sequential( 25 | nn.Linear(20000, self.feature_dim), 26 | nn.LayerNorm(self.feature_dim)) 27 | 28 | self.outputs = dict() 29 | 30 | def forward_conv(self, obs): 31 | self.outputs['obs'] = obs 32 | 33 | conv = torch.relu(self.convs[0](obs)) 34 | self.outputs['conv1'] = conv 35 | 36 | for i in range(1, self.num_layers): 37 | conv = torch.relu(self.convs[i](conv)) 38 | self.outputs['conv%s' % (i + 1)] = conv 39 | h = conv.view(conv.size(0), -1) 40 | return h 41 | 42 | def forward(self, obs, detach=False): 43 | h = self.forward_conv(obs) 44 | if detach: 45 | h = h.detach() 46 | 47 | out = self.head(h) 48 | if not self.output_logits: 49 | out = torch.tanh(out) 50 | 51 | self.outputs['out'] = out 52 | 53 | return out 54 | 55 | def log(self, logger, step): 56 | for k, v in self.outputs.items(): 57 | logger.log_histogram(f'train_encoder/{k}_hist', v, step) 58 | if len(v.shape) > 2: 59 | logger.log_image(f'train_encoder/{k}_img', v[0], step) 60 | 61 | for i in range(self.num_layers): 62 | logger.log_param(f'train_encoder/conv{i + 1}', self.convs[i], step) 63 | -------------------------------------------------------------------------------- /core/traj_opt/gen_init_target/cutrearrange_generator.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from core.traj_opt.gen_init_target.state_generator import StateGenerator 3 | import copy 4 | 5 | 6 | def cut(state, cut_loc, slide_a, slide_b): 7 | ns = copy.deepcopy(state) 8 | flag = ns['state'][0][:, 0] <= cut_loc 9 | ns['state'][0][flag, 0] -= slide_a 10 | ns['state'][0][np.logical_not(flag), 0] += slide_b 11 | return ns, flag 12 | 13 | 14 | def move_cluster(state, flag, dx, dz, dy=0): 15 | # need to move forward 16 | new_state = copy.deepcopy(state) 17 | mean = new_state['state'][0][flag].mean(axis=0) 18 | new_state['state'][0][flag] += np.array([dx, mean[1] + dy, dz]) - mean 19 | return new_state 20 | 21 | 22 | class CutrearrangeGenerator(StateGenerator): 23 | def __init__(self, *args, **kwargs): 24 | super(CutrearrangeGenerator, self).__init__(*args, **kwargs) 25 | self.env.reset() 26 | for i in range(50): 27 | self.env.step(np.array([0, 0, 0] + [0, 0, 0] + [0, 0, 0, 0])) 28 | self.initial_state = self.env.get_state() 29 | self.N = 1 30 | 31 | def _generate(self): 32 | for i in range(self.N): 33 | def rand(a, b): 34 | return np.random.random() * (b - a) + a 35 | 36 | state = copy.deepcopy(self.initial_state) 37 | width = np.array([rand(0.2, 0.24), 0.08, rand(0.04, 0.08)]) 38 | state['state'][0] = (np.random.random((len(state['state'][0]), 3)) * 2 - 1) * (0.5 * width) + np.array([0.5, 0.0669, 0.5]) 39 | 40 | da = db = 0. 41 | if np.random.randint(2): 42 | da = rand(0.1, 0.13) 43 | else: 44 | db = rand(0.1, 0.13) 45 | 46 | cut_loc = rand(0.48, 0.52) 47 | cutted, flag = cut(state, cut_loc, da, db) 48 | 49 | def sample_target(): 50 | # make sure that it is moved away 51 | y = rand(0.3, 0.35) * (np.random.randint(2) * 2 - 1) + 0.5 52 | x = rand(0.3, 0.7) 53 | return (x, y) 54 | 55 | while True: # sample two targets which are far away 56 | a = sample_target() 57 | b = sample_target() 58 | if np.linalg.norm(np.array(a) - np.array(b)) >= 0.3: 59 | break 60 | 61 | move1 = move_cluster(cutted, flag, *a) 62 | move2 = move_cluster(move1, np.logical_not(flag), *b) 63 | 64 | self.save_init(3 * i, state) 65 | self.save_init(3 * i + 1, cutted) 66 | self.save_init(3 * i + 2, move1) 67 | self.save_target(3 * i, cutted) 68 | self.save_target(3 * i + 1, move1) 69 | self.save_target(3 * i + 2, move2) 70 | -------------------------------------------------------------------------------- /core/traj_opt/gen_init_target/gathermove_generator.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from core.traj_opt.gen_init_target.state_generator import StateGenerator 3 | from functools import partial 4 | 5 | 6 | class GathermoveGenerator(StateGenerator): 7 | def __init__(self, *args, **kwargs): 8 | super(GathermoveGenerator, self).__init__(*args, **kwargs) 9 | self.N = 100 10 | self.xs = np.linspace(0.36, 0.4, self.N) 11 | self.rs = np.linspace(0.04, 0.07, self.N // 10) 12 | 13 | def case(self, cfg, i): # Sphere 14 | # Disable collision checking during generation. Need to disable both, or the collision will not be disabled 15 | cfg.PRIMITIVES[0]['collision_group'] = (0, 0, 0) 16 | cfg.PRIMITIVES[1]['collision_group'] = (0, 0, 0) 17 | cfg.SHAPES[0]['seed'] = i 18 | 19 | def case2(self, cfg, i): 20 | cfg.SHAPES[0]['shape'] = 'sphere' 21 | del cfg.SHAPES[0]['pos_max'], cfg.SHAPES[0]['pos_min'], cfg.SHAPES[0]['seed'] 22 | pos = (0.65, 0.08, 0.5) 23 | r = self.rs[i * 11117771 % 12837119 % self.N // 10] 24 | cfg.SHAPES[0]['radius'] = r 25 | new_pos = (self.xs[i], r + 0.08, pos[2]) 26 | cfg.SHAPES[0]['init_pos'] = new_pos 27 | 28 | def _generate(self): 29 | for i in range(self.N): 30 | self.env.reset(target_cfg_modifier=partial(self.case, i=i)) 31 | for _ in range(10): 32 | self.env.step(np.array([0] * 13)) # Wait for dough to drop 33 | self.save_init(2 * i) 34 | if i % 2 ==0: 35 | n = 25 36 | else: 37 | n = 10 38 | for j in range(n): # Grasp in one direction 39 | action = np.array([0] * 7 + [0.] * 6) 40 | action[6] = ((j < 15) - 0.5) * 2 41 | self.env.step(action) 42 | self.taichi_env.primitives[0].set_state(0, [0.7, 0.06, 0.5, 0.707, 0.707, 0., 0., 0.4]) 43 | for j in range(25): # Grasp in the other direction 44 | action = np.array([0] * 7 + [0.] * 6) 45 | action[6] = ((j < 15) - 0.5) * 2 46 | self.env.step(action) 47 | self.taichi_env.primitives[0].set_state(0, [0.7, 0.06, 0.5, 0.5, 0.5, -0.5, 0.5, 0.4]) 48 | self.save_init(2 * i + 1) 49 | self.save_target(i) 50 | 51 | for i in range(self.N): 52 | self.env.reset(target_cfg_modifier=partial(self.case2, i=i)) 53 | self.save_target(i + self.N) 54 | -------------------------------------------------------------------------------- /core/traj_opt/gen_init_target/liftspread_generator.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from diffskill.gen_init_target.state_generator import StateGenerator 3 | from functools import partial 4 | 5 | 6 | class LiftspreadGenerator(StateGenerator): 7 | def __init__(self, *args, **kwargs): 8 | super(LiftspreadGenerator, self).__init__(*args, **kwargs) 9 | self.N = 10 10 | self.xs = np.linspace(0.32, 0.38, self.N) # Dough location on the cutting board 11 | self.xs_target = np.linspace(0.25, 0.35, self.N) # Target location of the box dough 12 | self.rs = np.linspace(0.04, 0.07, self.N) 13 | 14 | def randomize_tool_state(self, cfg, pin_min_y): 15 | pos = eval(cfg.PRIMITIVES[0]['init_pos']) 16 | height = np.random.rand() * 0.1 + 0.04 17 | x_shift = np.random.rand() * 0.1 - 0.08 18 | cfg.PRIMITIVES[0]['init_pos'] = (pos[0] + x_shift, float(pin_min_y + height), pos[2]) 19 | 20 | def case1(self, cfg, i, j): # Sphere 21 | pos = eval(cfg.SHAPES[0]['init_pos']) 22 | cfg.SHAPES[0]['radius'] = self.rs[j] 23 | new_pos = (self.xs[i], self.rs[j] + 0.1, pos[2]) 24 | cfg.SHAPES[0]['init_pos'] = new_pos 25 | 26 | self.randomize_tool_state(cfg, self.rs[j] * 2 + 0.1) 27 | 28 | def case11(self, cfg, i, j): # Sphere on initial place 29 | pos = eval(cfg.SHAPES[0]['init_pos']) 30 | cfg.SHAPES[0]['shape'] = 'sphere' 31 | cfg.SHAPES[0]['radius'] = self.rs[j] 32 | new_pos = (pos[0], self.rs[j] + 0.03, pos[2]) 33 | cfg.SHAPES[0]['init_pos'] = new_pos 34 | 35 | self.randomize_tool_state(cfg, self.rs[j] * 2 + 0.1) 36 | 37 | def case3(self, cfg, i, j): # Target 38 | cfg.SHAPES[0]['shape'] = 'box' 39 | w = np.linspace(0.26, 0.3, self.N) 40 | h = np.linspace(0.025, 0.04, self.N) 41 | w_id, h_id = np.random.randint(0, self.N), np.random.randint(0, self.N) 42 | pos = eval(cfg.SHAPES[0]['init_pos']) 43 | new_pos = (self.xs_target[i], h[h_id] + 0.1, pos[2]) 44 | new_size = (w[w_id], h[h_id], 0.16) 45 | cfg.SHAPES[0]['init_pos'] = new_pos 46 | cfg.SHAPES[0]['width'] = new_size 47 | del cfg.SHAPES[0]['radius'] 48 | self.randomize_tool_state(cfg, 0.19) 49 | 50 | def _generate(self): 51 | for case_id, case in enumerate([self.case1, self.case11]): 52 | for i in range(self.N): 53 | for j in range(self.N): 54 | self.env.reset(target_cfg_modifier=partial(case, i=i, j=j)) 55 | self.save_init(idx=i * self.N + j + self.N * self.N * case_id) 56 | 57 | np.random.seed(0) 58 | for case_id, case in enumerate([self.case1, self.case3]): 59 | for i in range(self.N): 60 | for j in range(self.N): 61 | self.env.reset(target_cfg_modifier=partial(case, i=i, j=j)) 62 | self.save_target(idx=i * self.N + j + self.N * self.N * case_id) 63 | -------------------------------------------------------------------------------- /core/utils/camera_info.json: -------------------------------------------------------------------------------- 1 | { 2 | "class_name" : "PinholeCameraParameters", 3 | "extrinsic" : 4 | [ 5 | 0.99988534836532872, 6 | -0.011445440959168104, 7 | -0.0099142324763738954, 8 | 0.0, 9 | -0.012447164521824866, 10 | -0.24839850661516383, 11 | -0.96857795246780376, 12 | 0.0, 13 | 0.0086231212279553448, 14 | 0.96859030750498776, 15 | -0.24851249060697403, 16 | 0.0, 17 | -0.60497881910966111, 18 | -0.37683675069503242, 19 | 1.3058392152363218, 20 | 1.0 21 | ], 22 | "intrinsic" : 23 | { 24 | "height" : 1080, 25 | "intrinsic_matrix" : 26 | [ 27 | 935.30743608719376, 28 | 0.0, 29 | 0.0, 30 | 0.0, 31 | 935.30743608719376, 32 | 0.0, 33 | 959.5, 34 | 539.5, 35 | 1.0 36 | ], 37 | "width" : 1920 38 | }, 39 | "version_major" : 1, 40 | "version_minor" : 0 41 | } -------------------------------------------------------------------------------- /core/utils/core_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import random 4 | 5 | 6 | class VArgs(object): 7 | def __init__(self, vv): 8 | for key, val in vv.items(): 9 | setattr(self, key, val) 10 | 11 | def set_ipdb_debugger(): 12 | import sys 13 | import ipdb 14 | import traceback 15 | 16 | def info(t, value, tb): 17 | if t == KeyboardInterrupt: 18 | traceback.print_exception(t, value, tb) 19 | return 20 | else: 21 | traceback.print_exception(t, value, tb) 22 | ipdb.pm() 23 | 24 | sys.excepthook = info 25 | 26 | def set_random_seed(seed): 27 | random.seed(seed) 28 | np.random.seed(seed) 29 | torch.manual_seed(seed) 30 | torch.cuda.manual_seed(seed) 31 | torch.cuda.manual_seed_all(seed) 32 | torch.backends.cudnn.deterministic = True 33 | torch.backends.cudnn.benchmark = False 34 | 35 | 36 | def list_to_array(l, tensor=False, axis=0): 37 | """Concatenate list and save index""" 38 | if axis == 0: 39 | list_idx = [len(ele) for ele in l] 40 | else: 41 | list_idx = [ele.shape[1] for ele in l] 42 | 43 | if not tensor: 44 | l = np.concatenate(l, axis=axis) 45 | else: 46 | l = torch.cat(l, dim=axis) 47 | return l, list_idx 48 | 49 | 50 | def array_to_list(arr, list_idx, axis=0): 51 | """Concatenate array to list""" 52 | ret = [] 53 | cnt = 0 54 | for num in list_idx: 55 | if axis == 0: 56 | ret.append(arr[cnt:cnt + num]) 57 | elif axis == 1: 58 | ret.append(arr[:, cnt:cnt + num]) 59 | else: 60 | raise NotImplementedError 61 | cnt += num 62 | return ret 63 | -------------------------------------------------------------------------------- /core/utils/pc_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | 4 | def batch_resample_pc(pcs, N): 5 | """ Resample the pointcloud to have N points. pcs has shape B x M x 3 6 | Return pc shape: B x N x 3 7 | """ 8 | B, M = pcs.shape[:2] 9 | if M == N: 10 | return pcs 11 | if M < N: 12 | K = N - M 13 | idx = np.random.choice(M, K) 14 | return np.concatenate([pcs, pcs[:, idx]], axis=1) 15 | else: 16 | idx = np.random.choice(M, N, replace=False) 17 | return pcs[:, idx] 18 | 19 | 20 | if __name__ == '__main__': 21 | import numpy as np 22 | 23 | a = np.random.random([128, 500, 3]) 24 | print(batch_resample_pc(a, 1000).shape) 25 | 26 | 27 | def resample_pc(pc, N): 28 | """Resample the pointcloud to have N points""" 29 | if len(pc) == N: 30 | return pc.reshape(N, 3) 31 | if len(pc) < N: 32 | m = N - len(pc) 33 | idx = np.random.choice(len(pc), m) 34 | return np.vstack([pc, pc[idx]]) 35 | else: 36 | idx = np.random.choice(len(pc), N) 37 | return pc[idx].reshape(N, 3) 38 | 39 | 40 | def decompose_pc(pc, label, N=None): 41 | """Decompose a pointcloud into set based on dbscan label. 42 | If N is not None, resample each pc to have N points 43 | """ 44 | assert len(pc) == len(label) 45 | max_label = np.max(label) 46 | pcs = [] 47 | for l in range(max_label + 1): 48 | p = pc[np.where(label == l)] 49 | p = resample_pc(p, N) 50 | pcs.append(p) 51 | return pcs 52 | 53 | 54 | def combine_pc_list(list_pcs): 55 | pcs = [] 56 | for i in range(len(list_pcs)): 57 | pcs.append(np.vstack(list_pcs[i])) 58 | return pcs 59 | -------------------------------------------------------------------------------- /core/utils/simple_logger.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | 4 | class LogDict(object): 5 | def __init__(self): 6 | self.dict = None 7 | 8 | def log_dict(self, d): 9 | if self.dict is None: 10 | self.dict = {} 11 | for key, val in d.items(): 12 | self.dict[key] = [val] 13 | else: 14 | for key, val in d.items(): 15 | self.dict[key].append(val) 16 | 17 | def agg(self, reduction='sum', numpy=False): 18 | assert reduction in ['sum', 'mean'] 19 | ret = {key: sum(self.dict[key]) for key in self.dict.keys()} 20 | if reduction == 'mean': 21 | for key in self.dict.keys(): 22 | ret[key] /= len(self.dict[key]) 23 | if numpy: 24 | for key in ret.keys(): 25 | if isinstance(ret[key], torch.Tensor): 26 | ret[key] = ret[key].item() # Assume scalar 27 | return ret 28 | -------------------------------------------------------------------------------- /core/utils/torch_chamfer.py: -------------------------------------------------------------------------------- 1 | import torch 2 | 3 | def compute_chamfer_distance(x, y): 4 | """ 5 | Compute the chamfer distance loss between two point sets. 6 | :param x: (B, N, 3) torch tensor 7 | :param y: (B, M, 3) torch tensor 8 | :return: (B,) torch tensor 9 | """ 10 | x = x.unsqueeze(2).repeat(1, 1, y.shape[1], 1) 11 | y = y.unsqueeze(1).repeat(1, x.shape[1], 1, 1) 12 | dist = torch.sum((x - y) ** 2, dim=-1) 13 | dist1, _ = torch.min(dist, dim=1) 14 | dist2, _ = torch.min(dist, dim=2) 15 | return dist1.mean(dim=1) + dist2.mean(dim=1) 16 | 17 | if __name__ == '__main__': 18 | x = torch.rand(512, 500, 3) 19 | y = torch.rand(512, 200, 3) 20 | print(compute_chamfer_distance(x, y).shape) 21 | -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | name: plb 2 | channels: 3 | - conda-forge 4 | - defaults 5 | - pytorch 6 | dependencies: 7 | - python=3.6 8 | - numpy 9 | - pip 10 | - pip: 11 | - joblib 12 | - gym==0.14.0 13 | - moviepy 14 | - opencv-python 15 | - matplotlib 16 | - transforms3d 17 | - python-dateutil 18 | - cloudpickle 19 | - pyyaml 20 | - yacs 21 | - scikit-learn 22 | - scipy 23 | - taichi==0.7.26 24 | - pyquaternion 25 | - transformations 26 | - natsort 27 | - wandb 28 | - torch_geometric 29 | - pytorch3d 30 | - pillow==8.4.0 31 | - imageio==2.15.0 32 | - torchdiffeq 33 | - open3d 34 | -------------------------------------------------------------------------------- /media/CutRearrange-v1_PASTA.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/media/CutRearrange-v1_PASTA.gif -------------------------------------------------------------------------------- /media/CutRearrangeSpread-v1_PASTA.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/media/CutRearrangeSpread-v1_PASTA.gif -------------------------------------------------------------------------------- /media/LiftSpread-v1_PASTA.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/media/LiftSpread-v1_PASTA.gif -------------------------------------------------------------------------------- /media/pasta.jpeg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/media/pasta.jpeg -------------------------------------------------------------------------------- /plb/__init__.py: -------------------------------------------------------------------------------- 1 | from . import envs -------------------------------------------------------------------------------- /plb/config/__init__.py: -------------------------------------------------------------------------------- 1 | from .utils import load -------------------------------------------------------------------------------- /plb/config/default_config.py: -------------------------------------------------------------------------------- 1 | from yacs.config import CfgNode as CN 2 | 3 | _C = CN() 4 | 5 | cfg = _C 6 | 7 | # ---------------------------------------------------------------------------- # 8 | # Simulator 9 | # ---------------------------------------------------------------------------- # 10 | _C.SIMULATOR = CN() 11 | _C.SIMULATOR.dim = 3 12 | _C.SIMULATOR.quality = 1 # control the number of particles and size of the grids 13 | _C.SIMULATOR.quality_multiplier = 1. # hack to increase quality into a float value.. 14 | _C.SIMULATOR.yield_stress = 50. 15 | _C.SIMULATOR.dtype = "float32" 16 | _C.SIMULATOR.max_steps = 1024 17 | _C.SIMULATOR.n_particles = 30000 18 | _C.SIMULATOR.lower_bound = 0. 19 | _C.SIMULATOR.E = 5e3 20 | _C.SIMULATOR.nu = 0.15 # Young's modulus and Poisson's ratio 21 | _C.SIMULATOR.ground_friction = 1.5 22 | _C.SIMULATOR.gravity = (0, -1, 0) 23 | 24 | # ---------------------------------------------------------------------------- # 25 | # PRIMITIVES, i.e., Controller 26 | # ---------------------------------------------------------------------------- # 27 | _C.PRIMITIVES = list() 28 | 29 | # ---------------------------------------------------------------------------- # 30 | # Controller 31 | # ---------------------------------------------------------------------------- # 32 | _C.SHAPES = list() 33 | 34 | # ---------------------------------------------------------------------------- # 35 | # RENDERER 36 | # ---------------------------------------------------------------------------- # 37 | _C.RENDERER = RENDERER = CN() 38 | RENDERER.name = 'plb' 39 | RENDERER.spp = 50 40 | RENDERER.max_ray_depth = 2 41 | RENDERER.image_res = (128, 128) 42 | RENDERER.tina_img_res = 1024 43 | RENDERER.voxel_res = (168, 168, 168) 44 | RENDERER.target_res = (64, 64, 64) 45 | 46 | RENDERER.dx = 1. / 150 47 | RENDERER.sdf_threshold = 0.37 * 0.56 48 | RENDERER.max_ray_depth = 2 49 | RENDERER.bake_size = 6 50 | RENDERER.use_roulette = False 51 | 52 | RENDERER.light_direction = (2., 1., 0.7) 53 | RENDERER.camera_pos = (0.5, 1.2, 4.) 54 | RENDERER.camera_rot = (0.2, 0) 55 | RENDERER.cam_center = (0.5, 0., 0.5) 56 | RENDERER.cam_theta = 0. 57 | RENDERER.cam_phi = -0.8 58 | RENDERER.cam_radius = 0.88 59 | RENDERER.use_directional_light = False 60 | RENDERER.use_roulette = False 61 | RENDERER.max_num_particles = 1000000 62 | 63 | RENDERER.mesh = False 64 | 65 | # ---------------------------------------------------------------------------- # 66 | # ENV 67 | # ---------------------------------------------------------------------------- # 68 | _C.ENV = ENV = CN() 69 | 70 | loss = ENV.loss = CN() 71 | loss.soft_contact = False 72 | loss_weight = loss.weight = CN() 73 | loss_weight.sdf = 10 74 | loss_weight.density = 10 75 | loss_weight.contact = 1 76 | loss.target_path = '' 77 | 78 | ENV.n_observed_particles = 200 79 | ENV.cached_state_path = '' 80 | ENV.env_name = '' 81 | 82 | _C.VARIANTS = list() 83 | 84 | 85 | def get_cfg_defaults(): 86 | return _C.clone() 87 | -------------------------------------------------------------------------------- /plb/config/utils.py: -------------------------------------------------------------------------------- 1 | from .default_config import get_cfg_defaults 2 | from yacs.config import CfgNode 3 | 4 | def make_cls_config(self, cfg=None, **kwargs): 5 | _cfg = self.default_config() 6 | if cfg is not None: 7 | if isinstance(cfg, str): 8 | _cfg.merge_from_file(cfg) 9 | else: 10 | _cfg.merge_from_other_cfg(cfg) 11 | if len(kwargs) > 0: 12 | _cfg.merge_from_list(sum(list(kwargs.items()), ())) 13 | return _cfg 14 | 15 | 16 | def purge_cfg(cfg: CfgNode): 17 | """Purge configuration for clean logs and logical check. 18 | Notes: 19 | If a CfgNode has 'TYPE' attribute, 20 | its CfgNode children the key of which do not contain 'TYPE' will be removed. 21 | """ 22 | target_key = cfg.get('TYPE', None) 23 | removed_keys = [] 24 | for k, v in cfg.items(): 25 | if isinstance(v, CfgNode): 26 | if target_key is not None and (k != target_key): 27 | removed_keys.append(k) 28 | else: 29 | purge_cfg(v) 30 | for k in removed_keys: 31 | del cfg[k] 32 | 33 | def load(path=None, opts=None): 34 | cfg = get_cfg_defaults() 35 | if path is not None: 36 | cfg.merge_from_file(path) 37 | if opts is not None: 38 | cfg.merge_from_list(opts) 39 | purge_cfg(cfg) 40 | cfg.freeze() 41 | return cfg -------------------------------------------------------------------------------- /plb/engine/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/plb/engine/__init__.py -------------------------------------------------------------------------------- /plb/engine/primitive/__init__.py: -------------------------------------------------------------------------------- 1 | from .primitives import Primitives -------------------------------------------------------------------------------- /plb/engine/renderer/__init__.py: -------------------------------------------------------------------------------- 1 | from .renderer import Renderer -------------------------------------------------------------------------------- /plb/engine/renderer/renderer_tester.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import taichi as ti 4 | from plb.config import load 5 | from plb.engine.renderer.renderer import Renderer 6 | from plb.engine.primitive.primitives import Sphere 7 | 8 | 9 | def test_render(): 10 | cfg = load() 11 | sphere1 = Sphere(radius=0.2) 12 | sphere2 = Sphere(radius=0.1) 13 | sphere3 = Sphere(radius=0.1) 14 | 15 | render = Renderer(cfg.RENDERER, (sphere1, sphere2, sphere3)) 16 | 17 | sphere1.set_state(0, [0.1, 0.2, 0.2, 1, 0, 0, 0]) 18 | sphere2.set_state(0, [0.9, 0.2, 0.2, 1, 0, 0, 0]) 19 | sphere3.set_state(0, [0.5, 0.8, 0.6, 1, 0, 0, 0]) 20 | # sphere1.set_color([0.3, 0.7, 0.3]) 21 | # sphere2.set_color([0.3, 0.3, 0.7]) 22 | # sphere3.set_color([0.7, 0.3, 0.7]) 23 | 24 | x = np.random.random((10000, 3)) * 0.2 + np.array([0.5, 0.5, 0.5]) 25 | 26 | render.set_particles(x, np.zeros(10000, ) + (255 << 8) + 255) 27 | 28 | img = render.render_frame(40, shape=1, primitive=1, target=0) 29 | cv2.imshow('x', img[:, :, 3] / np.max(img[:, :, 3])) 30 | cv2.waitKey(0) 31 | 32 | 33 | if __name__ == '__main__': 34 | test_render() 35 | -------------------------------------------------------------------------------- /plb/engine/renderer/test.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # Name : test 3 | # Author : Zhiao Huang 4 | # Email : z2huang@eng.ucsd.edu 5 | # Date : 2/13/2021 6 | # Distributed under terms of the MIT license. 7 | 8 | import os 9 | import cv2 10 | import numpy as np 11 | import torch 12 | import copy 13 | import pickle 14 | from typing import Optional, TypeVar, Type, Union 15 | 16 | import pandas as pd 17 | import seaborn as sns 18 | from matplotlib import pyplot as plt 19 | 20 | import matplotlib.pyplot as plt 21 | from plb.envs import make 22 | env = make("Move-v1") 23 | env.reset() 24 | 25 | #%% 26 | print(env.taichi_env.renderer.target_density.to_numpy().sum()) 27 | print(env.taichi_env.renderer.visualize_target) 28 | plt.imshow(env.render(mode='rgb_array')[..., ::-1]) 29 | plt.show() 30 | 31 | -------------------------------------------------------------------------------- /plb/engine/shapes/__init__.py: -------------------------------------------------------------------------------- 1 | from .shape_maker import Shapes -------------------------------------------------------------------------------- /plb/envs/__init__.py: -------------------------------------------------------------------------------- 1 | import gym 2 | from .env import PlasticineEnv 3 | from gym import register 4 | 5 | ENVS = [] 6 | for env_name in ['Move', 'Torus', 'Rope', 'Writer', "Pinch", "Rollingpin", "Chopsticks", "Table", 'TripleMove', 'Assembly']: 7 | for id in range(5): 8 | register( 9 | id=f'{env_name}-v{id + 1}', 10 | entry_point=f"plb.envs.env:PlasticineEnv", 11 | kwargs={'cfg_path': f"{env_name.lower()}.yml", "version": id + 1}, 12 | max_episode_steps=50 13 | ) 14 | 15 | register(id='GatherMove-v1', 16 | entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 17 | kwargs={'cfg_path': "env_ymls/gather_move.yml", "version": 1}, 18 | max_episode_steps=50) 19 | 20 | register(id='LiftSpread-v1', 21 | entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 22 | kwargs={'cfg_path': "env_ymls/lift_spread.yml", "version": 1}, 23 | max_episode_steps=50) 24 | 25 | register(id='CutRearrange-v1', 26 | entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 27 | kwargs={'cfg_path': "env_ymls/cut_rearrange.yml", "version": 1}, 28 | max_episode_steps=50) 29 | 30 | register(id='CutRearrange-v2', 31 | entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 32 | kwargs={'cfg_path': "env_ymls/cut_rearrange_0528.yml", "version": 1}, 33 | max_episode_steps=50) 34 | 35 | register(id='CutRearrangeSpread-v1', 36 | entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 37 | kwargs={'cfg_path': "env_ymls/cut_rearrange_spread.yml", "version": 1}, 38 | max_episode_steps=100) 39 | 40 | # Single-stage task 41 | register(id='Lift-v1', entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 42 | kwargs={'cfg_path': "env_ymls/single_stage/lift.yml", "version": 1}, 43 | max_episode_steps=50) 44 | register(id='Spread-v1', entry_point=f"plb.envs.multitask_env:MultitaskPlasticineEnv", 45 | kwargs={'cfg_path': "env_ymls/single_stage/spread.yml", "version": 1}, 46 | max_episode_steps=50) 47 | 48 | 49 | def make(env_name, nn=False, return_dist=False, **kwargs): 50 | env: PlasticineEnv = gym.make(env_name, nn=nn, return_dist=return_dist, **kwargs) 51 | return env 52 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/assembly.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | yield_stress: 100. 3 | ground_friction: 100. 4 | SHAPES: 5 | - shape: box 6 | width: (0.16, 0.16, 0.16) 7 | init_pos: (0.6, 0.08, 0.5) 8 | n_particles: 6000 9 | color: ((150<<8)+ (150<<16)) 10 | - shape: sphere 11 | radius: 0.06 12 | init_pos: (0.3, 0.06, 0.5) 13 | n_particles: 4000 14 | color: ((0<<8)+ (150<<16) +150) 15 | PRIMITIVES: 16 | - shape: Sphere 17 | radius: 0.04 18 | init_pos: (0.38, 0.06, 0.5) 19 | color: (0.8, 0.8, 0.8) 20 | friction: 0.9 21 | action: 22 | dim: 3 23 | scale: (0.009, 0.009, 0.009) 24 | - shape: Sphere 25 | radius: 0.04 26 | init_pos: (0.22, 0.06, 0.5) 27 | color: (0.8, 0.8, 0.8) 28 | friction: 0.9 29 | action: 30 | dim: 3 31 | scale: (0.009, 0.009, 0.009) 32 | ENV: 33 | loss: 34 | target_path: envs/assets/Assembly3D-v1.npy 35 | RENDERER: 36 | camera_pos: (0.5, 0.5, 2.5) 37 | camera_rot: (0.1, 0.) 38 | VARIANTS: 39 | - ENV: 40 | loss: 41 | target_path: envs/assets/Assembly3D-v1.npy 42 | - ENV: 43 | loss: 44 | target_path: envs/assets/Assembly3D-v2.npy 45 | - ENV: 46 | loss: 47 | target_path: envs/assets/Assembly3D-v3.npy 48 | - ENV: 49 | loss: 50 | target_path: envs/assets/Assembly3D-v4.npy 51 | - ENV: 52 | loss: 53 | target_path: envs/assets/Assembly3D-v5.npy 54 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/chopsticks.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | n_particles: 10000 3 | yield_stress: 200. 4 | ground_friction: 0. 5 | gravity: (0, -5, 0) 6 | SHAPES: 7 | - shape: box 8 | width: (0.04, 0.04, 0.6) 9 | init_pos: (0.5, 0.02, 0.5) 10 | color: 100 11 | PRIMITIVES: 12 | - shape: Chopsticks 13 | h: 0.2 14 | r: 0.02 15 | init_pos: (0.5, 0.15, 0.5) 16 | init_rot: (1., 0., 0., 0.) 17 | init_gap: 0.06 18 | color: (0.8, 0.8, 0.8) 19 | friction: 10. 20 | action: 21 | dim: 7 22 | scale: (0.02, 0.02, 0.02, 0.04, 0.04, 0.04, 0.02) 23 | ENV: 24 | loss: 25 | target_path: envs/assets/Chopsticks3D-v1.npy 26 | RENDERER: 27 | use_directional_light: True 28 | VARIANTS: 29 | - ENV: 30 | loss: 31 | target_path: envs/assets/Chopsticks3D-v1.npy 32 | - ENV: 33 | loss: 34 | target_path: envs/assets/Chopsticks3D-v2.npy 35 | - ENV: 36 | loss: 37 | target_path: envs/assets/Chopsticks3D-v3.npy 38 | - ENV: 39 | loss: 40 | target_path: envs/assets/Chopsticks3D-v4.npy 41 | - ENV: 42 | loss: 43 | target_path: envs/assets/Chopsticks3D-v5.npy 44 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/cut_rearrange.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 30000 4 | quality_multiplier: 1.25 5 | yield_stress: 150. 6 | ground_friction: 0.5 7 | gravity: (0, -10, 0) 8 | quality: 1 9 | dtype: float32 10 | lower_bound: 1. 11 | SHAPES: 12 | - shape: box 13 | init_pos: (0.5, 0.12, 0.5) 14 | width: (0.2, 0.08, 0.08) # Computed based on the volume of the sphere 0.00143675 15 | color: 100 16 | n_particles: 5000 17 | PRIMITIVES: 18 | - shape: Knife 19 | h: (0.15, 0.15) 20 | size: (0.025, 0.2, 0.06) 21 | prot: (1.0, 0.0, 0.0, 0.58) 22 | init_pos: (0.5, 0.3, 0.5) 23 | color: (0.7568, 0.6039, 0.4196) 24 | friction: 0. 25 | action: 26 | dim: 3 27 | scale: (0.015, 0.015, 0.0) 28 | - shape: Gripper 29 | size: (0.015, 0.1, 0.06) 30 | init_pos: (0.5, 0.10, 0.5) 31 | init_gap: 0.18 32 | minimal_gap: 0.08 33 | maximal_gap: 0.2 34 | #h: 0.10 35 | init_rot: (0.707, 0.0, 0.707, 0.0) 36 | color: (0.7568, 0.6039, 0.4196) 37 | friction: 10. 38 | action: 39 | dim: 7 40 | scale: (0.015, 0.015, 0.015, 0., 0., 0., 0.015) # previous is 0.01 + 0.015 41 | 42 | ENV: 43 | cached_state_path: datasets/CutRearrange-v1/1215_cutrearrange 44 | env_name: CutRearrange-v1 45 | 46 | RENDERER: 47 | name: tina 48 | image_res: (512, 512) 49 | 50 | cam_center: (0.52, 0.3, 0.6) 51 | cam_theta: 0. 52 | cam_phi: -1.3 53 | cam_radius: 0.53 54 | 55 | #cam_center: (0.5, 0.3, 0.5) 56 | #cam_theta: 1.57 57 | #cam_phi: 0. 58 | #cam_radius: 0.75 59 | 60 | use_directional_light: True 61 | light_direction: (0., 1., 1.) -------------------------------------------------------------------------------- /plb/envs/env_ymls/cut_rearrange_spread.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 30000 4 | quality_multiplier: 1.25 5 | yield_stress: 150. 6 | ground_friction: 0.5 7 | gravity: (0, -20, 0) 8 | quality: 1 9 | dtype: float32 10 | lower_bound: 1. 11 | SHAPES: 12 | - shape: box 13 | init_pos: (0.5, 0.12, 0.5) 14 | width: (0.2, 0.08, 0.08) # Computed based on the volume of the sphere 0.00143675 15 | color: 100 16 | n_particles: 5000 17 | PRIMITIVES: 18 | - shape: Knife 19 | h: (0.15, 0.15) 20 | size: (0.025, 0.2, 0.06) 21 | prot: (1.0, 0.0, 0.0, 0.58) 22 | init_pos: (0.8, 0.3, 0.1) 23 | color: (0.7568, 0.6039, 0.4196) 24 | friction: 0. 25 | action: 26 | dim: 3 27 | scale: (0.015, 0.015, 0.0) 28 | - shape: Box 29 | size: (0.015, 0.05, 0.1) 30 | init_pos: (0.5, 0.25, 0.1) 31 | init_rot: (1., 0.0, 0.0, 0.0) 32 | color: (0.7568, 0.6039, 0.4196) 33 | friction: 10. 34 | action: 35 | dim: 6 36 | scale: (0.010, 0.010, 0.010, 0., 0.025, 0.) 37 | - shape: RollingPinExt 38 | h: 0.15 39 | r: 0.03 40 | init_pos: (0.25, 0.25, 0.1) # was 0.2 on the second dimension 41 | init_rot: (0.707, 0.707, 0., 0.) 42 | color: (0.7568, 0.6039, 0.4196) 43 | friction: 0.9 44 | action: 45 | dim: 6 46 | scale: (0.7, 0.005, 0.005, 0.005, 0., 0.) 47 | lower_bound: (0., 0.03, 0.) 48 | 49 | ENV: 50 | cached_state_path: datasets/CutRearrangeSpread-v1/0609_cutrearrangespread 51 | env_name: CutRearrangeSpread-v1 52 | 53 | RENDERER: 54 | name: tina 55 | image_res: (512, 512) 56 | 57 | cam_center: (0.5, 0.3, 0.6) 58 | cam_theta: 0. 59 | cam_phi: -1.0 60 | cam_radius: 0.7 61 | 62 | use_directional_light: True 63 | light_direction: (0., 1., 1.) -------------------------------------------------------------------------------- /plb/envs/env_ymls/gather_move.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 30000 4 | yield_stress: 200. 5 | ground_friction: 1.5 6 | gravity: (0, -20, 0) 7 | quality: 1 8 | SHAPES: 9 | - shape: scatter 10 | pos_min: (0.64, 0.02, 0.38) 11 | pos_max: (0.76, 0.035, 0.62) 12 | color: (0.7568, 0.6039, 0.4196) 13 | seed: 0 14 | PRIMITIVES: 15 | - shape: Gripper 16 | size: (0.015, 0.09, 0.05) # To make rendering realistic. Actual floor is higher than the rendered 17 | init_pos: (0.7, 0.06, 0.5) 18 | init_rot: (0.5, 0.5, -0.5, 0.5) 19 | init_gap: 0.4 20 | minimal_gap: 0.05 21 | color: (0.7568, 0.6039, 0.4196) 22 | friction: 1. 23 | action: 24 | dim: 7 25 | scale: (0.015, 0.0, 0.015, 0.0, 0.0, 0.1, 0.03) 26 | collision_group: (0, 0, 1) 27 | - shape: Box # Lifter 28 | size: (0.07, 0.07, 0.02) # To make rendering realistic. Actual floor is higher than the rendered 29 | init_pos: (0.7, 0.01, 0.5) # Was 0.7 on dim 3 30 | init_rot: (0.707, 0.707, 0., 0.) 31 | color: (0.7568, 0.6039, 0.4196) 32 | friction: 50. 33 | action: 34 | dim: 6 35 | scale: (0.01, 0.01, 0., 0.0, 0., 0.05) 36 | collision_group: (0, 0, 1) 37 | - shape: Box # Obstacle 38 | size: (0.2, 0.28, 0.04) # To make rendering realistic. Actual floor is higher than the rendered 39 | init_pos: (0.33, 0.05, 0.5) 40 | init_rot: (0.707, 0.707, 0., 0.) 41 | color: (0.5, 0.5, 0.5) 42 | friction: 5. # Make this larger later 43 | action: 44 | dim: 0 45 | ENV: 46 | # cached_state_path: datasets/0307_gathermove 47 | cached_state_path: datasets/0202_gathermove 48 | 49 | RENDERER: 50 | name: tina 51 | cam_center: (0.55, 0., 0.5) 52 | cam_theta: 0. 53 | cam_phi: -0.85 54 | cam_radius: 0.75 55 | tina_img_res: 512 56 | sdf_threshold: 0.4 57 | use_directional_light: True 58 | light_direction: (0., 1., 1.) -------------------------------------------------------------------------------- /plb/envs/env_ymls/lift_spread.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 30000 4 | yield_stress: 200. 5 | ground_friction: 1.5 6 | gravity: (0, -20, 0) 7 | quality: 1 8 | SHAPES: 9 | - shape: sphere 10 | init_pos: (0.65, 0.08, 0.5) 11 | radius: (0.05) 12 | color: 100 13 | # - shape: sphere 14 | # init_pos: (0.5, 0.23576, 0.5) 15 | # width: (0.2, 0.03575, 0.2) # Computed based on the volume of the sphere 0.00143675 16 | # color: 100 17 | 18 | PRIMITIVES: 19 | - shape: RollingPinExt 20 | h: 0.3 21 | r: 0.03 22 | init_pos: (0.3, 0.25, 0.5) # was 0.2 on the second dimension 23 | init_rot: (0.707, 0.707, 0., 0.) 24 | color: (0.7568, 0.6039, 0.4196) 25 | friction: 0.9 26 | action: 27 | dim: 6 28 | scale: (0.7, 0.005, 0.005, 0.005, 0., 0.) 29 | lower_bound: (0., 0.16, 0.) 30 | - shape: Box 31 | size: (0.1, 0.1, 0.02) # To make rendering realistilc. Actual floor is higher than the rendered 32 | init_pos: (0.65, 0.02, 0.5) 33 | init_rot: (0.707, 0.707, 0., 0.) 34 | color: (0.7568, 0.6039, 0.4196) 35 | friction: 50. 36 | action: 37 | dim: 6 38 | scale: (0.01, 0.01, 0., 0.0, 0., 0.05) 39 | collision_group: (0, 0, 1) 40 | # lower_bound: (.6, 0., 0.) # 0.15 is the height of the obstacle 41 | - shape: Box # Obstacle 42 | size: (0.2, 0.28, 0.07) # To make rendering realistic. Actual floor is higher than the rendered 43 | init_pos: (0.3, 0.05, 0.5) 44 | init_rot: (0.707, 0.707, 0., 0.) 45 | color: (0.5, 0.5, 0.5) 46 | friction: 5. # Make this larger later 47 | action: 48 | dim: 0 49 | ENV: 50 | cached_state_path: datasets/LiftSpread-v1/0202_liftspread 51 | 52 | RENDERER: 53 | name: tina 54 | tina_img_res: 512 55 | cam_center: (0.45, 0.07, 0.5) 56 | cam_theta: 0. 57 | cam_phi: -0.6 58 | cam_radius: 0.8 59 | use_directional_light: True 60 | light_direction: (0., 1., 1.) 61 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/move.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 10000 4 | yield_stress: 200. 5 | SHAPES: 6 | - shape: sphere 7 | radius: 0.2049069760770578/2 8 | init_pos: (0.6757143040494873, 0.5619162002773135, 0.7515980438048129) 9 | color: (127<<16) 10 | PRIMITIVES: 11 | - shape: Sphere 12 | radius: 0.03 13 | init_pos: (0.5757143040494873, 0.5619162002773135, 0.7515980438048129) 14 | color: (0.7, 0.7, 0.7) 15 | friction: 0.9 16 | action: 17 | dim: 3 18 | scale: (0.01, 0.01, 0.01) 19 | - shape: Sphere 20 | radius: 0.03 21 | init_pos: (0.7757143040494873, 0.5619162002773135, 0.7515980438048129) 22 | color: (0.7, 0.7, 0.7) 23 | friction: 0.9 24 | action: 25 | dim: 3 26 | scale: (0.01, 0.01, 0.01) 27 | VARIANTS: 28 | - SHAPES: 29 | - init_pos: (0.6757143040494873, 0.5619162002773135, 0.7515980438048129) 30 | radius: 0.2049069760770578/2 31 | PRIMITIVES: 32 | - init_pos: (0.5757143040494873, 0.5619162002773135, 0.7515980438048129) 33 | - init_pos: (0.7757143040494873, 0.5619162002773135, 0.7515980438048129) 34 | ENV: 35 | loss: 36 | target_path: envs/assets/Move3D-v1.npy 37 | - SHAPES: 38 | - init_pos: (0.4800617702933018, 0.6114161266624294, 0.2150469121879661) 39 | radius: 0.22128338675873624/2 40 | PRIMITIVES: 41 | - init_pos: (0.3800617702933018, 0.6114161266624294, 0.2150469121879661) 42 | - init_pos: (0.5800617702933018, 0.6114161266624294, 0.2150469121879661) 43 | ENV: 44 | loss: 45 | target_path: envs/assets/Move3D-v2.npy 46 | - SHAPES: 47 | - init_pos: (0.5953388885096601, 0.7803511669469463, 0.3652372561756634) 48 | radius: 0.21518886629207218/2 49 | PRIMITIVES: 50 | - init_pos: (0.4953388885096601, 0.7803511669469463, 0.3652372561756634) 51 | - init_pos: (0.6953388885096601, 0.7803511669469463, 0.3652372561756634) 52 | ENV: 53 | loss: 54 | target_path: envs/assets/Move3D-v3.npy 55 | - SHAPES: 56 | - init_pos: (0.5608152006865512, 0.5151402950552514, 0.4707541125135959) 57 | radius: 0.23144406058863135/2 58 | PRIMITIVES: 59 | - init_pos: (0.4608152006865512, 0.5151402950552514, 0.4707541125135959) 60 | - init_pos: (0.6608152006865512, 0.5151402950552514, 0.4707541125135959) 61 | ENV: 62 | loss: 63 | target_path: envs/assets/Move3D-v4.npy 64 | - SHAPES: 65 | - init_pos: (0.2958401778083163, 0.5385429137124296, 0.7461548784761765) 66 | radius: 0.23726089169300607/2 67 | PRIMITIVES: 68 | - init_pos: (0.1958401778083163, 0.5385429137124296, 0.7461548784761765) 69 | - init_pos: (0.3958401778083163, 0.5385429137124296, 0.7461548784761765) 70 | ENV: 71 | loss: 72 | target_path: envs/assets/Move3D-v5.npy 73 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/pinch.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | yield_stress: 50. 3 | ground_friction: 100. 4 | 5 | SHAPES: 6 | - shape: box 7 | width: (0.2, 0.2, 0.2) 8 | init_pos: (0.5, 0.1, 0.5) 9 | n_particles: 6000 10 | color: ((150<<8)+ (150<<16)) 11 | PRIMITIVES: 12 | - shape: Sphere 13 | radius: 0.04 14 | init_pos: (0.5, 0.35, 0.5) 15 | color: (0.8, 0.8, 0.8) 16 | friction: 0.9 17 | lower_bound: (0.1, 0.1, 0.1) 18 | upper_bound: (0.9, 0.9, 0.9) 19 | action: 20 | dim: 3 21 | scale: (0.02, 0.02, 0.02) 22 | ENV: 23 | loss: 24 | target_path: envs/assets/Pinch3D-v1.npy 25 | 26 | RENDERER: 27 | camera_pos: (0.5, 2.5, 2.2) 28 | camera_rot: (0.8, 0.) 29 | VARIANTS: 30 | - ENV: 31 | loss: 32 | target_path: envs/assets/Pinch3D-v1.npy 33 | - ENV: 34 | loss: 35 | target_path: envs/assets/Pinch3D-v2.npy 36 | - ENV: 37 | loss: 38 | target_path: envs/assets/Pinch3D-v3.npy 39 | - ENV: 40 | loss: 41 | target_path: envs/assets/Pinch3D-v4.npy 42 | - ENV: 43 | loss: 44 | target_path: envs/assets/Pinch3D-v5.npy 45 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/rollingpin.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 10000 4 | yield_stress: 50. 5 | ground_friction: 1.5 6 | SHAPES: 7 | - shape: box 8 | width: (0.3, 0.1, 0.3) 9 | init_pos: (0.5, 0.05, 0.5) 10 | color: 100 11 | PRIMITIVES: 12 | - shape: RollingPin 13 | h: 0.3 14 | r: 0.03 15 | init_pos: (0.5, 0.123, 0.5) 16 | init_rot: (0.707, 0.707, 0., 0.) 17 | color: (0.8, 0.8, 0.8) 18 | friction: 0.9 19 | action: 20 | dim: 3 21 | scale: (0.6666666666666667, 0.06666666666666668, 0.001) 22 | ENV: 23 | loss: 24 | target_path: envs/assets/Rollingpin3D-v1.npy 25 | RENDERER: 26 | camera_pos: (0.5, 2.5, 2.) 27 | camera_rot: (1.0, 0.) 28 | use_directional_light: True 29 | light_direction: (0., 1., 0.) 30 | VARIANTS: 31 | - ENV: 32 | loss: 33 | target_path: envs/assets/Rollingpin3D-v1.npy 34 | - ENV: 35 | loss: 36 | target_path: envs/assets/Rollingpin3D-v2.npy 37 | - ENV: 38 | loss: 39 | target_path: envs/assets/Rollingpin3D-v3.npy 40 | - ENV: 41 | loss: 42 | target_path: envs/assets/Rollingpin3D-v4.npy 43 | - ENV: 44 | loss: 45 | target_path: envs/assets/Rollingpin3D-v5.npy 46 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/rope.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | yield_stress: 50. 3 | ground_friction: 0.3 4 | SHAPES: 5 | - shape: box 6 | width: (0.6, 0.06, 0.06) 7 | init_pos: (0.5, 0.03, 0.73) 8 | color: (((0 << 8) + 150) << 8) 9 | PRIMITIVES: 10 | - shape: Sphere 11 | radius: 0.03 12 | init_pos: (0.22, 0.015, 0.82) 13 | color: (0.8, 0.8, 0.8) 14 | friction: 0.9 15 | action: 16 | dim: 3 17 | scale: (0.01, 0.01, 0.01) 18 | - shape: Sphere 19 | radius: 0.03 20 | init_pos: (0.78, 0.015, 0.82) 21 | color: (0.8, 0.8, 0.8) 22 | friction: 0.9 23 | action: 24 | dim: 3 25 | scale: (0.01, 0.01, 0.01) 26 | - shape: Cylinder 27 | h: 0.1 28 | r: 0.2 29 | init_pos: (0.3919300650726247, 0., 0.4990770359432596) 30 | color: (0.3, 0.3, 0.3) 31 | friction: 0.9 32 | ENV: 33 | loss: 34 | target_path: envs/data/Rope3D-v1.npy 35 | RENDERER: 36 | camera_pos: (0.5, 2.5, 2.) 37 | camera_rot: (1.0, 0.) 38 | VARIANTS: 39 | - PRIMITIVES: 40 | - 41 | - 42 | - init_pos: (0.3919300650726247, 0, 0.4990770359432596) 43 | ENV: 44 | loss: 45 | target_path: envs/assets/Rope3D-v1.npy 46 | - PRIMITIVES: 47 | - 48 | - 49 | - init_pos: (0.4827737598605798, 0, 0.572508568647028) 50 | ENV: 51 | loss: 52 | target_path: envs/assets/Rope3D-v2.npy 53 | - PRIMITIVES: 54 | - 55 | - 56 | - init_pos: (0.48953026610561057, 0, 0.5199459480962076) 57 | ENV: 58 | loss: 59 | target_path: envs/assets/Rope3D-v3.npy 60 | - PRIMITIVES: 61 | - 62 | - 63 | - init_pos: (0.46968068720064815, 0, 0.3868456769743354) 64 | ENV: 65 | loss: 66 | target_path: envs/assets/Rope3D-v4.npy 67 | - PRIMITIVES: 68 | - 69 | - 70 | - init_pos: (0.49333308965447087, 0, 0.5946055392248519) 71 | ENV: 72 | loss: 73 | target_path: envs/assets/Rope3D-v5.npy 74 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/single_stage/lift.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 10000 4 | yield_stress: 200. 5 | ground_friction: 1.5 6 | gravity: (0, -20, 0) 7 | quality: 1 8 | SHAPES: 9 | - shape: sphere 10 | init_pos: (0.65, 0.08, 0.5) 11 | radius: (0.05) 12 | color: 100 13 | # - shape: sphere 14 | # init_pos: (0.5, 0.23576, 0.5) 15 | # width: (0.2, 0.03575, 0.2) # Computed based on the volume of the sphere 0.00143675 16 | # color: 100 17 | 18 | PRIMITIVES: 19 | - shape: RollingPinExt 20 | h: 0.3 21 | r: 0.03 22 | init_pos: (0.3, 0.25, 0.5) # was 0.2 on the second dimension 23 | init_rot: (0.707, 0.707, 0., 0.) 24 | color: (0.7568, 0.6039, 0.4196) 25 | friction: 0.9 26 | action: 27 | dim: 6 28 | scale: (0.7, 0.005, 0.005, 0.005, 0., 0.) 29 | lower_bound: (0., 0.16, 0.) 30 | - shape: Box 31 | size: (0.1, 0.1, 0.02) # To make rendering realistilc. Actual floor is higher than the rendered 32 | init_pos: (0.65, 0.02, 0.5) 33 | init_rot: (0.707, 0.707, 0., 0.) 34 | color: (0.7568, 0.6039, 0.4196) 35 | friction: 50. 36 | action: 37 | dim: 6 38 | scale: (0.01, 0.01, 0., 0.0, 0., 0.05) 39 | collision_group: (0, 0, 1) 40 | # lower_bound: (.6, 0., 0.) # 0.15 is the height of the obstacle 41 | - shape: Box # Obstacle 42 | size: (0.2, 0.28, 0.07) # To make rendering realistic. Actual floor is higher than the rendered 43 | init_pos: (0.3, 0.05, 0.5) 44 | init_rot: (0.707, 0.707, 0., 0.) 45 | color: (0.5, 0.5, 0.5) 46 | friction: 5. # Make this larger later 47 | action: 48 | dim: 0 49 | ENV: 50 | cached_state_path: datasets/1116_Lift 51 | 52 | RENDERER: 53 | name: tina 54 | tina_img_res: 512 55 | cam_center: (0.55, 0.07, 0.5) 56 | cam_theta: 0. 57 | cam_phi: -0.5 58 | cam_radius: 0.75 59 | use_directional_light: True 60 | light_direction: (0., 1., 1.) 61 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/single_stage/spread.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 10000 4 | yield_stress: 200. 5 | ground_friction: 1.5 6 | gravity: (0, -20, 0) 7 | quality: 1 8 | SHAPES: 9 | - shape: sphere 10 | init_pos: (0.65, 0.08, 0.5) 11 | radius: (0.05) 12 | color: 100 13 | # - shape: sphere 14 | # init_pos: (0.5, 0.23576, 0.5) 15 | # width: (0.2, 0.03575, 0.2) # Computed based on the volume of the sphere 0.00143675 16 | # color: 100 17 | 18 | PRIMITIVES: 19 | - shape: RollingPinExt 20 | h: 0.3 21 | r: 0.03 22 | init_pos: (0.3, 0.25, 0.5) # was 0.2 on the second dimension 23 | init_rot: (0.707, 0.707, 0., 0.) 24 | color: (0.7568, 0.6039, 0.4196) 25 | friction: 0.9 26 | action: 27 | dim: 6 28 | scale: (0.7, 0.005, 0.005, 0.005, 0., 0.) 29 | lower_bound: (0., 0.16, 0.) 30 | - shape: Box 31 | size: (0.1, 0.1, 0.02) # To make rendering realistilc. Actual floor is higher than the rendered 32 | init_pos: (0.65, 0.02, 0.5) 33 | init_rot: (0.707, 0.707, 0., 0.) 34 | color: (0.7568, 0.6039, 0.4196) 35 | friction: 50. 36 | action: 37 | dim: 6 38 | scale: (0.01, 0.01, 0., 0.0, 0., 0.05) 39 | collision_group: (0, 0, 1) 40 | # lower_bound: (.6, 0., 0.) # 0.15 is the height of the obstacle 41 | - shape: Box # Obstacle 42 | size: (0.2, 0.28, 0.07) # To make rendering realistic. Actual floor is higher than the rendered 43 | init_pos: (0.3, 0.05, 0.5) 44 | init_rot: (0.707, 0.707, 0., 0.) 45 | color: (0.5, 0.5, 0.5) 46 | friction: 5. # Make this larger later 47 | action: 48 | dim: 0 49 | ENV: 50 | cached_state_path: datasets/1118_Spread_correct 51 | 52 | RENDERER: 53 | name: tina 54 | tina_img_res: 512 55 | cam_center: (0.55, 0.07, 0.5) 56 | cam_theta: 0. 57 | cam_phi: -0.5 58 | cam_radius: 0.75 59 | use_directional_light: True 60 | light_direction: (0., 1., 1.) 61 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/spread.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 10000 4 | yield_stress: 50. 5 | ground_friction: 1.5 6 | gravity: (0, -10, 0) 7 | quality: 1 8 | SHAPES: 9 | - shape: sphere 10 | radius: 0.15 11 | init_pos: (0.5, 0.15, 0.5) 12 | color: 100 13 | PRIMITIVES: 14 | - shape: RollingPin 15 | h: 0.3 16 | r: 0.03 17 | init_pos: (0.1, 0.18, 0.5) 18 | init_rot: (0.707, 0.707, 0., 0.) 19 | color: (0.8, 0.8, 0.8) 20 | friction: 0.9 21 | action: 22 | dim: 3 23 | scale: (1., 0.3333333333333, 0.005) 24 | - shape: Box 25 | size: (0.015, 0.1, 0.1) 26 | init_pos: (0.8, 0.18, 0.5) 27 | init_rot: (0.9305, 0, 0.366, 0.) 28 | color: (0.8, 0.8, 0.8) 29 | friction: 0.9 30 | action: 31 | dim: 6 32 | scale: (0.05,0.05,0.05,0.3,0.3,0.3) 33 | ENV: 34 | loss: 35 | target_path: envs/assets/Rollingpin3D-v1.npy 36 | RENDERER: 37 | camera_pos: (0.5, 2., 2.) 38 | camera_rot: (0.9, 0.) 39 | use_directional_light: True 40 | light_direction: (0., 1., 1.) 41 | VARIANTS: 42 | - ENV: 43 | loss: 44 | target_path: envs/assets/Rollingpin3D-v1.npy 45 | - ENV: 46 | loss: 47 | target_path: envs/assets/Rollingpin3D-v2.npy 48 | - ENV: 49 | loss: 50 | target_path: envs/assets/Rollingpin3D-v3.npy 51 | - ENV: 52 | loss: 53 | target_path: envs/assets/Rollingpin3D-v4.npy 54 | - ENV: 55 | loss: 56 | target_path: envs/assets/Rollingpin3D-v5.npy 57 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/table.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | yield_stress: 50. 3 | nu: 0.05 4 | ground_friction: 0.3 5 | SHAPES: 6 | - shape: box 7 | width: (0.04, 0.1, 0.04) 8 | init_pos: (0.5-0.075, 0.1, 0.5-0.075) 9 | n_particles: 2000 10 | - shape: box 11 | width: (0.04, 0.1, 0.04) 12 | init_pos: (0.5-0.075, 0.1, 0.5+0.075) 13 | n_particles: 2000 14 | - shape: box 15 | width: (0.04, 0.1, 0.04) 16 | init_pos: (0.5+0.075, 0.1, 0.5-0.075) 17 | n_particles: 2000 18 | - shape: box 19 | width: (0.04, 0.1, 0.04) 20 | init_pos: (0.5+0.075, 0.1, 0.5+0.075) 21 | n_particles: 2000 22 | - shape: box 23 | width: (0.2, 0.05, 0.2) 24 | init_pos: (0.5, 0.18, 0.5) 25 | color: (((200 << 8) + 200) << 8) 26 | n_particles: 2000 27 | PRIMITIVES: 28 | - shape: Sphere 29 | radius: 0.04 30 | init_pos: (0.5, 0.06, 0.5) 31 | friction: 0.9 32 | color: (0.8, 0.8, 0.8) 33 | action: 34 | dim: 3 35 | scale: (0.03, 0.0, 0.03) 36 | ENV: 37 | loss: 38 | target_path: envs/assets/Table3D-v1.npy 39 | RENDERER: 40 | camera_pos: (0.5, 0.3, 1.5) 41 | camera_rot: (0.1, 0.) 42 | use_directional_light: True 43 | 44 | VARIANTS: 45 | - ENV: 46 | loss: 47 | target_path: envs/assets/Table3D-v1.npy 48 | - ENV: 49 | loss: 50 | target_path: envs/assets/Table3D-v2.npy 51 | - ENV: 52 | loss: 53 | target_path: envs/assets/Table3D-v3.npy 54 | - ENV: 55 | loss: 56 | target_path: envs/assets/Table3D-v4.npy 57 | - ENV: 58 | loss: 59 | target_path: envs/assets/Table3D-v5.npy 60 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/torus.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | yield_stress: 50. 3 | ground_friction: 100. 4 | 5 | SHAPES: 6 | - shape: box 7 | width: (0.3, 0.1, 0.3) 8 | init_pos: (0.5, 0.05, 0.5) 9 | color: (((200 << 8) + 200) << 8) 10 | PRIMITIVES: 11 | - shape: Torus 12 | tx: 0.05 13 | ty: 0.03 14 | init_pos: (0.5, 0.2, 0.5) 15 | init_rot: (0., 0., 0., 1.) 16 | friction: 0.9 17 | color: (0.8, 0.8, 0.8) 18 | lower_bound: (0., 0.05, 0.) 19 | action: 20 | dim: 3 21 | scale: (0.004, 0.004, 0.004) 22 | ENV: 23 | loss: 24 | target_path: envs/assets/Torus3D-v1.npy 25 | 26 | RENDERER: 27 | camera_pos: (0.5, 2.5, 2.2) 28 | camera_rot: (1., 0.) 29 | use_directional_light: True 30 | 31 | VARIANTS: 32 | - ENV: 33 | loss: 34 | target_path: envs/assets/Torus3D-v1.npy 35 | - ENV: 36 | loss: 37 | target_path: envs/assets/Torus3D-v2.npy 38 | - ENV: 39 | loss: 40 | target_path: envs/assets/Torus3D-v3.npy 41 | - ENV: 42 | loss: 43 | target_path: envs/assets/Torus3D-v4.npy 44 | - ENV: 45 | loss: 46 | target_path: envs/assets/Torus3D-v5.npy 47 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/triplemove.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | yield_stress: 200. 3 | SHAPES: 4 | - shape: box 5 | width: (0.1, 0.1, 0.1) 6 | init_pos: (0.3, 0.05, 0.5) 7 | n_particles: 3333 8 | - shape: box 9 | width: (0.1, 0.1, 0.1) 10 | init_pos: (0.5, 0.05, 0.5) 11 | n_particles: 3333 12 | - shape: box 13 | width: (0.1, 0.1, 0.1) 14 | init_pos: (0.7, 0.05, 0.5) 15 | n_particles: 3333 16 | PRIMITIVES: 17 | - shape: Sphere 18 | radius: 0.03 19 | init_pos: (0.23, 0.05, 0.5) 20 | color: (0.8, 0.8, 0.8) 21 | friction: 0.9 22 | action: 23 | dim: 3 24 | scale: (0.01, 0.01, 0.01) 25 | - shape: Sphere 26 | radius: 0.03 27 | init_pos: (0.37, 0.05, 0.5) 28 | color: (0.8, 0.8, 0.8) 29 | friction: 0.9 30 | action: 31 | dim: 3 32 | scale: (0.01, 0.01, 0.01) 33 | - shape: Sphere 34 | radius: 0.03 35 | init_pos: (0.43, 0.05, 0.5) 36 | color: (0.8, 0.8, 0.8) 37 | friction: 0.9 38 | action: 39 | dim: 3 40 | scale: (0.01, 0.01, 0.01) 41 | - shape: Sphere 42 | radius: 0.03 43 | init_pos: (0.57, 0.05, 0.5) 44 | color: (0.8, 0.8, 0.8) 45 | friction: 0.9 46 | action: 47 | dim: 3 48 | scale: (0.01, 0.01, 0.01) 49 | - shape: Sphere 50 | radius: 0.03 51 | init_pos: (0.63, 0.05, 0.5) 52 | color: (0.8, 0.8, 0.8) 53 | friction: 0.9 54 | action: 55 | dim: 3 56 | scale: (0.01, 0.01, 0.01) 57 | - shape: Sphere 58 | radius: 0.03 59 | init_pos: (0.77, 0.05, 0.5) 60 | color: (0.8, 0.8, 0.8) 61 | friction: 0.9 62 | action: 63 | dim: 3 64 | scale: (0.01, 0.01, 0.01) 65 | ENV: 66 | loss: 67 | target_path: envs/assets/TripleMove3D-v1.npy 68 | RENDERER: 69 | camera_pos: (0.5, 2.5, 2.) 70 | camera_rot: (1.0, 0.) 71 | use_directional_light: True 72 | light_direction: (0.1, 1, 0.3) 73 | 74 | VARIANTS: 75 | - ENV: 76 | loss: 77 | target_path: envs/assets/TripleMove3D-v1.npy 78 | - ENV: 79 | loss: 80 | target_path: envs/assets/TripleMove3D-v2.npy 81 | - ENV: 82 | loss: 83 | target_path: envs/assets/TripleMove3D-v3.npy 84 | - ENV: 85 | loss: 86 | target_path: envs/assets/TripleMove3D-v4.npy 87 | - ENV: 88 | loss: 89 | target_path: envs/assets/TripleMove3D-v5.npy 90 | -------------------------------------------------------------------------------- /plb/envs/env_ymls/writer.yml: -------------------------------------------------------------------------------- 1 | SIMULATOR: 2 | E: 5000. 3 | n_particles: 10000 4 | yield_stress: 50. 5 | ground_friction: 100. 6 | SHAPES: 7 | - shape: box 8 | width: (0.3, 0.1, 0.3) 9 | init_pos: (0.5, 0.05, 0.5) 10 | color: ((((200 << 8) + 200) << 8) + 0) 11 | PRIMITIVES: 12 | - shape: Capsule 13 | h: 0.06 14 | r: 0.03 15 | init_pos: (0.5, 0.13, 0.5) 16 | init_rot: (0., 0., 0., 1.) 17 | lower_bound: (0., 0.05, 0.) 18 | color: (0.8, 0.8, 0.8) 19 | friction: 0. 20 | action: 21 | dim: 3 22 | scale: (0.01, 0.01, 0.01) 23 | ENV: 24 | loss: 25 | target_path: envs/assets/Writer3D-v1.npy 26 | RENDERER: 27 | camera_pos: (0.5, 2., 1.2) 28 | camera_rot: (1.3, 0.) 29 | use_directional_light: True 30 | 31 | VARIANTS: 32 | - ENV: 33 | loss: 34 | target_path: envs/assets/Writer3D-v1.npy 35 | - ENV: 36 | loss: 37 | target_path: envs/assets/Writer3D-v2.npy 38 | - ENV: 39 | loss: 40 | target_path: envs/assets/Writer3D-v3.npy 41 | - ENV: 42 | loss: 43 | target_path: envs/assets/Writer3D-v4.npy 44 | - ENV: 45 | loss: 46 | target_path: envs/assets/Writer3D-v5.npy 47 | -------------------------------------------------------------------------------- /plb/envs/utils.py: -------------------------------------------------------------------------------- 1 | from yacs.config import CfgNode 2 | 3 | def merge_dict(a, b): 4 | if b is None: 5 | return a 6 | import copy 7 | a = copy.deepcopy(a) 8 | for key in a: 9 | if key in b: 10 | if not isinstance(b[key], dict): 11 | a[key] = b[key] 12 | else: 13 | assert not isinstance(a[key], list) 14 | a[key] = merge_dict(a[key], b[key]) 15 | for key in b: 16 | if key not in a: 17 | raise ValueError("Key is not in dict A!") 18 | return a 19 | 20 | 21 | def merge_lists(a, b): 22 | outs = [] 23 | assert isinstance(a, list) and isinstance(b, list) 24 | for i in range(len(a)): 25 | assert isinstance(a[i], dict) 26 | x = a[i] 27 | if i < len(b): 28 | x = merge_dict(a[i], b[i]) 29 | outs.append(x) 30 | return outs 31 | 32 | -------------------------------------------------------------------------------- /prepare.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | conda activate plb 3 | export PYTHONPATH=$PWD:$PWD/taichi_three:$PWD/PointFlow:$PWD/setvae:$PYTHONPATH 4 | export PATH=$PWD/taichi_three:$PATH 5 | export LD_LIBRARY_PATH=$LD_LIBRARY_PATH -------------------------------------------------------------------------------- /run_scripts/diffskill_train_all.sh: -------------------------------------------------------------------------------- 1 | python core/diffskill/train_diffskill_rgb.py \ 2 | --env_name 'CutRearrange-v1' \ 3 | --input_mode 'rgbd' \ 4 | --dataset_path '' \ 5 | --train_modules 'vae' 'reward' 'policy' 'fea' \ 6 | --plan_step 3 \ 7 | --num_tools 2 \ 8 | --batch_size 256 \ 9 | --actor_latent_dim 64 \ 10 | --reward_latent_dim 1024 \ 11 | --fea_latent_dim 1024 \ 12 | --t_relative True \ 13 | --num_random_neg 2048 \ 14 | --num_buffer_neg 2048 \ 15 | --il_eval_freq 1 \ 16 | --il_num_epoch 50 \ 17 | --min_zlogl -4 \ 18 | --adam_iter 500 \ 19 | --adam_lr 5e-1 \ 20 | --adam_sample 1000 -------------------------------------------------------------------------------- /run_scripts/difskill_plan.sh: -------------------------------------------------------------------------------- 1 | python core/diffskill/train_diffskill_rgb.py \ 2 | --run_plan True \ 3 | --resume_path '' \ 4 | --adam_iter 10 \ 5 | --adam_sample 1280 \ 6 | --plan_bs 128 \ 7 | --plan_step 3 -------------------------------------------------------------------------------- /run_scripts/download.py: -------------------------------------------------------------------------------- 1 | # Use gdown to download file and unzip it 2 | import gdown 3 | import zipfile 4 | import os 5 | import shutil 6 | 7 | def download_and_unzip(url, output_dir): 8 | gdown.download(url, 'temp.zip', quiet=False) 9 | if not os.path.exists(output_dir): 10 | os.makedirs(output_dir) 11 | 12 | os.system('unzip temp.zip -d ' + output_dir) 13 | os.remove('temp.zip') 14 | 15 | def download_folder(file_id, output_dir): 16 | """Download folder from google drive""" 17 | os.system(f"gdown https://drive.google.com/drive/folders/{file_id} -O {output_dir} --folder") 18 | 19 | 20 | def download_pretrained_model(output_dir): 21 | """Download all pretrained models""" 22 | file_id = '18tmH0stc1z_TzfAHbQDu5HASNkaWFKk_' # (300M) 23 | download_folder(file_id, output_dir) 24 | 25 | def download_init_target(env_name, output_dir): 26 | """Download init and target for all environments""" 27 | urls = { 28 | 'LiftSpread-v1': 'https://drive.google.com/uc?id=1b4Qw6cbWbtEiP3MO7v6WOLkgCplFnDb1', 29 | 'CutRearrange-v1': 'https://drive.google.com/uc?id=1XFGKwngAX_4gVzmP38DMv0mo44a9e3Sd', 30 | 'CutRearrangeSpread-v1': 'https://drive.google.com/uc?id=1IEAev3VeCXAKVEeZmYlWjlT2MHLWyedl', 31 | } 32 | url = urls[env_name] 33 | download_and_unzip(url, output_dir) 34 | 35 | def download_demonstration(env_name, output_dir): 36 | """Download init and target for all environments""" 37 | urls = { 38 | 'LiftSpread-v1': 'https://drive.google.com/uc?id=1gRaaMEG6ytmjOOqG-K79kMteGhXW_QFB', # (6G) 39 | 'CutRearrange-v1': 'https://drive.google.com/uc?id=1wKgFLSkdueVGoy_vDGxPiOnVwENKsSZT', # (5G) 40 | 'CutRearrangeSpread-v1': 'https://drive.google.com/uc?id=1OFcC8po2ZBmgMxtJ5C7uiWPG7ESm5W2v', # (3G) 41 | } 42 | url = urls[env_name] 43 | download_and_unzip(url, output_dir) 44 | 45 | if __name__ == '__main__': 46 | import argparse 47 | 48 | parser = argparse.ArgumentParser() 49 | parser.add_argument('--command', type=str, choices=['init_target', 'demo', 'pretrained']) 50 | parser.add_argument( 51 | '--env_name', type=str, default='CutRearrangeSpread', help='The environment to download', 52 | choices=['LiftSpread-v1', 'CutRearrange-v1', 'CutRearrangeSpread-v1', 'all']) 53 | args = parser.parse_args() 54 | 55 | if args.env_name == 'all': 56 | envs = ['LiftSpread-v1', 'CutRearrange-v1', 'CutRearrangeSpread-v1'] 57 | else: 58 | envs = [args.env_name] 59 | 60 | if args.command == 'init_target': 61 | download_init_target(args.env_name, output_dir='./datasets/') 62 | elif args.command == 'demo': 63 | download_demonstration(args.env_name, output_dir='./data/released_datasets') 64 | elif args.command == 'pretrained': 65 | download_pretrained_model(output_dir='./data/') 66 | -------------------------------------------------------------------------------- /run_scripts/flat3d_train_all.sh: -------------------------------------------------------------------------------- 1 | python core/diffskill/train_flat3d.py \ 2 | --env_name 'CutRearrange-v1' \ 3 | --input_mode 'pc' \ 4 | --dataset_path '' \ 5 | --vae_resume_path '' \ 6 | --plan_step 3 \ 7 | --num_tools 2 \ 8 | --dimz 2 \ 9 | --batch_size 64 \ 10 | --actor_latent_dim 64 \ 11 | --reward_latent_dim 1024 \ 12 | --fea_latent_dim 1024 \ 13 | --t_relative True \ 14 | --num_random_neg 2048 \ 15 | --num_buffer_neg 2048 \ 16 | --fea_z_noise 0.01 \ 17 | --fea_t_noise 0.01 \ 18 | --il_eval_freq 1 \ 19 | --il_num_epoch 1 -------------------------------------------------------------------------------- /run_scripts/generate_init_target.sh: -------------------------------------------------------------------------------- 1 | python core/traj_opt/gen_init_target/state_generator.py \ 2 | --env_name 'CutRearrange-v1' \ 3 | --save_dir 'datasets/0001_test_dataset' \ 4 | --mode 'train' 5 | -------------------------------------------------------------------------------- /run_scripts/pasta_plan.sh: -------------------------------------------------------------------------------- 1 | # : pointflow model after trained with 'load_set_from_buffer=True'. 2 | # : pasta model with policy trained. 3 | # --vae_resume_path '' \ 4 | # --resume_path '' \ 5 | 6 | python core/pasta/train_pasta.py \ 7 | --env_name 'CutRearrangeSpread-v1' \ 8 | --input_mode 'pc' \ 9 | --run_plan True \ 10 | --vae_resume_path 'data/released_models/pointflow_vae/CutRearrangeSpread-v1/pointflow_pasta_crs.pt' \ 11 | --resume_path 'data/released_models/pasta_abstraction/CutRearrangeSpread-v1/fullmodel_pasta_crs.ckpt' \ 12 | --train_modules '' \ 13 | --load_modules 'reward' 'fea' 'policy' \ 14 | --plan_step 3 \ 15 | --num_tools 3 \ 16 | --dimz 2 \ 17 | --actor_arch 'pointnet' \ 18 | --actor_latent_dim 64 \ 19 | --reward_latent_dim 1024 \ 20 | --fea_arch 'v3' \ 21 | --fea_latent_dim 1024 \ 22 | --fea_center True \ 23 | --eval_plan True \ 24 | --plan_bs 5000 \ 25 | --adam_sample 5000 \ 26 | --adam_iter 100 \ 27 | --use_wandb 1 28 | -------------------------------------------------------------------------------- /run_scripts/pasta_train_abstraction.sh: -------------------------------------------------------------------------------- 1 | # : dataset path after generating dbscan labels for the dataset for pasta training 2 | # : pointflow model after trained with 'load_set_from_buffer=True'. 3 | # : pasta model with policy trained. 4 | # --dataset_path '' \ 5 | # --vae_resume_path '' \ 6 | # --resume_path '' \ 7 | 8 | python core/pasta/train_pasta.py \ 9 | --env_name 'CutRearrangeSpread-v1' \ 10 | --input_mode 'pc' \ 11 | --dataset_path 'data/released_datasets/CutRearrangeSpread-v1/0609_crs_dataset' \ 12 | --vae_resume_path 'data/released_models/pointflow_vae/CutRearrangeSpread-v1/pointflow_pasta_crs.pt' \ 13 | --resume_path 'data/released_models/pasta_policy/policy_pasta_crs.ckpt' \ 14 | --train_modules 'reward' 'fea' \ 15 | --load_modules 'policy' \ 16 | --plan_step 3 \ 17 | --num_tools 3 \ 18 | --dimz 2 \ 19 | --batch_size 256 \ 20 | --actor_batch_size 10 \ 21 | --actor_latent_dim 64 \ 22 | --reward_latent_dim 1024 \ 23 | --fea_latent_dim 1024 \ 24 | --num_random_neg 2048 \ 25 | --num_buffer_neg 2048 \ 26 | --fea_z_noise 0.02 \ 27 | --fea_t_noise 0.01 \ 28 | --il_lr 1e-4 \ 29 | --il_eval_freq 10 \ 30 | --il_num_epoch 300 \ 31 | --actor_arch 'pointnet' \ 32 | --fea_arch 'v3' \ 33 | --hard_negative_type 'obs_goal' \ 34 | --eval_plan True \ 35 | --plan_bs 5000 \ 36 | --adam_sample 5000 \ 37 | --adam_iter 100 \ 38 | --use_wandb 1 39 | -------------------------------------------------------------------------------- /run_scripts/pasta_train_policy.sh: -------------------------------------------------------------------------------- 1 | # : dataset path after generating dbscan labels for the dataset for pasta training 2 | # : pointflow model after trained with 'load_set_from_buffer=True' 3 | 4 | python core/pasta/train_pasta.py \ 5 | --env_name 'CutRearrangeSpread-v1' \ 6 | --input_mode 'pc' \ 7 | --dataset_path 'data/released_datasets/demonstrations/CutRearrangeSpread-v1/0609_crs_dataset' \ 8 | --vae_resume_path 'data/released_models/pointflow_vae/CutRearrangeSpread-v1/pointflow_pasta_crs.pt' \ 9 | --train_modules 'policy' \ 10 | --plan_step 3 \ 11 | --num_tools 3 \ 12 | --dimz 2 \ 13 | --batch_size 256 \ 14 | --actor_batch_size 10 \ 15 | --actor_latent_dim 64 \ 16 | --reward_latent_dim 1024 \ 17 | --fea_latent_dim 1024 \ 18 | --num_random_neg 2048 \ 19 | --num_buffer_neg 2048 \ 20 | --il_eval_freq 100 \ 21 | --il_num_epoch 2000 \ 22 | --eval_skill True \ 23 | --obs_noise 0.005 \ 24 | --il_lr 1e-4 \ 25 | --actor_arch 'pointnet' 26 | -------------------------------------------------------------------------------- /run_scripts/run_dbscan.sh: -------------------------------------------------------------------------------- 1 | # replace dataset_path to your '' to generate dbscan labels for the dataset for pasta training 2 | python core/pasta/generate_dbscan_label.py \ 3 | --dataset_path '' \ 4 | --dbscan_eps 0.03 \ 5 | --dbscan_min_samples 6 \ 6 | --dbscan_min_points 10 -------------------------------------------------------------------------------- /run_scripts/run_gbto.sh: -------------------------------------------------------------------------------- 1 | python core/traj_opt/gen_data.py \ 2 | --env_name 'CutRearrange-v1' \ 3 | --algo 'imitation' \ 4 | --adam_loss_type 'twoway_chamfer' \ 5 | --data_name 'demo' \ 6 | --gen_num_batch 1 \ 7 | --gen_batch_id 0 \ 8 | --gd_max_iter 5 -------------------------------------------------------------------------------- /taichi_three/LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 archibate 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /taichi_three/Tina-Dev.py: -------------------------------------------------------------------------------- 1 | bl_info = { 2 | 'name': 'Tina (dev mode)', 3 | 'description': 'A soft-renderer based on Taichi programming language', 4 | 'author': 'archibate <1931127624@qq.com>', 5 | 'version': (0, 0, 0), 6 | 'blender': (2, 81, 0), 7 | 'location': 'Render -> Tina', 8 | 'support': 'TESTING', 9 | 'wiki_url': 'https://github.com/archibate/tina/wiki', 10 | 'tracker_url': 'https://github.com/archibate/tina/issues', 11 | 'warning': 'Development mode', 12 | 'category': 'Render', 13 | } 14 | 15 | 16 | import sys 17 | sys.path.insert(0, '/home/bate/Develop/three_taichi') 18 | 19 | 20 | registered = False 21 | 22 | 23 | def register(): 24 | print('Tina-Dev register...') 25 | import tina_blend 26 | tina_blend.register() 27 | 28 | global registered 29 | registered = True 30 | print('...register done') 31 | 32 | 33 | def unregister(): 34 | print('Tina-Dev unregister...') 35 | import tina_blend 36 | tina_blend.unregister() 37 | 38 | global registered 39 | registered = False 40 | print('...unregister done') 41 | 42 | 43 | def reload_addon(): 44 | import tina 45 | import tina_blend 46 | if registered: 47 | tina_blend.unregister() 48 | tina.__lazyreload__() 49 | tina_blend.__lazyreload__() 50 | tina_blend.register() 51 | 52 | 53 | @eval('lambda x: x()') 54 | def _(): 55 | class Reload: 56 | def __repr__(self): 57 | import os 58 | import bpy 59 | os.system('clear') 60 | reload_addon() 61 | bpy.context.scene.frame_current = bpy.context.scene.frame_current 62 | return 'reloaded' 63 | 64 | __import__('bpy').a = Reload() 65 | -------------------------------------------------------------------------------- /taichi_three/assets/.gitignore: -------------------------------------------------------------------------------- 1 | bunny.* 2 | lens.* 3 | *.npy 4 | *.vdb 5 | *.blend* 6 | hp_* 7 | *.hdr 8 | *.npz 9 | semisphere.* 10 | monkey_cornell.* 11 | -------------------------------------------------------------------------------- /taichi_three/assets/atms.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/atms.png -------------------------------------------------------------------------------- /taichi_three/assets/cloth.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/cloth.jpg -------------------------------------------------------------------------------- /taichi_three/assets/cornell.mtl: -------------------------------------------------------------------------------- 1 | # Blender MTL File: 'cornell.blend' 2 | # Material Count: 2 3 | 4 | newmtl Material 5 | Ns 225.000000 6 | Ka 1.000000 1.000000 1.000000 7 | Kd 0.800000 0.800000 0.800000 8 | Ks 0.500000 0.500000 0.500000 9 | Ke 0.000000 0.000000 0.000000 10 | Ni 1.450000 11 | d 1.000000 12 | illum 2 13 | 14 | newmtl Material.001 15 | Ns 820.041320 16 | Ka 0.881818 0.881818 0.881818 17 | Kd 0.800000 0.800000 0.800000 18 | Ks 0.545455 0.545455 0.545455 19 | Ke 0.000000 0.000000 0.000000 20 | Ni 1.450000 21 | d 1.000000 22 | illum 3 23 | -------------------------------------------------------------------------------- /taichi_three/assets/cube.obj: -------------------------------------------------------------------------------- 1 | # Blender v2.90.0 OBJ File: '' 2 | # www.blender.org 3 | o Cube 4 | v 1.000000 1.000000 -1.000000 5 | v 1.000000 -1.000000 -1.000000 6 | v 1.000000 1.000000 1.000000 7 | v 1.000000 -1.000000 1.000000 8 | v -1.000000 1.000000 -1.000000 9 | v -1.000000 -1.000000 -1.000000 10 | v -1.000000 1.000000 1.000000 11 | v -1.000000 -1.000000 1.000000 12 | vt 0.625000 0.500000 13 | vt 0.875000 0.500000 14 | vt 0.875000 0.750000 15 | vt 0.625000 0.750000 16 | vt 0.375000 0.750000 17 | vt 0.625000 1.000000 18 | vt 0.375000 1.000000 19 | vt 0.375000 0.000000 20 | vt 0.625000 0.000000 21 | vt 0.625000 0.250000 22 | vt 0.375000 0.250000 23 | vt 0.125000 0.500000 24 | vt 0.375000 0.500000 25 | vt 0.125000 0.750000 26 | vn 0.0000 1.0000 0.0000 27 | vn 0.0000 0.0000 1.0000 28 | vn -1.0000 0.0000 0.0000 29 | vn 0.0000 -1.0000 0.0000 30 | vn 1.0000 0.0000 0.0000 31 | vn 0.0000 0.0000 -1.0000 32 | s off 33 | f 1/1/1 5/2/1 7/3/1 3/4/1 34 | f 4/5/2 3/4/2 7/6/2 8/7/2 35 | f 8/8/3 7/9/3 5/10/3 6/11/3 36 | f 6/12/4 2/13/4 4/5/4 8/14/4 37 | f 2/13/5 1/1/5 3/4/5 4/5/5 38 | f 6/11/6 5/10/6 1/1/6 2/13/6 39 | -------------------------------------------------------------------------------- /taichi_three/assets/lut.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/lut.jpg -------------------------------------------------------------------------------- /taichi_three/assets/multimtl.mtl: -------------------------------------------------------------------------------- 1 | # Blender MTL File: 'None' 2 | # Material Count: 2 3 | 4 | newmtl Material1 5 | Ns 225.000000 6 | Ka 1.000000 1.000000 1.000000 7 | Kd 0.149896 0.800000 0.202048 8 | Ks 0.500000 0.500000 0.500000 9 | Ke 0.000000 0.000000 0.000000 10 | Ni 1.450000 11 | d 1.000000 12 | illum 2 13 | 14 | newmtl Material2 15 | Ns 225.000000 16 | Ka 1.000000 1.000000 1.000000 17 | Kd 0.800000 0.093049 0.192718 18 | Ks 0.500000 0.500000 0.500000 19 | Ke 0.000000 0.000000 0.000000 20 | Ni 1.450000 21 | d 1.000000 22 | illum 2 23 | -------------------------------------------------------------------------------- /taichi_three/assets/normal.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/normal.png -------------------------------------------------------------------------------- /taichi_three/assets/pattern.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/pattern.jpg -------------------------------------------------------------------------------- /taichi_three/assets/plane.obj: -------------------------------------------------------------------------------- 1 | # Blender v2.90.0 OBJ File: '' 2 | # www.blender.org 3 | o Plane 4 | v -1.000000 0.000000 1.000000 5 | v 1.000000 0.000000 1.000000 6 | v -1.000000 0.000000 -1.000000 7 | v 1.000000 0.000000 -1.000000 8 | vt 0.000000 0.000000 9 | vt 1.000000 0.000000 10 | vt 1.000000 1.000000 11 | vt 0.000000 1.000000 12 | vn 0.0000 1.0000 0.0000 13 | s off 14 | f 1/1/1 2/2/1 4/3/1 3/4/1 15 | -------------------------------------------------------------------------------- /taichi_three/assets/shadow.mtl: -------------------------------------------------------------------------------- 1 | # Blender MTL File: 'shadow.blend' 2 | # Material Count: 2 3 | 4 | newmtl Material 5 | Ns 323.999994 6 | Ka 1.000000 1.000000 1.000000 7 | Kd 0.800000 0.800000 0.800000 8 | Ks 0.500000 0.500000 0.500000 9 | Ke 0.0 0.0 0.0 10 | Ni 1.450000 11 | d 1.000000 12 | illum 2 13 | 14 | newmtl None 15 | Ns 500 16 | Ka 0.8 0.8 0.8 17 | Kd 0.8 0.8 0.8 18 | Ks 0.8 0.8 0.8 19 | d 1 20 | illum 2 21 | -------------------------------------------------------------------------------- /taichi_three/assets/shadow.obj: -------------------------------------------------------------------------------- 1 | # Blender v2.80 (sub 75) OBJ File: 'shadow.blend' 2 | # www.blender.org 3 | mtllib shadow.mtl 4 | o Cube 5 | v 0.303343 -0.008899 -0.428899 6 | v 0.303343 -0.428899 0.008899 7 | v 0.303343 0.428899 -0.008899 8 | v 0.303343 0.008899 0.428899 9 | v -0.303343 -0.008899 -0.428899 10 | v -0.303343 -0.428899 0.008899 11 | v -0.303343 0.428899 -0.008899 12 | v -0.303343 0.008899 0.428899 13 | vt 0.375000 0.000000 14 | vt 0.625000 0.000000 15 | vt 0.625000 0.250000 16 | vt 0.375000 0.250000 17 | vt 0.375000 0.250000 18 | vt 0.625000 0.250000 19 | vt 0.625000 0.500000 20 | vt 0.375000 0.500000 21 | vt 0.625000 0.750000 22 | vt 0.375000 0.750000 23 | vt 0.625000 0.750000 24 | vt 0.625000 1.000000 25 | vt 0.375000 1.000000 26 | vt 0.125000 0.500000 27 | vt 0.375000 0.500000 28 | vt 0.375000 0.750000 29 | vt 0.125000 0.750000 30 | vt 0.625000 0.500000 31 | vt 0.875000 0.500000 32 | vt 0.875000 0.750000 33 | vn 0.0000 0.6923 -0.7216 34 | vn 0.0000 0.7216 0.6923 35 | vn -1.0000 0.0000 0.0000 36 | vn 0.0000 -0.6923 0.7216 37 | vn 1.0000 -0.0000 0.0000 38 | vn 0.0000 -0.7216 -0.6923 39 | usemtl Material 40 | s off 41 | f 1/1/1 5/2/1 7/3/1 3/4/1 42 | f 4/5/2 3/6/2 7/7/2 8/8/2 43 | f 8/8/3 7/7/3 5/9/3 6/10/3 44 | f 6/10/4 2/11/4 4/12/4 8/13/4 45 | f 2/14/5 1/15/5 3/16/5 4/17/5 46 | f 6/18/6 5/19/6 1/20/6 2/11/6 47 | o Plane 48 | v -2.000000 0.750960 2.106195 49 | v 2.000000 0.750960 2.106195 50 | v -2.000000 -2.135532 -0.662949 51 | v 2.000000 -2.135532 -0.662949 52 | vt 0.000000 0.000000 53 | vt 1.000000 0.000000 54 | vt 1.000000 1.000000 55 | vt 0.000000 1.000000 56 | vn 0.0000 0.6923 -0.7216 57 | usemtl None 58 | s off 59 | f 9/21/7 10/22/7 12/23/7 11/24/7 60 | -------------------------------------------------------------------------------- /taichi_three/assets/skybox.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/skybox.jpg -------------------------------------------------------------------------------- /taichi_three/assets/smallptwall.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Xingyu-Lin/PASTA/ab892a11c0a2de7be0226b60cb7fde48002e5845/taichi_three/assets/smallptwall.png -------------------------------------------------------------------------------- /taichi_three/bench.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | 4 | ti.init(ti.gpu) 5 | 6 | 7 | nx, ny = 512, 512 8 | 9 | 10 | x = ti.field(float, (2, nx, ny // 2)) 11 | r = ti.field(float, (2, nx, ny // 2)) 12 | 13 | 14 | @ti.pyfunc 15 | def at(i, j): 16 | return (i + j) % 2, i, j // 2 17 | 18 | 19 | @ti.kernel 20 | def solve(): 21 | for t in ti.static(range(2)): 22 | for i, j in ti.ndrange(nx, ny): 23 | if (i + j) % 2 != t: 24 | continue 25 | 26 | xl = x[at(i - 1, j)] 27 | xr = x[at(i + 1, j)] 28 | xb = x[at(i, j - 1)] 29 | xt = x[at(i, j + 1)] 30 | div = r[at(i, j)] 31 | xc = (xl + xr + xb + xt - div) * 0.25 32 | x[at(i, j)] = xc 33 | 34 | 35 | @ti.kernel 36 | def dump(out: ti.ext_arr()): 37 | for i, j in ti.ndrange(nx, ny): 38 | out[i, j] = x[at(i, j)] 39 | 40 | 41 | r[at(256, 256)] = -2 42 | gui = ti.GUI('jacobi', (nx, ny)) 43 | while gui.running and not gui.get_event(gui.ESCAPE): 44 | for i in range(200): 45 | solve() 46 | out = np.empty((nx, ny), dtype=np.float32) 47 | dump(out) 48 | gui.set_image(out) 49 | gui.show() 50 | -------------------------------------------------------------------------------- /taichi_three/conv.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | 4 | 5 | A = np.array([ 6 | [0, 1, 0], 7 | [1, 0, 1], 8 | [0, 1, 0], 9 | ]) 10 | 11 | 12 | def conv(A, B): 13 | m, n = A.shape 14 | s, t = B.shape 15 | C = np.zeros((m + s - 1, n + t - 1), dtype=A.dtype) 16 | for i in range(m): 17 | for j in range(n): 18 | for k in range(s): 19 | for l in range(t): 20 | C[i + k, j + l] += A[i, j] * B[k, l] 21 | return C 22 | 23 | 24 | B = A 25 | print(B) 26 | B = conv(B, A) 27 | print(B) 28 | B = conv(B, A) 29 | print(B) 30 | B = conv(B, A) 31 | print(B) 32 | -------------------------------------------------------------------------------- /taichi_three/docs/connect.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | verts, faces = tina.readobj('assets/monkey.obj', simple=True) 7 | 8 | scene = tina.Scene() 9 | mesh = tina.ConnectiveMesh() 10 | scene.add_object(mesh, tina.Classic()) 11 | 12 | mesh.set_vertices(verts) 13 | mesh.set_faces(faces) 14 | 15 | 16 | gui = ti.GUI() 17 | while gui.running: 18 | scene.input(gui) 19 | verts[100, 1] = ti.sin(gui.frame * 0.02) 20 | mesh.set_vertices(verts) 21 | scene.render() 22 | gui.set_image(scene.img) 23 | gui.show() 24 | -------------------------------------------------------------------------------- /taichi_three/docs/gltf.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.cpu) 5 | 6 | scene = tina.Scene(smoothing=True, texturing=True) 7 | scene.load_gltf('assets/cornell.gltf') 8 | #scene.load_gltf('/home/bate/Downloads/军用飞机/DamagedHelmet.gltf') 9 | 10 | gui = ti.GUI('gltf') 11 | scene.init_control(gui, center=(0, 2, 0), radius=6) 12 | 13 | while gui.running: 14 | scene.input(gui) 15 | scene.render() 16 | gui.set_image(scene.img) 17 | gui.show() 18 | -------------------------------------------------------------------------------- /taichi_three/docs/lighting.py: -------------------------------------------------------------------------------- 1 | # In this episode, you'll learn how to use lights and materials in Tina 2 | # 3 | # This tutorial is based on docs/monkey.py, make sure you check that first 4 | 5 | import taichi as ti 6 | import tina 7 | 8 | ti.init(ti.gpu) 9 | 10 | scene = tina.Scene() 11 | 12 | # 5. Material - for describing the material of an object 13 | material = tina.PBR(metallic=0.6, roughness=0.2) 14 | # parameters may also be specified by textures (add texturing=True to Scene) 15 | #material = tina.PBR(basecolor=tina.Texture('assets/cloth.jpg')) 16 | 17 | model = tina.MeshModel('assets/monkey.obj') 18 | # load our model into the scene with material specified: 19 | scene.add_object(model, material) 20 | 21 | gui = ti.GUI('lighting') 22 | 23 | # now, let's add some custom light sources into the scene for test 24 | # 25 | # first of all, remove the 'default light' from scene: 26 | scene.lighting.clear_lights() 27 | # adds a directional light with direction (0, 0, 1), with white color 28 | # the direction will be automatically normalized to obtain desired result 29 | scene.lighting.add_light(dir=[0, 0, 1], color=[1, 1, 1]) 30 | # adds a point light at position (1, 1.5, 0.3), with red color 31 | scene.lighting.add_light(pos=[1, 1.5, 0.3], color=[1, 0, 0]) 32 | # specifies the ambient color to be dark green 33 | scene.lighting.set_ambient_light([0, 0.06, 0]) 34 | 35 | while gui.running: 36 | scene.input(gui) 37 | scene.render() 38 | gui.set_image(scene.img) 39 | gui.show() 40 | -------------------------------------------------------------------------------- /taichi_three/docs/monkey.py: -------------------------------------------------------------------------------- 1 | # Tina is a real-time soft renderer based on Taichi for visualizing 3D scenes. 2 | # 3 | # To get started, let's try to load and display a monkey model in the GUI. 4 | 5 | import taichi as ti 6 | import tina 7 | 8 | ti.init(ti.gpu) # use GPU backend for better speed 9 | 10 | # to make tina actually display things, we need at least three things: 11 | # 12 | # 1. Scene - the top structure that manages all resources in the scene 13 | scene = tina.Scene() 14 | 15 | # 2. Model - the model to be displayed 16 | # 17 | # here we use `tina.MeshModel` which can load models from OBJ format files 18 | model = tina.MeshModel('assets/monkey.obj') 19 | # and, don't forget to add the model into the scene so that it gets displayed 20 | scene.add_object(model) 21 | 22 | # 3. GUI - we also need to create an window for display 23 | gui = ti.GUI('monkey') 24 | 25 | while gui.running: 26 | # update the camera transform from mouse events (will invoke gui.get_events) 27 | scene.input(gui) 28 | 29 | # render scene to image 30 | scene.render() 31 | 32 | # show the image in GUI 33 | gui.set_image(scene.img) 34 | gui.show() 35 | -------------------------------------------------------------------------------- /taichi_three/docs/options.py: -------------------------------------------------------------------------------- 1 | # In this episode, you'll learn some basic options to specify for a Tina scene. 2 | # 3 | # This tutorial is based on docs/monkey.py, make sure you check that first 4 | 5 | import taichi as ti 6 | import tina 7 | 8 | ti.init(ti.gpu) 9 | 10 | # There are some options you may specify to tina.Scene, try turn off some 11 | # of them and see what's the difference 12 | # 13 | # culling: enable face culling for better performance (default: on) 14 | # clipping: enable view space clipping for rid objects out of depth (default: off) 15 | # smoothing: enable smooth shading by interpolating normals (default: off) 16 | # texturing: enable texture coordinates, see docs/texture.py (default: off) 17 | # taa: temporal anti-aliasing, won't work well for dynamic scenes (default: off) 18 | scene = tina.Scene(smoothing=True, taa=True) 19 | 20 | # (make sure your OBJ model have normals to make smooth shading work) 21 | model = tina.MeshModel('assets/torus.obj') 22 | scene.add_object(model) 23 | 24 | gui = ti.GUI('options') 25 | 26 | while gui.running: 27 | scene.input(gui) 28 | scene.render() 29 | gui.set_image(scene.img) 30 | gui.show() 31 | -------------------------------------------------------------------------------- /taichi_three/docs/particles.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.Scene() 8 | 9 | pars = tina.SimpleParticles() 10 | material = tina.Classic() 11 | scene.add_object(pars, material) 12 | 13 | gui = ti.GUI('particles') 14 | 15 | pos = np.random.rand(1024, 3).astype(np.float32) * 2 - 1 16 | pars.set_particles(pos) 17 | radius = np.random.rand(1024).astype(np.float32) * 0.1 + 0.1 18 | pars.set_particle_radii(radius) 19 | color = np.random.rand(1024, 3).astype(np.float32) * 0.8 + 0.2 20 | pars.set_particle_colors(color) 21 | 22 | while gui.running: 23 | scene.input(gui) 24 | scene.render() 25 | gui.set_image(scene.img) 26 | gui.show() 27 | -------------------------------------------------------------------------------- /taichi_three/docs/pathtrace.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True, texturing=True) 8 | scene.lighting.skybox = tina.Atomsphere() 9 | material = tina.Phong() 10 | mesh = tina.MeshModel('assets/monkey.obj') 11 | scene.add_object(mesh, material) 12 | pars = tina.SimpleParticles() 13 | scene.add_object(pars, material) 14 | pars.set_particles(np.random.rand(2**10, 3) * 2 - 1) 15 | 16 | gui = ti.GUI('pathtrace', scene.res) 17 | 18 | scene.update() 19 | while gui.running: 20 | scene.input(gui) 21 | scene.render(nsteps=8) 22 | gui.set_image(scene.img) 23 | gui.show() 24 | -------------------------------------------------------------------------------- /taichi_three/docs/primitives.py: -------------------------------------------------------------------------------- 1 | import tina 2 | 3 | scene = tina.Scene(smoothing=True, texturing=True, taa=True) 4 | 5 | #mesh = tina.PrimitiveMesh.sphere() 6 | mesh = tina.PrimitiveMesh.cylinder() 7 | wire = tina.MeshToWire(mesh) 8 | scene.add_object(mesh) 9 | scene.add_object(wire) 10 | 11 | gui = tina.ti.GUI('primitives') 12 | 13 | while gui.running: 14 | scene.input(gui) 15 | scene.render() 16 | gui.set_image(scene.img) 17 | gui.show() 18 | -------------------------------------------------------------------------------- /taichi_three/docs/smooth.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | scene = tina.Scene(smoothing=True) 7 | 8 | model = tina.MeshSmoothNormal(tina.MeshModel('assets/monkey.obj')) 9 | scene.add_object(model) 10 | 11 | gui = ti.GUI('smooth') 12 | 13 | while gui.running: 14 | scene.input(gui) 15 | scene.render() 16 | gui.set_image(scene.img) 17 | gui.show() 18 | -------------------------------------------------------------------------------- /taichi_three/docs/transform.py: -------------------------------------------------------------------------------- 1 | # In this episode, you'll learn how to make use of nodes for transformation 2 | # 3 | # This tutorial is based on docs/monkey.py, make sure you check that first 4 | 5 | import taichi as ti 6 | import tina 7 | 8 | ti.init(ti.gpu) 9 | 10 | scene = tina.Scene() 11 | 12 | # load the monkey model with the `tina.MeshModel` node: 13 | model = tina.MeshModel('assets/monkey.obj') 14 | # transform the mesh using the `tina.Transform` node: 15 | tmodel = tina.MeshTransform(model) 16 | # load the desired node to be displayed: 17 | scene.add_object(tmodel) 18 | #scene.add_object(model) 19 | 20 | gui = ti.GUI('transform') 21 | 22 | while gui.running: 23 | scene.input(gui) 24 | 25 | # create a matrix representing translation along X-axis 26 | dx = ti.sin(gui.frame * 0.03) 27 | matrix = tina.translate([dx, 0, 0]) 28 | # set the model transformation matrix for `tina.Transform` node 29 | tmodel.set_transform(matrix) 30 | 31 | scene.render() 32 | gui.set_image(scene.img) 33 | gui.show() 34 | -------------------------------------------------------------------------------- /taichi_three/docs/triangle.py: -------------------------------------------------------------------------------- 1 | # Tina is a real-time soft renderer based on Taichi for visualizing 3D scenes. 2 | # 3 | # To get started, let's try to make a simple triangle and display it in the GUI. 4 | 5 | import taichi as ti 6 | import numpy as np 7 | import tina 8 | 9 | # Tina use a right-handed coordinate system in world space: 10 | # 11 | # +Z: back. 12 | # +X: right. 13 | # +Y: up. 14 | # 15 | # The camera looks from +Z to target by default. 16 | 17 | # make a simple triangle by specifying the vertices 18 | verts = np.array([[ 19 | [-1, -1, 0], # vertex 1 20 | [ 1, -1, 0], # vertex 2 21 | [ 0, 1, 0], # vertex 3 22 | ]]) 23 | # also note that face vertices needs to be **counter-clockwise** to be visible 24 | # you may disable such face culling policy by using tina.Scene(culling=False) 25 | 26 | # to make tina actually display things, we need at least three things: 27 | # 28 | # 1. Scene - the top structure that manages all resources in the scene 29 | scene = tina.Scene() 30 | 31 | # 2. Model - the model to be displayed 32 | # 33 | # here we use `tina.SimpleMesh` which allows use to specify the vertices manually 34 | mesh = tina.SimpleMesh() 35 | # and, don't forget to add the object into the scene so that it gets displayed 36 | scene.add_object(mesh) 37 | 38 | # 3. GUI - we also need to create an window for display 39 | gui = ti.GUI('triangle') 40 | 41 | while gui.running: 42 | # update the camera transform from mouse events (will invoke gui.get_events) 43 | scene.input(gui) 44 | 45 | # set face vertices by feeding a numpy array into it 46 | mesh.set_face_verts(verts) 47 | # render image with objects (a triangle in this case) in scene 48 | scene.render() 49 | 50 | # show the image in GUI 51 | gui.set_image(scene.img) 52 | gui.show() 53 | -------------------------------------------------------------------------------- /taichi_three/docs/volume.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | dens = np.load('assets/smoke.npy') 8 | scene = tina.Scene(N=dens.shape[0], taa=True, density=16) 9 | volume = tina.SimpleVolume(N=dens.shape[0]) 10 | #model = tina.MeshModel('assets/monkey.obj') 11 | #scene.add_object(model, tina.CookTorrance(metallic=0.8)) 12 | scene.add_object(volume) 13 | 14 | gui = ti.GUI('volume', scene.res) 15 | 16 | volume.set_volume_density(dens) 17 | while gui.running: 18 | scene.input(gui) 19 | scene.render() 20 | gui.set_image(scene.img) 21 | gui.show() 22 | -------------------------------------------------------------------------------- /taichi_three/docs/wireframe.py: -------------------------------------------------------------------------------- 1 | # In this episode, you'll learn how to render a wireframe model in Tina 2 | # 3 | # This tutorial is based on docs/monkey.py, make sure you check that first 4 | 5 | import taichi as ti 6 | import tina 7 | 8 | ti.init(ti.cpu) 9 | 10 | # you may specify the line width for rendering wireframes: 11 | # taa=True turns on Temporal Anti-Aliasing to make lines smoother 12 | scene = tina.Scene(linewidth=2, taa=True) 13 | 14 | # load the monkey using `tina.MeshModel` node (`tina.SimpleMesh` works too): 15 | model = tina.MeshModel('assets/monkey.obj') 16 | # convert the mesh to its wireframe using the `tina.MeshToWire` node: 17 | wiremodel = tina.MeshToWire(model) 18 | 19 | # add the wireframe model into scene: 20 | scene.add_object(wiremodel) 21 | 22 | # add the original model, with a tiny scale: 23 | model = tina.MeshTransform(model, tina.scale(0.9)) 24 | scene.add_object(model) 25 | 26 | gui = ti.GUI('wireframe') 27 | 28 | while gui.running: 29 | scene.input(gui) 30 | scene.render() 31 | gui.set_image(scene.img) 32 | gui.show() 33 | -------------------------------------------------------------------------------- /taichi_three/examples/cornell_box.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True, texturing=True) 8 | scene.load_gltf('assets/cornell.gltf') 9 | scene.add_object(tina.MeshTransform(tina.MeshModel('assets/plane.obj'), 10 | tina.translate([0, 3.98, 0]) @ tina.scale(0.1)), tina.Lamp(color=64)) 11 | 12 | if isinstance(scene, tina.PTScene): 13 | scene.update() 14 | 15 | gui = ti.GUI('cornell_box', scene.res) 16 | scene.init_control(gui, center=(0, 2, 0), radius=5) 17 | 18 | while gui.running: 19 | scene.input(gui) 20 | if isinstance(scene, tina.PTScene): 21 | scene.render(nsteps=8) 22 | else: 23 | scene.render() 24 | gui.set_image(scene.img) 25 | gui.show() 26 | 27 | ti.imwrite(scene.img, 'cornell.png') 28 | -------------------------------------------------------------------------------- /taichi_three/examples/ibl_matball.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | scene = tina.Scene(smoothing=True, taa=True, ibl=True) 7 | 8 | roughness = tina.Param(float, initial=0)#.15) 9 | metallic = tina.Param(float, initial=0)#.25) 10 | material = tina.PBR(metallic=metallic, roughness=roughness) 11 | 12 | #model = tina.PrimitiveMesh.sphere() 13 | model = tina.MeshModel('assets/bunny.obj') 14 | scene.add_object(model, material) 15 | 16 | gui = ti.GUI('matball') 17 | roughness.make_slider(gui, 'roughness') 18 | metallic.make_slider(gui, 'metallic') 19 | 20 | while gui.running: 21 | scene.input(gui) 22 | scene.render() 23 | gui.set_image(scene.img) 24 | gui.show() 25 | -------------------------------------------------------------------------------- /taichi_three/examples/matball.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.cpu) 5 | 6 | scene = tina.Scene(smoothing=True) 7 | 8 | metallic = tina.Param(float, initial=0.0) 9 | specular = tina.Param(float, initial=0.5) 10 | roughness = tina.Param(float, initial=0.4) 11 | material = tina.PBR(metallic=metallic, roughness=roughness, specular=specular) 12 | 13 | #shineness = tina.Param(float, initial=32) 14 | #specular = tina.Param(float, initial=0.5) 15 | #material = tina.Classic(shineness=shineness, specular=specular) 16 | 17 | model = tina.PrimitiveMesh.sphere() 18 | scene.add_object(model, material) 19 | 20 | gui = ti.GUI('matball') 21 | if 'roughness' in globals(): 22 | roughness.make_slider(gui, 'roughness') 23 | if 'metallic' in globals(): 24 | metallic.make_slider(gui, 'metallic') 25 | if 'shineness' in globals(): 26 | shineness.make_slider(gui, 'shineness', 1, 500, 1) 27 | if 'specular' in globals(): 28 | specular.make_slider(gui, 'specular') 29 | 30 | while gui.running: 31 | scene.input(gui) 32 | scene.render() 33 | gui.set_image(scene.img) 34 | gui.show() 35 | -------------------------------------------------------------------------------- /taichi_three/examples/meshgrid_wave.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import time 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | 8 | @ti.func 9 | def Z(xy, t): 10 | return 0.1 * ti.sin(10 * xy.norm() - ti.tau * t) 11 | 12 | 13 | @ti.kernel 14 | def deform_mesh(t: float): 15 | for i, j in mesh.pos: 16 | mesh.pos[i, j].z = Z(mesh.pos[i, j].xy, t) 17 | 18 | 19 | scene = tina.Scene(smoothing=True) 20 | 21 | mesh = tina.MeshNoCulling(tina.MeshGrid(64)) 22 | scene.add_object(mesh) 23 | 24 | gui = ti.GUI('meshgrid_wave', scene.res) 25 | 26 | while gui.running: 27 | scene.input(gui) 28 | 29 | deform_mesh(gui.frame * 0.01) 30 | 31 | scene.render() 32 | gui.set_image(scene.img) 33 | gui.show() 34 | -------------------------------------------------------------------------------- /taichi_three/examples/rtx_matball.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | scene = tina.PTScene(smoothing=True) 7 | scene.engine.skybox = tina.Atomsphere() 8 | #scene.engine.skybox = tina.Skybox('assets/skybox.jpg') 9 | 10 | roughness = tina.Param(float, initial=0.15) 11 | metallic = tina.Param(float, initial=1.0) 12 | specular = tina.Param(float, initial=0.5) 13 | material = tina.PBR(metallic=metallic, roughness=roughness, specular=specular) 14 | 15 | model = tina.PrimitiveMesh.sphere() 16 | scene.add_object(model, material) 17 | 18 | gui = ti.GUI('matball') 19 | roughness.make_slider(gui, 'roughness') 20 | metallic.make_slider(gui, 'metallic') 21 | specular.make_slider(gui, 'specular') 22 | 23 | scene.update() 24 | while gui.running: 25 | scene.input(gui) 26 | scene.render() 27 | gui.set_image(scene.img) 28 | gui.show() 29 | 30 | #tina.pfmwrite('/tmp/color.pfm', scene.img) 31 | -------------------------------------------------------------------------------- /taichi_three/melt/dump.py: -------------------------------------------------------------------------------- 1 | from tina.advans import * 2 | 3 | pars = np.load('/tmp/mpm.npy') 4 | 5 | scene = tina.Scene() 6 | model = tina.SimpleParticles() 7 | scene.add_object(model) 8 | 9 | gui = ti.GUI() 10 | while gui.running: 11 | scene.input(gui) 12 | model.set_particles(pars) 13 | scene.render() 14 | gui.set_image(scene.img) 15 | gui.show() 16 | -------------------------------------------------------------------------------- /taichi_three/melt/lbvh.py: -------------------------------------------------------------------------------- 1 | #from tina.common import * 2 | 3 | 4 | #verts, faces = tina.readobj('assets/bunny.obj', simple=True) 5 | #verts = verts[faces] 6 | import numpy as np 7 | pos = np.load('assets/fluid.npy')[:3] 8 | 9 | 10 | def expandBits(v): 11 | v = (v * 0x00010001) & 0xFF0000FF; 12 | v = (v * 0x00000101) & 0x0F00F00F; 13 | v = (v * 0x00000011) & 0xC30C30C3; 14 | v = (v * 0x00000005) & 0x49249249; 15 | return v; 16 | 17 | def morton3D(x, y, z): 18 | x = min(max(x * 1024, 0), 1023); 19 | y = min(max(y * 1024, 0), 1023); 20 | z = min(max(z * 1024, 0), 1023); 21 | xx = expandBits(int(x)); 22 | yy = expandBits(int(y)); 23 | zz = expandBits(int(z)); 24 | return xx * 4 + yy * 2 + zz; 25 | 26 | 27 | def generateHierarchy(ids, codes, first, last): 28 | if first == last: 29 | return ids[first] 30 | 31 | split = findSplit(codes, first, last) 32 | 33 | childA = generateHierarchy(ids, codes, first, split) 34 | childB = generateHierarchy(ids, codes, split + 1, last) 35 | 36 | return (childA, childB) 37 | 38 | 39 | def countLeadingZeros(x): 40 | n = 0 41 | while x != 0: 42 | x >>= 1 43 | n += 1 44 | return 32 - n 45 | 46 | 47 | def findSplit(codes, first, last): 48 | code_first = codes[first] 49 | code_last = codes[last] 50 | 51 | if code_first == code_last: 52 | return (first + last) >> 1 53 | 54 | common_prefix = countLeadingZeros(code_first ^ code_last) 55 | 56 | split = first 57 | step = last - first 58 | 59 | while True: 60 | step = (step + 1) >> 1 61 | new_split = split + step 62 | if new_split < last: 63 | code_split = codes[new_split] 64 | split_prefix = countLeadingZeros(code_first ^ code_split) 65 | if split_prefix > common_prefix: 66 | split = new_split 67 | 68 | if step <= 1: 69 | break 70 | 71 | return split 72 | 73 | 74 | ids = range(len(pos)) 75 | codes = [morton3D(*pos[i]) for i in ids] 76 | _ = sorted(zip(ids, codes), key=lambda x: x[1]) 77 | ids = [_ for _, __ in _] 78 | codes = [_ for __, _ in _] 79 | tree = generateHierarchy(ids, codes, 0, len(ids) - 1) 80 | print(pos) 81 | print(tree) 82 | exit(1) 83 | -------------------------------------------------------------------------------- /taichi_three/melt/vbrbd.py: -------------------------------------------------------------------------------- 1 | # https://developer.nvidia.com/gpugems/gpugems3/part-v-physics-simulation/chapter-29-real-time-rigid-body-simulation-gpus 2 | 3 | from tina.advans import * 4 | #from voxelizer import MeshVoxelizer 5 | 6 | NE = 64 7 | N = 128 8 | R = 0.1 9 | Ks = 600 10 | Kd = 6 11 | Kt = 0 12 | Dt = 0.001 13 | Grav = 2.5 14 | pos = ti.Vector.field(3, float, N) 15 | vel = ti.Vector.field(3, float, N) 16 | epos = ti.Vector.field(3, float, NE) 17 | erot = ti.Vector.field(4, float, NE) 18 | 19 | @ti.kernel 20 | def reset(): 21 | for i in pos: 22 | pos[i] = V(ti.random(), ti.random(), ti.random()) * 2 - 1 23 | vel[i] = 0 24 | 25 | @ti.kernel 26 | def substep(): 27 | for i in pos: 28 | acc = V(0., 0., 0.) 29 | for j in range(N): 30 | if i == j: 31 | continue 32 | r = pos[j] - pos[i] 33 | v = vel[j] - vel[i] 34 | rn = r.norm() 35 | if rn > R * 2: 36 | continue 37 | rnn = r / rn 38 | fs = -Ks * (R * 2 - rn) * rnn 39 | vd = v.dot(rnn) * rnn 40 | ft = Kt * (v - vd) 41 | fd = Kd * vd 42 | acc += fs + fd + ft 43 | acc += V(0., -Grav, 0.) 44 | vel[i] += acc * Dt 45 | for i in pos: 46 | cond = (pos[i] < -1 and vel[i] < 0) or (pos[i] > 1 and vel[i] > 0) 47 | vel[i] = -0.1 * vel[i] if cond else vel[i] 48 | pos[i] += vel[i] * Dt 49 | 50 | def step(): 51 | for i in range(32): 52 | substep() 53 | 54 | scene = tina.Scene() 55 | model = tina.SimpleParticles(radius=R) 56 | scene.add_object(model) 57 | 58 | reset() 59 | gui = ti.GUI() 60 | while gui.running: 61 | scene.input(gui) 62 | if not gui.is_pressed(gui.SPACE): step() 63 | if gui.is_pressed('r'): reset() 64 | model.set_particles(pos.to_numpy()) 65 | scene.render() 66 | gui.set_image(scene.img) 67 | gui.show() 68 | -------------------------------------------------------------------------------- /taichi_three/setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | import tina 3 | 4 | setuptools.setup( 5 | name='taichi-tina', 6 | version='.'.join(map(str, tina.__version__)), 7 | author=tina.__author__.split('<')[0][:-1], 8 | author_email=tina.__author__.split('<')[1][:-1], 9 | url='https://github.com/taichi-dev/taichi_three', 10 | description='A real-time soft renderer based on the Taichi programming language', 11 | long_description=tina.__doc__, 12 | license=tina.__license__, 13 | keywords=['graphics', 'renderer'], 14 | classifiers=[ 15 | 'Intended Audience :: Developers', 16 | 'Intended Audience :: Science/Research', 17 | 'Topic :: Multimedia :: Graphics', 18 | 'Topic :: Games/Entertainment :: Simulation', 19 | 'Operating System :: OS Independent', 20 | ], 21 | python_requires='>=3.6', 22 | install_requires=[ 23 | 'taichi', 24 | 'transformations' 25 | ], 26 | packages=setuptools.find_packages(), 27 | ) 28 | -------------------------------------------------------------------------------- /taichi_three/tests/bdpt.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True, texturing=True) 8 | 9 | scene.add_object(tina.PrimitiveMesh.sphere(), tina.Glass()) 10 | scene.add_object(tina.MeshTransform(tina.MeshModel('assets/cube.obj'), 11 | tina.translate([0, -3, 0]) @ tina.scale(2)), tina.Lambert()) 12 | 13 | scene.add_object(tina.MeshTransform(tina.MeshModel('assets/plane.obj'), 14 | tina.translate([0, 4, 0]) @ tina.scale(0.2)), tina.Lamp(color=tina.Texture("assets/cloth.jpg")) * 64) 15 | 16 | gui = ti.GUI('bdpt', scene.res) 17 | 18 | scene.update() 19 | while gui.running: 20 | if scene.input(gui): 21 | scene.clear() 22 | scene.render(nsteps=6) 23 | scene.render_light(nsteps=6) 24 | gui.set_image(scene.img) 25 | gui.show() 26 | -------------------------------------------------------------------------------- /taichi_three/tests/blooming.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | scene = tina.Scene(smoothing=True, blooming=True) 7 | 8 | material = tina.Diffuse() 9 | model = tina.MeshModel('assets/monkey.obj') 10 | scene.add_object(model, material) 11 | 12 | gui = ti.GUI(res=scene.res) 13 | light = gui.slider('light', 0, 32, 0.1) 14 | light.value = 6 15 | 16 | scene.lighting.clear_lights() 17 | 18 | while gui.running: 19 | scene.lighting.set_ambient_light([light.value] * 3) 20 | scene.input(gui) 21 | scene.render() 22 | gui.set_image(scene.img) 23 | gui.show() 24 | -------------------------------------------------------------------------------- /taichi_three/tests/brdf.py: -------------------------------------------------------------------------------- 1 | from tina.advans import * 2 | 3 | ti.init(ti.cpu) 4 | 5 | PBR = 1 6 | 7 | if PBR: 8 | roughness = tina.Param(float, initial=0.4) 9 | metallic = tina.Param(float, initial=1.0) 10 | material = tina.PBR(basecolor=[1, 1, 1], roughness=roughness, metallic=metallic, specular=0.0) 11 | else: 12 | shineness = tina.Param(float, initial=10) 13 | material = tina.Phong(diffuse=[0, 0, 0], specular=[1, 1, 1], shineness=shineness) 14 | 15 | iu = ti.field(float, ()) 16 | 17 | k = 1 18 | n = 400 19 | f = ti.Vector.field(3, float, (n * k, n)) 20 | 21 | @ti.kernel 22 | def render(): 23 | idir = spherical(iu[None], 0.) 24 | for i, j in f: 25 | ou, ov = 1 - (i + 0.5) / n, (j + 0.5) / n 26 | odir = spherical(ou, ov) 27 | brdf = material.brdf(V(0., 0., 1.), idir, odir) 28 | f[i, j] = brdf * ou 29 | 30 | gui = ti.GUI('brdf', (n * k, n)) 31 | if PBR: 32 | roughness.make_slider(gui, 'roughness') 33 | metallic.make_slider(gui, 'metallic') 34 | else: 35 | shineness.make_slider(gui, 'shineness', 1, 128, 1) 36 | iu_slider = gui.slider('U', 0, 1, 0.01) 37 | avg_label = gui.label('average') 38 | 39 | while gui.running: 40 | for e in gui.get_events(): 41 | if e.key == gui.ESCAPE: 42 | gui.running = False 43 | render() 44 | im = f.to_numpy() 45 | avg_label.value = np.sum(im) / (3 * k * n**2) 46 | gui.set_image(1 - np.exp(-im)) 47 | gui.show() 48 | iu[None] = iu_slider.value 49 | -------------------------------------------------------------------------------- /taichi_three/tests/cookibl.py: -------------------------------------------------------------------------------- 1 | import tina 2 | 3 | tina.ti.init(tina.ti.gpu) 4 | 5 | li = tina.SkyboxLighting('assets/ballroom.npy', precision=2048) 6 | li.save('assets/ballroom.npz') 7 | -------------------------------------------------------------------------------- /taichi_three/tests/emission.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | scene = tina.PTScene() 7 | #scene.engine.skybox = tina.Atomsphere() 8 | 9 | model = tina.MeshTransform(tina.MeshModel('assets/plane.obj'), 10 | tina.translate([0, 0, 4]) @ tina.eularXYZ([ti.pi / 2, 0, 0])) 11 | material = tina.Emission() * 2 12 | scene.add_object(model, material) 13 | 14 | metallic = tina.Param(float, initial=1.0) 15 | roughness = tina.Param(float, initial=0.01) 16 | model = tina.MeshModel('assets/monkey.obj') 17 | material = tina.PBR(metallic=metallic, roughness=roughness) 18 | scene.add_object(model, material) 19 | 20 | gui = ti.GUI(res=scene.res) 21 | metallic.make_slider(gui, 'metallic') 22 | roughness.make_slider(gui, 'roughness') 23 | 24 | scene.update() 25 | while gui.running: 26 | scene.input(gui) 27 | scene.render() 28 | gui.set_image(scene.img) 29 | gui.show() 30 | -------------------------------------------------------------------------------- /taichi_three/tests/fpe.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.cpu) 6 | 7 | scene = tina.PTScene(smoothing=True) 8 | scene.engine.skybox = tina.PlainSkybox() 9 | 10 | scene.add_object(tina.MeshModel('assets/monkey.obj'), tina.Lambert()) 11 | 12 | gui = ti.GUI('fpe', scene.res) 13 | 14 | scene.update() 15 | while gui.running: 16 | if scene.input(gui): 17 | scene.clear() 18 | scene.render() 19 | gui.set_image(scene.img) 20 | gui.show() 21 | -------------------------------------------------------------------------------- /taichi_three/tests/fxaa.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.opengl) 5 | 6 | 7 | scene = tina.Scene(fxaa=True) 8 | model = tina.MeshModel('assets/monkey.obj') 9 | scene.add_object(model, tina.PBR(metallic=0.5, roughness=0.3)) 10 | 11 | 12 | gui = ti.GUI(res=scene.res) 13 | 14 | abs_thresh = gui.slider('abs_thresh', 0, 0.1, 0.002) 15 | rel_thresh = gui.slider('rel_thresh', 0, 0.5, 0.01) 16 | factor = gui.slider('factor', 0, 1, 0.01) 17 | abs_thresh.value = 0.0625 18 | rel_thresh.value = 0.063 19 | factor.value = 1 20 | 21 | while gui.running: 22 | scene.fxaa.rel_thresh[None] = rel_thresh.value 23 | scene.fxaa.abs_thresh[None] = abs_thresh.value 24 | scene.fxaa.factor[None] = factor.value 25 | scene.input(gui) 26 | scene.render() 27 | gui.set_image(scene.img) 28 | gui.show() 29 | -------------------------------------------------------------------------------- /taichi_three/tests/glass.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True, texturing=True) 8 | scene.engine.skybox = tina.Skybox('assets/skybox.jpg', cubic=True) 9 | scene.add_object(tina.MeshModel('assets/cube.obj'), tina.Glass()) 10 | 11 | gui = ti.GUI('bdpt', scene.res) 12 | 13 | scene.update() 14 | while gui.running: 15 | if scene.input(gui): 16 | scene.clear() 17 | scene.render(nsteps=12) 18 | #scene.render_light(nsteps=6) 19 | gui.set_image(scene.img) 20 | gui.show() 21 | -------------------------------------------------------------------------------- /taichi_three/tests/inter.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.cpu) 5 | 6 | 7 | scene = tina.Scene() 8 | model = tina.MeshModel('assets/cube.obj') 9 | material = tina.Emission() * tina.Input('pos') 10 | scene.add_object(model, material) 11 | 12 | gui = ti.GUI(res=scene.res) 13 | 14 | while gui.running: 15 | scene.input(gui) 16 | scene.render() 17 | gui.set_image(scene.img) 18 | gui.show() 19 | -------------------------------------------------------------------------------- /taichi_three/tests/knn.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | 4 | img = ti.imread('noise.png') 5 | #img = ti.imread('/opt/cuda/samples/3_Imaging/imageDenoising/data/portrait_noise.bmp') 6 | img = np.float32(img / 255) 7 | w, h, chans = img.shape 8 | 9 | src = ti.Vector.field(chans, float, (w, h)) 10 | dst = ti.Vector.field(chans, float, (w, h)) 11 | 12 | src.from_numpy(img) 13 | 14 | noise = 1 / 0.32**2 15 | lerp_c = 0.2 16 | win_rad = 3 17 | blk_rad = 3 18 | win_area = (2 * win_rad + 1)**2 19 | wei_thres = 0.02 20 | lerp_thres = 0.79 21 | 22 | @ti.kernel 23 | def denoise(): 24 | for x, y in src: 25 | cnt = 0.0 26 | wei = 0.0 27 | clr = ti.Vector([0.0, 0.0, 0.0]) 28 | 29 | clr00 = src[x, y] 30 | for i, j in ti.ndrange((-win_rad, win_rad + 1), (-win_rad, win_rad + 1)): 31 | wij = 0.0 32 | 33 | clrij = src[x + i, y + j] 34 | disij = (clr00 - clrij).norm_sqr() 35 | 36 | wij = ti.exp(-(disij * noise + (i**2 + j**2) / win_area)) 37 | cnt += 1 / win_area if wij > wei_thres else 0 38 | clr += clrij * wij 39 | wei += wij 40 | 41 | clr /= wei 42 | lerp_q = lerp_c if cnt > lerp_thres else 1 - lerp_c 43 | dst[x, y] = clr * (1 - lerp_q) + src[x, y] * lerp_q 44 | 45 | 46 | denoise() 47 | gui1 = ti.GUI('before', (w, h)) 48 | gui2 = ti.GUI('after', (w, h)) 49 | gui1.fps_limit = 30 50 | gui2.fps_limit = 30 51 | while gui1.running and gui2.running: 52 | gui1.running = not gui1.get_event(ti.GUI.ESCAPE) 53 | gui2.running = not gui2.get_event(ti.GUI.ESCAPE) 54 | gui1.set_image(src) 55 | gui2.set_image(dst) 56 | gui1.show() 57 | gui2.show() 58 | -------------------------------------------------------------------------------- /taichi_three/tests/mis.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.opengl) 6 | 7 | scene = tina.PTScene(smoothing=True) 8 | 9 | roughness = tina.Param(float, initial=0.2) 10 | metallic = tina.Param(float, initial=1.0) 11 | scene.add_object(tina.MeshModel('assets/sphere.obj'), tina.PBR(metallic=metallic, roughness=roughness)) 12 | 13 | scene.add_object(tina.MeshTransform(tina.MeshModel('assets/plane.obj'), 14 | tina.translate([0, 0, 4]) @ tina.eularXYZ([ti.pi / 2, 0, 0]) 15 | #@ tina.scale(0.1)), tina.Lamp(color=32)) 16 | @ tina.scale(0.4)), tina.Lamp(color=1)) 17 | 18 | gui = ti.GUI('path', scene.res) 19 | roughness.make_slider(gui, 'roughness', 0, 1, 0.01) 20 | metallic.make_slider(gui, 'metallic', 0, 1, 0.01) 21 | 22 | scene.update() 23 | while gui.running: 24 | if scene.input(gui): 25 | scene.clear() 26 | scene.render(nsteps=6) 27 | gui.set_image(scene.img) 28 | gui.show() 29 | 30 | #ti.imwrite(scene.img, 'output.png') 31 | -------------------------------------------------------------------------------- /taichi_three/tests/mtlid.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene() 8 | verts, faces = tina.readobj('assets/sphere.obj', simple=True) 9 | verts = verts[faces] 10 | model = tina.SimpleMesh() 11 | scene.add_object(model) 12 | scene.engine.skybox = tina.Atomsphere() 13 | scene.materials = [tina.Lambert(), tina.Lambert() * [0, 1, 0]] 14 | 15 | mtlids = np.zeros(verts.shape[0]) 16 | mtlids[3] = 1 17 | mtlids[5] = 1 18 | model.set_face_verts(verts) 19 | model.set_face_mtlids(mtlids) 20 | 21 | scene.update() 22 | scene.visualize() 23 | -------------------------------------------------------------------------------- /taichi_three/tests/nlm.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | 4 | img = ti.imread('noise.png') 5 | #img = ti.imread('/opt/cuda/samples/3_Imaging/imageDenoising/data/portrait_noise.bmp') 6 | img = np.float32(img / 255) 7 | w, h, chans = img.shape 8 | 9 | src = ti.Vector.field(chans, float, (w, h)) 10 | dst = ti.Vector.field(chans, float, (w, h)) 11 | 12 | src.from_numpy(img) 13 | 14 | noise = 1 / 1.45**2 15 | lerp_c = 0.2 16 | win_rad = 3 17 | blk_rad = 3 18 | win_area = (2 * win_rad + 1)**2 19 | wei_thres = 0.1 20 | lerp_thres = 0.1 21 | 22 | @ti.kernel 23 | def denoise(): 24 | for x, y in src: 25 | cnt = 0.0 26 | wei = 0.0 27 | clr = ti.Vector([0.0, 0.0, 0.0]) 28 | 29 | for i, j in ti.ndrange((-win_rad, win_rad + 1), (-win_rad, win_rad + 1)): 30 | wij = 0.0 31 | 32 | for m, n in ti.ndrange((-blk_rad, blk_rad + 1), (-blk_rad, blk_rad + 1)): 33 | clr00 = src[x + m, y + n] 34 | clrij = src[x + i + m, y + j + n] 35 | wij += (clr00 - clrij).norm_sqr() 36 | 37 | wij = ti.exp(-(wij * noise + (i**2 + j**2) / win_area)) 38 | cnt += 1 / win_area if wij > wei_thres else 0 39 | clr += src[x + i, y + j] * wij 40 | wei += wij 41 | 42 | clr /= wei 43 | lerp_q = lerp_c if cnt > lerp_thres else 1 - lerp_c 44 | dst[x, y] = clr * (1 - lerp_q) + src[x, y] * lerp_q 45 | 46 | 47 | denoise() 48 | gui1 = ti.GUI('before', (w, h)) 49 | gui2 = ti.GUI('after', (w, h)) 50 | gui1.fps_limit = 30 51 | gui2.fps_limit = 30 52 | while gui1.running and gui2.running: 53 | gui1.running = not gui1.get_event(ti.GUI.ESCAPE) 54 | gui2.running = not gui2.get_event(ti.GUI.ESCAPE) 55 | gui1.set_image(src) 56 | gui2.set_image(dst) 57 | gui1.show() 58 | gui2.show() 59 | -------------------------------------------------------------------------------- /taichi_three/tests/noise.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True, texturing=True) 8 | #scene.lighting.skybox = tina.Skybox('assets/skybox.jpg', cubic=True) 9 | model = tina.MeshModel('assets/bunny.obj') 10 | #material = tina.PBR(roughness=0.0, metallic=0.0) 11 | material = tina.PBR(roughness=0.2, metallic=0.8) 12 | scene.add_object(model, material) 13 | denoise = tina.Denoise(scene.res) 14 | 15 | if isinstance(scene, tina.PTScene): 16 | scene.update() 17 | 18 | gui = ti.GUI('noise', scene.res) 19 | 20 | while gui.running: 21 | scene.input(gui) 22 | if isinstance(scene, tina.PTScene): 23 | scene.render(nsteps=5) 24 | else: 25 | scene.render() 26 | #gui.set_image(scene.img) 27 | denoise.src.from_numpy(scene.img) 28 | denoise.nlm(radius=2, noiseness=0.9) 29 | gui.set_image(denoise.dst) 30 | gui.show() 31 | -------------------------------------------------------------------------------- /taichi_three/tests/path.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True) 8 | #scene.lighting.skybox = tina.Atomsphere() 9 | 10 | #material = tina.Phong(color=[0.25, 0.5, 0.5]) 11 | #scene.add_object(tina.PrimitiveMesh.sphere(), tina.Glass()) 12 | roughness = tina.Param(float, initial=0.332) 13 | metallic = tina.Param(float, initial=0.0) 14 | scene.add_object(tina.MeshModel('assets/sphere.obj'), tina.PBR(metallic=metallic, roughness=roughness)) 15 | #scene.add_object(tina.PrimitiveMesh.sphere(), tina.Mirror()) 16 | #scene.add_object(tina.MeshModel('assets/bunny.obj'), tina.Glass()) 17 | #scene.add_object(tina.MeshModel('assets/monkey.obj'), tina.Mirror()) 18 | #scene.add_object(tina.MeshModel('assets/cube.obj'), tina.Classic()) 19 | #scene.add_object(tina.MeshTransform(tina.MeshModel('assets/cube.obj'), tina.translate([0, -1.2, 0]) @ tina.scale([2, 0.05, 2])), tina.Lambert()) 20 | 21 | scene.add_object(tina.MeshTransform(tina.MeshModel('assets/plane.obj'), 22 | tina.translate([0, 0, 4]) @ tina.eularXYZ([ti.pi / 2, 0, 0]) 23 | @ tina.scale(0.1)), tina.Lamp(color=128)) 24 | 25 | gui = ti.GUI('path', scene.res) 26 | roughness.make_slider(gui, 'roughness', 0, 1, 0.01) 27 | metallic.make_slider(gui, 'metallic', 0, 1, 0.01) 28 | 29 | scene.update() 30 | while gui.running: 31 | if scene.input(gui): 32 | scene.clear() 33 | scene.render(nsteps=6) 34 | gui.set_image(scene.img) 35 | gui.show() 36 | 37 | #ti.imwrite(scene.img, 'output.png') 38 | -------------------------------------------------------------------------------- /taichi_three/tests/pick.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import pickle 3 | import tina 4 | 5 | 6 | x = ti.Matrix.shield(3, 3, float, 4) 7 | x[3][2, 1] = 1 8 | print(x) 9 | x = pickle.dumps(x) 10 | print(x) 11 | ti.init() 12 | x = pickle.loads(x) 13 | print(x) 14 | exit(1) 15 | -------------------------------------------------------------------------------- /taichi_three/tests/probe.py: -------------------------------------------------------------------------------- 1 | import tina 2 | import taichi as ti 3 | 4 | 5 | scene = tina.Scene(texturing=True, prob=True) 6 | 7 | verts, faces = tina.readobj('assets/monkey.obj', simple=True) 8 | mesh = tina.SimpleMesh() 9 | texture = tina.LerpTexture(x0=[1.0, 1.0, 1.0], x1=[0.0, 0.0, 1.0]) 10 | material = tina.Diffuse(color=texture) 11 | scene.add_object(mesh, material) 12 | 13 | probe = tina.ProbeShader(scene.res) 14 | scene.post_shaders.append(probe) 15 | 16 | mesh.set_face_verts(verts[faces]) 17 | 18 | @ti.func 19 | def ontouch(probe, I, r): 20 | e = probe.elmid[I] 21 | for i in ti.static(range(3)): 22 | mesh.coors[e, i].x = 1.0 23 | 24 | gui = ti.GUI() 25 | while gui.running: 26 | scene.input(gui) 27 | if gui.is_pressed(gui.LMB): 28 | mx, my = gui.get_cursor_pos() 29 | probe.touch(ontouch, mx, my, 0) 30 | scene.render() 31 | gui.set_image(scene.img) 32 | gui.show() 33 | -------------------------------------------------------------------------------- /taichi_three/tests/ptlight.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True) 8 | 9 | roughness = tina.Param(float, initial=0.2) 10 | metallic = tina.Param(float, initial=1.0) 11 | scene.add_object(tina.MeshModel('assets/sphere.obj'), tina.PBR(metallic=metallic, roughness=roughness)) 12 | 13 | pars = tina.SimpleParticles() 14 | scene.add_object(pars, tina.Lamp(color=32)) 15 | 16 | gui = ti.GUI('path', scene.res) 17 | roughness.make_slider(gui, 'roughness', 0, 1, 0.01) 18 | metallic.make_slider(gui, 'metallic', 0, 1, 0.01) 19 | 20 | pars.set_particles(np.array([ 21 | [0, 0, 5], 22 | ])) 23 | 24 | scene.update() 25 | while gui.running: 26 | if scene.input(gui): 27 | scene.clear() 28 | scene.render(nsteps=6) 29 | gui.set_image(scene.img) 30 | gui.show() 31 | 32 | #ti.imwrite(scene.img, 'output.png') 33 | -------------------------------------------------------------------------------- /taichi_three/tests/raytrace.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import taichi_inject 4 | import tina 5 | 6 | ti.init(ti.gpu) 7 | 8 | scene = tina.Scene(smoothing=True, texturing=True, rtx=True, taa=True) 9 | material = tina.Phong(shineness=64) 10 | mesh = tina.PrimitiveMesh.sphere() 11 | #mesh = tina.MeshModel('assets/monkey.obj') 12 | scene.add_object(mesh, material) 13 | 14 | gui = ti.GUI('raytrace', scene.res) 15 | 16 | scene.update() 17 | while gui.running: 18 | scene.input(gui) 19 | scene.render() 20 | gui.set_image(scene.img) 21 | gui.show() 22 | -------------------------------------------------------------------------------- /taichi_three/tests/rtao.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True)#, texturing=True) 8 | scene.engine.skybox = tina.PlainSkybox() 9 | 10 | scene.add_object(tina.MeshModel('assets/monkey.obj'), tina.Lambert()) 11 | #scene.add_object(tina.MeshModel('assets/cube.obj'), tina.Emission() * [.9, .4, .9]) 12 | 13 | gui = ti.GUI('rtao', scene.res) 14 | 15 | scene.update() 16 | while gui.running: 17 | if scene.input(gui): 18 | scene.clear() 19 | scene.render() 20 | gui.set_image(scene.img) 21 | gui.show() 22 | -------------------------------------------------------------------------------- /taichi_three/tests/rtx.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | scene = tina.Scene(smoothing=True, rtx=True) 7 | 8 | material = tina.Phong(shineness=256) 9 | 10 | model = tina.PrimitiveMesh.sphere() 11 | scene.add_object(model, material) 12 | 13 | model2 = tina.MeshTransform(tina.PrimitiveMesh.sphere(), tina.translate([1.4, 0, 0]) @ tina.scale(0.5)) 14 | scene.add_object(model2, material) 15 | 16 | gui = ti.GUI('matball') 17 | 18 | scene.update() 19 | scene.init_control(gui) 20 | while gui.running: 21 | scene.input(gui) 22 | scene.render() 23 | gui.set_image(scene.img) 24 | gui.show() 25 | -------------------------------------------------------------------------------- /taichi_three/tests/skybox.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene() 8 | scene.add_object(tina.MeshModel('assets/shadow.obj')) 9 | scene.lighting.skybox = tina.Atomsphere() 10 | 11 | gui = ti.GUI('sky', scene.res) 12 | 13 | scene.update() 14 | while gui.running: 15 | scene.input(gui) 16 | scene.render() 17 | gui.set_image(scene.img) 18 | gui.show() 19 | -------------------------------------------------------------------------------- /taichi_three/tests/ssao.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.cuda) 5 | 6 | scene = tina.Scene(smoothing=True, ssao=True)#, ibl=True) 7 | #model = tina.MeshModel('assets/bunny.obj') 8 | model = tina.MeshModel('assets/monkey.obj') 9 | #model = tina.MeshModel('/home/bate/Documents/testssao.obj') 10 | scene.add_object(model) 11 | 12 | if not scene.ibl: 13 | scene.lighting.clear_lights() 14 | scene.lighting.set_ambient_light([1, 1, 1]) 15 | 16 | gui = ti.GUI(res=scene.res) 17 | 18 | radius = gui.slider('radius', 0, 2, 0.01) 19 | thresh = gui.slider('thresh', 0, 1, 0.01) 20 | factor = gui.slider('factor', 0, 2, 0.01) 21 | radius.value = 0.2 22 | thresh.value = 0.0 23 | factor.value = 1.0 24 | 25 | while gui.running: 26 | scene.ssao.radius[None] = radius.value 27 | scene.ssao.thresh[None] = thresh.value 28 | scene.ssao.factor[None] = factor.value 29 | 30 | scene.input(gui) 31 | scene.render() 32 | gui.set_image(scene.img) 33 | gui.show() 34 | -------------------------------------------------------------------------------- /taichi_three/tests/ssr.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | ti.init(ti.gpu) 5 | 6 | 7 | scene = tina.Scene((640, 480), smoothing=True, ssr=True, taa=True) 8 | monkey_material = tina.PBR(metallic=0.0, roughness=0.4) 9 | monkey = tina.MeshModel('assets/monkey.obj') 10 | scene.add_object(monkey, monkey_material) 11 | 12 | param_metallic = tina.Param() 13 | param_roughness = tina.Param() 14 | plane_material = tina.PBR(metallic=param_metallic, roughness=param_roughness) 15 | plane = tina.MeshTransform(tina.MeshGrid(32), 16 | tina.scale(2) @ tina.eularXYZ([-ti.pi / 2, 0, 0])) 17 | scene.add_object(plane, plane_material) 18 | 19 | gui = ti.GUI(res=scene.res) 20 | nsteps = gui.slider('nsteps', 1, 128, 1) 21 | nsamples = gui.slider('nsamples', 1, 128, 1) 22 | stepsize = gui.slider('stepsize', 0, 32, 0.1) 23 | tolerance = gui.slider('tolerance', 0, 64, 0.1) 24 | blurring = gui.slider('blurring', 1, 8, 1) 25 | metallic = gui.slider('metallic', 0, 1, 0.01) 26 | roughness = gui.slider('roughness', 0, 1, 0.01) 27 | nsteps.value = 64 28 | nsamples.value = 12 29 | blurring.value = 4 30 | stepsize.value = 2 31 | tolerance.value = 15 32 | metallic.value = 1.0 33 | roughness.value = 0.0 34 | 35 | while gui.running: 36 | scene.ssr.nsteps[None] = int(nsteps.value) 37 | scene.ssr.nsamples[None] = int(nsamples.value) 38 | scene.ssr.blurring[None] = int(blurring.value) 39 | scene.ssr.stepsize[None] = stepsize.value 40 | scene.ssr.tolerance[None] = tolerance.value 41 | param_metallic.value[None] = metallic.value 42 | param_roughness.value[None] = roughness.value 43 | scene.input(gui) 44 | scene.render() 45 | gui.set_image(scene.img) 46 | gui.show() 47 | -------------------------------------------------------------------------------- /taichi_three/tests/volume.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | dens = np.load('assets/smoke.npy') 8 | scene = tina.Scene(N=dens.shape[0], density=16)#, taa=True) 9 | volume = tina.SimpleVolume(N=dens.shape[0]) 10 | #model = tina.MeshModel('assets/monkey.obj') 11 | #scene.add_object(model, tina.CookTorrance(metallic=0.8)) 12 | scene.add_object(volume) 13 | 14 | gui = ti.GUI('volume', scene.res) 15 | 16 | volume.set_volume_density(dens) 17 | while gui.running: 18 | scene.input(gui) 19 | scene.render() 20 | gui.set_image(scene.img) 21 | gui.show() 22 | -------------------------------------------------------------------------------- /taichi_three/tests/voxl.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | dens = np.load('assets/smoke.npy')[::1, ::1, ::1] 8 | scene = tina.PTScene() 9 | scene.engine.skybox = tina.Atomsphere() 10 | volume = tina.VolumeScale(tina.SimpleVolume(N=dens.shape[0]), scale=5) 11 | scene.add_object(tina.MeshModel('assets/monkey.obj')) 12 | g = tina.Param(float, initial=0.76) 13 | scene.add_object(volume, tina.HenyeyGreenstein(g=g)) 14 | #scene.add_object(tina.MeshTransform(tina.MeshModel('assets/plane.obj'), 15 | # tina.translate([0, 0, 4]) @ tina.eularXYZ([ti.pi / 2, 0, 0])), 16 | # tina.Emission() * 4) 17 | 18 | gui = ti.GUI('volume', scene.res) 19 | g.make_slider(gui, 'g', -1, 1, 0.01) 20 | 21 | volume.set_volume_density(dens) 22 | scene.update() 23 | while gui.running: 24 | scene.input(gui) 25 | scene.render()#nsteps=32) 26 | gui.set_image(scene.img) 27 | gui.show() 28 | -------------------------------------------------------------------------------- /taichi_three/tests/wave.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | ti.init(ti.gpu) 6 | 7 | scene = tina.PTScene(smoothing=True) 8 | #scene.lighting.skybox = tina.Atomsphere() 9 | #scene.lighting.skybox = tina.Skybox('assets/grass.jpg') 10 | 11 | scene.add_object(tina.PrimitiveMesh.sphere(), tina.Glass()) 12 | scene.add_object(tina.MeshTransform(tina.MeshModel('assets/cube.obj'), 13 | tina.translate([0, -3, 0]) @ tina.scale(2)), tina.Lambert()) 14 | 15 | scene.lighting.set_lights(np.array([ 16 | [0, 5, 0], 17 | ])) 18 | scene.lighting.set_light_colors(np.array([ 19 | [16, 16, 16], 20 | ])) 21 | scene.lighting.set_light_radii(np.array([ 22 | 0.01, 23 | ])) 24 | 25 | gui = ti.GUI('wave', scene.res) 26 | 27 | scene.update() 28 | while gui.running: 29 | if scene.input(gui): 30 | scene.clear() 31 | scene.render_light(nsteps=6) 32 | scene.render(nsteps=6) 33 | gui.set_image(scene.img) 34 | gui.show() 35 | 36 | #tina.pfmwrite('/tmp/a.pfm', scene.img) 37 | -------------------------------------------------------------------------------- /taichi_three/tests/wire.py: -------------------------------------------------------------------------------- 1 | # In this episode, you'll learn how to render a wireframe model 2 | # 3 | # This tutorial is based on docs/monkey.py, make sure you check that first 4 | 5 | import taichi as ti 6 | import numpy as np 7 | import tina 8 | 9 | scene = tina.Scene() 10 | 11 | # load the monkey using `tina.MeshModel` node (`tina.SimpleMesh` works too): 12 | model = tina.SimpleMesh(npolygon=2) 13 | scene.add_object(model) 14 | 15 | 16 | gui = ti.GUI('wireframe') 17 | 18 | while gui.running: 19 | scene.input(gui) 20 | a = gui.frame * 0.03 21 | verts = np.array([ 22 | [[0, 0, 0], [np.cos(a), np.sin(a), 0]], 23 | ], dtype=np.float32) 24 | model.set_face_verts(verts) 25 | scene.render() 26 | gui.set_image(scene.img) 27 | gui.show() 28 | -------------------------------------------------------------------------------- /taichi_three/tina/__init__.py: -------------------------------------------------------------------------------- 1 | __version__ = (0, 1, 1) 2 | __author__ = 'archibate <1931127624@qq.com>' 3 | __license__ = 'MIT' 4 | 5 | print('[Tina] version', '.'.join(map(str, __version__))) 6 | 7 | 8 | def require_version(*ver): 9 | def tos(ver): 10 | return '.'.join(map(str, ver)) 11 | 12 | msg = f'This program requires Tina version {tos(ver)} to work.\n' 13 | msg += f'However your installed Tina version is {tos(__version__)}.\n' 14 | msg += f'Try run `pip install taichi-tina=={tos(ver)}` to upgrade/downgrade.' 15 | if __version__ > ver: 16 | print(msg) 17 | elif __version__ < ver: 18 | raise RuntimeError(msg) 19 | 20 | 21 | from .lazimp import * 22 | from .hacker import * 23 | from .common import * 24 | from .advans import * 25 | 26 | if __import__('tina').lazyguard: 27 | from .shield import * 28 | from .util import * 29 | from .matr import * 30 | from .core import * 31 | from .path import * 32 | from .assimp import * 33 | from .mesh import * 34 | from .pars import * 35 | from .voxl import * 36 | from .scene import * 37 | from .skybox import * 38 | from .random import * 39 | from .probe import * 40 | from .postp import * 41 | -------------------------------------------------------------------------------- /taichi_three/tina/__main__.py: -------------------------------------------------------------------------------- 1 | def main(*args): 2 | cmd, *args = args 3 | if cmd == 'mesh': 4 | from .cli.mesh import main 5 | return main(*args) 6 | if cmd == 'volume': 7 | from .cli.volume import main 8 | return main(*args) 9 | if cmd == 'particles': 10 | from .cli.particles import main 11 | return main(*args) 12 | else: 13 | print('bad command:', cmd) 14 | exit(1) 15 | 16 | 17 | if __name__ == '__main__': 18 | import sys 19 | main(*sys.argv[1:]) 20 | -------------------------------------------------------------------------------- /taichi_three/tina/advans.py: -------------------------------------------------------------------------------- 1 | from .common import * 2 | 3 | 4 | inf = 1e6 5 | eps = 1e-6 6 | 7 | 8 | def texture_as_field(filename): 9 | if isinstance(filename, str): 10 | img_np = ti.imread(filename) 11 | else: 12 | img_np = np.array(filename) 13 | if img_np.dtype == np.uint8: 14 | img_np = np.float32(img_np / 255) 15 | 16 | if len(img_np.shape) == 3: 17 | img = ti.Vector.field(img_np.shape[2], float, img_np.shape[:2]) 18 | img._dense_shape = img_np.shape[:2] 19 | else: 20 | assert len(img_np.shape) == 2 21 | img = ti.field(float, img_np.shape) 22 | img._dense_shape = img_np.shape 23 | 24 | @ti.materialize_callback 25 | def init_texture(): 26 | img.from_numpy(img_np) 27 | 28 | return img 29 | 30 | 31 | @ti.pyfunc 32 | def aces_tonemap(color): 33 | # https://zhuanlan.zhihu.com/p/21983679 34 | return color * (2.51 * color + 0.03) / (color * (2.43 * color + 0.59) + 0.14) 35 | 36 | 37 | @ti.pyfunc 38 | def _film(x): 39 | # https://zhuanlan.zhihu.com/p/21983679 40 | A = 0.22 41 | B = 0.30 42 | C = 0.10 43 | D = 0.20 44 | E = 0.01 45 | F = 0.30 46 | 47 | return (x * (A * x + C * B) + D * E) / (x * (A * x + B) + D * F) - E / F 48 | 49 | 50 | @ti.pyfunc 51 | def film_tonemap(color): 52 | return _film(1.6 * color) / _film(11.2) 53 | 54 | 55 | @ti.pyfunc 56 | def ce_tonemap(color): 57 | if ti.static(ti.inside_kernel()): 58 | return 1 - ti.exp(-color) 59 | else: 60 | return 1 - np.exp(-color) 61 | 62 | 63 | @ti.pyfunc 64 | def ce_untonemap(color): 65 | if ti.static(ti.inside_kernel()): 66 | # 255 / 256 ~= 0.996 67 | return -ti.log(1 - clamp(color * 0.996, 0, 1)) 68 | else: 69 | return -np.log(1 - np.clip(color * 0.996, 0, 1)) 70 | 71 | 72 | @eval('lambda x: x()') 73 | def get_tonemap_image(): 74 | @ti.kernel 75 | def _get_image(out: ti.ext_arr(), img: ti.template(), tonemap: ti.template()): 76 | for I in ti.grouped(img): 77 | val = tonemap(img[I]) 78 | if ti.static(isinstance(val, ti.Matrix)): 79 | for k in ti.static(range(val.n)): 80 | out[I, k] = val[k] 81 | else: 82 | out[I] = val 83 | 84 | 85 | def get_image(img, tonemap=lambda x: x): 86 | shape = img.shape 87 | if isinstance(img, ti.Matrix): 88 | shape = shape + (img.n,) 89 | out = np.empty(shape) 90 | _get_image(out, img, tonemap) 91 | return out 92 | 93 | return get_image 94 | 95 | 96 | @ti.func 97 | def tangentspace(nrm, up=V(233., 666., 512.)): 98 | bitan = nrm.cross(up).normalized() 99 | tan = bitan.cross(nrm) 100 | return ti.Matrix.cols([tan, bitan, nrm]) 101 | 102 | 103 | 104 | @ti.func 105 | def spherical(h, p): 106 | unit = V(ti.cos(p * ti.tau), ti.sin(p * ti.tau)) 107 | dir = V23(ti.sqrt(max(0, 1 - h**2)) * unit, h) 108 | return dir 109 | 110 | 111 | @ti.func 112 | def unspherical(dir): 113 | p = ti.atan2(dir.y, dir.x) / ti.tau 114 | return dir.z, p % 1 115 | -------------------------------------------------------------------------------- /taichi_three/tina/assimp/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .obj import * 3 | from .gltf import * 4 | from .tet import * 5 | from .pfm import * 6 | -------------------------------------------------------------------------------- /taichi_three/tina/assimp/pfm.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sys 3 | 4 | def pfmwrite(path, im): 5 | im = im.swapaxes(0, 1) 6 | scale = max(1e-10, -im.min(), im.max()) 7 | h, w = im.shape[:2] 8 | with open(path, 'wb') as f: 9 | f.write(b'PF\n' if len(im.shape) >= 3 else b'Pf\n') 10 | f.write(f'{w} {h}\n'.encode()) 11 | f.write(f'{scale if sys.byteorder == "big" else -scale}\n'.encode()) 12 | f.write((im / scale).astype(np.float32).tobytes()) 13 | -------------------------------------------------------------------------------- /taichi_three/tina/cli/__init__.py: -------------------------------------------------------------------------------- 1 | from . import mesh, volume, particles -------------------------------------------------------------------------------- /taichi_three/tina/cli/mesh.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import tina 3 | 4 | 5 | def main(filename): 6 | ti.init(ti.gpu) 7 | 8 | obj = tina.readobj(filename, scale='auto') 9 | scene = tina.Scene((1024, 768), maxfaces=len(obj['f']), smoothing=True) 10 | model = tina.MeshModel(obj) 11 | scene.add_object(model) 12 | 13 | gui = ti.GUI('mesh', scene.res, fast_gui=True) 14 | while gui.running: 15 | scene.input(gui) 16 | scene.render() 17 | gui.set_image(scene.img) 18 | gui.show() 19 | 20 | 21 | if __name__ == '__main__': 22 | import sys 23 | main(*sys.argv[1:]) 24 | -------------------------------------------------------------------------------- /taichi_three/tina/cli/particles.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | 6 | def main(filename, radius=0.05): 7 | ti.init(ti.gpu) 8 | 9 | pos = np.load(filename).astype(np.float32) 10 | scene = tina.Scene((1024, 768)) 11 | pars = tina.SimpleParticles(maxpars=len(pos)) 12 | scene.add_object(pars) 13 | 14 | gui = ti.GUI('particles', scene.res, fast_gui=True) 15 | pars.set_particles(pos) 16 | pars.set_particle_radii(np.ones(len(pos), dtype=np.float32) * radius) 17 | while gui.running: 18 | scene.input(gui) 19 | scene.render() 20 | gui.set_image(scene.img) 21 | gui.show() 22 | 23 | 24 | if __name__ == '__main__': 25 | import sys 26 | main(*sys.argv[1:]) 27 | -------------------------------------------------------------------------------- /taichi_three/tina/cli/volume.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import tina 4 | 5 | 6 | def main(filename, density=16): 7 | ti.init(ti.gpu) 8 | 9 | dens = np.load(filename).astype(np.float32) 10 | scene = tina.Scene((1024, 768), N=dens.shape[0], taa=True, density=density) 11 | volume = tina.SimpleVolume(N=dens.shape[0]) 12 | scene.add_object(volume) 13 | 14 | gui = ti.GUI('volume', scene.res, fast_gui=True) 15 | volume.set_volume_density(dens) 16 | while gui.running: 17 | scene.input(gui) 18 | scene.render() 19 | gui.set_image(scene.img) 20 | gui.show() 21 | 22 | 23 | 24 | if __name__ == '__main__': 25 | import sys 26 | main(*sys.argv[1:]) 27 | -------------------------------------------------------------------------------- /taichi_three/tina/core/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .engine import * 3 | from .triangle import * 4 | from .particle import * 5 | from .wireframe import * 6 | from .volume import * 7 | from .shader import * 8 | from .lighting import * 9 | from .particle_sdf import * -------------------------------------------------------------------------------- /taichi_three/tina/core/engine.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class Engine: 6 | def __init__(self, res_x=512, res_y=None): 7 | if res_y is None: 8 | res_y = res_x 9 | self.res = tovector((res_x, res_y) if isinstance(res_x, int) else res_x) 10 | 11 | self.depth = ti.field(float, self.res) 12 | self.maxdepth = 2**25 13 | 14 | self.W2V = ti.Matrix.field(4, 4, float, ()) 15 | self.V2W = ti.Matrix.field(4, 4, float, ()) 16 | 17 | self.bias = ti.Vector.field(2, float, ()) 18 | 19 | @ti.materialize_callback 20 | @ti.kernel 21 | def init_engine(): 22 | self.W2V[None] = ti.Matrix.identity(float, 4) 23 | self.W2V[None][2, 2] = -1 24 | self.V2W[None] = ti.Matrix.identity(float, 4) 25 | self.V2W[None][2, 2] = -1 26 | self.bias[None] = [0.5, 0.5] 27 | 28 | ti.materialize_callback(self.clear_depth) 29 | 30 | @ti.kernel 31 | def randomize_bias(self, center: ti.template()): 32 | if ti.static(center): 33 | self.bias[None] = [0.5, 0.5] 34 | else: 35 | #r = ti.sqrt(ti.random()) 36 | #a = ti.random() * ti.tau 37 | #x, y = r * ti.cos(a) * 0.5 + 0.5, r * ti.sin(a) * 0.5 + 0.5 38 | x, y = ti.random(), ti.random() 39 | self.bias[None] = [x, y] 40 | 41 | @ti.kernel 42 | def render_background(self, shader: ti.template()): 43 | for P in ti.grouped(ti.ndrange(*self.res)): 44 | if self.depth[P] >= self.maxdepth: 45 | uv = (float(P) + self.bias[None]) / self.res * 2 - 1 46 | ro = mapply_pos(self.V2W[None], V(uv.x, uv.y, -1.0)) 47 | ro1 = mapply_pos(self.V2W[None], V(uv.x, uv.y, +1.0)) 48 | rd = (ro1 - ro).normalized() 49 | shader.shade_background(P, rd) 50 | 51 | @ti.func 52 | def to_viewspace(self, p): 53 | return mapply_pos(self.W2V[None], p) 54 | 55 | @ti.pyfunc 56 | def from_viewspace(self, p): 57 | return mapply_pos(self.V2W[None], p) 58 | 59 | @ti.func 60 | def to_viewport(self, p): 61 | return (p.xy * 0.5 + 0.5) * self.res 62 | 63 | @ti.func 64 | def from_viewport(self, p): 65 | return p / self.res * 2 - 1 66 | 67 | @ti.kernel 68 | def clear_depth(self): 69 | for P in ti.grouped(self.depth): 70 | self.depth[P] = (self.maxdepth << 4) 71 | 72 | def set_camera(self, view, proj): 73 | W2V = proj @ view 74 | V2W = np.linalg.inv(W2V) 75 | self.W2V.from_numpy(np.array(W2V, dtype=np.float32)) 76 | self.V2W.from_numpy(np.array(V2W, dtype=np.float32)) 77 | -------------------------------------------------------------------------------- /taichi_three/tina/core/material.py: -------------------------------------------------------------------------------- 1 | '''mockup''' 2 | -------------------------------------------------------------------------------- /taichi_three/tina/inject.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import functools 3 | import atexit 4 | import time 5 | 6 | 7 | _cbs = [] 8 | 9 | 10 | @atexit.register 11 | def _exit_callback(): 12 | for func in _cbs: 13 | func() 14 | 15 | 16 | def _inject(module, name, enable=True): 17 | if not enable: 18 | return lambda x: x 19 | 20 | def decorator(hook): 21 | func = getattr(module, name) 22 | if hasattr(func, '_injected'): 23 | func = func._injected 24 | 25 | @functools.wraps(func) 26 | def wrapped(*args, **kwargs): 27 | _taichi_skip_traceback = 1 28 | clb = hook(*args, **kwargs) 29 | ret = func(*args, **kwargs) 30 | if clb is not None: 31 | clb(ret) 32 | return ret 33 | 34 | wrapped._injected = func 35 | setattr(module, name, wrapped) 36 | return hook 37 | 38 | return decorator 39 | 40 | 41 | @_inject(ti.Kernel, 'materialize') 42 | def _(self, key=None, args=None, arg_features=None): 43 | self._key = key 44 | 45 | if key is None: 46 | key = (self.func, 0) 47 | if not self.runtime.materialized: 48 | self.runtime.materialize() 49 | if key in self.compiled_functions: 50 | return 51 | grad_suffix = "" 52 | if self.is_grad: 53 | grad_suffix = "_grad" 54 | name = "{}_c{}_{}{}".format(self.func.__qualname__, 55 | self.kernel_counter, key[1], 56 | grad_suffix) 57 | self._kname[key] = name 58 | 59 | @_cbs.append 60 | def callback(): 61 | if self._profile.get(name): 62 | x = sorted(self._profile[name]) 63 | if len(x) % 2 == 0: 64 | dt = (x[len(x) // 2] + x[len(x) // 2 - 1]) / 2 65 | else: 66 | dt = x[len(x) // 2] 67 | print(f'[{max(x):8.05f} {dt:8.05f} {len(x):4d}] {name}') 68 | 69 | 70 | @_inject(ti.Kernel, '__call__') 71 | def _(self, *args, **kwargs): 72 | def callback(ret): 73 | t1 = time.time() 74 | dt = t1 - t0 75 | self._profile.setdefault(self._kname[self._key], []).append(dt) 76 | 77 | if not hasattr(self, '_profile'): 78 | self._profile = {} 79 | if not hasattr(self, '_kname'): 80 | self._kname = {} 81 | 82 | #ti.trace('calling [' + self.func.__qualname__ + ']') 83 | t0 = time.time() 84 | return callback 85 | 86 | 87 | print('[Tina] Taichi JIT injected!') 88 | 89 | __all__ = [] 90 | 91 | 92 | if __name__ == '__main__': 93 | @ti.data_oriented 94 | class ODOP: 95 | @ti.kernel 96 | def func(self): 97 | print(233) 98 | 99 | c = ODOP() 100 | c.func() 101 | exit(1) 102 | -------------------------------------------------------------------------------- /taichi_three/tina/matr/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .nodes import * 3 | from .material import * 4 | from .wavelen import * 5 | -------------------------------------------------------------------------------- /taichi_three/tina/matr/wavelen.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.func 5 | def blackbody(temp, wave): 6 | wave *= 1e-9 7 | HCC2 = 1.1910429723971884140794892e-29 8 | HKC = 1.438777085924334052222404423195819240925e-2 9 | return wave ** -5 * HCC2 / (ti.exp(HKC / (wave * temp)) - 1) 10 | 11 | 12 | @ti.func 13 | def wav_to_rgb(wav): 14 | # https://blog.csdn.net/tanmx219/article/details/91658415 15 | r = 0.0 16 | g = 0.0 17 | b = 0.0 18 | alpha = 0.0 19 | 20 | # 3, 3, 2.5 21 | if 380.0 <= wav < 440.0: 22 | # .5, 0, 1 23 | r = -1.0 * (wav - 440.0) / (440.0 - 380.0) 24 | g = 0.0 25 | b = 1.0 26 | elif 440.0 <= wav < 490.0: 27 | # 0, .5, 1 28 | r = 0.0 29 | g = (wav - 440.0) / (490.0 - 440.0) 30 | b = 1.0 31 | elif 490.0 <= wav < 510.0: 32 | # 0, 1, .5 33 | r = 0.0 34 | g = 1.0 35 | b = -1.0 * (wav - 510.0) / (510.0 - 490.0) 36 | elif 510.0 <= wav < 580.0: 37 | # .5, 1, 0 38 | r = (wav - 510.0) / (580.0 - 510.0) 39 | g = 1.0 40 | b = 0.0 41 | elif 580.0 <= wav < 645.0: 42 | # 1, .5, 0 43 | r = 1.0 44 | g = -1.0 * (wav - 645.0) / (645.0 - 580.0) 45 | b = 0.0 46 | elif 645.0 <= wav < 780.0: 47 | # 1, 0, 0 48 | r = 1.0 49 | g = 0.0 50 | b = 0.0 51 | else: 52 | r = 0.0 53 | g = 0.0 54 | b = 0.0 55 | 56 | # 在可见光谱的边缘处强度较低。 57 | if 380.0 <= wav < 420.0: 58 | alpha = 0.30 + 0.70 * (wav - 380.0) / (420.0 - 380.0) 59 | elif 420.0 <= wav < 701.0: 60 | alpha = 1.0 61 | elif 701.0 <= wav < 780.0: 62 | alpha = 0.30 + 0.70 * (780.0 - wav) / (780.0 - 700.0) 63 | else: 64 | alpha = 0.0 65 | 66 | r *= 1.0 / 0.43511572 67 | g *= 1.0 / 0.5000342 68 | b *= 3.0 / 2.5 / 0.45338875 69 | return V(r, g, b) * alpha 70 | 71 | 72 | @ti.func 73 | def rgb_at_wav(rgb, wav): 74 | if ti.static(not isinstance(rgb, ti.Matrix)): 75 | return rgb 76 | ret = 0.0 77 | r, g, b = rgb 78 | if 645 <= wav < 780: 79 | ret += r / (0.3041293 * 1.179588 * 1.0081447) 80 | if 500 <= wav < 540: 81 | ret += g / (0.3092271 * 1.0824629 * 0.9944099) 82 | if 420 <= wav < 460: 83 | ret += b / (0.3240771 * 1.1763208 * 1.0136887) 84 | return ret 85 | 86 | 87 | @ti.func 88 | def random_wav(cho): 89 | #cho = ti.random(int) % 6 90 | cho = cho % 6 91 | ret = ti.random() 92 | if cho == 0: 93 | ret = lerp(ret, 645, 780) 94 | elif cho == 1: 95 | ret = lerp(ret, 580, 645) 96 | elif cho == 2: 97 | ret = lerp(ret, 510, 580) 98 | elif cho == 3: 99 | ret = lerp(ret, 490, 510) 100 | elif cho == 4: 101 | ret = lerp(ret, 440, 490) 102 | elif cho == 5: 103 | ret = lerp(ret, 380, 440) 104 | return ret 105 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .grid import * 3 | from .cull import * 4 | from .norm import * 5 | from .wire import * 6 | from .trans import * 7 | from .model import * 8 | from .simple import * 9 | from .prim import * 10 | from .conn import * 11 | from .export import * 12 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/base.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class MeshEditBase: 6 | def __init__(self, mesh): 7 | self.mesh = mesh 8 | 9 | def __getattr__(self, attr): 10 | return getattr(self.mesh, attr) 11 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/conn.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class ConnectiveMesh: 6 | def __init__(self, maxfaces=MAX, maxverts=MAX, npolygon=3): 7 | self.faces = ti.Vector.field(npolygon, int, maxfaces) 8 | self.verts = ti.Vector.field(3, float, maxverts) 9 | self.coors = ti.Vector.field(2, float, maxverts) 10 | self.norms = ti.Vector.field(3, float, maxverts) 11 | self.nfaces = ti.field(int, ()) 12 | 13 | self.maxfaces = maxfaces 14 | self.maxverts = maxverts 15 | self.npolygon = npolygon 16 | 17 | def get_npolygon(self): 18 | return self.npolygon 19 | 20 | @ti.func 21 | def pre_compute(self): 22 | pass 23 | 24 | def get_max_vert_nindex(self): 25 | return self.maxverts 26 | 27 | @ti.func 28 | def get_indiced_vert(self, i): 29 | return self.verts[i] 30 | 31 | @ti.func 32 | def get_indiced_norm(self, i): 33 | return self.norms[i] 34 | 35 | @ti.func 36 | def get_indiced_coor(self, i): 37 | return self.coors[i] 38 | 39 | @ti.func 40 | def get_nfaces(self): 41 | return self.nfaces[None] 42 | 43 | @ti.func 44 | def get_face_vert_indices(self, n): 45 | i = self.faces[n][0] 46 | j = self.faces[n][1] 47 | k = self.faces[n][2] 48 | return i, j, k 49 | 50 | @ti.func 51 | def _get_face_props(self, prop, n): 52 | a = prop[self.faces[n][0]] 53 | b = prop[self.faces[n][1]] 54 | c = prop[self.faces[n][2]] 55 | return a, b, c 56 | 57 | @ti.func 58 | def get_face_verts(self, n): 59 | return self._get_face_props(self.verts, n) 60 | 61 | @ti.func 62 | def get_face_coors(self, n): 63 | return self._get_face_props(self.coors, n) 64 | 65 | @ti.func 66 | def get_face_norms(self, n): 67 | return self._get_face_props(self.norms, n) 68 | 69 | @ti.kernel 70 | def set_vertices(self, verts: ti.ext_arr()): 71 | nverts = min(verts.shape[0], self.verts.shape[0]) 72 | for i in range(nverts): 73 | for k in ti.static(range(3)): 74 | self.verts[i][k] = verts[i, k] 75 | 76 | @ti.kernel 77 | def set_vert_norms(self, norms: ti.ext_arr()): 78 | nverts = min(norms.shape[0], self.norms.shape[0]) 79 | for i in range(nverts): 80 | for k in ti.static(range(3)): 81 | self.norms[i][k] = norms[i, k] 82 | 83 | @ti.kernel 84 | def set_vert_coors(self, coors: ti.ext_arr()): 85 | nverts = min(coors.shape[0], self.coors.shape[0]) 86 | for i in range(nverts): 87 | for k in ti.static(range(3)): 88 | self.coors[i][k] = coors[i, k] 89 | 90 | @ti.kernel 91 | def set_faces(self, faces: ti.ext_arr()): 92 | self.nfaces[None] = min(faces.shape[0], self.faces.shape[0]) 93 | for i in range(self.nfaces[None]): 94 | for k in ti.static(range(self.npolygon)): 95 | self.faces[i][k] = faces[i, k] 96 | 97 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/cull.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | from .base import MeshEditBase 3 | 4 | 5 | class MeshFlipCulling(MeshEditBase): 6 | @ti.func 7 | def _flip(self, props: ti.template()): 8 | props_rev = list(reversed(props)) 9 | for i, prop in ti.static(enumerate(props_rev)): 10 | props[i] = prop 11 | 12 | @ti.func 13 | def get_face_verts(self, n): 14 | verts = self.mesh.get_face_verts(n) 15 | self._flip(verts) 16 | return verts 17 | 18 | @ti.func 19 | def get_face_norms(self, n): 20 | norms = self.mesh.get_face_norms(n) 21 | self._flip(norms) 22 | return norms 23 | 24 | @ti.func 25 | def get_face_coors(self, n): 26 | coors = self.mesh.get_face_coors(n) 27 | self._flip(coors) 28 | return coors 29 | 30 | 31 | class MeshNoCulling(MeshFlipCulling): 32 | @ti.func 33 | def get_nfaces(self): 34 | return self.mesh.get_nfaces() * 2 35 | 36 | @ti.func 37 | def get_face_verts(self, n): 38 | verts = self.mesh.get_face_verts(n // 2) 39 | if n % 2 != 0: 40 | self._flip(verts) 41 | return verts 42 | 43 | @ti.func 44 | def get_face_norms(self, n): 45 | norms = self.mesh.get_face_norms(n // 2) 46 | if n % 2 != 0: 47 | self._flip(norms) 48 | for i, norm in ti.static(enumerate(norms)): 49 | norms[i] = -norm 50 | return norms 51 | 52 | @ti.func 53 | def get_face_coors(self, n): 54 | coors = self.mesh.get_face_coors(n // 2) 55 | if n % 2 != 0: 56 | self._flip(coors) 57 | return coors 58 | 59 | 60 | class MeshFlipNormal(MeshEditBase): 61 | @ti.func 62 | def get_face_norms(self, n): 63 | norms = self.mesh.get_face_norms(n) 64 | for i, norm in ti.static(enumerate(norms)): 65 | norms[i] = -norm 66 | return norms 67 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/grid.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class MeshGrid: 6 | def __init__(self, res, as_quad=False): 7 | if isinstance(res, int): 8 | res = res, res 9 | self.res = V(*res) 10 | self.pos = ti.Vector.field(3, float, self.res) 11 | self.nrm = ti.Vector.field(3, float, self.res) 12 | self.tex = ti.Vector.field(2, float, self.res) 13 | self.as_quad = as_quad 14 | 15 | @ti.materialize_callback 16 | @ti.kernel 17 | def init_pos(): 18 | for I in ti.grouped(self.pos): 19 | u, v = I / (self.res - 1) 20 | self.tex[I] = V(u, v) 21 | self.pos[I] = V(u * 2 - 1, v * 2 - 1, 0) 22 | 23 | def get_npolygon(self): 24 | return 4 if self.as_quad else 3 25 | 26 | @ti.func 27 | def pre_compute(self): 28 | for i, j in self.pos: 29 | i2 = max(i - 1, 0) 30 | j2 = max(j - 1, 0) 31 | i1 = min(i + 1, self.res.x - 1) 32 | j1 = min(j + 1, self.res.y - 1) 33 | dy = self.pos[i, j1] - self.pos[i, j2] 34 | dx = self.pos[i1, j] - self.pos[i2, j] 35 | self.nrm[i, j] = dx.cross(dy).normalized() 36 | 37 | @ti.func 38 | def get_nfaces(self): 39 | n = (self.res.x - 1) * (self.res.y - 1) 40 | if ti.static(self.as_quad): 41 | return n 42 | else: 43 | return n * 2 44 | 45 | @ti.func 46 | def _get_face_props(self, prop, n): 47 | stride = self.res.x - 1 48 | m = n 49 | if ti.static(not self.as_quad): 50 | m = n // 2 51 | i, j = V(m // stride, m % stride) 52 | a, b = prop[i, j], prop[i + 1, j] 53 | c, d = prop[i + 1, j + 1], prop[i, j + 1] 54 | if ti.static(self.as_quad): 55 | return a, b, c, d 56 | if n % 2 != 0: 57 | a, b, c = a, c, d 58 | return a, b, c 59 | 60 | @ti.func 61 | def get_face_verts(self, n): 62 | return self._get_face_props(self.pos, n) 63 | 64 | @ti.func 65 | def get_face_norms(self, n): 66 | return self._get_face_props(self.nrm, n) 67 | 68 | @ti.func 69 | def get_face_coors(self, n): 70 | return self._get_face_props(self.tex, n) 71 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/model.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class MeshModel: 6 | def __init__(self, obj, *_, **__): 7 | ''' 8 | :param obj: (OBJ | str) the OBJ model as to be displayed, automatically invoke tina.readobj if a string is specified 9 | :return: (Mesh) the mesh object to add into scene 10 | ''' 11 | 12 | if isinstance(obj, str): 13 | obj = tina.readobj(obj, *_, **__) 14 | 15 | self.faces = ti.Matrix.field(3, 3, int, len(obj['f'])) 16 | self.verts = ti.Vector.field(3, float, len(obj['v'])) 17 | self.coors = ti.Vector.field(2, float, len(obj['vt'])) 18 | self.norms = ti.Vector.field(3, float, len(obj['vn'])) 19 | 20 | @ti.materialize_callback 21 | def init_mesh(): 22 | faces = obj['f'] 23 | if len(faces.shape) == 2: 24 | faces = np.stack([faces, faces, faces], axis=2) 25 | self.faces.from_numpy(faces.astype(np.uint32)) 26 | self.verts.from_numpy(obj['v']) 27 | self.coors.from_numpy(obj['vt']) 28 | self.norms.from_numpy(obj['vn']) 29 | 30 | self.maxfaces = len(obj['f']) 31 | self.maxverts = len(obj['v']) 32 | self.maxcoors = len(obj['vt']) 33 | self.maxnorms = len(obj['vn']) 34 | 35 | def get_npolygon(self): 36 | return 3 37 | 38 | @ti.func 39 | def pre_compute(self): 40 | pass 41 | 42 | def get_max_vert_nindex(self): 43 | return self.maxverts 44 | 45 | @ti.func 46 | def get_nfaces(self): 47 | return self.faces.shape[0] 48 | 49 | @ti.func 50 | def get_face_vert_indices(self, n): 51 | i = self.faces[n][0, 0] 52 | j = self.faces[n][1, 0] 53 | k = self.faces[n][2, 0] 54 | return i, j, k 55 | 56 | @ti.func 57 | def _get_face_props(self, prop, index: ti.template(), n): 58 | a = prop[self.faces[n][0, index]] 59 | b = prop[self.faces[n][1, index]] 60 | c = prop[self.faces[n][2, index]] 61 | return a, b, c 62 | 63 | @ti.func 64 | def get_face_verts(self, n): 65 | return self._get_face_props(self.verts, 0, n) 66 | 67 | @ti.func 68 | def get_face_coors(self, n): 69 | return self._get_face_props(self.coors, 1, n) 70 | 71 | @ti.func 72 | def get_face_norms(self, n): 73 | return self._get_face_props(self.norms, 2, n) 74 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/norm.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | from .base import MeshEditBase 3 | 4 | 5 | class MeshFlatNormal(MeshEditBase): 6 | @ti.func 7 | def _calc_norm(self, a, b, c): 8 | return (b - a).cross(c - a).normalized() 9 | 10 | @ti.func 11 | def get_face_norms(self, n): 12 | verts = self.mesh.get_face_verts(n) 13 | norm = self._calc_norm(verts[0], verts[1], verts[2]) 14 | return [norm for vert in verts] 15 | 16 | 17 | class MeshSmoothNormal(MeshEditBase): 18 | def __init__(self, mesh, cached=True): 19 | super().__init__(mesh) 20 | 21 | N = self.mesh.get_max_vert_nindex() 22 | self.norm = ti.Vector.field(3, float, N) 23 | 24 | self.cached = cached 25 | if self.cached: 26 | ti.materialize_callback(self.update_normal) 27 | 28 | @ti.func 29 | def pre_compute(self): 30 | self.mesh.pre_compute() 31 | if ti.static(not self.cached): 32 | self._update_normal() 33 | 34 | @ti.func 35 | def _update_normal(self): 36 | for i in self.norm: 37 | self.norm[i] = 0 38 | for n in range(self.mesh.get_nfaces()): 39 | i, j, k = self.mesh.get_face_vert_indices(n) 40 | a, b, c = self.mesh.get_face_verts(n) 41 | nrm = (b - a).cross(c - a).normalized() 42 | self.norm[i] += nrm 43 | self.norm[j] += nrm 44 | self.norm[k] += nrm 45 | for i in self.norm: 46 | self.norm[i] = self.norm[i].normalized() 47 | 48 | @ti.kernel 49 | def update_normal(self): 50 | self._update_normal() 51 | 52 | @ti.func 53 | def get_face_norms(self, n): 54 | i, j, k = self.mesh.get_face_vert_indices(n) 55 | return self.norm[i], self.norm[j], self.norm[k] 56 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/trans.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from ..common import * 4 | from .base import MeshEditBase 5 | 6 | 7 | class MeshTransform(MeshEditBase): 8 | def __init__(self, mesh, trans=None): 9 | super().__init__(mesh) 10 | 11 | self.trans = ti.Matrix.field(4, 4, float, ()) 12 | self.trans_normal = ti.Matrix.field(3, 3, float, ()) 13 | 14 | @ti.materialize_callback 15 | @ti.kernel 16 | def init_trans(): 17 | self.trans[None] = ti.Matrix.identity(float, 4) 18 | self.trans_normal[None] = ti.Matrix.identity(float, 3) 19 | 20 | if trans is not None: 21 | @ti.materialize_callback 22 | def init_trans_arg(): 23 | self.set_transform(trans) 24 | 25 | def set_transform(self, trans): 26 | trans_normal = np.transpose(np.linalg.inv(trans)) 27 | self.trans[None] = np.array(trans).tolist() 28 | self.trans_normal[None] = np.array(trans_normal).tolist() 29 | 30 | def get_verts(self): 31 | x = self.mesh.verts.to_numpy() 32 | x = np.concatenate((x, np.ones_like(x[:, :1])), -1) 33 | return (x @ self.trans.to_numpy().T)[:, :3] 34 | 35 | @ti.func 36 | def get_face_verts(self, n): 37 | verts = self.mesh.get_face_verts(n) 38 | for i, vert in ti.static(enumerate(verts)): 39 | verts[i] = mapply_pos(self.trans[None], vert) 40 | return verts 41 | 42 | @ti.func 43 | def get_face_norms(self, n): 44 | norms = self.mesh.get_face_norms(n) 45 | for i, norm in ti.static(enumerate(norms)): 46 | norms[i] = self.trans_normal[None] @ norm 47 | return norms 48 | -------------------------------------------------------------------------------- /taichi_three/tina/mesh/wire.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | from .base import MeshEditBase 3 | 4 | 5 | class MeshToWire(MeshEditBase): 6 | def __init__(self, mesh): 7 | super().__init__(mesh) 8 | if hasattr(self.mesh, 'get_npolygon'): 9 | self.src_npolygon = self.mesh.get_npolygon() 10 | else: 11 | self.src_npolygon = 3 12 | 13 | def get_npolygon(self): 14 | return 2 15 | 16 | @ti.func 17 | def get_nfaces(self): 18 | return self.mesh.get_nfaces() * self.src_npolygon 19 | 20 | @ti.func 21 | def _get_face_props(self, src_getter: ti.template(), n): 22 | props = src_getter(n // self.src_npolygon) 23 | ep1 = n % self.src_npolygon 24 | ep2 = (n + 1) % self.src_npolygon 25 | p1 = list_subscript(props, ep1) 26 | p2 = list_subscript(props, ep2) 27 | return p1, p2 28 | 29 | @ti.func 30 | def get_face_verts(self, n): 31 | return self._get_face_props(self.mesh.get_face_verts, n) 32 | 33 | @ti.func 34 | def get_face_coors(self, n): 35 | return self._get_face_props(self.mesh.get_face_coors, n) 36 | 37 | @ti.func 38 | def get_face_norms(self, n): 39 | return self._get_face_props(self.mesh.get_face_norms, n) 40 | -------------------------------------------------------------------------------- /taichi_three/tina/pars/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .simple import * 3 | from .trans import * 4 | -------------------------------------------------------------------------------- /taichi_three/tina/pars/base.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class ParsEditBase: 6 | def __init__(self, pars): 7 | self.pars = pars 8 | 9 | def __getattr__(self, attr): 10 | return getattr(self.pars, attr) 11 | -------------------------------------------------------------------------------- /taichi_three/tina/pars/export.py: -------------------------------------------------------------------------------- 1 | from ..advans import * 2 | 3 | 4 | @ti.kernel 5 | def _pre_compute_pars_npars(pars: ti.template()) -> int: 6 | pars.pre_compute() 7 | return pars.get_npars() 8 | 9 | 10 | @ti.kernel 11 | def _export_pars_verts(pars: ti.template(), out: ti.ext_arr()): 12 | for i in range(out.shape[0]): 13 | vert = pars.get_particle_position(i) 14 | for j, v in ti.static(enumerate(vert)): 15 | out[i, j] = v 16 | 17 | @ti.kernel 18 | def _export_pars_sizes(pars: ti.template(), out: ti.ext_arr()): 19 | for i in range(out.shape[0]): 20 | size = pars.get_particle_radius(i) 21 | out[i] = size 22 | 23 | @ti.kernel 24 | def _export_pars_colors(pars: ti.template(), out: ti.ext_arr()): 25 | for i in range(out.shape[0]): 26 | color = pars.get_particle_color(i) 27 | for j, c in ti.static(enumerate(color)): 28 | out[i, j] = c 29 | 30 | 31 | def export_simple_pars(pars): 32 | npars = _pre_compute_pars_npars(pars) 33 | ret = {} 34 | assert hasattr(pars, 'get_particle_position') 35 | out = np.empty((npars, 3), dtype=np.float32) 36 | _export_pars_verts(pars, out) 37 | ret['v'] = out 38 | if hasattr(pars, 'get_particle_radius'): 39 | out = np.empty(npars, dtype=np.float32) 40 | _export_pars_sizes(pars, out) 41 | ret['vr'] = out 42 | if hasattr(pars, 'get_particle_color'): 43 | out = np.empty((npars, 3), dtype=np.float32) 44 | _export_pars_colors(pars, out) 45 | ret['vc'] = out 46 | return ret 47 | -------------------------------------------------------------------------------- /taichi_three/tina/pars/simple.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class SimpleParticles: 6 | def __init__(self, maxpars=65536, radius=0.02): 7 | self.verts = ti.Vector.field(3, float, maxpars) 8 | self.sizes = ti.field(float, maxpars) 9 | self.colors = ti.Vector.field(3, float, maxpars) 10 | self.npars = ti.field(int, ()) 11 | 12 | @ti.materialize_callback 13 | def init_pars(): 14 | self.sizes.fill(radius) 15 | self.colors.fill(1) 16 | 17 | self.maxpars = maxpars 18 | 19 | @ti.func 20 | def pre_compute(self): 21 | pass 22 | 23 | @ti.func 24 | def get_npars(self): 25 | return min(self.npars[None], self.maxpars) 26 | 27 | @ti.func 28 | def get_particle_position(self, n): 29 | return self.verts[n] 30 | 31 | @ti.func 32 | def get_particle_radius(self, n): 33 | return self.sizes[n] 34 | 35 | @ti.func 36 | def get_particle_color(self, n): 37 | return self.colors[n] 38 | 39 | @ti.kernel 40 | def set_particles(self, verts: ti.ext_arr()): 41 | self.npars[None] = min(verts.shape[0], self.verts.shape[0]) 42 | for i in range(self.npars[None]): 43 | for k in ti.static(range(3)): 44 | self.verts[i][k] = verts[i, k] 45 | 46 | @ti.kernel 47 | def set_particle_radii(self, sizes: ti.ext_arr()): 48 | for i in range(self.npars[None]): 49 | self.sizes[i] = sizes[i] 50 | 51 | @ti.kernel 52 | def set_particle_colors(self, colors: ti.ext_arr()): 53 | for i in range(self.npars[None]): 54 | for k in ti.static(range(3)): 55 | self.colors[i][k] = colors[i, k] 56 | -------------------------------------------------------------------------------- /taichi_three/tina/pars/trans.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | from .base import ParsEditBase 3 | 4 | 5 | class ParsTransform(ParsEditBase): 6 | def __init__(self, pars): 7 | super().__init__(pars) 8 | 9 | self.trans = ti.Matrix.field(4, 4, float, ()) 10 | self.scale = ti.field(float, ()) 11 | 12 | @ti.materialize_callback 13 | @ti.kernel 14 | def init_trans(): 15 | self.trans[None] = ti.Matrix.identity(float, 4) 16 | self.scale[None] = 1 17 | 18 | def set_transform(self, trans, scale): 19 | self.trans[None] = np.array(trans).tolist() 20 | self.scale[None] = scale 21 | 22 | @ti.func 23 | def get_particle_position(self, n): 24 | vert = self.pars.get_particle_position(n) 25 | return mapply_pos(self.trans[None], vert) 26 | 27 | @ti.func 28 | def get_particle_radius(self, n): 29 | size = self.pars.get_particle_radius(n) 30 | return self.scale[None] * size 31 | -------------------------------------------------------------------------------- /taichi_three/tina/path/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .engine import * 3 | from .geometry import * 4 | from .particle import * 5 | from .triangle import * 6 | from .volume import * 7 | from .tree import * 8 | -------------------------------------------------------------------------------- /taichi_three/tina/path/volume.py: -------------------------------------------------------------------------------- 1 | from ..advans import * 2 | from .geometry import * 3 | 4 | 5 | @ti.data_oriented 6 | class VolumeTracer: 7 | is_dedicated_tracer = True 8 | 9 | def __init__(self, **extra_options): 10 | pass 11 | 12 | def clear_objects(self): 13 | pass 14 | 15 | def add_object(self, voxl, mtlid): 16 | self.voxl = voxl 17 | 18 | @ti.func 19 | def sample_density(self, pos): 20 | return self.voxl.sample_volume(pos) 21 | #return 2.6 if (pos // 0.5).sum() % 2 == 0 else 0.0 22 | 23 | @ti.func 24 | def hit(self, ro, rd, maxfar=inf): 25 | # travel volume 26 | bmin, bmax = self.voxl.get_bounding_box() 27 | near, far = tina.ray_aabb_hit(bmin, bmax, ro, rd) 28 | 29 | hitind = -1 30 | depth = inf 31 | if near <= far: 32 | near = max(near, 0) 33 | far = min(far, maxfar) 34 | 35 | t = near 36 | int_rho = 0.0 37 | ran = -ti.log(ti.random()) 38 | step = 0.01 39 | t += step * (.5 + .5 * ti.random()) 40 | while t < far: 41 | dt = min(step, far - t) 42 | pos = ro + t * rd 43 | rho = self.sample_density(pos) 44 | int_rho += rho * dt 45 | if ran < int_rho: # ti.random() > ti.exp(-Integrate rho*dt) 46 | depth = t 47 | hitind = 0 48 | break 49 | t += dt 50 | 51 | return depth, hitind, V(0., 0.) 52 | 53 | @ti.func 54 | def calc_geometry(self, near, ind, uv, ro, rd): 55 | nrm = V(0., 0., 0.) 56 | return nrm, V(0., 0.) 57 | 58 | @ti.func 59 | def sample_light_pos(self, org): 60 | return V(0., 0., 0.), 0, 0.0 61 | 62 | @ti.kernel 63 | def update_emission(self, mtltab: ti.template()): 64 | pass 65 | 66 | def update(self): 67 | pass 68 | 69 | @ti.func 70 | def get_material_id(self, ind): 71 | return 1 72 | -------------------------------------------------------------------------------- /taichi_three/tina/postp/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .tonemap import * 3 | from .denoise import * 4 | from .blooming import * 5 | from .fxaa import * 6 | from .ssao import * 7 | from .ssr import * 8 | -------------------------------------------------------------------------------- /taichi_three/tina/postp/blooming.py: -------------------------------------------------------------------------------- 1 | from ..advans import * 2 | 3 | 4 | @ti.data_oriented 5 | class Blooming: 6 | def __init__(self, res): 7 | self.res = tovector(res) 8 | self.img = ti.Vector.field(3, float, self.res // 2) 9 | self.tmp = ti.Vector.field(3, float, self.res // 2) 10 | self.gwei = ti.field(float, 128) 11 | self.thresh = ti.field(float, ()) 12 | self.factor = ti.field(float, ()) 13 | self.scale = ti.field(float, ()) 14 | self.sigma = ti.field(float, ()) 15 | self.radius = ti.field(int, ()) 16 | 17 | @ti.materialize_callback 18 | def init_params(): 19 | self.thresh[None] = 1 20 | self.factor[None] = 1 21 | self.radius[None] = min(*self.res) // 16 22 | self.sigma[None] = 1 23 | self.scale[None] = 0.25 24 | self.init_gwei() 25 | 26 | @ti.kernel 27 | def init_gwei(self): 28 | sum = -1.0 29 | radius = self.radius[None] 30 | sigma = self.sigma[None] 31 | for i in range(radius + 1): 32 | x = sigma * i / radius 33 | y = ti.exp(-x**2) 34 | self.gwei[i] = y 35 | sum += y * 2 36 | for i in range(radius + 1): 37 | self.gwei[i] /= sum 38 | 39 | @ti.func 40 | def filter(self, x): 41 | t = max(0, x - self.thresh[None]) 42 | t = 1 - 1 / (1 + self.scale[None] * t) 43 | return self.factor[None] * t 44 | 45 | @ti.kernel 46 | def apply(self, image: ti.template()): 47 | for I in ti.grouped(self.img): 48 | res = V(0., 0., 0.) 49 | for J in ti.grouped(ti.ndrange(2, 2)): 50 | res += self.filter(image[I * 2 + J]) 51 | self.img[I] = res / 4 52 | for I in ti.grouped(self.img): 53 | res = self.img[I] * self.gwei[0] 54 | for i in range(1, self.radius[None] + 1): 55 | val = self.img[max(0, I.x - i), I.y] 56 | val += self.img[min(self.img.shape[0] - 1, I.x + i), I.y] 57 | res += val * self.gwei[i] 58 | self.tmp[I] = res 59 | for I in ti.grouped(self.img): 60 | res = self.tmp[I] * self.gwei[0] 61 | for i in range(1, self.radius[None] + 1): 62 | dir = V(0, 1) 63 | val = self.tmp[I.x, max(0, I.y - i)] 64 | val += self.tmp[I.x, min(self.img.shape[1] - 1, I.y + i)] 65 | res += val * self.gwei[i] 66 | self.img[I] = res 67 | for I in ti.grouped(image): 68 | image[I] += bilerp(self.img, I / 2) 69 | -------------------------------------------------------------------------------- /taichi_three/tina/postp/fxaa.py: -------------------------------------------------------------------------------- 1 | from ..advans import * 2 | 3 | 4 | @ti.func 5 | def luminance(c): 6 | return V(0.2989, 0.587, 0.114).dot(c) 7 | 8 | 9 | # https://catlikecoding.com/unity/tutorials/advanced-rendering/fxaa/ 10 | @ti.data_oriented 11 | class FXAA: 12 | def __init__(self, res): 13 | self.res = res 14 | self.lumi = ti.field(float, self.res) 15 | self.img = ti.Vector.field(3, float, self.res) 16 | self.abs_thresh = ti.field(float, ()) 17 | self.rel_thresh = ti.field(float, ()) 18 | self.factor = ti.field(float, ()) 19 | 20 | @ti.materialize_callback 21 | def init_params(): 22 | self.abs_thresh[None] = 0.0625 23 | self.rel_thresh[None] = 0.063 24 | self.factor[None] = 1 25 | 26 | @ti.kernel 27 | def apply(self, image: ti.template()): 28 | for I in ti.grouped(image): 29 | self.lumi[I] = clamp(luminance(image[I]), 0, 1) 30 | self.img[I] = image[I] 31 | for I in ti.grouped(image): 32 | m = self.lumi[I] 33 | n = self.lumi[I + V(0, 1)] 34 | e = self.lumi[I + V(1, 0)] 35 | s = self.lumi[I + V(0, -1)] 36 | w = self.lumi[I + V(-1, 0)] 37 | ne = self.lumi[I + V(1, 1)] 38 | nw = self.lumi[I + V(-1, 1)] 39 | se = self.lumi[I + V(1, -1)] 40 | sw = self.lumi[I + V(-1, -1)] 41 | hi = max(m, n, e, s, w) 42 | lo = min(m, n, e, s, w) 43 | c = hi - lo 44 | if c < self.abs_thresh[None] or c < self.rel_thresh[None] * hi: 45 | continue 46 | filt = 2 * (n + e + s + w) 47 | filt += ne + nw + se + sw 48 | filt = abs(filt / 12 - m) 49 | filt = clamp(filt / c, 0, 1) 50 | blend = smoothstep(filt, 0, 1)**2 * self.factor[None] 51 | hori = abs(n + s - 2 * m) * 2 52 | hori += abs(ne + se - 2 * e) 53 | hori += abs(nw + sw - 2 * w) 54 | vert = abs(e + w - 2 * m) * 2 55 | vert += abs(ne + nw - 2 * n) 56 | vert += abs(se + sw - 2 * s) 57 | is_hori = hori >= vert 58 | 59 | plumi = n if is_hori else e 60 | nlumi = s if is_hori else w 61 | pgrad = abs(plumi - m) 62 | ngrad = abs(nlumi - m) 63 | 64 | if pgrad < ngrad: 65 | blend = -blend 66 | dir = V(0, 1) if is_hori else V(1, 0) 67 | 68 | image[I] = bilerp(self.img, I + blend * dir) 69 | -------------------------------------------------------------------------------- /taichi_three/tina/postp/tonemap.py: -------------------------------------------------------------------------------- 1 | from ..advans import * 2 | 3 | 4 | @ti.data_oriented 5 | class ToneMapping: 6 | def __init__(self, res): 7 | self.res = res 8 | 9 | @ti.kernel 10 | def apply(self, image: ti.template()): 11 | for I in ti.grouped(image): 12 | image[I] = aces_tonemap(image[I]) 13 | -------------------------------------------------------------------------------- /taichi_three/tina/probe.py: -------------------------------------------------------------------------------- 1 | from .advans import * 2 | 3 | 4 | @ti.data_oriented 5 | class ProbeShader: 6 | def __init__(self, res): 7 | self.res = res 8 | self.elmid = ti.field(int, res) 9 | self.texcoord = ti.Vector.field(2, float, res) 10 | 11 | @ti.kernel 12 | def clear_buffer(self): 13 | for I in ti.grouped(self.elmid): 14 | self.elmid[I] = -1 15 | self.texcoord[I] *= 0 16 | 17 | @ti.func 18 | def shade_color(self, engine, P, p, f, pos, normal, texcoord, color): 19 | self.elmid[P] = f 20 | self.texcoord[P] = texcoord 21 | 22 | @ti.kernel 23 | def touch(self, callback: ti.template(), mx: float, my: float, rad: float): 24 | p = V(mx, my) * self.res 25 | bot, top = ifloor(p - rad), iceil(p + rad) 26 | for I in ti.grouped(ti.ndrange((bot.x, top.x + 1), (bot.y, top.y + 1))): 27 | r = (I - p).norm() 28 | if r > rad: 29 | continue 30 | if self.elmid[I] == -1: 31 | continue 32 | callback(self, I, r) 33 | -------------------------------------------------------------------------------- /taichi_three/tina/scene/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .raster import * 3 | from .tracer import * 4 | -------------------------------------------------------------------------------- /taichi_three/tina/shield.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | import functools 4 | import pickle 5 | 6 | 7 | def _getstate(self): 8 | return self._pickable_decl, self.to_numpy() 9 | 10 | 11 | def _Expr_setstate(self, state): 12 | (a, b), data = state 13 | 14 | field = ti.field(*a, **b) 15 | op = ti.impl.pytaichi.global_vars.pop() 16 | assert op is field, (op, field) 17 | self.__init__(field.ptr) 18 | ti.impl.pytaichi.global_vars.append(self) 19 | 20 | @ti.materialize_callback 21 | def restore_state(): 22 | self.from_numpy(data) 23 | 24 | 25 | def _Matrix_setstate(self, state): 26 | (a, b), data = state 27 | 28 | olen = len(ti.impl.pytaichi.global_vars) 29 | field = ti.Matrix.field(*a, **b) 30 | self.__init__([[field(i, j) for j in range(field.m)] for i in range(field.n)]) 31 | 32 | @ti.materialize_callback 33 | def restore_state(): 34 | self.from_numpy(data) 35 | 36 | 37 | ti.Expr.__getstate__ = _getstate 38 | ti.Matrix.__getstate__ = _getstate 39 | ti.Expr.__setstate__ = _Expr_setstate 40 | ti.Matrix.__setstate__ = _Matrix_setstate 41 | 42 | 43 | def _mock(foo): 44 | @functools.wraps(foo) 45 | def wrapped(*a, **b): 46 | ret = foo(*a, **b) 47 | if foo == ti.Matrix._Vector_field: 48 | a = (a[0], 1) + a[1:] 49 | ret._pickable_decl = a, b 50 | return ret 51 | 52 | return wrapped 53 | 54 | 55 | ti.field = _mock(ti.field) 56 | ti.Vector.field = _mock(ti.Vector.field) 57 | ti.Matrix.field = _mock(ti.Matrix.field) 58 | 59 | 60 | print('[Tina] Taichi fields pickle hacked') 61 | 62 | 63 | __all__ = [] 64 | -------------------------------------------------------------------------------- /taichi_three/tina/util/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .accumator import * 3 | from .control import * 4 | from .matrix import * 5 | from .mciso import * 6 | from .stack import * 7 | -------------------------------------------------------------------------------- /taichi_three/tina/util/accumator.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class Accumator: 6 | def __init__(self, res): 7 | self.img = ti.Vector.field(3, float, res) 8 | self.count = ti.field(int, ()) 9 | 10 | @ti.kernel 11 | def clear(self): 12 | self.count[None] = 0 13 | for I in ti.grouped(self.img): 14 | self.img[I] *= 0 15 | 16 | @ti.kernel 17 | def update(self, src: ti.template()): 18 | self.count[None] += 1 19 | inv_count = 1 / self.count[None] 20 | for I in ti.grouped(self.img): 21 | color = src[I] 22 | self.img[I] *= 1 - inv_count 23 | self.img[I] += color * inv_count 24 | -------------------------------------------------------------------------------- /taichi_three/tina/util/stack.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class Stack: 6 | def __init__(self, N_mt=512 * 512, N_len=32, field=None): 7 | if ti.cfg.arch == ti.cpu and ti.cfg.cpu_max_num_threads == 1 or ti.cfg.arch == ti.cc: 8 | N_mt = 1 9 | print('[Tina] Using', N_mt, 'threads') 10 | self.N_mt = N_mt 11 | self.N_len = N_len 12 | self.val = ti.field(int) if field is None else field 13 | self.blk1 = ti.root.dense(ti.i, N_mt) 14 | self.blk2 = self.blk1.dense(ti.j, N_len) 15 | self.blk2.place(self.val) 16 | self.len = ti.field(int, N_mt) 17 | 18 | @staticmethod 19 | def instance(): 20 | return Stack.g_stack 21 | 22 | @ti.func 23 | def __iter__(self): 24 | for i in range(self.N_mt): 25 | stack = self.Proxy(self, i) 26 | setattr(Stack, 'g_stack', stack) 27 | yield i 28 | delattr(Stack, 'g_stack') 29 | 30 | @ti.data_oriented 31 | class Proxy: 32 | def __init__(self, stack, mtid): 33 | self.stack = stack 34 | self.mtid = mtid 35 | 36 | def __getattr__(self, attr): 37 | return getattr(self.stack, attr) 38 | 39 | @ti.func 40 | def size(self): 41 | return self.len[self.mtid] 42 | 43 | @ti.func 44 | def clear(self): 45 | self.len[self.mtid] = 0 46 | 47 | @ti.func 48 | def push(self, val): 49 | l = self.len[self.mtid] 50 | self.val[self.mtid, l] = val 51 | self.len[self.mtid] = l + 1 52 | 53 | @ti.func 54 | def pop(self): 55 | l = self.len[self.mtid] 56 | val = self.val[self.mtid, l - 1] 57 | self.len[self.mtid] = l - 1 58 | return val 59 | -------------------------------------------------------------------------------- /taichi_three/tina/voxl/__init__.py: -------------------------------------------------------------------------------- 1 | if __import__('tina').lazyguard: 2 | from .simple import * 3 | from .trans import * 4 | from .scale import * 5 | -------------------------------------------------------------------------------- /taichi_three/tina/voxl/base.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class VoxlEditBase: 6 | def __init__(self, voxl): 7 | self.voxl = voxl 8 | 9 | def __getattr__(self, attr): 10 | return getattr(self.voxl, attr) 11 | -------------------------------------------------------------------------------- /taichi_three/tina/voxl/scale.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | from .base import VoxlEditBase 3 | 4 | 5 | class VolumeScale(VoxlEditBase): 6 | def __init__(self, voxl, scale=1): 7 | super().__init__(voxl) 8 | 9 | self.scale = ti.field(float, ()) 10 | 11 | @ti.materialize_callback 12 | def init_scale(): 13 | self.scale[None] = scale 14 | 15 | def set_scale(self, scale): 16 | self.scale[None] = scale 17 | 18 | @ti.func 19 | def sample_volume(self, pos): 20 | return self.voxl.sample_volume(pos) * self.scale[None] 21 | -------------------------------------------------------------------------------- /taichi_three/tina/voxl/simple.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | 3 | 4 | @ti.data_oriented 5 | class SimpleVolume: 6 | def __init__(self, N): 7 | self.N = N 8 | self.dens = ti.field(float, (N, N, N)) 9 | 10 | @ti.materialize_callback 11 | def init_pars(): 12 | self.dens.fill(1) 13 | 14 | def set_volume_density(self, dens): 15 | self.dens.from_numpy(dens) 16 | 17 | @ti.func 18 | def pre_compute(self): 19 | pass 20 | 21 | @ti.func 22 | def sample_volume(self, pos): 23 | ret = 0.0 24 | if all(-1 <= pos <= 1): 25 | ret = trilerp(self.dens, (pos * 0.5 + 0.5) * self.N) 26 | return ret 27 | 28 | @ti.func 29 | def sample_gradient(self, pos): 30 | ret = ti.Vector.zero(float, 3) 31 | for i in ti.static(range(3)): 32 | dir = U3(i) * 0.5 / self.N 33 | hi = self.sample_volume(pos + dir) 34 | lo = self.sample_volume(pos - dir) 35 | ret[i] = (hi - lo) / 2 36 | return ret 37 | 38 | @ti.func 39 | def get_transform(self): 40 | return ti.Matrix.identity(float, 4) 41 | 42 | def get_bounding_box(self): 43 | return V3(-1.), V3(1.) 44 | -------------------------------------------------------------------------------- /taichi_three/tina/voxl/trans.py: -------------------------------------------------------------------------------- 1 | from ..common import * 2 | from .base import VoxlEditBase 3 | 4 | 5 | class VolumeTransform(VoxlEditBase): 6 | def __init__(self, voxl): 7 | super().__init__(voxl) 8 | 9 | self.trans = ti.Matrix.field(4, 4, float, ()) 10 | self.inv_trans = ti.Matrix.field(4, 4, float, ()) 11 | 12 | @ti.materialize_callback 13 | @ti.kernel 14 | def init_trans(): 15 | self.trans[None] = ti.Matrix.identity(float, 4) 16 | self.inv_trans[None] = ti.Matrix.identity(float, 4) 17 | 18 | @ti.func 19 | def get_transform(self): 20 | return self.trans[None] @ self.voxl.get_transform() 21 | 22 | def set_transform(self, trans): 23 | self.trans[None] = np.array(trans).tolist() 24 | self.inv_trans[None] = np.linalg.inv(trans).tolist() 25 | -------------------------------------------------------------------------------- /taichi_three/view.py: -------------------------------------------------------------------------------- 1 | import taichi as ti 2 | import numpy as np 3 | from matplotlib import cm 4 | import os 5 | 6 | cmap = cm.get_cmap('magma') 7 | 8 | 9 | res = 256, 64, 64 10 | #res = 1024, 256, 1 11 | rho = ti.field(float, res) 12 | vel = ti.Vector.field(3, float, res) 13 | img = ti.field(float, (res[0], res[1])) 14 | 15 | 16 | def load(frame): 17 | path = f'/tmp/{frame:06d}.npz' 18 | if not os.path.exists(path): 19 | return False 20 | with np.load(path) as data: 21 | rho.from_numpy(data['rho']) 22 | vel.from_numpy(data['vel']) 23 | return True 24 | 25 | 26 | @ti.func 27 | def color(x, y, z): 28 | return vel[x, y, z].norm() * 5 29 | 30 | 31 | @ti.kernel 32 | def render(): 33 | for x, y in img: 34 | ret = 0.0 35 | cnt = 0 36 | for z in range(res[2] // 4, max(1, res[2] * 3 // 4)): 37 | ret += color(x, y, z) 38 | cnt += 1 39 | img[x, y] = ret / cnt 40 | 41 | 42 | frame = 0 43 | while load(frame): 44 | print('render for', frame) 45 | render() 46 | im = cmap(img.to_numpy()) 47 | ti.imshow(im) 48 | ti.imwrite(im, f'/tmp/{frame:06d}.png') 49 | frame += 1 50 | --------------------------------------------------------------------------------