from pathlib import Path
from typing import Tuple

import numpy as np
import json
import torch
import torchvision.transforms as tf
from einops import rearrange, repeat
from jaxtyping import Float, Int64
from omegaconf import DictConfig
from PIL import Image
from torch import Tensor
from torch.utils.data import Dataset
from tqdm import tqdm

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-i','--input_dir',  type=str,default="",required=False,help="src dir")
parser.add_argument('-o','--output_dir', type=str,default="",required=False,help="where to save poses")
args = parser.parse_args()

extrinsics=torch.stack([torch.from_numpy(np.array(x["transform_matrix"])) for x in json.load(open(args.input_dir+"/transforms.json"))["frames"]])

torch.save(extrinsics.cpu(),"%s/poses.pt"%args.output_dir)