File size: 4,787 Bytes
55e58d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
from pathlib import Path
import shutil
import os
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import pickle
import cv2
from tqdm.contrib.concurrent import process_map
dust3r_path = Path("/home/jovyan/users/bulat/workspace/3drec/Indoor/DUSt3R/res/scannet_posed_35")
pcds_path = Path("/home/jovyan/users/bulat/workspace/3drec/Indoor/DUSt3R/res_scennet_posed_35_new")
def resize_intrinsics(K, base_shape, new_shape):
fx = K[0, 0]
fy = K[1, 1]
cx = K[0, 2]
cy = K[1, 2]
fx_new = fx * new_shape[0] / base_shape[0]
fy_new = fy * new_shape[1] / base_shape[1]
cx_new = cx * new_shape[0] / base_shape[0]
cy_new = cy * new_shape[1] / base_shape[1]
return np.array([[fx_new, 0, cx_new, 0], [0, fy_new, cy_new, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
def prepare_scene(scene_id):
scannet_processed_path = Path("/home/jovyan/users/bulat/workspace/3drec/Indoor/MaskClustering/data/scannet/processed") / scene_id
out_path = Path("data/scannet_dust3r_posed_35_bulat/processed") / scene_id
out_path.mkdir(parents=True, exist_ok=True)
path_to_scene = dust3r_path / scene_id
scene_params = torch.load(path_to_scene / 'scene_params.pt')
pcd_path = pcds_path / (scene_id + '.ply')
poses = scene_params['poses']
depths = scene_params['depths']
Ks = scene_params['Ks']
image_names = [Path(image_name) for image_name in scene_params['image_files']]
shutil.copy(pcd_path, out_path / f'{scene_id}_vh_clean_2.ply')
path_to_out_intrinsic = out_path / 'intrinsic'
path_to_out_intrinsic.mkdir(parents=True, exist_ok=True)
shutil.copy(scannet_processed_path / 'intrinsic' / 'intrinsic_color.txt', path_to_out_intrinsic / 'intrinsic_color.txt')
shutil.copy(scannet_processed_path / 'intrinsic' / 'intrinsic_depth.txt', path_to_out_intrinsic / 'intrinsic_depth.txt')
shutil.copy(scannet_processed_path / 'intrinsic' / 'extrinsic_color.txt', path_to_out_intrinsic / 'extrinsic_color.txt')
shutil.copy(scannet_processed_path / 'intrinsic' / 'extrinsic_depth.txt', path_to_out_intrinsic / 'extrinsic_depth.txt')
# base_intrinsic_path = f'/home/jovyan/users/lemeshko/mmdetection3d/data/scannet/posed_images/{scene_id}/intrinsic.txt'
# base_intrinsic = np.loadtxt(base_intrinsic_path)
# np.savetxt(path_to_out_intrinsic / 'intrinsic_color.txt', resize_intrinsics(Ks.mean(0), (512, 384), (1280, 960)))
# np.savetxt(path_to_out_intrinsic / 'intrinsic_depth.txt', resize_intrinsics(Ks.mean(0), (512, 384), (640, 480)))
# np.savetxt(path_to_out_intrinsic / 'extrinsic_color.txt', np.eye(4))
# np.savetxt(path_to_out_intrinsic / 'extrinsic_depth.txt', np.eye(4))
path_to_out_color = out_path / 'color'
path_to_out_color.mkdir(parents=True, exist_ok=True)
path_to_out_depth = out_path / 'depth'
path_to_out_depth.mkdir(parents=True, exist_ok=True)
path_to_out_pose = out_path / 'pose'
path_to_out_pose.mkdir(parents=True, exist_ok=True)
image_file = Path(image_names[0])
gt_depth = image_file.parent.parent / 'depth_maps' / (image_file.stem + ".png")
gt_pose = image_file.parent.parent / 'poses' / (image_file.stem + ".txt")
print(gt_depth)
# gt_depth = cv2.imread(str(gt_depth), -1) / 1000.0
# gt_pose = np.loadtxt(str(gt_pose))
# h, w = gt_depth.shape
# depth = depths[0].numpy().astype(np.float32)
# depth = cv2.resize(depth, (w, h), interpolation=cv2.INTER_NEAREST)
# valid_mask = (gt_depth > 1e-10) & (depth > 1e-10)
# depth = depth[valid_mask]
# gt_depth = gt_depth[valid_mask]
# scale = np.median(gt_depth / depth)
# first_pose = np.linalg.inv(poses[0].numpy())
for i, image_name in enumerate(image_names):
image_name = Path(image_name)
def get_image_name(image_name):
return str(int(image_name.stem))
shutil.copy(Path("/home/jovyan/users/lemeshko/mmdetection3d/data/scannet/posed_images") / scene_id / image_name.name, path_to_out_color / (get_image_name(image_name) + '.jpg'))
depth_map = depths[i].numpy() * 1000
depth_map = cv2.resize(depth_map, (640, 480))
depth_map = depth_map.astype(np.uint16)
cv2.imwrite(str(path_to_out_depth / (get_image_name(image_name) + '.png')), depth_map)
pose = poses[i].numpy()
np.savetxt(str(path_to_out_pose / (get_image_name(image_name) + '.txt')), pose)
if __name__ == "__main__":
pkl_path = '/home/jovyan/users/lemeshko/mmdetection3d/data/scannetOV20/scannet_oneformer3d_infos_val.pkl'
with open(pkl_path, 'rb') as file:
lst = pickle.load(file)['data_list']
scenes = sorted([val['lidar_points']['lidar_path'][:-4] for val in lst])
process_map(prepare_scene, scenes, chunksize=10) |