"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os, sys, cv2
import numpy as np
import torch
from torch.autograd import Variable



def init_depth_model():
    NeWCRFs_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/NeWCRFs')
    sys.path.append(os.path.join(NeWCRFs_ROOT, 'newcrfs'))
    from networks.NewCRFDepth import NewCRFDepth

    # model config
    class cfg: pass
    cfg.version = 'large07'
    cfg.checkpoint_path = os.path.join(NeWCRFs_ROOT, 'model_zoo', 'model_nyu.ckpt')
    cfg.inv_depth = False
    cfg.max_depth = 10  # meters
    cfg.input_size = (480, 640)  # (height, width)
    cfg.input_mean = (123.68, 116.78, 103.94)
    cfg.input_std = (0.017, 0.017, 0.017)
    
    model = NewCRFDepth(version=cfg.version, inv_depth=False, max_depth=cfg.max_depth)
    model = torch.nn.DataParallel(model)
    
    checkpoint = torch.load(cfg.checkpoint_path)
    model.load_state_dict(checkpoint['model'])
    model.eval()
    model.cuda()
    
    return model, cfg


def preprocess_rgb(frame, cfg):
    frame = cv2.resize(frame, (cfg.input_size[1], cfg.input_size[0]), interpolation=cv2.INTER_LINEAR)
    input_image = frame.astype(np.float32)
    # normalize
    for i in range(3):
        input_image[:, :, i] = (input_image[:, :, i] - cfg.input_mean[i]) * cfg.input_std[i]
    # transpose
    input_images = np.expand_dims(input_image, axis=0)  # (1, H, W, C)
    input_images = np.transpose(input_images, (0, 3, 1, 2))  # (1, C, H, W)
    image = Variable(torch.from_numpy(input_images)).cuda()
    return image


def postprocess_depth(depth, frame_size):
    """ depth: (h, w)
        frame_size: (H, W)
    """
    # resize depth map to the rgb frame size
    depth_resize = cv2.resize(depth, (frame_size[1], frame_size[0]))
    return depth_resize


def monocular_depth_estimation(video, model, cfg):
    """ depth estimation for each frame """
    depth_maps = []
    for frame in video:
        # preprocess video frame as the model input
        frame_input = preprocess_rgb(frame, cfg)
        with torch.no_grad(): # inference
            depth_est = model(frame_input)
            depth_est = depth_est.cpu().numpy().squeeze()  # (h, w)
        # post process
        depth_est = postprocess_depth(depth_est, frame.shape[:2])  # (H, W)
        depth_maps.append(depth_est)
    depth_maps = np.stack(depth_maps, axis=0)  # (T, H, W)
    return depth_maps

preprocess/EgoPAT3D/depth_estimation.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
import pyk4a
from pyk4a import PyK4APlayback
from tqdm import tqdm
import pickle
import argparse

from google.protobuf import text_format
from mediapipe.framework.formats import landmark_pb2
from mediapipe.python.solutions import drawing_utils
import mediapipe.python.solutions.drawing_styles as mp_style
import mediapipe.python.solutions.hands as mp_hands
from typing import Mapping, Optional, Tuple

_VISIBILITY_THRESHOLD = 0.5
_PRESENCE_THRESHOLD = 0.5
WHITE_COLOR = (224, 224, 224)



def read_landmarks(landmark_file, flip=False):
    landmarks = dict()
    with open(landmark_file, 'r') as f:
        pid = 0
        for line in f.readlines():
            line = line.strip()
            if line.split(':')[0].isnumeric():
                frame_id = int(line.split(':')[0])
                landmarks[frame_id] = np.zeros((21, 3))  # a new frame
                assert pid == 0 or pid == 20, 'invalid frame!'
                pid = -1
                continue
            if 'landmark' in line: 
                pid += 1  # a new 3D point
                continue
            if 'x:' in line: 
                xval = float(line.split(': ')[1])
                landmarks[frame_id][pid, 0] = 1.0 - xval if flip else xval
                continue
            if 'y:' in line:
                landmarks[frame_id][pid, 1] = float(line.split(': ')[1])
                continue
            if 'z:' in line:
                landmarks[frame_id][pid, 2] = float(line.split(': ')[1])
                continue
    return landmarks


def read_landmarks_text(landmark_file, flip=False):
    landmarks = dict()
    with open(landmark_file, 'r') as f:
        for line in f.readlines():
            line = line.strip()
            if line.split(':')[0].isnumeric():
                frame_id = int(line.split(':')[0])
                landmarks[frame_id] = ''
                pointer = 0
                continue
            if flip and 'x:' in line:
                xval = 1.0 - float(line.split(': ')[1])  # flip x coordinates horizontally
                line = '  x: {}'.format(xval)
            if pointer <= 105:  # 1 + 5 * 21 = 105 rows
                if ('x:' in line) or ('y:' in line) or ('}' in line):
                    line += ' '
                landmarks[frame_id] += line
                pointer += 1
    return landmarks


def read_clip_ranges(rng_file):
    ranges = dict()
    with open(rng_file, 'r') as f:
        for line in f.readlines():
            cid = int(line.strip().split('. ')[0])
            start = int(line.strip().split('. ')[1].split(', ')[0])
            end = int(line.strip().split('. ')[1].split(', ')[1])
            ranges[cid] = (start, end)
    return ranges


def get_2d_coordinates(landmarks_text, start, end, resolution):
    image_cols, image_rows = resolution[1], resolution[0]
    results = {}
    for i in range(start, end):
        if i not in landmarks_text:
            continue
        # parse landmarks
        landmark_list = text_format.Parse(landmarks_text[i], landmark_pb2.NormalizedLandmarkList())
        # get the 2D coordinates
        idx_to_coordinates = {}
        for idx, landmark in enumerate(landmark_list.landmark):
            if ((landmark.HasField('visibility') and
                landmark.visibility < _VISIBILITY_THRESHOLD) or
                (landmark.HasField('presence') and
                landmark.presence < _PRESENCE_THRESHOLD)):
                continue
            landmark_px = drawing_utils._normalized_to_pixel_coordinates(landmark.x, landmark.y, image_cols, image_rows)
            if landmark_px:
                idx_to_coordinates[idx] = landmark_px  # (x, y)
        # save to dict
        results[i] = idx_to_coordinates
    return results


def get_3d_coordinates(landmarks, landmarks_2d, frame_start, depth_clip=None):
    results = {}
    for frame_id, idx_to_coords in landmarks_2d.items():
        pxyz = np.copy(landmarks[frame_id])  # (21, 3)
        pxyz[1:, 2] += pxyz[0, 2]  # z values of top 20 are relative offsets to the 21-th z value
        idx_to_coords3d = {}
        # iterate each 2D point (u,v) to fetch the depth (z)
        for idx, puv in idx_to_coords.items():
            x_norm, y_norm, Z = pxyz[idx]
            # retrieve z from (u,v) location of depth map
            if depth_clip is not None:
                u, v = int(puv[1]), int(puv[0])
                Z = depth_clip[frame_id - frame_start, u, v]
            # from 2D (u,v) to 3D (X, Y, Z)
            Z = Z / 1000  # mm to m
            X = (x_norm * resolution[1] - 1.94228662e+03) * Z / 1.80820276e+03
            Y = (y_norm * resolution[0] - 1.12382178e+03) * Z / 1.80794556e+03
            idx_to_coords3d[idx] = [X, Y, Z]  # (m, m, m)
        results[frame_id] = idx_to_coords3d
    return results


def colorize(
    image: np.ndarray,
    clipping_range: Tuple[Optional[int], Optional[int]] = (None, None),
    colormap: int = cv2.COLORMAP_HSV,
) -> np.ndarray:
    if clipping_range[0] or clipping_range[1]:
        img = image.clip(clipping_range[0], clipping_range[1])
    else:
        img = image.copy()
    img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)
    img = cv2.applyColorMap(img, colormap)
    return img


def get_depth_vis(depth_mat, depth_range=[0, 1200], ratio=1.0):
    if np.max(depth_mat) > depth_range[1]:
        depth_range[1] = np.max(depth_mat)
    depth_vis = colorize(depth_mat, tuple(depth_range))
    # resize into smaller size
    dst_size = (int(depth_vis.shape[1] * ratio), int(depth_vis.shape[0] * ratio))
    depth_vis = cv2.resize(depth_vis, dst_size)
    return depth_vis


def get_depth(playback, start, end, pointer, ratio=1.0, vis_depth=False):
    
    depth_clip = np.zeros([end - start] + resolution)  # (N, H, W, 3)
    depth_vis_all = []
    # move pointer to the clip start frame
    while pointer < start:
        capture = playback.get_next_capture()
        pointer += 1
    # get the depth of the clip
    while pointer < end:
        capture = playback.get_next_capture()
        if (capture.color is None) or (capture.depth is None):
            pointer += 1  # move pointer
            continue
        # rgb = convert_to_bgra_if_required(playback.configuration["color_format"], capture.color)
        depth_mat = pyk4a.depth_image_to_color_camera(capture.depth, playback.calibration, playback.thread_safe)
        depth_clip[pointer - start] = depth_mat
        if vis_depth:
            # save visualized depth
            depth_vis = get_depth_vis(depth_mat, depth_range=depth_range, ratio=ratio)
            depth_vis_all.append(depth_vis)
        pointer += 1  # move pointer
    depth_vis_all = np.stack(depth_vis_all, axis=0) if vis_depth else None
    return depth_clip, depth_vis_all, pointer


def compute_3d_trajectory(result_file, record_name):

    # read hand landmarks
    landmark_file = os.path.join(record_path, 'hand_frames', 'hand_landmarks.txt')
    if not os.path.exists(landmark_file):
        return False
    landmarks = read_landmarks(landmark_file, flip=True)
    # get the landmark text with flipped x value
    landmarks_text = read_landmarks_text(landmark_file, flip=True)

    # mkv video
    playback = PyK4APlayback(mkvraw_file)
    playback.open()
    pointer = 0

    # read clip ranges
    cliprng_file = os.path.join(record_path, 'hand_frames', 'clip_ranges.txt')
    clip_ranges = read_clip_ranges(cliprng_file)

    # process each clip
    results = []
    for cid, (start, end) in tqdm(clip_ranges.items(), desc=f'Record {record_name}', total=len(clip_ranges)):
        assert start < end, "invalid clip!"
        if end - start > 5400: # more than ~3min (invalid clip):
            continue
        # get 2D landmarks
        landmarks_2d = get_2d_coordinates(landmarks_text, start, end, resolution)

        # get depth frames
        depth_clip, depth_vis, pointer = get_depth(playback, start, end, pointer, ratio=vis_ratio, vis_depth=(cid==list(clip_ranges.keys())[0]))
        if depth_vis is not None:
            write_video(depth_vis, depth_visfile, fps=5)
        
        # get 3D landmarks
        trajectory_3d = get_3d_coordinates(landmarks, landmarks_2d, start, depth_clip=depth_clip)
        results.append(trajectory_3d)

    if len(results) > 0:
        with open(result_file.format('local'), 'wb') as f:
            pickle.dump(results, f, protocol=pickle.HIGHEST_PROTOCOL)

    playback.close()
    return True


def read_video(video_file, ratio=0.5):
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    videos = []
    while success:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        dst_size = (int(frame.shape[1] * ratio), int(frame.shape[0] * ratio))
        frame = cv2.resize(frame, dst_size)
        videos.append(frame)
        success, frame = cap.read()
    videos = np.array(videos)
    return videos


def write_video(mat, video_file, fps=5):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)


def read_clip_ranges(rng_file):
    ranges = dict()
    with open(rng_file, 'r') as f:
        for line in f.readlines():
            cid = int(line.strip().split('. ')[0])
            start = int(line.strip().split('. ')[1].split(', ')[0])
            end = int(line.strip().split('. ')[1].split(', ')[1])
            ranges[cid] = (start, end)
    return ranges


def draw_on_frames(frame, idx_to_coords3d, ratio=0.25):
    # drawing config
    connection_drawing_spec = mp_style.get_default_hand_connections_style()
    landmark_drawing_spec = mp_style.get_default_hand_landmarks_style()
    
    def _XYZ_to_uv(X, Y, Z, ratio=1.0):
        u = X * 1.80820276e+03 / Z + 1.94228662e+03
        v = Y * 1.80794556e+03 / Z + 1.12382178e+03
        u, v = int(u * ratio), int(v * ratio)  # scaling
        return u, v

    # Draws the connections if the start and end landmarks are both visible.
    for connection in mp_hands.HAND_CONNECTIONS:
        start_idx = connection[0]
        end_idx = connection[1]
        if start_idx in idx_to_coords3d and end_idx in idx_to_coords3d:
            drawing_spec = connection_drawing_spec[connection] if isinstance(
                connection_drawing_spec, Mapping) else connection_drawing_spec
            # start point (XYZ to uv)
            X, Y, Z = idx_to_coords3d[start_idx]  # (X, Y, Z) in meters
            if Z == 0: continue
            px_start, py_start = _XYZ_to_uv(X, Y, Z, ratio=ratio)
            # end point (XYZ to uv)
            X, Y, Z = idx_to_coords3d[end_idx]  # (X, Y, Z) in meters
            if Z == 0: continue
            px_end, py_end = _XYZ_to_uv(X, Y, Z, ratio=ratio)
            # draw line
            cv2.line(frame, (px_start, py_start), (px_end, py_end), drawing_spec.color, drawing_spec.thickness)
    
    # draw keypoints
    for idx, landmark_pxyz in idx_to_coords3d.items():
        drawing_spec = landmark_drawing_spec[idx] if isinstance(
        landmark_drawing_spec, Mapping) else landmark_drawing_spec
        # White circle border
        circle_border_radius = max(drawing_spec.circle_radius + 1,
                                    int(drawing_spec.circle_radius * 1.2))
        # (XYZ to uv)
        X, Y, Z = landmark_pxyz  # (X, Y, Z) in meters
        if Z == 0: continue
        px, py = _XYZ_to_uv(X, Y, Z, ratio=ratio)
        cv2.circle(frame, (px, py), circle_border_radius, WHITE_COLOR, drawing_spec.thickness)
        # Fill color into the circle
        cv2.circle(frame, (px, py), drawing_spec.circle_radius, drawing_spec.color, drawing_spec.thickness)


def draw_trajectory_3d(landmarks_3d, start, end, fps=5, size=1000, fig_size=5, views=[30, -120], minmax=[(0, 0, -1), (1, 1, 2)], savefile='temp.mp4'):

    xmin, ymin, zmin = minmax[0]
    xmax, ymax, zmax = minmax[1]

    # init a 3D figure
    fig = plt.figure(figsize=(fig_size, fig_size))
    ax = fig.add_subplot(projection='3d')
    ax.view_init(views[0], views[1])  # elev: 30, azim: -120
    # Make a 3D quiver plot
    ox, oy, oz = np.zeros((3,3))
    dx, dy, dz = np.array([[xmax-xmin, 0, 0], [0, zmax-zmin, 0], [0, 0, ymax-ymin]]) * 0.2
    ax.quiver(ox, oy, oz, dx, dy, dz)
    # Setup coordinate system
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(zmin, zmax)
    ax.set_zlim(ymin, ymax)
    ax.invert_zaxis()
    ax.set_xlabel('X')  # x: rightward
    ax.set_ylabel('Z')  # z: innerward
    ax.set_zlabel('Y')  # y: downward
    plt.locator_params(axis='x', nbins=5)
    plt.locator_params(axis='y', nbins=5)
    plt.locator_params(axis='z', nbins=5)

    # init video writer
    traj3d_writer = FFMpegWriter(fps=fps, metadata=dict(title='3D Trajectory', artist='Matplotlib'))
    # process
    with traj3d_writer.saving(fig, savefile, dpi=int(size / fig_size)):
        traj3d = []
        pts3d_last = np.array(list(list(landmarks_3d.values())[0].values()))
        for i in range(start, end):
            if i in landmarks_3d:
                pts3d = np.array(list(landmarks_3d[i].values()))
                pts3d_last = np.copy(pts3d)
            else:
                pts3d = np.copy(pts3d_last)
            # compute the 3d center of the hand as the trajectory point
            pt_hand = np.mean(pts3d[pts3d[:, 2] > 0], axis=0)
            traj3d.append(pt_hand)
            xs, ys, zs = np.array(traj3d)[:, 0], np.array(traj3d)[:, 1], np.array(traj3d)[:, 2]
            ax.plot(xs, zs, ys, '-o')
            ax.set_title('Hand trajectory [Frame: {}, View: ({}, {})]'.format(i-start, views[0], views[1]))
            # grab a frame
            traj3d_writer.grab_frame()  # dpi * (fig_size, fig_size)
    plt.close()


def vis(results, ratio=0.25):
    # video file path
    print('Read video...')
    video_file = os.path.join(record_path, 'rgb_video.mp4')
    video_all = read_video(video_file, ratio=ratio)  # (N, 540, 960, 3)

    # read clip ranges
    cliprng_file = os.path.join(record_path, 'hand_frames', 'clip_ranges.txt')
    clip_ranges = read_clip_ranges(cliprng_file)

    # visualize for each clip
    for cid, (start, end) in tqdm(clip_ranges.items(), desc='Visualizing clips', total=len(clip_ranges)):
        landmarks_3d = results[cid - 1]  # clip id starts from 1
        vis_mat = np.copy(video_all[start: end])

        for frame_id, idx_to_coords3d in landmarks_3d.items():
            # draw on video frame
            draw_on_frames(vis_mat[frame_id - start], idx_to_coords3d, ratio=ratio)

        # draw 3D trajectory video
        draw_trajectory_3d(landmarks_3d, start, end, fps=5, size=int(resolution[0]*ratio), views=[10, -75], minmax=[(-0.5, -0.5, 0), (0.5, 0.5, 1.0)], savefile='temp.mp4')
        video_traj3d = read_video('temp.mp4', ratio=1.0)  # (M, 540, 540, 3)

        # visualize RGB and 3D trajectory side-by-side
        vis_mat = np.concatenate([vis_mat, video_traj3d], axis=2)  # (M, 540, 1500, 3)

        vis_file = os.path.join(vis_dir, 'vis_{}_clip{}.mp4'.format(record_name, cid))
        write_video(vis_mat, vis_file, fps=5)
    
    print(len(results))

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--vis', action='store_true', help='vis the results')
    parser.add_argument('--scene', default=1, help='the scene id')
    args = parser.parse_args()
    scene_id = str(args.scene)
    
    root_path = os.path.join(os.path.dirname(__file__), '../../data')
    resolution = [2160, 3840]  # height, width
    vis_ratio = 0.25
    depth_range = [0, 1200]  # for depth visualization
    visited = False
    
    for record_name in sorted(os.listdir(os.path.join(root_path, 'EgoPAT3D-complete', scene_id))):
        # scene_name = 'bathroomCabinet'
        # record_id = '2'
        # record_name = '{}_{}'.format(scene_name, record_id)
        record_path = os.path.join(root_path, 'EgoPAT3D-complete', scene_id, record_name)  # './1/bathroomCabinet_2'
        mkvraw_file = os.path.join(root_path, 'EgoPAT3D-mkv', scene_id, record_name + '.mkv')  # './1/bathroomCabinet_2.mkv'
        if not os.path.isdir(record_path):
            continue
        
        # result trajectory path
        traj_path = os.path.join(root_path, 'trajectory', scene_id)
        os.makedirs(traj_path, exist_ok=True)
        traj_file = os.path.join(traj_path, record_name + '_{}.pkl')  # local, global
        
        # result depth visualization path
        depth_vispath = os.path.join(root_path, 'depth_vis{}'.format(vis_ratio), scene_id)
        os.makedirs(depth_vispath, exist_ok=True)
        depth_visfile = os.path.join(depth_vispath, record_name + '_clip0.mp4')

        if not os.path.exists(traj_file.format("local")):
            success = compute_3d_trajectory(traj_file, record_name)
        
        if args.vis and not visited:
            vis_dir = os.path.join(traj_path, record_name + '_vis')
            os.makedirs(vis_dir, exist_ok=True)
            
            # saved landmarks
            with open(traj_file.format("local"), 'rb') as f:
                results = pickle.load(f)
            vis(results, ratio=vis_ratio)
            visited = True  # only visualize one recording for each scene

    if os.path.exists('temp.mp4'):
        os.remove('temp.mp4')
preprocess/EgoPAT3D/get_3dhand_trajectory.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import open3d as o3d
import numpy as np
import matplotlib.pyplot as plt


if __name__ == '__main__':
    pcd_file = '../../data/sequences/bathroomCabinet/bathroomCabinet_2/pointcloud/300.ply'

    pcd = o3d.io.read_point_cloud(pcd_file)
    pointxyz = np.asarray(pcd.points)
    pointcolor = np.asarray(pcd.colors)
    # pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]])

    # camera intrinsics and exterinsics
    cam_param = o3d.camera.PinholeCameraParameters()
    cam_param.extrinsic = np.array([[1,0,0,0], 
                                    [0,1,0,0],
                                    [0,0,1,0],
                                    [0,0,0,1]])
    cam_param.intrinsic.set_intrinsics(width=3840, height=2160, fx=1.80820276e+03, fy=1.80794556e+03, cx=1.94228662e+03, cy=1.12382178e+03)

    # get depth
    viewer = o3d.visualization.Visualizer()
    viewer.create_window()
    viewer.add_geometry(pcd)
    viewer.run()
    control = viewer.get_view_control()
    control.convert_from_pinhole_camera_parameters(cam_param)
    depth = viewer.capture_depth_float_buffer()

    plt.imshow(np.asarray(depth))
    plt.imsave("testing_depth.png", np.asarray(depth), dpi = 1)

preprocess/EgoPAT3D/pc2depth.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import sys
PYK4A_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/pyk4a')
sys.path.append(os.path.join(PYK4A_ROOT, 'example'))
from helpers import convert_to_bgra_if_required, colorize

import pyk4a
from pyk4a import PyK4APlayback
import cv2
import numpy as np


def read_by_whileloop(mkvraw_file, rgb_path, depth_path):
    playback = PyK4APlayback(mkvraw_file)
    playback.open()
    num=1
    depth_range = [0, 1000]
    while True:
        capture = playback.get_next_capture()
        if capture.color is not None:
            rgb = convert_to_bgra_if_required(playback.configuration["color_format"], capture.color)
            cv2.imwrite(os.path.join(rgb_path, '{}.jpg'.format(num)), rgb)
            if capture.depth is not None:
              depth = pyk4a.depth_image_to_color_camera(capture.depth, playback.calibration, playback.thread_safe)
              if depth is not None:
                if np.max(depth) > depth_range[1]:
                    depth_range[1] = np.max(depth)
                depth = colorize(depth, tuple(depth_range))
                cv2.imwrite(os.path.join(depth_path, '{}.png'.format(num)), depth)
                print('{}'.format(num))
            num += 1
            if num >= 1000:
                break
    cv2.destroyAllWindows()
    playback.close()


def get_specified_frame(playback, pointer, i):
    rgb, depth = None, None
    while True:
        capture = playback.get_next_capture()
        if (pointer == i) and (capture.color is not None) and (capture.depth is not None):
            rgb = convert_to_bgra_if_required(playback.configuration["color_format"], capture.color)
            depth = pyk4a.depth_image_to_color_camera(capture.depth, playback.calibration, playback.thread_safe)
            break
        if pointer < i:
            pointer += 1
        else:
            break
    return rgb, depth, pointer


def read_by_frameid(mkvraw_file, rgb_path, depth_path):
    playback = PyK4APlayback(mkvraw_file)
    playback.open()
    # fps_dict = {'FPS.FPS_5': 5, 'FPS.FPS_15': 15, 'FPS.FPS_30': 30}
    # fps = fps_dict[str(playback.configuration['camera_fps'])]
    # offset = playback.configuration["start_timestamp_offset_usec"]
    pointer = 0
    depth_range = [0, 1000]
    for i in range(20, 1001, 10):
        print('{}'.format(i))
        rgb, depth, pointer = get_specified_frame(playback, pointer, i)
        if (rgb is not None) and (depth is not None):
            cv2.imwrite(os.path.join(rgb_path, '{}.jpg'.format(i)), rgb)
            depth = colorize(depth, tuple(depth_range))
            cv2.imwrite(os.path.join(depth_path, '{}.png'.format(i)), depth)
    
        # pos = int(i / fps * 1000000)
        # playback.seek(pos, origin=SeekOrigin.BEGIN)
        # capture = playback.get_next_capture()
        # rgb = convert_to_bgra_if_required(playback.configuration["color_format"], capture.color)
        # depth = pyk4a.depth_image_to_color_camera(capture.depth, playback.calibration, playback.thread_safe)



def main():
    scene_id = '1'
    scene_name = 'bathroomCabinet'
    record_id = '2'
    record_name = '{}_{}'.format(scene_name, record_id)

    root_path = os.path.join(os.path.dirname(__file__), '../../data')
    record_path = os.path.join(root_path, 'EgoPAT3D-complete', scene_id, record_name)  # './1/bathroomCabinet_2'

    mkvraw_file = os.path.join(record_path, '{}.mkv'.format(record_name))
    # mkvraw_file = '../{}.mkv'.format(record_name)

    result_path = '../{}'.format(record_name)
    rgb_path = os.path.join(result_path, 'color')
    depth_path = os.path.join(result_path, 'depth')
    os.makedirs(rgb_path, exist_ok=True)
    os.makedirs(depth_path, exist_ok=True)

    read_by_whileloop(mkvraw_file, rgb_path, depth_path)

    # read_by_frameid(mkvraw_file, rgb_path, depth_path)



if __name__ == '__main__':

    main()

    
preprocess/EgoPAT3D/read_mkv_pyk4a.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import numpy as np
import pickle
import cv2
from tqdm import tqdm


def read_trajfile(trajfile):
    assert os.path.exists(trajfile), "File does not exist! {}".format(trajfile)
    with open(trajfile, 'rb') as f:
        traj_clips = pickle.load(f)
    return traj_clips


def get_all_clips(traj3d_all):
    all_clips = []
    for traj3d in traj3d_all:
        if len(traj3d) < 2: # at least two time points 
            continue
        start = int(list(traj3d.keys())[0])  # the start frame may not valid
        end = int(list(traj3d.keys())[-1])
        # find the first valid landmark
        for fid, landmarks in traj3d.items():
            pts3d = np.array(list(landmarks.values()))
            if len(pts3d[pts3d[:, 2] > 0]) > 0:
                start = fid  # move start forward
                break
        if start >= end:
            continue
        all_clips.append((start, end))
    return all_clips


def read_video(video_file, ratio=0.5):
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    videos = []
    while success:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        dst_size = (int(frame.shape[1] * ratio), int(frame.shape[0] * ratio))
        frame = cv2.resize(frame, dst_size)
        videos.append(frame)
        success, frame = cap.read()
    videos = np.array(videos)
    fps = cap.get(cv2.CAP_PROP_FPS) if len(videos) > 0 else 30
    return videos, fps


def write_video(mat, video_file, fps=30):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)


if __name__ == '__main__':
    
    ratio = 0.25  # we down-sampling the rgb video into 1/16 size for each clip (540 x 960)
    
    # trajectory directory
    root_dir = os.path.join(os.path.dirname(__file__), '../../data', 'EgoPAT3D')
    traj_dir = os.path.join(root_dir, 'EgoPAT3D-postproc', 'trajectory')
    
    # result directory
    videoclip_path = os.path.join(root_dir, 'EgoPAT3D-postproc', 'clip_videos')
    os.makedirs(videoclip_path, exist_ok=True)
    
    # get clip from trajectory annotation
    for scene_id in sorted(os.listdir(traj_dir)):
        scene_dir = os.path.join(traj_dir, scene_id)
        recordings = sorted(os.listdir(scene_dir))
        for filename in tqdm(recordings, total=len(recordings), desc='Process scene {}'.format(scene_id)):
            if not filename.endswith('pkl'):
                continue
            # read pkl trajectory file
            traj_file = os.path.join(scene_dir, filename)
            traj3d_all = read_trajfile(traj_file)
            
            # gather valid clips
            all_clips = get_all_clips(traj3d_all)
            
            record_name = filename.split('_')[0] + '_' + filename.split('_')[1]
            # create saving directory
            clip_savedir = os.path.join(videoclip_path, scene_id, record_name)
            os.makedirs(clip_savedir, exist_ok=True)
            
            # if some result clip files exist, check if there is a need to continue on the following codes
            need_continue = False
            for start, end in all_clips:
                # save each video clip
                clip_file = os.path.join(clip_savedir, 'clip_s{}_e{}.mp4'.format(start, end))
                if not os.path.exists(clip_file):
                    need_continue = True
                    break
                else:
                    clip_mat, _ = read_video(clip_file, ratio=1)
                    if len(clip_mat) == 0:  # Re-generate video clip
                        need_continue = True
                        break
            
            if need_continue:
                # read the whole mp4 video
                video_file = os.path.join(root_dir, 'EgoPAT3D-complete', scene_id, record_name, 'rgb_video.mp4')
                rgb_video, FPS = read_video(video_file, ratio=ratio)
                
                for start, end in all_clips:
                    # save each video clip
                    clip_file = os.path.join(clip_savedir, 'clip_s{}_e{}.mp4'.format(start, end))
                    if not os.path.exists(clip_file):
                        write_video(rgb_video[start: end+1], clip_file, fps=FPS)
preprocess/EgoPAT3D/split_videos.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

""" In this step, we use the provided manually annotated clip ranges and the hand landmarks
    to find the exact temporal region of each hand trajectory. Then, video clip of each trajectory is stored as mp4 file.
    At the same time, we resized the videos into 1/16 (1/4 x 1/4) of original 4K size.
"""
import os
from tqdm import tqdm
import cv2
import numpy as np


def read_clip_ranges(rng_file):
    ranges = dict()
    with open(rng_file, 'r') as f:
        for line in f.readlines():
            cid = int(line.strip().split('. ')[0])
            start = int(line.strip().split('. ')[1].split(', ')[0])
            end = int(line.strip().split('. ')[1].split(', ')[1])
            ranges[cid] = (start, end)
    return ranges



def read_landmarks(landmark_file, flip=False):
    landmarks = dict()
    with open(landmark_file, 'r') as f:
        pid = 0
        for line in f.readlines():
            line = line.strip()
            if line.split(':')[0].isnumeric():
                frame_id = int(line.split(':')[0])
                landmarks[frame_id] = np.zeros((21, 3))  # a new frame
                assert pid == 0 or pid == 20, 'invalid frame!'
                pid = -1
                continue
            if 'landmark' in line: 
                pid += 1  # a new 3D point
                continue
            if 'x:' in line: 
                xval = float(line.split(': ')[1])
                landmarks[frame_id][pid, 0] = 1.0 - xval if flip else xval
                continue
            if 'y:' in line:
                landmarks[frame_id][pid, 1] = float(line.split(': ')[1])
                continue
            if 'z:' in line:
                landmarks[frame_id][pid, 2] = float(line.split(': ')[1])
                continue
    return landmarks


def read_video(video_file, ratio=0.25):
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    videos = []
    while success:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        dst_size = (int(frame.shape[1] * ratio), int(frame.shape[0] * ratio))
        frame = cv2.resize(frame, dst_size)
        videos.append(frame)
        success, frame = cap.read()
    videos = np.array(videos)
    fps = cap.get(cv2.CAP_PROP_FPS) if len(videos) > 0 else 30
    return videos, fps


def get_hand_clip(start, end, landmarks):
    # get the first frame which has landmarks of hands
    first_frame = start
    while first_frame < end:
        if first_frame in landmarks:
            break
        first_frame += 1
    # get the last frame which has landmakrs of hands
    last_frame = end
    while last_frame > first_frame:
        if last_frame in landmarks:
            break
        last_frame -= 1
    return first_frame, last_frame


def compute_hand_center(landmarks):
    """ landmarks: (21.3) """
    pxyz = np.copy(landmarks)  # (21, 3)
    pxyz[1:, 2] += pxyz[0, 2]  # z values of top 20 are relative offsets to the 21-th z value
    center = np.mean(pxyz, axis=0)  # (3,)
    return center


def write_video(mat, video_file, fps=30):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)



if __name__ == '__main__':
    
    """ get all hand clips with landmarks in first frame """
    
    # input: hand landmarks directory
    root_path = os.path.join(os.path.dirname(__file__), '../../data/EgoPAT3D')
    resolution = [2160, 3840]  # height, width
    vis_ratio = 0.25
    CLIP_MIN, CLIP_MAX = 10, 5400  # more than ~3min or less than 10 frames (invalid clip):
    
    # output: first frame with hand center warpped by RAFT
    output_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'video_clips_hand')
    os.makedirs(output_dir, exist_ok=True)
    
    input_dir = os.path.join(root_path, 'EgoPAT3D-complete')
    for scene_id in sorted(os.listdir(input_dir)):
        # output path
        result_scene = os.path.join(output_dir, scene_id)
        os.makedirs(result_scene, exist_ok=True)
        
        for record_name in sorted(os.listdir(os.path.join(input_dir, scene_id))):
            record_path = os.path.join(input_dir, scene_id, record_name)
            if not os.path.isdir(record_path):
                continue
            
            # read hand landmarks
            landmark_file = os.path.join(record_path, 'hand_frames', 'hand_landmarks.txt')
            if not os.path.exists(landmark_file):
                continue  # discard record video without landmarks
            landmarks = read_landmarks(landmark_file, flip=True)
            
            # clip range file
            cliprng_file = os.path.join(record_path, 'hand_frames', 'clip_ranges.txt')
            clip_ranges = read_clip_ranges(cliprng_file)
            
            # output path
            result_record = os.path.join(result_scene, record_name)
            os.makedirs(result_record, exist_ok=True)
            
            # video file
            video_file = os.path.join(record_path, 'rgb_video.mp4')
            rgb_video, fps = read_video(video_file, vis_ratio)
            
            for cid, (start, end) in tqdm(clip_ranges.items(), desc=f'Record {record_name}', total=len(clip_ranges)):
                # get the first and last frames
                first_frame, last_frame = get_hand_clip(start, end, landmarks)
                
                # ignore too long or too short clips
                num_frames = last_frame - first_frame + 1
                if num_frames > CLIP_MAX or num_frames < CLIP_MIN:
                    continue
                
                # hand center (2D) of the start frame
                hand_center_1 = compute_hand_center(landmarks[first_frame])  # (u,v,z)
                px1 = int(resolution[1] * vis_ratio * hand_center_1[0])
                py1 = int(resolution[0] * vis_ratio * hand_center_1[1])
                
                # hand center (2D) of the end frame
                hand_center_2 = compute_hand_center(landmarks[last_frame])  # (u,v,z)
                px2 = int(resolution[1] * vis_ratio * hand_center_2[0])
                py2 = int(resolution[0] * vis_ratio * hand_center_2[1])
                
                # save frame
                result_clip_file = os.path.join(result_record, 'c{}_s{}_e{}_sx{}_sy{}_ex{}_ey{}.mp4'.format(cid, first_frame, last_frame, px1, py1, px2, py2))
                if not os.path.exists(result_clip_file):
                    write_video(rgb_video[first_frame: last_frame + 1], result_clip_file, fps=fps)
            
preprocess/EgoPAT3D/step1_split_videos_with_landmarks.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

""" In this step, we use RAFT model to get the forward and backward optical flows.
    The flows are used to warp the initial & the last hand location to get the forward & backward trajectories.
    Finally, the hand trajectory is the mean of the forward & backward trajectories.
"""
import os
from tqdm import tqdm
import cv2
import numpy as np
import torch
import argparse
import sys
import pickle

RAFT_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/RAFT')
sys.path.append(os.path.join(RAFT_ROOT, 'core'))
from raft import RAFT
from utils.utils import InputPadder



def init_raft(args, raft_root='./RAFT'):
    # class cfg: pass
    args.model = os.path.join(raft_root, 'models', 'raft-things.pth')
    args.small = False  # do not use small model
    args.mixed_precision = False  # do not use mixed precision inference
    args.alternate_corr = False  # do not use efficent correlation implementation
    args.iters = 32
    args.device = torch.device('cuda')
    # cfg = vars(cfg)  # class to dict
    # for k, v in cfg.items():
    #     parser.add_argument('--' + k, default=v)
    # args = parser.parse_args()  # dict to parser
    
    model = torch.nn.DataParallel(RAFT(args))
    model.load_state_dict(torch.load(args.model))
    model = model.module
    model.to(args.device)
    model.eval()
    return model, args


def read_video(video_file):
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    video = []
    while success:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        video.append(frame)
        success, frame = cap.read()
    video = np.array(video)
    return video


def write_video(mat, video_file, fps=30):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)
        
        
def coord_check(px, py, img_size=(540, 960)):
    px_new = max(min(px, img_size[1] - 1), 0)
    py_new = max(min(py, img_size[0] - 1), 0)
    return px_new, py_new


@torch.no_grad()
def flow_warpping(video, px, py, model, iters=120, device='cuda'):
    """ video: (T, H, W, C) """
    # toTensor() and toGPU()
    video_tensor = torch.from_numpy(video.astype(np.float32)).permute(0, 3, 1, 2).to(device)  # (T, C, H, W)
    padder = InputPadder(video_tensor.shape[1:])
    traj2d = [[px, py]]
    for image1, image2 in zip(video_tensor[:-1], video_tensor[1:]):
        # pad the frames
        image1, image2 = padder.pad(image1[None], image2[None])
        # flow inference
        _, flow_up = model(image1, image2, iters=iters, test_mode=True)  # (1, 2, 544, 960)
        flow_up = padder.unpad(flow_up)
        # warpping
        new_px = traj2d[-1][0] + float(flow_up[0, 0, int(traj2d[-1][1]), int(traj2d[-1][0])].cpu())  # u = u + du
        new_py = traj2d[-1][1] + float(flow_up[0, 1, int(traj2d[-1][1]), int(traj2d[-1][0])].cpu())  # v = v + dv
        # clipping
        new_px, new_py = coord_check(new_px, new_py, img_size=video_tensor.shape[2:])
        traj2d.append([new_px, new_py])
    traj2d = np.array(traj2d, dtype=np.float32)
    return traj2d


def run_warping():
    
    for scene_id in args.scenes:
        # output path
        result_scene = os.path.join(output_dir, scene_id)
        os.makedirs(result_scene, exist_ok=True)
        
        for record_name in sorted(os.listdir(os.path.join(clip_path, scene_id))):
            # input video folder
            video_path = os.path.join(clip_path, scene_id, record_name)
            
            # output path
            result_record = os.path.join(result_scene, record_name)
            os.makedirs(result_record, exist_ok=True)
            
            all_clips = os.listdir(video_path)
            for clip_name in tqdm(sorted(all_clips), desc=f'Record {record_name}', total=len(all_clips)):                
                # read video with hand location
                sx = int(clip_name.split('.mp4')[0].split('_')[-4][2:])
                sy = int(clip_name.split('.mp4')[0].split('_')[-3][2:])
                ex = int(clip_name.split('.mp4')[0].split('_')[-2][2:])
                ey = int(clip_name.split('.mp4')[0].split('_')[-1][2:])
                # read video
                video = read_video(os.path.join(video_path, clip_name))  # (N, H, W, C)
                # coordinate check
                sx, sy = coord_check(sx, sy, img_size=video.shape[1:3])
                ex, ey = coord_check(ex, ey, img_size=video.shape[1:3])
                
                with torch.no_grad():
                    # forward pass: warp the first hand location to the future frames by RAFT
                    traj2d_forward = flow_warpping(video, sx, sy, model, iters=args.iters, device=args.device)
                    # backward pass: warp the last hand location to the initial frames by RAFT
                    traj2d_backward = flow_warpping(video[::-1], ex, ey, model, iters=args.iters, device=args.device)
                    # take weighted mean
                    weights = args.conf + (1.0 - args.conf) / (1 + np.exp(np.arange(1, video.shape[0] + 1) - video.shape[0] / 2))
                    traj2d = traj2d_forward * weights[:, None] + traj2d_backward[::-1] * (1 - weights[:, None])
                
                # save the 2D trajectory
                result_traj_file = os.path.join(result_record, '{}.pkl'.format(clip_name.split('.mp4')[0]))
                with open(result_traj_file, 'wb') as f:
                    pickle.dump({'traj2d': traj2d}, f, protocol=pickle.HIGHEST_PROTOCOL)


def visualize(scene_id='1', record_name='bathroomCabinet_1', clip_name='c9_s439_e483_x438_y189'):
    
    # read trajectory file
    input_traj_file = os.path.join(output_dir, scene_id, record_name, clip_name + '.pkl')
    assert os.path.exists(input_traj_file), 'trajectory file does not exist!'
    with open(input_traj_file, 'rb') as f:
        data = pickle.load(f)
    traj2d = data['traj2d']  # (T, 2)
    
    # read video file
    video_file = os.path.join(clip_path, scene_id, record_name, clip_name + '.mp4')
    assert os.path.exists(video_file), 'video clip file does not exist!'
    video = read_video(video_file)  # (T, H, W, C)
    
    # visualize
    vis_video = []
    for hand_point, img in zip(traj2d, video):
        img_vis = cv2.circle(img, (int(hand_point[0]), int(hand_point[1])), radius=5, color=(0, 0, 255), thickness=-1)
        vis_video.append(img_vis)
    vis_video = np.array(vis_video)
    
    result_vis_file = os.path.join(output_dir, scene_id, record_name, 'vis_' + clip_name + '.mp4')
    write_video(vis_video, result_vis_file, fps=5)


if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_root', type=str, default='../../data/EgoPAT3D')
    parser.add_argument("--scenes", nargs='+', default=[], help='the scene ids.')
    parser.add_argument("--vis", action='store_true', help='visualize an example')
    parser.add_argument('--conf', type=float, default=0.3, help='the least confidence of forward warpping')
    args = parser.parse_args()
    
    root_path = os.path.join(os.path.dirname(__file__), args.data_root)
    
    # path to video clips with hand annotation of the first frame
    clip_path = os.path.join(root_path, 'EgoPAT3D-postproc', 'video_clips_hand')
    
    # result path
    output_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'trajectory_warpped')
    os.makedirs(output_dir, exist_ok=True)
    
    if args.vis:
        assert os.path.exists(output_dir)
        print('Visualize 2D trajectory...')
        visualize(scene_id='1', record_name='bathroomCabinet_1', clip_name='c100_s5241_e5270_sx563_sy454_ex722_ey355')
    else:
        # get input scenes
        all_scenes = os.listdir(clip_path)
        if len(args.scenes) > 0:
            for s in args.scenes:
                assert s in all_scenes
        else:
            args.scenes = all_scenes
        # initialize a RAFT flow model
        model, args = init_raft(args, raft_root=RAFT_ROOT)
    
        run_warping()
    
preprocess/EgoPAT3D/step2_flow_warped_2dhand_trajectory.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import pickle
import numpy as np
from tqdm import tqdm
import pyk4a
from pyk4a import PyK4APlayback
import argparse



def get_depth(playback, start, end, pointer):
    """get depth video clip"""
    depth_clip = np.zeros([end - start + 1] + resolution)  # (N, H, W)
    has_invalid = False
    # move pointer to the clip start frame
    while pointer < start:
        capture = playback.get_next_capture()
        pointer += 1
    # get the depth of the clip
    num_valid = 0
    while pointer < end + 1:
        capture = playback.get_next_capture()
        if (capture.color is None) or (capture.depth is None):
            pointer += 1  # move pointer
            has_invalid = True
            continue
        # assert (capture.color is not None) and (capture.depth is not None)
        depth_mat = pyk4a.depth_image_to_color_camera(capture.depth, playback.calibration, playback.thread_safe)
        depth_clip[pointer - start] = depth_mat
        pointer += 1  # move pointer
        num_valid += 1
    if num_valid < end - start + 1:
        has_invalid = True
    return depth_clip, pointer, has_invalid


def get_traj3d(traj2d, depth_clip):
    """ traj2d: (T, 2)
        depth_clip: (T, H, W)
    """
    traj3d = np.zeros((traj2d.shape[0], 3))
    for t, (coord, depth) in enumerate(zip(traj2d, depth_clip)):
        # scaling the 2D coordinates to raw 4K resolution
        u, v = int(coord[0] / vis_ratio), int(coord[1] / vis_ratio)
        # from 2D (u,v) to 3D (X, Y, Z) using the constant camera instrinsics
        Z = depth[v, u] / 1000  # mm to m
        X = (u - 1.94228662e+03) * Z / 1.80820276e+03
        Y = (v - 1.12382178e+03) * Z / 1.80794556e+03
        traj3d[t] = np.array([X, Y, Z])
    return traj3d
    


def fetch_trajectory_depth():
    
    for scene_id in sorted(args.scenes):
        for record_name in sorted(os.listdir(os.path.join(traj_dir, scene_id))):
            
            # get depth from raw MKV videos
            mkv_filepath = os.path.join(mkv_dir, scene_id, record_name + '.mkv')
            playback = PyK4APlayback(mkv_filepath)
            playback.open()
            pointer = 0
            
            # sort the filenames of clip trajectories in an ascending order
            all_trajfiles = os.listdir(os.path.join(traj_dir, scene_id, record_name))
            all_trajfiles = list(filter(lambda x: x.endswith('.pkl'), all_trajfiles))
            clip_ids = np.argsort([int(filename.split('_')[0][1:]) for filename in all_trajfiles])
            all_trajfiles = [all_trajfiles[i] for i in clip_ids]
            
            for traj_filename in tqdm(all_trajfiles, desc=f'Record {record_name}', total=len(all_trajfiles)):
                # read trajectory file
                traj_file = os.path.join(traj_dir, scene_id, record_name, traj_filename)
                with open(traj_file, 'rb') as f:
                    data = pickle.load(f)
                traj2d = data['traj2d']  # (T, 2)
                
                # get depth clip
                start_frame = int(traj_filename.split(".pkl")[0].split('_')[1][1:])
                end_frame = int(traj_filename.split(".pkl")[0].split('_')[2][1:])
                depth_clip, pointer, has_invalid = get_depth(playback, start_frame, end_frame, pointer)
                
                if has_invalid:
                    print("Invalid clip! Record: {}, Clip: {}".format(record_name, traj_filename))
                
                # get 3D trajectory (camera coordinate system)
                traj3d = get_traj3d(traj2d, depth_clip)
                data.update({'traj3d': traj3d})
                
                # write results to the pickle file
                with open(traj_file, 'wb') as f:
                    pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)

            playback.close()



if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_root', type=str, default='../../data/EgoPAT3D')
    parser.add_argument('--scenes', nargs='+', default=[])
    args = parser.parse_args()
    
    root_path = os.path.join(os.path.dirname(__file__), args.data_root)
    resolution = [2160, 3840]  # height, width
    vis_ratio = 0.25
    
    # result path
    traj_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'trajectory_warpped')
    assert os.path.exists(traj_dir), 'Path does not exist!'
    
    # mkv video path which contains depth
    mkv_dir = os.path.join(root_path, 'EgoPAT3D-mkv')
    
    if len(args.scenes) == 0:
        args.scenes = os.listdir(traj_dir)
    print("Start processing scenes: {}".format(args.scenes))
    
    fetch_trajectory_depth()
    
    print("Finished scenes: {}".format(args.scenes))
preprocess/EgoPAT3D/step3_fetch_depth_traj3d.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

""" 
    In this step, we repair the 3D trajectories whose depths are partially (or entirely) missing/invalid, and those contains outliers in tail part.
    For the missing depths, we use the valid hand depth and corresponding frame indices to fit a multinomial with sine model by least-square fitting. 
    And the missing depths are estimated. For move the tail outliers, we use first-order different to identify the cut-off location.
"""
import os
from tqdm import tqdm
import pickle
import numpy as np
from scipy.optimize import leastsq
import matplotlib.pyplot as plt
import shutil
    


def read_traj_file(filepath, target='3d'):
    assert os.path.exists(filepath), "File does not exist! {}".format(filepath)
    with open(filepath, 'rb') as f:
        data = pickle.load(f)
    traj2d = data['traj2d']
    traj3d = None
    if target == '3d':
        assert 'traj3d' in data, "3D trajectories do not exist in file: {}".format(filepath)
        traj3d = data['traj3d']
    return traj2d, traj3d


def repair_depth(traj3d, traj2d, invalids, vis_ratio, vis=False):
    """ repair the depth using least square f"""
    def model_func(p, x):
        f = np.poly1d(p[:-2])
        y = f(x) + p[-2] * np.sin(p[-1]*x)
        return y

    def error_func(p, x, y):
        error = y - model_func(p, x)
        return error
    
    # get the valid true depth
    fit_data, fit_labels = [], []
    for ids, coord in enumerate(traj2d):
        if ids not in invalids:  # valid segments
            u, v = int(coord[0]), int(coord[1])
            fit_data.append(ids)
            fit_labels.append(traj3d[ids, 2])

    # least square fitting by third-order multinomial + sin model
    ret = leastsq(error_func, [0, 0, 1, 0, 1, 0.05], args=(np.array(fit_data), np.array(fit_labels)), epsfcn=0.0001)  # start with straight line: y=0
    
    # repaire
    traj3d_repaire = np.copy(traj3d)
    test_data, test_labels = [], []
    for ids in invalids:
        # repaire 3D coordinates
        u, v = int(traj2d[ids, 0]), int(traj2d[ids, 1])  # (u,v) in (540, 960) frame
        Z = model_func(ret[0], ids)
        test_data.append(ids)
        test_labels.append(Z)
        # (u,v) + Z -> (X, Y)
        X = (u / vis_ratio - 1.94228662e+03) * Z / 1.80820276e+03
        Y = (v / vis_ratio - 1.12382178e+03) * Z / 1.80794556e+03
        # substitute (X, Y, Z)
        traj3d_repaire[ids] = np.array([X, Y, Z])
    
    if vis:
        # plot the fitting 
        plt.figure()
        plt.scatter(fit_data, fit_labels, label='train data')
        x_draw = np.arange(1, len(traj2d) + 1, 0.0001)
        plt.plot(x_draw, model_func(ret[0], x_draw), 'r-', label='model')
        plt.scatter(test_data, test_labels, marker='^', label='test data')
        plt.xlabel('Frame index')
        plt.ylabel('Hand depth (m)')
        plt.legend()
        plt.tight_layout()
    
    return traj3d_repaire


def repair_depth_by_model(traj2d, depth_maps, vis_ratio=0.25):
    """ Replace with the estimated depth """
    traj3d = np.zeros((traj2d.shape[0], 3))
    for ids, (coord, depth) in enumerate(zip(traj2d, depth_maps)):
        u, v = int(coord[0]), int(coord[1])
        Z = depth[v, u] * 0.5
        X = (u / vis_ratio - 1.94228662e+03) * Z / 1.80820276e+03
        Y = (v / vis_ratio - 1.12382178e+03) * Z / 1.80794556e+03
        traj3d[ids] = np.array([X, Y, Z])
    return traj3d



def count_invalid():
    
    total_invalid, total = 0, 0
    for scene_id in sorted(os.listdir(traj_dir)):
        for record_name in sorted(os.listdir(os.path.join(traj_dir, scene_id))):
            
            num_clips_invalid = 0
            all_trajfiles = os.listdir(os.path.join(traj_dir, scene_id, record_name))
            all_trajfiles = list(filter(lambda x: x.endswith('.pkl'), all_trajfiles))
            total += len(all_trajfiles)
            
            for traj_filename in sorted(all_trajfiles):
                # read trajectory file
                traj_file = os.path.join(traj_dir, scene_id, record_name, traj_filename)
                traj2d, traj3d = read_traj_file(traj_file)  # (N,2), (N,3)
                
                invalids = np.where(traj3d[:, 2] <= MIN_DEPTH)[0]
                if len(invalids) >= len(traj3d) - MIN_VALID_POINTS:  # all points are valid
                    num_clips_invalid += 1

            total_invalid += num_clips_invalid
            print("Record:{}, total clips: {}, invalid clips: {}".format(record_name, len(all_trajfiles), num_clips_invalid))
    
    print('Total invalid: {} / {}'.format(total_invalid, total))


def remove_tail_outliers(Zs, num_tails=6, thresh=0.15):
    """ Remove the outlier depth in the tail of a trajectory"""
    num_preserve = len(Zs)
    if len(Zs) >= num_tails:
        tails = Zs[-num_tails:]
        diffs = np.abs(tails[1:] - tails[:-1])
        if np.max(diffs) > thresh:
            num_preserve = len(Zs) - num_tails + np.argmax(diffs) + 1
    return num_preserve


def main():
    
    num_vis = 0
    for scene_id in sorted(os.listdir(traj_dir)):
        for record_name in sorted(os.listdir(os.path.join(traj_dir, scene_id))):
            # repaire results dir 
            traj_repair_path = os.path.join(traj_repair_dir, scene_id, record_name)
            os.makedirs(traj_repair_path, exist_ok=True)

            all_trajfiles = os.listdir(os.path.join(traj_dir, scene_id, record_name))
            all_trajfiles = list(filter(lambda x: x.endswith('.pkl'), all_trajfiles))
            for traj_filename in tqdm(all_trajfiles, desc=f'Record {record_name}', total=len(all_trajfiles)):
                # result file
                traj_repair_file = os.path.join(traj_repair_path, traj_filename)
                
                # read trajectory file
                traj_file = os.path.join(traj_dir, scene_id, record_name, traj_filename)
                traj2d, traj3d = read_traj_file(traj_file)  # (N,2), (N,3)
                
                # trim the trajectory rail part
                num_preserve = remove_tail_outliers(traj3d[:, 2])
                traj2d, traj3d = traj2d[:num_preserve], traj3d[:num_preserve]
                
                invalids = np.where(traj3d[:, 2] <= MIN_DEPTH)[0]
                if len(invalids) == 0:  # all points are valid
                    pass # no further action needed
                
                elif len(invalids) > 0 and len(invalids) <= len(traj3d) - MIN_VALID_POINTS: # there are invalid depths, and at least 5 valid depths
                    # repair invalid depth
                    traj3d = repair_depth(traj3d, traj2d, invalids, vis=VIS)
                    num_vis += 1
                    if VIS and num_vis <= 20:
                        plt.savefig('../output/temp_trim_fit/fit_{}.png'.format(num_vis))
                        plt.close()

                else:  # most points are invalid, we directly discard the trajectory
                    # print("Discard the trajectory: {}/{}/{}".format(scene_id, record_name, traj_filename))
                    continue

                # write results to the pickle file
                with open(traj_repair_file, 'wb') as f:
                    pickle.dump({'traj2d': traj2d, 'traj3d': traj3d, 'num_preserve': num_preserve}, f, protocol=pickle.HIGHEST_PROTOCOL)

    # delete empty trajectory dir
    for scene_id in sorted(os.listdir(traj_repair_dir)):
        for record_name in sorted(os.listdir(os.path.join(traj_repair_dir, scene_id))):
            traj_folder = os.path.join(traj_repair_dir, scene_id, record_name)
            all_trajfiles = list(filter(lambda x: x.endswith('.pkl'), os.listdir(traj_folder)))
            if len(all_trajfiles) == 0:
                shutil.rmtree(traj_folder)
                print("Removed the empty folder: {}".format(traj_folder))


if __name__ == '__main__':
    
    # root path
    root_path = os.path.join(os.path.dirname(__file__), '../../data/EgoPAT3D')
    vis_ratio = 0.25
    MIN_DEPTH = 0.001
    MIN_VALID_POINTS = 10  # the least number of valid 3D points for a valid trajectory
    NUM_TAIL_FRAMES = 6
    VIS=True
    
    # trajectory path
    traj_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'trajectory_warpped')
    assert os.path.exists(traj_dir), 'Path does not exist!'
    
    # video path
    rgb_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'video_clips_hand')
    assert os.path.exists(rgb_dir), 'Path does not exist!'
    
    # repaired trajectory path
    traj_repair_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'trajectory_repair')
    os.makedirs(traj_repair_dir, exist_ok=True)
    
    # count_invalid()
    
    main()

preprocess/EgoPAT3D/step4_repair_invalid_depth.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

""" 
    In this file, we compare the visualized 3D trajectories amon the true depth (from RGB-D camera recordings),
    the repaired depth by least square fitting, and the estimated depth (from NewCRFs model)
"""
import os
import numpy as np
import matplotlib.pyplot as plt
import cv2
import pickle
from matplotlib.animation import FFMpegWriter
from step4_repair_invalid_depth import read_traj_file, remove_tail_outliers, repair_depth
from depth_estimation import init_depth_model, monocular_depth_estimation



def read_video(video_file):
    assert os.path.exists(video_file), "File does not exist! {}".format(video_file)
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    video = []
    while success:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        video.append(frame)
        success, frame = cap.read()
    video = np.array(video)
    return video


def repair_depth_by_model(traj3d, traj2d, invalids, video, model, cfg, vis_ratio=0.25):
    """ Replace with the estimated depth for invalid depth frames """
    # compute depth scale
    all_scales = []
    for ids in range(len(traj3d)):
        if ids in invalids:
            continue
        # depth estimation
        depth_map = monocular_depth_estimation([video[ids]], model, cfg)
        # repaire 3D coordinates
        u, v = int(traj2d[ids, 0]), int(traj2d[ids, 1])  # (u,v) in (540, 960) frame
        Z = depth_map[0][v, u]
        # compute scale
        s = traj3d[ids, 2] / Z
        all_scales.append(s)
    scale = np.mean(np.array(all_scales))
    
    # repair depth with the scale and depth estimation model
    traj3d_repaire = np.copy(traj3d)
    for ids in invalids:
        # depth estimation
        depth_map = monocular_depth_estimation([video[ids]], model, cfg)
        # repaire 3D coordinates
        u, v = int(traj2d[ids, 0]), int(traj2d[ids, 1])  # (u,v) in (540, 960) frame
        Z = depth_map[0][v, u] * scale
        X = (u / vis_ratio - 1.94228662e+03) * Z / 1.80820276e+03
        Y = (v / vis_ratio - 1.12382178e+03) * Z / 1.80794556e+03
        traj3d_repaire[ids] = np.array([X, Y, Z])
    
    return traj3d_repaire


def find_vis_samples(model=None, cfg=None):
    # for each scene, find a complete trajectory as an example for visualization
    examples = dict()
    for scene_id in sorted(os.listdir(traj_dir)):
        found = False
        
        for record_name in sorted(os.listdir(os.path.join(traj_dir, scene_id))):
            # loop all trajectory files
            all_trajfiles = os.listdir(os.path.join(traj_dir, scene_id, record_name))
            all_trajfiles = list(filter(lambda x: x.endswith('.pkl'), all_trajfiles))
            clip_ids = np.argsort([int(filename.split('_')[0][1:]) for filename in all_trajfiles])
            all_trajfiles = [all_trajfiles[i] for i in clip_ids]
            
            for traj_filename in all_trajfiles:
                # read trajectory file
                traj_file = os.path.join(traj_dir, scene_id, record_name, traj_filename)
                traj2d, traj3d = read_traj_file(traj_file)  # (N,2), (N,3)
                
                # trim the trajectory rail part
                num_preserve = remove_tail_outliers(traj3d[:, 2])
                traj2d, traj3d = traj2d[:num_preserve], traj3d[:num_preserve]
                
                invalids = np.where(traj3d[:, 2] <= MIN_DEPTH)[0]
                if len(invalids) > 3 and len(invalids) < len(traj3d) - MIN_VALID_POINTS:
                    # repair by least square fitting
                    traj3d_lsq = repair_depth(traj3d, traj2d, invalids, vis_ratio)
                    
                    # read video
                    video_file = os.path.join(rgb_dir, scene_id, record_name, traj_filename[:-4] + '.mp4')
                    video = read_video(video_file)  # rgb(T, H, W, C)
                    video = video[:num_preserve]
                    
                    # repair by depth estimation model
                    traj3d_est = repair_depth_by_model(traj3d, traj2d, invalids, video, model, cfg)
                    
                    examples[scene_id] = {'scene_id': scene_id, 'record_name': record_name, 'clip_name': traj_filename[:-4],
                                          'traj2d': traj2d, 'traj3d_true': traj3d, 'traj3d_lsq': traj3d_lsq, 'traj3d_est': traj3d_est, 'num_preserve': num_preserve}
                    found = True
                    print('Found one video clip (len={}): {}/{}/{}'.format(len(traj3d), scene_id, record_name, traj_filename.split('_')[0]))
                
                if found:
                    break  # break the clip loop
            if found:
                break  # break the record loop
    
    return examples


def draw_uv_points(video, traj2d):
    """ draw uv location on rgb video """
    video_vis = np.copy(video)
    for frame, coord in zip(video_vis, traj2d):
        u, v = int(coord[0]), int(coord[1])
        cv2.circle(frame, (u, v), radius=5, color=(0, 255, 0), thickness=-1)
    return video_vis


def draw_depth_comparison(traj3d_true, traj3d_lsq, traj3d_est, canvas_height=540, fig_size=5, fps=5, savefile='temp.mp4'):
    
    # initialize a canvas for visualization
    fig = plt.figure(figsize=(fig_size, fig_size))
    ax = fig.add_subplot()
    ax.set_xlabel('frame id')
    ax.set_ylabel('hand depth (m)')
    plt.locator_params(axis='x', nbins=5)
    plt.locator_params(axis='y', nbins=5)
    
    Z_true, Z_lsq, Z_est = traj3d_true[:, 2], traj3d_lsq[:, 2], traj3d_est[:, 2]
    max_depth = max(np.max(Z_true), np.max(Z_est))
    plt.ylim(0, max_depth + 0.01)
    plt.xlim(0, traj3d_true.shape[0] + 1)
    plt.grid(axis='y')
    
    # draw with plt canvas
    traj_writer = FFMpegWriter(fps=fps, metadata=dict(title='depth comparison', artist='Matplotlib'))
    with traj_writer.saving(fig, savefile, dpi=int(canvas_height / fig_size)):
        for i in range(len(Z_true)):
            ax.scatter(np.arange(1, i+2), Z_est[:(i+1)], color='b', label='Repair by NewCRFs')
            ax.scatter(np.arange(1, i+2), Z_lsq[:(i+1)], color='k', label='Repair by LSF')
            ax.scatter(np.arange(1, i+2), Z_true[:(i+1)], color='r', label='RGB-D Camera')
            if i == 0:
                plt.legend()
            traj_writer.grab_frame()
    plt.close()
    
    # read the cached video
    traj_video = read_video(savefile)
    
    return traj_video


def write_video(mat, video_file, fps=5):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)


def vis_compare(examples):
    # vis
    for scene_id, data in examples.items():
        # read rgb video
        video_file = os.path.join(rgb_dir, data['scene_id'], data['record_name'], data['clip_name'] + '.mp4')
        rgb = read_video(video_file)  # rgb(T, H, W, C)
        rgb = rgb[:data['num_preserve']]
        
        # draw 2D hand locations on RGB
        vis_mat = draw_uv_points(rgb, data['traj2d'])

        # draw 3D trajectory on XoZ plane
        compare_vis = draw_depth_comparison(data['traj3d_true'], data['traj3d_lsq'], data['traj3d_est'], canvas_height=rgb.shape[1], fps=VIS_FPS)
        vis_mat = np.concatenate([vis_mat, compare_vis], axis=2) 
        
        # save visualized video
        filename = 'compare_{}_{}_{}'.format(data['scene_id'], data['record_name'], data['clip_name'].split("_")[0])
        vis_file = os.path.join(vis_path, filename + '.mp4')
        write_video(vis_mat, vis_file, fps=VIS_FPS)


def plot_compare(examples, fig_size=5, fontsize=20):
    
    for scene_id, data in examples.items():
        
        traj3d_true, traj3d_lsq, traj3d_est = data['traj3d_true'], data['traj3d_lsq'], data['traj3d_est']
        
        # draw on figure
        fig = plt.figure(figsize=(fig_size, fig_size))
        ax = fig.add_subplot()
        ax.set_xlabel('frame id', fontsize=fontsize)
        ax.set_ylabel('hand depth (m)', fontsize=fontsize)
        plt.locator_params(axis='x', nbins=5)
        plt.locator_params(axis='y', nbins=5)
        
        Z_true, Z_lsq, Z_est = traj3d_true[:, 2], traj3d_lsq[:, 2], traj3d_est[:, 2]
        max_depth = max(np.max(Z_true), np.max(Z_est))
        plt.ylim(0, max_depth + 0.01)
        plt.xlim(0, traj3d_true.shape[0] + 1)
        plt.grid(axis='y')
        
        ax.scatter(np.arange(1, len(Z_true) + 1), Z_est[:len(Z_true)], color='b', label='Repair by NewCRFs')
        ax.scatter(np.arange(1, len(Z_true) + 1), Z_lsq[:len(Z_true)], color='k', label='Repair by LSF')
        ax.scatter(np.arange(1, len(Z_true) + 1), Z_true[:len(Z_true)], color='r', label='RGB-D Camera')
        plt.legend(fontsize=fontsize)
        plt.tight_layout()
        
        # save
        fig_dir = os.path.join(vis_path, 'figures')
        os.makedirs(fig_dir, exist_ok=True)
        filename = 'compare_{}_{}_{}'.format(data['scene_id'], data['record_name'], data['clip_name'].split("_")[0])
        plt.savefig(os.path.join(fig_dir, filename + '.png'))
        plt.savefig(os.path.join(fig_dir, filename + '.pdf'))


if __name__ == '__main__':
    
    # root path
    root_path = os.path.join(os.path.dirname(__file__), '../../data/EgoPAT3D')
    MIN_DEPTH = 0.001
    MIN_VALID_POINTS = 10
    NUM_TAIL_FRAMES = 6
    VIS_FPS = 5
    vis_ratio = 0.25
    
    # trajectory path
    traj_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'trajectory_warpped')
    assert os.path.exists(traj_dir), 'Path does not exist!'
    
    # video path
    rgb_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'video_clips_hand')
    assert os.path.exists(rgb_dir), 'Path does not exist!'
    
    # result path
    vis_path = os.path.join(os.path.dirname(__file__), '../../output/EgoPAT3D/compare_gt_depth')
    os.makedirs(vis_path, exist_ok=True)
    vis_data_file = os.path.join(vis_path, 'vis_data.pkl')
    
    if not os.path.exists(vis_data_file):
        # initialize depth estimation model
        model, cfg = init_depth_model()
        
        # find examples for visualization
        examples = find_vis_samples(model, cfg)
        
        # save
        with open(vis_data_file, 'wb') as f:
            pickle.dump(examples, f, protocol=pickle.HIGHEST_PROTOCOL)
    else:
        with open(vis_data_file, 'rb') as f:
            examples = pickle.load(f)
    
    print("Visualizing the comparison...")
    # vis_compare(examples)
    
    plot_compare(examples)
preprocess/EgoPAT3D/step5_compare_depth.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import argparse
from tqdm import tqdm
import pickle
import numpy as np
try:
    # the pyk4a is only installed on monitor-connected Windows/Linux machine
    import sys
    PYK4A_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/pyk4a')
    sys.path.append(os.path.join(PYK4A_ROOT, 'example'))
    from helpers import convert_to_bgra_if_required
    import pyk4a
    from pyk4a import PyK4APlayback
    import open3d as o3d
except:
    pass
import cv2
from functools import reduce
import imageio


def read_video(video_file):
    assert os.path.exists(video_file), "File does not exist! {}".format(video_file)
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    video = []
    while success:
        video.append(frame)
        success, frame = cap.read()
    video = np.array(video)
    return video


def read_traj_file(filepath, target='3d'):
    assert os.path.exists(filepath), "File does not exist! {}".format(filepath)
    with open(filepath, 'rb') as f:
        data = pickle.load(f)
    traj2d = data['traj2d']
    traj3d = None
    if target == '3d':
        assert 'traj3d' in data, "3D trajectories do not exist in file: {}".format(filepath)
        traj3d = data['traj3d']
    num_preserve = data['num_preserve'] if 'num_preserve' in data else len(traj2d)
    return traj2d, traj3d, num_preserve


def read_odometry(odom_file):
    assert os.path.exists(odom_file), "File does not exist! {}".format(odom_file)
    all_transforms = np.load(odom_file)  # (T, 4, 4)
    return all_transforms


def get_rgb_depth(playback, start, end, pointer, ratio=1.0):
    """get depth video clip"""
    height = int(intrinsics['h'] * ratio)
    width = int(intrinsics['w'] * ratio)
    rgb_clip = np.zeros([end - start + 1, height, width, 3])  # (N, H, W, C)
    depth_clip = np.zeros([end - start + 1, height, width])  # (N, H, W)
    validities = np.zeros([end - start + 1], dtype=np.bool_)
    # move pointer to the clip start frame
    while pointer < start:
        capture = playback.get_next_capture()
        pointer += 1
    # get the depth of the clip
    while pointer < end + 1:
        capture = playback.get_next_capture()
        if (capture.color is None) or (capture.depth is None):
            pointer += 1  # move pointer
            has_invalid = True
            continue
        # get rgb frame
        rgb_mat = convert_to_bgra_if_required(playback.configuration["color_format"], capture.color)
        rgb_clip[pointer - start] = cv2.resize(rgb_mat, dsize=(width, height))
        # get depth frame
        depth_mat = pyk4a.depth_image_to_color_camera(capture.depth, playback.calibration, playback.thread_safe)
        depth_clip[pointer - start] = cv2.resize(depth_mat, dsize=(width, height))
        validities[pointer - start] = True
        pointer += 1  # move pointer
    return rgb_clip, depth_clip, pointer, validities


def compute_video_odometry(rgb_clip, depth_clip, validities):
    
    assert rgb_clip.shape[0] == depth_clip.shape[0]
    num_frames = rgb_clip.shape[0]
    
    cam_model = o3d.camera.PinholeCameraIntrinsic()
    cam_model.set_intrinsics(int(intrinsics['w']), int(intrinsics['h']), intrinsics['fx'], intrinsics['fy'], intrinsics['ox'], intrinsics['oy'])
    option = o3d.pipelines.odometry.OdometryOption()
    odo_method = o3d.pipelines.odometry.RGBDOdometryJacobianFromHybridTerm()
    odo_init = np.identity(4)
    
    # construct the rgb-depth pair
    rgbd_list = []
    for t in range(num_frames):
        rgb = o3d.geometry.Image(rgb_clip[t].astype(np.uint8))
        depth = o3d.geometry.Image(depth_clip[t].astype(np.uint8))
        rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(rgb, depth)
        rgbd_list.append(rgbd)
    
    # calculate odometry for each pair of adjacent RGBD frames
    all_transforms = [odo_init]
    for source, target, valid in zip(rgbd_list[1:], rgbd_list[:-1], validities[1:]):
        success, transform, info = o3d.pipelines.odometry.compute_rgbd_odometry(source, target, cam_model, odo_init, odo_method, option)
        if (not success) or (not valid):
            transform = np.identity(4)
        all_transforms.append(transform)
    
    return np.array(all_transforms)


def main():
    
    for scene_id in sorted(args.scenes):
        for record_name in sorted(os.listdir(os.path.join(traj_dir, scene_id))):
            # odometry results dir 
            odometry_result_path = os.path.join(odometry_dir, scene_id, record_name)
            os.makedirs(odometry_result_path, exist_ok=True)

            # trajectory file list must be sorted according to the clip order
            traj_folder = os.path.join(traj_dir, scene_id, record_name)
            all_trajfiles = list(filter(lambda x: x.endswith('.pkl'), os.listdir(traj_folder)))
            clip_ids = np.argsort([int(filename.split('_')[0][1:]) for filename in all_trajfiles])
            all_trajfiles = [all_trajfiles[i] for i in clip_ids]
            
            if len(os.listdir(odometry_result_path)) == len(all_trajfiles):
                continue
            
            # get depth from raw MKV videos
            mkv_filepath = os.path.join(mkv_dir, scene_id, record_name + '.mkv')
            playback = PyK4APlayback(mkv_filepath)
            playback.open()
            pointer = 0
            
            for traj_filename in tqdm(all_trajfiles, desc=f'Record {record_name}', total=len(all_trajfiles)):
                # result file
                odometry_file = os.path.join(odometry_result_path, traj_filename[:-4] + '.npy')

                # get depth clip
                start_frame = int(traj_filename.split(".pkl")[0].split('_')[1][1:])
                end_frame = int(traj_filename.split(".pkl")[0].split('_')[2][1:])
                rgb_clip, depth_clip, pointer, validities = get_rgb_depth(playback, start_frame, end_frame, pointer)
                
                # compute the odometry matrix for each frame
                all_transforms = compute_video_odometry(rgb_clip, depth_clip, validities)
                
                # save
                # write results to the pickle file
                with open(odometry_file, 'wb') as f:
                    np.save(f, all_transforms, allow_pickle=True)
            
            playback.close()


def XYZ_to_uv(traj3d, max_depth=3):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points 
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['ox'])
    traj2d[:, 1] = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['oy'])
    # clip the coordinates 
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['w'])
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['h'])
    traj2d = np.floor(traj2d).astype(np.int32)
    
    return traj2d


def vis_frame_traj(frame, traj2d):
    """ frame: (H, W, 3)
        traj2d: (T, 2)
    """
    cv2.polylines(frame, [traj2d], False, (0, 255, 255))
    for t, uv in enumerate(traj2d):
        cv2.circle(frame, uv, radius=5, color=(0, 0, 255), thickness=-1)


def video_to_gif(video, giffile):
    assert giffile.endswith('.gif')
    with imageio.get_writer(giffile, mode='I', duration=0.2) as writer:
        for frame in video:
            writer.append_data(frame)
    
            
def visualize_transformed_traj(video, traj3d, odometry):
    """ video: (T, H, W, 3)
        traj3d: (T, 3)
        odometry: (T, 4, 4)
    """
    length = len(odometry)
    traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
    for i in range(length):
        # compute the transformed points of all future points
        traj_list = [traj3d[i]]  # initial 3D point
        for j in range(i + 1, length):
            odom = reduce(np.dot, odometry[(i+1):(j+1)])  # (4, 4)
            future_point = odom.dot(traj3d_homo[j].T)  # (4,)
            traj_list.append(future_point[:3])
        # visualize current frame
        traj2d = XYZ_to_uv(np.array(traj_list))
        vis_frame_traj(video[i], traj2d)


def vis_odometry():
    
    # examples to visualize
    examples = [{'scene': '1', 'record': 'bathroomCabinet_1', 'clip': 'c9_s439_e480_sx438_sy189_ex605_ey309'},
                {'scene': '1', 'record': 'bathroomCabinet_9', 'clip': 'c1_s64_e103_sx718_sy420_ex631_ey247'},
                {'scene': '2', 'record': 'bathroomCounter_10', 'clip': 'c1_s44_e152_sx600_sy381_ex510_ey415'},
                {'scene': '3', 'record': 'bin_7', 'clip': 'c3_s168_e204_sx550_sy315_ex494_ey354'}
                ]
    for data in examples:
        # input files
        video_file = os.path.join(video_dir, data['scene'], data['record'], data['clip'] + '.mp4')
        traj_file = os.path.join(traj_dir, data['scene'], data['record'], data['clip'] + '.pkl')
        odom_file = os.path.join(odometry_dir, data['scene'], data['record'], data['clip'] + '.npy')
        # read trajectory
        traj2d, traj3d, num_preserve = read_traj_file(traj_file)
        # read video 
        video = read_video(video_file)
        video = video[:num_preserve]
        # read odometry
        odometry = read_odometry(odom_file)
        odometry = odometry[:num_preserve]
        
        visualize_transformed_traj(video, traj3d, odometry)

        # write to GIF
        giffile = os.path.join(vis_path, 'odom_{}_{}_{}.gif'.format(data['scene'], data['record'], data['clip'].split("_")[0]))
        video_to_gif(video, giffile)


if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_root', type=str, default='../../data/EgoPAT3D')
    parser.add_argument('--scenes', nargs='+', default=[])
    parser.add_argument('--vis', action='store_true')
    args = parser.parse_args()
    
    # root path
    root_path = os.path.join(os.path.dirname(__file__), args.data_root)
    vis_ratio = 0.25
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}
    for k, v in intrinsics.items():
        v *= vis_ratio  # scaling the camera intrinsics
        intrinsics[k] = v
    
    if not args.vis:  # not needed for visualization
        # mkv video path
        mkv_dir = os.path.join(root_path, 'EgoPAT3D-mkv')
        assert os.path.exists(mkv_dir), 'Path does not exist!'
    
    # (relative) trajectory path
    traj_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'trajectory_repair')
    assert os.path.exists(traj_dir), 'Path does not exist!'
    
    # result odometry path
    odometry_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'odometry')
    os.makedirs(odometry_dir, exist_ok=True)
    
    if len(args.scenes) == 0:
        args.scenes = os.listdir(traj_dir)
    print("Start processing scenes: {}".format(args.scenes))
    
    if args.vis:
        video_dir = os.path.join(root_path, 'EgoPAT3D-postproc', 'video_clips_hand')
        vis_path = os.path.join(os.path.dirname(__file__), '../../output/temp_odom_vis')
        os.makedirs(vis_path, exist_ok=True)
        # vis
        vis_odometry()
    else:
        main()
    
    print("Finished scenes: {}".format(args.scenes))

preprocess/EgoPAT3D/step6_compute_odometry.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from tqdm import tqdm

from google.protobuf import text_format
from mediapipe.framework.formats import landmark_pb2
from mediapipe.python.solutions import drawing_utils
import mediapipe.python.solutions.drawing_styles as mp_style
import mediapipe.python.solutions.hands as mp_hands


def read_video(video_file, ratio=0.5):
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    videos = []
    while success:
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        dst_size = (int(frame.shape[1] * ratio), int(frame.shape[0] * ratio))
        frame = cv2.resize(frame, dst_size)
        videos.append(frame)
        success, frame = cap.read()
    videos = np.array(videos)
    return videos


def read_landmarks(landmark_file):
    landmarks = dict()
    with open(landmark_file, 'r') as f:
        pid = 0
        for line in f.readlines():
            line = line.strip()
            if line.split(':')[0].isnumeric():
                frame_id = int(line.split(':')[0])
                landmarks[frame_id] = np.zeros((21, 3))  # a new frame
                assert pid == 0 or pid == 20, 'invalid frame!'
                pid = -1
                continue
            if 'landmark' in line: 
                pid += 1  # a new 3D point
                continue
            if 'x:' in line: 
                landmarks[frame_id][pid, 0] = 1.0 - float(line.split(': ')[1])
                continue
            if 'y:' in line:
                landmarks[frame_id][pid, 1] = float(line.split(': ')[1])
                continue
            if 'z:' in line:
                landmarks[frame_id][pid, 2] = float(line.split(': ')[1])
                continue
    return landmarks


def read_landmarks_text(landmark_file, flip=False):
    landmarks = dict()
    with open(landmark_file, 'r') as f:
        for line in f.readlines():
            line = line.strip()
            if line.split(':')[0].isnumeric():
                frame_id = int(line.split(':')[0])
                landmarks[frame_id] = ''
                pointer = 0
                continue
            if flip and 'x:' in line:
                xval = 1.0 - float(line.split(': ')[1])  # flip x coordinates horizontally
                line = '  x: {}'.format(xval)
            if pointer <= 105:  # 1 + 5 * 21 = 105 rows
                if ('x:' in line) or ('y:' in line) or ('}' in line):
                    line += ' '
                landmarks[frame_id] += line
                pointer += 1
    return landmarks


def generate_landmarks(start, end, origin=[0, 0, 0]):
    landmarks, landmarks_text = dict(), dict()
    # define a unit 3D cude, with pre-defined coordinate system
    # rightward: x, downward: y, innerward: z, cube center: origin
    cube = np.array([[1, -1, -1], [1, -1, 1], [-1, -1, 1], [-1, -1, -1],
                     [1, 1, -1], [1, 1, 1], [-1, 1, 1], [-1, 1, -1]], dtype=np.float32) * 0.5  # (8, 3)
    cube += origin
    step = int((end - start + 1) / 8)  # the number of frames for each of the 8 corners
    for s in range(8):
        t_start, t_end = start + s * step, start + (s+1) * step
        lm_points = np.repeat(np.expand_dims(cube[s], axis=0), 21, axis=0)  # (21, 3)
        lm_points[1:, 2] -= lm_points[0, 2]  # relative z
        for frame_id in range(t_start, t_end):
            landmarks[frame_id] = np.copy(lm_points)
            landmarks_text[frame_id] = ''.join(['landmark {' + 'x: {} y: {} z: {}'.format(pt[0], pt[1], pt[2]) + '} ' for pt in lm_points])
    return landmarks, landmarks_text



def read_clip_ranges(rng_file):
    ranges = dict()
    with open(rng_file, 'r') as f:
        for line in f.readlines():
            cid = int(line.strip().split('. ')[0])
            start = int(line.strip().split('. ')[1].split(', ')[0])
            end = int(line.strip().split('. ')[1].split(', ')[1])
            ranges[cid] = (start, end)
    return ranges


def write_video(mat, video_file, fps=5):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)


def draw_trajectory_3d(landmarks, start, end, pts3d_last, size=1000, fig_size=5, views=[30, -120], minmax=[(0, 0, -1), (1, 1, 1)], savefile='temp.mp4'):
    # find the range
    # pts_all = np.concatenate([landmarks[i] for i in range(start, end) if i in landmarks], axis=0)
    # pts_all = np.concatenate([pts_all, np.array([[0,0,0]])], axis=0)
    # xmin, ymin, zmin = np.min(pts_all, axis=0)
    # xmax, ymax, zmax = np.max(pts_all, axis=0)
    xmin, ymin, zmin = minmax[0]
    xmax, ymax, zmax = minmax[1]

    # init a 3D figure
    fig = plt.figure(figsize=(fig_size, fig_size))
    ax = fig.add_subplot(projection='3d')
    ax.view_init(views[0], views[1])  # elev: 30, azim: -120
    # Make a 3D quiver plot
    ox, oy, oz = np.zeros((3,3))
    dx, dy, dz = np.array([[xmax-xmin, 0, 0], [0, zmax-zmin, 0], [0, 0, ymax-ymin]]) * 0.2
    ax.quiver(ox, oy, oz, dx, dy, dz)
    # Setup coordinate system
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(zmin, zmax)
    ax.set_zlim(ymin, ymax)
    ax.invert_zaxis()
    ax.set_xlabel('X')  # x: rightward
    ax.set_ylabel('Z')  # z: innerward
    ax.set_zlabel('Y')  # y: downward
    plt.locator_params(axis='x', nbins=5)
    plt.locator_params(axis='y', nbins=5)
    plt.locator_params(axis='z', nbins=5)

    # init video writer
    traj3d_writer = FFMpegWriter(fps=5, metadata=dict(title='3D Trajectory', artist='Matplotlib'))
    # process
    with traj3d_writer.saving(fig, savefile, dpi=int(size / fig_size)):
        traj3d = []
        for i in range(start, end):
            if i in landmarks:
                pts3d = np.copy(landmarks[i])
                pts3d_last = np.copy(pts3d)
            else:
                pts3d = np.copy(pts3d_last)
            pts3d[1:, 2] += pts3d[0, 2]
            pt_hand = np.mean(pts3d, axis=0)
            traj3d.append(pt_hand)
            xs, ys, zs = np.array(traj3d)[:, 0], np.array(traj3d)[:, 1], np.array(traj3d)[:, 2]
            ax.plot(xs, zs, ys, '-o')
            ax.set_title('Hand trajectory [Frame: {}, View: ({}, {})]'.format(i-start, views[0], views[1]))
            # grab a frame
            traj3d_writer.grab_frame()  # dpi * (fig_size, fig_size)
    plt.close()


def draw_landmarks_2d(video, landmarks_text, start, end):
    video_anno = video[start: end]
    for i in range(start, end):
        if i not in landmarks_text:
            continue
        # draw landmarks
        landmark_list = text_format.Parse(landmarks_text[i], landmark_pb2.NormalizedLandmarkList())
        drawing_utils.draw_landmarks(
            video_anno[i - start],
            landmark_list,
            mp_hands.HAND_CONNECTIONS,
            mp_style.get_default_hand_landmarks_style(),
            mp_style.get_default_hand_connections_style())
    return video_anno


def main():
    # vis path
    vis_dir = os.path.join(root_path, './vis')
    os.makedirs(vis_dir, exist_ok=True)

    # read video
    print('Read video...')
    video_file = os.path.join(record_path, 'rgb_video.mp4')
    video_all = read_video(video_file, ratio=0.25)  # (N, 540, 960, 3)

    # read hand landmarks
    landmark_file = os.path.join(record_path, 'hand_frames', 'hand_landmarks.txt')
    landmarks = read_landmarks(landmark_file)
    pts3d_last = list(landmarks.items())[0][1]  # (21, 3)

    landmarks_text = read_landmarks_text(landmark_file, flip=True)

    # read clip ranges
    cliprng_file = os.path.join(record_path, 'hand_frames', 'clip_ranges.txt')
    clip_ranges = read_clip_ranges(cliprng_file)

    # visualize for each clip
    for cid, (start, end) in tqdm(clip_ranges.items(), desc='Visualizing clips', total=len(clip_ranges)):
        # draw trajectory and transform into video
        draw_trajectory_3d(landmarks, start, end, pts3d_last, size=video_all.shape[1], views=[10, -75], savefile='temp.mp4')
        video_traj3d = read_video('temp.mp4', ratio=1.0)  # (M, 540, 540, 3)

        # draw landmarks on RGB video
        # video_rgb = video_all[start: end]
        video_rgb = draw_landmarks_2d(video_all, landmarks_text, start, end)

        # visualize RGB and 3D trajectory side-by-side
        mat_vis = np.concatenate([video_rgb, video_traj3d], axis=2)  # (M, 540, 1500, 3)

        # write result video
        vis_file = os.path.join(vis_dir, 'vis_{}_{}_clip{}.mp4'.format(scene_name, record_id, cid))
        write_video(mat_vis, vis_file, fps=5)


def coord_check():

    # vis path
    vis_dir = os.path.join(root_path, './vis_check')
    os.makedirs(vis_dir, exist_ok=True)

    # read video
    print('Read video...')
    video_file = os.path.join(record_path, 'rgb_video.mp4')
    video_all = read_video(video_file, ratio=0.25)  # (N, 540, 960, 3)

    # read clip ranges
    cliprng_file = os.path.join(record_path, 'hand_frames', 'clip_ranges.txt')
    clip_ranges = read_clip_ranges(cliprng_file)
    cid = 1
    start, end = clip_ranges[cid]

    # generate fake landmarks
    landmarks, landmarks_text = generate_landmarks(start, end, origin=[0.5, 0.5, 0])
    pts3d_last = list(landmarks.items())[0][1]  # (21, 3)

    # draw trajectory and transform into video
    draw_trajectory_3d(landmarks, start, end, pts3d_last, size=video_all.shape[1], views=[10, -75], minmax=[(-1,-1,-1),(1,1,1)], savefile='temp.mp4')
    video_traj3d = read_video('temp.mp4', ratio=1.0)  # (M, 540, 540, 3)

    # # draw landmarks on RGB video
    # video_rgb = video_all[start: end]
    video_rgb = draw_landmarks_2d(video_all, landmarks_text, start, end)

    # visualize RGB and 3D trajectory side-by-side
    mat_vis = np.concatenate([video_rgb, video_traj3d], axis=2)  # (M, 540, 1500, 3)

    # write result video
    vis_file = os.path.join(vis_dir, 'vis_{}_{}_clip{}.mp4'.format(scene_name, record_id, cid))
    write_video(mat_vis, vis_file, fps=5)




if __name__ == '__main__':

    scene_id = '1'
    scene_name = 'bathroomCabinet'
    record_id = '2'

    root_path = os.path.dirname(__file__, "../../data/EgoPAT3D")
    record_path = os.path.join(root_path, 'complete', scene_id, '{}_{}'.format(scene_name, record_id))  # 'complete/1/bathroomCabinet_2'

    # # coordinate check
    # coord_check()

    main()

    if os.path.exists('temp.mp4'):
        os.remove('temp.mp4')
preprocess/EgoPAT3D/visualize_gt.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os


def get_split_info(split_file):
    sample_info = []
    with open(split_file, 'r') as f:
        for line in f.readlines():
            img_file = line.strip()
            subject, scene, obj = img_file.split('/')[:3]
            sub = subject.split('subject')[-1]
            info = 'sub{}_{}_{}'.format(sub, scene, obj)
            if info not in sample_info:
                sample_info.append(info)
    return sample_info


def write_split(info, filepath):
    with open(filepath, 'w') as f:
        for sample in info:
            f.writelines(sample + '\n')


if __name__ == '__main__':
    data_root = '../../data/H2O/ego'

    dest_dir = '../../data/H2O/Ego3DTraj/splits'
    os.makedirs(dest_dir, exist_ok=True)
    
    split_file = os.path.join(data_root, 'label_split', 'pose_train.txt')
    train_info = get_split_info(split_file)  # 114
    write_split(train_info, os.path.join(dest_dir, 'train.txt'))

    split_file = os.path.join(data_root, 'label_split', 'pose_val.txt')
    val_info = get_split_info(split_file)    # 24
    write_split(val_info, os.path.join(dest_dir, 'val.txt'))

    split_file = os.path.join(data_root, 'label_split', 'pose_test.txt')
    test_info = get_split_info(split_file)   # 46
    write_split(test_info, os.path.join(dest_dir, 'test.txt'))
    
    print("Train samples: {}, val samples: {}, test samples: {}".format(len(train_info), len(val_info), len(test_info)))

preprocess/H2O/get_splits.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import numpy as np
import cv2
import pickle
import sys
import copy



def frames_to_video(frames_path):
    filenames = os.listdir(frames_path)
    num_frames = len(filenames)
    video = []
    for i in range(num_frames):
        img_file = os.path.join(frames_path, '{:06d}.jpg'.format(i))
        assert os.path.exists(img_file), 'Frame {} does not exist!'.format(img_file)
        im = cv2.imread(img_file)
        video.append(im)  # H x W x 3
    video = np.array(video)
    return video


def write_video(mat, video_file, fps=30):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)


def read_depth_map(depth_file):
    assert os.path.exists(depth_file), "Depth file does not exist! {}".format(depth_file)
    depth_img = cv2.imread(depth_file, cv2.IMREAD_ANYDEPTH)  # (720, 1280), uint16, min=0, max=3168
    depth_img = depth_img.astype(np.float32) / 1000.0  # in meters
    return depth_img


def XYZ_to_uv(traj3d, intrinsics):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points in local coordinate system
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['cx'])
    traj2d[:, 1] = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['cy'])
    # clip the coordinates 
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['width']-1)
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['height']-1)
    traj2d = np.floor(traj2d).astype(np.int32)
    return traj2d

def world_to_uv(point, cam2world, intrinsics):
    # world to camera
    extrinsics = np.linalg.inv(cam2world)
    pt_homo = np.concatenate([point[:, np.newaxis], np.ones((1, 1))], axis=0)  # (4, 1)
    pt_local = extrinsics.dot(pt_homo)[:3, :]  # (3, 1)
    # project 3D coordinates onto pixel 2D coordinates
    pt_uv = XYZ_to_uv(pt_local.T, intrinsics)  # (1, 2)
    return pt_uv[0]


def uvd_to_world(uvd, cam2world, intrinsics):
    # uv_to_XYZ
    pt_local = np.zeros_like(uvd)
    pt_local[0] = (uvd[0] - intrinsics['cx']) * uvd[2] / intrinsics['fx']
    pt_local[1] = (uvd[1] - intrinsics['cy']) * uvd[2] / intrinsics['fy']
    pt_local[2] = uvd[2]
    # cam to world
    pt_homo = np.concatenate((pt_local[:, np.newaxis], np.array([[1]])), axis=0)  # (4,1)
    
    # E = np.linalg.inv(cam2world)
    # R = np.linalg.inv(E.T.dot(E)).dot(E.T)
    # pt_world = R.dot(pt_homo)

    # Note that this results is almost the same as the commented lines, 
    # which is an exact form of cam_to_world transformation.
    pt_world = cam2world.dot(pt_homo)
    return pt_world[:3, 0]

def replace_with_depth(point, depth_map, cam2world, intrinsics):
    """ point: 
    """ 
    uv = world_to_uv(point, cam2world, intrinsics)
    d = np.array([depth_map[uv[1], uv[0]]])
    uvd = np.concatenate([uv.astype(np.float32), d])
    pt_world = uvd_to_world(uvd, cam2world, intrinsics)
    return pt_world


def split_traj(traj, cam_poses, valids):
    """ traj: ndarray, (L, 3)
        cam_poses: (L, 4, 4)
        valids: list, (L,), binary
    """
    valids_extend = np.array([0] + valids + [0]).astype(int) # L+2
    starts = np.where(np.diff(valids_extend) == 1)[0]
    ends = np.where(np.diff(valids_extend) == -1)[0] - 1
    
    traj_data = []
    for s, e in zip(starts, ends):
        data = {
            'traj3d': traj[s: e+1],
            'cam2world': cam_poses[s: e+1],
            'start': s,
            'end': e
        }
        traj_data.append(data)
    return traj_data


def get_3d_traj(cam_dir, campose_folder, handpose_folder, depth_folder='depth', from_depth=False, intrinsics=None):
    # parse poses
    campose_dir = os.path.join(cam_dir, campose_folder)
    handpose_dir = os.path.join(cam_dir, handpose_folder)
    num_frames = len(os.listdir(handpose_dir))
    if from_depth:
        depth_dir = os.path.join(cam_dir, depth_folder)

    left_traj, right_traj, cam_poses = [], [], []
    left_valids, right_valids = [], []
    for i in range(num_frames):
        hpose_file = os.path.join(handpose_dir, '{:06d}.txt'.format(i))
        assert os.path.exists(hpose_file), 'Hand pose file {} does not exist!'.format(hpose_file)
        # parse hand pose file
        data = np.loadtxt(hpose_file)
        left_hasanno, left_hpose = data[0] == 1, data[1:64].reshape(-1, 3)
        right_hasanno, right_hpose = data[64] == 1, data[65:128].reshape(-1, 3)

        # compute the center of each hand by the 4 joints of fingers root (5, 9, 13, 7)
        left_ctr = np.mean(left_hpose[np.array([5, 9, 13, 17]), :], axis=0)
        right_ctr = np.mean(right_hpose[np.array([5, 9, 13, 17]), :], axis=0)

        # parse camera pose file
        cpose_file = os.path.join(campose_dir, '{:06d}.txt'.format(i))
        assert os.path.exists(cpose_file), 'Camera pose file {} does not exist!'.format(cpose_file)
        cam2world = np.loadtxt(cpose_file).reshape(4, 4)
        cam_poses.append(cam2world)

        if from_depth:
            # read depth data
            depth_map = read_depth_map(os.path.join(depth_dir, '{:06d}.png'.format(i)))
            # update the 3D coordinates by using depth from sensors
            left_ctr = replace_with_depth(left_ctr, depth_map, cam2world, intrinsics)
            right_ctr = replace_with_depth(right_ctr, depth_map, cam2world, intrinsics)

        # save
        left_traj.append(left_ctr)
        right_traj.append(right_ctr)
        left_valids.append(left_hasanno)
        right_valids.append(right_hasanno)
    
    # figure out which segment is a valid trajectory (hand exist)
    trajdata_left = split_traj(np.stack(left_traj, axis=0), np.stack(cam_poses, axis=0), left_valids)
    trajdata_right = split_traj(np.stack(right_traj, axis=0), np.stack(cam_poses, axis=0), right_valids)
    
    return trajdata_left, trajdata_right


def process(from_depth=False):
    for sub in subject_ids:
        subject_dir = os.path.join(data_root, 'subject{}_ego'.format(sub))
        if not os.path.isdir(subject_dir):
            continue

        for scene in sorted(os.listdir(subject_dir)):
            scene_dir = os.path.join(subject_dir, scene)
            if not os.path.isdir(scene_dir):
                continue 

            for obj in sorted(os.listdir(scene_dir)):
                obj_dir = os.path.join(scene_dir, obj)
                if not os.path.isdir(obj_dir):
                    continue

                cam_dir = os.path.join(obj_dir, cam)
                # dest name
                sample_name = 'sub{}_{}_{}'.format(sub, scene, obj)
                print("process sample: {}".format(sample_name))

                # read video frames
                vid_file = os.path.join(video_dir, sample_name + '.mp4')
                if not os.path.exists(vid_file):
                    # RGB video frames
                    video = frames_to_video(os.path.join(cam_dir, 'rgb256'))
                    # save as mp4 video file
                    write_video(video, vid_file)
                
                traj_file = os.path.join(traj_dir, sample_name + '.pkl')
                if not os.path.exists(traj_file):
                    # read camera intrinsics
                    data = np.loadtxt(os.path.join(cam_dir, 'cam_intrinsics.txt'))
                    intrinsics = {'fx': data[0], 'fy': data[1], 'cx': data[2], 
                                  'cy': data[3], 'width': data[4], 'height': data[5]}

                    # get 2D and 3D hand trajectory
                    trajdata_left, trajdata_right = get_3d_traj(cam_dir, 'cam_pose', 'hand_pose', 
                                                                depth_folder='depth', from_depth=from_depth, intrinsics=intrinsics)

                    # save results
                    with open(traj_file, 'wb') as f:
                        pickle.dump({'left_hand': trajdata_left, 
                                     'right_hand': trajdata_right,
                                     'intrinsics': intrinsics}, f, protocol=pickle.HIGHEST_PROTOCOL)


def read_traj(traj_file):
    with open(traj_file, 'rb') as f:
        traj_data = pickle.load(f)
        trajdata_left = traj_data['left_hand']
        trajdata_right = traj_data['right_hand']
        intrinsics = traj_data['intrinsics']
    return trajdata_left, trajdata_right, intrinsics


def get_stats(trajdata_ref, trajdata_target):
    # compute ADE in 3D world space
    displace_errors = np.sqrt(np.sum((trajdata_ref - trajdata_target)**2, axis=-1))
    de_mean = np.mean(displace_errors)
    de_max = np.max(displace_errors)
    # compute absolute depth error
    depth_errors = np.fabs(trajdata_ref[:, 2] - trajdata_target[:, 2])
    dz_mean = np.mean(depth_errors)
    dz_max = np.max(depth_errors)
    # ratio
    ratio = len(np.where(displace_errors > 0.2)[0]) / len(displace_errors)

    print("DE(mean): {:.3f} mm, DE(max): {:.3f} mm, DZ(mean): {:.3f} mm, DZ(max): {:.3f} mm, P(DE > 0.20) = {:.2f}%".format(
        de_mean*1000, de_max*1000, dz_mean*1000, dz_max*1000, ratio * 100))
    
    all_stats = {
        'displace_errors': displace_errors,
        'ratio': ratio,
        'de_mean': de_mean,
        'dz_mean': dz_mean
    }
    return all_stats


def plot_frequency(errors_before, errors_after, ratio_before, ratio_after):
    import matplotlib.pyplot as plt
    import seaborn as sns
    fig, ax = plt.subplots(figsize=(8, 4))
    fontsize = 18
    sns.kdeplot(errors_before, color='red', fill=True, label='raw depth, P(DE>0.2)={:.2f}%'.format(ratio_before * 100))
    sns.kdeplot(errors_after, color='blue', fill=True, label='depth repair, P(DE>0.2)={:.2f}%'.format(ratio_after * 100))
    # bins = np.linspace(0, 1, 1000)
    # ax.hist(errors_before, bins, lw=1, ec="red", fc="red", alpha=0.5, label='before repair')
    # ax.hist(errors_after, bins, lw=1, ec="blue", fc="blue", alpha=0.5, label='after repair')
    plt.xlim(0,0.8)
    plt.xticks(fontsize=fontsize)
    plt.yticks(fontsize=fontsize)
    plt.xlabel('Displacement Error (m)', fontsize=fontsize)
    plt.ylabel('Frequency', fontsize=fontsize)
    plt.legend(loc='upper right', fontsize=fontsize)
    plt.tight_layout()
    plt.savefig(os.path.join(dest_dir, 'compare_hist.png'))
    plt.savefig(os.path.join(dest_dir, 'compare_hist.pdf'))


def plot_stat_bars(stats_before, stats_after, fontsize=18):
    import matplotlib.pyplot as plt
    # plt.rcParams['font.family'] = 'DeJavu Serif'
    # plt.rcParams['font.serif'] = ['Times New Roman']
    measures = ('mDE', 'mDZ')
    values = {
        'H2O-DT w/o repair': (stats_before['de_mean'] * 100, stats_before['dz_mean'] * 100),
        'H2O-DT': (stats_after['de_mean'] * 100, stats_after['dz_mean'] * 100)
    }
    x = np.arange(len(measures))  # the label locations
    width = 0.3  # the width of the bars
    multiplier = 0.5
    fig, ax = plt.subplots(figsize=(5,4))

    for k, v in values.items():
        offset = width * multiplier
        rects = ax.bar(x + offset, v, width, label=k)
        ax.bar_label(rects, padding=4, fontsize=fontsize, fmt='%.1f')
        multiplier += 1
    
    plt.yticks(np.arange(0, 25, 4), fontsize=fontsize)
    ax.set_ylabel('distance error (cm)', fontsize=fontsize)
    ax.set_ylim(0, 25)
    ax.set_xticks(x + width, measures, fontsize=fontsize)
    ax.set_xlim(-0.45, 2)
    ax.legend(loc='upper right', fontsize=fontsize-3)
    plt.tight_layout()
    plt.savefig(os.path.join(dest_dir, 'compare_bars.png'))
    plt.savefig(os.path.join(dest_dir, 'compare_bars.pdf'))


def compare_trajectories():
    """ Compare the trajecotries obtained by methods with & without depth images
    """
    trajdata_all, trajdata_all_withdepth, trajdata_all_repair = [], [], []
    for sample in sorted(os.listdir(traj_dir)):
        # read trajectory data from 3D pose
        trajdata_left, trajdata_right, intrinsics = read_traj(os.path.join(traj_dir, sample))

        trajdata = np.concatenate([np.concatenate([traj['traj3d'] for traj in trajdata_left], axis=0),
                                   np.concatenate([traj['traj3d'] for traj in trajdata_right], axis=0)], axis=0)
        trajdata_all.append(trajdata)

        # read trajectory data from 3D pose + repair
        trajdata_left, trajdata_right, intrinsics = read_traj(os.path.join(traj_dir + '_depth', sample))
        trajdata_withdepth = np.concatenate([np.concatenate([traj['traj3d'] for traj in trajdata_left], axis=0),
                                             np.concatenate([traj['traj3d'] for traj in trajdata_right], axis=0)], axis=0)
        trajdata_all_withdepth.append(trajdata_withdepth)

        # read trajectory data from 3D pose + repair
        trajdata_left, trajdata_right, intrinsics = read_traj(os.path.join(traj_dir + '_repair', sample))
        trajdata_repair = np.concatenate([np.concatenate([traj['traj3d'] for traj in trajdata_left], axis=0),
                                          np.concatenate([traj['traj3d'] for traj in trajdata_right], axis=0)], axis=0)
        trajdata_all_repair.append(trajdata_repair)
    
    trajdata_all = np.concatenate(trajdata_all, axis=0)  # 682659
    trajdata_all_withdepth = np.concatenate(trajdata_all_withdepth, axis=0)
    trajdata_all_repair = np.concatenate(trajdata_all_repair, axis=0)

    print("==> Before repair:")
    stats_before = get_stats(trajdata_all, trajdata_all_withdepth)
    # DE(mean): 170.892 mm, DE(max): 1343.636 mm, DZ(mean): 160.702 mm, DZ(max): 1173.206 mm (before repair)

    print("==> After repair:")
    stats_after = get_stats(trajdata_all, trajdata_all_repair)
    # DE(mean): 133.106 mm, DE(max): 888.224 mm, DZ(mean): 115.336 mm, DZ(max): 700.663 mm. (after repair)
    
    plot_frequency(stats_before['displace_errors'], stats_after['displace_errors'], stats_before['ratio'], stats_after['ratio'])

    plot_stat_bars(stats_before, stats_after)



def repair_by_lsq(traj_raw, traj_ref, test_data, num_frames, threshold):
    from scipy.optimize import leastsq
    def model_func(p, x):
        f = np.poly1d(p[:-2])
        y = f(x) + p[-2] * np.sin(p[-1]*x)
        return y

    def error_func(p, x, y):
        error = y - model_func(p, x)
        return error
    
    fit_data = [i for i in range(num_frames) if i not in test_data]
    fit_labels = [traj_raw[i, 2] for i in fit_data]
    
    # least square fitting by third-order multinomial + sin model
    param_init = [0, 1, 1, 0, 1, 0.05]
    ret = leastsq(error_func, param_init, 
                    args=(np.array(fit_data), np.array(fit_labels)), epsfcn=0.0001)  # start with straight line: y=0

    traj_repaired = np.copy(traj_raw)
    for i in test_data:
        depth_pred = model_func(ret[0], i)  # repair by LSQ modelz
        if np.fabs(depth_pred - traj_ref[i, 2]) < 2 * threshold:
            traj_repaired[i, 2] = depth_pred
    
    return traj_repaired


def repair_hand_traj(trajdata_hand_raw, trajdata_hand_ref, threshold, MIN_VALID_POINTS):

    repaired = False
    trajdata_repaired = copy.deepcopy(trajdata_hand_raw)
    for i, (trajdata_raw, trajdata_ref) in enumerate(zip(trajdata_hand_raw, trajdata_hand_ref)):
        traj_raw = np.clip(trajdata_raw['traj3d'], 0.1, 1.0) if CLIP_VAL else trajdata_raw['traj3d']
        traj_ref = trajdata_ref['traj3d']
        depth_errors = np.fabs(traj_raw[:, 2] - traj_ref[:, 2])
        num_frames = len(traj_raw)

        # find invalid trajectory points
        invalids = np.where(depth_errors > threshold)[0].tolist() # input frame indices whose depths need to be repaired
        # repair
        if len(invalids) > 0 and len(invalids) <= num_frames - MIN_VALID_POINTS:
            traj_repaired = repair_by_lsq(traj_raw, traj_ref, invalids, num_frames, threshold)
            trajdata_repaired[i]['traj3d'] = traj_repaired
            repaired = True
        
        if CLIP_VAL:
            # clamp the depth into [0.1, 1]
            zs = trajdata_repaired[i]['traj3d'][:, 2]
            if (zs < 0.1).any() or (zs > 1.0).any():
                trajdata_repaired[i]['traj3d'][:, 2] = np.clip(zs, 0.1, 1.0)
    
    return trajdata_repaired, repaired



def repair_depth(threshold=0.20):
    """ Repair the depth values for those trajectories that have large depth error
    """
    for sample in sorted(os.listdir(traj_dir)):
        # read trajectory data from 3D pose
        trajdata_left_ref, trajdata_right_ref, _ = read_traj(os.path.join(traj_dir, sample))
        # read trajectory data from 3D pose + depth
        trajdata_left_raw, trajdata_right_raw, intrinsics = read_traj(os.path.join(traj_dir + '_depth', sample))

        # repairing
        trajdata_left_repaired, repaired1 = repair_hand_traj(trajdata_left_raw, trajdata_left_ref, threshold, MIN_VALID_POINTS)
        trajdata_right_repaired, repaired2 = repair_hand_traj(trajdata_right_raw, trajdata_right_ref, threshold, MIN_VALID_POINTS)
        
        print("repair sample: {}({})".format(sample[:-4], 'repaired' if repaired1 or repaired2 else ''))

        # save
        with open(os.path.join(traj_repair_dir, sample), 'wb') as f:
            pickle.dump({'left_hand': trajdata_left_repaired, 
                         'right_hand': trajdata_right_repaired,
                         'intrinsics': intrinsics}, f, protocol=pickle.HIGHEST_PROTOCOL)



if __name__ == '__main__':
    data_root = '../../data/H2O/ego'
    cam = 'cam4'  # egocentric camera
    subject_ids = ['1', '2', '3', '4']

    from_depth = True if len(sys.argv) > 1 and sys.argv[1] == 'from_depth' else False

    # destination dirs
    dest_dir = '../../data/H2O/Ego3DTraj'
    os.makedirs(dest_dir, exist_ok=True)

    video_dir = os.path.join(dest_dir, 'video')
    os.makedirs(video_dir, exist_ok=True)

    traj_dir = os.path.join(dest_dir, 'traj')
    if from_depth: traj_dir += '_depth'
    os.makedirs(traj_dir, exist_ok=True)

    if len(sys.argv) > 1 and sys.argv[1] == 'compare': 
        compare_trajectories()
    
    elif len(sys.argv) > 1 and sys.argv[1] == 'repair':
        traj_repair_dir = os.path.join(dest_dir, 'traj_repair')
        os.makedirs(traj_repair_dir, exist_ok=True)
        # MIN_VALID_POINTS=10
        MIN_VALID_POINTS=6
        CLIP_VAL=False
        repair_depth(threshold=0.20)
    
    else:
        process(from_depth=from_depth)

preprocess/H2O/preprocess.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import pickle
import cv2
import numpy as np
import os
import imageio
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter
from mpl_toolkits.mplot3d import Axes3D # <--- This is important for 3d plotting 

TRAJ_COLOR = (255, 0, 0)    # red
PAST_COLOR = (255, 255, 0)  # yellow
FUTURE_COLOR = (0, 0, 255)  # blue



def read_traj(traj_file):
    with open(traj_file, 'rb') as f:
        traj_data = pickle.load(f)
        trajdata_left = traj_data['left_hand']
        trajdata_right = traj_data['right_hand']
        intrinsics = traj_data['intrinsics']
    return trajdata_left, trajdata_right, intrinsics


def read_video(video_file, toRGB=True):
    assert os.path.exists(video_file), "File does not exist! {}".format(video_file)
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    video = []
    while success:
        if toRGB:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        video.append(frame)
        success, frame = cap.read()
    video = np.array(video)
    return video


def XYZ_to_uv(traj3d, intrinsics):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points in local coordinate system
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['cx'])
    traj2d[:, 1] = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['cy'])
    # clip the coordinates 
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['width'])
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['height'])
    traj2d = np.floor(traj2d).astype(np.int32)
    return traj2d


def resize_coords(traj2d, width, height, intrinsics):
    traj2d_new = np.copy(traj2d)
    traj2d_new[:, 0] = np.floor(traj2d[:, 0] / intrinsics['width'] * width)
    traj2d_new[:, 1] = np.floor(traj2d[:, 1] / intrinsics['height'] * height)
    return traj2d_new


def draw_traj2d_frame(frame, traj2d, traj_start=0, future_start=15, radius=5):
    cv2.polylines(frame, [traj2d], False, TRAJ_COLOR, thickness=2)
    for t, uv in enumerate(traj2d):
        color = PAST_COLOR if traj_start + t <= future_start else FUTURE_COLOR
        cv2.circle(frame, tuple(uv), radius=radius, color=color, thickness=-1)


def draw_traj3d_frame(frame, traj3d, temp_file='temp.mp4', compose='concat'):
    # draw 3D plot
    height, width = frame.shape[:2]
    scale = 0.8 if compose == 'overlap' else 1.0
    canvas_size = min(height, width) * scale
    traj3d_frame = draw_trajectory_3d(traj3d, (canvas_size, canvas_size), dpi=90, views=[30, -120], savefile=temp_file)
    if compose == 'overlap':
        # overlap the traj3d frame on the the original frame
        bottom_left = frame[height - traj3d_frame.shape[0]:, :traj3d_frame.shape[1], :]
        frame[height - traj3d_frame.shape[0]:, :traj3d_frame.shape[1], :] = cv2.addWeighted(bottom_left, 0.4, traj3d_frame, 0.6, 0)  # bottom left
    else:
        # concatenate horizontally
        traj3d_frame = cv2.resize(traj3d_frame, (int(height / traj3d_frame.shape[0] * traj3d_frame.shape[1]), height))
        frame = cv2.hconcat([frame, traj3d_frame])
    return frame


def draw_trajectory_3d(traj3d, video_size, dpi=90, views=[30, -120], minmax=None, savefile='temp.mp4'):

    if minmax is not None:
        xmin, ymin, zmin = minmax[0]
        xmax, ymax, zmax = minmax[1]
    else:
        xmin, xmax = np.min(traj3d[:, 0]) - 0.1, np.max(traj3d[:, 0]) + 0.1
        ymin, ymax = np.min(traj3d[:, 1]) - 0.1, np.max(traj3d[:, 1]) + 0.1
        zmin, zmax = np.min(traj3d[:, 2]) - 0.1, np.max(traj3d[:, 2]) + 0.1

    fig_w, fig_h = int(video_size[1] // dpi), int(video_size[0] // dpi)

    # init a 3D figure
    fig = plt.figure(figsize=(fig_w, fig_h))
    ax = fig.add_subplot(projection='3d')
    ax.view_init(views[0], views[1])  # elev: 30, azim: -120
    # Make a 3D quiver plot
    ox, oy, oz = np.zeros((3,3))
    dx, dy, dz = np.array([[xmax-xmin, 0, 0], [0, zmax-zmin, 0], [0, 0, ymax-ymin]]) * 0.2
    ax.quiver(ox, oy, oz, dx, dy, dz)
    # Setup coordinate system
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(zmin, zmax)
    ax.set_zlim(ymin, ymax)
    ax.invert_zaxis()
    ax.set_xlabel('X')  # x: rightward
    ax.set_ylabel('Z')  # z: innerward
    ax.set_zlabel('Y')  # y: downward
    plt.locator_params(axis='x', nbins=5)
    plt.locator_params(axis='y', nbins=5)
    plt.locator_params(axis='z', nbins=5)
    plt.tight_layout()

    traj3d_writer = FFMpegWriter(fps=30, metadata=dict(title='3D Trajectory', artist='Matplotlib'))
    with traj3d_writer.saving(fig, savefile, dpi=dpi):
        xs, ys, zs = np.array(traj3d)[:, 0], np.array(traj3d)[:, 1], np.array(traj3d)[:, 2]
        ax.plot(xs, zs, ys, '-o', color='red', markersize=2)
        # grab a single frame
        traj3d_writer.grab_frame()
    plt.close()

    frame_traj3d = read_video(savefile, toRGB=True)
    # remoe temporary file
    if os.path.exists(savefile):
        os.remove(savefile)
    
    return frame_traj3d[0]



def video_to_gif(video, giffile, fps=5.0, toBGR=False):
    assert giffile.endswith('.gif')
    with imageio.get_writer(giffile, mode='I', duration=1.0/fps) as writer:
        for frame in video:
            frame_vis = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) if toBGR else np.copy(frame)
            writer.append_data(frame_vis)


def visualize(trajdata, video, intrinsics, n_future=30, obs_ratio=1.0, radius=5, with_3d=False, fps=30):
    """ Visualize trajectory
    """
    height, width = video.shape[1:3]
    temp_file = os.path.join(result_dir, 'temp.mp4')
    result = []
    # visualize 2D trajectory
    for traj_data in trajdata:
        start, end = traj_data['start'], traj_data['end']
        clip = np.copy(video[start: end+1])
        traj3d = traj_data['traj3d']  # suppose this is in world coordinate system
        cam2world = traj_data['cam2world']
        future_start = int(len(traj3d) * obs_ratio)
        horizon = min(n_future, len(traj3d))

        video_vis = []
        for i, (frame, Rc2w) in enumerate(zip(clip, cam2world)):
            # get the global trajectory
            horizon_end = min(i+horizon, len(traj3d))
            traj3d_global = traj3d[i:horizon_end]

            # world to cam
            extrinsics = np.linalg.inv(Rc2w)
            global_homo = np.concatenate([traj3d_global, np.ones((horizon_end-i, 1))], axis=1)  # (30, 4)
            traj3d_local = extrinsics.dot(global_homo.T)[:3, :].T  # (30, 3)

            # project 3D coordinates onto pixel 2D coordinates
            traj2d_vis = XYZ_to_uv(traj3d_local, intrinsics)
            
            # scaling due to video resize
            traj2d_vis = resize_coords(traj2d_vis, width, height, intrinsics)

            # visualize 2D trajectory on the RGB frame
            draw_traj2d_frame(frame, traj2d_vis, traj_start=i, future_start=future_start, radius=radius)

            if with_3d:
                # visualize 3D trajectory
                frame = draw_traj3d_frame(frame, traj3d_global, temp_file='temp.mp4')

            video_vis.append(frame)
        result.append(np.array(video_vis))
    
    return result
    

def main(traj_file, vid_file, result_dir, longest=False):

    # read trajectory data
    trajdata_left, trajdata_right, intrinsics = read_traj(traj_file)
    
    # read video
    video = read_video(vid_file, toRGB=True)
    fps = 30

    # visualize left & right hand trajectories
    vis_left = visualize(trajdata_left, video, intrinsics, n_future=30, radius=3, with_3d=True, fps=fps)
    vis_right = visualize(trajdata_right, video, intrinsics, n_future=30, radius=3, with_3d=True, fps=fps)

    if longest:
        lens = [len(vis) for vis in vis_left]
        vis_left = vis_left[np.argmax(lens)]  # the longest left hand
        lens = [len(vis) for vis in vis_right]
        vis_right = vis_right[np.argmax(lens)]  # the longest left hand
        vis = vis_left if len(vis_left) > len(vis_right) else vis_right
        # save to GIF file
        video_to_gif(vis, os.path.join(result_dir, '{}.gif'.format(sample)), fps=fps)
    else:
        # save to GIF file
        for n, vis in enumerate(vis_left):
            video_to_gif(vis, os.path.join(result_dir, '{}_left{}.gif'.format(sample, n+1)), fps=fps)
        for n, vis in enumerate(vis_right):
            video_to_gif(vis, os.path.join(result_dir, '{}_right{}.gif'.format(sample, n+1)), fps=fps)



if __name__ == '__main__':

    data_root = '../../data/H2O/Ego3DTraj'
    samples = ['sub1_h1_4', 'sub1_h2_4', 'sub1_k1_4', 'sub1_k2_4', 'sub1_o1_4', 'sub1_o2_4',
               'sub2_h1_4', 'sub2_h2_4', 'sub2_k1_4', 'sub2_k2_4', 'sub2_o1_4', 'sub2_o2_4',
               'sub3_h1_4', 'sub3_h2_4', 'sub3_k1_4', 'sub3_k2_4', 'sub3_o1_4', 'sub3_o2_4',
               'sub4_h1_4', 'sub4_h2_4', 'sub4_k1_4', 'sub4_k2_4', 'sub4_o1_4', 'sub4_o2_4']

    for sample in samples:
        
        print("Visualizing sample: {}".format(sample))
        traj_file = os.path.join(data_root, 'traj/{}.pkl'.format(sample))
        vid_file = os.path.join(data_root, 'video/{}.mp4'.format(sample))

        result_dir = os.path.join(data_root, 'vis')
        os.makedirs(result_dir, exist_ok=True)

        main(traj_file, vid_file, result_dir, longest=True)


preprocess/H2O/visualize.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import torch
import os
import pickle
import cv2
import numpy as np
from tqdm import tqdm
from functools import reduce
from src.utils import denormalize, global_to_local, normalize, vis_frame_traj, write_video,video_to_gif, read_video



class EgoPAT3D(Dataset):
    def __init__(self, root_dir, phase='train', transform=None, data_cfg=None, model_cfg=None):
        self.root_dir = root_dir
        self.phase = phase
        self.transform = transform
        self._MAX_FRAMES = getattr(data_cfg, 'max_frames', -1) # the maximum number of input frames
        self._MAX_DEPTH = getattr(data_cfg, 'max_depth', 3)  # the max depth of hand trajectory
        self.target = getattr(model_cfg, 'target', '3d')
        self.scenes = getattr(data_cfg, 'scenes', None)
        self.modalities = getattr(model_cfg, 'modalities', ['rgb', 'loc'])
        self.tinyset = getattr(data_cfg, 'tinyset', False)
        self.load_all = getattr(data_cfg, 'load_all', False)
        self.use_odom = getattr(model_cfg, 'use_odom', False)
        if self.use_odom: assert self.target == '3d', "Do not support 2D target when using odometry!"
        self.centralize = getattr(model_cfg, 'centralize', False)
        # Our own dataset split
        self.scene_splits = {'train': {'1': ['1', '2', '3', '4', '5', '6', '7'], 
                                       '2': ['1', '2', '3', '4', '5', '6', '7'], 
                                       '3': ['1', '2', '3', '4', '5', '6'],
                                       '4': ['1', '2', '3', '4', '5', '6', '7'],
                                       '5': ['1', '2', '3', '4', '5', '6'], 
                                       '6': ['1', '2', '3', '4', '5', '6'], 
                                       '7': ['1', '2', '3', '4', '5', '6', '7'],
                                       '9': ['1', '2', '3', '4', '5', '6', '7'],
                                       '10': ['1', '2', '3', '4', '5', '6', '7'],
                                       '11': ['1', '2', '3', '4', '5', '6', '7'],
                                       '12': ['1', '2', '3', '4', '5', '6', '7']
                                        },
                             'val': {'1': ['8'], 
                                     '2': ['8'], '3': ['7'], '4': ['8'], '5': ['7'], '6': ['7'], 
                                     '7': ['8'], '9': ['8'], '10': ['8'], '11': ['8'], '12': ['8']
                                     },
                             'test': {'1': ['9', '10'], 
                                      '2': ['9', '10'], '3': ['9', '10'], '4': ['9', '10'], '5': ['8', '9'], '6': ['8', '9'], 
                                      '7': ['9', '10'], '9': ['9', '10'], '10': ['9', '10'], '11': ['9', '10'], '12': ['9', '10']
                                     },
                             'test_novel': {'13': ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10'], 
                                            '14': ['2', '3', '4', '5', '6', '7', '8', '9', '10'], 
                                            '15': ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10']}}
        assert self.phase in self.scene_splits, "Invalid dataset split: {}".format(self.phase)
        if self.scenes is not None:
            data = dict(filter(lambda elem: elem[0] in self.scenes, self.scene_splits[self.phase].items()))
            self.scene_splits[self.phase] = data
        # camera intrinsics
        self.intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                           'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                           'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
        self.vis_ratio = 0.25  # the ratio of rgb video size to the 4K size
        
        # the data path needed
        self.rgb_dir, self.traj_dir, self.odom_dir = self._init_data_path()
        if self.tinyset:
            selects = self._read_selected_list()  # 100 selected samples
            samples = selects[:64] if self.phase == 'train' else selects[64:]
            self.rgb_paths, self.traj_data, self.odom_data, self.preserves = self._read_selected_data(samples)
        else:    
            # read dataset lists
            self.rgb_paths, self.traj_data, self.odom_data, self.preserves = self._read_list()

        if self.load_all and 'rgb' in self.modalities:
            self.rgb_data = self._read_all_videos()
        
        
    def _read_traj_file(self, filepath):
        assert os.path.exists(filepath), "File does not exist! {}".format(filepath)
        with open(filepath, 'rb') as f:
            data = pickle.load(f)
        traj2d = data['traj2d']
        traj3d = data['traj3d']
        num_preserve = data['num_preserve'] if 'num_preserve' in data else len(traj2d)
        return traj2d, traj3d, num_preserve


    def _read_odom_file(self, filepath):
        assert os.path.exists(filepath), "File does not exist! {}".format(filepath)
        all_transforms = np.load(filepath)  # (T, 4, 4)
        return all_transforms

    
    def _init_data_path(self):
        # rgb video root
        rgb_dir = os.path.join(self.root_dir, 'EgoPAT3D-postproc', 'video_clips_hand')
        assert os.path.exists(rgb_dir), 'Path does not exist! {}'.format(rgb_dir)
        # trajectory root
        traj_dir = os.path.join(self.root_dir, 'EgoPAT3D-postproc', 'trajectory_repair')
        assert os.path.exists(traj_dir), 'Path does not exist! {}'.format(traj_dir)
        # visual odometry path (used only if self.use_odom=True)
        odom_dir = os.path.join(self.root_dir, 'EgoPAT3D-postproc', 'odometry')
        assert os.path.exists(odom_dir), 'Path does not exist! {}'.format(traj_dir)
        return rgb_dir, traj_dir, odom_dir


    def _read_selected_list(self):
        select_listfile = os.path.join(self.root_dir, 'EgoPAT3D-postproc', 'selected.txt')
        assert os.path.exists(select_listfile), 'Path does not exist! {}'.format(select_listfile)
        selects = []
        with open(select_listfile, 'r') as f:
            for line in f.readlines():
                selects.append(line.strip())
        return selects
    
    
    def _read_list(self):
        rgb_paths, traj_data, odom_data, preserves = [], [], [], []
        # read the paths
        for scene_id, record_splits in tqdm(self.scene_splits[self.phase].items(), ncols=0, desc='Read trajectories'):
            scene_rgb_path = os.path.join(self.rgb_dir, scene_id)
            scene_traj_path = os.path.join(self.traj_dir, scene_id)
            scene_odom_path = os.path.join(self.odom_dir, scene_id)
            
            record_names = list(filter(lambda x: x.split('_')[-1] in record_splits, os.listdir(scene_traj_path)))  # get the splits of records
            for record in record_names:
                record_rgb_path = os.path.join(scene_rgb_path, record)
                record_traj_path = os.path.join(scene_traj_path, record)
                record_odom_path = os.path.join(scene_odom_path, record)
                
                traj_files = list(filter(lambda x: x.endswith('.pkl'), os.listdir(record_traj_path)))  # a list of '*.pkl'
                for traj_name in traj_files:
                    # read trajectory
                    traj2d, traj3d, num_preserve = self._read_traj_file(os.path.join(record_traj_path, traj_name))
                    traj_data.append({'traj2d': traj2d, 'traj3d': traj3d})
                    # save rgb video path
                    rgb_paths.append(os.path.join(record_rgb_path, traj_name[:-4] + '.mp4'))
                    # read odometry data
                    if self.use_odom:
                        odom = self._read_odom_file(os.path.join(record_odom_path, traj_name[:-4] + '.npy'))  # (T, 4, 4)
                    else:
                        odom = np.eye(4, dtype=np.float32)[None, :, :].repeat(num_preserve, axis=0)  # identity matrix
                    odom_data.append(odom)
                    preserves.append(num_preserve)

        return rgb_paths, traj_data, odom_data, preserves


    def _read_selected_data(self, sample_list):
        rgb_paths, traj_data, odom_data, preserves = [], [], [], []
        # read the paths
        for sample in sample_list:
            scene_id, record, clip_name = sample.split('/')
            # read trajectory
            traj2d, traj3d, num_preserve = self._read_traj_file(os.path.join(self.traj_dir, scene_id, record, clip_name + '.pkl'))
            traj_data.append({'traj2d': traj2d, 'traj3d': traj3d})
            # save rgb video path
            rgb_paths.append(os.path.join(self.rgb_dir, scene_id, record, clip_name + '.mp4'))
            # read odometry data
            if self.use_odom:
                odom = self._read_odom_file(os.path.join(self.odom_dir, scene_id, record, clip_name + '.npy'))  # (T, 4, 4)
            else:
                odom = np.eye(4, dtype=np.float32)[None, :, :].repeat(num_preserve, axis=0)  # identity matrix
            odom_data.append(odom)
            preserves.append(num_preserve)

        return rgb_paths, traj_data, odom_data, preserves
    
    
    def _read_all_videos(self):
        """ Read all video data at once into CPU memory
        """
        video_data = []
        for filename, preserve in zip(self.rgb_paths, self.preserves):
            rgb = self._read_rgb(filename)
            rgb = rgb[:preserve]
            if self.transform:
                rgb = self.transform(rgb)  # (3, T, H=540, W=960)
            else:
                rgb = torch.from_numpy(np.transpose(rgb, (3, 0, 1, 2)))  # (3, T, H=540, W=960)
            video_data.append(rgb)
        return video_data


    def _read_rgb(self, video_file):
        cap = cv2.VideoCapture(video_file)
        success, frame = cap.read()
        videos = []
        while success:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            videos.append(frame)
            success, frame = cap.read()
        cap.release()
        cv2.destroyAllWindows()
        videos = np.array(videos)
        return videos
    

    def _XYZ_to_uv(self, traj3d):
        width = self.intrinsics['w'] * self.vis_ratio
        height = self.intrinsics['h'] * self.vis_ratio
        # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
        u = (traj3d[:, 0] * self.intrinsics['fx'] / traj3d[:, 2] + self.intrinsics['ox']) * self.vis_ratio
        v = (traj3d[:, 1] * self.intrinsics['fy'] / traj3d[:, 2] + self.intrinsics['oy']) * self.vis_ratio
        u = torch.clamp(u, min=0, max=width-1)
        v = torch.clamp(v, min=0, max=height-1)
        traj2d = torch.stack((u, v), dim=-1)
        return traj2d
    
    
    def _normalize(self, traj):
        traj_new = traj.clone()
        width = self.intrinsics['w'] * self.vis_ratio
        height = self.intrinsics['h'] * self.vis_ratio
        if self.target == '2d':
            # normalize the (u,v) into the range 0-1
            traj_new[:, 0] /= width
            traj_new[:, 1] /= height
        elif self.target == '3d':
            # transform the (X, Y, Z) into the normalized (u,v,z) by PINHOLE camera model
            traj2d = self._XYZ_to_uv(traj)
            traj_new[:, 0] = traj2d[:, 0] / width
            traj_new[:, 1] = traj2d[:, 1] / height
            traj_new[:, 2] = traj[:, 2] / self._MAX_DEPTH
        if self.centralize:
            traj_new -= 0.5
        return traj_new

    
    def _get_projected_traj3d(self, traj3d, odometry):
        """ traj3d: (T, 3)
            odometry: (T, 4, 4)
        """
        length = len(odometry)
        traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
        all_traj3d_proj = []
        for i in range(length):
            # compute the transformed points of all future points
            traj3d_proj = [traj3d[i]]  # initial 3D point
            for j in range(i + 1, length):
                odom = reduce(np.dot, odometry[(i+1):(j+1)])  # (4, 4)
                future_point = odom.dot(traj3d_homo[j].T)  # (4,)
                traj3d_proj.append(future_point[:3])
            all_traj3d_proj.append(np.array(traj3d_proj))
        all_traj3d_proj = np.concatenate(all_traj3d_proj, axis=0)  # (T*(T+1)/2, 3)
        return all_traj3d_proj
    
    
    def __len__(self):
        return len(self.rgb_paths)
    
    
    def __getitem__(self, index):
        
        traj2d = torch.from_numpy(self.traj_data[index]['traj2d']).to(torch.float32)  # (T, 2)
        traj3d = torch.from_numpy(self.traj_data[index]['traj3d']).to(torch.float32)  # (T, 3)
        traj = traj2d if self.target == '2d' else traj3d
        num_preserve = self.preserves[index]
        if 'rgb' in self.modalities:
            if self.load_all:
                rgb = self.rgb_data[index]  # (3, T, H=540, W=960)
            else:
                # read mp4 file
                rgb = self._read_rgb(self.rgb_paths[index])  # (T, H=540, W=960, 3)
                rgb = rgb[:num_preserve]
                # ------------  process video input  ------------
                # data transformation
                if self.transform:
                    rgb = self.transform(rgb)  # (3, T, H=540, W=960)
                else:
                    rgb = torch.from_numpy(np.transpose(rgb, (3, 0, 1, 2)))  # (3, T, H=540, W=960)
            assert rgb.shape[1] == traj.shape[0], "Trajectory is inconsistent with RGB video!"
            
        # clipping or padding with zero
        max_frames = traj.size(0) if self._MAX_FRAMES < 0 else self._MAX_FRAMES
        len_valid = min(traj.size(0), max_frames)
        
        input_data = torch.tensor([])  # a placeholder (not used)
        if 'rgb' in self.modalities:
            input_data = torch.zeros([rgb.size(0), max_frames, rgb.size(2), rgb.size(3)], dtype=rgb.dtype)
            input_data[:, :len_valid] = rgb[:, :len_valid]
        
        # ------------  process target trajectory output  ------------ 
        len_pad, len_real = max_frames, len_valid
        odometry = torch.eye(4)[None, :, :].repeat(len_pad, 1, 1).to(torch.float32)
        output_traj = torch.zeros([len_pad, traj.size(1)], dtype=traj.dtype)
        traj_valid = traj[:len_valid]  # (T, 3) or (T, 2): [3D or 2D local trajectory]
        if self.use_odom:
            odom = self.odom_data[index][:num_preserve]  # (T, 4, 4)
            odom_valid = odom[:len_valid]
            odometry[:len_valid] = torch.from_numpy(odom_valid)
            # project all future hand locations to each of current frames
            traj3d_valid_all = self._get_projected_traj3d(traj3d[:len_valid].numpy(), odom_valid)  # (T*(T+1)/2, 3)
            traj3d_valid = torch.from_numpy(traj3d_valid_all[:len_real])  # only use the trajectory in the first frame
            traj_valid = self._XYZ_to_uv(traj3d_valid) if self.target == '2d' else traj3d_valid  # (T, 3) or (T, 2): [3D or 2D global trajectory]
        
        # transformed to normalized 2d representation
        traj_norm_valid = self._normalize(traj_valid)
        output_traj[:len_real] = traj_norm_valid

        return self.rgb_paths[index], input_data, odometry, len_valid, output_traj



def build_dataloaders(cfg, phase='trainval'):
    """Loading the dataset"""
    from gluoncv.torch.data.transforms.videotransforms.video_transforms import Compose, Resize, Normalize
    from gluoncv.torch.data.transforms.videotransforms.volume_transforms import ClipToTensor

    
    data_root = os.path.join(cfg.DATA.data_path, cfg.DATA.dataset)
    transform_train, transform_test = None, None
    if cfg.DATA.transform is not None:
        input_size = cfg.DATA.transform.input_size
        transform_train = Compose([Resize(input_size), ClipToTensor(), Normalize(mean=cfg.DATA.transform.means, std=cfg.DATA.transform.stds)])
        transform_test = Compose([Resize(input_size), ClipToTensor(), Normalize(mean=cfg.DATA.transform.means, std=cfg.DATA.transform.stds)])
    
    if phase == 'trainval':
        # train set
        trainset = EgoPAT3D(data_root, phase='train', transform=transform_train, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        train_loader = DataLoader(trainset, batch_size=cfg.TRAIN.batch_size, shuffle=True, num_workers=cfg.num_workers, pin_memory=True)
        # validation set
        valset = EgoPAT3D(data_root, phase='val', transform=transform_test, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        val_loader = DataLoader(valset, batch_size=cfg.TEST.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True)
        # print
        print("Number of train/val: {}/{}".format(len(trainset), len(valset)))
        return train_loader, val_loader
    else:
        # test set
        testset = EgoPAT3D(data_root, phase='test', transform=transform_test, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        test_loader = DataLoader(testset, batch_size=cfg.TEST.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True)
        # test_novel set
        testnovel_set = EgoPAT3D(data_root, phase='test_novel', transform=transform_test, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        testnovel_loader = DataLoader(testnovel_set, batch_size=cfg.TEST.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True)
        # print
        print("Number of test/test_novel samples: {}/{}".format(len(testset), len(testnovel_set)))
        return test_loader, testnovel_loader


    
def vis_input_output(train_loader, trainset, target='3d', use_odom=False, save_gif=True):

    result_dir = '../output/selected_samples/'
    os.makedirs(result_dir, exist_ok=True)
    vis_dir = os.path.join(result_dir, 'vis')
    os.makedirs(vis_dir, exist_ok=True)
    filelist_fid = open(os.path.join(result_dir, 'filelist.txt'), 'w')
    
    for i, (filename, clip, odometry, nf, traj) in tqdm(enumerate(train_loader), total=len(trainset), desc='Train'):
        # get file info
        sample_info = filename[0][:-4].split('/')
        vis_filename = 'vis_{}_{}'.format(sample_info[-2], sample_info[-1].split('_')[0])

        # read video by filename
        video = read_video(filename[0])
        video_size = video.shape[1:3]  # (H, W)

        if trainset.centralize:
            traj += 0.5
        u = torch.floor(traj[0, :, 0] * video_size[1]).to(torch.long)  # (T,) or (T*(T+1)/2,)
        v = torch.floor(traj[0, :, 1] * video_size[0]).to(torch.long)  # (T,) or (T*(T+1)/2,)

        if use_odom:
            frame = cv2.cvtColor(video[0], cv2.COLOR_RGB2BGR)  # use the first frame
            traj2d = torch.hstack((u[:, None], v[:, None])).numpy()
            vis_frame_traj(frame, traj2d)
            vis_file = os.path.join(vis_dir, vis_filename + '_global.png')
            cv2.imwrite(vis_file, frame)
        
        # visualize
        num_frames = int(nf[0])
        vis_mat = []
        for t in range(num_frames):
            frame = video[t]
            if use_odom:
                # back-project future hand locations of the first frame to current frame
                if t > 0:
                    # normalized 3D to actual 3D in camera_0
                    traj3d = denormalize(traj[0, t:], target='3d', intrinsics=trainset.intrinsics, max_depth=trainset._MAX_DEPTH)
                    # camera_0 to camera_t (Note: all T-t points are in the same camera_0 ,thus their odometry are ignored by setting to identity matrix)
                    odom_new = np.copy(odometry[0].numpy())
                    odom_new[(t+1):] = np.eye(4)[None, :, :]
                    traj3d = global_to_local(traj3d, odom_new, start_idx=t+1)
                    # in camera_t: 3D to normalized 2D
                    traj2d = normalize(traj3d, target='3d', intrinsics=trainset.intrinsics, vis_ratio=trainset.vis_ratio, max_depth=trainset._MAX_DEPTH)
                    # normalized 2D to (u,v)
                    u = np.floor(traj2d[:, 0] * video_size[1]).astype(np.int32)
                    v = np.floor(traj2d[:, 1] * video_size[0]).astype(np.int32)
                    traj2d = np.stack((u, v), axis=1)
                else:
                    traj2d = torch.hstack((u[:, None], v[:, None])).numpy()
                vis_frame_traj(frame, traj2d)
            else:
                cv2.circle(frame, (int(u[t]), int(v[t])), radius=5, color=(0, 255, 0), thickness=-1)
            vis_mat.append(frame)
        vis_mat = np.array(vis_mat)  # (T, 224, 224, 3)
        if save_gif:
            video_to_gif(vis_mat, os.path.join(vis_dir, vis_filename) + '.gif', fps=5)
        else:
            write_video(vis_mat, os.path.join(vis_dir, vis_filename) + '.mp4', fps=5)
        
        filelist_fid.writelines('{}/{}/{}\n'.format(sample_info[-3], sample_info[-2], sample_info[-1]))
    filelist_fid.close()


def get_frame_stats():
    
    num_frames = []
    for i, (filename, clip, odometry, nf, traj) in tqdm(enumerate(train_loader), total=len(trainset), desc='Train'):
        num_frames.append(int(nf[0]))
    
    for i, (filename, clip, odometry, nf, traj) in tqdm(enumerate(val_loader), total=len(valset), desc='Val'):
        num_frames.append(int(nf[0]))
    
    print("# frames: max={}, min={}, mean={}".format(max(num_frames), min(num_frames), np.mean(num_frames)))
    plt.hist(num_frames, bins=100)
    plt.savefig('../data/EgoPAT3D/num_frames_tinyset.png')


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    from gluoncv.torch.data.transforms.videotransforms.video_transforms import Compose, Resize, Normalize
    from gluoncv.torch.data.transforms.videotransforms.volume_transforms import ClipToTensor
    import sys

    rgb_transform = Compose([Resize([64, 64]), ClipToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
    
    class data_cfg: pass
    data_cfg.scenes = ['1']
    data_cfg.tinyset = True
    data_cfg.load_all = True

    class model_cfg: pass
    model_cfg.target = '3d'
    model_cfg.use_odom = True
    
    trainset = EgoPAT3D('../data/EgoPAT3D', phase='train', transform=rgb_transform, data_cfg=data_cfg, model_cfg=model_cfg)
    train_loader = DataLoader(trainset, batch_size=1, shuffle=False, num_workers=4, pin_memory=True)

    valset = EgoPAT3D('../data/EgoPAT3D', phase='val', transform=rgb_transform, data_cfg=data_cfg, model_cfg=model_cfg)
    val_loader = DataLoader(valset, batch_size=1, shuffle=False, num_workers=0, pin_memory=False)

    print("Number of train/val samples: {}/{}".format(len(trainset), len(valset)))
    
    if sys.argv[1] == 'vis':
        # visualize data from DataLoader
        vis_input_output(train_loader, trainset, target=model_cfg.target, use_odom=model_cfg.use_odom)
    
    elif sys.argv[1] == 'stat':
        # get the stats of the number frames for all clips
        get_frame_stats()

    elif sys.argv[1] == 'check_val':
        for i, (clip, nframes, traj_gt) in tqdm(enumerate(val_loader),
                total=len(val_loader), ncols=0, desc='check valset'):
            if i == 0:
                print(clip.size(), traj_gt.size())  #  (B, 3, T, H, W),  (B, T, 3)
    else:
        raise NotImplementedError

    print("Done!")
src/EgoPAT3DLoader.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import torch
import os
import pickle
import cv2
import numpy as np


H2O_STATS = {'2d': {'min': [0, 0], 
                    'max': [1280,  720], 
                    'mean': [680.92605918, 349.150914]}, 
             '3dc': {'min': [-0.33688032, -0.3860082 ,  0.13619463], 
                     'max': [0.45182842, 0.34188776, 0.76754248], 
                     'mean': [ 0.02773121, -0.00888447,  0.4658583 ]}, 
             '3dw': {'min': [-0.32412812, -0.27855986,  0.16916465], 
                     'max': [0.37120016, 0.28456741, 0.74655667], 
                     'mean': [0.01147821, 0.01590901, 0.4859054 ]}}

INTRINSICS = {'fx': 636.6593017578125, 'fy': 636.251953125, 
              'cx': 635.283881879317, 'cy': 366.8740353496978, 
              'width': 1280.0, 'height': 720.0}  # all samples' intrinsics actuall are the same as this!!!


def normalize_traj(traj, target='3d', use_global=True):
    if target == '3d':
        # use the min-max normalization
        if use_global:
            minvals = np.array([H2O_STATS['3dw']['min']])  # (1, 3)
            maxvals = np.array([H2O_STATS['3dw']['max']])
        else:
            minvals = np.array([H2O_STATS['3dc']['min']])  # (1, 3)
            maxvals = np.array([H2O_STATS['3dc']['max']])
        traj_norm = (traj - minvals) / (maxvals - minvals + 1e-6)
    else:
        # use the intrinsic range for normalization
        traj_norm = np.copy(traj).astype(np.float32)
        traj_norm[:, 0] = traj[:, 0] / H2O_STATS['2d']['max'][0]  # 1280
        traj_norm[:, 1] = traj[:, 1] / H2O_STATS['2d']['max'][1]  # 720
    return traj_norm


def denormalize_traj(traj_norm, target='3d', use_global=True):
    if target == '3d':
        # recover from min-max normalization
        if use_global:
            minvals = np.array([H2O_STATS['3dw']['min']])  # (1, 3)
            maxvals = np.array([H2O_STATS['3dw']['max']])
        else:
            minvals = np.array([H2O_STATS['3dc']['min']])  # (1, 3)
            maxvals = np.array([H2O_STATS['3dc']['max']])
        traj = traj_norm * (maxvals - minvals + 1e-6) + minvals
    else:
        # recover from the intrinsic range for normalization
        traj = np.copy(traj_norm)
        traj[:, 0] = traj_norm[:, 0] * H2O_STATS['2d']['max'][0]  # 1280
        traj[:, 1] = traj_norm[:, 1] * H2O_STATS['2d']['max'][1]  # 720
    return traj


def world_to_camera(traj3d_global, cam2world):
    traj3d_local = []
    for xyz, Rc2w in zip(traj3d_global, cam2world):
        # world to cam
        extrinsics = np.linalg.inv(Rc2w) # 4 x 4
        xyz_homo = np.concatenate([xyz[:, np.newaxis], np.ones((1, 1))], axis=0)  # (4, 1)
        xyz_local = extrinsics.dot(xyz_homo)[:3, 0]  # (3,)
        traj3d_local.append(xyz_local)
    traj3d_local = np.array(traj3d_local)

    return traj3d_local


def XYZ_to_uv(traj3d, intrinsics):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points in local coordinate system
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = traj3d[:, 0] * intrinsics['fx'] / (traj3d[:, 2] + 1e-6) + intrinsics['cx']
    traj2d[:, 1] = traj3d[:, 1] * intrinsics['fy'] / (traj3d[:, 2] + 1e-6) + intrinsics['cy']

    # clip the coordinates if out of boundary
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['width']-1)
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['height']-1)
    traj2d = np.floor(traj2d).astype(np.int32)

    return traj2d


class H2O(Dataset):
    def __init__(self, root_dir, phase='train', transform=None, data_cfg=None, model_cfg=None):
        self.root_dir = os.path.join(root_dir, 'Ego3DTraj')  # ./data/H2O/Ego3DTraj
        self.phase = phase
        self.transform = transform

        # data config
        self._MAX_FRAMES = getattr(data_cfg, 'max_frames', -1) # the maximum number of input frames
        self._MAX_DEPTH = getattr(data_cfg, 'max_depth', 1.0)  # the max depth of hand trajectory
        self.load_all = getattr(data_cfg, 'load_all', False)

        # model config
        self.target = getattr(model_cfg, 'target', '3d')
        self.modalities = getattr(model_cfg, 'modalities', ['rgb', 'loc'])
        self.use_global = getattr(model_cfg, 'use_global', False)  # local camera 3D space by default
        self.centralize = getattr(model_cfg, 'centralize', False)
        self.normalize = getattr(model_cfg, 'normalize', True)

        # dataset split
        self.samples = self.get_split()

        # read all trajectory data
        self.traj_data, self.video_data = self.read_all_data()
        
        if isinstance(phase, list): # ['train', 'val', 'test']
            self.check_intrinsics()
            self.traj_stats = self.get_stats()


    def get_split(self):
        def read_split(filepath):
            assert os.path.exists(filepath), "File path does not exist! {}".format(filepath)
            samples = []
            with open(filepath, 'r') as f:
                for line in f.readlines():
                    samples.append(line.strip())
            return samples
        
        split_dir = os.path.join(self.root_dir, 'splits')
        if isinstance(self.phase, str):
            samples = read_split(os.path.join(split_dir, self.phase + '.txt'))
        else:  # read all splits
            samples = []
            for phase in self.phase:
                samples.extend(read_split(os.path.join(split_dir, phase + '.txt')))
        return samples
    

    def get_time_stamps(self, total_frames, clip_len, stepsize=15):
        time_stamps = []
        ps, pe = 0, clip_len-1
        while pe < total_frames:
            time_stamps.append([ps, pe])
            ps += stepsize
            pe += stepsize
        time_stamps = np.array(time_stamps)  # (N,2)
        return time_stamps
    

    def preprocess_trajdata(self, trajdata, intrinsics):
        """ traj3d: list of 3D trajectories in world space
            intrinsics: camera intrinsics
        """
        traj_2d_all, traj_3dc_all, traj_3dw_all, timestamps_all, campose_call = [], [], [], [], []
        for data in trajdata:
            start, end = data['start'], data['end']
            traj3d_world = data['traj3d']
            cam2world = data['cam2world']
            num_frames = end - start + 1
            
            # 3D transformation from world to camera system
            traj3d_camera = world_to_camera(traj3d_world, cam2world)
            # 3D to 2D
            traj2d = XYZ_to_uv(traj3d_camera, intrinsics)

            if self.normalize:
                # normalize trajectories
                traj2d = normalize_traj(traj2d, target='2d')
                traj3d_world = normalize_traj(traj3d_world, target='3d', use_global=True)
                traj3d_camera = normalize_traj(traj3d_camera, target='3d', use_global=False)

            # sampling trajectory clips
            max_frames = num_frames if self._MAX_FRAMES < 0 else self._MAX_FRAMES  # 40
            time_stamps = self.get_time_stamps(num_frames, max_frames)

            if time_stamps.shape[0] > 0:
                for ps, pe in time_stamps:
                    traj_2d_all.append(traj2d[ps: pe+1])
                    traj_3dc_all.append(traj3d_camera[ps: pe+1])
                    traj_3dw_all.append(traj3d_world[ps: pe+1])
                    timestamps_all.append(np.array([start + ps, start + pe]))  # offset by trajectory start
                    campose_call.append(cam2world[ps: pe+1])
            else:
                # self._MAX_FRAMES == -1 or self._MAX_FRAMES > num_frames
                traj_2d_all.append(traj2d)
                traj_3dc_all.append(traj3d_camera)
                traj_3dw_all.append(traj3d_world)
                timestamps_all.append(np.array([start, end]))
                campose_call.append(cam2world)
        
        return traj_2d_all, traj_3dc_all, traj_3dw_all, timestamps_all, campose_call


    def read_rgb(self, video_file):
        assert os.path.exists(video_file), "Video file does not exist! {}".format(video_file)
        cap = cv2.VideoCapture(video_file)
        success, frame = cap.read()
        videos = []
        while success:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            videos.append(frame)
            success, frame = cap.read()
        cap.release()
        cv2.destroyAllWindows()
        videos = np.array(videos)
        return videos
    

    def read_all_data(self):
        traj_data = {'2d': [], '3dc': [], '3dw': []}
        video_data = {'data': [], 'file': [], 'timestamps': [], 'campose': []}
        for sample in self.samples:
            traj_file = os.path.join(self.root_dir, 'traj', sample + '.pkl')
            assert os.path.exists(traj_file), 'Trajectory file does not exist! {}'.format(traj_file)
            with open(traj_file, 'rb') as f:
                data = pickle.load(f)
            traj2d_l, traj3d_lc, traj3d_lw, timestamps_l, campose_l = self.preprocess_trajdata(data['left_hand'], data['intrinsics'])
            traj2d_r, traj3d_rc, traj3d_rw, timestamps_r, campose_r = self.preprocess_trajdata(data['right_hand'], data['intrinsics'])

            # merge data
            traj_data['2d'].extend(traj2d_l + traj2d_r)
            traj_data['3dc'].extend(traj3d_lc + traj3d_rc)
            traj_data['3dw'].extend(traj3d_lw + traj3d_rw)
            timestamps = timestamps_l + timestamps_r
            campose = campose_l + campose_r

            # read video data
            vid_file = os.path.join(self.root_dir, 'video', sample + '.mp4')
            video_data['file'].extend([vid_file]*len(timestamps))
            video_data['timestamps'].extend(timestamps)
            video_data['campose'].extend(campose)
            if self.load_all:
                video = self.read_rgb(vid_file)  # (T, H, W, C)
                video = self.transform(video) if self.transform is not None else torch.from_numpy(np.transpose(video, (3, 0, 1, 2))) # [3, T, H, W]
                video_data['data'].extend([video[:, ps: pe+1] for ps, pe in timestamps])
        
        return traj_data, video_data


    def check_intrinsics(self):
        all_intrinsics = []
        for sample in self.samples:
            traj_file = os.path.join(self.root_dir, 'traj', sample + '.pkl')
            with open(traj_file, 'rb') as f:
                data = pickle.load(f)
            all_intrinsics.append(data['intrinsics'])
        
        for i in range(len(all_intrinsics)-1):
            if all_intrinsics[i] != all_intrinsics[i+1]:
                print("Not all intrinsics are equal!")


    def get_stats(self):
        stats = {k: dict() for k in self.traj_data.keys()}
        for k, traj_list in self.traj_data.items():
            dim = 2 if k == '2d' else 3
            sumvals = np.zeros((dim), dtype=np.float32)
            num_pt = 0
            minvals, maxvals = [], []
            for traj in traj_list:
                minvals.append(np.min(traj, axis=0))
                maxvals.append(np.max(traj, axis=0))
                sumvals += np.sum(traj, axis=0)
                num_pt += traj.shape[0]
            # compute stats
            stats[k]['min'] = np.min(np.array(minvals), axis=0)
            stats[k]['max'] = np.max(np.array(maxvals), axis=0)
            stats[k]['mean'] = sumvals / num_pt
        return stats


    def __len__(self):
        return len(self.traj_data['3dw'])

    
    def __getitem__(self, index):

        input_data = torch.tensor([])  # a placeholder (not used)
        vid_file = self.video_data['file'][index]
        ps, pe = self.video_data['timestamps'][index]
        len_valid = pe - ps + 1
        cam_pose = np.copy(self.video_data['campose'][index])  # camera pose (cam to world)

        # get normalized & centralized trajectory data
        if self.target == '3d':
            output_traj = np.copy(self.traj_data['3dw'][index]) if self.use_global else np.copy(self.traj_data['3dc'][index])
            dim_traj = 3
        else:
            output_traj = np.copy(self.traj_data['2d'][index])  # in [0, 1]
            dim_traj = 2
        if self.centralize:
            output_traj = output_traj * 2.0 - 1.0  # in [-1, 1]
        
        if len_valid < self._MAX_FRAMES:  # need to pad trajectory
            output_traj = np.concatenate([output_traj, np.zeros((self._MAX_FRAMES-len_valid, dim_traj))], axis=0)
            cam_pose = np.concatenate([cam_pose, np.tile(np.eye(4), (self._MAX_FRAMES-len_valid, 1, 1))], axis=0)
        # ndarray to torch.tensor
        output_traj = torch.from_numpy(output_traj).to(torch.float32)
        input_pose = torch.from_numpy(cam_pose).to(torch.float32)

        # get video data
        if 'rgb' in self.modalities:
            if self.load_all:
                input_data = self.video_data['data'][index].clone()  # already preprocessed
            else:
                rgb = self.read_rgb(vid_file)
                rgb = self.transform(rgb) if self.transform is not None else torch.from_numpy(np.transpose(rgb, (3, 0, 1, 2))) # [3, T, H, W]
                input_data = rgb[:, ps: pe+1]
            assert input_data.size(1) == len_valid

            if len_valid < self._MAX_FRAMES: # need to pad video
                c, t, h, w = input_data.size()
                input_data = torch.concat([input_data, torch.zeros([c, self._MAX_FRAMES-len_valid, h, w])], axis=1)
    
        return vid_file, input_data, input_pose, len_valid, output_traj



def build_dataloaders(cfg, phase='trainval'):
    """Loading the dataset"""
    from gluoncv.torch.data.transforms.videotransforms.video_transforms import Compose, Resize, Normalize
    from gluoncv.torch.data.transforms.videotransforms.volume_transforms import ClipToTensor
    
    data_root = os.path.join(cfg.DATA.data_path, cfg.DATA.dataset)
    transform_train, transform_test = None, None
    if cfg.DATA.transform is not None:
        input_size = cfg.DATA.transform.input_size
        transform_train = Compose([Resize(input_size), ClipToTensor(), Normalize(mean=cfg.DATA.transform.means, std=cfg.DATA.transform.stds)])
        transform_test = Compose([Resize(input_size), ClipToTensor(), Normalize(mean=cfg.DATA.transform.means, std=cfg.DATA.transform.stds)])
    
    if phase == 'trainval':
        # train set
        trainset = H2O(data_root, phase='train', transform=transform_train, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        train_loader = DataLoader(trainset, batch_size=cfg.TRAIN.batch_size, shuffle=True, num_workers=cfg.num_workers, pin_memory=True)
        # validation set
        valset = H2O(data_root, phase='val', transform=transform_test, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        val_loader = DataLoader(valset, batch_size=cfg.TEST.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True)
        # print
        print("Number of train/val: {}/{}".format(len(trainset), len(valset)))
        return train_loader, val_loader
    else:
        # test set
        testset = H2O(data_root, phase='test', transform=transform_test, data_cfg=cfg.DATA, model_cfg=cfg.MODEL)
        test_loader = DataLoader(testset, batch_size=cfg.TEST.batch_size, shuffle=False, num_workers=cfg.num_workers, pin_memory=True)
        # print
        print("Number of test samples: {}".format(len(testset)))
        return test_loader
    

if __name__ == '__main__':

    from gluoncv.torch.data.transforms.videotransforms.video_transforms import Compose, Resize, Normalize
    from gluoncv.torch.data.transforms.videotransforms.volume_transforms import ClipToTensor
    from tqdm import tqdm

    class data_cfg: pass
    data_cfg.max_frames = 64
    data_cfg.load_all = True

    class model_cfg: pass
    model_cfg.target = '2d'
    model_cfg.modalities = ['rgb', 'loc']
    model_cfg.use_global = True
    model_cfg.centralize = True
    model_cfg.normalize = True

    rgb_transform = Compose([Resize([64, 64]), ClipToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
    # allset = H2O('data/H2O', phase=['train', 'test', 'val'], transform=rgb_transform, data_cfg=data_cfg, model_cfg=model_cfg)

    trainset = H2O('data/H2O', phase='train', transform=rgb_transform, data_cfg=data_cfg, model_cfg=model_cfg)
    train_loader = DataLoader(trainset, batch_size=16, shuffle=True, num_workers=4, pin_memory=True)

    valset = H2O('data/H2O', phase='val', transform=rgb_transform, data_cfg=data_cfg, model_cfg=model_cfg)
    val_loader = DataLoader(valset, batch_size=1, shuffle=False, num_workers=0, pin_memory=False)

    testset = H2O('data/H2O', phase='val', transform=rgb_transform, data_cfg=data_cfg, model_cfg=model_cfg)
    test_loader = DataLoader(testset, batch_size=1, shuffle=False, num_workers=0, pin_memory=False)

    print("# train: {}, # val: {}, # test: {}".format(len(trainset), len(valset), len(testset)))

    for i, (filename, clip, odometry, nf, traj) in tqdm(enumerate(train_loader), total=len(train_loader), desc='Train'):
        # print("clip size: {}, traj size: {}".format(clip.size(), traj.size()))
        pass

    for i, data in tqdm(enumerate(val_loader), total=len(val_loader), desc='Val'):
        # print("clip size: {}, traj size: {}".format(clip.size(), traj.size()))
        pass

    for i, data in tqdm(enumerate(test_loader), total=len(test_loader), desc='Test'):
        # print("clip size: {}, traj size: {}".format(clip.size(), traj.size()))
        pass
src/H2OLoader.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F



class CoordLoss(nn.Module):
    def __init__(self, cfg, max_epoch=1500):
        """Simple Mean-squared Error as the loss function"""
        super(CoordLoss, self).__init__()
        if cfg.function == 'MSE':
            self.coord_loss = nn.MSELoss(reduction='none')
        elif cfg.function == 'Huber':
            self.coord_loss = nn.HuberLoss(reduction='none', delta=getattr(cfg, 'delta', 1.0))

        self.scale = getattr(cfg, 'scale', 1.0)
        self.robust_weighing = getattr(cfg, 'robust_weighing', False)
        self.tau = getattr(cfg, 'tau', 0.1)
        self.velo_coeff = getattr(cfg, 'velo_coeff', 0)
        self.velo_warp = getattr(cfg, 'velo_warp', 0)


    def get_robust_weights(self, z_values, tau=0.1):
        """ Compute the robust weights using temporal depth values
        """
        diffs = torch.cat([torch.tensor([0.0]).to(z_values.device), torch.abs(torch.diff(z_values))])
        horizon = z_values.size(0)
        weights = horizon * F.softmax(-diffs / tau, dim=-1)
        return weights
    

    def compute_loss(self, y_hat, y_true, logvar=torch.tensor([0.0]), weights=torch.tensor([1.0])):
        """ Loss function
        """
        # huber loss
        loss = self.coord_loss(y_hat, y_true)  # (Tu, 3)
        # uncertainty regularization
        beta_1 = torch.exp(-logvar).to(loss.device)
        beta_2 = logvar.to(loss.device)
        if logvar.size(-1) in [1, 3]:  # temporal or spacetime or None
            loss = beta_1 * loss + beta_2  # (Tu, 3)
        elif logvar.size(-1) == 2: # hybrid
            loss_xy = beta_1[:, 0:1] * loss[:, :-1] + beta_2[:, 0:1]
            loss_z = (beta_1[:, 1:2] * loss[:, -1:] + beta_2[:, 1:2]) * weights.unsqueeze(-1)
            loss = torch.cat([loss_xy, loss_z], dim=-1)
        # take sum over xyz, and mean over t
        loss = torch.mean(torch.sum(loss, dim=-1))
        return loss


    def forward(self, outputs, nframes, trajectories):
        """
            outputs: a list of dict() with obs_ratio (scalar) as key and unobs_pred (Tu, 3) as value
            nframes: (B,)
            trajectories: (B, T, 3)
        """
        outputs_unobs, outputs_obs = outputs
        unobs_avg_loss, obs_avg_loss, cnt = 0, 0, 0
        unobs_velo_loss, unobs_warp_loss, obs_velo_loss = 0, 0, 0
        
        for b, ratio_preds in enumerate(outputs_unobs):
            # get the number observed frames and full frames
            num_full = nframes[b]
            traj_gt = trajectories[b]  # (T,3)
            weights = self.get_robust_weights(traj_gt[:num_full, 2], tau=self.tau) if self.robust_weighing and traj_gt.size(-1) == 3 else torch.ones((num_full)).to(traj_gt.device)

            for i, (r, preds) in enumerate(ratio_preds.items()):
                num_obs = torch.floor(num_full * float(r)).to(torch.long)
                out_dim = preds['traj'].size(-1)

                # prepare the ground truth (from the last observed to the last valid frame)
                gts = traj_gt[num_obs: num_full, :]

                logvar = preds['unct'] if 'unct' in preds else torch.tensor([0.0])
                unobs_avg_loss += self.compute_loss(preds['traj'], gts[:, :out_dim], logvar=logvar, weights=weights[num_obs: num_full])
                
                if 'velo' in preds:
                    # velocity loss
                    gt_velo = torch.diff(torch.cat((traj_gt[num_obs-1: num_obs, :], gts), dim=0), dim=0)
                    unobs_velo_loss += self.velo_coeff * self.compute_loss(preds['velo'], gt_velo[:, :out_dim])
                    # accumulative warping loss
                    traj_warp = traj_gt[num_obs-1: num_obs, :out_dim] + torch.cumsum(preds['velo'], dim=0)
                    unobs_warp_loss += self.velo_warp * self.compute_loss(traj_warp, preds['traj'])
                
                # loss for the observed frames
                if outputs_obs is not None:
                    pred_obs = outputs_obs[b][r]
                    logvar = pred_obs['unct'] if 'unct' in pred_obs else torch.tensor([0.0])
                    obs_avg_loss += self.compute_loss(pred_obs['traj'], traj_gt[:num_obs, :out_dim], logvar=logvar, weights=weights[:num_obs])
                    
                    if 'velo' in pred_obs:
                        # velocity loss
                        gt_velo = torch.diff(torch.cat((torch.zeros((1, out_dim)).to(gts.device), traj_gt[:num_obs]), dim=0), dim=0)
                        obs_velo_loss += self.velo_coeff * self.compute_loss(pred_obs['velo'], gt_velo[:, :out_dim])     
                cnt += 1

        unobs_avg_loss /= cnt
        total_loss = unobs_avg_loss
        loss_out = {'total_loss': total_loss,
                    'unobs_loss': unobs_avg_loss}
        
        if self.velo_coeff > 0:
            unobs_velo_loss /= cnt
            total_loss += unobs_velo_loss
            loss_out.update({'total_loss': total_loss,
                             'unobs_velo_loss': unobs_velo_loss})
        
        if self.velo_warp > 0:
            unobs_warp_loss /= cnt
            total_loss += unobs_warp_loss
            loss_out.update({'total_loss': total_loss,
                             'unobs_warp_loss': unobs_warp_loss})
        
        if outputs_obs is not None:
            obs_avg_loss /= cnt
            total_loss += obs_avg_loss
            loss_out.update({'total_loss': total_loss,
                             'obs_loss': obs_avg_loss})
            
            if self.velo_coeff > 0:
                obs_velo_loss /= cnt
                total_loss += obs_velo_loss
                loss_out.update({'total_loss': total_loss,
                                'obs_velo_loss': obs_velo_loss})

        for k, v in loss_out.items():
            loss_out[k] = v * self.scale

        return loss_out
src/Losses.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import argparse
import yaml
from easydict import EasyDict
from pprint import pformat
import os
import torch



def parse_configs(phase='train'):
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='./config/baseline.yml',
                        help='The relative path of dataset.')
    parser.add_argument('--gpus', type=str, default="0", 
                        help="The delimited list of GPU IDs separated with comma. Default: '0'.")
    parser.add_argument('--num_workers', type=int, default=4, metavar='N',
                        help='The number of workers to load dataset. Default: 4')
    parser.add_argument('--eval', action='store_true', 
                        help='If specified, run the evaluation only.')
    parser.add_argument('--tag', type=str, default='debug',
                        help='The tag to save model results')
    args = parser.parse_args()
    with open(args.config, 'r') as f:
        cfg = EasyDict(yaml.safe_load(f))
    cfg.update(vars(args))
    # add device
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    cfg.update(device=device)
    cfg.MODEL.update(device=device)
    # add root path
    root_path = os.path.dirname(__file__)
    cfg.update(root_path=root_path)

    if cfg.tag == 'debug':
        cfg.TRAIN.epoch = 1
        if cfg.TRAIN.scheduler.type == 'cosine_warmup':
            cfg.TRAIN.epoch += cfg.TRAIN.scheduler.num_restarts # the minimum training epochs
        cfg.TRAIN.eval_interval = 1
        cfg.TRAIN.snapshot_interval = 1
        cfg.DATA.load_all = False
    
    if phase == 'test' and hasattr(cfg.DATA, 'load_all'):
        cfg.DATA.load_all = False
    
    # save configs to file
    exp_dir = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag)
    os.makedirs(exp_dir, exist_ok=True)
    with open(os.path.join(exp_dir, 'config_{}.yaml'.format(phase)), 'w') as f:
        f.writelines(pformat(vars(cfg)))

    return cfg
src/config.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch
import torch.nn as nn
from torchvision import models
from torch.hub import load_state_dict_from_url
import timm
import math
from functools import reduce
from operator import mul
import numpy as np
import os, sys
HoiForecast_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/hoi-forecast')
sys.path.append(HoiForecast_ROOT)
""" 
The followings are hoi-forecast classes
"""
from networks.layer import Mlp



class FullyConvResNet(models.ResNet):
    """ Remove the pooling layers of ResNet to make it fully convolutional
    """
    def __init__(self, base='resnet18', img_size=64, pretrained=False, **kwargs):
        # Start with standard resnet18 defined here 
        super().__init__(block = models.resnet.BasicBlock, layers=[2, 2, 2, 2], num_classes=1000, **kwargs)

        if pretrained:
            state_dict = load_state_dict_from_url(models.resnet.model_urls[base])
            self.load_state_dict(state_dict)

        # compute the feature dimension
        dim = float(img_size)
        for n in range(5):
            dim = np.ceil(dim / 2.0)
        self.num_features = int((16 * 32) * dim**2)


    # Reimplementing forward pass. 
    def forward_features(self, x):
        # Standard forward for resnet18
        x = self.conv1(x)  # (64, H/2, W/2)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)  # (64, H/4, W/4)

        x = self.layer1(x)
        x = self.layer2(x)  # (128, H/8, W/8)
        x = self.layer3(x)  # (256, H/16, W/16)
        x = self.layer4(x)  # (512, H/32, W/32)
        
        x = torch.flatten(x, start_dim=1)  # (self.num_features,)
        return x


class FrameEmbedding(nn.Module):
    """ Embed each frame into a low-dimensional feature"""
    def __init__(self, cfg, input_size=224):
        super(FrameEmbedding, self).__init__()
        self.input_size = input_size
        self.embed_dim = cfg.feat_dim
        self.freeze_backbone = getattr(cfg, 'freeze_backbone', False)
        self.backbone_name = cfg.backbone
        # visual prompt tuning
        self.vpt = None
        if hasattr(cfg, 'vpt'):
            self.freeze_backbone = True
            self.vpt = getattr(cfg, 'vpt')
        
        if self.backbone_name == 'vit/b16-224':
            self.backbone = timm.create_model('vit_base_patch16_224', img_size=self.input_size, pretrained=True, num_classes=0)
        elif self.backbone_name == 'resnet18':
            img_size = self.input_size + 2 * self.vpt['num_tokens'] if self.vpt else self.input_size
            pretrained = True if self.vpt else False  # if not vpt, we train resnet from scratch
            self.backbone = FullyConvResNet('resnet18', img_size=img_size, pretrained=pretrained)  # train from scratch
        else:
            raise NotImplementedError
        
        if self.freeze_backbone:
            for name, param in self.backbone.named_parameters():
                param.requires_grad = False  # fix the parameters
        
        if self.vpt and self.backbone_name == 'vit/b16-224':
            self.prompt_proj = nn.Identity()  # do not project the prompt
            self.prompt_dropout = nn.Dropout(self.vpt['dropout'])
            self.prompt_embeddings = nn.Parameter(torch.zeros(1, self.vpt['prompt_len'], self.vpt['prompt_dim']))
            # xavier_uniform initialization
            val = math.sqrt(6. / float(3 * reduce(mul, self.backbone.patch_embed.patch_size, 1) + self.vpt['prompt_dim']))
            nn.init.uniform_(self.prompt_embeddings.data, -val, val)
        
        if self.vpt and self.backbone_name == 'resnet18':
            self.prompt_norm = nn.Identity()
            self.prompt_embeddings_tb = nn.Parameter(torch.zeros(
                    1, 3, 2 * self.vpt['num_tokens'],
                    input_size + 2 * self.vpt['num_tokens']
            ))
            self.prompt_embeddings_lr = nn.Parameter(torch.zeros(
                    1, 3, input_size, 2 * self.vpt['num_tokens']
            ))
            nn.init.normal_(self.prompt_embeddings_tb.data)
            nn.init.normal_(self.prompt_embeddings_lr.data)
        
        # two-layers MLP
        self.mlp = Mlp(self.backbone.num_features, hidden_features=512, out_features=self.embed_dim)
    
    
    def _pos_embed(self, x):
        x = torch.cat((self.backbone.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
        x = x + self.backbone.pos_embed
        return self.backbone.pos_drop(x)
    
    
    def incorporate_vit_prompt(self, x):
        # combine prompt embeddings with image-patch embeddings
        B = x.shape[0]
        # after CLS token, all before image patches
        x = self.backbone.patch_embed(x)  # (batch_size, 1 + n_patches, hidden_dim)
        x = self._pos_embed(x)
        # insert the learnable prompts
        x = torch.cat((
                x[:, :1, :],
                self.prompt_dropout(self.prompt_proj(self.prompt_embeddings).expand(B, -1, -1)),
                x[:, 1:, :]
            ), dim=1)  # (batch_size, cls_token + n_prompt + n_patches, hidden_dim)
        return x
    
    
    def incorporate_conv_prompt(self, x):
        B = x.shape[0]
        prompt_emb_lr = self.prompt_norm(
            self.prompt_embeddings_lr).expand(B, -1, -1, -1)
        prompt_emb_tb = self.prompt_norm(
            self.prompt_embeddings_tb).expand(B, -1, -1, -1)

        x = torch.cat((
            prompt_emb_lr[:, :, :, :self.vpt['num_tokens']],
            x, prompt_emb_lr[:, :, :, self.vpt['num_tokens']:]
            ), dim=-1)
        x = torch.cat((
            prompt_emb_tb[:, :, :self.vpt['num_tokens'], :],
            x, prompt_emb_tb[:, :, self.vpt['num_tokens']:, :]
        ), dim=-2)  # (B, 3, crop_size + num_prompts, crop_size + num_prompts)
        return x


    def forward(self, x):
        if self.vpt and self.backbone_name == 'vit/b16-224':
            # prompt tuning
            x = self.incorporate_vit_prompt(x)
            x = self.backbone.blocks(x)
            x = self.backbone.norm(x)
            x = x[:, 0]
        elif self.vpt and self.backbone_name == 'resnet18':
            # prompt tuning
            x = self.incorporate_conv_prompt(x)
            x = self.backbone.forward_features(x)  # (B, D)
        else:
            # cnn backbone
            x = self.backbone.forward_features(x)  # (B, D)
        x = self.mlp(x)
        return x
src/models/EmbedFrame.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch.nn as nn

class MlpModel(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0., out_layer=None):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)
        self.out_layer = out_layer() if out_layer is not None else None

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        if self.out_layer is not None:
            x = self.out_layer(x)
        return x
src/models/EmbedTrajectory.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch.nn as nn


class MlpHead(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, bn_layer=None, act_layer=nn.GELU, drop=0., out_layer=None):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.bn1 = bn_layer(hidden_features) if bn_layer is not None else nn.Identity()
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.bn2 = bn_layer(out_features) if bn_layer is not None else nn.Identity()
        self.drop = nn.Dropout(drop)
        self.out_layer = out_layer() if out_layer is not None else None

    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.bn2(x)
        x = self.drop(x)
        if self.out_layer is not None:
            x = self.out_layer(x)
        return x


class UncertaintyHead(nn.Module):
    def __init__(self, in_features, hidden_features=None, act_layer=nn.GELU, cfg=None):
        super().__init__()
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features, bias=False)
        self.act = act_layer()
        
        # parse the configurations
        self.act_out = getattr(cfg, 'act_out', None)
        self.unct_type = getattr(cfg, 'type', 'spacetime')  # spacetime, temporal, hybrid
        
        out_dims = {'spacetime': 3, 'hybrid': 2, 'temporal': 1}
        assert self.unct_type in out_dims, "Invalid uncertainty type."
        
        self.fc2 = nn.Linear(hidden_features, out_dims[self.unct_type], bias=False)
        # output activation
        self.out_layer = nn.Softplus(beta=1) if self.act_out else None
        # initialize
        self.apply(self.weights_init)


    def weights_init(self, m):
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, mean=0.0, std=0.0001)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.fc2(x)
        if self.out_layer:
            x = self.out_layer(x)
        return x

src/models/Heads.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch
import torch.nn as nn
import os, sys
HoiForecast_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/hoi-forecast')
sys.path.append(HoiForecast_ROOT)
""" 
The followings are hoi-forecast classes
"""
from networks.embedding import PositionalEncoding
from networks.layer import MultiHeadAttention



class ProbTranV3(nn.Module):
    def __init__(self, vis_dim, loc_dim, modalities, seq_len=40, h_dim=128, z_dim=16, pred_link=False):
        super(ProbTranV3, self).__init__()
        self.modalities = modalities
        assert any(mod in modalities for mod in ['rgb', 'loc'])  # at least one of modalitties should exist
        # input dimension
        self.x_dim = vis_dim + loc_dim
        if 'rgb' not in modalities:
            self.x_dim = loc_dim
        if 'loc' not in modalities:
            self.x_dim = vis_dim
        # predictive link dimension
        self.loc_dim = loc_dim
        self.h_dim = h_dim
        self.pred_link = pred_link

        # xt --> ht
        self.phi_x = nn.Sequential(
            nn.Linear(self.x_dim, h_dim),
            nn.ReLU(),
            nn.Linear(h_dim, h_dim))
        self.time_embed_h = PositionalEncoding(h_dim, max_len=seq_len)
        self.layer_norm_h = nn.LayerNorm(h_dim)
        # w(1:t-1) --> w_bar_t
        self.attn_w = MultiHeadAttention(h_dim // 2, num_heads=8, qkv_bias=False,
                                       qk_scale=None, attn_drop=0, proj_drop=0)
        self.layer_norm_wb = nn.LayerNorm(h_dim)
        # h(1:C) --> w_hat_t
        self.attn_h = MultiHeadAttention(h_dim, num_heads=8, qkv_bias=False,
                                       qk_scale=None, attn_drop=0, proj_drop=0)
        self.layer_norm_wh = nn.LayerNorm(h_dim * 2)
        # Prior: w_hat_t --> zt
        self.prior_mlp = nn.Sequential(
            nn.Linear(h_dim * 2, h_dim),
            nn.ReLU(),
            nn.Linear(h_dim, h_dim))
        self.prior_mean = nn.Sequential(
            nn.Linear(h_dim, z_dim),)
        # zt --> wt
        self.phi_z = nn.Sequential(
            nn.Linear(h_dim * 2 + z_dim, h_dim),
            nn.ReLU(),
            nn.Linear(h_dim, h_dim // 2))
        self.time_embed_w = PositionalEncoding(h_dim // 2, max_len=seq_len)
        self.layer_norm_w = nn.LayerNorm(h_dim // 2)
        # w(1:T) --> x(1:T)
        dim_phi_w = h_dim // 2 + self.loc_dim if self.pred_link and 'loc' in self.modalities else h_dim // 2
        self.phi_w = nn.Sequential(
            nn.Linear(dim_phi_w, self.loc_dim),
            nn.ReLU(),
            nn.Linear(self.loc_dim, self.loc_dim))
    

    def forward(self, mask, vis_obs=None, traj_obs=None):
        """
            \log p(z1:T | x1:C) = \sum\log pθ(zt | z1:t-1, x1:C )
            vis_obs: (B*N, T, D)
            traj_obs: (B*N, T, D)
            mask: (B*N, T), mask of the observed time steps
        """
        batch_size, seq_len = mask.size()[:2]
        device = mask.device

        wt = torch.zeros([batch_size, 1, self.h_dim // 2], device=device)
        w_accum = [wt]  # [(B*N, 1, d/2)]

        x_obs = torch.cat([obs for obs in [vis_obs, traj_obs] if obs is not None], dim=-1)
        # ht = LayerNorm(MLP(xt) + Position(t))
        h = self.layer_norm_h(self.time_embed_h(self.phi_x(x_obs)))

        for t in range(seq_len):
            # w¯ t = LayerNorm([wt−1; Attention(wt−1, w1:t−1, w1:t−1)])
            attn_w = self.attn_w(q=wt, k=torch.cat(w_accum, dim=1), v=torch.cat(w_accum, dim=1), mask=mask[:, :(t+1)].unsqueeze(1))
            wbar_t = self.layer_norm_wb(torch.cat([wt, attn_w], dim=-1))  # (B*N, 1, d)

            # wˆ t = LayerNorm([w¯ t; Attention(w¯ t, h1:C , h1:C )])
            attn_h = self.attn_h(q=wbar_t, k=h, v=h, mask=mask.unsqueeze(1))
            what_t = self.layer_norm_wh(torch.cat([wbar_t, attn_h], dim=-1))  # (B*N, 1, 2d)

            # zt = MLP(wˆ t)
            prior_shared = self.prior_mlp(what_t)
            zt = self.prior_mean(prior_shared)  # (B*N, 1, dz)

            # wt = LayerNorm(MLP([wˆ t; zt]) + Position(t))
            fuse_wz = self.phi_z(torch.cat([what_t, zt], dim=-1))
            wt = self.layer_norm_w(fuse_wz + self.time_embed_w.pe[:, t:(t+1)])
            w_accum.append(wt)
        
        phi_w_in = torch.cat(w_accum[1:], dim=1)  # (B*N, T, d/2)
        if self.pred_link and 'loc' in self.modalities:
            loc_tm1 = torch.cat((torch.zeros(batch_size, 1, self.loc_dim).to(device), traj_obs[:, :-1, :]), dim=1)  # (B*N, T, d)
            phi_w_in = torch.cat([phi_w_in, loc_tm1], dim=-1)
        x_pred = self.phi_w(phi_w_in)  # (B*N, T, D)

        return x_pred, dict()

src/models/ProTranV3.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch
import torch.nn as nn
from einops import rearrange
from src.models.Heads import MlpHead, UncertaintyHead
from src.models.EmbedTrajectory import MlpModel
from src.models.EmbedFrame import FrameEmbedding
from src.models.transformer_layers import Encoder
from src.models.utils import get_masks, get_rgb_observed, get_traj_observed, gather_results
from src.models.ProTranV3 import ProbTranV3


class TransformerSSM(nn.Module):
    def __init__(self, cfg, seq_len=64, input_size=224):
        super(TransformerSSM, self).__init__()
        self.embed_dim = cfg.feat_dim
        self.loc_dim = cfg.loc_dim
        self.modalities = cfg.modalities
        assert any(mod in cfg.modalities for mod in ['rgb', 'loc'])  # at least one of modalitties should exist
        self.transformer_params = vars(cfg.transformer) if hasattr(cfg, 'transformer') else dict()
        self.use_odom = getattr(cfg, 'use_odom', False)
        self.sup_obs = getattr(cfg, 'sup_obs', False)
        self.ssm_params = vars(cfg.ssm)
        self.ssm_params.update(seq_len=seq_len)
        self.aleatoric_uncertainty = getattr(cfg, 'aleatoric_uncertainty', None)
        self.pred_velocity = getattr(cfg, 'pred_velocity', False)
        
        # output config
        self.target = getattr(cfg, 'target', '3d')  # ['2d', '3d']
        self.ignore_depth = getattr(cfg, 'ignore_depth', False)
        self.output_size = 3 if self.target == '3d' and (not self.ignore_depth) else 2
        self.centralize = getattr(cfg, 'centralize', False)
        
        # input feature extraction
        input_dim = 0
        if 'rgb' in self.modalities:
            # frame-level feature extraction
            self.global_feature = FrameEmbedding(cfg, input_size=input_size)
            input_dim += self.embed_dim
        if 'loc' in self.modalities:
            # location to feature
            self.location_feature = MlpModel(self.output_size, self.loc_dim // 2, self.loc_dim)
            input_dim += self.loc_dim
        
        # state transition model: pθ(z1:T | x1:C )
        self.transition = ProbTranV3(self.embed_dim, self.loc_dim, self.modalities, **self.ssm_params)
        # transformer encoders
        self.encoder_vis = Encoder(num_frames=seq_len, embed_dim=self.embed_dim, **self.transformer_params)
        self.encoder_traj = Encoder(num_frames=seq_len, embed_dim=self.loc_dim, **self.transformer_params)

        out_layer = nn.Tanh if self.centralize else nn.Sigmoid
        dim_pred = self.loc_dim
        self.predict = MlpHead(dim_pred, dim_pred // 2, self.output_size, out_layer=out_layer)

        if self.aleatoric_uncertainty:
            self.predict_uncertainty = UncertaintyHead(dim_pred, dim_pred //2, cfg=cfg.aleatoric_uncertainty)
        
        if self.pred_velocity:
            self.predict_velo = MlpHead(dim_pred, dim_pred // 2, self.output_size, bn_layer=nn.LayerNorm, out_layer=nn.Tanh)
    
    
    def _get_rgb_features(self, x, num_ratios, mask):
        """
            x: (B, C, T, H, W)
            num_ratios: ()
            mask: (B*N, T)
        """
        rgb_input = get_rgb_observed(x, num_ratios, mask)  # (B*N, C, T, H, W)
        batch_size, max_frames = rgb_input.size(0), x.size(2)
        # get global appearance features
        rgb_input = rearrange(rgb_input, 'b c t h w -> (b t) c h w')
        rgb_feat = self.global_feature(rgb_input)  # (B*T, D)
        rgb_feat = rearrange(rgb_feat, '(b t) d -> b t d', b=batch_size, t=max_frames)  # (B*N, T, D)
        return rgb_feat


    def _get_loc_features(self, traj, num_ratios, mask):
        """
            traj: (B, T, 3)
            num_ratios: ()
            mask: (B*N, T)
        """
        traj_input = get_traj_observed(traj, num_ratios, mask)  # (B*N, T, 3)
        loc_feat = self.location_feature(traj_input)  # (B*N, T, D)
        return loc_feat


    def get_input_features(self, x, traj_all, num_ratios, mask):
        # get rgb & trajectory features of full observation
        x_feat = self._get_rgb_features(x, num_ratios, mask)  # (B*N, T, D)
        if 'loc' in self.modalities:
            y_feat = self._get_loc_features(traj_all, num_ratios, mask)  # (B*N, T, D)
            x_feat = torch.cat([x_feat, y_feat], dim=-1)
        return x_feat
    

    def forward(self, x, nframes, ratios, traj):
        """
            x: (B, C, T, H, W)
            nframes: (B,)
            ratios: (B, N)
            traj: (B, T, 3) or (B, T*(T+1)/2, 3)
        """
        batch_size, max_frames, num_ratios = x.size(0), traj.size(1), ratios.size(1)
        device = x.device
        # get trajectory target
        traj_all = traj[:, :, :self.output_size]

        # get masks of observed and unobserved frames
        mask_o, mask_u, last_obs_frames = get_masks(batch_size, ratios, max_frames, nframes, device)  # (B, N, T)
        mask_o = rearrange(mask_o, 'b n t -> (b n) t')  # (B*N, T)

        vis_feat, traj_feat = None, None
        if 'rgb' in self.modalities:
            # visual feature
            vis_feat = self._get_rgb_features(x, num_ratios, mask_o)  # (B*N, T, D)
            vis_feat = self.encoder_vis(vis_feat, mask_o)
        if 'loc' in self.modalities:
            # trajectory feature
            traj_feat = self._get_loc_features(traj_all, num_ratios, mask_o)
            traj_feat = self.encoder_traj(traj_feat, mask_o)
        # state transition: x_obs (+ y_obs) --> x_future
        future_pred, pred_dict = self.transition(mask_o, vis_obs=vis_feat, traj_obs=traj_feat)

        # trajectory prediction
        preds = self.predict(future_pred)  # (B*N, T, 3)
        pred_dict.update({'traj': preds})

        # uncertainty prediction
        if self.aleatoric_uncertainty:
            uncts = self.predict_uncertainty(future_pred)  # (B*N, T, 3)
            pred_dict.update({'unct': uncts})
        
        # velocity prediction
        if self.pred_velocity:
            velos = self.predict_velo(future_pred.detach())  # (B*N, T, 3)
            pred_dict.update({'velo': velos})
        
        """ ----------- parse prediction results ----------- """
        outputs, outputs_obs = gather_results(pred_dict, nframes, ratios, last_obs_frames, self.sup_obs)
        
        return outputs, outputs_obs
    

    def inference(self, x, nframes, ratios, traj=None):
        return self(x, nframes, ratios, traj)

    

src/models/TransformerSSM.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
*
* NOTICE:
* Code slightly modified from original repo: https://github.com/stevenlsw/hoi-forecast/tree/master/networks,
* which is protected under the MIT License: https://github.com/stevenlsw/hoi-forecast/blob/master/LICENSE.
* We sincerely thank the contributors of the original repo.
*
"""

import torch
import torch.nn as nn

import os, sys
HoiForecast_ROOT = os.path.join(os.path.dirname(__file__), '../../third_party/hoi-forecast')
sys.path.append(HoiForecast_ROOT)
""" 
The followings are hoi-forecast classes
"""
from networks.net_utils import DropPath
from networks.embedding import PositionalEncoding
from networks.layer import MultiHeadAttention, Mlp



def get_pad_mask(seq, pad_idx=0):
    if not isinstance(pad_idx, int):
        raise TypeError("<pad_index> has to be an int!")
    if seq.dim() == 3:
        return seq != pad_idx  # (B, T, T)
    return (seq != pad_idx).unsqueeze(1) # (B, 1, T)


class EncoderBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,
                 drop_path=0.1, act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super(EncoderBlock, self).__init__()
        self.norm1 = norm_layer(dim)
        self.attn = MultiHeadAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias,
                                       qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)


    def forward(self, x, mask=None):
        src_mask = None
        if mask is not None:
            src_mask = get_pad_mask(mask, pad_idx=0)  # (B, 1, T)
        x_norm = self.norm1(x)
        x = x + self.drop_path(self.attn(q=x_norm, k=x_norm, v=x_norm, mask=src_mask))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class Encoder(nn.Module):
    """ Transformer Encoder """
    def __init__(self, num_frames, embed_dim=512, depth=6, num_heads=8, mlp_ratio=4.,
                 qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
                 drop_path_rate=0.1, norm_layer=nn.LayerNorm, dropout=0.):
        super(Encoder, self).__init__()
        # dropout layer
        self.dropout = nn.Dropout(dropout)
        # temporal positional encoding
        self.time_embed = PositionalEncoding(embed_dim, max_len=num_frames)
        self.time_drop = nn.Dropout(p=drop_rate)
        
        # construct encoder blocks with increasing drop path rates
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
        self.encoder_blocks = nn.ModuleList([EncoderBlock(
            dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
            drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer
        ) for i in range(depth)])
        
        # layer normalization
        self.norm = norm_layer(embed_dim)
        
    
    def forward(self, x, mask=None):
        """ x: (B, T, D)
            mask: # (B, T)
        """
        x = self.time_embed(x)
        x = self.time_drop(x)
        
        for blk in self.encoder_blocks:
            x = blk(x, mask=mask)  # (B, T, D)
        
        x = self.norm(x)
        return x

src/models/transformer_layers.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch
from einops import rearrange



def get_masks(batch_size, ratios, max_frames, nframes, device):
    """Get the multi-hot masks for observed and unobserved frames """
    num_ratios = ratios.size(1)
    mask_o = torch.zeros((batch_size, num_ratios, max_frames)).to(device, non_blocking=True)  # (B, N, T)
    mask_u = torch.zeros((batch_size, num_ratios, max_frames)).to(device, non_blocking=True)  # (B, N, T)
    last_frames = torch.zeros((batch_size, num_ratios)).long()
    for b in range(batch_size):  # loop for each video clip
        # get the number observed frames and full frames
        num_full = int(nframes[b])
        num_obs = torch.floor(num_full * ratios[b]).to(torch.long)  # (2,)
        last_frames[b] = num_obs - 1  # index
        # get masks
        for i, n_o in enumerate(num_obs):
            mask_o[b, i, :n_o] = 1
            mask_u[b, i, n_o: num_full] = 1
    return mask_o, mask_u, last_frames


def get_rgb_observed(rgb, num_ratios, mask_o):
    """ get the observed part of RGB frames"""
    rgb_input = rgb[:, None].repeat(1, num_ratios, 1, 1, 1, 1)
    rgb_input = rearrange(rgb_input, 'b n c t h w -> (b n) c t h w')  # (B*N, C, T, H, W)
    rgb_mask = mask_o[:, None, :].repeat(1, rgb.size(1), 1).unsqueeze(-1).unsqueeze(-1)  # (B*N, C, T, 1, 1)
    rgb_input = rgb_input * rgb_mask
    return rgb_input  # (B*N, C, T, H, W)


def get_traj_observed(traj_all, num_ratios, mask_o):
    """ get the observed part of trajectory"""
    traj_input = traj_all[:, None].repeat(1, num_ratios, 1, 1)  # (B, N, T, 3)
    traj_input = rearrange(traj_input, 'b n t c-> (b n) t c')  # (B*N, T, 3)
    traj_mask = mask_o.unsqueeze(-1)  # (B*N, T, 1)
    traj_input = traj_input * traj_mask
    return traj_input  # (B*N, T, 3)


def gather_results(pred_dict, nframes, ratios, last_obs_frames, sup_obs=False):
    """ pred_dict: 'traj': (B*N, T, 3), 'kld': (B*N, T)
        nframes: (B,)
        ratios: (B, N)
        last_obs_frames: (B, N)
    """
    def split_dict(pred_dict, start, end):
        result = dict()
        for k, v in pred_dict.items():
            result.update({k: v[p, start: end]})
        return result
        
    batch_size, num_ratios = ratios.size()
    outputs = [dict() for _ in range(batch_size)]
    outputs_obs = [dict() for _ in range(batch_size)] if sup_obs else None
    for b in range(batch_size):
        for n in range(num_ratios):
            p = b * num_ratios + n
            key = '{:.1f}'.format(float(ratios[b, n]))  # we use ratio string as the key
            # unobserved part
            outputs[b][key] = split_dict(pred_dict, start=last_obs_frames[b, n]+1, end=int(nframes[b]))
            # observed part
            if sup_obs:
                outputs_obs[b][key] = split_dict(pred_dict, start=0, end=last_obs_frames[b, n]+1)
    return outputs, outputs_obs
src/models/utils.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
*
* NOTICE:
* Function `CosineAnnealingWarmupRestarts_New` is from the original repo: https://github.com/sgvaze/osr_closed_set_all_you_need/blob/main/utils/schedulers.py,
* which is protected under the MIT License: https://github.com/sgvaze/osr_closed_set_all_you_need/blob/main/LICENSE.md
* We sincerely thank the contributors of the original repo.
*
"""

import torch
import math


class CosineAnnealingWarmupRestarts_New(torch.optim.lr_scheduler.CosineAnnealingWarmRestarts):

    def __init__(self, warmup_epochs, lr_decay=1.0, lr_decay_epochs=None, *args, **kwargs):

        super(CosineAnnealingWarmupRestarts_New, self).__init__(*args, **kwargs)

        # Init optimizer with low learning rate
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.eta_min

        self.warmup_epochs = warmup_epochs
        self.lr_decay = lr_decay
        self.lr_decay_epochs = lr_decay_epochs

        self.warmup_lrs = []
        for base_lr in self.base_lrs:
            # Get target LR after warmup is complete
            target_lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * warmup_epochs / self.T_i)) / 2

            # Linearly interpolate between minimum lr and target_lr
            linear_step = (target_lr - self.eta_min) / self.warmup_epochs
            warmup_lrs = [self.eta_min + linear_step * (n + 1) for n in range(warmup_epochs)]
            self.warmup_lrs.append(warmup_lrs)

    def step(self, epoch=None):

        # Called on super class init
        if epoch is None:
            super(CosineAnnealingWarmupRestarts_New, self).step(epoch=epoch)

        else:
            if epoch < self.warmup_epochs:
                for i, param_group in enumerate(self.optimizer.param_groups):
                    param_group['lr'] = self.warmup_lrs[i][epoch]

                # Fulfill misc super() funcs
                self.last_epoch = math.floor(epoch)
                self.T_cur = epoch
                self._last_lr = [group['lr'] for group in self.optimizer.param_groups]

            else:
                # decay the base_lr which will be used in the step() call for cosine restart
                if self.lr_decay_epochs and epoch + 1 in self.lr_decay_epochs:
                    self.base_lrs = [lr * self.lr_decay for lr in self.base_lrs]
                super(CosineAnnealingWarmupRestarts_New, self).step(epoch=epoch)


def get_optimizer(cfg, params):
    # get optimizer
    if cfg.optimizer == 'SGD':
        optimizer = torch.optim.SGD(params, lr=cfg.base_lr, momentum=cfg.momentum)
    elif cfg.optimizer == 'Adam':
        optimizer = torch.optim.Adam(params, lr=cfg.base_lr, weight_decay=cfg.weight_decay)
    elif cfg.optimizer == 'Adan':
        from adan_pytorch import Adan
        optimizer = Adan(params, lr=cfg.base_lr, weight_decay=cfg.weight_decay, betas=cfg.betas)
    else:
        raise NotImplementedError
    
    return optimizer


def get_scheduler(cfg, optimizer):
    # get learning rate scheduler
    if cfg.scheduler.type == 'step':
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, cfg.scheduler.step_size, gamma=cfg.scheduler.gamma)
    elif cfg.scheduler.type == 'cosine':
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=cfg.epoch, last_epoch=-1, eta_min=0)
    elif cfg.scheduler.type == 'cosine_warmup':
        scheduler = CosineAnnealingWarmupRestarts_New(warmup_epochs=cfg.scheduler.warmup_epoch,
                                                      lr_decay=cfg.scheduler.lr_decay,
                                                      lr_decay_epochs=cfg.scheduler.lr_decay_epoch, 
                                                      optimizer=optimizer, 
                                                      T_0=int(cfg.epoch / (cfg.scheduler.num_restarts + 1)),
                                                      eta_min=cfg.base_lr * cfg.scheduler.min_lr_factor)
    else:
        raise NotImplementedError
    return scheduler
    
src/optimizers.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

from tqdm import tqdm
import numpy as np
from src.utils import random_ratios_batch, get_test_ratios, send_to_gpu
import src.utils as utils
import src.utils_h2o as utils_h2o



def train_epoch(cfg, model, traindata_loader, criterion, optimizer, writer, epoch):

    model.train()
    loss = 0
    pbar = tqdm(total=len(traindata_loader), ncols=0, desc='train epoch {}/{}'.format(epoch + 1, cfg.TRAIN.epoch))
    for batch_id, batch_data in enumerate(traindata_loader):

        filename, clip, odometry, nframes, traj_gt = batch_data
        # send data to device 
        clip, traj_gt = send_to_gpu((clip, traj_gt), cfg.device, non_blocking=True)
        
        # generate a batch of random observation ratios
        ratios = random_ratios_batch(nframes, cfg.TRAIN)

        # run inference
        outputs = model(clip, nframes, ratios, traj=traj_gt)

        # compute losses
        losses = criterion(outputs, nframes, traj_gt)
        loss += losses['total_loss']

        # backward
        optimizer.zero_grad()
        losses['total_loss'].backward()
        optimizer.step()

        # write the lr and losses
        iter_cur = epoch * len(traindata_loader) + batch_id
        writer.add_scalar('misc/lr', optimizer.state_dict()['param_groups'][0]['lr'], iter_cur)
        for k, v in losses.items():
            writer.add_scalars('train/{}'.format(k), {k: v}, iter_cur)  # draw loss curves in different figures
        
        pbar.set_postfix({"train loss": losses['total_loss'].item()})
        pbar.update()
    
    loss /= len(traindata_loader)
    pbar.close()

    return loss


def eval_epoch(cfg, model, valdata_loader, criterion, writer, epoch):
    
    model.eval()

    all_preds, all_gts = {}, {}
    pbar = tqdm(total=len(valdata_loader), ncols=0, desc='eval epoch {}'.format(epoch + 1))
    for batch_id, batch_data in enumerate(valdata_loader):
        
        _, clip, _, nframes, traj_gt = batch_data
        # send data to device 
        clip, traj_gt = send_to_gpu((clip, traj_gt), cfg.device, non_blocking=True)
        
        # fixed set of observation ratios
        ratios = get_test_ratios(cfg.TEST.ratios, nframes)
        
        # run inference
        outputs = model(clip, nframes, ratios, traj=traj_gt)

        # compute losses
        losses = criterion(outputs, nframes, traj_gt)

        # write the lr and losses
        iter_cur = epoch * len(valdata_loader) + batch_id
        for k, v in losses.items():
            writer.add_scalars('test/{}'.format(k), {k: v}, iter_cur)

        # gather unobserved predictions and ground truths 
        preds, gts = utils.gather_eval_results(outputs[0], nframes, traj_gt, 
                                                ignore_depth=model.ignore_depth, centralize=model.centralize)
        if batch_id == 0:
            all_preds.update(preds)
            all_gts.update(gts)
        else:
            for r in list(preds.keys()):
                all_preds[r].extend(preds[r])
                all_gts[r].extend(gts[r])
        
        pbar.set_postfix({"val loss": losses['total_loss'].item()})
        pbar.update()
    
    for r in list(all_preds.keys()):
        all_preds[r] = np.vstack(all_preds[r])  # (N, 3)
        all_gts[r] = np.vstack(all_gts[r])  # (N, 3)

    pbar.close()
    
    return all_preds, all_gts


def eval_h2o_epoch(cfg, model, valdata_loader, criterion, writer, epoch):

    model.eval()

    all_preds, all_gts, all_poses = {}, {}, {}
    pbar = tqdm(total=len(valdata_loader), ncols=0, desc='eval epoch {}'.format(epoch + 1))
    for batch_id, batch_data in enumerate(valdata_loader):
        
        filename, clip, campose, nframes, traj_gt = batch_data
        # send data to device 
        clip, traj_gt = send_to_gpu((clip, traj_gt), cfg.device, non_blocking=True)
        
        # fixed set of observation ratios
        ratios = get_test_ratios(cfg.TEST.ratios, nframes)
        
        # run inference
        outputs = model(clip, nframes, ratios, traj=traj_gt)

        # compute losses
        losses = criterion(outputs, nframes, traj_gt)

        # write the lr and losses
        iter_cur = epoch * len(valdata_loader) + batch_id
        for k, v in losses.items():
            writer.add_scalars('test/{}'.format(k), {k: v}, iter_cur)

        # gather unobserved predictions and ground truths 
        preds, gts, poses = utils_h2o.gather_eval_results(outputs[0], nframes, traj_gt, campose, ignore_depth=model.ignore_depth, 
                                                    use_global=cfg.MODEL.use_global, centralize=model.centralize)

        if batch_id == 0:
            all_preds.update(preds)
            all_gts.update(gts)
            all_poses.update(poses)
        else:
            for r in list(preds.keys()):
                all_preds[r].extend(preds[r])
                all_gts[r].extend(gts[r])
                all_poses[r].extend(poses[r])
        
        pbar.set_postfix({"val loss": losses['total_loss'].item()})
        pbar.update()
    
    for r in list(all_preds.keys()):
        all_preds[r] = np.vstack(all_preds[r])  # (N, 3)
        all_gts[r] = np.vstack(all_gts[r])  # (N, 3)
        all_poses[r] = np.concatenate(all_poses[r])  # (N, 4, 4)
    
    pbar.close()

    return all_preds, all_gts, all_poses



def run_test(cfg, model, data_loader):

    all_preds, all_gts = {}, {}
    for batch_id, batch_data in tqdm(enumerate(data_loader), total=len(data_loader), desc='Run testing'):
        
        filename, clip, odometry, nframes, traj_gt = batch_data
        # send data to device 
        clip, traj_gt = send_to_gpu((clip, traj_gt), cfg.device, non_blocking=True)
        
        # fixed set of observation ratios
        ratios = get_test_ratios(cfg.TEST.ratios, nframes)
        
        # run inference
        outputs = model.inference(clip, nframes, ratios, traj=traj_gt)

        # gather unobserved predictions and ground truths 
        # Note: the preds and gts are in Global 3D Space for target == '3d, or in pixel space of the 1st frame
        preds, gts = utils.gather_eval_results(outputs[0], nframes, traj_gt, 
                                            ignore_depth=model.ignore_depth, centralize=model.centralize)
        if batch_id == 0:
            all_preds.update(preds)
            all_gts.update(gts)
        else:
            for r in list(preds.keys()):
                all_preds[r].extend(preds[r])
                all_gts[r].extend(gts[r])
    
    for r in list(all_preds.keys()):
        all_preds[r] = np.vstack(all_preds[r])  # (N, 3)
        all_gts[r] = np.vstack(all_gts[r])  # (N, 3)
        
    return all_preds, all_gts


def run_h2o_test(cfg, model, data_loader):

    all_preds, all_gts, all_poses = {}, {}, {}
    for batch_id, batch_data in tqdm(enumerate(data_loader), total=len(data_loader), desc='Run testing'):
        
        filename, clip, campose, nframes, traj_gt = batch_data
        # send data to device 
        clip, traj_gt = send_to_gpu((clip, traj_gt), cfg.device, non_blocking=True)
        
        # fixed set of observation ratios
        ratios = get_test_ratios(cfg.TEST.ratios, nframes)
        
        # run inference
        outputs = model.inference(clip, nframes, ratios, traj=traj_gt)

        # gather unobserved predictions and ground truths 
        # Note: the preds and gts are in Global 3D Space for target == '3d, or in pixel space of the 1st frame
        preds, gts, poses = utils_h2o.gather_eval_results(outputs[0], nframes, traj_gt, campose, ignore_depth=model.ignore_depth, 
                                                    use_global=cfg.MODEL.use_global, centralize=model.centralize)
        if batch_id == 0:
            all_preds.update(preds)
            all_gts.update(gts)
            all_poses.update(poses)
        else:
            for r in list(preds.keys()):
                all_preds[r].extend(preds[r])
                all_gts[r].extend(gts[r])
                all_poses[r].extend(poses[r])
    
    for r in list(all_preds.keys()):
        all_preds[r] = np.vstack(all_preds[r])  # (N, 3)
        all_gts[r] = np.vstack(all_gts[r])  # (N, 3)
        all_poses[r] = np.concatenate(all_poses[r])  # (N, 4, 4)
        
    return all_preds, all_gts, all_poses


def print_eval_results(writer, all_ades_val, all_fdes_val, all_ades_train=None, all_fdes_train=None, epoch=0, loss_train=0):
    mADE_val = np.mean(list(all_ades_val.values()))
    mFDE_val = np.mean(list(all_fdes_val.values()))
    writer.add_scalars('test/DE_val', {'mADE': mADE_val, 'mFDE': mFDE_val}, epoch)
    info = "==> [Epoch {}]: Loss(train) = {:.9f}, mADE(val) = {:.3f}, mFDE(val) = {:.3f}".format(epoch + 1, loss_train, mADE_val, mFDE_val)
    if all_ades_train and all_fdes_train:
        mADE_train = np.mean(list(all_ades_train.values()))
        mFDE_train = np.mean(list(all_fdes_train.values()))
        writer.add_scalars('test/DE_train', {'mADE': mADE_train, 'mFDE': mFDE_train}, epoch)
        info += ', mADE(train) = {:.3f}, mFDE(train)'.format(mADE_train, mFDE_train)
    print(info)

src/runtime.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import random
import os
import numpy as np
import torch
from copy import deepcopy
from functools import reduce
import cv2
import imageio
import matplotlib.pyplot as plt
from matplotlib.animation import FFMpegWriter



def set_deterministic(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if you are using multi-GPU.
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True


def send_to_gpu(data_tuple, device, non_blocking=False):
    gpu_tensors = []
    for item in data_tuple:
        gpu_tensors.append(item.to(device, non_blocking=non_blocking))
    return tuple(gpu_tensors)


def get_depth_anchors(traj_gt, num_obs, num_full, offset_anchor='Zc', time_orders=None):
    if offset_anchor == 'Zc':
        # use the last observed depth as the offset anchor
        anchors = traj_gt[num_obs-1: num_obs, 2:3]
    else:
        zs_obs = traj_gt[:num_obs, 2:3]
        if isinstance(zs_obs, np.ndarray):
            zs_obs = torch.from_numpy(zs_obs).to(time_orders.device)
        # use historical (z, t) values to fit a four-order multinomial model
        solution = torch.linalg.lstsq(time_orders[:num_obs, :], zs_obs).solution  # (5,1)
        anchors = torch.matmul(time_orders[num_obs:num_full, :], solution)
        if isinstance(traj_gt, np.ndarray):
            anchors = anchors.cpu().numpy()
    return anchors


def gather_eval_results(outputs, nframes, trajectories, ignore_depth=False, centralize=False, return_unct=False):
    """ gather predictions and gts for evaluation purpose
        return: all_preds (N, 9, 3)
                all_gts (N, 9, 3)
    """
    # constant camera parameters
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    max_depth = 3
    num_order = 3
    
    trajectories = trajectories.cpu().numpy()
    all_preds, all_gts = {}, {}
    if return_unct: all_uncts = {}
    target = '3d' if trajectories.shape[-1] == 3 else '2d'
    # parse results
    for b, ratio_preds in enumerate(outputs):
        # get the number observed frames and full frames
        num_full = nframes[b]
        traj_gt = trajectories[b]  # (T,3)
        if ignore_depth:
            traj_gt = traj_gt[:, :2]
        if centralize:
            traj_gt += 0.5
        traj_gt = denormalize(traj_gt, target, intrinsics, max_depth)
        
        for i, (r, results) in enumerate(ratio_preds.items()):
            num_obs = torch.floor(num_full * float(r)).to(torch.long)
            preds = results['traj'].cpu().numpy()
            gts = traj_gt[num_obs:num_full, :]
            # denormalize
            if centralize:
                preds += 0.5
            preds = denormalize(preds, target, intrinsics, max_depth)
            # get uncertainty if any
            if return_unct:
                logvars = results['unct'].cpu().numpy() if 'unct' in results else np.ones_like(preds)
                uncts = np.exp(logvars)
            
            # gather results
            if r in all_preds and r in all_gts:
                all_preds[r].append(preds)  # (M, 3)
                all_gts[r].append(gts)  # (M, 3)
                if return_unct: all_uncts[r].append(uncts)
            else:
                all_preds[r] = [preds]
                all_gts[r] = [gts]
                if return_unct: all_uncts[r] = [uncts]
    if return_unct: 
        return all_preds, all_uncts, all_gts
    return all_preds, all_gts


def XYZ_to_uv(traj3d, intrinsics):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points 
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['ox'])
    traj2d[:, 1] = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['oy'])
    # clip the coordinates 
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['w'])
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['h'])
    traj2d = np.floor(traj2d).astype(np.int32)
    return traj2d


def compute_cle(all_preds, all_gts):
    """Compute the Center Location Error (CLE)"""
    all_cles = dict()
    for r in list(all_gts.keys()):
        average_cle = np.mean(np.sum((all_preds[r] - all_gts[r])**2, axis=-1))  # (9,)
        all_cles[r] = average_cle
    
    # compute the mean of results from 9 observation ratios
    mCLE = np.mean(list(all_cles.values()))
    
    return all_cles, mCLE


def normalize_2d(uv, intrinsics):
    uv_norm = np.copy(uv).astype(np.float32)
    uv_norm[:, 0] /= intrinsics['w']
    uv_norm[:, 1] /= intrinsics['h']
    return uv_norm


def compute_displacement_errors(all_preds, all_gts, target='3d', eval_space='3d', intrinsics=None):
    """Compute the Displacement Errors (ADE and FDE)"""
    all_ades, all_fdes = dict(), dict()
    for r in list(all_gts.keys()):
        preds, gts = all_preds[r], all_gts[r]
        
        if target == '3d' and eval_space == '2d':
            preds = XYZ_to_uv(preds, intrinsics)
            gts = XYZ_to_uv(gts, intrinsics)
        
        if target == '3d' and eval_space == 'norm2d':
            preds = normalize_2d(XYZ_to_uv(preds, intrinsics), intrinsics)
            gts = normalize_2d(XYZ_to_uv(gts, intrinsics), intrinsics)
        
        if target == '2d' and eval_space == 'norm2d':
            preds = normalize_2d(preds, intrinsics)
            gts = normalize_2d(gts, intrinsics)

        displace_errors = np.sqrt(np.sum((preds - gts)**2, axis=-1))  # (Tu,)
        # ADE score
        all_ades[r] = np.mean(displace_errors)
        # FDE score
        all_fdes[r] = displace_errors[-1]
    
    return all_ades, all_fdes


def compute_block_distances(all_preds, all_gts, target='3d', eval_space='3d', intrinsics=None):
    """Compute the block distances along x, y, and z dimensions"""
    all_dxs, all_dys, all_dzs = dict(), dict(), dict()
    for r in list(all_gts.keys()):
        preds, gts = all_preds[r], all_gts[r]
        
        if target == '3d' and eval_space == '2d':
            preds = XYZ_to_uv(preds, intrinsics)
            gts = XYZ_to_uv(gts, intrinsics)
        
        if target == '3d' and eval_space == 'norm2d':
            preds = normalize_2d(XYZ_to_uv(preds, intrinsics), intrinsics)
            gts = normalize_2d(XYZ_to_uv(gts, intrinsics), intrinsics)
        
        if target == '2d' and eval_space == 'norm2d':
            preds = normalize_2d(preds, intrinsics)
            gts = normalize_2d(gts, intrinsics)
        
        # delta X
        all_dxs[r] = np.mean(np.fabs(preds[:, 0] - gts[:, 0]))
        # delta Y
        all_dys[r] = np.mean(np.fabs(preds[:, 1] - gts[:, 1]))
        if preds.shape[-1] == 3:
            # delta Z
            all_dzs[r] = np.mean(np.fabs(preds[:, 2] - gts[:, 2]))
    return all_dxs, all_dys, all_dzs


def first_last_nonzero_indices(mask, dim=1):
    """ Find the first and the last non-zero indices 
        mask: (B, T)
    """
    assert mask.sum(dim=dim).all(), "mask should not contain all-zero rows!"
    mask_cumsum = mask.cumsum(dim=dim)
    first_nz = (mask_cumsum == 0).sum(dim=dim)
    r2l_cumsum = mask + torch.sum(mask, dim=dim, keepdims=True) - torch.cumsum(mask, dim=dim)
    last_nz = mask.size(dim) - (r2l_cumsum == 0).sum(dim=dim) - 1
    return first_nz, last_nz


def get_subsequent_mask(seq_len, device, diagonal=1):
    """ Get the sequential mask for transformer decoder
        seq_len: int
        return: (1, seq_len, seq_len) bottom-triangular bool matrix
    """
    mask = (1 - torch.triu(torch.ones((1, seq_len, seq_len), device=device), diagonal=diagonal)).bool()
    return mask


def random_ratios_batch(nframes, cfg, min_obs=2):
    """ Generate a batch of random ratios
    """
    if hasattr(cfg, 'ratios'):  # fixed ratio list
        ratios_batch = torch.tensor([cfg.ratios], dtype=torch.float32).repeat(nframes.size(0), 1)
        return ratios_batch
    # random selection 
    assert cfg.num_ratios <= 9 and cfg.num_ratios >= 1, "invalid number of ratios!"
    ratios_batch = torch.zeros((nframes.size(0), cfg.num_ratios), dtype=torch.float32)
    for b, len_valid in enumerate(nframes):
        ratios = []
        while len(ratios) != cfg.num_ratios:
            r = np.random.choice(np.arange(1, 10)/10, 1, replace=False)[0]
            if r * int(len_valid) >= min_obs and r not in ratios:  # ignore ratios that produce too short clips
                ratios.append(r)
        ratios_batch[b] = torch.tensor(ratios, dtype=torch.float32)
    return ratios_batch


def get_test_ratios(test_ratios, nframes):
    """ test_ratios: list of decimal numbers
        nframes: (B,)
    """
    valid_ratios = [r for r in test_ratios if all(nframes * r >= 1)]
    ratios = torch.tensor([valid_ratios]).repeat(nframes.size(0), 1)  # (B, M)
    return ratios


def save_the_latest(data, ckpt_file, topK=3, ignores=[]):
    """ Only keeping the latest topK checkpoints.
    """
    # find the existing checkpoints in a sorted list
    folder = os.path.dirname(ckpt_file)
    num_exist = len(os.listdir(folder))
    if num_exist >= topK + len(ignores):
        # remove the old checkpoints
        ext = ckpt_file.split('.')[-1]
        all_ckpts = list(filter(lambda x: x.endswith('.' + ext), os.listdir(folder)))
        all_epochs = [int(filename.split('.')[-2].split('_')[-1]) for filename in all_ckpts]
        fids = np.argsort(all_epochs)  # transformer_90.pth
        # iteratively remove
        for i in fids[:(num_exist - topK + 1)]:
            if all_epochs[i] in ignores:
                continue
            file_to_remove = os.path.join(folder, all_ckpts[i])
            if os.path.isfile(file_to_remove):
                os.remove(file_to_remove)
    torch.save(data, ckpt_file)


def read_video(video_file, toRGB=True):
    assert os.path.exists(video_file), "File does not exist! {}".format(video_file)
    cap = cv2.VideoCapture(video_file)
    success, frame = cap.read()
    video = []
    while success:
        if toRGB:
            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        video.append(frame)
        success, frame = cap.read()
    video = np.array(video)
    return video


def read_odometry(odom_file):
    assert os.path.exists(odom_file), "File does not exist! {}".format(odom_file)
    all_transforms = np.load(odom_file)  # (T, 4, 4)
    return all_transforms


def len_quad(max_len, i):
    return i * max_len - int(i * (i - 1) / 2)


def write_video(mat, video_file, fps=5):
    video_writer = cv2.VideoWriter(video_file, cv2.VideoWriter_fourcc(*'mp4v'), fps, (mat.shape[2], mat.shape[1]))
    for frame in mat:
        video_writer.write(frame)

def vis_frame_traj(frame, traj2d):
    cv2.polylines(frame, [traj2d], False, (0, 255, 255))
    for t, uv in enumerate(traj2d):
        cv2.circle(frame, uv, radius=5, color=(0, 0, 255), thickness=-1)

def video_to_gif(video, giffile, fps=5.0, toBGR=False):
    assert giffile.endswith('.gif')
    with imageio.get_writer(giffile, mode='I', duration=1.0/fps) as writer:
        for frame in video:
            frame_vis = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) if toBGR else np.copy(frame)
            writer.append_data(frame_vis)


def global_to_local(traj3d, odometry, refer_idx=0, start_idx=1):
    """ Transform 3D trajectory of global frame to local frame
        traj3d: (Tu, 3)
    """
    assert start_idx >= refer_idx
    length = traj3d.shape[0]
    traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
    traj3d_local = []
    for i, pt_homo in enumerate(traj3d_homo):
        # collect transformation matrices from the first frame to the current fram
        trans = reduce(np.dot, odometry[refer_idx: i + start_idx - refer_idx])
        trans_inv = np.dot(np.linalg.inv(np.dot(trans.T, trans)), trans.T)  # (A^TxA)^-1xA, (4, 4)
        pt_local = trans_inv.dot(pt_homo)[:3]
        traj3d_local.append(pt_local)
    traj3d_local = np.array(traj3d_local)
    return traj3d_local


def transform_to_local(traj3d_global, odometry, num_obs=0):
    traj3d_local = []
    for i, loc3d in enumerate(traj3d_global):
        t = num_obs + i
        # camera_0 to camera_t (Note: all T-t points are in the same camera_0 ,thus their odometry are ignored by setting to identity matrix)
        odom_new = np.copy(odometry)
        odom_new[(t+1):] = np.repeat(np.eye(4)[None, :, :], odometry.shape[0]-t-1, axis=0)
        # global to local transform
        loc3d_local = global_to_local(loc3d[None, :], odom_new, start_idx=t+1)
        traj3d_local.append(loc3d_local)
    traj3d_local = np.concatenate(traj3d_local, axis=0)
    return traj3d_local


def normalize(traj, target='3d', intrinsics=None, vis_ratio=0.25, max_depth=3.0):
    width = intrinsics['w'] * vis_ratio
    height = intrinsics['h'] * vis_ratio
    if target == '2d':
        traj[:, 0] /= width
        traj[:, 1] /= height
    elif target == '3d':
        # transform the (X, Y, Z) into the normalized (u,v,z) by PINHOLE camera model
        traj[:, 0] = (traj[:, 0] * intrinsics['fx'] / traj[:, 2] + intrinsics['ox']) / intrinsics['w']
        traj[:, 1] = (traj[:, 1] * intrinsics['fy'] / traj[:, 2] + intrinsics['oy']) / intrinsics['h']
        traj[:, 2] = traj[:, 2] / max_depth
    traj[:, 0] = np.clip(traj[:, 0], a_min=0, a_max=(width-1)/width)
    traj[:, 1] = np.clip(traj[:, 1], a_min=0, a_max=(height-1)/height)
    return traj


def denormalize(traj, target='3d', intrinsics=None, max_depth=3.0):
    """ transform the normalized (u,v,depth) to (X,Y,Z)
    """
    u = traj[:, 0] * intrinsics['w']
    v = traj[:, 1] * intrinsics['h']
    traj2d = np.stack((u, v), axis=1)
    if target == '3d':
        Z = traj[:, 2] * max_depth
        X = (u - intrinsics['ox']) * Z / intrinsics['fx']
        Y = (v - intrinsics['oy']) * Z / intrinsics['fy']
        traj3d = np.stack((X, Y, Z), axis=1)
        return traj3d
    return traj2d


def output_transform(outputs, trajectories, nframes, odometry, 
        intrinsics=None, vis_ratio=0.25, max_depth=3.0, centralize=False, use_odom=False, ignore_depth=False):
    """ outputs: predictions of observed and unobserved (torch cuda tensor)
        trajectories: (B, T, 3), ndarray, normalized
    """
    outputs_unobs, outputs_obs = deepcopy(outputs[0]), deepcopy(outputs[1])
    trajectory_local = np.copy(trajectories)
    if centralize:
        trajectory_local += 0.5  # in range [0,1]
    for b, ratio_preds in enumerate(outputs[0]):
        # get the number observed frames and full frames
        num_full = int(nframes[b])
        traj_gt = trajectory_local[b, :num_full]  # (T,3)
        
        for i, (r, pred_dict) in enumerate(ratio_preds.items()):
            num_obs = int(num_full * float(r))
            preds = pred_dict['traj'].cpu().numpy()
            if centralize:
                preds += 0.5
            if ignore_depth and traj_gt.shape[1] == 3:
                preds = np.hstack((preds, traj_gt[num_obs:, 2:3]))  # (Tu, 3), compensate the depth for prediction
            
            # global to local
            if use_odom:
                # de-normalize & 2D-3D
                preds = denormalize(preds, target='3d', intrinsics=intrinsics, max_depth=max_depth)
                # transform from global (camera_0) to local (camera_t) in 3D space
                preds = transform_to_local(preds, odometry[b, :num_full], num_obs)
                # normalize back
                preds = normalize(preds, target='3d', intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth)
            outputs_unobs[b][r]['traj'] = preds  # in [0,1]

            if outputs_obs is not None:
                preds_obs = outputs_obs[b][r]['traj'].cpu().numpy()
                if centralize:
                    preds_obs += 0.5
                if ignore_depth and traj_gt.shape[1] == 3:
                    preds_obs = np.hstack((preds_obs, traj_gt[:num_obs, 2:3]))  # (To, 3)
                if use_odom:
                    preds_obs = denormalize(preds_obs, target='3d', intrinsics=intrinsics, max_depth=max_depth)
                    preds_obs = transform_to_local(preds_obs, odometry[b, :num_full])
                    preds_obs = normalize(preds_obs, target='3d', intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth)
                outputs_obs[b][r]['traj'] = preds_obs

        if use_odom:
            traj_gt = denormalize(traj_gt, target='3d', intrinsics=intrinsics, max_depth=max_depth)
            traj_gt = transform_to_local(traj_gt, odometry[b, :num_full])
            trajectory_local[b, :num_full] = normalize(traj_gt, target='3d', intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth)

    return (outputs_unobs, outputs_obs), trajectory_local


def draw_trajectory_3d_with_gt(pred3d, gt3d, video_size=[270, 480], \
    ratio=0.5, fps=5, dpi=30, views=[30, -120], minmax=None, savefile='temp.mp4'):

    if minmax is None:
        all_points = np.concatenate([pred3d, gt3d], axis=0)
        xmin, ymin, zmin = np.min(all_points, axis=0).tolist()
        xmax, ymax, zmax = np.max(all_points, axis=0).tolist()
    else:
        xmin, ymin, zmin = minmax[0]
        xmax, ymax, zmax = minmax[1]
    fig_w, fig_h = int(video_size[1] // dpi), int(video_size[0] // dpi)

    # init a 3D figure
    fig = plt.figure(figsize=(fig_w, fig_h))
    ax = fig.add_subplot(projection='3d')
    ax.view_init(views[0], views[1])  # elev: 30, azim: -120
    # Make a 3D quiver plot
    ox, oy, oz = np.zeros((3,3))
    dx, dy, dz = np.array([[xmax-xmin, 0, 0], [0, zmax-zmin, 0], [0, 0, ymax-ymin]]) * 0.2
    ax.quiver(ox, oy, oz, dx, dy, dz)
    # Setup coordinate system
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(zmin, zmax)
    ax.set_zlim(ymin, ymax)
    ax.invert_zaxis()
    ax.set_xlabel('X')  # x: rightward
    ax.set_ylabel('Z')  # z: innerward
    ax.set_zlabel('Y')  # y: downward
    plt.locator_params(axis='x', nbins=5)
    plt.locator_params(axis='y', nbins=5)
    plt.locator_params(axis='z', nbins=5)
    plt.tight_layout()

    # init video writer
    writer = FFMpegWriter(fps=fps, metadata=dict(title='3D Trajectory', artist='Matplotlib'))
    future_start = int(len(pred3d) * ratio)
    # process
    with writer.saving(fig, savefile, dpi=dpi):
        pred_past, pred_future = [], []
        gt_past, gt_future = [], []
        for t, (pred, gt) in enumerate(zip(pred3d, gt3d)):
            if t <= future_start:
                pred_past.append(pred)
                gt_past.append(gt)
            else:
                pred_future.append(pred)
                gt_future.append(gt)
            
            if t == future_start:  # connect the future and past
                pred_future.append(pred)
                gt_future.append(gt)

            # draw past ground truth trajectory
            xs, ys, zs = np.array(gt_past)[:, 0], np.array(gt_past)[:, 1], np.array(gt_past)[:, 2]
            ax.plot(xs, zs, ys, '-o', mfc='green', c='blue', fillstyle='full', mec='green', mew=0.0)
            # draw past predicted trajectory
            xs, ys, zs = np.array(pred_past)[:, 0], np.array(pred_past)[:, 1], np.array(pred_past)[:, 2]
            ax.plot(xs, zs, ys, '-o', mfc='blue', c='blue', fillstyle='full', mec='blue', mew=0.0)

            if t > future_start:
                # draw future ground truth trajectory
                xs, ys, zs = np.array(gt_future)[:, 0], np.array(gt_future)[:, 1], np.array(gt_future)[:, 2]
                ax.plot(xs, zs, ys, '-o', mfc='limegreen', c='limegreen', fillstyle='full', mec='limegreen', mew=0.0)
                # draw future predicted trajectory
                xs, ys, zs = np.array(pred_future)[:, 0], np.array(pred_future)[:, 1], np.array(pred_future)[:, 2]
                ax.plot(xs, zs, ys, '-o', mfc='red', c='red', fillstyle='full', mec='red', mew=0.0)
            # grab a frame
            writer.grab_frame()
    plt.close()

src/utils.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import torch
from src.H2OLoader import normalize_traj, denormalize_traj, world_to_camera, XYZ_to_uv, INTRINSICS
import numpy as np


def gather_eval_results(outputs, nframes, trajectories, cam_poses, ignore_depth=False, use_global=False, centralize=False, return_unct=False):
    """ gather predictions and gts for evaluation purpose
        return: all_preds (N, 9, 3)
                all_gts (N, 9, 3)
    """
    trajectories = trajectories.cpu().numpy()
    cam_poses = cam_poses.numpy()
    all_preds, all_gts, all_poses = {}, {}, {}
    if return_unct: all_uncts = {}
    target = '3d' if trajectories.shape[-1] == 3 else '2d'
    # parse results
    for b, ratio_preds in enumerate(outputs):
        # get the number observed frames and full frames
        num_full = nframes[b]
        traj_gt = trajectories[b]  # (T,3)
        campose = cam_poses[b]
        if ignore_depth: traj_gt = traj_gt[:, :2]

        # de-centralize and de-normalize
        if centralize: traj_gt = (traj_gt + 1.0) * 0.5  # in [0, 1]
        traj_gt = denormalize_traj(traj_gt, target=target, use_global=use_global)
        
        for i, (r, results) in enumerate(ratio_preds.items()):
            num_obs = torch.floor(num_full * float(r)).to(torch.long)
            preds = results['traj'].cpu().numpy()
            gts = traj_gt[num_obs:num_full, :]
            poses = campose[num_obs:num_full]  # (N,4,4)

            # de-centralize and de-normalize
            if centralize: preds = (preds + 1.0) * 0.5  # in [0, 1]
            preds = denormalize_traj(preds, target=target, use_global=use_global)
            # get uncertainty if any
            if return_unct:
                logvars = results['unct'].cpu().numpy() if 'unct' in results else np.ones_like(preds)
                uncts = np.exp(logvars)
            
            # gather results
            if r in all_preds and r in all_gts:
                all_preds[r].append(preds)  # (M, 3)
                all_gts[r].append(gts)  # (M, 3)
                all_poses[r].append(poses)
                if return_unct: all_uncts[r].append(uncts)
            else:
                all_preds[r] = [preds]
                all_gts[r] = [gts]
                all_poses[r] = [poses]
                if return_unct: all_uncts[r] = [uncts]
    if return_unct: 
        return all_preds, all_uncts, all_gts, all_poses
    return all_preds, all_gts, all_poses


def traj_transform(traj, cam2world=None, target='3d', eval_space='3d', use_global=True):
    """ transform trajectory space
    """
    if target == '3d' and eval_space in ['2d', 'norm2d']:
        if use_global:
            traj = world_to_camera(traj, cam2world)
        # camera 3D to 2D
        traj = XYZ_to_uv(traj, INTRINSICS)

        if eval_space == 'norm2d':
            traj = normalize_traj(traj, target='2d')
    
    if target == '2d' and eval_space == 'norm2d':
        traj = normalize_traj(traj, target='2d')
    
    return traj



def compute_displacement_errors(all_preds, all_gts, all_camposes, target='3d', eval_space='3d', use_global=True):
    """Compute the Displacement Errors (ADE and FDE)"""
    all_ades, all_fdes = dict(), dict()
    for r in list(all_gts.keys()):
        preds, gts, cam2world = all_preds[r], all_gts[r], all_camposes[r]
        
        # transform trajectory
        preds = traj_transform(preds, cam2world, target, eval_space, use_global)
        gts = traj_transform(gts, cam2world, target, eval_space, use_global)

        displace_errors = np.sqrt(np.sum((preds - gts)**2, axis=-1))  # (Tu,)
        # ADE score
        all_ades[r] = np.mean(displace_errors)
        # FDE score
        all_fdes[r] = displace_errors[-1]
    
    return all_ades, all_fdes


def compute_block_distances(all_preds, all_gts, all_camposes, target='3d', eval_space='3d', use_global=True):
    """Compute the block distances along x, y, and z dimensions"""
    all_dxs, all_dys, all_dzs = dict(), dict(), dict()
    for r in list(all_gts.keys()):
        preds, gts, cam2world = all_preds[r], all_gts[r], all_camposes[r]
        
        # transform trajectory
        preds = traj_transform(preds, cam2world, target, eval_space, use_global)
        gts = traj_transform(gts, cam2world, target, eval_space, use_global)
        
        # delta X
        all_dxs[r] = np.mean(np.fabs(preds[:, 0] - gts[:, 0]))
        # delta Y
        all_dys[r] = np.mean(np.fabs(preds[:, 1] - gts[:, 1]))
        if preds.shape[-1] == 3:
            # delta Z
            all_dzs[r] = np.mean(np.fabs(preds[:, 2] - gts[:, 2]))
    return all_dxs, all_dys, all_dzs
src/utils_h2o.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import numpy as np
import torch
from terminaltables import AsciiTable


def load_checkpoint(cfg, model):
    ckpt_dir = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'snapshot')
    if cfg.TEST.test_epoch is None:  # the lastest epoch
        all_ckpts = list(filter(lambda x: x.endswith('.pth'), os.listdir(ckpt_dir)))
        all_epochs = [int(filename.split('.')[-2].split('_')[-1]) for filename in all_ckpts]  # transformer_90.pth
        fids = np.argsort(all_epochs)
        ckpt_file = os.path.join(ckpt_dir, all_ckpts[fids[-1]])
        test_epoch = all_epochs[fids[-1]]
    else:
        assert isinstance(cfg.TEST.test_epoch, int)
        ckpt_file = os.path.join(ckpt_dir, cfg.TRAIN.snapshot_prefix + '%02d.pth'%(cfg.TEST.test_epoch))
        test_epoch = cfg.TEST.test_epoch
    print('Loading the model checkpoint: {}'.format(ckpt_file))
    checkpoint = torch.load(ckpt_file)
    model.load_state_dict(checkpoint['model'])
    return model, test_epoch


def print_de_table(all_ades, all_fdes, subset):
    mADE = np.mean(list(all_ades.values()))
    mFDE = np.mean(list(all_fdes.values()))
    print("mADE({}) = {:.3f}, mFDE({}) = {:.3f}".format(subset, mADE, subset, mFDE))
    # print as table
    display = [list(all_ades.keys()), ["{:.3f} / {:.3f}".format(v1, v2) for v1, v2 in zip(all_ades.values(), all_fdes.values())]]
    table = AsciiTable(display)
    table.inner_footing_row_border = True
    table.justify_columns = {i: 'center' for i in range(len(all_ades))}
    print(table.table)


def print_delta_table(all_dxs, all_dys, all_dzs, subset):
    mDX = np.mean(list(all_dxs.values()))
    mDY = np.mean(list(all_dys.values()))
    print_str = "mDX({}) = {:.3f}, mDY({}) = {:.3f}".format(subset, mDX, subset, mDY)
    if len(all_dzs) > 0:
        mDZ = np.mean(list(all_dzs.values()))
        print_str += ", mDZ({}) = {:.3f}".format(subset, mDZ)
    print(print_str)
    # print as table
    if len(all_dzs) > 0:
        display = [list(all_dxs.keys()), ["{:.3f} / {:.3f} / {:.3f}".format(v1, v2, v3) 
                                      for v1, v2, v3 in zip(all_dxs.values(), all_dys.values(), all_dzs.values())]]
    else:  # 2D targets
        display = [list(all_dxs.keys()), ["{:.3f} / {:.3f}".format(v1, v2) 
                                      for v1, v2 in zip(all_dxs.values(), all_dys.values())]]
    table = AsciiTable(display)
    table.inner_footing_row_border = True
    table.justify_columns = {i: 'center' for i in range(len(all_dxs))}
    print(table.table)
src/utils_io.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import numpy as np
import torch
import cv2
import imageio
from src.utils import output_transform, read_video, draw_trajectory_3d_with_gt



def draw_uv_points(video, preds, start=0, color=(0, 255, 0)):
    """ video: (T, H, W, C)
        preds: (Tm, 2/3)
    """
    height, width = video.shape[1:3]
    for t, loc in enumerate(preds):
        u, v = int(loc[0] * width), int(loc[1] * height)  # normalized loc to real loc
        cv2.circle(video[start + t], (u, v), radius=5, color=color, thickness=-1)
    return video


def video_to_gif(video, giffile, toBGR=False):
    assert giffile.endswith('.gif')
    with imageio.get_writer(giffile, mode='I', duration=0.2) as writer:
        for frame in video:
            if toBGR:
                frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            writer.append_data(frame)


def remove_the_old(result_path, topk=5, ignores=[]):
    files = [name for name in os.listdir(result_path) if name.endswith('.gif')]
    if len(files) >= topk + len(ignores):
        all_epochs = [int(filename.split('_')[0].split('epoch')[-1]) for filename in files]  # epoch20_bathroomCabinet_7_c40r0.6.gif
        fids = np.argsort(all_epochs)
        # iteratively remove
        for i in fids[:(len(files) - topk + 1)]:
            if all_epochs[i] in ignores:
                continue
            file_to_remove = os.path.join(result_path, files[i])
            if os.path.isfile(file_to_remove):
                os.remove(file_to_remove)


def vis_demo(video, outputs, nframes, trajectories, result_prefix=None, obs_ratios=['0.2', '0.5', '0.8'], toBGR=True, return_video=False):
    outputs_unobs, outputs_obs = outputs
    vis_video = {}
    for r in obs_ratios:
        num_full = int(nframes[0])  # only one video
        traj_gt = trajectories[0, :num_full]  # (Tm,3)
        video_vis = np.copy(video[:num_full])
        
        # visualize ground truth
        video_vis = draw_uv_points(video_vis, traj_gt, start=0, color=(0, 0, 255))  # blue
        
        # visualize predictions of the observed part
        if outputs_obs is not None:
            preds = outputs_obs[0][str(r)]['traj']
            video_vis = draw_uv_points(video_vis, preds, start=0, color=(255, 255, 0))  # yellow
        
        # visualize predictions of the future part
        preds = outputs_unobs[0][str(r)]['traj']
        num_obs = int(num_full * float(r))
        video_vis = draw_uv_points(video_vis, preds, start=num_obs-1, color=(0, 0, 0))  # black
        
        if not return_video:
            # save as GIF file
            video_to_gif(video_vis, result_prefix + 'r{}.gif'.format(r), toBGR=toBGR)
        else:
            vis_video[r] = video_vis
    return vis_video


def vis_traj3d(outputs, traj3d_gt, obs_ratios, nframes, canvas_size, result_prefix=None, show_past_pred=False, return_video=False):
    """ Visualize results in 3D space
    """
    outputs_unobs, outputs_obs = outputs
    vis_video = {}
    result_dir = os.path.dirname(result_prefix) if result_prefix is not None else './'
    for r in obs_ratios:
        preds_future = outputs_unobs[0][str(r)]['traj']
        # get the observed part
        num_full = int(nframes[0])  # only one video
        num_obs = int(num_full * float(r))
        # preds_past = outputs_obs[0][r]['traj'] if outputs_obs is not None else traj3d_gt[0, :num_obs]
        preds_past = outputs_obs[0][str(r)]['traj'] if outputs_obs is not None and show_past_pred else traj3d_gt[0, :num_obs]
        # complete prediction
        traj3d_pred = np.concatenate([preds_past, preds_future], axis=0)

        # draw 3D
        temp_file = os.path.join(result_dir, 'temp_pred.mp4')
        draw_trajectory_3d_with_gt(traj3d_pred, traj3d_gt[0, :num_full], canvas_size, ratio=float(r),
            fps=5, dpi=90, views=[30, -120], minmax=None, savefile=temp_file)
        video_traj3d = read_video(temp_file)
        
        if not return_video:
            video_to_gif(video_traj3d, result_prefix + 'r{}.gif'.format(r))
        else:
            vis_video[r] = video_traj3d

        if os.path.exists(temp_file):
            os.remove(temp_file)
    return vis_video


def visualize_example(train_loader, result_dir, epoch, num=1, target='3d'):
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}
    means = np.array([[[0.5, 0.5, 0.5]]])  # (1,1,3)
    stds = np.array([[[0.5, 0.5, 0.5]]])  # (1,1,3)
    
    for i, (clip, nf, traj) in enumerate(train_loader):
        if i >= num:
            break  # only visualize 1 batches (clips)
        u = torch.floor(traj[0, :, 0] * clip.size(-1)).to(torch.long)  # (T,)
        v = torch.floor(traj[0, :, 1] * clip.size(-2)).to(torch.long)  # (T,)
        # visualize
        vis_mat = []
        for t in range(nf[0]):
            frame = clip[0, :, t].permute(1, 2, 0).contiguous().numpy()  # (224, 224, 3)
            frame = ((frame * stds + means) * 255).astype(np.uint8)  # broadcast add & mul 
            cv2.circle(frame, (int(u[t]), int(v[t])), radius=5, color=(0, 255, 0), thickness=-1)
            vis_mat.append(frame)
        vis_mat = np.array(vis_mat)  # (T, 224, 224, 3)
        # save to GIF figure
        video_to_gif(vis_mat, os.path.join(result_dir, 'vis_epoch{}_clip{}.gif'.format(epoch, i)))


def run_runtime_demo(cfg, model, data_batch, obs_ratios=['0.2', '0.5', '0.8'], result_folder='demo_train', prefix_name='epoch0_', topk=5):
    # input data (a training sample)
    filename, input_data, odometry, nframes, gt_traj = data_batch[0][0], data_batch[1][0:1], data_batch[2][0:1], data_batch[3][0:1], data_batch[4][0:1]
    scene_record = filename[:-4].split('/')[-2]
    clip_name = filename[:-4].split('/')[-1].split('_')[0]
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    vis_ratio = 0.25
    max_depth = 3
    
    # directly read the original video
    video = read_video(filename)
    
    # result path
    result_path = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, result_folder)
    os.makedirs(result_path, exist_ok=True)
    # remove the old figs
    remove_the_old(result_path, topk=topk, ignores=getattr(cfg.TRAIN.scheduler, 'lr_decay_epoch', []))  # only save demos from the latest topK models
    
    # inference
    with torch.no_grad():
        ratios = torch.tensor([[float(r) for r in obs_ratios]]).repeat(nframes.size(0), 1)  # (B, M)
        outputs = model.inference(input_data.cuda(), nframes, ratios, traj=gt_traj.cuda())
    
    # transform from global to local coordinate frame
    outputs, gt_traj = output_transform(outputs, gt_traj.numpy(), nframes.numpy(), odometry.numpy(), 
                                        intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth, 
                                        centralize=model.centralize, use_odom=model.use_odom, ignore_depth=model.ignore_depth)
    
    # visualize
    result_prefix = os.path.join(result_path, prefix_name + '{}_{}'.format(scene_record, clip_name))
    vis_demo(video, outputs, nframes, gt_traj, result_prefix, obs_ratios=obs_ratios)

    if model.target == '3d':
        # visualize 3D trajectory
        result_3d_path = os.path.join(result_path, 'traj3d')
        os.makedirs(result_3d_path, exist_ok=True)
        remove_the_old(result_3d_path, topk=topk, ignores=getattr(cfg.TRAIN.scheduler, 'lr_decay_epoch', []))  # only save demos from the latest topK models

        canvas_size = (intrinsics['h'] * vis_ratio * 0.5, intrinsics['w'] * vis_ratio * 0.5)
        result_prefix = os.path.join(result_3d_path, prefix_name + '{}_{}'.format(scene_record, clip_name))
        vis_traj3d(outputs, gt_traj, obs_ratios, nframes, canvas_size, result_prefix)
src/utils_vis.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import numpy as np
import os
import torch
import importlib

from src.config import parse_configs
from src.EgoPAT3DLoader import build_dataloaders
from src.utils import set_deterministic, compute_displacement_errors, compute_block_distances 
from src.runtime import run_test
from src.utils_io import load_checkpoint, print_de_table, print_delta_table


def get_test_results(model, test_loader, result_file):
    if not os.path.exists(result_file):
        # run test inference
        with torch.no_grad():
            all_preds, all_gt = run_test(cfg, model, test_loader)
        # save predictions
        np.savez(result_file[:-4], pred=all_preds, gt=all_gt)
    else:
        print("Result file exists. Loaded from file: %s."%(result_file))
        all_results = np.load(result_file, allow_pickle=True)
        all_preds, all_gt = all_results['pred'][()], all_results['gt'][()]
    return all_preds, all_gt


def test(cfg, intrinsics):

    # build test dataloaders
    print("Loading dataset...")
    test_loader, testnovel_loader = build_dataloaders(cfg, phase='test')

    # build the model
    model_module = importlib.import_module('src.models.{}'.format(cfg.MODEL.arch))
    model = getattr(model_module, cfg.MODEL.arch)(cfg.MODEL, seq_len=cfg.DATA.max_frames, input_size=cfg.DATA.transform.input_size[0])
    model = model.to(device=cfg.device)
    # load checkpoints
    model, test_epoch = load_checkpoint(cfg, model)
    model = model.eval()
    
    # result folder
    result_path = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'test-e{}'.format(test_epoch))
    os.makedirs(result_path, exist_ok=True)
    eval_space = getattr(cfg.TEST, 'eval_space', '3d')

    ### test on the seen scenes
    print("Seen secenes...")
    all_preds, all_gt = get_test_results(model, test_loader, os.path.join(result_path, 'test_seen.npz'))
    # evaluation
    all_ades_seen, all_fdes_seen = compute_displacement_errors(all_preds, all_gt, 
                                                               target=model.target, eval_space=eval_space, intrinsics=intrinsics)
    all_dxs_seen, all_dys_seen, all_dzs_seen = compute_block_distances(all_preds, all_gt, 
                                                               target=model.target, eval_space=eval_space, intrinsics=intrinsics)

    # print tables
    print_de_table(all_ades_seen, all_fdes_seen, subset='Seen')
    print_delta_table(all_dxs_seen, all_dys_seen, all_dzs_seen, subset='Seen')


    ### test on the unseen scenes
    print("\nUnseen secenes...")
    all_preds, all_gt = get_test_results(model, testnovel_loader, os.path.join(result_path, 'test_unseen.npz'))
    # evaluation
    all_ades_unseen, all_fdes_unseen = compute_displacement_errors(all_preds, all_gt, 
                                                               target=model.target, eval_space=eval_space, intrinsics=intrinsics)
    all_dxs_unseen, all_dys_unseen, all_dzs_unseen = compute_block_distances(all_preds, all_gt, 
                                                               target=model.target, eval_space=eval_space, intrinsics=intrinsics)
    
    print_de_table(all_ades_unseen, all_fdes_unseen, subset='Unseen')
    print_delta_table(all_dxs_unseen, all_dys_unseen, all_dzs_unseen, subset='Unseen')

    print("\nDone!")


if __name__ == '__main__':
    # parse input arguments
    cfg = parse_configs(phase='test')

    # fix random seed 
    set_deterministic(cfg.seed)
    
    # constant camera parameters
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code

    # test
    test(cfg, intrinsics)
test.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import argparse
import yaml
from easydict import EasyDict
import numpy as np
import cv2
import pickle
import torch
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from gluoncv.torch.data.transforms.videotransforms.video_transforms import Compose, Resize, Normalize
from gluoncv.torch.data.transforms.videotransforms.volume_transforms import ClipToTensor
from src.utils import set_deterministic, output_transform, read_video
from src.utils_vis import vis_demo, vis_traj3d, video_to_gif
import importlib
import random
from functools import reduce
from src.utils_io import load_checkpoint



def read_traj_file(filepath, target='3d'):
    assert os.path.exists(filepath), "File does not exist! {}".format(filepath)
    with open(filepath, 'rb') as f:
        data = pickle.load(f)
    traj2d = data['traj2d']
    traj3d = None
    if target == '3d':
        assert 'traj3d' in data, "3D trajectories do not exist in file: {}".format(filepath)
        traj3d = data['traj3d']
    num_preserve = data['num_preserve'] if 'num_preserve' in data else len(traj2d)
    return traj2d, traj3d, num_preserve


def read_odometry(filepath):
    assert os.path.exists(filepath), "File does not exist! {}".format(filepath)
    all_transforms = np.load(filepath)  # (T, 4, 4)
    return all_transforms


def get_projected_traj3d(traj3d, odometry):
    """ traj3d: (T, 3)
        odometry: (T, 4, 4)
    """
    length = len(odometry)
    traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
    all_traj3d_proj = []
    for i in range(length):
        # compute the transformed points of all future points
        traj3d_proj = [traj3d[i]]  # initial 3D point
        for j in range(i + 1, length):
            odom = reduce(np.dot, odometry[(i+1):(j+1)])  # (4, 4)
            future_point = odom.dot(traj3d_homo[j].T)  # (4,)
            traj3d_proj.append(future_point[:3])
        all_traj3d_proj.append(np.array(traj3d_proj))
    all_traj3d_proj = np.concatenate(all_traj3d_proj, axis=0)  # (T*(T+1)/2, 3)
    return all_traj3d_proj


def _XYZ_to_uv(traj3d, intrinsics=None, vis_ratio=0.25):
    width = intrinsics['w'] * vis_ratio
    height = intrinsics['h'] * vis_ratio
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    u = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['ox']) * vis_ratio
    v = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['oy']) * vis_ratio
    u = np.clip(u, a_min=0, a_max=width-1)
    v = np.clip(v, a_min=0, a_max=height-1)
    traj2d = torch.stack((u, v), dim=-1)
    return traj2d


def _normalize(traj, intrinsics=None, vis_ratio=0.25, max_depth=3.0, target='3d', centralize=False):
    traj_new = traj.clone()
    width = intrinsics['w'] * vis_ratio
    height = intrinsics['h'] * vis_ratio
    if target == '2d':
        traj_new[:, 0] = traj[:, 0] / width
        traj_new[:, 1] = traj[:, 1] / height
    elif target == '3d':
        # transform the (X, Y, Z) into the normalized (u,v,z) by PINHOLE camera model
        traj2d = _XYZ_to_uv(traj, intrinsics, vis_ratio)
        traj_new[:, 0] = traj2d[:, 0] / width
        traj_new[:, 1] = traj2d[:, 1] / height
        traj_new[:, 2] = traj[:, 2] / max_depth
    if centralize:
        traj_new -= 0.5
    return traj_new


def random_select_data(cfg):
    ### randomly select one clip as an example
    data_root = os.path.join(cfg.DATA.data_path, cfg.DATA.dataset)
    scene_test_splits = {'1': ['9', '10'], '2': ['9', '10'], '3': ['9', '10'], '4': ['9', '10'], '5': ['8', '9'], '6': ['8', '9'], 
                            '7': ['9', '10'], '9': ['9', '10'], '10': ['9', '10'], '11': ['9', '10'], '12': ['9', '10']}
    if hasattr(cfg.DATA, 'scenes'):
        scene_test_splits = dict(filter(lambda elem: elem[0] in cfg.DATA.scenes, scene_test_splits.items()))
    # random select a scene
    scene_id = random.choice(list(scene_test_splits))
    # random select a record
    record_id = random.choice(scene_test_splits[scene_id])
    record_path = os.path.join(data_root, 'EgoPAT3D-postproc', 'trajectory_repair', scene_id)
    record_name = list(filter(lambda x: x.split('_')[-1] == record_id, os.listdir(record_path)))[0]
    # random select a clip
    clip_path = os.path.join(record_path, record_name)
    all_clips = list(filter(lambda x: x.endswith('.pkl'), os.listdir(clip_path)))
    clip_name = random.choice(all_clips)[:-4]
    return {'scene': scene_id, 'record': record_name, 'clip': clip_name}


def prepare_demo_data(example, cfg, intrinsics=None, vis_ratio=0.25, max_depth=3.0):
    
    data_root = os.path.join(cfg.DATA.data_path, cfg.DATA.dataset)
    path_template = os.path.join(data_root, 'EgoPAT3D-postproc', '{}', example['scene'], example['record'], example['clip'] + '.{}')

    # prepare ground truth
    traj_file = path_template.format('trajectory_repair', 'pkl')
    traj2d, traj3d, num_preserve = read_traj_file(traj_file, target=cfg.MODEL.target)
    traj2d = torch.from_numpy(traj2d).to(torch.float32)
    traj3d = torch.from_numpy(traj3d).to(torch.float32)
    traj = traj3d if cfg.MODEL.target == '3d' else traj2d
    
    # read video
    video_file = path_template.format('video_clips_hand', 'mp4')
    video = read_video(video_file)  # (T, H, W, C), RGB data
    video = video[:num_preserve]

    # read odometry
    odom_data = np.eye(4, dtype=np.float32)[None, :, :].repeat(num_preserve, axis=0)  # identity matrix
    use_odom = getattr(cfg.MODEL, 'use_odom', False)
    if use_odom:
        odom_file = path_template.format('odometry', 'npy')
        odom_data = read_odometry(odom_file)
        odom_data = odom_data[:num_preserve]  # (T, 4, 4)
    
    # preprocessing
    transform_test = Compose([Resize(cfg.DATA.transform.input_size), 
                              ClipToTensor(), 
                              Normalize(mean=cfg.DATA.transform.means, std=cfg.DATA.transform.stds)])
    input_data = transform_test(video)   # (C, T, h, w)
    max_frames = traj.size(0) if cfg.DATA.max_frames < 0 else cfg.DATA.max_frames
    len_valid = min(input_data.size(1), max_frames)

    input_data = torch.zeros([input_data.size(0), max_frames, input_data.size(2), input_data.size(3)], dtype=input_data.dtype)
    input_data[:, :len_valid] = input_data[:, :len_valid]
    input_data = input_data.unsqueeze(0)
    
    nframes = torch.tensor(len_valid).unsqueeze(0)
    
    odometry = torch.eye(4)[None, :, :].repeat(max_frames, 1, 1).to(torch.float32)
    gt_traj = torch.zeros([max_frames, traj.size(1)], dtype=traj.dtype)
    traj_valid = traj[:len_valid]  # (T, 3)
    if use_odom:
        odom_valid = odom_data[:len_valid]
        odometry[:len_valid] = torch.from_numpy(odom_valid)
        # project all future hand locations to each of current frames
        traj3d_valid_all = get_projected_traj3d(traj3d[:len_valid].numpy(), odom_valid)  # (T*(T+1)/2, 3)
        traj3d_valid = torch.from_numpy(traj3d_valid_all[:len_valid])  # only use the trajectory in the first frame
        traj_valid = _XYZ_to_uv(traj3d_valid, intrinsics, vis_ratio) if cfg.MODEL.target == '2d' else traj3d_valid

    traj_valid = _normalize(traj_valid, intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth,
                target=cfg.MODEL.target, centralize=getattr(cfg.MODEL, 'centralize', False))
    gt_traj[:len_valid] = traj_valid
    gt_traj = gt_traj.unsqueeze(0)
    odometry = odometry.unsqueeze(0)
    
    return video, input_data, odometry, nframes, gt_traj


def run_demo(cfg, sample=None, ratios=None):

    # result path
    result_path = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'demo')
    os.makedirs(result_path, exist_ok=True)

    # constant camera parameters
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    vis_ratio = 0.25
    max_depth = 3

    # input data
    if sample is None:
        sample = random_select_data(cfg=cfg)
    video, input_data, odometry, nframes, gt_traj = prepare_demo_data(sample, cfg, intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth)
    
    # input observation ratios
    if ratios is None:
        ratios = cfg.TEST.ratios
    test_ratios = torch.tensor([[float(r) for r in ratios]]).repeat(nframes.size(0), 1)  # (B, M)
    
    # load model
    model_module = importlib.import_module('src.models.{}'.format(cfg.MODEL.arch))
    model = getattr(model_module, cfg.MODEL.arch)(cfg.MODEL, seq_len=cfg.DATA.max_frames, input_size=cfg.DATA.transform.input_size[0])
    model = model.to(device=cfg.device)
    model, epoch = load_checkpoint(cfg, model)
    model = model.eval()

    # inference
    with torch.no_grad():
        outputs = model.inference(input_data.cuda(), nframes, test_ratios, traj=gt_traj.cuda())
    
    # transform from global to local coordinate frame
    outputs, gt_traj = output_transform(outputs, gt_traj.numpy(), nframes.numpy(), odometry.numpy(), 
                                        intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth, 
                                        centralize=model.centralize, use_odom=model.use_odom, ignore_depth=model.ignore_depth)
    # visualize
    prefix_name = '{}_{}_{}_epoch{}_'.format(sample['scene'], sample['record'], sample['clip'].split('_')[0], epoch)
    result_prefix = os.path.join(result_path, prefix_name)
    vis_demo(video, outputs, nframes, gt_traj, result_prefix, obs_ratios=ratios)

    # visualize 3D trajectory
    result_3d_path = os.path.join(result_path, 'traj3d')
    os.makedirs(result_3d_path, exist_ok=True)

    canvas_size = (intrinsics['h'] * vis_ratio * 0.5, intrinsics['w'] * vis_ratio * 0.5)
    result_prefix = os.path.join(result_3d_path, prefix_name)
    vis_traj3d(outputs, gt_traj, ratios, nframes, canvas_size, result_prefix)



def recover_video(input_data, cfg, video_size=(540, 960)):
    """ input_data: (C, T, H, W)
    """
    means = np.array([[cfg.DATA.transform.means]])  # (1, 1, 3)
    stds = np.array([[cfg.DATA.transform.stds]])  # (1, 1, 3)
    num_frames = input_data.size(1)
    video = []
    for t in range(num_frames):
        frame = input_data[:, t].permute(1, 2, 0).contiguous().numpy()  # (H, W, 3)
        frame = (frame * stds + means) * 255  # broadcast add & mul 
        frame = cv2.resize(frame, dsize=(video_size[1], video_size[0]))
        frame = cv2.cvtColor(frame.astype(np.uint8), cv2.COLOR_RGB2BGR)
        video.append(frame)
    video = np.array(video, dtype=np.uint8)  # (T, 224, 224, 3)
    return video
    

def parse_configs():
    parser = argparse.ArgumentParser()
    parser.add_argument('--config', type=str, default='./config/baseline.yml',
                        help='The relative path of dataset.')
    parser.add_argument('--gpus', type=str, default="0", 
                        help="The delimited list of GPU IDs separated with comma. Default: '0'.")
    parser.add_argument('--num_workers', type=int, default=4, metavar='N',
                        help='The number of workers to load dataset. Default: 4')
    parser.add_argument('--tag', type=str, default='default',
                        help='The tag to save model results')
    args = parser.parse_args()
    with open(args.config, 'r') as f:
        cfg = EasyDict(yaml.safe_load(f))
    cfg.update(vars(args))
    # add device
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    cfg.update(device=device)
    cfg.MODEL.update(device=device)
    # add root path
    root_path = os.path.dirname(__file__)
    cfg.update(root_path=root_path)

    cfg.DATA.load_all = False

    return cfg
    

if __name__ == '__main__':
    
    # parse input arguments
    cfg = parse_configs()
    
    # fix random seed 
    set_deterministic(cfg.seed)
    
    # demo_sample = {'scene': '1', 'record': 'bathroomCabinet_9', 'clip': 'c1_s64_e103_sx718_sy420_ex631_ey247'}
    demo_sample = {'scene': '1', 'record': 'bathroomCabinet_8', 'clip': 'c100_s5336_e5369_sx586_sy367_ex518_ey346'}
    run_demo(cfg, sample=demo_sample)
tools/demo.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import numpy as np
import os
import torch
import importlib
from src.config import parse_configs
from src.EgoPAT3DLoader import build_dataloaders
from src.utils import set_deterministic, output_transform, read_video, get_test_ratios, video_to_gif, denormalize
from src.utils_vis import vis_traj3d
from src.utils_io import load_checkpoint
from tqdm import tqdm
import cv2
import pickle
from functools import reduce
from copy import deepcopy



def overlay_video(vis2d, vis3d=None):
    if vis3d is None:
        return vis2d
    video_vis = np.copy(vis2d)
    for i, (frame2d, frame3d) in enumerate(zip(vis2d, vis3d)):
        # overlay
        bottom_left = frame2d[frame2d.shape[0] - frame3d.shape[0]:, :frame3d.shape[1], :]
        frame2d[frame2d.shape[0] - frame3d.shape[0]:, :frame3d.shape[1], :] = cv2.addWeighted(bottom_left, 0.4, frame3d, 0.6, 0)  # bottom left
        video_vis[i] = frame2d
    return video_vis


def overlay_frame(frame2d, frame3d):
    frame = np.copy(frame2d)
    bottom_left = frame2d[frame2d.shape[0] - frame3d.shape[0]:, :frame3d.shape[1], :]
    frame[frame2d.shape[0] - frame3d.shape[0]:, :frame3d.shape[1], :] = cv2.addWeighted(bottom_left, 0.4, frame3d, 0.6, 0)  # bottom left
    return frame


def compute_mse(outputs, traj_gt, num_full, future=False):
    mse = 0
    for i, (r, preds) in enumerate(outputs.items()):
        num_obs = torch.floor(num_full * float(r)).to(torch.long).numpy()
        gts = traj_gt[num_obs: num_full] if future else traj_gt[:num_obs]
        mse += np.mean(np.sqrt(np.sum((preds['traj'] - gts)**2, axis=-1)))
    mse /= (i+1)
    return float(mse)




def XYZ_to_uv(traj3d, intrinsics):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points 
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['ox'])
    traj2d[:, 1] = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['oy'])
    # clip the coordinates 
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['w'])
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['h'])
    traj2d = np.floor(traj2d).astype(np.int32)
    return traj2d


def vis_frame_traj(frame, traj2d, traj_start=0, future_start=15, 
                   TRAJ_COLOR=(255, 255, 0), PAST_COLOR=(255, 0, 0), FUTURE_COLOR=(0, 0, 255)):
    cv2.polylines(frame, [traj2d], False, TRAJ_COLOR, thickness=2)
    for t, uv in enumerate(traj2d):
        color = PAST_COLOR if traj_start + t <= future_start else FUTURE_COLOR
        cv2.circle(frame, uv, radius=5, color=color, thickness=-1)


def visualize_transformed_traj(video, traj3d, odometry, intrinsics, ratio=0.5,
                               TRAJ_COLOR=(255, 255, 0), PAST_COLOR=(255, 0, 0), FUTURE_COLOR=(0, 0, 255)):
    """ video: (T, H, W, 3)
        traj3d: (T, 3)
        odometry: (T, 4, 4)
    """
    length = len(odometry)
    traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
    future_start = int(length * ratio)
    for i in range(length):
        # compute the transformed points of all future points
        traj_list = [traj3d[i]]  # initial 3D point
        for j in range(i + 1, length):
            odom = reduce(np.dot, odometry[(i+1):(j+1)])  # (4, 4)
            future_point = odom.dot(traj3d_homo[j].T)  # (4,)
            traj_list.append(future_point[:3])
        # visualize current frame
        traj2d = XYZ_to_uv(np.array(traj_list), intrinsics)
        vis_frame_traj(video[i], traj2d, traj_start=i, future_start=future_start, 
                       TRAJ_COLOR=TRAJ_COLOR, PAST_COLOR=PAST_COLOR, FUTURE_COLOR=FUTURE_COLOR)


def run_demo_test(cfg, model, data_loader, candidates=[]):

    all_files, all_nframes, all_odoms = [], [], []
    local_results = {'outputs': [], 'gt_traj': [], 'mse': []}
    global_results = {'outputs': [], 'gt_traj': [], 'mse': []}
    for batch_id, batch_data in tqdm(enumerate(data_loader), total=len(data_loader), desc='Run testing'):
        
        filename, clip, odometry, nframes, traj_gt = batch_data
        filename, input_data, odometry, nframes, gt_traj = filename[0], clip[0:1], odometry[0:1], nframes[0:1], traj_gt[0:1]
        
        # filtering
        scene_record = filename[:-4].split('/')[-2]
        clip_name = filename[:-4].split('/')[-1].split('_')[0]
        if '{}_{}'.format(scene_record, clip_name) not in candidates:
            continue
        
        # inference
        with torch.no_grad():
            ratios = get_test_ratios(cfg.TEST.ratios, nframes)
            outputs = model.inference(input_data.cuda(), nframes, ratios, traj=gt_traj.cuda())
        
        # transform from global to local coordinate frame (for 2D visualize)
        outputs_local, gt_traj_local = output_transform(outputs, gt_traj.numpy(), nframes.numpy(), odometry.numpy(), 
                                            intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth, 
                                            centralize=model.centralize, use_odom=True, ignore_depth=model.ignore_depth)

        outputs_global, gt_traj_global = output_transform(outputs, gt_traj.numpy(), nframes.numpy(), odometry.numpy(), 
                                        intrinsics=intrinsics, vis_ratio=vis_ratio, max_depth=max_depth, 
                                        centralize=model.centralize, use_odom=False, ignore_depth=model.ignore_depth)
        
        mse_global = compute_mse(outputs_global[0][0], gt_traj_global[0], nframes[0], future=True)
        mse_local = compute_mse(outputs_local[0][0], gt_traj_local[0], nframes[0], future=True)
        
        all_files.append(filename)
        all_nframes.append(nframes.numpy())
        all_odoms.append(odometry.numpy())
        local_results['outputs'].append(outputs_local)
        local_results['gt_traj'].append(gt_traj_local)
        local_results['mse'].append(mse_local)
        global_results['outputs'].append(outputs_global)
        global_results['gt_traj'].append(gt_traj_global)
        global_results['mse'].append(mse_global)
        
    return all_files, all_nframes, all_odoms, local_results, global_results


def scale_intrinsics():
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    vis_ratio = 0.25
    max_depth = 3
    for k, v in intrinsics.items():
        v *= vis_ratio  # scaling the camera intrinsics
        intrinsics[k] = v
    return 

def vis_2d(video, outputs, trajectories, nframes, odoms, obs_ratios=[0.6]):
    outputs_unobs, outputs_obs = outputs
    vis_video = {}
    for r in obs_ratios:
        num_full = int(nframes[0])  # only one video
        num_obs = int(num_full * float(r))
        
        traj_gt_norm = trajectories[0, :num_full]  # (Tm,3)
        vis_video[r] = np.copy(video[:num_full])
        odometry = odoms[0, :num_full]
        
        # visualize ground truth 2D trajectory
        traj_gt = denormalize(traj_gt_norm, target='3d', intrinsics=intrinsics, max_depth=max_depth)
        visualize_transformed_traj(vis_video[r], traj_gt, odometry, scaled_intrinsics, ratio=r,
                                   TRAJ_COLOR=(255, 255, 0), PAST_COLOR=(0, 255, 0), FUTURE_COLOR=(0, 255, 0))  # green -> green
        
        # visualize predicted 2D trajectory
        # preds_obs = outputs_obs[0][str(r)]['traj'] if outputs_obs is not None else np.copy(traj_gt_norm[:num_obs])
        preds_obs = np.copy(traj_gt_norm[:num_obs])
        preds_norm = np.concatenate((preds_obs, outputs_unobs[0][str(r)]['traj']))
        preds = denormalize(preds_norm, target='3d', intrinsics=intrinsics, max_depth=max_depth)
        visualize_transformed_traj(vis_video[r], preds, odometry, scaled_intrinsics, ratio=r, 
                                   TRAJ_COLOR=(0, 255, 255), PAST_COLOR=(0, 0, 255), FUTURE_COLOR=(255, 0, 0))  # blue -> red
    return vis_video


def visualize(all_files, all_nframes, all_odoms, local_results, global_results, result_path, target='3d', saveframe=False, saveGIF=False):
    
    for filename, nframes, odoms, outputs_local, gt_traj_local, outputs_global, gt_traj_global in zip(all_files, all_nframes, all_odoms,
                                                                                             local_results['outputs'], local_results['gt_traj'],
                                                                                             global_results['outputs'], global_results['gt_traj'],):
        scene_record = filename[:-4].split('/')[-2]
        clip_name = filename[:-4].split('/')[-1].split('_')[0]
        # directly read the original video
        video = read_video(filename, toRGB=False)
        
        # visualize 2D
        vis_video = vis_2d(video, outputs_local, gt_traj_local, nframes, odoms, obs_ratios=cfg.TEST.ratios)
        # visualize 3D
        vis3d_video = {r: None for r in cfg.TEST.ratios}
        if target == '3d':
            canvas_size = (intrinsics['h'] * vis_ratio * 0.5, intrinsics['w'] * vis_ratio * 0.5)
            vis3d_video = vis_traj3d(outputs_global, gt_traj_global, cfg.TEST.ratios, nframes, canvas_size, result_prefix=None, return_video=True)
        
        # save gif/png
        for r in cfg.TEST.ratios:

            if saveGIF:
                gif_dir = os.path.join(result_path, 'gif_{}'.format(r))
                os.makedirs(gif_dir, exist_ok=True)
                result_prefix = os.path.join(gif_dir, '{}_{}'.format(scene_record, clip_name))
                # vis
                vis = overlay_video(vis_video[r], vis3d=vis3d_video[r])
                video_to_gif(vis, result_prefix + '.gif')

            if saveframe:
                img_dir = os.path.join(result_path, 'img_{}'.format(r))
                os.makedirs(img_dir, exist_ok=True)
                result_prefix = os.path.join(img_dir, '{}_{}'.format(scene_record, clip_name))
                # vis
                vis = overlay_frame(vis_video[r][0], vis3d_video[r][-1])
                vis = cv2.cvtColor(vis, cv2.COLOR_BGR2RGB)
                cv2.imwrite(result_prefix + '.png', vis)


def run_visualization(model, test_loader, result_path, candidates):
    os.makedirs(result_path, exist_ok=True)
    
    cache_file = os.path.join(result_path, 'results.pkl')
    if not os.path.exists(cache_file):
        # run test inference
        all_files, all_nframes, all_odoms, local_results, global_results = run_demo_test(cfg, model, test_loader, candidates)
        with open(cache_file, 'wb') as f:
            pickle.dump([all_files, all_nframes, all_odoms, local_results, global_results], f, protocol=pickle.HIGHEST_PROTOCOL)
    else:
        with open(cache_file, 'rb') as f:
            print('Loading the cached results: {}'.format(cache_file))
            all_files, all_nframes, all_odoms, local_results, global_results = pickle.load(f)
        # filtering
        inds = [i for i, filename in enumerate(all_files) if '{}_{}'.format(filename[:-4].split('/')[-2], 
                                                                            filename[:-4].split('/')[-1].split('_')[0]) in candidates]
        all_files = [all_files[i] for i in inds]
        all_nframes = [all_nframes[i] for i in inds]
        all_odoms = [all_odoms[i] for i in inds]
        local_results = {k: [v[i] for i in inds] for k, v in local_results.items()}
        global_results = {k: [v[i] for i in inds] for k, v in global_results.items()}

    # visualization
    visualize(all_files, all_nframes, all_odoms, local_results, global_results, result_path, model.target, saveframe=False, saveGIF=True)
    


def demo():
    
    # build test dataloaders
    test_loader, testnovel_loader = build_dataloaders(cfg, phase='test')

    # build the model
    model_module = importlib.import_module('src.models.{}'.format(cfg.MODEL.arch))
    model = getattr(model_module, cfg.MODEL.arch)(cfg.MODEL, seq_len=cfg.DATA.max_frames, input_size=cfg.DATA.transform.input_size[0])
    model = model.to(device=cfg.device)
    # load checkpoints
    model, test_epoch = load_checkpoint(cfg, model)
    model = model.eval()
    
    # run on see test data
    run_visualization(model, test_loader, 
                      os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'demo_paper', 'seen'),
                      candidates=['microwave_9_c100', 'pantryShelf_9_c36'])
    
    # run on unsee test data
    run_visualization(model, testnovel_loader, 
                      os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'demo_paper', 'unseen'),
                      candidates=['stoveTop_10_c100', 'windowsillAC_8_c29'])


if __name__ == '__main__':
    # parse input arguments
    cfg = parse_configs(phase='test')
    cfg.TEST.batch_size = 1
    cfg.DATA.load_all = False

    # fix random seed 
    set_deterministic(cfg.seed)
    
    # constant camera parameters
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    max_depth = 3
    vis_ratio = 0.25
    
    scaled_intrinsics = deepcopy(intrinsics)
    for k, v in intrinsics.items():
        v *= vis_ratio  # scaling the camera intrinsics
        scaled_intrinsics[k] = v
    
    # test
    demo()
tools/demo_paper.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
from demo import read_traj_file, read_video, read_odometry
from src.utils import video_to_gif
import numpy as np
import cv2
import matplotlib.pyplot as plt
from functools import reduce
from matplotlib.animation import FFMpegWriter

TRAJ_COLOR = (255, 255, 0)  # yellow
PAST_COLOR = (255, 0, 0)    # red
FUTURE_COLOR = (0, 0, 255)  # blue



def XYZ_to_uv(traj3d, intrinsics):
    """ traj3d: (T, 3), a list of 3D (X, Y,Z) points 
    """
    # transform the (X, Y, Z) into the (u,v) by PINHOLE camera model
    traj2d = np.zeros((traj3d.shape[0], 2), dtype=np.float32)
    traj2d[:, 0] = (traj3d[:, 0] * intrinsics['fx'] / traj3d[:, 2] + intrinsics['ox'])
    traj2d[:, 1] = (traj3d[:, 1] * intrinsics['fy'] / traj3d[:, 2] + intrinsics['oy'])
    # clip the coordinates 
    traj2d[:, 0] = np.clip(traj2d[:, 0], 0, intrinsics['w'])
    traj2d[:, 1] = np.clip(traj2d[:, 1], 0, intrinsics['h'])
    traj2d = np.floor(traj2d).astype(np.int32)
    return traj2d


def vis_frame_traj(frame, traj2d, traj_start=0, future_start=15):
    cv2.polylines(frame, [traj2d], False, TRAJ_COLOR, thickness=2)
    for t, uv in enumerate(traj2d):
        color = PAST_COLOR if traj_start + t <= future_start else FUTURE_COLOR
        cv2.circle(frame, uv, radius=5, color=color, thickness=-1)


def visualize_transformed_traj(video, traj3d, odometry, intrinsics, ratio=0.5):
    """ video: (T, H, W, 3)
        traj3d: (T, 3)
        odometry: (T, 4, 4)
    """
    length = len(odometry)
    traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
    future_start = int(length * ratio)
    for i in range(length):
        # compute the transformed points of all future points
        traj_list = [traj3d[i]]  # initial 3D point
        for j in range(i + 1, length):
            odom = reduce(np.dot, odometry[(i+1):(j+1)])  # (4, 4)
            future_point = odom.dot(traj3d_homo[j].T)  # (4,)
            traj_list.append(future_point[:3])
        # visualize current frame
        traj2d = XYZ_to_uv(np.array(traj_list), intrinsics)
        vis_frame_traj(video[i], traj2d, traj_start=i, future_start=future_start)


def local_to_global(traj3d, odometry):
    """ Transform local 3D coordinates into the Global 3D coordinates
        traj3d: (T, 3)
        odometry: (T, 4, 4)
    """
    length = len(odometry)
    traj3d_homo = np.hstack((traj3d, np.ones((length, 1))))  # (T, 4)
    traj3d_global = []
    for i in range(length):
        # compute the transformed points of all future points
        traj_list = [traj3d[i]]  # initial 3D point
        for j in range(i + 1, length):
            odom = reduce(np.dot, odometry[(i+1):(j+1)])  # (4, 4)
            future_point = odom.dot(traj3d_homo[j].T)  # (4,)
            traj_list.append(future_point[:3])
        # record
        traj3d_global.append(np.array(traj_list))
    return traj3d_global



def draw_trajectory_3d(traj3d, video_size, ratio=0.5, fps=5, dpi=30, views=[30, -120], minmax=[(0, 0, -1), (1, 1, 2)], savefile='temp.mp4'):

    xmin, ymin, zmin = minmax[0]
    xmax, ymax, zmax = minmax[1]
    fig_w, fig_h = int(video_size[1] // dpi), int(video_size[0] // dpi)

    # init a 3D figure
    fig = plt.figure(figsize=(fig_w, fig_h))
    ax = fig.add_subplot(projection='3d')
    ax.view_init(views[0], views[1])  # elev: 30, azim: -120
    # Make a 3D quiver plot
    ox, oy, oz = np.zeros((3,3))
    dx, dy, dz = np.array([[xmax-xmin, 0, 0], [0, zmax-zmin, 0], [0, 0, ymax-ymin]]) * 0.2
    ax.quiver(ox, oy, oz, dx, dy, dz)
    # Setup coordinate system
    ax.set_xlim(xmin, xmax)
    ax.set_ylim(zmin, zmax)
    ax.set_zlim(ymin, ymax)
    ax.invert_zaxis()
    ax.set_xlabel('X')  # x: rightward
    ax.set_ylabel('Z')  # z: innerward
    ax.set_zlabel('Y')  # y: downward
    plt.locator_params(axis='x', nbins=5)
    plt.locator_params(axis='y', nbins=5)
    plt.locator_params(axis='z', nbins=5)
    plt.tight_layout()

    # init video writer
    traj3d_writer = FFMpegWriter(fps=fps, metadata=dict(title='3D Trajectory', artist='Matplotlib'))
    future_start = int(len(traj3d) * ratio)
    # process
    with traj3d_writer.saving(fig, savefile, dpi=dpi):
        traj3d_past, traj_future = [], []
        for t, pt in enumerate(traj3d):
            if t <= future_start:
                traj3d_past.append(pt)
            else:
                traj_future.append(pt)
            # draw past trajectory
            xs, ys, zs = np.array(traj3d_past)[:, 0], np.array(traj3d_past)[:, 1], np.array(traj3d_past)[:, 2]
            ax.plot(xs, zs, ys, '-o', color='red')
            if t > future_start:
                # draw future trajectory
                xs, ys, zs = np.array(traj_future)[:, 0], np.array(traj_future)[:, 1], np.array(traj_future)[:, 2]
                ax.plot(xs, zs, ys, '-o', color='blue')
            # grab a frame
            traj3d_writer.grab_frame()
    plt.close()


if __name__ == '__main__':

    data_root = os.path.join(os.path.dirname(__file__), '../data', 'EgoPAT3D')
    video_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'video_clips_hand')
    traj_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'trajectory_repair')
    odometry_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'odometry')
    
    data = {'scene': '1', 'record': 'bathroomCabinet_1', 'clip': 'c9_s439_e480_sx438_sy189_ex605_ey309'}
    # data = {'scene': '1', 'record': 'bathroomCabinet_2', 'clip': 'c61_s3163_e3201_sx567_sy385_ex487_ey207'}
    # data = {'scene': '1', 'record': 'bathroomCabinet_3', 'clip': 'c53_s2632_e2662_sx603_sy389_ex558_ey301'}
    
    folder = 'scene{}_record{}_{}'.format(data['scene'], data['record'].split('_')[-1], data['clip'].split('_')[0])
    result_dir = 'output/demo_task/{}'.format(folder)
    os.makedirs(result_dir, exist_ok=True)

    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    vis_ratio = 0.25
    max_depth = 3
    for k, v in intrinsics.items():
        v *= vis_ratio  # scaling the camera intrinsics
        intrinsics[k] = v
    obs_ratio = 0.5

    # input files
    video_file = os.path.join(video_dir, data['scene'], data['record'], data['clip'] + '.mp4')
    traj_file = os.path.join(traj_dir, data['scene'], data['record'], data['clip'] + '.pkl')
    odom_file = os.path.join(odometry_dir, data['scene'], data['record'], data['clip'] + '.npy')
    # read trajectory
    traj2d, traj3d, num_preserve = read_traj_file(traj_file)
    # read video 
    video = read_video(video_file, toRGB=True)  # RGB video
    video = video[:num_preserve]
    # read odometry
    odometry = read_odometry(odom_file)
    odometry = odometry[:num_preserve]
    num_obs = int(num_preserve * obs_ratio)
    
    # visualize 2D trajectory on video
    visualize_transformed_traj(video, traj3d, odometry, intrinsics, ratio=obs_ratio)

    # write to GIF
    filename = 'vis_scene{}_record{}_clip.gif'.format(data['scene'], data['record'].split('_')[-1], data['clip'].split('_')[0])
    giffile = os.path.join(result_dir, filename)
    video_to_gif(video, giffile, toBGR=True)
    
    traj3d_global = local_to_global(traj3d, odometry)
    traj3d = traj3d_global[0]  # only refer to the camera of the first frame 

    print('x_min={:.6f}, x_max={:.6f} \ny_min={:.6f}, y_max={:.6f} \nz_min={:.6f}, z_max={:.6f}'.format(
        np.min(traj3d[:, 0]), np.max(traj3d[:, 0]),
        np.min(traj3d[:, 1]), np.max(traj3d[:, 1]),
        np.min(traj3d[:, 2]), np.max(traj3d[:, 2])))

    temp_file = os.path.join(result_dir, 'temp.mp4')

    # visualize 3D trajectory
    canvas_size = min(intrinsics['h'], intrinsics['w']) * 0.82
    draw_trajectory_3d(traj3d, (canvas_size, canvas_size), ratio=obs_ratio,
        fps=5, dpi=90, views=[30, -120], minmax=[(-0.1, -0.2, 0.1), (0.3, 0.1, 0.7)], savefile=temp_file)
    video_traj3d = read_video(temp_file, toRGB=True)
    video_to_gif(video_traj3d, os.path.join(result_dir, 'traj3d.gif'), toBGR=True)
    
    # overlap two frames
    for t, name in zip([0, num_obs, -1], ['start', 'current', 'end']):
        frame_vis = np.copy(video[t])
        frame3d = np.copy(video_traj3d[t])
        bottom_left = frame_vis[frame_vis.shape[0] - frame3d.shape[0]:, :frame3d.shape[1], :]
        frame_vis[frame_vis.shape[0] - frame3d.shape[0]:, :frame3d.shape[1], :] = cv2.addWeighted(bottom_left, 0.4, frame3d, 0.6, 0)  # bottom left
        cv2.imwrite(os.path.join(result_dir, 'task_{}.png'.format(name)), frame_vis)

    # overlap the two videos
    video_vis = np.copy(video)
    for i, (traj, rgb) in enumerate(zip(video_traj3d, video)):
        bottom_left = rgb[rgb.shape[0] - traj.shape[0]:, :traj.shape[1], :]
        rgb[rgb.shape[0] - traj.shape[0]:, :traj.shape[1], :] = cv2.addWeighted(bottom_left, 0.4, traj, 0.6, 0)  # bottom left
        video_vis[i] = rgb
    video_to_gif(video_vis, os.path.join(result_dir, 'demo.gif'), toBGR=True)

    # save frames
    frame_dir = os.path.join(result_dir, 'frames')
    os.makedirs(frame_dir, exist_ok=True)
    for t, frame in enumerate(video_vis):
        # add frame number
        cv2.putText(frame, "#{}".format(t+1).zfill(2), (frame.shape[1] - 250, 150), cv2.FONT_HERSHEY_SIMPLEX, 4, color=(0,255,255), thickness=5)
        cv2.imwrite(os.path.join(frame_dir, '{}.png'.format(t + 1)), frame)
    
    if os.path.exists(temp_file):
        os.remove(temp_file)


tools/illustrate_task.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import random
from tqdm import tqdm
import numpy as np
import cv2
from demo import read_traj_file, read_video, read_odometry
from src.utils import video_to_gif
from illustrate_task import visualize_transformed_traj, local_to_global, draw_trajectory_3d


scene_splits = {'train': {'1': ['1', '2', '3', '4', '5', '6', '7'], 
                          '2': ['1', '2', '3', '4', '5', '6', '7'], 
                          '3': ['1', '2', '3', '4', '5', '6'],
                          '4': ['1', '2', '3', '4', '5', '6', '7'],
                          '5': ['1', '2', '3', '4', '5', '6'], 
                          '6': ['1', '2', '3', '4', '5', '6'], 
                          '7': ['1', '2', '3', '4', '5', '6', '7'],
                          '9': ['1', '2', '3', '4', '5', '6', '7'],
                          '10': ['1', '2', '3', '4', '5', '6', '7'],
                          '11': ['1', '2', '3', '4', '5', '6', '7'],
                          '12': ['1', '2', '3', '4', '5', '6', '7']},
                'val': {'1': ['8'], 
                        '2': ['8'], '3': ['7'], '4': ['8'], '5': ['7'], '6': ['7'], 
                        '7': ['8'], '9': ['8'], '10': ['8'], '11': ['8'], '12': ['8']},
                'test': {'1': ['9', '10'], 
                        '2': ['9', '10'], '3': ['9', '10'], '4': ['9', '10'], '5': ['8', '9'], '6': ['8', '9'], 
                        '7': ['9', '10'], '9': ['9', '10'], '10': ['9', '10'], '11': ['9', '10'], '12': ['9', '10']},
                'test_novel': {'13': ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10'], 
                            '14': ['2', '3', '4', '5', '6', '7', '8', '9', '10'], 
                            '15': ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10']}}
scene_names = {'1': 'bathroomCabinet', '2': 'bathroomCounter', '3': 'bin', '4': 'desk', '5': 'drawer', '6': 'kitchenCounter', '7': 'kitchenCupboard',
               '9': 'microwave', '10': 'nightstand', '11': 'pantryShelf', '12': 'smallBins', '13': 'stoveTop', '14': 'windowsillAC', '15': 'woodenTable'}



if __name__ == '__main__':

    data_root = os.path.join(os.path.dirname(__file__), '../data', 'EgoPAT3D')
    video_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'video_clips_hand')
    traj_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'trajectory_repair')
    odometry_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'odometry')
    # intrinsics and extra info
    intrinsics = {'fx': 1.80820276e+03, 'fy': 1.80794556e+03, 
                  'ox': 1.94228662e+03, 'oy': 1.12382178e+03,
                  'w': 3840, 'h': 2160}  # from EgoPAT3D preprocessing code
    vis_ratio = 0.25
    max_depth = 3
    for k, v in intrinsics.items():
        v *= vis_ratio  # scaling the camera intrinsics
        intrinsics[k] = v
    obs_ratio = 0.5
    
    num_record_vis = 3  # number of records to visualize for each scene (<= 3)
    num_clip_vis = 3  # number of samples to visualize for each record (<= 6)
    subset = 'test_novel'
    random.seed(42)
    
    # result path
    result_dir = os.path.join(data_root, 'EgoPAT3D-postproc', 'visualize', subset)
    os.makedirs(result_dir, exist_ok=True)
    temp_file = os.path.join(result_dir, 'temp.mp4')
    
    for scene_id, record_splits in tqdm(scene_splits[subset].items(), ncols=0, desc=subset):
        
        scene_rgb_path = os.path.join(video_dir, scene_id)
        scene_traj_path = os.path.join(traj_dir, scene_id)
        scene_odom_path = os.path.join(odometry_dir, scene_id)
        
        record_names = list(filter(lambda x: x.split('_')[-1] in record_splits, os.listdir(scene_traj_path)))  # get the splits of records
        record_names_selected = random.choices(record_names, k=num_record_vis)
        for record in record_names_selected:
            record_rgb_path = os.path.join(scene_rgb_path, record)
            record_traj_path = os.path.join(scene_traj_path, record)
            record_odom_path = os.path.join(scene_odom_path, record)
            
            traj_files = list(filter(lambda x: x.endswith('.pkl'), os.listdir(record_traj_path)))  # a list of '*.pkl'
            traj_files_selected = random.choices(traj_files, k=num_clip_vis)
            for traj_name in traj_files_selected:
                # input files
                video_file = os.path.join(record_rgb_path, traj_name[:-4] + '.mp4')
                traj_file = os.path.join(record_traj_path, traj_name)
                odom_file = os.path.join(record_odom_path, traj_name[:-4] + '.npy')
                # read trajectory
                traj2d, traj3d, num_preserve = read_traj_file(traj_file)
                # read video 
                video = read_video(video_file, toRGB=True)  # RGB video
                video = video[:num_preserve]
                # read odometry
                odometry = read_odometry(odom_file)
                odometry = odometry[:num_preserve]
                
                # visualize 2D trajectory on video
                visualize_transformed_traj(video, traj3d, odometry, intrinsics, ratio=obs_ratio)
                
                # local 3D to global 3D
                traj3d_global = local_to_global(traj3d, odometry)
                traj3d = traj3d_global[0]  # only refer to the camera of the first frame 
                
                # visualize 3D trajectory
                canvas_size = min(intrinsics['h'], intrinsics['w']) * 0.82
                draw_trajectory_3d(traj3d, (canvas_size, canvas_size), ratio=obs_ratio,
                    fps=5, dpi=90, views=[30, -120], minmax=[(-0.1, -0.2, 0.1), (0.3, 0.1, 0.7)], savefile=temp_file)
                video_traj3d = read_video(temp_file, toRGB=True)
                
                # overlap the two videos
                video_vis = np.copy(video)
                for i, (traj, rgb) in enumerate(zip(video_traj3d, video)):
                    bottom_left = rgb[rgb.shape[0] - traj.shape[0]:, :traj.shape[1], :]
                    rgb[rgb.shape[0] - traj.shape[0]:, :traj.shape[1], :] = cv2.addWeighted(bottom_left, 0.4, traj, 0.6, 0)  # bottom left
                    video_vis[i] = rgb
                
                # save to GIF file
                gif_name = '{}_{}.gif'.format(record, traj_name.split('_')[0])
                video_to_gif(video_vis, os.path.join(result_dir, gif_name), toBGR=True)
                
                if os.path.exists(temp_file):
                    os.remove(temp_file)
tools/visualize_dataset.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import os
import torch
from tensorboardX import SummaryWriter
import importlib

from src.config import parse_configs
from src.EgoPAT3DLoader import build_dataloaders
from src.utils import set_deterministic, compute_displacement_errors, save_the_latest
from src.runtime import train_epoch, eval_epoch, print_eval_results
import src.Losses as loss_module
from src.optimizers import get_optimizer, get_scheduler



def train(cfg):
    # model snapshots
    model_dir = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'snapshot')
    os.makedirs(model_dir, exist_ok=True)

    # tensorboard logging
    logs_dir = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'logs')
    os.makedirs(logs_dir, exist_ok=True)
    writer = SummaryWriter(logs_dir)

    # build data loaders
    print("Loading dataset...")
    traindata_loader, valdata_loader = build_dataloaders(cfg, phase='trainval')
    
    # build the model
    model_module = importlib.import_module('src.models.{}'.format(cfg.MODEL.arch))
    model = getattr(model_module, cfg.MODEL.arch)(cfg.MODEL, seq_len=cfg.DATA.max_frames, input_size=cfg.DATA.transform.input_size[0]).train()
    model = model.to(device=cfg.device)

    # build the loss criterion
    args_loss = {'cfg': cfg.TRAIN.loss}
    criterion = getattr(loss_module, cfg.TRAIN.loss.type)(**args_loss)
    criterion = criterion.to(device=cfg.device)
    
    # optimizer & lr scheduler
    optimizer = get_optimizer(cfg.TRAIN, model.parameters())
    scheduler = get_scheduler(cfg.TRAIN, optimizer)
    
    # training loop
    for epoch in range(cfg.TRAIN.epoch):
        # train one epoch
        loss_train = train_epoch(cfg, model, traindata_loader, criterion, optimizer, writer, epoch)

        if (epoch + 1) % cfg.TRAIN.eval_interval == 0 or epoch + 1 == cfg.TRAIN.epoch:
            # test a model snapshot
            with torch.no_grad():
                all_preds_val, all_gt_val = eval_epoch(cfg, model, valdata_loader, criterion, writer, epoch)
                all_ades_val, all_fdes_val = compute_displacement_errors(all_preds_val, all_gt_val)

            print_eval_results(writer, all_ades_val, all_fdes_val, epoch=epoch, loss_train=loss_train)

        if (epoch + 1) % cfg.TRAIN.snapshot_interval == 0 or epoch + 1 == cfg.TRAIN.epoch:
            # save model snapshot
            save_dict = {'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict()}
            model_file = os.path.join(model_dir, cfg.TRAIN.snapshot_prefix + '%02d.pth'%(epoch + 1))
            save_the_latest(save_dict, model_file, topK=1, ignores=getattr(cfg.TRAIN.scheduler, 'lr_decay_epoch', []))
            print('Model has been saved as: %s'%(model_file))

        # update learning rate
        scheduler.step(epoch=epoch)

    writer.close()


if __name__ == '__main__':
    # parse input arguments
    cfg = parse_configs()

    # fix random seed 
    set_deterministic(cfg.seed)

    # train
    train(cfg)
train.py
oppo-us-research/USST
"""
* Copyright (c) 2023 OPPO. All rights reserved.
*
*
* Licensed under the Apache License, Version 2.0 (the "License"):
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and 
* limitations under the License.
"""

import numpy as np
import os
import torch
from tensorboardX import SummaryWriter
import importlib

from src.config import parse_configs
from src.H2OLoader import build_dataloaders
from src.utils import set_deterministic, save_the_latest
from src.utils_h2o import compute_displacement_errors, compute_block_distances
from src.runtime import train_epoch, eval_h2o_epoch, run_h2o_test, print_eval_results
import src.Losses as loss_module
from src.optimizers import get_optimizer, get_scheduler
from src.utils_io import load_checkpoint, print_de_table, print_delta_table



def get_test_results(model, test_loader, result_file):
    if not os.path.exists(result_file):
        # run test inference
        with torch.no_grad():
            all_preds, all_gt, all_cam_poses = run_h2o_test(cfg, model, test_loader)
        # save predictions
        np.savez(result_file[:-4], pred=all_preds, gt=all_gt, campose=all_cam_poses)
    else:
        print("Result file exists. Loaded from file: %s."%(result_file))
        all_results = np.load(result_file, allow_pickle=True)
        all_preds, all_gt, all_cam_poses = all_results['pred'][()], all_results['gt'][()], all_results['campose'][()]
    return all_preds, all_gt, all_cam_poses


def test(cfg):

    # build test dataloaders
    print("Loading dataset...")
    test_loader = build_dataloaders(cfg, phase='test')

    # build the model
    model_module = importlib.import_module('src.models.{}'.format(cfg.MODEL.arch))
    model = getattr(model_module, cfg.MODEL.arch)(cfg.MODEL, seq_len=cfg.DATA.max_frames, input_size=cfg.DATA.transform.input_size[0])
    model = model.to(device=cfg.device)
    # load checkpoints
    model, test_epoch = load_checkpoint(cfg, model)
    model = model.eval()
    
    # result folder
    result_path = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'test-e{}'.format(test_epoch))
    os.makedirs(result_path, exist_ok=True)
    eval_space = getattr(cfg.TEST, 'eval_space', '3d')

    ### test on the seen scenes
    all_preds, all_gt, all_cam_poses = get_test_results(model, test_loader, os.path.join(result_path, 'test_results.npz'))
    # evaluation
    all_ades, all_fdes = compute_displacement_errors(all_preds, all_gt, all_cam_poses,
                                                               target=model.target, eval_space=eval_space, use_global=cfg.MODEL.use_global)
    all_dxs, all_dys, all_dzs = compute_block_distances(all_preds, all_gt, all_cam_poses,
                                                               target=model.target, eval_space=eval_space, use_global=cfg.MODEL.use_global)
    
    # print tables
    print_de_table(all_ades, all_fdes, subset='test')
    print_delta_table(all_dxs, all_dys, all_dzs, subset='test')

    print("\nDone!")


def train(cfg):
    # model snapshots
    model_dir = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'snapshot')
    os.makedirs(model_dir, exist_ok=True)

    # tensorboard logging
    logs_dir = os.path.join(cfg.output_dir, cfg.DATA.dataset, cfg.tag, 'logs')
    os.makedirs(logs_dir, exist_ok=True)
    writer = SummaryWriter(logs_dir)

    # build data loaders
    print("Loading dataset...")
    traindata_loader, valdata_loader = build_dataloaders(cfg, phase='trainval')
    
    # build the model
    model_module = importlib.import_module('src.models.{}'.format(cfg.MODEL.arch))
    model = getattr(model_module, cfg.MODEL.arch)(cfg.MODEL, seq_len=cfg.DATA.max_frames, input_size=cfg.DATA.transform.input_size[0]).train()
    model = model.to(device=cfg.device)

    # build the loss criterion
    criterion = getattr(loss_module, cfg.TRAIN.loss.type)(cfg=cfg.TRAIN.loss)
    criterion = criterion.to(device=cfg.device)
    
    # optimizer & lr scheduler
    optimizer = get_optimizer(cfg.TRAIN, model.parameters())
    scheduler = get_scheduler(cfg.TRAIN, optimizer)
    
    # training loop
    for epoch in range(cfg.TRAIN.epoch):
        # train one epoch
        loss_train = train_epoch(cfg, model, traindata_loader, criterion, optimizer, writer, epoch)

        if (epoch + 1) % cfg.TRAIN.eval_interval == 0 or epoch + 1 == cfg.TRAIN.epoch:
            # test a model snapshot
            with torch.no_grad():
                all_preds_val, all_gt_val, all_cam_poses = eval_h2o_epoch(cfg, model, valdata_loader, criterion, writer, epoch)
                all_ades_val, all_fdes_val = compute_displacement_errors(all_preds_val, all_gt_val, all_cam_poses, use_global=cfg.MODEL.use_global)

            print_eval_results(writer, all_ades_val, all_fdes_val, epoch=epoch, loss_train=loss_train)

        if (epoch + 1) % cfg.TRAIN.snapshot_interval == 0 or epoch + 1 == cfg.TRAIN.epoch:
            # save model snapshot
            save_dict = {'epoch': epoch, 'model': model.state_dict(), 'optimizer': optimizer.state_dict()}
            model_file = os.path.join(model_dir, cfg.TRAIN.snapshot_prefix + '%02d.pth'%(epoch + 1))
            save_the_latest(save_dict, model_file, topK=1, ignores=getattr(cfg.TRAIN.scheduler, 'lr_decay_epoch', []))
            print('Model has been saved as: %s'%(model_file))
            
        # update learning rate
        scheduler.step(epoch=epoch)

    writer.close()
    


if __name__ == '__main__':
    # parse input arguments
    cfg = parse_configs()

    # fix random seed 
    set_deterministic(cfg.seed)

    if cfg.eval:
        test(cfg)
    else:
        train(cfg)
trainval_h2o.py
oppo-us-research/USST