This repository contains the code for Semantically-enhanced Variational Autoencoders for Deep Collision Prediction
from setuptools import find_packages
from distutils.core import setup

setup(
    name='sevae',
    version='0.0.1',
    author='Mihir Kulkarni',
    license="BSD-3-Clause",
    packages=find_packages(),
    author_email='mihir.kulkarni@ntnu.no',
    description='Package for training and evaluating Semantically-enhanced VAEs for collision prediction',
    install_requires=['torch>=1.13',
                      'numpy']
)
setup.py
ntnu-arl/sevae
from scipy.fft import fft2, ifft2
import numpy as np

class FFT:
    def __init__(self):
        self.image = None
        self.fft = None
        self.reconstructed_image = None
    
    def get_image_reconstruction_with_compressed_dimensions(self, image, latent_dims=128):
        pixels = image.shape[0]*image.shape[1]
        self.fft = fft2(image)

        self.thresh_fft = self.fft.copy()
        self.fft_magnitude = np.abs(self.fft).flatten()
        self.fft_magnitude.sort()
        threshold = self.fft_magnitude[-latent_dims]
        self.thresh_fft[np.abs(self.fft) < threshold] = 0
        self.reconstructed_image = ifft2(self.thresh_fft).real # Check if real is to be used or abs
        return self.fft, self.reconstructed_image
    
    def forward(self, image, latent_dims=128):
        return self.get_image_reconstruction_with_compressed_dimensions(image, latent_dims)[1]
sevae/baselines/fft.py
ntnu-arl/sevae
import numpy as np
import pywt

class WaveletTransforms:
    def __init__(self, name='db1', level=7):
        self.image = None
        self.coeffs = None
        self.reconstructed_image = None
        self.name = name
        self.level = level

    def get_wavelet_coefficients(self, image, latent_dims=128):
        self.image = image
        self.coeffs = pywt.wavedec2(image, wavelet=self.name, level=self.level)
        # select top latent_dims coefficients after sorting
        coeff_array, coeff_slices = pywt.coeffs_to_array(self.coeffs)
        Csort = np.sort(np.abs(coeff_array.reshape(-1)))
        threshold = Csort[-latent_dims]
        coeff_array[np.abs(coeff_array) < threshold] = 0
        coeffs_filt_array = pywt.array_to_coeffs(coeff_array, coeff_slices, output_format='wavedec2')
        self.reconstructed_image = pywt.waverec2(coeffs_filt_array, wavelet=self.name)
        return coeff_array, self.reconstructed_image

    def forward_with_latent_dims(self, image, latent_dims=128):
        return self.get_wavelet_coefficients(image, latent_dims)
    
    def get_reconstruction_with_latent_dims(self, latent_space, slices):
        coeffs_filt_array = pywt.array_to_coeffs(latent_space, slices, output_format='wavedec2')
        reconstructed_image = pywt.waverec2(coeffs_filt_array, wavelet=self.name)
        return reconstructed_image
sevae/baselines/wavelets.py
ntnu-arl/sevae
import torch
import tensorflow as tf
import numpy as np
import os
import sys
from math import ceil

sys.path.append('.')
# from data_processing.utilities import get_files_ending_with



# TFRecord is used here to read from a serialized dataset. 

# There can really be a better way to do this, but need to find out what and how. Explore: using pickle to store the data.
# It is usually very painful to have to use both tensorflow and pytorch in the same project. Better to stick to one.

IMAGES_PER_TF_RECORD = 1000

class DepthSemanticImageDataset(torch.utils.data.IterableDataset):
    """DepthSemanticImageDataset wraps a torch.nn.IterableDataset around a tf.data.TFRecordDataset. 

    This allows the the depth images to fit in memory as the TFRecordDataset retrieves depth images sequencially during training.
    """

    def __init__(self, tfrecord_folder, batch_size=32, shuffle=False, one_tfrecord=False, test_dataset=False):
        super(DepthSemanticImageDataset).__init__()
        self.tfrecord_folder = tfrecord_folder
        self.dataset, self.data_len = self.load_tfrecords(
            is_shuffle_and_repeat=shuffle, batch_size=batch_size, one_tfrecord=one_tfrecord, test_dataset=test_dataset)

    def read_tfrecord(self, serialized_example):
        print("Reading dataset")
        feature_description_v2 = {
            'depth_image_raw': tf.io.FixedLenFeature([], tf.string),
            'depth_image_filtered': tf.io.FixedLenFeature([], tf.string),
            'semantic_image_raw': tf.io.FixedLenFeature([], tf.string),
            'num_envs': tf.io.FixedLenFeature([], tf.int64),
            'num_images_per_env': tf.io.FixedLenFeature([], tf.int64),}
        
        feature_description_old = {
            'depth_image_raw': tf.io.FixedLenFeature([], tf.string),
            'semantic_image_raw': tf.io.FixedLenFeature([], tf.string),
            'num_envs': tf.io.FixedLenFeature([], tf.int64),
            'num_images_per_env': tf.io.FixedLenFeature([], tf.int64),}
        try:
            example = tf.io.parse_single_example(
                serialized_example, feature_description_v2)
        except:
            example = tf.io.parse_single_example(
                serialized_example, feature_description_old)
        print("parsing depth image")
        parsed_depth_tensor = tf.cast(tf.io.parse_tensor(
            example['depth_image_raw'], out_type=np.float), tf.float32) / 1000.0
        print("parsing depth image filtered")
        parsed_depth_tensor_filtered = tf.cast(tf.io.parse_tensor(
            example['depth_image_filtered'], out_type=np.float), tf.float32) / 1000.0
        print("parsing semantic image")
        parsed_semantic_tensor = tf.cast(tf.io.parse_tensor(
            example['semantic_image_raw'], out_type=np.float), tf.int64)
        # If semantic information is not aviailable, set it to a matrix of zeros
        if parsed_semantic_tensor.shape[0] == 1:
            parsed_semantic_tensor = tf.zeros_like(parsed_depth_tensor)
        print("parsing num envs")
        parsed_num_envs = example['num_envs']
        print("parsing images per env")
        parsed_num_images_per_env = example['num_images_per_env']
        return parsed_num_envs, parsed_num_images_per_env, parsed_depth_tensor, parsed_depth_tensor_filtered, parsed_semantic_tensor

    def load_tfrecords(self, is_shuffle_and_repeat=True, shuffle_buffer_size=5000, prefetch_buffer_size_multiplier=2, batch_size=32, one_tfrecord=False, test_dataset=False):
        print('Loading tfrecords from folder ', self.tfrecord_folder)
        assert os.path.exists(self.tfrecord_folder), "TF Record folder path: {} does not exist".format(self.tfrecord_folder)
        assert os.path.isdir(self.tfrecord_folder), "TF Record folder path: {} is not a directory".format(self.tfrecord_folder)

        self.tfrecord_fnames = [os.path.join(self.tfrecord_folder, f) for f in os.listdir(self.tfrecord_folder) if f.endswith('.tfrecord')]
        
        print("Found {} tfrecords".format(len(self.tfrecord_fnames)))
        print("TFRecord names: ", self.tfrecord_fnames)
        
        assert len(self.tfrecord_fnames) > 0
        if is_shuffle_and_repeat:
            np.random.shuffle(self.tfrecord_fnames)
        else:
            # 176 tfrecords for train, 20 for test
            self.tfrecord_fnames = sorted(self.tfrecord_fnames)

        if one_tfrecord:
            self.tfrecord_fnames = self.tfrecord_fnames[:1]
            print(self.tfrecord_fnames)

        dataset = tf.data.TFRecordDataset(self.tfrecord_fnames)
        dataset = dataset.map(self.read_tfrecord,
                              num_parallel_calls=tf.data.experimental.AUTOTUNE)
        if is_shuffle_and_repeat:
            dataset = dataset.shuffle(
                buffer_size=shuffle_buffer_size, reshuffle_each_iteration=True)
        dataset = dataset.batch(batch_size)
        dataset = dataset.prefetch(
            buffer_size=prefetch_buffer_size_multiplier * batch_size)

        print('Iterating length... ', end="\t")
        return dataset, IMAGES_PER_TF_RECORD *len(self.tfrecord_fnames)

    def __iter__(self):
        print("gotcha")
        return self.dataset.__iter__()

    def __len__(self):
        return self.data_len


def collate_batch(batch):
    """Used as the 'collate_fn' when creating a torch.nn.DataLoader."""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    num_envs = torch.Tensor(np.array(batch[0][0]))
    num_images_per_env = torch.Tensor(np.array(batch[0][1]))
    depth_images = torch.Tensor(np.array(batch[0][2]))
    filtered_images = torch.Tensor(np.array(batch[0][3]))
    semantic_images = torch.Tensor(np.array(batch[0][4]))
    

    return num_envs, num_images_per_env, depth_images, filtered_images, semantic_images

sevae/datasets/depth_semantic_image_dataset.py
ntnu-arl/sevae
import time
import torch

import sevae
from sevae.networks.VAE.vae import VAE

# collision_repreesntation path
COLL_REP_PATH = sevae.__path__[0]
# weights path
# weights/seVAE_weights/dragvoll_dataset_test_6_LD_128_epoch_39.pth
# WEIGHTS_PATH = COLL_REP_PATH + "/weights/seVAE_weights/seVAE_LD_128.pth"

WEIGHTS_PATH = COLL_REP_PATH + "/weights/seVAE_weights/seVAE_evo_LD_128.pth"


class VAENetworkInterface():
    def __init__(self, latent_space_dims=128, device="cuda:0"):
        self.latent_space = latent_space_dims
        self.device = device

        # Load networks in constructor
        try:
            self.vae = VAE(input_dim=1, latent_dim=self.latent_space)
            dict = torch.load(WEIGHTS_PATH)
            self.vae.load_state_dict(dict, strict=True)
            self.vae = self.vae.to(self.device)
        except:
            print("Could not load networks")
            raise Exception("Could not load networks")
        print("Loading: " + WEIGHTS_PATH)
        print("VAE MODEL LOADED TO DEVICE: ", self.device)

    def forward(self, image_numpy, calculate_reconstruction=False):
        self.start_time = time.time()
        with torch.no_grad():
            torch_image = torch.from_numpy(image_numpy).float().to(self.device)
            torch_image = torch.clamp(torch_image, 0.0, 1.0)
            z_sampled, means, *_ = self.vae.encode(torch_image.view(1, 1, 270, 480))
            reconstructed_image = 0
            if calculate_reconstruction:
                reconstructed_image = self.vae.decode(means).cpu().numpy()
        self.end_time = time.time()
        return means.cpu().numpy(), reconstructed_image, 1000*(self.end_time - self.start_time)
    
    def get_compute_time(self):
        return (self.end_time - self.start_time)*1000

sevae/inference/scripts/VAENetworkInterface.py
ntnu-arl/sevae
import rospy
from sensor_msgs.msg import Image
from std_msgs.msg import Float32MultiArray


import numpy as np
import torch


from sevae.inference.scripts.VAENetworkInterface import VAENetworkInterface

import argparse

# Use argparser to distinguish between simulation and real world
parser = argparse.ArgumentParser()
parser.add_argument('--sim', default="False")
parser.add_argument('--show_cv', default="False")
parser.add_argument('--sim_use_flightmare', default="False")
args = parser.parse_args()

LATENT_SPACE = 128
IMAGE_WIDTH = 480
IMAGE_HEIGHT = 270
IMAGE_DIMENSIONS = IMAGE_WIDTH*IMAGE_HEIGHT

IMAGE_MAX_DEPTH = 10.0 #7.50

LATENT_SPACE_TOPIC_NAME = "/latent_space"

if args.sim == "False":
    IS_SIM = False
    IMAGE_TOPIC = "/d455/depth/image_rect_raw"
    USE_CV_BRIDGE = False
    CALCULATE_RECONSTRUCTION = True

else:
    IS_SIM = True
    print('args.sim_use_flightmare:', args.sim_use_flightmare)
    if args.sim_use_flightmare == "True":
        IMAGE_TOPIC = "/delta/agile_autonomy/unity_depth" # '/delta/agile_autonomy/sgm_depth'
    else: # Gazebo
        IMAGE_TOPIC = "/delta/rgbd/camera_depth/depth"
    print('IMAGE_TOPIC:', IMAGE_TOPIC)
    USE_CV_BRIDGE = True
    CALCULATE_RECONSTRUCTION = True
from cv_bridge import CvBridge, CvBridgeError


class VAENavInterface():
    def __init__(self):

        self.obs_tensor = torch.zeros(
            IMAGE_DIMENSIONS, device="cuda:0")
        self.target_x = 0
        self.target_y = 0
        self.target_z = 0
        self.skip_next = False
        self.net_interface = VAENetworkInterface(LATENT_SPACE, "cuda:0")
        if IS_SIM == True:
            self.bridge = CvBridge()
        # Publish filtered image
        self.decode_img_publisher = rospy.Publisher('/decoded_image', Image, queue_size=1)
        # Subscribe to image
        self.image_sub = rospy.Subscriber(
            IMAGE_TOPIC, Image, self.image_callback, queue_size=1)
        print("Subscribed to {}".format(IMAGE_TOPIC))
        
        # Publish latent space as Float32MultiArray
        self.latent_space_publisher = rospy.Publisher(LATENT_SPACE_TOPIC_NAME, Float32MultiArray, queue_size=2)
        print("Publishing latent space to {}", LATENT_SPACE_TOPIC_NAME)


    def image_callback(self, data):
        if self.skip_next:
            self.skip_next = False
            return

        # # convert from ROS image to torch image
        if IS_SIM:
            try:
                cv_image = self.bridge.imgmsg_to_cv2(data, "passthrough")
            except CvBridgeError as e:
                print(e)
            cv_image = np.copy(cv_image)
            cv_image = cv_image.reshape((data.height,data.width))
            if args.sim_use_flightmare == "True":
                cv_image = cv_image[5:275,:] # flightmare crashes if we use the resolution of 480x270, hence we use 480x280! 
                cv_image = cv_image.astype('float32') * 0.001 # convert pixel value to meter
            cv_image[np.isnan(cv_image)] = IMAGE_MAX_DEPTH
        else:
            cv_image = np.ndarray((data.height, data.width), '<H', data.data, 0)
            cv_image = cv_image.astype('float32') * 0.001 # convert pixel value to meter
            cv_image[np.isnan(cv_image)] = IMAGE_MAX_DEPTH # max_range


        # send image to GPU on torch
        np_image = np.asarray(cv_image).copy()
        np_image[np_image > IMAGE_MAX_DEPTH] = IMAGE_MAX_DEPTH
        np_image[np_image < 0.20] = -1.0
        
        np_image.setflags(write=1)
        np_image = np_image/IMAGE_MAX_DEPTH
        np_image[np_image < 0.2/IMAGE_MAX_DEPTH] = -1.0

        input_image = np_image.copy()
        if np_image.shape != (IMAGE_HEIGHT,IMAGE_WIDTH):
            torch_img = torch.from_numpy(np_image).float()
            rescaled_image = torch.nn.functional.interpolate(torch_img.unsqueeze(0).unsqueeze(0), (IMAGE_HEIGHT, IMAGE_WIDTH)).squeeze(0).squeeze(0)
            input_image = rescaled_image.numpy()


        means, reconstruction, compute_time = self.net_interface.forward(input_image, calculate_reconstruction=CALCULATE_RECONSTRUCTION)
        
        if CALCULATE_RECONSTRUCTION == True:
            img_filtered_uint8 = (reconstruction * 255).astype(np.uint8)
            msg_filtered = Image()
            msg_filtered.height = IMAGE_HEIGHT # hardcoded and not data.height
            msg_filtered.width = IMAGE_WIDTH # hardcoded and not data.width
            msg_filtered.encoding = "8UC1"
            msg_filtered.is_bigendian = 0
            msg_filtered.step = IMAGE_WIDTH # 1 byte for each pixel
            msg_filtered.data = np.reshape(img_filtered_uint8, (IMAGE_WIDTH*IMAGE_HEIGHT,)).tolist()
            self.decode_img_publisher.publish(msg_filtered)


        latent_space_msg = Float32MultiArray()
        latent_space_msg.data = means.flatten().tolist()
        self.latent_space_publisher.publish(latent_space_msg)
        print('Compute time:', compute_time)


if __name__ == "__main__":
    rospy.init_node("vae_interface")
    print("Node Initialized. Loading weights.")
    nav_interface = VAENavInterface()
    print("Loaded weights. Lets gooooooooooo.......")
    rospy.spin()

sevae/inference/src/vae_node.py
ntnu-arl/sevae
import torch
import torch.nn as nn

def unweighted_non_semantic_reconstruction_with_MSE_KLD_Loss_unweighted_for_invalid_pixels(x, x_logit, semantic_image, gcam_map, mean, logvar, beta_coeff=1.0, latent_dims=128):
    invalid_pixel_mask = torch.where(x > 0, torch.ones_like(x), torch.zeros_like(x))
    MSE_LOSS = nn.MSELoss(reduction="none")
    cross_ent = MSE_LOSS(torch.sigmoid(x_logit), x) * invalid_pixel_mask
    reconstruction_loss = torch.mean(torch.sum(cross_ent, dim=[1, 2, 3]))
    kld_loss = -0.5 * torch.mean(torch.sum(1 + logvar - mean.pow(2) - logvar.exp(), dim=1))
    beta_norm = (beta_coeff*latent_dims)/(480.0*270.0)
    kld_loss = kld_loss * beta_norm
    return reconstruction_loss + kld_loss, kld_loss


def unweighted_semantic_reconstruction_with_MSE_KLD_Loss_unweighted_for_invalid_pixels(x, x_logit, semantic_image, gcam_map, mean, logvar):
    invalid_pixel_mask = torch.where(x > 0, torch.ones_like(x), torch.zeros_like(x))
    MSE_LOSS = nn.MSELoss(reduction="none")

    # Get unique semantic labels
    unique_semantic_labels, unique_counts = torch.unique(semantic_image[semantic_image > 10], return_counts=True)
    weight_matrix = torch.ones_like(semantic_image)

    # for each unique semantic label, get the number of pixels in the semantic mask and mask the invalid pixels
    for i in range(len(unique_semantic_labels)):
        if unique_semantic_labels[i] < 10:
            continue
        counts = unique_counts[i]
        weight_per_pixel = 1.0
        if counts > 40:
            weight_per_pixel = max(6000.0 / counts, 15.0)
        
        # add weight to weight matrix based on the number of pixels per semantic label
        weight_matrix[semantic_image == unique_semantic_labels[i]] *= weight_per_pixel
    cross_ent = MSE_LOSS(torch.sigmoid(x_logit), x) * invalid_pixel_mask * weight_matrix
    reconstruction_loss = torch.mean(torch.sum(cross_ent, dim=[1, 2, 3]))
    kld_loss = -0.5 * torch.mean(torch.sum(1 + logvar - mean.pow(2) - logvar.exp(), dim=1))
    beta_norm = (1.0*128)/(480.0*270.0)
    kld_loss = kld_loss * beta_norm
    return reconstruction_loss + kld_loss, kld_loss

def vanilla_image_MSE_Loss(x, x_hat):
    invalid_pixel_mask = torch.where(x_hat > 0, torch.ones_like(x_hat), torch.zeros_like(x_hat))
    MSE_LOSS = nn.MSELoss(reduction="none")
    cross_ent = MSE_LOSS(x_hat, x)*invalid_pixel_mask
    reconstruction_loss = torch.mean(torch.sum(cross_ent, dim=[1, 2, 3]))
    return reconstruction_loss
sevae/networks/Loss/loss_functions.py
ntnu-arl/sevae
import math

class RunningLoss:
    def __init__(self, batch_size):
        self.batch_size = batch_size
        self.aggregated_loss = 0.0
        self.num_samples = 0
        self.reset()
    
    def update(self, loss):
        self.aggregated_loss += loss/self.batch_size
        self.num_samples += 1
        self.avg = self.aggregated_loss / (self.num_samples)
    
    def reset(self):
        self.aggregated_loss = 0.0
        self.num_samples = 0
        self.avg = 0.0
    

sevae/networks/Loss/running_loss.py
ntnu-arl/sevae
import torch
import torch.nn as nn

class ImgDecoder(nn.Module):
    """Generates an image reconstruction using ConvTranspose2d layers. 
    
    Adapted from https://github.com/microsoft/AirSim-Drone-Racing-VAE-Imitation/blob/master/racing_models/cmvae.py
    """

    def __init__(self, input_dim=1, latent_dim=64, with_logits=False):
        """
        Parameters
        ----------
        latent_dim: int
            The latent dimension.
        """

        super(ImgDecoder, self).__init__()
        print('[ImgDecoder] Starting create_model')
        self.with_logits = with_logits
        self.n_channels = input_dim
        self.dense = nn.Linear(latent_dim, 512)
        self.dense1 = nn.Linear(512, 9*15*128)
        
        # Pytorch docs: output_padding is only used to find output shape, but does not actually add zero-padding to output
        self.deconv1 = nn.ConvTranspose2d(128, 128, kernel_size=3, stride=1, padding=1)
        self.deconv2 = nn.ConvTranspose2d(128, 64, kernel_size=5, stride=2, padding=(2,2), output_padding=(0,1), dilation=1)
        self.deconv4 = nn.ConvTranspose2d(64, 32, kernel_size=6, stride=4, padding=(2,2), output_padding=(0,0), dilation=1)
        self.deconv6 = nn.ConvTranspose2d(32, 16, kernel_size=6, stride=2, padding=(0,0), output_padding=(0,1))
        self.deconv7 = nn.ConvTranspose2d(16, self.n_channels, kernel_size=4, stride=2, padding=2) # tanh activation or sigmoid

        print('[ImgDecoder] Done with create_model')

        print("Defined decoder.")

    def forward(self, z):
        return self.decode(z)
    
    def decode(self, z):
        x = self.dense(z)
        x = torch.relu(x)
        x = self.dense1(x)
        x = x.view(x.size(0), 128, 9, 15)

        x = self.deconv1(x)
        x = torch.relu(x)

        x = self.deconv2(x)
        x = torch.relu(x)

        x = self.deconv4(x)
        x = torch.relu(x)


        x = self.deconv6(x)
        x = torch.relu(x)

        x = self.deconv7(x)
        # print(f"- After deconv 7, mean: {x.mean():.3f} var: {x.var():.3f}")
        if self.with_logits:
            return x

        x = torch.sigmoid(x)
        # print(f"- After sigmoid, mean: {x.mean():.3f} var: {x.var():.3f}")
        return x


if __name__== '__main__':
    from torchsummary import summary
    
    latent_dim = 64
    img_decoder = ImgDecoder(latent_dim=latent_dim, input_dim=1).to("cuda")
    summary(img_decoder, (1,latent_dim), device="cuda")
sevae/networks/VAE/condensed_decoder.py
ntnu-arl/sevae
import torch
import torch.nn as nn

class Encoder(nn.Module):
    """
    ResNet8 architecture as encoder.
    """
    def __init__(self, input_dim, latent_dim):
        """
        Parameters:
        ----------
        input_dim: int
            Number of input channels in the image.
        latent_dim: int
            Number of latent dimensions
        """
        super(Encoder, self).__init__()
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.define_encoder()
        self.elu = nn.ELU()
        print("Defined Encoder encoder.")
    
    def define_encoder(self):
        
        # define batch norm functions
        # self.bn0 = nn.BatchNorm2d(32)
        # self.bn1 = nn.BatchNorm2d(64)
        # self.bn2 = nn.BatchNorm2d(128)
        # self.bn3 = nn.BatchNorm2d(64)
        # self.bn4 = nn.BatchNorm2d(64)
        # self.bn5 = nn.BatchNorm2d(128)

        # self.max0 = nn.MaxPool2d(kernel_size=2, stride=2)

        # define conv functions
        self.conv0 = nn.Conv2d(self.input_dim, 32, kernel_size=5, stride=2, padding=2)
        self.conv0_1 = nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=2)
        nn.init.xavier_uniform_(self.conv0_1.weight, gain=nn.init.calculate_gain('linear'))
        nn.init.zeros_(self.conv0_1.bias)
        
        
        self.conv1_0 = nn.Conv2d(32, 32, kernel_size=5, stride=2, padding=1)
        self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        nn.init.xavier_uniform_(self.conv1_1.weight, gain=nn.init.calculate_gain('linear'))
        nn.init.zeros_(self.conv1_1.bias)

        self.conv2_0 = nn.Conv2d(64, 64, kernel_size=5, stride=2, padding=2)
        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
        nn.init.xavier_uniform_(self.conv2_1.weight, gain=nn.init.calculate_gain('linear'))
        nn.init.zeros_(self.conv2_1.bias)

        self.conv3_0 = nn.Conv2d(128, 128, kernel_size=5, stride=2)


        self.conv0_jump_2 = nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1)
        self.conv1_jump_3 = nn.Conv2d(64, 128, kernel_size=5, stride=4, padding=(2, 1))



        



        
        # self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1)
        # self.conv2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
        # self.conv3 = nn.Conv2d(32, 32, kernel_size=1, stride=2)
        # nn.init.xavier_uniform_(self.conv3.weight, gain=nn.init.calculate_gain('linear'))
        # nn.init.zeros_(self.conv3.bias)

        # self.conv4 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        # self.conv5 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
        # self.conv6 = nn.Conv2d(32, 64, kernel_size=1, stride=2)
        # nn.init.xavier_uniform_(self.conv6.weight, gain=nn.init.calculate_gain('linear'))
        # nn.init.zeros_(self.conv6.bias)

        # self.conv7 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
        # self.conv8 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
        # self.conv9 = nn.Conv2d(64, 128, kernel_size=1, stride=2)
        # nn.init.xavier_uniform_(self.conv9.weight, gain=nn.init.calculate_gain('linear'))
        # nn.init.zeros_(self.conv9.bias)

        
        self.dense0 = nn.Linear(3*6*128, 512)
        self.dense1 = nn.Linear(512, 2*self.latent_dim)
        # self.dense2 = nn.Linear(2*self.latent_dim, 2*self.latent_dim)

        print("[Encoder] Encoder network initialized.")

    def forward(self, img):
        return self.encode(img)
    
    def encode(self, img):
        """
        Encodes the input image.
        """

        # conv0
        x0_0 = self.conv0(img)
        x0_1 = self.conv0_1(x0_0)
        x0_1 = self.elu(x0_1)

        x1_0 = self.conv1_0(x0_1)
        x1_1 = self.conv1_1(x1_0)

        x0_jump_2 = self.conv0_jump_2(x0_1)

        x1_1 = x1_1 + x0_jump_2

        x1_1 = self.elu(x1_1)

        x2_0 = self.conv2_0(x1_1)
        x2_1 = self.conv2_1(x2_0)

        x1_jump3 = self.conv1_jump_3(x1_1)

        x2_1 = x2_1 + x1_jump3

        x2_1 = self.elu(x2_1)

        x3_0 = self.conv3_0(x2_1)

        x = x3_0.view(x3_0.size(0), -1)

        x = self.dense0(x)
        x = self.elu(x)
        x = self.dense1(x)
        return x





        




if __name__== '__main__':
    from torchsummary import summary
    latent_dim = 64
    encoder = Encoder(latent_dim=latent_dim, input_dim=1).to("cuda")
    summary(encoder, (1, 270, 480), device="cuda")   

sevae/networks/VAE/condensed_encoder.py
ntnu-arl/sevae
import torch
import torch.nn as nn

from .condensed_encoder import Encoder
from .condensed_decoder import ImgDecoder




import sys

class Lambda(nn.Module):
    """Lambda function that accepts tensors as input."""
    
    def __init__(self, func):
        super(Lambda, self).__init__()
        self.func = func
        
    def forward(self, x):
        return self.func(x)


class VAE(nn.Module):
    """Variational Autoencoder for reconstruction of depth images."""

    def __init__(self, input_dim=1, latent_dim=64, with_logits=False, inference_mode = False):
        """
        Parameters
        ----------
        input_dim: int
            The number of input channels in an image.
        latent_dim: int
            The latent dimension.
        """

        super(VAE, self).__init__()
        
        self.with_logits = with_logits
        self.input_dim = input_dim
        self.latent_dim = latent_dim
        self.inference_mode = inference_mode
        self.encoder = Encoder(input_dim=self.input_dim, latent_dim=self.latent_dim)
        self.img_decoder = ImgDecoder(input_dim=1, latent_dim=self.latent_dim, with_logits=self.with_logits)
        
        self.mean_params = Lambda(lambda x: x[:, :self.latent_dim]) # mean parameters
        self.logvar_params = Lambda(lambda x: x[:, self.latent_dim:]) # log variance parameters
        
    
    def forward(self, img):
        """ Do a forward pass of the VAE. Generates a reconstructed image based on img
        Parameters
        ----------
        img: torch.Tensor
            The input image.
        """
        
        # encode
        z = self.encoder(img)

        # reparametrization trick
        mean = self.mean_params(z)
        logvar = self.logvar_params(z)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        if self.inference_mode:
            eps = torch.zeros_like(eps)
        z_sampled = mean + eps * std

        # decode
        img_recon = self.img_decoder(z_sampled)
        return img_recon, mean, logvar, z_sampled

    def forward_test(self, img):
        """ Do a forward pass of the VAE. Generates a reconstructed image based on img
        Parameters
        ----------
        img: torch.Tensor
            The input image.
        """
        
        # encode
        z = self.encoder(img)

        # reparametrization trick
        mean = self.mean_params(z)
        logvar = self.logvar_params(z)
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        if self.inference_mode:
            eps = torch.zeros_like(eps)
        z_sampled = mean + eps * std

        # decode
        img_recon = self.img_decoder(z_sampled)
        return img_recon, mean, logvar, z_sampled

    
    def encode(self, img):
        """ Do a forward pass of the VAE. Generates a latent vector based on img
        Parameters
        ----------
        img: torch.Tensor
            The input image.
        """
        z = self.encoder(img)
        
        means = self.mean_params(z)
        logvars = self.logvar_params(z)
        std = torch.exp(0.5 * logvars)
        eps = torch.randn_like(logvars)
        if self.inference_mode:
            eps = torch.zeros_like(eps)
        z_sampled = means + eps * std

        return z_sampled, means, std
    

    def decode(self, z):
        """ Do a forward pass of the VAE. Generates a reconstructed image based on z
        Parameters
        ----------
        z: torch.Tensor
            The latent vector.
        """
        img_recon = self.img_decoder(z)
        if self.with_logits:
            return torch.sigmoid(img_recon)
        return img_recon
    
    def set_inference_mode(self, mode):
        self.inference_mode = mode


if __name__ == "__main__":
    from torchsummary import summary

    LATENT_DIM = 64
    device = torch.device("cpu")
    encoder = Encoder(input_dim=1, latent_dim=LATENT_DIM).to(device)
    summary(encoder, input_size=(1, 270, 480), batch_size=-1, device=device.type)

    decoder = ImgDecoder(input_dim=1, latent_dim=LATENT_DIM).to(device)
    summary(decoder, input_size=(1, LATENT_DIM), batch_size=-1, device=device.type)

    vae = VAE(input_dim=1, latent_dim=LATENT_DIM, with_logits=False)
    summary(vae, input_size=(1, 270, 480), batch_size=-1, device=device.type)
sevae/networks/VAE/vae.py
ntnu-arl/sevae
import os, sys
from random import shuffle

import math
import time
import numpy as np
import matplotlib.pyplot as plt
from collections import deque

from networks.VAE.vae import VAE
from networks.Loss.running_loss import RunningLoss
from networks.Loss.loss_functions import *
# from datasets.depth_image_dataset import DepthImageDataset, collate_batch
from datasets.depth_semantic_image_dataset import IMAGES_PER_TF_RECORD, DepthSemanticImageDataset, collate_batch

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchsummary import summary
from torch.utils.tensorboard import SummaryWriter
import torchvision


# import module for random sampling
import random


import tensorflow as tf
import yaml
import inspect




#  Use argument parser to set arguments of experiment name
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--experiment_name", type=str, default="default")
parser.add_argument("--load_model", type=str, default=None)
parser.add_argument("--load_model_path", type=str, default=None)


gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        # Currently, memory growth needs to be the same across GPUs
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.experimental.list_logical_devices('GPU')
        print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
    except RuntimeError as e:
        # Memory growth must be set before GPUs have been initialized
        print("GPU error")
        print(e)


device = torch.device("cuda")
device0 = torch.device("cuda:0")


# VAE Hyperparams
LATENT_DIM = 128
NUM_EPOCHS = 40
ONE_TFRECORD = False
BATCH_SIZE = 64
LEARNING_RATE = 1e-4

SAVE_MODEL = False
LOAD_MODEL = False
SAVE_INTERVAL = 10  # Save the model every 10 batches

# Model Paths
BASE_PATH = "/home/arl/MihirK/DATA"
EXPERIMENT_NAME = "default"
EXPERIMENT_BASE_PATH = os.path.join(BASE_PATH, EXPERIMENT_NAME)
SAVE_MODEL_FILE = os.path.join(BASE_PATH, "models") # "_epochxx.pth" appended in training
LOAD_MODEL_FILE = os.path.join(BASE_PATH, "vae_models/final_BCE__weighted_100_canny_mae2__weighted_10_depth_100.pth")

# Data Path
TFRECORD_FOLDER = os.path.join(BASE_PATH, "datasets")
TFRECORD_TEST_FOLDER = os.path.join(TFRECORD_FOLDER, "test")

MULTI_GPU = True
FILL_UNDEFINED_PIXELS_WITH_NEGATIVE_VALUES = True

ADD_NOISE_TO_INPUT = False



MAX_DEPTH = 10.0
MIN_DEPTH = 0.15


def make_grid_for_tensorboard(images_list, n_grids=2):
    joined_images = []
    [joined_images.extend(images[:n_grids]) for images in images_list]
    return torchvision.utils.make_grid(joined_images, nrow=n_grids, padding=5)


def get_noise(means, std_dev, const_multiplier):
    '''
    
    '''
    return const_multiplier*torch.normal(means, std_dev)

def process_for_training(input_image, filled_input_image, semantic_input_image):
    '''
    Function to process the input image for training
    '''
    processed_input_image = input_image.clone()
    processed_filled_input_image = filled_input_image.clone()
    processed_semantic_input_image = semantic_input_image.clone()

    processed_input_image[processed_input_image > MAX_DEPTH] = MAX_DEPTH
    processed_input_image[processed_input_image < MIN_DEPTH] = -1.0
    processed_input_image = processed_input_image / MAX_DEPTH
    processed_input_image[processed_input_image < 0] = -1.0

    processed_filled_input_image = torch.clamp(processed_filled_input_image, min=0, max=MAX_DEPTH)
    processed_filled_input_image[processed_filled_input_image < MIN_DEPTH] = MAX_DEPTH
    processed_filled_input_image = processed_filled_input_image / MAX_DEPTH

    processed_semantic_input_image = torch.where(input_image < 0.99*MAX_DEPTH, semantic_input_image, torch.zeros_like(semantic_input_image))

    # Do not have any semantic less thatsemaantic number 9
    processed_semantic_input_image[processed_semantic_input_image < 9] = 0

    if FILL_UNDEFINED_PIXELS_WITH_NEGATIVE_VALUES:
        image_to_reconstruct = torch.where(input_image > MIN_DEPTH, processed_input_image, -1.0*torch.ones_like(input_image))
    
    processed_input_image_with_noise = processed_input_image.clone()
    if ADD_NOISE_TO_INPUT:
        std_dev = torch.zeros_like(input_image)
        std_dev[:] = input_image*MIN_DEPTH/MAX_DEPTH # interpret this as: std_dev at max depth = 0.15m. std_dev at min depth = 0.0m. linearly increasing in between.
        processed_input_image_with_noise = image_to_reconstruct + get_noise(torch.zeros_like(image_to_reconstruct), torch.ones_like(image_to_reconstruct), std_dev)
        processed_input_image_with_noise[processed_input_image_with_noise > 1.0] = 1.0
        processed_input_image_with_noise[input_image < 0] = -1.0

    return processed_input_image_with_noise, image_to_reconstruct, processed_input_image, processed_filled_input_image, processed_semantic_input_image


def train_model(model, train_dataset_loader, test_dataset_loader, epochs, writer, batch_size, loss_fn, optimizer, save_interval=10):
    '''
    Function to train the given input model based on the given data and update the tensorboard
    '''
    torch.autograd.set_detect_anomaly(True)
    # Initialize the model
    model.train()
    # Initialize the loss and the optimizer
    loss_meter = RunningLoss(batch_size)
    # Initialize the number of batches
    num_batches = int(len(train_dataset_loader)/ batch_size)

    num_test_batches = int(len(test_dataset_loader)/ batch_size)
    # Initialize the number of epochs
    for epoch in range(epochs):
        if MULTI_GPU:
            model.module.set_inference_mode(False)
        else:
            model.set_inference_mode(False)
        model.train()
        # Initialize the time counter
        epoch_start_time = time.time()
        # Initialize the number of batches
        for batch_idx, (num_envs, images_per_env, depth_data, filtered_data, semantic_data) in enumerate(train_dataset_loader):
            batch_start_time = time.time()
            model.zero_grad()
            optimizer.zero_grad()

            depth_data = depth_data.to(device).unsqueeze(1)
            filtered_data = filtered_data.to(device).unsqueeze(1)
            semantic_data = semantic_data.to(device).unsqueeze(1)

            noisy_image, depth_data_to_reconstruct, filtered_data, filled_filtered_data, semantic_data = process_for_training(depth_data, filtered_data, semantic_data)

            # Forward pass
            reconstructed_image, means, log_vars, sampled_latent_vars  = model(noisy_image)
            clamped_g_cam_map = None
            loss, kld_loss = loss_fn(depth_data_to_reconstruct, reconstructed_image, semantic_data, clamped_g_cam_map, means, log_vars)
            
            # Update the loss meter
            loss_meter.update(loss.item())
            avg_iter_time = (time.time() - epoch_start_time) / (batch_idx + 1)

            # Backward pass
            loss.backward()
            optimizer.step()

            # Update the tensorboard
            if batch_idx % save_interval == 0 and batch_idx is not 0:
                writer.add_scalar('Train/Loss', loss.item()/BATCH_SIZE, epoch * num_batches + batch_idx)
                writer.add_scalar('Train/KLD Loss', kld_loss.item()/BATCH_SIZE, epoch * num_batches + batch_idx)
                print(f"[TRAINING] Epoch: {epoch}/{epochs} Batch: {batch_idx}/{num_batches} Avg. Train Loss: {loss.item()/BATCH_SIZE:.4f}, KL Div Loss.: {kld_loss.item()/BATCH_SIZE:.4f}"\
                    f"Time: {time.time() - batch_start_time:.2f}s, Est. time remaining: {(num_batches - batch_idx)*avg_iter_time :.2f}s")

                # add image to the tensorboard
                grid = make_grid_for_tensorboard([filled_filtered_data, depth_data_to_reconstruct, noisy_image, torch.sigmoid(reconstructed_image), semantic_data], n_grids=4)
                writer.add_image('training/images', grid, global_step=epoch*num_batches + batch_idx)
                if batch_idx % (5*save_interval) == 0:
                    torchvision.utils.save_image(grid, EXPERIMENT_BASE_PATH + "/training_images" + "/" + EXPERIMENT_NAME + "_epoch_" + str(epoch) + "_batch_" + str(batch_idx) + ".png")

        # Print the statistics
        print('Epoch: %d, Loss: %.4f, Time: %.4f' %
              (epoch, loss_meter.avg, time.time() - epoch_start_time))
        # Reset the loss meter
        loss_meter.reset()
        print("Saving model...")
        torch.save(model.state_dict(), os.path.join(
            EXPERIMENT_BASE_PATH+"/models", '%s_LD_%d_epoch_%d.pth' % (EXPERIMENT_NAME, LATENT_DIM,  epoch)))
        print("[DONE] Savng model at ", EXPERIMENT_BASE_PATH+"/models")

        # # Evaluate the model
        model.eval()
        if MULTI_GPU:
            model.module.set_inference_mode(True)
        else:
            model.set_inference_mode(True)
        for batch_idx, (num_envs, images_per_env, depth_data, filtered_data, semantic_data) in enumerate(test_dataset_loader):
            model.zero_grad()
            optimizer.zero_grad()
            # Process for testing
            
            depth_data = depth_data.to(device).unsqueeze(1)
            filtered_data = filtered_data.to(device).unsqueeze(1)
            semantic_data = semantic_data.to(device).unsqueeze(1)

            noisy_image, depth_data_to_reconstruct, filtered_data, filled_filtered_data, semantic_data = process_for_training(depth_data, filtered_data, semantic_data)

            # Forward pass
            reconstructed_image, means, log_vars, sampled_latent_vars  = model(noisy_image)
            clamped_g_cam_map = None
            loss, kld_loss = loss_fn(depth_data_to_reconstruct, reconstructed_image, semantic_data, clamped_g_cam_map, means, log_vars)
            
            # Update the loss meter
            loss_meter.update(loss.item())
            avg_iter_time = (time.time() - epoch_start_time) / (batch_idx + 1)

            grid = make_grid_for_tensorboard([filled_filtered_data, depth_data_to_reconstruct, noisy_image, torch.sigmoid(reconstructed_image), semantic_data], n_grids=4)
            writer.add_image('testing/images', grid, global_step=epoch)
            if batch_idx % (save_interval) == 0 and batch_idx is not 0:
                print(f"[TESTING] Epoch: {epoch}/{epochs} Batch: {batch_idx}/{num_test_batches} Avg. Train Loss: {loss.item()/BATCH_SIZE:.4f}, KL Div Loss.: {kld_loss.item()/BATCH_SIZE:.4f}")
                torchvision.utils.save_image(grid, EXPERIMENT_BASE_PATH + "/testing_images" + "/" + EXPERIMENT_NAME + "_epoch_" + str(epoch) + "_batch_" + str(batch_idx) + ".png")
                # Update the tensorboard
                writer.add_scalar('Test/Loss', loss.item()/BATCH_SIZE,
                                epoch * num_test_batches + batch_idx)
                writer.add_scalar('Test/KL Div Loss', kld_loss.item()/BATCH_SIZE,
                                    epoch * num_test_batches + batch_idx)
                              
        # Print the statistics
        print('Test Loss:', loss_meter.avg)
        loss_meter.reset()

    return model


def main():
    global USE_GUIDED_ATTENTION_LOSS
    global EXPERIMENT_NAME
    global EXPERIMENT_BASE_PATH

    # Load Dataset
    print("Loading train dataset from ", TFRECORD_FOLDER)
    train_dataset = DepthSemanticImageDataset(tfrecord_folder=TFRECORD_FOLDER, shuffle=True, batch_size=BATCH_SIZE, one_tfrecord=ONE_TFRECORD)
    print("Loading test dataset", TFRECORD_TEST_FOLDER)
    test_dataset = DepthSemanticImageDataset(tfrecord_folder=TFRECORD_TEST_FOLDER, shuffle=False, batch_size=BATCH_SIZE, one_tfrecord=ONE_TFRECORD, test_dataset=True)

    # Parse arguments from argparser 
    args = parser.parse_args()

    # Check if load model is true , and if so, get the model path
    if args.load_model:
        LOAD_MODEL = True
        model_path = args.model_path
        LOAD_MODEL_FILE = args.model_filesss
        print("Loading model from ", model_path)
    else:
        LOAD_MODEL = False
        LOAD_MODEL_FILE = None
        model_path = None
    
    EXPERIMENT_NAME = args.experiment_name
    EXPERIMENT_BASE_PATH = os.path.join(BASE_PATH, EXPERIMENT_NAME)

    # Define the data loaders
    train_loader = DataLoader(dataset=train_dataset,
                              batch_size=1, collate_fn=collate_batch)
    test_loader = DataLoader(
        dataset=test_dataset, batch_size=1, collate_fn=collate_batch, shuffle=False)
    print("Loaded data loaders")

    print("Number of training samples:", len(train_dataset))
    print("Number of testing samples:", len(test_dataset))
    
    EXPERIMENT_BASE_PATH = os.path.join(BASE_PATH, "experiments", EXPERIMENT_NAME)
    # initialize the folder for saving model, params, config, and training and testing images
    if not os.path.exists(EXPERIMENT_BASE_PATH):
        os.makedirs(os.path.join(EXPERIMENT_BASE_PATH))
        os.makedirs(os.path.join(EXPERIMENT_BASE_PATH, "models"))
        os.makedirs(os.path.join(EXPERIMENT_BASE_PATH, "training_images"))
        os.makedirs(os.path.join(EXPERIMENT_BASE_PATH, "testing_images"))
        os.makedirs(os.path.join(EXPERIMENT_BASE_PATH, "tensorboard"))
        writer = SummaryWriter(log_dir=os.path.join(EXPERIMENT_BASE_PATH, "tensorboard"))

    # elif EXPERIMENT_NAME == "default":
    #     pass
    else:
        print("Experiment name already exists. Please choose another name.")
        sys.exit(0)

    # Define and initialize the model
    model = VAE(latent_dim=LATENT_DIM, with_logits=True)
    model = model.to(device)
    summary(model, (1, 270, 480))
    if MULTI_GPU:
        model = nn.DataParallel(model)
        model = model.to(device0)
    if LOAD_MODEL:
        if MULTI_GPU:
            model.module.load_state_dict(torch.load(LOAD_MODEL_FILE))
        else:
            model.load_state_dict(torch.load(LOAD_MODEL_FILE))

    loss = unweighted_semantic_reconstruction_with_MSE_KLD_Loss_unweighted_for_invalid_pixels

    # Define the optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)

    TF_RECORDS_TRAINING = train_dataset.tfrecord_fnames
    TF_RECORDS_TESTING = test_dataset.tfrecord_fnames


    # populate a file detailing the config as a yaml file and save them to a file
    # the details should include batch sizes, tf_record file names for training and testing, latent spaces, loss function name, etc.
    config = {
        "base_path": BASE_PATH,
        "experiment_base_path": EXPERIMENT_BASE_PATH,
        "experiment name": EXPERIMENT_NAME,
        "batch_size": BATCH_SIZE,
        "tfrecord_folder": TFRECORD_FOLDER,
        "tfrecord_test_folder": TFRECORD_TEST_FOLDER,
        "latent_dim": LATENT_DIM,
        "loss": loss.__name__,
        "loss_function_definition": inspect.getsource(loss),
        "learning_rate": LEARNING_RATE,
        "num_epochs": NUM_EPOCHS,
        "save_interval": SAVE_INTERVAL,
        "experiment_name": EXPERIMENT_NAME,
        "load_model": LOAD_MODEL,
        "load_model_file": LOAD_MODEL_FILE,
        "model_path": model_path,
        "multi_gpu": MULTI_GPU,
        "use_filtered_data_to_fill_in_undefined_pixels": FILL_UNDEFINED_PIXELS_WITH_FILTERED_DATA,
        "use_negative_values_to_fill_in_undefined_pixels": FILL_UNDEFINED_PIXELS_WITH_NEGATIVE_VALUES,
        "tf_records_training": TF_RECORDS_TRAINING,
        "tf_records_testing": TF_RECORDS_TESTING
        }
    
    # write config to a yaml file and save it to the experiment folder
    with open(os.path.join(EXPERIMENT_BASE_PATH, "config.yaml"), "w") as f:
        yaml.dump(config, f)
    
        
    # Print the current torch seed
    print("Current seed for training: ", torch.seed())

    # Train the model
    model = train_model(model, train_loader, test_loader, NUM_EPOCHS, writer, BATCH_SIZE,
                        loss, optimizer, save_interval=SAVE_INTERVAL)


if __name__ == "__main__":
    main()

sevae/train_seVAE.py
ntnu-arl/sevae

import os
import glob
import pickle
import sys
import torch


from sevae.networks.VAE.vae import VAE




LOAD_MODEL= True


# Data Path
BASE_PATH = "/home/arl/MihirK/DATA"

# VAE Hyperparams
LATENT_DIM = 128
LOAD_MODEL = True

MAX_DEPTH = 10.0
MIN_DEPTH = 0.2

global_pickle_counter = 0

def read_pickle_files(depth_pickle_file, append_filename):
    print("Reading pickle files: ", depth_pickle_file)
    if not os.path.isfile(depth_pickle_file):
        print("Depth file does not exist")
        return

    with open(depth_pickle_file, 'rb') as f:
        depth_img_list = pickle.load(f)

    num_env = 1
    num_timesteps = 0

    depth_latents = []
    depth_flipped_latents = []

    for episode in range(len(depth_img_list)):
        depth_latents.append([])
        depth_flipped_latents.append([])
        for timestep in range(len(depth_img_list[episode])):
            depth_image = depth_img_list[episode][timestep].copy()
            depth_image_torch = torch.from_numpy(depth_image).squeeze(-1).to("cuda")
            depth_image_torch = depth_image_torch.unsqueeze(0).unsqueeze(0)
            depth_image_flipped = torch.flip(depth_image_torch, [3])

            depth_image_torch = torch.clamp(depth_image_torch, 0, MAX_DEPTH)
            depth_image_flipped = torch.clamp(depth_image_flipped, 0, MAX_DEPTH)
            depth_image_torch = depth_image_torch / MAX_DEPTH
            depth_image_flipped = depth_image_flipped / MAX_DEPTH

            depth_image[depth_image < MIN_DEPTH/MAX_DEPTH] = -1.0
            depth_image_flipped[depth_image_flipped < MIN_DEPTH/MAX_DEPTH] = -1.0

            # get the latent vector
            latent_vector_straight_image = model.encode(depth_image_torch)[1].detach().cpu().numpy().squeeze(0)
            latent_vector_flipped_image = model.encode(depth_image_flipped)[1].detach().cpu().numpy().squeeze(0)

            depth_latents[episode].append(latent_vector_straight_image)
            depth_flipped_latents[episode].append(latent_vector_flipped_image)
            num_timesteps += 1

    # print(len(depth_img_list), num_timesteps)
    # save pickle
    depth_latent_filename = depth_pickle_file.replace("di_dump", "di_latent_"+append_filename)
    depth_flipped_latent_filename = depth_pickle_file.replace("di_dump", "di_flipped_latent_"+append_filename)
    print("Saving pickles to: ", depth_latent_filename, depth_flipped_latent_filename)
    with open(depth_latent_filename, 'wb') as f:
        pickle.dump(depth_latents, f)
    with open(depth_flipped_latent_filename, 'wb') as f:
        pickle.dump(depth_flipped_latents, f)
    print("Saved pickle to VAE compatible format")



if __name__ == "__main__":
    # parse over folder and get all di_dump*.p pickles and rgb_dump*.p pickles
    # read experiment name to load weights from 
    if len(sys.argv) < 3:
        print("Using IROS 2023 submission model weights")
        experiment_name  = "NTNU_thin_ft_labeled_invald_px_test"
    else:
        experiment_name = sys.argv[2]

    LOAD_MODEL_FILE = os.path.join(BASE_PATH, "experiments/"+experiment_name+"/models/"+experiment_name+"_LD_128_epoch_39.pth")
    

    # read folder name from first argument 
    folder_name = sys.argv[1]

    print("Loading di_dump pickles from folder: ", folder_name)

    # print("Checking folder: ", os.join(folder_name))

    if len(sys.argv) < 4:
        print("Appending filename with \"test\"")
        append_filename = "test"
    else:
        append_filename = sys.argv[3]

    print("Appending filename of latent pickles with: ", append_filename)

    print("Number of pickle pairs to process: ", len(glob.glob(folder_name + "/di_dump*.p")))

    model = VAE(latent_dim=LATENT_DIM, with_logits=True)

    print("Loading model from file: ", LOAD_MODEL_FILE)
    if LOAD_MODEL:
        print("Loaded state dict")
        loaded_dict = torch.load(LOAD_MODEL_FILE)
        for key in list(loaded_dict.keys()):
            if "module." in key:
                loaded_dict[key.replace("module.", "")] = loaded_dict[key]
                del loaded_dict[key]
        print(loaded_dict.keys())
        model.load_state_dict(loaded_dict)
    
    model.to("cuda")
    print("[DONE] Loading model from file: ", LOAD_MODEL_FILE)

    files = glob.glob(folder_name + "/di_dump*.p")
    pickles_list = []
    # get all files in the folder having rgb_dump in the name
    for file in glob.glob(folder_name + "/di_dump*.p"):
        depth_filename = file
        print("Reading files: ", depth_filename)
        # check if files exist 
        if not os.path.isfile(depth_filename):
            print("Depth file does not exist")
            continue
        pickles_list.append(depth_filename)
    
    pickles_list.sort()
    print("Number of pickles to process: ", len(pickles_list))
    for pickle_name in pickles_list:
        read_pickle_files(pickle_name, append_filename)
sevae/utils/create_latent_pickles_from_ORACLE_datasets.py
ntnu-arl/sevae
from collections import OrderedDict

# define function to clean up dataparallel config for loading
def clean_state_dict(state_dict):
    '''
    Removes the "module." from the state_dict keys
    '''
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        name = k
        if "module." in k:
            name = k.replace("module.", "")
        if "dronet." in name:
            name = name.replace("dronet.", "encoder.")
        new_state_dict[name] = v
    return new_state_dict


sevae/utils/utils.py
ntnu-arl/sevae