semantic counts

Author: csguestp

Diff for: datasets/scannet_dataset.py

@@ -29,13 +29,14 @@ def g(x_):
 
 ## ScanNet dataset class
 class ScanNetDataset(Dataset):
-    def __init__(self, options, split, random=True):
+    def __init__(self, options, split, load_confidence=False, random=True):
         self.options = options
         self.split = split
         self.random = random
         self.imagePaths = []
-        self.dataFolder = '/gruvi/Data/chenliu/ScanNet/scans/'
-
+        self.dataFolder = options.dataFolder
+        self.load_confidence = load_confidence
+        
         with open('split_' + split + '.txt', 'r') as f:
             for line in f:
                 scene_id = line.strip()
@@ -134,32 +135,7 @@ def __getitem__(self, index):
             colors[:, :3] = colors[:, :3] + np.random.randn(3) * 0.1            
             pass
 
-        if self.options.trainingMode == 'semantic':
-            unique_instances, indices, instances = np.unique(instances, return_index=True, return_inverse=True)
-            labels = labels[indices]
-            labels[labels == -100] = 20
-            new_coords = np.zeros(coords.shape, dtype=coords.dtype)
-            for instance in range(len(unique_instances)):
-                instance_mask = instances == instance
-                instance_coords = coords[instance_mask]
-                mins = instance_coords.min(0)
-                maxs = instance_coords.max(0)
-                max_range = (maxs - mins).max()
-                if self.split == 'train':
-                    padding = (maxs - mins) * np.random.random(3) * 0.1
-                else:
-                    padding = max_range * 0.05
-                    pass
-                max_range += padding * 2
-                mins = (mins + maxs) / 2 - max_range / 2
-                instance_coords = np.clip(np.round((instance_coords - mins) / max_range * full_scale), 0, full_scale - 1)
-                new_coords[instance_mask] = instance_coords
-                continue
-            coords = np.concatenate([new_coords, np.expand_dims(instances, -1)], axis=-1)
-            sample = [coords.astype(np.int64), colors.astype(np.float32), faces.astype(np.int64), labels.astype(np.int64), instances.astype(np.int64), self.imagePaths[index]]
-            return sample
-
-        if self.options.trainingMode == 'confidence':
+        if self.load_confidence:
             scene_id = self.imagePaths[index].split('/')[-1].split('_vh_clean_2')[0]
             info = torch.load('test/output_normal_augment_2_' + self.split + '/cache/' + scene_id + '.pth')
             if len(info) == 2:

Diff for: datasets/semantic_counts_pixelwise.npy

@@ -35,34 +35,10 @@ def parse_args():
                         default=2.5e-4, type=float)
     parser.add_argument('--numEpochs', dest='numEpochs',
                         help='the number of epochs',
-                        default=1000, type=int)
+                        default=50, type=int)
     parser.add_argument('--startEpoch', dest='startEpoch',
                         help='starting epoch index',
                         default=-1, type=int)
-    parser.add_argument('--modelType', dest='modelType',
-                        help='model type',
-                        default='', type=str)
-    parser.add_argument('--heatmapThreshold', dest='heatmapThreshold',
-                        help='heatmap threshold for positive predictions',
-                        default=0.5, type=float)
-    parser.add_argument('--distanceThreshold3D', dest='distanceThreshold3D',
-                        help='distance threshold 3D',
-                        default=0.2, type=float)
-    parser.add_argument('--distanceThreshold2D', dest='distanceThreshold2D',
-                        help='distance threshold 2D',
-                        default=20, type=float)
-    parser.add_argument('--numInputPlanes', dest='numInputPlanes',
-                        help='the number of input planes',
-                        default=1024, type=int)
-    parser.add_argument('--numOutputPlanes', dest='numOutputPlanes',
-                        help='the number of output planes',
-                        default=10, type=int)
-    parser.add_argument('--numInputClasses', dest='numInputClasses',
-                        help='the number of input classes',
-                        default=0, type=int)
-    parser.add_argument('--numOutputClasses', dest='numOutputClasses',
-                        help='the number of output classes',
-                        default=0, type=int)    
     parser.add_argument('--width', dest='width',
                         help='input width',
                         default=256, type=int)
@@ -77,32 +53,29 @@ def parse_args():
                         default=50, type=int)    
     parser.add_argument('--numScales', dest='numScales',
                         help='the number of scales',
-                        default=1, type=int)
+                        default=2, type=int)
     parser.add_argument('--numCrossScales', dest='numCrossScales',
                         help='the number of cross scales',
                         default=0, type=int)            
     parser.add_argument('--numNeighbors', dest='numNeighbors',
                         help='the number of neighbors',
                         default=6, type=int)
-    parser.add_argument('--outputScale', dest='outputScale',
-                        help='output scale',
-                        default=256, type=int)
-    parser.add_argument('--negativeWeights', dest='negativeWeights',
-                        help='negative weights',
-                        default='531111', type=str)
-    parser.add_argument('--trainingMode', dest='trainingMode',
-                        help='training mode',
-                        default='all', type=str)
     ## Flags
     parser.add_argument('--visualizeMode', dest='visualizeMode',
                         help='visualization mode',
                         default='', type=str)    
     parser.add_argument('--suffix', dest='suffix',
                         help='suffix to distinguish experiments',
-                        default='', type=str)    
+                        default='normal_augment', type=str)    
     parser.add_argument('--useCache', dest='useCache',
                         help='use cache instead of re-computing existing examples',
                         default=0, type=int)    
+    parser.add_argument('--dataFolder', dest='dataFolder',
+                        help='data folder',
+                        default='/gruvi/Data/chenliu/ScanNet/scans/', type=str)
+    parser.add_argument('--labelFile', dest='labelFile',
+                        help='path to scannetv2-labels.combined.tsv',
+                        default='/gruvi/Data/chenliu/ScanNet/tasks/scannetv2-labels.combined.tsv', type=str)
 
     args = parser.parse_args()
     return args

Diff for: models/instance.py

@@ -7,7 +7,7 @@
 
 # python imports
 import math
-import os, sys, argparse
+import os, sys
 import inspect
 import json
 import glob
@@ -18,29 +18,15 @@
     print("Failed to import numpy package.")
     sys.exit(-1)
 
-currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
-parentdir = os.path.dirname(currentdir)
-sys.path.insert(0,parentdir)
-import util
-import util_3d
+# currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
+# parentdir = os.path.dirname(currentdir)
+# sys.path.insert(0,parentdir)
+from scripts.util import read_label_mapping
+from scripts.util_3d import read_mesh_vertices
 import torch
-import multiprocessing as mp    
-from utils import write_ply_label
-
-TASK_TYPES = {'label', 'instance'}
-
-parser = argparse.ArgumentParser()
-#parser.add_argument('--scan_path', required=True, help='path to scannet scene (e.g., data/ScanNet/v2/scene0000_00')
-parser.add_argument('--label_map_file', default='/gruvi/Data/chenliu/ScanNet/tasks/scannetv2-labels.combined.tsv', help='path to scannetv2-labels.combined.tsv')
-parser.add_argument('--type', default='instance', help='task type [label or instance]')
-opt = parser.parse_args()
-assert opt.type in TASK_TYPES
-
-label_map = util.read_label_mapping(opt.label_map_file, label_from='raw_category', label_to='nyu40id')
-# remapper=np.ones(150)*(-100)
-# for i,x in enumerate([1,2,3,4,5,6,7,8,9,10,11,12,14,16,24,28,33,34,36,39]):
-#     remapper[x]=i
-
+import multiprocessing as mp
+import functools
+#from utils import write_ply_label
 
 def read_aggregation(filename):
     assert os.path.isfile(filename)
@@ -76,44 +62,23 @@ def read_segmentation(filename):
     return seg_to_verts, num_verts
 
 
-def export(filename):
+def export(filename, label_map):
     scan_name = filename.split('_vh')[0]
     mesh_file = os.path.join(scan_name + '_vh_clean_2.ply')
     agg_file = os.path.join(scan_name + '.aggregation.json')
     seg_file = os.path.join(scan_name + '_vh_clean_2.0.010000.segs.json')
 
-    if os.path.exists(mesh_file[:-4] + '.pth') and len(torch.load(mesh_file[:-4] + '.pth')) == 5 and False:
+    print(filename)        
+    if os.path.exists(mesh_file[:-4] + '.pth'):
         return
-    print(filename)    
 
-    #mesh_vertices, mesh_colors, faces = util_3d.read_mesh_vertices(mesh_file)
+    mesh_vertices, mesh_colors, faces = read_mesh_vertices(mesh_file)
     if os.path.exists(agg_file):
         object_id_to_segs, label_to_segs = read_aggregation(agg_file)
         seg_to_verts, num_verts = read_segmentation(seg_file)
         label_ids = np.zeros(shape=(num_verts), dtype=np.uint32)     # 0: unannotated
         instance_ids = np.zeros(shape=(num_verts), dtype=np.uint32)  # 0: unannotated
 
-        invalid_instance_ids = np.zeros(shape=(num_verts), dtype=np.uint32)  # 0: unannotated
-        # print(len(seg_to_verts))
-        for object_id, segs in object_id_to_segs.items():
-            object_verts = []
-            for seg in segs:
-                verts = seg_to_verts[seg]
-                object_verts.append(verts)
-                continue
-            nums = np.array([len(verts) for verts in object_verts])
-            invalid_indices = np.logical_and(nums < (0.5 * nums.sum()), nums >= 100)
-            invalid_indices = invalid_indices.nonzero()[0]
-            if len(invalid_indices) == 0:
-                continue                
-            seg = segs[np.random.choice(invalid_indices)]
-            verts = seg_to_verts[seg]
-            
-            invalid_instance_ids[verts] = object_id            
-            continue
-        torch.save((invalid_instance_ids, ), mesh_file[:-4] + '_invalid.pth')
-        return
-    
         # write_ply_label('test/mesh.ply', mesh_vertices, faces, label_ids)
         # exit(1)
 
@@ -132,46 +97,19 @@ def export(filename):
         label_ids = np.zeros(shape=(num_verts), dtype=np.uint32)     # 0: unannotated
         instance_ids = np.zeros(shape=(num_verts), dtype=np.uint32)  # 0: unannotated
         pass
-    #point_cloud = torch.load(mesh_file[:-4] + '.pth')
-    #print(point_cloud)
-    #print([(v.shape, v.min(0), v.max(0), v.dtype) for v in point_cloud])
-    #print(np.abs(remapper[label_ids] - point_cloud[2]).max())
-    #print([(v.shape, v.min(), v.max()) for v in [mesh_vertices, mesh_colors, label_ids, instance_ids]])
-    #exit(1)
     mesh_vertices = np.ascontiguousarray(mesh_vertices - mesh_vertices.mean(0))
     mesh_colors = np.ascontiguousarray(mesh_colors) / 127.5 - 1
-    # print(np.abs(mesh_vertices - point_cloud[0]).max())
-    # print(np.abs(mesh_colors - point_cloud[1]).max())            
-    # print(np.abs(remapper[label_ids] - point_cloud[2]).max())
-    # exit(1)    
     torch.save((mesh_vertices, mesh_colors, label_ids, instance_ids, faces), mesh_file[:-4] + '.pth')
     return
 
-def main():
-    ROOT_FOLDER = '/gruvi/Data/chenliu/ScanNet/scans/'
-    files = sorted(glob.glob(ROOT_FOLDER + '*/*_vh_clean_2.ply'))
-    #print(files)
-    #exit(1)
+def prepare_data(options):
+    ROOT_FOLDER = options.dataFolder
+    files = sorted(glob.glob(options.dataFolder + '*/*_vh_clean_2.ply'))
+    p = mp.Pool(processes=mp.cpu_count())
 
-    # files = [filename for filename in files if 'scene0568_00' in filename]
-    # print(files)
-    # export(files[0])
-    # exit(1)
+    label_map = read_label_mapping(options.labelFile, label_from='raw_category', label_to='nyu40id')
 
-    
-    #print(mp.cpu_count())
-    # for filename in files:
-    #     export(filename)
-    #     continue
-    # exit(1)
-    p = mp.Pool(processes=mp.cpu_count())
-    p.map(export, files)
+    p.map(functools.partial(export, label_map=label_map), files)
     p.close()
     p.join()
-    
-    #for filename in files:
-    #export(mesh_file, agg_file, seg_file, opt.label_map_file)
-
-
-if __name__ == '__main__':
-    main()
+    return

Diff for: options.py

@@ -14,7 +14,7 @@
     print("pip install plyfile")
     sys.exit(-1)
 
-import util
+import scripts.util
 
 
 # matrix: 4x4 np array