datasets_preprocess/scannet_traj_seq_gen.py

import sys, os
from copy import deepcopy

import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__))))

if 'META_INTERNAL' in os.environ.keys() and os.environ['META_INTERNAL'] == "False":
    generate_html = None
    from dust3r.dummy_io import *
else:
    from meta_internal.io import *

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--div", default=1, type=int)
parser.add_argument("--node-no", default=0, type=int)
parser.add_argument("--hardness", type=str)
parser.add_argument("--n-v", type = int)
parser.add_argument("--n-render", type = int)
parser.add_argument("--data-type", type = str)
parser.add_argument("--split", type = str, default = "all")
parser.add_argument("--render-overlap", type = float, default = 0.95)
parser.add_argument("--n-tuple-per-scene", type = int, default = 1000)
args = parser.parse_args()
print('args', args)
node_no = args.node_no

import torch.distributed as dist

def init_distributed():
    if not dist.is_initialized():
        dist.init_process_group(backend='gloo', init_method='env://')

def get_rank():
    if not dist.is_initialized():
        return 0
    return dist.get_rank()

init_distributed()
rank = get_rank()


cuda_id = int(rank) % 8
device = f"cuda:{cuda_id}"
print('cuda id', cuda_id)

import PIL
import numpy as np
import torch
import glob, sys
import json
import imageio
import cv2

from dust3r.utils.geometry import depthmap_to_absolute_camera_coordinates

train_name_list_path = g_pathmgr.get_local_path(global_variable.train_name_list_path)
train_name_list = json.load(open(global_variable.train_name_list_path, 'r'))

n_v = args.n_v
n_render = args.n_render
# n_v = 4
# n_render = 1
# data_type = "scannet"
data_type = args.data_type
hardness = args.hardness
split = args.split

n_tuple_per_scene = args.n_tuple_per_scene
n_try = 1000
target_n = 100
n_try_scene = 100000

render_range = [
    [0, 4 * (i + 1)]
    for i in range(6)
]

cover_hardness = {
    'hard': [0.1, 0.4],
    'easy': [0.3, 0.7],
    'easier': [0.3, 1.0],
}
if "_" in hardness:
    mi, ma = hardness.split('_')
    cover_hardness[hardness] = [float(mi), float(ma)]

n_inference = n_v - n_render


# dataset_name = "scannet_large_easy_12"
# dataset_name = "scannet_large_easier_12"
# dataset_name = "scannet_large_12"
# dataset_name = f"scannetpp_large_easier_{n_v}"
# dataset_name = f"scannetpp_large_easy_{n_v}"
dataset_name = f"{data_type}_{hardness}_{n_v}_seq_{split}"

if data_type == "scannet":
    nn_shrink_rate = 3
elif data_type == "scannetpp":
    nn_shrink_rate = 9

def compare(a, b):
    a = int(os.path.basename(a)[:-4])
    b = int(os.path.basename(b)[:-4])
    return a < b

def key(a):
    return int(os.path.basename(a)[:-4])

def min_dis(A, B):

    A = torch.from_numpy(A).reshape(-1, 3).to(device)
    B = torch.from_numpy(B).reshape(-1, 3).to(device)
    from pytorch3d.ops import knn_points
    dis, _, _ = knn_points(B[None], A[None]) # B querying in A, dis: [1, B.shape[0], 1]
    return dis[0,:,0]

def cover(pc1_, pc2_): # querying pc2 in pc1
    import numpy as np
    pc1 = pc1_.reshape(-1, 3)
    pc2 = pc2_.reshape(-1, 3)
    
    distances = min_dis(pc1, pc2)
    
    thres = 0.015 * nn_shrink_rate
    return distances[(distances > 0) * (distances < thres)].shape[0] / distances.shape[0]

def get_score(pc1, pc2):
    return (cover(pc1, pc2) + cover(pc2, pc1)) / 2

def extract_valid_frames(valid_frames_ss):
    valid_frames_name = []
    for x in valid_frames_ss:
        x = x.split(' ')
        if len(x) <= 1:
            break
        if int(x[2]) != 0:
            print('get bad frame', x)
            # exit(0)
        valid_frames_name.append(int(x[1]))
    return valid_frames_name

def qvec2rotmat(qvec):
    return np.array([
        [1 - 2 * qvec[2]**2 - 2 * qvec[3]**2,
         2 * qvec[1] * qvec[2] - 2 * qvec[0] * qvec[3],
         2 * qvec[3] * qvec[1] + 2 * qvec[0] * qvec[2]],
        [2 * qvec[1] * qvec[2] + 2 * qvec[0] * qvec[3],
         1 - 2 * qvec[1]**2 - 2 * qvec[3]**2,
         2 * qvec[2] * qvec[3] - 2 * qvec[0] * qvec[1]],
        [2 * qvec[3] * qvec[1] - 2 * qvec[0] * qvec[2],
         2 * qvec[2] * qvec[3] + 2 * qvec[0] * qvec[1],
         1 - 2 * qvec[1]**2 - 2 * qvec[2]**2]])

def qt2w2c(q, t) -> np.ndarray:
    R = qvec2rotmat(q)
    world2cam = np.eye(4)
    world2cam[:3, :3] = R
    world2cam[:3, 3] = t
    return world2cam

def extract_image_txt(lines):
    c2ws = []
    frame_names = []
    for x in lines:
        if ".jpg" in x:
            x = x.split(' ')
            q = [float(y) for y in x[1:5]]
            t = [float(y) for y in x[5:8]]
            frame = int(x[9].replace('frame_', '').replace('.jpg', ''))
            c2ws.append(np.linalg.inv(qt2w2c(q, t)))
            frame_names.append(frame)
            # print('extract', x, q, t, frame)
    
    return c2ws, frame_names

def extract_K(lines):
    ss = lines[3].split(' ')
    K = np.eye(3)
    K[0,0] = float(ss[4])
    K[1,1] = float(ss[5])
    K[0,2] = float(ss[6])
    K[1,2] = float(ss[7])
    return K

def get_scene_list():
    if data_type == "scannet":
        scan_list = glob.glob("/vol22/zt15/scannet/scans/*")
        scan_list.sort()
        # scan_list = np.load(g_pathmgr.get_local_path(global_variable.scannet_scan_list))
    elif data_type == "scannetpp":
        scan_list = np.load(g_pathmgr.get_local_path(global_variable.scannetpp_scan_list))
        scan_list = [os.path.join(global_variable.scannetpp_data_dir + "/data", x.split('/')[-1]) for x in scan_list]
    return scan_list

def filter_test(scene_list):
    scene_list_new = []
    for x in scene_list:
        scene_name = os.path.basename(x)
        if scene_name not in train_name_list:
            scene_list_new.append(x)
    return scene_list_new

def main():
    
    scene_list = get_scene_list()
    if args.split == "test":
        scene_list = filter_test(scene_list)
    # exit(0)
    tuple_done = 0
    for scene_id, scene_name in enumerate(scene_list[::1]):
        if scene_id % args.div != cuda_id + args.node_no * 8:
            continue
        print('trying', scene_id, cuda_id, cuda_id + args.node_no * 8, scene_name)
        skip_scene = False
        
        
        target_scene_name = f"{global_variable.metadata_dir}/{dataset_name}/{os.path.basename(scene_name)}"
        last_name = f"{target_scene_name}/{str(n_tuple_per_scene - 1).zfill(6)}_extra.pt"
        if g_pathmgr.exists(last_name):
            print('exist', last_name)
            continue
        print('doing', scene_id, cuda_id, cuda_id + args.node_no * 8, scene_name)
        
        g_pathmgr.mkdirs(target_scene_name)
        
        print('scene_name', scene_name)
        pose_raw_list = []
        pose_list = []
        if data_type == "scannet":
            valid_frame_name = os.path.join(scene_name, "valid_frames.txt")
            if g_pathmgr.exists(valid_frame_name):
                with g_pathmgr.open(valid_frame_name, 'r') as f:
                    valid_frames_ss = f.readlines()
                valid_frames_name = extract_valid_frames(valid_frames_ss)
                rgb_suff = "jpg"
                if not g_pathmgr.exists(os.path.join(scene_name, "frames", "color", f"{valid_frames_name[0]}.{rgb_suff}")):
                    print('not exist', os.path.join(scene_name, "frames", "color", f"{valid_frames_name[0]}.{rgb_suff}"))
                    rgb_suff = "png"
            else:
                valid_frames_name = glob.glob(f"{scene_name}/frames/color/*")
                valid_frames_name.sort()
                rgb_suff = valid_frames_name[0].split('.')[-1]
                valid_frames_name = [os.path.basename(x).split('.')[0] for x in valid_frames_name]
            rgb_list = [os.path.join(scene_name, "frames", "color", f"{valid_frame}.{rgb_suff}") for valid_frame in valid_frames_name]
            depth_list = [os.path.join(scene_name, "frames", "depth", f"{valid_frame}.png") for valid_frame in valid_frames_name]
            pose_list = [os.path.join(scene_name, "frames", "pose", f"{valid_frame}.txt") for valid_frame in valid_frames_name]
            intrinsic_depth = np.loadtxt(g_pathmgr.get_local_path(os.path.join(scene_name, "frames", "intrinsic", "intrinsic_depth.txt")))
            intrinsic_depth_original = deepcopy(intrinsic_depth)

        elif data_type == "scannetpp":
            poses, frame_names = extract_image_txt(open(g_pathmgr.get_local_path(os.path.join(scene_name, "iphone", "colmap", "images.txt")), "r").readlines())
            intrinsic_depth = extract_K(open(g_pathmgr.get_local_path(os.path.join(scene_name, "iphone", "colmap", "cameras.txt")), "r").readlines())
            intrinsic_depth_original = deepcopy(intrinsic_depth)
            intrinsic_depth[:2] *= 1 / nn_shrink_rate
            rgb_list = []
            depth_list = []
            pose_raw_list = [x for x in poses]
            depth_dir = os.path.join(f"{global_variable.scannetpp_dir}/render", scene_name.split('/')[-1])
            for frame in frame_names:
                rgb_name   = os.path.join(scene_name, "iphone", "rgb", f"frame_{str(frame).zfill(6)}.jpg")
                depth_name = os.path.join(depth_dir, "iphone", "render_depth", f"frame_{str(frame).zfill(6)}.png")
                rgb_list.append(rgb_name)
                depth_list.append(depth_name)

        # step = max(len(rgb_list) // target_n, 1)
        id_list = [int(i / target_n * len(rgb_list)) for i in range(target_n)]
        rgb_list = [rgb_list[id] for id in id_list]
        depth_list = [depth_list[id] for id in id_list]
        if pose_list:
            pose_list = [pose_list[id] for id in id_list]
        if pose_raw_list:
            pose_raw_list = [pose_raw_list[id] for id in id_list]
        
        print('loading:', len(rgb_list))
        # rgb_preload = []
        depth_preload = []
        pose_preload = []
        valid_set = []
        for id in range(len(rgb_list)):
            try:
                rgb = imageio.imread(g_pathmgr.get_local_path(rgb_list[id])).astype(np.float32) / 255
                depth = imageio.imread(g_pathmgr.get_local_path(depth_list[id])).astype(np.float32) / 1000
                if len(pose_list) > 0:
                    pose = np.loadtxt(g_pathmgr.get_local_path(pose_list[id]))
                else:
                    pose = pose_raw_list[id]
                # rgb_preload.append(rgb)
                depth_preload.append(depth)
                pose_preload.append(pose)
                valid_set.append(id)
            except:
                pass
            print('loading', id)
        if len(depth_preload) < len(rgb_list) * 0.7:
            continue
        rgb_list = [rgb_list[x] for x in valid_set]
        depth_list = [depth_list[x] for x in valid_set]
        if len(pose_list):
            pose_list = [pose_list[x] for x in valid_set]
        if len(pose_raw_list):
            pose_raw_list = [pose_raw_list[x] for x in valid_set]
        print('load Done')
        for cnt in range(n_tuple_per_scene):

            if skip_scene:
                break
            
            cnt_all_test = 0
            current_name = f"{target_scene_name}/{str(cnt).zfill(6)}_extra.pt"
            if g_pathmgr.exists(current_name):
                print('secEx', current_name)
                continue
            # generate n_v views randomly from rgb_list and depth_list, do not replace.

            id_list = [None for i in range(n_v)]
            focal = (intrinsic_depth[0,0] + intrinsic_depth[1, 1]) / 2
            pcd_all = []
            C = np.zeros((n_v, n_v))
            i = 0
            rgbs = []
            cover_mask_all = [None for i in range(n_render)]
            start_id = np.random.randint(target_n - n_v)
            while i < n_v:
                            
                if skip_scene:
                    break

                try_cnt = 0

                while 1:

                    try_cnt += 1
                    id_list[i] = start_id + i # np.random.choice(len(depth_list))
                    cnt_all_test += 1

                    if try_cnt > n_try:
                        print('failed')
                        try_cnt = 0
                        i = -1
                        rgbs = []
                        pcd_all = []
                        C = np.zeros((n_v, n_v))
                        id_list = [None for i in range(n_v)]
                        break
                    
                    if id_list[i] in id_list[:i]:
                        continue

                    print('test', 'scene_id', scene_id, 'cnt', cnt, 'cnt_all_test', cnt_all_test, 'i', i, 'id_list[i]', id_list[i], 'try_cnt', try_cnt) # test 57 608 206 9 184 130
                    if cnt_all_test > n_try_scene:
                        print('failed scene')
                        skip_scene = True
                        break

                    import time
                    t = time.time()
                    id = id_list[i]
                    # rgb = rgb_preload[id]
                    depth = depth_preload[id]
                    depth = cv2.resize(depth, (depth.shape[1] // nn_shrink_rate, depth.shape[0] // nn_shrink_rate))
                    pose = pose_preload[id]
                    print('load time', time.time() - t)
                    t = time.time()
                    pcd, valid_mask = depthmap_to_absolute_camera_coordinates(depth, intrinsic_depth, pose) # pcd (480, 640, 3) (480, 640) 0.9917024739583333
                    # pcd = pcd[::nn_shrink_rate,::nn_shrink_rate]
                    # valid_mask = valid_mask[::nn_shrink_rate,::nn_shrink_rate]
                    pcd_valid = pcd[valid_mask]
                    print('depth calc time', time.time() - t)
                    # print('pcd', pcd.shape, valid_mask.shape, valid_mask.mean()) # pcd (480, 640, 3) (480, 640) 0.9874479166666666
                    if i < n_inference:
                        score_list = []
                        t = time.time()
                        for j in range(i):
                            score_list.append(get_score(pcd_all[j], pcd_valid))
                        print('nn time', time.time() - t)
                        print(score_list)

                        # if len(score_list) == 0 or (np.max(score_list) > 0.3 and np.max(score_list) < 0.7):
                        # if len(score_list) == 0 or (np.max(score_list) > 0.3 and np.max(score_list) < 1.0):
                        if 1:
                            pcd_all.append(pcd_valid)
                            try_cnt = 0
                            break
                    else:
                        i_render = i - n_inference
                        score_list = []
                        t = time.time()
                        pcd_all_combined = np.concatenate(pcd_all[render_range[i_render][0] : render_range[i_render][1]], 0)
                        # def cover(pc1_, pc2_): # querying pc2 in pc1
                        score = cover(pcd_all_combined, pcd_valid)
                        if score > args.render_overlap:
                            # cover_mask = cover(pcd_all_combined, pcd_original, return_mask = True)
                            # cover_mask = cover_mask.reshape(-1).cpu() * torch.from_numpy(valid_mask_original.reshape(-1)).float()
                            # cover_mask = cover_mask.reshape(depth.shape[0], depth.shape[1])
                            # cover_mask_all[i_render] = cover_mask
                            pcd_all.append(pcd_valid)
                            try_cnt = 0
                            break
                i += 1
                
                # rgb = cv2.resize(rgb, (depth.shape[1], depth.shape[0]))
                # np.save(f"{target_scene_name}/{str(cnt).zfill(6)}_rgb_{i}.npy", rgb)
                # np.save(f"{target_scene_name}/{str(cnt).zfill(6)}_pcd_{i}.npy", pcd)
                # np.save(f"{target_scene_name}/{str(cnt).zfill(6)}_valid_{i}.npy", valid_mask)
                # np.save(f"{target_scene_name}/{str(cnt).zfill(6)}_focal_{i}.npy", focal)
                # np.save(f"{target_scene_name}/{str(cnt).zfill(6)}_pose_{i}.npy", pose)

                # rgb = cv2.resize(rgb, (depth.shape[1], depth.shape[0]))
                # rgb = (rgb * 255).astype(np.uint8)
                # rgbs.append(rgb)
            
            tuple_done += 1
            # rgbs = np.concatenate(rgbs, axis=1)
            # imageio.imwrite(f"{target_scene_name}/{str(cnt).zfill(6)}_rgb.png", rgbs)

            if skip_scene:
                break

            C_avg = []
            for i in range(n_v):
                for j in range(n_v):
                    if i != j:
                        C[i,j] = get_score(pcd_all[i], pcd_all[j])
                        C_avg.append(C[i,j])
                    # print(i,j, C[i,j])
            
            rgb_list_ = [rgb_list[id] for id in id_list]
            depth_list_ = [depth_list[id] for id in id_list]

            if len(pose_list):
                pose_list_ = [pose_list[id] for id in id_list] # from here tomorrow
                pose_info = {"pose_list": pose_list_}
            else:
                pose_raw_list_ = [pose_raw_list[id].tolist() for id in id_list]
                pose_info = {"pose_raw_list": pose_raw_list_}

            extra_info = {
                'C': C.tolist(),
                'C_avg': np.mean(np.array(C_avg)).item(),
                'id_list': id_list,
                'rgb_list': rgb_list_,
                'depth_list': depth_list_,
                'intrinsic_raw': intrinsic_depth_original.tolist(),
                **pose_info,
                }
            print('extra_info', scene_id, scene_name, extra_info)
            torch.save(extra_info, f"{target_scene_name}/{str(cnt).zfill(6)}_extra.pt")
            # with open(f"{target_scene_name}/{str(cnt).zfill(6)}_extra.json", "w") as f:
            #     json.dump(extra_info, f, indent=4)
            
            print('C')
            print(C)
            print('tuple done', tuple_done, cuda_id, id_list)


main()