import torch
from tqdm import tqdm
from torch_scatter import scatter
from utilities import *
dev = "cuda" if torch.cuda.is_available() else "cpu"
# load the graph
graph = torch.load("./data/girl_skate.pt")
graph = graph.to(dev)
# To visualize the noisy pointcloud :
displaySur(**dict(position=graph.x.cpu().numpy(), texture=graph.tex.cpu().numpy()))
class Jacobi_iterative(object):
"""
To apply Jacobi's fixed point iteration scheme
"""
def __init__(self, graph, **kwargs):
self.edge_index = graph.edge_index
self.edge_attr = graph.edge_attr
self.dim_size = kwargs["dim_size"]
self.f0 = kwargs["noisy_sig"]
self.lamb = kwargs["lamb"]
def run(self, itr, signal):
tol = 1e-15
for i in tqdm(range(itr)):
num = self.f0 + (2*self.lamb) * (scatter(self.edge_attr * signal[self.edge_index[0]], self.edge_index[1], dim=0, dim_size=self.dim_size, reduce="add"))
den = 1.0 + (2*self.lamb) * (scatter(self.edge_attr, self.edge_index[1], dim=0, dim_size=self.dim_size, reduce="add"))
tmp_signal = num/den
eps = torch.nn.MSELoss()(signal,tmp_signal)
signal = tmp_signal
if eps < tol :
print(f"Tolerance reached at itr:{i}!")
return signal
return signal
lamb = 0.2
itr = 100
hp = dict(noisy_sig=graph.tex, lamb=lamb, dim_size=graph.tex.shape[0])
j_i = Jacobi_iterative(graph, **hp)
new_sig = j_i.run(itr, graph.tex)
29%|██▉ | 29/100 [00:00<00:00, 429.23it/s]
Tolerance reached at itr:29!
# Display the processed pointcloud
displaySur(**dict(position=graph.x.cpu().numpy(), texture=new_sig.cpu().numpy()))
# You should see something like following:
# (Left) noisy pointcloud
# (Right) processed pointcloud
from matplotlib import image as mpimg
from matplotlib import pyplot as plt
res = mpimg.imread("./data/out_color.png")
f = plt.figure(figsize=(12,12))
plt.imshow(res, cmap="gray")
plt.show()