utils for meshing sdf pcs

Author: Ritvik Vasan

pointcloudutils/__init__.py

@@ -11,6 +11,3 @@
 
 def get_module_version():
     return __version__
-
-
-from .example import Example  # noqa: F401

pointcloudutils/datamodules/shapenet.py

@@ -16,6 +16,7 @@ def __init__(
         self,
         dataset_folder: str = "/allen/aics/modeling/ritvik/projects/occupancy_networks/data/ShapeNet",
         method: str = "shapenet_dfnet",
+        x_label: str = 'pcloud',
         dataset_type: str = "partial_pointcloud",
         train_split: str = "train",
         val_split: str = "val",
@@ -47,6 +48,7 @@ def __init__(
         self.categories = categories
         self.train_split = train_split
         self.val_split = val_split
+        self.x_label = x_label
         self.test_split = test_split
         self.points_subsample = points_subsample
         self.input_type = input_type
@@ -91,7 +93,7 @@ def _get_dataset(self, mode):
         )
 
         if inputs_field is not None:
-            fields["pcloud"] = inputs_field
+            fields[self.x_label] = inputs_field
 
         if self.return_idx:
             fields["idx"] = IndexField()
@@ -110,6 +112,7 @@ def _get_dataset(self, mode):
             self.splits[mode],
             self.categories,
             transform,
+            self.x_label,
         )
         return dataset
 

pointcloudutils/datamodules/shapenet_dataset/core.py

@@ -41,6 +41,7 @@ def __init__(
         split=None,
         categories=None,
         transform=None,
+        x_label='pcloud',
     ):
         """Initialization of the the 3D shape dataset.
 
@@ -55,6 +56,7 @@ def __init__(
         # Attributes
         self.dataset_folder = dataset_folder
         self.fields = fields
+        self.x_label = x_label
         self.transform = transform
         self.split = split
         # If categories is None, use all subfolders
@@ -141,15 +143,15 @@ def __getitem__(self, idx):
 
     def transforms(self, data, transform_type=None):
         if "rotate" in transform_type:
-            data["pcloud"], data["points"], data["points_iou"] = rotate_pointcloud(
-                pointcloud=data["pcloud"],
+            data[self.x_label], data["points"], data["points_iou"] = rotate_pointcloud(
+                pointcloud=data[self.x_label],
                 points=data["points"],
                 points_iou=data.get("points_iou"),
             )
 
         if "translate" in transform_type:
-            data["pcloud"], data["points"], data["points_iou"] = translate_pointcloud(
-                pointcloud=data["pcloud"],
+            data[self.x_label], data["points"], data["points_iou"] = translate_pointcloud(
+                pointcloud=data[self.x_label],
                 points=data["points"],
                 points_iou=data.get("points_iou"),
             )
@@ -164,13 +166,13 @@ def transforms(self, data, transform_type=None):
 
             if data.get("points_iou") is not None:
                 (
-                    data["pcloud"],
-                    data["pcloud"],
+                    data[self.x_label],
+                    data[self.x_label],
                     data["points_iou"],
                     points_df,
                     points_iou_df,
                 ) = single_translate_pointcloud(
-                    pointcloud=data["pcloud"],
+                    pointcloud=data[self.x_label],
                     points=data["points"],
                     points_iou=data["points_iou"],
                     points_df=points_df,
@@ -181,8 +183,8 @@ def transforms(self, data, transform_type=None):
                 if points_iou_df is not None:
                     data["points_iou.df"] = points_iou_df
             else:
-                data["pcloud"], data["points"], points_df = single_translate_pointcloud(
-                    pointcloud=data["pcloud"],
+                data[self.x_label], data["points"], points_df = single_translate_pointcloud(
+                    pointcloud=data[self.x_label],
                     points=data["points"],
                     points_df=points_df,
                 )

pointcloudutils/example.py

@@ -2,7 +2,7 @@
 from torch import nn
 
 from cyto_dl import utils
-from .vnn import VNLinear, VNRotationMatrix
+from cyto_dl.nn.point_cloud.vnn import VNLinear, VNRotationMatrix
 
 log = utils.get_pylogger(__name__)