import torch
import os
import argparse
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from lib.LandmarkDataset import LandmarkDataset
from lib.CalibRecorder import CalibRecorder
from lib.Calibrator import Calibrator
from lib.FVMModule import FVMModule
from lib.CameraModule import CameraModule


device = torch.device('cuda:0')

parser = argparse.ArgumentParser()
parser.add_argument('--data_source', type=str, default='./data/input')
args = parser.parse_args()
DATA_SOURCE = args.data_source

data_folder = os.path.join(DATA_SOURCE, 'images')
param_folder = os.path.join(DATA_SOURCE, 'params')
dataset = LandmarkDataset(data_folder=data_folder)
bfm = FVMModule(batch_size=len(dataset)).to(device)
camera = CameraModule(2048).to(device).eval()
recorder = CalibRecorder(data_folder=data_folder, param_folder=param_folder)

calibrator = Calibrator(dataset, bfm, camera, recorder)
calibrator.calibrate()