Merge pull request graspnet#51 from graspnet/1.2.11

GouMinghao · web-flow · commit c57e6a1b365a · 2023-12-31T11:54:25.000+08:00
1.2.11
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -22,7 +22,7 @@
 author = 'graspnet'
 
 # The full version, including alpha/beta/rc tags
-release = '1.2.10'
+release = '1.2.11'
 
 
 # -- General configuration ---------------------------------------------------
diff --git a/graspnetAPI/__init__.py b/graspnetAPI/__init__.py
@@ -1,5 +1,5 @@
 __author__ = 'mhgou'
-__version__ = '1.2.10'
+__version__ = '1.2.11'
 
 from .graspnet import GraspNet
 from .graspnet_eval import GraspNetEval
diff --git a/graspnetAPI/graspnet.py b/graspnetAPI/graspnet.py
@@ -67,7 +67,7 @@ def _isArrayLike(obj):
 
 
 class GraspNet():
-    def __init__(self, root, camera='kinect', split='train'):
+    def __init__(self, root, camera='kinect', split='train', sceneIds=[]):
         '''
 
         graspnetAPI main class.
@@ -76,10 +76,12 @@ def __init__(self, root, camera='kinect', split='train'):
 
         - camera: string of type of camera: "kinect" or "realsense"
 
-        - split: string of type of split of dataset: "all", "train", "test", "test_seen", "test_similar" or "test_novel"
+        - split: string of type of split of dataset: "all", "train", "test", "test_seen", "test_similar", "test_novel" or "custom"
+
+        - sceneIds: list of custom scene ids.
         '''
         assert camera in ['kinect', 'realsense'], 'camera should be kinect or realsense'
-        assert split in ['all', 'train', 'test', 'test_seen', 'test_similar', 'test_novel'], 'split should be all/train/test/test_seen/test_similar/test_novel'
+        assert split in ['all', 'train', 'test', 'test_seen', 'test_similar', 'test_novel', "custom"], 'split should be all/train/test/test_seen/test_similar/test_novel'
         self.root = root
         self.camera = camera
         self.split = split
@@ -97,6 +99,8 @@ def __init__(self, root, camera='kinect', split='train'):
             self.sceneIds = list(range(130, 160))
         elif split == 'test_novel':
             self.sceneIds = list(range(160, 190))
+        elif split == "custom":
+            self.sceneIds = sceneIds
 
         self.rgbPath = []
         self.depthPath = []
diff --git a/setup.py b/setup.py
@@ -5,14 +5,14 @@
 
 setup(
     name='graspnetAPI',
-    version='1.2.10',
+    version='1.2.11',
     description='graspnet API',
     author='Hao-Shu Fang, Chenxi Wang, Minghao Gou',
     author_email='gouminghao@gmail.com',
     url='https://graspnet.net',
     packages=find_packages(),
     install_requires=[
-        'numpy',
+        'numpy==1.20.3',
         'scipy',
         'transforms3d==0.3.1',
         'open3d>=0.8.0.0',
@@ -30,7 +30,7 @@
         'cvxopt',
         'dill',
         'h5py',
-        'sklearn',
+        'scikit-learn',
         'grasp_nms'
     ]
 )
