3D metric reconstruction using Droid-slam and Metric3D
import sys
sys.path.append('droid_slam')

from tqdm import tqdm
import numpy as np
import torch
import cv2
import os
import json
import argparse
import glob
from PIL import Image as PILImage
from PIL.ExifTags import GPSTAGS, TAGS
from pathlib import Path
from droid import Droid
from scipy.spatial.transform import Rotation

def pose_matrix_from_quaternion(pvec):
    """ convert 4x4 pose matrix to (t, q) """
    pose = np.eye(4)
    pose[:3,:3] = Rotation.from_quat(pvec[3:]).as_matrix()
    pose[:3, 3] = pvec[:3]
    return pose

def show_image(image):
    image = image.permute(1, 2, 0).cpu().numpy()
    cv2.imshow('image', image / 255.0)
    cv2.waitKey(1)

IMG_EXTENSIONS = ["png", "PNG", "jpg", "JPG"]
def get_sorted_image_names_in_dir(dir_path: str):
    image_paths = []
    for extension in IMG_EXTENSIONS:
        search_path = os.path.join(dir_path, f"*.{extension}")
        image_paths.extend(glob.glob(search_path))

    return sorted(image_paths)


def __compute_sensor_width_from_exif(exif_data) -> float:
        """Compute sensor_width_mm from `ExifImageWidth` tag,

        Equation: sensor_width = pixel_x_dim / focal_plane_x_res * unit_conversion_factor

        Returns:
            sensor_width_mm.
        """

        sensor_width_mm = 0.0

        # Read `ExifImageWidth` and `FocalPlaneXResolution`.
        pixel_x_dim = exif_data.get("ExifImageWidth")
        focal_plane_x_res = exif_data.get("FocalPlaneXResolution")
        focal_plane_res_unit = exif_data.get("FocalPlaneResolutionUnit")
        if (
            pixel_x_dim is not None
            and pixel_x_dim > 0
            and focal_plane_x_res is not None
            and focal_plane_x_res > 0
            and focal_plane_res_unit is not None
            and focal_plane_res_unit > 0
        ):
            ccd_width = pixel_x_dim / focal_plane_x_res

            if focal_plane_res_unit == CENTIMETERS_FOCAL_PLANE_RES_UNIT:
                sensor_width_mm = ccd_width * 10.0  # convert cm to mm
            elif focal_plane_res_unit == INCHES_FOCAL_PLANE_RES_UNIT:
                sensor_width_mm = ccd_width * MILLIMETERS_PER_INCH  # convert inch to mm

        return sensor_width_mm
    
def get_intrinsics_from_exif(exif_data, img_w_px, img_h_px):
    """Constructs the camera intrinsics from exif tag.

    Equation: focal_px=max(w_px,h_px)∗focal_mm / ccdw_mm

    Ref:
    - https://www.awaresystems.be/imaging/tiff/tifftags/privateifd/exif.html
    - https://github.com/colmap/colmap/blob/e3948b2098b73ae080b97901c3a1f9065b976a45/src/util/bitmap.cc#L282
    - https://openmvg.readthedocs.io/en/latest/software/SfM/SfMInit_ImageListing/
    - https://photo.stackexchange.com/questions/40865/how-can-i-get-the-image-sensor-dimensions-in-mm-to-get-circle-of-confusion-from # noqa: E501

    Returns:
        intrinsics matrix (3x3).
    """
    max_size = max(img_w_px, img_h_px)

    # Initialize principal point.
    center_x = img_w_px / 2
    center_y = img_h_px / 2

    # Initialize focal length as None.
    focal_length_px = None

    # Read from `FocalLengthIn35mmFilm`.
    focal_length_35_mm = exif_data.get("FocalLengthIn35mmFilm")
    if focal_length_35_mm is not None and focal_length_35_mm > 0:
        focal_length_px = focal_length_35_mm / 35.0 * max_size
    else:
        # Read from `FocalLength` mm.
        focal_length_mm = exif_data.get("FocalLength")
        if focal_length_mm is None or focal_length_mm <= 0:
            return None

        # Compute sensor width, either from database or from EXIF.
        sensor_width_mm = __compute_sensor_width_from_exif(exif_data)
        if sensor_width_mm > 0.0:
            focal_length_px = focal_length_mm / sensor_width_mm * max_size

    if focal_length_px is None or focal_length_px <= 0.0:
        return None

    return focal_length_px, focal_length_px, float(center_x), float(center_y)

def load_image(img_path: str, return_intrinsics=False):
    """Load the image from disk.

    Notes: EXIF is read as a map from (tag_id, value) where tag_id is an integer.
    In order to extract human-readable names, we use the lookup table TAGS or GPSTAGS.
    Images will be converted to RGB if in a different format.

    Args:
        img_path (str): the path of image to load.

    Returns:
        loaded image in RGB format.
    """
    original_image = PILImage.open(img_path)

    exif_data = original_image._getexif()
    if exif_data is not None:
        parsed_data = {}
        for tag_id, value in exif_data.items():
            # extract the human readable tag name
            if tag_id in TAGS:
                tag_name = TAGS.get(tag_id)
            elif tag_id in GPSTAGS:
                tag_name = GPSTAGS.get(tag_id)
            else:
                tag_name = tag_id
            parsed_data[tag_name] = value

        exif_data = parsed_data
        
    img_fname = Path(img_path).name
    image_array = np.array(original_image)
    bgr_image_array = image_array[:, :, ::-1]
    if return_intrinsics:
        return bgr_image_array, get_intrinsics_from_exif(exif_data, original_image.size[0], original_image.size[1])
    else:
        return bgr_image_array

def image_stream(data, stride, use_depth=False):
    """ image generator """
    all_image_paths = get_sorted_image_names_in_dir(data)
    num_all_imgs = len(all_image_paths)
    _, (fx, fy, cx, cy) = load_image(all_image_paths[0], return_intrinsics=True)
    
    for t, imfile in enumerate(all_image_paths):
        image = load_image(imfile)
        h0, w0, _ = image.shape
        h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))
        w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))

        image = cv2.resize(image, (w1, h1))
        image = image[:h1-h1%8, :w1-w1%8]
        if use_depth:
            depth = depth_model.infer_pil(image)

        image = torch.as_tensor(image).permute(2, 0, 1)

        intrinsics = torch.as_tensor([fx, fy, cx, cy])
        intrinsics[0::2] *= (w1 / w0)
        intrinsics[1::2] *= (h1 / h0)

        if use_depth:
            depth = torch.as_tensor(depth)
            yield t, image[None], intrinsics, depth
        else:
            yield t, image[None], intrinsics, None

def save_reconstruction(droid, output_dir):
    t = droid.video.counter.value
    tstamps = droid.video.tstamp[:t].cpu().numpy()
    images = droid.video.images[:t].cpu().numpy() # [N, 3, 384, 512]
    disps = droid.video.disps_up[:t].cpu().numpy() # [N, 384, 512]
    poses = droid.video.poses[:t].cpu().numpy() # [N, 7]
    intrinsics = droid.video.intrinsics[:t].cpu().numpy() # [N, 4]
    
    # transfer [fx fy cx cy] to 4x4 intrinsics matrix
    intrinsics = intrinsics[0] # [fx fy cx cy]
    intrinsics_4x4 = np.eye(4)
    intrinsics_4x4[0, 0] = intrinsics[0]  # fx
    intrinsics_4x4[1, 1] = intrinsics[1]  # fy
    intrinsics_4x4[0, 2] = intrinsics[2]  # cx
    intrinsics_4x4[1, 2] = intrinsics[3]  # cy

    os.makedirs(f'{output_dir}/pose', exist_ok=True)
    os.makedirs(f'{output_dir}/color', exist_ok=True)
    os.makedirs(f'{output_dir}/depth', exist_ok=True)
    os.makedirs(f'{output_dir}/intrinsic', exist_ok=True)

    transformed_data = {
        "w": int(images[0].shape[2]),
        "h": int(images[0].shape[1]),
        "fl_x": float(intrinsics[0]) * 8.0,
        "fl_y": float(intrinsics[1]) * 8.0,
        "cx": float(intrinsics[2]) * 8.0,
        "cy": float(intrinsics[3]) * 8.0,
        "k1": 0,
        "k2": 0,
        "p1": 0,
        "p2": 0,
        "camera_model": "OPENCV",
        "frames": []
    }

    # save images, poses, intrinsics
    for _id  in tqdm(range(t)):
        pose = np.linalg.inv(pose_matrix_from_quaternion(poses[_id]))
        pose[:, 2] *= -1
        pose[:, 1] *= -1
        image = images[_id]
        image = image.transpose(1, 2, 0)
        depth = (disps[_id] * 1000).astype(np.uint16)
        np.savetxt(f'{output_dir}/pose/{_id}.txt', pose)
        cv2.imwrite(f'{output_dir}/color/{_id}.jpg', image)
        cv2.imwrite(f'{output_dir}/depth/{_id}.png', depth)
        np.savetxt(f'{output_dir}/intrinsic/intrinsic_color.txt', intrinsics_4x4)
        np.savetxt(f'{output_dir}/intrinsic/intrinsic_depth.txt', intrinsics_4x4)

        transformed_frame = {
            "file_path": f"color/{_id}.png",
            "transform_matrix": pose.tolist(),
        }
        transformed_data['frames'].append(transformed_frame)
        output_data_path = os.path.join(output_dir, 'transforms.json')
        with open(output_data_path, 'w') as f:
            f.write(json.dumps(transformed_data, indent=4))
        
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=str, help="path to data directory")
    parser.add_argument("--output_dir", help="path to saved reconstruction")
    parser.add_argument("--t0", default=0, type=int, help="starting frame")
    parser.add_argument("--stride", default=3, type=int, help="frame stride")

    parser.add_argument("--weights", default="droid.pth")
    parser.add_argument("--buffer", type=int, default=512)
    parser.add_argument("--image_size", default=[240, 320])
    parser.add_argument("--disable_vis", action="store_true")

    parser.add_argument("--beta", type=float, default=0.3, help="weight for translation / rotation components of flow")
    parser.add_argument("--filter_thresh", type=float, default=2.4, help="how much motion before considering new keyframe")
    parser.add_argument("--warmup", type=int, default=8, help="number of warmup frames")
    parser.add_argument("--keyframe_thresh", type=float, default=4.0, help="threshold to create a new keyframe")
    parser.add_argument("--frontend_thresh", type=float, default=16.0, help="add edges between frames whithin this distance")
    parser.add_argument("--frontend_window", type=int, default=25, help="frontend optimization window")
    parser.add_argument("--frontend_radius", type=int, default=2, help="force edges between frames within radius")
    parser.add_argument("--frontend_nms", type=int, default=1, help="non-maximal supression of edges")

    parser.add_argument("--backend_thresh", type=float, default=22.0)
    parser.add_argument("--backend_radius", type=int, default=2)
    parser.add_argument("--backend_nms", type=int, default=3)
    parser.add_argument("--upsample", action="store_true")
    parser.add_argument("--use_depth", action="store_true")
    args = parser.parse_args()

    args.stereo = False
    torch.multiprocessing.set_start_method('spawn')

    if args.use_depth:
        DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
        depth_model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to(DEVICE).eval()

    droid = None

    # need high resolution depths
    if args.output_dir is not None:
        args.upsample = True

    tstamps = []
    for (t, image, intrinsics, depth) in tqdm(image_stream(args.data, args.stride, args.use_depth)):
        if t < args.t0:
            continue

        if not args.disable_vis:
            show_image(image[0])

        if droid is None:
            args.image_size = [image.shape[2], image.shape[3]]
            droid = Droid(args)
        
        droid.track(t, image, depth, intrinsics=intrinsics)

    if args.output_dir is not None:
        save_reconstruction(droid, args.output_dir)


demo.py
hugoycj/DROID-SLAM-Metric3D
import sys
sys.path.append('droid_slam')

from tqdm import tqdm
import numpy as np
import torch
import cv2
import os
import json
import argparse
from evo.core import sync

import evo
from evo.core.trajectory import PoseTrajectory3D
from evo.tools import file_interface
from evo.core import sync
import evo.main_ape as main_ape
from evo.core.metrics import PoseRelation

from droid import Droid
from scipy.spatial.transform import Rotation

def pose_matrix_from_quaternion(pvec):
    """ convert 4x4 pose matrix to (t, q) """
    pose = np.eye(4)
    pose[:3,:3] = Rotation.from_quat(pvec[3:]).as_matrix()
    pose[:3, 3] = pvec[:3]
    return pose

def show_image(image):
    image = image.permute(1, 2, 0).cpu().numpy()
    cv2.imshow('image', image / 255.0)
    cv2.waitKey(1)

def image_stream(data, use_depth=False, stride=0):
    """ image generator """
    # Read json from {data}/transforms.json
    with open(os.path.join(data, 'transforms.json'), 'r') as f:
        transforms = json.load(f)
    fx, fy = transforms['fl_x'], transforms['fl_y']
    cx, cy = transforms['cx'], transforms['cy']
    
    image_list = []
    for _frame in transforms['frames']:
        image_list.append(_frame['file_path'])
    
    for t, imfile in enumerate(image_list):
        image = cv2.imread(os.path.join(data, imfile))
        h0, w0, _ = image.shape
        h1 = int(h0 * np.sqrt((384 * 512) / (h0 * w0)))
        w1 = int(w0 * np.sqrt((384 * 512) / (h0 * w0)))

        image = cv2.resize(image, (w1, h1))
        image = image[:h1-h1%8, :w1-w1%8]
        if use_depth:
            depth = depth_model.infer_pil(image)

        image = torch.as_tensor(image).permute(2, 0, 1)

        intrinsics = torch.as_tensor([fx, fy, cx, cy])
        intrinsics[0::2] *= (w1 / w0)
        intrinsics[1::2] *= (h1 / h0)

        if use_depth:
            depth = torch.as_tensor(depth)
            yield t, image[None], intrinsics, depth
        else:
            yield t, image[None], intrinsics, None

def save_reconstruction(droid, output_dir):
    t = droid.video.counter.value
    tstamps = droid.video.tstamp[:t].cpu().numpy()
    images = droid.video.images[:t].cpu().numpy() # [N, 3, 384, 512]
    disps = droid.video.disps_up[:t].cpu().numpy() # [N, 384, 512]
    poses = droid.video.poses[:t].cpu().numpy() # [N, 7]
    intrinsics = droid.video.intrinsics[:t].cpu().numpy() # [N, 4]
    
    # transfer [fx fy cx cy] to 4x4 intrinsics matrix
    intrinsics = intrinsics[0] # [fx fy cx cy]
    intrinsics_4x4 = np.eye(4)
    intrinsics_4x4[0, 0] = intrinsics[0]  # fx
    intrinsics_4x4[1, 1] = intrinsics[1]  # fy
    intrinsics_4x4[0, 2] = intrinsics[2]  # cx
    intrinsics_4x4[1, 2] = intrinsics[3]  # cy

    os.makedirs(f'{output_dir}/pose', exist_ok=True)
    os.makedirs(f'{output_dir}/color', exist_ok=True)
    os.makedirs(f'{output_dir}/depth', exist_ok=True)
    os.makedirs(f'{output_dir}/intrinsic', exist_ok=True)

    transformed_data = {
        "w": int(images[0].shape[2]),
        "h": int(images[0].shape[1]),
        "fl_x": float(intrinsics[0]) * 8.0,
        "fl_y": float(intrinsics[1]) * 8.0,
        "cx": float(intrinsics[2]) * 8.0,
        "cy": float(intrinsics[3]) * 8.0,
        "k1": 0,
        "k2": 0,
        "p1": 0,
        "p2": 0,
        "camera_model": "OPENCV",
        "frames": []
    }

    # save images, poses, intrinsics
    for _id  in tqdm(range(t)):
        pose = np.linalg.inv(pose_matrix_from_quaternion(poses[_id]))
        pose[:, 2] *= -1
        pose[:, 1] *= -1
        image = images[_id]
        image = image.transpose(1, 2, 0)
        depth = (disps[_id] * 1000).astype(np.uint16)
        np.savetxt(f'{output_dir}/pose/{_id}.txt', pose)
        cv2.imwrite(f'{output_dir}/color/{_id}.jpg', image)
        cv2.imwrite(f'{output_dir}/depth/{_id}.png', depth)
        np.savetxt(f'{output_dir}/intrinsic/intrinsic_color.txt', intrinsics_4x4)
        np.savetxt(f'{output_dir}/intrinsic/intrinsic_depth.txt', intrinsics_4x4)

        
        transformed_frame = {
            "file_path": f"color/{_id}.png",
            "transform_matrix": pose.tolist(),
        }
        transformed_data['frames'].append(transformed_frame)
        output_data_path = os.path.join(output_dir, 'transforms.json')
        with open(output_data_path, 'w') as f:
            f.write(json.dumps(transformed_data, indent=4))
        
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--data", type=str, help="path to data directory")
    parser.add_argument("--output_dir", help="path to saved reconstruction")
    parser.add_argument("--t0", default=0, type=int, help="starting frame")
    parser.add_argument("--stride", default=1, type=int, help="frame stride")

    parser.add_argument("--weights", default="droid.pth")
    parser.add_argument("--buffer", type=int, default=512)
    parser.add_argument("--image_size", default=[240, 320])
    parser.add_argument("--disable_vis", action="store_true")

    parser.add_argument("--beta", type=float, default=0.3, help="weight for translation / rotation components of flow")
    parser.add_argument("--filter_thresh", type=float, default=0.1, help="how much motion before considering new keyframe")
    parser.add_argument("--warmup", type=int, default=8, help="number of warmup frames")
    parser.add_argument("--keyframe_thresh", type=float, default=1.0, help="threshold to create a new keyframe")
    parser.add_argument("--frontend_thresh", type=float, default=16.0, help="add edges between frames whithin this distance")
    parser.add_argument("--frontend_window", type=int, default=50, help="frontend optimization window")
    parser.add_argument("--frontend_radius", type=int, default=2, help="force edges between frames within radius")
    parser.add_argument("--frontend_nms", type=int, default=1, help="non-maximal supression of edges")

    parser.add_argument("--backend_thresh", type=float, default=22.0)
    parser.add_argument("--backend_radius", type=int, default=2)
    parser.add_argument("--backend_nms", type=int, default=3)
    parser.add_argument("--upsample", action="store_true")
    parser.add_argument("--use_depth", action="store_true")
    args = parser.parse_args()

    args.stereo = False
    torch.multiprocessing.set_start_method('spawn')

    if args.use_depth:
        DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
        depth_model = torch.hub.load('isl-org/ZoeDepth', "ZoeD_N", pretrained=True).to(DEVICE).eval()

    droid = None

    # need high resolution depths
    if args.output_dir is not None:
        args.upsample = True

    tstamps = []
    for (t, image, intrinsics, depth) in tqdm(image_stream(args.data, args.use_depth, args.stride)):
        if t < args.t0:
            continue

        if not args.disable_vis:
            show_image(image[0])

        if droid is None:
            args.image_size = [image.shape[2], image.shape[3]]
            droid = Droid(args)
        
        droid.track(t, image, depth, intrinsics=intrinsics)
        tstamps.append(t)

    
    if args.output_dir is not None:
        save_reconstruction(droid, args.output_dir)
demo_for_ngp_format.py
hugoycj/DROID-SLAM-Metric3D
import torch
import lietorch
import numpy as np

from lietorch import SE3
from factor_graph import FactorGraph


class DroidBackend:
    def __init__(self, net, video, args):
        self.video = video
        self.update_op = net.update

        # global optimization window
        self.t0 = 0
        self.t1 = 0

        self.upsample = args.upsample
        self.beta = args.beta
        self.backend_thresh = args.backend_thresh
        self.backend_radius = args.backend_radius
        self.backend_nms = args.backend_nms
        
    @torch.no_grad()
    def __call__(self, steps=12):
        """ main update """

        t = self.video.counter.value
        if not self.video.stereo and not torch.any(self.video.disps_sens):
             self.video.normalize()

        graph = FactorGraph(self.video, self.update_op, corr_impl="alt", max_factors=16*t, upsample=self.upsample)

        graph.add_proximity_factors(rad=self.backend_radius, 
                                    nms=self.backend_nms, 
                                    thresh=self.backend_thresh, 
                                    beta=self.beta)

        graph.update_lowmem(steps=steps)
        graph.clear_edges()
        self.video.dirty[:t] = True

droid_slam/droid_backend.py
hugoycj/DROID-SLAM-Metric3D
import torch
import lietorch
import numpy as np

from lietorch import SE3
from factor_graph import FactorGraph


class DroidFrontend:
    def __init__(self, net, video, args):
        self.video = video
        self.update_op = net.update
        self.graph = FactorGraph(video, net.update, max_factors=48, upsample=args.upsample)

        # local optimization window
        self.t0 = 0
        self.t1 = 0

        # frontent variables
        self.is_initialized = False
        self.count = 0

        self.max_age = 25
        self.iters1 = 4
        self.iters2 = 2

        self.warmup = args.warmup
        self.beta = args.beta
        self.frontend_nms = args.frontend_nms
        self.keyframe_thresh = args.keyframe_thresh
        self.frontend_window = args.frontend_window
        self.frontend_thresh = args.frontend_thresh
        self.frontend_radius = args.frontend_radius

    def __update(self):
        """ add edges, perform update """

        self.count += 1
        self.t1 += 1

        if self.graph.corr is not None:
            self.graph.rm_factors(self.graph.age > self.max_age, store=True)

        self.graph.add_proximity_factors(self.t1-5, max(self.t1-self.frontend_window, 0), 
            rad=self.frontend_radius, nms=self.frontend_nms, thresh=self.frontend_thresh, beta=self.beta, remove=True)

        self.video.disps[self.t1-1] = torch.where(self.video.disps_sens[self.t1-1] > 0, 
           self.video.disps_sens[self.t1-1], self.video.disps[self.t1-1])

        for itr in range(self.iters1):
            self.graph.update(None, None, use_inactive=True)

        # set initial pose for next frame
        poses = SE3(self.video.poses)
        d = self.video.distance([self.t1-3], [self.t1-2], beta=self.beta, bidirectional=True)

        if d.item() < self.keyframe_thresh:
            self.graph.rm_keyframe(self.t1 - 2)
            
            with self.video.get_lock():
                self.video.counter.value -= 1
                self.t1 -= 1

        else:
            for itr in range(self.iters2):
                self.graph.update(None, None, use_inactive=True)

        # set pose for next itration
        self.video.poses[self.t1] = self.video.poses[self.t1-1]
        self.video.disps[self.t1] = self.video.disps[self.t1-1].mean()

        # update visualization
        self.video.dirty[self.graph.ii.min():self.t1] = True

    def __initialize(self):
        """ initialize the SLAM system """

        self.t0 = 0
        self.t1 = self.video.counter.value

        self.graph.add_neighborhood_factors(self.t0, self.t1, r=3)

        for itr in range(8):
            self.graph.update(1, use_inactive=True)

        self.graph.add_proximity_factors(0, 0, rad=2, nms=2, thresh=self.frontend_thresh, remove=False)

        for itr in range(8):
            self.graph.update(1, use_inactive=True)


        # self.video.normalize()
        self.video.poses[self.t1] = self.video.poses[self.t1-1].clone()
        self.video.disps[self.t1] = self.video.disps[self.t1-4:self.t1].mean()

        # initialization complete
        self.is_initialized = True
        self.last_pose = self.video.poses[self.t1-1].clone()
        self.last_disp = self.video.disps[self.t1-1].clone()
        self.last_time = self.video.tstamp[self.t1-1].clone()

        with self.video.get_lock():
            self.video.ready.value = 1
            self.video.dirty[:self.t1] = True

        self.graph.rm_factors(self.graph.ii < self.warmup-4, store=True)

    def __call__(self):
        """ main update """

        # do initialization
        if not self.is_initialized and self.video.counter.value == self.warmup:
            self.__initialize()
            
        # do update
        elif self.is_initialized and self.t1 < self.video.counter.value:
            self.__update()

        

droid_slam/droid_frontend.py
hugoycj/DROID-SLAM-Metric3D
import torch
import lietorch
import numpy as np

import matplotlib.pyplot as plt
from lietorch import SE3
from modules.corr import CorrBlock, AltCorrBlock
import geom.projective_ops as pops


class FactorGraph:
    def __init__(self, video, update_op, device="cuda:0", corr_impl="volume", max_factors=-1, upsample=False):
        self.video = video
        self.update_op = update_op
        self.device = device
        self.max_factors = max_factors
        self.corr_impl = corr_impl
        self.upsample = upsample

        # operator at 1/8 resolution
        self.ht = ht = video.ht // 8
        self.wd = wd = video.wd // 8

        self.coords0 = pops.coords_grid(ht, wd, device=device)
        self.ii = torch.as_tensor([], dtype=torch.long, device=device)
        self.jj = torch.as_tensor([], dtype=torch.long, device=device)
        self.age = torch.as_tensor([], dtype=torch.long, device=device)

        self.corr, self.net, self.inp = None, None, None
        self.damping = 1e-6 * torch.ones_like(self.video.disps)

        self.target = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)
        self.weight = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)

        # inactive factors
        self.ii_inac = torch.as_tensor([], dtype=torch.long, device=device)
        self.jj_inac = torch.as_tensor([], dtype=torch.long, device=device)
        self.ii_bad = torch.as_tensor([], dtype=torch.long, device=device)
        self.jj_bad = torch.as_tensor([], dtype=torch.long, device=device)

        self.target_inac = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)
        self.weight_inac = torch.zeros([1, 0, ht, wd, 2], device=device, dtype=torch.float)

    def __filter_repeated_edges(self, ii, jj):
        """ remove duplicate edges """

        keep = torch.zeros(ii.shape[0], dtype=torch.bool, device=ii.device)
        eset = set(
            [(i.item(), j.item()) for i, j in zip(self.ii, self.jj)] +
            [(i.item(), j.item()) for i, j in zip(self.ii_inac, self.jj_inac)])

        for k, (i, j) in enumerate(zip(ii, jj)):
            keep[k] = (i.item(), j.item()) not in eset

        return ii[keep], jj[keep]

    def print_edges(self):
        ii = self.ii.cpu().numpy()
        jj = self.jj.cpu().numpy()

        ix = np.argsort(ii)
        ii = ii[ix]
        jj = jj[ix]

        w = torch.mean(self.weight, dim=[0,2,3,4]).cpu().numpy()
        w = w[ix]
        for e in zip(ii, jj, w):
            print(e)
        print()

    def filter_edges(self):
        """ remove bad edges """
        conf = torch.mean(self.weight, dim=[0,2,3,4])
        mask = (torch.abs(self.ii-self.jj) > 2) & (conf < 0.001)

        self.ii_bad = torch.cat([self.ii_bad, self.ii[mask]])
        self.jj_bad = torch.cat([self.jj_bad, self.jj[mask]])
        self.rm_factors(mask, store=False)

    def clear_edges(self):
        self.rm_factors(self.ii >= 0)
        self.net = None
        self.inp = None

    @torch.cuda.amp.autocast(enabled=True)
    def add_factors(self, ii, jj, remove=False):
        """ add edges to factor graph """

        if not isinstance(ii, torch.Tensor):
            ii = torch.as_tensor(ii, dtype=torch.long, device=self.device)

        if not isinstance(jj, torch.Tensor):
            jj = torch.as_tensor(jj, dtype=torch.long, device=self.device)

        # remove duplicate edges
        ii, jj = self.__filter_repeated_edges(ii, jj)


        if ii.shape[0] == 0:
            return

        # place limit on number of factors
        if self.max_factors > 0 and self.ii.shape[0] + ii.shape[0] > self.max_factors \
                and self.corr is not None and remove:
            
            ix = torch.arange(len(self.age))[torch.argsort(self.age).cpu()]
            self.rm_factors(ix >= self.max_factors - ii.shape[0], store=True)

        net = self.video.nets[ii].to(self.device).unsqueeze(0)

        # correlation volume for new edges
        if self.corr_impl == "volume":
            c = (ii == jj).long()
            fmap1 = self.video.fmaps[ii,0].to(self.device).unsqueeze(0)
            fmap2 = self.video.fmaps[jj,c].to(self.device).unsqueeze(0)
            corr = CorrBlock(fmap1, fmap2)
            self.corr = corr if self.corr is None else self.corr.cat(corr)

            inp = self.video.inps[ii].to(self.device).unsqueeze(0)
            self.inp = inp if self.inp is None else torch.cat([self.inp, inp], 1)

        with torch.cuda.amp.autocast(enabled=False):
            target, _ = self.video.reproject(ii, jj)
            weight = torch.zeros_like(target)

        self.ii = torch.cat([self.ii, ii], 0)
        self.jj = torch.cat([self.jj, jj], 0)
        self.age = torch.cat([self.age, torch.zeros_like(ii)], 0)

        # reprojection factors
        self.net = net if self.net is None else torch.cat([self.net, net], 1)

        self.target = torch.cat([self.target, target], 1)
        self.weight = torch.cat([self.weight, weight], 1)

    @torch.cuda.amp.autocast(enabled=True)
    def rm_factors(self, mask, store=False):
        """ drop edges from factor graph """

        # store estimated factors
        if store:
            self.ii_inac = torch.cat([self.ii_inac, self.ii[mask]], 0)
            self.jj_inac = torch.cat([self.jj_inac, self.jj[mask]], 0)
            self.target_inac = torch.cat([self.target_inac, self.target[:,mask]], 1)
            self.weight_inac = torch.cat([self.weight_inac, self.weight[:,mask]], 1)

        self.ii = self.ii[~mask]
        self.jj = self.jj[~mask]
        self.age = self.age[~mask]
        
        if self.corr_impl == "volume":
            self.corr = self.corr[~mask]

        if self.net is not None:
            self.net = self.net[:,~mask]

        if self.inp is not None:
            self.inp = self.inp[:,~mask]

        self.target = self.target[:,~mask]
        self.weight = self.weight[:,~mask]


    @torch.cuda.amp.autocast(enabled=True)
    def rm_keyframe(self, ix):
        """ drop edges from factor graph """


        with self.video.get_lock():
            self.video.images[ix] = self.video.images[ix+1]
            self.video.poses[ix] = self.video.poses[ix+1]
            self.video.disps[ix] = self.video.disps[ix+1]
            self.video.disps_sens[ix] = self.video.disps_sens[ix+1]
            self.video.intrinsics[ix] = self.video.intrinsics[ix+1]

            self.video.nets[ix] = self.video.nets[ix+1]
            self.video.inps[ix] = self.video.inps[ix+1]
            self.video.fmaps[ix] = self.video.fmaps[ix+1]

        m = (self.ii_inac == ix) | (self.jj_inac == ix)
        self.ii_inac[self.ii_inac >= ix] -= 1
        self.jj_inac[self.jj_inac >= ix] -= 1

        if torch.any(m):
            self.ii_inac = self.ii_inac[~m]
            self.jj_inac = self.jj_inac[~m]
            self.target_inac = self.target_inac[:,~m]
            self.weight_inac = self.weight_inac[:,~m]

        m = (self.ii == ix) | (self.jj == ix)

        self.ii[self.ii >= ix] -= 1
        self.jj[self.jj >= ix] -= 1
        self.rm_factors(m, store=False)


    @torch.cuda.amp.autocast(enabled=True)
    def update(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, motion_only=False):
        """ run update operator on factor graph """

        # motion features
        with torch.cuda.amp.autocast(enabled=False):
            coords1, mask = self.video.reproject(self.ii, self.jj)
            motn = torch.cat([coords1 - self.coords0, self.target - coords1], dim=-1)
            motn = motn.permute(0,1,4,2,3).clamp(-64.0, 64.0)
        
        # correlation features
        corr = self.corr(coords1)

        self.net, delta, weight, damping, upmask = \
            self.update_op(self.net, self.inp, corr, motn, self.ii, self.jj)

        if t0 is None:
            t0 = max(1, self.ii.min().item()+1)

        with torch.cuda.amp.autocast(enabled=False):
            self.target = coords1 + delta.to(dtype=torch.float)
            self.weight = weight.to(dtype=torch.float)

            ht, wd = self.coords0.shape[0:2]
            self.damping[torch.unique(self.ii)] = damping

            if use_inactive:
                m = (self.ii_inac >= t0 - 3) & (self.jj_inac >= t0 - 3)
                ii = torch.cat([self.ii_inac[m], self.ii], 0)
                jj = torch.cat([self.jj_inac[m], self.jj], 0)
                target = torch.cat([self.target_inac[:,m], self.target], 1)
                weight = torch.cat([self.weight_inac[:,m], self.weight], 1)

            else:
                ii, jj, target, weight = self.ii, self.jj, self.target, self.weight


            damping = .2 * self.damping[torch.unique(ii)].contiguous() + EP

            target = target.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()
            weight = weight.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()

            # dense bundle adjustment
            self.video.ba(target, weight, damping, ii, jj, t0, t1, 
                itrs=itrs, lm=1e-4, ep=0.1, motion_only=motion_only)
        
            if self.upsample:
                self.video.upsample(torch.unique(self.ii), upmask)

        self.age += 1


    @torch.cuda.amp.autocast(enabled=False)
    def update_lowmem(self, t0=None, t1=None, itrs=2, use_inactive=False, EP=1e-7, steps=8):
        """ run update operator on factor graph - reduced memory implementation """

        # alternate corr implementation
        t = self.video.counter.value

        num, rig, ch, ht, wd = self.video.fmaps.shape
        corr_op = AltCorrBlock(self.video.fmaps.view(1, num*rig, ch, ht, wd))

        for step in range(steps):
            print("Global BA Iteration #{}".format(step+1))
            with torch.cuda.amp.autocast(enabled=False):
                coords1, mask = self.video.reproject(self.ii, self.jj)
                motn = torch.cat([coords1 - self.coords0, self.target - coords1], dim=-1)
                motn = motn.permute(0,1,4,2,3).clamp(-64.0, 64.0)

            s = 8
            for i in range(0, self.jj.max()+1, s):
                v = (self.ii >= i) & (self.ii < i + s)
                iis = self.ii[v]
                jjs = self.jj[v]

                ht, wd = self.coords0.shape[0:2]
                corr1 = corr_op(coords1[:,v], rig * iis, rig * jjs + (iis == jjs).long())

                with torch.cuda.amp.autocast(enabled=True):
                 
                    net, delta, weight, damping, upmask = \
                        self.update_op(self.net[:,v], self.video.inps[None,iis], corr1, motn[:,v], iis, jjs)

                    if self.upsample:
                        self.video.upsample(torch.unique(iis), upmask)

                self.net[:,v] = net
                self.target[:,v] = coords1[:,v] + delta.float()
                self.weight[:,v] = weight.float()
                self.damping[torch.unique(iis)] = damping

            damping = .2 * self.damping[torch.unique(self.ii)].contiguous() + EP
            target = self.target.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()
            weight = self.weight.view(-1, ht, wd, 2).permute(0,3,1,2).contiguous()

            # dense bundle adjustment
            self.video.ba(target, weight, damping, self.ii, self.jj, 1, t, 
                itrs=itrs, lm=1e-5, ep=1e-2, motion_only=False)

            self.video.dirty[:t] = True

    def add_neighborhood_factors(self, t0, t1, r=3):
        """ add edges between neighboring frames within radius r """

        ii, jj = torch.meshgrid(torch.arange(t0,t1), torch.arange(t0,t1))
        ii = ii.reshape(-1).to(dtype=torch.long, device=self.device)
        jj = jj.reshape(-1).to(dtype=torch.long, device=self.device)

        c = 1 if self.video.stereo else 0

        keep = ((ii - jj).abs() > c) & ((ii - jj).abs() <= r)
        self.add_factors(ii[keep], jj[keep])

    
    def add_proximity_factors(self, t0=0, t1=0, rad=2, nms=2, beta=0.25, thresh=16.0, remove=False):
        """ add edges to the factor graph based on distance """

        t = self.video.counter.value
        ix = torch.arange(t0, t)
        jx = torch.arange(t1, t)

        ii, jj = torch.meshgrid(ix, jx)
        ii = ii.reshape(-1)
        jj = jj.reshape(-1)

        d = self.video.distance(ii, jj, beta=beta)
        d[ii - rad < jj] = np.inf
        d[d > 100] = np.inf

        ii1 = torch.cat([self.ii, self.ii_bad, self.ii_inac], 0)
        jj1 = torch.cat([self.jj, self.jj_bad, self.jj_inac], 0)
        for i, j in zip(ii1.cpu().numpy(), jj1.cpu().numpy()):
            for di in range(-nms, nms+1):
                for dj in range(-nms, nms+1):
                    if abs(di) + abs(dj) <= max(min(abs(i-j)-2, nms), 0):
                        i1 = i + di
                        j1 = j + dj

                        if (t0 <= i1 < t) and (t1 <= j1 < t):
                            d[(i1-t0)*(t-t1) + (j1-t1)] = np.inf


        es = []
        for i in range(t0, t):
            if self.video.stereo:
                es.append((i, i))
                d[(i-t0)*(t-t1) + (i-t1)] = np.inf

            for j in range(max(i-rad-1,0), i):
                es.append((i,j))
                es.append((j,i))
                d[(i-t0)*(t-t1) + (j-t1)] = np.inf

        ix = torch.argsort(d)
        for k in ix:
            if d[k].item() > thresh:
                continue

            if len(es) > self.max_factors:
                break

            i = ii[k]
            j = jj[k]
            
            # bidirectional
            es.append((i, j))
            es.append((j, i))

            for di in range(-nms, nms+1):
                for dj in range(-nms, nms+1):
                    if abs(di) + abs(dj) <= max(min(abs(i-j)-2, nms), 0):
                        i1 = i + di
                        j1 = j + dj

                        if (t0 <= i1 < t) and (t1 <= j1 < t):
                            d[(i1-t0)*(t-t1) + (j1-t1)] = np.inf

        ii, jj = torch.as_tensor(es, device=self.device).unbind(dim=-1)
        self.add_factors(ii, jj, remove)

droid_slam/factor_graph.py
hugoycj/DROID-SLAM-Metric3D