I have a scripts that can generate slam using rgb video and depth estimation. now, i can create linear maps on straight motion on a corridor. but whatever reason, despite the footage camera just going straight on corridor no up or down movement, as you can see in the camera trajectory, the camera movement goes down like -25 linearly. also when there is a turn on the video, slam process cant handle it and after the turn the map is nonsense. also when there is a up or down motion on the camera despite the camera is came back to the original position before movement, the mapping is ridiculous and for example if the camera moves upwards the map starts to be created there. Also point clouds are overlapped and a smooth transition is not provided. Actually I need to see an almost rectangular corridor scene but the connections between the frames are quite weak. Find solutions for all these problems and give me code suggestions. Think as much as you can.

Camera Trajectory:

X,Y,Z
1.335631506522596545e-01,3.909969910141680022e-02,-9.959320549241844622e-01
2.761117046358165794e-01,1.176601407773113228e-01,-1.980384714111716793e+00
4.281914740062272573e-01,1.780915967998213667e-01,-2.968405250036573850e+00
5.586917653470301337e-01,2.416417789205533295e-01,-3.979270763077455086e+00
7.016325554856915891e-01,3.175934404215720042e-01,-4.957247252016697914e+00
8.446621575753525679e-01,3.600348655048209445e-01,-5.957281790773906849e+00
9.702873497041186734e-01,4.284337848014099226e-01,-6.955621319910493483e+00
1.117943145340354105e+00,4.998321046638229959e-01,-7.917363007725041868e+00
1.205634397748885034e+00,5.741987592677019192e-01,-8.872652812414941081e+00
1.356548769800091980e+00,6.588043250803754525e-01,-9.871264850139315428e+00
1.508230938456357473e+00,7.916599306272588121e-01,-1.089008352130205637e+01
1.648874653280515412e+00,8.615478862857782749e-01,-1.186057918295001734e+01
1.818987532489576342e+00,9.448405968359098184e-01,-1.282653058551184699e+01
1.941855145696292473e+00,1.026555043892105612e+00,-1.381544676244391390e+01
2.102781730244725367e+00,1.111608871350449945e+00,-1.482044789382075045e+01
2.251254445669872872e+00,1.204879018976154237e+00,-1.584449606741855199e+01
2.386185449614589338e+00,1.282392421340864397e+00,-1.681277929402066462e+01
2.486260018165475483e+00,1.324215770780428514e+00,-1.776956774430320607e+01
2.599885009170571415e+00,1.409341839209690139e+00,-1.875779861052231112e+01
2.724745446558732187e+00,1.477661842153065841e+00,-1.974403402743631020e+01
2.863257056085137808e+00,1.529056013067355169e+00,-2.074313367447696876e+01
2.980636758614750192e+00,1.606902883343632737e+00,-2.170657974430658399e+01
3.167865365888362206e+00,1.677330936066333988e+00,-2.268562409987585937e+01
3.275303069861895455e+00,1.744036513807038613e+00,-2.368702804027229547e+01
3.399178907080772749e+00,1.821181215379479523e+00,-2.468028479542448750e+01
3.531284374779977586e+00,1.896387140784229208e+00,-2.562966536826266406e+01
3.647020334667235630e+00,1.976675213637114181e+00,-2.663157852693044703e+01
3.776872918148584279e+00,2.059823544574268528e+00,-2.762187645231425748e+01
3.940807000104240121e+00,2.132621486128423172e+00,-2.863622066309268632e+01


slam.py

import os
import sys
import cv2
import numpy as np
import open3d as o3d
import argparse
import threading
import time
import matplotlib.pyplot as plt

# Add the current directory to the Python path
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

# Add the 'zoe' and 'Metric3D' directories to the Python path
zoe_dir = os.path.join(current_dir, 'zoe')
metric3d_dir = os.path.join(current_dir, 'metric3dd')
sys.path.append(zoe_dir)
sys.path.append(metric3d_dir)

# Import point cloud generators
from point_cloud_generator import PointCloudGenerator as MidasPointCloudGenerator
from zoe.point_cloud_generator import PointCloudGenerator as ZoePointCloudGenerator
from metric3dd.point_cloud_generator import PointCloudGenerator as Metric3DPointCloudGenerator
from pose_estimator import PoseEstimator

def register_point_clouds(source_pcd, target_pcd):
    icp_threshold = 0.02
    trans_init = np.eye(4)
    reg_p2p = o3d.pipelines.registration.registration_icp(
        source_pcd, target_pcd, icp_threshold, trans_init,
        o3d.pipelines.registration.TransformationEstimationPointToPoint())
    return reg_p2p.transformation

class SLAMSystem:
    def __init__(self, video_path, calibration_file, real_time_vis=False, model='zoe', save_point_clouds=False, frame_interval=1):
        self.calibration_file = calibration_file
        self.camera_matrix, self.dist_coeffs = self.load_camera_parameters()
        self.frame_interval = frame_interval
        
        if model == 'midas':
            self.generator = MidasPointCloudGenerator(calibration_path=calibration_file)
            self.pcd_folder = 'pcd_midas'
        elif model == 'metric3d':
            self.generator = Metric3DPointCloudGenerator(calibration_path=calibration_file)
            self.pcd_folder = 'pcd_metric'
        else:  # default to zoe
            self.generator = ZoePointCloudGenerator(calibration_path=calibration_file)
            self.pcd_folder = 'pcd_zoe'
        
        focal_length = (self.camera_matrix[0, 0] + self.camera_matrix[1, 1]) / 2
        principal_point = (self.camera_matrix[0, 2], self.camera_matrix[1, 2])
        self.estimator = PoseEstimator(focal_length, principal_point)
        
        self.cap = cv2.VideoCapture(video_path)
        self.global_map = o3d.geometry.PointCloud()
        self.prev_frame = None
        self.pose = np.eye(4)
        self.real_time_vis = real_time_vis
        self.save_point_clouds = save_point_clouds
        self.vis = None
        self.is_running = True
        self.camera_positions = []
        self.vis_updated = threading.Event()
        if self.real_time_vis:
            self.vis = o3d.visualization.Visualizer()
            self.vis.create_window()
            self.vis.add_geometry(self.global_map)
            self.vis_thread = threading.Thread(target=self.visualization_loop)
            self.vis_thread.start()

    def load_camera_parameters(self):
        try:
            data = np.load(self.calibration_file)
            camera_matrix = data['Camera_matrix']
            dist_coeffs = data['distCoeff'].ravel()
            print(f"Camera matrix:\n{camera_matrix}")
            print(f"Distortion coefficients: {dist_coeffs}")
            return camera_matrix, dist_coeffs
        except Exception as e:
            print(f"Error loading camera parameters: {str(e)}")
            print("Using default camera parameters")
            return np.array([[700, 0, 320], [0, 700, 240], [0, 0, 1]]), np.zeros(5)

    def calculate_scale(self, depth_prev, depth_curr):
        depth_prev_nonzero = depth_prev[depth_prev > 0]
        depth_curr_nonzero = depth_curr[depth_curr > 0]

        if len(depth_prev_nonzero) == 0 or len(depth_curr_nonzero) == 0:
            print("Insufficient depth information for scale calculation. Using default scale.")
            return 1.0

        # Use percentile-based scaling
        prev_depth = np.percentile(depth_prev_nonzero, 50)
        curr_depth = np.percentile(depth_curr_nonzero, 50)

        if curr_depth == 0:
            print("Current depth is zero, cannot calculate scale.")
            return 1.0

        scale = prev_depth / curr_depth

        # Limit scale factor to prevent extreme values
        max_scale = 1.1
        min_scale = 0.9
        scale = np.clip(scale, min_scale, max_scale)

        print(f"Calculated scale factor: {scale}")
        return scale

    def update_visualization(self):
        if self.real_time_vis:
            self.vis.update_geometry(self.global_map)
            self.vis.poll_events()
            self.vis.update_renderer()
            self.vis_updated.set()

    def visualization_loop(self):
        while self.is_running:
            self.vis_updated.wait()
            self.update_visualization()
            self.vis_updated.clear()

    def run(self):
        frame_count = 0
        depth_prev = None

        if self.save_point_clouds:
            os.makedirs(self.pcd_folder, exist_ok=True)

        try:
            while True:
                ret, frame = self.cap.read()
                if not ret:
                    print("End of video stream.")
                    break

                if frame_count % self.frame_interval == 0:
                    print(f"Processing frame {frame_count}")

                    point_cloud_np, depth_map = self.generator.process_frame_with_depth(frame)

                    if self.save_point_clouds:
                        pcd_filename = os.path.join(self.pcd_folder, f"frame_{frame_count:04d}.txt")
                        self.generator.save_point_cloud_txt(point_cloud_np, pcd_filename)
                        print(f"Saved point cloud to {pcd_filename}")
                    else:
                        current_pcd = o3d.geometry.PointCloud()
                        current_pcd.points = o3d.utility.Vector3dVector(point_cloud_np[:, :3])
                        current_pcd.colors = o3d.utility.Vector3dVector(point_cloud_np[:, 3:] / 255.0)

                        scale = None

                        if self.prev_frame is not None and depth_prev is not None:
                            gray_prev = cv2.cvtColor(self.prev_frame, cv2.COLOR_BGR2GRAY)
                            gray_current = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                            kp1, desc1 = self.estimator.detect_and_compute(gray_prev)
                            kp2, desc2 = self.estimator.detect_and_compute(gray_current)

                            matches = self.estimator.match_features(desc1, desc2)
                            if len(matches) == 0:
                                print(f"Frame {frame_count}: No valid matches found.")
                            else:
                                scale = self.calculate_scale(depth_prev, depth_map)
                                R, t = self.estimator.estimate_pose(kp1, kp2, matches, scale=scale)
                                if R is not None and t is not None:
                                    T = np.eye(4)
                                    T[:3, :3] = R
                                    T[:3, 3] = t.squeeze()
                                    self.pose = self.pose @ T
                                    print(f"Updated pose:\n{self.pose}")
                                    # Store the camera position
                                    camera_position = self.pose[:3, 3]
                                    self.camera_positions.append(camera_position)
                                else:
                                    print(f"Frame {frame_count}: Pose estimation failed.")

                        current_pcd.transform(self.pose)
                        
                        self.global_map += current_pcd
                        self.global_map = self.global_map.voxel_down_sample(voxel_size=0.05)

                        if self.real_time_vis:
                            self.vis_updated.set()

                        self.prev_frame = frame.copy()
                        depth_prev = depth_map.copy()

                frame_count += 1

        except Exception as e:
            print(f"An error occurred during SLAM processing: {str(e)}")

        finally:
            self.cap.release()

            if not self.save_point_clouds:
                o3d.io.write_point_cloud("global_map.pcd", self.global_map)
                print("Saved global map to 'global_map.pcd'.")

                if self.real_time_vis:
                    self.is_running = False
                    self.vis_thread.join()
                    self.vis.destroy_window()
                else:
                    o3d.visualization.draw_geometries([self.global_map])

            if len(self.camera_positions) > 0:
                camera_positions = np.array(self.camera_positions)
                fig = plt.figure()
                ax = fig.add_subplot(111, projection='3d')
                ax.plot(camera_positions[:, 0], camera_positions[:, 1], camera_positions[:, 2], 'r*-')
                ax.set_xlabel('X')
                ax.set_ylabel('Y')
                ax.set_zlabel('Z')
                ax.set_title('Camera Trajectory')
                plt.show()

                # Save camera trajectory to file
                np.savetxt("camera_trajectory.txt", camera_positions, delimiter=",", header="X,Y,Z", comments="")
                print("Saved camera trajectory to 'camera_trajectory.txt'.")

def main():
    parser = argparse.ArgumentParser(description="Run SLAM system with optional real-time visualization or save point clouds")
    parser.add_argument("--video_path", type=str, default="videos/kutupkisa.mp4", help="Path to input video")
    parser.add_argument("--calibration_file", type=str, required=True, help="Path to camera calibration NPZ file")
    parser.add_argument("--real_time_vis", action="store_true", help="Enable real-time visualization")
    parser.add_argument("--model", type=str, choices=['midas', 'zoe', 'metric3d'], default='zoe', help="Choose the point cloud generation model")
    parser.add_argument("--save_point_clouds", action="store_true", help="Save point clouds as txt files without running SLAM")
    parser.add_argument("--frame_interval", type=int, default=1, help="Process every Nth frame (default: 1)")
    args = parser.parse_args()

    slam = SLAMSystem(args.video_path, args.calibration_file, args.real_time_vis, args.model, args.save_point_clouds, args.frame_interval)
    slam.run()

if __name__ == "__main__":
    main()

point_cloud_generator.py

import cv2
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
import open3d as o3d

class PointCloudGenerator:
    def __init__(self, calibration_path='zoe/CalibrationMatrix_college_cpt.npz', transformation=None, voxel_size=0.02):
        # Metric3D setup
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.metric3d_model = torch.hub.load('yvanyin/metric3d', 'metric3d_vit_large', pretrain=True)
        self.metric3d_model.to(self.device)
        self.metric3d_model.eval()

        # Load camera parameters from NPZ file
        if calibration_path:
            self.load_camera_parameters(calibration_path)
        else:
            self.set_default_camera_parameters()

        # Compute Q matrix
        self.compute_Q_matrix()

        # Rotation matrix (keeping this as it was in the original code)
        self.R = np.array([
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 0]
        ])

        # Transformation matrix
        if transformation is None:
            self.transformation = np.array([
            [1, 0, 0, 0],
            [0, 1, 0, 0],  # Flipping Y-axis
            [0, 0, 1, 0],  # Flipping Z-axis (may cause issues)
            [0, 0, 0, 0]
        ])
        else:
            self.transformation = transformation

        self.voxel_size = voxel_size

    def set_default_camera_parameters(self):
        self.FX = 350.0  # Default focal length
        self.FY = 350.0
        self.CX = 160.0  # Default principal point
        self.CY = 120.0
        self.dist_coeffs = np.zeros(5)
        print("Using default camera parameters")

    def load_camera_parameters(self, calibration_path):
        try:
            calibration_data = np.load(calibration_path)
            if 'Camera_matrix' in calibration_data:
                camera_matrix = calibration_data['Camera_matrix']
                self.FX = camera_matrix[0, 0]
                self.FY = camera_matrix[1, 1]
                self.CX = camera_matrix[0, 2]
                self.CY = camera_matrix[1, 2]
                print(f"Camera parameters loaded: FX: {self.FX}, FY: {self.FY}, CX: {self.CX}, CY: {self.CY}")
            else:
                raise KeyError("Camera matrix not found in calibration file")

            if 'distCoeff' in calibration_data:
                self.dist_coeffs = calibration_data['distCoeff'].ravel()
                print(f"Distortion coefficients loaded: {self.dist_coeffs}")
            else:
                print("Distortion coefficients not found, assuming no distortion")
                self.dist_coeffs = np.zeros(5)

        except Exception as e:
            print(f"Error loading camera parameters: {e}")
            self.set_default_camera_parameters()

    def compute_Q_matrix(self):
        self.Q = np.array([[1.0, 0.0, 0.0, -self.CX],
                           [0.0, 1.0, 0.0, -self.CY],
                           [0.0, 0.0, 0.0, self.FX],
                           [0.0, 0.0, 1.0/self.FY, 0.0]], dtype=np.float32)

    def process_frame_with_depth(self, frame):
        # Undistort the frame
        h, w = frame.shape[:2]
        new_camera_matrix, _ = cv2.getOptimalNewCameraMatrix(
            np.array([[self.FX, 0, self.CX], [0, self.FY, self.CY], [0, 0, 1]]), 
            self.dist_coeffs, (w, h), 1, (w, h)
        )
        undistorted_frame = cv2.undistort(frame, 
                                          np.array([[self.FX, 0, self.CX], [0, self.FY, self.CY], [0, 0, 1]]), 
                                          self.dist_coeffs, 
                                          None, 
                                          new_camera_matrix)

        # RGB to Depth
        color_image = cv2.cvtColor(undistorted_frame, cv2.COLOR_BGR2RGB)
        image_tensor = transforms.ToTensor()(color_image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            pred_depth, confidence, _ = self.metric3d_model.inference({'input': image_tensor})
            depth_map = pred_depth.squeeze().cpu().numpy()

        # Ensure depth_map matches the original frame size
        depth_map = cv2.resize(depth_map, (w, h), interpolation=cv2.INTER_LINEAR)

        # Scale depth map to real-world units (e.g., meters)
        depth_scale = 1.0  # Adjust this value based on your depth model's output scale
        scaled_depth_map = depth_map * depth_scale

        # Create 3D points
        rows, cols = scaled_depth_map.shape
        c, r = np.meshgrid(np.arange(cols), np.arange(rows))
        points = np.zeros((rows * cols, 3), dtype=np.float32)
        points[:, 0] = (c.flatten() - self.CX) * scaled_depth_map.flatten() / self.FX
        points[:, 1] = (r.flatten() - self.CY) * scaled_depth_map.flatten() / self.FY
        points[:, 2] = scaled_depth_map.flatten()

        # Improved filtering
        min_depth = 0.1  # Adjust based on your scene
        max_depth = 5.0  # Adjust based on your scene
        min_brightness = 10  # Adjust to filter out very dark points

        mask = (scaled_depth_map.flatten() > min_depth) & \
               (scaled_depth_map.flatten() < max_depth) & \
               (np.mean(color_image.reshape(-1, 3), axis=1) > min_brightness)

        valid_points = points[mask]
        valid_colors = color_image.reshape(-1, 3)[mask]

        # Apply transformations
        #transformed_points = np.dot(valid_points, self.R.T)
        homogeneous_points = np.hstack((valid_points, np.ones((valid_points.shape[0], 1))))
        transformed_points = np.dot(homogeneous_points, self.transformation.T)[:, :3]
        # Downsample
        if self.voxel_size > 0:
            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(transformed_points)
            pcd.colors = o3d.utility.Vector3dVector(valid_colors.astype(np.float64) / 255.0)
            downsampled_pcd = pcd.voxel_down_sample(voxel_size=self.voxel_size)
            transformed_points = np.asarray(downsampled_pcd.points)
            valid_colors = np.asarray(downsampled_pcd.colors) * 255

        # Combine points and colors for final point cloud
        point_cloud = np.hstack((transformed_points, valid_colors.astype(np.uint8)))

        return point_cloud, depth_map


    def save_point_cloud_txt(self, point_cloud, filename):
        np.savetxt(filename, point_cloud, fmt='%.6f %.6f %.6f %d %d %d')

def save_point_cloud_pcd(point_cloud, filename):
    with open(filename, 'w') as f:
        f.write("# .PCD v0.7 - Point Cloud Data file format\n")
        f.write("VERSION 0.7\n")
        f.write("FIELDS x y z rgb\n")
        f.write("SIZE 4 4 4 4\n")
        f.write("TYPE F F F U\n")
        f.write("COUNT 1 1 1 1\n")
        f.write(f"WIDTH {len(point_cloud)}\n")
        f.write("HEIGHT 1\n")
        f.write("VIEWPOINT 0 0 0 1 0 0 0 0\n")
        f.write(f"POINTS {len(point_cloud)}\n")
        f.write("DATA ascii\n")
        for point in point_cloud:
            x, y, z, r, g, b = point
            rgb = int(r) << 16 | int(g) << 8 | int(b)
            f.write(f"{x} {y} {z} {rgb}\n")

pose_estimator.py

import cv2
import numpy as np

class PoseEstimator:
    def __init__(self, focal_length, principal_point):
        self.orb = cv2.ORB_create(5000)
        self.matcher = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)
        self.K = np.array([
            [focal_length, 0, principal_point[0]],
            [0, focal_length, principal_point[1]],
            [0, 0, 1]
        ])

    def detect_and_compute(self, img):
        keypoints, descriptors = self.orb.detectAndCompute(img, None)
        return keypoints, descriptors

    def match_features(self, desc1, desc2):
        if desc1 is None or desc2 is None:
            print("One of the descriptors is None.")
            return []

        # Use KNN matcher with k=2
        matcher = cv2.BFMatcher(cv2.NORM_HAMMING)
        knn_matches = matcher.knnMatch(desc1, desc2, k=2)
        
        # Apply ratio test
        good_matches = []
        for m, n in knn_matches:
            if m.distance < 0.999 * n.distance:
                good_matches.append(m)
                
        return good_matches

    def estimate_pose(self, kp1, kp2, matches, scale=1.0):
        if len(matches) < 8:
            print("Not enough matches to compute pose.")
            return None, None

        try:
            pts1 = np.float32([kp1[m.queryIdx].pt for m in matches])
            pts2 = np.float32([kp2[m.trainIdx].pt for m in matches])
            
            E, mask = cv2.findEssentialMat(pts1, pts2, self.K, method=cv2.RANSAC, prob=0.999, threshold=1.0)
            if E is None:
                print("Essential matrix computation failed.")
                return None, None
            
            _, R, t, mask_pose = cv2.recoverPose(E, pts1, pts2, self.K)

            # Apply scale to the translation vector
            t = t * scale
            print(f"Applied scale factor: {scale}")
            print(f"Estimated rotation R:\n{R}")
            print(f"Estimated translation t:\n{t.T}")
            return R, t
        except Exception as e:
            print(f"Exception during pose estimation: {e}")
            return None, None