Feature/decimation

.dvc/.gitignore

@@ -0,0 +1,3 @@
+/config.local
+/tmp
+/cache

.dvc/config

@@ -0,0 +1,4 @@
+[core]
+    remote = storage
+['remote "storage"']
+    url = s3://machinelearning-assets/roofmeasurements/datasets/roof_segmentation

.dvcignore

@@ -0,0 +1,3 @@
+# Add patterns of files dvc should ignore, which could improve
+# the performance. Learn more at
+# https://dvc.org/doc/user-guide/dvcignore

.gitignore

@@ -5,5 +5,5 @@
 # data files
 *.obj
 checkpoints
-datasets
+#datasets
 runs

data/blender_scripts/extract_vertex_labels.py

@@ -0,0 +1,16 @@
+
+import bpy
+
+ob = bpy.context.object
+obdata = bpy.context.object.data
+
+label = []
+for v in obdata.vertices:
+    if bpy.context.object.vertex_groups['roof'].index in [i.group for i in v.groups]:
+        label.append(str(1))
+    else:
+        label.append(str(0))
+
+with open('/home/ihahanov/Projects/roof-measurements/dl_roof_extraction/meshcnn'
+          '/datasets/roof_seg/vseg/model12.eseg', 'w') as f:
+    f.write('\n'.join(label))

data/make_annotation_from_vertex_labels.py

@@ -0,0 +1,196 @@
+import numpy as np
+import os
+import glob
+import filecmp
+import sys
+
+
+'''
+Creates esseg files for accuracy with smooth transitions between classes 
+Requires Objects and corresponding labels per edge
+Author: Rana Hanocka / Lisa Schneider
+
+@input: 
+    <input_path> path where seg, sseg, train, test folders are placed 
+
+@output:
+    esseg files for all objects
+    to run it from cmd line:
+    python create_sseg.py /home/user/MedMeshCNN/datasets/human_seg/
+'''
+
+def compute_face_normals_and_areas(vs, faces):
+    face_normals = np.cross(vs[faces[:, 1]] - vs[faces[:, 0]],
+                            vs[faces[:, 2]] - vs[faces[:, 1]])
+    face_areas = np.sqrt((face_normals ** 2).sum(axis=1))
+    face_normals /= face_areas[:, np.newaxis]
+    assert (not np.any(face_areas[:, np.newaxis] == 0)), 'has zero area face'
+    face_areas *= 0.5
+    return face_normals, face_areas
+
+
+def remove_non_manifolds(vs, faces):
+    edges_set = set()
+    mask = np.ones(len(faces), dtype=bool)
+    _, face_areas = compute_face_normals_and_areas(vs, faces)
+    for face_id, face in enumerate(faces):
+        if face_areas[face_id] == 0:
+            mask[face_id] = False
+            continue
+        faces_edges = []
+        is_manifold = False
+        for i in range(3):
+            cur_edge = (face[i], face[(i + 1) % 3])
+            if cur_edge in edges_set:
+                is_manifold = True
+                break
+            else:
+                faces_edges.append(cur_edge)
+        if is_manifold:
+            mask[face_id] = False
+        else:
+            for idx, edge in enumerate(faces_edges):
+                edges_set.add(edge)
+    return faces[mask], face_areas[mask]
+
+def get_gemm_edges(faces, export_name_edges):
+    """
+    gemm_edges: array (#E x 4) of the 4 one-ring neighbors for each edge
+    sides: array (#E x 4) indices (values of: 0,1,2,3) indicating where an edge is in the gemm_edge entry of the 4 neighboring edges
+    for example edge i -> gemm_edges[gemm_edges[i], sides[i]] == [i, i, i, i]
+    """
+    edge_nb = []
+    sides = []
+    edge2key = dict()
+    edges = []
+    edges_count = 0
+    nb_count = []
+    for face_id, face in enumerate(faces):
+        faces_edges = []
+        for i in range(3):
+            cur_edge = (face[i], face[(i + 1) % 3])
+            faces_edges.append(cur_edge)
+        for idx, edge in enumerate(faces_edges):
+            edge = tuple(sorted(list(edge)))
+            faces_edges[idx] = edge
+            if edge not in edge2key:
+                edge2key[edge] = edges_count
+                edges.append(list(edge))
+                edge_nb.append([-1, -1, -1, -1])
+                sides.append([-1, -1, -1, -1])
+                nb_count.append(0)
+                edges_count += 1
+        for idx, edge in enumerate(faces_edges):
+            edge_key = edge2key[edge]
+            edge_nb[edge_key][nb_count[edge_key]] = edge2key[faces_edges[(idx + 1) % 3]]
+            edge_nb[edge_key][nb_count[edge_key] + 1] = edge2key[faces_edges[(idx + 2) % 3]]
+            nb_count[edge_key] += 2
+        for idx, edge in enumerate(faces_edges):
+            edge_key = edge2key[edge]
+            sides[edge_key][nb_count[edge_key] - 2] = nb_count[edge2key[faces_edges[(idx + 1) % 3]]] - 1
+            sides[edge_key][nb_count[edge_key] - 1] = nb_count[edge2key[faces_edges[(idx + 2) % 3]]] - 2
+    edges = np.array(edges, dtype=np.int32)
+    np.savetxt(export_name_edges, edges, fmt='%i')
+    return edge_nb, edges
+
+
+def load_labels(path):
+    with open(path, 'r') as f:
+        content = f.read().splitlines()
+        return content
+
+def create_sseg_file(gemms, labels, export_name_seseg):
+    gemmlabels = {}
+    classes = len(np.unique(labels))
+    class_to_idx = {v: i for i, v in enumerate(np.unique(labels))}
+    totaledges = len(gemms)
+    sseg = np.zeros([ totaledges, classes])
+    for i, edges in enumerate(gemms):
+        alllabels = []
+        for edge in range(len(edges)):
+            lookupEdge = edges[edge]
+            label = labels[lookupEdge]
+            alllabels.append(label)
+        gemmlabels[i] = alllabels
+
+    for i, edges in enumerate(gemms):
+            gemmlab = gemmlabels[i]
+            uniqueValues, counts = np.unique(gemmlab, return_counts=True)
+            for j, label in enumerate(uniqueValues):
+                weight = 0.125*counts[j]
+                sseg[i][class_to_idx[label]] = weight
+    np.savetxt(export_name_seseg, sseg,  fmt='%1.6f')
+
+def get_obj(file):
+    vs, faces = [], []
+    f = open(file)
+    for line in f:
+        line = line.strip()
+        splitted_line = line.split()
+        if not splitted_line:
+            continue
+        elif splitted_line[0] == 'v':
+            vs.append([float(v) for v in splitted_line[1:4]])
+        elif splitted_line[0] == 'f':
+            face_vertex_ids = [int(c.split('/')[0]) for c in splitted_line[1:]]
+            assert len(face_vertex_ids) == 3
+            face_vertex_ids = [(ind - 1) if (ind >= 0) else (len(vs) + ind)
+                               for ind in face_vertex_ids]
+            faces.append(face_vertex_ids)
+    f.close()
+    vs = np.asarray(vs)
+    faces = np.asarray(faces, dtype=int)
+    assert np.logical_and(faces >= 0, faces < len(vs)).all()
+    return faces, vs
+
+
+import trimesh as tm
+def edges_to_path(edges, color=tm.visual.color.random_color()):
+    lines = np.asarray(edges)
+    args = tm.path.exchange.misc.lines_to_path(lines)
+    colors = [color for _ in range(len(args['entities']))]
+    path = tm.path.Path3D(**args, colors=colors)
+    return path
+
+
+def show_mesh(edges, vs, label,  colors=[[0,0,0,255], [120,120,120,255]]):
+    colors = np.array(colors)
+    edges = vs[edges]
+    tm.Scene([edges_to_path(e, colors[int(l)]) for e, l in zip(edges, label)]).show()
+
+
+
+def create_files(path):
+    for filename in glob.glob(os.path.join(path, 'obj/*.obj')):
+        print(filename)
+        basename = os.path.splitext(os.path.basename(filename))[0]
+        v_label_name = os.path.join(os.path.join(path, 'vseg'), basename + '.eseg')
+        label_name = os.path.join(os.path.join(path, 'seg'), basename + '.eseg')
+        export_name_seseg = os.path.join(os.path.join(path, 'sseg'), basename + '.seseg')
+        export_name_edges = os.path.join(os.path.join(path, 'edges'), basename + '.edges')
+
+        faces, vs = get_obj(filename)
+        faces, face_areas = remove_non_manifolds(vs, faces)
+        gemms, edges = get_gemm_edges(faces, export_name_edges)
+        with open(v_label_name) as f:
+            v_label = np.array(f.readlines(), dtype=int)
+
+        edge_label = []
+        for e in edges:
+            if v_label[e[0]] == 1 and v_label[e[1]] == 1:
+                edge_label.append(str(2))
+            else:
+                edge_label.append(str(1))
+
+        with open(label_name, 'w') as f:
+            f.write('\n'.join(edge_label))
+        print(len(edge_label))
+        if os.path.isfile(label_name):
+
+            create_sseg_file(gemms, edge_label, export_name_seseg)
+        else:
+            print(label_name, "is no directory")
+
+
+if __name__ == '__main__':
+    create_files(sys.argv[1])

data/segmentation_data.py

@@ -5,6 +5,27 @@
 import numpy as np
 from models.layers.mesh import Mesh
 
+
+import trimesh as tm
+def edges_to_path(edges, color=tm.visual.color.random_color()):
+    lines = np.asarray(edges)
+    args = tm.path.exchange.misc.lines_to_path(lines)
+    colors = [color for _ in range(len(args['entities']))]
+    path = tm.path.Path3D(**args, colors=colors)
+    return path
+
+
+def show_mesh(mesh, label,  colors=[[0,0,0,255], [120,120,120,255]]):
+    colors = np.array(colors)
+    edges = mesh.vs[mesh.edges]
+    tm.Scene([edges_to_path(e, colors[int(l)]) for e, l in zip(edges, label)]).show()
+
+
+def show_vertices(mesh, label,  colors=[[0,0,0,255], [120,120,120,255]]):
+    colors = np.array(colors)
+    tm.PointCloud(mesh.vs, colors=np.array(colors)[label]).show()
+
+
 class SegmentationData(BaseDataset):
 
     def __init__(self, opt):
@@ -29,6 +50,7 @@ def __getitem__(self, index):
         mesh = Mesh(file=path, opt=self.opt, hold_history=True, export_folder=self.opt.export_folder)
         meta = {}
         meta['mesh'] = mesh
+        meta['path'] = path
         label = read_seg(self.seg_paths[index]) - self.offset
         label = pad(label, self.opt.ninput_edges, val=-1, dim=0)
         meta['label'] = label

datasets/.gitignore

@@ -0,0 +1 @@
+/roof_seg

datasets/roof_seg.dvc

@@ -0,0 +1,5 @@
+outs:
+- md5: 5ad27ad6392a64493c0976a5ad1a8296.dir
+  size: 176735174
+  nfiles: 404
+  path: roof_seg

decimate.py

@@ -0,0 +1,108 @@
+from copy import deepcopy
+import os
+from collections import OrderedDict
+
+import trimesh as tm
+import numpy as np
+import torch
+
+from options.pl_options import PLOptions
+from data import DataLoader
+from data.segmentation_data import Mesh
+from util.util import pad
+from train_pl import MeshSegmenter
+
+
+def show_mesh(mesh, label):
+    edges = mesh.edges
+    vertices = mesh.vs
+    vertex_label = np.zeros(len(vertices))
+    for e_l, e in zip(label[0], edges):
+        if e_l == 1:
+            vertex_label[e] = 1
+    faces = mesh.faces
+    vertex_colors = np.array([[255, 100, 0, 255], [0, 100, 255, 255]])[vertex_label.astype(int)]
+    trimesh = tm.Trimesh(faces=faces, vertices=vertices, vertex_colors=vertex_colors)
+    trimesh.show()
+    return trimesh
+
+
+def simplify_rooftop(roof_segment: tm.Trimesh, n_triangles) -> tm.Trimesh:
+    """
+    Perform mesh simplificaiton based on desired triangles number
+    :param roof_segment: Trimesh - submesh of the roof
+    :param n_triangles: int - number of triangles the simplified mesh would contain
+    :return: tm.Trimesh - Simplified mesh
+    """
+    n_triangles = max([n_triangles, 5])
+    segment = roof_segment.simplify_quadratic_decimation(n_triangles)
+
+    return segment
+
+
+def load_obj(path, opt, mean, std):
+    mesh = Mesh(file=path, opt=opt, hold_history=True, export_folder=opt.export_folder)
+    meta = {}
+    meta['mesh'] = [mesh]
+    meta['path'] = [path]
+    edge_features = mesh.extract_features()
+    edge_features = pad(edge_features, opt.ninput_edges)
+    edge_features = (edge_features - mean) / std
+    meta['edge_features'] = np.expand_dims(edge_features, 0)
+    meta['label'] = np.array([])
+    meta['soft_label'] = np.array([])
+    return meta
+
+
+def run_decimation(epoch=-1):
+    opt = PLOptions().parse()
+    opt.serial_batches = True  # no shuffle
+    dataset = DataLoader(opt)
+    model = MeshSegmenter(opt)
+
+    device = torch.device('cuda:{}'.format(opt.gpu_ids[0])) if opt.gpu_ids else torch.device('cpu')
+    checkpoint = torch.load(opt.model_path, map_location=device)
+
+    state_dict = checkpoint['state_dict']
+    new_state_dict = OrderedDict()
+    for k, v in state_dict.items():
+        new_key = k.replace('.module', '')
+        new_state_dict[new_key] = v
+    model.load_state_dict(new_state_dict)
+
+    for i, data in enumerate(dataset):
+        if i != 21 and i != 22:
+            continue
+        print(i, data['path'])
+        obj_name = os.path.basename(data['path'][0])
+        print(obj_name)
+
+        torch.cuda.empty_cache()
+
+        mesh = deepcopy(data['mesh'][0])
+        pred_class = model.forward(data).max(1)[1]
+        tm_mesh = show_mesh(mesh, label=pred_class)
+
+        torch.cuda.empty_cache()
+
+        for desired_triangle_area in [0.5, 0.8, 1, 1.2, 1.4, 1.7, 2, 2.5, 3.5, 5, 7]:
+            print(desired_triangle_area)
+
+            tm_mesh_new = simplify_rooftop(tm_mesh, int((tm_mesh.area / desired_triangle_area)))
+
+            new_obj_name = obj_name[:-4] + '_' + str(desired_triangle_area) + '.obj'
+            obj_path = os.path.join(opt.decimation_dir, new_obj_name)
+
+            with open(obj_path, mode='w') as f:
+                f.write(tm.exchange.obj.export_obj(tm_mesh_new))
+                data_new = load_obj(obj_path, opt, dataset.dataset.mean, dataset.dataset.std)
+
+                mesh_new = deepcopy(data_new['mesh'][0])
+                pred_class_new = model.forward(data_new).max(1)[1]
+                show_mesh(mesh_new, label=pred_class_new)
+
+                torch.cuda.empty_cache()
+            # os.unlink(f.name)
+
+if __name__ == '__main__':
+    run_decimation()