import os
import sys
from pathlib import Path
from GPUtil import showUtilization as gpu_usage
from tensorboardX import SummaryWriter


base_dir = Path(__file__).absolute().parent
sys.path.append(os.path.join(base_dir, 'monodepth2'))
sys.path.append(os.path.join(base_dir, 'manydepth'))

from collections import deque
import numpy as np
import torch
from src.agents import RandomAgent, HeuristicRandomAgent, ChamferAgent, DepthMAEAgent, MyAgent, FrontAgent, ForwardAgent
from src.config import get_config

from src.env_wrapper import SimEnv
from src.my_utils import MapStat, Timings 
from src.options import Options

def find_nearest(num, num_list):
    num_array = np.array(num_list)
    closest_num = num_array[np.abs(num_array - num).argmin()]
    return closest_num


def main():
    options = Options()
    opt = options.parse()

    if opt.model_type == 'flow2depth':
        from src.models import MyTrainer as ModelTrainer
    else:
        raise NotImplementedError('Model {} is not implemented'.format(opt.model_type))

    model_trainer = ModelTrainer(opt)
    cfg = get_config(opt.cfg_path)

    model_type = 'manydepth' if opt.model_type in ['manydepth', 'flow2depth'] else 'monodepth'

    env = SimEnv(cfg, opt.scene_id)
    if opt.agent_type == 'random':
        agent = RandomAgent(env, cfg.AGENT)
    elif opt.agent_type == 'forward':
        agent = ForwardAgent(env, cfg.AGENT)
    elif opt.agent_type == 'frontier':
        agent = FrontAgent(env, cfg.AGENT, model_type, opt.video_interval)
    elif opt.agent_type == 'mix':
        agent = MyAgent(env, cfg.AGENT, model_type, opt.video_interval)
    
    if opt.agent_type not in ['random', 'forward']:
        dummy = HeuristicRandomAgent(env, cfg.AGENT)
    else: 
        dummy = agent

    if type(agent) in [FrontAgent, MyAgent]:
        print('With frontier explor!')  
    else:
        print('No explor' )

    logger = SummaryWriter(model_trainer.log_path)
    map_logger = MapStat(env)
    timer = Timings()
    timer.reset()

    ## Init test as the starting of curve
    data_count = 0
    APPROX_RANGE = [i for i in range(cfg.TRAIN.WARMUP_DATA, cfg.TRAIN.FULL_DATA_LEN, cfg.COLLECT_STEPS)]
    print(APPROX_RANGE)
    metrics = model_trainer.test(logger, model_trainer.step)
    acc = metrics['test/a1']()
    logger.add_scalar('Explor/a1_wrt_data_num', acc,  data_count)
    ##

    count_steps = 0
    action = ['move_forward'] * 20
    observations, _ = env.step(action)  # noqa: F841

    training_obs = deque(maxlen=cfg.TRAIN.Q_MAXLEN)

    while len(training_obs) < cfg.TRAIN.WARMUP_DATA:
        training_obs += observations
        observations = dummy.act(model_trainer.model, observations)

    data_count = len(training_obs)

    logger.add_scalar('Explor/data_num', len(list(training_obs)), model_trainer.step,)

    # Old version is True False False
    _ = model_trainer.train_one_step(list(training_obs), logger, count_steps)
    torch.cuda.empty_cache()

    reset_done_flag = False
    #while count_steps <= cfg.COLLECT_NUM:
    while len(training_obs) < cfg.TRAIN.FULL_DATA_LEN: 
        if len(training_obs) > cfg.TRAIN.Q_MAXLEN / 5:

            if opt.reset_halfway and len(training_obs) > 18000 and (not reset_done_flag):
                loop_num = 3
                reset_done_flag = True
                model_trainer.reset_model()
            else:
                loop_num = 2

            for _ in range(loop_num):
                acc = model_trainer.train_one_step(list(training_obs), logger, count_steps)
                torch.cuda.empty_cache()
            model_trainer.log_trajectory(env, list(training_obs), logger, count_steps)

            log_data_approx = find_nearest(data_count, APPROX_RANGE)
            logger.add_scalar('Explor/data_num', log_data_approx, model_trainer.step,)
            logger.add_scalar('Explor/a1_wrt_data_num', acc,  log_data_approx)
            print('Training one iter, data count: {}'.format(data_count))
        timer.time("Model train")

        obs_len = 0
        print('Collect data')
        while obs_len < cfg.COLLECT_STEPS:
            observations = agent.act(model_trainer.model, list(training_obs))
            obs_len += len(observations)
            training_obs += observations
            data_count += len(observations)

        map_logger.step([d['pos'] for d in list(training_obs)])
        timer.time("Env step")


        tgt_index = opt.frame_ids[1]
        if (count_steps % (opt.log_frequency // 20) == 0):
            buf, percent = map_logger.summary()
            logger.add_image('Explor/map', buf, count_steps)
            logger.add_scalar('Explor/percent', percent, count_steps)
            if type(agent) in [FrontAgent, MyAgent]:
                buf = agent.map.visualize(None,None, 'buf')
                logger.add_image('Explor/costmap', buf, count_steps)
                if hasattr(agent, 'cur_costmap') and (agent.cur_costmap is not None):
                    buf = agent.cur_costmap.visualize(None, None, 'buf')
                    logger.add_image('Explor/insist_cur', buf, count_steps)
                    buf = agent.past_costmap.visualize(None, None, 'buf')
                    logger.add_image('Explor/insist_pst', buf, count_steps)
                agent.map.reset_viz()
            timer.time("Model eval")

            print('Steps = {}'.format(count_steps), timer.summary())
        count_steps += 1

    # Collect data end, 
    buf, percent = map_logger.summary()
    logger.add_image('Explor/map', buf, count_steps+1)
    logger.add_scalar('Explor/percent', percent, count_steps+1)
    if type(agent) in [FrontAgent, MyAgent]:
        buf = agent.map.visualize(None,None, 'buf')
        logger.add_image('Explor/costmap', buf, count_steps+1)
        agent.map.reset_viz()

    for i in range(4): 
        _ = model_trainer.train_one_step(list(training_obs), logger, count_steps)
        model_trainer.test(logger, model_trainer.step)

    model_trainer.test(logger, model_trainer.step, final=True)

    del training_obs
    del model_trainer
    del env

if __name__ == '__main__':
    main()
        



embodied_depth.py
embodied-depth/Embodied-Depth-Prediction
from src.agents.baseline_agent import RandomAgent, HeuristicRandomAgent, ForwardAgent
from src.agents.frontagent import FrontAgent
from src.agents.combo_agent import MyAgent

src/agents/__init__.py
embodied-depth/Embodied-Depth-Prediction
import random
import numpy as np
import cv2

class RandomAgent:
    def __init__(self, env, cfg):
        self.env = env
        self.steps = cfg.FRAMENUM_PER_STEP
        
        self.action_space = {
            0: 'turn_left', 
            1: 'move_forward', 
            2: 'turn_right'
        }

    def act(self, model, observations):
        actions = self._random_sample(num=self.steps)
        return self._collect_data(actions)

    def _random_sample(self, num=1):
        out_ = [self.action_space[np.random.randint(0,3)] for _ in range(num)]
        return out_

    def _collect_data(self, policy, get_flag=False):
        if get_flag:
            return self.env.step(policy)
        else:
            return self.env.step(policy)[0]

class HeuristicRandomAgent:
    def __init__(self, env, cfg):
        self.env = env
        self.steps = cfg.FRAMENUM_PER_STEP
        
        self.policy = {
            0: ['turn_left'] * 8 + ['move_forward'] * (self.steps - 8 ),
            1: ['turn_right'] * 8 + ['move_forward'] * (self.steps - 8 ),
            2: ['move_forward'] * self.steps
        }


    def act(self, model, observations):
        actions = self._random_sample(num=self.steps)
        return self._collect_data(actions)

    def _random_sample(self, num=1):
        out_ = self.policy[np.random.randint(0,3)] 
        return out_

    def _collect_data(self, policy, get_flag=False):
        if get_flag:
            return self.env.step(policy)
        else:
            return self.env.step(policy)[0]

class ForwardAgent:
    def __init__(self, env, cfg):
        self.env = env
        self.steps = cfg.FRAMENUM_PER_STEP
        self.policy = {
            0: ['turn_left'] * 8 + ['move_forward'] * (self.steps - 8 ),
            1: ['turn_right'] * 8 + ['move_forward'] * (self.steps - 8 ),
            2: ['move_forward'] * self.steps
        }

        self.random_obs_count = 0
        self.explore_prob = cfg.EXPLORE_PROB

    def act(self, model, observations):

        if np.random.rand() < self.explore_prob:
            obs_ = self._random_walk()
        else:
            obs_ = []
            policy = ['move_forward'] * self.steps

            obs, collision_flag = self._collect_data(policy)
            obs_ += obs

            while collision_flag:
                policy = ['turn_right'] * 5
                obs, collision_flag = self._collect_data(policy)
                obs_ += obs

        return obs_

    def _random_walk(self):
        observations = []
        for _ in range(5):
            policy = self.policy[np.random.randint(0,3)]
            obs = self._collect_data(policy=policy)
            observations += obs

        self.random_obs_count += len(observations)
        return observations 

    def _collect_data(self, policy):
        return self.env.step(policy)
src/agents/baseline_agent.py
embodied-depth/Embodied-Depth-Prediction
import numpy as np
from skimage.measure import label

from src.agents.frontagent import FrontAgent, Map


class MyAgent(FrontAgent):
    def __init__(self, env, cfg, model_name='monodepth', interval=1):
        super(MyAgent, self).__init__(env, cfg, model_name, interval)

        self.steps = cfg.FRAMENUM_PER_STEP
        self.ref_frame_index = cfg.REF_FRAME_INDEX

        self.exploit_counter = 9000
        self.explore_counter = 0

        self.last_policy = None

    def act(self, model, history_observations):
        if self.exploit_counter > 400:
            print('Front Explore')
            history = history_observations[-400:]
            self._update_states(model=model, history_observations=history)
            obs = self._explore_act(model, history)

            self.explore_counter += len(obs)
            if self.explore_counter > 800:
                self.exploit_counter = 0
                self.explore_counter = 0
            self.last_policy = 'frontier'
        else:
            obs = self._inconsist_act(model, history_observations)
            self.last_policy = 'inconsist'

            self.exploit_counter += len(obs)

        return obs

    def _inconsist_act(self, model, observations):
        
        if self.model_name == 'monodepth':
            self.cur_costmap = self.create_singlemap(model, observations[-1:])
            self.past_costmap = self.create_singlemap(model, observations[self.ref_frame_index:self.ref_frame_index+1])
        elif self.model_name == 'manydepth':
            self.cur_costmap = self.create_singlemap(model, observations[-self.interval-1:])
            self.past_costmap = self.create_singlemap(model, observations[self.ref_frame_index-self.interval:self.ref_frame_index+1])


        if (self.cur_costmap is not None) and (self.past_costmap is not None):
            inconsist_center, flag = self.find_inconsistent_region(self.cur_costmap, self.past_costmap)
            if flag:
                goal_idx = self.find_navigable_point(inconsist_center)
                goal = self.from_grid(goal_idx[0], goal_idx[1], self.map.shape, self.env.physical_bounds)

                self.map.visualize(goal=goal)

                return self.navigate_get_obs(model=model, current_pose=self.history_obs[-1]['Ext'], goal=goal)

        return self._random_walk()

    def find_inconsistent_region(self, map1, map2):
        occ1 = map1.get_map() > 0
        occ2 = map2.get_map() > 0
        mask1 = map1.get_map() != 0
        mask2 = map2.get_map() != 0
        inconsist = np.logical_and(mask1, mask2) * (occ1 != occ2)

        groups = label(inconsist)
        group_num = groups.max()
        max_group = []

        for i in range(1, group_num+1):
            cur_group = np.argwhere(groups == i)
            if len(cur_group) > len(max_group):
                max_group = cur_group

        if max_group != []:
            c =  np.array(max_group).mean(axis=0)
            return c.astype(int), True 
        else:
            return None, False

    def find_navigable_point(self, goal):
        p_ = np.array(goal).copy()
        while not self.map.is_navigable(p_):
            noise = 4 * np.random.normal(size=(2))
            if self.map.is_legal( np.floor(p_ + noise)):
                p_ = (p_ + noise).astype(int)
        return p_
        
    def create_singlemap(self, model, observation):
        pose = observation[-1]['Ext'] 
        if model is None:
            depth = observation[-1]['depth_sensor'] 
        else:
            depth = model.pred_depth(observation)
        points = self.backproj(depth).squeeze()

        # filter some points in terms of height
        points = points[:, points[1,:] > self.cliplow]
        points = points[:, points[1,:] < self.cliphigh]

        points = points[:, (points[2] != 0)]
        # We do not trust the farther distance estimation
        ## here z's direction is backward
        occ_points = points[:, (points[2] > self.mark_range)]

        empty_points = points.clone()
        empty_points[:3, (points[2] < self.empty_range)] *= self.empty_range /  empty_points[2, (points[2] < self.empty_range)]

        occ_proj_points = (pose[:3, :] @ occ_points.numpy()).T
        empty_proj_points = (pose[:3, :] @ empty_points.numpy()).T

        if occ_proj_points.shape[1] == 0 or occ_proj_points.shape[0] ==0:
            return

        costmap = Map(self.env, self.cfg)
        costmap.update(occ_pts=occ_proj_points, empty_pts=empty_proj_points, current_pose=pose)
        return costmap

src/agents/combo_agent.py
embodied-depth/Embodied-Depth-Prediction
import io
import cv2
import torch
import torch.nn.functional as F
from collections import deque
from PIL import Image
import numpy as np
from tqdm import tqdm
from src.my_utils import *
try:
    from habitat.utils.visualizations import maps
except ImportError:
    print('No habitat in the environment')

from matplotlib import pyplot as plt
from multiprocessing import Pool

from src.agents.planners import Planner
from skimage.measure import label

def euclidean(a, b):
    return np.sqrt( ((a-b)**2).mean() )

def dest_point_from_ang_dist(current_pos, angle, distance):
    z, x = current_pos
    delta_x = distance * np.cos(angle)
    delta_z = distance * np.sin(angle)
    return z+delta_z, x+delta_x

class FrontAgent:
    def __init__(self, env, cfg=None, model_name='monodepth', interval=1):
        self.cfg = cfg
        self.map = Map(env, cfg)
        self.env = env
        self.model_name = model_name

        self.interval = interval
        ## 
        self.planner = Planner(env.physical_bounds, cfg.FRONTIER.PLANNER.DILATION_SIZE, env.type)
        self.history_obs = deque(maxlen=200)
        self.reliable_traj_dist = cfg.FRONTIER.RELIABLE_TRAJ_DIST

        self.hist_frontier = deque(maxlen=3)
        self.steps = cfg.FRAMENUM_PER_STEP
        self.policy = {
            0: ['turn_left'] * 8 + ['move_forward'] * (self.steps - 8 ),
            1: ['move_forward'] * self.steps,
            2: ['turn_right'] * 8 + ['move_forward'] * (self.steps - 8 ),
        }

        if env.type == 'sim':
            #self.K = np.array([[320, 0, 320, 0],
            #                [0,  96, 96, 0],
            #                [0,   0,  1, 0],
            #                [0, 0, 0, 1]], dtype=np.float32)
            self.K = np.array([[1, 0, 0, 0],
                            [0,  1, 0, 0],
                            [0,   0,  1, 0],
                            [0, 0, 0, 1]], dtype=np.float32)
            self.backproj = self.sim_backproj
        else:
            self.K = np.array([[904.62 , 0,     640., 0],
                            [0,      904.62 , 360., 0],
                            [0, 0, 1, 0],
                            [0, 0, 0, 1]], dtype=np.float32)
            self.full_res_shape = (640, 192) # (1024, 320) #(640, 192)
            self.K[0] = self.K[0]   * self.full_res_shape[0] / 1280.
            self.K[1] = self.K[1]   * self.full_res_shape[1] / 720.
            self.backproj = self.real_backproj

        self.HEIGHT = 192
        self.WIDTH = 640

        self.explore_prob = cfg.EXPLORE_PROB
        self.mark_range = self.cfg.FRONTIER.MAP_MARK_RANGE
        self.empty_range = self.cfg.FRONTIER.MAP_CLEAR_RANGE
        self.cliplow  = self.cfg.FRONTIER.PCL_HEIGHT_CLIP_LOW
        self.cliphigh = self.cfg.FRONTIER.PCL_HEIGHT_CLIP_HIGH

        # TODO:
        self.counter = 0

    def act(self, model, history_observations):
        #history = history_observations[-2000::15] + [history_observations[-1]]
        # We only assume the recent 30 frames is convincing, catering to the real world egomotion cases
        if np.random.rand() < self.explore_prob:
            obs_ = self._random_walk()
        else:
            history = history_observations[-400:]
            self._update_states(model=model, history_observations=history)
            obs_ = self._explore_act(model, history)
            #obs += self._round_look()
        return obs_


    def _update_states(self, model, history_observations):
        # We only assume the recent 30 frames is convincing, catering to the real world egomotion cases
        self.update_historyobs(history_observations)

        # Remember the place the robot went as empty [deprecated]
        #self.map._update_empty_map([obs['Ext'] for obs in history_observations])

        self.map.reset_costmap()
        sample_interval  = max(1, len(self.history_obs) // 60)
        self.update_map(model, list(self.history_obs)[::sample_interval])

        # before planning to the new area, update nearby local map
        #nearby_obs = self._round_look()
        #self.update_map(model, nearby_obs)

    def _explore_act(self, model, history_observations):

        current_pos = self.history_obs[-1]['Ext'][(2, 0), 3]

        # Find next goal
        costmap = self.map.get_map()
        frontiers = self._compute_frontiers(costmap)
        grid_idx, goal = self._pick_goal_from_frontiers(frontiers, current_pos)

        if grid_idx is not None:
            self.hist_frontier.append(grid_idx.astype(np.int))
            if len(self.hist_frontier) >2 and (self.hist_frontier[-1] == self.hist_frontier[-2]).any() and (self.hist_frontier[-2] == self.hist_frontier[-3]).any():
                print("RESET COSTMAP cuz the center does not change{}".format(self.hist_frontier))
                self.map.reset_costmap()
                return self._random_walk()


            self.map.visualize(grid_idx ) 
            return self.navigate_get_obs(model=model, current_pose=self.history_obs[-1]['Ext'], goal=goal)
        else:
            return self._random_walk()

    def navigate_get_obs(self, model, current_pose, goal):
        observations_  = []
        collide_flag = True
        counter = 0

        while collide_flag and (counter <= 5):
            if counter > 0:
                # use random to walk out of the corner
                while collide_flag:
                    #print('Collide, turn right and seek to be out')
                    policy = ['turn_right'] * 10
                    obs, collide_flag = self._collect_data(policy, get_flag=True)
                    observations_ += obs
                    policy = ['move_forward'] * 10
                    obs, collide_flag = self._collect_data(policy, get_flag=True)
                    observations_ += obs
                current_pose = observations_[-1]['Ext']

            actions, grid_path = self.planner.plan(self.map.get_map() , current_pose, goal) 
            if observations_ != []:
                self.update_map(model, observations_[-50::5])
            self.map._vis_plan(grid_path)
            obs, collide_flag = self._collect_data(actions, get_flag=True)
            observations_ += obs
            current_pose = observations_[-1]['Ext']
            counter += 1

        return observations_

    def _label_collision(self, cur_ext_matrix):
        cur_pos = [cur_ext_matrix[2, 3], cur_ext_matrix[0, 3]]
        cur_ang = np.arctan2(-cur_ext_matrix[0, 0], -cur_ext_matrix[0, 2]) 

        half_fov = np.pi / 4
        distance = 0.1 * np.sqrt(2)

        point1 = dest_point_from_ang_dist(current_pos=cur_pos, angle=cur_ang-half_fov, distance=distance)
        point2 = dest_point_from_ang_dist(current_pos=cur_pos, angle=cur_ang+half_fov, distance=distance)

        points = np.linspace(point1, point2, num=50)
        self.map._label_collision_occ(points)


    def update_historyobs(self, new_obs):
        if euclidean(new_obs[0]['Ext'][(0, 2), 3] , new_obs[-1]['Ext'][(0, 2), 3]) > self.reliable_traj_dist:
            self.history_obs = deque(new_obs)
        else:
            self.history_obs += deque(new_obs)

        while True:
            if euclidean(self.history_obs[0]['Ext'][(0, 2), 3] , self.history_obs[-1]['Ext'][(0, 2), 3]) > self.reliable_traj_dist:
                self.history_obs.popleft()
            else:
                break

        # Some statistics
        traj = np.array([obs['Ext'][(0, 2), 3] for obs in self.history_obs ])
        print('Step={}, the mean = {}, \t var = {} '.format(self.counter, traj.mean(axis=0), traj.var(axis=0)))
        self.counter += 1


    def update_map(self, model, history_observations):
        if self.model_name == 'monodepth':
            for i, obs in enumerate(tqdm(history_observations)):
                self._update_map(model, [obs])
        elif self.model_name == 'manydepth':
            for i in tqdm(range(self.interval + 1, len(history_observations))):
                self._update_map(model, history_observations[i-self.interval-1:i] )

    def _update_map(self, model, observation):
        pose = observation[-1]['Ext'] 
        if model is None:
            depth = observation[-1]['depth_sensor'] 
        else:
            depth = model.pred_depth(observation)
        points = self.backproj(depth).squeeze()

        # filter some points in terms of height
        points = points[:, points[1,:] > self.cliplow]
        points = points[:, points[1,:] < self.cliphigh]

        points = points[:, (points[2] != 0)]
        # We do not trust the farther distance estimation
        ## here z's direction is backward
        occ_points = points[:, (points[2] > self.mark_range)]

        empty_points = points.clone()
        empty_points[:3, (points[2] < self.empty_range)] *= self.empty_range /  empty_points[2, (points[2] < self.empty_range)]

        occ_proj_points = (pose[:3, :] @ occ_points.numpy()).T
        empty_proj_points = (pose[:3, :] @ empty_points.numpy()).T

        if occ_proj_points.shape[1] == 0 or occ_proj_points.shape[0] ==0:
            return
        self.map.update(occ_pts=occ_proj_points, empty_pts=empty_proj_points, current_pose=pose)


    def real_backproj(self, depth, batch_size=1):
        if depth.ndim == 2:
            depth = np.expand_dims(depth, axis=0)
        depth = torch.from_numpy(depth)

        inv_K = torch.from_numpy(np.linalg.pinv(self.K))
        
        meshgrid = np.meshgrid(range(self.WIDTH), range(self.HEIGHT), indexing='xy')
        
        id_coords = np.stack(meshgrid, axis=0).astype(np.float32)
        id_coords = torch.from_numpy(id_coords)
                                    
        ones = torch.ones(batch_size, 1, self.HEIGHT * self.WIDTH)

        pix_coords = torch.unsqueeze(torch.stack(
                [id_coords[0].view(-1), id_coords[1].view(-1)], 0), 0)
        pix_coords = pix_coords.repeat(batch_size, 1, 1)
        pix_coords = torch.cat([pix_coords, ones], 1)

        # forward(self, depth, inv_K):
        cam_points = torch.matmul(inv_K[:3, :3], pix_coords)
        cam_points = depth.view(batch_size, 1, -1) * cam_points
        cam_points = torch.cat([cam_points, ones], 1)

        return cam_points


    def sim_backproj(self, depth, batch_size=1):
        if depth.ndim == 2:
            depth = np.expand_dims(depth, axis=0)
        depth = torch.from_numpy(depth)

        inv_K = torch.from_numpy(np.linalg.pinv(self.K))
        
        #meshgrid = np.meshgrid(range(self.WIDTH), range(self.HEIGHT), indexing='xy')
        meshgrid = np.meshgrid(np.linspace(-1,1,self.WIDTH), np.linspace(1,-1,self.HEIGHT), indexing='xy')
        
        id_coords = np.stack(meshgrid, axis=0).astype(np.float32)
        id_coords = torch.from_numpy(id_coords)
                                    
        ones = torch.ones(batch_size, 1, self.HEIGHT * self.WIDTH)

        pix_coords = torch.unsqueeze(torch.stack(
                [id_coords[0].view(-1), id_coords[1].view(-1)], 0), 0)
        pix_coords = pix_coords.repeat(batch_size, 1, 1)
        pix_coords = torch.cat([pix_coords, -ones], 1)

        # forward(self, depth, inv_K):
        cam_points = torch.matmul(inv_K[:3, :3], pix_coords)
        cam_points = depth.view(batch_size, 1, -1) * cam_points
        cam_points = torch.cat([cam_points, ones], 1)

        return cam_points

    def _compute_frontiers(self, map):
        frontiers = np.zeros_like(map)
        for i in np.argwhere(map == 0):
            if i[0] == 0 or i[1] == 0 or i[0] == map.shape[0]-1 or i[1] == map.shape[1]-1:
                continue
            if (map[i[0]-1:i[0]+2, i[1]-1:i[1]+2] < 0).any():
                frontiers[i[0], i[1]] = 1

        fronts_group = label(frontiers)
        group_num = fronts_group.max()
        groups = []
        max_group = []

        out_ = []
        for i in range(1, group_num+1):
            cur_group = np.argwhere(fronts_group == i)
            if len(cur_group) > 20:
                groups += [f for f in cur_group]
                out_.append(cur_group)
            if len(cur_group) > len(max_group):
                max_group = cur_group

        if groups != []:
            self.map.visualize(None, groups, None)
        return out_


    def _pick_goal_from_frontiers(self, frontiers, current_pos):
        ## Pick the frontiers w/ max Euclidean distance 
        if frontiers != []:
            dists = []
            groups = []
            for _, group in enumerate(frontiers):
                grid_idx, goal = self._get_frontier_center(group)
                dist =  euclidean(current_pos, goal)
                if dist >= 1.:
                    dists.append(dist)
                    groups.append(group)

            if dists != []:
                dists = np.array(dists)
                dists /= dists.sum()
                grid_idx, goal = self._get_frontier_center(
                                        groups[
                                            np.random.choice(len(dists), p=dists)
                                            ]
                                        )
                return grid_idx, goal

        return None, None

    def from_grid(self, grid_x, grid_y, grid_resolution, bounds):
        lower_bound, upper_bound = bounds

        grid_size = (
            abs(upper_bound[2] - lower_bound[2]) / grid_resolution[0],
            abs(upper_bound[0] - lower_bound[0]) / grid_resolution[1],
        )
        realworld_x = lower_bound[2] + grid_x * grid_size[0]
        realworld_y = lower_bound[0] + grid_y * grid_size[1]
        return realworld_x, realworld_y

    def _get_frontier_center(self, frontiers):
        c =  np.array(frontiers).mean(axis=0)
        goal = self.from_grid(int(c[0]), int(c[1]), self.map.shape, self.env.physical_bounds)
        #print('Frontier Goal is {} (grid), {}'.format(c, goal))
        return c, goal

    def _collect_data(self, policy, get_flag=False):
        if get_flag:
            obs, collide_flag = self.env.step(policy)
            if collide_flag:
                self._label_collision(obs[-1]['Ext'])
            return obs, collide_flag
        else:
            return self.env.step(policy)[0]

    def _random_walk(self):
        observations = []
        for _ in range(5):
            policy = self.policy[np.random.randint(0,3)]
            obs = self._collect_data(policy=policy)
            observations += obs
        return observations

    def _round_look(self):
        observations = []
        policy = ['turn_left']  * 180
        obs = self._collect_data(policy=policy)
        observations += obs[1::4]
        return observations


# CostMap utils
def _bresenhamline_nslope(slope):
    """
    Normalize slope for Bresenham's line algorithm.
    >>> s = np.array([[-2, -2, -2, 0]])
    >>> _bresenhamline_nslope(s)
    array([[-1., -1., -1.,  0.]])
    >>> s = np.array([[0, 0, 0, 0]])
    >>> _bresenhamline_nslope(s)
    array([[ 0.,  0.,  0.,  0.]])
    >>> s = np.array([[0, 0, 9, 0]])
    >>> _bresenhamline_nslope(s)
    array([[ 0.,  0.,  1.,  0.]])
    """
    scale = np.amax(np.abs(slope), axis=1).reshape(-1, 1)
    zeroslope = (scale == 0).all(1)
    scale[zeroslope] = np.ones(1)
    normalizedslope = np.array(slope, dtype=np.double) / scale
    normalizedslope[zeroslope] = np.zeros(slope[0].shape)
    return normalizedslope

def _bresenhamlines(start, end, max_iter):
    """
    Returns npts lines of length max_iter each. (npts x max_iter x dimension) 
    >>> s = np.array([[3, 1, 9, 0],[0, 0, 3, 0]])
    >>> _bresenhamlines(s, np.zeros(s.shape[1]), max_iter=-1)
    array([[[ 3,  1,  8,  0],
            [ 2,  1,  7,  0],
            [ 2,  1,  6,  0],
            [ 2,  1,  5,  0],
            [ 1,  0,  4,  0],
            [ 1,  0,  3,  0],
            [ 1,  0,  2,  0],
            [ 0,  0,  1,  0],
            [ 0,  0,  0,  0]],
    <BLANKLINE>
           [[ 0,  0,  2,  0],
            [ 0,  0,  1,  0],
            [ 0,  0,  0,  0],
            [ 0,  0, -1,  0],
            [ 0,  0, -2,  0],
            [ 0,  0, -3,  0],
            [ 0,  0, -4,  0],
            [ 0,  0, -5,  0],
            [ 0,  0, -6,  0]]])
    """
    if max_iter == -1:
        max_iter = np.amax(np.amax(np.abs(end - start), axis=1)) - 1
        max_iter_each = np.amax(np.abs(end - start), axis=1) - 1

    npts, dim = start.shape
    nslope = _bresenhamline_nslope(end - start)

    # steps to iterate on
    stepseq = np.arange(1, max_iter + 1)
    stepmat = np.tile(stepseq, (dim, 1)).T

    # some hacks for broadcasting properly
    bline = start[:, np.newaxis, :] + nslope[:, np.newaxis, :] * stepmat

    bline = np.concatenate([bline[i, range(max_iter_each[i])] for i in range(npts)])

    # Approximate to nearest int
    return np.array(np.rint(bline), dtype=start.dtype)

def bresenhamline(start, end, max_iter=5):
    """
    Returns a list of points from (start, end] by ray tracing a line b/w the
    points.
    Parameters:
        start: An array of start points (number of points x dimension)
        end:   An end points (1 x dimension)
            or An array of end point corresponding to each start point
                (number of points x dimension)
        max_iter: Max points to traverse. if -1, maximum number of required
                  points are traversed
    Returns:
        linevox (n x dimension) A cumulative array of all points traversed by
        all the lines so far.
    >>> s = np.array([[3, 1, 9, 0],[0, 0, 3, 0]])
    >>> bresenhamline(s, np.zeros(s.shape[1]), max_iter=-1)
    array([[ 3,  1,  8,  0],
           [ 2,  1,  7,  0],
           [ 2,  1,  6,  0],
           [ 2,  1,  5,  0],
           [ 1,  0,  4,  0],
           [ 1,  0,  3,  0],
           [ 1,  0,  2,  0],
           [ 0,  0,  1,  0],
           [ 0,  0,  0,  0],
           [ 0,  0,  2,  0],
           [ 0,  0,  1,  0],
           [ 0,  0,  0,  0],
           [ 0,  0, -1,  0],
           [ 0,  0, -2,  0],
           [ 0,  0, -3,  0],
           [ 0,  0, -4,  0],
           [ 0,  0, -5,  0],
           [ 0,  0, -6,  0]])
    """
    # Return the points as a single array
    return _bresenhamlines(start, end, max_iter).reshape(-1, start.shape[-1])

class Map():
    def __init__(self, env, cfg):
        self.env = env

        self.log_occupied = cfg.FRONTIER.MAP_LOG_OCCUPIED
        self.log_free = cfg.FRONTIER.MAP_LOG_FREE
        self.decay_factor = cfg.FRONTIER.MAP_DECAY

        if self.env.type == 'sim':
            top_down_map = maps.get_topdown_map_from_sim(
                env._sim, map_resolution=512
                ) 
            self.total = (top_down_map == 1).sum()
            self.top_down_map = maps.colorize_topdown_map(top_down_map)
        else:
            # Resolution is 0.03 meter per pixel
            self.top_down_map = np.zeros((1000, 1000, 3))

        self.shape = self.top_down_map.shape[:2]
        self.log_odds_prob = np.zeros(self.top_down_map.shape[:2])
        self.confident_empty_map = np.zeros(self.top_down_map.shape[:2])

        self.real_bounds = self.env.physical_bounds

    def is_legal(self, point):
        return (point >= 0).all() and (point < self.shape).all()

    def is_navigable(self, point):
        return self.log_odds_prob[point[0], point[1]] <= 0

    def reset_viz(self):
        if self.env.type == 'sim':
            top_down_map = maps.get_topdown_map_from_sim(
                self.env._sim, map_resolution=512
                ) 
            self.top_down_map = maps.colorize_topdown_map(top_down_map)
        else:
            self.top_down_map = np.zeros((1000, 1000, 3))

    def to_3dgrid(self, points, grid_resolution):
        '''
        points with shape (N, 2),  
        '''
        lower_bound, upper_bound = self.real_bounds

        grid_size = (
            abs(upper_bound[2] - lower_bound[2]) / grid_resolution[0],
            abs(upper_bound[0] - lower_bound[0]) / grid_resolution[1],
            abs(upper_bound[1] - lower_bound[1]) / grid_resolution[1],
        )
        grid_z = ((points[:, 2] - lower_bound[2]) / grid_size[0]).astype(np.int)
        grid_x = ((points[:, 0] - lower_bound[0]) / grid_size[1]).astype(np.int)
        grid_y = ((points[:, 1] - lower_bound[1]) / grid_size[1]).astype(np.int)

        return np.stack( [grid_z, grid_x, grid_y], axis=1)

    def to_grid(self, points, grid_resolution):
        '''
        points with shape (N, 2),  
        '''
        lower_bound, upper_bound = self.real_bounds

        grid_size = (
            abs(upper_bound[2] - lower_bound[2]) / grid_resolution[0],
            abs(upper_bound[0] - lower_bound[0]) / grid_resolution[1],
        )
        grid_x = ((points[:, 0] - lower_bound[2]) / grid_size[0]).astype(np.int)
        grid_y = ((points[:, 1] - lower_bound[0]) / grid_size[1]).astype(np.int)
        return np.stack( [grid_x, grid_y], axis=1)

    def _vis_costmap(self):
        for p in np.array((self.log_odds_prob>0).nonzero()).T:
            overlay = self.top_down_map.copy()
            cv2.circle(overlay, (p[1], p[0]), radius=1, color=(255, 0, 0), thickness=1)
            alpha = 0.5
            self.top_down_map = cv2.addWeighted(overlay, alpha, self.top_down_map, 1 - alpha, 0)

        for p in np.array((self.log_odds_prob<0).nonzero()).T:
            overlay = self.top_down_map.copy()
            cv2.circle(overlay, (p[1], p[0]), radius=1, color=(0, 255, 0), thickness=1)
            alpha = 0.1
            self.top_down_map = cv2.addWeighted(overlay, alpha, self.top_down_map, 1 - alpha, 0)

    def _vis_plan(self, traj):
        p = traj[0]
        cv2.circle(self.top_down_map, (p[1], p[0]), radius=7, color=(100, 250, 0), thickness=4)
        for p in traj[1:]:
            cv2.circle(self.top_down_map, (p[1], p[0]), radius=1, color=(200, 200, 0), thickness=1)
        p = traj[-1]
        cv2.circle(self.top_down_map, (p[1], p[0]), radius=7, color=(0, 0 ,250), thickness=4)

    def visualize(self, goal=None, fronts=None, name=None):
        if goal is not None:
            goal = (int(goal[0]), int(goal[1]))
            cv2.circle(self.top_down_map, (goal[1], goal[0]), radius=7, color=(0, 255, 0), thickness=3)

        if fronts is not None:
            f = np.array(fronts)
            for p in f:
                cv2.circle(self.top_down_map, (p[1], p[0]), radius=1, color=(0, 0, 255), thickness=1)

        if name is not None:
            self._vis_costmap()
            plt.axis('off')
            plt.imshow(self.top_down_map)
            if name =='buf':
                buf = io.BytesIO() 
                plt.savefig(buf, format='jpeg')
                buf.seek(0)
                image = Image.open(buf)
                image = ToTensor()(image)

                self.reset_viz()
                return image
            else:
                self.reset_viz()
                plt.savefig(name)

    def reset_costmap(self):
        self.log_odds_prob = np.zeros_like(self.log_odds_prob)
        self.reset_viz()
    
    def _update_empty_map(self, _traj):
        traj = np.stack(_traj)
        points = np.stack([traj[:, 2, 3], traj[:, 0, 3]], axis=1)
        trajs = self.to_grid(points, self.log_odds_prob.shape)
        self.confident_empty_map[trajs[:, 0], trajs[:, 1]] = -1

    def _update_3dlog_odds(self, pcls, current_pose):
        # pcls (Num, 3)


        grid_idxs = self.to_3dgrid(pcls, self.log_odds_prob.shape)

        boundz, boundx = self.log_odds_prob.shape 
        filter_grid_idxs = grid_idxs[grid_idxs[:, 0]<boundz] 
        filter_grid_idxs = filter_grid_idxs[0<= filter_grid_idxs[:, 0]] 
        filter_grid_idxs = filter_grid_idxs[filter_grid_idxs[:, 1]<boundx] 
        filter_grid_idxs = filter_grid_idxs[0<= filter_grid_idxs[:, 1]] 

        current_pose = np.expand_dims(current_pose, 0)
        start_points = self.to_3dgrid(current_pose, self.log_odds_prob.shape) #- np.array([mins]) 

        mins, maxs = [], []
        for i in range(3):
            mins.append(min(filter_grid_idxs[:, i].min(), int(start_points[:, i])))
            maxs.append(max(filter_grid_idxs[:, i].max(), int(start_points[:, i])))
        
        #occ3d = np.zeros((boundz, boundx, boundx))
        occ3d = np.zeros([maxs[i] - mins[i] + 1 for i in range(3)])

        filter_grid_idxs -= np.array([mins])
        start_points -= np.array([mins]) 
        starts = start_points.repeat(filter_grid_idxs.shape[0], axis=0)
        ## Occupied points
        occ3d[filter_grid_idxs[:, 0], filter_grid_idxs[:, 1], filter_grid_idxs[:, 2]] = 1

        ## Empty points
        bresenham_path = bresenhamline( starts,
                                        filter_grid_idxs,
                                        max_iter=-1)
        bresenham_path = bresenham_path[bresenham_path[:, 0]<boundz] 
        bresenham_path = bresenham_path[0<=bresenham_path[:, 0]] 
        bresenham_path = bresenham_path[bresenham_path[:, 1]<boundx] 
        bresenham_path = bresenham_path[0<=bresenham_path[:, 1]] 
        bresenham_path = bresenham_path[bresenham_path[:, 2]<boundx] 
        bresenham_path = bresenham_path[0<=bresenham_path[:, 2]] 
        occ3d[bresenham_path[:, 0], bresenham_path[:, 1], bresenham_path[:, 2]] = 0.5

        ## Project 3d to 2d
        occ2d = occ3d.max(axis=2)
        occupied_idx = np.argwhere(occ2d == 1) + np.array([mins[:2]]) #        
        empty_idx = np.argwhere(occ2d == 0.5) + np.array([mins[:2]])#        

        ## Update log odds 
        self.log_odds_prob *= self.decay_factor
        self.log_odds_prob[occupied_idx[:, 0], occupied_idx[:, 1]] += self.log_occupied
        self.log_odds_prob[empty_idx[:, 0], empty_idx[:, 1]] -= self.log_free

        self.log_odds_prob = np.clip(self.log_odds_prob, -1000, 1000)
        #self.log_odds_prob[self.confident_empty_map < 0] = -1000

        for p in np.array((self.log_odds_prob>0).nonzero()).T:
            cv2.circle(self.top_down_map, (p[1], p[0]), radius=1, color=(255, 0, 0), thickness=1)

    def _clip_grid(self, idxs, boundx, boundy):
        filter_grid_idxs = idxs[idxs[:, 1]<boundy] 
        filter_grid_idxs = filter_grid_idxs[0<= filter_grid_idxs[:, 1]] 
        filter_grid_idxs = filter_grid_idxs[filter_grid_idxs[:, 0]<boundx] 
        filter_grid_idxs = filter_grid_idxs[0<= filter_grid_idxs[:, 0]] 

        return filter_grid_idxs

    def _update_log_odds(self, occ_pcls, empty_pcls, current_pose):
        # pcls (Num, 3)
        self.log_odds_prob *= self.decay_factor

        boundx, boundy = self.log_odds_prob.shape 
        # Occupied
        occ_points = np.stack([occ_pcls[:, 2], occ_pcls[:, 0]], axis=1)
        occ_grid_idxs = np.unique(self.to_grid(occ_points, self.log_odds_prob.shape), axis=0)
        filter_occ_grids = self._clip_grid(occ_grid_idxs, boundx, boundy)
        self.log_odds_prob[filter_occ_grids[:, 0], filter_occ_grids[:, 1]] += self.log_occupied

        # Empty
        empty_points = np.stack([empty_pcls[:, 2], empty_pcls[:, 0]], axis=1)
        empty_grid_idxs = np.unique(self.to_grid(empty_points, self.log_odds_prob.shape), axis=0)
        filter_empty_grids = self._clip_grid(empty_grid_idxs, boundx, boundy)
        if filter_empty_grids.shape[0] > 0:
            start_points = self.to_grid(
                            np.array([current_pose[[2,0]]]), 
                            self.log_odds_prob.shape,
                            )
            starts = start_points.repeat(filter_empty_grids.shape[0], axis=0)
            bresenham_path = bresenhamline( starts,
                                            filter_empty_grids,
                                            max_iter=-1)
            bresenham_path = self._clip_grid(bresenham_path, boundx, boundy)
            self.log_odds_prob[bresenham_path[:, 0], bresenham_path[:, 1]] -= self.log_free

        # Saturate the log odds
        self.log_odds_prob = np.clip(self.log_odds_prob, -1000, 1000)

        for p in np.array((self.log_odds_prob>0).nonzero()).T:
            cv2.circle(self.top_down_map, (p[1], p[0]), radius=1, color=(255, 0, 0), thickness=1)

    def _update_log_oddsv2(self, occ_pcls, empty_pcls, current_pose):
        # pcls (Num, 3)
        self.log_odds_prob *= self.decay_factor

        occupied_map = np.zeros_like(self.log_odds_prob)
        empty_map = np.zeros_like(self.log_odds_prob)

        boundx, boundy = self.log_odds_prob.shape 
        # Occupied
        occ_points = np.stack([occ_pcls[:, 2], occ_pcls[:, 0]], axis=1)
        occ_grid_idxs = np.unique(self.to_grid(occ_points, self.log_odds_prob.shape), axis=0)
        filter_occ_grids = self._clip_grid(occ_grid_idxs, boundx, boundy)


        occupied_map[filter_occ_grids[:, 0], filter_occ_grids[:, 1]] = 1

        # Empty
        empty_points = np.stack([empty_pcls[:, 2], empty_pcls[:, 0]], axis=1)
        empty_grid_idxs = np.unique(self.to_grid(empty_points, self.log_odds_prob.shape), axis=0)
        filter_empty_grids = self._clip_grid(empty_grid_idxs, boundx, boundy)
        if filter_empty_grids.shape[0] > 0:
            start_points = self.to_grid(
                            np.array([current_pose[[2,0]]]), 
                            self.log_odds_prob.shape,
                            )
            starts = start_points.repeat(filter_empty_grids.shape[0], axis=0)
            bresenham_path = bresenhamline( starts,
                                            filter_empty_grids,
                                            max_iter=-1)
            bresenham_path = self._clip_grid(bresenham_path, boundx, boundy)

            empty_map[bresenham_path[:, 0], bresenham_path[:, 1]] = 1
  
        # Merge Occupied map and empty map
        # Keep occupied points if it is also labeled as empty
        overlap = np.logical_and(occupied_map, empty_map)
        empty_map[overlap] = 0
        self.log_odds_prob += self.log_occupied * occupied_map
        self.log_odds_prob -= self.log_free * empty_map

        # Saturate the log odds
        self.log_odds_prob = np.clip(self.log_odds_prob, -1000, 1000)

    def _label_collision_occ(self, points):
        '''
        Input:
            points: [N, 2]  in continuous space
        '''
        boundx, boundy = self.log_odds_prob.shape 
        occ_grid_idxs = np.unique(self.to_grid(points, self.log_odds_prob.shape), axis=0)
        filter_occ_grids = self._clip_grid(occ_grid_idxs, boundx, boundy)

        self.log_odds_prob[filter_occ_grids[:, 0], filter_occ_grids[:, 1]] = 1000
        for p in zip(filter_occ_grids[:, 0],filter_occ_grids[:, 1]) :
            overlay = self.top_down_map.copy()
            cv2.circle(overlay, (p[1], p[0]), radius=2, color=(10, 10, 10), thickness=1)
            alpha = 0.3
            self.top_down_map = cv2.addWeighted(overlay, alpha, self.top_down_map, 1 - alpha, 0)

    def update(self, occ_pts, empty_pts, current_pose):
        self._update_log_oddsv2(occ_pts, empty_pts, current_pose[:3, 3])
            
    def get_map(self):
        return self.log_odds_prob

    def summary(self):
        #map = (self.map / self.map.max() * 255).astype(np.uint8)
        #vis = np.stack([map, np.zeros_like(map), np.zeros_like(map)], dim=-1)
        ##vis = cv2.cvtColor(map ,cv2.COLOR_GRAY2RGB)
        #vis_ = cv2.addWeighted(top_down_map, 0.5, vis, 0.5, 0)
        plt.imshow(self.top_down_map)
        print(len(self.map.nonzero()[0]) / self.total)



src/agents/frontagent.py
embodied-depth/Embodied-Depth-Prediction
# Modified from the work https://github.com/richardos/occupancy-grid-a-star
import math
import numpy as np
from heapq import heappush, heappop
import matplotlib.pyplot as plt

from scipy import ndimage


def dist2d(point1, point2):
    x1, y1 = point1[0:2]
    x2, y2 = point2[0:2]

    dist2 = (x1 - x2)**2 + (y1 - y2)**2

    return math.sqrt(dist2)

def _get_movements_4n():
    """
    Get all possible 4-connectivity movements.
    :return: list of movements with cost [(dx, dy, movement_cost)]
    """
    return [(1, 0, 1.0),
            (0, 1, 1.0),
            (-1, 0, 1.0),
            (0, -1, 1.0)]


def _get_movements_8n():
    """
    Get all possible 8-connectivity movements. Equivalent to get_movements_in_radius(1).
    :return: list of movements with cost [(dx, dy, movement_cost)]
    """
    s2 = math.sqrt(2)
    return [(1, 0, 1.0),
            (0, 1, 1.0),
            (-1, 0, 1.0),
            (0, -1, 1.0),
            (1, 1, s2),
            (-1, 1, s2),
            (-1, -1, s2),
            (1, -1, s2)]


def a_star(start_m, goal_m, gmap, movement='8N', occupancy_cost_factor=100):
    """
    A* for 2D occupancy grid.
    :param start_m: start node (x, y) in meters
    :param goal_m: goal node (x, y) in meters
    :param gmap: the grid map
    :param movement: select between 4-connectivity ('4N') and 8-connectivity ('8N', default)
    :param occupancy_cost_factor: a number the will be multiplied by the occupancy probability
        of a grid map cell to give the additional movement cost to this cell (default: 3).
    :return: a tuple that contains: (the resulting path in meters, the resulting path in data array indices)
    """

    # get array indices of start and goal
    start = gmap.get_index_from_coordinates(start_m[0], start_m[1])
    goal = gmap.get_index_from_coordinates(goal_m[0], goal_m[1])

    # check if start and goal nodes correspond to free spaces
    #if gmap.is_occupied_idx(start):
    #    raise Exception('Start node is not traversable')

    #if gmap.is_occupied_idx(goal):
    #    raise Exception('Goal node is not traversable')

    # add start node to front
    # front is a list of (total estimated cost to goal, total cost from start to node, node, previous node)
    start_node_cost = 0
    start_node_estimated_cost_to_goal = dist2d(start, goal) + start_node_cost
    front = [(start_node_estimated_cost_to_goal, start_node_cost, start, None)]

    # use a dictionary to remember where we came from in order to reconstruct the path later on
    came_from = {}

    # get possible movements
    if movement == '4N':
        movements = _get_movements_4n()
    elif movement == '8N':
        movements = _get_movements_8n()
    else:
        raise ValueError('Unknown movement')

    # while there are elements to investigate in our front.
    while front:
        # get smallest item and remove from front.
        element = heappop(front)

        # if this has been visited already, skip it
        total_cost, cost, pos, previous = element
        if gmap.is_visited_idx(pos):
            continue

        # now it has been visited, mark with cost
        gmap.mark_visited_idx(pos)

        # set its previous node
        came_from[pos] = previous

        # if the goal has been reached, we are done!
        if pos == goal:
            break

        # check all neighbors
        for dx, dy, deltacost in movements:
            # determine new position
            new_x = pos[0] + dx
            new_y = pos[1] + dy
            new_pos = (new_x, new_y)

            # check whether new position is inside the map
            # if not, skip node
            if not gmap.is_inside_idx(new_pos):
                continue

            # add node to front if it was not visited before and is not an obstacle
            if (not gmap.is_visited_idx(new_pos)): #and (not gmap.is_occupied_idx(new_pos)):
                potential_function_cost = gmap.get_data_idx(new_pos)*occupancy_cost_factor
                new_cost = cost + deltacost + potential_function_cost
                new_total_cost_to_goal = new_cost + dist2d(new_pos, goal) + potential_function_cost

                heappush(front, (new_total_cost_to_goal, new_cost, new_pos, pos))

    # reconstruct path backwards (only if we reached the goal)
    path = []
    path_idx = []
    if pos == goal:
        while pos:
            path_idx.append(pos)
            # transform array indices to meters
            pos_m_x, pos_m_y = gmap.get_coordinates_from_index(pos[0], pos[1])
            path.append((pos_m_x, pos_m_y))
            pos = came_from[pos]

        # reverse so that path is from start to goal.
        path.reverse()
        path_idx.reverse()

    return path, path_idx


class OccupancyGridMap:
    def __init__(self, data_array, bounds, occupancy_threshold=0.8):
        """
        Creates a grid map
        :param data_array: a 2D array with a value of occupancy per cell (values from 0 - 1)
        :param cell_size: cell size in meters
        :param occupancy_threshold: A threshold to determine whether a cell is occupied or free.
        A cell is considered occupied if its value >= occupancy_threshold, free otherwise.
        """

        self.data = data_array
        self.dim_cells = data_array.shape

        self.lower_bound, self.upper_bound = bounds

        cell_size = min(abs(self.upper_bound[coord] - self.lower_bound[coord]) / 512 for coord in [0, 2])
        self.dim_meters = (self.dim_cells[0] * cell_size, self.dim_cells[1] * cell_size)
        self.cell_size = cell_size
        self.occupancy_threshold = occupancy_threshold
        # 2D array to mark visited nodes (in the beginning, no node has been visited)
        self.visited = np.zeros(self.dim_cells, dtype=np.float32)

    def mark_visited_idx(self, point_idx):
        """
        Mark a point as visited.
        :param point_idx: a point (x, y) in data array
        """
        x_index, y_index = point_idx
        if x_index < 0 or y_index < 0 or x_index >= self.dim_cells[0] or y_index >= self.dim_cells[1]:
            raise Exception('Point is outside map boundary')

        self.visited[x_index][y_index] = 1.0

    def mark_visited(self, point):
        """
        Mark a point as visited.
        :param point: a 2D point (x, y) in meters
        """
        x, y = point
        x_index, y_index = self.get_index_from_coordinates(x, y)

        return self.mark_visited_idx((x_index, y_index))

    def is_visited_idx(self, point_idx):
        """
        Check whether the given point is visited.
        :param point_idx: a point (x, y) in data array
        :return: True if the given point is visited, false otherwise
        """
        x_index, y_index = point_idx
        if x_index < 0 or y_index < 0 or x_index >= self.dim_cells[0] or y_index >= self.dim_cells[1]:
            raise Exception('Point is outside map boundary')

        if self.visited[x_index][y_index] == 1.0:
            return True
        else:
            return False

    def is_visited(self, point):
        """
        Check whether the given point is visited.
        :param point: a 2D point (x, y) in meters
        :return: True if the given point is visited, false otherwise
        """
        x, y = point
        x_index, y_index = self.get_index_from_coordinates(x, y)

        return self.is_visited_idx((x_index, y_index))

    def get_data_idx(self, point_idx):
        """
        Get the occupancy value of the given point.
        :param point_idx: a point (x, y) in data array
        :return: the occupancy value of the given point
        """
        x_index, y_index = point_idx
        if x_index < 0 or y_index < 0 or x_index >= self.dim_cells[0] or y_index >= self.dim_cells[1]:
            raise Exception('Point is outside map boundary')

        return self.data[x_index][y_index]

    def get_data(self, point):
        """
        Get the occupancy value of the given point.
        :param point: a 2D point (x, y) in meters
        :return: the occupancy value of the given point
        """
        x, y = point
        x_index, y_index = self.get_index_from_coordinates(x, y)

        return self.get_data_idx((x_index, y_index))

    def set_data_idx(self, point_idx, new_value):
        """
        Set the occupancy value of the given point.
        :param point_idx: a point (x, y) in data array
        :param new_value: the new occupancy values
        """
        x_index, y_index = point_idx
        if x_index < 0 or y_index < 0 or x_index >= self.dim_cells[0] or y_index >= self.dim_cells[1]:
            raise Exception('Point is outside map boundary')

        self.data[x_index][y_index] = new_value

    def set_data(self, point, new_value):
        """
        Set the occupancy value of the given point.
        :param point: a 2D point (x, y) in meters
        :param new_value: the new occupancy value
        """
        x, y = point
        x_index, y_index = self.get_index_from_coordinates(x, y)

        self.set_data_idx((x_index, y_index), new_value)

    def is_inside_idx(self, point_idx):
        """
        Check whether the given point is inside the map.
        :param point_idx: a point (x, y) in data array
        :return: True if the given point is inside the map, false otherwise
        """
        x_index, y_index = point_idx
        if x_index < 0 or y_index < 0 or x_index >= self.dim_cells[0] or y_index >= self.dim_cells[1]:
            return False
        else:
            return True

    def is_inside(self, point):
        """
        Check whether the given point is inside the map.
        :param point: a 2D point (x, y) in meters
        :return: True if the given point is inside the map, false otherwise
        """
        x, y = point
        x_index, y_index = self.get_index_from_coordinates(x, y)

        return self.is_inside_idx((x_index, y_index))

    def is_occupied_idx(self, point_idx):
        """
        Check whether the given point is occupied according the the occupancy threshold.
        :param point_idx: a point (x, y) in data array
        :return: True if the given point is occupied, false otherwise
        """
        x_index, y_index = point_idx
        if self.get_data_idx((x_index, y_index)) >= self.occupancy_threshold:
            return True
        else:
            return False

    def is_occupied(self, point):
        """
        Check whether the given point is occupied according the the occupancy threshold.
        :param point: a 2D point (x, y) in meters
        :return: True if the given point is occupied, false otherwise
        """
        x, y = point
        x_index, y_index = self.get_index_from_coordinates(x, y)

        return self.is_occupied_idx((x_index, y_index))

    def get_index_from_coordinates(self, x, y):
        """
        Get the array indices of the given point.
        :param x: the point's x-coordinate in meters
        :param y: the point's y-coordinate in meters
        :return: the corresponding array indices as a (x, y) tuple
        """
        grid_size = (
            abs(self.upper_bound[2] - self.lower_bound[2]) / self.dim_cells[0],
            abs(self.upper_bound[0] - self.lower_bound[0]) / self.dim_cells[1],
            )
        grid_x = int((x - self.lower_bound[2]) / grid_size[0])
        grid_y = int((y - self.lower_bound[0]) / grid_size[1])
        return grid_x, grid_y

    def get_coordinates_from_index(self, x_index, y_index):
        """
        Get the coordinates of the given array point in meters.
        :param x_index: the point's x index
        :param y_index: the point's y index
        :return: the corresponding point in meters as a (x, y) tuple
        """
        grid_size = (
            abs(self.upper_bound[2] - self.lower_bound[2]) / self.dim_cells[0],
            abs(self.upper_bound[0] - self.lower_bound[0]) / self.dim_cells[1],
        )
        realworld_x = self.lower_bound[2] + x_index * grid_size[0]
        realworld_y = self.lower_bound[0] + y_index * grid_size[1]
        return realworld_x, realworld_y

    def plot(self, alpha=1, min_val=0, origin='lower'):
        """
        plot the grid map
        """
        plt.imshow(self.data, vmin=min_val, vmax=1, origin=origin, interpolation='none', alpha=alpha)
        plt.draw()



class Planner:
    def __init__(self, bounds, dilation_size, mode='sim'):
        self.bounds = bounds
        self.dilation_size = dilation_size
        if mode=='sim':
            self.turn_angle = 2.
            self.move_meter = 0.03
        elif mode=='real':
            self.turn_angle = 4.
            self.move_meter = 0.057


    def plan(self, costmap, start, goal):
        phy_path, grid_path = self.get_waypoints(costmap, start, goal)
        actions = self.path2action(phy_path, start)
        return actions, grid_path

    def get_waypoints(self, _costmap, start, goal,):
        '''
        costmap
        start: Extrinsic matrix with shape (4, 4)
        goal: [z, x]   
        '''
        start_pos = [start[2, 3], start[0, 3]]

        costmap = (_costmap > 0).astype(np.float32)
        costmap[costmap > 0] = 1000  # occupied area
        costmap[_costmap == 0] = 100 # unknown area

        # Dilate the occupied points, in case of collision
        costmap = ndimage.grey_dilation(costmap,
                                        size=self.dilation_size)

        gmap = OccupancyGridMap(
                    data_array=costmap,
                    bounds=self.bounds,
        )

        phy_path, grid_path = \
                a_star( start_pos ,
                goal_m=goal,
                gmap=gmap,
                movement='8N', 
                occupancy_cost_factor=100  )

        return phy_path, grid_path

    def path2action(self, waypoints, start):
        start_pos = [start[2, 3], start[0, 3]]
        start_ang = np.arctan2(-start[0, 0], -start[0, 2]) * 180 / np.pi

        poses = np.concatenate([np.array([start_pos]), np.array(waypoints)], axis=0)
        ang = np.arctan2(poses[1:, 0] - poses[:-1, 0],
                         poses[1:, 1] - poses[:-1, 1]) * 180 / np.pi

        angs = np.concatenate([np.expand_dims(start_ang, 0), ang], axis=0)
        rots = angs[1:] - angs[:-1]
        rots[np.abs(rots) > 180] -= np.sign(rots[np.abs(rots) > 180]) * 360
        rots = np.around(rots, decimals=1)

        dists = np.sqrt(((poses[1:, ] - poses[:-1]) ** 2).sum(axis=1))

        actions = []
        forward_counter = 0
        for r, dist in zip(rots, dists):
            if r == 0:
                forward_counter += dist
            else:
                actions +=  ['move_forward'] * int(forward_counter / self.move_meter)  
                if r < 0:
                    actions += ['turn_left'] * int( np.abs(r) // self.turn_angle)  #+ ['move_forward'] * int(np.ceil(dist))
                elif r > 0:
                    actions += ['turn_right'] * int( np.abs(r) // self.turn_angle) #+ ['move_forward'] * int(np.ceil(dist))
                forward_counter = dist

        if forward_counter != 0:
            actions +=  ['move_forward'] * int(forward_counter / self.move_meter)  

        return actions








src/agents/planners.py
embodied-depth/Embodied-Depth-Prediction
#!/usr/bin/env python3

from typing import List, Optional, Union

import yacs.config

# Default Habitat config node
class Config(yacs.config.CfgNode):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs, new_allowed=True)


CN = Config

DEFAULT_CONFIG_DIR = "configs/"
CONFIG_FILE_SEPARATOR = ","

# -----------------------------------------------------------------------------
# Config definition
# -----------------------------------------------------------------------------
_C = CN()
_C.SEED = 100


def get_config(
    config_paths: Optional[Union[List[str], str]] = None,
    opts: Optional[list] = None,
) -> CN:
    r"""Create a unified config with default values overwritten by values from
    :p:`config_paths` and overwritten by options from :p:`opts`.

    :param config_paths: List of config paths or string that contains comma
        separated list of config paths.
    :param opts: Config options (keys, values) in a list (e.g., passed from
        command line into the config. For example,
        :py:`opts = ['FOO.BAR', 0.5]`. Argument can be used for parameter
        sweeping or quick tests.
    """
    config = _C.clone()
    if config_paths:
        if isinstance(config_paths, str):
            if CONFIG_FILE_SEPARATOR in config_paths:
                config_paths = config_paths.split(CONFIG_FILE_SEPARATOR)
            else:
                config_paths = [config_paths]

        for config_path in config_paths:
            config.merge_from_file(config_path)

    if opts:
        config.merge_from_list(opts)

    config.freeze()
    return config
src/config.py
embodied-depth/Embodied-Depth-Prediction
from .indoor_dataset import MHabitatRAWDataset
from .robo_dataset import MRealRAWDataset
from .embodied_dataset import EmDataset
from .embodied_realdataset import EmRealDataset
src/datasets/__init__.py
embodied-depth/Embodied-Depth-Prediction
from __future__ import absolute_import, division, print_function

import os
import random
import numpy as np
import copy
from PIL import Image  # using pillow-simd for increased speed

import torch
import torch.utils.data as data
import torch.nn.functional as F
from torchvision import transforms

def pil_loader(path):
    # open path as file to avoid ResourceWarning
    # (https://github.com/python-pillow/Pillow/issues/835)
    with open(path, 'rb') as f:
        with Image.open(f) as img:
            return img.convert('RGB')

def depth_to_disp(_depth, min_depth, max_depth):
    mask = _depth != 0

    depth = torch.clamp(_depth, min_depth, max_depth)
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = 1 / depth
    disp_ = (scaled_disp - min_disp) / (max_disp - min_disp) * mask
    return disp_

class MultiMonoDataset(data.Dataset):
    """Superclass for monocular dataloaders

    Args:
        data_path
        filenames
        height
        width
        frame_idxs
        num_scales
        is_train
        img_ext
    """
    def __init__(self,
                 data_path,
                 filenames,
                 height,
                 width,
                 frame_idxs,
                 num_scales,
                 is_train=False,
                 img_ext='.jpg',
                 load_pose=False,
                 input_depth=False,
                 num_past_frame_input=1,
                 video_interval=1,
                 input_T_dim=None):
        super(MultiMonoDataset, self).__init__()

        self.data_path = data_path
        self.filenames = filenames
        self.height = height
        self.width = width
        self.num_scales = num_scales
        self.interp = Image.ANTIALIAS

        self.frame_idxs = frame_idxs
        self.num_past_frame_input = num_past_frame_input
        self.video_interval = video_interval
        self.input_T_dim = input_T_dim

        self.read_center_id = -min([min(frame_idxs), -video_interval * num_past_frame_input]) \
                        if len(self.input_T_dim) > 1 else \
                       -min(frame_idxs) + video_interval * num_past_frame_input
        self.read_len = self.read_center_id + max(self.frame_idxs) + 1

        self.is_train = is_train
        self.img_ext = img_ext

        self.loader = pil_loader
        self.to_tensor = transforms.ToTensor()
        self.input_depth = input_depth

        # We need to specify augmentations differently in newer versions of torchvision.
        # We first try the newer tuple version; if this fails we fall back to scalars
        try:
            self.brightness = (0.8, 1.2)
            self.contrast = (0.8, 1.2)
            self.saturation = (0.8, 1.2)
            self.hue = (-0.1, 0.1)
            transforms.ColorJitter.get_params(
                self.brightness, self.contrast, self.saturation, self.hue)
        except TypeError:
            self.brightness = 0.2
            self.contrast = 0.2
            self.saturation = 0.2
            self.hue = 0.1

        self.resize = {}
        for i in range(self.num_scales):
            s = 2 ** i
            self.resize[i] = transforms.Resize((self.height // s, self.width // s),
                                               interpolation=self.interp)

        self.load_depth = self.check_depth()
        self.load_pose  = load_pose

    def preprocess(self, _colors, _depths, color_aug):
        """Resize colour images to the required scales and augment if required

        We create the color_aug object in advance and apply the same augmentation to all
        images in this item. This ensures that all images input to the pose network receive the
        same augmentation.
        """
        outputs = {}
        outputs['depth_gt'] = torch.tensor(_depths[self.read_center_id]).unsqueeze(0)

        for s in range(self.num_scales):
            colors = [self.to_tensor(self.resize[s](color)) for color in _colors]
            depths = torch.tensor(np.stack(_depths)).unsqueeze(1)
            #depths = [self.to_tensor(self.resize[s](depth)) for depth in _depths]
            for idx in self.frame_idxs:
                # GT color and depth for warpping and eval
                outputs[('color', idx, s)] = colors[self.read_center_id + idx]
                outputs[('color_aug', idx, s)] = torch.stack([color_aug(colors[self.read_center_id + idx])])

            for idx in self.input_T_dim:
                # Color as input
                read_idxs = slice(self.read_center_id - self.num_past_frame_input * self.video_interval,
                                 self.read_center_id + 1,
                                  self.video_interval)
                outputs[('color_aug', idx, s)] = torch.stack([color_aug(c) for c in colors[read_idxs]])
                outputs[('color', 0)] = torch.stack([c for c in colors[read_idxs]])[:-1]

                # Depth as input
                if self.input_depth:
                    depth_input = torch.stack([d for d in depths[read_idxs]]) 
                    if self.depth_noise:
                        depth_input = depth_input + 0.05 * depth_input / 2. * torch.randn_like(depth_input)
                        depth_input[depth_input > 4.] = 0
                        outputs[('depth_input', idx, s)] = depth_to_disp(depth_input, 0.1, 40)
                    else:
                        depth_input[depth_input > 4.] = 0
                        outputs[('depth_input', idx, s)] = depth_to_disp(depth_input, 0.1, 40)

        return outputs

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        """Returns a single training item from the dataset as a dictionary.

        Values correspond to torch tensors.
        Keys in the dictionary are either strings or tuples:

            ("color", <frame_id>, <scale>)          for raw colour images, with shape (B, 1, 3, H, W)
            ("color_aug", <frame_id>, <scale>)      for augmented colour images with shape (B, num_past_frame_input, 3, H, W),
            ("K", scale) or ("inv_K", scale)        for camera intrinsics,
            "stereo_T"                              for camera extrinsics, and
            "depth_gt"                              for ground truth depth maps.

        <frame_id> is either:
            an integer (e.g. 0, -1, or 1) representing the temporal step relative to 'index',
        or
            "s" for the opposite image in the stereo pair.

        <scale> is an integer representing the scale of the image relative to the fullsize image:
            -1      images at native resolution as loaded from disk
            0       images resized to (self.width,      self.height     )
            1       images resized to (self.width // 2, self.height // 2)
            2       images resized to (self.width // 4, self.height // 4)
            3       images resized to (self.width // 8, self.height // 8)
        """
        inputs = {}
        # adjusting intrinsics to match each scale in the pyramid
        for scale in range(self.num_scales):
            K = self.K.copy()
            K[0, :] *= 1 // (2 ** scale)
            K[1, :] *= 1 // (2 ** scale)
            inv_K = np.linalg.pinv(K)

            inputs[("K", scale)] = torch.from_numpy(K)
            inputs[("inv_K", scale)] = torch.from_numpy(inv_K)

        # Read Color and Depth
        do_color_aug = self.is_train and random.random() > 0.5
        #do_color_aug = False
        color_aug = transforms.ColorJitter(
                        self.brightness, 
                        self.contrast, 
                        self.saturation, 
                        self.hue) if do_color_aug else (lambda x: x)
        ### do_flip = self.is_train and random.random() > 0.5
        do_flip = False

        line = self.filenames[index].split()
        folder = line[0]

        if len(line) == 3:
            frame_index = int(line[1].split('_')[0])

        colors = self.get_color(folder, frame_index)
        depths = self.get_depth(folder, frame_index)

        inputs.update(self.preprocess(colors, depths, color_aug))

        # Read pose for warpping and input
        warp_pose, warp_pose_inv, v_pose, poses = self.get_pose(folder, frame_index)
        for i, f in enumerate(self.frame_idxs):
            inputs[("cam_T_cam", 0, f)] = warp_pose[i]
            if warp_pose_inv is not None:
                inputs[("cam_T_cam_inv", 0, f)] = warp_pose_inv[i]

        inputs['v_poses'] = torch.tensor(np.stack(v_pose))

        # For optical flow input
        inputs['flowpose'] = poses

        return inputs

    def get_color(self, folder, frame_index):
        raise NotImplementedError

    def check_depth(self):
        raise NotImplementedError

    def get_depth(self, folder, frame_index):
        raise NotImplementedError

    def get_pose(self, folder, frame_index):
        raise NotImplementedError


src/datasets/base_dataset.py
embodied-depth/Embodied-Depth-Prediction
from __future__ import absolute_import, division, print_function

import os
import random
import numpy as np
import copy
from PIL import Image  # using pillow-simd for increased speed

import torch
import torch.utils.data as data
import torch.nn.functional as F
from torchvision import transforms
#from wrapper.my_utils import cam1_to_cam0
from src.my_utils import correct_yz, cam1_to_cam0



def traj_loader(path, start_idx):
    with open(path, 'r') as f:
        lines = f.readlines()
        out_ = np.zeros((len(lines), 4 ,4))
        for i, l in enumerate(lines):
            num = l.split(' ')[1:]
            out_[i, :3, :] = np.array([float(n) for n in num ]).reshape(-1, 4)
            out_[i, 3, 3] = 1

    return out_

class EmDataset(data.Dataset):
    """Superclass for monocular dataloaders

    Args:
        height
        width
        num_scales
        is_train
        img_ext
    """
    def __init__(self,
                data,
                frame_idxs,
                 height,
                 width,
                 num_scales,
                 is_train=True,
                 load_pose=False,
                 input_depth=False):
        super(EmDataset, self).__init__()

        images = np.stack([d['color_sensor'] for d in data])
        depths = np.stack([d['depth_sensor'] for d in data])
        poses  = [ correct_yz(d['Ext']) for d in data]  

        self.images = images
        self.depths = depths
        self.poses = poses
        self.height = height
        self.width = width
        self.num_scales = num_scales
        self.interp = Image.ANTIALIAS

        self.frame_idxs = frame_idxs
        self.frame_interval = max(frame_idxs)

        self.to_tensor = transforms.ToTensor()
        self.input_depth = input_depth
        self.is_train = is_train

        # We need to specify augmentations differently in newer versions of torchvision.
        # We first try the newer tuple version; if this fails we fall back to scalars
        try:
            self.brightness = (0.8, 1.2)
            self.contrast = (0.8, 1.2)
            self.saturation = (0.8, 1.2)
            self.hue = (-0.1, 0.1)
            transforms.ColorJitter.get_params(
                self.brightness, self.contrast, self.saturation, self.hue)
        except TypeError:
            self.brightness = 0.2
            self.contrast = 0.2
            self.saturation = 0.2
            self.hue = 0.1

        self.resize = {}
        for i in range(self.num_scales):
            s = 2 ** i
            self.resize[i] = transforms.Resize((self.height // s, self.width // s),
                                               interpolation=self.interp)

        self.load_depth = True
        self.load_pose  = load_pose

        self.K = np.array([[320, 0, 320, 0],
                           [0,  96, 96, 0],
                           [0,   0,  1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)

        self.inputs_memo = {}
        for scale in range(self.num_scales):
            K = self.K.copy()

            K[0, :] = K[0, :] // (2 ** scale)
            K[1, :] = K[1, :]  // (2 ** scale)

            inv_K = np.linalg.pinv(K)

            self.inputs_memo[("K", scale)] = torch.from_numpy(K)
            self.inputs_memo[("inv_K", scale)] = torch.from_numpy(inv_K)

    def preprocess(self, inputs, color_aug):
        """Resize colour images to the required scales and augment if required

        We create the color_aug object in advance and apply the same augmentation to all
        images in this item. This ensures that all images input to the pose network receive the
        same augmentation.
        """
        for k in list(inputs):
            frame = inputs[k]
            if "color" in k:
                n, im, i = k
                for i in range(self.num_scales):
                    inputs[(n, im, i)] = self.resize[i](inputs[(n, im, i - 1)])

        for k in list(inputs):
            f = inputs[k]
            if "color" in k:
                n, im, i = k
                inputs[(n, im, i)] = self.to_tensor(f)
                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f)).unsqueeze(0)
            elif "color_pass" in k:
                n, im, i = k
                inputs[(n, im, i)] = self.to_tensor(f)
                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f))


    def __len__(self):
        assert self.images.shape[0] == self.depths.shape[0]
        return self.images.shape[0] - 2 * self.frame_interval

    def __getitem__(self, index):
        """Returns a single training item from the dataset as a dictionary.

        Values correspond to torch tensors.
        Keys in the dictionary are either strings or tuples:

            ("color", <frame_id>, <scale>)          for raw colour images,
            ("color_aug", <frame_id>, <scale>)      for augmented colour images,
            ("K", scale) or ("inv_K", scale)        for camera intrinsics,
            "stereo_T"                              for camera extrinsics, and
            "depth_gt"                              for ground truth depth maps.

        <frame_id> is either:
            an integer (e.g. 0, -1, or 1) representing the temporal step relative to 'index',
        or
            "s" for the opposite image in the stereo pair.

        <scale> is an integer representing the scale of the image relative to the fullsize image:
            -1      images at native resolution as loaded from disk
            0       images resized to (self.width,      self.height     )
            1       images resized to (self.width // 2, self.height // 2)
            2       images resized to (self.width // 4, self.height // 4)
            3       images resized to (self.width // 8, self.height // 8)
        """
        frame_index = index + self.frame_interval # skip index=0 and -1, cuz where no 
        inputs = {}

        do_color_aug = self.is_train and random.random() > 0.5
        #do_flip = self.is_train and random.random() > 0.5
        do_flip = False

        frame_idxs = self.frame_idxs #[0, -self.frame_interval, self.frame_interval]
        for i in frame_idxs:
            color = Image.fromarray(self.images[frame_index + i]).convert('RGB')
            if do_flip:
                color = color.transpose(Image.FLIP_LEFT_RIGHT)
            inputs[("color", i, -1)] = color

        for scale in range(self.num_scales):
            inputs[("K", scale)] = self.inputs_memo[("K", scale)].clone()
            inputs[("inv_K", scale)] = self.inputs_memo[("inv_K", scale)].clone()


        if do_color_aug:
            color_aug = transforms.ColorJitter(
                self.brightness, self.contrast, self.saturation, self.hue)
        else:
            color_aug = (lambda x: x)

        depth_gt = self.depths[frame_index]
        inputs["depth_gt"] = np.expand_dims(depth_gt, 0)
        inputs["depth_gt"] = torch.from_numpy(inputs["depth_gt"].astype(np.float32))

        self.preprocess(inputs, color_aug)
        for i in frame_idxs:
            del inputs[("color", i, -1)]
            del inputs[("color_aug", i, -1)]

        if self.load_pose:
            for i in frame_idxs[1:]:
                inputs[("cam_T_cam", 0, i)] = torch.from_numpy(cam1_to_cam0(self.poses[frame_index + i],
                                                        self.poses[frame_index]))
                inputs[("cam_T_cam_inv", 0, i)] = torch.from_numpy(cam1_to_cam0(self.poses[frame_index], self.poses[frame_index+i]))

            if len(self.frame_idxs) > 3: # hack: for video input
                poses = [ self.poses[frame_index + f] for f in self.frame_idxs]
                poses_ref_in_other = np.stack([(np.matmul(np.linalg.pinv(pose), poses[0])) for pose in poses[3:]]) # hack: skip 0 -i i 
                inputs['flowpose'] = poses_ref_in_other


        return inputs


src/datasets/embodied_dataset.py
embodied-depth/Embodied-Depth-Prediction
from __future__ import absolute_import, division, print_function

import os
import random
import numpy as np
import copy
from PIL import Image  # using pillow-simd for increased speed

import torch
import torch.utils.data as data
import torch.nn.functional as F
from torchvision import transforms
#from wrapper.my_utils import cam1_to_cam0

def cam1_to_cam0(T_world_camera1, T_world_camera0):
    T_camera1_world = np.linalg.pinv(T_world_camera1)
    T_camera1_camera0 = np.matmul(T_camera1_world, T_world_camera0)
    return T_camera1_camera0


def traj_loader(path, start_idx):
    with open(path, 'r') as f:
        lines = f.readlines()
        out_ = np.zeros((len(lines), 4 ,4))
        for i, l in enumerate(lines):
            num = l.split(' ')[1:]
            out_[i, :3, :] = np.array([float(n) for n in num ]).reshape(-1, 4)
            out_[i, 3, 3] = 1

    return out_

class EmRealDataset(data.Dataset):
    """Superclass for monocular dataloaders

    Args:
        height
        width
        num_scales
        is_train
        img_ext
    """
    def __init__(self,
                data,
                frame_idxs,
                 height,
                 width,
                 num_scales,
                 is_train=True,
                 load_pose=False,
                 input_depth=False):
        super(EmRealDataset, self).__init__()

        images = data['color']  
        depths = data['depth'] 
        poses  = data['Ext'] 

        self.images = images
        self.depths = depths
        self.poses = poses
        self.height = height
        self.width = width
        self.num_scales = num_scales
        self.interp = Image.ANTIALIAS

        self.frame_idxs = frame_idxs
        self.frame_interval = max(frame_idxs)

        self.to_tensor = transforms.ToTensor()
        self.input_depth = input_depth
        self.is_train = is_train

        # We need to specify augmentations differently in newer versions of torchvision.
        # We first try the newer tuple version; if this fails we fall back to scalars
        try:
            self.brightness = (0.8, 1.2)
            self.contrast = (0.8, 1.2)
            self.saturation = (0.8, 1.2)
            self.hue = (-0.1, 0.1)
            transforms.ColorJitter.get_params(
                self.brightness, self.contrast, self.saturation, self.hue)
        except TypeError:
            self.brightness = 0.2
            self.contrast = 0.2
            self.saturation = 0.2
            self.hue = 0.1

        self.resize = {}
        for i in range(self.num_scales):
            s = 2 ** i
            self.resize[i] = transforms.Resize((self.height // s, self.width // s),
                                               interpolation=self.interp)

        self.load_depth = True
        self.load_pose  = load_pose

        self.K = np.array([[1, 0, 0, 0],
                           [0, 1, 0, 0],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)

        self.inputs_memo = {}
        for scale in range(self.num_scales):
            K = self.K.copy()

            #K[0, :] *= self.width // (2 ** scale)
            #K[1, :] *= self.height // (2 ** scale)
            K[0, :] *= 1 // (2 ** scale)
            K[1, :] *= 1 // (2 ** scale)

            inv_K = np.linalg.pinv(K)

            self.inputs_memo[("K", scale)] = torch.from_numpy(K)
            self.inputs_memo[("inv_K", scale)] = torch.from_numpy(inv_K)
        
        self.length = None
        self.length = self.__len__()

    def preprocess(self, inputs, color_aug):
        """Resize colour images to the required scales and augment if required

        We create the color_aug object in advance and apply the same augmentation to all
        images in this item. This ensures that all images input to the pose network receive the
        same augmentation.
        """
        for k in list(inputs):
            frame = inputs[k]
            if "color" in k:
                n, im, i = k
                for i in range(self.num_scales):
                    inputs[(n, im, i)] = self.resize[i](inputs[(n, im, i - 1)])

        for k in list(inputs):
            f = inputs[k]
            if "color" in k:
                n, im, i = k
                inputs[(n, im, i)] = self.to_tensor(f)
                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f)).unsqueeze(0)
            elif "color_pass" in k:
                n, im, i = k
                inputs[(n, im, i)] = self.to_tensor(f)
                inputs[(n + "_aug", im, i)] = self.to_tensor(color_aug(f))


    def __len__(self):
        if self.length == None: 
            self.length = 0
            self.len_mark = []
            for img_list, depth_list, pose_list in zip(self.images, self.depths, self.poses):
                assert img_list.shape[0] == depth_list.shape[0]
                assert img_list.shape[0] == pose_list.shape[0]
                self.length += img_list.shape[0] - 2 * self.frame_interval
                self.len_mark.append(self.length)

        return self.length

    def __getitem__(self, index):
        """Returns a single training item from the dataset as a dictionary.

        Values correspond to torch tensors.
        Keys in the dictionary are either strings or tuples:

            ("color", <frame_id>, <scale>)          for raw colour images,
            ("color_aug", <frame_id>, <scale>)      for augmented colour images,
            ("K", scale) or ("inv_K", scale)        for camera intrinsics,
            "stereo_T"                              for camera extrinsics, and
            "depth_gt"                              for ground truth depth maps.

        <frame_id> is either:
            an integer (e.g. 0, -1, or 1) representing the temporal step relative to 'index',
        or
            "s" for the opposite image in the stereo pair.

        <scale> is an integer representing the scale of the image relative to the fullsize image:
            -1      images at native resolution as loaded from disk
            0       images resized to (self.width,      self.height     )
            1       images resized to (self.width // 2, self.height // 2)
            2       images resized to (self.width // 4, self.height // 4)
            3       images resized to (self.width // 8, self.height // 8)
        """
        group_index = 0
        for i in range(len(self.len_mark) - 1):
            if self.len_mark[i+1] > index:
                frame_index = index - self.len_mark[i]
            else:
                group_index += 1
                if group_index == len(self.len_mark) - 1:
                    frame_index = index - self.len_mark[-1]

        frame_index += self.frame_interval # skip index=0 and -1, cuz where no 
        inputs = {}

        do_color_aug = self.is_train and random.random() > 0.5
        #do_flip = self.is_train and random.random() > 0.5
        do_flip = False

        frame_idxs = self.frame_idxs #[0, -self.frame_interval, self.frame_interval]
        for i in frame_idxs:
            color = Image.fromarray(self.images[group_index][frame_index + i]).convert('RGB')
            if do_flip:
                color = color.transpose(Image.FLIP_LEFT_RIGHT)
            inputs[("color", i, -1)] = color

        for scale in range(self.num_scales):
            inputs[("K", scale)] = self.inputs_memo[("K", scale)].clone()
            inputs[("inv_K", scale)] = self.inputs_memo[("inv_K", scale)].clone()


        if do_color_aug:
            color_aug = transforms.ColorJitter(
                self.brightness, self.contrast, self.saturation, self.hue)
        else:
            color_aug = (lambda x: x)

        depth_gt = self.depths[group_index][frame_index]
        inputs["depth_gt"] = np.expand_dims(depth_gt, 0)
        inputs["depth_gt"] = torch.from_numpy(inputs["depth_gt"].astype(np.float32))

        self.preprocess(inputs, color_aug)
        for i in frame_idxs:
            del inputs[("color", i, -1)]
            del inputs[("color_aug", i, -1)]

        if self.load_pose:
            for i in frame_idxs[1:]:
                inputs[("cam_T_cam", 0, i)] = torch.from_numpy(cam1_to_cam0(self.poses[frame_index + i],
                                                        self.poses[frame_index]))

            poses = [ self.poses[frame_index + f]  for f in self.frame_idxs]
            poses_ref_in_other = np.stack([(np.matmul(np.linalg.inv(pose), poses[0])) for pose in poses[3:]])
            inputs['flowpose'] = poses_ref_in_other


        return inputs


src/datasets/embodied_realdataset.py
embodied-depth/Embodied-Depth-Prediction
from __future__ import absolute_import, division, print_function
import sys
#from pathlib import Path
#base_dir = Path(__file__).absolute().parent.parent
#sys.path.append(str(base_dir))

import os
import numpy as np

from src.datasets.base_dataset import MultiMonoDataset
from src.my_utils import correct_yz, cam1_to_cam0


class MHabitatRAWDataset(MultiMonoDataset):
    """Superclass for different types of KITTI dataset loaders
    """
    def __init__(self, *args, **kwargs):
        super(MHabitatRAWDataset, self).__init__(*args, **kwargs)

        # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size.
        # To normalize you need to scale the first row by 1 / image_width and the second row
        # by 1 / image_height. Monodepth2 assumes a principal point to be exactly centered.
        # If your principal point is far from the center you might need to disable the horizontal
        # flip augmentation.
        # Takes from this issue https://github.com/facebookresearch/habitat-sim/issues/80
        self.K = np.array([[320, 0, 320, 0],
                           [0,  96, 96, 0],
                           [0,   0,  1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)

        #self.K[0, :] *= self.width
        #self.K[1, :] *= self.height
        self.full_res_shape = (640, 192)
        self.side_map = {"2": 2, "3": 3, "l": 2, "r": 3}
        self.depth_noise = True
        self.exts = {}
        self.color_memo = {}
        self.dep_memo = {}

    def check_depth(self):
        return True

    def get_image_path(self, folder, frame_index ):
        f_str = "color/{:08d}{}".format(frame_index, self.img_ext)
        image_path = os.path.join(
            self.data_path, folder, f_str)
        return image_path

    def get_color(self, folder, frame_index):
        colors = []
        for i in range(self.read_len):
            idx = frame_index + i - self.read_center_id
            color = self.loader(self.get_image_path(folder, idx))
            colors.append(color)
        return colors

    def get_depth(self, folder, frame_index):
        depths = []
        for i in range(self.read_len):
            idx = frame_index + i - self.read_center_id
            f_str = "depth/{:08d}.npy".format(idx)
            depth_path = os.path.join(
                        self.data_path, folder, f_str)
            depth = np.load(depth_path)
            depths.append(depth.astype(np.float32))

        return depths

    def get_pose(self, folder, frame_index):
        '''
        return: 
        warp_pose: relative pose used in the loss computation
        video_pose: relative pose used when inputting video and its corresponding poses
        poses_ref_in_other: relative pose used for deriving depth from optical flows

        '''

        if folder not in self.exts.keys():
            pose_dir = os.path.join(self.data_path, folder, 'poses.npy')
            self.exts[folder] = correct_yz(np.load(pose_dir))

        warp_pose = []
        warp_pose_inv = []
        for f in self.frame_idxs:
            warp_pose.append(cam1_to_cam0(
                self.exts[folder][frame_index + f ],
                 self.exts[folder][frame_index ],
            ))
            warp_pose_inv.append(cam1_to_cam0(
                self.exts[folder][frame_index], 
                self.exts[folder][frame_index+f]
                ))
           


        video_pose = []
        for i in self.input_T_dim:
            buf = []
            for f in range(-self.num_past_frame_input * self.video_interval, 1, self.video_interval):
                buf.append(cam1_to_cam0(
                    self.exts[folder][frame_index + i  ] ,
                    self.exts[folder][frame_index + i + f]
                ))

            video_pose.append(buf)

        poses = [ (self.exts[folder][frame_index + f]) for f in self.frame_idxs]
        if len(poses) > 3:
            poses_ref_in_other = np.stack([(np.matmul(np.linalg.pinv(pose), poses[0])) for pose in poses[3:]]) # (#ofRefFrames, 4, 4)
        else:
            poses_ref_in_other = np.stack(poses)

        return warp_pose, warp_pose_inv, video_pose, poses_ref_in_other



if __name__ == '__main__':
    from utils import readlines
    from torch.utils.data import DataLoader

    data_path = '/Users/mac/Desktop/Habitat/interactivity_output'
    fpath = '/Users/mac/Desktop/Habitat/all_files.txt'
    train_filenames = readlines(fpath)

    train_dataset = IndoorRAWDataset(
            data_path, train_filenames, 544, 720,
            [0, -1, 1], 4, is_train=True, img_ext='.png')
    train_loader = DataLoader(
        train_dataset, 2, True,
        num_workers=0, pin_memory=True, drop_last=True)

    for i in train_loader:
        print(i.keys())
        break

src/datasets/indoor_dataset.py
embodied-depth/Embodied-Depth-Prediction
#import sys
#from pathlib import Path
#base_dir = Path(__file__).absolute().parent.parent
#sys.path.append(str(base_dir))

import os
import skimage.transform
import numpy as np
import torch
import PIL.Image as pil
from scipy.signal import medfilt2d

from src.datasets.base_dataset import MultiMonoDataset

def cam1_to_cam0(_world_camera1, _world_camera0):
    T_camera1_world = np.linalg.pinv(_world_camera1)
    T_camera1_camera0 = np.matmul(T_camera1_world, _world_camera0)
    return T_camera1_camera0

def traj_loader(path, start_idx):
    with open(path, 'r') as f:
        lines = f.readlines()
        out_ = np.zeros((len(lines), 4 ,4), dtype=np.float32)
        for i, l in enumerate(lines):
            num = l.split(' ')[1:]
            out_[i, :3, :] = np.array([float(n) for n in num ]).reshape(-1, 4)
            out_[i, 3, 3] = 1

    return out_

class MRealRAWDataset(MultiMonoDataset):
    """Superclass for different types of KITTI dataset loaders
    """
    def __init__(self, *args, **kwargs):
        super(MRealRAWDataset, self).__init__(*args, **kwargs)

        self.use_odom = False
        # focal length 904.62 mm
        # NOTE: Make sure your intrinsics matrix is *normalized* by the original image size.
        # To normalize you need to scale the first row by 1 / image_width and the second row
        # by 1 / image_height. Monodepth2 assumes a principal point to be exactly centered.
        # If your principal point is far from the center you might need to disable the horizontal
        # flip augmentation.
        # Takes from this issue https://github.com/facebookresearch/habitat-sim/issues/80
        self.K = np.array([[904.62 , 0,     640., 0],
                           [0,      904.62 , 360., 0],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)

        self.full_res_shape = (640, 192) # (1024, 320) #(640, 192)
        self.K[0] = self.K[0]   * self.full_res_shape[0] / 1280.
        self.K[1] = self.K[1]   * self.full_res_shape[1] / 720.

        self.side_map = {"2": 2, "3": 3, "l": 2, "r": 3}
        self.exts = {}
        self.color_memo = {}
        self.dep_memo = {}
        self.depth_noise = False

    def check_depth(self):
        return True

    def get_image_path(self, folder, frame_index):
        f_str = "color/frame{}{}".format(frame_index, self.img_ext)
        image_path = os.path.join(
            self.data_path, folder, f_str)
        return image_path

    def get_color(self, folder, frame_index):
        colors = []
        for i in range(self.read_len):
            idx = frame_index + i - self.read_center_id
            color = self.loader(self.get_image_path(folder, idx))
            colors.append(color)

        return colors

    def get_depth(self, folder, frame_index):
        depths = []
        for i in range(self.read_len):
            idx = frame_index + i - self.read_center_id
            f_str = "depth/dframe{}.npy".format(frame_index + i - self.read_center_id)
            depth_path = os.path.join(
                   self.data_path, folder, f_str)
            depth = np.load(depth_path) / 1000.
            depth = skimage.transform.resize(
                    depth, self.full_res_shape[::-1], order=0, preserve_range=True, mode='constant')
            depths.append(depth.astype(np.float32))

        return depths

    def get_pose(self, folder, frame_index):
        if folder not in self.exts.keys():
            if self.use_odom:
                pose_dir = os.path.join(self.data_path, folder, 'odom.npy')
            else:
                pose_dir = os.path.join(self.data_path, folder, 'poses.npy')
            self.exts[folder] = np.load(pose_dir).astype(np.float32)

        warp_pose = []
        video_pose = []
        for f in self.frame_idxs:
            warp_pose.append(cam1_to_cam0(
                self.exts[folder][frame_index + f ],
                 self.exts[folder][frame_index ],
            ))

        for i in self.input_T_dim:
            buf = []
            for f in range(-self.num_past_frame_input * self.video_interval, 1, self.video_interval):
                buf.append(cam1_to_cam0(
                    self.exts[folder][frame_index + i + f ],
                    self.exts[folder][frame_index + i],
                ))
            video_pose.append(buf)

        poses = [ self.exts[folder][i + frame_index] for i in self.frame_idxs]
        if len(poses) > 3:
            poses_ref_in_other = np.stack([(np.matmul(np.linalg.inv(pose), poses[0])) for pose in poses[3:]]) # (#ofRefFrames, 4, 4)
        else:
            #print("warning")
            poses_ref_in_other = np.stack(poses)

        return warp_pose, None, video_pose, poses_ref_in_other


  

src/datasets/robo_dataset.py
embodied-depth/Embodied-Depth-Prediction
import os
import numpy as np
from matplotlib import pyplot as plt
import cv2

#import imageio.v2 as imageio
try:
    import habitat_sim
    import magnum as mn
    from habitat.utils.visualizations import maps
except ImportError:
    pass

try:
    from src.ros import DataListener, Controller
except ImportError:
    print('no rospy in lib')

class CollisionError(Exception):
    def __init__(self, error_info=None):
        super().__init__(self)
        self.error_info = error_info

    def __str__(self):
        return self.error_info


class SimEnv:
    def __init__(self, config, scene_id):
        self.config = config
        base_dir = self.config.SCENE_PATH
        self.SCENE_MAP = \
            {
                0: base_dir + 'mp3d/JeFG25nYj2p/JeFG25nYj2p.glb',
                1: base_dir + 'gibson/Herricks.glb', 
                2: base_dir + 'mp3d/17DRP5sb8fy/17DRP5sb8fy.glb',
                3: base_dir + 'gibson/Eastville.glb',
                4: base_dir + 'mp3d/2t7WUuJeko7/2t7WUuJeko7.glb',
                5: base_dir + 'gibson/Corder.glb',
            }

        self.type = 'sim'
        # Create the simulator
        self.SCENE_PATH = self.SCENE_MAP[scene_id]
        cfg = self._get_env_cfg()
        self._sim = habitat_sim.Simulator(cfg)
        self._sim.seed(self.config.SEED)

        self.physical_bounds = self._sim.pathfinder.get_bounds()

        # Create the agent to the sim
        self.agent = self._sim.initialize_agent(self.config.ENV.DEFAULT_AGENT)

        self.vel_control = habitat_sim.physics.VelocityControl()
        self.vel_control.controlling_lin_vel = True
        self.vel_control.lin_vel_is_local = True
        self.vel_control.controlling_ang_vel = True
        self.vel_control.ang_vel_is_local = True

        self._hist_turn = 'turn_right'

    def _check_collision(self, _depth):
        H, W = _depth.shape
        return (_depth[H // 2 - 10 : H // 2 + 10, W//2 - 10 : W // 2 + 10] < 0.1).all()
        
    def _step(self, action):
        if action in ['turn_left', 'turn_right']:
            observations = self._sim.step(action)
            #observations = self._sim.step('move_forward')
        else:
            observations = self._sim.step(action)
        state = self.agent.get_state()
        observations['Ext'] = np.array(self.agent.scene_node.transformation_matrix())
        observations['pos'] = state.position
        observations['angle'] = state.rotation
        return [observations]

    def step(self, action):
        '''
        Input: lists of actions 
        Return: 
            obs_: lists of observations
            collided_flag: whether collision occurs
        '''
        obs_ = []
        if type(action) == list:
            for i, a in enumerate(action):
                obs_ += self.step(a)
                if self._check_collision(obs_[-1]['depth_sensor']) or obs_[-1]['collided']:
                    return obs_, True
            return obs_, False
        else:
            obs_ += self._step(action)
            if action in ['turn_left', 'turn_right']:
                self._hist_turn = action
            return obs_

    def reset(self):
        pass
    
    def _get_env_cfg(self):
        sim_cfg = habitat_sim.SimulatorConfiguration()
        sim_cfg.scene_id = self.SCENE_PATH
        sim_cfg.gpu_device_id = self.config.ENV.CUDA_ID

        # a RGB visual sensor 
        rgb_sensor_spec = habitat_sim.CameraSensorSpec()
        rgb_sensor_spec.uuid = "color_sensor"
        rgb_sensor_spec.sensor_type = habitat_sim.SensorType.COLOR
        rgb_sensor_spec.resolution = [self.config.SENSOR.HEIGHT, self.config.SENSOR.WIDTH]  
        rgb_sensor_spec.position = [0.0, self.config.SENSOR.POS_HEIGHT, 0.0]
        rgb_sensor_spec.orientation = [self.config.SENSOR.PITCH, 0.0, 0.0,]
        rgb_sensor_spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE

        # a Depth sensor
        depth_sensor_1st_person_spec = habitat_sim.CameraSensorSpec()
        depth_sensor_1st_person_spec.uuid = "depth_sensor"
        depth_sensor_1st_person_spec.sensor_type = habitat_sim.SensorType.DEPTH
        depth_sensor_1st_person_spec.resolution = [self.config.SENSOR.HEIGHT, self.config.SENSOR.WIDTH]  
        depth_sensor_1st_person_spec.position = [0.0, self.config.SENSOR.POS_HEIGHT, 0.0]
        depth_sensor_1st_person_spec.orientation = [self.config.SENSOR.PITCH, 0.0, 0.0,]
        depth_sensor_1st_person_spec.sensor_subtype = habitat_sim.SensorSubType.PINHOLE

        # agent
        agent_cfg = habitat_sim.agent.AgentConfiguration()
        agent_cfg.sensor_specifications = [rgb_sensor_spec, depth_sensor_1st_person_spec]
        agent_cfg.action_space = {
        "move_forward": habitat_sim.agent.ActionSpec(
            "move_forward", habitat_sim.agent.ActuationSpec(amount=self.config.AGENT.FORWARD_MOUNT)
        ),
        "move_backward":habitat_sim.agent.ActionSpec(
            "move_backward", habitat_sim.agent.ActuationSpec(amount=-self.config.AGENT.FORWARD_MOUNT)
        ),

        "turn_left": habitat_sim.agent.ActionSpec(
            "turn_left", habitat_sim.agent.ActuationSpec(amount=self.config.AGENT.LR_ANGLE)
        ),
        "turn_right": habitat_sim.agent.ActionSpec(
            "turn_right", habitat_sim.agent.ActuationSpec(amount=self.config.AGENT.LR_ANGLE)
        ),
        "look_up": habitat_sim.agent.ActionSpec(
            "look_up", habitat_sim.agent.ActuationSpec(amount=self.config.AGENT.UD_ANGLE)
        ),
        "look_down": habitat_sim.agent.ActionSpec(
            "look_down", habitat_sim.agent.ActuationSpec(amount=self.config.AGENT.UD_ANGLE)
        ),
        }

        '''
        action_list = [
            "move_left",
            "turn_left",
            "move_right",
            "turn_right",
            "move_backward",
            "look_up",
            "move_forward",
            "look_down",
            "move_down",
            "move_up",
        ] 
        '''
        return habitat_sim.Configuration(sim_cfg, [agent_cfg])


class RealEnv:
    def __init__(self, no_return=False):
        #self.config = config

        self.no_return = no_return
        self.type = 'real'
        # Create the simulator

        self.control = Controller()
        self.datalistener = DataListener()

        self.physical_bounds = [(-15, -15, -15), (15, 15, 15) ]
        #self.physical_bounds = self._sim.pathfinder.get_bounds()


    def _check_collision(self, _depth):
        H, W = _depth.shape
        return (_depth[H // 2 - 10 : H // 2 + 10, W//2 - 10 : W // 2 + 10] < 0.1).all()
        
    def step(self, actions):
        '''
        Input: lists of actions 
        Return: 
            obs_: lists of observations
            collided_flag: whether collision occurs
        '''
        self.control.act(actions)
        if not self.no_return:
            obs_ = self.datalistener.output_data()

            return obs_
        else:
            return [None]

def test():
    cfg = get_config('configs/default_cfg.yaml')
    env = SimEnv(cfg)
    actions = ['turn_right', 'turn_right','move_forward', 'move_forward', 'move_forward', 'move_forward', 'move_forward',  'turn_right',] * 50

    obss = []
    trace = []
    for a in actions:
        obs = env.step(a)[0]
        obss.append(obs['color_sensor'])
        trace.append( obs['pos'] )

    print('*' * 20)
    print(obss[0].shape)

    m = MapStat(env)
    m.step(trace=trace)
    m.summary()

if __name__ == '__main__':
    from config import get_config
    from tensorboardX import SummaryWriter
    from torchvision.transforms import ToTensor
    from PIL import Image
    test()
src/env_wrapper.py
embodied-depth/Embodied-Depth-Prediction


from src.models.flow2depth.trainer import MyTrainer
from src.models.flow2depth.options import flow2depthOptions

#from src.models.monodepth.options import monodepthOptions
#from src.models.manydepth.options import manydepthOptions
#from src.models.scdepth.options import scdepthOptions
#from src.models.LeReS.options import leresOptions
#
#from src.models.monodepth.trainer import MonoTrainer
#from src.models.manydepth.trainer import ManyTrainer
#from src.models.scdepth.trainer import ScTrainer

src/models/__init__.py
embodied-depth/Embodied-Depth-Prediction
import os
import sys
import json
import glob
from PIL import Image  # using pillow-simd for increased speed
from pathlib import Path
from collections import deque

import time
from tqdm import tqdm
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader, Subset
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler, WeightedRandomSampler
from tensorboardX import SummaryWriter


from src.models.layers import *
from src.models.utils import *
from src.datasets  import MHabitatRAWDataset, MRealRAWDataset, EmDataset 
from src.my_utils import *


class BaseTrainer:
    def __init__(self, options):
        '''
        Compulsory-specified attribute:
            self.model
                must have two attributes 
                `frames_to_load_train` decide the number of frames from Dataset input 
                `frames_to_load_test` decide the number of frames when the model online preprocesses the data from environment
        Attention: Model should be defined before super().__init__()
        '''

        self.opt = options
        log_path = os.path.join(self.opt.log_dir, self.opt.model_type + '_' + self.opt.agent_type +  '_' + self.opt.model_name )
        if not self.opt.test:
            times = len(glob.glob( os.path.join(log_path, 'log*'), recursive=False))
            self.log_path = os.path.join(log_path,  'log' + str(times))
        else:
            self.log_path = log_path
            

        ## Val and test dataset
        self.fpath = os.path.join(str(Path(__file__).absolute().parent.parent.parent), "splits", self.opt.split, "{}_files.txt")
        val_filenames = readlines(self.fpath.format("val{}".format(self.opt.test_scene_id)))
        test_filenames = readlines(self.fpath.format("test{}".format(self.opt.test_scene_id)))

        datasets_dict = {
                         "habitat": MHabitatRAWDataset,
                         'real': MRealRAWDataset
                         }
        self.dataset = datasets_dict[self.opt.dataset]

        num_scales = 4 #if self.opt.model_type == 'manydepth' else 1

        img_ext = '.png' if self.opt.png else '.jpg'
        val_dataset = self.dataset(
            data_path=self.opt.data_path,
            filenames=val_filenames,
            height=self.opt.height,
            width=self.opt.width,
            frame_idxs=self.model.frames_to_load_train,
            num_scales=num_scales,
            is_train=False,
            img_ext=img_ext,
            input_depth=False,
            num_past_frame_input=0,
            video_interval=self.opt.video_interval,
            input_T_dim=self.opt.input_T_dim )
        self.val_loader = DataLoader(
            val_dataset, batch_size=self.opt.batch_size, shuffle=False,
            num_workers=8, pin_memory=True, drop_last=True)

        test_dataset = self.dataset(
            data_path=self.opt.data_path,
            filenames=test_filenames,
            height=self.opt.height,
            width=self.opt.width,
            frame_idxs=self.model.frames_to_load_train,
            num_scales=num_scales,
            is_train=False,
            img_ext=img_ext,
            input_depth=False,
            num_past_frame_input=0,
            video_interval=self.opt.video_interval,
            input_T_dim=self.opt.input_T_dim )
        self.test_loader = DataLoader(
            test_dataset, batch_size=self.opt.batch_size, shuffle=False,
            num_workers=12, pin_memory=True, drop_last=True)

        ## Log variable
        self.step = 0
        self.epoch = 0
        self.val_history_best = 0

    def train_offline(self):
        pass

    def set_train(self):
        pass

    def set_eval(self):
        pass

    def process_batch(self, inputs):
        '''
        return:
            outputs
            losses 
        '''
        pass

    def log(self, logger, inputs, outputs, loss):
        pass

    def save_model(self, path, name):
        '''
        Save all model ckpts and optimizer in the following path
         `[tensorboard_dir]/[self.log_path]/models/weights_{epoch}_val_{valacc}/`
        '''
        #save_folder = os.path.join(self.log_path, "models", "weights_{}".format(cur_step))
        #save_path = os.path.join(save_folder, "{}.pth".format("adam"))
        #torch.save(self.model_optimizer.state_dict(), save_path)
        pass

    def load_model(self, path):
        pass



    # Above need to specify; below has been written

    def train_one_step(self, data, logger=None, cur_step=0, dataweight=None):
        self.set_train()

        dataset = EmDataset(
                            data, 
                            self.model.frames_to_load_train,
                            self.opt.height,
                            self.opt.width,
                            self.num_scales,
                            is_train=True,
                            load_pose=(self.opt.pose_model_type == 'gt')
                            )

        dataloader = DataLoader(
            dataset, self.opt.batch_size, True,
            num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)

        loss_buf = []
        for i in range(self.opt.update_times):
            for inputs in dataloader:
                outputs, losses = self.process_batch(inputs)

                self.model_optimizer.zero_grad()
                losses["loss"].backward()
                self.model_optimizer.step()
                loss_buf.append(losses["loss"].cpu().data)

                if self.step % 1024 == 0:
                    logger.add_scalar('loss/loss', losses["loss"].cpu().data, self.step)
                    logger.add_scalar('lr_rate', self.model_lr_scheduler.get_last_lr()[0], self.step)

                early_phase = self.step % 400 == 0 and self.step < 5000
                late_phase = self.step % 4000 == 0

                if (early_phase or late_phase) and (len(outputs.keys()) > 2):
                    self.log(logger, inputs, outputs, losses)

                self.step += 1
                if self.step % 2048 == 0:
                    self.set_eval()
                    self.val(logger, self.step)
                    #self.test(logger, self.step)
                    self.set_train()

                if (self.step ) % 8000 == 0:
                    self.save_model(self.log_path, 'latest')

        self.model_lr_scheduler.step()
        self.set_eval()
        metrics = self.test(logger, self.step)
        

        acc_ = metrics['test/a1']()
        self.set_eval()
        print("Action Step {}, training step {}, lr = {}".format(cur_step, self.step, self.model_lr_scheduler.get_last_lr()))


        del dataset
        del dataloader

        # For manydepth 
        if hasattr(self.opt, "freeze_teacher_epoch" ):
            self.epoch += 1
            if self.epoch == self.opt.freeze_teacher_epoch:
                self.freeze_teacher()


        return acc_

    def cal_acc(self, b):
        with torch.no_grad():
            outputs, losses = self.process_batch(b)

            depth_pred  = outputs[('depth', 0, 0)].cpu()
            depth_gt = b['depth_gt'].cpu()

            mask = depth_gt > 0
            #scale_coef  = 1 if self.opt.dataset == 'real' else torch.median(depth_gt[mask]) / torch.median(depth_pred[mask])
            scale_coef  =  torch.median(depth_gt[mask]) / torch.median(depth_pred[mask])

            depth_pred *= scale_coef  #torch.median(depth_gt) / torch.median(depth_pred)
            #depth_pred *= torch.median(depth_gt) / torch.median(depth_pred)
            depth_pred = torch.clamp(depth_pred, min=self.opt.min_depth, max=self.opt.max_depth)
            depth_pred_ = depth_pred.clone()

            depth_gt = depth_gt[mask]
            depth_pred = depth_pred[mask]

            depth_errors_ = compute_depth_errors(depth_gt, depth_pred)
        return depth_pred_, depth_errors_, losses

    def test(self, logger, cur_step, final=False, load_best=True):

        d_videos = []
        loss_buf = AvgMeter()
        depth_metric_names = [
            "test/abs_rel", "test/sq_rel", "test/rms", "test/log_rms", "test/a1", "test/a2", "test/a3"] + \
            ["test/>5m_abs_rel", "test/>5m_sq_rel", "test/>5m_rms", "test/>5m_log_rms", "test/>5m_a1", "test/>5m_a2", "test/>5m_a3"]
        metrics = {k: AvgMeter() for k in depth_metric_names}

        if final:
            if load_best:
                self.load_best_model()

        with torch.no_grad():
            for b in tqdm(self.test_loader):
                #videos.append(b[('color', 0, 0)].cpu())
                depth_pred, depth_errors, losses = self.cal_acc(b)
                loss_buf += losses["loss"].cpu().data
                d_videos.append(depth_pred)
                for k in depth_errors.keys():
                    metric_name = k.replace('de/', 'test/').replace('da/', 'test/')
                    metrics[metric_name] += depth_errors[k].cpu().data


        d_videos = torch.cat(d_videos).unsqueeze(0)

        if final:
            if self.opt.save_pred_npy:
                np.save(os.path.join(self.log_path, 
                                     self.opt.model_name.split('/')[0] + '_'+self.opt.model_type + '_' + self.opt.agent_type + '_scene{}.npy'.format(self.opt.test_scene_id)), 
                        d_videos.numpy() )
            print('Final Metircs')
            main_table_path = os.path.join(Path(self.log_path).parent, 'table.csv')
            write_metrics = (not os.path.exists(main_table_path))
            with open(main_table_path, 'a') as f:
                if write_metrics:
                    for k in metrics.keys():
                        f.write(k + ',')
                    f.write('\n')

                for k, v in metrics.items():
                    print("{}: {} \t".format(k, v()))
                    f.write( '{},'.format(v()) )
                f.write('\n')

            # calculate mean and std so far
            summary_path = os.path.join(Path(self.log_path).parent, 'summary.csv')
            tb = np.loadtxt(main_table_path, delimiter=",", dtype=str)
            data = tb[1:, :-1].astype(float)
            name = tb[0]
            mean = data.mean(0)
            std = data.std(0)
            with open(summary_path, 'w') as f:
                for k in name:
                    f.write(k + ',')
                f.write('\n')
                for k in mean:
                    f.write('{:.03f},'.format(k))
                f.write('\n')
                for k in std:
                    f.write('{:.03f},'.format(k))
                f.write('\n')

        else:
            logger.add_scalar("Eval_loss", loss_buf(), cur_step)
            if cur_step % 7 == 0:
                logger.add_video('Eval/depth', depth_to_rgb(d_videos.numpy(), self.opt.max_depth), fps=60, global_step=cur_step)
            if self.opt.save_pred_npy:
                np.save(os.path.join(self.log_path, 
                                     '/our_pred_depth_scene{}_e{}.npy'.format(self.opt.test_scene_id, self.epoch)), 
                        d_videos.numpy() )
            for k, v in metrics.items():
                logger.add_scalar("{}".format(k), 
                                    v(),
                                    cur_step
                                    )

            return metrics

    def val(self, logger, cur_step):
        loss_buf = AvgMeter()
        depth_metric_names = [
            "val/abs_rel", "val/sq_rel", "val/rms", "val/log_rms", "val/a1", "val/a2", "val/a3"] + \
            ["val/>5m_abs_rel", "val/>5m_sq_rel", "val/>5m_rms", "val/>5m_log_rms", "val/>5m_a1", "val/>5m_a2", "val/>5m_a3"]
        metrics = {k: AvgMeter() for k in depth_metric_names}

        with torch.no_grad():
            for b in self.val_loader:
                depth_pred, depth_errors, losses = self.cal_acc(b)
                for k in depth_errors.keys():
                    metric_name = k.replace('de/', 'val/').replace('da/', 'val/')
                    metrics[metric_name] += depth_errors[k].cpu().data
                    loss_buf += losses["loss"].cpu().data

            logger.add_scalar("val_loss", loss_buf(), cur_step)
            for k, v in metrics.items():
                logger.add_scalar("{}".format(k), 
                                    v(),
                                    cur_step
                                    )

            if (self.val_history_best < metrics['val/a1']()):
                self.val_history_best = metrics['val/a1']()
                self.save_model(self.log_path, 'e_{}_val_{}'.format(cur_step, self.val_history_best ))
                self.test(logger, cur_step, final=False)
    
    def load_best_model(self):
        weights_names = glob.glob(os.path.join(self.log_path, 'models', 'weights*val*'), recursive=False)
        print('Choose from {} ...'.format(weights_names))
        best = 0
        best_name = None
        for weight_name in weights_names:
            acc = float(weight_name.split('val_')[-1])
            if best < acc:
                best = acc
                best_name = weight_name

        print('Load {} ...'.format(best_name))
        self.load_model(best_name)

        for n in weights_names:
            if n != best_name:
                for root, dirs, files in os.walk(n, topdown=False):
                    for name in files:
                        os.remove(os.path.join(root, name))
        
    def save_opts(self):
        """Save options to disk so we know what we ran this experiment with
        """
        models_dir = os.path.join(self.log_path, "models")
        if not os.path.exists(models_dir):
            os.makedirs(models_dir)
        to_save = self.opt.__dict__.copy()

        with open(os.path.join(models_dir, 'opt.json'), 'w') as f:
            json.dump(to_save, f, indent=2)

    def log_trajectory(self, env, data, logger, cur_step):
        
        gt_poses = [  d['pos'] for d in data]

        INTERVAL = self.opt.frame_ids[-1]
        vis_gt_poses = gt_poses[0 ::INTERVAL]

        img = vis_map(env, vis_gt_poses, None)

        logger.add_image('Train/Trajectory', img, cur_step)



src/models/base_trainer.py
embodied-depth/Embodied-Depth-Prediction
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

from .update import BasicUpdateBlock, SmallUpdateBlock
from .extractor import BasicEncoder, SmallEncoder
from .corr import CorrBlock, AlternateCorrBlock
from .utils.utils import bilinear_sampler, coords_grid, upflow8

try:
    autocast = torch.cuda.amp.autocast
except:
    # dummy autocast for PyTorch < 1.6
    class autocast:
        def __init__(self, enabled):
            pass
        def __enter__(self):
            pass
        def __exit__(self, *args):
            pass


class RAFT(nn.Module):
    def __init__(self, args):
        super(RAFT, self).__init__()
        self.args = args

        if args.small:
            self.hidden_dim = hdim = 96
            self.context_dim = cdim = 64
            args.corr_levels = 4
            args.corr_radius = 3
        
        else:
            self.hidden_dim = hdim = 128
            self.context_dim = cdim = 128
            args.corr_levels = 4
            args.corr_radius = 4

        if 'dropout' not in self.args:
            self.args.dropout = 0

        if 'alternate_corr' not in self.args:
            self.args.alternate_corr = False

        # feature network, context network, and update block
        if args.small:
            self.fnet = SmallEncoder(output_dim=128, norm_fn='instance', dropout=args.dropout)        
            self.cnet = SmallEncoder(output_dim=hdim+cdim, norm_fn='none', dropout=args.dropout)
            self.update_block = SmallUpdateBlock(self.args, hidden_dim=hdim)

        else:
            self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', dropout=args.dropout)        
            self.cnet = BasicEncoder(output_dim=hdim+cdim, norm_fn='batch', dropout=args.dropout)
            self.update_block = BasicUpdateBlock(self.args, hidden_dim=hdim)

    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, nn.BatchNorm2d):
                m.eval()

    def initialize_flow(self, img):
        """ Flow is represented as difference between two coordinate grids flow = coords1 - coords0"""
        N, C, H, W = img.shape
        coords0 = coords_grid(N, H//8, W//8, device=img.device)
        coords1 = coords_grid(N, H//8, W//8, device=img.device)

        # optical flow computed as difference: flow = coords1 - coords0
        return coords0, coords1

    def upsample_flow(self, flow, mask):
        """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
        N, _, H, W = flow.shape
        mask = mask.view(N, 1, 9, 8, 8, H, W)
        mask = torch.softmax(mask, dim=2)

        up_flow = F.unfold(8 * flow, [3,3], padding=1)
        up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)

        up_flow = torch.sum(mask * up_flow, dim=2)
        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
        return up_flow.reshape(N, 2, 8*H, 8*W)


    def forward(self, image1, image2, iters=12, flow_init=None, upsample=True, test_mode=False):
        """ Estimate optical flow between pair of frames """


        hdim = self.hidden_dim
        cdim = self.context_dim

        # run the feature network
        with autocast(enabled=self.args.mixed_precision):
            fmap1, fmap2 = self.fnet([image1, image2])        
        
        fmap1 = fmap1.float()
        fmap2 = fmap2.float()
        if self.args.alternate_corr:
            corr_fn = AlternateCorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
        else:
            corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)

        # run the context network
        with autocast(enabled=self.args.mixed_precision):
            cnet = self.cnet(image1)
            net, inp = torch.split(cnet, [hdim, cdim], dim=1)
            net = torch.tanh(net)
            inp = torch.relu(inp)

        coords0, coords1 = self.initialize_flow(image1)

        if flow_init is not None:
            coords1 = coords1 + flow_init

        flow_predictions = []
        for itr in range(iters):
            coords1 = coords1.detach()
            corr = corr_fn(coords1) # index correlation volume

            flow = coords1 - coords0
            with autocast(enabled=self.args.mixed_precision):
                net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)

            # F(t+1) = F(t) + \Delta(t)
            coords1 = coords1 + delta_flow

            # upsample predictions
            if up_mask is None:
                flow_up = upflow8(coords1 - coords0)
            else:
                flow_up = self.upsample_flow(coords1 - coords0, up_mask)
            
            flow_predictions.append(flow_up)

        if test_mode:
            return coords1 - coords0, flow_up
            
        return flow_predictions

src/models/flow2depth/RAFT/core/raft.py
embodied-depth/Embodied-Depth-Prediction
import sys
sys.path.append('core')

import argparse
import os
import cv2
import glob
import numpy as np
import torch
from PIL import Image

from core.raft import RAFT
from core.utils import flow_viz
from core.utils.utils import InputPadder

from imageio import get_writer
from flow2depth import generate_image_homogeneous_coordinates, rot_bearing_mul, triangulation

import matplotlib.pyplot as plt


DEVICE = 'cuda'

def load_image(imfile):
    img = np.array(Image.open(imfile)).astype(np.uint8)
    img = torch.from_numpy(img).permute(2, 0, 1).float()
    if 'png' in imfile:
        img = img[:3]
    return img[None].to(DEVICE)


def viz(img, flo, pred_depth):
    img = img[0].permute(1,2,0).cpu().numpy().astype(np.uint8)
    flo = flo[0].permute(1,2,0).cpu().numpy()
    
    # map flow to rgb image
    flo = flow_viz.flow_to_image(flo)
    img_flo = np.concatenate([img, flo], axis=0)

    #img_flo = np.concatenate([img_flo, display_d(pred_depth).flatten(0,1).permute(1,2,0).repeat(1,1,3).numpy()], axis=0)
    #img_flo = np.concatenate([img_flo, colormap(pred_depth, torch_transpose=False).flatten(0,1).permute(1,2,0)], axis=0)
    #d_ = 
    depth = (colormap(pred_depth, torch_transpose=False).reshape((192, 640, 3)) * 255 ).astype(np.uint8)
    img_flo = np.concatenate([img_flo, depth], axis=0)
    # import matplotlib.pyplot as plt
    # plt.imshow(img_flo / 255.0)
    # plt.show()

    return img_flo
    return img_flo[:, :, [2,1,0]]
    cv2.imshow('image', img_flo[:, :, [2,1,0]]/255.0)
    
#_DEPTH_COLORMAP = plt.get_cmap('plasma', 256)  # for plotting
_DEPTH_COLORMAP = plt.get_cmap('Spectral', 256)  # for plotting
def colormap(inputs, normalize=True, torch_transpose=True):
    if isinstance(inputs, torch.Tensor):
        inputs = inputs.detach().cpu().numpy()

    vis = inputs
    if normalize:
        ma = float(vis.max())
        mi = float(vis.min())
        d = ma - mi if ma != mi else 1e5
        vis = (vis - mi) / d

    if vis.ndim == 4:
        vis = vis.transpose([0, 2, 3, 1])
        vis = _DEPTH_COLORMAP(vis)
        vis = vis[:, :, :, 0, :3]
        if torch_transpose:
            vis = vis.transpose(0, 3, 1, 2)
    elif vis.ndim == 3:
        vis = _DEPTH_COLORMAP(vis)
        vis = vis[:, :, :, :3]
        if torch_transpose:
            vis = vis.transpose(0, 3, 1, 2)
    elif vis.ndim == 2:
        vis = _DEPTH_COLORMAP(vis)
        vis = vis[..., :3]
        if torch_transpose:
            vis = vis.transpose(2, 0, 1)

    return vis

def display_d(depth, ax=None):
    d_max = 10.
    d_min = 0.1
    depth[depth<d_min] = d_min
    depth[depth>d_max] = d_max

    depth = depth / (d_max - d_min)
    return depth * 255.

def traj_loader(path):
    with open(path, 'r') as f:
        lines = f.readlines()
        out_ = np.zeros((len(lines), 4 ,4), dtype=np.float32)
        for i, l in enumerate(lines):
            num = l.split(' ')[1:]
            out_[i, :3, :] = np.array([float(n) for n in num ]).reshape(-1, 4)
            out_[i, 3, 3] = 1
    return out_


def demo(args):
    model = torch.nn.DataParallel(RAFT(args))
    model.load_state_dict(torch.load(args.model))

    model = model.module
    model.to(DEVICE)
    model.eval()

    log = get_writer('/data/scratch/qqbao/DEBUG/flow_colormap_2flow.mp4', fps=45)

    poses_all = np.load(args.path + '../poses.npy')
    y_correct = np.array([[1., 0, 0, 0],
                    [0, -1., 0, 0],
                    [0, 0 , -1., 0],
                    [0, 0 , 0, 1.]], dtype=np.float32)

    fc = np.array([320., 96.])
    cc = np.array([320., 96.])
    height =  192
    width  = 640

    with torch.no_grad():
        images = glob.glob(os.path.join(args.path, '*.png')) + \
                 glob.glob(os.path.join(args.path, '*.jpg'))
        
        images = sorted(images)
        for i, (imfile1, imfile2, imfile3) in enumerate(zip(images[:-6], images[3:-3], images[6:])):
            image1 = load_image(imfile1)
            image2 = load_image(imfile2)
            image3 = load_image(imfile3)

            padder = InputPadder(image1.shape)
            image1, image2, image3 = padder.pad(image1, image2, image3)


            image1 = 2 * (image1 / 255.0) - 1.0
            image2 = 2 * (image2 / 255.0) - 1.0
            image3 = 2 * (image3 / 255.0) - 1.0

            image1 = image1.contiguous()
            image2 = image2.contiguous()
            image3 = image3.contiguous()

            flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
            flow_low, flow_up2 = model(image1, image3, iters=20, test_mode=True)

            flow_up = flow_up.cpu()
            flow_up2 = flow_up2.cpu()

            poses = [
                np.matmul(poses_all[i], y_correct),
                np.matmul(poses_all[i+3], y_correct),
                np.matmul(poses_all[i+6], y_correct)]

            poses_ref_in_other = [torch.from_numpy(np.matmul(np.linalg.inv(pose), poses[0])) for pose in poses[1:]]
            rots_ref_in_other = torch.from_numpy(np.stack([pose[0:3, 0:3] for pose in poses_ref_in_other])).unsqueeze(0)
            ts = torch.from_numpy(np.stack([pose[0:3, 3] for pose in poses_ref_in_other])).unsqueeze(0)
            homo =  generate_image_homogeneous_coordinates(fc, cc, width, height).permute(2, 0, 1).unsqueeze(0)
            bearings_ref_in_other = [rot_bearing_mul(rots_ref_in_other[:, k, ...], homo) for k in range(rots_ref_in_other.shape[1])]
            pred_depth, _ = triangulation(bearings_ref_in_other, ts, [flow_up, flow_up2], residual=True)


            log.append_data(viz(image1, flow_up, pred_depth))

    log.close()

def demo2(args):
    model = torch.nn.DataParallel(RAFT(args))
    model.load_state_dict(torch.load(args.model))

    model = model.module
    model.to(DEVICE)
    model.eval()

    log = get_writer('/data/scratch/qqbao/DEBUG/REALDATA_colormap_2flow_interval16_new.mp4', fps=45)

    poses_all = traj_loader(os.path.join(args.path, 'new_trajectory.txt'))

    height = 480 #720 
    width  = 640
    fc = np.array([904.62, 904.62]).astype(float) 
    cc = np.array([640., 360.]).astype(float)
    fc[0] *= width / 1280.
    cc[0] *= width / 1280.
    fc[1] *= height / 720.
    cc[1] *= height / 720.
    print('FC, cc', self.fc, self.cc)
    #fc = np.array([904.62, 640.]) * width / 1280
    #cc = np.array([904.62, 360.]) * height / 720

    with torch.no_grad():
        images = glob.glob(os.path.join(args.path, 'color', '*.png')) + \
                 glob.glob(os.path.join(args.path, 'color', '*.jpg'))
        
        images = sorted(images)
        images = [os.path.join(args.path,'color', 'frame{}.jpg'.format(i)) for i in range(1, 1000) ]
        for i, (imfile1, imfile2, imfile3) in enumerate(zip(images[:-16], images[13:-3], images[16:])):
            image1 = load_image(imfile1)
            image2 = load_image(imfile2)
            image3 = load_image(imfile3)

            padder = InputPadder(image1.shape)
            image1, image2, image3 = padder.pad(image1, image2, image3)

            flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
            flow_low, flow_up2 = model(image1, image3, iters=20, test_mode=True)

            flow_up = flow_up.cpu()
            flow_up2 = flow_up2.cpu()

            poses = [
                poses_all[i],
                poses_all[i+13],
                poses_all[i+16]]

            poses_ref_in_other = [torch.from_numpy(np.matmul(np.linalg.inv(pose), poses[0])) for pose in poses[1:]]
            rots_ref_in_other = torch.from_numpy(np.stack([pose[0:3, 0:3] for pose in poses_ref_in_other])).unsqueeze(0)
            ts = torch.from_numpy(np.stack([pose[0:3, 3] for pose in poses_ref_in_other])).unsqueeze(0)
            homo =  generate_image_homogeneous_coordinates(fc, cc, width, height).permute(2, 0, 1).unsqueeze(0)
            bearings_ref_in_other = [rot_bearing_mul(rots_ref_in_other[:, k, ...], homo) for k in range(rots_ref_in_other.shape[1])]
            pred_depth, _ = triangulation(bearings_ref_in_other, ts, [flow_up, flow_up2], residual=True)


            log.append_data(viz(image1, flow_up, pred_depth))

    log.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', help="restore checkpoint")
    parser.add_argument('--path', help="dataset for evaluation")
    parser.add_argument('--small', action='store_true', help='use small model')
    parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
    parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation')
    args = parser.parse_args()

    demo(args)

src/models/flow2depth/RAFT/demo.py
embodied-depth/Embodied-Depth-Prediction
#import sys
#sys.path.append('core')

import argparse
import os
import cv2
import glob
import numpy as np
import torch
from PIL import Image
from torchvision.transforms import Grayscale

from src.models.flow2depth.RAFT.core.raft import RAFT
from src.models.flow2depth.RAFT.core.utils import flow_viz
from src.models.flow2depth.RAFT.core.utils.utils import InputPadder
import multiprocessing

#from core.raft import RAFT
#from core.utils import flow_viz
#from core.utils.utils import InputPadder

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

class FlowNet(nn.Module):
    def __init__(self, args):
        super(FlowNet, self).__init__()
        self.args = args

        raft = torch.nn.DataParallel(RAFT(args))
        raft.load_state_dict(torch.load(args.flow_model_path))
        self.flownet = raft.module
        self.flownet = self.flownet.to(args.device)
        self.flownet.eval()

    def forward(self, img0, img_refs):
        with torch.no_grad():
            flow_ups = []
            for img_ref in img_refs:
                _, flow_up = self.flownet(img0, img_ref, iters=20, test_mode=True)
                flow_ups.append(flow_up)

        return torch.stack(flow_ups, dim=1)

def get_flow(imgs):
	out_ = cv2.calcOpticalFlowFarneback(imgs[0], imgs[1], None, 0.5, 10, 80, 3, 7, 1.5, 0)
	return torch.tensor(out_).permute(2, 0, 1)


class FlowBaseline(nn.Module):
    def __init__(self, args):
        super(FlowBaseline, self).__init__()
        self.args = args
        self.grayscale = Grayscale(num_output_channels=1)

    def forward(self, _img0, _img_refs, iters=20, test_mode=True):                                                                                  
        flow_ups = []                                                                                                                               
        img0 = (self.grayscale(_img0.cpu()).squeeze().numpy() * 255).astype(np.uint8)                                                               
        img_refs = (self.grayscale(_img_refs.cpu()).squeeze().numpy() * 255).astype(np.uint8)                                                      
        if len(_img0) > 1:
            slices = [(img0[i], img_refs[i]) for i in range(len(img0))]
            # Create a multiprocessing pool
            pool = multiprocessing.Pool(processes=len(img0))
            # Apply the process_slice function to each slice in parallel
            results = pool.map(get_flow, slices)
            # Close the pool to free up resources
            pool.close()
            # Wait for all processes to finish
            pool.join()
            flow_ups = results
        else:
            flow_ups = [get_flow((img0, img_refs))]
			
        #for i0, img_ref in zip(img0, img_refs):                                                                                                     
        #    flow_up = cv2.calcOpticalFlowFarneback(i0, img_ref, None, 0.5, 10, 80, 10, 7, 1.5, 0)
        #    flow_ups.append(torch.tensor(flow_up).to(_img0.device).permute(2, 0, 1))
                                     
        return None, torch.stack(flow_ups, dim=0).to(_img0.device)


class Flow2Depth(nn.Module):
    def __init__(self, args,):
        super(Flow2Depth, self).__init__()
        self.args = args

        if args.use_opencv_flow:
            self.flownet = FlowBaseline(args)
        else:
            raft = torch.nn.DataParallel(RAFT(args))
            raft.load_state_dict(torch.load(args.flow_model_path))
            self.flownet = raft.module
            self.flownet = self.flownet.to(args.device)
            self.flownet.eval()

        #self.height = 480 #720 
        #self.width  = 640
        self.height = args.height
        self.width  = args.width
        if args.dataset == 'real':
            #self.fc = np.array([904.62, 640.]).astype(float) * self.width / 1280
            #self.cc = np.array([904.62, 360.]).astype(float)* self.height / 720
            self.fc = np.array([904.62, 904.62]).astype(float) 
            self.cc = np.array([640., 360.]).astype(float)
            self.fc[0] *= self.width / 1280.
            self.cc[0] *= self.width / 1280.
            self.fc[1] *= self.height / 720.
            self.cc[1] *= self.height / 720.
            print('FC, cc', self.fc, self.cc)

        elif args.dataset == 'habitat':
            self.fc = np.array([320., 96.])
            self.cc = np.array([320., 96.])
        else:
            raise NotImplementedError('No support for {} dataset. Need to know its Camera Intrinsics'.format(self.opt.dataset))

        self.homo = generate_image_homogeneous_coordinates(self.fc, self.cc, self.width, self.height).permute(2, 0, 1).unsqueeze(0)

    def forward(self, img0, img_refs, poses):
        with torch.no_grad():

            flow_ups = []
            for img_ref in img_refs:
                _, flow_up = self.flownet(img0, img_ref, iters=20, test_mode=True)
                flow_ups.append(flow_up)

            batch_size = img0.shape[0]
            device = img0.device

            #poses_ref_in_other = [torch.from_numpy(np.matmul(np.linalg.inv(pose), poses[0])).to(device) for pose in poses[1:]]
            #rots_ref_in_other = torch.from_numpy(np.stack([pose[0:3, 0:3] for pose in poses])).to(device)
            #ts = torch.from_numpy(np.stack([pose[0:3, 3] for pose in poses])).to(device)
            rots = poses[:,:,  :3, :3]
            ts = poses[:, :, :3, 3]
            homo = self.homo.repeat(batch_size, 1, 1, 1).to(device)
            bearings_ref_in_other = [rot_bearing_mul(rots[:, k, ...], homo) for k in range(rots.shape[1])]

            pred_depth, _ = self.triangulation(bearings_ref_in_other, ts, flow_ups, residual=True)

        return depth_to_disp(pred_depth, min_depth=self.args.min_depth, max_depth=self.args.max_depth)
        #return pred_depth

    def triangulation(self, bearings_ref_in_other, t_ref_in_other, flows, residual=False):
        rs, ss = self.pre_triangulation(bearings_ref_in_other, t_ref_in_other, flows, concat=False)
        # get output = (z, residual, hessian)
        outputs = [ls_2view(*r_s) for r_s in zip(rs, ss)]
        # weighted sum of z with weight hessian
        
        hessian = sum([output[2] for output in outputs])
        pred_depths = sum([output[0] * output[2] for output in outputs]) / (hessian + 1e-12)

        if residual:
            # hessian*(z* - z)^2 + residual
            error = torch.sqrt(
                sum([output[2] * (pred_depths - output[0]) ** 2 + output[1] for output in outputs]).clamp_min(0))
            sqrt_hessian = torch.sqrt(hessian)
            return pred_depths, (error, sqrt_hessian)
        else:
            return pred_depths

    def flow2bearing(self, flow, fc, cc, normalize=True):
        assert len(flow.shape) == 4
        height, width = flow.shape[2:4]
        device = flow.device
        xx, yy = np.meshgrid(range(width), range(height))
        pixel = torch.zeros_like(flow)
        match = [flow[:, 0, ...] + torch.from_numpy(xx).to(device), flow[:, 1] + torch.from_numpy(yy).to(device)]
        pixel[:, 0] = (match[0] - cc[0]) / fc[0]
        pixel[:, 1] = (match[1] - cc[1]) / fc[1]
        pixel = torch.cat((pixel, torch.ones_like(pixel[:, 0:1])), dim=1)

        if normalize:
            pixel = F.normalize(pixel)
        return pixel

    def pre_triangulation(self, bearings_ref_in_other, t_ref_in_other, flows, concat=True):
        device = flows[0].device
        fc = torch.from_numpy(self.fc).to(device)
        cc = torch.from_numpy(self.cc).to(device)

        #fc = torch.from_numpy(np.array([320., 96.])).to(device) #/ resize
        #cc = torch.from_numpy(np.array([320., 96.])).to(device) #/ resize
        #fc = torch.from_numpy(np.array([904.62, 640.])).to(device) / 2.#/ resize
        #cc = torch.from_numpy(np.array([904.62, 360.])).to(device) * 480. / 720. #/ resize

        bearings_other = [self.flow2bearing(flow, fc, cc, normalize=True) for flow in flows]
        ss = [torch.cross(bearings_other[k], bearings_ref_in_other[k], dim=1) for k in
                range(len(bearings_other))]
        rs = [torch.cross(bearings_other[k], t_ref_in_other[:, k, :, None, None].expand_as(bearings_other[k]), dim=1)
                for k in range(len(bearings_other))]

        if concat:
            s = torch.cat(ss, dim=1)
            r = torch.cat(rs, dim=1)
            return r, s
        else:
            return rs, ss

def depth_to_disp(_depth, min_depth, max_depth):
    #mask1 = ( (_depth == 0).mean((-1, -2), dtype=torch.float) < 0.5).reshape(_depth.shape[0], 1, 1, 1).repeat(1, *_depth.shape[1:])
    #mask2 = _depth != 0
    #mask = mask1 * mask2
    
    mask = _depth != 0
    depth = torch.clamp(_depth, min_depth, max_depth)
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = 1 / depth
    disp_ = (scaled_disp - min_disp) / (max_disp - min_disp)* mask
    return disp_

def generate_image_homogeneous_coordinates(fc, cc, image_width, image_height):
    homogeneous = np.zeros((image_height, image_width, 3))
    homogeneous[:, :, 2] = 1

    xx, yy = np.meshgrid([i for i in range(0, image_width)], [i for i in range(0, image_height)])
    homogeneous[:, :, 0] = (xx - cc[0]) / fc[0]
    homogeneous[:, :, 1] = (yy - cc[1]) / fc[1]

    return torch.from_numpy(homogeneous.astype(np.float32))

def rot_bearing_mul(rot, bearing):
    # rot: B x 3 x 3, bearing: B x 3 x H x W
    product = torch.bmm(rot, bearing.view(bearing.shape[0], 3, -1))
    return product.view(bearing.shape)

def ls_2view(r, s):
    hessian = (s * s).sum(dim=1, keepdims=True)
    z = -(s * r).sum(dim=1, keepdims=True) / (hessian + 1e-30)
    e = (r * r).sum(dim=1, keepdims=True) - hessian * (z ** 2)

    invalid_mask = (z <= 0.1)
    invalid_mask |= (z >= 40)
    # invalid_mask |= (e > 0.015 ** 2)
    z[invalid_mask] = 0
    e[invalid_mask] = 0
    hessian[invalid_mask] = 0
    return z, e, hessian


def triangulation(bearings_ref_in_other, t_ref_in_other, flows, residual=False):
    rs, ss = pre_triangulation(bearings_ref_in_other, t_ref_in_other, flows, concat=False)
    # get output = (z, residual, hessian)
    outputs = [ls_2view(*r_s) for r_s in zip(rs, ss)]
    # weighted sum of z with weight hessian
    
    hessian = sum([output[2] for output in outputs])
    pred_depths = sum([output[0] * output[2] for output in outputs]) / (hessian + 1e-12)

    if residual:
        # hessian*(z* - z)^2 + residual
        error = torch.sqrt(
            sum([output[2] * (pred_depths - output[0]) ** 2 + output[1] for output in outputs]).clamp_min(0))
        sqrt_hessian = torch.sqrt(hessian)
        return pred_depths, (error, sqrt_hessian)
    else:
        return pred_depths

def flow2bearing(flow, fc, cc, normalize=True):
    assert len(flow.shape) == 4
    height, width = flow.shape[2:4]
    device = flow.device
    xx, yy = np.meshgrid(range(width), range(height))
    pixel = torch.zeros_like(flow)
    match = [flow[:, 0, ...] + torch.from_numpy(xx).to(device), flow[:, 1] + torch.from_numpy(yy).to(device)]
    pixel[:, 0] = (match[0] - cc[0]) / fc[0]
    pixel[:, 1] = (match[1] - cc[1]) / fc[1]
    pixel = torch.cat((pixel, torch.ones_like(pixel[:, 0:1])), dim=1)

    if normalize:
        pixel = F.normalize(pixel)
    return pixel

def pre_triangulation(bearings_ref_in_other, t_ref_in_other, flows, concat=True):
    device = flows[0].device

    fc = torch.from_numpy(np.array([320., 96.])).to(device) #/ resize
    cc = torch.from_numpy(np.array([320., 96.])).to(device) #/ resize
    #fc = torch.from_numpy(np.array([904.62, 640.])).to(device) / 2.#/ resize
    #cc = torch.from_numpy(np.array([904.62, 360.])).to(device) * 480. / 720. #/ resize

    bearings_other = [flow2bearing(flow, fc, cc, normalize=True) for flow in flows]
    ss = [torch.cross(bearings_other[k], bearings_ref_in_other[k], dim=1) for k in
            range(len(bearings_other))]
    rs = [torch.cross(bearings_other[k], t_ref_in_other[:, k, :, None, None].expand_as(bearings_other[k]), dim=1)
            for k in range(len(bearings_other))]

    if concat:
        s = torch.cat(ss, dim=1)
        r = torch.cat(rs, dim=1)
        return r, s
    else:
        return rs, ss

src/models/flow2depth/RAFT/flow2depth.py
embodied-depth/Embodied-Depth-Prediction
from src.models.flow2depth.resnet_encoder import RGBDEncoderv2,  ResnetEncoder, RGBDEncoder
from src.models.flow2depth.depth_decoder import  DepthDecoderT, DepthDecoderv5 
src/models/flow2depth/__init__.py
embodied-depth/Embodied-Depth-Prediction
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

from scipy.spatial.transform import Rotation as R
import numpy as np
import torch
import torch.nn as nn

from collections import OrderedDict
from src.models.layers import *


def dist_T(t):
    I = torch.eye(3).repeat(t.shape[0], t.shape[1], 1, 1).to(t.device)
    dist = torch.sqrt(
            (t[:, :, :3, 3].abs()).sum(dim=(-1))  + \
            2. / 3. * ( I - t[:,:, :3, :3])[:,:, (0,1,2), (0,1,2)].sum(dim=-1)
    )
    dist[dist.isnan()] = 0
    return dist

class DepthDecoder(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True, bayes=False):
        super(DepthDecoder, self).__init__()

        self.num_output_channels = num_output_channels
        self.use_skips = use_skips
        self.upsample_mode = 'nearest'
        self.scales = scales

        self.num_ch_enc = num_ch_enc
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])

        # decoder
        self.convs = OrderedDict()
        for i in range(4, -1, -1):
            # upconv_0
            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out, bayes=bayes)

            # upconv_1
            num_ch_in = self.num_ch_dec[i]
            if self.use_skips and i > 0:
                num_ch_in += self.num_ch_enc[i - 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out, bayes=bayes)

        for s in self.scales:
            self.convs[("dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)

        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_features):
        self.outputs = {}

        # decoder
        x = input_features[-1]
        for i in range(4, -1, -1):
            x = self.convs[("upconv", i, 0)](x)
            x = [upsample(x)]
            if self.use_skips and i > 0:
                x += [input_features[i - 1]]
            x = torch.cat(x, 1)
            x = self.convs[("upconv", i, 1)](x)
            if i in self.scales:
                self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))

        return self.outputs

'''
V5: Tailor to the Multi-input (including Single) Resnet encoders, use pose_dist as decay coef for the embeddings
Receive multi-frame images (may including depth) in every layer with shape (B, T, c, h, w)
'''
class DepthDecoderv5(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True, 
                    bayes=False, frame_num=1, input_pose=False, multi_pose=False, pose_dist_info=False, pose_rot=False, pose_decay=False):
        super(DepthDecoderv5, self).__init__()

        self.num_output_channels = num_output_channels
        self.use_skips = use_skips
        self.upsample_mode = 'nearest'
        self.scales = scales
        self.input_pose = input_pose
        self.multi_pose = multi_pose
        self.pose_rot = pose_rot
        self.pose_dist_info = pose_dist_info
        self.pose_decay = pose_decay
        self.pose_dim = 6 if pose_rot else 2
        if self.pose_dist_info:
            self.pose_dim += 1 

        self.decayT = 0.5

        self.num_ch_enc = num_ch_enc 
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])
        if self.input_pose:
            if self.multi_pose:
                for i in range(len(self.num_ch_enc)):
                    self.num_ch_enc[i] += self.pose_dim * frame_num
            else:
                self.num_ch_enc[-1] += self.pose_dim * frame_num

        # decoder
        self.convs = OrderedDict()
        for i in range(4, -1, -1):
            # upconv_0
            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out, bayes=bayes)

            # upconv_1
            num_ch_in = self.num_ch_dec[i]
            if self.use_skips and i > 0:
                num_ch_in += self.num_ch_enc[i - 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out, bayes=bayes)

        for s in self.scales:
            self.convs[("dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)

        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_features, input_poses=None):
        self.outputs = {}

        B, T, _, _, _ = input_features[-1].shape
        # pose_dist_decay 
        if (input_poses is not None) and self.pose_decay:
            pose_dist = dist_T(input_poses).reshape(B, T, 1, 1, 1)
            pose_dist_decay = torch.exp(-pose_dist / self.decayT)
            for i in range(len(input_features)):
                x = input_features[i]
                B, T, C, H, W = x.shape
                input_features[i] = (x * pose_dist_decay).reshape(B, -1, H, W)
        else:
            for i in range(len(input_features)):
                x = input_features[i]
                input_features[i] = x.flatten(1, 2)

        x = input_features[-1]
        # add pose info
        if self.input_pose:
            pose_info = input_poses[:, :, (0,2), 3].reshape(B, -1, 1, 1)
            if self.pose_rot:
                roat = R.from_matrix(input_poses.flatten(0,1)[:, :3, :3].cpu().numpy())  # (B,T,4,4) -> (B,4)
                quat = roat.as_quat()
                quat = torch.from_numpy(quat).reshape(B, -1, 1, 1).to(x.device).to(x.dtype)
                pose_info = torch.cat([pose_info, quat], dim=1)

            x = torch.cat([x, pose_info.repeat((1, 1,x.shape[-2],x.shape[-1]))], dim=1)

        # decode
        for i in range(4, -1, -1):
            x = self.convs[("upconv", i, 0)](x)
            x = [upsample(x)]
            if self.use_skips and i > 0:
                x += [input_features[i - 1]]
            x = torch.cat(x, 1)
            # Add pose info in every layer
            if i > 0 and self.input_pose and self.multi_pose:
                x = torch.cat([x, pose_info.repeat((1, 1,x.shape[-2],x.shape[-1]))], dim=1)
            x = self.convs[("upconv", i, 1)](x)
            if i in self.scales:
                self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))

        return self.outputs


'''
T: Tailor to the Multi-input (including Single) Resnet encoders
Delete Pose

Receive multi-frame images (may including depth) in every layer with shape (B, T, c, h, w)
'''
class DepthDecoderT(nn.Module):
    def __init__(self, num_ch_enc, scales=range(4), num_output_channels=1, use_skips=True, 
                    bayes=False, frame_num=1, input_pose=False, multi_pose=False, pose_dist_info=False, pose_rot=False, pose_decay=False):
        super(DepthDecoderT, self).__init__()

        self.num_output_channels = num_output_channels
        self.use_skips = use_skips
        self.upsample_mode = 'nearest'
        self.scales = scales
        self.input_pose = input_pose
        self.multi_pose = multi_pose
        self.pose_rot = pose_rot
        self.pose_dist_info = pose_dist_info
        self.pose_decay = pose_decay
        self.pose_dim = 6 if pose_rot else 2
        if self.pose_dist_info:
            self.pose_dim += 1 

        self.decayT = 0.5

        self.num_ch_enc = num_ch_enc // frame_num
        self.num_ch_dec = np.array([16, 32, 64, 128, 256])

        # decoder
        self.convs = OrderedDict()
        for i in range(4, -1, -1):
            # upconv_0
            num_ch_in = self.num_ch_enc[-1] if i == 4 else self.num_ch_dec[i + 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 0)] = ConvBlock(num_ch_in, num_ch_out, bayes=bayes)

            # upconv_1
            num_ch_in = self.num_ch_dec[i]
            if self.use_skips and i > 0:
                num_ch_in += self.num_ch_enc[i - 1]
            num_ch_out = self.num_ch_dec[i]
            self.convs[("upconv", i, 1)] = ConvBlock(num_ch_in, num_ch_out, bayes=bayes)

        for s in self.scales:
            self.convs[("dispconv", s)] = Conv3x3(self.num_ch_dec[s], self.num_output_channels)

        self.decoder = nn.ModuleList(list(self.convs.values()))
        self.sigmoid = nn.Sigmoid()

    def forward(self, input_features, input_poses=None):
        self.outputs = {}

        B, T, _, _, _ = input_features[-1].shape
        # pose_dist_decay 
        for i in range(len(input_features)):
            x = input_features[i]
            input_features[i] = x.flatten(0, 1)

        x = input_features[-1]

        # decode
        for i in range(4, -1, -1):
            x = self.convs[("upconv", i, 0)](x)
            x = [upsample(x)]
            if self.use_skips and i > 0:
                x += [input_features[i - 1]]
            x = torch.cat(x, 1)
            # Add pose info in every layer
            x = self.convs[("upconv", i, 1)](x)
            if i in self.scales:
                self.outputs[("disp", i)] = self.sigmoid(self.convs[("dispconv", i)](x))

        _, _, H, W = self.outputs[("disp", i)].shape
        self.outputs[("disp", i)] = self.outputs[("disp", i)].reshape(B, T, 1, H, W)
        return self.outputs

src/models/flow2depth/depth_decoder.py
embodied-depth/Embodied-Depth-Prediction
import os
import json
from PIL import Image  # using pillow-simd for increased speed
from pathlib import Path
from collections import deque

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler

from src.my_utils import *
from src.models.utils import *
from src.models.layers import *

from src.models.flow2depth import RGBDEncoder, RGBDEncoderv2, DepthDecoderv5, DepthDecoderT #, UpdateLayerv2
from src.models.flow2depth.RAFT.flow2depth import Flow2Depth

def cam1_to_cam0(_world_camera1, _world_camera0):
    T_camera1_world = np.linalg.pinv(_world_camera1)
    T_camera1_camera0 = np.matmul(T_camera1_world, _world_camera0)
    return T_camera1_camera0

def depth_to_disp(_depth, min_depth, max_depth):
    mask = _depth != 0

    depth = torch.clamp(_depth, min_depth, max_depth)
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = 1 / depth
    disp_ = (scaled_disp - min_disp) / (max_disp - min_disp) * mask
    return disp_

class DepthModel:
    def __init__(self, options):
        self.opt = options

        # checking height and width are multiples of 32
        assert self.opt.height % 32 == 0, "'height' must be a multiple of 32"
        assert self.opt.width % 32 == 0, "'width' must be a multiple of 32"

        self.models = {}
        self.parameters_to_train = []

        self.device = torch.device("cpu" if self.opt.no_cuda else "cuda")

        self.num_scales = len(self.opt.scales)
        self.num_input_frames = len(self.opt.frame_ids)
        self.num_pose_frames = 2 if self.opt.pose_model_input == "pairs" else self.num_input_frames

        assert self.opt.frame_ids[0] == 0, "frame_ids must start with 0"

        self.use_pose_net = self.opt.pose_model_type != 'gt'
        self.input_pose = self.opt.input_pose

        self.frames_to_load_test = [0]
        self.frames_to_load_train = self.opt.frame_ids.copy()
        for idx in range(-self.opt.num_past_frame_input * self.opt.video_interval, 1, self.opt.video_interval):
            if idx not in self.frames_to_load_test:
                self.frames_to_load_test.append(idx)
                self.frames_to_load_train.append(idx)

        print(self.frames_to_load_train)
        print(self.frames_to_load_test)
        # TODO: clarify this GUY!!
        if self.opt.rgbdcat:
            self.models["encoder"] = RGBDEncoder(
        			self.opt.num_layers,
        			self.opt.weights_init == "pretrained", 
        			num_input_images=1,  #self.opt.num_past_frame_input+1 ,
        			#num_input_images=self.opt.num_past_frame_input+1 ,
        			input_depth=self.opt.depth_encoder) 
        else:
            self.models["encoder"] = RGBDEncoderv2(
                    self.opt.num_layers,
        			self.opt.weights_init == "pretrained", 
        			num_input_images=1,  #self.opt.num_past_frame_input+1 ,
        			#num_input_images=self.opt.num_past_frame_input+1 ,
        			input_depth=self.opt.depth_encoder) 

        self.models["depth"] = DepthDecoderv5(
                self.models["encoder"].num_ch_enc, 
                self.opt.scales, 
                bayes=self.opt.bayes, 
                frame_num=self.opt.num_past_frame_input + 1 ,
                input_pose=self.opt.input_pose,
                multi_pose=self.opt.multi_pose,
                pose_dist_info=self.opt.pose_dist_info,
                pose_rot=self.opt.pose_rot)


        for k in self.models.keys():
            self.models[k].to(self.device)
            if k != 'refine':
                self.parameters_to_train += list(self.models[k].parameters())

        if self.opt.flow2depth:
            self.models["flow2depth"] = Flow2Depth(self.opt)
            self.models["flow2depth"].to(self.device)
            if self.opt.flow2depth_train:
                self.parameters_to_train += list(self.models["flow2depth"].parameters())

        if self.use_pose_net:
            if self.opt.pose_model_type == "separate_resnet":
                self.models["pose_encoder"] = ResnetEncoder(
                    self.opt.num_layers,
                    self.opt.weights_init == "pretrained",
                    num_input_images=self.num_pose_frames,
                    input_depth=self.opt.input_depth)

                self.models["pose_encoder"].to(self.device)
                self.parameters_to_train += list(self.models["pose_encoder"].parameters())

                self.models["pose"] = PoseDecoder(
                    self.models["pose_encoder"].num_ch_enc,
                    num_input_features=1,
                    num_frames_to_predict_for=2)
            else:
                raise NotImplementedError


        if self.opt.load_weights_folder is not None:
            self.load_model()

        print("Training model named:\n  ", self.opt.model_name)
        print("Models and tensorboard events files are saved to:\n  ", self.opt.log_dir)
        print("Training is using:\n  ", self.device)


        self.backproject_depth = {}
        self.project_3d = {}
        for scale in self.opt.scales:
            h = self.opt.height // (2 ** scale)
            w = self.opt.width // (2 ** scale)

            self.backproject_depth[scale] = BackprojectDepth(self.opt.batch_size, h, w) #if self.opt.dataset != 'habitat' else HabBackprojectDepth(self.opt.batch_size, h, w)
            self.backproject_depth[scale].to(self.device)

            self.project_3d[scale] = Project3D(self.opt.batch_size, h, w) #if self.opt.dataset != 'habitat' else  HabProject3D(self.opt.batch_size, h, w)
            self.project_3d[scale].to(self.device)

        if not self.opt.no_ssim:
            self.ssim = SSIM()
            self.ssim.to(self.device)

    def forward(self, inputs):
        outputs, losses = self._process_batch(inputs)
        return outputs, losses

    def _preprocess(self, obs, ):
        '''
        Input: 
            obs: Dict{'color_sensor': np.ndarray, 'depth_sensor': np.ndarray} 
        Output:
            img_: torch.Tensor (B,  3, H, W)
            depth_: numpy.ndarray  (B,  1, H, W)
        '''
        correct = torch.tensor(np.array([[1., 0, 0, 0],
                          [0, -1., 0, 0],
                          [0, 0 , -1., 0],
                          [0, 0 , 0, 1.]], dtype=np.float32))
        correct_f = lambda x: torch.matmul(x, correct)

        if type(obs) in [list, deque] :
            inputs = {}
            series = [self._preprocess(o) for o in obs]

            minus_len = -min(self.frames_to_load_test)
            data_len = len(obs) 
            for i in self.frames_to_load_test:
                inputs[('color', i, 0)] = torch.stack([obs['rgb'] for obs in series[minus_len+i: data_len+i]]).to(self.device)
                inputs[('flowpose', i)] = torch.stack(
                                                        [
                                                           cam1_to_cam0(
                                                                        correct_f(series[idx +  i]['pose']),
                                                                        correct_f(series[idx     ]['pose'])) 
                                                            for idx in range(minus_len+i, data_len+i)
                                                        ]
                                                    ).to(self.device)

            inputs['gt_depth'] = np.stack([obs['gt_depth'] for obs in series[minus_len: data_len]])
             
        else:
            to_tensor = transforms.ToTensor()
            img_ = to_tensor(Image.fromarray(obs['color_sensor']).convert('RGB')).unsqueeze(0).to(torch.float32)
            depth_ = np.expand_dims(obs['depth_sensor'], 0)
            pose_ = torch.tensor(obs['Ext'])

            inputs = {
                'rgb': img_,
                'gt_depth': depth_,
                'pose': pose_
            }

        return inputs

    def pred_depth(self, obs):
        #input_color, gt_depth = self._preprocess(obs)
        inputs = self._preprocess(obs)
        gt_depth = inputs['gt_depth']
        with torch.no_grad():

            input_rgb = inputs[('color', 0, 0)]

            if self.opt.depth_encoder:
                if self.opt.flow2depth:
                    img0 = inputs["color", 0, 0].clone().flatten(0,1)
                    img_refs = [inputs["color", i, 0].clone().flatten(0,1) for i in self.frames_to_load_test[1:]]
                    poses = torch.stack([inputs[('flowpose',i)] for i in self.frames_to_load_test[1:]], dim=1)

                    input_depth = self.models['flow2depth'](img0, img_refs, poses).unsqueeze(1)
                    inputs[('depth_input', 0, 0)]  = input_depth
                else:
                    input_depth = inputs[('depth_input', 0, 0)] 
            else:
                input_depth = None

            features = self.models["encoder"](input_rgb, input_depth) 
            output = self.models["depth"](features, inputs["v_poses"].squeeze()) if self.input_pose else self.models["depth"](features)

            pred_disp, _ = disp_to_depth(output[("disp", 0)], self.opt.min_depth, self.opt.max_depth)
            pred_disp = pred_disp.cpu()[:, 0].numpy()

            if self.opt.post_process:
                N = pred_disp.shape[0] // 2
                pred_disp = batch_post_process_disparity(pred_disp[:N], pred_disp[N:, :, ::-1])

            # Disp -> depth
            #pred_disp = cv2.resize(pred_disp, (gt_width, gt_height))
            pred_depth = 1 / pred_disp

            #mask = gt_depth > 0 
            pred_depth *= self.opt.pred_depth_scale_factor
            if not self.opt.disable_median_scaling:
                ratio = np.median(gt_depth) / np.median(pred_depth)
                pred_depth *= ratio


            MIN_DEPTH = self.opt.min_depth
            MAX_DEPTH = self.opt.max_depth
            pred_depth[pred_depth < MIN_DEPTH] = MIN_DEPTH
            pred_depth[pred_depth > MAX_DEPTH] = MAX_DEPTH

        return pred_depth.squeeze()

    def pred_pose(self, obs):
        input_color, gt_depth = self._preprocess(obs)
        with torch.no_grad():
        
            all_color_aug = torch.cat([input_color[i:i+1] for i in range(input_color.shape[0])], 1)

            features = [self.models['pose_encoder'](all_color_aug)]
            axisangle, translation = self.models['pose'](features)

            pred_poses = transformation_from_parameters(axisangle[:, 0], translation[:, 0]).cpu().numpy()

        return pred_poses

    def _process_batch(self, inputs):
        """Pass a minibatch through the network and generate images and losses
        """
        for key, ipt in inputs.items():
            inputs[key] = ipt.to(self.device)

        if self.opt.geo_consist:
            input_rgb = torch.cat([inputs["color_aug", i, 0].squeeze() for i in self.opt.frame_ids]).unsqueeze(1)
            input_depth = torch.cat([inputs[('depth_input', i, 0)] for i in self.opt.frame_ids ]) if self.opt.depth_encoder else None
        else:
            input_rgb = inputs["color_aug", 0, 0]

            if self.opt.depth_encoder:
                if self.opt.flow2depth:
                    img0 = inputs["color", 0, 0].clone().squeeze()
                    img_refs = [inputs["color", i, 0].clone().squeeze() for i in self.frames_to_load_test[1:]]
                    poses = inputs['flowpose']

                    input_depth = self.models['flow2depth'](img0, img_refs, poses).unsqueeze(1)
                    inputs[('depth_input', 0, 0)]  = input_depth
                else:
                    input_depth = inputs[('depth_input', 0, 0)] 
            else:
                input_depth = None

        features = self.models["encoder"](input_rgb, input_depth) 

        outputs = self.models["depth"](features, inputs["v_poses"].squeeze()) if self.input_pose else self.models["depth"](features)

        outputs.update(self._predict_poses(inputs, features))

        self._generate_images_pred(inputs, outputs)
        losses = self._compute_losses(inputs, outputs)

        return outputs, losses

    def _predict_poses(self, inputs, features):
        """Predict poses between input frames for monocular sequences.
        """
        outputs = {}
        if self.opt.pose_model_type == 'gt':
            for f_i in self.opt.frame_ids[1:]:
                outputs[("cam_T_cam", 0, f_i)] = inputs[("cam_T_cam", 0, f_i)].to(torch.float32)
            return outputs

        if self.num_pose_frames == 2:
            # In this setting, we compute the pose to each source frame via a
            # separate forward pass through the pose network.

            # select what features the pose network takes as input
            if self.opt.pose_model_type == "shared":
                pose_feats = {f_i: features[f_i] for f_i in self.opt.frame_ids}
            else:
                pose_feats = {f_i: inputs["color_aug", f_i, 0] for f_i in self.opt.frame_ids}

            for f_i in self.opt.frame_ids[1:]:
                if f_i != "s":
                    # To maintain ordering we always pass frames in temporal order
                    if f_i < 0:
                        pose_inputs = [pose_feats[f_i], pose_feats[0]]
                    else:
                        pose_inputs = [pose_feats[0], pose_feats[f_i]]

                    if self.opt.pose_model_type == "separate_resnet":
                        pose_inputs = [self.models["pose_encoder"](torch.cat(pose_inputs, 1))]
                    elif self.opt.pose_model_type == "posecnn":
                        pose_inputs = torch.cat(pose_inputs, 1)

                    axisangle, translation = self.models["pose"](pose_inputs)
                    outputs[("axisangle", 0, f_i)] = axisangle
                    outputs[("translation", 0, f_i)] = translation

                    # Invert the matrix if the frame id is negative
                    outputs[("cam_T_cam", 0, f_i)] = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0], invert=(f_i < 0))

        else:
            raise NotImplementedError

        return outputs

    def _generate_images_pred(self, inputs, outputs):
        """Generate the warped (reprojected) color images for a minibatch.
        Generated images are saved into the `outputs` dictionary.
        """
        for scale in self.opt.scales:
            disp = outputs[("disp", scale)]
            if self.opt.geo_consist:
                outputs[("disp", scale)] = outputs[("disp", scale)].chunk(len(self.opt.frame_ids), 0)[0]
            else:
                outputs[("disp", scale)] = outputs[("disp", scale)]

            if self.opt.v1_multiscale:
                source_scale = scale
            else:
                disp = F.interpolate(
                    disp, [self.opt.height, self.opt.width], mode="bilinear", align_corners=False)
                source_scale = 0

            _, depth = disp_to_depth(disp, self.opt.min_depth, self.opt.max_depth)
            

            if self.opt.geo_consist:
                depth = depth.chunk(len(self.opt.frame_ids), 0)
            else:
                depth = [depth]

            outputs[("depth", 0, scale)] = depth[0]

            for i, frame_id in enumerate(self.opt.frame_ids[1:]):

                if frame_id == "s":
                    T = inputs["stereo_T"]
                else:
                    T = outputs[("cam_T_cam", 0, frame_id)]

                # from the authors of https://arxiv.org/abs/1712.00175
                if self.opt.pose_model_type == "posecnn":

                    axisangle = outputs[("axisangle", 0, frame_id)]
                    translation = outputs[("translation", 0, frame_id)]

                    inv_depth = 1 / depth
                    mean_inv_depth = inv_depth.mean(3, True).mean(2, True)

                    T = transformation_from_parameters(
                        axisangle[:, 0], translation[:, 0] * mean_inv_depth[:, 0], frame_id < 0)

                cam_points = self.backproject_depth[source_scale](
                    depth[0], inputs[("inv_K", source_scale)])
                pix_coords, warp_depth = self.project_3d[source_scale](
                    cam_points, inputs[("K", source_scale)], T)

                outputs[("sample", frame_id, scale)] = pix_coords
                outputs[('warp_depth', frame_id, scale)] = warp_depth.unsqueeze(1)

                outputs[("color", frame_id, scale)] = F.grid_sample(
                    inputs[("color", frame_id, source_scale)],
                    outputs[("sample", frame_id, scale)],
                    padding_mode="border")

                if self.opt.geo_consist:
                    outputs[("depth",frame_id, scale)] = F.grid_sample(
                        depth[i+1],
                        outputs[("sample", frame_id, scale)],
                        padding_mode="zeros")

                if not self.opt.disable_automasking:
                    outputs[("color_identity", frame_id, scale)] = \
                        inputs[("color", frame_id, source_scale)]

    def _compute_reprojection_loss(self, pred, target):
        """Computes reprojection loss between a batch of predicted and target images
        """
        abs_diff = torch.abs(target - pred)
        l1_loss = abs_diff.mean(1, True)

        if self.opt.no_ssim:
            reprojection_loss = l1_loss
        else:
            ssim_loss = self.ssim(pred, target).mean(1, True)
            reprojection_loss = 0.85 * ssim_loss + 0.15 * l1_loss

        return reprojection_loss

    def _compute_depth_loss(self, disp, _gt_depth):
        # scale the ground truth
        assert not _gt_depth.isnan().any(), '  '.format(_gt_depth.isnan.sum())
        #mask = (F.interpolate(_gt_depth, disp.shape[-2:], mode='bilinear') > 0 ).detach()

        #gt_disp = 1 / (_gt_depth + 0.01)
        #depth = torch.clamp(F.interpolate(gt_disp, disp.shape[-2:], mode='bilinear'),
        #                    0.,
        #                    100.)
        
        _gt_depth   = _gt_depth.unsqueeze(1)
        mask = torch.logical_and(_gt_depth > 0, disp>0)
        _loss = F.l1_loss(torch.log(disp[mask]), torch.log(_gt_depth[mask]))
        return _loss

    def _compute_depth_consistency(self, outputs):
        loss_ = 0
        for i in self.opt.frame_ids[1:]:
            mask = outputs[("depth", i, 0)] > 0
            diff = (torch.log(outputs[("depth", i, 0)][mask]) - torch.log(outputs[('warp_depth', i, 0)][mask])).abs()
            loss_ += diff.mean()

        return loss_    

    def _compute_losses(self, inputs, outputs):
        """Compute the reprojection and smoothness losses for a minibatch
        """
        losses = {}
        total_loss = 0

        for scale in self.opt.scales:
            loss = 0
            reprojection_losses = []

            if self.opt.v1_multiscale:
                source_scale = scale
            else:
                source_scale = 0

            disp = outputs[("disp", scale)]
            color = inputs[("color", 0, scale)]
            target = inputs[("color", 0, source_scale)]

            for frame_id in self.opt.frame_ids[1:]:
                pred = outputs[("color", frame_id, scale)]
                reprojection_losses.append(self._compute_reprojection_loss(pred, target))

            reprojection_losses = torch.cat(reprojection_losses, 1)

            if not self.opt.disable_automasking:
                identity_reprojection_losses = []
                for frame_id in self.opt.frame_ids[1:]:
                    pred = inputs[("color", frame_id, source_scale)]
                    identity_reprojection_losses.append(
                        self._compute_reprojection_loss(pred, target))

                identity_reprojection_losses = torch.cat(identity_reprojection_losses, 1)

                if self.opt.avg_reprojection:
                    identity_reprojection_loss = identity_reprojection_losses.mean(1, keepdim=True)
                else:
                    # save both images, and do min all at once below
                    identity_reprojection_loss = identity_reprojection_losses

            elif self.opt.predictive_mask:
                # use the predicted mask
                mask = outputs["predictive_mask"]["disp", scale]
                if not self.opt.v1_multiscale:
                    mask = F.interpolate(
                        mask, [self.opt.height, self.opt.width],
                        mode="bilinear", align_corners=False)

                reprojection_losses *= mask

                # add a loss pushing mask to 1 (using nn.BCELoss for stability)
                weighting_loss = 0.2 * nn.BCELoss()(mask, torch.ones(mask.shape).cuda())
                loss += weighting_loss.mean()

            if self.opt.avg_reprojection:
                reprojection_loss = reprojection_losses.mean(1, keepdim=True)
            else:
                reprojection_loss = reprojection_losses

            if not self.opt.disable_automasking:
                # add random numbers to break ties
                identity_reprojection_loss += torch.randn(
                    identity_reprojection_loss.shape, device=self.device) * 0.00001

                combined = torch.cat((identity_reprojection_loss, reprojection_loss), dim=1)
            else:
                combined = reprojection_loss

            if combined.shape[1] == 1:
                to_optimise = combined
            else:
                to_optimise, idxs = torch.min(combined, dim=1)

            if not self.opt.disable_automasking:
                outputs["identity_selection/{}".format(scale)] = (
                    idxs > identity_reprojection_loss.shape[1] - 1).float()

            loss += to_optimise.mean()

            mean_disp = disp.mean(2, True).mean(3, True)
            norm_disp = disp / (mean_disp + 1e-7)
            smooth_loss = get_smooth_loss(norm_disp, color)
            
            #if ('depth_input', 0, 0) in inputs.keys():
            #    disp_loss = self._compute_depth_loss(disp, inputs[('depth_input', 0, 0)][:, -1])
            #    loss += disp_loss * self.opt.disp_loss_weight / (2 ** scale)
            #    losses["loss/disp_loss{}".format(scale)] = disp_loss

            if scale == 0 and ('depth_input', 0, 0) in inputs.keys():
                gt_disp = depth_to_disp(inputs[('depth_gt')][:, -1], self.opt.min_depth, self.opt.max_depth)
                disp_loss = self._compute_depth_loss(disp, gt_disp)
                loss += disp_loss * self.opt.disp_loss_weight / (2 ** scale)
                losses["loss/disp_loss{}".format(scale)] = disp_loss

            loss += self.opt.disparity_smoothness * smooth_loss / (2 ** scale)

            if ('disp_old', 0) in outputs.keys():
                idloss = F.mse_loss(disp , outputs[('disp_old', 0)][:, -1]) * 0.01
                loss += idloss
                losses['loss/id_loss']= idloss.detach().cpu().data

            if self.opt.geo_consist:
                geo_consist_loss = self._compute_depth_consistency(outputs)
                loss += self.opt.geo_consist * geo_consist_loss
                losses["loss/geo_loss{}".format(scale)] = geo_consist_loss

            total_loss += loss
            losses["loss/{}".format(scale)] = loss
            losses['loss/smooth_loss']= smooth_loss.detach().cpu().data


        total_loss /= self.num_scales
        losses["loss"] = total_loss
        return losses



    def train(self):
        """Convert all models to training mode
        """
        for m in self.models.values():
            m.train()

    def eval(self):
        """Convert all models to testing/evaluation mode
        """
        for m in self.models.values():
            m.eval()

    def save_model(self, path, e):
        """Save model weights to disk
        """
        save_folder = os.path.join(path, "models", "weights_{}".format(e))
        if not os.path.exists(save_folder):
            os.makedirs(save_folder)

        for model_name, model in self.models.items():
            save_path = os.path.join(save_folder, "{}.pth".format(model_name))
            to_save = model.state_dict()
            if model_name == 'encoder':
                # save the sizes - these are needed at prediction time
                to_save['height'] = self.opt.height
                to_save['width'] = self.opt.width
                to_save['use_stereo'] = self.opt.use_stereo
            torch.save(to_save, save_path)

    def load_model(self, path=None):
        """Load model(s) from disk
        """
        if path is None:
            self.opt.load_weights_folder = os.path.expanduser(self.opt.load_weights_folder)
            folder = self.opt.load_weights_folder
        else:
            folder = path

        assert os.path.isdir(folder), \
            "Cannot find folder {}".format(folder)
        print("loading model from folder {}".format(folder))

        for n in self.opt.models_to_load:
            print("Loading {} weights...".format(n))
            path = os.path.join(folder, "{}.pth".format(n))
            model_dict = self.models[n].state_dict()
            pretrained_dict = torch.load(path)
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
            model_dict.update(pretrained_dict)
            self.models[n].load_state_dict(model_dict)

src/models/flow2depth/model.py
embodied-depth/Embodied-Depth-Prediction
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import os
import argparse

file_dir = os.path.dirname(__file__)  # the directory that options.py resides in


def flow2depthOptions(parser):
     # PATHS
     parser.add_argument("--data_path",
                              type=str,
                              help="path to the training data",
                              default=os.path.join(file_dir, "kitti_data"))
     parser.add_argument("--log_dir",
                              type=str,
                              help="log directory",
                              default=os.path.join(os.path.expanduser("~"), "tmp"))

     # TRAINING options
     parser.add_argument("--model_name",
                              type=str,
                              help="the name of the folder to save the model in",
                              default="mdp")
     parser.add_argument("--split",
                              type=str,
                              help="which training split to use",
                              #choices=["eigen_zhou", "eigen_full", "odom", "benchmark", 'habitat','habitat2', 'habitat_pose', 'minikitti', 'real'],
                              default="eigen_zhou")
     parser.add_argument("--num_layers",
                              type=int,
                              help="number of resnet layers",
                              default=18,
                              choices=[18, 34, 50, 101, 152])
     parser.add_argument("--dataset",
                              type=str,
                              help="dataset to train on",
                              default="kitti",
                              choices=["kitti", "kitti_odom", "kitti_depth", "kitti_test", 'habitat', 'minikitti', 'real'])
     parser.add_argument("--height",
                              type=int,
                              help="input image height",
                              default=192)
     parser.add_argument("--width",
                              type=int,
                              help="input image width",
                              default=640)
     parser.add_argument("--disparity_smoothness",
                              type=float,
                              help="disparity smoothness weight",
                              default=1e-3)

     parser.add_argument("--disp_loss_weight",
                         help="disparity loss with gt depth",
                              type=float,
                              default=0.)
     parser.add_argument("--geo_consist",
                              help="geometry consistency loss weight",
                              type=float,
                              default=0.)

     parser.add_argument("--scales",
                              nargs="+",
                              type=int,
                              help="scales used in the loss",
                              default=[0, 1, 2, 3])
     parser.add_argument("--min_depth",
                              type=float,
                              help="minimum depth",
                              default=0.1)
     parser.add_argument("--max_depth",
                              type=float,
                              help="maximum depth",
                              default=100.0)
     parser.add_argument("--use_stereo",
                              help="if set, uses stereo pair for training",
                              action="store_true")
     parser.add_argument("--frame_ids",
                              nargs="+",
                              type=int,
                              help="frames to load",
                              default=[0, -1, 1])
     parser.add_argument("--warp_ids",
                              nargs="+",
                              type=int,
                              help="frames to load",
                              default=[0, -1, 1])

     # OPTIMIZATION options
     parser.add_argument("--batch_size",
                              type=int,
                              help="batch size",
                              default=12)
     parser.add_argument("--learning_rate",
                              type=float,
                              help="learning rate",
                              default=1e-4)

     parser.add_argument('--wdecay', type=float, default=0.0001)
     parser.add_argument('--epsilon', type=float, default=1e-8)
     parser.add_argument('--clip', type=float, default=-1.0)
     

     parser.add_argument('--num_steps', type=int, default=10000)
     parser.add_argument("--num_epochs",
                              type=int,
                              help="number of epochs",
                              default=20)
     parser.add_argument("--scheduler_step_size",
                              type=int,
                              help="step size of the scheduler",
                              default=15)

     # ABLATION options
     parser.add_argument("--warmup_orinet",
                         action="store_true",
                         default=False)

     parser.add_argument("--transform_in_refine",
                         action="store_true",
                         default=False)

     parser.add_argument("--flow2depth",
                         action="store_true",
                         default=False)

     parser.add_argument("--rgbdcat",
                         action="store_true",
                         default=False)

     parser.add_argument("--use_opencv_flow",
                              help="if set, and flow2depth is set meanwhile, use opencv func to generate optical flows",
                              action="store_true")

     ## RAFT options
     parser.add_argument("--flow_model_path",
                              default='')

     parser.add_argument("--small",
                         action="store_true",
                              default=False)
                              
     parser.add_argument("--device",
                              default='cuda')

     parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision')
     parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation')


     ### Depth decoder ablation
     parser.add_argument("--stack_video",
                         action="store_true",
                         default=False)

     parser.add_argument("--multi_pose",
                         action="store_true",
                         default=False)

     parser.add_argument("--pose_dist_info",
                         action="store_true",
                         default=False)

     parser.add_argument("--pose_rot",
                         action="store_true",
                         default=False)

     parser.add_argument("--pose_decay",
                         action="store_true",
                         default=False)

     parser.add_argument("--depth_encoder",
                         action="store_true",
                         default=False)


     parser.add_argument("--bayes",
                              action="store_true")
     parser.add_argument("--input_depth",
                              help="if set, input noisy depth map",
                              action="store_true")
     parser.add_argument("--input_pose",
                              help="if set, input noisy depth map",
                              action="store_true")



     parser.add_argument("--v1_multiscale",
                              help="if set, uses monodepth v1 multiscale",
                              action="store_true")
     parser.add_argument("--avg_reprojection",
                              help="if set, uses average reprojection loss",
                              action="store_true")
     parser.add_argument("--disable_automasking",
                              help="if set, doesn't do auto-masking",
                              action="store_true")
     parser.add_argument("--predictive_mask",
                              help="if set, uses a predictive masking scheme as in Zhou et al",
                              action="store_true")
     parser.add_argument("--no_ssim",
                              help="if set, disables ssim in the loss",
                              action="store_true")
     parser.add_argument("--weights_init",
                              type=str,
                              help="pretrained or scratch",
                              default="pretrained",
                              choices=["pretrained", "scratch"])
     parser.add_argument("--pose_model_input",
                              type=str,
                              help="how many images the pose network gets",
                              default="pairs",
                              choices=["pairs", "all"])
     parser.add_argument("--pose_model_type",
                              type=str,
                              help="normal or shared",
                              default="separate_resnet",
                              choices=["posecnn", "separate_resnet", "shared", 'gt'])

     # SYSTEM options
     parser.add_argument("--no_cuda",
                              help="if set disables CUDA",
                              action="store_true")

     # LOADING options
     parser.add_argument("--load_weights_folder",
                              type=str,
                              help="name of model to load")
     parser.add_argument("--models_to_load",
                              nargs="+",
                              type=str,
                              help="models to load",
                              default=["encoder", "depth", "pose_encoder", "pose"])

     # LOGGING options
     parser.add_argument("--log_frequency",
                              type=int,
                              help="number of batches between each tensorboard log",
                              default=250)
     parser.add_argument("--save_frequency",
                              type=int,
                              help="number of epochs between each save",
                              default=12)

     # EVALUATION options
     parser.add_argument("--eval_stereo",
                              help="if set evaluates in stereo mode",
                              action="store_true")
     parser.add_argument("--eval_mono",
                              help="if set evaluates in mono mode",
                              action="store_true")
     parser.add_argument("--disable_median_scaling",
                              help="if set disables median scaling in evaluation",
                              action="store_true")
     parser.add_argument("--pred_depth_scale_factor",
                              help="if set multiplies predictions by this number",
                              type=float,
                              default=1)
     parser.add_argument("--ext_disp_to_eval",
                              type=str,
                              help="optional path to a .npy disparities file to evaluate")
     parser.add_argument("--eval_split",
                              type=str,
                              default="eigen",
                              choices=[
                                   "eigen", "eigen_benchmark", "benchmark", "odom_9", "odom_10", "habitat", 'habitat_pose'],
                              help="which split to run eval on")
     parser.add_argument("--save_pred_disps",
                              help="if set saves predicted disparities",
                              action="store_true")
     parser.add_argument("--no_eval",
                              help="if set disables evaluation",
                              action="store_true")
     parser.add_argument("--eval_eigen_to_benchmark",
                              help="if set assume we are loading eigen results from npy but "
                                   "we want to evaluate using the new benchmark.",
                              action="store_true")
     parser.add_argument("--eval_out_dir",
                              help="if set will output the disparities to this folder",
                              type=str)
     parser.add_argument("--post_process",
                              help="if set will perform the flipping post processing "
                                   "from the original monodepth paper",
                              action="store_true")

     return parser 


src/models/flow2depth/options.py
embodied-depth/Embodied-Depth-Prediction
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import numpy as np

import torch
import torch.nn as nn
import torchvision.models as models
import torch.utils.model_zoo as model_zoo

def depth_to_disp(_depth, min_depth, max_depth):
    depth = torch.clamp(_depth, min_depth, max_depth)
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = 1 / depth
    disp_ = (scaled_disp - min_disp) / (max_disp - min_disp)
    return disp_

class ResNetMultiImageInput(models.ResNet):
    """Constructs a resnet model with varying number of input images.
    Adapted from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
    """
    def __init__(self, block, layers, num_classes=1000, num_input_images=1, input_depth=False):
        super(ResNetMultiImageInput, self).__init__(block, layers)
        self.inplanes = 64
        self.input_depth = input_depth
        if input_depth:
            self.conv1 = nn.Conv2d(
                num_input_images * 4, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        else:
            self.conv1 = nn.Conv2d(
                num_input_images * 3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(self.inplanes)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)


def resnet_multiimage_input(num_layers, pretrained=False, num_input_images=1, input_depth=False):
    """Constructs a ResNet model.
    Args:
        num_layers (int): Number of resnet layers. Must be 18 or 50
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        num_input_images (int): Number of frames stacked as input
    """
    assert num_layers in [18, 50], "Can only run with 18 or 50 layer resnet"
    blocks = {18: [2, 2, 2, 2], 50: [3, 4, 6, 3]}[num_layers]
    block_type = {18: models.resnet.BasicBlock, 50: models.resnet.Bottleneck}[num_layers]
    model = ResNetMultiImageInput(block_type, blocks, num_input_images=num_input_images, input_depth=input_depth)

    if pretrained:
        loaded = model_zoo.load_url(models.resnet.model_urls['resnet{}'.format(num_layers)])
        if input_depth:
            loaded['conv1.weight'] = torch.cat([loaded['conv1.weight'], torch.randn((64, 1, 7, 7 ))], dim=1)
        loaded['conv1.weight'] = torch.cat(
            [loaded['conv1.weight']] * num_input_images, 1) / num_input_images
        model.load_state_dict(loaded)
    return model


class ResnetEncoder(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self, num_layers, pretrained, num_input_images=1, input_depth=False):
        super(ResnetEncoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512])
        self.input_depth = input_depth

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))

        if num_input_images > 1:
            self.encoder = resnet_multiimage_input(num_layers, pretrained, num_input_images, input_depth=input_depth)
        else:
            self.encoder = resnets[num_layers](pretrained)
            if self.input_depth:
                self.encoder.conv1 =  nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        if self.input_depth:
            bias = torch.tensor([0.485, 0.456, 0.406, 1.565]).reshape(1,4, 1,1).repeat(1, num_input_images, 1, 1)
            sigma = torch.tensor([0.229, 0.224, 0.225, 1.761]).reshape(1,4, 1,1).repeat(1,num_input_images , 1, 1)

        else:
            bias = torch.tensor([0.485, 0.456, 0.406]).reshape(1,3, 1,1).repeat(1, num_input_images, 1, 1)
            sigma = torch.tensor([0.229, 0.224, 0.225]).reshape(1,3, 1,1).repeat(1,num_input_images , 1, 1)

        #self.register_buffer('bias', bias)
        #self.register_buffer('sigma', sigma)

        self.bias = bias
        self.sigma = sigma

    def forward(self, input_image):
        self.features = []
        #if self.bias == None:
        #    if self.input_depth:
        #        self.bias = torch.tensor([0.485, 0.456, 0.406, 1.565]).reshape(1,4, 1,1).to(input_image.device)
        #        self.sigma = torch.tensor([0.229, 0.224, 0.225, 1.761]).reshape(1,4, 1,1).to(input_image.device)
        #    else:
        #        self.bias = torch.tensor([0.485, 0.456, 0.406 ]).reshape(1,3, 1,1).to(input_image.device)
        #        self.sigma = torch.tensor([0.229, 0.224, 0.225 ]).reshape(1,3, 1,1).to(input_image.device)

        x = (input_image - self.bias.to(input_image.device)) / self.sigma.to(input_image.device)

        x = self.encoder.conv1(x)
        x = self.encoder.bn1(x)
        self.features.append(self.encoder.relu(x))
        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))
        self.features.append(self.encoder.layer2(self.features[-1]))
        self.features.append(self.encoder.layer3(self.features[-1]))
        self.features.append(self.encoder.layer4(self.features[-1]))

        return self.features

class MultiInputResnetEncoder(nn.Module):
    """Pytorch module for a resnet encoder
    """
    def __init__(self, num_layers, pretrained, num_input_images=1):
        super(MultiInputResnetEncoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512]) * num_input_images

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))


        self.encoder = resnets[num_layers](pretrained)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        bias = torch.tensor([0.485, 0.456, 0.406]).reshape(1,3, 1,1)
        sigma = torch.tensor([0.229, 0.224, 0.225]).reshape(1,3, 1,1)

        self.bias = bias
        self.sigma = sigma

    def forward(self, input_image):
        '''
        Input:
            input_image [B, T, C, H, W] 
        Output:
            features [B, T, c, h, w]
        '''
        assert input_image.dim() == 5, 'Wrong input dims should be 5, but received {}'.format(input_image.dim())
        self.features = []

        B, T, _, _, _ = input_image.shape
        input_image = input_image.flatten(0,1)
        x = (input_image - self.bias.to(input_image.device)) / self.sigma.to(input_image.device)

        x = self.encoder.conv1(x)
        x = self.encoder.bn1(x)
        self.features.append(self.encoder.relu(x))
        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))
        self.features.append(self.encoder.layer2(self.features[-1]))
        self.features.append(self.encoder.layer3(self.features[-1]))
        self.features.append(self.encoder.layer4(self.features[-1]))

        for i in range(len(self.features)):
            _, c, h, w = self.features[i].shape
            self.features[i] = self.features[i].reshape(B, T, c, h, w)

        return self.features

class MultiInputDepthEncoder(nn.Module):
    def __init__(self, num_layers, num_input_images=1):
        super(MultiInputDepthEncoder, self).__init__()

        self.num_ch_enc = np.array([64, 64, 128, 256, 512]) * num_input_images

        resnets = {18: models.resnet18,
                   34: models.resnet34,
                   50: models.resnet50,
                   101: models.resnet101,
                   152: models.resnet152}

        if num_layers not in resnets:
            raise ValueError("{} is not a valid number of resnet layers".format(num_layers))

        pretrained = False

        self.encoder = resnets[num_layers](pretrained)
        self.encoder.conv1 =  nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)

        if num_layers > 34:
            self.num_ch_enc[1:] *= 4

        bias  = torch.tensor([ 0.]).reshape(1,1, 1,1)
        sigma = torch.tensor([ 1.]).reshape(1,1, 1,1)
        self.bias = bias
        self.sigma = sigma

    def forward(self, input_image):
        '''
        Input:
            input_image [B, T, C, H, W] 
        Output:
            features [B, T, c, h, w]
        '''
        assert input_image.dim() == 5, 'Wrong input dims should be 5, but received {}'.format(input_image.dim())
        self.features = []

        B, T, _, _, _ = input_image.shape
        input_image = input_image.flatten(0,1)
        x = (input_image - self.bias.to(input_image.device)) / self.sigma.to(input_image.device)

        x = self.encoder.conv1(x)
        x = self.encoder.bn1(x)
        self.features.append(self.encoder.relu(x))
        self.features.append(self.encoder.layer1(self.encoder.maxpool(self.features[-1])))
        self.features.append(self.encoder.layer2(self.features[-1]))
        self.features.append(self.encoder.layer3(self.features[-1]))
        self.features.append(self.encoder.layer4(self.features[-1]))

        for i in range(len(self.features)):
            _, c, h, w = self.features[i].shape
            self.features[i] = self.features[i].reshape(B, T, c, h, w)

        return self.features

# Cat RGB embedding and Depth embedding
class RGBDEncoder(nn.Module):
    def __init__(self, num_layers, pretrained, num_input_images=1, input_depth=False):
        super(RGBDEncoder, self).__init__()
        self.input_depth = input_depth
        self.rgb_encoder = MultiInputResnetEncoder(num_layers, pretrained, num_input_images)
        if self.input_depth:
            self.d_encoder  = MultiInputDepthEncoder(num_layers, num_input_images)
            self.num_ch_enc = [ i+j for i, j in zip(self.rgb_encoder.num_ch_enc, self.d_encoder.num_ch_enc)]
        else:
            self.num_ch_enc = self.rgb_encoder.num_ch_enc

    def forward(self, rgb, d=None):
        rgb_f = self.rgb_encoder(rgb)
        if self.input_depth:
            d_f   = self.d_encoder(d)
            out_ = []
            for f1, f2 in zip(rgb_f, d_f):
                out_.append(torch.cat([f1, f2], dim=2))
        else:
            out_ = rgb_f

        return out_

# Add RGB embedding and Depth embedding
class RGBDEncoderv2(nn.Module):
    def __init__(self, num_layers, pretrained, num_input_images=1, input_depth=False):
        super(RGBDEncoderv2, self).__init__()
        self.input_depth = input_depth
        self.rgb_encoder = MultiInputResnetEncoder(num_layers, pretrained, num_input_images)
        if self.input_depth:
            self.d_encoder  = MultiInputDepthEncoder(num_layers, num_input_images)
        self.num_ch_enc = self.rgb_encoder.num_ch_enc

    def forward(self, rgb, d=None):
        rgb_f = self.rgb_encoder(rgb)
        if self.input_depth:
            d_f   = self.d_encoder(d)
            out_ = []
            for f1, f2 in zip(rgb_f, d_f):
                out_.append(f1 + f2)
        else:
            out_ = rgb_f

        return out_

if __name__ == '__main__':
    net = ResnetEncoder(18, True, 5, False)
src/models/flow2depth/resnet_encoder.py
embodied-depth/Embodied-Depth-Prediction
import os
import sys
import json
from PIL import Image  # using pillow-simd for increased speed
from pathlib import Path
from collections import deque

import time
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from torch.utils.data import DataLoader, Subset
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler, WeightedRandomSampler
from tensorboardX import SummaryWriter


from src.models.layers import *
from src.models.utils import *
from src.datasets  import MHabitatRAWDataset, MRealRAWDataset, EmDataset 
from src.my_utils import *

from src.models.base_trainer import BaseTrainer


class MyTrainer(BaseTrainer):
    def __init__(self, opt):
        from src.models.flow2depth.model import DepthModel

        self.model = DepthModel(opt)

        super().__init__(opt)

        self.num_scales = len(self.opt.scales)

        param_to_train = self.model.parameters_to_train #if self.opt.model_type != 'ours' else \
        #[
        #        {'params': self.model.parameters_to_train},
        #        {'params': self.model.models['flow2depth'].parameters(), 'lr': 1e-7}
        #    ]
        self.model_optimizer = optim.Adam(
                                    param_to_train, 
                                    self.opt.learning_rate, 
                                    weight_decay=0,
                                    eps=self.opt.epsilon)
        self.model_lr_scheduler = optim.lr_scheduler.StepLR(
            self.model_optimizer, self.opt.scheduler_step_size, 0.1)

        self.save_opts()

    def reset_model(self):
        if self.opt.model_type == 'monodepth':
            from src.models.monodepth.model import DepthModel
        elif self.opt.model_type == 'ours':
            from src.models.flow2depth.model import DepthModel
        self.model = DepthModel(self.opt)
        self.model_optimizer = optim.Adam(self.model.parameters_to_train, self.opt.learning_rate, weight_decay=0, eps=self.opt.epsilon)
        self.model_lr_scheduler = optim.lr_scheduler.StepLR(
            self.model_optimizer, self.opt.scheduler_step_size, 0.1)

    def train_offline(self):
        """Run the entire training pipeline
        """
        img_ext = '.png' if self.opt.png else '.jpg'
        train_filenames = readlines(self.fpath.format("train{}".format(self.opt.train_scene_id)))
        num_train_samples = len(train_filenames)
        self.num_total_steps = num_train_samples // self.opt.batch_size * self.opt.num_epochs

        train_dataset = self.dataset(
            self.opt.data_path, train_filenames, self.opt.height, self.opt.width,
            self.model.frames_to_load_train, 
            num_scales=self.num_scales, is_train=True, img_ext=img_ext, 
            load_pose=self.opt.pose_model_type=='gt',
            input_depth=self.opt.depth_encoder,
            num_past_frame_input=0,
            video_interval=self.opt.video_interval,
            input_T_dim=self.opt.input_T_dim 
            )
        self.train_loader = DataLoader(
            train_dataset, self.opt.batch_size, True,
            num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)

        self.writers = {}
        for mode in ["train", "test"]:
            self.writers[mode] = SummaryWriter(os.path.join(self.log_path, mode))

        self.epoch = 0
        self.start_time = time.time()
        for self.epoch in range(self.opt.num_epochs):
            self._run_epoch()
            print("Epoch {}, training step {}, lr = {}".format(self.epoch, self.step, self.model_lr_scheduler.get_last_lr()))

    def _run_epoch(self):
        """Run a single epoch of training and validation
        """

        self.model_lr_scheduler.step()
        print("Training")
        self.model.train()

        for batch_idx, inputs in enumerate(self.train_loader):

            before_op_time = time.time()

            outputs, losses = self.model.forward(inputs)

            self.model_optimizer.zero_grad()
            losses["loss"].backward()

            if self.opt.clip > 0 :
                torch.nn.utils.clip_grad_norm_(self.model.parameters_to_train, self.opt.clip)
            self.model_optimizer.step()
            #self.model_lr_scheduler.step()

            duration = time.time() - before_op_time

            # log less frequently after the first 2000 steps to save time & disk space
            early_phase = batch_idx % self.opt.log_frequency == 0 and self.step < 3000
            late_phase = self.step % 3000 == 0

            if early_phase or late_phase:
                self.log_time(batch_idx, duration, losses["loss"].cpu().data)

                if "depth_gt" in inputs:
                    depth_pred = outputs[("depth", 0, 0)]
                    depth_gt = inputs["depth_gt"]
                    mask = depth_gt > 0
                    depth_gt = depth_gt[mask]
                    depth_pred = depth_pred[mask]

                    depth_errors = compute_depth_errors(depth_gt, depth_pred)
                    for metric, val in depth_errors.items():
                        losses[metric] = np.array(depth_errors[metric].cpu().data)

                self.log("train", inputs, outputs, losses)

            #if self.step % 2048 == 10:
            #    self.model.eval()
            #    self.test(self.writers['test'], self.step)
            #    self.model.train()

            self.step += 1

        if self.epoch % 2 == 0:
            self.set_eval()
            self.val(self.writers['train'], self.step)
            if self.epoch % 4 == 0:
                self.save_model(self.log_path, 'e_{}'.format(self.epoch ))
            self.test(self.writers['test'], self.step)
            self.set_train()


    def set_train(self):
        self.model.train()

    def set_eval(self):
        self.model.eval()

    def process_batch(self, inputs):
        '''
        return:
            outputs
            losses 
        '''
        return self.model.forward(inputs)


    def save_model(self, path, name):
        '''
        Save all model ckpts and optimizer in the following path
         `[tensorboard_dir]/[self.log_path]/models/weights_{epoch}_val_{valacc}/`
        '''
        self.model.save_model(path, name)
        #save_path = os.path.join(save_folder, "{}.pth".format("adam"))
        #torch.save(self.model_optimizer.state_dict(), save_path)


    def load_model(self, path):
        self.model.load_model(path)


    def log_time(self, batch_idx, duration, loss):
        """Print a logging statement to the terminal
        """
        samples_per_sec = self.opt.batch_size / duration
        time_sofar = time.time() - self.start_time
        training_time_left = (
            self.num_total_steps / self.step - 1.0) * time_sofar if self.step > 0 else 0
        print_string = "epoch {:>3} | batch {:>6} | examples/s: {:5.1f}" + \
            " | loss: {:.5f} | time elapsed: {} | time left: {}"
        print(print_string.format(self.epoch, batch_idx, samples_per_sec, loss,
                                  sec_to_hm_str(time_sofar), sec_to_hm_str(training_time_left)))

    def log(self, mode, inputs, outputs, losses):
        """Write an event to the tensorboard events file
        """
        if type(mode) == str:
            writer = self.writers[mode]
        else:
            writer = mode

        if type(outputs) == list:
            sigmas = self.get_uncertainty(outputs)
            outputs = outputs[0]
        else:
            sigmas = None

        for l, v in losses.items():
            writer.add_scalar("{}".format(l), v, self.step)

        for j in range(min(4, self.opt.batch_size)):  # write a maxmimum of four images
            writer.add_image(
                "gt_disp/{}".format(j),
                colorize(depth_to_disp(inputs["depth_gt"][j].cpu().detach(), 
                                self.opt.min_depth,
                                self.opt.max_depth)), 
                 self.step)
            if ('depth_input', 0, 0) in inputs.keys():
                writer.add_image(
                    "input_disp/{}".format(j),
                    colorize(inputs[('depth_input', 0, 0)][j, -1].cpu().detach()), 
                    self.step)

            if sigmas is not None:
                writer.add_image(
                    "depth_Nsigma/{}".format(j),
                    normalize_image(sigmas[j]), self.step)

            if not(("color", self.opt.frame_ids[1], 0) in inputs.keys()):
                break
            for s in self.opt.scales:
                for frame_id in self.opt.frame_ids:
                    writer.add_image(
                        "color_{}_{}/{}".format(frame_id, s, j),
                        inputs[("color", frame_id, s)][j,:3].cpu().data, self.step)
                    if s == 0 and frame_id != 0:
                        writer.add_image(
                            "color_pred_{}_{}/{}".format(frame_id, s, j),
                            outputs[("color", frame_id, s)][j,:3].cpu().data, self.step)
                    if self.opt.geo_consist and frame_id!= 0:
                        writer.add_image(
                            "computed_depth_{}/{}".format(frame_id, j),
                            colorize(normalize_image(outputs[("warp_depth", frame_id, 0)][j]).cpu().detach()), self.step)
                        writer.add_image(
                            "warp_depth_{}/{}".format(frame_id, j),
                            colorize(normalize_image(outputs[("depth", frame_id, 0)][j].cpu().detach())), self.step)

                writer.add_image(
                    "disp_{}/{}".format(s, j),
                    colorize(outputs[("disp", s)][j].cpu().detach()), self.step)

                if ("disp_old", s) in outputs.keys():
                    writer.add_image(
                    "disp_old_{}/{}".format(s, j),
                    colorize(outputs[("disp_old", s)][j, -1].cpu().detach()), self.step)
                

                if self.opt.predictive_mask:
                    for f_idx, frame_id in enumerate(self.opt.frame_ids[1:]):
                        writer.add_image(
                            "predictive_mask_{}_{}/{}".format(frame_id, s, j),
                            outputs["predictive_mask"][("disp", s)][j, f_idx][None, ...],
                            self.step)

                #elif not self.opt.disable_automasking:
                #    writer.add_image(
                #        "automask_{}/{}".format(s, j),
                #        outputs["identity_selection/{}".format(s)][j][None, ...], self.step)

    def special_train_offline(self):
        """Run the entire training pipeline
        """
        img_ext = '.png' if self.opt.png else '.jpg'
        train_filenames = readlines(self.fpath.format("train{}".format(self.opt.train_scene_id)))
        num_train_samples = len(train_filenames)
        self.num_total_steps = num_train_samples // self.opt.batch_size * self.opt.num_epochs

        train_dataset = self.dataset(
            self.opt.data_path, train_filenames, self.opt.height, self.opt.width,
            self.model.frames_to_load_train, 
            num_scales=1, is_train=True, img_ext=img_ext, 
            load_pose=self.opt.pose_model_type=='gt',
            input_depth=self.opt.depth_encoder,
            num_past_frame_input=0,
            video_interval=self.opt.video_interval,
            input_T_dim=self.opt.input_T_dim 
            )
        
        whole_length = len(train_dataset)

        self.writers = {}
        for mode in ["train", "test"]:
            self.writers[mode] = SummaryWriter(os.path.join(self.log_path, mode))

        self.epoch = 0
        self.start_time = time.time()

        start_len = 6000
        for self.epoch in range(self.opt.num_epochs):
            if self.epoch % 1 == 0:
                current_train_set = Subset(train_dataset, [i for i in range(start_len)])
                self.writers['train'].add_scalar('Explor/data_num', start_len, self.step)
                self.train_loader = DataLoader(
                    current_train_set , self.opt.batch_size, True,
                    num_workers=self.opt.num_workers, pin_memory=True, drop_last=True)
                if self.epoch >= 0:
                    start_len += 3000
                    start_len = min(start_len, whole_length)

            self._run_epoch()
            print("Epoch {}, training step {}, len_data={}, lr = {}".format(self.epoch, self.step, len(current_train_set), self.model_lr_scheduler.get_last_lr()))



if __name__ == "__main__":
    from monodepth2.options import MonodepthOptions

    options = MonodepthOptions()
    opts = options.parse()

    trainer = ModelTrainer(opts)
    trainer.train_offline()
    #trainer.test()

src/models/flow2depth/trainer.py
embodied-depth/Embodied-Depth-Prediction
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function

import numpy as np

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

        
def cam1_to_cam0(T_world_camera1, T_world_camera0):
    T_camera1_world = np.linalg.pinv(T_world_camera1)
    T_camera1_camera0 = np.matmul(T_camera1_world, T_world_camera0)
    return T_camera1_camera0

def disp_to_depth(disp, min_depth, max_depth):
    """Convert network's sigmoid output into depth prediction
    The formula for this conversion is given in the 'additional considerations'
    section of the paper.
    """
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = min_disp + (max_disp - min_disp) * disp
    depth = 1 / scaled_disp
    return scaled_disp, depth

def depth_to_disp(_depth, min_depth, max_depth):
    depth = torch.clamp(_depth, min_depth, max_depth)
    min_disp = 1 / max_depth
    max_disp = 1 / min_depth
    scaled_disp = 1 / depth
    disp_ = (scaled_disp - min_disp) / (max_disp - min_disp)
    return disp_


def transformation_from_parameters(axisangle, translation, invert=False):
    """Convert the network's (axisangle, translation) output into a 4x4 matrix
    """
    R = rot_from_axisangle(axisangle)
    t = translation.clone()

    if invert:
        R = R.transpose(1, 2)
        t *= -1

    T = get_translation_matrix(t)

    if invert:
        M = torch.matmul(R, T)
    else:
        M = torch.matmul(T, R)

    return M


def get_translation_matrix(translation_vector):
    """Convert a translation vector into a 4x4 transformation matrix
    """
    T = torch.zeros(translation_vector.shape[0], 4, 4).to(device=translation_vector.device)

    t = translation_vector.contiguous().view(-1, 3, 1)

    T[:, 0, 0] = 1
    T[:, 1, 1] = 1
    T[:, 2, 2] = 1
    T[:, 3, 3] = 1
    T[:, :3, 3, None] = t

    return T


def rot_from_axisangle(vec):
    """Convert an axisangle rotation into a 4x4 transformation matrix
    (adapted from https://github.com/Wallacoloo/printipi)
    Input 'vec' has to be Bx1x3
    """
    angle = torch.norm(vec, 2, 2, True)
    axis = vec / (angle + 1e-7)

    ca = torch.cos(angle)
    sa = torch.sin(angle)
    C = 1 - ca

    x = axis[..., 0].unsqueeze(1)
    y = axis[..., 1].unsqueeze(1)
    z = axis[..., 2].unsqueeze(1)

    xs = x * sa
    ys = y * sa
    zs = z * sa
    xC = x * C
    yC = y * C
    zC = z * C
    xyC = x * yC
    yzC = y * zC
    zxC = z * xC

    rot = torch.zeros((vec.shape[0], 4, 4)).to(device=vec.device)

    rot[:, 0, 0] = torch.squeeze(x * xC + ca)
    rot[:, 0, 1] = torch.squeeze(xyC - zs)
    rot[:, 0, 2] = torch.squeeze(zxC + ys)
    rot[:, 1, 0] = torch.squeeze(xyC + zs)
    rot[:, 1, 1] = torch.squeeze(y * yC + ca)
    rot[:, 1, 2] = torch.squeeze(yzC - xs)
    rot[:, 2, 0] = torch.squeeze(zxC - ys)
    rot[:, 2, 1] = torch.squeeze(yzC + xs)
    rot[:, 2, 2] = torch.squeeze(z * zC + ca)
    rot[:, 3, 3] = 1

    return rot


class ConvBlock(nn.Module):
    """Layer to perform a convolution followed by ELU
    """
    def __init__(self, in_channels, out_channels, bayes=False):
        super(ConvBlock, self).__init__()

        if bayes:
            self.conv = BayesianConv3x3(in_channels, out_channels)
        else:
            self.conv = Conv3x3(in_channels, out_channels)
        self.nonlin = nn.ELU(inplace=True)

    def forward(self, x):
        out = self.conv(x)
        out = self.nonlin(out)
        return out


class Conv3x3(nn.Module):
    """Layer to pad and convolve input
    """
    def __init__(self, in_channels, out_channels, use_refl=True):
        super(Conv3x3, self).__init__()

        if use_refl:
            self.pad = nn.ReflectionPad2d(1)
        else:
            self.pad = nn.ZeroPad2d(1)
        self.conv = nn.Conv2d(int(in_channels), int(out_channels), 3)

    def forward(self, x):
        out = self.pad(x)
        out = self.conv(out)
        return out

class BayesianConv3x3(nn.Module):

    def __init__(self, in_ch, out_ch, kernel_size=(3, 3), drop_rate=0.1, use_refl=True):
        super(BayesianConv3x3, self).__init__()
        self.p = 1 - drop_rate
        assert(0 <= self.p <= 1)

        if use_refl:
            self.pad = nn.ReflectionPad2d(1)
        else:
            self.pad = nn.ZeroPad2d(1)

        self.conv_op = F.conv2d
        self.conv_layer = torch.nn.Conv2d(in_channels=in_ch, out_channels=out_ch, kernel_size=kernel_size,
                                          )

    def forward(self, x):
        out = self.pad(x)

        mask = torch.ones(self.conv_layer.weight.shape).to(self.conv_layer.weight.device) * self.p
        mask = torch.bernoulli(mask)

        masked_kernels = self.conv_layer.weight * mask
        return self.conv_op(out, masked_kernels, self.conv_layer.bias)
    

class BackprojectDepth(nn.Module):
    """Layer to transform a depth image into a point cloud
    """
    def __init__(self, batch_size, height, width):
        super(BackprojectDepth, self).__init__()

        self.batch_size = batch_size
        self.height = height
        self.width = width

        meshgrid = np.meshgrid(range(self.width), range(self.height), indexing='xy')
        self.id_coords = np.stack(meshgrid, axis=0).astype(np.float32)
        self.id_coords = nn.Parameter(torch.from_numpy(self.id_coords),
                                      requires_grad=False)

        self.ones = nn.Parameter(torch.ones(self.batch_size, 1, self.height * self.width),
                                 requires_grad=False)

        self.pix_coords = torch.unsqueeze(torch.stack(
            [self.id_coords[0].view(-1), self.id_coords[1].view(-1)], 0), 0)
        self.pix_coords = self.pix_coords.repeat(batch_size, 1, 1)
        self.pix_coords = nn.Parameter(torch.cat([self.pix_coords, self.ones], 1),
                                       requires_grad=False)

    def forward(self, depth, inv_K):
        cam_points = torch.matmul(inv_K[:, :3, :3], self.pix_coords)
        cam_points = depth.view(self.batch_size, 1, -1) * cam_points
        cam_points = torch.cat([cam_points, self.ones], 1)

        return cam_points


class Project3D(nn.Module):
    """Layer which projects 3D points into a camera with intrinsics K and at position T
    """
    def __init__(self, batch_size, height, width, eps=1e-7):
        super(Project3D, self).__init__()

        self.batch_size = batch_size
        self.height = height
        self.width = width
        self.eps = eps

    def forward(self, points, K, T):
        P = torch.matmul(K, T)[:, :3, :]

        cam_points = torch.matmul(P, points)

        # TODO
        pix_coords = cam_points[:, :2, :] / (cam_points[:, 2, :].unsqueeze(1) + self.eps)
        pix_coords = pix_coords.view(self.batch_size, 2, self.height, self.width)
        pix_coords = pix_coords.permute(0, 2, 3, 1)
        pix_coords[..., 0] /= self.width - 1
        pix_coords[..., 1] /= self.height - 1
        pix_coords = (pix_coords - 0.5) * 2

        return pix_coords, cam_points[:, 2, :].reshape(-1, self.height, self.width)

class HabBackprojectDepth(nn.Module):
    """Layer to transform a depth image into a point cloud
    """
    def __init__(self, batch_size, height, width):
        super(HabBackprojectDepth, self).__init__()

        self.batch_size = batch_size
        self.height = height
        self.width = width

        meshgrid = np.meshgrid(np.linspace(-1,1,self.width), np.linspace(1,-1,self.height), indexing='xy')
        self.id_coords = np.stack(meshgrid, axis=0).astype(np.float32)
        self.id_coords = nn.Parameter(torch.from_numpy(self.id_coords),
                                      requires_grad=False)

        self.ones = nn.Parameter(torch.ones(self.batch_size, 1, self.height * self.width),
                                 requires_grad=False)

        self.pix_coords = torch.unsqueeze(torch.stack(
            [self.id_coords[0].view(-1), self.id_coords[1].view(-1)], 0), 0)
        self.pix_coords = self.pix_coords.repeat(batch_size, 1, 1)
        self.pix_coords = nn.Parameter(torch.cat([self.pix_coords, -self.ones], 1),
                                       requires_grad=False)

    def forward(self, depth, inv_K):
        cam_points = torch.matmul(inv_K[:, :3, :3], self.pix_coords)
        cam_points = depth.view(self.batch_size, 1, -1) * cam_points
        cam_points = torch.cat([cam_points, self.ones], 1)

        return cam_points


class HabProject3D(nn.Module):
    """Layer which projects 3D points into a camera with intrinsics K and at position T
    """
    def __init__(self, batch_size, height, width, eps=1e-7):
        super(HabProject3D, self).__init__()

        self.batch_size = batch_size
        self.height = height
        self.width = width
        self.eps = eps

    def forward(self, points, K, T):
        P = torch.matmul(K, T)[:, :3, :]

        cam_points = torch.matmul(P, points)

        # TODO
        pix_coords = cam_points[:, :2, :] / -(cam_points[:, 2, :].unsqueeze(1) + self.eps)
        pix_coords = pix_coords.view(self.batch_size, 2, self.height, self.width)
        pix_coords = pix_coords.permute(0, 2, 3, 1)

        pix_coords[:, :, :, 1] *= -1
        return pix_coords, -cam_points[:, 2, :].reshape(-1, self.height, self.width)


def upsample(x):
    """Upsample input tensor by a factor of 2
    """
    return F.interpolate(x, scale_factor=2, mode="nearest")


def get_smooth_loss(disp, img):
    """Computes the smoothness loss for a disparity image
    The color image is used for edge-aware smoothness
    """
    grad_disp_x = torch.abs(disp[:, :, :, :-1] - disp[:, :, :, 1:])
    grad_disp_y = torch.abs(disp[:, :, :-1, :] - disp[:, :, 1:, :])

    grad_img_x = torch.mean(torch.abs(img[:, :, :, :-1] - img[:, :, :, 1:]), 1, keepdim=True)
    grad_img_y = torch.mean(torch.abs(img[:, :, :-1, :] - img[:, :, 1:, :]), 1, keepdim=True)

    grad_disp_x *= torch.exp(-grad_img_x)
    grad_disp_y *= torch.exp(-grad_img_y)

    return grad_disp_x.mean() + grad_disp_y.mean()


class SSIM(nn.Module):
    """Layer to compute the SSIM loss between a pair of images
    """
    def __init__(self):
        super(SSIM, self).__init__()
        self.mu_x_pool   = nn.AvgPool2d(3, 1)
        self.mu_y_pool   = nn.AvgPool2d(3, 1)
        self.sig_x_pool  = nn.AvgPool2d(3, 1)
        self.sig_y_pool  = nn.AvgPool2d(3, 1)
        self.sig_xy_pool = nn.AvgPool2d(3, 1)

        self.refl = nn.ReflectionPad2d(1)

        self.C1 = 0.01 ** 2
        self.C2 = 0.03 ** 2

    def forward(self, x, y):
        x = self.refl(x)
        y = self.refl(y)

        mu_x = self.mu_x_pool(x)
        mu_y = self.mu_y_pool(y)

        sigma_x  = self.sig_x_pool(x ** 2) - mu_x ** 2
        sigma_y  = self.sig_y_pool(y ** 2) - mu_y ** 2
        sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y

        SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)
        SSIM_d = (mu_x ** 2 + mu_y ** 2 + self.C1) * (sigma_x + sigma_y + self.C2)

        return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)


def compute_depth_errors(gt, pred):
    out_ = {}
    out_.update(_compute_depth_errors(gt, pred, ''))

    if (gt > 5).sum() > 0 :
        gt_5m = gt[gt > 5]
        pred_5m = pred[gt>5]
        out_.update(_compute_depth_errors(gt_5m, pred_5m, '>5m_'))
    return out_


def _compute_depth_errors(gt, pred, name):
    """Computation of error metrics between predicted and ground truth depths
    """
    thresh = torch.max((gt / pred), (pred / gt))
    a1 = (thresh < 1.25     ).float().mean()
    a2 = (thresh < 1.25 ** 2).float().mean()
    a3 = (thresh < 1.25 ** 3).float().mean()

    rmse = (gt - pred) ** 2
    rmse = torch.sqrt(rmse.mean())

    rmse_log = (torch.log(gt) - torch.log(pred)) ** 2
    rmse_log = torch.sqrt(rmse_log.mean())

    abs_rel = torch.mean(torch.abs(gt - pred) / gt)

    sq_rel = torch.mean((gt - pred) ** 2 / gt)

    out_ = {
         "de/" + name + "abs_rel": abs_rel,
         "de/" + name + "sq_rel" : sq_rel,
         "de/" + name + "rms":rmse, 
         "de/" + name + "log_rms":rmse_log, 
         "da/" + name + "a1":a1, 
         "da/" + name + "a2":a2, 
         "da/" + name + "a3":a3
    }
    return out_

src/models/layers.py
embodied-depth/Embodied-Depth-Prediction
# Copyright Niantic 2019. Patent Pending. All rights reserved.
#
# This software is licensed under the terms of the Monodepth2 licence
# which allows for non-commercial use only, the full terms of which are made
# available in the LICENSE file.

from __future__ import absolute_import, division, print_function
import os
import torch
import numpy as np
import hashlib
import zipfile
from six.moves import urllib
import matplotlib

class AvgMeter(object):
    def __init__(self):
        self.value = 0.
        self.count = 0
    def __add__(self, other:float):
        self.value = (self.value * self.count + other) / (self.count + 1)
        self.count += 1
        return self
    def __call__(self):
        return self.value

def readlines(filename):
    """Read all the lines in a text file and return as a list
    """
    with open(filename, 'r') as f:
        lines = f.read().splitlines()
    return lines


def normalize_image(x):
    """Rescale image pixels to span range [0, 1]
    """
    if type(x) == np.ndarray:
        ma = float(x.max())
        mi = float(x.min())
        d = ma - mi if ma != mi else 1e5
    elif type(x) == torch.Tensor:
        ma = float(x.max().cpu().data)
        mi = float(x.min().cpu().data)
        d = ma - mi if ma != mi else 1e5
    return (x - mi) / d

def depth_to_rgb(depth_image: np.ndarray, clip_max: float = 10.0) -> np.ndarray:
    """Normalize depth image into [0, 1] and convert to grayscale rgb

    :param depth_image: Raw depth observation image from sensor output.
    :param clip_max: Max depth distance for clipping and normalization.

    :return: Clipped grayscale depth image data.
    """
    d_im = np.clip(depth_image, 0, clip_max)
    d_im /= clip_max
    rgb_d_im = (d_im * 255).astype(np.uint8)
    #cmapper = matplotlib.cm.get_cmap('Spectral')
    #value = cmapper(d_im,bytes=True).transpose(( 0 ,1 ,2,  5, 3, 4))[:, :, 0, :3] # (nxmx4)
    return d_im


def sec_to_hm(t):
    """Convert time in seconds to time in hours, minutes and seconds
    e.g. 10239 -> (2, 50, 39)
    """
    t = int(t)
    s = t % 60
    t //= 60
    m = t % 60
    t //= 60
    return t, m, s


def sec_to_hm_str(t):
    """Convert time in seconds to a nice string
    e.g. 10239 -> '02h50m39s'
    """
    h, m, s = sec_to_hm(t)
    return "{:02d}h{:02d}m{:02d}s".format(h, m, s)


def download_model_if_doesnt_exist(model_name):
    """If pretrained kitti model doesn't exist, download and unzip it
    """
    # values are tuples of (<google cloud URL>, <md5 checksum>)
    download_paths = {
        "mono_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_640x192.zip",
             "a964b8356e08a02d009609d9e3928f7c"),
        "stereo_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_640x192.zip",
             "3dfb76bcff0786e4ec07ac00f658dd07"),
        "mono+stereo_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_640x192.zip",
             "c024d69012485ed05d7eaa9617a96b81"),
        "mono_no_pt_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_no_pt_640x192.zip",
             "9c2f071e35027c895a4728358ffc913a"),
        "stereo_no_pt_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_no_pt_640x192.zip",
             "41ec2de112905f85541ac33a854742d1"),
        "mono+stereo_no_pt_640x192":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_no_pt_640x192.zip",
             "46c3b824f541d143a45c37df65fbab0a"),
        "mono_1024x320":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono_1024x320.zip",
             "0ab0766efdfeea89a0d9ea8ba90e1e63"),
        "stereo_1024x320":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/stereo_1024x320.zip",
             "afc2f2126d70cf3fdf26b550898b501a"),
        "mono+stereo_1024x320":
            ("https://storage.googleapis.com/niantic-lon-static/research/monodepth2/mono%2Bstereo_1024x320.zip",
             "cdc5fc9b23513c07d5b19235d9ef08f7"),
        }

    if not os.path.exists("models"):
        os.makedirs("models")

    model_path = os.path.join("models", model_name)

    def check_file_matches_md5(checksum, fpath):
        if not os.path.exists(fpath):
            return False
        with open(fpath, 'rb') as f:
            current_md5checksum = hashlib.md5(f.read()).hexdigest()
        return current_md5checksum == checksum

    # see if we have the model already downloaded...
    if not os.path.exists(os.path.join(model_path, "encoder.pth")):

        model_url, required_md5checksum = download_paths[model_name]

        if not check_file_matches_md5(required_md5checksum, model_path + ".zip"):
            print("-> Downloading pretrained model to {}".format(model_path + ".zip"))
            urllib.request.urlretrieve(model_url, model_path + ".zip")

        if not check_file_matches_md5(required_md5checksum, model_path + ".zip"):
            print("   Failed to download a file which matches the checksum - quitting")
            quit()

        print("   Unzipping model...")
        with zipfile.ZipFile(model_path + ".zip", 'r') as f:
            f.extractall(model_path)

        print("   Model unzipped to {}".format(model_path))

src/models/utils.py
embodied-depth/Embodied-Depth-Prediction
import collections
import timeit
import io
import cv2
import torch
import numpy as np
#import quaternion
try:
    from habitat.utils.visualizations import maps
except:
    pass
from torchvision.transforms import ToTensor
from PIL import Image  
import matplotlib
import matplotlib.cm
from matplotlib import pyplot as plt

PI = 3.1415926
D2PI = lambda x: x /  180 * PI
K = np.array([[1 , 0, 0., 0],
                           [0, 1, 0., 0],
                           [0, 0, 1, 0],
                           [0, 0, 0, 1]], dtype=np.float32)


_DEPTH_COLORMAP = plt.get_cmap('plasma', 256)  # for plotting
def colormap(inputs, normalize=True, torch_transpose=True):
    if isinstance(inputs, torch.Tensor):
        inputs = inputs.detach().cpu().numpy()

    vis = inputs
    if normalize:
        ma = float(vis.max())
        mi = float(vis.min())
        d = ma - mi if ma != mi else 1e5
        vis = (vis - mi) / d

    if vis.ndim == 4:
        vis = vis.transpose([0, 2, 3, 1])
        vis = _DEPTH_COLORMAP(vis)
        vis = vis[:, :, :, 0, :3]
        if torch_transpose:
            vis = vis.transpose(0, 3, 1, 2)
    elif vis.ndim == 3:
        vis = _DEPTH_COLORMAP(vis)
        vis = vis[:, :, :, :3]
        if torch_transpose:
            vis = vis.transpose(0, 3, 1, 2)
    elif vis.ndim == 2:
        vis = _DEPTH_COLORMAP(vis)
        vis = vis[..., :3]
        if torch_transpose:
            vis = vis.transpose(2, 0, 1)

    return vis



def colorize(value, vmin=None, vmax=None, cmap='Spectral'):
    """
    A utility function for Torch/Numpy that maps a grayscale image to a matplotlib
    colormap for use with TensorBoard image summaries.
    By default it will normalize the input value to the range 0..1 before mapping
    to a grayscale colormap.
    Arguments:
      - value: 2D Tensor of shape [height, width] or 3D Tensor of shape
        [height, width, 1].
      - vmin: the minimum value of the range used for normalization.
        (Default: value minimum)
      - vmax: the maximum value of the range used for normalization.
        (Default: value maximum)
      - cmap: a valid cmap named for use with matplotlib's `get_cmap`.
        (Default: Matplotlib default colormap)
    
    Returns a 4D uint8 tensor of shape [height, width, 4].
    """

    # normalize
    vmin = value.min() if vmin is None else vmin
    vmax = value.max() if vmax is None else vmax
    if vmin!=vmax:
        value = (value - vmin) / (vmax - vmin) # vmin..vmax
    #else:
    #    # Avoid 0-division
    #    value = value*0.
    # squeeze last dim if it exists

    value = value.squeeze()

    cmapper = matplotlib.cm.get_cmap(cmap)
    value = cmapper(value,bytes=True).transpose((2, 0, 1))[:3] # (nxmx4)
    return value

    #rgb_d_im = np.uint8(value * 255)
    #dep_ = cv2.applyColorMap(rgb_d_im, cv2.COLORMAP_MAGMA)
    #return dep_


def cam1_to_cam0(T_world_camera1, T_world_camera0):
    T_camera1_world = np.linalg.pinv(T_world_camera1)
    T_camera1_camera0 = np.matmul(T_camera1_world, T_world_camera0)
    return T_camera1_camera0


# from https://github.com/tinghuiz/SfMLearner
def compute_ate(gtruth_xyz, pred_xyz_o):
    # Make sure that the first matched frames align (no need for rotational alignment as
    # all the predicted/ground-truth snippets have been converted to use the same coordinate
    # system with the first frame of the snippet being the origin).
    offset = gtruth_xyz[0] - pred_xyz_o[0]
    pred_xyz = pred_xyz_o + offset[None, :]

    # Optimize the scaling factor
    scale = np.sum(gtruth_xyz * pred_xyz) / np.sum(pred_xyz ** 2)
    alignment_error = pred_xyz * scale - gtruth_xyz
    rmse = np.sqrt(np.sum(alignment_error ** 2)) / gtruth_xyz.shape[0]
    return rmse

class MapStat:
    def __init__(self, env):
        self.env = env
        top_down_map = maps.get_topdown_map_from_sim(
            env._sim, map_resolution=512
            ) 
        self.total = (top_down_map == 1).sum()
        self.top_down_map = maps.colorize_topdown_map(top_down_map)
        self.map = np.zeros(self.top_down_map.shape[:2])
    
    def step(self, trace):
        po = [
            maps.to_grid(
                    p[2], p[0],
                    self.map.shape,
                    sim=self.env._sim,
                )
                for p in trace
            ]
        maps.draw_path(self.top_down_map, po,  maps.MAP_SHORTEST_PATH_COLOR, thickness=2)
        for p in po:
            self.map[p] += 1

    def summary(self, path=None):
        plt.imshow(self.top_down_map)
        if path == None:
            buf = io.BytesIO()
            plt.savefig(buf, format='jpeg')
            buf.seek(0)
            image = Image.open(buf)
            image = ToTensor()(image)
            percent = (len(self.map.nonzero()[0]) / self.total)
            return image, percent
        else:
            plt.savefig(path, format='jpeg')

def vis_map(env, gt_trace, pred_trace):
    top_down_map = maps.get_topdown_map_from_sim(
            env._sim, map_resolution=1024
        )

    gt_po = [
            maps.to_grid(
                    p[2],
                    p[0],
                    (top_down_map.shape[0], top_down_map.shape[1]),
                    sim=env._sim,
                )
                for p in gt_trace
    ]

    top_down_map = maps.colorize_topdown_map(top_down_map)
    maps.draw_path(top_down_map, gt_po, (255, 0, 0) , thickness=4)

    if pred_trace is not None:
        pred_po = [
                maps.to_grid(
                        p[2],
                        p[0],
                        (top_down_map.shape[0], top_down_map.shape[1]),
                        sim=env._sim,
                    )
                    for p in pred_trace
        ]
        maps.draw_path(top_down_map, pred_po,  (0, 255,0), thickness=4)
        for p in pred_po:
            cv2.circle(top_down_map, (p[1], p[0]), radius=5, color=(0, 200, 0), thickness=3)

    for i, p in enumerate(gt_po):
        if i == 0:
            cv2.circle(top_down_map, (p[1], p[0]), radius=8, color=(200, 0, 0), thickness=5)
        else:
            cv2.circle(top_down_map, (p[1], p[0]), radius=5, color=(200, 0, 0), thickness=3)

    plt.imshow(top_down_map)

    buf = io.BytesIO()
    plt.savefig(buf, format='jpeg')
    buf.seek(0)

    image = Image.open(buf)
    image = ToTensor()(image)
    return image


def dist(p1, p2):
    '''
    Considering the different point numbers in different areas,
    we normalize the distance according to the point numbers
    p1 : (1, 3, N_1, ) 
    p2 : (1, 3, N_2, ) 
    '''
    n_1 = p1.shape[2]
    n_2 = p2.shape[2]
    if n_1 == 0 or n_2 ==0:
        return 1000
    error = (
                (p1.unsqueeze(3).repeat(1, 1, 1, n_2) - \
                p2.unsqueeze(2).repeat(1, 1, n_1, 1, )) ** 2
            ).sum(dim=1)
    
    dist1 = error.min(dim=-1).values.mean(-1)  
    dist2 = error.min(dim=-2).values.mean(-1)
    return dist1 + dist2

def backproj(depth, K, batch_size=1):
    '''
    Input:
        depth: (*, W, H) 
        K: (4, 4)
    
    '''
    #TODO: check the 3d points' scale
    height = 192
    width = 640
    
    inv_K = torch.from_numpy(np.linalg.pinv(K))
    
    meshgrid = np.meshgrid(range(width), range(height), indexing='xy')
    
    id_coords = np.stack(meshgrid, axis=0).astype(np.float32)
    id_coords = torch.from_numpy(id_coords)
                                   
    ones = torch.ones(batch_size, 1, height * width)

    pix_coords = torch.unsqueeze(torch.stack(
            [id_coords[0].view(-1), id_coords[1].view(-1)], 0), 0)
    pix_coords = pix_coords.repeat(batch_size, 1, 1)
    pix_coords = torch.cat([pix_coords, ones], 1)

    cam_points = torch.matmul(inv_K[:3, :3], pix_coords)
    cam_points = depth.view(batch_size, 1, -1) * cam_points
    cam_points = torch.cat([cam_points, ones], 1)

    return cam_points

def correct_yz(pose):
    '''correct pose in Habitat to the standard xyz 
    pose with shape (*, 4, 4)
    '''
    correct = np.array([[1., 0, 0, 0],
                        [0, -1., 0, 0],
                        [0, 0 , -1., 0],
                        [0, 0 , 0, 1.]], dtype=np.float32)
    return  np.matmul(pose, correct)

def split(_points):
    '''
    points: (1, 3, N) 
    '''
    points = _points[:, :, _points[0, 1, :] > 0]
    angles = torch.atan2(-points[..., 2, :], points[...,0, :]).squeeze()
    l_ = points[:,:, D2PI(75) > angles].clone()
    r_ = points[:,:, D2PI(105) < angles].clone()
    m = points[:,:, D2PI(75) < angles].clone()
    m_angles = torch.atan2(-m[..., 2, :], m[...,0, :]).squeeze()
    m_ = m[:,:, D2PI(105) > m_angles ].clone()

    return l_, m_, r_

class AvgMeter(object):
    def __init__(self):
        self.value = 0.
        self.count = 0
    def __add__(self, other:float):
        self.value = (self.value * self.count + other) / (self.count + 1)
        self.count += 1
        return self
    def __call__(self):
        return self.value 

    
class Timings:
    """Not thread-safe."""

    def __init__(self):
        self._means = collections.defaultdict(int)
        self._vars = collections.defaultdict(int)
        self._counts = collections.defaultdict(int)
        self.reset()

    def reset(self):
        self.last_time = timeit.default_timer()

    def time(self, name):
        """Save an update for event `name`.
        Nerd alarm: We could just store a
            collections.defaultdict(list)
        and compute means and standard deviations at the end. But thanks to the
        clever math in Sutton-Barto
        (http://www.incompleteideas.net/book/first/ebook/node19.html) and
        https://math.stackexchange.com/a/103025/5051 we can update both the
        means and the stds online. O(1) FTW!
        """
        now = timeit.default_timer()
        x = now - self.last_time
        self.last_time = now

        n = self._counts[name]

        mean = self._means[name] + (x - self._means[name]) / (n + 1)
        var = (
            n * self._vars[name] + n * (self._means[name] - mean) ** 2 + (x - mean) ** 2
        ) / (n + 1)

        self._means[name] = mean
        self._vars[name] = var
        self._counts[name] += 1

    def means(self):
        return self._means

    def vars(self):
        return self._vars

    def stds(self):
        return {k: v ** 0.5 for k, v in self._vars.items()}

    def summary(self, prefix=""):
        means = self.means()
        stds = self.stds()
        total = sum(means.values())

        result = prefix
        for k in sorted(means, key=means.get, reverse=True):
            result += f"\n    %s: %.6fs +- %.6fs (%.2f%%) " % (
                k,
                means[k],
                 stds[k],
                100 * means[k] / total,
            )
        result += "\nTotal: %.6fs" % ( total)
        return result

src/my_utils.py
embodied-depth/Embodied-Depth-Prediction
from __future__ import absolute_import, division, print_function

import os
import argparse

class Options:
    def __init__(self):
        self.parser = argparse.ArgumentParser(description="Monodepth options")

        # Embodied env setting
        self.parser.add_argument("--scene_id",
                            type=int,
                            default=4)
        self.parser.add_argument("--train_scene_id",
                            type=int,
                            help='id for offline data training',
                            default=4)
        self.parser.add_argument("--test_scene_id",
                            type=int,
                            default=4)
        self.parser.add_argument("--cfg_path",
                        type=str,
                        default='configs/default_cfg.yaml')

        # Visalization setting
        self.parser.add_argument("--video_prefix",
                                        type=str,
                                        default='')
        self.parser.add_argument("--clipmax",
                                type=int,
                                default=80 )
        
        self.parser.add_argument("--max_frame_interval",
                                   type=int,
                                   default=6)
        self.parser.add_argument("--min_frame_interval",
                                   type=int,
                                   default=1)

        self.parser.add_argument("--update_times",
                                   type=int,
                                   default=1,
                                   help='number of Data Training for one iteration data collection')

       # Dataset setting 
        self.parser.add_argument("--png",
                              help="if set, trains from raw KITTI png files (instead of jpgs)",
                              action="store_true")
        self.parser.add_argument("--video_interval",
                            type=int,
                            default=1)

        self.parser.add_argument("--input_T_dim",
                            nargs="+",
                            type=int,
                            help="time index used in consistency loss",
                            default=[0])

        
        self.parser.add_argument("--num_past_frame_input",
                                type=int,
                                help="frames to load",
                                default=0)
        
        self.parser.add_argument("--num_workers",
                              type=int,
                              help="number of dataloader workers",
                              default=12)

        # Ablation
        self.parser.add_argument("--model_type",
                            choices=['monodepth', 'scdepth', 'manydepth', 'flow2depth', 'leres'])
        self.parser.add_argument("--agent_type",
                            choices=['random', 'chamfer', 'mae', 'frontier', 'mix', 'edge', 'forward'],
                            default='random')
        self.parser.add_argument("--reset_halfway",
                                help="Whether only test in the offline mode",
                                default=False,
                                action="store_true")

        self.parser.add_argument("--full_trainset",
                                help="Whether use full trainset directly in the offline mode",
                                default=False,
                                action="store_true")

        self.parser.add_argument("--test",
                                help="Whether only test in the offline mode",
                                default=False,
                                action="store_true")

        self.parser.add_argument("--save_pred_npy",
                                help="Whether only test in the offline mode",
                                default=False,
                                action="store_true")

    def parse(self):
        subparser = self.parser.add_subparsers()

        from src.models import flow2depthOptions 
        #from src.models import monodepthOptions 
        #from src.models import manydepthOptions 
        #from src.models import scdepthOptions 
        #from src.models import leresOptions 

        flow2depthOptions(subparser.add_parser('flow2depth'))
        #monodepthOptions(subparser.add_parser('monodepth'))
        #manydepthOptions(subparser.add_parser('manydepth'))
        #scdepthOptions(subparser.add_parser('scdepth'))
        #leresOptions(subparser.add_parser('leres'))

        self.options = self.parser.parse_args()
        return self.options

src/options.py
embodied-depth/Embodied-Depth-Prediction
from src.ros.roscontrol import Controller
from src.ros.nplistener import DataListener

src/ros/__init__.py
embodied-depth/Embodied-Depth-Prediction
#!/usr/bin/env python

import time
import rospy
from rospy_tutorials.msg import Floats
from rospy.numpy_msg import numpy_msg
import numpy as np
import skimage.transform

from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError

cvb=CvBridge()

def callback(data):
    try:
        array = cvb.imgmsg_to_cv2(data)
        print(array.shape)
        #print(rospy.get_name(), "I heard %s"%(array.shape))
    except CvBridgeError as e:
        print(e)

def listener():
    rospy.init_node('listener')
    rospy.Subscriber("floats", Image, callback)
    rospy.spin()


class DataListener:
    def __init__(self):
        self.rgb = []
        self.dep = []
        self.traj = None
        self.history_idx = 0

        rospy.Subscriber('data/rgb', Image, self.grab_rgb)
        rospy.Subscriber('data/depth', Image, self.grab_dep)
        rospy.Subscriber('data/traj', Image, self.grab_traj)

    def output_data(self):
        while self.traj is None:
            time.sleep(1)
            print('Waiting for the first batch data')

        out_len = min(len(self.rgb), len(self.dep), self.traj.shape[0])
        while out_len - self.history_idx < 5:
            time.sleep(1)
            print('Waiting for the batch data')
            out_len = min(len(self.rgb), len(self.dep), self.traj.shape[0])

        out_ = []
        for i in range(self.history_idx, out_len):
            out = {
                'color_sensor': self.rgb[i],
                'depth_sensor': self.dep[i] / 1000.,
                'Ext': self.traj[i]
            }
            out_.append(out)
        self.history_idx = out_len

        return out_, False


    def grab_rgb(self, data):
        try:
            array = cvb.imgmsg_to_cv2(data)
            self.rgb.append(array)
        except CvBridgeError as e:
            print(e)

    def grab_dep(self, data):
        try:
            array = cvb.imgmsg_to_cv2(data)
            depth = skimage.transform.resize(
                    array, (192, 640), order=0, preserve_range=True, mode='constant')
            self.dep.append(depth)
        except CvBridgeError as e:
            print(e)

    def grab_traj(self, data):
        try:
            array = cvb.imgmsg_to_cv2(data)
            if self.traj is  None:
                self.traj = array
            else:
                self.traj = np.concatenate([self.traj, array], axis=0)
        except CvBridgeError as e:
            print(e)

if __name__ == '__main__':
    listener()

src/ros/nplistener.py
embodied-depth/Embodied-Depth-Prediction
import os
import glob
import rospy
import numpy as np
import cv2
import sys
from nav_msgs.msg import Odometry
from sensor_msgs.msg import Image, CameraInfo
from cv_bridge import CvBridge, CvBridgeError
import orbslam2
import time
from collections import deque
from scipy.spatial.transform import Rotation as R

cvb=CvBridge()


input_length = 0
def slaminput(system):
    global input_length
    #if len(system.depth_buffer) > input_length and len(system.depth_buffer) == len(system.rgb_buffer):
    if len(system.depth_buffer) > input_length and len(system.rgb_buffer) > input_length:
        system.slam.process_image_rgbd(system.rgb_buffer[input_length],
                system.depth_buffer[input_length] / 1000.,
                len(system.rgb_buffer))

        input_length += 1

        if (str(system.slam.get_tracking_state()) == 'LOST'):
            system.slam_lost_count += 1
            if system.slam_lost_count > 60:
                #system.print_traj()
                print("*" * 20)
                print('RESET cuz lost too long')
                print("*" * 20)
                system.write_to_disk()
                system.reset()
                input_length = 0
                time.sleep(1)
    #if (str(system.slam.get_tracking_state()) == 'LOST'):
    #    system.slam_lost_count += 1
    #    if system.slam_lost_count > 60:
    #        #system.print_traj()
    #        print("*" * 20)
    #        print('RESET cuz lost too long')
    #        print("*" * 20)
    #        system.write_to_disk()
    #        system.reset()
    #        input_length = 0
    #        time.sleep(5)



def imnormalize(image):
    """
    Normalize a list of sample image data in the range of 0 to 1
    : image:image data.
    : return: Numpy array of normalize data
    """
    xmax = np.max(image)
    xmin = 0
    a = 0
    b = 255

    return ((np.array(image,dtype=np.float32) - xmin) * (b - a)) / (xmax - xmin)


class ImageListener:
    def __init__(self, vocab_path, settings_path, disk_path, send_data_flag=True):
        self.disk_path = disk_path
        if not os.path.exists(self.disk_path):
            os.makedirs(self.disk_path)
        self.rgb_buffer = []
        self.depth_buffer = []
        self.odom = []
        self.i_rgb=0
        self.i_depth=0
        self.i_odom=0
        self.rgb_mem=None
        self.depth_mem=None
        self.first_flag_depth=True
        self.first_flag_rgb=True

        self.last_trans_index = 0
        self.pause_receive_flag = False
        self.static_flag = False
        self.last_static_flag = False
        self.flag_smoother = deque(maxlen=80)



        self.traj_num_idx = 0
        self.endflag = False

        self.send_data_flag = send_data_flag

        ## Receiver from realseanse
        rospy.init_node("grabrgb")
        rospy.loginfo("Running RGB Grabber")
        rospy.Subscriber("/camera/color/image_raw",Image,self.grabrgb)
        #rospy.Subscriber("/camera/depth/image_rect_raw",Image,self.checkdepth)
        rospy.Subscriber("/camera/aligned_depth_to_color/image_raw",Image,self.checkdepth)
        #rospy.spin()


        ## Open or close the RGBD save by the odometry
        rospy.Subscriber('/odometry/filtered', Odometry, self.listen_odom)

        ## Slam setting
        self.slam = orbslam2.System(vocab_path, settings_path, orbslam2.Sensor.RGBD)
        self.slam.set_use_viewer(True)
        self.slam.initialize()
        self.slam_lost_count = 0

        ## Send to pytorch module
        self.pubimg = rospy.Publisher('data/rgb', Image,queue_size=10)
        self.pubdepth = rospy.Publisher('data/depth', Image,queue_size=10)
        self.pubtraj = rospy.Publisher('data/traj', Image,queue_size=10)
        #rospy.init_node('talker', anonymous=True)
        self.rate = rospy.Rate(20) # 10hz

        ## 

    def reset(self):
        self.traj_num_idx += 1

        self.rgb_buffer = []
        self.depth_buffer = []
        self.odom = []
        self.i_rgb=0
        self.i_depth=0
        self.i_odom=0
        self.rgb_mem=None
        self.depth_mem=None
        self.first_flag_depth=True
        self.first_flag_rgb=True

        self.static_flag = False
        self.last_static_flag = False
        self.slam.reset()
        self.slam_lost_count = 0

    def listen_odom(self, msg):
        self.last_static_flag = self.static_flag
        self.static_flag = (abs(msg.twist.twist.linear.x) < 8e-3)  and \
                (abs(msg.twist.twist.angular.z) < 8e-3 )
        self.flag_smoother.append(self.static_flag)

        if (self.i_odom <= self.i_depth) and (not self.static_flag) and (not self.pause_receive_flag):
            rmatrix = R.from_quat([msg.pose.pose.orientation.x,
                                    msg.pose.pose.orientation.y,
                                    msg.pose.pose.orientation.z,
                                    msg.pose.pose.orientation.w]).as_matrix()
            t_vec = np.array([[-msg.pose.pose.position.x],
                              [-msg.pose.pose.position.y],
                              [ msg.pose.pose.position.z]])
            zero = np.array([[0, 0, 0, 1]])
            matrix = np.concatenate([rmatrix, t_vec], axis=1)
            matrix = np.concatenate([matrix, zero], axis=0)

            self.i_odom += 1
            self.odom.append(matrix)


        if (self.last_static_flag != self.static_flag):
            print('Static flag changes from {} to {} '.format(self.last_static_flag, self.static_flag))

        if self.send_data_flag and np.array(self.flag_smoother).all() and (len(self.rgb_buffer) - self.last_trans_index) > 10:
            # If remain still for a little while, start to transmit the data
            self.send_data()


    def grabrgb(self, msg):
        #print('flag rgb :', self.pause_receive_flag)
        if (self.i_rgb <= self.i_depth) and (not self.static_flag) and (not self.pause_receive_flag):
            try:
                cv_image = cvb.imgmsg_to_cv2(msg,"bgr8")
            except CvBridgeError as e:
                print(e)

            #print(msg.step)
            image_normal= np.array(cv_image)
            if self.first_flag_rgb == True :
                self.rgb_mem = np.copy(image_normal)
                self.first_flag_rgb=False
            elif np.array_equal(self.rgb_mem,image_normal) :
                return None
            else :
                self.rgb_mem = np.copy(image_normal)

            self.i_rgb+=1
            self.rgb_buffer.append(image_normal)


            #if len(self.depth_buffer) > 0:
            #    self.slam.process_image_rgbd(self.rgb_buffer[-1],
            #            self.depth_buffer[-1] / 1000.,
            #            len(self.rgb_buffer))

            #    if (str(self.slam.get_tracking_state()) == 'LOST'):
            #        self.slam_lost_count += 1
            #        if self.slam_lost_count > 60:
            #            self.print_traj()
            #            print('RESET cuz lost too long')
            #            self.slam.reset()
            #            self.slam_lost_count = 0


    def checkdepth(self, msg):
        if (not self.static_flag) and (not self.pause_receive_flag):
            try:
                cv_image = cvb.imgmsg_to_cv2(msg,msg.encoding)
            except CvBridgeError as e:
                print(e)
            
            image_normal= np.array(imnormalize(cv_image),dtype=np.uint8)
            numpy_image= np.array(cv_image)
            if self.first_flag_depth == True:
                self.depth_mem = np.copy(numpy_image)
                self.first_flag_depth=False
            if (self.depth_mem==numpy_image).all() :
                return
            else:
                self.depth_mem = np.copy(numpy_image)
                
            self.i_depth+=1
            self.depth_buffer.append(numpy_image)

    def get_recent_traj(self):
        out = []
        trajs = self.slam.get_trajectory_points()

        print('Len of trajs: ', len(trajs))
        print('Len of rgb: ', len(self.rgb_buffer))

        for i in range(self.last_trans_index, len(trajs)):
            traj = trajs[i]
            stamp, r00, r01, r02, t0, r10, r11, r12, t1, r20, r21, r22, t2 = traj
            ext_matrix = np.array([[r00, r01, r02, t0],
                                   [r10, r11, r12, t1],
                                   [r20, r21, r22, t2],
                                   [0.,  0.,  0.,  1]])
            out.append(ext_matrix)

        out = np.stack(out)
        return out

    def get_all_traj(self):
        out = []
        trajs = self.slam.get_trajectory_points()

        for i in range(0, len(trajs)):
            traj = trajs[i]
            stamp, r00, r01, r02, t0, r10, r11, r12, t1, r20, r21, r22, t2 = traj
            ext_matrix = np.array([[r00, r01, r02, t0],
                                   [r10, r11, r12, t1],
                                   [r20, r21, r22, t2],
                                   [0.,  0.,  0.,  1]])
            out.append(ext_matrix)

        out = np.stack(out)
        return out

    def end(self):
        self.slam.shutdown()

        out_file_name = os.path.join( self.disk_path, 'train1_files.txt')
        shootings = glob.glob( os.path.join(self.disk_path, 'traj*'), recursive=False)

        for shoot in shootings:
            trajs = np.load(os.path.join(shoot, 'poses.npy'))
            num = trajs.shape[0]


            with open(out_file_name, 'a') as f:
                for i in range(20, num-20):
                    f.write(shoot + ' ' + str(i) + ' l\n')

        self.endflag = True

    def send_data(self):
        self.pause_receive_flag = True
        self.input_rest_rgbd2slam()
        self.align_rgb_d_len()


        if len(self.rgb_buffer) - self.last_trans_index > 0:
            traj = self.get_recent_traj()

            send_len = min(traj.shape[0] + self.last_trans_index, len(self.rgb_buffer))
            traj = traj[: send_len - self.last_trans_index]

            #for i in range(self.last_trans_index, len(self.rgb_buffer)):
            for i in range(self.last_trans_index, send_len):
                img = self.rgb_buffer[i]
                dep = self.depth_buffer[i]

                self.pubimg.publish(cvb.cv2_to_imgmsg(img))
                self.pubdepth.publish(cvb.cv2_to_imgmsg(dep))
                self.rate.sleep()


            self.last_trans_index = len(self.rgb_buffer)
            self.pubtraj.publish(cvb.cv2_to_imgmsg(traj))
            self.rate.sleep()

            print('Send Over')

        self.pause_receive_flag = False

    def write_to_disk(self):
        self.input_rest_rgbd2slam()
        self.align_rgb_d_len()
        path = os.path.join(self.disk_path, 'traj{}'.format(self.traj_num_idx))
        if not os.path.exists(path):
            os.makedirs(path)
        if not os.path.exists(path + '/color'):
            os.makedirs(path+ '/color')
            os.makedirs(path+ '/depth')
            os.makedirs(path+ '/depthmap')

        print('*'*20)
        print('Save into traj{}'.format(self.traj_num_idx) )
        print('Image Length: ', len(self.rgb_buffer))
        print('Depth Length: ', len(self.depth_buffer))
        traj = self.get_all_traj()
        print('Traj Length: ', traj.shape[0])
        np.save(path + '/poses.npy', traj)

        odoms = np.stack(self.odom)
        np.save(path + '/odom.npy', odoms)

        for i, (im, d) in enumerate(zip(self.rgb_buffer, self.depth_buffer)):
            #print('RGB  shape,', im.shape)
            #print('DE, shape, ', d.shape)
            cv2.imwrite(path + '/color/frame{}.jpg'.format(i), im)
            np.save(path + '/depth/dframe{}.npy'.format(i), d)
            d_ = imnormalize(d)
            cv2.imwrite(path + '/depthmap/dframe{}.jpg'.format(i), d_)

        self.pause_receive_flag = False

    # Only used for online test
    def print_traj(self):
        out = []
        trajs = self.slam.get_trajectory_points()
        i = 0
        for traj in trajs:
            stamp, r00, r01, r02, t0, r10, r11, r12, t1, r20, r21, r22, t2 = traj
            ext_matrix = np.array([[r00, r01, r02, t0],
                                   [r10, r11, r12, t1],
                                   [r20, r21, r22, t2],
                                   [0.,  0.,  0.,  1]])
            i += 1
            if i % 20 == 0:
                print('Time: {}, matrix: {}'.format(stamp, ext_matrix))

    def input_rest_rgbd2slam(self):
        self.pause_receive_flag = True
        self.align_rgb_d_len()
        global input_length
        while input_length  < len(self.rgb_buffer):
            slaminput(self)

        time.sleep(2)

    def align_rgb_d_len(self):
        while len(self.rgb_buffer) > len(self.depth_buffer):
            self.rgb_buffer.pop()

        while len(self.rgb_buffer) < len(self.depth_buffer):
            self.depth_buffer.pop()

        while len(self.odom) > len(self.rgb_buffer):
            self.odom.pop()
        

def write(img, depth, path):
    print('Image Length: ', len(img))
    print('Depth Length: ', len(depth))
    for i, (im, d) in enumerate(zip(img, depth)):
        #print('RGB  shape,', im.shape)
        #print('DE, shape, ', d.shape)
        cv2.imwrite(path + '/frame{}.jpg'.format(i), im)
        #np.save(path + '/depth/dframe{}.npy'.format(i), d)
        d_ = imnormalize(d)
        cv2.imwrite(path + '/dframe{}.jpg'.format(i), d_)


def test():
    if len(sys.argv) != 4:
        print('Usage: python ./roscam.py path_to_vocabulary path_to_settings  save_rgbd_dirname' )
    lis = ImageListener(
                        vocab_path=sys.argv[1],
                        settings_path=sys.argv[2],
                        disk_path='/home/yilundu/qingquantmp/tmp/' + str(sys.argv[3])
                        )


    start_time = time.time()
    while time.time() - start_time < 40:
        slaminput(lis)

    #lis.print_traj()
    lis.write_to_disk()
    #write(lis.rgb_buffer, lis.depth_buffer, '/home/yilundu/qingquantmp/tmp/testfig')
    print('end')
    lis.end()

def main():
    if len(sys.argv) < 4:
        print('Usage: python ./roscam.py path_to_vocabulary path_to_settings  save_rgbd_dirname send_data_flag' )

    if len(sys.argv) >4:
        send_data_flag = bool(sys.argv[4])
    else:
        send_data_flag = False

    lis = ImageListener(
                        vocab_path=sys.argv[1],
                        settings_path=sys.argv[2],
                        disk_path='/home/yilundu/qingquantmp/tmp/' + str(sys.argv[3]),
                        send_data_flag=send_data_flag
                        )


    start_time = time.time()
    while time.time() - start_time < 120:
        slaminput(lis)

    lis.write_to_disk()
    lis.end()


if __name__ == '__main__':
    main()

src/ros/roscam.py
embodied-depth/Embodied-Depth-Prediction
import numpy as np
import rospy
from geometry_msgs.msg import Twist, Vector3
from sensor_msgs.msg import LaserScan
from random import randint


class CollisionError(Exception):
    def __init__(self, error_info=None):
        super().__init__(self)
        self.error_info = error_info

    def __str__(self):
        return self.error_info


class Controller:
    def __init__(self):
        rospy.init_node('myagent')
        self.rate = rospy.Rate(20)
        
        self.ACTION_SPACE = {
                'move_forward': Twist(
                    Vector3(-0.1, 0, 0),
                    Vector3(0, 0, 0),            
                    ),
                'back': Twist(
                    Vector3(0.1, 0, 0),
                    Vector3(0, 0, 0),            
                    ),
                'turn_left': Twist(
                    Vector3(0., 0, 0),
                    Vector3(0, 0, 0.1),            
                    ),
                'turn_right': Twist(
                    Vector3(0., 0, 0),
                    Vector3(0, 0, -0.1),            
                    ),
        }

        rospy.Subscriber("/scan", LaserScan, self.lasercallback)
        self.pub = rospy.Publisher('/jackal_velocity_controller/cmd_vel', Twist, queue_size=10)
        self.collision_flag = False
        
    def act(self ,policy:list):
        if not self.collision_flag:
            for a in policy:
                if not self.collision_flag:
                    self._act(a)
                else:
                    break
        if self.collision_flag:
            for a in ['back']*5:
                self._act(a)

    def _act(self, action:str):
        rospy.loginfo(action)
        if not self.collision_flag:
            steps = 10 if action == 'move_forward' else 15
            for i in range(steps): # enlong the move time
                self.pub.publish(self.ACTION_SPACE[action])
                self.rate.sleep()
        else: 
            self.pub.publish(self.ACTION_SPACE['back'])
    
    def lasercallback(self, msg):
        arr =  np.array(msg.ranges[100:600]) # 120 degree
        arr = arr[arr > 0.03] # filter numbers below the min range
        self.collision_flag = (arr < 0.25).any()

def test_control():
    pub = Controller()
    #while not rospy.is_shutdown():
    for _ in range(10):
        action = ['move_forward'] * 5
        pub.act(action)

def lasercallback(msg):
    arr =  np.array(msg.ranges[180:540]) # 90 degree
    arr = arr[arr > 0.03] # filter numbers below the min range
    return (arr < 0.2).any()
    
def baselidar_test():
    rospy.init_node('scan_values')
    rospy.Subscriber("/scan",LaserScan, lasercallback)
    rospy.spin()

if __name__ == '__main__':
    #baselidar_test()
    test_control()
src/ros/roscontrol.py
embodied-depth/Embodied-Depth-Prediction
# Testing the orbslam2 python binding
import os
import numpy as np
import cv2
import sys
import orbslam2
import time
from tqdm import tqdm
from pathlib import Path

def main(vocab_path, settings_path, sequence_path, association_path):

    rgb_filenames, depth_filenames, timestamps = load_images(association_path)
    num_images = len(timestamps)

    slam = orbslam2.System(vocab_path, settings_path, orbslam2.Sensor.RGBD)
    slam.set_use_viewer(True)
    slam.initialize()

    times_track = [0 for _ in range(num_images)]
    print('-----')
    print('Start processing sequence ...')
    print('Images in the sequence: {0}'.format(num_images))

    for idx in tqdm(range(num_images)):
        rgb_image = cv2.imread(os.path.join(sequence_path, rgb_filenames[idx]), cv2.IMREAD_UNCHANGED)
        #depth_image = cv2.imread(os.path.join(sequence_path, depth_filenames[idx]), cv2.IMREAD_UNCHANGED)
        depth_image = np.load(os.path.join(sequence_path, depth_filenames[idx])) / 1000.
        tframe = timestamps[idx]

        if rgb_image is None:
            print("failed to load image at {0}".format(rgb_filenames[idx]))
            return 1
        if depth_image is None:
            print("failed to depth at {0}".format(depth_filenames[idx]))
            return 1

        t1 = time.time()
        slam.process_image_rgbd(rgb_image, depth_image, tframe)
        t2 = time.time()

        ttrack = t2 - t1
        times_track[idx] = ttrack

        t = 0
        if idx < num_images - 1:
            t = timestamps[idx + 1] - tframe
        elif idx > 0:
            t = tframe - timestamps[idx - 1]

        #if ttrack < t:
        #    time.sleep(t - ttrack)

    save_path = Path(association_path).parent.absolute()
    save_trajectory(slam.get_trajectory_points(), os.path.join(save_path, 'new_trajectory.txt'))

    slam.shutdown()

    times_track = sorted(times_track)
    total_time = sum(times_track)
    print('-----')
    print('median tracking time: {0}'.format(times_track[num_images // 2]))
    print('mean tracking time: {0}'.format(total_time / num_images))

    return 0


def load_images(path_to_association):
    rgb_filenames = []
    depth_filenames = []
    timestamps = []
    with open(path_to_association) as times_file:
        for line in times_file:
            if len(line) > 0 and not line.startswith('#'):
                t, rgb, _, depth = line.rstrip().split(' ')[0:4]
                rgb_filenames.append(rgb)
                depth_filenames.append(depth)
                timestamps.append(float(t))
    return rgb_filenames, depth_filenames, timestamps


def save_trajectory(trajectory, filename):
    out = []
    for traj in trajectory:
        stamp, r00, r01, r02, t0, r10, r11, r12, t1, r20, r21, r22, t2 = traj
        ext_matrix = np.array([[r00, r01, r02, t0],
                               [r10, r11, r12, t1],
                               [r20, r21, r22, t2],
                               [0.,  0.,  0.,  1]])
        out.append(ext_matrix)

    out = np.stack(out)
    print(out.shape)
    np.save('poses0.npy' , out)


if __name__ == '__main__':
    if len(sys.argv) != 5:
        print('Usage: ./orbslam_rgbd_tum path_to_vocabulary path_to_settings path_to_sequence path_to_association')
    main(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])


src/ros/slamoffline.py
embodied-depth/Embodied-Depth-Prediction