main.py

import cv2
import numpy as np
import argparse
from src.dataset import KittiDataset, VideoDataset
from src.vo import VisualOdometry

def main():
    parser = argparse.ArgumentParser(description='Python Visual Odometry')
    parser.add_argument('--video', type=str, help='Path to video file')
    parser.add_argument('--kitti', type=str, help='Path to KITTI sequence folder (containing image_0 or image_1)')
    args = parser.parse_args()

    # Initialize Dataset
    if args.kitti:
        dataset = KittiDataset(args.kitti)
    elif args.video:
        dataset = VideoDataset(args.video)
    else:
        print("Please provide --video or --kitti path")
        return

    # Initialize VO
    # Note: Focal length and principal point should ideally come from calibration.
    # For KITTI, these are known. For random video, we might guess.
    cam_focal = 718.8560  # calibration for KITTI sequence 00
    cam_pp = (607.1928, 185.2157) # calibration for KITTI sequence 00
    
    vo = VisualOdometry(cam_focal, cam_pp)

    # Trajectory drawing
    traj_img = np.zeros((600, 600, 3), dtype=np.uint8)
    
    # Main Loop
    for i, img in enumerate(dataset):
        vo.update(img, i)

        # Visualization
        cur_t = vo.cur_t
        if i > 2:
            x, y, z = cur_t[0][0], cur_t[1][0], cur_t[2][0]
        else:
            x, y, z = 0., 0., 0.
            
        # Draw Trajectory
        # Map x, z to image coordinates
        # Adjust scale and center based on your video behavior
        draw_x = int(x) + 300
        draw_y = int(z) + 100
        cv2.circle(traj_img, (draw_x, draw_y), 1, (i*255/1000, 255, 0), 1)

        # Display
        cv2.rectangle(traj_img, (10, 20), (600, 60), (0,0,0), -1)
        text = f"Coordinates: x={x:.2f}m y={y:.2f}m z={z:.2f}m"
        cv2.putText(traj_img, text, (20, 40), cv2.FONT_HERSHEY_PLAIN, 1, (255,255,255), 1, 8)

        cv2.imshow('Road Facing Camera', img)
        cv2.imshow('Trajectory', traj_img)
        
        if cv2.waitKey(1) == 27:
            break

    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()