-
Notifications
You must be signed in to change notification settings - Fork 1
/
view_psd_data.py
69 lines (56 loc) · 2.04 KB
/
view_psd_data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
import matplotlib
from matplotlib import cm
import numpy as np
import trimesh
import os
from glob import glob
import json
import argparse
class NpEncoder(json.JSONEncoder):
def default(self, obj):
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.ndarray):
return obj.tolist()
def read_psd_shape(file):
file = glob(file)[0]
mesh = trimesh.load(file, process=False, maintain_order=True)
return mesh
def read_psd_segmask(file):
file = glob(file)[0]
with open(file, 'r') as f:
lines = f.readlines()
mask = []
for i, line in enumerate(lines):
if i % 2 == 1:
mask.append([int(x)-1 for x in line.strip().split()])
return mask
def visualize_psd_shape(shape_file, segmask_file, textured=False):
mesh = read_psd_shape(shape_file)
mask = read_psd_segmask(segmask_file)
if textured:
norm = matplotlib.colors.Normalize(0, 20, clip=True)
mapper = cm.ScalarMappable(norm=norm, cmap=cm.tab20)
for cid, cmask in enumerate(mask):
# print(cid, len(cmask))
color = mapper.to_rgba(cid%20)
color = np.array(color)*255
color = color.astype(np.uint8)
# color = trimesh.visual.random_color()
# print(color)
mesh.visual.face_colors[cmask] = color
return mesh, mask
if __name__ == "__main__":
folder = "/home/lyang/yl_code/dataset/labeledDb/LabeledDB_new/*"
parser = argparse.ArgumentParser('Segmentation GUI')
parser.add_argument('--input', type=str, default='167', help='Input mesh path.')
args = parser.parse_args()
## load mesh
shape_id = args.input
fpath = f"./data/segmentation_data/*/{shape_id}.off"
mesh, mask = visualize_psd_shape(fpath, fpath.replace(".off", "_labels.txt"), textured=True)
# mask_path = f"{shape_id}_mask.json"
# with open(mask_path, 'w') as f:
# json.dump(mask, f, cls=NpEncoder, ensure_ascii=False, indent=4)
# mesh.export(f"{shape_id}.obj")
mesh.show()