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 __init__(self,
root:str,
device:torch.device=torch.device("cpu"),
train:bool=True,
test:bool=True,
transform_data:bool=True):
self.url = 'https://drive.google.com/file/d/1dp4sMvZ8cmIIITE-qj6zYpZb0-v-4Kgf/view?usp=sharing'
self.categories = ["big_cats","cows","dogs","hippos","horses"]
# center each mesh into its centroid
pre_transform = transforms.Center()
# transform
if transform_data:
# rotate and move
transform = transforms.Compose([
Move(mean=[0,0,0], std=[0.05,0.05,0.05]),
Rotate(dims=[0,1,2]),
ToDevice(device)])
else:
transform=ToDevice(device)
super().__init__(root=root, transform=transform, pre_transform=pre_transform)
self.data, self.slices = torch.load(self.processed_paths[0])
self.downscaler = dscale.Downscaler(
filename=join(self.processed_dir,"ds"), mesh=self.get(0), factor=2)
if train and not test:
self.data, self.slices = self.collate([self.get(i) for i in range(len(self)) if self.get(i).pose < 16])
elif not train and test:
self.data, self.slices = self.collate([self.get(i) for i in range(len(self)) if self.get(i).pose >= 16])
def __init__(self,
root: str,
device: torch.device = torch.device("cpu"),
train: bool = True,
test: bool = True,
transform_data: bool = True):
self.url = 'http://www.cs.cf.ac.uk/shaperetrieval/shrec14/'
if transform_data:
# rotate and move
transform = transforms.Compose([
transforms.Center(),
# transforms.RandomScale((0.8,1.2)),
Rotate(dims=[1]),
Move(mean=[0, 0, 0], std=[0.05, 0.05, 0.05]),
transforms.RandomTranslate(0.01),
ToDevice(device)
])
else:
transform = ToDevice(device)
# center each mesh into its centroid
super().__init__(root=root,
transform=transform,
pre_transform=transforms.Center())
self.data, self.slices = torch.load(self.processed_paths[0])
testset_slice, trainset_slice = list(range(0, 40)) + list(
range(200, 240)), list(range(40, 200)) + list(range(240, 400))
if train and not test:
self.data, self.slices = self.collate(
[self[i] for i in trainset_slice])
elif not train and test:
self.data, self.slices = self.collate(
[self[i] for i in testset_slice])
def test_compose():
transform = T.Compose([T.Center(), T.AddSelfLoops()])
assert str(transform) == ('Compose([\n'
' Center(),\n'
' AddSelfLoops()\n'
'])')
pos = torch.Tensor([[0, 0], [2, 0], [4, 0]])
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
data = Data(edge_index=edge_index, pos=pos)
data = transform(data)
assert len(data) == 2
assert data.pos.tolist() == [[-2, 0], [0, 0], [2, 0]]
assert data.edge_index.size() == (2, 7)
def test_compose():
transform = T.Compose([T.Center(), T.AddSelfLoops()])
assert transform.__repr__() == ('Compose([\n'
' Center(),\n'
' AddSelfLoops(),\n'
'])')
pos = torch.tensor([[0, 0], [2, 0], [4, 0]], dtype=torch.float)
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
data = Data(edge_index=edge_index, pos=pos)
out = transform(data)
assert out.pos.tolist() == [[-2, 0], [0, 0], [2, 0]]
assert out.edge_index.tolist() == [[0, 0, 1, 1, 1, 2, 2],
[0, 1, 0, 1, 2, 1, 2]]
import torch_geometric
#from torch_geometric.transforms import GridSampling
import torch_geometric.transforms as T
from human_corres.data import ImgData, ImgBatch
from human_corres.config import PATH_TO_SURREAL
import human_corres.transforms as H
num_views = 20
IDlist = np.stack([
np.arange(100000 * num_views),
np.arange(115000 * num_views, 215000 * num_views)
],
axis=0)
num_test = 5000
DefaultTransform = T.Compose([
T.Center(),
T.RandomRotate(30, axis=0),
T.RandomRotate(30, axis=1),
T.RandomRotate(30, axis=2),
])
class SurrealFEPts5k(Dataset):
"""Surreal 3D points for Feature Extraction (FE).
Samples a fixed number of points.
Output: dictionary with keys {points3d, correspondence}
Data Format:
points3d: [num_points, 3] real numbers.
correspondence: [num_points] integers in range [6890].
"""
def __init__(self, center):
self.is_center = center
if center:
self.center = T.Center()
else:
self.center = None
import scipy.io as sio
from torch.utils.data import DataLoader
from torch.utils.data import Dataset
import numpy as np
from human_corres.utils import helper
from human_corres.data import Data
import torch_geometric
import torch_geometric.transforms as T
import human_corres as hc
import human_corres.transforms as H
num_views = 20
IDlist = np.arange(5000 * num_views)
num_test = 50
TrainTransform = T.Compose([
T.Center(),
T.RandomRotate(30, axis=0),
T.RandomRotate(30, axis=1),
T.RandomRotate(30, axis=2),
H.GridSampling(0.01)
])
TestTransform = T.Compose([T.Center(), H.GridSampling(0.01)])
class DGFSurrealFEPts(Dataset):
"""Surreal 3D points for Feature Extraction (FE).
Output: dictionary with keys {points3d, correspondence}
Data Format:
points3d: [num_points, 3] real numbers.
correspondence: [num_points] integers in range [6890].
def test(model, loader, args, epoch, REG):
model.train()
handler = MultiGPUHandler(loader, args, training=False, reg=True)
CORRES = REG.replace('REG_', 'CORRES_')
start = 0
#import ipdb; ipdb.set_trace()
with progressbar.ProgressBar(max_value=len(loader),
widgets=handler.widgets) as bar:
for i, data_list in enumerate(loader):
if (start % 100 != 0) and ((start + 1) % 100 != 0) and (
(start + 2) % 100 != 0):
start += len(data_list)
continue
#print(data_list)
centered = [T.Center()(data).to(args.device) for data in data_list]
saved_pos = [data.pos.cpu().numpy() for data in centered]
saved_y = [data.y.cpu().numpy() for data in centered]
saved_pos = [
torch.as_tensor(pos, dtype=torch.float).to(args.device)
for pos in saved_pos
]
saved_y = [
torch.as_tensor(y, dtype=torch.float).to(args.device)
for y in saved_y
]
sampled = [
H.GridSampling(0.01)(data).to(args.device) for data in centered
]
out_dict = model(sampled)
corres = correspondence(out_dict['feats'], handler.template_feats)
pred_before_reg = []
pred_after_reg = []
offset = 0
x_before_reg = handler.template_points[corres, :]
x_after_reg = out_dict['x_out0']
for idx, data in enumerate(sampled):
length = data.pos.shape[0]
pred0 = Prediction(data.pos,
x_before_reg[offset:(offset + length)])
pred1 = Prediction(data.pos,
x_after_reg[offset:(offset + length)])
pred0 = pred0.knn_interpolate(saved_pos[idx], 3)
pred1 = pred1.knn_interpolate(saved_pos[idx], 3)
offset += length
#error_smpl = np.loadtxt(CORRES.replace('.reg', '.errors').format(start+idx)).reshape(-1, 2)
pred0.evaluate_errors(saved_y[idx][:, -3:])
pred1.evaluate_errors(saved_y[idx][:, -3:])
pred0.save_to_mat(CORRES.format(start + idx))
pred1.save_to_mat(REG.format(start + idx))
#print(errors0.mean(), errors1.mean(), error_smpl.mean(0))
#errors = torch.stack([errors0, errors1], dim=-1)
#np.savetxt(CORRES.format(start+idx), errors.cpu().numpy(), fmt='%.6f %.6f')
start += len(data_list)
#if (i % 10 == 0) or (i == len(loader)-1):
# handler.visualize(bar)
torch.cuda.empty_cache()
return {}
