def get_neighbours(point_map, p, shape):
i = util.R(p.x)
j = util.R(p.y)
if i < 0 or j < 0 or i >= shape[0] or j >= shape[1]: return []
neighbours = []
neighbours.extend(point_map[(i, j)])
if j != shape[1] - 1: neighbours.extend(point_map[(i, j + 1)])
if j != 0: neighbours.extend(point_map[(i, j - 1)])
if i != shape[0] - 1: neighbours.extend(point_map[(i + 1, j)])
if i != shape[0] - 1 and j != shape[1] - 1:
neighbours.extend(point_map[(i + 1, j + 1)])
if i != shape[0] - 1 and j != 0:
neighbours.extend(point_map[(i + 1, j - 1)])
if i != 0: neighbours.extend(point_map[(i - 1, j)])
if i != 0 and j != shape[1] - 1:
neighbours.extend(point_map[(i - 1, j + 1)])
if i != 0 and j != 0: neighbours.extend(point_map[(i - 1, j - 1)])
return neighbours
def get_cluster_points(f, img, shape, img_name, artist, dir="pixelClustering"):
print("IMAGE SHAPE: (%s, %s)" % (shape[0], shape[1]))
f.write("Speed factor: %f\n" % SPEED_FACTOR)
f.write("Stopping Point: %f\n" % STOPPING_POINT)
f.write("Min Speed: %f\n" % MIN_SPEED)
f.write("Min Neighbours: %f\n" % MIN_NEIGHBOURS)
f.write("Neighbourhood Radius: %f\n\n" % NBHD_RADIUS)
blur = cv2.GaussianBlur(img, (7, 7), 1)
util.display_img(blur, "Blurred", DISPLAY_IMG)
sobelx = cv2.Sobel(blur, cv2.CV_64F, 1, 0, ksize=5)
sobely = cv2.Sobel(blur, cv2.CV_64F, 0, 1, ksize=5)
scaled_x = cv2.convertScaleAbs(sobelx)
scaled_y = cv2.convertScaleAbs(sobely)
cv2.imwrite('%s/%s/%s/%s_sobelx.png' % (dir, artist, img_name, img_name),
scaled_x)
cv2.imwrite('%s/%s/%s/%s_sobely.png' % (dir, artist, img_name, img_name),
scaled_y)
moving_points, point_map = get_moving_pixels(f, img, sobelx, sobely)
print("Number of moving points: %s" % len(moving_points))
f.write("Initial number of moving points: %d\n" % len(moving_points))
newImg = util.mapToImage(moving_points, img.shape)
cv2.imwrite(
'%s/%s/%s/%s_COLOR_%d.png' % (dir, artist, img_name, img_name, 1),
newImg)
#util.display_img(newImg, "Mapped from Point Set!", DISPLAY_IMG)
# MOVEMENT PHASE
start_time = time.time()
movement_phase = True
counter = 0
while movement_phase:
counter += 1
bad_pix = []
# movement phase
for i in xrange(len(moving_points)):
idx = (util.R(moving_points[i].x), util.R(moving_points[i].y))
x_speed = getMinSpeed(SPEED_FACTOR * moving_points[i].gradx)
y_speed = getMinSpeed(SPEED_FACTOR * moving_points[i].grady)
if moving_points[i].stopped: continue
moving_points[i].x = moving_points[i].x - x_speed
moving_points[i].y = moving_points[i].y - y_speed
idx2 = (util.R(moving_points[i].x), util.R(moving_points[i].y))
# update space in bin
if idx != idx2:
point_map[idx].remove(moving_points[i])
if idx2 in point_map:
point_map[idx2].append(moving_points[i])
else:
print("BAD POINT FOUND: (%.3f, %.3f)" %
(moving_points[i].x, moving_points[i].y))
bad_pix.append(moving_points[i])
for p in bad_pix:
moving_points.remove(p)
# evaluation phase
for p in moving_points:
stop = stop_pixel(point_map, p, shape)
if stop:
p.stopped = stop
print("PHASE COMPLETE.")
newImg2 = util.mapToImage_moving(moving_points, img.shape)
cv2.imwrite(
'%s/%s/%s/%s_COLOR_%d.png' % (dir, artist, img_name, img_name,
(counter + 1)), newImg2)
#util.display_img(newImg2, "Mapped from Point Set %s!" % (counter+1), DISPLAY_IMG
# Determine whether to break the loop
stp = get_percent_stopped_pixels(moving_points)
print("Percentage stopped: %f" % stp)
f.write("Rep %d. Percentage stopped: %f\n" % (counter, stp))
if stp > STOPPING_POINT:
print("STOPPING THE LOOP")
movement_phase = False
end_time = time.time()
moving_points = clean_up(moving_points)
f.write("\nTOTAL Number of iterations: %d\n" % counter)
f.write("Time taken (in seconds): %f\n" % (end_time - start_time))
f.write("Final number of moving points: %d\n" % len(moving_points))
newImg2 = util.mapToImage(moving_points, img.shape)
util.display_img(newImg2, "THE FINAL RESULT", DISPLAY_IMG)
cv2.imwrite(
'%s/%s/%s/%s_COLOR_%s.png' %
(dir, artist, img_name, img_name, "FINAL"), newImg2)
cv2.destroyAllWindows()
calculate_local_stroke_thickness(f, point_map, moving_points, shape)
util.save_point_set("%s.txt" % img_name, artist, "pointSet", moving_points)
return moving_points
def train():
# define model, dataloader, 3dmm eigenvectors, optimization method
model = densenet
model.load_state_dict(torch.load("model/chkpoint_000.pt"))
model.cuda()
#optimizer2 = torch.optim.Adam(model.parameters(),lr=1e-4)
adjustmentnet.load_state_dict(torch.load("model/chkpoint_adj_000.pt"))
adjustmentnet.cuda()
adjustmentnet.train()
loader = dataloader.BatchLoader
face3dmm = dataloader.Face3DMM()
optimizer = torch.optim.Adam(list(adjustmentnet.parameters()) +
list(model.parameters()),
lr=1e-5)
# main training loop
for epoch in itertools.count():
for i, batch in enumerate(loader):
optimizer.zero_grad()
x = batch['image'].cuda()
y = batch['lm2d'].cuda()
y_pred = model(x)
batchsize = x.shape[0]
alphas = y_pred[:, :199]
betas = y_pred[:, 199:228]
s = y_pred[:, 228]
t = y_pred[:, 229:231]
t = torch.tanh(t)
r = y_pred[:, 231:235]
r = torch.tanh(r) * (3.14 / 4)
# apply 3DMM model from predicted parameters
alpha_matrix = alphas.unsqueeze(2).expand(
*alphas.size(), alphas.size(1)) * torch.eye(
alphas.size(1)).cuda()
beta_matrix = betas.unsqueeze(2).expand(*betas.size(),
betas.size(1)) * torch.eye(
betas.size(1)).cuda()
shape_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.shape_eigenvec]),
alpha_matrix)
exp_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.exp_eigenvec]), beta_matrix)
shape_cov = shape_cov.sum(2)
exp_cov = exp_cov.sum(2)
# alignment
shape = (face3dmm.mu_shape.unsqueeze(0) + shape_cov.view(
(batchsize, 53215, 3))) + exp_cov.view((batchsize, 53215, 3))
lm = shape[:, face3dmm.lm, :]
R = util.R(r).cuda()
scaledshape = s.unsqueeze(1).unsqueeze(1) * torch.bmm(lm, R)
alignedshape = t.unsqueeze(1) + scaledshape[:, :, :2]
# adjustment network applied onto the landmarks of 3dMM taking input image
adjustment = torch.tanh(adjustmentnet(x).view(-1, 68, 2))
pred = (alignedshape + adjustment) * 112 + 112
gt = y * 112 + 112
# weight update
loss = torch.mean(torch.norm(gt - pred, p=2, dim=2))
loss.backward()
optimizer.step()
gt = gt[0].cpu().data.numpy()
pred = pred[0].cpu().data.numpy()
sample = x[0].cpu().permute(1, 2, 0).data.numpy()
sample = sample * 255
util.viewLM(sample.astype(np.uint8), pred)
#io.imsave(f"example_{i:04d}.png",sample)
print(f"epoch/batch {epoch}/{i} | Loss: {loss:.4f}")
print("saving!")
torch.save(model.state_dict(), f"model/chkpoint_{epoch:03d}.pt")
torch.save(adjustmentnet.state_dict(),
f"model/chkpoint_adj_{epoch:03d}.pt")
def train():
# define model, dataloader, 3dmm eigenvectors, optimization method
model = densenet
model.cuda()
adjustmentnet.cuda()
loader = dataloader.BatchLoader
face3dmm = dataloader.Face3DMM()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
# main training loop
for epoch in itertools.count():
for i, batch in enumerate(loader):
optimizer.zero_grad()
x = batch['image'].cuda()
y = batch['lm2d'].cuda()
y_pred = model(x)
batchsize = x.shape[0]
alphas = y_pred[:, :199]
betas = y_pred[:, 199:228]
s = y_pred[:, 228]
t = y_pred[:, 229:231]
t = torch.tanh(t)
r = y_pred[:, 231:235]
r = torch.tanh(r) * (3.14 / 4)
# apply 3DMM model from predicted parameters
alpha_matrix = alphas.unsqueeze(2).expand(
*alphas.size(), alphas.size(1)) * torch.eye(
alphas.size(1)).cuda()
beta_matrix = betas.unsqueeze(2).expand(*betas.size(),
betas.size(1)) * torch.eye(
betas.size(1)).cuda()
shape_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.shape_eigenvec]),
alpha_matrix)
exp_cov = torch.bmm(
torch.stack(batchsize * [face3dmm.exp_eigenvec]), beta_matrix)
shape_cov = shape_cov.sum(2)
exp_cov = exp_cov.sum(2)
# alignment
shape = (face3dmm.mu_shape.unsqueeze(0) + shape_cov.view(
(batchsize, 53215, 3))) + exp_cov.view((batchsize, 53215, 3))
lm = shape[:, face3dmm.lm, :]
R = util.R(r).cuda()
scaledshape = s.unsqueeze(1).unsqueeze(1) * torch.bmm(lm, R)
alignedshape = t.unsqueeze(1) + scaledshape[:, :, :2]
# adjustment network applied onto the landmarks of 3dMM taking input image
# adjustment = adjustmentnet(x).view(-1,68,2)
# weight update
loss = torch.mean(torch.abs(y - (alignedshape)))
loss.backward()
optimizer.step()
print(f"epoch/batch {epoch}/{i} | Loss: {loss:.4f}")
print("saving!")
torch.save(model.state_dict(), f"model/chkpoint_{epoch:03d}.pt")
def extract_topology(f, moving_points, point_map, img, corners, artist,
img_name):
# encode some constants
f.write("SCALE FACTOR: %d\n" % SCALE_FACTOR)
f.write("PRUNING RADIUS: %.4f\n" % PRUNING_RADIUS)
G = create_cluster_graph(f, point_map, moving_points, img.shape)
""" Compute Minimum Spanning Tree """
start_time = time.time()
mst = nx.minimum_spanning_tree(G)
end_time = time.time()
print("NUMBER OF MST NODES: %d" % mst.number_of_nodes())
print("NUMBER OF MST EDGES: %d" % mst.number_of_edges())
f.write("NUMBER OF MST NODES: %d\n" % mst.number_of_nodes())
f.write("NUMBER OF MST EDGES: %d\n" % mst.number_of_edges())
f.write("Time to compute mst (s): %s\n\n" % (end_time - start_time))
newImg = util.mapToImage(mst.nodes(), img.shape)
util.thicken_line(newImg)
util.display_img(newImg, "Post MST", DISPLAY_IMG)
cv2.imwrite(
"topology/%s/%s/%s_CompleteGraph.jpg" % (artist, img_name, img_name),
newImg)
""" Iterative Pruning """
f.write("--- ITERATIVE PRUNING ---\n")
num_spc_pts = len(moving_points)
prev_spc_pts = num_spc_pts # avoid infinite loops
streak_len = 0
counter = 0
endpoints = []
junctions = []
itr_pruning_start = time.time()
while num_spc_pts >= len(moving_points) // SCALE_FACTOR:
counter += 1
leaf_nodes = [x for x in mst.nodes_iter() if mst.degree(x) == 1]
for leaf in leaf_nodes:
total_weight = 0
curr_node = leaf
while total_weight < max(PRUNING_RADIUS, curr_node.lsw):
n = mst.neighbors(curr_node)
if len(n) != 1: break # used to be == 0
e = mst.get_edge_data(curr_node, n[0])
if total_weight + e['weight'] < max(PRUNING_RADIUS,
curr_node.lsw):
total_weight += e['weight']
mst.remove_node(curr_node)
curr_node = n[0]
else:
break
print("NUM LEAF NODES: %d" % len(leaf_nodes))
print("NUMBER OF PRUNED MST NODES: %d" % mst.number_of_nodes())
print("NUMBER OF PRUNED MST EDGES: %d" % mst.number_of_edges())
endpoints = [x for x in mst.nodes_iter() if mst.degree(x) == 1]
junctions = [x for x in mst.nodes_iter() if mst.degree(x) >= 3]
num_spc_pts = len(endpoints) + len(junctions)
print("REP %d: %d points, %d special points" %
(counter, len(moving_points), num_spc_pts))
f.write("REP %d: %d points, %d special points\n" %
(counter, len(moving_points), num_spc_pts))
mstImg = util.mapToImage(mst.nodes(), img.shape)
colorImg = mstImg.copy()
colorImg = cv2.cvtColor(colorImg, cv2.COLOR_GRAY2BGR)
for endpoint in endpoints:
cv2.circle(colorImg, (util.R(endpoint.y), util.R(endpoint.x)), 3,
(0, 0, 255), -1)
for junction in junctions:
cv2.circle(colorImg, (util.R(junction.y), util.R(junction.x)), 3,
(255, 0, 0), -1)
util.display_img(colorImg, "Post MST Pruning %d" % counter,
DISPLAY_IMG)
cv2.imwrite(
"topology/%s/%s/%s_Pruned_%d.jpg" %
(artist, img_name, img_name, counter), colorImg)
# AVOID INFINITE LOOPS
if prev_spc_pts == num_spc_pts:
if streak_len < 2:
prev_spc_pts = num_spc_pts
streak_len += 1
else:
break
else:
prev_spc_pts = num_spc_pts
itr_pruning_end = time.time()
print("\n--- FINAL POINTSET ---")
print("%d points, %d special points\n" %
(len(moving_points), (len(endpoints) + len(junctions))))
f.write("\n--- FINAL POINTSET ---\n")
f.write("%d points, %d special points\n" %
(len(moving_points), (len(endpoints) + len(junctions))))
f.write("Pruning time (in s): %s\n\n" %
(itr_pruning_end - itr_pruning_start))
mstImg = util.mapToImage(mst.nodes(), img.shape)
colorImg = mstImg.copy()
colorImg = cv2.cvtColor(colorImg, cv2.COLOR_GRAY2BGR)
for corner in corners:
x, y = corner.ravel()
cv2.circle(colorImg, (x, y), 3, (0, 255, 0), -1)
cv2.imwrite(
"topology/%s/%s/%s_CORNERS_mst.jpg" % (artist, img_name, img_name),
colorImg)
f.write("NUMBER OF Connected Components: %d\n" %
nx.number_connected_components(mst))
loner_nodes = [x for x in mst.nodes_iter() if mst.degree(x) == 0]
mst.remove_nodes_from(loner_nodes)
mst.remove_nodes_from(endpoints)
mst.remove_nodes_from(junctions)
junction_endpoints = endpoints + junctions
all_pairs = np.transpose([
np.tile(junction_endpoints, len(junction_endpoints)),
np.repeat(junction_endpoints, len(junction_endpoints))
])
print("Number of pairs: %d" % len(all_pairs))
f.write("Number of endpoint-junction pairs: %d\n" % len(all_pairs))
print("NUMBER OF CC's: %d" % nx.number_connected_components(mst))
f.write("NEW NUMBER OF Connected Components: %d\n" %
nx.number_connected_components(mst))
components_subgraphs = nx.connected_component_subgraphs(mst)
point_pairs = []
counter = 0
cc_start_time = time.time()
for comp in components_subgraphs:
counter += 1
endpoints = [x for x in comp.nodes_iter() if comp.degree(x) == 1]
if len(endpoints) != 2: continue
min_dist = 1000
pair = []
for p0, p1 in all_pairs:
dist1 = util.pointDist(p0, endpoints[0]) + util.pointDist(
p1, endpoints[1])
dist2 = util.pointDist(p0, endpoints[1]) + util.pointDist(
p1, endpoints[0])
if dist1 < min_dist:
min_dist = dist1
pair = (p0, p1)
if dist2 < min_dist:
min_dist = dist2
pair = (p1, p0)
point_pairs.append(pair)
testImg = img.copy()
testImg = cv2.cvtColor(testImg, cv2.COLOR_GRAY2BGR)
for node in comp.nodes_iter():
cv2.circle(testImg, (util.R(node.y), util.R(node.x)), 3,
(0, 255, 0), -1)
cv2.circle(testImg, (util.R(pair[0].y), util.R(pair[0].x)), 3,
(0, 0, 255), -1)
cv2.circle(testImg, (util.R(pair[1].y), util.R(pair[1].x)), 3,
(0, 0, 255), -1)
util.display_img(testImg, "Test", False)
if len(comp.nodes()) > MIN_NODES_DETECTED:
cv2.imwrite(
"topology/%s/%s/%s_DetectedCurve_%d.jpg" %
(artist, img_name, img_name, counter), testImg)
cc_end_time = time.time()
print("NUMBER OF POINT PAIRS: %d\n" % len(point_pairs))
f.write("NUMBER OF POINT PAIRS: %d\n" % len(point_pairs))
f.write("Time to Find Point Pairs (in s): %s\n\n" %
(cc_end_time - cc_start_time))
util.save_buddy_points("%s.txt" % img_name, artist, "buddyPoints",
point_pairs)
f.write(" --- Calculating Paths --- \n")
curves = []
counter = 0
curve_start = time.time()
for pair in point_pairs:
counter += 1
path = nx.dijkstra_path(G, pair[0], pair[1])
curves.append(path)
print("Length of path: %d" % len(path))
f.write("Length of path: %d\n" % len(path))
testImg = img.copy()
testImg = cv2.cvtColor(testImg, cv2.COLOR_GRAY2BGR)
for node in path:
cv2.circle(testImg, (util.R(node.y), util.R(node.x)), 3,
(0, 0, 255), -1)
util.display_img(testImg, "Test", False)
if len(path) > MIN_NODES_ACTUAL:
cv2.imwrite(
"topology/%s/%s/%s_ActualCurve_%d.jpg" %
(artist, img_name, img_name, counter), testImg)
curve_end = time.time()
util.save_curves("%s.txt" % img_name, artist, "curves", curves)
f.write("\nTime to Seperate Curves (in s): %s\n\n" %
(curve_end - curve_start))
alpha_matrix = alphas.unsqueeze(2).expand(
*alphas.size(), alphas.size(1)) * torch.eye(alphas.size(1)).cuda()
beta_matrix = betas.unsqueeze(2).expand(
*betas.size(), betas.size(1)) * torch.eye(betas.size(1)).cuda()
shape_cov = torch.bmm(torch.stack(batchsize * [face3dmm.shape_eigenvec]),
alpha_matrix)
exp_cov = torch.bmm(torch.stack(batchsize * [face3dmm.exp_eigenvec]),
beta_matrix)
shape_cov = shape_cov.sum(2)
exp_cov = exp_cov.sum(2)
# alignment
shape = (face3dmm.mu_shape.unsqueeze(0) + shape_cov.view(
(batchsize, 53215, 3))) + exp_cov.view((batchsize, 53215, 3))
lm = shape[:, face3dmm.lm, :]
R = util.R(r).cuda()
scaledshape = s.unsqueeze(1).unsqueeze(1) * torch.bmm(lm, R)
alignedshape = t.unsqueeze(1) + scaledshape[:, :, :2]
# adjustment network applied onto the landmarks of 3dMM taking input image
adjustment = torch.tanh(adjustmentnet(x).view(-1, 68, 2))
pred = adjustment * 112 + 112
gt = y * 112 + 112
# weight update
gt = gt[0].cpu().data.numpy()
pred = pred[0].cpu().data.numpy()
le = gt[36, :]
re = gt[45, :]
d = np.linalg.norm(le - re)