|
|
import os |
|
|
import json |
|
|
import argparse |
|
|
import torch |
|
|
from tqdm import tqdm |
|
|
|
|
|
from scripts.camera.geometry.camera import SimpleRadial |
|
|
from scripts.camera.geometry.gravity import Gravity |
|
|
from scripts.camera.geometry.perspective_fields import get_perspective_field |
|
|
from scripts.camera.utils.conversions import fov2focal |
|
|
from scripts.camera.utils.text import parse_camera_params |
|
|
|
|
|
class Cam_Generator: |
|
|
def __init__(self, mode="base"): |
|
|
self.mode = mode |
|
|
|
|
|
def _load_text(self, caption, h=512, w=512, k1=0, k2=0): |
|
|
|
|
|
roll, pitch, vfov = parse_camera_params(caption, self.mode) |
|
|
|
|
|
|
|
|
f = fov2focal(torch.tensor(vfov), h) |
|
|
px, py = w / 2, h / 2 |
|
|
params = torch.tensor([w, h, f, f, px, py, k1, k2]).float() |
|
|
gravity = torch.tensor([roll, pitch]).float() |
|
|
return params, gravity |
|
|
|
|
|
def _read_param(self, parameters, gravity): |
|
|
|
|
|
camera = SimpleRadial(parameters).float() |
|
|
roll, pitch = gravity.unbind(-1) |
|
|
gravity_obj = Gravity.from_rp(roll, pitch) |
|
|
camera = camera.scale(torch.Tensor([1, 1])) |
|
|
return {"camera": camera, "gravity": gravity_obj} |
|
|
|
|
|
def _get_perspective(self, data): |
|
|
|
|
|
camera = data["camera"] |
|
|
gravity_obj = data["gravity"] |
|
|
up_field, lat_field = get_perspective_field( |
|
|
camera, gravity_obj, use_up=True, use_latitude=True |
|
|
) |
|
|
del camera, gravity_obj |
|
|
return torch.cat([up_field[0], lat_field[0]], dim=0) |
|
|
|
|
|
def get_cam(self, caption): |
|
|
params, gravity = self._load_text(caption) |
|
|
data = self._read_param(params, gravity) |
|
|
return self._get_perspective(data) |
|
|
|
|
|
def process_folders(input_root, output_root, start_idx=0, num_folders=None, mode="base"): |
|
|
gen = Cam_Generator(mode=mode) |
|
|
all_dirs = sorted([ |
|
|
d for d in os.listdir(input_root) |
|
|
if os.path.isdir(os.path.join(input_root, d)) |
|
|
]) |
|
|
if num_folders is None: |
|
|
num_folders = len(all_dirs) - start_idx |
|
|
selected = all_dirs[start_idx:start_idx + num_folders] |
|
|
|
|
|
for sub in tqdm(selected, desc="Subfolders"): |
|
|
in_sub = os.path.join(input_root, sub) |
|
|
out_sub = os.path.join(output_root, sub) |
|
|
os.makedirs(out_sub, exist_ok=True) |
|
|
|
|
|
json_files = sorted([ |
|
|
f for f in os.listdir(in_sub) |
|
|
if f.lower().endswith('.json') |
|
|
]) |
|
|
|
|
|
for jf in tqdm(json_files, desc=f"Processing {sub}", leave=False): |
|
|
in_path = os.path.join(in_sub, jf) |
|
|
with open(in_path, 'r', encoding='utf-8') as f: |
|
|
data = json.load(f) |
|
|
caption = data.get('caption', '') |
|
|
cam = gen.get_cam(caption) |
|
|
out_name = os.path.splitext(jf)[0] + '.pt' |
|
|
out_path = os.path.join(out_sub, out_name) |
|
|
torch.save(cam, out_path) |
|
|
|
|
|
def main(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description="Batch process the captions to the camera maps and save as .pt" |
|
|
) |
|
|
parser.add_argument('--input_root', type=str, |
|
|
help='Root directory of JSON subfolders') |
|
|
parser.add_argument('--output_root', type=str, |
|
|
help='Root directory to save .pt files') |
|
|
parser.add_argument('--start_idx', type=int, default=0, |
|
|
help='Start index of subfolders (0-based, default=0)') |
|
|
parser.add_argument('--num_folders', type=int, default=None, |
|
|
help='Number of subfolders to process (default: all)') |
|
|
parser.add_argument('--mode', type=str, default='base', |
|
|
help='parse_camera_params mode') |
|
|
args = parser.parse_args() |
|
|
|
|
|
process_folders( |
|
|
args.input_root, |
|
|
args.output_root, |
|
|
start_idx=args.start_idx, |
|
|
num_folders=args.num_folders, |
|
|
mode=args.mode |
|
|
) |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
main() |
|
|
|
