two_image_3D_pose.py

import numpy as np
import cv2 as cv

from superpointclass import SuperPointDescriptors
from class_Undistort import UndistortCamera
from ImageFolderVideoReader import ImageFolderVideoReader
from class_Undistort import UndistortCamera

# CREATE VIDEO OBJECT FROM FOLDER PATH
folder_path = "/Users/sumitsarkar/Code/python_codes/MH_03_medium/mav0/cam0/data"
calibration_path  = "/Users/sumitsarkar/Code/python_codes/MH_03_medium/mav0/cam0/sensor.yaml"

def resize_image(image,height, width):
    
    img_size = [height, width]
    interp = cv.INTER_AREA
    resize_image = cv.resize(image, (img_size[1], img_size[0]), interpolation=interp)
    return resize_image
    #return image

def undistort_keypoints(keypoints, camera_matrix, dist_coeffs):
    """
    Undistort keypoints using camera calibration
    
    Args:
        keypoints: List of cv2.KeyPoint objects
        camera_matrix: 3x3 camera matrix
        dist_coeffs: Distortion coefficients [k1, k2, p1, p2, k3]
    
    Returns:
        List of undistorted keypoints
    """
    if not keypoints:
        return []
    
    # Extract point coordinates
    points = np.array([[kp.pt[0], kp.pt[1]] for kp in keypoints], dtype=np.float32)
    points = points.reshape(-1, 1, 2)
    
    # Undistort points
    undistorted_points = cv.undistortPoints(points, camera_matrix, dist_coeffs, P=camera_matrix)
    undistorted_points = undistorted_points.reshape(-1, 2)
    
    # Create new keypoints with undistorted coordinates
    undistorted_kp = []
    for i, kp in enumerate(keypoints):
        new_kp = cv.KeyPoint(
            x=undistorted_points[i][0],
            y=undistorted_points[i][1],
            size=kp.size,
            angle=kp.angle,
            response=kp.response,
            octave=kp.octave
        )
        undistorted_kp.append(new_kp)
    
    return undistorted_kp

try:
    # Create video reader from image folder
    video_reader = ImageFolderVideoReader(folder_path)
    undistort_image = UndistortCamera(calibration_path)
    undistort_image.load_camera_params()
    H = 240
    W = 326
    
    _, frame_t_0 = video_reader.read()
    frame_t_0_gray = cv.cvtColor(frame_t_0, cv.COLOR_BGR2GRAY)

    # 1. Undistort
    # 2. Resize
    undistort_frame_t_0_gray = undistort_image.get_undistortImage(frame_t_0_gray)
    undistort_frame_t_0_gray_resize = resize_image(undistort_frame_t_0_gray, H, W)

    # INITIALIZE SUPERPOINT
    _superpoint = SuperPointDescriptors()
    keypoints0, desc0 = _superpoint.get_keypoints(undistort_frame_t_0_gray_resize)

    # COLLECT ORIGINAL CAMERA PARAMETERS, assuming same camera for both images
    camera_K, camera_D = undistort_image.get_camera_params()

    # COLLECT CAMERA PARAMS AFTER RESIZE 
    #camera_K_resize, camera_D_resize = undistort_image.get_resized_camera_params(H, W)

   # Process frames like a regular video
    while True:

        _, frame_t_1 = video_reader.read()
        frame_t_1_gray = cv.cvtColor(frame_t_1, cv.COLOR_BGR2GRAY)

        # 1. Undistort
        # 2. Resize
        undistort_frame_t_1_gray = undistort_image.get_undistortImage(frame_t_1_gray)
        undistort_frame_t_1_gray_resize = resize_image(undistort_frame_t_1_gray, H, W)

        # DETECT KEYPOINTS AND GENERATE DESCRIPTORS
        keypoints1, desc1 = _superpoint.get_keypoints(undistort_frame_t_1_gray_resize)

        # GET THE MATCHES 
        matches = _superpoint.get_matches(desc0, desc1)

        # Extract matched points
        pts0 = np.float32([keypoints0[m.queryIdx].pt for m in matches]).reshape(-1, 2)
        pts1 = np.float32([keypoints1[m.trainIdx].pt for m in matches]).reshape(-1, 2)

        # Use fundamental matrix for robust filtering
        F, mask = cv.findFundamentalMat(
            pts0, pts1,
            method=cv.FM_RANSAC,
            ransacReprojThreshold=0.5,
            confidence=0.99
        )
        
        # Filter matches based on inlier mask
        good_matches = [matches[i] for i in range(len(matches)) if mask[i] == 1]

        # calculate percentage inliers
        print("number of inlier points = ", len(good_matches))
        print("Percentage Inlier = ", 100*len(good_matches)/len(matches))
        print("")

        # DRAW THE MATCHES 
        img_matches = cv.drawMatches(
            undistort_frame_t_0_gray_resize, keypoints0,           # Query image and keypoints
            undistort_frame_t_1_gray_resize, keypoints1,           # Train image and keypoints
            good_matches[:200],                 # Matches to draw
            None,                         # Output image (None = create new)
            flags=cv.DrawMatchesFlags_NOT_DRAW_SINGLE_POINTS

        )
        
        # Display the frame
        cv.imshow('MATCHES', img_matches)
        #cv.imshow('FRAMES', resize_frame_t_1_gray)

        # Copy current frame to prevs frame
        desc0 = desc1.copy()
        keypoints0 = keypoints1
        undistort_frame_t_0_gray_resize = undistort_frame_t_1_gray_resize.copy()

        # Press 'q' to quit
        if cv.waitKey(1) & 0xFF == ord('q'):
            break
    
    video_reader.release()
    cv.destroyAllWindows()
    
except ValueError as e:
    print(f"Error: {e}")