I’m working on a SLAM system that processes every Nth frame from a video, where N is controlled by a variable called frame_interval. The system uses point clouds generated by various models (Zoe, Midas, or Metric3D) and performs pose estimation and keyframe selection based on frame intervals, translation, and rotation between poses. However, I'm encountering an issue with keyframe selection when using larger frame_interval values.

Problem:

For example, when I set frame_interval=30, I expect to process frames 0, 30, 60, etc. However, when processing the 150th frame, the keyframe selected might be as early as frame 11, which doesn’t seem appropriate given the large frame interval. I suspect the issue is related to how keyframe selection is based on both frame intervals and distance/pose criteria, and it's not accounting for the larger intervals correctly.

Keyframe selection is done using the following logic:

def is_keyframe(self, frame_index, current_pose):
    if not self.keyframes:
        return True  # Always select the first frame as a keyframe

    frames_since_last_keyframe = frame_index - self.last_keyframe_index
    if frames_since_last_keyframe < self.min_keyframe_distance:
        return False

    last_keyframe = self.keyframes[-1]
    last_pose = last_keyframe['pose']

    # Calculate translation and rotation between current pose and last keyframe
    translation = np.linalg.norm(current_pose[:3, 3] - last_pose[:3, 3])
    rotation = np.degrees(np.arccos((np.trace(current_pose[:3, :3] @ last_pose[:3, :3].T) - 1) / 2))

    return translation > self.min_translation or rotation > self.min_rotation


I have a few parameters that control the keyframe selection:

min_keyframe_distance = 10 (minimum number of processed frames between keyframes)
min_translation = 2 meters (minimum translation required to select a new keyframe)
min_rotation = 20 degrees (minimum rotation required for a new keyframe)
My SLAM system also includes pose estimation between frames, point cloud generation, and visualization. However, this issue with the keyframe selection logic is causing problems when I increase the frame interval, which suggests that the logic for choosing keyframes may not be scaling well with different frame intervals.

Key Considerations:
Frame Interval Handling:

When using a larger frame interval (e.g., frame_interval=30), the system should still choose keyframes that are appropriately spaced out, but instead, it sometimes selects keyframes from early frames (e.g., frame 11 while processing frame 150). I suspect that the condition based on frames_since_last_keyframe and min_keyframe_distance doesn’t account for varying frame_interval values properly.

Additionally, when the interval is larger, the system should adjust the criteria for selecting keyframes to ensure that keyframes are not too close together in terms of frame numbers.

Translation and Rotation Thresholds:

The translation and rotation thresholds (min_translation and min_rotation) may not scale well with larger frame intervals. With larger intervals, the difference in motion between frames is naturally larger, so the thresholds might need dynamic adjustment to ensure keyframes are chosen appropriately.
Keyframe Addition:

In the method add_keyframe, I add keyframes to the list without further checks about the frame interval. I suspect there might be an additional check needed here to prevent adding keyframes that are spaced too far apart or too close together based on frame intervals.
Pose Estimation with Different Frame Rates:

With larger frame intervals, the system may skip important pose updates or select poses based on sparse data. The current pose estimation is based on ORB feature matching, but when using larger frame intervals, the matches may become sparse, affecting the quality of pose estimates.
Visualization Issues:

The real-time visualization (update_visualization) might not reflect the changes appropriately when working with larger intervals or faster/slower frame rates, especially if keyframe selection is incorrect.
Question: How can I modify the keyframe selection process to correctly handle different frame intervals, ensuring that keyframes are selected appropriately based on both frame index and pose differences (translation and rotation)? Additionally, are there any improvements I can make to the pose estimation or point cloud handling to ensure robust performance with varying frame intervals?

Below is a portion of the code for better context:

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()

        self.frame_interval = frame_interval
        self.keyframes = []
        self.last_keyframe_index = -1
        self.min_keyframe_distance = 10  # Minimum number of processed frames between keyframes
        self.min_translation = 2  # Minimum translation to consider a new keyframe (in meters)
        self.min_rotation = 20  # Minimum rotation to consider a new keyframe (in degrees)

    def is_keyframe(self, frame_index, current_pose):
        if not self.keyframes:
            return True  # Always select the first frame as a keyframe

        frames_since_last_keyframe = frame_index - self.last_keyframe_index
        if frames_since_last_keyframe < self.min_keyframe_distance:
            return False

        last_keyframe = self.keyframes[-1]
        last_pose = last_keyframe['pose']

        # Calculate translation and rotation between current pose and last keyframe
        translation = np.linalg.norm(current_pose[:3, 3] - last_pose[:3, 3])
        rotation = np.degrees(np.arccos((np.trace(current_pose[:3, :3] @ last_pose[:3, :3].T) - 1) / 2))

        return translation > self.min_translation or rotation > self.min_rotation

    def add_keyframe(self, frame, depth_map, pose, frame_index):
        self.keyframes.append({
            'frame': frame,
            'depth_map': depth_map,
            'pose': pose,
            'index': frame_index
        })
        self.last_keyframe_index = frame_index
        print(f"Added keyframe at frame {frame_index}")

    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
        processed_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) < 8:
                                print(f"Frame {frame_count}: Insufficient matches found ({len(matches)}). Skipping pose estimation.")
                            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}")
                                    
                                    # Add keyframe if necessary
                                    if self.is_keyframe(processed_frame_count, self.pose):
                                        self.add_keyframe(frame, depth_map, self.pose, processed_frame_count)
                                    
                                    # 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()
                        
                        processed_frame_count +=1

                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()

Can you help me adjust the keyframe selection logic to scale properly with different frame_interval values and ensure that the system selects keyframes in a consistent, frame-appropriate manner? If you identify other issues with the code, suggestions for improvements would be appreciated!