diff --git a/README.md b/README.md index cfde5e2..25eb06f 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# kiui kit +# kiuikit -A toolkit for personal use. +A toolkit for computer vision (especially 3D vision) tasks. **Features**: * Collection of *maintained, reusable and trustworthy* code snippets. diff --git a/kiui/nn.py b/kiui/nn.py new file mode 100644 index 0000000..40ff1ee --- /dev/null +++ b/kiui/nn.py @@ -0,0 +1,57 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +import numpy as np + +class MLP(nn.Module): + def __init__(self, dim_in, dim_out, dim_hidden, num_layers, bias=True): + super().__init__() + self.dim_in = dim_in + self.dim_out = dim_out + self.dim_hidden = dim_hidden + self.num_layers = num_layers + + net = [] + for l in range(num_layers): + net.append(nn.Linear(self.dim_in if l == 0 else self.dim_hidden, self.dim_out if l == num_layers - 1 else self.dim_hidden, bias=bias)) + + self.net = nn.ModuleList(net) + + def forward(self, x): + for l in range(self.num_layers): + x = self.net[l](x) + if l != self.num_layers - 1: + x = F.relu(x, inplace=True) + return x + +class HashGrid(nn.Module): + def __init__(self, + input_dim=3, + num_levels=16, + level_dim=2, + log2_hashmap_size=19, + base_resolution=16, + desired_resolution=2048, + interpolation='linear' + ): + super().__init__() + import tinycudann as tcnn + self.encoder = tcnn.Encoding( + n_input_dims=input_dim, + encoding_config={ + "otype": "HashGrid", + "n_levels": num_levels, + "n_features_per_level": level_dim, + "log2_hashmap_size": log2_hashmap_size, + "base_resolution": base_resolution, + "per_level_scale": np.exp2(np.log2(desired_resolution / num_levels) / (num_levels - 1)), + "interpolation": "Smoothstep" if interpolation == 'smoothstep' else "Linear", + }, + dtype=torch.float32, + ) + self.input_dim = input_dim + self.output_dim = self.encoder.n_output_dims # patch + + def forward(self, x, bound=1): + return self.encoder((x + bound) / (2 * bound)) \ No newline at end of file diff --git a/setup.py b/setup.py index c3d3ecf..9abdc41 100644 --- a/setup.py +++ b/setup.py @@ -3,7 +3,7 @@ if __name__ == "__main__": setup( name="kiui", - version="0.1.9", + version="0.1.10", description="self-use toolkits", long_description=open("README.md", encoding="utf-8").read(), long_description_content_type="text/markdown", @@ -27,6 +27,8 @@ "rich", "numpy", "scipy", + "scikit-image", + "scikit-learn", "pandas", "trimesh", "pygltflib", @@ -35,7 +37,6 @@ "opencv-python", "imageio", "imageio-ffmpeg", - "scikit-image", "rembg[gpu,cli]", "dearpygui", 'huggingface_hub',