repo_description: [CVPR 2021, Oral] PREDATOR: Registration of 3D Point Clouds with Low Overlap. 
repo_topics: 3dvision attention-mechanism point-cloud registration transformer
code: """
  General utility functions
  
  Author: Shengyu Huang
  Last modified: 30.11.2020
  """
  
  import os,re,sys,json,yaml,random, argparse, torch, pickle
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  import numpy as np
  from scipy.spatial.transform import Rotation
  
  from sklearn.neighbors import NearestNeighbors
  from scipy.spatial.distance import minkowski
  _EPS = 1e-7  # To prevent division by zero
  
  
  class Logger:
      def __init__(self, path):
          self.path = path
          self.fw = open(self.path+'/log','a')
  
      def write(self, text):
          self.fw.write(text)
          self.fw.flush()
  
      def close(self):
          self.fw.close()
  
  def save_obj(obj, path ):
      """
      save a dictionary to a pickle file
      """
      with open(path, 'wb') as f:
          pickle.dump(obj, f)
  
  def load_obj(path):
      """
      read a dictionary from a pickle file
      """
      with open(path, 'rb') as f:
          return pickle.load(f)
  
  def load_config(path):
      """
      Loads config file:
  
      Args:
          path (str): path to the config file
  
      Returns: 
          config (dict): dictionary of the configuration parameters, merge sub_dicts
  
      """
      with open(path,'r') as f:
          cfg = yaml.safe_load(f)
      
      config = dict()
      for key, value in cfg.items():
          for k,v in value.items():
              config[k] = v
  
      return config
  
  
  def setup_seed(seed):
      """
      fix random seed for deterministic training
      """
      torch.manual_seed(seed)
      torch.cuda.manual_seed_all(seed)
      np.random.seed(seed)
      random.seed(seed)
      torch.backends.cudnn.deterministic = True
  
  def square_distance(src, dst, normalised = False):
      """
      Calculate Euclid distance between each two points.
      Args:
          src: source points, [B, N, C]
          dst: target points, [B, M, C]
      Returns:
          dist: per-point square distance, [B, N, M]
      """
      B, N, _ = src.shape
      _, M, _ = dst.shape
      dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
      if(normalised):
          dist += 2
      else:
          dist += torch.sum(src ** 2, dim=-1)[:, :, None]
          dist += torch.sum(dst ** 2, dim=-1)[:, None, :]
  
      dist = torch.clamp(dist, min=1e-12, max=None)
      return dist
      
  
  def validate_gradient(model):
      """
      Confirm all the gradients are non-nan and non-inf
      """
      for name, param in model.named_parameters():
          if param.grad is not None:
              if torch.any(torch.isnan(param.grad)):
                  return False
              if torch.any(torch.isinf(param.grad)):
                  return False
      return True
  
  
  def natural_key(string_):
      """
      Sort strings by numbers in the name
      """
      return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_)]