4DGaussians/scripts/blender2colmap.py
2023-12-02 14:13:12 +08:00

89 lines
2.8 KiB
Python

import os
import numpy as np
import glob
import sys
import json
from PIL import Image
from tqdm import tqdm
import shutil
import math
def fov2focal(fov, pixels):
return pixels / (2 * math.tan(fov / 2))
def rotmat2qvec(R):
Rxx, Ryx, Rzx, Rxy, Ryy, Rzy, Rxz, Ryz, Rzz = R.flat
K = np.array([
[Rxx - Ryy - Rzz, 0, 0, 0],
[Ryx + Rxy, Ryy - Rxx - Rzz, 0, 0],
[Rzx + Rxz, Rzy + Ryz, Rzz - Rxx - Ryy, 0],
[Ryz - Rzy, Rzx - Rxz, Rxy - Ryx, Rxx + Ryy + Rzz]]) / 3.0
eigvals, eigvecs = np.linalg.eigh(K)
qvec = eigvecs[[3, 0, 1, 2], np.argmax(eigvals)]
if qvec[0] < 0:
qvec *= -1
return qvec
root_dir = sys.argv[1]
colmap_dir = os.path.join(root_dir,"sparse_")
if not os.path.exists(colmap_dir):
os.makedirs(colmap_dir)
imagecolmap_dir = os.path.join(root_dir,"image_colmap")
if not os.path.exists(imagecolmap_dir):
os.makedirs(imagecolmap_dir)
image_dir = os.path.join(root_dir)
images = os.listdir(image_dir)
images.sort()
camera_json = os.path.join(root_dir,"transforms_train.json")
with open (camera_json) as f:
meta = json.load(f)
try:
image_size = meta['w'], meta['h']
focal = [meta['fl_x'],meta['fl_y']]
except:
try:
image_size = meta['frames'][0]['w'], meta['frames'][0]['h']
focal = [meta['frames'][0]['fl_x'],meta['frames'][0]['fl_y']]
except:
image_size = 800,800
focal = fov2focal(meta['camera_angle_x'], 800)
focal = [focal,focal]
# size = image.size
# breakpoint()
object_images_file = open(os.path.join(colmap_dir,"images.txt"),"w")
object_cameras_file = open(os.path.join(colmap_dir,"cameras.txt"),"w")
idx=0
sizes=1
cnt=0
while len(meta['frames'])//sizes > 200:
sizes += 1
for frame in meta['frames']:
cnt+=1
if cnt % sizes != 0:
continue
matrix = np.linalg.inv(np.array(frame["transform_matrix"]))
R = -np.transpose(matrix[:3,:3])
R[:,0] = -R[:,0]
T = -matrix[:3, 3]
T = -np.matmul(R,T)
T = [str(i) for i in T]
qevc = [str(i) for i in rotmat2qvec(np.transpose(R))]
print(idx+1," ".join(qevc)," ".join(T),1,frame['file_path'].split('/')[-1]+".png","\n",file=object_images_file)
print(idx,"SIMPLE_PINHOLE",image_size[0],image_size[1],focal[0],image_size[0]/2,image_size[1]/2,file=object_cameras_file)
idx+=1
# breakpoint()
print(os.path.join(image_dir,frame['file_path']),os.path.join(imagecolmap_dir,frame['file_path'].split('/')[-1]+".png"))
shutil.copy(os.path.join(image_dir,frame['file_path']+".png"),os.path.join(imagecolmap_dir,frame['file_path'].split('/')[-1]+".png"))
# write camera infomation.
# print(1,"SIMPLE_PINHOLE",image_size[0],image_size[1],focal[0],image_sizep0/2,image_size[1]/2,file=object_cameras_file)
object_point_file = open(os.path.join(colmap_dir,"points3D.txt"),"w")
object_cameras_file.close()
object_images_file.close()
object_point_file.close()