trying to figx data generation

This commit is contained in:
2023-02-18 00:53:41 +00:00
parent a94be773d3
commit e0d1916901
6 changed files with 317 additions and 222 deletions
+82 -79
View File
@@ -103,37 +103,43 @@ class Cache:
def read(self, main_id):
frames = [frame for frame in self.frames if frame.split()[0] == main_id]
for frame in frames:
id = frame.split()[0]
garment = frame.split()[2]
if id in self.cache.keys():
continue
# frames = [frame for frame in self.frames if (frame.split()[0] == main_id.split()[0] and frame.split()[2] == main_id.split()[1])]
# for frame in frames:
id = main_id.split()[0]
garment = main_id.split()[1]
self.cache[id] = Data()
self.cache[f'{id} {garment}'] = Data()
mask = read_garment_mask(self.folder, id, garment)
_, faces = _maskFaces(np.zeros((14475,3)), F, mask)
v = read_posed_garment_frame(folder, id, garment, 0)[mask,:]
m = tm.Trimesh(v, faces, process=False, validate=False)
outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
outliers = np.nonzero(outliers)[0]
outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
faces = np.concatenate((faces, outliers_faces), axis=0)
m = tm.Trimesh(v, faces, process=False, validate=False)
is_skirt = (garment == 'Dress' or garment == 'Skirt')
if self.include_neighbours:
mask = read_garment_mask(self.folder, id, garment)
_, faces = _maskFaces(np.zeros((14475,3)), F if not is_skirt else F_skirt, mask)
v = read_posed_garment_frame(folder, id, garment, 0)[mask,:]
m = tm.Trimesh(v, faces, process=False, validate=False)
outliers = np.array([len(v) for v in m.vertex_neighbors]) == 0
outliers = np.nonzero(outliers)[0]
outliers_n_neighbours = get_closest_n_points(v, outliers, 2)
outliers_faces = np.concatenate((outliers[:, np.newaxis], outliers_n_neighbours), axis=-1)
# sorted_faces = [sorted(f) for f in faces.astype(np.int64)]
# for outlier_face in outliers_faces:
# assert sorted(outlier_face) not in sorted_faces
faces = np.concatenate((faces, outliers_faces), axis=0)
m = tm.Trimesh(v, faces, process=False, validate=False)
if self.include_neighbours:
if not is_skirt:
k_neighbors = np.argsort(dist[mask, :][:, mask], axis=-1)[:, :self.neighbourhood]
self.cache[id].k_neighbors = k_neighbors
# G = tm.graph.vertex_adjacency_graph(m)
# for i in range(v.shape[0]):
# for k in range(25):
# assert nx.has_path(G, i, k_neighbors[i, k]), f'No path between {i} and {k_neighbors[i, k]}'
else:
k_neighbors = np.argsort(dist_skirt[mask, :][:, mask], axis=-1)[:, :self.neighbourhood]
self.cache[f'{id} {garment}'].k_neighbors = k_neighbors
G = tm.graph.vertex_adjacency_graph(m)
for i in range(v.shape[0]):
for k in range(25):
assert nx.has_path(G, i, k_neighbors[i, k]), f'No path between {i} and {k_neighbors[i, k]}'
self.cache[id].mask = mask
self.cache[id].faces = faces
self.cache[f'{id} {garment}'].mask = mask
self.cache[f'{id} {garment}'].faces = faces
def __getitem__(self, id):
if id not in self.cache.keys():
@@ -164,10 +170,10 @@ def calculate_dist(folder, neighbourhood, frames):
def generate_curvatures(folder, neighbourhood, frames):
# ids = os.listdir(folder)[int(args.start):int(args.end)]
# frames = frames[int(args.start):int(args.end)]
os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw'), exist_ok=True)
os.makedirs(os.path.join(folder + '_extras', f'mesh_curvature_raw'), exist_ok=True)
cache = Cache(folder, neighbourhood, frames)
filters = tb.Filters(complevel=9, complib='blosc2')
# filters = tb.Filters(complevel=9, complib='blosc2')
for frame in tqdm.tqdm(frames):
@@ -175,29 +181,29 @@ def generate_curvatures(folder, neighbourhood, frames):
frame_nr = int(frame.split()[1])
garment = frame.split()[2]
mask = cache[id].mask
faces = cache[id].faces
mask = cache[f'{id} {garment}'].mask
faces = cache[f'{id} {garment}'].faces
v = read_posed_garment_frame(folder, id, garment, frame_nr)[mask,:]
m = tm.Trimesh(v, faces, process=False, validate=False)
export_to_ply(m, temp_file)
cloud = PyntCloud.from_file(temp_file)
k_neighbors = cache[id].k_neighbors
k_neighbors = cache[f'{id} {garment}'].k_neighbors
ev = cloud.add_scalar_field("eigen_values", k_neighbors=k_neighbors)
cloud.add_scalar_field("curvature", ev=ev)
curvature = np.array(cloud.points["curvature(" + str(neighbourhood+1) + ")"])
out = np.zeros((14475), dtype=np.float32)
out[mask] = curvature
out[mask] = curvature.real if str(curvature.dtype) == 'complex64' else curvature
np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
if f'curvature_raw' in group:
array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw')
array[:] = out
else:
carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters)
carray[:] = out
np.savez_compressed(os.path.join(folder + '_extras', f'mesh_curvature_raw', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
# group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
# if f'curvature_raw' in group:
# array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_raw')
# array[:] = out
# else:
# carray = h5file.create_carray(group, "curvature_raw", atom=tb.Float32Atom(), shape=out.shape, filters=filters)
# carray[:] = out
return None
@@ -214,34 +220,34 @@ def get_bins(values, nr_of_bins):
def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
cache = Cache(folder, neighbourhood, frames, include_neighbours=False)
os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_global'), exist_ok=True)
os.makedirs(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_garment'), exist_ok=True)
os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_global'), exist_ok=True)
os.makedirs(os.path.join(folder + '_extras', f'curvature_mesh_garment'), exist_ok=True)
labels = np.linspace(0, bins.shape[0] - 2, bins.shape[0] - 1)
filters = tb.Filters(complevel=9, complib='blosc2')
# filters = tb.Filters(complevel=9, complib='blosc2')
for i, frame in enumerate(tqdm.tqdm(frames)):
id = frame.split()[0]
frame_nr = int(frame.split()[1])
garment = frame.split()[2]
curvature = curvatures[i][cache[id].mask]
curvature = curvatures[i][cache[f'{id} {garment}'].mask]
frame_labels = np.zeros((curvature.shape[0]), dtype=np.uint8)
for j in range(bins.shape[0] - 2):
frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
frame_labels[curvature >= bins[-2]] = labels[-1]
out = np.zeros((14475), dtype=np.uint8)
out[cache[id].mask] = frame_labels + 1
out[cache[f'{id} {garment}'].mask] = frame_labels + 1
group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
if f'curvature_mesh_global' in group:
array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global')
array[:] = out
else:
carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
carray[:] = out
# group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
# if f'curvature_mesh_global' in group:
# array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_global')
# array[:] = out
# else:
# carray = h5file.create_carray(group, "curvature_mesh_global", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
# carray[:] = out
np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_global', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
# writePC2(os.path.join(garment_dir, "curvature_label_" + str(nr_of_labels) + "_" + str(neighbourhood) + ".pc16"), unmasked_label, True)
# redo it with local bins
@@ -251,17 +257,17 @@ def generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood):
frame_labels[(curvature >= bins[j]) * (curvature < bins[j+1])] = labels[j]
frame_labels[curvature >= bins[-2]] = labels[-1]
out = np.zeros((14475), dtype=np.uint8)
out[cache[id].mask] = frame_labels + 1
out[cache[f'{id} {garment}'].mask] = frame_labels + 1
group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
if f'curvature_mesh_garment' in group:
array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment')
array[:] = out
else:
carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
carray[:] = out
# group = h5file.get_node(f'/{id}/{frame_nr}/{garment}')
# if f'curvature_mesh_garment' in group:
# array = h5file.get_node(f'/{id}/{frame_nr}/{garment}/curvature_mesh_garment')
# array[:] = out
# else:
# carray = h5file.create_carray(group, "curvature_mesh_garment", atom=tb.UInt8Atom(), shape=out.shape, filters=filters)
# carray[:] = out
np.savez_compressed(os.path.join('/'.join(folder.split('/')[:-1]), 'backup', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
np.savez_compressed(os.path.join(folder + '_extras', f'curvature_mesh_garment', f'{id}_{frame_nr}_{garment}.npz'), mesh_curvature_raw=out)
def visualize_curvature_labels(folder, frames):
@@ -319,6 +325,7 @@ def read_raw_curvatures(folder, frames, neighbourhood):
return curvatures_all
def read_raw_curvatures_2(folder, frames, neighbourhood):
curvatures_all = []
for frame in tqdm.tqdm(frames):
@@ -335,14 +342,13 @@ def read_raw_curvatures_2(folder, frames, neighbourhood):
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("-f", "--folder", help = "process the given folder", default=r'/home/workspace/data/dataset/cloth3d/train')
parser.add_argument("-s", "--start", help = "start of range", default=0)
parser.add_argument("-e", "--end", help = "end of range", default=600000)
parser.add_argument("-s", "--start", help = "start of range", default=300)
parser.add_argument("-e", "--end", help = "end of range", default=600)
parser.add_argument("-n", "--name", help = "end of range", default="lhunor")
args = parser.parse_args()
is_skirt = False
skirt_modifier = '_skirt' if is_skirt else ''
F = np.load("custom_utils/smpl_faces" + skirt_modifier + ".npy").astype(np.int32)
F = np.load("custom_utils/smpl_faces.npy").astype(np.int32)
F_skirt = np.load("custom_utils/smpl_faces_skirt.npy").astype(np.int32)
# k_neighbours(k=25)
@@ -353,13 +359,10 @@ if __name__ == '__main__':
neighbourhood = 25
temp_file = '/home/workspace/local-data/temp/' + ''.join(random.choice(string.ascii_lowercase) for i in range(10)) + '.ply'
if not is_skirt:
dist = np.load('custom_data/dijkstra.npy')
else:
dist = np.load('custom_data/dijkstra_skirt.npy')
dist = np.load('custom_data/dijkstra.npy')
dist_skirt = np.load('custom_data/dijkstra_skirt.npy')
print('SKIRT' if is_skirt else 'NOT SKIRT')
print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {is_skirt} -neigh {neighbourhood}')
print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -neigh {neighbourhood}')
# frames = sorted(os.listdir(os.path.join(folder, 'mesh_posed')))[int(args.start):int(args.end)]
# frames = [frame.split('.')[0] for frame in frames]
@@ -372,21 +375,21 @@ if __name__ == '__main__':
root = '/'.join(folder.split('/')[:-1])
type = folder.split('/')[-1]
h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
# h5file = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='a')
start = time.perf_counter()
# frames = [f'05235 {i} Tshirt' for i in range(300)]
# generate_curvatures(folder, neighbourhood, frames)
curvatures = read_raw_curvatures(folder, frames, neighbourhood)
generate_curvatures(folder, neighbourhood, frames)
# curvatures = read_raw_curvatures_2(folder, frames, neighbourhood)
# bins = get_bins(curvatures, nr_of_labels)
# np.save(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy', bins)
# print(bins)
bins = np.load(f'custom_data/{type}_bins{skirt_modifier}_{neighbourhood}.npy')
generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
# bins = np.load(f'custom_data/{type}_bins_{neighbourhood}.npy')
# generate_curvature_labels(folder, bins, curvatures, frames, neighbourhood)
# visualize_curvature_labels(folder, frames)
print(time.perf_counter() - start)
h5file.close()
pass