-
Notifications
You must be signed in to change notification settings - Fork 2
/
Copy pathmain.py
394 lines (352 loc) · 18.5 KB
/
main.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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
from typing import List
import heapq
from queue import Queue
import numpy as np
from numpy.linalg import norm
from tqdm import tqdm
import time
import functools
def timed(func):
# 函数运行计时
@functools.wraps(func)
def timed_wrapper(*args, **kwargs):
print(f'{func.__name__}', end=' ')
start_time = time.perf_counter()
result = func(*args, **kwargs)
end_time = time.perf_counter()
print(f'{end_time - start_time:.2f}s')
return result
return timed_wrapper
class NeighborInfo:
def __init__(self, e_vids, fid, angle, ang_dis, geo_dis):
self.vids = e_vids # 邻边的两个顶点的id
self.fid = fid # 邻面的id
self.angle = angle # 与邻面的夹角
self.ang_dis = ang_dis # 与邻面的角距离
self.geo_dis = geo_dis # 与邻面的测地距离
self.dis = 0 # 与邻面的距离度量
class Face:
def __init__(self, vs, vids): # vs: 三个顶点的位置,vid: 三个顶点的id
self.vids = vids
self.center = sum(vs) / 3 # 三个顶点的质心位置
n = np.cross(vs[1] - vs[0], vs[2] - vs[0]) # 法向量
norm_len = norm(n)
self.norm = n if norm_len < 1e-12 else n / norm_len # 单位法向量
self.label = 0 # 分割标签
self.nbrs: List[NeighborInfo] = [] # 所有邻面的信息
class Model:
@staticmethod
def read_ply(ply_path):
vertices, faces, v_num, f_num = [], [], 0, 0
with open(ply_path, 'r') as f:
lines = [line.strip() for line in f.readlines()]
for i, line in enumerate(lines):
if line.startswith('element vertex'):
v_num = int(line.split(' ')[-1])
if line.startswith('element face'):
f_num = int(line.split(' ')[-1])
if line == 'endheader':
break
for line in lines[-(v_num + f_num):-f_num]:
x, y, z = line.split(' ')[:3]
vertices.append([float(x), float(y), float(z)])
for line in lines[-f_num:]:
v1, v2, v3 = line.split(' ')[1:4]
faces.append([int(v1), int(v2), int(v3)])
return np.array(vertices), np.array(faces)
def __init__(self, ply_path):
self.vs, fs = Model.read_ply(ply_path) # 所有顶点位置
self.fs = [Face(self.vs[f], f) for f in fs] # 所有面片
# STEP 1:计算所有邻居的信息
self.avg_ang_dis = 0
self.compute_neighbor()
# STEP 2:计算任意两个面片间的最短路径
self.f_dis = np.full((len(self.fs), len(self.fs)), np.inf)
self.compute_shortest()
self.label_nums = 0 # 目前所有面片已分割出多少种类
@staticmethod
def compute_dis(f0: Face, f1: Face, e_vs): # e_vs是邻边的两个顶点位置
# 计算两个面片之间的角距离
angle = np.arccos(np.dot(f0.norm, f1.norm))
is_convex = np.dot(f0.norm, f1.center - f0.center) < 1e-12
eta = 0.2 if is_convex else 1.0 # 根据是否是凸的决定eta
ang_dis = eta * (1 - np.dot(f0.norm, f1.norm))
# 计算两个面片之间的测地距离
# 计算方法:如果将两个面片展平,测地距离就是两质心相连的直线段,构成一个三角形,夹角就是angle0+angle1,用余弦定理就能算出距离
axis, d0, d1 = e_vs[1] - e_vs[0], f0.center - e_vs[0], f1.center - e_vs[0] # 3个共起点的向量
axis_len, d0_len, d1_len = norm(axis), norm(d0), norm(d1)
angle0 = np.arccos(np.dot(d0, axis) / d0_len / axis_len)
angle1 = np.arccos(np.dot(d1, axis) / d1_len / axis_len)
geo_dis = d0_len * d0_len + d1_len * d1_len - 2 * d0_len * d1_len * np.cos(angle0 + angle1)
return angle, ang_dis, geo_dis
def compute_neighbor(self): # 将所有面片的邻边信息计算出来
class Edge:
def __init__(self, ev, fid):
self.vids = (ev[0], ev[1]) if ev[0] < ev[1] else (ev[1], ev[0]) # 两个顶点的id,升序
self.fid = fid # 面片id
es = [] # 所有棱边
for i, vids in enumerate([f.vids for f in self.fs]):
es.extend([Edge((vids[0], vids[1]), i), Edge((vids[1], vids[2]), i), Edge((vids[2], vids[0]), i)])
# 为所有面片找到与之相邻的面片,并计算相邻面片的角距离ang、测地距离geo和总距离
visited_es = {}
for e in es:
if e.vids not in visited_es:
visited_es[e.vids] = e.fid
else: # 遇到第二次才进行计算
f0, f1 = visited_es[e.vids], e.fid
angle, ang_dis, geo_dis = Model.compute_dis(self.fs[f0], self.fs[f1], self.vs[list(e.vids)])
self.fs[f0].nbrs.append(NeighborInfo(e.vids, f1, angle, ang_dis, geo_dis))
self.fs[f1].nbrs.append(NeighborInfo(e.vids, f0, angle, ang_dis, geo_dis))
count = sum([len(f.nbrs) for f in self.fs])
self.avg_ang_dis = sum([sum([n.ang_dis for n in f.nbrs]) for f in self.fs]) / count
avg_geo_dis = sum([sum([n.geo_dis for n in f.nbrs]) for f in self.fs]) / count
delta = 0.8
for f in self.fs:
for n in f.nbrs:
n.dis = (1 - delta) * n.ang_dis / self.avg_ang_dis + delta * n.geo_dis / avg_geo_dis
@timed
def compute_shortest(self):
import multiprocessing
import functools
import dijkstra
f_nbrs_id = [[n.fid for n in f.nbrs] for f in self.fs]
f_nbrs_dis = [[n.dis for n in f.nbrs] for f in self.fs]
assert all([len(x) == 3 for x in f_nbrs_id]) # 要求每个面片都有3个邻居
f_nbrs_id, f_nbrs_dis = np.array(f_nbrs_id), np.array(f_nbrs_dis)
num_proc = 6 # 最短路算法并行,并行进程数
with multiprocessing.Pool(num_proc) as p:
results = [res for res in p.imap(functools.partial(dijkstra.dijkstra_c, f_nbrs_id, f_nbrs_dis),
np.array_split(list(range(len(self.fs))), num_proc))]
self.f_dis = np.concatenate(results, axis=0)
def compute_flow(self, f_types):
# f_types是所有面片目前的类型:无关区域0 边界区域1,2 模糊区域3。
# 对f_types不为0的这些面片求最大流,函数返回时,f_types的3都会变成1或2
# 参考:
# Ford-Fulkerson增广路算法-EK算法。无向图等价于就直接用两个有向边就行。
# https://www.desgard.com/2020/03/03/max-flow-ford-fulkerson.html
# https://oi-wiki.org/graph/flow/max-flow/
# 初始化流图,用于计算最大流
class FlowEdge:
def __init__(self, fro, to, cap, flow):
self.fro = fro
self.to = to
self.cap = cap
self.flow = flow
# es[x],点x的所有邻边信息(起始点为x),最后两个是用于的源点和汇点
es: List[List[FlowEdge]] = [[] for _ in range(len(self.fs) + 2)]
for i, f in enumerate(self.fs):
for n in f.nbrs:
es[i].append(FlowEdge(i, n.fid, 1 / (1 + n.ang_dis / self.avg_ang_dis), 0))
start, target = len(es) - 2, len(es) - 1
f_types = np.array(list(f_types) + [4, 4])
for i, f_type in enumerate(f_types):
if f_type == 1:
es[start].append(FlowEdge(start, i, float('inf'), 0))
elif f_type == 2:
es[i].append(FlowEdge(i, target, float('inf'), 0))
for i in range(len(es)):
if f_types[i]:
for e in es[i]:
e.flow = 0
sum_flow = 0 # 总最大流
while True:
# STEP1: 从源点一直BFS,碰到汇点就停
p = [-1] * (len(self.fs) + 2) # p[x], BFS中x的父结点
a = [0.0] * (len(self.fs) + 2) # a[x]: BFS中x的父结点给到的流量
a[start] = float('inf')
q, q_history = Queue(), [] # 用q_history存储历史
q.put(start)
while not q.empty():
cur = q.get()
for e in es[cur]: # 遍历cur的所有边e,如果e没有被搜索到且还有残余流量,就流
if f_types[e.to] != 0 and not a[e.to] and e.cap > e.flow:
p[e.to] = cur # 设置父边
a[e.to] = min(a[cur], e.cap - e.flow) # 设置流量
q.put(e.to) # 继续搜
q_history.append(e.to)
if a[target]:
break
flow_add = a[target]
if not flow_add: # 已经搜索不到了,结束,BFS涉及的区域,都是1
for i in q_history:
f_types[i] = 1
for i in range(len(self.fs)):
if f_types[i] == 3:
f_types[i] = 2
break
# STEP2: 反向传播流量
cur = target
while cur != start:
for e in es[p[cur]]:
if e.to == cur:
e.flow += flow_add # 增加路径的flow值
for e in es[cur]:
if e.to == p[cur]:
e.flow -= flow_add # 减小反向路径的flow值
cur = p[cur]
sum_flow += flow_add
return f_types
def write_ply(self, ply_path):
with open(ply_path, 'w') as f:
f.write(f"ply\nformat ascii 1.0\n"
f"element vertex {len(self.vs)}\nproperty float x\nproperty float y\nproperty float z\n"
f"element face {len(self.fs)}\nproperty list uchar int vertex_indices\n"
f"property uint8 red\nproperty uint8 green\nproperty uint8 blue\n"
f"end_header\n")
for v in self.vs:
f.write(f'{v[0]} {v[1]} {v[2]}\n')
for face in self.fs:
f.write(f'3 {face.vids[0]} {face.vids[1]} {face.vids[2]} ')
label = face.label
f.write(f'{60 * (label % 4 + 1)} {80 * ((label + 1) % 3 + 1)} {50 * ((label + 2) % 5 + 1)}\n')
print(f'save to {ply_path}')
class Segment:
def __init__(self, model, level, fids=None):
# global
self.model = model
self.fids = fids if fids else list(range(len(model.fs))) # 只考虑哪些面片
self.level = level # 层次化分解的深度
# local
self.fs = [model.fs[fid] for fid in self.fids]
self.f_dis = model.f_dis[self.fids][:, self.fids]
def average(arr):
return np.sum(arr) / (len(arr) * (len(arr) - 1))
# 计算一些所需的最大距离/平均距离
local_max_dis_fids = np.unravel_index(self.f_dis.argmax(), self.f_dis.shape) # 最远的一对面片
self.global_max_dis = self.model.f_dis.max() - self.model.f_dis.min()
assert self.global_max_dis != np.inf # 要求是连通图
self.global_avg_dis, self.local_avg_dis = average(self.model.f_dis), average(self.f_dis)
def k_way_reps():
# 选到其他各点距离之和最小的点初始点
reps, G = [np.argmin(np.sum(self.f_dis, axis=1))], []
for i in range(20): # 20个试验点,每次新加入一个种子使最近的种子最远
rep, max_dis = 0, 0.0
for j in range(len(self.f_dis)):
min_dis = np.min([self.f_dis[j][reps]])
if min_dis > max_dis:
max_dis = min_dis
rep = j
reps.append(rep)
G.append(max_dis)
num = np.argmax([G[num] - G[num + 1] for num in range(len(G) - 2)]) + 2 # 最大化G[num]-G[num+1]
# NOTE: 如果想进行二路分解,只需要直接在此处设置num=2
return num, reps[:num] # 种子数和代表点
self.num, self.reps = k_way_reps() # 注意: reps是local变量,即reps[i]的值作为self.f_dis的索引而不是self.model.f_dis
# 用uniques记录不重复的reps的索引,有可能出现种子重复的情况,只考虑不重复的种子
self.uniques = np.sort(np.unique(self.reps, return_index=True)[1])
if self.num == 2:
rep0, rep1 = sorted(local_max_dis_fids)
self.reps[0], self.reps[1] = rep0, rep1
# 计算同类相邻面片夹角的极差
max_ang, min_ang = 0, np.pi
for f in self.fs:
for n in f.nbrs:
if self.model.fs[n.fid].label == f.label:
min_ang, max_ang = min(n.angle, min_ang), max(n.angle, max_ang)
self.ang_diff = max_ang - min_ang
def seg(self):
prob = np.zeros((self.num, len(self.f_dis)))
OFFSET, FUZZY = self.model.label_nums, self.model.label_nums + self.num # OFFSET之前共有多少种类,FUZZY标志模糊区域
def compute_prob():
for fid in range(len(self.f_dis)):
if fid in self.reps:
prob[self.reps.index(fid)][fid] = 1
continue
sum_prob = sum([1 / self.f_dis[fid][self.reps[u]] for u in self.uniques]) # 只考虑不重合的点
prob[:, fid] = 1 / self.f_dis[fid][self.reps] / sum_prob # 计算平均
def assign(): # 给面片打标签
eps = 0.04 if self.num <= 3 else 0.02 # 确定清晰区域的阈值, 这个参数需要调,对分割结果影响较大
counts = np.zeros(self.num) # 每个类别的数量
prob[[i for i in range(self.num) if i not in self.uniques]] = 0
for fid in range(len(self.f_dis)):
if len(self.uniques) > 1:
label1, label2 = heapq.nlargest(2, range(len(self.uniques)), prob[self.uniques, fid].take)
prob1, prob2 = prob[label1][fid], prob[label2][fid]
else:
label1, label2, prob1, prob2 = self.uniques[0], -1, 1.0, 0.0
if prob1 - prob2 > eps:
self.fs[fid].label = OFFSET + label1
counts[label1] += 1
else:
self.fs[fid].label = FUZZY + label1 * self.num + label2
def recompute_reps():
# STEP1: 用论文3.3节最后一段的改进方法来计算P(fi∈Sj)
assign() # 先粗分类一下
# rep_dis[k][i]表示第k类种子到第i个面片的距离,用面片i到所有第k类面片的的平均距离来表示
rep_dis = np.zeros((self.num, len(self.f_dis)))
counts = np.zeros(self.num) # 每个类别的数量
for k_f in range(len(self.f_dis)): # 计算总距离累加
k = self.fs[k_f].label - OFFSET
if k < self.num:
counts[k] += 1
# k_f是一个新发现的第k类的面片,每个面片i累加上k_f到i(也就是i到k_f)的距离。最终得到i到所有k类面片的距离累加
rep_dis[k] += self.f_dis[k_f]
for k in range(self.num):
rep_dis[k] = rep_dis[k] / counts[k] if counts[k] else np.inf # 求平均距离
prob[:] = 1 / (rep_dis + 1e-12) / np.sum(1 / (rep_dis + 1e-12), axis=0) # 计算概率P
# STEP2: 用论文3.3节第2部分计算新的种子
rep_cost = np.dot(prob, self.f_dis) # rep_cost[k][i]表示第k个类用i做种子的开销
reps = list(np.argmin(rep_cost, axis=1))
return reps, rep_cost
def assign_fuzzy():
for i in self.uniques: # 两片模糊区域间的面片两两分割
for j in self.uniques:
if j <= i:
continue
f_types = np.zeros(len(self.model.fs))
# STEP1: 确定具体分割的面片,找到哪些是模糊区域 3, 哪些是边界区域1,2, 哪些是无关区域0
for fid, f in zip(self.fids, self.fs):
if f.label == FUZZY + i * self.num + j or f.label == FUZZY + j * self.num + i:
f_types[fid] = 3
for n in f.nbrs:
nf = self.model.fs[n.fid]
if nf.label == OFFSET + i:
f_types[n.fid] = 1
elif nf.label == OFFSET + j:
f_types[n.fid] = 2
# STEP2: 计算分割
f_types = self.model.compute_flow(f_types)
# STEP3: 执行分割结果
for fid in self.fids:
if f_types[fid] == 1:
self.model.fs[fid].label = OFFSET + i
elif f_types[fid] == 2:
self.model.fs[fid].label = OFFSET + j
for _ in tqdm(range(20), desc=f'{self.level} step'):
# 论文3.3中的STEP1: 计算概率
compute_prob()
# 论文3.3中的STEP2: 重新计算reps
new_reps, cost = recompute_reps()
# STEP3: 判断是否更新
new_cost = [cost[i][rep] for i, rep in enumerate(new_reps)]
old_cost = [cost[i][rep] for i, rep in enumerate(self.reps)]
changed = any([(c1 < c0 - 1e-12 and r1 != r0)
for r1, r0, c1, c0 in zip(new_reps, self.reps, new_cost, old_cost)])
if changed:
self.reps = new_reps
self.uniques = np.sort(np.unique(self.reps, return_index=True)[1])
else:
break
recompute_reps()
assign() # 清晰部分
assign_fuzzy() # 模糊部分
self.model.label_nums += self.num
reps_f_dis = self.f_dis[self.reps][:, self.reps]
local_max_patch_dis = np.max(reps_f_dis)
if self.level > 0 or local_max_patch_dis / self.global_max_dis < 0.1: # 最多只递归一层,可调
return
# 递归
segments = []
for sid in range(self.num):
if sid in self.uniques:
# 先统一建好建所有segment,再统一seg,不然seg导致label被换了标记,取模运算会有问题
fids = [fid for fid in self.fids if self.model.fs[fid].label % self.num == sid]
segments.append(Segment(self.model, self.level + 1, fids))
for segment in segments:
if segment.ang_diff > 0.3 and segment.local_avg_dis / segment.global_avg_dis > 0.2:
segment.seg()
if __name__ == '__main__':
ply = 'horse' # 'dino' 'horse'
mesh_model = Model(f'data/{ply}.ply')
Segment(mesh_model, 0).seg()
mesh_model.write_ply(f'data/{ply}-output.ply')