[IEEE T-RO 2023] Source code of RING and RING++ for loop closure detection in LiDAR SLAM.
import os
import sys
import glob
import random
import numpy as np
import pandas as pd
import utils.config as cfg
from typing import List, Tuple
from utils.poses import xyz_ypr2m
from utils.tools import find_nearest_ndx, check_in_train_set, check_in_test_set
import matplotlib.pyplot as plt
from sklearn.neighbors import KDTree
from scipy.spatial import distance_matrix
from torch.utils.data import Dataset, ConcatDataset, DataLoader


# calibration matrix
def read_calib_file(filepath):
    """Read in a calibration file and parse into a dictionary."""
    data = {}

    with open(filepath, 'r') as f:
        for line in f.readlines():
            key, value = line.split(':', 1)
            # The only non-float values in these files are dates, which
            # we don't care about anyway
            try:
                data[key] = np.array([float(x) for x in value.split()])
            except ValueError:
                pass

    cam_to_velo = np.reshape(data['Tr'], (3, 4))
    cam_to_velo = np.vstack([cam_to_velo, [0, 0, 0, 1]])
    return cam_to_velo


def load_pc_kitti(file_pathname: str):
    # Load point cloud from file
    pc = np.fromfile(file_pathname, dtype=np.float32)
    pc = np.reshape(pc, (-1, 4)) # x,y,z,intensity
    pc = pc[:,:3]

    return pc


class KITTIPointCloudLoader:
    def __init__(self, remove_zero_points: bool = True, remove_ground_plane: bool = True, ground_plane_level: float = -1.0):
        # remove_zero_points: remove points with all zero coordinates
        # remove_ground_plane: remove points on ground plane level and below
        # ground_plane_level: ground plane level
        self.remove_zero_points = remove_zero_points
        self.remove_ground_plane = remove_ground_plane
        self.ground_plane_level = ground_plane_level

    def read_pc(self, file_pathname: str):
        # Reads the point cloud and preprocess (optional removal of zero points and points on the ground plane
        pc = load_pc_kitti(file_pathname)
        if self.remove_zero_points:
            mask = np.all(np.isclose(pc, 0), axis=1)
            pc = pc[~mask]

        if self.remove_ground_plane:
            mask = pc[:, 2] > self.ground_plane_level
            pc = pc[mask]        
        
        pc = np.unique(pc, axis=0) # remove duplicate points  

        return pc
    
    def normalize_pc(self, pc, xbound: float = cfg.point_cloud["x_bound"][1], ybound: float = cfg.point_cloud["y_bound"][1], zbound: float = 10.):
        # Normalize the point cloud
        pc = np.array(pc, dtype=np.float32)
        pc = pc[:,:3]

        mask = (np.abs(pc[:,0]) < xbound) * (np.abs(pc[:,1]) < ybound) * (pc[:,2] > self.ground_plane_level) * (pc[:,2] < zbound)

        pc = pc[mask]
        pc[:,0] = pc[:,0] / xbound
        pc[:,1] = pc[:,1] / ybound
        pc[:,2] = pc[:,2] / zbound 

        return pc   
    

class KITTISequence(Dataset):
    """
    Single KITTI sequence indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_name: str, split: str = 'train', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        self.dataset_root = dataset_root
        self.sequence_name = sequence_name
        self.sequence_path = os.path.join(self.dataset_root, self.sequence_name)
        assert os.path.exists(self.sequence_path), f'Cannot access sequence: {self.sequence_path}'
        self.split = split
        self.sampling_distance = sampling_distance
        # Maximum discrepancy between timestamps of LiDAR scan and global pose in seconds
        self.pose_time_tolerance = 1.
        self.pose_file = os.path.join(self.sequence_path, 'poses.txt')
        assert os.path.exists(self.pose_file), f'Cannot access global pose file: {self.pose_file}'
        self.lidar_path = os.path.join(self.sequence_path, 'velodyne')
        assert os.path.exists(self.lidar_path), f'Cannot access lidar scans: {self.lidar_path}'
        # self.lidar_files = sorted(glob.glob(os.path.join(self.lidar_path, '*.bin')))
        self.pc_loader = KITTIPointCloudLoader()
        self.times_file = os.path.join(self.sequence_path, 'times.txt')    
        assert os.path.exists(self.times_file), f'Cannot access sequence times file: {self.times_file}'
        self.calib_file = os.path.join(self.sequence_path, 'calib.txt')     
        assert os.path.exists(self.calib_file), f'Cannot access sequence calib file: {self.calib_file}'
        self.timestamps, self.filepaths, self.poses, self.xys = self.get_scan_poses()
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_name}-{split}')
        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.filepaths)

    def __getitem__(self, ndx):
        pc_filepath = self.filepaths[ndx]
        pc = self.load_pc(pc_filepath)
        return {'ts': self.timestamps[ndx], 'pc': pc, 'position': self.xys[ndx], 'pose': self.poses[ndx]}

    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs
        
    def get_scan_poses(self):
        # Read global poses from .txt file and link each lidar_scan with the nearest pose
        # threshold: timestamp threshold in seconds
        # Returns a dictionary with (4, 4) pose matrix indexed by a timestamp (as integer)
        calib = read_calib_file(self.calib_file)   
        with open(self.pose_file, "r") as h:
            pose_data = h.readlines()

        n = len(pose_data)
        print(f'Number of global poses: {n}')
        system_timestamps = np.genfromtxt(self.times_file)
        poses = np.zeros((n, 4, 4), dtype=np.float32)       # 4x4 pose matrix

        for ndx, pose in enumerate(pose_data):
            # Split by comma and remove whitespaces
            temp = [e.strip() for e in pose.split(' ')]
            assert len(temp) == 12, f'Invalid line in global poses file: {temp}'
            # poses in kitti are in cam0 reference
            poses[ndx] = np.array([[float(temp[0]), float(temp[1]), float(temp[2]), float(temp[3])],
                                    [float(temp[4]), float(temp[5]), float(temp[6]), float(temp[7])],
                                    [float(temp[8]), float(temp[9]), float(temp[10]), float(temp[11])],
                                    [0., 0., 0., 1.]])
            poses[ndx] = np.linalg.inv(calib) @ (poses[ndx] @ calib)

        # List LiDAR scan timestamps
        all_lidar_timestamps = [int(os.path.splitext(f)[0]) for f in os.listdir(self.lidar_path) if
                                os.path.splitext(f)[1] == '.bin']
        all_lidar_timestamps.sort()        
        print(f'Number of global scans: {len(all_lidar_timestamps)}')

        lidar_timestamps = []
        lidar_filepaths = []
        lidar_poses = []
        count_rejected = 0

        for ndx, lidar_ts in enumerate(all_lidar_timestamps):
            # Find index of the closest timestamp
            closest_ts_ndx = find_nearest_ndx(lidar_ts, system_timestamps)
            delta = abs(system_timestamps[closest_ts_ndx] - lidar_ts)
            # Timestamp is in nanoseconds = 1e-9 second
            if delta > self.pose_time_tolerance * 1000000000:
                # Reject point cloud without corresponding pose
                count_rejected += 1
                continue

            lidar_timestamps.append(lidar_ts)
            lidar_filepaths.append(os.path.join(self.lidar_path, f'{lidar_ts}.bin'))
            lidar_poses.append(poses[closest_ts_ndx])

        lidar_timestamps = np.array(lidar_timestamps, dtype=np.int64)
        lidar_filepaths = np.array(lidar_filepaths)
        lidar_poses = np.array(lidar_poses, dtype=np.float32)     # 4x4 pose matrix
        lidar_xys = lidar_poses[:, :2, 3]                         # 2D position
        
        # Split data into train / test set
        if self.split != 'all':      
            if self.split == 'train':
                mask = check_in_train_set(lidar_xys, dataset='kitti')
            elif self.split == 'test':
                mask = check_in_test_set(lidar_xys, dataset='kitti')
            lidar_timestamps = lidar_timestamps[mask]
            lidar_filepaths = lidar_filepaths[mask]
            lidar_poses = lidar_poses[mask]
            lidar_xys = lidar_xys[mask]

        # Sample lidar scans 
        prev_position = None
        mask = []
        for ndx, position in enumerate(lidar_xys):
            if prev_position is None:
                mask.append(ndx)
                prev_position = position
            else:
                displacement = np.linalg.norm(prev_position - position)
                if displacement > self.sampling_distance:
                    mask.append(ndx)
                    prev_position = position

        lidar_timestamps = lidar_timestamps[mask]
        lidar_filepaths = lidar_filepaths[mask]
        lidar_poses = lidar_poses[mask]
        lidar_xys = lidar_xys[mask]

        print(f'{len(lidar_timestamps)} scans with valid pose, {count_rejected} rejected due to unknown pose')
        return lidar_timestamps, lidar_filepaths, lidar_poses, lidar_xys

    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


class KITTISequences(Dataset):
    """
    Multiple NCLT sequences indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_names: List[str], split: str = 'train', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        self.dataset_root = dataset_root
        self.sequence_names = sequence_names
        self.split = split
        self.sampling_distance = sampling_distance
        
        # Load all sequences
        sequences = []
        for sequence_name in sequence_names:
            sequence = KITTISequence(dataset_root, sequence_name, split, sampling_distance)
            sequences.append(sequence)
        self.dataset = ConcatDataset(sequences)
        self.pc_loader = sequences[0].pc_loader

        # Concatenate all sequences
        self.timestamps = np.concatenate([s.timestamps for s in sequences])
        self.filepaths = np.concatenate([s.filepaths for s in sequences])
        self.poses = np.concatenate([s.poses for s in sequences])
        self.xys = np.concatenate([s.xys for s in sequences])
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_names}-{split}')

        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, ndx):
        return self.dataset[ndx]

    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs
        
    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


if __name__ == "__main__":
    # load dataset
    base_path = './data/KITTI/'
    folder = '08' # sequence folder for training   
    train_dataset = KITTISequence(base_path, folder, split='train')
    test_dataset = KITTISequence(base_path, folder, split='test')
    train_positions = train_dataset.xys
    test_positions = test_dataset.xys
    # plot the splited results
    plt.scatter(train_positions[:,0], train_positions[:,1], c='r', s=0.1)
    plt.scatter(test_positions[:,0], test_positions[:,1], c='b', s=0.1)
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.legend(['train', 'test'])
    plt.title(f'{folder} Trajectory Split')    
    plt.savefig(f'{folder}_split.png')
    plt.show()
    
datasets/KITTIDataset.py
lus6-Jenny/RING
import os
import sys
import glob
import random
import numpy as np
import pandas as pd
import utils.config as cfg
from typing import List, Tuple
from utils.poses import xyz_ypr2m
from utils.tools import find_nearest_ndx, check_in_train_set, check_in_test_set
import matplotlib.pyplot as plt
from sklearn.neighbors import KDTree
from scipy.spatial.transform import Rotation
from torch.utils.data import Dataset, ConcatDataset, DataLoader


# calibrate the lidar data of mulran dataset
def get_mulran_calib():
    calib = np.eye(4)
    calib[0,3] = 1.7042
    calib[1,3] = -0.021
    calib[2,3] = 1.8047
    calib[:3,:3] = Rotation.from_euler('xyz',[0.0001,0.0003,179.6654],degrees=True).as_matrix()
    return calib


def load_pc_mulran(file_pathname: str):
    # Load point cloud from file
    n_vec = 4
    pc = np.fromfile(file_pathname, dtype=np.float32)
    pc = np.reshape(pc, (-1, n_vec))
    pc = pc[:,:3]
    return pc


class MulRanPointCloudLoader:
    def __init__(self, remove_zero_points: bool = True, remove_ground_plane: bool = True, ground_plane_level: float = -0.9):
        # remove_zero_points: remove points with all zero coordinates
        # remove_ground_plane: remove points on ground plane level and below
        # ground_plane_level: ground plane level
        self.remove_zero_points = remove_zero_points
        self.remove_ground_plane = remove_ground_plane
        self.ground_plane_level = ground_plane_level

    def read_pc(self, file_pathname: str):
        # Reads the point cloud and preprocess (optional removal of zero points and points on the ground plane
        pc = load_pc_mulran(file_pathname)

        if self.remove_zero_points:
            mask = np.all(np.isclose(pc, 0), axis=1)
            pc = pc[~mask]

        if self.remove_ground_plane:
            mask = pc[:, 2] > self.ground_plane_level
            pc = pc[mask]        

        pc = np.unique(pc, axis=0) # remove duplicate points  

        return pc
    
    def normalize_pc(self, pc, xbound: float = cfg.point_cloud["x_bound"][1], ybound: float = cfg.point_cloud["y_bound"][1], zbound: float = 19.):
        # Normalize the point cloud
        pc = np.array(pc, dtype=np.float32)
        pc = pc[:,:3]

        mask = (np.abs(pc[:,0]) < xbound) * (np.abs(pc[:,1]) < ybound) * (pc[:,2] > self.ground_plane_level) * (pc[:,2] < zbound)

        pc = pc[mask]
        pc[:,0] = pc[:,0] / xbound
        pc[:,1] = pc[:,1] / ybound
        pc[:,2] = (pc[:,2] - self.ground_plane_level) / (zbound - self.ground_plane_level) 

        return pc   
    

class MulRanSequence(Dataset):
    """
    Single MulRan sequence indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_name: str, split: str = 'train', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        self.dataset_root = dataset_root
        self.sequence_name = sequence_name
        self.sequence_path = os.path.join(self.dataset_root, self.sequence_name)
        assert os.path.exists(self.sequence_path), f'Cannot access sequence: {self.sequence_path}'
        self.split = split
        self.sampling_distance = sampling_distance
        # Maximum discrepancy between timestamps of LiDAR scan and global pose in seconds
        self.pose_time_tolerance = 1.
        self.pose_file = os.path.join(self.sequence_path, 'global_pose.csv')
        assert os.path.exists(self.pose_file), f'Cannot access global pose file: {self.pose_file}'
        self.lidar_path = os.path.join(self.sequence_path, 'Ouster')
        assert os.path.exists(self.lidar_path), f'Cannot access lidar scans: {self.lidar_path}'
        # self.lidar_files = sorted(glob.glob(os.path.join(self.lidar_path, '*.bin')))
        self.pc_loader = MulRanPointCloudLoader()
        self.timestamps, self.filepaths, self.poses, self.xys = self.get_scan_poses()
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_name}-{split}')
        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.filepaths)

    def __getitem__(self, ndx):
        pc_filepath = self.filepaths[ndx]
        pc = self.load_pc(pc_filepath)
        return {'ts': self.timestamps[ndx], 'pc': pc, 'position': self.xys[ndx], 'pose': self.poses[ndx]}

    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs
        
    def get_scan_poses(self):
        # Read global poses from .csv file and link each lidar_scan with the nearest pose
        # threshold: timestamp threshold in seconds
        # Returns a dictionary with (4, 4) pose matrix indexed by a timestamp (as integer)
        pose_data = pd.read_csv(self.pose_file, header=0, \
            names=['timestamp','x1','x2','x3','northing','y1','y2','y3','easting','z1','z2','z3','height'], low_memory=False)

        n = len(pose_data)
        print(f'Number of global poses: {n}')
        system_timestamps = np.zeros((n,), dtype=np.int64)
        poses = np.zeros((n, 4, 4), dtype=np.float32)       # 4x4 pose matrix
        calib = get_mulran_calib()
        calib = np.linalg.inv(calib)

        for ndx in range(n):
            row = pose_data.iloc[ndx]
            assert len(row) == 13, f'Invalid line in global poses file: {row}'
            ts = int(row['timestamp'])
            # transform to se3 matrix
            se3 = np.eye(4, dtype=np.float32)            
            se3[0,0] = float(row['x1'])
            se3[0,1] = float(row['x2'])
            se3[0,2] = float(row['x3'])
            se3[0,3] = float(row['northing'])
            se3[1,0] = float(row['y1'])
            se3[1,1] = float(row['y2'])
            se3[1,2] = float(row['y3'])
            se3[1,3] = float(row['easting'])
            se3[2,0] = float(row['z1'])
            se3[2,1] = float(row['z2'])
            se3[2,2] = float(row['z3'])
            se3[2,3] = float(row['height'])
            # add to system_timestamps and poses
            system_timestamps[ndx] = int(ts)
            poses[ndx] = se3 @ calib

        # Ensure timestamps and poses are sorted in ascending order
        sorted_ndx = np.argsort(system_timestamps, axis=0)
        system_timestamps = system_timestamps[sorted_ndx]
        poses = poses[sorted_ndx]

        # List LiDAR scan timestamps
        all_lidar_timestamps = [int(os.path.splitext(f)[0]) for f in os.listdir(self.lidar_path) if
                                os.path.splitext(f)[1] == '.bin']
        all_lidar_timestamps.sort()
        print(f'Number of global scans: {len(all_lidar_timestamps)}')

        lidar_timestamps = []
        lidar_filepaths = []
        lidar_poses = []
        count_rejected = 0

        for ndx, lidar_ts in enumerate(all_lidar_timestamps):
            # Find index of the closest timestamp
            closest_ts_ndx = find_nearest_ndx(lidar_ts, system_timestamps)
            delta = abs(system_timestamps[closest_ts_ndx] - lidar_ts)
            # Timestamp is in nanoseconds = 1e-9 second
            if delta > self.pose_time_tolerance * 1000000000:
                # Reject point cloud without corresponding pose
                count_rejected += 1
                continue

            lidar_timestamps.append(lidar_ts)
            lidar_filepaths.append(os.path.join(self.lidar_path, f'{lidar_ts}.bin'))
            lidar_poses.append(poses[closest_ts_ndx])
        
        lidar_timestamps = np.array(lidar_timestamps, dtype=np.int64)
        lidar_filepaths = np.array(lidar_filepaths)
        lidar_poses = np.array(lidar_poses, dtype=np.float32)     # 4x4 pose matrix
        lidar_xys = lidar_poses[:, :2, 3]                         # 2D position

        # split data into train / test set
        if self.split != 'all':      
            if self.split == 'train':
                mask = check_in_train_set(lidar_xys, dataset='mulran')
            elif self.split == 'test':
                mask = check_in_test_set(lidar_xys, dataset='mulran')
            lidar_timestamps = lidar_timestamps[mask]
            lidar_filepaths = lidar_filepaths[mask]
            lidar_poses = lidar_poses[mask]
            lidar_xys = lidar_xys[mask]

        # Sample lidar scans             
        prev_position = None
        mask = []
        for ndx, position in enumerate(lidar_xys):
            if prev_position is None:
                mask.append(ndx)
                prev_position = position
            else:
                displacement = np.linalg.norm(prev_position - position)
                if displacement > self.sampling_distance:
                    mask.append(ndx)
                    prev_position = position

        lidar_timestamps = lidar_timestamps[mask]
        lidar_filepaths = lidar_filepaths[mask]
        lidar_poses = lidar_poses[mask]
        lidar_xys = lidar_xys[mask]

        print(f'{len(lidar_timestamps)} scans with valid pose, {count_rejected} rejected due to unknown pose')
        return lidar_timestamps, lidar_filepaths, lidar_poses, lidar_xys

    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


class MulRanSequences(Dataset):
    """
    Multiple MulRan sequences indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_names: List[str], split: str = 'train', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        self.dataset_root = dataset_root
        self.sequence_names = sequence_names
        self.split = split
        self.sampling_distance = sampling_distance
        
        # Load all sequences
        sequences = []
        for sequence_name in sequence_names:
            sequence = MulRanSequence(dataset_root, sequence_name, split, sampling_distance)
            sequences.append(sequence)
        self.dataset = ConcatDataset(sequences)
        self.pc_loader = sequences[0].pc_loader

        # Concatenate all sequences
        self.timestamps = np.concatenate([s.timestamps for s in sequences])
        self.filepaths = np.concatenate([s.filepaths for s in sequences])
        self.poses = np.concatenate([s.poses for s in sequences])
        self.xys = np.concatenate([s.xys for s in sequences])
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_names}-{split}')

        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, ndx):
        return self.dataset[ndx]

    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs

    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


if __name__ == "__main__":
    # load dataset
    base_path = './data/MulRan/'
    folder = 'Sejong01' # sequence folder for training    
    train_dataset = MulRanSequence(base_path, folder, split='train')
    test_dataset = MulRanSequence(base_path, folder, split='test')
    train_positions = train_dataset.xys
    test_positions = test_dataset.xys
    # plot the splited results
    plt.scatter(train_positions[:,0], train_positions[:,1], c='r', s=0.1)
    plt.scatter(test_positions[:,0], test_positions[:,1], c='b', s=0.1)
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.legend(['train', 'test'])
    plt.title(f'{folder} Trajectory Split')
    plt.savefig(f'{folder}_split.png')
    plt.show()
       
datasets/MulRanDataset.py
lus6-Jenny/RING
import os
import sys
import glob
import random
import struct
import numpy as np
import pandas as pd
import utils.config as cfg
from typing import List, Tuple
from utils.poses import xyz_ypr2m
from utils.tools import find_nearest_ndx, check_in_train_set, check_in_test_set
import matplotlib.pyplot as plt
from sklearn.neighbors import KDTree
from scipy.spatial import distance_matrix
from torch.utils.data import Dataset, ConcatDataset, DataLoader


# nclt pointcloud utils
def convert(x_s, y_s, z_s):
    scaling = 0.005 # 5 mm
    offset = -100.0

    x = x_s * scaling + offset
    y = y_s * scaling + offset
    z = z_s * scaling + offset

    return x, y, z    


def load_pc_nclt(file_pathname: str):
    # Load point cloud from file
    hits = []
    with open(file_pathname,'rb') as f_bin:
        while True:
            x_str = f_bin.read(2)
            if x_str == b"": # eof
                break
            x = struct.unpack('<H', x_str)[0]
            y = struct.unpack('<H', f_bin.read(2))[0]
            z = struct.unpack('<H', f_bin.read(2))[0]
            i = struct.unpack('B', f_bin.read(1))[0]
            l = struct.unpack('B', f_bin.read(1))[0]
            x, y, z = convert(x, y, z)
            s = "%5.3f, %5.3f, %5.3f, %d, %d" % (x, y, z, i, l)
            hits += [[x, y, -z]] # flip z axis
        hits = np.asarray(hits)
    return hits


class NCLTPointCloudLoader:
    def __init__(self, remove_zero_points: bool = True, remove_ground_plane: bool = True, ground_plane_level: float = 1.0):
        # remove_zero_points: remove points with all zero coordinates
        # remove_ground_plane: remove points on ground plane level and below
        # ground_plane_level: ground plane level
        self.remove_zero_points = remove_zero_points
        self.remove_ground_plane = remove_ground_plane
        self.ground_plane_level = ground_plane_level

    def read_pc(self, file_pathname: str):
        # Reads the point cloud and preprocess (optional removal of zero points and points on the ground plane
        pc = load_pc_nclt(file_pathname)
        if self.remove_zero_points:
            mask = np.all(np.isclose(pc, 0), axis=1)
            pc = pc[~mask]

        if self.remove_ground_plane:
            mask = pc[:, 2] > self.ground_plane_level
            pc = pc[mask]    

        pc = np.unique(pc, axis=0) # remove duplicate points    

        return pc
    
    def normalize_pc(self, pc, xbound: float = cfg.point_cloud["x_bound"][1], ybound: float = cfg.point_cloud["y_bound"][1], zbound: float = 20.):
        # Normalize the point cloud
        pc = np.array(pc, dtype=np.float32)
        pc = pc[:,:3]

        mask = (np.abs(pc[:,0]) < xbound) * (np.abs(pc[:,1]) < ybound) * (pc[:,2] > self.ground_plane_level) * (pc[:,2] < zbound)

        pc = pc[mask]
        pc[:,0] = pc[:,0] / xbound
        pc[:,1] = pc[:,1] / ybound
        pc[:,2] = (pc[:,2] - self.ground_plane_level) / (zbound - self.ground_plane_level)  

        return pc   
    

class NCLTSequence(Dataset):
    """
    Single NCLT sequence indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_name: str, split: str = 'train', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        self.dataset_root = dataset_root
        self.sequence_name = sequence_name
        self.sequence_path = os.path.join(self.dataset_root, self.sequence_name)
        assert os.path.exists(self.sequence_path), f'Cannot access sequence: {self.sequence_path}'
        self.split = split
        self.sampling_distance = sampling_distance
        # Maximum discrepancy between timestamps of LiDAR scan and global pose in seconds
        self.pose_time_tolerance = 1.
        self.pose_file = os.path.join(self.sequence_path, 'ground_truth/groundtruth_' + sequence_name + '.csv')
        assert os.path.exists(self.pose_file), f'Cannot access global pose file: {self.pose_file}'
        self.lidar_path = os.path.join(self.sequence_path, 'velodyne_data/velodyne_sync')
        assert os.path.exists(self.lidar_path), f'Cannot access lidar scans: {self.lidar_path}'
        # self.lidar_files = sorted(glob.glob(os.path.join(self.lidar_path, '*.bin')))
        self.pc_loader = NCLTPointCloudLoader()
        self.timestamps, self.filepaths, self.poses, self.xys = self.get_scan_poses()
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_name}-{split}')
        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.filepaths)

    def __getitem__(self, ndx):
        pc_filepath = self.filepaths[ndx]
        pc = self.load_pc(pc_filepath)
        return {'ts': self.timestamps[ndx], 'pc': pc, 'position': self.xys[ndx], 'pose': self.poses[ndx]}
    
    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs
        
    def get_scan_poses(self):
        # Read global poses from .csv file and link each lidar_scan with the nearest pose
        # threshold: timestamp threshold in seconds
        # Returns a dictionary with (4, 4) pose matrix indexed by a timestamp (as integer)
        pose_data = pd.read_csv(self.pose_file, header=0, \
            names=['timestamp','northing','easting','height','roll','pitch','yaw'], low_memory=False)

        n = len(pose_data)
        print(f'Number of global poses: {n}')
        system_timestamps = np.zeros((n,), dtype=np.int64)
        poses = np.zeros((n, 4, 4), dtype=np.float32)       # 4x4 pose matrix

        for ndx in range(n):
            row = pose_data.iloc[ndx]
            assert len(row) == 7, f'Invalid line in global poses file: {row}'
            ts = int(row['timestamp'])
            x = float(row['northing'])
            y = float(row['easting'])
            z = float(row['height'])
            roll = float(row['roll'])
            pitch = float(row['pitch'])
            yaw = float(row['yaw'])
            # transform to se3 matrix
            se3 = xyz_ypr2m(x, y, z, yaw, pitch, roll)
            # add to system_timestamps and poses
            system_timestamps[ndx] = int(ts)
            poses[ndx] = se3

        # Ensure timestamps and poses are sorted in ascending order
        sorted_ndx = np.argsort(system_timestamps, axis=0)
        system_timestamps = system_timestamps[sorted_ndx]
        poses = poses[sorted_ndx]

        # List LiDAR scan timestamps
        all_lidar_timestamps = [int(os.path.splitext(f)[0]) for f in os.listdir(self.lidar_path) if
                                os.path.splitext(f)[1] == '.bin']
        all_lidar_timestamps.sort()
        print(f'Number of global scans: {len(all_lidar_timestamps)}')

        lidar_timestamps = []
        lidar_filepaths = []
        lidar_poses = []
        count_rejected = 0

        for ndx, lidar_ts in enumerate(all_lidar_timestamps):
            # Find index of the closest timestamp
            closest_ts_ndx = find_nearest_ndx(lidar_ts, system_timestamps)
            delta = abs(system_timestamps[closest_ts_ndx] - lidar_ts)
            # Timestamp is in nanoseconds = 1e-9 second
            if delta > self.pose_time_tolerance * 1000000000:
                # Reject point cloud without corresponding pose
                count_rejected += 1
                continue

            lidar_timestamps.append(lidar_ts)
            lidar_filepaths.append(os.path.join(self.lidar_path, f'{lidar_ts}.bin'))
            lidar_poses.append(poses[closest_ts_ndx])

        lidar_timestamps = np.array(lidar_timestamps, dtype=np.int64)
        lidar_filepaths = np.array(lidar_filepaths)
        lidar_poses = np.array(lidar_poses, dtype=np.float32)     # 4x4 pose matrix
        lidar_xys = lidar_poses[:, :2, 3]                         # 2D position
        
        # Split data into train / test set
        if self.split != 'all':      
            if self.split == 'train':
                mask = check_in_train_set(lidar_xys, dataset='nclt')
            elif self.split == 'test':
                mask = check_in_test_set(lidar_xys, dataset='nclt')
            lidar_timestamps = lidar_timestamps[mask]
            lidar_filepaths = lidar_filepaths[mask]
            lidar_poses = lidar_poses[mask]
            lidar_xys = lidar_xys[mask]

        # Sample lidar scans 
        prev_position = None
        mask = []
        for ndx, position in enumerate(lidar_xys):
            if prev_position is None:
                mask.append(ndx)
                prev_position = position
            else:
                displacement = np.linalg.norm(prev_position - position)
                if displacement > self.sampling_distance:
                    mask.append(ndx)
                    prev_position = position

        lidar_timestamps = lidar_timestamps[mask]
        lidar_filepaths = lidar_filepaths[mask]
        lidar_poses = lidar_poses[mask]
        lidar_xys = lidar_xys[mask]

        print(f'{len(lidar_timestamps)} scans with valid pose, {count_rejected} rejected due to unknown pose')
        return lidar_timestamps, lidar_filepaths, lidar_poses, lidar_xys

    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


class NCLTSequences(Dataset):
    """
    Multiple NCLT sequences indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_names: List[str], split: str = 'train', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        self.dataset_root = dataset_root
        self.sequence_names = sequence_names
        self.split = split
        self.sampling_distance = sampling_distance
        
        # Load all sequences
        sequences = []
        for sequence_name in sequence_names:
            sequence = NCLTSequence(dataset_root, sequence_name, split, sampling_distance)
            sequences.append(sequence)
        self.dataset = ConcatDataset(sequences)
        self.pc_loader = sequences[0].pc_loader

        # Concatenate all sequences
        self.timestamps = np.concatenate([s.timestamps for s in sequences])
        self.filepaths = np.concatenate([s.filepaths for s in sequences])
        self.poses = np.concatenate([s.poses for s in sequences])
        self.xys = np.concatenate([s.xys for s in sequences])
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_names}-{split}')

        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, ndx):
        return self.dataset[ndx]

    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs
        
    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


if __name__ == "__main__":
    # load dataset
    base_path = './data/NCLT/'
    folder = '2012-02-04' # sequence folder for training   
    train_dataset = NCLTSequence(base_path, folder, split='train')
    test_dataset = NCLTSequence(base_path, folder, split='test')
    train_positions = train_dataset.xys
    test_positions = test_dataset.xys
    # plot the splited results
    plt.scatter(train_positions[:,0], train_positions[:,1], c='r', s=0.1)
    plt.scatter(test_positions[:,0], test_positions[:,1], c='b', s=0.1)
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.legend(['train', 'test'])
    plt.title(f'{folder} Trajectory Split')    
    plt.savefig(f'{folder}_split.png')
    plt.show()
    
datasets/NCLTDataset.py
lus6-Jenny/RING
import os
import sys
import copy
import glob
import random
import numpy as np
import pandas as pd
import utils.config as cfg
from typing import List, Tuple
from utils.poses import xyz_ypr2m
import datasets.velodyne as velodyne
from utils.tools import find_nearest_ndx, check_in_train_set, check_in_test_set
import matplotlib.pyplot as plt
from sklearn.neighbors import KDTree
from scipy.spatial.transform import Rotation
from torch.utils.data import Dataset, ConcatDataset, DataLoader


# calibrate the lidar data of oxford radar robotcar dataset
def get_oxford_radar_calib():
    calib = np.eye(4)
    calib[0,3] = -1.7132
    calib[1,3] = 0.1181
    calib[2,3] = 1.1948
    calib[:3,:3] = Rotation.from_euler('xyz',[-0.0125,0.0400,0.0050],degrees=True).as_matrix()
    return calib


# concantate point clouds generated by the left and right lidar
def pc_concantate(pc_left, pc_right, extrinsics_dir):
    """
    input pc_left: Nx4
    input pc_right: Mx4
    input extrinsics_dir: path to extrinsics
    return pc_cat: (N+M)x4
    """
    # transpose pc_left and pc_right
    pc_left = pc_left.transpose(1, 0)
    pc_right = pc_right.transpose(1, 0)
    # load extrinsics of left lidar
    left = np.loadtxt(os.path.join(extrinsics_dir, 'velodyne_left.txt'), delimiter=' ', dtype=np.float32)
    roll, pitch, yaw = left[3], left[4], left[5]
    # convert euler angles to rotation matrix
    left_T = np.eye(4)
    R_x = np.array([[1, 0, 0],
                    [0, np.cos(roll), -np.sin(roll)],
                    [0, np.sin(roll), np.cos(roll)]])
    R_y = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                    [0, 1, 0],
                    [-np.sin(pitch), 0, np.cos(pitch)]])
    R_z = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                    [np.sin(yaw), np.cos(yaw), 0],
                    [0, 0, 1]])
    R = np.dot(R_z, np.dot(R_y, R_x))
    # translation
    t = np.array(left[:3])
    left_T[:3, :3] = R
    left_T[:3, 3] = t
    intensity_array = copy.deepcopy(pc_left[3,:])
    pc_left[3,:] = 1
    pc_left_ed = np.dot(left_T, pc_left)  
    pc_left[3,:] = intensity_array
    pc_left_ed[3,:] = intensity_array

    # load extrinsics of right lidar
    right = np.loadtxt(os.path.join(extrinsics_dir, 'velodyne_right.txt'), delimiter=' ', dtype=np.float32)
    roll, pitch, yaw = right[3], right[4], right[5]
    # convert euler angles to rotation matrix
    right_T = np.eye(4)
    R_x = np.array([[1, 0, 0],
                    [0, np.cos(roll), -np.sin(roll)],
                    [0, np.sin(roll), np.cos(roll)]])
    R_y = np.array([[np.cos(pitch), 0, np.sin(pitch)],
                    [0, 1, 0],
                    [-np.sin(pitch), 0, np.cos(pitch)]])
    R_z = np.array([[np.cos(yaw), -np.sin(yaw), 0],
                    [np.sin(yaw), np.cos(yaw), 0],
                    [0, 0, 1]])
    R = np.dot(R_z, np.dot(R_y, R_x))
    # translation
    t = np.array(right[:3])
    right_T[:3, :3] = R
    right_T[:3, 3] = t
    intensity_array = copy.deepcopy(pc_right[3,:])
    pc_right[3,:] = 1
    pc_right_ed = np.dot(right_T, pc_right)  
    pc_right[3,:] = intensity_array
    pc_right_ed[3,:] = intensity_array
    pc_cat = np.concatenate((pc_left_ed, pc_right_ed), axis=1)
    pc_cat = pc_cat.transpose(1, 0)
    pc_cat = pc_cat.astype(np.float32)
    return pc_cat


# load point cloud from bin or raw png file
def load_pc_oxford_radar(scan_path):
    n_vec = 4
    if scan_path.endswith('.bin'):
        pc = np.fromfile(scan_path, dtype=np.float32)
        pc = np.reshape(pc, (-1, n_vec))
        # print(pc.shape)
        # print('------ debug ------', np.min(pc[:,0]), np.max(pc[:,0]), np.min(pc[:,1]), np.max(pc[:,1]), np.min(pc[:,2]), np.max(pc[:,2]))
    elif scan_path.endswith('.png'):
        ranges, intensities, angles, approximate_timestamps = velodyne.load_velodyne_raw(scan_path)
        pc = velodyne.velodyne_raw_to_pointcloud(ranges, intensities, angles)
        pc = pc.reshape((-1, n_vec))
    pc = pc.astype(np.float32)
    return pc


class OxfordRadarPointCloudLoader:
    def __init__(self, remove_zero_points: bool = True, remove_ground_plane: bool = True, ground_plane_level: float = -13.0):
        # remove_zero_points: remove points with all zero coordinates
        # remove_ground_plane: remove points on ground plane level and below
        # ground_plane_level: ground plane level
        self.remove_zero_points = remove_zero_points
        self.remove_ground_plane = remove_ground_plane
        self.ground_plane_level = ground_plane_level

    def read_pc(self, file_pathname: str):
        # Reads the point cloud and preprocess (optional removal of zero points and points on the ground plane
        pc = load_pc_oxford_radar(file_pathname)
        if self.remove_zero_points:
            mask = np.all(np.isclose(pc, 0), axis=1)
            pc = pc[~mask]

        if self.remove_ground_plane:
            mask = pc[:, 2] > self.ground_plane_level
            pc = pc[mask]        
        
        pc = np.unique(pc, axis=0) # remove duplicate points  

        return pc

    def read_pc_cat(self, left_file_pathname: str, right_file_pathname: str, extrinsics_dir: str):
        # Reads the point cloud and preprocess (optional removal of zero points and points on the ground plane
        pc_left = load_pc_oxford_radar(left_file_pathname)
        pc_right = load_pc_oxford_radar(right_file_pathname)
        pc = pc_concantate(pc_left, pc_right, extrinsics_dir)
        if self.remove_zero_points:
            mask = np.all(np.isclose(pc, 0), axis=1)
            pc = pc[~mask]

        if self.remove_ground_plane:
            mask = pc[:, 2] > self.ground_plane_level
            pc = pc[mask]    

        pc = np.unique(pc, axis=0) # remove duplicate points      

        return pc
                
    def normalize_pc(self, pc, xbound: float = cfg.point_cloud["x_bound"][1], ybound: float = cfg.point_cloud["y_bound"][1], zbound: float = 15.):
        # Normalize the point cloud
        pc = np.array(pc, dtype=np.float32)
        pc = pc[:,:3]

        mask = (np.abs(pc[:,0]) < xbound) * (np.abs(pc[:,1]) < ybound) * (pc[:,2] > self.ground_plane_level) * (pc[:,2] < zbound)

        pc = pc[mask]
        pc[:,0] = pc[:,0] / xbound
        pc[:,1] = pc[:,1] / ybound
        pc[:,2] = pc[:,2] / zbound 

        return pc   


class OxfordRadarSequence(Dataset):
    """
    Single Oxford Radar sequence indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_name: str, split: str = 'train', lidar: str = 'all', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        assert lidar in ['left', 'right', 'all']
        self.dataset_root = dataset_root
        self.sequence_name = sequence_name
        self.sequence_path = os.path.join(self.dataset_root, self.sequence_name)   
        assert os.path.exists(self.sequence_path), f'Cannot access sequence: {self.sequence_path}'          
        self.split = split
        self.lidar = lidar
        self.sampling_distance = sampling_distance
        # Maximum discrepancy between timestamps of LiDAR scan and global pose in seconds
        self.pose_time_tolerance = 1.
        self.pose_file = os.path.join(self.sequence_path, 'ins.csv')
        assert os.path.exists(self.pose_file), f'Cannot access global pose file: {self.pose_file}'
        self.lidar_path = os.path.join(self.sequence_path, 'velodyne_' + self.lidar)
        # assert os.path.exists(self.lidar_path), f'Cannot access lidar scans: {self.lidar_path}'      
        # self.lidar_files = sorted(glob.glob(os.path.join(self.lidar_path, '*.bin')))
        self.pc_loader = OxfordRadarPointCloudLoader()
        self.timestamps, self.filepaths, self.poses, self.xys = self.get_scan_poses()
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_name}-{split}-{lidar}')
        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.filepaths)

    def __getitem__(self, ndx):
        pc_filepath = self.filepaths[ndx]
        pc = self.load_pc(pc_filepath)
        return {'ts': self.timestamps[ndx], 'pc': pc, 'position': self.xys[ndx], 'pose': self.poses[ndx]}
    
    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs
        
    def get_scan_poses(self):
        # Read global poses from .csv file and link each lidar_scan with the nearest pose
        # threshold: timestamp threshold in seconds
        # Returns a dictionary with (4, 4) pose matrix indexed by a timestamp (as integer)
        pose_data = pd.read_csv(self.pose_file, header=0, low_memory=False)

        n = len(pose_data)
        print(f'Number of global poses: {n}')
        system_timestamps = np.zeros((n,), dtype=np.int64)
        poses = np.zeros((n, 4, 4), dtype=np.float32)       # 4x4 pose matrix
        calib = get_oxford_radar_calib()

        for ndx in range(n):
            row = pose_data.iloc[ndx]
            assert len(row) == 15, f'Invalid line in global poses file: {row}'
            ts = int(row['timestamp'])
            x = float(row['northing'])
            y = float(row['easting'])
            z = float(row['down'])
            roll = float(row['roll'])
            pitch = float(row['pitch'])
            yaw = float(row['yaw'])            
            # transform to se3 matrix
            se3 = xyz_ypr2m(x, y, z, yaw, pitch, roll)
            # add to system_timestamps and poses
            system_timestamps[ndx] = int(ts)
            poses[ndx] = se3 @ calib
            # poses[ndx] = se3


        # Ensure timestamps and poses are sorted in ascending order
        sorted_ndx = np.argsort(system_timestamps, axis=0)
        system_timestamps = system_timestamps[sorted_ndx]
        poses = poses[sorted_ndx]

        # List LiDAR scan timestamps
        if self.lidar == 'all':
            if not os.path.exists(self.lidar_path):
                os.makedirs(self.lidar_path)
                print(f'Created directory: {self.lidar_path}')
            left_lidar_path = os.path.join(self.sequence_path, 'velodyne_left')
            right_lidar_path = os.path.join(self.sequence_path, 'velodyne_right')
            left_lidar_timestamps = [int(os.path.splitext(f)[0]) for f in os.listdir(left_lidar_path)]
            right_lidar_timestamps = [int(os.path.splitext(f)[0]) for f in os.listdir(right_lidar_path)]
            left_lidar_timestamps.sort()
            right_lidar_timestamps.sort()
            print(f'Number of left lidar scans: {len(left_lidar_timestamps)}')
            print(f'Number of right lidar scans: {len(right_lidar_timestamps)}')
            for ndx, lidar_ts in enumerate(left_lidar_timestamps):
                # Find index of the closest timestamp
                closest_ts_ndx = find_nearest_ndx(lidar_ts, right_lidar_timestamps)
                delta = abs(right_lidar_timestamps[closest_ts_ndx] - lidar_ts)
                # Timestamp is in nanoseconds = 1e-9 second
                if delta > self.pose_time_tolerance * 1000000000:
                    continue                
                # Concatenate left and right lidar scans and save to .bin file
                bin_filepath = os.path.join(self.lidar_path, f'{lidar_ts}.bin')
                if not os.path.exists(bin_filepath):
                    left_lidar_filepath = os.path.join(left_lidar_path, f'{lidar_ts}.png')
                    right_lidar_filepath = os.path.join(right_lidar_path, f'{right_lidar_timestamps[closest_ts_ndx]}.png')
                    pc_cat = self.pc_loader.read_pc_cat(left_lidar_filepath, right_lidar_filepath, self.dataset_root)
                    pc_cat.tofile(bin_filepath)
                    print(f'Save the concatenated lidar scan from {left_lidar_filepath} and {right_lidar_filepath} to {bin_filepath}')
                
        all_lidar_timestamps = [int(os.path.splitext(f)[0]) for f in os.listdir(self.lidar_path)]
        all_lidar_timestamps.sort()
        print(f'Number of global scans: {len(all_lidar_timestamps)}')
        
        lidar_timestamps = []
        lidar_filepaths = []
        lidar_poses = []
        count_rejected = 0

        for ndx, lidar_ts in enumerate(all_lidar_timestamps):
            # Find index of the closest timestamp
            closest_ts_ndx = find_nearest_ndx(lidar_ts, system_timestamps)
            delta = abs(system_timestamps[closest_ts_ndx] - lidar_ts)
            # Timestamp is in nanoseconds = 1e-9 second
            if delta > self.pose_time_tolerance * 1000000000:
                # Reject point cloud without corresponding pose
                count_rejected += 1
                continue
            
            if self.lidar == 'all':
                lidar_filepaths.append(os.path.join(self.lidar_path, f'{lidar_ts}.bin'))
            else:
                lidar_filepaths.append(os.path.join(self.lidar_path, f'{lidar_ts}.png'))
            lidar_timestamps.append(lidar_ts)
            lidar_poses.append(poses[closest_ts_ndx])
        
        lidar_timestamps = np.array(lidar_timestamps, dtype=np.int64)
        lidar_filepaths = np.array(lidar_filepaths)
        lidar_poses = np.array(lidar_poses, dtype=np.float32)     # 4x4 pose matrix
        lidar_xys = lidar_poses[:, :2, 3]                         # 2D position

        # split data into train / test set
        if self.split != 'all':      
            if self.split == 'train':
                mask = check_in_train_set(lidar_xys, dataset='oxford_radar')
            elif self.split == 'test':
                mask = check_in_test_set(lidar_xys, dataset='oxford_radar')
            lidar_timestamps = lidar_timestamps[mask]
            lidar_filepaths = lidar_filepaths[mask]
            lidar_poses = lidar_poses[mask]
            lidar_xys = lidar_xys[mask]

        # Sample lidar scans             
        prev_position = None
        mask = []
        for ndx, position in enumerate(lidar_xys):
            if prev_position is None:
                mask.append(ndx)
                prev_position = position
            else:
                displacement = np.linalg.norm(prev_position - position)
                if displacement > self.sampling_distance:
                    mask.append(ndx)
                    prev_position = position
        
        lidar_timestamps = lidar_timestamps[mask]
        lidar_filepaths = lidar_filepaths[mask]
        lidar_poses = lidar_poses[mask]
        lidar_xys = lidar_xys[mask]

        print(f'{len(lidar_timestamps)} scans with valid pose, {count_rejected} rejected due to unknown pose')
        return lidar_timestamps, lidar_filepaths, lidar_poses, lidar_xys

    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


class OxfordRadarSequences(Dataset):
    """
    Multiple Oxford Radar sequences indexed as a single dataset containing point clouds and poses
    """
    def __init__(self, dataset_root: str, sequence_names: List[str], split: str = 'train', lidar: str = 'all', sampling_distance: float = 0.2):
        assert os.path.exists(dataset_root), f'Cannot access dataset root: {dataset_root}'
        assert split in ['train', 'test', 'all']
        assert lidar in ['left', 'right', 'all']
        self.dataset_root = dataset_root
        self.sequence_names = sequence_names
        self.split = split
        self.lidar = lidar
        self.sampling_distance = sampling_distance
        
        # Load all sequences
        sequences = []
        for sequence_name in sequence_names:
            sequence = OxfordRadarSequence(dataset_root, sequence_name, split, lidar, sampling_distance)
            sequences.append(sequence)
        self.dataset = ConcatDataset(sequences)
        self.pc_loader = sequences[0].pc_loader

        # Concatenate all sequences
        self.timestamps = np.concatenate([s.timestamps for s in sequences])
        self.filepaths = np.concatenate([s.filepaths for s in sequences])
        self.poses = np.concatenate([s.poses for s in sequences])
        self.xys = np.concatenate([s.xys for s in sequences])
        assert len(self.filepaths) == len(self.poses), f'Number of lidar file paths and poses do not match: {len(self.filepaths)} vs {len(self.poses)}'
        print(f'{len(self.filepaths)} scans in {sequence_names}-{split}-{lidar}')

        # Build a kdtree based on X, Y position
        self.kdtree = KDTree(self.xys)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, ndx):
        return self.dataset[ndx]

    def load_pc(self, scan_path):
        # Load point cloud from file
        pc = self.pc_loader.read_pc(scan_path)
        return pc
    
    def load_pcs(self, scan_paths):
        # Load point clouds from files
        pcs = []
        for scan_path in scan_paths:
            pc = self.load_pc(scan_path)
            if len(pc) == 0:
                continue
            pcs.append(pc)
        pcs = np.array(pcs)
        return pcs  
    
    def normalize_pc(self, pc):
        # Normalize the point cloud
        pc = self.pc_loader.normalize_pc(pc)
        return pc
    
    def normalize_pcs(self, pcs):
        # Normalize point clouds
        pcs = [self.normalize_pc(pc) for pc in pcs]
        return pcs

    def find_neighbours_ndx(self, position, radius):
        # Returns indices of neighbourhood point clouds for a given position
        assert position.ndim == 1
        assert position.shape[0] == 2
        # Reshape into (1, 2) axis
        position = position.reshape(1, -1)
        neighbours = self.kdtree.query_radius(position, radius)[0]
        return neighbours.astype(np.int32)


if __name__ == "__main__":
    # load dataset
    base_path = './data/OxfordRadar/'
    folder = '2019-01-11-13-24-51' # sequence folder for training    
    train_dataset = OxfordRadarSequence(base_path, folder, split='train')
    test_dataset = OxfordRadarSequence(base_path, folder, split='test')
    train_positions = train_dataset.xys
    test_positions = test_dataset.xys
    # plot the splited results
    plt.scatter(train_positions[:,0], train_positions[:,1], c='r', s=0.1)
    plt.scatter(test_positions[:,0], test_positions[:,1], c='b', s=0.1)
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.legend(['train', 'test'])
    plt.title(f'{folder} Trajectory Split')
    plt.savefig(f'{folder}_split.png')
    plt.show()

datasets/OxfordRadarDataset.py
lus6-Jenny/RING
################################################################################
#
# Copyright (c) 2017 University of Oxford
# Authors:
#  Dan Barnes (dbarnes@robots.ox.ac.uk)
#
# This work is licensed under the Creative Commons
# Attribution-NonCommercial-ShareAlike 4.0 International License.
# To view a copy of this license, visit
# http://creativecommons.org/licenses/by-nc-sa/4.0/ or send a letter to
# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
#
###############################################################################

from typing import AnyStr
import numpy as np
import os
import cv2

# Hard coded configuration to simplify parsing code
hdl32e_range_resolution = 0.002  # m / pixel
hdl32e_minimum_range = 3.0
hdl32e_elevations = np.array([-0.1862, -0.1628, -0.1396, -0.1164, -0.0930,
                              -0.0698, -0.0466, -0.0232, 0., 0.0232, 0.0466, 0.0698,
                              0.0930, 0.1164, 0.1396, 0.1628, 0.1862, 0.2094, 0.2327,
                              0.2560, 0.2793, 0.3025, 0.3259, 0.3491, 0.3723, 0.3957,
                              0.4189, 0.4421, 0.4655, 0.4887, 0.5119, 0.5353])[:, np.newaxis]
hdl32e_base_to_fire_height = 0.090805
hdl32e_cos_elevations = np.cos(hdl32e_elevations)
hdl32e_sin_elevations = np.sin(hdl32e_elevations)


def load_velodyne_binary(velodyne_bin_path: AnyStr):
    """Decode a binary Velodyne example (of the form '<timestamp>.bin')
    Args:
        example_path (AnyStr): Oxford Radar RobotCar Dataset binary Velodyne pointcloud example path
    Returns:
        ptcld (np.ndarray): XYZI pointcloud from the binary Velodyne data Nx4
    Notes:
        - The pre computed points are *NOT* motion compensated.
        - Converting a raw velodyne scan to pointcloud can be done using the
            `velodyne_ranges_intensities_angles_to_pointcloud` function.
    """
    ext = os.path.splitext(velodyne_bin_path)[1]
    if ext != ".bin":
        raise RuntimeError("Velodyne binary pointcloud file should have `.bin` extension but had: {}".format(ext))
    if not os.path.isfile(velodyne_bin_path):
        raise FileNotFoundError("Could not find velodyne bin example: {}".format(velodyne_bin_path))
    data = np.fromfile(velodyne_bin_path, dtype=np.float32)
    ptcld = data.reshape((4, -1))
    return ptcld


def load_velodyne_raw(velodyne_raw_path: AnyStr):
    """Decode a raw Velodyne example. (of the form '<timestamp>.png')
    Args:
        example_path (AnyStr): Oxford Radar RobotCar Dataset raw Velodyne example path
    Returns:
        ranges (np.ndarray): Range of each measurement in meters where 0 == invalid, (32 x N)
        intensities (np.ndarray): Intensity of each measurement where 0 == invalid, (32 x N)
        angles (np.ndarray): Angle of each measurement in radians (1 x N)
        approximate_timestamps (np.ndarray): Approximate linearly interpolated timestamps of each mesaurement (1 x N).
            Approximate as we only receive timestamps for each packet. The timestamp of the next frame will was used to
            interpolate the last packet timestamps. If there was no next frame, the last packet timestamps was
            extrapolated. The original packet timestamps can be recovered with:
                approximate_timestamps(:, 1:12:end) (12 is the number of azimuth returns in each packet)
     Notes:
       Reference: https://velodynelidar.com/lidar/products/manual/63-9113%20HDL-32E%20manual_Rev%20E_NOV2012.pdf
    """
    ext = os.path.splitext(velodyne_raw_path)[1]
    if ext != ".png":
        raise RuntimeError("Velodyne raw file should have `.png` extension but had: {}".format(ext))
    if not os.path.isfile(velodyne_raw_path):
        raise FileNotFoundError("Could not find velodyne raw example: {}".format(velodyne_raw_path))
    example = cv2.imread(velodyne_raw_path, cv2.IMREAD_GRAYSCALE)
    intensities, ranges_raw, angles_raw, timestamps_raw = np.array_split(example, [32, 96, 98], 0)
    ranges = np.ascontiguousarray(ranges_raw.transpose()).view(np.uint16).transpose()
    ranges = ranges * hdl32e_range_resolution
    angles = np.ascontiguousarray(angles_raw.transpose()).view(np.uint16).transpose()
    angles = angles * (2. * np.pi) / 36000
    approximate_timestamps = np.ascontiguousarray(timestamps_raw.transpose()).view(np.int64).transpose()
    return ranges, intensities, angles, approximate_timestamps


def velodyne_raw_to_pointcloud(ranges: np.ndarray, intensities: np.ndarray, angles: np.ndarray):
    """ Convert raw Velodyne data (from load_velodyne_raw) into a pointcloud
    Args:
        ranges (np.ndarray): Raw Velodyne range readings
        intensities (np.ndarray): Raw Velodyne intensity readings
        angles (np.ndarray): Raw Velodyne angles
    Returns:
        pointcloud (np.ndarray): XYZI pointcloud generated from the raw Velodyne data Nx4

    Notes:
        - This implementation does *NOT* perform motion compensation on the generated pointcloud.
        - Accessing the pointclouds in binary form via `load_velodyne_pointcloud` is approximately 2x faster at the cost
            of 8x the storage space
    """
    valid = ranges > hdl32e_minimum_range
    z = hdl32e_sin_elevations * ranges - hdl32e_base_to_fire_height
    xy = hdl32e_cos_elevations * ranges
    x = np.sin(angles) * xy
    y = -np.cos(angles) * xy

    xf = x[valid].reshape(-1)
    yf = y[valid].reshape(-1)
    zf = z[valid].reshape(-1)
    intensityf = intensities[valid].reshape(-1).astype(np.float32)
    ptcld = np.stack((xf, yf, zf, intensityf), 1)

    return ptcld

# ranges, intensity, angles, approximate_timestamps = load_velodyne_raw("/media/mav-lab/1T/Datasets/bypy/01.10/2019-01-10-12-32-52-radar-oxford-10k/velodyne_left/1547123579957366.png")
# ptcld = velodyne_raw_to_pointcloud(ranges, intensity, angles)
datasets/velodyne.py
lus6-Jenny/RING
import os
import time
import torch
import argparse
import numpy as np
from tqdm import tqdm
from typing import List, Dict
from sklearn.neighbors import KDTree

from utils.core import *
import utils.config as cfg
from utils.tools import imshow
from utils.poses import relative_pose, xyz_ypr2m, m2xyz_ypr, angle_clip, cal_pose_error
from utils.point_clouds import o3d_icp, fast_gicp, draw_pc, draw_pc_pair, draw_registration_result

from datasets.NCLTDataset import NCLTPointCloudLoader
from datasets.KITTIDataset import KITTIPointCloudLoader
from datasets.MulRanDataset import MulRanPointCloudLoader
from datasets.OxfordRadarDataset import OxfordRadarPointCloudLoader

from evaluation.plot_pose_errors import plot_cdf
from evaluation.plot_PR_curve import compute_PR_pairs
from evaluation.generate_evaluation_sets import EvaluationSet, EvaluationTuple


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('device', device)


def generate_representations(dataset: str, eval_subset: List[EvaluationTuple], bev_type: str = "occ"):
    # Generate BEV & RING representations for the evaluation set
    # Each BEV / RING representation is a numpy array of shape (C, H, W)
    # H: height, W: width, C: number of channels
    # C = 1 for occupancy BEV
    # C = 6 for feature BEV
    
    pcs = []
    pc_bevs = []
    pc_RINGs = []
    pc_TIRINGs = []
    if dataset == 'nclt':
        pc_loader = NCLTPointCloudLoader()
    elif dataset == 'kitti':
        pc_loader = KITTIPointCloudLoader()
    elif dataset == 'mulran':
        pc_loader = MulRanPointCloudLoader()
    elif dataset == 'oxford_radar':
        pc_loader = OxfordRadarPointCloudLoader()
    else:
        raise ValueError(f"Unknown dataset: {dataset}")
    
    for ndx, e in tqdm(enumerate(eval_subset)):
        pc = pc_loader.read_pc(e.rel_scan_filepath)
        pc = pc[:,:3]
        pc_norm = pc_loader.normalize_pc(pc)
        
        # Generate BEV features
        if device == torch.device("cpu"):
            pc_bev = generate_bev(pc) # BEV encoded by occupancy information
            pc_RING, pc_TIRING = generate_RING_cpu(pc_bev)
        else:      
            if bev_type == "occ":
                pc_bev = generate_bev_occ(pc_norm) # BEV encoded by occupancy information
            elif bev_type == "feat":
                pc_bev = generate_bev_feats(pc_norm) # BEV encoded by point features
            else:
                raise ValueError('bev_type is not "occ" or "feat"')
            pc_RING, pc_TIRING = generate_RING(pc_bev)
            
        pcs.append(pc)
        pc_bevs.append(pc_bev)
        pc_RINGs.append(pc_RING)
        pc_TIRINGs.append(pc_TIRING)
    
    return pcs, pc_bevs, pc_RINGs, pc_TIRINGs


def evaluate(dataset, eval_set_filepath: str, revisit_thresholds: List[float] = [5.0, 10.0, 15.0, 20.0], num_k: int = 25, bev_type: str = "occ"):
    
    assert os.path.exists(eval_set_filepath), f'Cannot access evaluation set pickle: {eval_set_filepath}'
    eval_set = EvaluationSet()
    eval_set.load(eval_set_filepath)
    
    map_set = eval_set.map_set
    query_set = eval_set.query_set

    map_pcs, map_bevs, map_RINGs, map_TIRINGs = generate_representations(dataset, map_set, bev_type)
    query_pcs, query_bevs, query_RINGs, query_TIRINGs = generate_representations(dataset, query_set, bev_type)
    
    C, H, W = map_bevs[0].shape

    map_xys = eval_set.get_map_positions()
    map_poses = eval_set.get_map_poses()
    num_maps = len(map_xys)
    print(f"No. maps: {num_maps}")
    
    tree = KDTree(map_xys)

    query_xys = eval_set.get_query_positions()
    query_poses = eval_set.get_query_poses()
    num_queries = len(query_xys)
    print(f"No. queries: {num_queries}")
    
    recalls_k = np.zeros((len(revisit_thresholds), num_k)) # Recall@k
    recalls_one_percent = np.zeros(len(revisit_thresholds)) # Recall@1% (recall when 1% of the database is retrieved)
    threshold = max(int(round(num_maps/100.0)), 1) # Recall@1% threshold
    pair_dists = np.zeros((num_queries, num_maps))
    rot_errors = np.zeros((len(revisit_thresholds), num_queries)) # rotation errors
    trans_errors = np.zeros((len(revisit_thresholds), num_queries))  # translation errors
    icp_rot_errors = np.zeros((len(revisit_thresholds), num_queries))  # rotation errors after ICP
    icp_trans_errors = np.zeros((len(revisit_thresholds), num_queries))  # translation errors after ICP    
    for query_ndx in tqdm(range(num_queries)):
        query_xy = query_xys[query_ndx]
        query_pose = query_poses[query_ndx]
        query_pc = query_pcs[query_ndx]
        query_bev = query_bevs[query_ndx]
        query_RING = query_RINGs[query_ndx]
        query_TIRING = query_TIRINGs[query_ndx]
        query_TIRING_repeated = query_TIRING.repeat(cfg.batch_size, 1, 1, 1)
        
        # ------------ Place Recognition ------------
        if num_maps % cfg.batch_size == 0:
            batch_num = num_maps // cfg.batch_size
        else:
            batch_num = num_maps // cfg.batch_size + 1      
        for i in range(batch_num):
            if i == batch_num - 1:
                map_TIRING = [map_TIRINGs[k] for k in range(i*cfg.batch_size, num_maps)]
                map_TIRING = torch.stack(map_TIRING, dim=0).reshape((-1, C, H, W))
                query_TIRING_repeated = query_TIRING.repeat(map_TIRING.shape[0], 1, 1, 1)
            else:
                map_TIRING = [map_TIRINGs[k] for k in range(i*cfg.batch_size, (i+1)*cfg.batch_size)]
                map_TIRING = torch.stack(map_TIRING, dim=0).reshape((-1, C, H, W))
            # batch_dists, batch_angles = fast_corr(query_TIRING_repeated, map_TIRING)
            batch_dists, batch_angles = batch_circorr(query_TIRING_repeated, map_TIRING)
            if i == 0:
                dists = batch_dists
                angles = batch_angles
            else:
                dists = np.concatenate((dists, batch_dists), axis=-1)
                angles = np.concatenate((angles, batch_angles), axis=-1)

        dists = dists.squeeze()
        angles = angles.squeeze()          
        pair_dists[query_ndx] = dists
        dists_sorted = np.sort(dists)
        idxs_sorted = np.argsort(dists)
        idx_top1 = idxs_sorted[0]
        for j, revisit_threshold in enumerate(revisit_thresholds):
            true_neighbors = tree.query_radius(query_xy.reshape(1,-1), revisit_threshold)[0]
            # Recall@k
            for k in range(num_k):
                # if np.linalg.norm(map_xys[idxs_sorted[k]] - query_xy) < revisit_threshold:
                if idxs_sorted[k] in true_neighbors:
                    recalls_k[j, k] += 1
                    break
            # Recall@1%
            if len(list(set(idxs_sorted[0:threshold]).intersection(set(true_neighbors)))) > 0:
                recalls_one_percent[j] += 1     

            # ------------ Pose Estimation ------------
             # Perform pose estimation for the top 1 match within the revisit threshold
            if idx_top1 in true_neighbors:
                map_pose = map_poses[idx_top1]
                map_pc = map_pcs[idx_top1]
                map_bev = map_bevs[idx_top1]
                angle_matched = angles[idx_top1]
                rel_pose = relative_pose(query_pose, map_pose)
                gt_x, gt_y, gt_z, gt_yaw, gt_pitch, gt_roll = m2xyz_ypr(rel_pose)
                print(f"-------- Query {query_ndx+1} matched with map {idx_top1+1} --------")
                print("Ground truth translation: x: {}, y: {}, rotation: {}".format(gt_x, gt_y, gt_yaw))

                ang_res = 2 * np.pi / cfg.num_ring # angular resolution
                # angle between the two matched RINGs in grids
                angle_matched_extra = angle_matched - cfg.num_ring // 2
                # convert the matched angle from grids to radians
                angle_matched_rad = angle_matched * ang_res 
                angle_matched_extra_rad = angle_matched_extra * ang_res        

                bev_rotated = rotate_bev(query_bev, angle_matched_rad)
                bev_rotated_extra = rotate_bev(query_bev, angle_matched_extra_rad)

                # solve the translation between the two matched bevs
                x, y, error = solve_translation(bev_rotated, map_bev)
                x_extra, y_extra, error_extra = solve_translation(bev_rotated_extra, map_bev)

                if error < error_extra:
                    pred_x = x / cfg.num_sector * (cfg.point_cloud["x_bound"][1] - cfg.point_cloud["x_bound"][0])  # in meters
                    pred_y = y / cfg.num_ring * (cfg.point_cloud["y_bound"][1] - cfg.point_cloud["y_bound"][0])   # in meters
                    pred_yaw = angle_matched_rad  # in radians
                else:
                    pred_x = x_extra / cfg.num_sector * (cfg.point_cloud["x_bound"][1] - cfg.point_cloud["x_bound"][0])  # in meters
                    pred_y = y_extra / cfg.num_ring * (cfg.point_cloud["y_bound"][1] - cfg.point_cloud["y_bound"][0])  # in meters 
                    pred_yaw = angle_matched_extra_rad  # in radians
                
                print("Estimated translation: x: {}, y: {}, rotation: {}".format(pred_x, pred_y, pred_yaw))

                yaw_err = np.abs(angle_clip(gt_yaw - pred_yaw)) * 180 / np.pi # in degrees
                x_err = np.abs(gt_x - pred_x) # in meters
                y_err = np.abs(gt_y - pred_y) # in meters
                trans_err = np.sqrt(x_err**2 + y_err**2) # in meters
                rot_errors[j, query_ndx] = yaw_err
                trans_errors[j, query_ndx] = trans_err

                # ------------ Pose Refinement ------------
                init_pose = xyz_ypr2m(pred_x, pred_y, 0, pred_yaw, 0, 0)
                times = time.time()
                icp_fitness_score, loop_transform = fast_gicp(query_pc, map_pc, max_correspondence_distance=cfg.icp_max_distance, init_pose=init_pose)
                # icp_fitness_score, loop_transform, _ = o3d_icp(query_pc, map_pc, transform=init_pose, point2plane=True, inlier_dist_threshold=cfg.icp_max_distance)
                timee = time.time()    
                # print("ICP processed time:", timee - times, 's')

                x, y, z, yaw, pitch, roll = m2xyz_ypr(loop_transform)
                print("Refined translation: x: {}, y: {}, rotation: {}".format(x, y, yaw))

                icp_rte, icp_rre = cal_pose_error(loop_transform, rel_pose)
                icp_rot_errors[j, query_ndx] = icp_rre
                icp_trans_errors[j, query_ndx] = icp_rte


    thresholds1 = np.linspace(0, 0.35, 10) 
    thresholds2 = np.linspace(0.35, 0.55, 80)
    thresholds3 = np.linspace(0.55, 1, 10)
    thresholds = np.concatenate((thresholds1, thresholds2, thresholds3))   
    precisions_all = np.zeros((len(revisit_thresholds), len(thresholds)))
    recalls_all = np.zeros((len(revisit_thresholds), len(thresholds)))
    f1s_all = np.zeros((len(revisit_thresholds), len(thresholds))) 
    quantiles = [0.25, 0.5, 0.75, 0.95]
    
    eval_setting = eval_set_filepath.split('/')[-1].split('.pickle')[0]
    folder = f"./results/{dataset}/{eval_setting}/revisit"
    for j, revisit_threshold in enumerate(revisit_thresholds):
        folder = f"{folder}_{revisit_threshold}"
    if not os.path.exists(folder):
        os.makedirs(folder)

    for j, revisit_threshold in enumerate(revisit_thresholds):
        # ------------ Recall@k and Recall@1% ------------
        recalls_k[j] = np.cumsum(recalls_k[j])/num_queries
        recalls_one_percent[j] = recalls_one_percent[j]/num_queries
        print(f"-------- Revisit threshold: {revisit_threshold} m --------")
        print(f"Recall@{num_k} at {revisit_threshold} m: {recalls_k[j]}")
        print(f"Recall@1% at {revisit_threshold} m: {recalls_one_percent[j]}")

        # ------------ PR Curve ------------
        save_path = f'{folder}/precision_recall_curve_{bev_type}_{revisit_threshold}m.pdf'
        precisions, recalls, f1s = compute_PR_pairs(pair_dists, query_xys, map_xys, thresholds=thresholds, save_path=save_path, revisit_threshold=revisit_threshold)
        precisions_all[j] = precisions
        recalls_all[j] = recalls
        f1s_all[j] = f1s

        # ------------ Pose Error ------------ 
        mean_rot_error = np.mean(rot_errors[j])
        mean_trans_error = np.mean(trans_errors[j])
        mean_icp_rot_error = np.mean(icp_rot_errors[j])
        mean_icp_trans_error = np.mean(icp_trans_errors[j])    
        rot_error_quantiles = np.quantile(rot_errors[j], quantiles)
        trans_error_quantiles = np.quantile(trans_errors[j], quantiles)
        icp_rot_error_quantiles = np.quantile(icp_rot_errors[j], quantiles)
        icp_trans_error_quantiles = np.quantile(icp_trans_errors[j], quantiles)    
        print(f"Mean rotation error at {revisit_threshold} m: {mean_rot_error}")    
        print(f"Mean translation error at {revisit_threshold} m: {mean_trans_error}")
        print(f"Mean icp rotation error at {revisit_threshold} m: {mean_icp_rot_error}")    
        print(f"Mean icp translation error at {revisit_threshold} m: {mean_icp_trans_error}")    
        print(f"Rotation error quantiles at {revisit_threshold} m: {rot_error_quantiles}")
        print(f"Translation error quantiles at {revisit_threshold} m: {trans_error_quantiles}")
        print(f"ICP rotation error quantiles at {revisit_threshold} m: {icp_rot_error_quantiles}")
        print(f"ICP translation error quantiles at {revisit_threshold} m: {icp_trans_error_quantiles}")

        save_path = f'{folder}/rot_error_cdf_{bev_type}_{revisit_threshold}m.pdf'
        plot_cdf(rot_errors[j], save_path=save_path, xlabel='Rotation Error (degrees)', ylabel='CDF', title='Rotation error CDF')

    np.savetxt(f'{folder}/recalls_k_{bev_type}.txt', recalls_k)
    np.savetxt(f'{folder}/recalls_one_percent_{bev_type}.txt', recalls_one_percent)
    np.savetxt(f'{folder}/precisions_{bev_type}.txt', precisions_all)
    np.savetxt(f'{folder}/recalls_{bev_type}.txt', recalls_all)
    np.savetxt(f'{folder}/f1s_{bev_type}.txt', f1s_all)
    np.savetxt(f'{folder}/rot_errors_{bev_type}.txt', rot_errors)
    np.savetxt(f'{folder}/trans_errors_{bev_type}.txt', trans_errors)
    np.savetxt(f'{folder}/icp_rot_errors_{bev_type}.txt', icp_rot_errors)
    np.savetxt(f'{folder}/icp_trans_errors_{bev_type}.txt', icp_trans_errors)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, default="nclt", help="dataset type (nclt / mulran / kitti / oxford_radar)")
    parser.add_argument('--eval_set_filepath', type=str, required=True, help='File path of the evaluation pickle')
    parser.add_argument('--revisit_thresholds', type=float, nargs='+', default=[5.0, 10.0, 15.0, 20.0], help='Revisit thresholds in meters')
    parser.add_argument('--num_k', type=int, default=25, help='Number of nearest neighbors for recall@k')
    parser.add_argument('--bev_type', type=str, default="occ", help='BEV type (occ / feat)')
    args = parser.parse_args()

    print(f'Dataset: {args.dataset}')
    print(f'Evaluation set path: {args.eval_set_filepath}')
    print(f'Revisit thresholds [m]: {args.revisit_thresholds}')
    print(f'Number of nearest neighbors for recall@k: {args.num_k}')
    print(f'BEV type: {args.bev_type}')

    evaluate(args.dataset, args.eval_set_filepath, args.revisit_thresholds, num_k=args.num_k, bev_type=args.bev_type)

evaluation/evaluate.py
lus6-Jenny/RING
import os
import pickle
import argparse
import numpy as np
from typing import List, Dict
from sklearn.neighbors import KDTree

from datasets.NCLTDataset import NCLTSequence, NCLTSequences
from datasets.KITTIDataset import KITTISequence, KITTISequences
from datasets.MulRanDataset import MulRanSequence, MulRanSequences
from datasets.OxfordRadarDataset import OxfordRadarSequence, OxfordRadarSequences


class EvaluationTuple:
    # Tuple describing an evaluation set element
    def __init__(self, timestamp: int, rel_scan_filepath: str, position: np.array, pose: np.array):
        # position: x, y position in meters
        # pose: 6 DoF pose (as 4x4 pose matrix)
        assert position.shape == (2,)
        assert pose is None or pose.shape == (4, 4)
        self.timestamp = timestamp
        self.rel_scan_filepath = rel_scan_filepath
        self.position = position
        self.pose = pose

    def to_tuple(self):
        return self.timestamp, self.rel_scan_filepath, self.position, self.pose


class EvaluationSet:
    # Evaluation set consisting of map and query elements
    def __init__(self, query_set: List[EvaluationTuple] = None, map_set: List[EvaluationTuple] = None):
        self.query_set = query_set
        self.map_set = map_set

    def save(self, pickle_filepath: str):
        # Pickle the evaluation set

        # Convert data to tuples and save as tuples
        query_l = []
        for e in self.query_set:
            query_l.append(e.to_tuple())

        map_l = []
        for e in self.map_set:
            map_l.append(e.to_tuple())
        pickle.dump([query_l, map_l], open(pickle_filepath, 'wb'))

    def load(self, pickle_filepath: str):
        # Load evaluation set from the pickle
        query_l, map_l = pickle.load(open(pickle_filepath, 'rb'))

        self.query_set = []
        for e in query_l:
            self.query_set.append(EvaluationTuple(e[0], e[1], e[2], e[3]))

        self.map_set = []
        for e in map_l:
            self.map_set.append(EvaluationTuple(e[0], e[1], e[2], e[3]))

    def get_map_positions(self):
        # Get map positions as (N, 2) array
        positions = np.zeros((len(self.map_set), 2), dtype=self.map_set[0].position.dtype)
        for ndx, pos in enumerate(self.map_set):
            positions[ndx] = pos.position
        return positions

    def get_query_positions(self):
        # Get query positions as (N, 2) array
        positions = np.zeros((len(self.query_set), 2), dtype=self.query_set[0].position.dtype)
        for ndx, pos in enumerate(self.query_set):
            positions[ndx] = pos.position
        return positions

    def get_map_poses(self):
        # Get map positions as (N, 2) array
        poses = np.zeros((len(self.map_set), 4, 4), dtype=self.map_set[0].position.dtype)
        for ndx, pos in enumerate(self.map_set):
            poses[ndx] = pos.pose
        return poses

    def get_query_poses(self):
        # Get query positions as (N, 2) array
        poses = np.zeros((len(self.query_set), 4, 4), dtype=self.query_set[0].position.dtype)
        for ndx, pos in enumerate(self.query_set):
            poses[ndx] = pos.pose
        return poses
    

def get_scans(sequence) -> List[EvaluationTuple]:
    # Get a list of all readings from the test area in the sequence
    elems = []
    for ndx in range(len(sequence)):
        pose = sequence.poses[ndx]
        position = sequence.xys[ndx]
        item = EvaluationTuple(sequence.timestamps[ndx], sequence.filepaths[ndx], position=position, pose=pose)
        elems.append(item)
    return elems


def filter_query_elements(query_set: List[EvaluationTuple], map_set: List[EvaluationTuple],
                          dist_threshold: float) -> List[EvaluationTuple]:
    # Function used in evaluation dataset generation
    # Filters out query elements without a corresponding map element within dist_threshold threshold
    map_pos = np.zeros((len(map_set), 2), dtype=np.float32)
    for ndx, e in enumerate(map_set):
        map_pos[ndx] = e.position

    # Build a kdtree
    kdtree = KDTree(map_pos)

    filtered_query_set = []
    count_ignored = 0
    for ndx, e in enumerate(query_set):
        position = e.position.reshape(1, -1)
        nn = kdtree.query_radius(position, dist_threshold, count_only=True)[0]
        if nn > 0:
            filtered_query_set.append(e)
        else:
            count_ignored += 1

    print(f"{count_ignored} query elements ignored - not having corresponding map element within {dist_threshold} [m] radius")
    return filtered_query_set


def generate_evaluation_set(dataset: str, dataset_root: str, map_sequence: str, query_sequence: str, split: str = 'test', 
map_sampling_distance: float = 20.0, query_sampling_distance: float = 5.0, dist_threshold: float = 20.0):
    assert split in ['test', 'all']
    if dataset == "nclt":
        map = NCLTSequence(dataset_root, sequence_name=map_sequence, split=split, sampling_distance=map_sampling_distance)
        query = NCLTSequence(dataset_root, sequence_name=query_sequence, split=split, sampling_distance=query_sampling_distance)
    elif dataset == "mulran":
        map = MulRanSequence(dataset_root, sequence_name=map_sequence, split=split, sampling_distance=map_sampling_distance)
        query = MulRanSequence(dataset_root, sequence_name=query_sequence, split=split, sampling_distance=query_sampling_distance)
    elif dataset == "kitti":
        map = KITTISequence(dataset_root, sequence_name=map_sequence, split=split, sampling_distance=map_sampling_distance)
        query = KITTISequence(dataset_root, sequence_name=query_sequence, split=split, sampling_distance=query_sampling_distance)
    elif dataset == "oxford_radar":
        map = OxfordRadarSequence(dataset_root, sequence_name=map_sequence, split=split, sampling_distance=map_sampling_distance)
        query = OxfordRadarSequence(dataset_root, sequence_name=query_sequence, split=split, sampling_distance=query_sampling_distance)
    else:
        raise ValueError('dataset is not "nclt" or "mulran" or "kitti" or "oxford_radar"')

    map_set = get_scans(map)
    query_set = get_scans(query)

    # Filters out query elements without a corresponding map element within dist_threshold threshold
    query_set = filter_query_elements(query_set, map_set, dist_threshold)

    print(f'{len(map_set)} database elements, {len(query_set)} query elements')
    return EvaluationSet(query_set, map_set)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, default="nclt", help="dataset type (nclt / mulran / kitti / oxford_radar)")
    parser.add_argument("--dataset_root", type=str, default="./data/NCLT", help="path to the dataset root")
    parser.add_argument("--map_sequence", type=str, default="2012-02-04", help="map sequence for loop closure detection")
    parser.add_argument("--query_sequence", type=str, default="2012-03-17", help="query sequence for loop closure detection")
    parser.add_argument("--map_sampling_distance", type=float, default=20.0, help="map sampling distance in meters")
    parser.add_argument("--query_sampling_distance", type=float, default=5.0, help="query sampling distance in meters")
    parser.add_argument("--dist_threshold", type=float, default=20.0, help="revisit threshold in meters")
    args = parser.parse_args()
    
    print(f'Dataset: {args.dataset}')
    print(f'Dataset root: {args.dataset_root}')
    print(f'Map sequence: {args.map_sequence}')
    print(f'Query sequence: {args.query_sequence}')
    print(f'Map sampling displacement between consecutive anchors: {args.map_sampling_distance}')
    print(f'Query sampling displacement between consecutive anchors: {args.query_sampling_distance}')
    print(f'Ignore query elements without a corresponding map element within a threshold [m]: {args.dist_threshold}')

    test_set = generate_evaluation_set(args.dataset, args.dataset_root, map_sequence=args.map_sequence, query_sequence=args.query_sequence, 
        map_sampling_distance=args.map_sampling_distance, query_sampling_distance=args.query_sampling_distance, dist_threshold=args.dist_threshold)

    pickle_name = f'test_{args.map_sequence}_{args.query_sequence}_{args.map_sampling_distance}_{args.query_sampling_distance}_{args.dist_threshold}.pickle'
    file_path_name = os.path.join(args.dataset_root, pickle_name)
    test_set.save(file_path_name)

evaluation/generate_evaluation_sets.py
lus6-Jenny/RING
import argparse
import matplotlib
import matplotlib.pyplot as plt
matplotlib.rcParams.update({'font.size': 14})
matplotlib.rcParams.update({'font.weight': "bold"})
matplotlib.rcParams.update({'axes.labelweight': "bold"})
matplotlib.rcParams.update({'axes.titleweight': "bold"})
import numpy as np
from sklearn.metrics import precision_recall_curve, auc, average_precision_score
from sklearn.neighbors import KDTree


def calculate_dist(pose1, pose2):
    dist = np.sqrt((pose1[0] - pose2[0])**2 + (pose1[1] - pose2[1])**2)
    return dist


def compute_PR_pairs(pair_dists, query_positions, map_positions, thresholds, save_path, revisit_threshold: float = 10.0):

    tree = KDTree(map_positions)
    num_thresholds = len(thresholds)
    num_hits = np.zeros(num_thresholds) 
    num_false_alarms = np.zeros(num_thresholds) 
    num_correct_rejections = np.zeros(num_thresholds) 
    num_misses = np.zeros(num_thresholds)     
    
    for i in range(pair_dists.shape[0]):
        dist = pair_dists[i]
        query_position = query_positions[i]
        indices = tree.query_radius(query_position.reshape(1,-1), revisit_threshold)[0]
        for j, thre in enumerate(thresholds):
            if(min(dist) < thre):
                # if under the theshold, it is considered seen.
                # and then check the correctness
                real_dist = calculate_dist(query_position, map_positions[np.argmin(dist)])
                if real_dist < revisit_threshold:
                    # TP: Hit
                    num_hits[j] = num_hits[j] + 1
                else:
                    # FP: False Alarm 
                    num_false_alarms[j] = num_false_alarms[j] + 1                         
            
            else:
                if len(indices) == 0:
                    # TN: Correct rejection
                    num_correct_rejections[j] = num_correct_rejections[j] + 1
                else:           
                    # FN: MISS
                    num_misses[j] = num_misses[j] + 1 
    
    precisions = num_hits / (num_hits + num_false_alarms + 1e-10)
    recalls = num_hits / (num_hits + num_misses + 1e-10)
    precisions[np.isnan(precisions)] = 1.0
    recalls[np.isnan(recalls)] = 0.0
    f1s = 2 * (precisions * recalls) / (precisions + recalls + 1e-10)
    
    # plot PR curve
    marker = 'x'
    markevery = 0.03
    plt.clf()
    plt.rcParams.update({'font.size': 16})
    fig = plt.figure()
    plt.plot(recalls, precisions, label='DeepRING', marker=marker, markevery=markevery)
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.xlabel("Recall [%]")
    plt.ylabel("Precision [%]")
    plt.xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0], ["0", "20", "40", "60", "80", "100"])
    plt.yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0], ["0", "20", "40", "60", "80", "100"])
    # plt.show()
    fig.savefig(save_path, bbox_inches='tight', pad_inches=0)
    plt.close()
    return precisions, recalls, f1s


def compute_AP(precisions, recalls):
    ap = 0.0
    for i in range(1, len(precisions)):
        ap += (recalls[i] - recalls[i-1])*precisions[i]
    return ap


def compute_AUC(precisions, recalls):
    auc = 0.0
    for i in range(len(precisions) - 1):
        auc += (recalls[i+1] - recalls[i])*(precisions[i+1] + precisions[i]) / 2.0
    return auc


if __name__ == '__main__':
    # settings
    bev_type = 'occ' # 'occ' or 'feat'
    method = 'RING' if bev_type == 'occ' else 'RING++'
    results_path = f'./results/nclt/test_2012-02-04_2012-03-17_20.0_5.0_10.0/revisit_5.0_10.0_15.0_20.0'
    revisit_thresholds = results_path.split('/')[-1].split('_')[1:]
    
    # load results
    precisions = np.loadtxt(f'{results_path}/precisions_{bev_type}.txt')
    recalls = np.loadtxt(f'{results_path}/recalls_{bev_type}.txt')
    f1s = np.loadtxt(f'{results_path}/f1s_{bev_type}.txt')

    # draw PR curve
    markers = ['o', 's', 'x', 'v', 'd', 'p', 'h']
    colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k']
    fig = plt.figure()

    for idx, revisit_threshold in enumerate(revisit_thresholds):
        # compute AP
        AP = compute_AP(precisions[idx], recalls[idx])
        print(f'AP: {AP}')
        
        # compute AUC
        AUC = compute_AUC(precisions[idx], recalls[idx])
        print(f'AUC: {AUC}')
        
        # plot
        plt.plot(recalls[idx], precisions[idx], marker = markers[idx], color = colors[idx], linewidth = 2, label = f'{method}({revisit_threshold}m)')
    
    plt.legend(loc = 'lower left')
    plt.xlim([0, 1])
    plt.ylim([0, 1])
    plt.xticks([0, 0.2, 0.4, 0.6, 0.8, 1.0], ["0", "20", "40", "60", "80", "100"])
    plt.yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0], ["0", "20", "40", "60", "80", "100"])
    # plot minor gridlines
    plt.minorticks_on()
    plt.grid(which='minor', linestyle='--', alpha=0.2)
    plt.grid(which='major', linestyle='-', alpha=0.5)
    # convert major tick lines from out direction to in direction
    plt.tick_params(axis='both', which='major', direction='in', top=True, right=True)
    # remove minor tick lines and labels
    plt.tick_params(axis='both', which='minor', bottom=False, top=False, labelbottom=False, right=False, left=False, labelleft=False)
    plt.xlabel('Recall [%]')
    plt.ylabel('Precision [%]')
    plt.savefig(f'{results_path}/precision_recall_curve_{bev_type}.pdf', dpi=600, bbox_inches='tight')
    plt.show()
    plt.close()        
evaluation/plot_PR_curve.py
lus6-Jenny/RING
import os
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

import matplotlib
matplotlib.rcParams.update({'font.size': 16})
matplotlib.rcParams.update({'font.weight': "bold"})
matplotlib.rcParams.update({'axes.labelweight': "bold"})
matplotlib.rcParams.update({'axes.titleweight': "bold"})


# plot cdf
def plot_cdf(errors, save_path, xlabel, ylabel, title):
    errors = np.sort(errors)
    y = np.arange(len(errors))/float(len(errors))
    plt.plot(errors, y)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(title)
    plt.savefig(save_path)
    # plt.show()
    plt.close()


# convert numpy to dataframe
def data_to_df(data, label, method):
    num = len(data)
    data_new = np.empty((num, 3), dtype=object)

    for i in range(num):
        data_new[i, 0] = data[i]
        data_new[i, 1] = label
        data_new[i, 2] = method
        
    data_df = pd.DataFrame(data_new)
    data_df.columns = ['Error', 'Threshold', 'Method']
    
    return data_df


if __name__ == "__main__":
    # settings
    icp = False
    bev_type = 'occ' # 'occ' or 'feat'
    method = 'RING' if bev_type == 'occ' else 'RING++'
    results_path = f'./results/nclt/test_2012-02-04_2012-03-17_20.0_5.0_10.0/revisit_5.0_10.0_15.0_20.0'
    revisit_thresholds = results_path.split('/')[-1].split('_')[1:]
    quantiles = [0.25, 0.5, 0.75, 0.95]
    
    # load results
    if icp:
        rot_error = np.loadtxt(f'{results_path}/icp_rot_errors_{bev_type}.txt') 
        trans_error = np.loadtxt(f'{results_path}/icp_trans_errors_{bev_type}.txt') 
    else:
        rot_error = np.loadtxt(f'{results_path}/rot_errors_{bev_type}.txt') 
        trans_error = np.loadtxt(f'{results_path}/trans_errors_{bev_type}.txt')         

    # convert to dataframe
    df_rot_total = pd.DataFrame()
    df_trans_total = pd.DataFrame()
    for idx, revisit_threshold in enumerate(revisit_thresholds):
        rot_err = rot_error[idx]
        trans_err = trans_error[idx]
        df_rot = data_to_df(rot_err, label=revisit_threshold, method=method)
        df_trans = data_to_df(trans_err, label=revisit_threshold, method=method)
        df_rot_total = pd.concat([df_rot, df_rot_total])
        df_trans_total = pd.concat([df_trans, df_trans_total])
        
        rot_error_quantiles = np.quantile(rot_err, quantiles)
        trans_error_quantiles = np.quantile(trans_err, quantiles)
        print(f"Rotation error quantiles at {revisit_threshold} m: {rot_error_quantiles}")
        print(f"Translation error quantiles at {revisit_threshold} m: {trans_error_quantiles}")        
        rot_success = rot_err[rot_err<=5]
        trans_success = trans_err[trans_err<=2]
        num = 0
        for j in range(len(rot_err)):
            if rot_err[j] <= 5 and trans_err[j] <= 2:
                num += 1
        print(f"Rotation success rate at {revisit_threshold} m: ", len(rot_success)/len(rot_err))
        print(f"Translation success rate at {revisit_threshold} m: ", len(trans_success)/len(trans_err))
        print(f"Pose success rate at {revisit_threshold} m: ", num/len(rot_err))

    markers = ['o', 's', 'x', 'v', 'd', 'p', 'h']
    colors = ['r', 'g', 'b', 'c', 'm', 'y', 'k']

    # draw boxplot for rotation and translation error
    fig, axs = plt.subplots(1,2, figsize=(12, 4))
    plt.subplots_adjust(wspace=0.3, hspace=0.3)
    my_pal = {"RING": colors[0], "RING++": colors[1]}
    sns.boxplot(data=df_rot_total, x='Threshold', y='Error', hue='Method', width=0.5, fliersize=2, linewidth=1.5, ax=axs[0], palette=my_pal)
    axs[0].set_yscale('log')
    axs[0].set_ylim(1e-1, max(df_rot_total['Error'].max(), df_rot_total['Error'].max()))
    axs[0].set_xlabel("Revisit Threshold (m)")
    axs[0].set_ylabel("Rotation Error (deg)")
    axs[0].legend([],[], frameon=False)

    sns.boxplot(data=df_trans_total, x='Threshold', y='Error', hue='Method', width=0.5, fliersize=2, linewidth=1.5, ax=axs[1], palette=my_pal)
    axs[1].set_ylim([0, 20])
    axs[1].set_xlabel("Revisit Threshold (m)")
    axs[1].set_ylabel("Translation Error (m)")
    axs[1].legend([],[], frameon=False)
    plt.legend(loc='upper right')
    plt.rcParams['xtick.direction'] = 'in'
    plt.rcParams['ytick.direction'] = 'in'
    plt.savefig(os.path.join(results_path, f"pose_estimation_error_{bev_type}.pdf"), dpi=600, bbox_inches='tight')
    plt.show()
    plt.close()

evaluation/plot_pose_errors.py
lus6-Jenny/RING
import os
import sys
import time
import torch
import numpy as np
from utils.core import *
from utils.icp import icp
import utils.config as cfg
from utils.tools import imshow
from utils.poses import xyz_ypr2m, m2xyz_ypr
from utils.point_clouds import load_pc_infer, fast_gicp, draw_pc, draw_pc_pair, draw_registration_result

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print('device', device)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
print('BASE_DIR', BASE_DIR)
sys.path.append(BASE_DIR)


# Apply 4x4 transformation matrix to point cloud
def apply_transform(pc, transform): 
    pc_tranformed = np.dot(pc, transform[:3, :3].T) + transform[:3, 3]
    pc_tranformed = pc_tranformed.astype(np.float32)

    return pc_tranformed


if __name__ == '__main__':
    pc1 = np.load('./test.npy')
    gt_x = np.random.uniform(-10, 10)
    gt_y = np.random.uniform(-10, 10)
    gt_yaw = np.random.uniform(0, 2*np.pi)
    transform = xyz_ypr2m(gt_x, gt_y, 0, gt_yaw, 0, 0)
    pc2 = apply_transform(pc1, transform)

    pc1, pc1_norm = load_pc_infer(pc1)
    pc2, pc2_norm = load_pc_infer(pc2)

    # draw_pc(pc1)
    # draw_pc(pc2)
    # draw_pc_pair(pc1, pc2)
    
    # ------------ RING Generation ------------
    bev_type = "occ" # "occ" or "feat"
    # generate BEV features
    times = time.time()
    if device == torch.device("cpu"):
        pc1_bev = generate_bev(pc1) # BEV encoded by occupancy information
        pc2_bev = generate_bev(pc2) # BEV encoded by occupancy information
    else:      
        if bev_type == "occ":
            pc1_bev = generate_bev_occ(pc1_norm) # BEV encoded by occupancy information
            pc2_bev = generate_bev_occ(pc2_norm) # BEV encoded by occupancy information
        elif bev_type == "feat":
            pc1_bev = generate_bev_feats(pc1_norm) # BEV encoded by point features
            pc2_bev = generate_bev_feats(pc2_norm) # BEV encoded by point features
        else:
            raise ValueError('bev_type is not "occ" or "feat"')
    timee = time.time()
    # print("Time of BEV generation:", timee - times, 's')
    print("BEV shapes:", pc1_bev.shape, pc2_bev.shape)

    # generate RING and TIRING descriptors
    times = time.time()
    if device == torch.device("cuda:0"):
        pc1_RING, pc1_TIRING = generate_RING(pc1_bev)
        pc2_RING, pc2_TIRING = generate_RING(pc2_bev)        
    elif device == torch.device("cpu"):
        pc1_RING, pc1_TIRING = generate_RING_cpu(pc1_bev)
        pc2_RING, pc2_TIRING = generate_RING_cpu(pc2_bev)
    timee = time.time()
    # print("Time of RING generation:", timee - times, 's')
    # ------------ RING Generation ------------

    # ------------ Pose Estimation ------------
    ang_res = 2 * np.pi / cfg.num_ring # angular resolution
    dist, angle_matched = fast_corr(pc1_TIRING, pc2_TIRING)

    # angle between the two matched RINGs in grids
    angle_matched_extra = angle_matched - cfg.num_ring // 2
    # convert the matched angle from grids to radians
    angle_matched_rad = angle_matched * ang_res 
    angle_matched_extra_rad = angle_matched_extra * ang_res        

    bev_rotated = rotate_bev(pc1_bev, angle_matched_rad)
    bev_rotated_extra = rotate_bev(pc1_bev, angle_matched_extra_rad)
    
    # imshow(pc1_bev[0], 'bev1.png')
    # imshow(pc1_RING[0], 'sinogram1.png')
    # imshow(pc1_TIRING[0], 'ting1.png')
    # imshow(pc2_bev[0], 'bev2.png')
    # imshow(pc2_RING[0], 'sinogram2.png')
    # imshow(pc2_TIRING[0], 'ting2.png')
    # imshow(bev_rotated[0], 'bev_rotated.png')
    # imshow(bev_rotated_extra[0], 'bev_rotated_extra.png')

    # solve the translation between the two matched bevs
    x, y, error = solve_translation(bev_rotated, pc2_bev)
    x_extra, y_extra, error_extra = solve_translation(bev_rotated_extra, pc2_bev)

    if error < error_extra:
        trans_x = x / cfg.num_sector * (cfg.point_cloud["x_bound"][1] - cfg.point_cloud["x_bound"][0])  # in meters
        trans_y = y / cfg.num_ring * (cfg.point_cloud["y_bound"][1] - cfg.point_cloud["y_bound"][0])   # in meters
        rot_yaw = angle_matched_rad  # in radians
    else:
        trans_x = x_extra / cfg.num_sector * (cfg.point_cloud["x_bound"][1] - cfg.point_cloud["x_bound"][0])  # in meters
        trans_y = y_extra / cfg.num_ring * (cfg.point_cloud["y_bound"][1] - cfg.point_cloud["y_bound"][0])  # in meters 
        rot_yaw = angle_matched_extra_rad  # in radians
    
    init_pose = xyz_ypr2m(trans_x, trans_y, 0, rot_yaw, 0, 0)
    print("Ground truth translation: x: {}, y: {}, rotation: {}".format(gt_x, gt_y, gt_yaw))
    print("Estimated translation: x: {}, y: {}, rotation: {}".format(trans_x, trans_y, rot_yaw))
    draw_registration_result(pc1, pc2, init_pose)
    # ------------ Pose Estimation ------------

    # ------------ Pose Refinement ------------
    times = time.time()
    icp_fitness_score, loop_transform = fast_gicp(pc1, pc2, max_correspondence_distance=cfg.icp_max_distance, init_pose=init_pose)
    timee = time.time() 

    # print("ICP fitness score:", icp_fitness_score)       
    # print("ICP processed time:", timee - times, 's')
    x, y, z, yaw, pitch, roll = m2xyz_ypr(loop_transform)
    print("Refined translation: x: {}, y: {}, rotation: {}".format(x, y, yaw))
    draw_registration_result(pc1, pc2, loop_transform)
    # ------------ Pose Refinement ------------
test.py
lus6-Jenny/RING
# circular correlation between two tensors

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms.functional as fn
import matplotlib.pyplot as plt

# compute the max correlation value and the corresponding circular shift 
# calculate the correlation and relative rotation between a and b (from a to b, point cloud transformation) using fft method

class circorr2(nn.Module):

    '''
    Circularly Correlates two images. Images must be of the same size.
    '''

    def __init__(self, is_circular=True, zero_mean_normalize=False):
        super(circorr2, self).__init__()
        self.InstanceNorm = nn.InstanceNorm2d(1, affine=False, track_running_stats=False)
        self.is_circular = is_circular
        self.zero_mean_normalize = zero_mean_normalize


    def forward(self, x1, x2, scale=1):
        """
        Args:
            x1 (torch.Tensor): Batch of Img1 of dimensions [B, C, H, W].
            x2 (torch.Tensor): Batch of Img2 of dimensions [B, C, H, W].
        Returns:
            scores (torch.Tensor): The circular cross correlation scores for the pairs. The output shape is [B, 1].
        """
        
        assert x1.shape[2:] == x2.shape[2:], "Images must be of the same size."

        if self.is_circular:
            corr, scores, shifts = self.circular_corr(x1, x2, scale)
            return corr, scores, shifts

        else:
            scores = self.corr(x1, x2)   
            return scores


    def corr(self, x1, x2):
        b, c, h, w = x2.shape

        if self.zero_mean_normalize:
            x1 = self.InstanceNorm(x1)
            x2 = self.InstanceNorm(x2)
        else:
            x1 = F.normalize(x1, dim=(-2,-1))
            x2 = F.normalize(x2, dim=(-2,-1))

        scores = torch.matmul(x1.view(b, 1, c*h*w), x2.view(b, c*h*w, 1))
        scores /= (h*w*c)
        
        return scores


    def circular_corr(self, x1, x2, scale=1):
        b, c, h, w = x2.shape

        if self.zero_mean_normalize:
            x1 = self.InstanceNorm(x1)
            x2 = self.InstanceNorm(x2)
        else:
            x1 = F.normalize(x1, dim=(-2,-1))
            x2 = F.normalize(x2, dim=(-2,-1))

        a_fft = torch.fft.fft2(x1, dim=-2, norm="ortho")
        b_fft = torch.fft.fft2(x2, dim=-2, norm="ortho")
        corr = torch.fft.ifft2(a_fft*b_fft.conj(), dim=-2, norm="ortho")  
        corr = torch.sqrt(corr.real**2 + corr.imag**2 + 1e-15)
        # sum the correlation over the channels
        corr = torch.sum(corr, dim=-3)
        # sum the correlation over the width
        corr = torch.sum(corr, dim=-1)
        # self.plot_corr(corr[0].detach().cpu().numpy(), 'corr.png')
        if scale != 1:
            # apply curve smoothing
            corr = self.curve_smoothing(corr, scale)
            # self.plot_corr(corr[0].detach().cpu().numpy(), 'corr_smooth.png')
        # find the max correlation value as score and the corresponding circular shift
        score = torch.max(corr, dim=-1)[0]/(0.15*c*h*w)
        # shift the correlation to the center 
        corr = torch.fft.fftshift(corr, dim=-1)
        shift = corr.shape[-1]//2 - torch.argmax(corr, dim=-1)
        return corr, score, shift


    def curve_smoothing(self, F, S):
        # width of the cross-correlation curve F
        B = F.shape[0]
        W = F.shape[1]
        # perform FFT to get the frequency domain representation of the curve
        F = torch.fft.fft2(F, dim=-1, norm="ortho")
        # zero-pad the curve
        padded_F = torch.cat((F[:,:int(W/2)], torch.zeros(B,(S-1)*W).to(F.device), F[:,int(W/2):]), dim=-1)
        # perform inverse FFT to get the smoothed curve
        smoothed_F = torch.fft.ifft2(padded_F, dim=-1, norm="ortho")
        # take the real part of the smoothed curve
        # smoothed_F_mag = torch.sqrt(smoothed_F.real**2 + smoothed_F.imag**2 + 1e-15)
        smoothed_F = smoothed_F.real
        return smoothed_F


    def plot_corr(self, a, save_path=None):
        # plot the correlation curve
        plt.plot(a)
        # plt.show()
        if save_path is not None:
           plt.savefig(save_path)
        plt.close()
        
        
if __name__ == "__main__":
    BATCH_SIZE = 10
    NUM_CHANNELS = 3
    IMG_H = 64
    IMG_W = 64
    x1 = torch.randn(BATCH_SIZE, NUM_CHANNELS, IMG_H, IMG_W)
    x2 = torch.randn(BATCH_SIZE, NUM_CHANNELS, IMG_H, IMG_W)
    circorr = circorr2()
    scores,shifts = circorr(x1, x2)
    print('output shape', scores.shape, shifts.shape)
    print('scores', scores)
    print('shifts', shifts)
utils/circorr2.py
lus6-Jenny/RING
'''
Parameters of RING for Loop Detection
'''

# Sampling Gap
train_sampling_gap = 10 # in meters
test_sampling_gap = 10 # in meters

# Top k Loop Detection 
top_k = 10

# Formation Loop Detection Window Size
window_size = 5

# Batch Size of Correlation Calculation
batch_size = 16

# Point Cloud Process
point_cloud = {
    'x_bound': [-70.0, 70.0],
    'y_bound': [-70.0, 70.0],
    'z_bound': [1.0, 20.0]
}

# BEV Resolution
num_ring = 120
num_sector = 120
num_height = 1
max_length = 1
max_height = 1

# Place Recognition Parameters
search_ratio = 0.1
num_candidates = 10
exclude_recent_nodes = 30
dist_threshold = 0.48 # descriptor distance (correlation) threshold for place recognition
revisit_criteria = 10 # in meters

# ICP Parameters
max_icp_iter = 100
icp_tolerance = 0.001
icp_max_distance = 5.0
num_icp_points = 6000
icp_fitness_score = 0.12
utils/config.py
lus6-Jenny/RING
import os
import sys
import cv2
import time
import torch

import voxelocc
import voxelfeat
import numpy as np
import open3d as o3d
import utils.config as cfg
from utils.tools import imshow
import utils.vox_utils.vox as vox
from utils.circorr2 import circorr2
from utils.extract_local_descriptor import build_neighbors_NN

import torch.nn as nn
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torchvision.transforms.functional as fn
from torch_radon import Radon, ParallelBeam, RadonFanbeam
from skimage.transform import radon, iradon, rescale

np.seterr(divide='ignore',invalid='ignore')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
InstanceNorm = nn.InstanceNorm2d(1, affine=False, track_running_stats=False)


############ FFT Operations ############
# Apply 2D fourier transform to the input data
def forward_fft(input):
    median_output = torch.fft.fft2(input, dim=(-2, -1), norm="ortho")
    median_output_r = median_output.real
    median_output_i = median_output.imag
    output = torch.sqrt(median_output_r ** 2 + median_output_i ** 2)
    # output = torch.fft.fftshift(output)
    return output, median_output


# Apply 1D fourier transform to the row of input data
def forward_row_fft(input):
    median_output = torch.fft.fft2(input, dim=-1, norm="ortho")
    median_output_r = median_output.real
    median_output_i = median_output.imag
    output = torch.sqrt(median_output_r ** 2 + median_output_i ** 2)
    return output, median_output


# Apply 1D fourier transform to the column of input data
def forward_column_fft(input):
    median_output = torch.fft.fft2(input, dim=-2, norm="ortho")
    median_output_r = median_output.real
    median_output_i = median_output.imag
    output = torch.sqrt(median_output_r ** 2 + median_output_i ** 2)
    return output, median_output
############ FFT Operations ############


############ BEV Features Generation ############
# Generate BEV encoded by occupancy information
def generate_bev(pc):
    '''
    pc: original point cloud
    '''
    # returns cfg.num_height * cfg.num_ring * cfg.num_sector matrix 
    pc = pc[:,:3]
    bounds = (cfg.point_cloud["x_bound"][0], cfg.point_cloud["x_bound"][1], cfg.point_cloud["y_bound"][0], \
              cfg.point_cloud["y_bound"][1], cfg.point_cloud["z_bound"][0], cfg.point_cloud["z_bound"][1])
    scene_centroid = [0.0, 0.0, 0.0]
    scene_centroid = torch.from_numpy(np.array(scene_centroid).reshape([1, 3])).float()
    vox_util = vox.Vox_util(cfg.num_height, cfg.num_ring, cfg.num_sector,
                            scene_centroid=scene_centroid.cuda(),
                            bounds=bounds,
                            assert_cube=False)
    pc = torch.from_numpy(pc).unsqueeze(0)
    occ_mem = vox_util.voxelize_xyz(pc, cfg.num_height, cfg.num_ring, cfg.num_sector, assert_cube=False)
    pc_bev = occ_mem.permute(0,1,3,2,4).squeeze()
    dim = pc_bev.ndim
    if dim == 2:
        return pc_bev.permute(1,0).unsqueeze(0).to(device)
    elif dim == 3:
        return pc_bev.permute(1,2,0).to(device)
    else:
        ValueError('pc_bev dim is not 2 or 3')


# Generate BEV encoded by occupancy information with CUDA accelerate 
def generate_bev_occ(pc):
    '''
    pc: normalized point cloud
    '''     
    # returns cfg.num_height * cfg.num_ring * cfg.num_sector matrix
    size = pc.shape[0]
    pc = pc[:,:3]
    pc_bev = np.zeros([cfg.num_height * cfg.num_ring * cfg.num_sector])
    pc = pc.transpose().flatten().astype(np.float32)

    transer_bev = voxelocc.GPUTransformer(pc, size, cfg.max_length, cfg.max_height, cfg.num_ring, cfg.num_sector, cfg.num_height, 1)
    transer_bev.transform()
    point_t_bev = transer_bev.retreive()
    point_t_bev = point_t_bev.reshape(-1, 3)
    point_t_bev = point_t_bev[...,2]

    pc_bev = point_t_bev.reshape(cfg.num_height, cfg.num_ring, cfg.num_sector)
    pc_bev = torch.from_numpy(pc_bev).to(device)
    return pc_bev


# Generate BEV encoded by point features with CUDA accelerate 
def generate_bev_feats(pc):
    '''
    pc: normalized point cloud
    ''' 
    # returns cfg.num_height * num_feats * cfg.num_ring * cfg.num_sector matrix    
    size = pc.shape[0]
    pc = pc[:,:3]

    ###### KNN ######
    k = 30  # num of neighbors
    k_indices, k_covs, k_entropy, k_eigens_, k_vectors_  = build_neighbors_NN(pc, k)

    ###### LPD-Net Features ######
    k_indices = k_indices.flatten().astype(np.int32)
    k_eigens_ = k_eigens_.flatten()
    pc_flatten = pc.flatten()

    featureExtractor = voxelfeat.GPUFeatureExtractor(pc_flatten, size, 13, 30, k_indices, k_eigens_)
    GPUfeat = featureExtractor.get_features()
    GPUfeat = GPUfeat.reshape(-1,13)
    GPUfeat = GPUfeat.take([0,1,3,10,11,12], axis=1)  # [0,1,3,10,11,12] x 5,7,8,9
    pc = np.concatenate((pc, GPUfeat), axis=1)  # num_points * 13 
    featsize = pc.shape[1]

    pc = pc.transpose().flatten().astype(np.float32)

    transer_bev = voxelfeat.GPUTransformer(pc, size, cfg.max_length, cfg.max_height, cfg.num_ring, cfg.num_sector, cfg.num_height, featsize)
    transer_bev.transform()
    point_t_bev = transer_bev.retreive()
    point_t_bev = point_t_bev.reshape(-1, featsize)
    point_t_bev = point_t_bev[...,3:]
    pc_bev = point_t_bev.reshape(cfg.num_height, cfg.num_ring, cfg.num_sector, featsize-3)
    
    pc_bev = torch.from_numpy(pc_bev).to(device)
    pc_bev = pc_bev.squeeze(0)
    pc_bev = pc_bev.permute(2,0,1)
    return pc_bev
############ BEV Features Generation ############


############ RING Descriptor Generation ############
#  GPU version of RING descriptor generation
def generate_RING(pc_bev):
    angles = torch.FloatTensor(np.linspace(0, 2*np.pi, cfg.num_ring).astype(np.float32))
    radon = ParallelBeam(cfg.num_sector, angles)

    pc_RING = radon.forward(pc_bev)
    pc_TIRING, _ = forward_row_fft(pc_RING)

    return pc_RING, pc_TIRING


#  CPU version of RING descriptor generation
def generate_RING_cpu(pc_bev):
    angles = np.linspace(0., 360., cfg.num_ring, endpoint=False).astype(np.float32)
    pc_RING = torch.from_numpy(radon(pc_bev.numpy().squeeze(0), theta=angles)).to(device).unsqueeze(0)
    pc_RING = pc_RING.transpose(1, 2)

    pc_TIRING, _ = forward_row_fft(pc_RING)

    return pc_RING, pc_TIRING
############ RING Descriptor Generation ############



############ Pose Estimation ############
# Compute the max correlation value and the corresponding circular shift 
def fast_corr(a, b, zero_mean_normalize=True):
    if zero_mean_normalize:
        a = fn.normalize(a, mean=a.mean(), std=a.std())
        b = fn.normalize(b, mean=b.mean(), std=b.std())
    else:
        a = F.normalize(a, dim=(-2,-1))
        b = F.normalize(b, dim=(-2,-1))    
    a_fft = torch.fft.fft2(a, dim=-2, norm="ortho")
    b_fft = torch.fft.fft2(b, dim=-2, norm="ortho")
    corr = torch.fft.ifft2(a_fft*b_fft.conj(), dim=-2, norm="ortho")  
    corr = torch.sqrt(corr.real**2 + corr.imag**2)
    corr = torch.sum(corr, dim=-3) # add this line for multi feature channels
    corr = torch.sum(corr, dim=-1).view(-1, cfg.num_ring) 
    
    corr = torch.fft.fftshift(corr, dim=-1)
    angle = cfg.num_ring//2 - torch.argmax(corr, dim=-1)
    max_corr = torch.max(corr, dim=-1)[0]
    dist = 1 - max_corr/(0.15*a.shape[-3]*a.shape[-2]*a.shape[-1])
    dist = dist.cpu().numpy().squeeze()
    angle = angle.cpu().numpy().squeeze()

    return dist, angle


def batch_circorr(a, b, scale=1):
    circorr = circorr2(is_circular=True, zero_mean_normalize=True)
    corr, score, angle = circorr(a, b, scale)
    dist = 1 - score
    dist = dist.cpu().numpy()
    angle = angle.cpu().numpy()

    return dist, angle


# Calculate the circular shift to compensate the yaw difference between two RINGs
def calculate_row_shift(shift):
    if shift < cfg.num_ring // 2:
        shift = -shift 
    else:
        shift = shift - cfg.num_ring
    return shift


# Rotate tensor image with torch
def rotate_bev(image: torch.Tensor, angle: float):
    # convert angle to degree
    angle = angle * 180 / np.pi
    image_rotated = fn.rotate(image, angle)
    return image_rotated


# Rotate numpy image with cv2
def rotate_image(image: np.ndarray, angle: float):
    col, row = cfg.num_ring, cfg.num_sector
    image_center = (col // 2, row // 2)
    rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
    # result = cv2.warpAffine(image, rot_mat,  (col, row), flags=cv2.INTER_LINEAR)
    result = cv2.warpAffine(image, rot_mat, (col, row))
    # plt.imshow(result)
    # plt.show()
    return result


# Exhaustive search for the best translation
def solve_translation(a, b, zero_mean_normalize=True):
    if zero_mean_normalize:
        a = fn.normalize(a, mean=a.mean(), std=a.std())
        b = fn.normalize(b, mean=b.mean(), std=b.std())
    else:
        a = F.normalize(a, dim=(-2,-1))
        b = F.normalize(b, dim=(-2,-1))  
    # 2D cross correlation
    a_fft = torch.fft.fft2(a, dim=(-2,-1), norm="ortho")
    b_fft = torch.fft.fft2(b, dim=(-2,-1), norm="ortho")
    corr = torch.fft.ifft2(a_fft*b_fft.conj(), dim=(-2,-1), norm="ortho")  
    corr = torch.sqrt(corr.real**2 + corr.imag**2)
    corr = torch.sum(corr, dim=-3) # add this line for multi feature channels
    # print((corr==torch.max(corr)).nonzero())

    # shift the correlation to the center
    corr = torch.fft.fftshift(corr, dim=(-2,-1))

    # get the index of the max correlation
    idx_x = (corr==torch.max(corr)).nonzero()[0][0]
    idx_y = (corr==torch.max(corr)).nonzero()[0][1]
    # print(idx_x, idx_y)    

    # get the translation value
    x = cfg.num_ring//2 - idx_x
    y = cfg.num_sector//2 - idx_y

    return x.cpu().numpy(), y.cpu().numpy(), -torch.max(corr).cpu().numpy()
############ Pose Estimation ############
utils/core.py
lus6-Jenny/RING
'''
	Pre-processing: prepare_data in LPD-Net
	generate KNN neighborhoods and calculate feature as the feature matrix of point
	Reference: LPD-Net: 3D Point Cloud Learning for Large-Scale Place Recognition and Environment Analysis, ICCV 2019
	author: Chuanzhe Suo(suo_ivy@foxmail.com)
	created: 10/26/18
'''
                        
import os
import sys
import time 
import torch
import numpy as np
import pandas as pd
from sklearn.neighbors import NearestNeighbors, KDTree

np.seterr(divide='ignore',invalid='ignore')

def calculate_entropy_array(eigen):
    L_ = (eigen[:,0] - eigen[:,1]) / eigen[:,0]
    P_ = (eigen[:,1] - eigen[:,2]) / eigen[:,0]
    S_ = eigen[:,2] / eigen[:,0]
    Entropy = -L_*np.log(L_)-P_*np.log(P_)-S_*np.log(S_)
    return Entropy


def covariation_eigenvalue(pointcloud, neighborhood_index):
    ### calculate covariation and eigenvalue of 3D and 2D
    # prepare neighborhood
    neighborhoods = pointcloud[neighborhood_index]

    # 3D cov and eigen by matrix
    Ex = np.average(neighborhoods, axis=1)
    Ex = np.reshape(np.tile(Ex,[neighborhoods.shape[1]]), neighborhoods.shape)
    P = neighborhoods-Ex
    cov_ = np.matmul(P.transpose((0,2,1)),P)/(neighborhoods.shape[1]-1)

    eigen_ = torch.linalg.eigvalsh(torch.from_numpy(cov_))
    eigen_ = eigen_.cpu().numpy()
    indices = np.argsort(eigen_)
    indices = indices[:,::-1]
    pcs_num_ = eigen_.shape[0]
    indx = np.reshape(np.arange(pcs_num_), [-1, 1])
    eig_ind = indices + indx*3
    vec_ind = np.reshape(eig_ind*3, [-1,1]) + np.full((pcs_num_*3,3), [0,1,2])
    vec_ind = np.reshape(vec_ind, [-1,3,3])
    eigen3d_ = np.take(eigen_, eig_ind)
    vectors_ = []
    entropy_ = calculate_entropy_array(eigen3d_)

    # 2d cov and eigen
    cov2d_ = cov_[:,:2,:2]
    eigen2d = torch.linalg.eigvalsh(torch.from_numpy(cov2d_))
    eigen2d = eigen2d.cpu().numpy()
    indices = np.argsort(eigen2d)
    indices = indices[:, ::-1]
    pcs_num_ = eigen2d.shape[0]
    indx = np.reshape(np.arange(pcs_num_), [-1, 1])
    eig_ind = indices + indx * 2
    eigen2d_ = np.take(eigen2d, eig_ind)

    eigens_ = np.append(eigen3d_,eigen2d_,axis=1)

    return cov_, entropy_, eigens_, vectors_


def build_neighbors_NN(pointcloud, k):
	### using KNN NearestNeighbors cluster according k
    start = time.time()
    nbrs = NearestNeighbors(n_neighbors=k, algorithm='kd_tree', n_jobs=-1).fit(pointcloud)
    distances, indices = nbrs.kneighbors(pointcloud)
    end = time.time()
    # print("neighbors time: ", end-start)
    
    start2 = time.time()
    covs, entropy, eigens_, vectors_ = covariation_eigenvalue(pointcloud, indices)
    end2 = time.time()
    # print("covariation_eigenvalue time: ", end2-start2)

    return indices, covs, entropy, eigens_, vectors_


def calculate_features(pointcloud, nbrs_index, eigens_, vectors_):
    ### calculate handcraft feature with eigens and statistics data

    # features using eigens
    eig3d = eigens_[:3]
    eig2d = eigens_[3:5]

    # 3d
    C_ = eig3d[2] / (eig3d.sum()) # change of curvature
    O_ = np.power((eig3d.prod() / np.power(eig3d.sum(), 3)), 1.0 / 3) # omni-variance
    # L_ = (eig3d[0] - eig3d[1]) / eig3d[0]
    E_ = -((eig3d / eig3d.sum()) * np.log(eig3d / eig3d.sum())).sum() # eigenvalue-entropy
    # P_ = (eig3d[1] - eig3d[2]) / eig3d[0]
    # S_ = eig3d[2] / eig3d[0]
    # A_ = (eig3d[0] - eig3d[2]) / eig3d[0]
    # X_ = eig3d.sum()
    # D_ = 3 * nbrs_index.shape[0] / (4 * math.pi * eig3d.prod())
    # 2d
    # S_2 = eig2d.sum() # 2D scattering
    L_2 = eig2d[1] / eig2d[0] # 2D linearity
    # features using statistics data
    neighborhood = pointcloud[nbrs_index]
    nbr_dz = neighborhood[:, 2] - neighborhood[:, 2].min()
    dZ_ = nbr_dz.max() # Fz features - maximum height difference
    vZ_ = np.var(nbr_dz) # Fz features - height variance
    # V_ = vectors_[2][2]
    # not ok: L_, D_, S_2, V_
    features = np.asarray([C_, O_, E_, L_2, dZ_, vZ_]) #([C_,O_,L_,E_,D_,S_2,L_2,dZ_,vZ_,V_])
    # features = np.asarray([dZ_])
    return features


def get_pointfeat(pointcloud):
    k = 30
    k_indices, k_covs, k_entropy, k_eigens_, k_vectors_  = build_neighbors_NN(pointcloud, k)
    start2 = time.time()

    points_feature = []
    for index in range(pointcloud.shape[0]):
        neighborhood = k_indices[index]
        eigens_ = k_eigens_[index]
        vectors_ = k_vectors_[index]

        # calculate point feature
        feature = calculate_features(pointcloud, neighborhood, eigens_, vectors_)
        points_feature.append(feature)

    end2 = time.time()
    print('calculate_features time:', end2-start2)

    points_feature = np.asarray(points_feature)

    return points_feature
utils/extract_local_descriptor.py
lus6-Jenny/RING
# from future.utils import iteritems
import os
from os.path import join as pjoin
from setuptools import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
import numpy


def find_in_path(name, path):
    """Find a file in a search path"""

    # Adapted fom http://code.activestate.com/recipes/52224
    for dir in path.split(os.pathsep):
        binpath = pjoin(dir, name)
        if os.path.exists(binpath):
            return os.path.abspath(binpath)
    return None


def locate_cuda():
    """Locate the CUDA environment on the system

    Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
    and values giving the absolute path to each directory.

    Starts by looking for the CUDAHOME env variable. If not found,
    everything is based on finding 'nvcc' in the PATH.
    """

    # First check if the CUDAHOME env variable is in use
    if 'CUDAHOME' in os.environ:
        home = os.environ['CUDAHOME']
        nvcc = pjoin(home, 'bin', 'nvcc')
    else:
        # Otherwise, search the PATH for NVCC
        nvcc = find_in_path('nvcc', os.environ['PATH'])
        if nvcc is None:
            raise EnvironmentError('The nvcc binary could not be '
                'located in your $PATH. Either add it to your path, '
                'or set $CUDAHOME')
        home = os.path.dirname(os.path.dirname(nvcc))

    cudaconfig = {'home': home, 'nvcc': nvcc,
                  'include': pjoin(home, 'include'),
                  'lib64': pjoin(home, 'lib64')}
    for k, v in iter(cudaconfig.items()):
        if not os.path.exists(v):
            raise EnvironmentError('The CUDA %s path could not be '
                                   'located in %s' % (k, v))

    return cudaconfig


def customize_compiler_for_nvcc(self):
    """Inject deep into distutils to customize how the dispatch
    to gcc/nvcc works.

    If you subclass UnixCCompiler, it's not trivial to get your subclass
    injected in, and still have the right customizations (i.e.
    distutils.sysconfig.customize_compiler) run on it. So instead of going
    the OO route, I have this. Note, it's kindof like a wierd functional
    subclassing going on.
    """

    # Tell the compiler it can processes .cu
    self.src_extensions.append('.cu')

    # Save references to the default compiler_so and _comple methods
    default_compiler_so = self.compiler_so
    super = self._compile

    # Now redefine the _compile method. This gets executed for each
    # object but distutils doesn't have the ability to change compilers
    # based on source extension: we add it.
    def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
        if os.path.splitext(src)[1] == '.cu':
            # use the cuda for .cu files
            self.set_executable('compiler_so', CUDA['nvcc'])
            # use only a subset of the extra_postargs, which are 1-1
            # translated from the extra_compile_args in the Extension class
            postargs = extra_postargs['nvcc']
        else:
            postargs = extra_postargs['gcc']

        super(obj, src, ext, cc_args, postargs, pp_opts)
        # Reset the default compiler_so, which we might have changed for cuda
        self.compiler_so = default_compiler_so

    # Inject our redefined _compile method into the class
    self._compile = _compile


# Run the customize_compiler
class custom_build_ext(build_ext):
    def build_extensions(self):
        customize_compiler_for_nvcc(self.compiler)
        build_ext.build_extensions(self)


CUDA = locate_cuda()

# Obtain the numpy include directory. This logic works across numpy versions.
try:
    numpy_include = numpy.get_include()
except AttributeError:
    numpy_include = numpy.get_numpy_include()


ext = Extension('voxelocc',
        sources = ['src/manager.cu', 'wrapper.pyx'],
        library_dirs = [CUDA['lib64']],
        libraries = ['cudart'],
        language = 'c++',
        runtime_library_dirs = [CUDA['lib64']],
        # This syntax is specific to this build system
        # we're only going to use certain compiler args with nvcc
        # and not with gcc the implementation of this trick is in
        # customize_compiler()
        extra_compile_args= {
            'gcc': [],
            'nvcc': [
                '--ptxas-options=-v', '-c',
                '--compiler-options', "'-fPIC'", "-std=c++11"
                ]
            },
            include_dirs = [numpy_include, CUDA['include'], 'src']
        )


setup(name = 'voxelocc',
      author = 'Xuecheng Xu',
      version = '0.1',
      ext_modules = [ext],
      cmdclass = {'build_ext': custom_build_ext},
      # Since the package has c code, the egg cannot be zipped
      zip_safe = False)

utils/generate_bev_occ_cython/setup.py
lus6-Jenny/RING
import voxelocc
import numpy as np
import numpy.testing as npt
import time
import os
import numpy.testing as npt
import matplotlib.pyplot as plt
from skimage.transform import radon
from mpl_toolkits.mplot3d import Axes3D

def load_pc_file(filename):
    # returns Nx3 matrix
    pc = np.fromfile(os.path.join("./", filename), dtype=np.float64)

    if(pc.shape[0] != 4096*3):
        print("Error in pointcloud shape")
        return np.array([])

    pc = np.reshape(pc,(pc.shape[0]//3, 3))    
    return pc

def project_2_bev(pointcloud, edge, range = 0):
    if range > 0:
        p_range = np.linalg.norm(pointcloud[:,:2], axis=1)
        inside = (p_range < range)
        pointcloud = pointcloud[inside,:]

    p_x = np.round(np.clip(pointcloud[:,0]*2, -edge, edge-1)+edge).astype(int)
    p_y = np.round(np.clip(pointcloud[:,1]*2, -edge, edge-1)+edge).astype(int)

    image = np.zeros((edge*2, edge*2)).astype(np.uint8)
    image[p_y, p_x] = 255
    return image

#### Test
test_data = load_pc_file("test.bin")
x = test_data[...,0]
y = test_data[...,1]
z = test_data[...,2]

fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(x, y, z)
plt.show()
plt.pause(0.1)
plt.close()

sc_image = project_2_bev(test_data,40,120)
plt.imshow(sc_image)
plt.show()

test_data = test_data.astype(np.float32)
test_data = test_data[np.newaxis,:,...]

#### Settings
num_points = 4096
size = num_points
num_y = 120
num_x = 40
num_height = 20
max_height = 1  # max height
max_length = 1  # max_length in xy direction (same for x and y)
num_in_voxel = 1 # Num points in a voxel, no use for occupied information

#### Usage for voxelization
test_data = test_data.transpose()
test_data = test_data.flatten()
transer = voxelocc.GPUTransformer(test_data, size, max_length, max_height, num_x, num_y, num_height, num_in_voxel)
transer.transform()
point_t = transer.retreive()
point_t = point_t.reshape(-1,3)
point_t = point_t.reshape(num_height, num_x, num_y, 3)



#### Visualization
for i in range(num_height):
    plt.imshow(point_t[i,:,:,2])
    plt.show()
    radon_image = radon(point_t[i,:,:,2], theta=np.arange(0,180,1))
    plt.imshow(radon_image)
    plt.show()
    plt.pause(0.3)

utils/generate_bev_occ_cython/test.py
lus6-Jenny/RING
# from future.utils import iteritems
import os
from os.path import join as pjoin
from setuptools import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext
import numpy


def find_in_path(name, path):
    """Find a file in a search path"""

    # Adapted fom http://code.activestate.com/recipes/52224
    for dir in path.split(os.pathsep):
        binpath = pjoin(dir, name)
        if os.path.exists(binpath):
            return os.path.abspath(binpath)
    return None


def locate_cuda():
    """Locate the CUDA environment on the system

    Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64'
    and values giving the absolute path to each directory.

    Starts by looking for the CUDAHOME env variable. If not found,
    everything is based on finding 'nvcc' in the PATH.
    """

    # First check if the CUDAHOME env variable is in use
    if 'CUDAHOME' in os.environ:
        home = os.environ['CUDAHOME']
        nvcc = pjoin(home, 'bin', 'nvcc')
    else:
        # Otherwise, search the PATH for NVCC
        nvcc = find_in_path('nvcc', os.environ['PATH'])
        if nvcc is None:
            raise EnvironmentError('The nvcc binary could not be '
                'located in your $PATH. Either add it to your path, '
                'or set $CUDAHOME')
        home = os.path.dirname(os.path.dirname(nvcc))

    cudaconfig = {'home': home, 'nvcc': nvcc,
                  'include': pjoin(home, 'include'),
                  'lib64': pjoin(home, 'lib64')}
    for k, v in iter(cudaconfig.items()):
        if not os.path.exists(v):
            raise EnvironmentError('The CUDA %s path could not be '
                                   'located in %s' % (k, v))

    return cudaconfig


def customize_compiler_for_nvcc(self):
    """Inject deep into distutils to customize how the dispatch
    to gcc/nvcc works.

    If you subclass UnixCCompiler, it's not trivial to get your subclass
    injected in, and still have the right customizations (i.e.
    distutils.sysconfig.customize_compiler) run on it. So instead of going
    the OO route, I have this. Note, it's kindof like a wierd functional
    subclassing going on.
    """

    # Tell the compiler it can processes .cu
    self.src_extensions.append('.cu')

    # Save references to the default compiler_so and _comple methods
    default_compiler_so = self.compiler_so
    super = self._compile

    # Now redefine the _compile method. This gets executed for each
    # object but distutils doesn't have the ability to change compilers
    # based on source extension: we add it.
    def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts):
        if os.path.splitext(src)[1] == '.cu':
            # use the cuda for .cu files
            self.set_executable('compiler_so', CUDA['nvcc'])
            # use only a subset of the extra_postargs, which are 1-1
            # translated from the extra_compile_args in the Extension class
            postargs = extra_postargs['nvcc']
        else:
            postargs = extra_postargs['gcc']

        super(obj, src, ext, cc_args, postargs, pp_opts)
        # Reset the default compiler_so, which we might have changed for cuda
        self.compiler_so = default_compiler_so

    # Inject our redefined _compile method into the class
    self._compile = _compile


# Run the customize_compiler
class custom_build_ext(build_ext):
    def build_extensions(self):
        customize_compiler_for_nvcc(self.compiler)
        build_ext.build_extensions(self)


CUDA = locate_cuda()

# Obtain the numpy include directory. This logic works across numpy versions.
try:
    numpy_include = numpy.get_include()
except AttributeError:
    numpy_include = numpy.get_numpy_include()


ext = Extension('voxelfeat',
        sources = ['src/manager.cu', 'wrapper.pyx'],
        library_dirs = [CUDA['lib64']],
        libraries = ['cudart'],
        language = 'c++',
        runtime_library_dirs = [CUDA['lib64']],
        # This syntax is specific to this build system
        # we're only going to use certain compiler args with nvcc
        # and not with gcc the implementation of this trick is in
        # customize_compiler()
        extra_compile_args= {
            'gcc': [],
            'nvcc': [
                '--ptxas-options=-v', '-c',
                '--compiler-options', "'-fPIC'", "-std=c++11"
                ]
            },
            include_dirs = [numpy_include, CUDA['include'], 'src']
        )


setup(name = 'voxelfeat',
      author = 'Xuecheng Xu',
      version = '0.1',
      ext_modules = [ext],
      cmdclass = {'build_ext': custom_build_ext},
      # Since the package has c code, the egg cannot be zipped
      zip_safe = False)

utils/generate_bev_pointfeat_cython/setup.py
lus6-Jenny/RING
import voxelfeat
import numpy as np
import numpy.testing as npt
import time
import os
import numpy.testing as npt
import matplotlib.pyplot as plt
from skimage.transform import radon
from mpl_toolkits.mplot3d import Axes3D
from sklearn.neighbors import NearestNeighbors, KDTree
import math
from knn_cuda import KNN
import torch

def calculate_entropy_array(eigen):
    L_ = (eigen[:,0] - eigen[:,1]) / eigen[:,0]
    P_ = (eigen[:,1] - eigen[:,2]) / eigen[:,0]
    S_ = eigen[:,2] / eigen[:,0]
    Entropy = -L_*np.log(L_)-P_*np.log(P_)-S_*np.log(S_)
    return Entropy


def covariation_eigenvalue(pointcloud, neighborhood_index):
    ### calculate covariation and eigenvalue of 3D and 2D
    # prepare neighborhood
    neighborhoods = pointcloud[neighborhood_index]

    # 3D cov and eigen by matrix
    Ex = np.average(neighborhoods, axis=1)
    Ex = np.reshape(np.tile(Ex,[neighborhoods.shape[1]]), neighborhoods.shape)
    P = neighborhoods-Ex
    cov_ = np.matmul(P.transpose((0,2,1)),P)/(neighborhoods.shape[1]-1)
    eigen_, vec_ = np.linalg.eig(cov_)
    indices = np.argsort(eigen_)
    indices = indices[:,::-1]
    pcs_num_ = eigen_.shape[0]
    indx = np.reshape(np.arange(pcs_num_), [-1, 1])
    eig_ind = indices + indx*3
    vec_ind = np.reshape(eig_ind*3, [-1,1]) + np.full((pcs_num_*3,3), [0,1,2])
    vec_ind = np.reshape(vec_ind, [-1,3,3])
    eigen3d_ = np.take(eigen_, eig_ind)
    vectors_ = np.take(vec_, vec_ind)
    entropy_ = calculate_entropy_array(eigen3d_)

    # 2D cov and eigen
    cov2d_ = cov_[:,:2,:2]
    eigen2d, vec_2d = np.linalg.eig(cov2d_)
    indices = np.argsort(eigen2d)
    indices = indices[:, ::-1]
    pcs_num_ = eigen2d.shape[0]
    indx = np.reshape(np.arange(pcs_num_), [-1, 1])
    eig_ind = indices + indx * 2
    eigen2d_ = np.take(eigen2d, eig_ind)

    eigens_ = np.append(eigen3d_,eigen2d_,axis=1)

    return cov_, entropy_, eigens_.astype(np.float32), vectors_

def build_neighbors_NN(pointcloud, k):
	### using KNN NearestNeighbors cluster according k
    nbrs = NearestNeighbors(n_neighbors=k).fit(pointcloud)
    distances, indices = nbrs.kneighbors(pointcloud)
    print("knn index shape: ", indices.shape)
    covs, entropy, eigens_, vectors_ = covariation_eigenvalue(pointcloud, indices)

    return indices, covs, entropy, eigens_, vectors_

def calculate_features(pointcloud, nbrs_index, eigens_, vectors_):
    ### calculate handcraft feature with eigens and statistics data

    # features using eigens
    eig3d = eigens_[:3]
    eig2d = eigens_[3:5]

    # 3d
    # if eig3d.sum() == 0:
    #     print("eig3d.sum() 0 ", eig3d.sum())
    C_ = eig3d[2] / (eig3d.sum())
    O_ = np.power((eig3d.prod() / np.power(eig3d.sum(), 3)), 1.0 / 3)
    L_ = (eig3d[0] - eig3d[1]) / eig3d[0]
    E_ = -((eig3d / eig3d.sum()) * np.log(eig3d / eig3d.sum())).sum()
    P_ = (eig3d[1] - eig3d[2]) / eig3d[0]
    S_ = eig3d[2] / eig3d[0]
    A_ = (eig3d[0] - eig3d[2]) / eig3d[0]
    X_ = eig3d.sum()
    D_ = 3 * nbrs_index.shape[0] / (4 * math.pi * eig3d.prod())
    # 2d
    S_2 = eig2d.sum()
    L_2 = eig2d[1] / eig2d[0]
    # features using statistics data
    neighborhood = pointcloud[nbrs_index]
    nbr_dz = neighborhood[:, 2] - neighborhood[:, 2].min()
    dZ_ = nbr_dz.max()
    vZ_ = np.var(nbr_dz)
    V_ = vectors_[2][2]

    features = np.asarray([C_, O_, L_, E_, P_, S_, A_, X_, D_, S_2, L_2, dZ_, vZ_])#([C_,O_,L_,E_,D_,S_2,L_2,dZ_,vZ_,V_])
    return features, vZ_

def get_pointfeat(pointcloud):
    # prepare KNN cluster number k
    k = 30
    k_indices, k_covs, k_entropy, k_eigens_, k_vectors_  = build_neighbors_NN(pointcloud, k)
    
    ref = torch.from_numpy(pointcloud).unsqueeze(0).float().cuda()
    query = torch.from_numpy(pointcloud).unsqueeze(0).float().cuda()
    knn = KNN(k, transpose_mode=True)
    d, indx = knn(ref, query)
    print("indx.shape ",indx.shape)

    print("test torch knn ", indx==k_indices)
    points_feature = []
    for index in range(pointcloud.shape[0]):
        ### per point
        neighborhood = k_indices[index].astype(np.int32)
        eigens_ = k_eigens_[index].astype(np.float32)
        vectors_ = k_vectors_[index]

        # calculate point feature
        feature, vZ_ = calculate_features(pointcloud, neighborhood, eigens_, vectors_)
        points_feature.append(feature)

    points_feature = np.asarray(points_feature)

    return points_feature, indx.detach().cpu().numpy(), k_eigens_


def load_pc_file(filename):
    # returns Nx3 matrix
    pc = np.fromfile(os.path.join("./", filename), dtype=np.float64)

    if(pc.shape[0] != 4096*3):
        print("Error in pointcloud shape")
        return np.array([])

    pc = np.reshape(pc,(pc.shape[0]//3, 3))    
    return pc


def project_2_bev(pointcloud, edge, range = 0):
    if range > 0:
        p_range = np.linalg.norm(pointcloud[:,:2], axis=1)
        inside = (p_range < range)
        pointcloud = pointcloud[inside,:]

    p_x = np.round(np.clip(pointcloud[:,0]*2, -edge, edge-1)+edge).astype(int)
    p_y = np.round(np.clip(pointcloud[:,1]*2, -edge, edge-1)+edge).astype(int)

    image = np.zeros((edge*2, edge*2)).astype(np.uint8)
    image[p_y, p_x] = 255
    return image


#### Test
test_data = load_pc_file("test.bin")
x = test_data[...,0]
y = test_data[...,1]
z = test_data[...,2]

# fig = plt.figure()
# ax = Axes3D(fig)
# ax.scatter(x, y, z)
# plt.show()
# plt.pause(0.1)
# plt.close()

sc_image = project_2_bev(test_data,40,120)
# plt.imshow(sc_image)
# plt.show()

test_data = test_data.astype(np.float32)
# test_data = test_data[np.newaxis,:,...]

#### Settings
num_points = 4096
size = num_points
num_y = 120
num_x = 120
num_height = 1
max_height = 1  # max height
max_length = 1  # max_length in xy direction (same for x and y)
num_in_voxel = 1 # Num points in a voxel, no use for occupied information
test_data = test_data.astype(np.float32)

#### Usage for voxelization
# print("test_data shape: ", test_data.shape)
pointfeat, k_indices, k_eigens_ = get_pointfeat(test_data)   # num_points * 10
CPUFeat = np.concatenate((test_data, pointfeat), axis=1)      # num_points * 13 

# #### trans CPU feat
CPUFeat_flatten = CPUFeat.transpose().flatten().astype(np.float32)
CPUtranser = voxelocc.GPUTransformer(CPUFeat_flatten, size, max_length, max_height, num_x, num_y, num_height, 16)
CPUtranser.transform()
CPUBEV = CPUtranser.retreive()
CPUBEV = CPUBEV.reshape(-1,16)
CPUBEV = CPUBEV[...,3:]
CPUBEV = CPUBEV.reshape(num_height, num_x, num_y, 13)
CPUBEV = CPUBEV.take([0,1,3,10,11,12], axis=3)

### Visualization
# for i in range(13):
#     plt.imshow(CPUBEV[0,:,:,i])
#     plt.show()
#     plt.pause(0.3)

print("test_data shape: ", test_data.shape)
#### GPU
test_data_flatten = test_data.flatten()
start = time.time()
k_indices = k_indices.flatten().astype(np.int32)
k_eigens_ = k_eigens_.flatten()

featureExtractor = voxelocc.GPUFeatureExtractor(test_data_flatten, size, 13, 30, k_indices, k_eigens_)
GPUfeat = featureExtractor.get_features()
GPUfeat = GPUfeat.reshape(-1,13)
GPUdataFeat = np.concatenate((test_data, GPUfeat), axis=1)      # num_points * 13 
end = time.time()
print("featureExtractor: ", end - start)

#### trans GPU feat
GPUdataFeat_flatten = GPUdataFeat.transpose().flatten()
GPUtranser = voxelocc.GPUTransformer(GPUdataFeat_flatten, size, max_length, max_height, num_x, num_y, num_height, 16)
GPUtranser.transform()
GPUBEV = GPUtranser.retreive()
print("GPUBEV shape: ", GPUBEV.shape)

GPUBEV = GPUBEV.reshape(-1,16)
GPUBEV = GPUBEV[...,3:]
GPUBEV = GPUBEV.reshape(num_height, num_x, num_y, 13)

GPUBEV = GPUBEV.take([0,1,3,10,11,12], axis=3)

print(np.all((CPUBEV - GPUBEV)==0))
#### Visualization
for i in range(6):
    plt.imshow(CPUBEV[0,:,:,i])
    plt.show()
    print(np.max(CPUBEV[0,:,:,i]))
    plt.imshow(GPUBEV[0,:,:,i])
    plt.show()
    print(np.max(GPUBEV[0,:,:,i]))
    plt.pause(0.3)

utils/generate_bev_pointfeat_cython/test.py
lus6-Jenny/RING
# utils for point cloud processing

import copy
import torch
import random
import pygicp
import numpy as np
import open3d as o3d
import utils.config as cfg


# Preprocess and normalize the point cloud 
def load_pc_infer(pc):
    # returns Nx3 matrix
    pc = np.array(pc, dtype=np.float32)
    idx = (pc[:,0] > cfg.point_cloud["x_bound"][0]) * (pc[:,0] < cfg.point_cloud["x_bound"][1]) * (pc[:,1] > cfg.point_cloud["y_bound"][0]) * \
          (pc[:,1] < cfg.point_cloud["y_bound"][1]) * (pc[:,2] > cfg.point_cloud["z_bound"][0]) * (pc[:,2] < cfg.point_cloud["z_bound"][1])

    pc = pc[idx]
    pc_filtered = copy.deepcopy(pc)
    pc[:,0] = pc[:,0] / cfg.point_cloud["x_bound"][1]
    pc[:,1] = pc[:,1] / cfg.point_cloud["y_bound"][1]
    pc[:,2] = (pc[:,2] - cfg.point_cloud["z_bound"][0]) / (cfg.point_cloud["z_bound"][1] - cfg.point_cloud["z_bound"][0])
    return pc_filtered, pc   


# Randomly downsample the pointcloud to a certain number of points
def random_sampling(orig_points, num_points):
    assert orig_points.shape[0] > num_points
    points_down_idx = random.sample(range(orig_points.shape[0]), num_points)
    down_points = orig_points[points_down_idx, :]

    return down_points


# Apply random rotation along z axis
def random_rotation(xyz, angle_range=(-np.pi, np.pi)):
    angle = np.random.uniform(*angle_range)
    rotation_matrix = np.array([[np.cos(angle), -np.sin(angle), 0],
                                [np.sin(angle), np.cos(angle), 0],
                                [0, 0, 1]]).transpose()
    return np.dot(xyz, rotation_matrix)


# Apply 4x4 SE(3) transformation matrix on (N, 3) point cloud or 3x3 transformation on (N, 2) point cloud
def apply_transform(pc: torch.Tensor, m: torch.Tensor):
    assert pc.ndim == 2
    n_dim = pc.shape[1]
    assert n_dim == 2 or n_dim == 3
    assert m.shape == (n_dim + 1, n_dim + 1)
    # (m @ pc.t).t = pc @ m.t
    pc = pc @ m[:n_dim, :n_dim].transpose(1, 0) + m[:n_dim, -1]

    return pc


# Make up feature with open3d
def make_open3d_feature(data, dim, npts):
    feature = o3d.pipelines.registration.Feature()
    feature.resize(dim, npts)
    feature.data = data.cpu().numpy().astype('d').transpose()
    return feature


# Make up point cloud with open3d
def make_open3d_point_cloud(xyz, color=None):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(xyz)
    if color is not None:
        pcd.colors = o3d.utility.Vector3dVector(color)
    return pcd


def draw_pc(pc):
    '''
    pc: np.ndarray
    '''
    pc = copy.deepcopy(pc)
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(pc)    
    pcd1.paint_uniform_color([1, 0.706, 0])
    coord = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2, origin=[0,0,0])
    o3d.visualization.draw_geometries([pcd1, coord],
                                      zoom=0.4459,
                                      front=[0.9288, -0.2951, -0.2242],
                                      lookat=[1.6784, 2.0612, 1.4451],
                                      up=[-0.3402, -0.9189, -0.1996])


def draw_pc_pair(pc1, pc2):
    '''
    pc1 & pc2: np.ndarray
    '''    
    pc1 = copy.deepcopy(pc1)
    pc2 = copy.deepcopy(pc2)    
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(pc1)
    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(pc2)    
    pcd1.paint_uniform_color([1, 0.706, 0])
    pcd2.paint_uniform_color([0, 0.651, 0.929])
    coord = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2, origin=[0,0,0])
    o3d.visualization.draw_geometries([pcd1, pcd2, coord],
                                      zoom=0.4459,
                                      front=[0.9288, -0.2951, -0.2242],
                                      lookat=[1.6784, 2.0612, 1.4451],
                                      up=[-0.3402, -0.9189, -0.1996])   


def draw_registration_result(source, target, transformation):
    '''
    source & target: np.ndarray
    '''        
    source = copy.deepcopy(source)
    target = copy.deepcopy(target)    
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(source)
    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(target)    
    pcd1.paint_uniform_color([1, 0.706, 0])
    pcd2.paint_uniform_color([0, 0.651, 0.929])
    pcd1.transform(transformation)
    coord = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2, origin=[0,0,0])    
    o3d.visualization.draw_geometries([pcd1, pcd2, coord],
                                      zoom=0.4459,
                                      front=[0.9288, -0.2951, -0.2242],
                                      lookat=[1.6784, 2.0612, 1.4451],
                                      up=[-0.3402, -0.9189, -0.1996])                              


# fast_gicp (https://github.com/SMRT-AIST/fast_gicp)
def fast_gicp(source, target, max_correspondence_distance=1.0, init_pose=np.eye(4)):
    # downsample the point cloud before registration
    source = pygicp.downsample(source, 0.25)
    target = pygicp.downsample(target, 0.25)

    # pygicp.FastGICP has more or less the same interfaces as the C++ version
    gicp = pygicp.FastGICP()
    gicp.set_input_target(target)
    gicp.set_input_source(source)

    # optional arguments
    gicp.set_num_threads(4)
    gicp.set_max_correspondence_distance(max_correspondence_distance)

    # align the point cloud using the initial pose calculated by RING
    T_matrix = gicp.align(initial_guess=init_pose)

    # get the fitness score
    fitness = gicp.get_fitness_score(max_range=1.0)
    # get the transformation matrix
    T_matrix = gicp.get_final_transformation()

    return fitness, T_matrix


# open3d icp
def o3d_icp(source, target, transform: np.ndarray = None, point2plane: bool = False,
        inlier_dist_threshold: float = 1.0, max_iteration: int = 200):
    # transform: initial alignment transform
    if transform is not None:
        transform = transform.astype(float)

    voxel_size = 0.25
    pcd1 = o3d.geometry.PointCloud()
    pcd1.points = o3d.utility.Vector3dVector(source)
    pcd1 = pcd1.voxel_down_sample(voxel_size=voxel_size)

    pcd2 = o3d.geometry.PointCloud()
    pcd2.points = o3d.utility.Vector3dVector(target)
    pcd2 = pcd2.voxel_down_sample(voxel_size=voxel_size)

    if point2plane:
        pcd1.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=20))
        pcd2.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=20))
        transform_estimation = o3d.pipelines.registration.TransformationEstimationPointToPlane()
    else:
        transform_estimation = o3d.pipelines.registration.TransformationEstimationPointToPoint()

    if transform is not None:
        reg_p2p = o3d.pipelines.registration.registration_icp(pcd1, pcd2, inlier_dist_threshold, transform,
                                                              estimation_method=transform_estimation,
                                                              criteria=o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=max_iteration))
    else:
        reg_p2p = o3d.pipelines.registration.registration_icp(pcd1, pcd2, inlier_dist_threshold,
                                                              estimation_method=transform_estimation,
                                                              criteria=o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=max_iteration))

    return reg_p2p.fitness, reg_p2p.transformation, reg_p2p.inlier_rmse


class RandomOcclude:
    def __init__(self, ang):
        assert 0 <= ang <= 360, f'occluded angle must be in [0, 360] range, is {ang}'
        self.ang = ang

    def __call__(self, pc):
        # centroid = np.mean(pc, axis=0)
        # pc -= centroid
        ang = np.deg2rad(self.ang)
        start_ang = np.random.rand(1) * 2 * np.pi
        start_ang = np.where(start_ang < 0, start_ang + 2 * np.pi, start_ang)
        end_ang = start_ang + ang
        end_ang = np.where(end_ang < 0, end_ang + 2 * np.pi, end_ang)
        angles = np.arctan2(pc[:, 1], pc[:, 0])
        angles = np.where(angles < 0, angles + 2 * np.pi, angles)
        mask = np.logical_or(angles < start_ang, angles > end_ang)
        pc = pc[mask, :]
        # pc += centroid
        return pc
utils/point_clouds.py
lus6-Jenny/RING
# utils for pose transformation

import numpy as np
import torch


def trans2hom(R, t):
    # Get 4x4 transformation matrix from rotation matrix and translation vector
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = t
    return T


def m2ypr(m):
    # Get yaw, pitch, roll angles from 4x4 transformation matrix
    # Based on formulas in Section 2.5.1 in:
    # A tutorial on SE(3) transformation parameterizations and on-manifold optimization
    # https://ingmec.ual.es/~jlblanco/papers/jlblanco2010geometry3D_techrep.pdf
    assert m.shape == (4, 4)
    pitch = np.arctan2(-m[2][0], np.sqrt(m[0][0]**2 + m[1][0]**2))
    # We do not handle degenerate case, when pitch is 90 degrees a.k.a. gimball lock
    assert not np.isclose(np.abs(pitch), np.pi/2)
    yaw = np.arctan2(m[1][0], m[0][0])
    roll = np.arctan2(m[2][1], m[2][2])
    return yaw, pitch, roll


def m2xyz_ypr(m):
    # Get yaw, pitch, roll angles from 4x4 transformation matrix
    # Based on formulas in Section 2.5.1 in:
    # A tutorial on SE(3) transformation parameterizations and on-manifold optimization
    # https://ingmec.ual.es/~jlblanco/papers/jlblanco2010geometry3D_techrep.pdf
    assert m.shape == (4, 4)
    yaw, pitch, roll = m2ypr(m)
    return m[0, 3], m[1, 3], m[2, 3], yaw, pitch, roll


def ypr2m(yaw, pitch, roll):
    # Construct 4x4 transformation matrix with rotation part set to given yaw, pitch, roll. Translation is set to 0.
    # Based on formulas in Section 2.2.1
    m = np.array([[np.cos(yaw) * np.cos(pitch), np.cos(yaw) * np.sin(pitch) * np.sin(roll) - np.sin(yaw) * np.cos(roll),
                   np.cos(yaw) * np.sin(pitch) * np.cos(roll) + np.sin(yaw) * np.sin(roll), 0.],
                  [np.sin(yaw) * np.cos(pitch), np.sin(roll) * np.sin(pitch) * np.sin(roll) + np.cos(yaw) * np.cos(roll),
                   np.sin(yaw) * np.sin(pitch) * np.cos(roll) - np.cos(yaw) * np.sin(roll), 0.],
                  [-np.sin(pitch), np.cos(pitch) * np.sin(roll), np.cos(pitch) * np.cos(roll), 0.],
                  [0., 0., 0., 1.]], dtype=np.float32)
    return m


def xyz_ypr2m(x, y, z, yaw, pitch, roll):
    # Construct 4x4 transformation matrix with given translation and rotation part set to given yaw, pitch, roll.
    # Based on formulas in Section 2.2.1
    m = ypr2m(yaw, pitch, roll)
    m[0, 3] = x
    m[1, 3] = y
    m[2, 3] = z
    return m


def relative_pose(m1, m2):
    # !!! DO NOT USE THIS FUNCTION FOR MULRAN POSES !!!
    # SE(3) pose is 4x4 matrix, such that
    # Pw = [R | T] @ [P]
    #      [0 | 1]   [1]
    # where Pw are coordinates in the world reference frame and P are coordinates in the camera frame
    # m1: coords in camera/lidar1 reference frame -> world coordinate frame
    # m2: coords in camera/lidar2 coords -> world coordinate frame
    # returns: coords in camera/lidar1 reference frame -> coords in camera/lidar2 reference frame
    #          relative pose of the first camera with respect to the second camera
    return np.linalg.inv(m2) @ m1


def angle_clip(angle):
    # Convert angle to [-pi, pi]
    if angle > 0:
        if angle > np.pi:
            return angle - 2*np.pi
    else:
        if angle < -np.pi:
            return angle + 2*np.pi
    return angle


def cal_pose_error(T_estimated, T_gt):
    # Calculate RTE (Relative Translation Error) and RRE (Relative Rotation Error)
    r'''
    RTE = \lVert t - \bar{t} \rVert_2
    RRE = acos((trace(R^T \cdot \bar{R}) - 1) / 2)
    '''
    rte = np.linalg.norm(T_estimated[:3, 3] - T_gt[:3, 3])
    cos_rre = (np.trace(T_estimated[:3, :3].transpose(1, 0) @ T_gt[:3, :3]) - 1.) / 2.
    rre = np.arccos(np.clip(cos_rre, a_min=-1., a_max=1.)) * 180. / np.pi  

    return rte, rre
utils/poses.py
lus6-Jenny/RING
# utils for dataset processing

import os
import cv2
import copy
import errno
import torch
import pickle
import numpy as np
import open3d as o3d
import utils.config as cfg
import matplotlib.pyplot as plt
from scipy.spatial import distance_matrix
import torchvision.transforms.functional as fn
from sklearn.neighbors import NearestNeighbors, KDTree

# Coordinates of test region centers
# nclt_test_region_centers = np.array([[-50,-250],[-50,150],[150,-250],[150,150]])
nclt_test_region_centers = np.array([[50, -50]])
mulran_test_region_centers = np.array([[345090.0743, 4037591.323], [345090.483, 4044700.04],
                                [350552.0308, 4041000.71], [349252.0308, 4044800.71]])
kitti_test_region_centers = np.array([[50, -50]])

oxford_radar_test_region_centers = np.array([[5735500, 620000], [5735500, 620500], [5735000, 620500], [5735000, 620000]])

# Radius of the test region
nclt_test_region_radius = 220
mulran_test_region_radius = 500
kitti_test_region_radius = 100
oxford_radar_test_region_radius = 500

# Boundary between the train and test region
test_region_boundary = 50


# Make directory
def mkdir(dirname):
    assert dirname is not None
    if dirname == '' or os.path.isdir(dirname):
        return
    try:
        os.makedirs(dirname)
    except OSError as e:
        if e.errno != errno.EEXIST:
            raise e


# Load pickle file
def load_pickle(filename):
    with open(filename, 'rb') as handle:
        data = pickle.load(handle)
        print("Pickle Data Loaded.")
        return data


# Find the closest timestamp of the target point cloud
def find_nearest_ndx(ts, timestamps):
    ndx = np.searchsorted(timestamps, ts)
    if ndx == 0:
        return ndx
    elif ndx == len(timestamps):
        return ndx - 1
    else:
        assert timestamps[ndx-1] <= ts <= timestamps[ndx]
        if ts - timestamps[ndx-1] < timestamps[ndx] - ts:
            return ndx - 1
        else:
            return ndx


# Sample point clouds at equidistant intervals
def sample_at_intervals(northing, easting, prev_northing, prev_easting, sampling_gap):
    is_submap = False
    euclidean = np.abs(np.sqrt((prev_northing-northing)**2 + (prev_easting-easting)**2))
    if euclidean >= sampling_gap:
        is_submap = True
    return is_submap


# Check if the location is in the train set
def check_in_train_set(pos, dataset):
    # returns true if pos is in train split
    assert pos.ndim == 2
    assert pos.shape[1] == 2    
    if dataset == "nclt":
        test_region_centers = nclt_test_region_centers
        test_region_radius = nclt_test_region_radius
    elif dataset == "mulran":
        test_region_centers = mulran_test_region_centers
        test_region_radius = mulran_test_region_radius
    elif dataset == "kitti":
        test_region_centers = kitti_test_region_centers
        test_region_radius = kitti_test_region_radius
    elif dataset == "oxford_radar":
        test_region_centers = oxford_radar_test_region_centers
        test_region_radius = oxford_radar_test_region_radius
    dist = distance_matrix(pos, test_region_centers)
    mask = (dist > test_region_radius + test_region_boundary).all(axis=1)
    return mask


# Check if the location is in the test set
def check_in_test_set(pos, dataset):
    # returns true if position is in evaluation split
    assert pos.ndim == 2
    assert pos.shape[1] == 2
    if dataset == "nclt":
        test_region_centers = nclt_test_region_centers
        test_region_radius = nclt_test_region_radius
    elif dataset == "mulran":
        test_region_centers = mulran_test_region_centers
        test_region_radius = mulran_test_region_radius 
    elif dataset == "kitti":
        test_region_centers = kitti_test_region_centers
        test_region_radius = kitti_test_region_radius        
    elif dataset == "oxford_radar":
        test_region_centers = oxford_radar_test_region_centers
        test_region_radius = oxford_radar_test_region_radius           
    dist = distance_matrix(pos, test_region_centers)
    mask = (dist < test_region_radius).any(axis=1)
    return mask


# Check if it's a new place (class) in a trajectory of train dataset
def check_train_class(pose, prev_poses, threshold=cfg.train_sampling_gap):
    assert pose.ndim == 2
    assert pose.shape[1] == 2
    dist = distance_matrix(pose, prev_poses)
    is_new_class = (dist >= threshold).all(axis=1)
    # print("is_new_class_train", is_new_class)
    return is_new_class


# Check if it's a new place (class) in a trajectory of test dataset
def check_test_class(pose, prev_poses, threshold=cfg.test_sampling_gap):
    assert pose.ndim == 2
    assert pose.shape[1] == 2
    dist = distance_matrix(pose, prev_poses)
    is_new_class = (dist >= threshold).all(axis=1)
    # print("is_new_class_test", is_new_class)
    return is_new_class


# Calculate the distance between two poses
def calculate_dist(pose1, pose2):
    dist = np.sqrt((pose1[0] - pose2[0])**2 + (pose1[1] - pose2[1])**2)
    return dist


# Check if there is a loop (revisitness) in the map
def is_revisited(query_pose, map_poses, revisit_threshold):
    tree = KDTree(map_poses)
    # get the nearest neighbor
    dist, idx = tree.query(np.array([query_pose]), k=1)
    # dist, idx = tree.query(query_pose.reshape(1, -1) , k=1)
    if dist[0] < revisit_threshold:
        revisited = True
    else:
        revisited = False
    
    return revisited, dist, idx


# Convert robot id and index to key for multi-robot optimization
def robotid_to_key(robotid):
    char_a = 97
    keyBits = 64
    chrBits = 8
    indexBits = keyBits - chrBits
    outkey = char_a + robotid
    print("robotid: ", robotid, " outkey: ",outkey)
    return outkey << indexBits


def imshow(tensor, title=None):
    # print(tensor.type())
    image = tensor.cpu().clone()  # we clone the tensor to not do changes on it
    image = image.squeeze(0)  # remove the fake batch dimension
    plt.imshow(image, cmap='jet')
    # plt.colorbar()
    plt.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
    plt.savefig(title, bbox_inches='tight', pad_inches=0)
    plt.show()


# Plot trajectory of both map and query
def plot_trajectory(positions, split, save_path):
    positions = np.array(positions)
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.plot(positions[:, 1], positions[:, 0], 'b-', label=split)
    plt.xlabel('Easting (m)')
    plt.ylabel('Northing (m)')
    plt.title('Trajectory')
    plt.savefig(save_path)
    # plt.show()
    plt.close()
utils/tools.py
lus6-Jenny/RING
import os
import sys

version = "2.0.0"
base_url = "https://rosh-public.s3-eu-west-1.amazonaws.com/radon-v2"

green = "\u001b[32m"
red = "\u001b[31m"
blue = "\u001b[34m"
reset = "\u001b[0m"

ok = f"{green}OK{reset}"
error = f"{red}ERROR{reset}"
not_found = f"{red}Not found{reset}"

print(f"{blue}Checking requirements{reset}")

# OS
platform = sys.platform
print("Operating System:", platform, end=" ")
if platform == "linux":
    print(ok)
else:
    print(error)
    print("Precompiled packages are only available for Linux")
    sys.exit(1)

# Python
python_version = sys.version_info.major * 10 + sys.version_info.minor
python_version_str = f"{sys.version_info.major}.{sys.version_info.minor}"
print("Python version:", python_version_str, end=" ")
if python_version in [36, 37, 38]:
    print(ok)
else:
    print(error)
    print("Precompiled packages are only available for Python 3.6, 3.7 and 3.8")
    sys.exit(1)

# PyTorch
try:
    import torch

    torch_available = True
except:
    torch_available = False

print("PyTorch:", end=" ")
if torch_available:
    torch_version = torch.__version__[:3]
    print(torch_version, end=" ")
    if torch_version in ["1.8", "1.7", "1.6", "1.5"]:
        print(ok)
    else:
        print(error)
        print("Precompiled packages are build for PyTorch 1.5 to 1.8")
        print("Consider manually compiling torch-radon")
        sys.exit(1)
else:
    print(not_found)
    print("You need to have PyTorch installed")
    sys.exit(1)

# CUDA
cuda_version = torch.version.cuda
print("CUDA:", cuda_version, end=" ")

if cuda_version in ["10.1", "10.2", "11.0", "11.1"]:
    print(ok)
else:
    print(error)
    print("Precompiled packages are build for CUDA 10.1, 10.2 and 11.0")
    print("Consider manually compiling torch-radon")
    sys.exit(1)

python_version = f"cp{python_version}-cp{python_version}m" if python_version < 38 else f"cp{python_version}-cp{python_version}"

package_url = f"{base_url}/cuda-{cuda_version}/torch-{torch_version}/torch_radon-{version}-{python_version}-linux_x86_64.whl"
install_command = f"pip install --force-reinstall --no-cache-dir {package_url}"
print(f"{blue}Executing: {install_command}{reset}")
os.system(install_command)

utils/torch-radon/auto_install.py
lus6-Jenny/RING
import datetime
import sphinx_rtd_theme
import doctest
import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), "../..")))

import torch_radon

extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.doctest',
    'sphinx.ext.intersphinx',
    'sphinx.ext.mathjax',
    'sphinx.ext.napoleon',
    'sphinx.ext.viewcode',
    'sphinx.ext.githubpages',
]

source_suffix = '.rst'
master_doc = 'index'

author = 'Matteo Ronchetti'
project = 'torch_radon'
copyright = '{}, {}'.format(datetime.datetime.now().year, author)

version = torch_radon.__version__
release = torch_radon.__version__

html_theme = 'sphinx_rtd_theme'
html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]

doctest_default_flags = doctest.NORMALIZE_WHITESPACE
intersphinx_mapping = {'python': ('https://docs.python.org/', None)}

utils/torch-radon/docs/source/conf.py
lus6-Jenny/RING
import torch
import torch.nn as nn
import numpy as np
from torch_radon import Radon

batch_size = 8
n_angles = 64
image_size = 128
channels = 4

device = torch.device('cuda')
criterion = nn.L1Loss()

# Instantiate a model for the sinogram and one for the image
sino_model = nn.Conv2d(1, channels, 5, padding=2).to(device)
image_model = nn.Conv2d(channels, 1, 3, padding=1).to(device)

# create empty images
x = torch.FloatTensor(batch_size, 1, image_size, image_size).to(device)

# instantiate Radon transform
angles = np.linspace(0, np.pi, n_angles)
radon = Radon(image_size, angles)

# forward projection
sinogram = radon.forward(x)

# apply sino_model to sinograms
filtered_sinogram = sino_model(sinogram)

# backprojection
backprojected = radon.backprojection(filtered_sinogram)

# apply image_model to backprojected images
y = image_model(backprojected)

# backward works as usual
loss = criterion(y, x)
loss.backward()
utils/torch-radon/examples/end_to_end.py
lus6-Jenny/RING
import matplotlib.pyplot as plt
import numpy as np
import torch
from utils import show_images

from torch_radon import Radon

device = torch.device('cuda')

img = np.load("phantom.npy")
print(img.shape)
image_size = img.shape[0]
n_angles = image_size

# Instantiate Radon transform. clip_to_circle should be True when using filtered backprojection.
angles = np.linspace(0, np.pi, n_angles, endpoint=False)
radon = Radon(image_size, angles, clip_to_circle=True)

with torch.no_grad():
    x = torch.FloatTensor(img).to(device)

    sinogram = radon.forward(x)
    filtered_sinogram = radon.filter_sinogram(sinogram)
    fbp = radon.backprojection(filtered_sinogram)

print("FBP Error", torch.norm(x - fbp).item())


# Show results
titles = ["Original Image", "Sinogram", "Filtered Sinogram", "Filtered Backprojection"]
show_images([x, sinogram, filtered_sinogram, fbp], titles, keep_range=False)
plt.show()
utils/torch-radon/examples/fbp.py
lus6-Jenny/RING
import matplotlib.pyplot as plt
import numpy as np
import torch
import time

from torch_radon import Radon, RadonFanbeam
from torch_radon.solvers import cg

from torch_radon.shearlet import ShearletTransform


def circle_mask(size, radius):
    center = (size - 1) / 2
    c0, c1 = np.ogrid[0:size, 0:size]
    return ((c0 - center) ** 2 + (c1 - center) ** 2) <= radius ** 2


def shrink(a, b):
    return (torch.abs(a) - b).clamp_min(0) * torch.sign(a)


def main():
    n_angles = 100
    image_size = 512
    circle_radius = 100
    source_dist = 1.5 * image_size
    batch_size = 1
    n_scales = 5

    angles = (np.linspace(0., 100., n_angles, endpoint=False) - 50.0) / 180.0 * np.pi

    x = np.zeros((image_size, image_size), dtype=np.float32)
    x[circle_mask(image_size, circle_radius)] = 1.0

    radon = Radon(image_size, angles)  # RadonFanbeam(image_size, angles, source_dist)
    shearlet = ShearletTransform(image_size, image_size, [0.5] * n_scales)

    torch_x = torch.from_numpy(x).cuda()
    torch_x = torch_x.view(1, image_size, image_size).repeat(batch_size, 1, 1)
    sinogram = radon.forward(torch_x)

    bp = radon.backward(sinogram)
    sc = shearlet.forward(bp)

    p_0 = 0.02
    p_1 = 0.1
    w = 3 ** shearlet.scales / 400
    w = w.view(1, -1, 1, 1).cuda()

    u_2 = torch.zeros_like(bp)
    z_2 = torch.zeros_like(bp)
    u_1 = torch.zeros_like(sc)
    z_1 = torch.zeros_like(sc)
    f = torch.zeros_like(bp)

    relative_error = []
    start_time = time.time()
    for i in range(100):
        cg_y = p_0 * bp + p_1 * shearlet.backward(z_1 - u_1) + (z_2 - u_2)
        f = cg(lambda x: p_0 * radon.backward(radon.forward(x)) + (1 + p_1) * x, f.clone(), cg_y, max_iter=50)
        sh_f = shearlet.forward(f)

        z_1 = shrink(sh_f + u_1, p_0 / p_1 * w)
        z_2 = (f + u_2).clamp_min(0)
        u_1 = u_1 + sh_f - z_1
        u_2 = u_2 + f - z_2

        relative_error.append((torch.norm(torch_x[0] - f[0]) / torch.norm(torch_x[0])).item())

    runtime = time.time() - start_time
    print("Running time:", runtime)
    print("Running time per image:", runtime / batch_size)
    print("Relative error: ", 100 * relative_error[-1])


with torch.no_grad():
    main()

utils/torch-radon/examples/invisible.py
lus6-Jenny/RING
import matplotlib.pyplot as plt
import numpy as np
import torch

from torch_radon import Radon
from torch_radon.solvers import Landweber
from utils import show_images

batch_size = 1
n_angles = 512
image_size = 512

img = np.load("phantom.npy")
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# instantiate Radon transform
angles = np.linspace(0, np.pi, n_angles, endpoint=False)
radon = Radon(image_size, angles)

landweber = Landweber(radon)

# estimate step size
alpha = landweber.estimate_alpha(image_size, device)

with torch.no_grad():
    x = torch.FloatTensor(img).reshape(1, 1, image_size, image_size).to(device)
    sinogram = radon.forward(x)

    # use landweber iteration to reconstruct the image
    # values returned by 'callback' are stored inside 'progress'
    reconstruction, progress = landweber.run(torch.zeros(x.size(), device=device), sinogram, alpha, iterations=500,
                                             callback=lambda xx: torch.norm(xx - x).item())

plt.plot(progress)
plt.xlabel("Iteration")
plt.ylabel("Reconstruction error")
plt.title("Landweber Reconstruction Error")
print("Landweber Error", torch.norm(x - reconstruction).item() / torch.norm(x).item())

titles = ["Original Image", "Sinogram", "Reconstruction"]
show_images([x, sinogram, reconstruction], titles, keep_range=False)
plt.show()

utils/torch-radon/examples/landweber.py
lus6-Jenny/RING
import numpy as np
import matplotlib.pyplot as plt


def show_images(imgs, titles=None, keep_range=True, shape=None, figsize=(8, 8.5)):
    imgs = [x.cpu().numpy() for x in imgs]
    combined_data = np.array(imgs)

    if titles is None:
        titles = [str(i) for i in range(combined_data.shape[0])]

    # Get the min and max of all images
    if keep_range:
        _min, _max = np.amin(combined_data), np.amax(combined_data)
    else:
        _min, _max = None, None

    if shape is None:
        shape = (1, len(imgs))

    fig, axes = plt.subplots(*shape, figsize=figsize, sharex=True, sharey=True)
    ax = axes.ravel()
    for i, (img, title) in enumerate(zip(imgs, titles)):
        ax[i].imshow(img, cmap=plt.cm.Greys_r, vmin=_min, vmax=_max)
        ax[i].set_title(title)

utils/torch-radon/examples/utils.py
lus6-Jenny/RING
import matplotlib.pyplot as plt
import numpy as np
import torch
import time

from torch_radon import Radon
from torch_radon.solvers import Landweber
from utils import show_images

batch_size = 1
n_angles = 512
image_size = 512

img = np.load("phantom.npy")
device = torch.device('cuda')

# instantiate Radon transform
angles = np.linspace(0, np.pi, n_angles, endpoint=False)
radon = Radon(image_size, angles)

landweber = Landweber(radon)
# estimate Landweber step size
alpha = landweber.estimate_alpha(image_size, device)

with torch.no_grad():
    x = torch.FloatTensor(img).reshape(1, 1, image_size, image_size).to(device)
    sinogram = radon.forward(x)
    filtered_sinogram = radon.filter_sinogram(sinogram)
    backprojection = radon.backprojection(sinogram, extend=True)

    # reconstruct using Filtered Backprojection
    fbp = radon.backprojection(filtered_sinogram, extend=False) * np.pi / n_angles

    # reconstruct using Landweber iteration
    land = landweber.run(torch.zeros(x.size(), device=device), sinogram, alpha, iterations=500)

print("FBP Error", torch.norm(x - fbp).item())
print("Landweber Error", torch.norm(x - land).item())

titles = ["Original Image", "Sinogram", "Filtered Sinogram", "Backprojection", "Filtered Backprojection", "Landweber"]
show_images([x, sinogram, filtered_sinogram, backprojection, fbp, land], titles, keep_range=False, shape=(2, 3))

plt.show()

utils/torch-radon/examples/visual_sample.py
lus6-Jenny/RING
import sys
from glob import glob
import os
import shutil
import subprocess
import re

from ptx_annotation import annotate_ptx


def mapper(src, dst):
    paths = glob(src)

    src_pre = src.index("*")
    src_post = src_pre - len(src) + 1
    pre = dst[:dst.index("*")]
    post = dst[dst.index("*") + 1:]

    return [(path, pre + path[src_pre:src_post] + post) for path in paths]


def run(command):
    print(f"\u001b[34m{command}\u001b[0m")
    if os.system(command) != 0:
        print("\u001b[31mERROR IN COMPILATION\u001b[0m")
        exit(-1)


def run_compilation(files, f):
    for src, dst in files:
        if not os.path.exists(dst) or os.path.getmtime(src) > os.path.getmtime(dst):
            print(f"\u001b[32mCompiling {src}\u001b[0m")
            command = f(src, dst)
            run(command)
        else:
            print(f"\u001b[32mAlready compiled {src}\u001b[0m")


def get_cuda_version(cuda_home):
    print("")
    print(cuda_home)

    nvcc_out = subprocess.run([f"{cuda_home}/bin/nvcc", "--version"], stdout=subprocess.PIPE).stdout.decode('utf-8')
    m = re.search(r"V[0-9]+.[0-9]+", nvcc_out)
    str_version = m.group(0)[1:]

    return int(str_version.replace(".", ""))


def build(compute_capabilities=(60, 70, 75, 80, 86), debug=False, cuda_home="/usr/local/cuda", cxx="g++",
          keep_intermediate=False):

    cuda_version = get_cuda_version(cuda_home)
    nvcc = f"{cuda_home}/bin/nvcc"
    include_dirs = ["./include"]
    intermediate_dir = "intermediates"

    # compute capabilities >= 80 are only for cuda >= 11
    if cuda_version <= 110:
        compute_capabilities = [x for x in compute_capabilities if x < 80]

    cu_files = mapper("src/*.cu", "objs/cuda/*.o")
    cpp_files = mapper("src/*.cpp", "objs/*.o")
    cpp_files = [x for x in cpp_files if x[0] != "src/pytorch.cpp"]

    all_objects = [y for x, y in cu_files + cpp_files]

    opt_flags = ["-g"] if debug else ["-DNDEBUG", "-O3"]

    include_flags = [f"-I{x}" for x in include_dirs]
    # -Wl,-Bstatic -lm
    cxx_flags = ["-std=gnu++11 -fPIC -D_GLIBCXX_USE_CXX11_ABI=0 -Wall"] + include_flags + opt_flags
    nvcc_base_flags = ["-std=c++11", f"-ccbin={cxx}", "-Xcompiler", "-fPIC",
                       "-Xcompiler -D_GLIBCXX_USE_CXX11_ABI=0"] + include_flags + opt_flags + [
                           "--generate-line-info --compiler-options -Wall", '-Xcudafe "--diag_suppress=unrecognized_gcc_pragma"']
    nvcc_flags = nvcc_base_flags + [f"-gencode arch=compute_{x},code=sm_{x}" for x in compute_capabilities]

    if keep_intermediate:
        if not os.path.exists(intermediate_dir):
            os.mkdir(intermediate_dir)
        nvcc_flags.append(f"-keep --keep-dir {intermediate_dir}")

    cxx_flags = " ".join(cxx_flags)
    nvcc_flags = " ".join(nvcc_flags)

    # create output directory
    if not os.path.exists("objs/cuda"):
        os.makedirs("objs/cuda")

    # compile
    run_compilation(cu_files, lambda src, dst: f"{nvcc} {nvcc_flags} -c {src} -o {dst}")
    run_compilation(cpp_files, lambda src, dst: f"{cxx} {cxx_flags} -c {src} -o {dst}")

    run(f"ar rc objs/libradon.a {' '.join(all_objects)}")

    if keep_intermediate:
        for path in os.listdir(intermediate_dir):
            if path.endswith(".ptx"):
                annotate_ptx(os.path.join(intermediate_dir, path))
            else:
                os.remove(os.path.join(intermediate_dir, path))


def clean():
    print(f"\u001b[32mCleaning\u001b[0m")
    for f in ["objs", "build", "dist", "intermediates", "torch_radon.egg-info"]:
        if os.path.exists(f):
            print(f"Removing '{f}'")
            shutil.rmtree(f)


def detect_compute_capabilities():
    import torch
    ccs = []
    for i in range(torch.cuda.device_count()):
        compute_capability = torch.cuda.get_device_capability(i)
        ccs.append(compute_capability[0]*10 + compute_capability[1])

    return ccs


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "clean":
        clean()
    else:
        cuda_home = os.getenv("CUDA_HOME", "/usr/local/cuda")
        compute_capabilities = (60, 70, 75, 80, 86)
        debug = "--debug" in sys.argv
        cxx = os.getenv("CXX", "g++")
        keep_intermediate = "--keep-intermediate" in sys.argv

        if "--local" in sys.argv:
            print("Detecting compute capabilities of local GPUs")
            compute_capabilities = detect_compute_capabilities()
            print("Compiling for compute capability", " and ".join([str(x) for x in compute_capabilities]))

        build(compute_capabilities, debug, cuda_home, cxx, keep_intermediate)

utils/torch-radon/make.py
lus6-Jenny/RING
import re
from demangler import demangle
import struct

def hex_to_float(h):
    return struct.unpack('!f', bytes.fromhex(h))[0]

def remove_tabs(txt):
    for i in range(len(txt)):
        if txt[i] not in ["\t", " "]:
            return txt[i:], i


def readlines(path):
    with open(path, "r") as f:
        return f.readlines()    

def annotate_ptx(path):
    lines = readlines(path)

    # extract file annotations
    files = dict()
    for line in lines:
        line = line.strip(" \t\n")
        
        m = re.match(r'\.file\s*([0-9]+)\s*"([^"]+)".*', line)
        if m:
            files[m.group(1)] = readlines(m.group(2))

    kernels = []
    dst_lines = []

    for line in lines:
        line = line.strip(" \t\n")

        # add comments with demangled kernel names
        m = re.match(r'\.visible\s*\.entry\s*([^(]+)', line)
        if m:
            kernel = demangle(m.group(1))
            line = f"// {kernel}\n{line}"
            kernels.append(kernel)
        
        # make .loc annotations explicit adding the corresponding source file line 
        m = re.match(r'\.loc\s*([0-9]+)\s*([0-9]+)\s*([0-9]+)', line)
        if m:
            file_id = m.group(1)
            line_n = int(m.group(2))
            column = int(m.group(3))
            src_txt, n_tabs = remove_tabs(files[file_id][line_n-1])
            column -= n_tabs + 1
            line = f"\n// {src_txt}// {' '*column}^\n{line}"
            
        # annotate hex float writing them in decimal format
        m = re.match(r'.*0f([0-9a-fA-F]{8})', line)
        if m:
            hexv = m.group(1)
            line = f"// {hexv} = {hex_to_float(hexv)}f\n{line}"
        
        dst_lines.append(line)

    with open(path, "w") as f:
        f.write("\n".join(dst_lines))
utils/torch-radon/ptx_annotation.py
lus6-Jenny/RING
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os
from make import build

with open("README.md", "r") as fh:
    long_description = fh.read()

cuda_home = os.getenv("CUDA_HOME", "/usr/local/cuda")
print(f"Using CUDA_HOME={cuda_home}")
build(cuda_home=cuda_home)

setup(name='torch_radon',
      version="2.0.0",
      author="Matteo Ronchetti",
      author_email="mttronchetti@gmail.com",
      description="Radon transform in PyTorch",
      long_description=long_description,
      long_description_content_type="text/markdown",
      url="https://github.com/matteo-ronchetti/torch-radon",

      packages=['torch_radon'],
      package_dir={
          'torch_radon': './torch_radon',
      },
      ext_modules=[
          CUDAExtension('torch_radon_cuda', [os.path.abspath('src/pytorch.cpp')],
                        include_dirs=[os.path.abspath('include')],
                        library_dirs=[os.path.abspath("objs")],
                        libraries=["m", "c", "gcc", "stdc++", "cufft", "radon"],
                        # extra_compile_args=["-static", "-static-libgcc", "-static-libstdc++"],
                        # strip debug symbols
                        extra_link_args=["-Wl,--strip-all"] #, "-static-libgcc", "-static-libstdc++"]
                        )
      ],
      cmdclass={'build_ext': BuildExtension},
      zip_safe=False,
      classifiers=[
          "Programming Language :: Python :: 3",
          "Operating System :: POSIX :: Linux",
          "License :: OSI Approved :: GNU General Public License v3 (GPLv3)"
      ],
      install_requires=[
          "scipy",
          "alpha-transform"
      ],
      )


"""
g++ -pthread -B /home/matteo/miniconda3/compiler_compat -L/home/matteo/miniconda3/lib -Wl,-rpath=/home/matteo/miniconda3/lib -Wl,--no-as-needed -Wl,--sysroot=/ /home/matteo/projects/torch-radon/build/temp.linux-x86_64-3.7/home/matteo/projects/torch-radon/src/pytorch.o -L/home/matteo/projects/torch-radon/objs -L/home/matteo/miniconda3/lib/python3.7/site-packages/torch/lib -L/usr/local/cuda-11/lib64 -lradon -lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda_cu -ltorch_cuda_cpp -o build/lib.linux-x86_64-3.7/torch_radon_cuda.cpython-37m-x86_64-linux-gnu.so -Wl,--strip-all -static-libgcc -static-libstdc++ -Wl,-static
g++ -shared /home/matteo/projects/torch-radon/build/temp.linux-x86_64-3.7/home/matteo/projects/torch-radon/src/pytorch.o -o build/lib.linux-x86_64-3.7/torch_radon_cuda.cpython-37m-x86_64-linux-gnu.so
g++ -pthread -B /home/matteo/miniconda3/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/matteo/projects/torch-radon/include -I/home/matteo/miniconda3/lib/python3.7/site-packages/torch/include -I/home/matteo/miniconda3/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I/home/matteo/miniconda3/lib/python3.7/site-packages/torch/include/TH -I/home/matteo/miniconda3/lib/python3.7/site-packages/torch/include/THC -I/usr/local/cuda-11/include -I/home/matteo/miniconda3/include/python3.7m -c -c /home/matteo/projects/torch-radon/src/pytorch.cpp -o /home/matteo/projects/torch-radon/build/temp.linux-x86_64-3.7/home/matteo/projects/torch-radon/src/pytorch.o -static -static-libgcc -static-libstdc++ -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=torch_radon_cuda -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14

"""
utils/torch-radon/setup.py
lus6-Jenny/RING
import torch
import numpy as np
from alpha_transform import AlphaShearletTransform
from torch_radon.shearlet import ShearletTransform
from nose.tools import assert_less
from .utils import relative_error


def test_shearlet():
    img_size = 128
    device = torch.device("cuda")
    scales = [0.5] * 2

    x = np.random.uniform(0, 1, (img_size, img_size))

    shearlet = ShearletTransform(img_size, img_size, scales)
    alpha_shearlet = AlphaShearletTransform(img_size, img_size, scales, real=True, parseval=True)
    print(alpha_shearlet.spectrograms[0].dtype)

    coeff_ = alpha_shearlet.transform(x, do_norm=False)
    rec_ = alpha_shearlet.adjoint_transform(coeff_, do_norm=False)


    with torch.no_grad():
        # check with double precision
        dx = torch.DoubleTensor(x).to(device)
        coeff = shearlet.forward(dx)
        rec = shearlet.backward(coeff)

        x_err_d = (torch.norm(rec - dx) / torch.norm(dx)).item()
        coef_err_d = relative_error(coeff_, coeff.cpu().numpy())
        rec_err_d = relative_error(rec_, rec.cpu().numpy())
        print(x_err_d, coef_err_d, rec_err_d)

        # check with single precision
        dx = torch.FloatTensor(x).to(device)
        coeff = shearlet.forward(dx)
        rec = shearlet.backward(coeff)

        x_err = (torch.norm(rec - dx) / torch.norm(dx)).item()
        coef_err = relative_error(coeff_, coeff.cpu().numpy())
        rec_err = relative_error(rec_, rec.cpu().numpy())
        print(x_err, coef_err, rec_err)

        assert_less(x_err_d, 1e-15)
        assert_less(coef_err_d, 1e-15)
        assert_less(rec_err_d, 1e-15)

        assert_less(x_err, 1e-6)
        assert_less(coef_err, 1e-6)
        assert_less(rec_err, 1e-6)

    # print(np.allclose(coeff_, coeff.cpu().numpy()))
    # print(np.allclose(rec_, rec.cpu().numpy()))
    #
    # print("Alpha forward", timeit.timeit(lambda: alpha_shearlet.transform(im, do_norm=False), number=10))
    # print("Alpha backward", timeit.timeit(lambda: alpha_shearlet.adjoint_transform(coeff_, do_norm=False), number=10))
    # print("My forward", timeit.timeit(lambda: shearlet.forward(x), number=30))
    # print("My backward", timeit.timeit(lambda: shearlet.backward(coeff), number=30))

utils/torch-radon/tests/_test_shearlet.py
lus6-Jenny/RING
import astra
import numpy as np


class AstraWrapper:
    def __init__(self, angles):
        self.angles = angles

        self.projectors = []
        self.algorithms = []
        self.data2d = []
        self.data3d = []

    def forward(self, x, spacing=1.0, det_count=-1):
        if det_count < 0:
            det_count = x.shape[1]

        vol_geom = astra.create_vol_geom(x.shape[1], x.shape[2], x.shape[0])
        phantom_id = astra.data3d.create('-vol', vol_geom, data=x)
        proj_geom = astra.create_proj_geom('parallel3d', spacing, 1.0, x.shape[0], det_count, -self.angles)

        proj_id, y = astra.creators.create_sino3d_gpu(phantom_id, proj_geom, vol_geom)

        self.projectors.append(proj_id)
        self.data3d.append(phantom_id)

        return proj_id, y
    # def forward(self, x):
    #     vol_geom = astra.create_vol_geom(x.shape[1], x.shape[2])
    #     proj_geom = astra.create_proj_geom('parallel', 1.0, x.shape[1], self.angles)
    #     proj_id = astra.create_projector('cuda', proj_geom, vol_geom)
    #
    #     self.projectors.append(proj_id)
    #
    #     ys = []
    #     for i in range(x.shape[0]):
    #         _, y = astra.create_sino(x[i], proj_id)
    #         ys.append(y.reshape(1, y.shape[0], y.shape[1]))
    #
    #     y = np.vstack(ys)
    #     return y

    def backproject(self, proj_id, s, bs):
        vol_geom = astra.create_vol_geom(s, s, bs)
        rec_id = astra.data3d.create('-vol', vol_geom)

        # Set up the parameters for a reconstruction algorithm using the GPU
        cfg = astra.astra_dict('BP3D_CUDA')
        cfg['ReconstructionDataId'] = rec_id
        cfg['ProjectionDataId'] = proj_id

        # Create the algorithm object from the configuration structure
        alg_id = astra.algorithm.create(cfg)
        astra.algorithm.run(alg_id, 1)

        self.algorithms.append(alg_id)
        self.data3d.append(rec_id)

        return astra.data3d.get(rec_id)

    def forward_single(self, x):
        vol_geom = astra.create_vol_geom(x.shape[0], x.shape[1])
        proj_geom = astra.create_proj_geom('parallel', 1.0, x.shape[0], self.angles)
        proj_id = astra.create_projector('cuda', proj_geom, vol_geom)

        self.projectors.append(proj_id)

        return astra.create_sino(x, proj_id)

    def fbp(self, x):
        s = x.shape[0]
        proj_id, _ = self.forward_single(x)
        vol_geom = astra.create_vol_geom(s, s)
        rec_id = astra.data2d.create('-vol', vol_geom)

        # create configuration
        cfg = astra.astra_dict('FBP_CUDA')
        cfg['ReconstructionDataId'] = rec_id
        cfg['ProjectionDataId'] = proj_id
        cfg['option'] = {'FilterType': 'Ram-Lak'}

        alg_id = astra.algorithm.create(cfg)
        astra.algorithm.run(alg_id)

        self.projectors.append(proj_id)
        self.algorithms.append(alg_id)
        self.data2d.append(rec_id)

        return astra.data2d.get(rec_id)

    def clean(self):
        # clean all astra stuff
        for pid in self.projectors:
            astra.projector.delete(pid)

        for pid in self.algorithms:
            astra.algorithm.delete(pid)

        for pid in self.data2d:
            astra.data2d.delete(pid)

        for pid in self.data3d:
            astra.data3d.delete(pid)

    def __del__(self):
        self.clean()
utils/torch-radon/tests/astra_wrapper.py
lus6-Jenny/RING
from .utils import generate_random_images, relative_error
import astra
from nose.tools import assert_less, assert_equal
import torch
import numpy as np
import torch_radon as tr
from parameterized import parameterized
import matplotlib.pyplot as plt

device = torch.device('cuda')

full_angles = np.linspace(0, 2*np.pi, 128).astype(np.float32)
many_angles = np.linspace(0, np.pi, 90).astype(np.float32)

params = []
for batch_size in [1, 8]:
    for volume_size in [64, 81]:
        for angles in [full_angles, many_angles]:
            for spacing in [1.0, 0.5, 1.3, 2.0]:
                for distances in [(1.5, 1.5), (2.0, 2.0), (1.2, 3.0)]:
                    for det_count in [1.0, 1.5]:
                        params.append((device, batch_size, volume_size, angles, spacing, distances, det_count))

half_params = [x for x in params if x[1] % 4 == 0]


@parameterized(params)
def test_fanflat_error(device, batch_size, volume_size, angles, det_spacing, distances, det_count):
    # generate random images
    det_count = int(det_count * volume_size)
    x = np.random.uniform(0.0, 1.0, (volume_size, volume_size, volume_size)).astype(np.float32)

    s_dist, d_dist = distances
    s_dist *= volume_size
    d_dist *= volume_size

    # astra
    vol_geom = astra.create_vol_geom(x.shape[1], x.shape[2], x.shape[0])
    proj_geom = astra.create_proj_geom('cone', det_spacing, det_spacing, det_count, det_count, angles, s_dist, d_dist)

    proj_id, astra_y = astra.create_sino3d_gpu(x, proj_geom, vol_geom)

    rec_id = astra.data3d.create('-vol', vol_geom)

    cfg = astra.astra_dict('BP3D_CUDA')
    cfg['ReconstructionDataId'] = rec_id
    cfg['ProjectionDataId'] = proj_id
    alg_id = astra.algorithm.create(cfg)
    astra.algorithm.run(alg_id, 1)

    astra_y = astra_y.transpose(1, 0, 2)
    astra_bp = astra.data3d.get(rec_id)

    # TODO clean astra structures

    # our implementation
    radon = tr.ConeBeam(det_count, angles, s_dist, d_dist, det_spacing_u=det_spacing, volume=volume_size)
    x = torch.FloatTensor(x).view(1, x.shape[0], x.shape[1], x.shape[2]).repeat(batch_size, 1, 1, 1).to(device)

    our_fp = radon.forward(x)
    our_bp = radon.backprojection(our_fp)

    our_fp = our_fp.cpu().numpy()
    batch_error = max([relative_error(our_fp[0], our_fp[i]) for i in range(1, batch_size)] + [0])
    forward_error = relative_error(astra_y, our_fp[0])

    our_bp = our_bp.cpu().numpy()
    batch_error_back = max([relative_error(our_bp[0], our_bp[i]) for i in range(1, batch_size)] + [0])
    back_error = relative_error(astra_bp, our_bp[0])

    # if forward_error > 2e-2:
    #     fig, ax = plt.subplots(3, 3)
    #     ax = ax.ravel()
    #     ax[0].imshow(astra_y[0])
    #     ax[1].imshow(our_fp[0, 0])
    #     ax[2].imshow(np.abs(our_fp[0, 0] - astra_y[0]))
    #     ax[3].imshow(astra_y[len(angles)//2])
    #     ax[4].imshow(our_fp[0, len(angles)//2])
    #     ax[5].imshow(np.abs(our_fp[0, len(angles)//2] - astra_y[len(angles)//2]))
    #     ax[6].imshow(astra_y[-1])
    #     ax[7].imshow(our_fp[0, -1])
    #     ax[8].imshow(np.abs(our_fp[0, -1] - astra_y[-1]))
    #     plt.show()

    print(f"batch: {batch_size}, size: {volume_size}, angles: {len(angles)}, spacing: {det_spacing}, distances: {distances}, det_count:{det_count}, forward: {forward_error}, back: {back_error}")
    
    # TODO better checks
    assert_less(batch_error, 1e-6)
    assert_less(forward_error, 2e-2)
    assert_less(batch_error_back, 1e-6)
    assert_less(back_error, 3e-3)


@parameterized(half_params)
def test_half(device, batch_size, volume_size, angles, det_spacing, distances, det_count):
    # generate random images
    det_count = int(det_count * volume_size)
    x = np.random.uniform(0.0, 1.0, (batch_size, volume_size, volume_size, volume_size)).astype(np.float32)

    s_dist, d_dist = distances
    s_dist *= volume_size
    d_dist *= volume_size

    radon = tr.ConeBeam(det_count, angles, s_dist, d_dist, det_spacing_u=det_spacing, volume=volume_size)
    x = torch.FloatTensor(x).to(device)

    single_fp = radon.forward(x) / len(angles)
    single_bp = radon.backprojection(single_fp)

    half_fp = radon.forward(x.half()) / len(angles)
    half_bp = radon.backprojection(half_fp)

    forward_error = relative_error(single_fp.cpu().numpy(), half_fp.float().cpu().numpy())
    back_error = relative_error(single_bp.cpu().numpy(), half_bp.float().cpu().numpy())

    print(f"batch: {batch_size}, size: {volume_size}, angles: {len(angles)}, spacing: {det_spacing}, distances: {distances}, det_count:{det_count}, forward: {forward_error}, back: {back_error}")
    
    # TODO better checks
    assert_less(forward_error, 3e-3)
    assert_less(back_error, 3e-3)

utils/torch-radon/tests/test_coneflat.py
lus6-Jenny/RING
import numpy as np
import torch
from nose.tools import assert_equal
from parameterized import parameterized
import random

from .utils import random_symbolic_function, symbolic_discretize, symbolic_forward, TestHelper

from torch_radon.volumes import Volume2D
import torch_radon as tr

random.seed(42)
device = torch.device('cuda')
test_helper = TestHelper("fanbeam")

# (batch_size, angles, volume, spacing, det_count, src_dist, det_dist)
params = []

# check different batch sizes
for batch_size in [1, 3, 17, 32]:
    params.append((batch_size, (0, 2*np.pi, 128), 128, 2.0, 128, 128, 128))

# check few and many angles which are not multiples of 16
for angles in [(0, 2*np.pi, 19), (0, 2*np.pi, 803)]:
    params.append((4, angles, 128, 2.0, 128, 128, 128))

# change volume size
for height, width in [(128, 256), (256, 128), (75, 149), (81, 81)]:
    s = max(height, width)
    params.append((4, (0, 2*np.pi, 64), Volume2D(height, width), 2.0, s, s, s))

# change volume scale and center
for center in [(0, 0), (17, -25), (53, 49)]:
    for voxel_size in [(1, 1), (0.75, 0.75), (1.5, 1.5), (0.7, 1.3), (1.3, 0.7)]:
        det_count = int(179 * max(voxel_size[0], 1) * max(voxel_size[1], 1) * np.sqrt(2))
        params.append((4, (0, 2*np.pi, 128), Volume2D(179, 123, center, voxel_size), 2.0, det_count, det_count, det_count))

for spacing in [1.0, 0.5, 1.3, 2.0]:
    for det_count in [79, 128, 243]:
        for src_dist, det_dist in [(128, 128), (64, 128), (128, 64), (503, 503)]:
            params.append((4, (0, 2*np.pi, 128), 128, spacing, det_count, src_dist, det_dist))


@parameterized(params)
def test_error(batch_size, angles, volume, spacing, det_count, src_dist, det_dist):
    radon = tr.FanBeam(det_count, angles, src_dist, det_dist, spacing, volume)

    f = random_symbolic_function(radon.volume.height, radon.volume.width)
    x = symbolic_discretize(f, radon.volume.height, radon.volume.width)

    f.scale(*radon.volume.voxel_size)
    f.move(*radon.volume.center)

    tx = torch.FloatTensor(x).unsqueeze(0).repeat(batch_size, 1, 1).to(device)

    y = symbolic_forward(f, radon.angles.cpu(), radon.projection.cfg).cpu().numpy()
    ty = radon.forward(tx)
    assert_equal(ty.size(0), batch_size)

    max_error = 2e-3 * (512 / y.shape[0]) * (512 / y.shape[1])

    description = f"Angles: {angles}\nVolume: {volume}\nSpacing: {spacing}, Count: {det_count}, Precision: float"
    test_helper.compare_images(y, ty, max_error, description)

    back_max_error = 1e-3
    test_helper.backward_check(tx, ty, radon, description, back_max_error)

    if batch_size % 4 == 0:
        ty = radon.forward(tx.half())

        description = f"Angles: {angles}\nVolume: {volume}\nSpacing: {spacing}, Count: {det_count}, Precision: half"
        test_helper.compare_images(y, ty, max_error, description)


# full_angles = np.linspace(0, 2 * np.pi, 180).astype(np.float32)
# limited_angles = np.linspace(0.2 * np.pi, 0.5 * np.pi, 50).astype(np.float32)
# sparse_angles = np.linspace(0, 2 * np.pi, 60).astype(np.float32)
# many_angles = np.linspace(0, 2 * np.pi, 800).astype(np.float32)

# params = []
# for batch_size in [1, 8]:
#     for image_size in [128, 151]:
#         for angles in [full_angles, limited_angles, sparse_angles, many_angles]:
#             for spacing in [1.0, 0.5, 1.3, 2.0]:
#                 for distances in [(1.2, 1.2), (2.0, 2.0), (1.2, 3.0)]:
#                     for det_count in [1.0, 1.5]:
#                         params.append((device, batch_size, image_size, angles, spacing, distances, det_count))

# half_params = [x for x in params if x[1] % 4 == 0]


# @parameterized(params)
# def test_fanbeam_error(device, batch_size, image_size, angles, spacing, distances, det_count):
#     # generate random images
#     # generate random images
#     det_count = int(det_count * image_size)
#     x = generate_random_images(1, image_size)[0]

#     s_dist, d_dist = distances
#     s_dist *= image_size
#     d_dist *= image_size

#     # astra
#     vol_geom = astra.create_vol_geom(x.shape[0], x.shape[1])
#     proj_geom = astra.create_proj_geom('fanflat', spacing, det_count, angles, s_dist, d_dist)
#     proj_id = astra.create_projector('cuda', proj_geom, vol_geom)

#     id, astra_y = astra.create_sino(x, proj_id)
#     _, astra_bp = astra.create_backprojection(astra_y, proj_id)

#     # TODO clean astra structures

#     # our implementation
#     radon = tr.FanBeam(det_count=det_count, det_spacing=spacing, angles=angles,
#                        src_dist=s_dist, det_dist=d_dist, volume=image_size)
#     x = torch.FloatTensor(x).to(device).view(1, x.shape[0], x.shape[1])
#     # repeat data to fill batch size
#     x = torch.cat([x] * batch_size, dim=0)

#     our_fp = radon.forward(x)
#     our_bp = radon.backprojection(our_fp)

#     forward_error = relative_error(astra_y, our_fp[0].cpu().numpy())
#     back_error = relative_error(astra_bp, our_bp[0].cpu().numpy())

#     # if back_error > 5e-3:
#     #     plt.imshow(astra_bp)
#     #     plt.figure()
#     #     plt.imshow(our_bp[0].cpu().numpy())
#     #     plt.show()
#     print(np.max(our_fp.cpu().numpy()), np.max(our_bp.cpu().numpy()))

#     print(
#         f"batch: {batch_size}, size: {image_size}, angles: {len(angles)}, spacing: {spacing}, distances: {distances}, forward: {forward_error}, back: {back_error}")
#     # TODO better checks
#     assert_less(forward_error, 1e-2)
#     assert_less(back_error, 5e-3)


# @parameterized(half_params)
# def test_half(device, batch_size, image_size, angles, spacing, distances, det_count):
#     # generate random images
#     det_count = int(det_count * image_size)
#     x = generate_random_images(batch_size, image_size)

#     s_dist, d_dist = distances
#     s_dist *= image_size
#     d_dist *= image_size

#     # our implementation
#     radon = tr.FanBeam(det_count=det_count, det_spacing=spacing, angles=angles,
#                        src_dist=s_dist, det_dist=d_dist, volume=image_size)
#     x = torch.FloatTensor(x).to(device)

#     # divide by len(angles) to avoid half-precision overflow
#     sinogram = radon.forward(x) / len(angles)
#     single_precision = radon.backprojection(sinogram)

#     h_sino = radon.forward(x.half()) / len(angles)
#     half_precision = radon.backprojection(h_sino)
#     print(torch.min(half_precision).item(), torch.max(half_precision).item())

#     forward_error = relative_error(sinogram.cpu().numpy(), h_sino.cpu().numpy())
#     back_error = relative_error(single_precision.cpu().numpy(), half_precision.cpu().numpy())

#     print(
#         f"batch: {batch_size}, size: {image_size}, angles: {len(angles)}, spacing: {spacing}, forward: {forward_error}, back: {back_error}")

#     assert_less(forward_error, 1e-3)
#     assert_less(back_error, 1e-3)

utils/torch-radon/tests/test_fanbeam.py
lus6-Jenny/RING
import numpy as np
import torch
from nose.tools import assert_equal

from torch_radon import Radon


# def test_noise():
#     device = torch.device('cuda')
#
#     x = torch.FloatTensor(3, 5, 64, 64).to(device)
#     lookup_table = torch.FloatTensor(128, 64).to(device)
#     x.requires_grad = True
#     angles = torch.FloatTensor(np.linspace(0, 2 * np.pi, 10).astype(np.float32))
#
#     radon = Radon(angles, 64)
#
#     sinogram = radon.forward(x)
#     assert_equal(sinogram.size(), (3, 5, 10, 64))
#
#     readings = radon.emulate_readings(sinogram, 5, 10.0)
#     assert_equal(readings.size(), (3, 5, 10, 64))
#     assert_equal(readings.dtype, torch.int32)
#
#     y = radon.readings_lookup(readings, lookup_table)
#     assert_equal(y.size(), (3, 5, 10, 64))
#     assert_equal(y.dtype, torch.float32)

utils/torch-radon/tests/test_noise.py
lus6-Jenny/RING
import numpy as np
import torch
from nose.tools import assert_equal
from parameterized import parameterized
import random

from .utils import random_symbolic_function, symbolic_discretize, symbolic_forward, TestHelper

from torch_radon.volumes import Volume2D
import torch_radon as tr

random.seed(42)
device = torch.device('cuda')
test_helper = TestHelper("parallel_beam")

# (batch_size, angles, volume, spacing, det_count)
params = []

# check different batch sizes
for batch_size in [1, 3, 17, 32]:
    params.append((batch_size, (0, np.pi, 128), 128, 1.0, 128))

# check few and many angles which are not multiples of 16
for angles in [(0, np.pi, 19), (0, np.pi, 803)]:
    params.append((4, angles, 128, 1.0, 128))

# change volume size
for height, width in [(128, 256), (256, 128), (75, 149), (81, 81)]:
    params.append((4, (0, np.pi, 64), Volume2D(height, width), 1.0, max(height, width)))

# change volume scale and center
for center in [(0, 0), (17, -25), (53, 49)]:
    for voxel_size in [(1, 1), (0.75, 0.75), (1.5, 1.5), (0.7, 1.3), (1.3, 0.7)]:
        det_count = int(179 * max(voxel_size[0], 1) * max(voxel_size[1], 1) * np.sqrt(2))
        params.append((4, (0, np.pi, 128), Volume2D(179, 123, center, voxel_size), 1.0, det_count))

for spacing in [1.0, 0.5, 1.3, 2.0]:
    for det_count in [79, 128, 243]:
        params.append((4, (0, np.pi, 128), 128, spacing, det_count))


@parameterized(params)
def test_error(batch_size, angles, volume, spacing, det_count):
    radon = tr.ParallelBeam(det_count, angles, spacing, volume)

    f = random_symbolic_function(radon.volume.height, radon.volume.width)
    x = symbolic_discretize(f, radon.volume.height, radon.volume.width)

    f.scale(*radon.volume.voxel_size)
    f.move(*radon.volume.center)

    tx = torch.FloatTensor(x).unsqueeze(0).repeat(batch_size, 1, 1).to(device)

    y = symbolic_forward(f, radon.angles.cpu(), radon.projection.cfg).cpu().numpy()
    ty = radon.forward(tx)
    assert_equal(ty.size(0), batch_size)

    max_error = 2e-3 * (512 / y.shape[0]) * (512 / y.shape[1])

    description = f"Angles: {angles}\nVolume: {volume}\nSpacing: {spacing}, Count: {det_count}, Precision: float"
    test_helper.compare_images(y, ty, max_error, description)

    back_max_error = 1e-3
    test_helper.backward_check(tx, ty, radon, description, back_max_error)

    if batch_size % 4 == 0:
        ty = radon.forward(tx.half())

        description = f"Angles: {angles}\nVolume: {volume}\nSpacing: {spacing}, Count: {det_count}, Precision: half"
        test_helper.compare_images(y, ty, max_error, description)

utils/torch-radon/tests/test_parallel_beam.py
lus6-Jenny/RING
from nose.tools import assert_equal
from torch_radon import Volume2D, Volume3D


def test_volume_2d():
    volume = Volume2D(64)

    assert_equal(volume.cfg.width, 64)
    assert_equal(volume.cfg.height, 64)
    assert_equal(volume.cfg.depth, 0)
    assert_equal(volume.cfg.is_3d, False)
    assert_equal(volume.cfg.dx, 0.0)
    assert_equal(volume.cfg.dy, 0.0)
    assert_equal(volume.cfg.dz, 0.0)

    volume = Volume2D(64, 128, (1.0, 2.0))
    assert_equal(volume.cfg.width, 128)
    assert_equal(volume.cfg.height, 64)
    assert_equal(volume.cfg.depth, 0)
    assert_equal(volume.cfg.is_3d, False)
    assert_equal(volume.cfg.dx, 1.0)
    assert_equal(volume.cfg.dy, 2.0)
    assert_equal(volume.cfg.dz, 0.0)


def test_volume_3d():
    volume = Volume3D(64)
    assert_equal(volume.cfg.width, 64)
    assert_equal(volume.cfg.height, 64)
    assert_equal(volume.cfg.depth, 64)
    assert_equal(volume.cfg.is_3d, True)
    assert_equal(volume.cfg.dx, 0.0)
    assert_equal(volume.cfg.dy, 0.0)
    assert_equal(volume.cfg.dz, 0.0)

    volume = Volume3D(64, 128)
    assert_equal(volume.cfg.width, 128)
    assert_equal(volume.cfg.height, 128)
    assert_equal(volume.cfg.depth, 64)
    assert_equal(volume.cfg.is_3d, True)
    assert_equal(volume.cfg.dx, 0.0)
    assert_equal(volume.cfg.dy, 0.0)
    assert_equal(volume.cfg.dz, 0.0)

    volume = Volume3D(64, 128, 256, (1.0, 2.0, 3.0))
    assert_equal(volume.cfg.width, 256)
    assert_equal(volume.cfg.height, 128)
    assert_equal(volume.cfg.depth, 64)
    assert_equal(volume.cfg.is_3d, True)
    assert_equal(volume.cfg.dx, 1.0)
    assert_equal(volume.cfg.dy, 2.0)
    assert_equal(volume.cfg.dz, 3.0)
utils/torch-radon/tests/test_parameter_classes.py
lus6-Jenny/RING
import torch
import numpy as np
from torch.autograd import gradcheck
import torch_radon as tr
from unittest import TestCase
from .utils import generate_random_images


class TestTorch(TestCase):
    def test_differentiation(self):
        device = torch.device('cuda')
        x = torch.FloatTensor(1, 64, 64).to(device)
        x.requires_grad = True
        angles = torch.FloatTensor(np.linspace(0, 2 * np.pi, 10).astype(np.float32)).to(device)

        radon = tr.ParallelBeam(64, angles)

        # check that backward is implemented for fp and bp
        y = radon.forward(x)
        z = torch.mean(radon.backprojection(y))
        z.backward()
        self.assertIsNotNone(x.grad)

    def test_shapes(self):
        """
        Check using channels is ok
        """
        device = torch.device('cuda')
        angles = torch.FloatTensor(np.linspace(0, 2 * np.pi, 10).astype(np.float32)).to(device)
        radon = tr.ParallelBeam(64, angles)

        # test with 2 batch dimensions
        x = torch.FloatTensor(2, 3, 64, 64).to(device)
        y = radon.forward(x)
        self.assertEqual(y.size(), (2, 3, 10, 64))
        z = radon.backprojection(y)
        self.assertEqual(z.size(), (2, 3, 64, 64))

        # no batch dimensions
        x = torch.FloatTensor(64, 64).to(device)
        y = radon.forward(x)
        self.assertEqual(y.size(), (10, 64))
        z = radon.backprojection(y)
        self.assertEqual(z.size(), (64, 64))

#     def test_gradients(self):
#         device = torch.device('cuda')
#         radon = Radon(64).to(device)
#         x = torch.FloatTensor(generate_random_images(1, 64)).to(device)
#         x.requires_grad = True
#         angles = torch.FloatTensor(np.linspace(0, 2 * np.pi, 10).astype(np.float32)).to(device)

#         def f(xx):
#             return radon.backprojection(radon.forward(xx, angles), angles)
        
#         self.assertEqual(gradcheck(f, x, 1e-1, 1e-3, 1e-2), True)

utils/torch-radon/tests/test_torch.py
lus6-Jenny/RING
import torch
import numpy as np
import torch_radon as tr
from unittest import TestCase
from .utils import generate_random_images


class TestVolume(TestCase):
    def test_non_square(self):
        width = 128
        height = 256
        n_angles = 128
        device = torch.device("cuda")

        img = np.random.uniform(0.0, 1.0, (height, width)).astype(np.float32)
        image_size = max(width, height)
        pad_y = (image_size - height) // 2
        pad_x = (image_size - width) // 2
        padded = np.pad(img, ((pad_y, pad_y), (pad_x,pad_x)))

        angles = np.linspace(0, np.pi, n_angles, endpoint=False)

        #projection = Projection.parallel_beam(image_size)
        projection = Projection.fanbeam(image_size, image_size, image_size)

        radon = Radon(angles, img.shape, projection)
        radon_p = Radon(angles, padded.shape, projection)

        with torch.no_grad():
            x = torch.FloatTensor(img).to(device).to(device) #.half()
            xp = torch.FloatTensor(padded).to(device).to(device) #.half()

            sinogram = radon.forward(x)
            sinogram_p = radon_p.forward(xp)

            bp = radon.backprojection(sinogram)
            bp_p = radon_p.backprojection(sinogram_p)

        self.assertLess(torch.norm(sinogram_p - sinogram) / torch.norm(sinogram), 6e-3)
        bp_p_cropped = bp_p[pad_y:pad_y+height, pad_x:pad_x+width]
        self.assertLess(torch.norm(bp_p_cropped - bp) / torch.norm(bp), 6e-4)


    def test_translation(self):
        dy = 16
        dx = 32
        device = torch.device("cuda")
        
        img = np.zeros((256, 256), dtype=np.float32)
        img[128-16+dy:128+16+dy, 128-16+dx:128+16+dx] = 1.0
        vol = Volume.create_2d(*img.shape)

        img_d = np.zeros((256, 256), dtype=np.float32)
        img_d[128-16:128+16, 128-16:128+16] = 1.0
        vol_d = Volume.create_2d(*img_d.shape, dy, dx)

        # compute croppings that align the images
        cropped_h = img.shape[0] - abs(dy)
        cropped_w = img.shape[1] - abs(dx)
        cropped_img = img[max(dy, 0):max(dy, 0) + cropped_h, max(dx, 0):max(dx, 0) + cropped_w] 
        cropped_img_d = img_d[max(-dy, 0):max(-dy, 0) + cropped_h, max(-dx, 0):max(-dx, 0) + cropped_w]
        # if this fails then there is an error in cropping
        self.assertLess(np.linalg.norm(cropped_img-cropped_img_d), 1e-6)

        angles = np.linspace(0, np.pi, 128, endpoint=False)
        projection = Projection.parallel_beam(256)

        radon = Radon(angles, vol, projection)
        radon_d = Radon(angles, vol_d, projection)

        with torch.no_grad():
            x = torch.FloatTensor(img).to(device)
            x_d = torch.FloatTensor(img_d).to(device)

            sinogram = radon.forward(x)
            sinogram_d = radon_d.forward(x_d)

            bp = radon.backprojection(sinogram)
            bp_d = radon_d.backprojection(sinogram_d)

        self.assertLess(torch.norm(sinogram_d - sinogram) / torch.norm(sinogram), 2e-3)

        cropped_bp = bp[max(dy, 0):max(dy, 0) + cropped_h, max(dx, 0):max(dx, 0) + cropped_w] 
        cropped_bp_d = bp_d[max(-dy, 0):max(-dy, 0) + cropped_h, max(-dx, 0):max(-dx, 0) + cropped_w]
        self.assertLess(torch.norm(cropped_bp - cropped_bp_d) / torch.norm(cropped_bp), 2e-4)

    def test_scale(self):
        sy = 1
        sx = 2
        device = torch.device("cuda")
        
        img = np.zeros((256, 256), dtype=np.float32)
        img[128-16:128+16, 128-16:128+16] = 1.0
        vol = Volume.create_2d(*img.shape)

        img_d = img[::sy, ::sx]
        vol_d = Volume.create_2d(*img_d.shape, sy=sy, sx=sx)

        angles = np.linspace(0, np.pi, 128, endpoint=False)
        projection = Projection.parallel_beam(256)

        radon = Radon(angles, vol, projection)
        radon_d = Radon(angles, vol_d, projection)

        with torch.no_grad():
            x = torch.FloatTensor(img).to(device)
            x_d = torch.FloatTensor(img_d).to(device)

            sinogram = radon.forward(x)
            sinogram_d = radon_d.forward(x_d)

            bp = radon.backprojection(sinogram)
            bp_d = radon_d.backprojection(sinogram_d)

        self.assertLess(torch.norm(sinogram_d - sinogram) / torch.norm(sinogram), 2e-3)

        # cropped_bp = bp[max(dy, 0):max(dy, 0) + cropped_h, max(dx, 0):max(dx, 0) + cropped_w] 
        # cropped_bp_d = bp_d[max(-dy, 0):max(-dy, 0) + cropped_h, max(-dx, 0):max(-dx, 0) + cropped_w]
        # self.assertLess(torch.norm(cropped_bp - cropped_bp_d) / torch.norm(cropped_bp), 2e-4)
utils/torch-radon/tests/test_volume.py
lus6-Jenny/RING
import numpy as np
import random
import os
import shutil
from nose.tools import assert_less, assert_equal
import matplotlib.pyplot as plt
import random
import torch

from torch_radon_cuda import SymbolicFunction, symbolic_discretize, symbolic_forward


def relative_error(ref, x):
    return np.linalg.norm(ref - x) / (np.linalg.norm(ref) + 1e-6)


def random_symbolic_function(height, width):
    f = SymbolicFunction(height, width)

    for i in range(30):
        w = random.uniform(0, 1)
        cx = random.uniform(-width/2, width/2)
        cy = random.uniform(-height/2, height/2)
        rx = random.uniform(width/32, width/4)
        ry = random.uniform(height/32, height/4)

        if random.randint(0, 1) < 0.25:
            f.add_ellipse(w, cx, cy, rx, ry)
        else:
            f.add_gaussian(5*w, cx, cy, 0.5/rx**2, 0.5/ry**2)

    return f


class TestHelper:
    def __init__(self, name):
        self.out_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")
        self.name = name
        self.count = 0

        if os.path.exists(self.out_folder):
            shutil.rmtree(self.out_folder)
        os.makedirs(self.out_folder)
        print(self.out_folder, self.name)

    def save_output_image(self, task):
        path = os.path.join(self.out_folder, f"{self.name}_{task}_{self.count}.png")
        self.count += 1
        plt.savefig(path, dpi=100, bbox_inches='tight')
        plt.close()

    def compare_images(self, target, res, max_error, description):
        is_half = res.dtype == torch.float16
        res = res.float().cpu().numpy()
        batch_size = res.shape[0]

        for i in range(batch_size):
            forward_error = relative_error(target, res[i])

            desc = f"Batch: {i}/{batch_size}, {description}\nError: {forward_error:.2e}, Max error: {max_error:.2e}"

            if forward_error > max_error:
                fig, ax = plt.subplots(1, 3, figsize=(10, 7))
                ax = ax.ravel()
                ax[0].imshow(target)
                ax[0].axis("off")
                ax[0].set_title("Target")
                ax[1].imshow(res[i])
                ax[1].axis("off")
                ax[1].set_title("Result")
                ax[2].imshow(np.abs(target - res[i]))
                ax[2].axis("off")
                ax[2].set_title("Difference")
                fig.suptitle(desc)
                fig.tight_layout()
                if is_half:
                    self.save_output_image("forward_half")
                else:
                    self.save_output_image("forward")
                break

            assert_less(forward_error, max_error)

    # makes sure that <y, Ax> = < A^T y, x>
    def backward_check(self, x, y, radon, description, back_max_error):
        batch_size = x.shape[0]

        test_sino = y / torch.max(y)
        target = torch.sum(test_sino * y, dim=(1,2)) / (radon.volume.voxel_size[0] * radon.volume.voxel_size[1])
        bp = radon.backward(test_sino)
        value = torch.sum(bp * x, dim=(1,2))

        for i in range(batch_size):
            backward_error = abs(target[i] - value[i]) / abs(target[i])
            desc = f"Batch: {i}/{batch_size}, {description}\nError: {backward_error:.2e}, Max error: {back_max_error:.2e}"

            if backward_error > back_max_error:
                fig, ax = plt.subplots(1, 2)
                ax = ax.ravel()
                ax[0].imshow(x[i].cpu().numpy())
                ax[0].axis("off")
                ax[1].imshow(bp[i].cpu().numpy())
                ax[1].axis("off")
                fig.suptitle(desc)
                fig.tight_layout()
                self.save_output_image("backward")
                break

        assert_less(backward_error, back_max_error)

utils/torch-radon/tests/utils.py
lus6-Jenny/RING
# select what to "export"

from .volumes import Volume2D, Volume3D
from .radon import FanBeam, ParallelBeam, ConeBeam
from .log_levels import *

# deprecated classes
from .radon import Radon, RadonFanbeam


__version__ = "2.0"
utils/torch-radon/torch_radon/__init__.py
lus6-Jenny/RING
import torch
import os

if os.getenv("TORCH_RADON_DOC_BUILD"):
    print("Not importing cuda backend because this is just the doc build")
    VolumeCfg = None
    ProjectionCfg = None

    class ExecCfg:
        pass
else:
    from torch_radon_cuda import *

utils/torch-radon/torch_radon/cuda_backend.py
lus6-Jenny/RING
import torch
from torch.autograd import Function

from . import cuda_backend


class RadonForward(Function):
    @staticmethod
    def forward(ctx, x, angles, tex_cache, vol_cfg, proj_cfg, exec_cfg_generator, exec_cfg=None):
        exec_cfg = exec_cfg_generator(vol_cfg, proj_cfg,  x.dtype == torch.half) if exec_cfg is None else exec_cfg
        sinogram = cuda_backend.forward(x, angles, tex_cache, vol_cfg, proj_cfg, exec_cfg)
        ctx.tex_cache = tex_cache
        ctx.vol_cfg = vol_cfg
        ctx.proj_cfg = proj_cfg.copy()
        ctx.exec_cfg_generator = exec_cfg_generator
        ctx.save_for_backward(angles)

        return sinogram

    @staticmethod
    def backward(ctx, grad_x):
        angles, = ctx.saved_tensors
        exec_cfg = ctx.exec_cfg_generator(ctx.vol_cfg, ctx.proj_cfg, grad_x.dtype == torch.half)
        grad = cuda_backend.backward(grad_x, angles, ctx.tex_cache, ctx.vol_cfg, ctx.proj_cfg, exec_cfg)
        return grad, None, None, None, None, None, None


class RadonBackprojection(Function):
    @staticmethod
    def forward(ctx, x, angles, tex_cache, vol_cfg, proj_cfg, exec_cfg_generator, exec_cfg=None):
        exec_cfg = exec_cfg_generator(vol_cfg, proj_cfg,  x.dtype == torch.half) if exec_cfg is None else exec_cfg
        image = cuda_backend.backward(x, angles, tex_cache,  vol_cfg, proj_cfg, exec_cfg)
        ctx.tex_cache = tex_cache
        ctx.vol_cfg = vol_cfg
        ctx.proj_cfg = proj_cfg.copy()
        ctx.exec_cfg_generator = exec_cfg_generator
        ctx.save_for_backward(angles)

        return image

    @staticmethod
    def backward(ctx, grad_x):
        angles, = ctx.saved_tensors
        exec_cfg = ctx.exec_cfg_generator(ctx.vol_cfg, ctx.proj_cfg, grad_x.dtype == torch.half)
        grad = cuda_backend.forward(grad_x, angles, ctx.tex_cache, ctx.vol_cfg, ctx.proj_cfg, exec_cfg)
        return grad, None, None, None, None, None, None

utils/torch-radon/torch_radon/differentiable_functions.py
lus6-Jenny/RING
import numpy as np
import torch
import torch.fft

try:
    import scipy.fft

    fftmodule = scipy.fft
except ImportError:
    import numpy.fft

    fftmodule = numpy.fft


class FourierFilters:
    def __init__(self):
        self.cache = dict()

    def get(self, size: int, filter_name: str, device):
        key = (size, filter_name)

        if key not in self.cache:
            ff = torch.FloatTensor(self.construct_fourier_filter(size, filter_name)).view(1, 1, -1, 1).to(device)
            self.cache[key] = ff

        return self.cache[key].to(device)

    @staticmethod
    def construct_fourier_filter(size, filter_name):
        """Construct the Fourier filter.

        This computation lessens artifacts and removes a small bias as
        explained in [1], Chap 3. Equation 61.

        Parameters
        ----------
        size: int
            filter size. Must be even.
        filter_name: str
            Filter used in frequency domain filtering. Filters available:
            ram-lak (ramp), shepp-logan, cosine, hamming, hann.

        Returns
        -------
        fourier_filter: ndarray
            The computed Fourier filter.

        References
        ----------
        .. [1] AC Kak, M Slaney, "Principles of Computerized Tomographic
               Imaging", IEEE Press 1988.

        """
        filter_name = filter_name.lower()

        n = np.concatenate((np.arange(1, size / 2 + 1, 2, dtype=np.int32),
                            np.arange(size / 2 - 1, 0, -2, dtype=np.int32)))
        f = np.zeros(size)
        f[0] = 0.25
        f[1::2] = -1 / (np.pi * n) ** 2

        # Computing the ramp filter from the fourier transform of its
        # frequency domain representation lessens artifacts and removes a
        # small bias as explained in [1], Chap 3. Equation 61
        fourier_filter = 2 * np.real(fftmodule.fft(f))  # ramp filter
        if filter_name == "ramp" or filter_name == "ram-lak":
            pass
        elif filter_name == "shepp-logan":
            # Start from first element to avoid divide by zero
            omega = np.pi * fftmodule.fftfreq(size)[1:]
            fourier_filter[1:] *= np.sin(omega) / omega
        elif filter_name == "cosine":
            freq = np.linspace(0, np.pi, size, endpoint=False)
            cosine_filter = fftmodule.fftshift(np.sin(freq))
            fourier_filter *= cosine_filter
        elif filter_name == "hamming":
            fourier_filter *= fftmodule.fftshift(np.hamming(size))
        elif filter_name == "hann":
            fourier_filter *= fftmodule.fftshift(np.hanning(size))
        else:
            print(
                f"[TorchRadon] Error, unknown filter type '{filter_name}', available filters are: 'ramp', 'shepp-logan', 'cosine', 'hamming', 'hann'")

        return fourier_filter[:size//2+1]

utils/torch-radon/torch_radon/filtering.py
lus6-Jenny/RING
from . import cuda_backend

DEBUG = 0
INFO = 1
WARN = 2
WARNING = 2
ERROR = 3

def set_log_level(level):
    level = max(0, min(int(level), ERROR))
    cuda_backend.set_log_level(level)
utils/torch-radon/torch_radon/log_levels.py
lus6-Jenny/RING
from .cuda_backend import ProjectionCfg

class Projection:
    PARALLEL = 0
    FANBEAM = 1
    CONE_FLAT = 2

    def __init__(self, cfg: ProjectionCfg):
        self.cfg = cfg

    @staticmethod
    def parallel_beam(det_count, det_spacing=1.0):
        return Projection(ProjectionCfg(det_count, det_spacing))

    @staticmethod
    def fanbeam(src_dist, det_dist, det_count, det_spacing=1.0):
        return Projection(ProjectionCfg(
            det_count, det_spacing,
            0, 1.0,
            src_dist, det_dist,
            0.0, 0.0,
            Projection.FANBEAM
        ))

    @staticmethod
    def coneflat(src_dist, det_dist, det_count_u, det_spacing_u=1.0, det_count_v=-1, det_spacing_v=-1.0, pitch=0.0, base_z=0.0):
        return Projection(ProjectionCfg(
            det_count_u, det_spacing_u,
            det_count_v, det_spacing_v,
            src_dist, det_dist,
            pitch, base_z,
            Projection.CONE_FLAT
        ))

    def is_2d(self):
        return self.cfg.is_2d()

utils/torch-radon/torch_radon/projection.py
lus6-Jenny/RING
from .volumes import Volume2D, Volume3D
from .projection import Projection
from .filtering import FourierFilters
from .utils import normalize_shape, ShapeNormalizer, expose_projection_attributes
from .differentiable_functions import RadonForward, RadonBackprojection
from . import cuda_backend
import numpy as np
import torch
import torch.nn.functional as F
from typing import Union
import warnings
warnings.simplefilter('default')


class ExecCfgGeneratorBase:
    def __init__(self):
        pass

    def __call__(self, vol_cfg, proj_cfg, is_half):
        if proj_cfg.projection_type == 2:
            ch = 4 if is_half else 1
            return cuda_backend.ExecCfg(8, 16, 8, ch)

        return cuda_backend.ExecCfg(16, 16, 1, 4)


class BaseRadon:
    def __init__(self, angles, volume, projection):
        # allows angles to be specified as (start_angle, end_angle, n_angles)
        if isinstance(angles, tuple) and len(angles) == 3:
            start_angle, end_angle, n_angles = angles
            angles = np.linspace(start_angle, end_angle, n_angles, endpoint=False)

        # make sure that angles are a PyTorch tensor
        if not isinstance(angles, torch.Tensor):
            angles = torch.FloatTensor(angles)

        self.angles = angles
        self.volume = volume
        self.projection = projection
        self.exec_cfg_generator = ExecCfgGeneratorBase()

        # caches used to avoid reallocation of resources
        self.tex_cache = cuda_backend.TextureCache(8)
        self.fft_cache = cuda_backend.FFTCache(8)
        self.fourier_filters = FourierFilters()

    def _move_parameters_to_device(self, device):
        if device != self.angles.device:
            self.angles = self.angles.to(device)

    def _check_input(self, x):
        if not x.is_contiguous():
            x = x.contiguous()

        if x.dtype == torch.float16:
            assert x.size(
                0) % 4 == 0, f"Batch size must be multiple of 4 when using half precision. Got batch size {x.size(0)}"

        return x

    def forward(self, x: torch.Tensor, angles: torch.Tensor = None, exec_cfg: cuda_backend.ExecCfg = None):
        r"""Radon forward projection.

        :param x: PyTorch GPU tensor.
        :param angles: PyTorch GPU tensor indicating the measuring angles, if None the angles given to the constructor are used
        :returns: PyTorch GPU tensor containing sinograms.
        """
        x = self._check_input(x)
        self._move_parameters_to_device(x.device)

        angles = angles if angles is not None else self.angles

        shape_normalizer = ShapeNormalizer(self.volume.num_dimensions())
        x = shape_normalizer.normalize(x)

        self.projection.cfg.n_angles = len(angles)

        y = RadonForward.apply(x, self.angles, self.tex_cache, self.volume.cfg, self.projection.cfg,
                               self.exec_cfg_generator, exec_cfg)

        return shape_normalizer.unnormalize(y)

    def backprojection(self, sinogram, angles: torch.Tensor = None, exec_cfg: cuda_backend.ExecCfg = None):
        r"""Radon backward projection.

        :param sinogram: PyTorch GPU tensor containing sinograms.
        :param angles: PyTorch GPU tensor indicating the measuring angles, if None the angles given to the constructor
            are used
        :returns: PyTorch GPU tensor containing backprojected volume.
        """
        sinogram = self._check_input(sinogram)
        self._move_parameters_to_device(sinogram.device)

        angles = angles if angles is not None else self.angles

        shape_normalizer = ShapeNormalizer(self.volume.num_dimensions())
        sinogram = shape_normalizer.normalize(sinogram)

        self.projection.cfg.n_angles = len(angles)

        y = RadonBackprojection.apply(sinogram, self.angles, self.tex_cache, self.volume.cfg, self.projection.cfg,
                                      self.exec_cfg_generator, exec_cfg)

        return shape_normalizer.unnormalize(y)

    def backward(self, sinogram, angles: torch.Tensor = None, exec_cfg: cuda_backend.ExecCfg = None):
        r"""Radon backward projection.

        :param sinogram: PyTorch GPU tensor containing sinograms.
        :param angles: PyTorch GPU tensor indicating the measuring angles, if None the angles given to the constructor
            are used
        :returns: PyTorch GPU tensor containing backprojected volume.
        """
        return self.backprojection(sinogram, angles, exec_cfg)

    @normalize_shape(2)
    def filter_sinogram(self, sinogram, filter_name="ramp"):
        size = sinogram.size(2)
        n_angles = sinogram.size(1)

        # Pad sinogram to improve accuracy
        padded_size = max(64, int(2 ** np.ceil(np.log2(2 * size))))
        pad = padded_size - size
        padded_sinogram = F.pad(sinogram.float(), (0, pad, 0, 0))

        sino_fft = cuda_backend.rfft(padded_sinogram, self.fft_cache) / np.sqrt(padded_size)

        # get filter and apply
        f = self.fourier_filters.get(padded_size, filter_name, sinogram.device)
        filtered_sino_fft = sino_fft * f

        # Inverse fft
        filtered_sinogram = cuda_backend.irfft(filtered_sino_fft, self.fft_cache) / np.sqrt(padded_size)
        filtered_sinogram = filtered_sinogram[:, :, :-pad] * (np.pi / (2 * n_angles))

        return filtered_sinogram.to(dtype=sinogram.dtype)


class ParallelBeam(BaseRadon):
    r"""
    |
    .. image:: https://raw.githubusercontent.com/matteo-ronchetti/torch-radon/
            master/pictures/parallelbeam.svg?sanitize=true
        :align: center
        :width: 400px
    |

    Class that implements Radon projection for the Parallel Beam geometry.

    :param det_count: *Required*. Number of rays that will be projected.
    :param angles: *Required*. Array containing the list of measuring angles. Can be a Numpy array, a PyTorch tensor or a tuple
        `(start, end, num_angles)` defining a range.
    :param det_spacing: Distance between two contiguous rays. By default is `1.0`.
    :param volume: Specifies the volume position and scale. By default a square uniform volume 
        of size :attr:`det_count` is used. To create a non-uniform volume specify an instance of :class:`torch_radon.Volume2D`.

    """

    def __init__(self, det_count: int, angles: Union[list, np.array, torch.Tensor, tuple],
                 det_spacing: float = 1.0, volume: Union[Volume2D, None, int, tuple] = None):

        if volume is None:
            volume = Volume2D(det_count, det_count)
        elif isinstance(volume, int):
            volume = Volume2D(volume, volume)
        elif isinstance(volume, tuple):
            volume = Volume2D(*volume)

        projection = Projection.parallel_beam(det_count, det_spacing)

        super().__init__(angles, volume, projection)


expose_projection_attributes(ParallelBeam, [
    ("det_count", "det_count_u"),
    ("det_spacing", "det_spacing_u")
])


class FanBeam(BaseRadon):
    r"""
    |
    .. image:: https://raw.githubusercontent.com/matteo-ronchetti/torch-radon/
            master/pictures/fanbeam.svg?sanitize=true
        :align: center
        :width: 400px
    |

    Class that implements Radon projection for the Fanbeam geometry.

    :param det_count: *Required*. Number of rays that will be projected.
    :param angles: *Required*. Array containing the list of measuring angles. Can be a Numpy array, a PyTorch tensor or a tuple
        `(start, end, num_angles)` defining a range.
    :param src_dist: Distance between the source of rays and the origin. If not specified is set equals to :attr:`det_count`. 
    :param det_dist: Distance between the detector plane and the origin. If not specified is set equals to :attr:`det_dist`.
    :param det_spacing: Distance between two contiguous rays. By default is `(src_dist + det_dist) / src_dist`.
    :param volume: Specifies the volume position and scale. By default a square uniform volume 
        of size :attr:`det_count` is used. To create a non-uniform volume specify an instance of :class:`torch_radon.Volume2D`.

    """

    def __init__(self, det_count: int, angles: Union[list, np.array, torch.Tensor, tuple],
                 src_dist: float = None, det_dist: float = None, det_spacing: float = None,
                 volume: Union[Volume2D, None, int, tuple] = None):

        if src_dist is None:
            src_dist = det_count

        if det_dist is None:
            det_dist = src_dist

        if det_spacing is None:
            det_spacing = (src_dist + det_dist) / src_dist

        if volume is None:
            volume = Volume2D(det_count, det_count)
        elif isinstance(volume, int):
            volume = Volume2D(volume, volume)
        elif isinstance(volume, tuple):
            volume = Volume2D(*volume)

        projection = Projection.fanbeam(src_dist, det_dist, det_count, det_spacing)

        super().__init__(angles, volume, projection)


class ConeBeam(BaseRadon):
    def __init__(self, det_count_u: int, angles: Union[list, np.array, torch.Tensor, tuple],
                 src_dist: float = None, det_dist: float = None, det_count_v: int = -1, det_spacing_u: float = 1.0,
                 det_spacing_v: float = -1.0, pitch: float = 0.0, base_z: float = 0.0,
                 volume: Union[Volume3D, None, int, tuple] = None):

        if src_dist is None:
            src_dist = det_count_u

        if det_dist is None:
            det_dist = src_dist

        det_count_v = det_count_v if det_count_v > 0 else det_count_u
        det_spacing_v = det_spacing_v if det_spacing_v > 0 else det_spacing_u

        if volume is None:
            volume = Volume3D(det_count_u, det_count_u, det_count_u)
        elif isinstance(volume, int):
            volume = Volume3D(volume, volume, volume)
        elif isinstance(volume, tuple):
            volume = Volume3D(*volume)

        projection = Projection.coneflat(src_dist, det_dist, det_count_u, det_spacing_u,
                                         det_count_v, det_spacing_v, pitch, base_z)

        super().__init__(angles, volume, projection)


expose_projection_attributes(ConeBeam, [
    "det_count_u", "det_count_v",
    "det_spacing_u", "det_spacing_v",
    ("src_dist", "s_dist"), ("det_dist", "d_dist"),
    "pitch", ("base_z", "initial_z")
])


# Deprecated classes kept for retrocompatibility

class Radon(BaseRadon):
    def __init__(self, resolution: int, angles, det_count=-1, det_spacing=1.0, clip_to_circle=False):
        warnings.warn(
            "Radon() class is deprecated, use ParallelBeam instead",
            DeprecationWarning
        )

        if det_count <= 0:
            det_count = resolution

        if clip_to_circle:
            print("[TORCH RADON] Clip to circle is deprecated and not considered anymore")

        volume = Volume2D(resolution)
        projection = Projection.parallel_beam(det_count, det_spacing)

        super().__init__(angles, volume, projection)


class RadonFanbeam(BaseRadon):
    def __init__(self, resolution: int, angles, source_distance: float, det_distance: float = -1, det_count: int = -1,
                 det_spacing: float = -1, clip_to_circle=False):
        warnings.warn(
            "RadonFanbeam() class is deprecated, use FanBeam instead",
            DeprecationWarning
        )

        if det_count <= 0:
            det_count = resolution

        if det_distance < 0:
            det_distance = source_distance
            det_spacing = 2.0
        if det_spacing < 0:
            det_spacing = (source_distance + det_distance) / source_distance

        if clip_to_circle:
            print("[TORCH RADON] Clip to circle is deprecated and not considered anymore")

        volume = Volume2D(resolution)

        projection = Projection.fanbeam(source_distance, det_distance, det_count, det_spacing)

        super().__init__(angles, volume, projection)


# class Radon:
#     def __init__(self, angles: Union[list, np.array, torch.Tensor], volume: Union[int, tuple, Volume],
#                  projection: Projection = None):
#         # make sure that angles are a PyTorch tensor
#         if not isinstance(angles, torch.Tensor):
#             angles = torch.FloatTensor(angles)

#         if isinstance(volume, int):
#             volume = Volume.create_2d(volume, volume)
#         elif isinstance(volume, tuple):
#             volume = Volume.create_2d(*volume) if len(volume) == 2 else Volume.create_3d(*volume)

#         if projection is None:
#             projection = Projection.parallel_beam(volume.max_dim())

#         self.angles = angles
#         self.volume = volume
#         self.projection = projection
#         self.exec_cfg_generator = ExecCfgGeneratorBase()

#         # caches used to avoid reallocation of resources
#         self.tex_cache = torch_radon_cuda.TextureCache(8)
#         self.fft_cache = torch_radon_cuda.FFTCache(8)
#         self.fourier_filters = FourierFilters()

#         seed = np.random.get_state()[1][0]
#         self.noise_generator = torch_radon_cuda.RadonNoiseGenerator(seed)

#     def _move_parameters_to_device(self, device):
#         if device != self.angles.device:
#             self.angles = self.angles.to(device)

#     def _check_input(self, x):
#         if not x.is_contiguous():
#             x = x.contiguous()

#         if x.dtype == torch.float16:
#             assert x.size(
#                 0) % 4 == 0, f"Batch size must be multiple of 4 when using half precision. Got batch size {x.size(0)}"

#         return x

#     def forward(self, x: torch.Tensor, angles: torch.Tensor = None, exec_cfg: ExecCfg = None):
#         r"""Radon forward projection.

#         :param x: PyTorch GPU tensor.
#         :param angles: PyTorch GPU tensor indicating the measuring angles, if None the angles given to the constructor
#         are used
#         :returns: PyTorch GPU tensor containing sinograms.
#         """
#         x = self._check_input(x)
#         self._move_parameters_to_device(x.device)

#         angles = angles if angles is not None else self.angles

#         shape_normalizer = ShapeNormalizer(self.volume.num_dimensions())
#         x = shape_normalizer.normalize(x)

#         self.projection.cfg.n_angles = len(angles)

#         y = RadonForward.apply(x, self.angles, self.tex_cache, self.volume.cfg, self.projection.cfg,
#                                self.exec_cfg_generator, exec_cfg)

#         return shape_normalizer.unnormalize(y)

#     def backprojection(self, sinogram, angles: torch.Tensor = None, exec_cfg: ExecCfg = None):
#         r"""Radon backward projection.

#         :param sinogram: PyTorch GPU tensor containing sinograms.
#         :param angles: PyTorch GPU tensor indicating the measuring angles, if None the angles given to the constructor
#         are used
#         :returns: PyTorch GPU tensor containing backprojected volume.
#         """
#         sinogram = self._check_input(sinogram)
#         self._move_parameters_to_device(sinogram.device)

#         angles = angles if angles is not None else self.angles

#         shape_normalizer = ShapeNormalizer(self.volume.num_dimensions())
#         sinogram = shape_normalizer.normalize(sinogram)

#         self.projection.cfg.n_angles = len(angles)

#         y = RadonBackprojection.apply(sinogram, self.angles, self.tex_cache, self.volume.cfg, self.projection.cfg,
#                                       self.exec_cfg_generator, exec_cfg)

#         return shape_normalizer.unnormalize(y)

#     def backward(self, sinogram):
#         r"""Radon backward projection.

#         :param sinogram: PyTorch GPU tensor containing sinograms.
#         :returns: PyTorch GPU tensor containing backprojected volume.
#         """
#         return self.backprojection(sinogram)

#     @normalize_shape(2)
#     def filter_sinogram(self, sinogram, filter_name="ramp"):
#         size = sinogram.size(2)
#         n_angles = sinogram.size(1)

#         # Pad sinogram to improve accuracy
#         padded_size = max(64, int(2 ** np.ceil(np.log2(2 * size))))
#         pad = padded_size - size
#         padded_sinogram = F.pad(sinogram.float(), (0, pad, 0, 0))
#         print(padded_sinogram.size())

#         # sino_fft = torch.fft.rfft(padded_sinogram, norm="ortho")
#         sino_fft = torch_radon_cuda.rfft(padded_sinogram, self.fft_cache)

#         # get filter and apply
#         f = self.fourier_filters.get(padded_size, "ramp", sinogram.device)
#         print("fft", sino_fft.size(), f.size())
#         filtered_sino_fft = sino_fft * f
#         print(filtered_sino_fft.size())

#         # Inverse fft
#         filtered_sinogram = torch_radon_cuda.irfft(filtered_sino_fft, self.fft_cache)
#         filtered_sinogram = filtered_sinogram[:, :, :-pad] * (np.pi / (2 * n_angles * padded_size))

#         return filtered_sinogram.to(dtype=sinogram.dtype)

#     @normalize_shape(2)
#     def add_noise(self, x, signal, density_normalization=1.0, approximate=False):
#         # print("WARN Radon.add_noise is deprecated")

#         torch_radon_cuda.add_noise(x, self.noise_generator, signal, density_normalization, approximate)
#         return x

#     @normalize_shape(2)
#     def emulate_readings(self, x, signal, density_normalization=1.0):
#         return torch_radon_cuda.emulate_sensor_readings(x, self.noise_generator, signal, density_normalization)

#     @normalize_shape(2)
#     def emulate_readings_new(self, x, signal, normal_std, k, bins):
#         return torch_radon_cuda.emulate_readings_new(x, self.noise_generator, signal, normal_std, k, bins)

#     @normalize_shape(2)
#     def readings_lookup(self, sensor_readings, lookup_table):
#         return torch_radon_cuda.readings_lookup(sensor_readings, lookup_table)

#     def set_seed(self, seed=-1):
#         if seed < 0:
#             seed = np.random.get_state()[1][0]

#         self.noise_generator.set_seed(seed)

#     def __del__(self):
#         self.noise_generator.free()


# def compute_lookup_table(sinogram, signal, normal_std, bins=4096, eps=0.01, eps_prob=0.99, eps_k=0.01, verbose=False):
#     s = sinogram.view(-1)
#     device = s.device

#     eps = np.quantile(sinogram.cpu().numpy(), eps) + eps_k

#     # Compute readings normalization value
#     if verbose:
#         print("Computing readings normalization value")
#     k = 0
#     for i in range(1, 5000):
#         a, b = torch_radon_cuda.compute_ab(s, signal, eps, bins * i)
#         if verbose:
#             print(a, b)
#         if a >= (a + b) * eps_prob:
#             k = bins * i
#             break
#     print("Readings normalization value = ", k // bins)

#     # Compute weights for Gaussian error
#     scale = k // bins
#     weights = []
#     for i in range(0, 64):
#         t = scipy.stats.norm.cdf((scale - i - 0.5) / normal_std) - scipy.stats.norm.cdf((- i - 0.5) / normal_std)
#         if t < 0.005:
#             break

#         weights.append(t)

#     weights = weights[scale:][::-1] + weights
#     weights = np.array(weights)

#     border_w = np.asarray([scipy.stats.norm.cdf((-x - 0.5) / normal_std) for x in range(scale)])
#     border_w = torch.FloatTensor(border_w).to(device)

#     log_factorial = np.arange(k + len(weights))
#     log_factorial[0] = 1
#     log_factorial = np.cumsum(np.log(log_factorial).astype(np.float64)).astype(np.float32)
#     log_factorial = torch.Tensor(log_factorial).to(device)

#     weights = torch.FloatTensor(weights).to(device)

#     lookup, lookup_var = torch_radon_cuda.compute_lookup_table(s, weights, signal, bins, scale, log_factorial, border_w)

#     return lookup, lookup_var, scale


# class ReadingsLookup:
#     def __init__(self, radon, bins=4096, mu=None, sigma=None, ks=None, signals=None, normal_stds=None):
#         self.radon = radon
#         self.bins = bins

#         self.mu = [] if mu is None else mu
#         self.sigma = [] if sigma is None else sigma
#         self.ks = [] if ks is None else ks

#         self.signals = [] if signals is None else signals
#         self.normal_stds = [] if normal_stds is None else normal_stds

#         self._mu = None
#         self._sigma = None
#         self._ks = None
#         self._signals = None
#         self._normal_stds = None
#         self._need_repacking = True

#     def repack(self, device):
#         self._mu = torch.FloatTensor(self.mu).to(device)
#         self._sigma = torch.FloatTensor(self.sigma).to(device)
#         self._ks = torch.IntTensor(self.ks).to(device)
#         self._signals = torch.FloatTensor(self.signals).to(device)
#         self._normal_stds = torch.FloatTensor(self.normal_stds).to(device)

#     @staticmethod
#     def from_file(path, radon):
#         obj = np.load(path)

#         bins = int(obj["bins"])

#         return ReadingsLookup(radon, bins, list(obj["mu"]), list(obj["sigma"]), list(obj["ks"]), list(obj["signals"]),
#                               list(obj["normal_stds"]))

#     def save(self, path):
#         self.repack("cpu")
#         np.savez(path, mu=self._mu, sigma=self._sigma, ks=self._ks, signals=self._signals,
#                  normal_stds=self._normal_stds, bins=self.bins)

#     def add_lookup_table(self, sinogram, signal, normal_std, eps=0.01, eps_prob=0.99, eps_k=0.01, verbose=True):
#         lookup, lookup_var, k = compute_lookup_table(sinogram, signal, normal_std, self.bins, eps, eps_prob, eps_k,
#                                                      verbose)

#         self.mu.append(lookup.cpu().numpy())
#         self.sigma.append(lookup_var.cpu().numpy())
#         self.ks.append(k)
#         self.signals.append(signal)
#         self.normal_stds.append(normal_std)
#         self._need_repacking = True

#     @normalize_shape(2)
#     def emulate_readings(self, sinogram, level):
#         if self._need_repacking or self._mu.device != sinogram.device:
#             self.repack(sinogram.device)

#         if isinstance(level, torch.Tensor):
#             return torch_radon_cuda.emulate_readings_multilevel(sinogram, self.radon.noise_generator, self._signals,
#                                                                 self._normal_stds, self._ks, level, self.bins)
#         else:
#             return torch_radon_cuda.emulate_readings_new(sinogram, self.radon.noise_generator, self.signals[level],
#                                                          self.normal_stds[level], self.ks[level], self.bins)

#     @normalize_shape(2)
#     def lookup(self, readings, level):
#         if self._need_repacking or self._mu.device != readings.device:
#             self.repack(readings.device)

#         if isinstance(level, torch.Tensor):
#             mu = torch_radon_cuda.readings_lookup_multilevel(readings, self._mu, level)
#             sigma = torch_radon_cuda.readings_lookup_multilevel(readings, self._sigma, level)
#         else:
#             mu = torch_radon_cuda.readings_lookup(readings, self._mu[level])
#             sigma = torch_radon_cuda.readings_lookup(readings, self._sigma[level])

#         return mu, sigma

#     def random_levels(self, size, device):
#         return torch.randint(0, len(self.mu), (size,), device=device, dtype=torch.int32)

utils/torch-radon/torch_radon/radon.py
lus6-Jenny/RING
from alpha_transform import AlphaShearletTransform
from alpha_transform.fourier_util import my_ifft_shift

from .utils import normalize_shape
import numpy as np
import torch
import os


class ShearletTransform:
    """
    Implementation of Alpha-Shearlet transform based on https://github.com/dedale-fet/alpha-transform/tree/master/alpha_transform.

    Once the shearlet spectrograms are computed all the computations are done on the GPU.

    :param width: Width of the images
    :param height: Height of the images
    :param alphas: List of alpha coefficients that will be used to generate shearlets
    :param cache: If specified it should be a path to a directory that will be used to cache shearlet coefficients in
        order to avoid recomputing them at each instantiation of this class.

    .. note::
        Support both float and double precision.
    """

    def __init__(self, width, height, alphas, cache=None):
        cache_name = f"{width}_{height}_{alphas}.npy"
        if cache is not None:
            if not os.path.exists(cache):
                os.makedirs(cache)

            cache_file = os.path.join(cache, cache_name)
            if os.path.exists(cache_file):
                shifted_spectrograms = np.load(cache_file)
            else:
                alpha_shearlet = AlphaShearletTransform(width, height, alphas, real=True, parseval=True)
                shifted_spectrograms = np.asarray([my_ifft_shift(spec) for spec in alpha_shearlet.spectrograms])
                np.save(cache_file, shifted_spectrograms)
        else:
            alpha_shearlet = AlphaShearletTransform(width, height, alphas, real=True, parseval=True)
            shifted_spectrograms = np.asarray([my_ifft_shift(spec) for spec in alpha_shearlet.spectrograms])

        self.shifted_spectrograms = torch.FloatTensor(shifted_spectrograms)

        self.shifted_spectrograms_d = torch.DoubleTensor(shifted_spectrograms)

    def _move_parameters_to_device(self, device):
        if device != self.shifted_spectrograms.device:
            self.shifted_spectrograms = self.shifted_spectrograms.to(device)
            self.shifted_spectrograms_d = self.shifted_spectrograms_d.to(device)

    @normalize_shape(2)
    def forward(self, x):
        """
        Do shearlet transform of a batch of images.

        :param x: PyTorch GPU tensor with shape :math:`(d_1, \\dots, d_n, h, w)`.
        :returns: PyTorch GPU tensor containing shearlet coefficients.
            Has shape :math:`(d_1, \\dots, d_n, \\text{n_shearlets}, h, w)`.
        """
        self._move_parameters_to_device(x.device)

        c = torch.fft.rfft(x, 2, norm="forward")
        print(self.shifted_spectrograms.size(), c.size())

        if x.dtype == torch.float64:
            cs = torch.einsum("fij,bijc->bfijc", self.shifted_spectrograms_d, c)
        else:
            cs = torch.einsum("fij,bijc->bfijc", self.shifted_spectrograms, c)

        return torch.fft.irfft(cs, 2, norm="forward")

    @normalize_shape(3)
    def backward(self, cs):
        """
        Do inverse shearlet transform.

        :param cs: PyTorch GPU tensor containing shearlet coefficients,
            with shape :math:`(d_1, \\dots, d_n, \\text{n_shearlets}, h, w)`.
        :returns: PyTorch GPU tensor containing reconstructed images.
            Has shape :math:`(d_1, \\dots, d_n, h, w)`.
        """

        cs_fft = torch.fft.rfft(cs, 2, norm="forward")
        print(self.shifted_spectrograms.size(), cs_fft.size())

        if cs.dtype == torch.float64:
            res = torch.einsum("fij,bfijc->bijc", self.shifted_spectrograms_d, cs_fft)
        else:
            print(self.shifted_spectrograms.size(), cs_fft.size())
            res = torch.einsum("fij,bfijc->bijc", self.shifted_spectrograms, cs_fft)

        return torch.fft.irfft(res, 2, norm="forward")

utils/torch-radon/torch_radon/shearlet.py
lus6-Jenny/RING
import torch
from .utils import _normalize_shape, _unnormalize_shape


def normalize(x):
    size = x.size()
    x = x.view(size[0], -1)
    norm = torch.norm(x, dim=1)
    x /= norm.view(-1, 1)
    return x.view(*size), torch.max(norm).item()


class Landweber:
    """
    Class that implements Landweber iteration to solve :math:`\\min_{x \\in C}  \\|Ax-y\\|_2^2`
    (see `Wikipedia page <https://en.wikipedia.org/wiki/Landweber_iteration>`_).

    The iteration used is :math:`x_{n+1} = \\mathcal{P}_C(x - \\alpha A^T A x_n)` where :math:`\\mathcal{P}_C` is the
    projection onto :math:`C`.

    :param operator: Instance of a class that implements products :math:`A x` (``operator.forward(x)``) and :math:`A^T y`
        (``operator.backward(y)``).
    :param projection: Function that implements :math:`\\mathcal{P}_C(\\cdot)`, if not specified no projection is used.
    :param grad: If true gradient will be enabled, more memory will be used but it will be possible to backpropagate.
    """

    def __init__(self, operator, projection=None, grad=False):
        self.operator = operator
        self.projection = lambda x: x if projection is None else projection
        self.grad = grad

    def estimate_alpha(self, img_size, device, n_iter=50, batch_size=8):
        """
        Use power iteration on :math:`A^T A` to estimate the maximum step size that still guarantees convergence.

        .. note::
            Because this computation is not exact it is advised to use a value of alpha lower that the one estimated by
            this method (for example multiplying the estimate by 0.95).

        :param img_size: Size of the image
        :param device: GPU device that will be used for computation
        :param n_iter: Number of iterations
        :param batch_size: Number of vectors used in the power iteration.
        :return: Estimated value for alpha
        """
        with torch.no_grad():
            x = torch.randn((batch_size, img_size, img_size), device=device)
            x, _ = normalize(x)
            for i in range(n_iter):
                next_x = self.operator.backward(self.operator.forward(x))
                x, v = normalize(next_x)

            return 2.0 / v

    def run(self, x_zero, y, alpha, iterations=100, callback=None):
        """
        Execute Landweber iterations.

        :param x_zero: Initial solution guess used as a starting point for the iteration
        :param y: Value of y in :math:`\\min_{x \\in C}  \\|Ax-y\\|_2^2`
        :param alpha: Step size, can be estimated using :attr:`estimate_alpha`
        :param iterations: Number of iterations
        :param callback: Optional function that will be called at each iteration with :math:`x_n` as argument.
            Values returned by ``callback`` will be stored in a list and returned together with the computed solution
        :return: If ``callback`` is specified returns ``x, values`` where ``x`` is the solution computed by the
            Landweber iteration and ``values`` is the list of values returned by ``callback`` at each iteration. If
            ``callback`` is not specified returns only ``x``
        """
        res = []
        with torch.set_grad_enabled(self.grad):
            x = x_zero.clone()
            for i in range(iterations):
                x = self.projection(x - alpha * self.operator.backward(self.operator.forward(x) - y))
                if callback is not None:
                    res.append(callback(x))

        if callback is not None:
            return x, res
        else:
            return x


def cg(forward, x, y, callback=None, max_iter=500, tol=1e-5):
    """
    Implements Conjugate Gradient algorithm for solving :math:`\\min_x  \\|Ax-y\\|_2^2`.

    .. note::
        For conjugate gradient to work the matrix :math:`A` must be symmetric positive definite. Otherwise use other
        solvers.

    :param forward: function that implements products :math:`A x` (``forward(x)``).
    :param x: Initial solution guess used as a starting point for the iteration
    :param y: Value of y in :math:`\\min_{x \\in C}  \\|Ax-y\\|_2^2`
    :param callback: Optional function that will be called at each iteration with :math:`x_n` and the residual as
            arguments. Values returned by ``callback`` will be stored in a list and returned together with the computed
            solution.
    :param max_iter: Maximum number of iterations.
    :param tol: Algorithm is stopped when :math:`\\frac{\\| Ax_n - y \\|}{\\| y \\|} \\leq \\text{tol}`
    :return: If ``callback`` is specified returns ``x, values`` where ``x`` is the solution computed by the
        Landweber iteration and ``values`` is the list of values returned by ``callback`` at each iteration. If
        ``callback`` is not specified returns only ``x``.
    """

    # normalize shape to be (batch, height, width)
    x, old_shape = _normalize_shape(x, 2)
    y, _ = _normalize_shape(y, 2)

    y_norm = torch.norm(y, dim=(1, 2))

    r = y - forward(x)
    p = r.clone()
    r_n = torch.sum(r ** 2, dim=(1, 2))

    values = []
    for i in range(max_iter):
        Ap = forward(p)
        alpha = (r_n / torch.sum(p * Ap, dim=(1, 2)).clamp_min(1e-8)).view(-1, 1, 1)
        x += alpha * p
        r_next = r - alpha * Ap

        if callback is not None:
            values.append(callback(x, r_next))

        r_next_n = torch.sum(r_next ** 2, dim=(1, 2))
        rs = torch.min(torch.sqrt(r_next_n) / y_norm).item()
        if rs <= tol:
            break

        beta = (r_next_n / r_n).view(-1, 1, 1)
        r_n = r_next_n
        p = r_next + beta * p
        r = r_next.clone()

    x = _unnormalize_shape(x, old_shape)

    if callback is None:
        return x
    else:
        return x, values


def cgne(operator, x, y, callback=None, max_iter=5000, tol=1e-5):
    """
    Implements Conjugate Gradient on the Normal Equations, an algorithm for solving :math:`\\min_x  \\|Ax-y\\|_2^2`.

    :param operator: Instance of a class that implements products :math:`A x` (``operator.forward(x)``) and :math:`A^T y`
        (``operator.backward(y)``).
    :param x: Initial solution guess used as a starting point for the iteration
        :param y: Value of y in :math:`\\min_{x \\in C}  \\|Ax-y\\|_2^2`
    :param callback: Optional function that will be called at each iteration with :math:`x_n` as argument.
            Values returned by ``callback`` will be stored in a list and returned together with the computed solution
    :param max_iter: Maximum number of iterations
    :param tol: Algorithm is stopped when :math:`\\frac{\\| s \\|}{\\| y \\|} \\leq \\text{tol}`
    :return: If ``callback`` is specified returns ``x, values`` where ``x`` is the solution computed by the
        Landweber iteration and ``values`` is the list of values returned by ``callback`` at each iteration. If
        ``callback`` is not specified returns only ``x``
    """
    # normalize shape to be (batch, height, width)
    x, old_shape = _normalize_shape(x, 2)
    y, _ = _normalize_shape(y, 2)

    y_norm = torch.norm(y, dim=(1, 2))
    d = y - operator.forward(x)

    s = operator.backprojection(d)
    p = s.clone()
    s_norm = torch.sum(s ** 2, dim=(1, 2))

    values = []

    for i in range(max_iter):
        q = operator.forward(p)
        alpha_den = torch.sum(q ** 2, dim=(1, 2)).clamp_min(1e-8)
        alpha = (s_norm / alpha_den).view(-1, 1, 1)

        x = x + alpha * p
        d = d - alpha * q
        s_next = operator.backprojection(d)
        s_next_norm = torch.sum(s_next ** 2, dim=(1, 2))

        if callback is not None:
            values.append(callback(x))

        # if s_next_norm is small break
        sn = torch.max(torch.sqrt(s_next_norm) / y_norm).item()
        if sn <= tol:
            break

        beta = (s_next_norm / s_norm).view(-1, 1, 1)
        s = s_next
        s_norm = s_next_norm
        p = s + beta * p

    x = _unnormalize_shape(x, old_shape)

    if callback is None:
        return x
    else:
        return x, values

utils/torch-radon/torch_radon/solvers.py
lus6-Jenny/RING
import torch


def _normalize_shape(x, d):
    old_shape = x.size()[:-d]
    x = x.view(-1, *(x.size()[-d:]))
    return x, old_shape


def _unnormalize_shape(y, old_shape):
    if isinstance(y, torch.Tensor):
        y = y.view(*old_shape, *(y.size()[1:]))
    elif isinstance(y, tuple):
        y = [yy.view(*old_shape, *(yy.size()[1:])) for yy in y]

    return y


class ShapeNormalizer:
    def __init__(self, d):
        self.d = d
        self.old_shape = None

    def normalize(self, x):
        x, self.old_shape = _normalize_shape(x, self.d)
        return x

    def unnormalize(self, y):
        if self.old_shape is None:
            raise Exception("Calling `unnormalize` before `normalize` ")

        return _unnormalize_shape(y, self.old_shape)


def normalize_shape(d):
    """
    Input with shape (batch_1, ..., batch_n, s_1, ..., s_d) is reshaped to (batch, s_1, s_2, ...., s_d)
    fed to f and output is reshaped to (batch_1, ..., batch_n, s_1, ..., s_o).
    :param d: Number of non-batch dimensions
    """

    def wrap(f):
        def wrapped(self, x, *args, **kwargs):
            x, old_shape = _normalize_shape(x, d)

            y = f(self, x, *args, **kwargs)

            return _unnormalize_shape(y, old_shape)

        wrapped.__doc__ = f.__doc__
        return wrapped

    return wrap


def projection_property_maker(name):
    @property
    def prop(self):
        return getattr(self.projection.cfg, name)

    @prop.setter
    def prop(self, value):
        setattr(self.projection.cfg, name, value)

    return prop


def expose_projection_attributes(pyclass, attributes: list):
    """Exposes the attributes of the projection directly from the class.
    For example, after exposing "det_spacing_u" (internal name) as "det_spacing" (exposed name), setting radon.det_spacing = 32
    is equivalent to setting  radon.projection.cfg.det_spacing_u = 32

    Args:
        pyclass: A python class (not instance of a class but the actual class)
        attributes: List of attributes, each element can be a string (then exposed_name = internal_name) or a tuple (exposed_name, internal_name)
    """
    for x in attributes:
        exposed_name, internal_name = x if isinstance(x, tuple) else (x, x)
        setattr(pyclass, exposed_name, projection_property_maker(internal_name))

utils/torch-radon/torch_radon/utils.py
lus6-Jenny/RING
from .cuda_backend import VolumeCfg


class Volume2D:
    def __init__(self, height, width=-1, center=(0.0, 0.0), voxel_size=(1.0, 1.0)):
        if width <= 0:
            width = height

        self.height = height
        self.width = width

        self.center = center
        self.voxel_size = voxel_size

        self.cfg = self.to_cfg()

    def to_cfg(self) -> VolumeCfg:
        return VolumeCfg(
            0, self.height, self.width,
            0.0, self.center[1], self.center[0],
            1.0, self.voxel_size[1], self.voxel_size[0],
            False
        )

    def num_dimensions(self):
        return 2

    def __str__(self) -> str:
        return f"Volume2D(height={self.height}, width={self.width}, center={self.center}, voxel_size={self.voxel_size})"


class Volume3D:
    def __init__(self, depth, height=-1, width=-1, center=(0.0, 0.0, 0.0), voxel_size=(1.0, 1.0, 1.0)):
        if height <= 0:
            height = depth

        if width <= 0:
            width = height

        self.depth = depth
        self.height = height
        self.width = width

        self.center = center
        self.voxel_size = voxel_size

        self.cfg = self.to_cfg()

    def shape(self):
        return (self.depth, self.height, self.width)

    def min(self):
        dx, dy, dz = self.center
        sx, sy, sz = self.voxel_size
        return [-self.width*sx / 2 + dx, -self.height*sy/2 + dy, -self.depth*sz/2 + dz]

    def max(self):
        dx, dy, dz = self.center
        sx, sy, sz = self.voxel_size
        return [self.width*sx / 2 + dx, self.height*sy/2 + dy, self.depth*sz/2 + dz]

    def to_cfg(self) -> VolumeCfg:
        return VolumeCfg(
            self.depth, self.height, self.width,
            self.center[2], self.center[1], self.center[0],
            self.voxel_size[2], self.voxel_size[1], self.voxel_size[0],
            True
        )

    def num_dimensions(self):
        return 3

utils/torch-radon/torch_radon/volumes.py
lus6-Jenny/RING
import re
import os
import urllib.request
import urllib.parse


url = 'https://download.pytorch.org/whl/torch_stable.html'
f = urllib.request.urlopen(url)
txt = f.read().decode('utf-8')

regex = r'<a\s*href\s*=\s*"[^"]*"\s*>cu([0-9]+)\/torch-([0-9.]+)[^-]*-cp([0-9]{2})'

configs = []
for cu, pt, py in re.findall(regex, txt):
    cuda = int(cu)
    python = int(py)
    pytorch = int(pt[:3].replace(".", ""))

    if cuda >= 100 and 38 >= python >= 36 and pytorch >= 16 and python == int(os.getenv("PYTHON_VERSION")):
        configs.append((cuda, python, pytorch))

configs = list(set(configs))

print(f"Need to compile {len(configs)} packages")

script = [
    "cd /code",
    'eval "$(/root/miniconda3/bin/conda shell.bash hook)"',
    "mkdir output",
    "export CXX=g++"
]

for i, (cuda, python, torch) in enumerate(configs):
    cuda_full = f"{cuda // 10}.{cuda % 10}"
    python_full = f"{python // 10}.{python % 10}"
    torch_full = f"{torch // 10}.{torch % 10}"
    env = f"py{python}cu{cuda}pt{torch}"

    script += [
        f"# Python {python}, PyTorch {torch_full}, CUDA {cuda_full}",
        f"printf 'Python {python}, PyTorch {torch_full}, CUDA {cuda_full}\n'",
        f"printf 'Progress {int(100*(i+1) / len(configs))}%% ({i+1}/{len(configs)})\n'",
        f"mkdir -p output/cuda-{cuda_full}/torch-{torch_full}",
        f"conda create -n {env} python={python_full}",
        f"conda install -n {env} pytorch={torch_full} cudatoolkit={cuda_full} -c pytorch -c conda-forge",
        f"source /root/miniconda3/bin/activate {env}",
        "pip install Demangler",
        "python --version",
        'python -c "import torch; print(torch.version.cuda)"',
        "python make.py clean",
        f"export CUDA_HOME=/usr/local/cuda-{cuda_full}",
        # force recompilation otherwise will reuse builds even if CUDA version changes
        "python setup.py build_ext --force",
        "python setup.py bdist_wheel",
        f"mv dist/*.whl output/cuda-{cuda_full}/torch-{torch_full}/",
        "conda deactivate",
        f"conda env remove -n {env}",
    ]

    if (i+1) % 5 == 0:
        script.append("conda clean -ayf")

    script += [
        "printf 'Disk free\n'",
        "df -h"
    ]

with open("/code/travis/do_build.sh", "w") as f:
    f.write("\n".join(script))

utils/torch-radon/travis/create_build_script.py
lus6-Jenny/RING
import numpy as np
import torch
EPS = 1e-6

def print_(name, tensor):
    tensor = tensor.detach().cpu().numpy()
    print(name, tensor, tensor.shape)

def print_stats(name, tensor):
    shape = tensor.shape
    tensor = tensor.detach().cpu().numpy()
    print('%s (%s) min = %.2f, mean = %.2f, max = %.2f' % (name, tensor.dtype, np.min(tensor), np.mean(tensor), np.max(tensor)), shape)
    
def strnum(x):
    s = '%g' % x
    if '.' in s:
        if x < 1.0:
            s = s[s.index('.'):]
    return s

def matmul2(mat1, mat2):
    return torch.matmul(mat1, mat2)

def pack_seqdim(tensor, B):
    shapelist = list(tensor.shape)
    B_, S = shapelist[:2]
    assert(B==B_)
    otherdims = shapelist[2:]
    tensor = torch.reshape(tensor, [B*S]+otherdims)
    return tensor

def unpack_seqdim(tensor, B):
    shapelist = list(tensor.shape)
    BS = shapelist[0]
    assert(BS%B==0)
    otherdims = shapelist[1:]
    S = int(BS/B)
    tensor = torch.reshape(tensor, [B,S]+otherdims)
    return tensor

def normalize_single(d):
    # d is a whatever shape torch tensor
    dmin = torch.min(d)
    dmax = torch.max(d)
    d = (d-dmin)/(EPS+(dmax-dmin))
    return d

def normalize(d):
    # d is B x whatever. normalize within each element of the batch
    out = torch.zeros(d.size())
    if d.is_cuda:
        out = out.cuda()
    B = list(d.size())[0]
    for b in list(range(B)):
        out[b] = normalize_single(d[b])
    return out

def reduce_masked_mean(x, mask, dim=None, keepdim=False):
    # x and mask are the same shape, or at least broadcastably so < actually it's safer if you disallow broadcasting
    # returns shape-1
    # axis can be a list of axes
    for (a,b) in zip(x.size(), mask.size()):
        # if not b==1:
        assert(a==b) # some shape mismatch!
    # assert(x.size() == mask.size())
    prod = x*mask
    if dim is None:
        numer = torch.sum(prod)
        denom = EPS+torch.sum(mask)
    else:
        numer = torch.sum(prod, dim=dim, keepdim=keepdim)
        denom = EPS+torch.sum(mask, dim=dim, keepdim=keepdim)

    mean = numer/denom
    return mean

def meshgrid2d(B, Y, X, stack=False, norm=False, device='cuda'):
    # returns a meshgrid sized B x Y x X

    grid_y = torch.linspace(0.0, Y-1, Y, device=torch.device(device))
    grid_y = torch.reshape(grid_y, [1, Y, 1])
    grid_y = grid_y.repeat(B, 1, X)

    grid_x = torch.linspace(0.0, X-1, X, device=torch.device(device))
    grid_x = torch.reshape(grid_x, [1, 1, X])
    grid_x = grid_x.repeat(B, Y, 1)

    if norm:
        grid_y, grid_x = normalize_grid2d(
            grid_y, grid_x, Y, X)

    if stack:
        # note we stack in xy order
        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
        grid = torch.stack([grid_x, grid_y], dim=-1)
        return grid
    else:
        return grid_y, grid_x
    
def meshgrid3d(B, Z, Y, X, stack=False, norm=False, device='cuda'):
    # returns a meshgrid sized B x Z x Y x X

    grid_z = torch.linspace(0.0, Z-1, Z, device=device)
    grid_z = torch.reshape(grid_z, [1, Z, 1, 1])
    grid_z = grid_z.repeat(B, 1, Y, X)

    grid_y = torch.linspace(0.0, Y-1, Y, device=device)
    grid_y = torch.reshape(grid_y, [1, 1, Y, 1])
    grid_y = grid_y.repeat(B, Z, 1, X)

    grid_x = torch.linspace(0.0, X-1, X, device=device)
    grid_x = torch.reshape(grid_x, [1, 1, 1, X])
    grid_x = grid_x.repeat(B, Z, Y, 1)

    if norm:
        grid_z, grid_y, grid_x = normalize_grid3d(
            grid_z, grid_y, grid_x, Z, Y, X)

    if stack:
        # note we stack in xyz order
        # (see https://pytorch.org/docs/stable/nn.functional.html#torch.nn.functional.grid_sample)
        grid = torch.stack([grid_x, grid_y, grid_z], dim=-1)
        return grid
    else:
        return grid_z, grid_y, grid_x

def gridcloud3d(B, Z, Y, X, norm=False, device='cuda'):
    # we want to sample for each location in the grid
    grid_z, grid_y, grid_x = meshgrid3d(B, Z, Y, X, norm=norm, device=device)
    x = torch.reshape(grid_x, [B, -1])
    y = torch.reshape(grid_y, [B, -1])
    z = torch.reshape(grid_z, [B, -1])
    # these are B x N
    xyz = torch.stack([x, y, z], dim=2)
    # this is B x N x 3
    return xyz

def normalize_grid2d(grid_y, grid_x, Y, X, clamp_extreme=True):
    # make things in [-1,1]
    grid_y = 2.0*(grid_y / float(Y-1)) - 1.0
    grid_x = 2.0*(grid_x / float(X-1)) - 1.0

    if clamp_extreme:
        grid_y = torch.clamp(grid_y, min=-2.0, max=2.0)
        grid_x = torch.clamp(grid_x, min=-2.0, max=2.0)

    return grid_y, grid_x

def normalize_grid3d(grid_z, grid_y, grid_x, Z, Y, X, clamp_extreme=True):
    # make things in [-1,1]
    grid_z = 2.0*(grid_z / float(Z-1)) - 1.0
    grid_y = 2.0*(grid_y / float(Y-1)) - 1.0
    grid_x = 2.0*(grid_x / float(X-1)) - 1.0

    if clamp_extreme:
        grid_z = torch.clamp(grid_z, min=-2.0, max=2.0)
        grid_y = torch.clamp(grid_y, min=-2.0, max=2.0)
        grid_x = torch.clamp(grid_x, min=-2.0, max=2.0)

    return grid_z, grid_y, grid_x

utils/vox_utils/basic.py
lus6-Jenny/RING
import torch
import utils.vox_utils.basic as basic

import numpy as np

def eye_4x4(B, device='cuda'):
    rt = torch.eye(4, device=torch.device(device)).view(1,4,4).repeat([B, 1, 1])
    return rt

def safe_inverse(a): #parallel version
    B, _, _ = list(a.shape)
    inv = a.clone()
    r_transpose = a[:, :3, :3].transpose(1,2) #inverse of rotation matrix

    inv[:, :3, :3] = r_transpose
    inv[:, :3, 3:4] = -torch.matmul(r_transpose, a[:, :3, 3:4])

    return inv

def safe_inverse_single(a):
    r, t = split_rt_single(a)
    t = t.view(3,1)
    r_transpose = r.t()
    inv = torch.cat([r_transpose, -torch.matmul(r_transpose, t)], 1)
    bottom_row = a[3:4, :] # this is [0, 0, 0, 1]
    # bottom_row = torch.tensor([0.,0.,0.,1.]).view(1,4)
    inv = torch.cat([inv, bottom_row], 0)
    return inv

def apply_4x4(RT, xyz):
    B, N, _ = list(xyz.shape)
    ones = torch.ones_like(xyz[:,:,0:1])
    xyz1 = torch.cat([xyz, ones], 2)
    xyz1_t = torch.transpose(xyz1, 1, 2)
    # this is B x 4 x N
    xyz2_t = torch.matmul(RT, xyz1_t)
    xyz2 = torch.transpose(xyz2_t, 1, 2)
    # xyz2 = xyz2 / xyz2[:,:,3:4]
    xyz2 = xyz2[:,:,:3]
    return xyz2


def apply_r4x4(xyz, RT):
    B, N, _ = list(xyz.shape)
    ones = torch.ones_like(xyz[:,:,0:1])
    xyz1 = torch.cat([xyz, ones], 2)
    # this is B x N x 4
    xyz2 = torch.matmul(xyz1, RT)
    # xyz2 = xyz2 / xyz2[:,:,3:4]
    xyz2 = xyz2[:,:,:3]
    return xyz2

def get_camM_T_camXs(origin_T_camXs, ind=0):
    B, S = list(origin_T_camXs.shape)[0:2]
    camM_T_camXs = torch.zeros_like(origin_T_camXs)
    for b in list(range(B)):
        camM_T_origin = safe_inverse_single(origin_T_camXs[b,ind])
        for s in list(range(S)):
            camM_T_camXs[b,s] = torch.matmul(camM_T_origin, origin_T_camXs[b,s])
    return camM_T_camXs

def split_rt_single(rt):
    r = rt[:3, :3]
    t = rt[:3, 3].view(3)
    return r, t

def split_rt(rt):
    r = rt[:, :3, :3]
    t = rt[:, :3, 3].view(-1, 3)
    return r, t

def merge_rt(r, t):
    # r is B x 3 x 3
    # t is B x 3
    B, C, D = list(r.shape)
    B2, D2 = list(t.shape)
    assert(C==3)
    assert(D==3)
    assert(B==B2)
    assert(D2==3)
    t = t.view(B, 3)
    rt = eye_4x4(B, device=t.device)
    rt[:,:3,:3] = r
    rt[:,:3,3] = t
    return rt

def scale_intrinsics(K, sx, sy):
    fx, fy, x0, y0 = split_intrinsics(K)
    fx = fx*sx
    fy = fy*sy
    x0 = x0*sx
    y0 = y0*sy
    K = merge_intrinsics(fx, fy, x0, y0)
    return K

def split_intrinsics(K):
    # K is B x 3 x 3 or B x 4 x 4
    fx = K[:,0,0]
    fy = K[:,1,1]
    x0 = K[:,0,2]
    y0 = K[:,1,2]
    return fx, fy, x0, y0

def merge_intrinsics(fx, fy, x0, y0):
    B = list(fx.shape)[0]
    K = torch.zeros(B, 4, 4, dtype=torch.float32, device=fx.device)
    K[:,0,0] = fx
    K[:,1,1] = fy
    K[:,0,2] = x0
    K[:,1,2] = y0
    K[:,2,2] = 1.0
    K[:,3,3] = 1.0
    return K

def merge_rtlist(rlist, tlist):
    B, N, D, E = list(rlist.shape)
    assert(D==3)
    assert(E==3)
    B, N, F = list(tlist.shape)
    assert(F==3)

    __p = lambda x: utils.basic.pack_seqdim(x, B)
    __u = lambda x: utils.basic.unpack_seqdim(x, B)
    rlist_, tlist_ = __p(rlist), __p(tlist)
    rtlist_ = merge_rt(rlist_, tlist_)
    rtlist = __u(rtlist_)
    return rtlist

def split_lrtlist(lrtlist):
    # splits a B x N x 19 tensor
    # into B x N x 3 (lens)
    # and B x N x 4 x 4 (rts)
    B, N, D = list(lrtlist.shape)
    assert(D==19)
    lenlist = lrtlist[:,:,:3].reshape(B, N, 3)
    ref_T_objs_list = lrtlist[:,:,3:].reshape(B, N, 4, 4)
    return lenlist, ref_T_objs_list

def merge_lrtlist(lenlist, rtlist):
    # lenlist is B x N x 3
    # rtlist is B x N x 4 x 4
    # merges these into a B x N x 19 tensor
    B, N, D = list(lenlist.shape)
    assert(D==3)
    B2, N2, E, F = list(rtlist.shape)
    assert(B==B2)
    assert(N==N2)
    assert(E==4 and F==4)
    rtlist = rtlist.reshape(B, N, 16)
    lrtlist = torch.cat([lenlist, rtlist], axis=2)
    return lrtlist

def apply_4x4_to_lrtlist(Y_T_X, lrtlist_X):
    B, N, D = list(lrtlist_X.shape)
    assert(D==19)
    B2, E, F = list(Y_T_X.shape)
    assert(B2==B)
    assert(E==4 and F==4)

    lenlist, rtlist_X = split_lrtlist(lrtlist_X)

    Y_T_Xs = Y_T_X.unsqueeze(1).repeat(1, N, 1, 1)
    Y_T_Xs_ = Y_T_Xs.view(B*N, 4, 4)
    rtlist_X_ = rtlist_X.reshape(B*N, 4, 4)
    rtlist_Y_ = utils.basic.matmul2(Y_T_Xs_, rtlist_X_)
    rtlist_Y = rtlist_Y_.reshape(B, N, 4, 4)
    lrtlist_Y = merge_lrtlist(lenlist, rtlist_Y)
    return lrtlist_Y

def apply_4x4_to_lrt(Y_T_X, lrt_X):
    B, D = list(lrt_X.shape)
    assert(D==19)
    B2, E, F = list(Y_T_X.shape)
    assert(B2==B)
    assert(E==4 and F==4)

    return apply_4x4_to_lrtlist(Y_T_X, lrt_X.unsqueeze(1)).squeeze(1)

def get_xyzlist_from_lenlist(lenlist):
    B, N, D = list(lenlist.shape)
    assert(D==3)
    lx, ly, lz = torch.unbind(lenlist, axis=2)

    xs = torch.stack([lx/2., lx/2., -lx/2., -lx/2., lx/2., lx/2., -lx/2., -lx/2.], axis=2)
    ys = torch.stack([ly/2., ly/2., ly/2., ly/2., -ly/2., -ly/2., -ly/2., -ly/2.], axis=2)
    zs = torch.stack([lz/2., -lz/2., -lz/2., lz/2., lz/2., -lz/2., -lz/2., lz/2.], axis=2)

    # these are B x N x 8
    xyzlist = torch.stack([xs, ys, zs], axis=3)
    # this is B x N x 8 x 3
    return xyzlist

def get_xyzlist_from_lrtlist(lrtlist, include_clist=False):
    B, N, D = list(lrtlist.shape)
    assert(D==19)

    lenlist, rtlist = split_lrtlist(lrtlist)
    # lenlist is B x N x 3
    # rtlist is B x N x 4 x 4

    xyzlist_obj = get_xyzlist_from_lenlist(lenlist)
    # xyzlist_obj is B x N x 8 x 3

    rtlist_ = rtlist.reshape(B*N, 4, 4)
    xyzlist_obj_ = xyzlist_obj.reshape(B*N, 8, 3)
    xyzlist_cam_ = apply_4x4(rtlist_, xyzlist_obj_)
    xyzlist_cam = xyzlist_cam_.reshape(B, N, 8, 3)

    if include_clist:
        clist_cam = get_clist_from_lrtlist(lrtlist).unsqueeze(2)
        xyzlist_cam = torch.cat([xyzlist_cam, clist_cam], dim=2)
    return xyzlist_cam

def get_clist_from_lrtlist(lrtlist):
    B, N, D = list(lrtlist.shape)
    assert(D==19)

    lenlist, rtlist = split_lrtlist(lrtlist)
    # lenlist is B x N x 3
    # rtlist is B x N x 4 x 4

    xyzlist_obj = torch.zeros((B, N, 1, 3), device=lrtlist.device)
    # xyzlist_obj is B x N x 8 x 3

    rtlist_ = rtlist.reshape(B*N, 4, 4)
    xyzlist_obj_ = xyzlist_obj.reshape(B*N, 1, 3)
    xyzlist_cam_ = apply_4x4(rtlist_, xyzlist_obj_)
    xyzlist_cam = xyzlist_cam_.reshape(B, N, 3)
    return xyzlist_cam

def wrap2pi(rad_angle):
    # puts the angle into the range [-pi, pi]
    return torch.atan2(torch.sin(rad_angle), torch.cos(rad_angle))

def xyd2pointcloud(xyd, pix_T_cam):
    # xyd is like a pointcloud but in pixel coordinates;
    # this means xy comes from a meshgrid with bounds H, W, 
    # and d comes from a depth map
    B, N, C = list(xyd.shape)
    assert(C==3)
    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
    xyz = pixels2camera(xyd[:,:,0], xyd[:,:,1], xyd[:,:,2], fx, fy, x0, y0)
    return xyz

def pixels2camera(x,y,z,fx,fy,x0,y0):
    # x and y are locations in pixel coordinates, z is a depth in meters
    # they can be images or pointclouds
    # fx, fy, x0, y0 are camera intrinsics
    # returns xyz, sized B x N x 3

    B = x.shape[0]
    
    fx = torch.reshape(fx, [B,1])
    fy = torch.reshape(fy, [B,1])
    x0 = torch.reshape(x0, [B,1])
    y0 = torch.reshape(y0, [B,1])

    x = torch.reshape(x, [B,-1])
    y = torch.reshape(y, [B,-1])
    z = torch.reshape(z, [B,-1])
    
    # unproject
    x = (z/fx)*(x-x0)
    y = (z/fy)*(y-y0)
    
    xyz = torch.stack([x,y,z], dim=2)
    # B x N x 3
    return xyz

def camera2pixels(xyz, pix_T_cam):
    # xyz is shaped B x H*W x 3
    # returns xy, shaped B x H*W x 2
    
    fx, fy, x0, y0 = split_intrinsics(pix_T_cam)
    x, y, z = torch.unbind(xyz, dim=-1)
    B = list(z.shape)[0]

    fx = torch.reshape(fx, [B,1])
    fy = torch.reshape(fy, [B,1])
    x0 = torch.reshape(x0, [B,1])
    y0 = torch.reshape(y0, [B,1])
    x = torch.reshape(x, [B,-1])
    y = torch.reshape(y, [B,-1])
    z = torch.reshape(z, [B,-1])

    EPS = 1e-4
    z = torch.clamp(z, min=EPS)
    x = (x*fx)/z + x0
    y = (y*fy)/z + y0
    xy = torch.stack([x, y], dim=-1)
    return xy

utils/vox_utils/geom.py
lus6-Jenny/RING
import numpy as np
import torch
import cv2

import basic
# color conversion libs, for flow vis
from skimage.color import (
    rgb2lab, rgb2yuv, rgb2ycbcr, lab2rgb, yuv2rgb, ycbcr2rgb,
    rgb2hsv, hsv2rgb, rgb2xyz, xyz2rgb, rgb2hed, hed2rgb)

def _convert(input_, type_):
    return {
        'float': input_.float(),
        'double': input_.double(),
    }.get(type_, input_)

def _generic_transform_sk_3d(transform, in_type='', out_type=''):
    def apply_transform_individual(input_):
        device = input_.device
        input_ = input_.cpu()
        input_ = _convert(input_, in_type)

        input_ = input_.permute(1, 2, 0).detach().numpy()
        transformed = transform(input_)
        output = torch.from_numpy(transformed).float().permute(2, 0, 1)
        output = _convert(output, out_type)
        return output.to(device)

    def apply_transform(input_):
        to_stack = []
        for image in input_:
            to_stack.append(apply_transform_individual(image))
        return torch.stack(to_stack)
    return apply_transform

hsv_to_rgb = _generic_transform_sk_3d(hsv2rgb)

def back2color(i, blacken_zeros=False):
    if blacken_zeros:
        const = torch.tensor([-0.5])
        i = torch.where(i==0.0, const.cuda() if i.is_cuda else const, i)
        return back2color(i)
    else:
        return ((i+0.5)*255).type(torch.ByteTensor)

def colorize(d):
    if d.ndim==2:
        d = d.unsqueeze(dim=0)
    else:
        assert(d.ndim==3)
    d = d.repeat(3, 1, 1)
    return d

def oned2inferno(d, norm=True):
    # convert a 1chan input to a 3chan image output

    # if it's just B x H x W, add a C dim
    if d.ndim==3:
        d = d.unsqueeze(dim=1)
    # d should be B x C x H x W, where C=1
    B, C, H, W = list(d.shape)
    assert(C==1)

    if norm:
        d = basic.normalize(d)
        
    rgb = torch.zeros(B, 3, H, W)
    for b in list(range(B)):
        rgb[b] = colorize(d[b])

    rgb = (255.0*rgb).type(torch.ByteTensor)

    return rgb

def convert_occ_to_height(occ, reduce_axis=3):
    B, C, D, H, W = list(occ.shape)
    assert(C==1)
    # note that height increases DOWNWARD in the tensor
    # (like pixel/camera coordinates)
    
    G = list(occ.shape)[reduce_axis]
    values = torch.linspace(float(G), 1.0, steps=G, dtype=torch.float32, device=occ.device)
    if reduce_axis==2:
        # fro view
        values = values.view(1, 1, G, 1, 1)
    elif reduce_axis==3:
        # top view
        values = values.view(1, 1, 1, G, 1)
    elif reduce_axis==4:
        # lateral view
        values = values.view(1, 1, 1, 1, G)
    else:
        assert(False) # you have to reduce one of the spatial dims (2-4)
    values = torch.max(occ*values, dim=reduce_axis)[0]/float(G)
    # values = values.view([B, C, D, W])
    return values

def draw_frame_id_on_vis(vis, frame_id, scale=0.5, left=5, top=20):

    rgb = vis.detach().cpu().numpy()[0]
    rgb = np.transpose(rgb, [1, 2, 0]) # put channels last
    rgb = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    color = (255, 255, 255)

    frame_str = basic.strnum(frame_id)

    cv2.putText(
        rgb,
        frame_str,
        (left, top), # from left, from top
        cv2.FONT_HERSHEY_SIMPLEX,
        scale, # font scale (float)
        color,
        1) # font thickness (int)
    rgb = cv2.cvtColor(rgb.astype(np.uint8), cv2.COLOR_BGR2RGB)
    vis = torch.from_numpy(rgb).permute(2, 0, 1).unsqueeze(0)
    return vis

class Summ_writer(object):
    def __init__(self, writer, global_step, log_freq=10, fps=8, scalar_freq=100, just_gif=False):
        self.writer = writer
        self.global_step = global_step
        self.log_freq = log_freq
        self.fps = fps
        self.just_gif = just_gif
        self.maxwidth = 10000
        self.save_this = (self.global_step % self.log_freq == 0)
        self.scalar_freq = max(scalar_freq,1)

    def summ_gif(self, name, tensor, blacken_zeros=False):
        # tensor should be in B x S x C x H x W

        assert tensor.dtype in {torch.uint8,torch.float32}
        shape = list(tensor.shape)

        if tensor.dtype == torch.float32:
            tensor = back2color(tensor, blacken_zeros=blacken_zeros)

        video_to_write = tensor[0:1]

        S = video_to_write.shape[1]
        if S==1:
            # video_to_write is 1 x 1 x C x H x W
            self.writer.add_image(name, video_to_write[0,0], global_step=self.global_step)
        else:
            self.writer.add_video(name, video_to_write, fps=self.fps, global_step=self.global_step)

        return video_to_write

    def summ_scalar(self, name, value):
        if (not (isinstance(value, int) or isinstance(value, float) or isinstance(value, np.float32))) and ('torch' in value.type()):
            value = value.detach().cpu().numpy()
        if not np.isnan(value):
            if (self.log_freq == 1):
                self.writer.add_scalar(name, value, global_step=self.global_step)
            elif self.save_this or np.mod(self.global_step, self.scalar_freq)==0:
                self.writer.add_scalar(name, value, global_step=self.global_step)

    def summ_oned(self, name, im, bev=False, fro=False, logvis=False, max_val=0, max_along_y=False, norm=True, frame_id=None, only_return=False):
        if self.save_this:
            if bev:
                B, C, H, _, W = list(im.shape)
                if max_along_y:
                    im = torch.max(im, dim=3)[0]
                else:
                    im = torch.mean(im, dim=3)
            elif fro:
                B, C, _, H, W = list(im.shape)
                if max_along_y:
                    im = torch.max(im, dim=2)[0]
                else:
                    im = torch.mean(im, dim=2)
            else:
                B, C, H, W = list(im.shape)

            im = im[0:1] # just the first one
            assert(C==1)

            if logvis and max_val:
                max_val = np.log(max_val)
                im = torch.log(im)
                im = torch.clamp(im, 0, max_val)
                im = im/max_val
                norm = False
            elif max_val:
                im = torch.clamp(im, 0, max_val)/max_val
                norm = False

            vis = oned2inferno(im, norm=norm)
            if W > self.maxwidth:
                vis = vis[...,:self.maxwidth]
            return self.summ_rgb(name, vis, blacken_zeros=False, frame_id=frame_id, only_return=only_return)

    def summ_rgb(self, name, ims, blacken_zeros=False, frame_id=None, only_return=False, halfres=False):
        if self.save_this:
            assert ims.dtype in {torch.uint8,torch.float32}

            if ims.dtype == torch.float32:
                ims = back2color(ims, blacken_zeros)

            #ims is B x C x H x W
            vis = ims[0:1] # just the first one
            B, C, H, W = list(vis.shape)

            if halfres:
                vis = F.interpolate(vis, scale_factor=0.5)

            if frame_id is not None:
                vis = draw_frame_id_on_vis(vis, frame_id)

            if int(W) > self.maxwidth:
                vis = vis[:,:,:,:self.maxwidth]

            if only_return:
                return vis
            else:
                return self.summ_gif(name, vis.unsqueeze(1), blacken_zeros)

    def summ_occ(self, name, occ, reduce_axes=[3], bev=False, fro=False, pro=False, frame_id=None, only_return=False):
        if self.save_this:
            B, C, D, H, W = list(occ.shape)
            if bev:
                reduce_axes = [3]
            elif fro:
                reduce_axes = [2]
            elif pro:
                reduce_axes = [4]
            for reduce_axis in reduce_axes:
                height = convert_occ_to_height(occ, reduce_axis=reduce_axis)
                if reduce_axis == reduce_axes[-1]:
                    return self.summ_oned(name=('%s_ax%d' % (name, reduce_axis)), im=height, norm=False, frame_id=frame_id, only_return=only_return)
                else:
                    self.summ_oned(name=('%s_ax%d' % (name, reduce_axis)), im=height, norm=False, frame_id=frame_id, only_return=only_return)

    def flow2color(self, flow, clip=50.0):
        """
        :param flow: Optical flow tensor.
        :return: RGB image normalized between 0 and 1.
        """

        # flow is B x C x H x W

        B, C, H, W = list(flow.size())

        flow = flow.clone().detach()

        abs_image = torch.abs(flow)
        flow_mean = abs_image.mean(dim=[1,2,3])
        flow_std = abs_image.std(dim=[1,2,3])

        if clip:
            flow = torch.clamp(flow, -clip, clip)/clip
        else:
            # Apply some kind of normalization. Divide by the perceived maximum (mean + std*2)
            flow_max = flow_mean + flow_std*2 + 1e-10
            for b in range(B):
                flow[b] = flow[b].clamp(-flow_max[b].item(), flow_max[b].item()) / flow_max[b].clamp(min=1)

        radius = torch.sqrt(torch.sum(flow**2, dim=1, keepdim=True)) #B x 1 x H x W
        radius_clipped = torch.clamp(radius, 0.0, 1.0)

        angle = torch.atan2(flow[:, 1:], flow[:, 0:1]) / np.pi #B x 1 x H x W

        hue = torch.clamp((angle + 1.0) / 2.0, 0.0, 1.0)
        saturation = torch.ones_like(hue) * 0.75
        value = radius_clipped
        hsv = torch.cat([hue, saturation, value], dim=1) #B x 3 x H x W

        #flow = tf.image.hsv_to_rgb(hsv)
        flow = hsv_to_rgb(hsv)
        flow = (flow*255.0).type(torch.ByteTensor)
        return flow

    def summ_flow(self, name, im, clip=0.0, only_return=False, frame_id=None):
        # flow is B x C x D x W
        if self.save_this:
            return self.summ_rgb(name, self.flow2color(im, clip=clip), only_return=only_return, frame_id=frame_id)
        else:
            return None

utils/vox_utils/improc.py
lus6-Jenny/RING
import numpy as np
import torch

class SimplePool():
    def __init__(self, pool_size, version='pt'):
        self.pool_size = pool_size
        self.version = version
        # random.seed(125)
        if self.pool_size > 0:
            self.num = 0
            self.items = []
        if not (version=='pt' or version=='np'):
            print('version = %s; please choose pt or np')
            assert(False) # please choose pt or np
            
    def __len__(self):
        return len(self.items)
    
    def mean(self, min_size='none'):
        if min_size=='half':
            pool_size_thresh = self.pool_size/2
        else:
            pool_size_thresh = 1
            
        if self.version=='np':
            if len(self.items) >= pool_size_thresh:
                return np.sum(self.items)/len(self.items)
            else:
                return np.nan
        if self.version=='pt':
            if len(self.items) >= pool_size_thresh:
                return torch.sum(self.items)/len(self.items)
            else:
                return torch.from_numpy(np.nan)
    
    def sample(self):
        idx = np.random.randint(len(self.items))
        return self.items[idx]
    
    def fetch(self, num=None):
        if self.version=='pt':
            item_array = torch.stack(self.items)
        elif self.version=='np':
            item_array = np.stack(self.items)
        if num is not None:
            # there better be some items
            assert(len(self.items) >= num)
                
            # if there are not that many elements just return however many there are
            if len(self.items) < num:
                return item_array
            else:
                idxs = np.random.randint(len(self.items), size=num)
                return item_array[idxs]
        else:
            return item_array
            
    def is_full(self):
        full = self.num==self.pool_size
        # print 'num = %d; full = %s' % (self.num, full)
        return full
    
    def empty(self):
        self.items = []
        self.num = 0
            
    def update(self, items):
        for item in items:
            if self.num < self.pool_size:
                # the pool is not full, so let's add this in
                self.num = self.num + 1
            else:
                # the pool is full
                # pop from the front
                self.items.pop(0)
            # add to the back
            self.items.append(item)
        return self.items

utils/vox_utils/misc.py
lus6-Jenny/RING
import numpy as np

def merge_rt(r, t):
    # r is 3 x 3
    # t is 3 or maybe 3 x 1
    t = np.reshape(t, [3, 1])
    rt = np.concatenate((r,t), axis=1)
    # rt is 3 x 4
    br = np.reshape(np.array([0,0,0,1], np.float32), [1, 4])
    # br is 1 x 4
    rt = np.concatenate((rt, br), axis=0)
    # rt is 4 x 4
    return rt

def merge_lrt(l, rt):
    # l is 3
    # rt is 4 x 4
    # merges these into a 19 vector
    D = len(l)
    assert(D==3)
    E, F = list(rt.shape)
    assert(E==4 and F==4)
    rt = rt.reshape(16)
    lrt = np.concatenate([l, rt], axis=0)
    return lrt

utils/vox_utils/py.py
lus6-Jenny/RING
import numpy as np
import torch
import torch.nn.functional as F
import utils.vox_utils.geom as geom
import utils.vox_utils.basic as basic
import matplotlib.pyplot as plt

class Vox_util(object):
    def __init__(self, Z, Y, X, scene_centroid, bounds, pad=None, assert_cube=False):
        self.XMIN, self.XMAX, self.YMIN, self.YMAX, self.ZMIN, self.ZMAX = bounds
        B, D = list(scene_centroid.shape)
        self.Z, self.Y, self.X = Z, Y, X

        scene_centroid = scene_centroid.detach().cpu().numpy()
        x_centroid, y_centroid, z_centroid = scene_centroid[0]
        self.XMIN += x_centroid
        self.XMAX += x_centroid
        self.YMIN += y_centroid
        self.YMAX += y_centroid
        self.ZMIN += z_centroid
        self.ZMAX += z_centroid

        self.default_vox_size_X = (self.XMAX-self.XMIN)/float(X)
        self.default_vox_size_Y = (self.YMAX-self.YMIN)/float(Y)
        self.default_vox_size_Z = (self.ZMAX-self.ZMIN)/float(Z)

        if pad:
            Z_pad, Y_pad, X_pad = pad
            self.ZMIN -= self.default_vox_size_Z * Z_pad
            self.ZMAX += self.default_vox_size_Z * Z_pad
            self.YMIN -= self.default_vox_size_Y * Y_pad
            self.YMAX += self.default_vox_size_Y * Y_pad
            self.XMIN -= self.default_vox_size_X * X_pad
            self.XMAX += self.default_vox_size_X * X_pad

        if assert_cube:
            # we assume cube voxels
            if (not np.isclose(self.default_vox_size_X, self.default_vox_size_Y)) or (not np.isclose(self.default_vox_size_X, self.default_vox_size_Z)):
                print('Z, Y, X', Z, Y, X)
                print('bounds for this iter:',
                      'X = %.2f to %.2f' % (self.XMIN, self.XMAX),
                      'Y = %.2f to %.2f' % (self.YMIN, self.YMAX),
                      'Z = %.2f to %.2f' % (self.ZMIN, self.ZMAX),
                )
                print('self.default_vox_size_X', self.default_vox_size_X)
                print('self.default_vox_size_Y', self.default_vox_size_Y)
                print('self.default_vox_size_Z', self.default_vox_size_Z)
            assert(np.isclose(self.default_vox_size_X, self.default_vox_size_Y))
            assert(np.isclose(self.default_vox_size_X, self.default_vox_size_Z))

    def Ref2Mem(self, xyz, Z, Y, X, assert_cube=False):
        # xyz is B x N x 3, in ref coordinates
        # transforms ref coordinates into mem coordinates
        B, N, C = list(xyz.shape)
        device = xyz.device
        assert(C==3)
        mem_T_ref = self.get_mem_T_ref(B, Z, Y, X, assert_cube=assert_cube, device=device)
        xyz = geom.apply_4x4(mem_T_ref, xyz)
        return xyz

    def Mem2Ref(self, xyz_mem, Z, Y, X, assert_cube=False):
        # xyz is B x N x 3, in mem coordinates
        # transforms mem coordinates into ref coordinates
        B, N, C = list(xyz_mem.shape)
        ref_T_mem = self.get_ref_T_mem(B, Z, Y, X, assert_cube=assert_cube, device=xyz_mem.device)
        xyz_ref = geom.apply_4x4(ref_T_mem, xyz_mem)
        return xyz_ref

    def get_mem_T_ref(self, B, Z, Y, X, assert_cube=False, device='cuda'):
        vox_size_X = (self.XMAX-self.XMIN)/float(X)
        vox_size_Y = (self.YMAX-self.YMIN)/float(Y)
        vox_size_Z = (self.ZMAX-self.ZMIN)/float(Z)

        if assert_cube:
            if (not np.isclose(vox_size_X, vox_size_Y)) or (not np.isclose(vox_size_X, vox_size_Z)):
                print('Z, Y, X', Z, Y, X)
                print('bounds for this iter:',
                      'X = %.2f to %.2f' % (self.XMIN, self.XMAX),
                      'Y = %.2f to %.2f' % (self.YMIN, self.YMAX),
                      'Z = %.2f to %.2f' % (self.ZMIN, self.ZMAX),
                )
                print('vox_size_X', vox_size_X)
                print('vox_size_Y', vox_size_Y)
                print('vox_size_Z', vox_size_Z)
            assert(np.isclose(vox_size_X, vox_size_Y))
            assert(np.isclose(vox_size_X, vox_size_Z))

        # translation
        # (this makes the left edge of the leftmost voxel correspond to XMIN)
        center_T_ref = geom.eye_4x4(B, device=device)
        center_T_ref[:,0,3] = -self.XMIN-vox_size_X/2.0
        center_T_ref[:,1,3] = -self.YMIN-vox_size_Y/2.0
        center_T_ref[:,2,3] = -self.ZMIN-vox_size_Z/2.0

        # scaling
        # (this makes the right edge of the rightmost voxel correspond to XMAX)
        mem_T_center = geom.eye_4x4(B, device=device)
        mem_T_center[:,0,0] = 1./vox_size_X
        mem_T_center[:,1,1] = 1./vox_size_Y
        mem_T_center[:,2,2] = 1./vox_size_Z
        mem_T_ref = basic.matmul2(mem_T_center, center_T_ref)

        return mem_T_ref

    def get_ref_T_mem(self, B, Z, Y, X, assert_cube=False, device='cuda'):
        mem_T_ref = self.get_mem_T_ref(B, Z, Y, X, assert_cube=assert_cube, device=device)
        # note safe_inverse is inapplicable here,
        # since the transform is nonrigid
        ref_T_mem = mem_T_ref.inverse()
        return ref_T_mem

    def get_inbounds(self, xyz, Z, Y, X, already_mem=False, padding=0.0, assert_cube=False):
        # xyz is B x N x 3
        # padding should be 0 unless you are trying to account for some later cropping
        if not already_mem:
            xyz = self.Ref2Mem(xyz, Z, Y, X, assert_cube=assert_cube)

        x = xyz[:,:,0]
        y = xyz[:,:,1]
        z = xyz[:,:,2]

        x_valid = ((x-padding)>-0.5).byte() & ((x+padding)<float(X-0.5)).byte()
        y_valid = ((y-padding)>-0.5).byte() & ((y+padding)<float(Y-0.5)).byte()
        z_valid = ((z-padding)>-0.5).byte() & ((z+padding)<float(Z-0.5)).byte()
        nonzero = (~(z==0.0)).byte()

        inbounds = x_valid & y_valid & z_valid & nonzero
        return inbounds.bool()

    def voxelize_xyz(self, xyz_ref, Z, Y, X, already_mem=False, assert_cube=False, clean_eps=0):
        B, N, D = list(xyz_ref.shape)
        assert(D==3)
        if already_mem:
            xyz_mem = xyz_ref
        else:
            xyz_mem = self.Ref2Mem(xyz_ref, Z, Y, X, assert_cube=assert_cube)
            xyz_zero = self.Ref2Mem(xyz_ref[:,0:1]*0, Z, Y, X, assert_cube=assert_cube)
        vox = self.get_occupancy(xyz_mem, Z, Y, X, clean_eps=clean_eps, xyz_zero=xyz_zero)
        return vox

    def voxelize_xyz_and_feats(self, xyz_ref, feats, Z, Y, X, already_mem=False, assert_cube=False, clean_eps=0):
        B, N, D = list(xyz_ref.shape)
        B2, N2, D2 = list(feats.shape)
        assert(D==3)
        assert(B==B2)
        assert(N==N2)
        if already_mem:
            xyz_mem = xyz_ref
        else:
            xyz_mem = self.Ref2Mem(xyz_ref, Z, Y, X, assert_cube=assert_cube)
            xyz_zero = self.Ref2Mem(xyz_ref[:,0:1]*0, Z, Y, X, assert_cube=assert_cube)
        feats = self.get_feat_occupancy(xyz_mem, feats, Z, Y, X, clean_eps=clean_eps, xyz_zero=xyz_zero)
        return feats

    def get_occupancy(self, xyz, Z, Y, X, clean_eps=0, xyz_zero=None):
        # xyz is B x N x 3 and in mem coords
        # we want to fill a voxel tensor with 1's at these inds
        B, N, C = list(xyz.shape)
        assert(C==3)

        # these papers say simple 1/0 occupancy is ok:
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Yang_PIXOR_Real-Time_3d_CVPR_2018_paper.pdf
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Luo_Fast_and_Furious_CVPR_2018_paper.pdf
        # cont fusion says they do 8-neighbor interp
        # voxelnet does occupancy but with a bit of randomness in terms of the reflectance value i think

        inbounds = self.get_inbounds(xyz, Z, Y, X, already_mem=True)
        x, y, z = xyz[:,:,0], xyz[:,:,1], xyz[:,:,2]
        mask = torch.zeros_like(x)
        mask[inbounds] = 1.0

        if xyz_zero is not None:
            # only take points that are beyond a thresh of zero
            dist = torch.norm(xyz_zero-xyz, dim=2)
            mask[dist < 0.1] = 0

        if clean_eps > 0:
            # only take points that are already near centers
            xyz_round = torch.round(xyz) # B, N, 3
            dist = torch.norm(xyz_round - xyz, dim=2)
            mask[dist > clean_eps] = 0

        # set the invalid guys to zero
        # we then need to zero out 0,0,0
        # (this method seems a bit clumsy)
        x = x*mask
        y = y*mask
        z = z*mask

        x = torch.round(x)
        y = torch.round(y)
        z = torch.round(z)
        x = torch.clamp(x, 0, X-1).int()
        y = torch.clamp(y, 0, Y-1).int()
        z = torch.clamp(z, 0, Z-1).int()

        x = x.view(B*N)
        y = y.view(B*N)
        z = z.view(B*N)

        dim3 = X
        dim2 = X * Y
        dim1 = X * Y * Z

        base = torch.arange(0, B, dtype=torch.int32, device=xyz.device)*dim1
        base = torch.reshape(base, [B, 1]).repeat([1, N]).view(B*N)

        vox_inds = base + z * dim2 + y * dim3 + x
        voxels = torch.zeros(B*Z*Y*X, device=xyz.device).float()
        voxels[vox_inds.long()] = 1.0
        # zero out the singularity
        voxels[base.long()] = 0.0
        voxels = voxels.reshape(B, 1, Z, Y, X)
        # B x 1 x Z x Y x X
        return voxels

    def get_feat_occupancy(self, xyz, feat, Z, Y, X, clean_eps=0, xyz_zero=None):
        # xyz is B x N x 3 and in mem coords
        # feat is B x N x D
        # we want to fill a voxel tensor with 1's at these inds
        B, N, C = list(xyz.shape)
        B2, N2, D2 = list(feat.shape)
        assert(C==3)
        assert(B==B2)
        assert(N==N2)

        # these papers say simple 1/0 occupancy is ok:
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Yang_PIXOR_Real-Time_3d_CVPR_2018_paper.pdf
        #  http://openaccess.thecvf.com/content_cvpr_2018/papers/Luo_Fast_and_Furious_CVPR_2018_paper.pdf
        # cont fusion says they do 8-neighbor interp
        # voxelnet does occupancy but with a bit of randomness in terms of the reflectance value i think

        inbounds = self.get_inbounds(xyz, Z, Y, X, already_mem=True)
        x, y, z = xyz[:,:,0], xyz[:,:,1], xyz[:,:,2]
        mask = torch.zeros_like(x)
        mask[inbounds] = 1.0

        if xyz_zero is not None:
            # only take points that are beyond a thresh of zero
            dist = torch.norm(xyz_zero-xyz, dim=2)
            mask[dist < 0.1] = 0

        if clean_eps > 0:
            # only take points that are already near centers
            xyz_round = torch.round(xyz) # B, N, 3
            dist = torch.norm(xyz_round - xyz, dim=2)
            mask[dist > clean_eps] = 0

        # set the invalid guys to zero
        # we then need to zero out 0,0,0
        # (this method seems a bit clumsy)
        x = x*mask # B, N
        y = y*mask
        z = z*mask
        feat = feat*mask.unsqueeze(-1) # B, N, D

        x = torch.round(x)
        y = torch.round(y)
        z = torch.round(z)
        x = torch.clamp(x, 0, X-1).int()
        y = torch.clamp(y, 0, Y-1).int()
        z = torch.clamp(z, 0, Z-1).int()

        # permute point orders
        perm = torch.randperm(N)
        x = x[:, perm]
        y = y[:, perm]
        z = z[:, perm]
        feat = feat[:, perm]

        x = x.view(B*N)
        y = y.view(B*N)
        z = z.view(B*N)
        feat = feat.view(B*N, -1)

        dim3 = X
        dim2 = X * Y
        dim1 = X * Y * Z

        base = torch.arange(0, B, dtype=torch.int32, device=xyz.device)*dim1
        base = torch.reshape(base, [B, 1]).repeat([1, N]).view(B*N)

        vox_inds = base + z * dim2 + y * dim3 + x
        feat_voxels = torch.zeros((B*Z*Y*X, D2), device=xyz.device).float()
        feat_voxels[vox_inds.long()] = feat
        # zero out the singularity
        feat_voxels[base.long()] = 0.0
        feat_voxels = feat_voxels.reshape(B, Z, Y, X, D2).permute(0, 4, 1, 2, 3)
        # B x C x Z x Y x X
        return feat_voxels

    def unproject_image_to_mem(self, rgb_camB, pixB_T_camA, camB_T_camA, Z, Y, X, assert_cube=False, xyz_camA=None):
        # rgb_camB is B x C x H x W
        # pixB_T_camA is B x 4 x 4

        # rgb lives in B pixel coords
        # we want everything in A memory coords

        # this puts each C-dim pixel in the rgb_camB
        # along a ray in the voxelgrid
        B, C, H, W = list(rgb_camB.shape)

        if xyz_camA is None:
            xyz_memA = basic.gridcloud3d(B, Z, Y, X, norm=False, device=pixB_T_camA.device)
            xyz_camA = self.Mem2Ref(xyz_memA, Z, Y, X, assert_cube=assert_cube)

        xyz_camB = geom.apply_4x4(camB_T_camA, xyz_camA)
        z = xyz_camB[:,:,2]

        xyz_pixB = geom.apply_4x4(pixB_T_camA, xyz_camA)
        normalizer = torch.unsqueeze(xyz_pixB[:,:,2], 2)

        EPS=1e-6
        # z = xyz_pixB[:,:,2]
        xy_pixB = xyz_pixB[:,:,:2]/torch.clamp(normalizer, min=EPS)
        # this is B x N x 2
        # this is the (floating point) pixel coordinate of each voxel
        x, y = xy_pixB[:,:,0], xy_pixB[:,:,1]
        # these are B x N

        # for i in range(5):
        #     canvas = np.ones((240, 360, 1))
        #     plt.figure(1)
        #     plt.scatter(x[i].clone().detach().cpu().numpy(), y[i].clone().detach().cpu().numpy(), c=np.ones_like(x[i].clone().cpu().numpy())-1, cmap='jet', alpha=0.1, s=1)
        #     plt.xlim(0, 240)
        #     plt.ylim(0, 360)
        #     plt.show()
        #     plt.savefig("/mnt/workspace/workgroup/xxc/proj"+str(i)+".jpg")


        x_valid = (x>-0.5).bool() & (x<float(W-0.5)).bool()
        y_valid = (y>-0.5).bool() & (y<float(H-0.5)).bool()
        z_valid = (z>-0.0).bool()
        valid_mem = (x_valid & y_valid & z_valid).reshape(B, 1, Z, Y, X).float()

        if (0):
            # handwritten version
            values = torch.zeros([B, C, Z*Y*X], dtype=torch.float32)
            for b in list(range(B)):
                values[b] = samp.bilinear_sample_single(rgb_camB[b], x_pixB[b], y_pixB[b])
        else:
            # native pytorch version
            y_pixB, x_pixB = basic.normalize_grid2d(y, x, H, W)
            # since we want a 3d output, we need 5d tensors
            z_pixB = torch.zeros_like(x)
            xyz_pixB = torch.stack([x_pixB, y_pixB, z_pixB], axis=2)
            rgb_camB = rgb_camB.unsqueeze(2)
            xyz_pixB = torch.reshape(xyz_pixB, [B, Z, Y, X, 3])
            values = F.grid_sample(rgb_camB, xyz_pixB, align_corners=False)

        values = torch.reshape(values, (B, C, Z, Y, X))
        values = values  * valid_mem
        return values

    def warp_tiled_to_mem(self, rgb_tileB, pixB_T_camA, camB_T_camA, Z, Y, X, DMIN, DMAX, assert_cube=False):
        # rgb_tileB is B,C,D,H,W
        # pixB_T_camA is B,4,4
        # camB_T_camA is B,4,4

        # rgb_tileB lives in B pixel coords but it has been tiled across the Z dimension
        # we want everything in A memory coords

        # this resamples the so that each C-dim pixel in rgb_tilB
        # is put into its correct place in the voxelgrid
        # (using the pinhole camera model)
        
        B, C, D, H, W = list(rgb_tileB.shape)

        xyz_memA = basic.gridcloud3d(B, Z, Y, X, norm=False, device=pixB_T_camA.device)

        xyz_camA = self.Mem2Ref(xyz_memA, Z, Y, X, assert_cube=assert_cube)

        xyz_camB = geom.apply_4x4(camB_T_camA, xyz_camA)
        z_camB = xyz_camB[:,:,2]

        # rgb_tileB has depth=DMIN in tile 0, and depth=DMAX in tile D-1
        z_tileB = (D-1.0) * (z_camB-float(DMIN)) / float(DMAX-DMIN)

        xyz_pixB = geom.apply_4x4(pixB_T_camA, xyz_camA)
        normalizer = torch.unsqueeze(xyz_pixB[:,:,2], 2)
        EPS=1e-6
        # z = xyz_pixB[:,:,2]
        xy_pixB = xyz_pixB[:,:,:2]/torch.clamp(normalizer, min=EPS)
        # this is B x N x 2
        # this is the (floating point) pixel coordinate of each voxel
        x, y = xy_pixB[:,:,0], xy_pixB[:,:,1]
        # these are B x N

        x_valid = (x>-0.5).bool() & (x<float(W-0.5)).bool()
        y_valid = (y>-0.5).bool() & (y<float(H-0.5)).bool()
        z_valid = (z_camB>0.0).bool()
        valid_mem = (x_valid & y_valid & z_valid).reshape(B, 1, Z, Y, X).float()

        z_tileB, y_pixB, x_pixB = basic.normalize_grid3d(z_tileB, y, x, D, H, W)
        xyz_pixB = torch.stack([x_pixB, y_pixB, z_tileB], axis=2)
        xyz_pixB = torch.reshape(xyz_pixB, [B, Z, Y, X, 3])
        values = F.grid_sample(rgb_tileB, xyz_pixB, align_corners=False)

        values = torch.reshape(values, (B, C, Z, Y, X))
        values = values * valid_mem
        return values
    

    def apply_mem_T_ref_to_lrtlist(self, lrtlist_cam, Z, Y, X, assert_cube=False):
        # lrtlist is B x N x 19, in cam coordinates
        # transforms them into mem coordinates, including a scale change for the lengths
        B, N, C = list(lrtlist_cam.shape)
        assert(C==19)
        mem_T_cam = self.get_mem_T_ref(B, Z, Y, X, assert_cube=assert_cube, device=lrtlist_cam.device)

    def xyz2circles(self, xyz, radius, Z, Y, X, soft=True, already_mem=True, also_offset=False, grid=None):
        # xyz is B x N x 3
        # radius is B x N or broadcastably so
        # output is B x N x Z x Y x X
        B, N, D = list(xyz.shape)
        assert(D==3)
        if not already_mem:
            xyz = self.Ref2Mem(xyz, Z, Y, X)

        if grid is None:
            grid_z, grid_y, grid_x = basic.meshgrid3d(B, Z, Y, X, stack=False, norm=False, device=xyz.device)
            # note the default stack is on -1
            grid = torch.stack([grid_x, grid_y, grid_z], dim=1)
            # this is B x 3 x Z x Y x X
            
        xyz = xyz.reshape(B, N, 3, 1, 1, 1)
        grid = grid.reshape(B, 1, 3, Z, Y, X)
        # this is B x N x Z x Y x X

        # round the xyzs, so that at least one value matches the grid perfectly,
        # and we get a value of 1 there (since exp(0)==1)
        xyz = xyz.round()

        if torch.is_tensor(radius):
            radius = radius.clamp(min=0.01)

        if soft:
            off = grid - xyz # B,N,3,Z,Y,X
            # interpret radius as sigma
            dist_grid = torch.sum(off**2, dim=2, keepdim=False)
            # this is B x N x Z x Y x X
            if torch.is_tensor(radius):
                radius = radius.reshape(B, N, 1, 1, 1)
            mask = torch.exp(-dist_grid/(2*radius*radius))
            # zero out near zero
            mask[mask < 0.001] = 0.0
            # h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
            # h[h < np.finfo(h.dtype).eps * h.max()] = 0
            # return h
            if also_offset:
                return mask, off
            else:
                return mask
        else:
            assert(False) # something is wrong with this. come back later to debug

            dist_grid = torch.norm(grid - xyz, dim=2, keepdim=False)
            # this is 0 at/near the xyz, and increases by 1 for each voxel away

            radius = radius.reshape(B, N, 1, 1, 1)

            within_radius_mask = (dist_grid < radius).float()
            within_radius_mask = torch.sum(within_radius_mask, dim=1, keepdim=True).clamp(0, 1)
            return within_radius_mask

    def xyz2circles_bev(self, xyz, radius, Z, Y, X, already_mem=True, also_offset=False):
        # xyz is B x N x 3
        # radius is B x N or broadcastably so
        # output is B x N x Z x Y x X
        B, N, D = list(xyz.shape)
        assert(D==3)
        if not already_mem:
            xyz = self.Ref2Mem(xyz, Z, Y, X)

        xz = torch.stack([xyz[:,:,0], xyz[:,:,2]], dim=2)

        grid_z, grid_x = basic.meshgrid2d(B, Z, X, stack=False, norm=False, device=xyz.device)
        # note the default stack is on -1
        grid = torch.stack([grid_x, grid_z], dim=1)
        # this is B x 2 x Z x X

        xz = xz.reshape(B, N, 2, 1, 1)
        grid = grid.reshape(B, 1, 2, Z, X)
        # these are ready to broadcast to B x N x Z x X

        # round the points, so that at least one value matches the grid perfectly,
        # and we get a value of 1 there (since exp(0)==1)
        xz = xz.round()

        if torch.is_tensor(radius):
            radius = radius.clamp(min=0.01)

        off = grid - xz # B,N,2,Z,X
        # interpret radius as sigma
        dist_grid = torch.sum(off**2, dim=2, keepdim=False)
        # this is B x N x Z x X
        if torch.is_tensor(radius):
            radius = radius.reshape(B, N, 1, 1, 1)
        mask = torch.exp(-dist_grid/(2*radius*radius))
        # zero out near zero
        mask[mask < 0.001] = 0.0
        
        # add a Y dim
        mask = mask.unsqueeze(-2)
        off = off.unsqueeze(-2)
        # # B,N,2,Z,1,X
        
        if also_offset:
            return mask, off
        else:
            return mask
        

utils/vox_utils/vox.py
lus6-Jenny/RING