datasets_preprocess/sintel_get_dynamics.py

import numpy as np
import cv2
import os
from tqdm import tqdm
import argparse
import torch

TAG_FLOAT = 202021.25
def flow_read(filename):
    """ Read optical flow from file, return (U,V) tuple.
    
    Original code by Deqing Sun, adapted from Daniel Scharstein.
    """
    f = open(filename,'rb')
    check = np.fromfile(f,dtype=np.float32,count=1)[0]
    assert check == TAG_FLOAT, 'flow_read:: Wrong tag in flow file (should be: {0}, is: {1}). Big-endian machine?'.format(TAG_FLOAT,check)
    width = np.fromfile(f,dtype=np.int32,count=1)[0]
    height = np.fromfile(f,dtype=np.int32,count=1)[0]
    size = width*height
    assert width > 0 and height > 0 and size > 1 and size < 100000000, 'flow_read:: Invalid input size (width = {0}, height = {1}).'.format(width,height)
    tmp = np.fromfile(f,dtype=np.float32,count=-1).reshape((height,width*2))
    u = tmp[:,np.arange(width)*2]
    v = tmp[:,np.arange(width)*2 + 1]
    return u,v

def cam_read(filename):
    """ Read camera data, return (M,N) tuple.
    
    M is the intrinsic matrix, N is the extrinsic matrix, so that

    x = M*N*X,
    where x is a point in homogeneous image pixel coordinates, and X is a
    point in homogeneous world coordinates.
    """
    f = open(filename,'rb')
    check = np.fromfile(f,dtype=np.float32,count=1)[0]
    assert check == TAG_FLOAT, 'cam_read:: Wrong tag in cam file (should be: {0}, is: {1}). Big-endian machine?'.format(TAG_FLOAT,check)
    M = np.fromfile(f,dtype='float64',count=9).reshape((3,3))
    N = np.fromfile(f,dtype='float64',count=12).reshape((3,4))
    return M,N

def depth_read(filename):
    """ Read depth data from file, return as numpy array. """
    f = open(filename,'rb')
    check = np.fromfile(f,dtype=np.float32,count=1)[0]
    assert check == TAG_FLOAT, 'depth_read:: Wrong tag in depth file (should be: {0}, is: {1}). Big-endian machine?'.format(TAG_FLOAT,check)
    width = np.fromfile(f,dtype=np.int32,count=1)[0]
    height = np.fromfile(f,dtype=np.int32,count=1)[0]
    size = width*height
    assert width > 0 and height > 0 and size > 1 and size < 100000000, 'depth_read:: Invalid input size (width = {0}, height = {1}).'.format(width,height)
    depth = np.fromfile(f,dtype=np.float32,count=-1).reshape((height,width))
    return depth

def RT_to_extrinsic_matrix(R, T):
    extrinsic_matrix = np.concatenate([R, T], axis=-1)
    extrinsic_matrix = np.concatenate([extrinsic_matrix, np.array([[0, 0, 0, 1]])], axis=0)
    return np.linalg.inv(extrinsic_matrix)

def depth_to_3d(depth_map, intrinsic_matrix):
    height, width = depth_map.shape
    i, j = np.meshgrid(np.arange(width), np.arange(height))
    
    # Convert pixel coordinates and depth values to 3D points
    x = (i - intrinsic_matrix[0, 2]) * depth_map / intrinsic_matrix[0, 0]
    y = (j - intrinsic_matrix[1, 2]) * depth_map / intrinsic_matrix[1, 1]
    z = depth_map
    
    points_3d = np.stack([x, y, z], axis=-1)
    return points_3d

def project_3d_to_2d(points_3d, intrinsic_matrix):
    # Convert 3D points to homogeneous coordinates
    projected_2d_hom = intrinsic_matrix @ points_3d.T
    # Convert from homogeneous coordinates to 2D image coordinates
    projected_2d = projected_2d_hom[:2, :] / projected_2d_hom[2, :]
    return projected_2d.T

def compute_optical_flow(depth1, depth2, pose1, pose2, intrinsic_matrix1, intrinsic_matrix2):
    # Input: All inputs as numpy arrays; convert torch tensors to numpy arrays if needed
    if isinstance(depth1, torch.Tensor):
        depth1 = depth1.cpu().numpy()
    if isinstance(depth2, torch.Tensor):
        depth2 = depth2.cpu().numpy()
    if isinstance(pose1, torch.Tensor):
        pose1 = pose1.cpu().numpy()
    if isinstance(pose2, torch.Tensor):
        pose2 = pose2.cpu().numpy()
    if isinstance(intrinsic_matrix1, torch.Tensor):
        intrinsic_matrix1 = intrinsic_matrix1.cpu().numpy()
    if isinstance(intrinsic_matrix2, torch.Tensor):
        intrinsic_matrix2 = intrinsic_matrix2.cpu().numpy()

    points_3d_frame1 = depth_to_3d(depth1, intrinsic_matrix1).reshape(-1, 3)
    points_3d_frame1_hom = np.concatenate([points_3d_frame1, np.ones((points_3d_frame1.shape[0], 1))], axis=1).T
    
    # Calculate the transformation matrix from frame 1 to frame 2
    transformation_matrix = (pose2) @ np.linalg.inv(pose1)
    points_3d_frame2_hom = transformation_matrix @ points_3d_frame1_hom
    points_3d_frame2 = (points_3d_frame2_hom[:3, :]).T

    points_2d_frame1 = project_3d_to_2d(points_3d_frame1, intrinsic_matrix1)
    points_2d_frame2 = project_3d_to_2d(points_3d_frame2, intrinsic_matrix2)

    # Compute optical flow vectors
    optical_flow = points_2d_frame2 - points_2d_frame1
    return optical_flow

def get_dynamic_label(base_dir, seq, continuous=False, threshold=13.75, save_dir='dynamic_label'):
    depth_dir = os.path.join(base_dir, 'depth', seq)
    cam_dir = os.path.join(base_dir, 'camdata_left', seq)
    flow_dir = os.path.join(base_dir, 'flow', seq)
    dynamic_label_dir = os.path.join(base_dir, save_dir, seq)
    os.makedirs(dynamic_label_dir, exist_ok=True)
    
    frames = sorted([f for f in os.listdir(depth_dir) if f.endswith('.dpt')])
    for i, frame1 in enumerate(frames):
        if i == len(frames) - 1:
            continue
        frame2 = frames[i + 1]
        
        frame1_id = frame1.split('.')[0]
        frame2_id = frame2.split('.')[0]

        # Load depth maps
        depth_map_frame1 = depth_read(os.path.join(depth_dir, frame1))
        depth_map_frame2 = depth_read(os.path.join(depth_dir, frame2))
        
        # Load camera intrinsics and poses
        intrinsic_matrix1, pose_frame1 = cam_read(os.path.join(cam_dir, f'{frame1_id}.cam'))
        intrinsic_matrix2, pose_frame2 = cam_read(os.path.join(cam_dir, f'{frame2_id}.cam'))
        
        # Pad pose with [0,0,0,1]
        pose_frame1 = np.concatenate([pose_frame1, np.array([[0, 0, 0, 1]])], axis=0)
        pose_frame2 = np.concatenate([pose_frame2, np.array([[0, 0, 0, 1]])], axis=0)
        
        # Compute optical flow
        optical_flow = compute_optical_flow(depth_map_frame1, depth_map_frame2, pose_frame1, pose_frame2, intrinsic_matrix1, intrinsic_matrix2)
        
        # Reshape the optical flow to the image dimensions
        height, width = depth_map_frame1.shape
        optical_flow_image = optical_flow.reshape(height, width, 2)
        
        # Load ground truth optical flow
        u, v = flow_read(os.path.join(flow_dir, f'{frame1_id}.flo'))
        gt_flow = np.stack([u, v], axis=-1)
        
        # Compute the error map
        error_map = np.linalg.norm(gt_flow - optical_flow_image, axis=-1)
        if not continuous:
            binary_error_map = error_map > threshold
            
            # Save the binary error map
            cv2.imwrite(os.path.join(dynamic_label_dir, f'{frame1_id}.png'), binary_error_map.astype(np.uint8) * 255)
        else:
            # Normalize the error map
            error_map = error_map / error_map.max()
            cv2.imwrite(os.path.join(dynamic_label_dir, f'{frame1_id}.png'), (error_map * 255).astype(np.uint8))

if __name__ == '__main__':
    # Process all sequences
    sequences = sorted(os.listdir('data/sintel/training/depth'))
    base_dir = 'data/sintel/training'
    parser = argparse.ArgumentParser()
    parser.add_argument('--continuous', action='store_true')
    parser.add_argument('--threshold', type=float, default=13.75)
    parser.add_argument('--save_dir', type=str, default='dynamic_label')
    args = parser.parse_args()
    for seq in tqdm(sequences):
        get_dynamic_label(base_dir, seq, continuous=args.continuous, threshold=args.threshold, save_dir=args.save_dir)
        print(f'Finished processing {seq}')