mast3r_expt.py

from mast3r.model import AsymmetricMASt3R
from mast3r.fast_nn import fast_reciprocal_NNs
import mast3r.utils.path_to_dust3r
from dust3r.inference import inference
from dust3r.utils.image import load_images
import os
import numpy as np
import torch
import torchvision.transforms.functional as F
import torch.nn.functional as tnf
from matplotlib import pyplot as plt
from PIL import Image
import cv2
import roma 
from scipy.spatial.transform import Rotation as R
from transforms3d.quaternions import qinverse, rotate_vector, qmult
from transforms3d.quaternions import quat2mat
from typing import List, Tuple
import re
import copy
from tqdm import tqdm

def extract_frame_number(img_path):
    """Extract frame number from image path using regex."""
    match = re.search(r'frame_(\d+)\.jpg', img_path)
    if match:
        return match.group(1)
    else:
        raise ValueError("Frame number not found in image path")

def get_mask_paths(frame_number, seq_folder):
    """Construct paths for the masks based on the frame number."""
    mask1_path = os.path.join(base_dir, f'/home2/swayam_421/indoor_loc/expts/dataset/upd_dataset/mapfree_subset/{seq_folder}/output_v2/seg_pix_confidences/query_pix_confidences_rframe_{frame_number}.pt')
    mask2_path = os.path.join(base_dir, f'/home2/swayam_421/indoor_loc/expts/dataset/upd_dataset/mapfree_subset/{seq_folder}/output_v2/seg_pix_confidences/ref_pix_confidences_rframe_{frame_number}.pt')
    # mask1 -> seq0, mask2 -> seq1 frames
    return mask1_path, mask2_path

def get_2d_3d_correspondences(matches_im0, matches_im1, pointmap1, pointmap2):
    pts2d = []
    pts3d = []
    
    for (x0, y0), (x1, y1) in zip(matches_im0, matches_im1):
        pts2d.append([x1, y1]) # The corresponding 2D point in image 2 -> query image
        pts3d.append(pointmap1[y0, x0]) # 3D point from the pointmap of image 1 -> reference image
        
    pts2d = np.array(pts2d, dtype=np.float32)
    pts3d = np.array(pts3d, dtype=np.float32)
    return pts2d, pts3d

# Conversion functions
def quaternion_to_matrix(qw, qx, qy, qz):
    r = R.from_quat([qx, qy, qz, qw])
    rotation_matrix = r.as_matrix()
    return rotation_matrix

def pose_to_matrix(qw, qx, qy, qz, tx, ty, tz):
    rotation_matrix = quaternion_to_matrix(qw, qx, qy, qz)
    translation_vector = np.array([tx, ty, tz])
    transformation_matrix = np.eye(4)
    transformation_matrix[:3, :3] = rotation_matrix
    transformation_matrix[:3, 3] = translation_vector
    return transformation_matrix

def compute_ground_truth_pose(pose2):
    ground_truth_pose = pose_to_matrix(*pose2)
    return ground_truth_pose

def pose_error_torch(R, t, Tgt, reduce=None):
    """Compute angular, scale and euclidean error of translation vector (metric). Compute angular rotation error."""

    Rgt = Tgt[:, :3, :3]                  # [B, 3, 3]
    tgt = Tgt[:, :3, 3:].transpose(1, 2)  # [B, 1, 3]

    scale_t = torch.linalg.norm(t, dim=-1)
    scale_tgt = torch.linalg.norm(tgt, dim=-1)

    cosine = (t @ tgt.transpose(1, 2)).squeeze(-1) / (scale_t * scale_tgt + 1e-9)
    cosine = torch.clip(cosine, -1.0, 1.0)    # handle numerical errors
    t_ang_err = torch.rad2deg(torch.acos(cosine))
    t_ang_err = torch.minimum(t_ang_err, 180 - t_ang_err)

    t_scale_err = scale_t / scale_tgt
    t_scale_err_sym = torch.maximum(scale_t / scale_tgt, scale_tgt / scale_t)
    t_euclidean_err = torch.linalg.norm(t - tgt, dim=-1)

    residual = R.transpose(1, 2) @ Rgt
    trace = torch.diagonal(residual, dim1=-2, dim2=-1).sum(-1)
    cosine = (trace - 1) / 2
    cosine = torch.clip(cosine, -1., 1.)  # handle numerical errors
    R_err = torch.rad2deg(torch.acos(cosine))

    if reduce is None:
        def fn(x): return x
    elif reduce == 'mean':
        fn = torch.mean
    elif reduce == 'median':
        fn = torch.median

    t_ang_err = fn(t_ang_err)
    t_scale_err = fn(t_scale_err)
    t_euclidean_err = fn(t_euclidean_err)
    R_err = fn(R_err)

    errors = {'t_err_ang': t_ang_err,
              't_err_scale': t_scale_err,
              't_err_scale_sym': t_scale_err_sym,
              't_err_euc': t_euclidean_err,
              'R_err': R_err}
    return errors

def rvec_tvec_to_matrix(rvec, tvec):
    rotation_matrix, _ = cv2.Rodrigues(rvec)
    transformation_matrix = np.eye(4)
    transformation_matrix[:3, :3] = rotation_matrix
    transformation_matrix[:3, 3] = tvec.flatten()
    return transformation_matrix

def convert_world2cam_to_cam2world(q, t):
    """Converts a world2cam pose to cam2world pose."""
    qinv = qinverse(q)
    tinv = -rotate_vector(t, qinv)
    return qinv, tinv

def project(pts: np.ndarray, K: np.ndarray, img_size: List[int] or Tuple[int] = None) -> np.ndarray:
    """Projects 3D points to image plane.

    Args:
        - pts [N, 3/4]: points in camera coordinates (homogeneous or non-homogeneous)
        - K [3, 3]: intrinsic matrix
        - img_size (width, height): optional, clamp projection to image borders
        Outputs:
        - uv [N, 2]: coordinates of projected points
    """

    assert len(pts.shape) == 2, 'incorrect number of dimensions'
    assert pts.shape[1] in [3, 4], 'invalid dimension size'
    assert K.shape == (3, 3), 'incorrect intrinsic shape'

    uv_h = (K @ pts[:, :3].T).T
    uv = uv_h[:, :2] / uv_h[:, -1:]

    if img_size is not None:
        uv[:, 0] = np.clip(uv[:, 0], 0, img_size[0])
        uv[:, 1] = np.clip(uv[:, 1], 0, img_size[1])

    return uv

def get_grid_multipleheight() -> np.ndarray:
    # create grid of points
    ar_grid_step = 0.3
    ar_grid_num_x = 7
    ar_grid_num_y = 4
    ar_grid_num_z = 7
    ar_grid_z_offset = 1.8
    ar_grid_y_offset = 0

    ar_grid_x_pos = np.arange(0, ar_grid_num_x)-(ar_grid_num_x-1)/2
    ar_grid_x_pos *= ar_grid_step

    ar_grid_y_pos = np.arange(0, ar_grid_num_y)-(ar_grid_num_y-1)/2
    ar_grid_y_pos *= ar_grid_step
    ar_grid_y_pos += ar_grid_y_offset

    ar_grid_z_pos = np.arange(0, ar_grid_num_z).astype(float)
    ar_grid_z_pos *= ar_grid_step
    ar_grid_z_pos += ar_grid_z_offset

    xx, yy, zz = np.meshgrid(ar_grid_x_pos, ar_grid_y_pos, ar_grid_z_pos)
    ones = np.ones(xx.shape[0]*xx.shape[1]*xx.shape[2])
    eye_coords = np.concatenate([c.reshape(-1, 1)
                                for c in (xx, yy, zz, ones)], axis=-1)
    return eye_coords

def read_intrinsics(intrinsics_file, image_path):
    with open(intrinsics_file, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if parts[0] == image_path:
                fx = float(parts[1])
                fy = float(parts[2])
                cx = float(parts[3])
                cy = float(parts[4])
                K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
                return K, parts[5], parts[6]
    return None, None, None

def read_pose(pose_file, image_path):
    with open(pose_file, 'r') as f:
        for line in f:
            parts = line.strip().split()
            if parts[0] == image_path:
                pose = list(map(float, parts[1:]))
                return pose
    return None

# global variable, avoids creating it again
eye_coords_glob = get_grid_multipleheight()

def reprojection_error(
        q_est: np.ndarray, t_est: np.ndarray, q_gt: np.ndarray, t_gt: np.ndarray, K: np.ndarray,
        W: int, H: int) -> float:
    eye_coords = eye_coords_glob

    # obtain ground-truth position of projected points
    uv_gt = project(eye_coords, K, (W, H))

    # residual transformation
    cam2w_est = np.eye(4)
    cam2w_est[:3, :3] = quat2mat(q_est)
    cam2w_est[:3, -1] = t_est
    cam2w_gt = np.eye(4)
    cam2w_gt[:3, :3] = quat2mat(q_gt)
    cam2w_gt[:3, -1] = t_gt

    # residual reprojection
    eyes_residual = (np.linalg.inv(cam2w_est) @ cam2w_gt @ eye_coords.T).T
    uv_pred = project(eyes_residual, K, (W, H))

    # get reprojection error
    repr_err = np.linalg.norm(uv_gt - uv_pred, ord=2, axis=1)
    mean_repr_err = float(repr_err.mean().item())
    return mean_repr_err

def load_mask(mask_path: str, size: tuple):
    mask_tensor = torch.load(mask_path).unsqueeze(0).unsqueeze(0)
    # print("Mask tensor shape: ", mask_tensor.shape)
    mask = tnf.interpolate(mask_tensor, size, mode='bilinear')
    # print("Mask shape after interpolation: ", mask.shape)
    mask = mask.squeeze()
    return mask.numpy()

def is_point_in_mask(point, mask):
    return mask[point[1], point[0]] > 0
    
def filter_matches_within_mask(matches_im0, matches_im1, mask1, mask2):
    """Filter matches that are within the mask given matches and masks."""
    matches_in_mask = []
    for m1, m2 in zip(matches_im0, matches_im1):
        if is_point_in_mask(m1, mask1) and is_point_in_mask(m2, mask2):
            matches_in_mask.append((m1, m2))

    return np.array(matches_in_mask)

def compute_metrics(pts2d, pts3d, K, gt_pose, img0, output_dir):
    success, r_pose, t_pose, _ = cv2.solvePnPRansac(pts3d, pts2d, K, None, flags=cv2.SOLVEPNP_SQPNP,
                                                    confidence=0.9999)
    
    r_pose = cv2.Rodrigues(r_pose)[0]  # world2cam == world2cam2
    RT = np.r_[np.c_[r_pose, t_pose], [(0,0,0,1)]] # world2cam2

    computed_pose = np.linalg.inv(RT)  # cam2toworld

    if success:
        # print(f"PnP with RANSAC succeeded")
        q = np.array((gt_pose[0], gt_pose[1], gt_pose[2], gt_pose[3]), dtype=np.float64)
        t = np.array((gt_pose[4], gt_pose[5], gt_pose[6]), dtype=np.float64)

        assert len(q) == 4 and len(t) == 3

        q, t = convert_world2cam_to_cam2world(q, t)

        R_est = computed_pose[:3, :3]
        t_est = computed_pose[:3, 3]
        errors = pose_error_torch(torch.from_numpy(R_est).unsqueeze(0), torch.from_numpy(t_est).unsqueeze(0),
                                  torch.from_numpy(pose_to_matrix(*q, *t)).unsqueeze(0))

        np.savetxt(os.path.join(output_dir, 'computed_pose.txt'), computed_pose)

        ground_truth_pose = pose_to_matrix(*q, *t)

        np.savetxt(os.path.join(output_dir, 'ground_truth_pose.txt'), ground_truth_pose)
        
        abs_transl_error = errors['t_err_euc'].item()
        abs_angular_error = errors['R_err'].item()
        abs_transl_error_deg = errors['t_err_ang'].item()
        abs_transl_error_scale = errors['t_err_scale'].item()
        abs_transl_error_scale_sym = errors['t_err_scale_sym'].item()

        # get q_est and t_est in qw, qx, qy, qz format
        q_est = R.from_matrix(R_est).as_quat()
        t_est = t_est

        width = img0.shape[1]
        height = img0.shape[0]
        reproj_err = reprojection_error(q_est, t_est, q, t, K, width, height)

        with open(os.path.join(output_dir, 'errors.txt'), 'w') as f:
            f.write(f"Absolute translation error: {abs_transl_error}\n")
            f.write(f"Absolute angular error: {abs_angular_error}\n")
            f.write(f"Absolute translation error in degrees: {abs_transl_error_deg}\n")
            f.write(f"Absolute translation error in scale: {abs_transl_error_scale}\n")
            f.write(f"Absolute translation error in scale symmetric: {abs_transl_error_scale_sym}\n")
            f.write(f"Reprojection error: {reproj_err}\n")

        errors = {}
        errors['abs_transl_error'] = abs_transl_error
        errors['abs_angular_error'] = abs_angular_error
        errors['abs_transl_error_deg'] = abs_transl_error_deg
        errors['abs_transl_error_scale'] = abs_transl_error_scale
        errors['abs_transl_error_scale_sym'] = abs_transl_error_scale_sym
        errors['reproj_err'] = reproj_err

        return errors
    else:
        print("PnP with RANSAC failed")
        return None

def match_images(img1_path, img2_path, mask1_path, mask2_path, output_dir):
    device = 'cuda'
    model_name = "naver/MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric"
    model = AsymmetricMASt3R.from_pretrained(model_name).to(device)
    images = load_images([img1_path, img2_path], size=512)

    output = inference([tuple(images)], model, device, batch_size=1, verbose=False)

    view1, pred1 = output['view1'], output['pred1']
    view2, pred2 = output['view2'], output['pred2']

    desc1, desc2 = pred1['desc'].squeeze(0).detach(), pred2['desc'].squeeze(0).detach()

    print(f"Desc1 shape: {desc1.shape}")
    print(f"Desc2 shape: {desc2.shape}")

    matches_im0, matches_im1 = fast_reciprocal_NNs(desc1, desc2, subsample_or_initxy1=8,
                                                   device=device, dist='dot', block_size=2**13)

    H0, W0 = view1['true_shape'][0]
    valid_matches_im0 = (matches_im0[:, 0] >= 3) & (matches_im0[:, 0] < int(W0) - 3) & (
        matches_im0[:, 1] >= 3) & (matches_im0[:, 1] < int(H0) - 3)

    H1, W1 = view2['true_shape'][0]
    valid_matches_im1 = (matches_im1[:, 0] >= 3) & (matches_im1[:, 0] < int(W1) - 3) & (
        matches_im1[:, 1] >= 3) & (matches_im1[:, 1] < int(H1) - 3)

    valid_matches = valid_matches_im0 & valid_matches_im1
    matches_im0, matches_im1 = matches_im0[valid_matches], matches_im1[valid_matches]
    sav_matches_im0 = copy.deepcopy(matches_im0)
    sav_matches_im1 = copy.deepcopy(matches_im1)

    conf_list = [pred1['desc_conf'].squeeze(0).cpu().numpy(), pred2['desc_conf'].squeeze(0).cpu().numpy()]
    matches_confs = np.minimum(
        conf_list[1][matches_im0[:, 1], matches_im0[:, 0]],
        conf_list[0][matches_im1[:, 1], matches_im1[:, 0]]
    )

    # if len(matches_confs) > 0:
    #     mask = matches_confs > 0.7
    #     matches_im0, matches_im1 = matches_im0[mask], matches_im1[mask]    
    # if matches_im0.shape[0] == 0:
    #     matches_im0 = sav_matches_im0
    #     matches_im1 = sav_matches_im1

    mask1, mask2 = load_images([mask1_path, mask2_path], size=(512))

    mask1_img = F.to_pil_image(mask1['img'])
    mask1_img.save(os.path.join(output_dir, 'mask1VIZ.png'))
    mask2_img = F.to_pil_image(mask2['img'])
    mask2_img.save(os.path.join(output_dir, 'mask2VIZ.png'))
    
    assert mask1 is not None and mask2 is not None
    
    matches_in_mask = filter_matches_within_mask(matches_im0, matches_im1, mask1['img'], mask2['img'])

    image_mean = torch.as_tensor([0.5, 0.5, 0.5], device='cpu').reshape(1, 3, 1, 1)
    image_std = torch.as_tensor([0.5, 0.5, 0.5], device='cpu').reshape(1, 3, 1, 1)

    viz_imgs = []
    for i, view in enumerate([view1, view2]):
        rgb_tensor = view['img'] * image_std + image_mean
        viz_imgs.append(rgb_tensor.squeeze(0).permute(1, 2, 0).cpu().numpy())

    H0, W0, H1, W1 = *viz_imgs[0].shape[:2], *viz_imgs[1].shape[:2]
    img0 = np.pad(viz_imgs[0], ((0, max(H1 - H0, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)
    img1 = np.pad(viz_imgs[1], ((0, max(H0 - H1, 0)), (0, 0), (0, 0)), 'constant', constant_values=0)
    img = np.concatenate((img0, img1), axis=1)
    plt.figure()
    plt.imshow(img)
    cmap = plt.get_cmap('jet')
    for m0, m1 in matches_in_mask:
        (x0, y0), (x1, y1) = m0.T, m1.T
        plt.plot([x0, x1 + W0], [y0, y1], '-+', scalex=False, scaley=False)

    plt.savefig(os.path.join(output_dir, 'matched_images.png'))  # Save the image to a file
    plt.close()  

    pointmap1 = pred1['pts3d'].squeeze(0).detach().cpu().numpy() # 3D points of image 1 in the coordinate system of camera 1
    pointmap2 = pred2['pts3d_in_other_view'].squeeze(0).detach().cpu().numpy() # 3D points of image 2 in the coordinate system of camera 1

    assert pointmap1.shape[2] == 3 and pointmap2.shape[2] == 3

    matches_im0 = matches_in_mask[:, 0]
    matches_im1 = matches_in_mask[:, 1]

    pts2d, pts3d = get_2d_3d_correspondences(matches_im0, matches_im1, pointmap1, pointmap2)

    pts2d_vanilla, pts3d_vanilla = get_2d_3d_correspondences(sav_matches_im0, sav_matches_im1, pointmap1, pointmap2)

    assert len(pts2d) == len(pts3d) and len(pts2d) >= 4

    img2_modif_path = os.path.join(*img2_path.split('/')[-2:])

    intrinsics_file = os.path.join(os.path.dirname(os.path.dirname(img2_path)), 'intrinsics.txt')

    K1, width, height = read_intrinsics(intrinsics_file, img2_modif_path)

    poses_file = os.path.join(os.path.dirname(os.path.dirname(img2_path)), 'poses.txt')
    
    pose2 = read_pose(poses_file, img2_modif_path)

    if pose2 is None:
        raise ValueError(f"Pose for {img2_path} not found in {poses_file}")
    
    if K1 is None:
        raise ValueError(f"Intrinsics not found in {intrinsics_file}")

    mask_results_dir = os.path.join(output_dir, 'mask_results')
    os.makedirs(mask_results_dir, exist_ok=True)

    vanilla_results_dir = os.path.join(output_dir, 'vanilla_results')
    os.makedirs(vanilla_results_dir, exist_ok=True)

    metrics_with_mask = compute_metrics(pts2d, pts3d, K1, pose2, viz_imgs[0], mask_results_dir)

    metrics_vanilla = compute_metrics(pts2d_vanilla, pts3d_vanilla, K1, pose2, viz_imgs[0], vanilla_results_dir)

    return metrics_with_mask, metrics_vanilla

def process_sequences(base_dir):
    global_results_abs_transl = []
    global_results_abs_angular = []
    global_results_abs_transl_deg = []
    global_results_abs_transl_scale = []
    global_results_abs_transl_scale_sym = []
    global_results_reproj = []

    global_results_abs_transl_vanilla = []
    global_results_abs_angular_vanilla = []
    global_results_abs_transl_deg_vanilla = []
    global_results_abs_transl_scale_vanilla = []
    global_results_abs_transl_scale_sym_vanilla = []
    global_results_reproj_vanilla = []

    sequences = [sequence_folder for sequence_folder in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, sequence_folder))]
    sequences.sort(key=lambda x: int(x[1:]))

    for sequence_folder in tqdm(sequences, desc="Processing sequences"):
        sequence_path = os.path.join(base_dir, sequence_folder)
        if os.path.isdir(sequence_path):
            img1_path = os.path.join(sequence_path, 'seq0', 'frame_00000.jpg')
            seq1_folder = os.path.join(sequence_path, 'seq1')
            
            if not os.path.exists(img1_path) or not os.path.isdir(seq1_folder):
                print("Warning: Required files/folders do not exist!!!")
                continue 
            
            results = []
            results_vanilla = []

            for img2_file in os.listdir(seq1_folder):
                img2_path = os.path.join(seq1_folder, img2_file)
                if img2_file.endswith('.jpg'):
                    output_dir = os.path.join(sequence_path, 'results', os.path.splitext(img2_file)[0])
                    os.makedirs(output_dir, exist_ok=True)
                    frame_number = extract_frame_number(img2_file)
                    mask1_path, mask2_path = get_mask_paths(frame_number, sequence_folder)
                    
                    errors_with_mask, errors_vanilla = match_images(img1_path, img2_path, mask1_path, mask2_path, output_dir)
                    results.append(errors_with_mask)
                    results_vanilla.append(errors_vanilla)

            abs_transl_errors = [errors['abs_transl_error'] for errors in results]
            abs_angular_errors = [errors['abs_angular_error'] for errors in results]
            abs_transl_errors_deg = [errors['abs_transl_error_deg'] for errors in results]
            abs_transl_errors_scale = [errors['abs_transl_error_scale'] for errors in results]
            abs_transl_errors_scale_sym = [errors['abs_transl_error_scale_sym'] for errors in results]
            reproj_errors = [errors['reproj_err'] for errors in results]

            abs_transl_errors_vanilla = [errors['abs_transl_error'] for errors in results_vanilla]
            abs_angular_errors_vanilla = [errors['abs_angular_error'] for errors in results_vanilla]
            abs_transl_errors_deg_vanilla = [errors['abs_transl_error_deg'] for errors in results_vanilla]
            abs_transl_errors_scale_vanilla = [errors['abs_transl_error_scale'] for errors in results_vanilla]
            abs_transl_errors_scale_sym_vanilla = [errors['abs_transl_error_scale_sym'] for errors in results_vanilla]
            reproj_errors_vanilla = [errors['reproj_err'] for errors in results_vanilla]

            global_results_abs_transl.append(abs_transl_errors)
            global_results_abs_angular.append(abs_angular_errors)
            global_results_abs_transl_deg.append(abs_transl_errors_deg)
            global_results_abs_transl_scale.append(abs_transl_errors_scale)
            global_results_abs_transl_scale_sym.append(abs_transl_errors_scale_sym)
            global_results_reproj.append(reproj_errors)

            global_results_abs_transl_vanilla.append(abs_transl_errors_vanilla)
            global_results_abs_angular_vanilla.append(abs_angular_errors_vanilla)
            global_results_abs_transl_deg_vanilla.append(abs_transl_errors_deg_vanilla)
            global_results_abs_transl_scale_vanilla.append(abs_transl_errors_scale_vanilla)
            global_results_abs_transl_scale_sym_vanilla.append(abs_transl_errors_scale_sym_vanilla)
            global_results_reproj_vanilla.append(reproj_errors_vanilla)

            with open(os.path.join(sequence_path, 'median_errors.txt'), 'w') as f:
                f.write(f"Median absolute translation error: {np.median(abs_transl_errors)}\n")
                f.write(f"Median absolute angular error: {np.median(abs_angular_errors)}\n")
                f.write(f"Median absolute translation error in degrees: {np.median(abs_transl_errors_deg)}\n")
                f.write(f"Median absolute translation error in scale: {np.median(abs_transl_errors_scale)}\n")
                f.write(f"Median absolute translation error in scale symmetric: {np.median(abs_transl_errors_scale_sym)}\n")
                f.write(f"Median reprojection error: {np.median(reproj_errors)}\n")

            with open(os.path.join(sequence_path, 'mean_errors.txt'), 'w') as f:
                f.write(f"Mean absolute translation error: {np.mean(abs_transl_errors)}\n")
                f.write(f"Mean absolute angular error: {np.mean(abs_angular_errors)}\n")
                f.write(f"Mean absolute translation error in degrees: {np.mean(abs_transl_errors_deg)}\n")
                f.write(f"Mean absolute translation error in scale: {np.mean(abs_transl_errors_scale)}\n")
                f.write(f"Mean absolute translation error in scale symmetric: {np.mean(abs_transl_errors_scale_sym)}\n")
                f.write(f"Mean reprojection error: {np.mean(reproj_errors)}\n")

            with open(os.path.join(sequence_path, 'median_errors_vanilla.txt'), 'w') as f:
                f.write(f"Median absolute translation error: {np.median(abs_transl_errors_vanilla)}\n")
                f.write(f"Median absolute angular error: {np.median(abs_angular_errors_vanilla)}\n")
                f.write(f"Median absolute translation error in degrees: {np.median(abs_transl_errors_deg_vanilla)}\n")
                f.write(f"Median absolute translation error in scale: {np.median(abs_transl_errors_scale_vanilla)}\n")
                f.write(f"Median absolute translation error in scale symmetric: {np.median(abs_transl_errors_scale_sym_vanilla)}\n")
                f.write(f"Median reprojection error: {np.median(reproj_errors_vanilla)}\n")

            with open(os.path.join(sequence_path, 'mean_errors_vanilla.txt'), 'w') as f:
                f.write(f"Mean absolute translation error: {np.mean(abs_transl_errors_vanilla)}\n")
                f.write(f"Mean absolute angular error: {np.mean(abs_angular_errors_vanilla)}\n")
                f.write(f"Mean absolute translation error in degrees: {np.mean(abs_transl_errors_deg_vanilla)}\n")
                f.write(f"Mean absolute translation error in scale: {np.mean(abs_transl_errors_scale_vanilla)}\n")
                f.write(f"Mean absolute translation error in scale symmetric: {np.mean(abs_transl_errors_scale_sym_vanilla)}\n")
                f.write(f"Mean reprojection error: {np.mean(reproj_errors_vanilla)}\n")

    with open(os.path.join(base_dir, 'global_median_errors.txt'), 'w') as f:
        f.write(f"Median absolute translation error: {np.median(np.concatenate(global_results_abs_transl))}\n")
        f.write(f"Median absolute angular error: {np.median(np.concatenate(global_results_abs_angular))}\n")
        f.write(f"Median absolute translation error in degrees: {np.median(np.concatenate(global_results_abs_transl_deg))}\n")
        f.write(f"Median absolute translation error in scale: {np.median(np.concatenate(global_results_abs_transl_scale))}\n")
        f.write(f"Median absolute translation error in scale symmetric: {np.median(np.concatenate(global_results_abs_transl_scale_sym))}\n")
        f.write(f"Median reprojection error: {np.median(np.concatenate(global_results_reproj))}\n")

    with open(os.path.join(base_dir, 'global_mean_errors.txt'), 'w') as f:
        f.write(f"Mean absolute translation error: {np.mean(np.concatenate(global_results_abs_transl))}\n")
        f.write(f"Mean absolute angular error: {np.mean(np.concatenate(global_results_abs_angular))}\n")
        f.write(f"Mean absolute translation error in degrees: {np.mean(np.concatenate(global_results_abs_transl_deg))}\n")   
        f.write(f"Mean absolute translation error in scale: {np.mean(np.concatenate(global_results_abs_transl_scale))}\n")
        f.write(f"Mean absolute translation error in scale symmetric: {np.mean(np.concatenate(global_results_abs_transl_scale_sym))}\n")
        f.write(f"Mean reprojection error: {np.mean(np.concatenate(global_results_reproj))}\n")

    with open(os.path.join(base_dir, 'global_median_errors_vanilla.txt'), 'w') as f:
        f.write(f"Median absolute translation error: {np.median(np.concatenate(global_results_abs_transl_vanilla))}\n")
        f.write(f"Median absolute angular error: {np.median(np.concatenate(global_results_abs_angular_vanilla))}\n")
        f.write(f"Median absolute translation error in degrees: {np.median(np.concatenate(global_results_abs_transl_deg_vanilla))}\n")
        f.write(f"Median absolute translation error in scale: {np.median(np.concatenate(global_results_abs_transl_scale_vanilla))}\n")
        f.write(f"Median absolute translation error in scale symmetric: {np.median(np.concatenate(global_results_abs_transl_scale_sym_vanilla))}\n")
        f.write(f"Median reprojection error: {np.median(np.concatenate(global_results_reproj_vanilla))}\n")

    with open(os.path.join(base_dir, 'global_mean_errors_vanilla.txt'), 'w') as f:
        f.write(f"Mean absolute translation error: {np.mean(np.concatenate(global_results_abs_transl_vanilla))}\n")
        f.write(f"Mean absolute angular error: {np.mean(np.concatenate(global_results_abs_angular_vanilla))}\n")
        f.write(f"Mean absolute translation error in degrees: {np.mean(np.concatenate(global_results_abs_transl_deg_vanilla))}\n")
        f.write(f"Mean absolute translation error in scale: {np.mean(np.concatenate(global_results_abs_transl_scale_vanilla))}\n")
        f.write(f"Mean absolute translation error in scale symmetric: {np.mean(np.concatenate(global_results_abs_transl_scale_sym_vanilla))}\n")
        f.write(f"Mean reprojection error: {np.mean(np.concatenate(global_results_reproj_vanilla))}\n")

    print("All Processing Done!!!")        
    
if __name__ == "__main__":
    base_dir = "/home2/swayam_421/indoor_loc/expts/dataset/upd_dataset/mapfree_subset"
    process_sequences(base_dir)