def transform_setup(graph_u=False, graph_gcn=False, rotation=180, samplePoints=1024, mesh=False, node_translation=0.01): if not graph_u and not graph_gcn: # Default transformation for scale noralization, centering, point sampling and rotating pretransform = T.Compose([T.NormalizeScale(), T.Center()]) transform = T.Compose([ T.SamplePoints(samplePoints), T.RandomRotate(rotation[0], rotation[1]) ]) print("pointnet rotation {}".format(rotation)) elif graph_u: pretransform = T.Compose([T.NormalizeScale(), T.Center()]) transform = T.Compose([ T.NormalizeScale(), T.Center(), T.SamplePoints(samplePoints, True, True), T.RandomRotate(rotation[0], rotation[1]), T.KNNGraph(k=graph_u) ]) elif graph_gcn: pretransform = T.Compose([T.NormalizeScale(), T.Center()]) if mesh: if mesh == "extraFeatures": transform = T.Compose([ T.RandomRotate(rotation[0], rotation[1]), T.GenerateMeshNormals(), T.FaceToEdge(True), T.Distance(norm=True), T.TargetIndegree(cat=True) ]) # , else: transform = T.Compose([ T.RandomRotate(rotation[0], rotation[1]), T.GenerateMeshNormals(), T.FaceToEdge(True), T.Distance(norm=True), T.TargetIndegree(cat=True) ]) else: transform = T.Compose([ T.SamplePoints(samplePoints, True, True), T.KNNGraph(k=graph_gcn), T.Distance(norm=True) ]) print("no mesh") print("Rotation {}".format(rotation)) print("Meshing {}".format(mesh)) else: print('no transfom') return transform, pretransform
def get_adj_list(data, max_adj=None): data.edge_index = None data = T.FaceToEdge(remove_faces=False)(data) edge, _ = add_self_loops(data.edge_index) data = T.ToDense()(data) adj_mat = data.adj num_list = adj_mat.sum(1).long().unsqueeze(1) if max_adj is None: max_adj = num_list.max().item() else: max_list = torch.full_like(num_list, max_adj).long() num_list = torch.where(num_list > max_adj, max_list, num_list) adj_list = torch.full_like(adj_mat, -1).long() N = data.pos.shape[0] for n in range(N): adj = adj_mat[n].nonzero() num = num_list[n] adj_list[n, :num] = adj.t()[:, :num] adj_list = torch.cat([adj_list[:, :max_adj], num_list], dim=1) # N * max_adj return adj_list.type_as(edge), edge
def process(self): data_list = [] f2e = transforms.FaceToEdge(remove_faces=False) for pindex, path in enumerate(tqdm.tqdm(self.raw_paths)): mesh = torch_geometric.io.read_ply(path) f2e(mesh) tmp = split(path)[1].split(".")[0].split("_") model_str, pose_str = tmp[-2], tmp[-1] category = "_".join(tmp[:-2]) mesh.model = int(model_str[5:]) mesh.pose = int(pose_str[4:]) mesh.y = self.categories.index(category) if self.pre_filter is not None and not self.pre_filter(mesh) : continue if self.pre_transform is not None: mesh = self.pre_transform(mesh) data_list.append(mesh) data, slices = self.collate(data_list) torch.save( (data, slices), self.processed_paths[0])
def process(self): # Read data into huge `Data` list. face2edge = transforms.FaceToEdge(remove_faces=False) data_list = [] num_classes = 10 for path in tqdm.tqdm(self.raw_paths): i = int(basename(path)[:-4]) mesh = torch_geometric.io.read_obj(path) mesh.y = i % num_classes mesh.subject = int(i / num_classes) face2edge(mesh) mesh.idx = i if self.pre_filter is not None and not self.pre_filter(mesh): continue if self.pre_transform is not None: mesh = self.pre_transform(mesh) data_list.append(mesh) data, slices = self.collate(data_list) torch.save((data, slices), self.processed_paths[0])
def __init__(self, batch_size=32, num_point_samples=4096): """ This code runs the preprocessing pipeline ones and stores the standardised coordinates. It seems overkill to run the entire pipeline just to normalise the landmarks but it is required to find the bounding box. """ root = '../data/CurvatureComplete_normalisation_only/' prep_dir = os.path.join(root, 'prep_parameters') # only perform preprocessing upto b-norm of the landmarks pre_transform = local_transforms.InvertibleCompose( [ # convert mesh faces to graph edges ptg_transforms.FaceToEdge(remove_faces=False), local_transforms.LabelCloner('x', 'continuous_curvature'), # standardise x so that x ∈ [0,1] local_transforms.UnitNormaliseScalar(), # propagate curvature feature to local extreme local_transforms.PropageteFeaturesToLocalExtremes(c=80), # remove vertices based on curvature local_transforms.FilterVerts( operator=lambda data: data.x[:, 0] < 0.85), # b-normalise local_transforms.NormalizeScale(invertible=True), local_transforms.LabelCloner('x', 'discrete_curvature'), ], invertible=True, skip_non_invertible=True, store_directory=prep_dir).load_parameters() super().__init__(root, batch_size=batch_size, transform=None, pre_transform=pre_transform)
parser.add_argument('--num_layers', type=int, default=2) parser.add_argument('--num_steps', type=int, default=10) parser.add_argument('--lr', type=float, default=0.001) parser.add_argument('--batch_size', type=int, default=512) parser.add_argument('--pre_epochs', type=int, default=15) parser.add_argument('--epochs', type=int, default=15) parser.add_argument('--runs', type=int, default=20) parser.add_argument('--test_samples', type=int, default=100) args = parser.parse_args() pre_filter1 = lambda d: d.num_nodes > 0 # noqa pre_filter2 = lambda d: d.num_nodes > 0 and d.name[:4] != '2007' # noqa transform = T.Compose([ T.Delaunay(), T.FaceToEdge(), T.Distance() if args.isotropic else T.Cartesian(), ]) path = osp.join('..', 'data', 'PascalVOC-WILLOW') pretrain_datasets = [] for category in PascalVOC.categories: dataset = PascalVOC(path, category, train=True, transform=transform, pre_filter=pre_filter2 if category in ['car', 'motorbike'] else pre_filter1) pretrain_datasets += [ValidPairDataset(dataset, dataset, sample=True)] pretrain_dataset = torch.utils.data.ConcatDataset(pretrain_datasets) pretrain_loader = DataLoader(pretrain_dataset,
import os.path as osp import torch import torch.nn.functional as F from torch_geometric.datasets import FAUST import torch_geometric.transforms as T from torch_geometric.loader import DataLoader from torch_geometric.nn import SplineConv path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'FAUST') pre_transform = T.Compose([T.FaceToEdge(), T.Constant(value=1)]) train_dataset = FAUST(path, True, T.Cartesian(), pre_transform) test_dataset = FAUST(path, False, T.Cartesian(), pre_transform) train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1) d = train_dataset[0] class Net(torch.nn.Module): def __init__(self): super().__init__() self.conv1 = SplineConv(1, 32, dim=3, kernel_size=5, aggr='add') self.conv2 = SplineConv(32, 64, dim=3, kernel_size=5, aggr='add') self.conv3 = SplineConv(64, 64, dim=3, kernel_size=5, aggr='add') self.conv4 = SplineConv(64, 64, dim=3, kernel_size=5, aggr='add') self.conv5 = SplineConv(64, 64, dim=3, kernel_size=5, aggr='add') self.conv6 = SplineConv(64, 64, dim=3, kernel_size=5, aggr='add') self.lin1 = torch.nn.Linear(64, 256) self.lin2 = torch.nn.Linear(256, d.num_nodes) def forward(self, data):
def __call__(self, data): data.x = data.pos data = T.FaceToEdge()(data) return data
return False class MyTransform(object): def __call__(self, data): data.face, data.x = None, torch.ones(data.num_nodes, 1) return data def norm(x, edge_index): deg = degree(edge_index[0], x.size(0), x.dtype, x.device) + 1 return x / deg.unsqueeze(-1) path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'FAUST') pre_transform = T.Compose([T.FaceToEdge(), MyTransform()]) train_dataset = FAUST(path, True, T.Cartesian(), pre_transform) train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) d = train_dataset[0] class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = SplineConv(1, 32, dim=3, kernel_size=5, norm=False) self.conv2 = SplineConv(32, 64, dim=3, kernel_size=5, norm=False) self.conv3 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False) self.conv4 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False) self.conv5 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False) self.conv6 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False) self.fc1 = torch.nn.Linear(64, 256)
from torch_geometric.utils import degree class MyTransform(object): def __call__(self, data): data.face, data.x = None, torch.ones(data.num_nodes, 1) return data def norm(x, edge_index): deg = degree(edge_index[0], x.size(0), x.dtype) + 1 return x / deg.unsqueeze(-1) path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'FAUST') pre_transform = T.Compose([T.FaceToEdge(), MyTransform()]) train_dataset = FAUST(path, True, T.Cartesian(), pre_transform) test_dataset = FAUST(path, False, T.Cartesian(), pre_transform) train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1) d = train_dataset[0] class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.conv1 = SplineConv(1, 32, dim=3, kernel_size=5, norm=False) self.conv2 = SplineConv(32, 64, dim=3, kernel_size=5, norm=False) self.conv3 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False) self.conv4 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False) self.conv5 = SplineConv(64, 64, dim=3, kernel_size=5, norm=False)
TRIPLET = False EPOCH_softmax = args.epoch_softmax EPOCH_hardloss = args.epoch_hardloss K = args.wavelet_scales + 1 SAVE_NAME = args.saving_name CPOINT_NAME = args.loading_name LEARNING_RATE = args.learning_rate_softmax WEIGHT_DECAY = args.weight_decay_softmax path = osp.join(osp.dirname(osp.realpath(__file__)), 'datasets', 'Faust') path_output = osp.join(osp.abspath('.'), 'outputs', SAVE_NAME) if not os.path.exists(path_output): os.makedirs(path_output) LOG_FOUT = open(path_output + '/log.out', 'w') pre_transform = T.FaceToEdge() train_dataset = FAUST_WAVELET(path, True, None, pre_transform) test_dataset = FAUST_WAVELET(path, False, None, pre_transform) train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True) train_loader_tri = DataListLoader(train_dataset, batch_size=2, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False) d = test_dataset[0] d.num_nodes = args.n_corr_points class L2Norm(torch.nn.Module): def __init__(self): super(L2Norm, self).__init__() # self.eps = 1e-10 def forward(self, x): norm = torch.sqrt(torch.sum(x * x, dim=1)) # + self.eps
def set_mesh_graph(self): mesh_transform = T.Compose([T.FaceToEdge(), T.Cartesian()]) mesh_graph = mesh_transform(self.mesh_graph) return mesh_graph.x, mesh_graph.edge_index, mesh_graph.edge_attr
def mesh_edge_subdiv(data): # face unpooling # v0 # /\ # n0 /__\ n2 # /\ /\ # /__\/__\ # v1 n1 v2 # [v0, v1, v2] -> [v0, n0, n2], [n0, v1, n1], [n1, v2, n2], [n0, n1, n2] vert = data.pos # V * 3 V = vert.shape[0] # use egde for index (since tensor cannot used for dict keys) data.edge_index = None data = T.FaceToEdge(remove_faces=False)(data) edge_list, _ = data.edge_index.sort(0) # 2 * E edge_sort_list, _ = coalesce(edge_list, None, V, V) # 2 * E' edge_sort_list = edge_sort_list.t() # E' * 2 face = data.face.t() # F * 3 F = face.shape[0] verts = [face[:, i] for i in range(3)] # [v0 v1 v2] edges = [face[:, :2], face[:, 1:], face[:, [-1, 0]]] # [e0(v0, v1), e1(v1, v2), e2(v2, v0)] edges_sort = [edge.sort(1)[0] for edge in edges] edge_dict = {} start_index = vert.shape[0] new_verts = [] pool_idx = torch.zeros(1, 2).type_as(edge_list) for f in range(F): ns = [] for e in range(3): edge_sort = edges_sort[e][f] idx = find_idx(edge_sort, edge_sort_list) if idx in edge_dict.keys(): n = edge_dict[idx] else: n = start_index edge_dict[idx] = torch.tensor([n]).type_as(edge_list) pool_idx = torch.cat([pool_idx, edge_sort.unsqueeze(0)], 0) start_index += 1 ns.append(edge_dict[idx]) if len(new_verts) == 0: new_verts.extend(ns) else: new_verts = [torch.cat([new_verts[i], ns[i]], 0) for i in range(3)] pool_idx = pool_idx[1:] # E" * 2 v0, v1, v2 = verts n0, n1, n2 = new_verts new_face0 = torch.stack([v0, n0, n2], 1) new_face1 = torch.stack([n0, v1, n1], 1) new_face2 = torch.stack([n1, v2, n2], 1) new_face3 = torch.stack([n0, n1, n2], 1) new_faces = torch.cat([new_face0, new_face1, new_face2, new_face3], 0) # F" * 3 return pool_idx, new_faces
def __init__(self, landmark='IP', batch_size=64, num_point_samples=2048, patch_size=8.): root = '../data/PatchBasedCompleteDentalDataModule_{}/'.format( landmark) self.prep_dir = os.path.join(root, 'prep_parameters') labels = [landmark + '_X', landmark + '_Y', landmark + '_Z'] gt_labels = [ 'gt_' + landmark + '_X', 'gt_' + landmark + '_Y', 'gt_' + landmark + '_Z' ] # read ids used in training training_ids = pd.read_csv(os.path.join( root, 'normalised_train_data_landmarks.csv'), header=0, index_col=0)['StudyID'] # read predicted landmarks for training examples landmarks = pd.concat([ pd.read_csv(f, header=0, index_col=0) for f in glob.glob(os.path.join(root, 'v0*.csv')) ]) landmarks = landmarks[landmarks['StudyID'].isin(training_ids.values)] assert len(landmarks) == len(training_ids) # get ground truth and predicted labels gt_landmarks = landmarks[gt_labels].astype(float) pred_landmarks = landmarks[labels].astype(float) # sample weights are proportional to the error, since these are less likely to occur self.sample_weights = torch.from_numpy( (pred_landmarks - gt_landmarks.values).abs().sum(1).values) # setup preprocessing steps, load them if they were executed before pre_transform = local_transforms.InvertibleCompose( [ # convert mesh faces to graph edges ptg_transforms.FaceToEdge(remove_faces=False), # # add edge_attr containing relative euclidean distance ptg_transforms.Distance(norm=False, cat=False), # extract patch around local_transforms.ExtractGeodesicPatch(patch_size, key='patch_center'), # b-normalise local_transforms.NormalizeScale(invertible=True), # save vertex features local_transforms.LabelCloner('x', 'continuous_curvature'), ptg_transforms.GenerateMeshNormals(), # remove labels used in preprocessing local_transforms.LabelCloner('pos', 'mesh_vert'), local_transforms.LabelCleaner(['edge_index', 'edge_attr']) ], invertible=True, skip_non_invertible=True, store_directory=self.prep_dir).load_parameters() pre_process = self.set_patch_centers transform = ptg_transforms.Compose([ local_transforms.LabelCloner('continuous_curvature', 'x'), local_transforms.SamplePoints(num_point_samples, remove_faces=False, include_normals=True, include_features=True), local_transforms.MergeLabels('norm'), local_transforms.ZNormalise('x'), local_transforms.LabelCleaner([ 'mesh_vert', 'norm', 'continuous_curvature', 'patch_center', 'face', 'mesh_norm', 'mesh_color' ]) ]) super().__init__(root, batch_size=batch_size, pre_process=pre_process, transform=transform, labels=labels, pre_transform=pre_transform)
def __init__(self, batch_size=32, num_point_samples=4096): root = '../data/CurvatureComplete/' prep_dir = os.path.join(root, 'prep_parameters') # Generate standardised landmarks to setup pca training_set_landmarks = torch.tensor( pd.read_csv(os.path.join(root, 'normalised_train_data_landmarks.csv'), header=0, index_col=0).values[:, 1:].astype('float32')).view( -1, 7, 3) pre_transform = local_transforms.InvertibleCompose( [ # convert mesh faces to graph edges ptg_transforms.FaceToEdge(remove_faces=False), local_transforms.LabelCloner('x', 'continuous_curvature'), # standardise x so that x ∈ [0,1] local_transforms.UnitNormaliseScalar(), # propagate curvature feature to local extreme local_transforms.PropageteFeaturesToLocalExtremes(c=80), # remove vertices based on curvature local_transforms.FilterVerts( operator=lambda data: data.x[:, 0] < 0.85), # b-normalise local_transforms.NormalizeScale(invertible=True), local_transforms.LabelCloner('x', 'discrete_curvature'), # normalise landmarks local_transforms.PCAProjectLandmarks( training_set_landmarks, n_components=3, invertible=True), # generate vertex features ptg_transforms.GenerateMeshNormals(), # remove labels used in preprocessing local_transforms.LabelCloner('pos', 'mesh_vert'), local_transforms.LabelCleaner(['edge_index']) ], invertible=True, skip_non_invertible=True, store_directory=prep_dir).load_parameters() transform = local_transforms.InvertibleCompose( [ local_transforms.LabelCloner('continuous_curvature', 'x'), local_transforms.SamplePoints(num_point_samples, remove_faces=False, include_normals=True, include_features=True), local_transforms.MergeLabels('norm'), local_transforms.ZNormalise('x'), local_transforms.LabelCleaner([ 'mesh_vert', 'norm', 'continuous_curvature', 'face', 'mesh_norm', 'discrete_curvature' ]) ], skip_non_invertible=True, invertible=True) super().__init__(root, batch_size=batch_size, transform=transform, pre_transform=pre_transform)