vismatch 1.1.1__py3-none-any.whl

vismatch/third_party/mast3r/dust3r/datasets_preprocess/preprocess_co3d.py

@@ -0,0 +1,295 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Script to pre-process the CO3D dataset.
+# Usage:
+# python3 datasets_preprocess/preprocess_co3d.py --co3d_dir /path/to/co3d
+# --------------------------------------------------------
+
+import argparse
+import random
+import gzip
+import json
+import os
+import os.path as osp
+
+import torch
+import PIL.Image
+import numpy as np
+import cv2
+
+from tqdm.auto import tqdm
+import matplotlib.pyplot as plt
+
+import path_to_root  # noqa
+import dust3r.datasets.utils.cropping as cropping  # noqa
+
+
+CATEGORIES = [
+    "apple", "backpack", "ball", "banana", "baseballbat", "baseballglove",
+    "bench", "bicycle", "book", "bottle", "bowl", "broccoli", "cake", "car", "carrot",
+    "cellphone", "chair", "couch", "cup", "donut", "frisbee", "hairdryer", "handbag",
+    "hotdog", "hydrant", "keyboard", "kite", "laptop", "microwave",
+    "motorcycle",
+    "mouse", "orange", "parkingmeter", "pizza", "plant", "remote", "sandwich",
+    "skateboard", "stopsign",
+    "suitcase", "teddybear", "toaster", "toilet", "toybus",
+    "toyplane", "toytrain", "toytruck", "tv",
+    "umbrella", "vase", "wineglass",
+]
+CATEGORIES_IDX = {cat: i for i, cat in enumerate(CATEGORIES)}  # for seeding
+
+SINGLE_SEQUENCE_CATEGORIES = sorted(set(CATEGORIES) - set(["microwave", "stopsign", "tv"]))
+
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--category", type=str, default=None)
+    parser.add_argument('--single_sequence_subset', default=False, action='store_true',
+                        help="prepare the single_sequence_subset instead.")
+    parser.add_argument("--output_dir", type=str, default="data/co3d_processed")
+    parser.add_argument("--co3d_dir", type=str, required=True)
+    parser.add_argument("--num_sequences_per_object", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--min_quality", type=float, default=0.5, help="Minimum viewpoint quality score.")
+
+    parser.add_argument("--img_size", type=int, default=512,
+                        help=("lower dimension will be >= img_size * 3/4, and max dimension will be >= img_size"))
+    return parser
+
+
+def convert_ndc_to_pinhole(focal_length, principal_point, image_size):
+    focal_length = np.array(focal_length)
+    principal_point = np.array(principal_point)
+    image_size_wh = np.array([image_size[1], image_size[0]])
+    half_image_size = image_size_wh / 2
+    rescale = half_image_size.min()
+    principal_point_px = half_image_size - principal_point * rescale
+    focal_length_px = focal_length * rescale
+    fx, fy = focal_length_px[0], focal_length_px[1]
+    cx, cy = principal_point_px[0], principal_point_px[1]
+    K = np.array([[fx, 0.0, cx], [0.0, fy, cy], [0.0, 0.0, 1.0]], dtype=np.float32)
+    return K
+
+
+def opencv_from_cameras_projection(R, T, focal, p0, image_size):
+    R = torch.from_numpy(R)[None, :, :]
+    T = torch.from_numpy(T)[None, :]
+    focal = torch.from_numpy(focal)[None, :]
+    p0 = torch.from_numpy(p0)[None, :]
+    image_size = torch.from_numpy(image_size)[None, :]
+
+    R_pytorch3d = R.clone()
+    T_pytorch3d = T.clone()
+    focal_pytorch3d = focal
+    p0_pytorch3d = p0
+    T_pytorch3d[:, :2] *= -1
+    R_pytorch3d[:, :, :2] *= -1
+    tvec = T_pytorch3d
+    R = R_pytorch3d.permute(0, 2, 1)
+
+    # Retype the image_size correctly and flip to width, height.
+    image_size_wh = image_size.to(R).flip(dims=(1,))
+
+    # NDC to screen conversion.
+    scale = image_size_wh.to(R).min(dim=1, keepdim=True)[0] / 2.0
+    scale = scale.expand(-1, 2)
+    c0 = image_size_wh / 2.0
+
+    principal_point = -p0_pytorch3d * scale + c0
+    focal_length = focal_pytorch3d * scale
+
+    camera_matrix = torch.zeros_like(R)
+    camera_matrix[:, :2, 2] = principal_point
+    camera_matrix[:, 2, 2] = 1.0
+    camera_matrix[:, 0, 0] = focal_length[:, 0]
+    camera_matrix[:, 1, 1] = focal_length[:, 1]
+    return R[0], tvec[0], camera_matrix[0]
+
+
+def get_set_list(category_dir, split, is_single_sequence_subset=False):
+    listfiles = os.listdir(osp.join(category_dir, "set_lists"))
+    if is_single_sequence_subset:
+        # not all objects have manyview_dev
+        subset_list_files = [f for f in listfiles if "manyview_dev" in f]
+    else:
+        subset_list_files = [f for f in listfiles if f"fewview_train" in f]
+
+    sequences_all = []
+    for subset_list_file in subset_list_files:
+        with open(osp.join(category_dir, "set_lists", subset_list_file)) as f:
+            subset_lists_data = json.load(f)
+            sequences_all.extend(subset_lists_data[split])
+
+    return sequences_all
+
+
+def prepare_sequences(category, co3d_dir, output_dir, img_size, split, min_quality, max_num_sequences_per_object,
+                      seed, is_single_sequence_subset=False):
+    random.seed(seed)
+    category_dir = osp.join(co3d_dir, category)
+    category_output_dir = osp.join(output_dir, category)
+    sequences_all = get_set_list(category_dir, split, is_single_sequence_subset)
+    sequences_numbers = sorted(set(seq_name for seq_name, _, _ in sequences_all))
+
+    frame_file = osp.join(category_dir, "frame_annotations.jgz")
+    sequence_file = osp.join(category_dir, "sequence_annotations.jgz")
+
+    with gzip.open(frame_file, "r") as fin:
+        frame_data = json.loads(fin.read())
+    with gzip.open(sequence_file, "r") as fin:
+        sequence_data = json.loads(fin.read())
+
+    frame_data_processed = {}
+    for f_data in frame_data:
+        sequence_name = f_data["sequence_name"]
+        frame_data_processed.setdefault(sequence_name, {})[f_data["frame_number"]] = f_data
+
+    good_quality_sequences = set()
+    for seq_data in sequence_data:
+        if seq_data["viewpoint_quality_score"] > min_quality:
+            good_quality_sequences.add(seq_data["sequence_name"])
+
+    sequences_numbers = [seq_name for seq_name in sequences_numbers if seq_name in good_quality_sequences]
+    if len(sequences_numbers) < max_num_sequences_per_object:
+        selected_sequences_numbers = sequences_numbers
+    else:
+        selected_sequences_numbers = random.sample(sequences_numbers, max_num_sequences_per_object)
+
+    selected_sequences_numbers_dict = {seq_name: [] for seq_name in selected_sequences_numbers}
+    sequences_all = [(seq_name, frame_number, filepath)
+                     for seq_name, frame_number, filepath in sequences_all
+                     if seq_name in selected_sequences_numbers_dict]
+
+    for seq_name, frame_number, filepath in tqdm(sequences_all):
+        frame_idx = int(filepath.split('/')[-1][5:-4])
+        selected_sequences_numbers_dict[seq_name].append(frame_idx)
+        mask_path = filepath.replace("images", "masks").replace(".jpg", ".png")
+        frame_data = frame_data_processed[seq_name][frame_number]
+        focal_length = frame_data["viewpoint"]["focal_length"]
+        principal_point = frame_data["viewpoint"]["principal_point"]
+        image_size = frame_data["image"]["size"]
+        K = convert_ndc_to_pinhole(focal_length, principal_point, image_size)
+        R, tvec, camera_intrinsics = opencv_from_cameras_projection(np.array(frame_data["viewpoint"]["R"]),
+                                                                    np.array(frame_data["viewpoint"]["T"]),
+                                                                    np.array(focal_length),
+                                                                    np.array(principal_point),
+                                                                    np.array(image_size))
+
+        frame_data = frame_data_processed[seq_name][frame_number]
+        depth_path = os.path.join(co3d_dir, frame_data["depth"]["path"])
+        assert frame_data["depth"]["scale_adjustment"] == 1.0
+        image_path = os.path.join(co3d_dir, filepath)
+        mask_path_full = os.path.join(co3d_dir, mask_path)
+
+        input_rgb_image = PIL.Image.open(image_path).convert('RGB')
+        input_mask = plt.imread(mask_path_full)
+
+        with PIL.Image.open(depth_path) as depth_pil:
+            # the image is stored with 16-bit depth but PIL reads it as I (32 bit).
+            # we cast it to uint16, then reinterpret as float16, then cast to float32
+            input_depthmap = (
+                np.frombuffer(np.array(depth_pil, dtype=np.uint16), dtype=np.float16)
+                .astype(np.float32)
+                .reshape((depth_pil.size[1], depth_pil.size[0])))
+        depth_mask = np.stack((input_depthmap, input_mask), axis=-1)
+        H, W = input_depthmap.shape
+
+        camera_intrinsics = camera_intrinsics.numpy()
+        cx, cy = camera_intrinsics[:2, 2].round().astype(int)
+        min_margin_x = min(cx, W - cx)
+        min_margin_y = min(cy, H - cy)
+
+        # the new window will be a rectangle of size (2*min_margin_x, 2*min_margin_y) centered on (cx,cy)
+        l, t = cx - min_margin_x, cy - min_margin_y
+        r, b = cx + min_margin_x, cy + min_margin_y
+        crop_bbox = (l, t, r, b)
+        input_rgb_image, depth_mask, input_camera_intrinsics = cropping.crop_image_depthmap(
+            input_rgb_image, depth_mask, camera_intrinsics, crop_bbox)
+
+        # try to set the lower dimension to img_size * 3/4 -> img_size=512 => 384
+        scale_final = ((img_size * 3 // 4) / min(H, W)) + 1e-8
+        output_resolution = np.floor(np.array([W, H]) * scale_final).astype(int)
+        if max(output_resolution) < img_size:
+            # let's put the max dimension to img_size
+            scale_final = (img_size / max(H, W)) + 1e-8
+            output_resolution = np.floor(np.array([W, H]) * scale_final).astype(int)
+
+        input_rgb_image, depth_mask, input_camera_intrinsics = cropping.rescale_image_depthmap(
+            input_rgb_image, depth_mask, input_camera_intrinsics, output_resolution)
+        input_depthmap = depth_mask[:, :, 0]
+        input_mask = depth_mask[:, :, 1]
+
+        # generate and adjust camera pose
+        camera_pose = np.eye(4, dtype=np.float32)
+        camera_pose[:3, :3] = R
+        camera_pose[:3, 3] = tvec
+        camera_pose = np.linalg.inv(camera_pose)
+
+        # save crop images and depth, metadata
+        save_img_path = os.path.join(output_dir, filepath)
+        save_depth_path = os.path.join(output_dir, frame_data["depth"]["path"])
+        save_mask_path = os.path.join(output_dir, mask_path)
+        os.makedirs(os.path.split(save_img_path)[0], exist_ok=True)
+        os.makedirs(os.path.split(save_depth_path)[0], exist_ok=True)
+        os.makedirs(os.path.split(save_mask_path)[0], exist_ok=True)
+
+        input_rgb_image.save(save_img_path)
+        scaled_depth_map = (input_depthmap / np.max(input_depthmap) * 65535).astype(np.uint16)
+        cv2.imwrite(save_depth_path, scaled_depth_map)
+        cv2.imwrite(save_mask_path, (input_mask * 255).astype(np.uint8))
+
+        save_meta_path = save_img_path.replace('jpg', 'npz')
+        np.savez(save_meta_path, camera_intrinsics=input_camera_intrinsics,
+                 camera_pose=camera_pose, maximum_depth=np.max(input_depthmap))
+
+    return selected_sequences_numbers_dict
+
+
+if __name__ == "__main__":
+    parser = get_parser()
+    args = parser.parse_args()
+    assert args.co3d_dir != args.output_dir
+    if args.category is None:
+        if args.single_sequence_subset:
+            categories = SINGLE_SEQUENCE_CATEGORIES
+        else:
+            categories = CATEGORIES
+    else:
+        categories = [args.category]
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    for split in ['train', 'test']:
+        selected_sequences_path = os.path.join(args.output_dir, f'selected_seqs_{split}.json')
+        if os.path.isfile(selected_sequences_path):
+            continue
+
+        all_selected_sequences = {}
+        for category in categories:
+            category_output_dir = osp.join(args.output_dir, category)
+            os.makedirs(category_output_dir, exist_ok=True)
+            category_selected_sequences_path = os.path.join(category_output_dir, f'selected_seqs_{split}.json')
+            if os.path.isfile(category_selected_sequences_path):
+                with open(category_selected_sequences_path, 'r') as fid:
+                    category_selected_sequences = json.load(fid)
+            else:
+                print(f"Processing {split} - category = {category}")
+                category_selected_sequences = prepare_sequences(
+                    category=category,
+                    co3d_dir=args.co3d_dir,
+                    output_dir=args.output_dir,
+                    img_size=args.img_size,
+                    split=split,
+                    min_quality=args.min_quality,
+                    max_num_sequences_per_object=args.num_sequences_per_object,
+                    seed=args.seed + CATEGORIES_IDX[category],
+                    is_single_sequence_subset=args.single_sequence_subset
+                )
+                with open(category_selected_sequences_path, 'w') as file:
+                    json.dump(category_selected_sequences, file)
+
+            all_selected_sequences[category] = category_selected_sequences
+        with open(selected_sequences_path, 'w') as file:
+            json.dump(all_selected_sequences, file)

vismatch/third_party/mast3r/dust3r/datasets_preprocess/preprocess_megadepth.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# Preprocessing code for the MegaDepth dataset
+# dataset at https://www.cs.cornell.edu/projects/megadepth/
+# --------------------------------------------------------
+import os
+import os.path as osp
+import collections
+from tqdm import tqdm
+import numpy as np
+os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
+import cv2
+import h5py
+
+import path_to_root  # noqa
+from dust3r.utils.parallel import parallel_threads
+from dust3r.datasets.utils import cropping  # noqa
+
+
+def get_parser():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--megadepth_dir', required=True)
+    parser.add_argument('--precomputed_pairs', required=True)
+    parser.add_argument('--output_dir', default='data/megadepth_processed')
+    return parser
+
+
+def main(db_root, pairs_path, output_dir):
+    os.makedirs(output_dir, exist_ok=True)
+
+    # load all pairs
+    data = np.load(pairs_path, allow_pickle=True)
+    scenes = data['scenes']
+    images = data['images']
+    pairs = data['pairs']
+
+    # enumerate all unique images
+    todo = collections.defaultdict(set)
+    for scene, im1, im2, score in pairs:
+        todo[scene].add(im1)
+        todo[scene].add(im2)
+
+    # for each scene, load intrinsics and then parallel crops
+    for scene, im_idxs in tqdm(todo.items(), desc='Overall'):
+        scene, subscene = scenes[scene].split()
+        out_dir = osp.join(output_dir, scene, subscene)
+        os.makedirs(out_dir, exist_ok=True)
+
+        # load all camera params
+        _, pose_w2cam, intrinsics = _load_kpts_and_poses(db_root, scene, subscene, intrinsics=True)
+
+        in_dir = osp.join(db_root, scene, 'dense' + subscene)
+        args = [(in_dir, img, intrinsics[img], pose_w2cam[img], out_dir)
+                for img in [images[im_id] for im_id in im_idxs]]
+        parallel_threads(resize_one_image, args, star_args=True, front_num=0, leave=False, desc=f'{scene}/{subscene}')
+
+    # save pairs
+    print('Done! prepared all pairs in', output_dir)
+
+
+def resize_one_image(root, tag, K_pre_rectif, pose_w2cam, out_dir):
+    if osp.isfile(osp.join(out_dir, tag + '.npz')):
+        return
+
+    # load image
+    img = cv2.cvtColor(cv2.imread(osp.join(root, 'imgs', tag), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB)
+    H, W = img.shape[:2]
+
+    # load depth
+    with h5py.File(osp.join(root, 'depths', osp.splitext(tag)[0] + '.h5'), 'r') as hd5:
+        depthmap = np.asarray(hd5['depth'])
+
+    # rectify = undistort the intrinsics
+    imsize_pre, K_pre, distortion = K_pre_rectif
+    imsize_post = img.shape[1::-1]
+    K_post = cv2.getOptimalNewCameraMatrix(K_pre, distortion, imsize_pre, alpha=0,
+                                           newImgSize=imsize_post, centerPrincipalPoint=True)[0]
+
+    # downscale
+    img_out, depthmap_out, intrinsics_out, R_in2out = _downscale_image(K_post, img, depthmap, resolution_out=(800, 600))
+
+    # write everything
+    img_out.save(osp.join(out_dir, tag + '.jpg'), quality=90)
+    cv2.imwrite(osp.join(out_dir, tag + '.exr'), depthmap_out)
+
+    camout2world = np.linalg.inv(pose_w2cam)
+    camout2world[:3, :3] = camout2world[:3, :3] @ R_in2out.T
+    np.savez(osp.join(out_dir, tag + '.npz'), intrinsics=intrinsics_out, cam2world=camout2world)
+
+
+def _downscale_image(camera_intrinsics, image, depthmap, resolution_out=(512, 384)):
+    H, W = image.shape[:2]
+    resolution_out = sorted(resolution_out)[::+1 if W < H else -1]
+
+    image, depthmap, intrinsics_out = cropping.rescale_image_depthmap(
+        image, depthmap, camera_intrinsics, resolution_out, force=False)
+    R_in2out = np.eye(3)
+
+    return image, depthmap, intrinsics_out, R_in2out
+
+
+def _load_kpts_and_poses(root, scene_id, subscene, z_only=False, intrinsics=False):
+    if intrinsics:
+        with open(os.path.join(root, scene_id, 'sparse', 'manhattan', subscene, 'cameras.txt'), 'r') as f:
+            raw = f.readlines()[3:]  # skip the header
+
+        camera_intrinsics = {}
+        for camera in raw:
+            camera = camera.split(' ')
+            width, height, focal, cx, cy, k0 = [float(elem) for elem in camera[2:]]
+            K = np.eye(3)
+            K[0, 0] = focal
+            K[1, 1] = focal
+            K[0, 2] = cx
+            K[1, 2] = cy
+            camera_intrinsics[int(camera[0])] = ((int(width), int(height)), K, (k0, 0, 0, 0))
+
+    with open(os.path.join(root, scene_id, 'sparse', 'manhattan', subscene, 'images.txt'), 'r') as f:
+        raw = f.read().splitlines()[4:]  # skip the header
+
+    extract_pose = colmap_raw_pose_to_principal_axis if z_only else colmap_raw_pose_to_RT
+
+    poses = {}
+    points3D_idxs = {}
+    camera = []
+
+    for image, points in zip(raw[:: 2], raw[1:: 2]):
+        image = image.split(' ')
+        points = points.split(' ')
+
+        image_id = image[-1]
+        camera.append(int(image[-2]))
+
+        # find the principal axis
+        raw_pose = [float(elem) for elem in image[1: -2]]
+        poses[image_id] = extract_pose(raw_pose)
+
+        current_points3D_idxs = {int(i) for i in points[2:: 3] if i != '-1'}
+        assert -1 not in current_points3D_idxs, bb()
+        points3D_idxs[image_id] = current_points3D_idxs
+
+    if intrinsics:
+        image_intrinsics = {im_id: camera_intrinsics[cam] for im_id, cam in zip(poses, camera)}
+        return points3D_idxs, poses, image_intrinsics
+    else:
+        return points3D_idxs, poses
+
+
+def colmap_raw_pose_to_principal_axis(image_pose):
+    qvec = image_pose[: 4]
+    qvec = qvec / np.linalg.norm(qvec)
+    w, x, y, z = qvec
+    z_axis = np.float32([
+        2 * x * z - 2 * y * w,
+        2 * y * z + 2 * x * w,
+        1 - 2 * x * x - 2 * y * y
+    ])
+    return z_axis
+
+
+def colmap_raw_pose_to_RT(image_pose):
+    qvec = image_pose[: 4]
+    qvec = qvec / np.linalg.norm(qvec)
+    w, x, y, z = qvec
+    R = np.array([
+        [
+            1 - 2 * y * y - 2 * z * z,
+            2 * x * y - 2 * z * w,
+            2 * x * z + 2 * y * w
+        ],
+        [
+            2 * x * y + 2 * z * w,
+            1 - 2 * x * x - 2 * z * z,
+            2 * y * z - 2 * x * w
+        ],
+        [
+            2 * x * z - 2 * y * w,
+            2 * y * z + 2 * x * w,
+            1 - 2 * x * x - 2 * y * y
+        ]
+    ])
+    # principal_axis.append(R[2, :])
+    t = image_pose[4: 7]
+    # World-to-Camera pose
+    current_pose = np.eye(4)
+    current_pose[: 3, : 3] = R
+    current_pose[: 3, 3] = t
+    return current_pose
+
+
+if __name__ == '__main__':
+    parser = get_parser()
+    args = parser.parse_args()
+    main(args.megadepth_dir, args.precomputed_pairs, args.output_dir)