Update codes.py

codes.py

@@ -1,66 +1,173 @@
 import os, sys
 import numpy as np
-from PIL import Image
-import itertools
-import glob
 import random
+import time
+import json
+import matplotlib.pyplot as plt
+import tqdm
 import torch
-import torchvision
-import torchvision.transforms as transforms
-import torch.nn as nn
-import torch.optim as optim
-import torch.nn.functional as F
-from torch.nn.functional import relu as RLU
 
-registration_method = 'Additive_Recurence' #{'Rawblock', 'matching_points', 'Additive_Recurence', 'Multiplicative_Recurence'} #'recurrent_matrix', 
-imposed_point = 0
-Arch = 'ResNet'
-Fix_Torch_Wrap = False
-BW_Position = False
-dim = 128
-dim0 =224
-crop_ratio = dim/dim0
 
 
+def generate_standard_elips(N_samples = 10, a= 1,b = 1):
+    radius = 0.5
+    center = 0
+    N_samples1 = int(N_samples/2 - 1)
+    N_samples2 = N_samples - N_samples1
+    x1 = np.random.uniform((center-radius)*a,(center+radius)*a, size = N_samples1)
+    x1_ordered = np.sort(x1)
+    y1 = center + b*np.sqrt(radius**2 - ((x1_ordered-center)/a)**2)
+    x2 =  np.random.uniform((center-radius)*a,(center+radius)*a, size = N_samples - N_samples1)
+    x2_ordered = -np.sort(-x2) #the minus sign to sort descindingly
+    y2 = center - b*np.sqrt(radius**2 - ((x2_ordered-center)/a)**2)
+    x = np.concatenate([x1_ordered, x2_ordered], axis=0)
+    y = np.concatenate([y1, y2], axis=0)
+    return x, y
 
+def destandarize_point(d, dim=128):
+    return dim*(d + 0.5)
 
-class Identity(nn.Module):
+def To_pointcloud(x,y,z=0):
+    N_points = x.shape[0]
+    point_cloud = np.zeros([N_points,3])
+    point_cloud[:,0] = x
+    point_cloud[:,1] = y
+    if not z==0:
+        point_cloud[:,2] = z
+    return point_cloud
+
+def To_xyz(point_cloud):
+    x = point_cloud[:,0]
+    y = point_cloud[:,1]
+    z = point_cloud[:,2]
+    return x,y,z
+
+
+def random_rigid_transformation(dim=2):
+    #dim = 4
+    rotation_x = 0
+    rotation_y = 0
+    rotation_z = random.uniform(0, 2)*np.pi
+    translation_x = random.uniform(-1, 1)*dim
+    translation_y = random.uniform(-1, 1)*dim
+    translation_z = 0
+    reflection_x = random.sample([-1,1],1)[0]
+    reflection_y = random.sample([-1,1],1)[0]
+    reflection_z = 1
+    Rotx = np.array([[1,0,0],
+                     [0,np.cos(rotation_x),-np.sin(rotation_x)],
+                     [0,np.sin(rotation_x),np.cos(rotation_x)]])
+    Roty = np.array([[np.cos(rotation_y),0,np.sin(rotation_y)],
+                     [0,1,0],
+                     [-np.sin(rotation_y),0,np.cos(rotation_y)]])
+    Rotz = np.array([[np.cos(rotation_z),-np.sin(rotation_z),0],
+                     [np.sin(rotation_z),np.cos(rotation_z),0],
+                     [0,0,1]])
+    Rotation = np.matmul(Rotz, np.matmul(Roty,Rotx))
+    Reflection = np.array([[reflection_x,0,0],[0,reflection_y,0],[0,0,reflection_z]])
+    Translation = np.array([translation_x,translation_y,translation_z])
+    RefRotation = np.matmul(Reflection,Rotation)
+    return RefRotation, Translation
+
+
+
+def rigid_2Dtransformation(prdction):
+    #prediction = [rotationz, reflectionx, reflectiony, translationx, translationy]
+    N_examples = prdction['rotation'].shape[0]
+    Translation = prdction['translation']
+    Reflection = torch.zeros([N_examples,3,3])#
+    Reflection[:,0,0] = prdction['reflection'][:,0]
+    Reflection[:,1,1] = prdction['reflection'][:,1]
+    Reflection[:,2,2] = 1.
+    rotation_z = prdction['rotation'][:,2]
+    Rotation = torch.zeros([N_examples,3,3])#np.repeat(np.eye(3)[None,:,:],N_examples, axis=0))
+    Rotation[:,0,0] = torch.cos(rotation_z)
+    Rotation[:,1,1] = torch.cos(rotation_z)
+    Rotation[:,0,1] = -torch.sin(rotation_z)
+    Rotation[:,2,2] = 1.
+    Rotation[:,1,0] = torch.sin(rotation_z)
+    RefRotation = torch.matmul(Reflection,Rotation)
+    return RefRotation, Translation
+
+def batch_hausdorff_prcnt_distance(batch_point_cloud1, point_cloud2, percentile = 0.95):
+    assert point_cloud2.shape[0]==3
+    assert batch_point_cloud1.shape[1]==3
+    distances = torch.norm(batch_point_cloud1[:, :, None,:] - point_cloud2[None, :, :,None], dim=1)
+    dists1 = torch.min(distances, dim=1).values
+    dists2 = torch.min(distances, dim=2).values
+    # Calculate the 95th percentile distance
+    percentile_95 = torch.quantile(torch.cat([dists1, dists2],axis=1), percentile, interpolation='linear', dim=1)
+    return percentile_95
+
+def HDloss(prd, pointcloud_source_norm_torch,pointcloud_target_norm_torch, percentile = 0.95):
+    A, b = rigid_2Dtransformation(prd)
+    point_cloud_wrapped = torch.matmul(A, pointcloud_source_norm_torch.T) + b[:,:,None]
+    loss = batch_hausdorff_prcnt_distance(point_cloud_wrapped, pointcloud_target_norm_torch.T, percentile)
+    return loss
+
+def Mean_HDloss(prd, pointcloud_source_norm_torch, pointcloud_target_norm_torch, percentile = 0.95):
+    loss = HDloss(prd, pointcloud_source_norm_torch, pointcloud_target_norm_torch, percentile = 0.95)
+    return torch.mean(loss)
+
+
+def wrap_pointcloud(record, pointcloud_source):
+    #normalize first
+    PC1_mean = np.mean(pointcloud_source, axis=0)
+    pointcloud_source_norm = pointcloud_source - PC1_mean
+    pointcloud_source_norm_torch = torch.tensor(pointcloud_source_norm, requires_grad=False).to(torch.float32)
+    # find Tx
+    A, b = rigid_2Dtransformation(record)
+    point_cloud_wrapped = torch.matmul(A, pointcloud_source_norm_torch.T) + b[:,:,None]
+    return point_cloud_wrapped
+
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+def move_dict2device(dictionary,device):
+    for key in list(dictionary.keys()):
+        dictionary[key] = dictionary[key].to(device)
+    return dictionary
+
+eps = 0.000001948
+class Optimization_model(torch.nn.Module):
     def __init__(self):
-        super(Identity, self).__init__()
-    def forward(self, x):
-        return x
-
-class Build_IRmodel_Resnet(nn.Module):
-    def __init__(self, resnet_model, registration_method = 'Additive_Recurence', BW_Position=False):
-        super(Build_IRmodel_Resnet, self).__init__()
-        self.resnet_model = resnet_model
-        self.BW_Position = BW_Position
-        self.N_parameters = 6
-        self.registration_method = registration_method
-        self.fc1 =nn.Linear(6, 64)
-        self.fc2 =nn.Linear(64, 128*3)
-        self.fc3 =nn.Linear(512, self.N_parameters)
+        super(Optimization_model, self).__init__()
+        self.alpha = torch.nn.Parameter(torch.tensor(0.5, requires_grad=True))
+        self.rotation = torch.nn.Parameter(torch.tensor([0.0, 0.0, 0.0], requires_grad=True))
+        self.translation = torch.nn.Parameter(torch.tensor([0.01,-0.01, 0.0], requires_grad=True))
+        self.reflection = torch.nn.Parameter(torch.sign(torch.tensor([0.01,-0.01, 1], requires_grad=True)))
+        #self.rigid = torch.nn.Parameter(torch.tensor([0.0,1.0,1.0,0.1,0.1], requires_grad=True))
     def forward(self, input_X_batch):
-        source = input_X_batch['source']
-        target = input_X_batch['target']
-        if 'Recurence' in self.registration_method:
-            M_i = input_X_batch['M_i'].view(-1, 6)
-            M_rep = F.relu(self.fc1(M_i))
-            M_rep = F.relu(self.fc2(M_rep)).view(-1,3,1,128)
-            concatenated_input = torch.cat((source,target,M_rep), dim=2)
-        else:
-            concatenated_input = torch.cat((source,target), dim=2)
-        resnet_output = self.resnet_model(concatenated_input)
-        predicted_line = self.fc3(resnet_output)
-        if 'Recurence' in self.registration_method:
-            predicted_part_mtrx = predicted_line.view(-1, 2, 3)
-            Prd_Affine_mtrx = predicted_part_mtrx + input_X_batch['M_i']
-            predction = {'predicted_part_mtrx':predicted_part_mtrx,
-                            'Affine_mtrx': Prd_Affine_mtrx}
-        else:
-            Prd_Affine_mtrx = predicted_line.view(-1, 2, 3)
-            predction = {'Affine_mtrx': Prd_Affine_mtrx}
-        return predction
+        predicted_rotation = self.alpha*self.rotation + (1-self.alpha)*input_X_batch['rotation']
+        predicted_translation = self.alpha*self.translation + (1-self.alpha)*input_X_batch['translation']
+        predicted_reflection = torch.sign(self.alpha*self.reflection +
+                                          (1-self.alpha)*input_X_batch['reflection']+eps)
+        return {'rotation':predicted_rotation,
+                'translation':predicted_translation,
+                'reflection':predicted_reflection}
+
+class Dataset(torch.utils.data.Dataset):
+  def __init__(self, dataset_size, N_dim = 2):
+        self.dataset_size = dataset_size
+        self.N_dim = 2
+  def __len__(self):
+        return int(self.dataset_size)
+  def __getitem__(self, index):
+        rotation = np.pi*(-1 + 2*torch.rand([3]))
+        translation = -0.1 + 0.2*torch.rand([3])
+        reflection  = torch.sign(torch.rand([3]) - 0.5)
+        if self.N_dim == 2:
+            rotation[0:2]=0
+            translation[2]=0
+            reflection[2]=1
+        random_solution = {'rotation':rotation,
+                'translation':translation,
+                'reflection':reflection}
+        return random_solution
+
+
+
+
+
 
 def pil_to_numpy(im):
     im.load()
@@ -80,93 +187,27 @@ def pil_to_numpy(im):
         raise RuntimeError("encoder error %d in tobytes" % s)
     return data
 
-def load_image_pil_accelerated(image_path, dim=128):
-    image = Image.open(image_path).convert("RGB")
-    array = pil_to_numpy(image)
-    tensor = torch.from_numpy(np.rollaxis(array,2,0)/255).to(torch.float32)
-    tensor = torchvision.transforms.Resize((dim,dim))(tensor)
-    return tensor
 
 
-def workaround_matrix(Affine_mtrx0, acc = 2):
-    # To find the equivalent torch-compatible matrix from a correct matrix set acc=2 #This will be needed for transforming an image
-    # To find the correct Affine matrix from Torch compatible matrix set acc=0.5
-    Affine_mtrx_adj = inv_AM(Affine_mtrx0)
-    Affine_mtrx_adj[:,:,2]*=acc
-    return Affine_mtrx_adj
-
-def inv_AM(Affine_mtrx):
-    AM3 = mtrx3(Affine_mtrx)
-    AM_inv = torch.linalg.inv(AM3)
-    return AM_inv[:,0:2,:]
-
-def mtrx3(Affine_mtrx):
-    mtrx_shape = Affine_mtrx.shape
-    if len(mtrx_shape)==3:
-        N_Mbatches = mtrx_shape[0]
-        AM3 = torch.zeros( [N_Mbatches,3,3])#.to(device)
-        AM3[:,0:2,:] = Affine_mtrx
-        AM3[:,2,2] = 1
-    elif len(mtrx_shape)==2:
-        N_Mbatches = 1
-        AM3 = torch.zeros([3,3])#.to(device)
-        AM3[0:2,:] = Affine_mtrx
-        AM3[2,2] = 1 
-    return AM3
-
-def standarize_point(d, dim=128, flip = False):
-    if flip:
-        d = -d
-    return d/dim - 0.5
-
-def destandarize_point(d, dim=128, flip = False):
-    if flip:
-        d = -d
-    return dim*(d + 0.5)
 
-def generate_standard_elips(N_samples = 100, a= 1,b = 1):
-    radius = 0.25
-    center = 0
-    N_samples1 = int(N_samples/2 - 1)
-    N_samples2 = N_samples - N_samples1
-    x1 = torch.distributions.uniform.Uniform(center-radius,center + radius).sample([N_samples1])
-    x1_ordered = torch.sort(x1).values
-    y1 = center + b*torch.sqrt(radius**2 - ((x1_ordered-center)/a)**2)
-    x2 = torch.distributions.uniform.Uniform(center-radius,center + radius).sample([N_samples2])
-    x2_ordered = torch.sort(x2, descending=True).values
-    y2 = center - b*torch.sqrt(radius**2 - ((x2_ordered-center)/a)**2)
-    x = torch.cat([x1_ordered, x2_ordered])
-    y = torch.cat([y1, y2])
-    return x, y
 
-def transform_standard_points(Affine_mat, x,y):
-    XY = torch.ones([3,x.shape[0]])
-    XY[0,:]= x
-    XY[1,:]= y
-    XYt = torch.matmul(Affine_mat.to('cpu').detach(),XY)
-    xt0 = XYt[0]
-    yt0 = XYt[1]
-    return xt0, yt0
-
-def wrap_points(img, x_source, y_source, l=1, DIM =dim):
-    for i in range(len(y_source)):
-        x0 = x_source[i].int()
-        y0 = y_source[i].int()
-        if (x0<DIM) and (x0>0) and (y0<DIM) and (y0>0):
-            img[:,:,y0-l:y0+l,x0-l:x0+l] = 0
-    return img
 
 
-def wrap_imge_cropped(Affine_mtrx, source_img, dim1=224, dim2=128):
-  source_img224 = torch.nn.ZeroPad2d(int((dim1-dim2)/2))(source_img)
-  grd = torch.nn.functional.affine_grid(Affine_mtrx, size=source_img224.shape,align_corners=False)
-  wrapped_img = torch.nn.functional.grid_sample(source_img224, grid=grd,
-                                                mode='bilinear', padding_mode='zeros', align_corners=False)
-  wrapped_img = torchvision.transforms.CenterCrop((dim2, dim2))(wrapped_img)
-  return wrapped_img
 
 
 
+
+
+
+
+def load_image_pil_accelerated(image_path, dim=128):
+    image = Image.open(image_path).convert("RGB")
+    array = pil_to_numpy(image)
+    tensor = torch.from_numpy(np.rollaxis(array,2,0)/255).to(torch.float32)
+    tensor = torchvision.transforms.Resize((dim,dim))(tensor)
+    return tensor
+
+
 def preprocess_image(image_path, dim = 128):
     img = torch.zeros([1,3,dim,dim])
     img[0] = load_image_pil_accelerated(image_path, dim)