Implement Rignet

rigging/rignet/README.md

@@ -0,0 +1,66 @@
+# RigNet
+
+## Input
+
+![Input](input.png)
+
+(File from https://github.com/zhan-xu/RigNet/blob/master/quick_start/17872_remesh.obj)
+
+## Output
+
+![Output](output_17872.png)
+
+## Setup
+
+The following commands installs pytorch1.7.1 and pytorch_geometric1.7.1.
+
+``` bash
+$ source setup.sh
+```
+
+For Windows,
+run `pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html -f https://pytorch-geometric.com/whl/torch-1.7.0+cpu.html`
+. Then you need to download Windows-compiled Rtree from [here](https://www.lfd.uci.edu/~gohlke/pythonlibs/#rtree), and
+install it by `pip install Rtree‑0.9.4‑cp37‑cp37m‑win_amd64.whl` (64-bit system)
+or `pip install Rtree‑0.9.4‑cp37‑cp37m‑win32.whl` (32-bit system).
+
+## Usage
+
+Automatically downloads the onnx and prototxt files on the first run. It is necessary to be connected to the Internet
+while downloading.
+
+For the sample object,
+
+``` bash
+$ python3 rignet.py
+```
+
+The predicted rigs are saved as *_rig.txt. Due to randomness, the results might be slightly different among each run.
+
+If you want to specify the input object, put the object file path after the `--input` option.  
+You can use `--savepath` option to change the name of the output file to save.
+
+```bash
+$ python3 rignet.py --input FILE_PATH --savepath SAVE_IMAGE_PATH
+```
+
+If you want to try your own models, remember to simplify the meshes so that the remeshed ones have vertices between 1K
+to 5K.
+
+You can specify bandwidth and threshold parameter.
+
+``` bash
+$ python3 rignet.py --threshold 0.75e-5 --bandwidth 0.045
+```
+
+## Reference
+
+- [RigNet](https://github.com/zhan-xu/RigNet)
+
+## Framework
+
+Pytorch
+
+## Model Format
+
+ONNX opset=11

rigging/rignet/binvox_rw.py

@@ -0,0 +1,183 @@
+#  Copyright (C) 2012 Daniel Maturana
+#  This file is part of binvox-rw-py.
+#
+#  binvox-rw-py is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  binvox-rw-py is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with binvox-rw-py. If not, see <http://www.gnu.org/licenses/>.
+#
+
+import numpy as np
+import struct
+
+
+class Voxels(object):
+    """ Holds a binvox model.
+    data is either a three-dimensional numpy boolean array (dense representation)
+    or a two-dimensional numpy float array (coordinate representation).
+
+    dims, translate and scale are the model metadata.
+
+    dims are the voxel dimensions, e.g. [32, 32, 32] for a 32x32x32 model.
+
+    scale and translate relate the voxels to the original model coordinates.
+
+    To translate voxel coordinates i, j, k to original coordinates x, y, z:
+
+    x_n = (i+.5)/dims[0]
+    y_n = (j+.5)/dims[1]
+    z_n = (k+.5)/dims[2]
+    x = scale*x_n + translate[0]
+    y = scale*y_n + translate[1]
+    z = scale*z_n + translate[2]
+
+    """
+
+    def __init__(self, data, dims, translate, scale, axis_order):
+        self.data = data
+        self.dims = dims
+        self.translate = translate
+        self.scale = scale
+        assert (axis_order in ('xzy', 'xyz'))
+        self.axis_order = axis_order
+
+    def clone(self):
+        data = self.data.copy()
+        dims = self.dims[:]
+        translate = self.translate[:]
+        return Voxels(data, dims, translate, self.scale, self.axis_order)
+
+    def write(self, fp):
+        write(self, fp)
+
+
+def read_header(fp):
+    """ Read binvox header. Mostly meant for internal use.
+    """
+    line = fp.readline().strip()
+    if not line.startswith(b'#binvox'):
+        raise IOError('Not a binvox file')
+    dims = list(map(int, fp.readline().strip().split(b' ')[1:]))
+    translate = list(map(float, fp.readline().strip().split(b' ')[1:]))
+    scale = list(map(float, fp.readline().strip().split(b' ')[1:]))[0]
+    _ = fp.readline()
+
+    return dims, translate, scale
+
+
+def read_as_3d_array(fp, fix_coords=True):
+    """ Read binary binvox format as array.
+
+    Returns the model with accompanying metadata.
+
+    Voxels are stored in a three-dimensional numpy array, which is simple and
+    direct, but may use a lot of memory for large models. (Storage requirements
+    are 8*(d^3) bytes, where d is the dimensions of the binvox model. Numpy
+    boolean arrays use a byte per element).
+
+    Doesn't do any checks on input except for the '#binvox' line.
+    """
+    dims, translate, scale = read_header(fp)
+    raw_data = np.frombuffer(fp.read(), dtype=np.uint8)
+    # if just using reshape() on the raw data:
+    # indexing the array as array[i,j,k], the indices map into the
+    # coords as:
+    # i -> x
+    # j -> z
+    # k -> y
+    # if fix_coords is true, then data is rearranged so that
+    # mapping is
+    # i -> x
+    # j -> y
+    # k -> z
+    values, counts = raw_data[::2], raw_data[1::2]
+    data = np.repeat(values, counts).astype(np.bool)
+    data = data.reshape(dims)
+    if fix_coords:
+        # xzy to xyz TODO the right thing
+        data = np.transpose(data, (0, 2, 1))
+        axis_order = 'xyz'
+    else:
+        axis_order = 'xzy'
+    return Voxels(data, dims, translate, scale, axis_order)
+
+
+def _sparse_to_dense(voxel_data, dims, dtype=np.bool):
+    if voxel_data.ndim != 2 or voxel_data.shape[0] != 3:
+        raise ValueError('voxel_data is wrong shape; should be 3xN array.')
+    if np.isscalar(dims):
+        dims = [dims] * 3
+    dims = np.atleast_2d(dims).T
+    # truncate to integers
+    xyz = voxel_data.astype(np.int)
+    # discard voxels that fall outside dims
+    valid_ix = ~np.any((xyz < 0) | (xyz >= dims), 0)
+    xyz = xyz[:, valid_ix]
+    out = np.zeros(dims.flatten(), dtype=dtype)
+    out[tuple(xyz)] = True
+    return out
+
+
+def _bwrite(fp, s):
+    fp.write(s.encode())
+
+
+def _write_pair(fp, state, ctr):
+    fp.write(struct.pack('B', state))
+    fp.write(struct.pack('B', ctr))
+
+
+def write(voxel_model, fp):
+    """ Write binary binvox format.
+
+    Note that when saving a model in sparse (coordinate) format, it is first
+    converted to dense format.
+
+    Doesn't check if the model is 'sane'.
+
+    """
+    if voxel_model.data.ndim == 2:
+        # TODO avoid conversion to dense
+        dense_voxel_data = _sparse_to_dense(voxel_model.data, voxel_model.dims)
+    else:
+        dense_voxel_data = voxel_model.data
+
+    _bwrite(fp, '#binvox 1\n')
+    _bwrite(fp, 'dim ' + ' '.join(map(str, voxel_model.dims)) + '\n')
+    _bwrite(fp, 'translate ' + ' '.join(map(str, voxel_model.translate)) + '\n')
+    _bwrite(fp, 'scale ' + str(voxel_model.scale) + '\n')
+    _bwrite(fp, 'data\n')
+    if not voxel_model.axis_order in ('xzy', 'xyz'):
+        raise ValueError('Unsupported voxel model axis order')
+
+    if voxel_model.axis_order == 'xzy':
+        voxels_flat = dense_voxel_data.flatten()
+    elif voxel_model.axis_order == 'xyz':
+        voxels_flat = np.transpose(dense_voxel_data, (0, 2, 1)).flatten()
+
+    # keep a sort of state machine for writing run length encoding
+    state = voxels_flat[0]
+    ctr = 0
+    for c in voxels_flat:
+        if c == state:
+            ctr += 1
+            # if ctr hits max, dump
+            if ctr == 255:
+                _write_pair(fp, state, ctr)
+                ctr = 0
+        else:
+            # if switch state, dump
+            _write_pair(fp, state, ctr)
+            state = c
+            ctr = 1
+    # flush out remainders
+    if ctr > 0:
+        _write_pair(fp, state, ctr)

rigging/rignet/requirements.txt

@@ -0,0 +1,7 @@
+open3d==0.9.0
+trimesh[easy]
+torch==1.7.1+cpu
+torch-scatter
+torch-sparse
+torch-cluster
+torch-geometric