Talk2BEV: Language-Enhanced Bird's Eye View Maps (Submitted to ICRA'24)
import argparse
import os
import random

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import gradio as gr

from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from llava.conversation import conv_templates, SeparatorStyle
from llava.model.builder import load_pretrained_model
from llava.utils import disable_torch_init
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria

from transformers import TextStreamer
# Load model directly
# Use a pipeline as a high-level helper
from transformers import pipeline

from PIL import Image

from minigpt4.common.config import Config
from minigpt4.common.dist_utils import get_rank
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Chat, CONV_VISION

# imports modules for registration
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *

from nuscenes.utils.data_classes import PointCloud, LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix

def parse_args():
    parser = argparse.ArgumentParser(description="Demo")
    parser.add_argument("--sam-checkpoint", type=str, default="sam_vit_h_4b8939.pth", help="path to sam weights.")
    parser.add_argument("--model-path", type=str, default="/raid/t1/scratch/vikrant.dewangan/LLaVA/ckpt-old/", help="path to llava weights.")
    parser.add_argument("--gpu-id", type=int, default=0, help="specify the gpu to load the model.")
    parser.add_argument("--conv-mode", type=int, default="llava_v1", help="specify the gpu to load the model.")
    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )
    args = parser.parse_args()
    return args

def setup_seeds(config):
    seed = config.run_cfg.seed + get_rank()

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    cudnn.benchmark = False
    cudnn.deterministic = True

def distance(x1, y1, z1, x2, y2, z2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + 0 * (z1 - z2) ** 2) ** ( 1 / 2)
 
# Function to calculate K closest points
def kClosest(points, target, K):
    pts = []
    n = len(points)
    d = []

    for i in range(n):
        d.append({
            "first": distance(points[i][0], points[i][1], points[i][2], target[0], target[1], target[2]),
            "second": i
        })
     
    d = sorted(d, key=lambda l:l["first"])
 
    for i in range(K):
        pt = []
        pt.append(points[d[i]["second"]][0])
        pt.append(points[d[i]["second"]][1])
        pt.append(points[d[i]["second"]][2])
        pt.append(points[d[i]["second"]][3])
        pts.append(pt)
 
    return pts

def get_image_projected_points(cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, arr):
    """
        arr: points to be back-projected: (K closest points)
    """
    cam_imgs = [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right]
    cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']
    calibration_data = {
        'CAM_FRONT_LEFT':{
            'translation': [-1.57525595, -0.50051938, -1.50696033],
            'rotation': [[ 0.82254604, -0.56868433, -0.00401771], [ 0.00647832,  0.01643407, -0.99984396], [ 0.56866162,  0.82239167,  0.01720189]],
            'camera_intrinsic': [[1257.8625342125129, 0.0, 827.2410631095686], [0.0, 1257.8625342125129, 450.915498205774], [0.0, 0.0, 1.0]]
        },
        'CAM_FRONT':{
            'translation': [-1.72200568, -0.00475453, -1.49491292],
            'rotation': [[ 0.01026021, -0.99987258, -0.01222952], [ 0.00843345,  0.01231626, -0.99988859], [ 0.9999118 ,  0.01015593,  0.00855874]],
            'camera_intrinsic': [[1252.8131021185304, 0.0, 826.588114781398], [0.0, 1252.8131021185304, 469.9846626224581], [0.0, 0.0, 1.0]]
        },
        'CAM_FRONT_RIGHT':{
            'translation': [-1.58082566,  0.49907871, -1.51749368],
            'rotation': [[-0.84397973, -0.53614138, -0.01583178], [ 0.01645551,  0.00362107, -0.99985804], [ 0.5361226 , -0.84412044,  0.00576637]],
            'camera_intrinsic': [[1256.7485116440405, 0.0, 817.7887570959712], [0.0, 1256.7485116440403, 451.9541780095127], [0.0, 0.0, 1.0]]
        },
        'CAM_BACK_LEFT':{
            'translation': [-1.035691  , -0.48479503, -1.59097015],
            'rotation': [[ 0.94776036,  0.31896113,  0.00375564], [ 0.00866572, -0.0139763 , -0.99986478], [-0.31886551,  0.94766474, -0.01601021]],
            'camera_intrinsic': [[1256.7414812095406, 0.0, 792.1125740759628], [0.0, 1256.7414812095406, 492.7757465151356], [0.0, 0.0, 1.0]]
        },
        'CAM_BACK':{
            'translation': [-0.02832603, -0.00345137, -1.57910346],
            'rotation': [[ 0.00242171,  0.99998907, -0.00400023], [-0.01675361, -0.00395911, -0.99985181], [-0.99985672,  0.00248837,  0.01674384]],
            'camera_intrinsic': [[809.2209905677063, 0.0, 829.2196003259838], [0.0, 809.2209905677063, 481.77842384512485], [0.0, 0.0, 1.0]]
        },
        'CAM_BACK_RIGHT':{
            'translation': [-1.0148781 ,  0.48056822, -1.56239545],
            'rotation': [[-0.93477554,  0.35507456, -0.01080503], [ 0.01587584,  0.0113705 , -0.99980932], [-0.35488399, -0.93476883, -0.01626597]],
            'camera_intrinsic': [[1259.5137405846733, 0.0, 807.2529053838625], [0.0, 1259.5137405846733, 501.19579884916527], [0.0, 0.0, 1.0]]
        }
    }

    min_dist = 1.0
    flag = 0
    cam_img = None
    for cam_ind, cam_key in enumerate(cam_keys):
        barr = np.copy(arr)
        ppc = LidarPointCloud(barr.T);
        ppc.translate(np.array(calibration_data[cam_key]['translation']))
        ppc.rotate(np.array(calibration_data[cam_key]['rotation']))
        ddepths_img = ppc.points[2, :]
        points = view_points(ppc.points[:3, :], np.array(calibration_data[cam_key]['camera_intrinsic']), normalize=True)
        print(points)
        mask = np.ones(ddepths_img.shape[0], dtype=bool)
        mask = np.logical_and(mask, ddepths_img > min_dist)
        mask = np.logical_and(mask, points[0, :] > 1)
        mask = np.logical_and(mask, points[0, :] < 1600 - 1)
        mask = np.logical_and(mask, points[1, :] > 1)
        mask = np.logical_and(mask, points[1, :] < 900 - 1)

        print(mask)
        if mask.sum() > 0:
            # found
            flag = 1
            points = points[:, mask]
            ddepths_img = ddepths_img[mask]
            cam_img = cam_imgs[cam_ind]
            print(points.shape, " points")
            break

        if flag == 0:
            # no point able to back-project, just use front cam
            print("AAAAAAAAA no point found")
            cam_img = cam_front
    return cam_img

# ========================================
#             Model Initialization
# ========================================


print('Initializing Chat')
args = parse_args()
cfg = Config(args)

args.image_file = "/home/t1/vikrant.dewangan/llm-bev/MiniGPT-4/sample.jpg"
args.model_base = None

################ LLAVA #############
# model_name = get_model_name_from_path(args.model_path)
model_name = "llava"
tokenizer, model, image_processor, context_len = load_pretrained_model(args.model_path, args.model_base, model_name, True, True)

if 'llama-2' in model_name.lower():
    conv_mode = "llava_llama_2"
elif "v1" in model_name.lower():
    conv_mode = "llava_v1"
elif "mpt" in model_name.lower():
    conv_mode = "mpt"
else:
    conv_mode = "llava_v0"
print('Initialization Finished')

if args.conv_mode is not None and conv_mode != args.conv_mode:
    print('[WARNING] the auto inferred conversation mode is {}, while `--conv-mode` is {}, using {}'.format(conv_mode, args.conv_mode, args.conv_mode))
else:
    args.conv_mode = conv_mode
conv = conv_templates[args.conv_mode].copy()
if "mpt" in model_name.lower():
    roles = ('user', 'assistant')
else:
    roles = conv.roles

####################################


# ========================================
#             Sam Initialization
# ========================================

from segment_anything import sam_model_registry, SamPredictor

model_type = "vit_h"
device = "cuda"

sam = sam_model_registry[model_type](checkpoint=args.sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)

# ========================================
#             Gradio Setting
# ========================================

def gradio_reset(img_list):
    return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first', interactive=False),gr.update(value="Upload & Start Chat", interactive=True)

def upload_img(gr_img, text_input):
    if gr_img is None:
        return None, None, gr.update(interactive=True), None
    image = gr_img
    image_tensor = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].half().to(device)
    image = None
    return image_tensor, gr.update(interactive=False), gr.update(interactive=True, placeholder='Type and press Enter'), gr.update(value="Start Chatting", interactive=False)

def gradio_ask(user_message, image, chatbot):
    if len(user_message) == 0:
        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
    if model.config.mm_use_im_start_end:
        inp = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + user_message
    else:
        inp = DEFAULT_IMAGE_TOKEN + '\n' + user_message
    conv.append_message(conv.roles[0], inp)
    chatbot = chatbot + [[user_message, None]]
    prompt = conv.get_prompt()
    input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).cuda()
    stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
    keywords = [stop_str]
    stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids)
    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    return image, streamer, stopping_criteria, input_ids, chatbot

def gradio_answer(input_ids, streamer, stopping_criteria, image_tensor, chatbot):
    output_ids = model.generate(
        input_ids,
        images=image_tensor,
        do_sample=True,
        temperature=1.0,
        max_new_tokens=1024,
        streamer=streamer,
        use_cache=True,
        stopping_criteria=[stopping_criteria])

    outputs = tokenizer.decode(output_ids[0, input_ids.shape[1]:]).strip()
    chatbot[-1][1] = outputs
    return chatbot

title = """<h1 align="center">Talk2BEV: Demo</h1>"""
description = """<h3>Upload BEV, perspective images, LiDAR point cloud and start chatting!</h3>"""
article = """<p><a href='https://minigpt-4.github.io'><img src='https://img.shields.io/badge/Project-Page-Green'></a></p><p><a href='https://github.com/Vision-CAIR/MiniGPT-4'><img src='https://img.shields.io/badge/Github-Code-blue'></a></p><p><a href='https://raw.githubusercontent.com/Vision-CAIR/MiniGPT-4/main/MiniGPT_4.pdf'><img src='https://img.shields.io/badge/Paper-PDF-red'></a></p>
"""

#TODO show examples below

def find_bounding_box(image):
    # Find the coordinates of non-zero (foreground) pixels along each channel
    foreground_pixels = np.array(np.where(np.any(image != 0, axis=2)))

    # Calculate the bounding box coordinates
    min_y, min_x = np.min(foreground_pixels, axis=1)
    max_y, max_x = np.max(foreground_pixels, axis=1)

    return min_y, min_x, max_y, max_x

def crop_around_bounding_box(image):
    min_y, min_x, max_y, max_x = find_bounding_box(image)

    # Crop the image using the bounding box coordinates
    cropped_image = image[min_y:max_y+1, min_x:max_x+1, :]

    return cropped_image

with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    # gr.Markdown(article)

    image_tensor = gr.State()
    streamer = gr.State()
    stopping_criteria = gr.State()
    img = gr.State()
    input_ids = gr.State()

    tolerance = gr.Slider(label="Tolerance", info="How different colors can be in a segment.", minimum=0, maximum=256*3, value=50)

    with gr.Row():
        cam_left = gr.Image(label="CAM_LEFT - Upload")
        cam_front = gr.Image(label="CAM_FRONT - Upload")
        cam_right = gr.Image(label="CAM_RIGHT - Upload")
    
    with gr.Row():
        cam_rear_left = gr.Image(label="CAM_REAR_LEFT - Upload")
        cam_rear = gr.Image(label="CAM_REAR - Upload")
        cam_rear_right = gr.Image(label="CAM_REAR_RIGHT - Upload")

    with gr.Row():
        bev_img = gr.Image(label="BEV Image - Click a segment")
        image = gr.Image(label="Selected Segment", type="pil")                    
        lidar_data = gr.File(label="LiDAR Data - Upload a file")

    with gr.Row():
        upload_button = gr.Button(value="Submit Mask & Start Chat", interactive=True, variant="primary")
        clear = gr.Button("Clear All")

    with gr.Row():
        num_beams = gr.Slider(
            minimum=1,
            maximum=10,
            value=1,
            step=1,
            interactive=True,
            label="beam search numbers)",
        )
        
        temperature = gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=1.0,
            step=0.1,
            interactive=True,
            label="Temperature",
        )

    with gr.Row():
        with gr.Column():
            chat_state = gr.State()
            img_list = gr.State()
            chatbot = gr.Chatbot(label='LLava')
            text_input = gr.Textbox(label='User', placeholder='Please upload your image first', interactive=False)

    def get_select_coords(cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, bev, lidar, tolerance, evt: gr.SelectData):
        lidardata = np.load(lidar.name)
        # import pdb; pdb.set_trace()
        print(evt.index[1], evt.index[0])
        print(bev[evt.index[1], evt.index[0]])

        labels_allowed_seg = [14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
        labels_allowed_hdmap = [24, 25, 26, 27]
        print(np.array(bev[evt.index[1], evt.index[0]]).shape, np.array([0,0,255]).shape)
        if bev[evt.index[1], evt.index[0]][0] == 0 and bev[evt.index[1], evt.index[0]][1] == 0 and bev[evt.index[1], evt.index[0]][2] == 255:
            labels_allowed = labels_allowed_seg
        else:
            labels_allowed = labels_allowed_hdmap
        print(labels_allowed)
        input_point = np.array([[evt.index[0], evt.index[1]]])
        input_label = np.array([1])

        print(lidardata.shape)
        target = np.array([(evt.index[1] - 100)/2, (evt.index[0] - 100)/2, 0])
        pts = []
        for pt in lidardata.T:
            if pt[-1] in labels_allowed:
                pts.append(pt)
        print("final shape lidar: ", np.array(pts).shape)

        arr = np.array(kClosest(np.array(pts), target, 1))
    
        # CAM_BACK
        translation = np.array([-0.02832603, -0.00345137, -1.57910346])
        rotation = np.array([[ 0.00242171,  0.99998907, -0.00400023], [-0.01675361, -0.00395911, -0.99985181], [-0.99985672,  0.00248837,  0.01674384]])
        camera_intrinsic = np.array([[809.2209905677063, 0.0, 829.2196003259838], [0.0, 809.2209905677063, 481.77842384512485], [0.0, 0.0, 1.0]])

        cam_imgs = [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right]
        cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']
        calibration_data = {
            'CAM_FRONT_LEFT':{
                'translation': [-1.57525595, -0.50051938, -1.50696033],
                'rotation': [[ 0.82254604, -0.56868433, -0.00401771], [ 0.00647832,  0.01643407, -0.99984396], [ 0.56866162,  0.82239167,  0.01720189]],
                'camera_intrinsic': [[1257.8625342125129, 0.0, 827.2410631095686], [0.0, 1257.8625342125129, 450.915498205774], [0.0, 0.0, 1.0]]
            },
            'CAM_FRONT':{
                'translation': [-1.72200568, -0.00475453, -1.49491292],
                'rotation': [[ 0.01026021, -0.99987258, -0.01222952], [ 0.00843345,  0.01231626, -0.99988859], [ 0.9999118 ,  0.01015593,  0.00855874]],
                'camera_intrinsic': [[1252.8131021185304, 0.0, 826.588114781398], [0.0, 1252.8131021185304, 469.9846626224581], [0.0, 0.0, 1.0]]
            },
            'CAM_FRONT_RIGHT':{
                'translation': [-1.58082566,  0.49907871, -1.51749368],
                'rotation': [[-0.84397973, -0.53614138, -0.01583178], [ 0.01645551,  0.00362107, -0.99985804], [ 0.5361226 , -0.84412044,  0.00576637]],
                'camera_intrinsic': [[1256.7485116440405, 0.0, 817.7887570959712], [0.0, 1256.7485116440403, 451.9541780095127], [0.0, 0.0, 1.0]]
            },
            'CAM_BACK_LEFT':{
                'translation': [-1.035691  , -0.48479503, -1.59097015],
                'rotation': [[ 0.94776036,  0.31896113,  0.00375564], [ 0.00866572, -0.0139763 , -0.99986478], [-0.31886551,  0.94766474, -0.01601021]],
                'camera_intrinsic': [[1256.7414812095406, 0.0, 792.1125740759628], [0.0, 1256.7414812095406, 492.7757465151356], [0.0, 0.0, 1.0]]
            },
            'CAM_BACK':{
                'translation': [-0.02832603, -0.00345137, -1.57910346],
                'rotation': [[ 0.00242171,  0.99998907, -0.00400023], [-0.01675361, -0.00395911, -0.99985181], [-0.99985672,  0.00248837,  0.01674384]],
                'camera_intrinsic': [[809.2209905677063, 0.0, 829.2196003259838], [0.0, 809.2209905677063, 481.77842384512485], [0.0, 0.0, 1.0]]
            },
            'CAM_BACK_RIGHT':{
                'translation': [-1.0148781 ,  0.48056822, -1.56239545],
                'rotation': [[-0.93477554,  0.35507456, -0.01080503], [ 0.01587584,  0.0113705 , -0.99980932], [-0.35488399, -0.93476883, -0.01626597]],
                'camera_intrinsic': [[1259.5137405846733, 0.0, 807.2529053838625], [0.0, 1259.5137405846733, 501.19579884916527], [0.0, 0.0, 1.0]]
            }
        }

        min_dist = 1.0
        flag = 0
        cam_img = None
        for cam_ind, cam_key in enumerate(cam_keys):
            barr = np.copy(arr)
            ppc = LidarPointCloud(barr.T);
            ppc.translate(np.array(calibration_data[cam_key]['translation']))
            ppc.rotate(np.array(calibration_data[cam_key]['rotation']))
            ddepths_img = ppc.points[2, :]
            points = view_points(ppc.points[:3, :], np.array(calibration_data[cam_key]['camera_intrinsic']), normalize=True)
            print(points)
            mask = np.ones(ddepths_img.shape[0], dtype=bool)
            mask = np.logical_and(mask, ddepths_img > min_dist)
            mask = np.logical_and(mask, points[0, :] > 1)
            mask = np.logical_and(mask, points[0, :] < 1600 - 1)
            mask = np.logical_and(mask, points[1, :] > 1)
            mask = np.logical_and(mask, points[1, :] < 900 - 1)

            print(mask)
            if mask.sum() > 0:
                # found
                flag = 1
                points = points[:, mask]
                ddepths_img = ddepths_img[mask]
                cam_img = cam_imgs[cam_ind]
                print(points.shape, " points")
                break

        if flag == 0:
            # no point able to back-project, just use front cam
            print("AAAAAAAAA no point found")
            cam_img = cam_front

        predictor.set_image(cam_img)
        input_point = np.array(points.astype(np.int32).T[:, :2])
        input_label = np.array([1 for i in range(len(points[0]))])
        print(input_point, input_label)
        masks, scores, logits = predictor.predict(point_coords=input_point,point_labels=input_label,multimask_output=True)
        print(masks.shape, " masks")
        idxs = (masks[-1].astype(np.uint8) * 200) > 0
        img_copy = np.copy(cam_img)
        img_copy[~idxs] = 0
        img_copy = crop_around_bounding_box(img_copy)
        return img_copy

    bev_img.select(get_select_coords, [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, bev_img, lidar_data, tolerance], image)
    
    upload_button.click(upload_img, [image, text_input], [image_tensor, image, text_input, upload_button])
    
    text_input.submit(gradio_ask, [text_input, img, chatbot], [img, streamer, stopping_criteria, input_ids, chatbot]).then(
        gradio_answer, [input_ids, streamer, stopping_criteria, image_tensor, chatbot], [chatbot]
    )
    clear.click(gradio_reset, [img], [chatbot, image, text_input, upload_button], queue=False)

demo.launch(share=True, enable_queue=True)

click2chat/click2chat_llava.py
llmbev/talk2bev
import argparse
import os
import random

import numpy as np
import torch
import torch.backends.cudnn as cudnn
import gradio as gr

from minigpt4.common.config import Config
from minigpt4.common.dist_utils import get_rank
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Chat, CONV_VISION

# imports modules for registration
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *

from nuscenes.utils.data_classes import PointCloud, LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix

def parse_args():
    parser = argparse.ArgumentParser(description="Demo")
    parser.add_argument("--cfg-path", required=True, default="eval_configs/minigpt4_eval.yaml", help="path to configuration file.")
    parser.add_argument("--gpu-id", type=int, default=0, help="specify the gpu to load the model.")
    parser.add_argument("--sam-checkpoint", type=str, default="sam_vit_h_4b8939.pth", help="path to sam weights.")
    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )
    args = parser.parse_args()
    return args

import numpy as np
from scipy.ndimage.measurements import label

def mean_blob_coordinates(image, x, y):
    image = (image[:,:,2]>0)
    labeled_array, num_features = label(image)
    blob_label = labeled_array[y, x]
    blob_coordinates = np.argwhere(labeled_array == blob_label)
    mean_x = np.mean(blob_coordinates[:, 1]).astype(np.int8)
    mean_y = np.mean(blob_coordinates[:, 0]).astype(np.int8)
    return mean_x, mean_y


def setup_seeds(config):
    seed = config.run_cfg.seed + get_rank()

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    cudnn.benchmark = False
    cudnn.deterministic = True

def distance(x1, y1, z1, x2, y2, z2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + 0 * (z1 - z2) ** 2) ** ( 1 / 2)
 
# Function to calculate K closest points
def kClosest(points, target, K):
    pts = []
    n = len(points)
    d = []

    for i in range(n):
        d.append({
            "first": distance(points[i][0], points[i][1], points[i][2], target[0], target[1], target[2]),
            "second": i
        })
     
    d = sorted(d, key=lambda l:l["first"])
 
    for i in range(K):
        pt = []
        pt.append(points[d[i]["second"]][0])
        pt.append(points[d[i]["second"]][1])
        pt.append(points[d[i]["second"]][2])
        pt.append(points[d[i]["second"]][3])
        pts.append(pt)
 
    return pts

def get_image_projected_points(cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, arr):
    """
        arr: points to be back-projected: (K closest points)
    """
    cam_imgs = [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right]
    cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']
    calibration_data = {
        'CAM_FRONT_LEFT':{
            'translation': [-1.57525595, -0.50051938, -1.50696033],
            'rotation': [[ 0.82254604, -0.56868433, -0.00401771], [ 0.00647832,  0.01643407, -0.99984396], [ 0.56866162,  0.82239167,  0.01720189]],
            'camera_intrinsic': [[1257.8625342125129, 0.0, 827.2410631095686], [0.0, 1257.8625342125129, 450.915498205774], [0.0, 0.0, 1.0]]
        },
        'CAM_FRONT':{
            'translation': [-1.72200568, -0.00475453, -1.49491292],
            'rotation': [[ 0.01026021, -0.99987258, -0.01222952], [ 0.00843345,  0.01231626, -0.99988859], [ 0.9999118 ,  0.01015593,  0.00855874]],
            'camera_intrinsic': [[1252.8131021185304, 0.0, 826.588114781398], [0.0, 1252.8131021185304, 469.9846626224581], [0.0, 0.0, 1.0]]
        },
        'CAM_FRONT_RIGHT':{
            'translation': [-1.58082566,  0.49907871, -1.51749368],
            'rotation': [[-0.84397973, -0.53614138, -0.01583178], [ 0.01645551,  0.00362107, -0.99985804], [ 0.5361226 , -0.84412044,  0.00576637]],
            'camera_intrinsic': [[1256.7485116440405, 0.0, 817.7887570959712], [0.0, 1256.7485116440403, 451.9541780095127], [0.0, 0.0, 1.0]]
        },
        'CAM_BACK_LEFT':{
            'translation': [-1.035691  , -0.48479503, -1.59097015],
            'rotation': [[ 0.94776036,  0.31896113,  0.00375564], [ 0.00866572, -0.0139763 , -0.99986478], [-0.31886551,  0.94766474, -0.01601021]],
            'camera_intrinsic': [[1256.7414812095406, 0.0, 792.1125740759628], [0.0, 1256.7414812095406, 492.7757465151356], [0.0, 0.0, 1.0]]
        },
        'CAM_BACK':{
            'translation': [-0.02832603, -0.00345137, -1.57910346],
            'rotation': [[ 0.00242171,  0.99998907, -0.00400023], [-0.01675361, -0.00395911, -0.99985181], [-0.99985672,  0.00248837,  0.01674384]],
            'camera_intrinsic': [[809.2209905677063, 0.0, 829.2196003259838], [0.0, 809.2209905677063, 481.77842384512485], [0.0, 0.0, 1.0]]
        },
        'CAM_BACK_RIGHT':{
            'translation': [-1.0148781 ,  0.48056822, -1.56239545],
            'rotation': [[-0.93477554,  0.35507456, -0.01080503], [ 0.01587584,  0.0113705 , -0.99980932], [-0.35488399, -0.93476883, -0.01626597]],
            'camera_intrinsic': [[1259.5137405846733, 0.0, 807.2529053838625], [0.0, 1259.5137405846733, 501.19579884916527], [0.0, 0.0, 1.0]]
        }
    }

    min_dist = 1.0
    flag = 0
    cam_img = None
    for cam_ind, cam_key in enumerate(cam_keys):
        barr = np.copy(arr)
        ppc = LidarPointCloud(barr.T);
        ppc.translate(np.array(calibration_data[cam_key]['translation']))
        ppc.rotate(np.array(calibration_data[cam_key]['rotation']))
        ddepths_img = ppc.points[2, :]
        points = view_points(ppc.points[:3, :], np.array(calibration_data[cam_key]['camera_intrinsic']), normalize=True)
        mask = np.ones(ddepths_img.shape[0], dtype=bool)
        mask = np.logical_and(mask, ddepths_img > min_dist)
        mask = np.logical_and(mask, points[0, :] > 1)
        mask = np.logical_and(mask, points[0, :] < 1600 - 1)
        mask = np.logical_and(mask, points[1, :] > 1)
        mask = np.logical_and(mask, points[1, :] < 900 - 1)

        if mask.sum() > 0:
            # found
            flag = 1
            points = points[:, mask]
            ddepths_img = ddepths_img[mask]
            cam_img = cam_imgs[cam_ind]
            break

        if flag == 0:
            # no point able to back-project, just use front cam
            cam_img = cam_front
    return cam_img

# ========================================
#             Model Initialization
# ========================================

args = parse_args()
cfg = Config(args)

model_config = cfg.model_cfg
model_config.device_8bit = args.gpu_id
model_cls = registry.get_model_class(model_config.arch)
model = model_cls.from_config(model_config).to('cuda:{}'.format(args.gpu_id))

vis_processor_cfg = cfg.datasets_cfg.cc_sbu_align.vis_processor.train
vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
chat = Chat(model, vis_processor, device='cuda:{}'.format(args.gpu_id))

# ========================================
#             Sam Initialization
# ========================================

from segment_anything import sam_model_registry, SamPredictor

model_type = "vit_h"
device = "cuda"

sam = sam_model_registry[model_type](checkpoint=args.sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)

# ========================================
#             Gradio Setting
# ========================================

def gradio_reset(chat_state, img_list):
    if chat_state is not None:
        chat_state.messages = []
    if img_list is not None:
        img_list = []
    return None, gr.update(value=None, interactive=True), gr.update(placeholder='Please upload your image first', interactive=False),gr.update(value="Upload & Start Chat", interactive=True), chat_state, img_list

def upload_img(gr_img, text_input, chat_state):
    if gr_img is None:
        return None, None, gr.update(interactive=True), chat_state, None
    chat_state = CONV_VISION.copy()
    img_list = []
    llm_message = chat.upload_img(gr_img, chat_state, img_list)
    return gr.update(interactive=False), gr.update(interactive=True, placeholder='Type and press Enter'), gr.update(value="Start Chatting", interactive=False), chat_state, img_list

def gradio_ask(user_message, chatbot, chat_state):
    if len(user_message) == 0:
        return gr.update(interactive=True, placeholder='Input should not be empty!'), chatbot, chat_state
    chat.ask(user_message, chat_state)
    chatbot = chatbot + [[user_message, None]]
    return '', chatbot, chat_state


def gradio_answer(chatbot, chat_state, img_list, num_beams, temperature):
    llm_message = chat.answer(conv=chat_state,
                              img_list=img_list,
                              num_beams=num_beams,
                              temperature=temperature,
                              max_new_tokens=300,
                              max_length=2000)[0]
    chatbot[-1][1] = llm_message
    return chatbot, chat_state, img_list

title = """<h1 align="center">Talk2BEV: Demo</h1>"""
description = """<h3>Upload BEV, perspective images, LiDAR point cloud and start chatting!</h3>"""
article = """<p><a href='https://minigpt-4.github.io'><img src='https://img.shields.io/badge/Project-Page-Green'></a></p><p><a href='https://github.com/Vision-CAIR/MiniGPT-4'><img src='https://img.shields.io/badge/Github-Code-blue'></a></p><p><a href='https://raw.githubusercontent.com/Vision-CAIR/MiniGPT-4/main/MiniGPT_4.pdf'><img src='https://img.shields.io/badge/Paper-PDF-red'></a></p>
"""

#TODO show examples below

def find_bounding_box(image):
    # Find the coordinates of non-zero (foreground) pixels along each channel
    foreground_pixels = np.array(np.where(np.any(image != 0, axis=2)))

    # Calculate the bounding box coordinates
    min_y, min_x = np.min(foreground_pixels, axis=1)
    max_y, max_x = np.max(foreground_pixels, axis=1)

    return min_y, min_x, max_y, max_x

def crop_around_bounding_box(image):
    min_y, min_x, max_y, max_x = find_bounding_box(image)

    # Crop the image using the bounding box coordinates
    cropped_image = image[min_y:max_y+1, min_x:max_x+1, :]

    return cropped_image

with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown(description)
    # gr.Markdown(article)

    tolerance = gr.Slider(label="Tolerance", info="How different colors can be in a segment.", minimum=0, maximum=256*3, value=50)

    with gr.Row():
        cam_left = gr.Image(label="CAM_LEFT - Upload")
        cam_front = gr.Image(label="CAM_FRONT - Upload")
        cam_right = gr.Image(label="CAM_RIGHT - Upload")
    
    with gr.Row():
        cam_rear_left = gr.Image(label="CAM_REAR_LEFT - Upload")
        cam_rear = gr.Image(label="CAM_REAR - Upload")
        cam_rear_right = gr.Image(label="CAM_REAR_RIGHT - Upload")

    with gr.Row():
        bev_img = gr.Image(label="BEV Image - Click a segment")
        image = gr.Image(label="Selected Segment", type="pil")                    
        lidar_data = gr.File(label="LiDAR Data - Upload a file")

    with gr.Row():
        upload_button = gr.Button(value="Submit Mask & Start Chat", interactive=True, variant="primary")
        clear = gr.Button("Clear All")

    with gr.Row():
        num_beams = gr.Slider(
            minimum=1,
            maximum=10,
            value=1,
            step=1,
            interactive=True,
            label="beam search numbers)",
        )
        
        temperature = gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=1.0,
            step=0.1,
            interactive=True,
            label="Temperature",
        )

    with gr.Row():
        with gr.Column():
            chat_state = gr.State()
            img_list = gr.State()
            chatbot = gr.Chatbot(label='MiniGPT-4')
            text_input = gr.Textbox(label='User', placeholder='Please upload your image first', interactive=False)

    def get_select_coords(cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, bev, lidar, tolerance, evt: gr.SelectData):
        lidardata = np.load(lidar.name)
        # import pdb; pdb.set_trace()

        labels_allowed_seg = [14, 15, 16, 17, 18, 19, 20, 21, 22, 23]
        labels_allowed_hdmap = [24, 25, 26, 27]
        if bev[evt.index[1], evt.index[0]][0] == 0 and bev[evt.index[1], evt.index[0]][1] == 0 and bev[evt.index[1], evt.index[0]][2] == 255:
            labels_allowed = labels_allowed_seg
        else:
            labels_allowed = labels_allowed_hdmap
        input_point = np.array([[evt.index[0], evt.index[1]]])
        input_label = np.array([1])

        target = np.array([(evt.index[1] - 100)/2, (evt.index[0] - 100)/2, 0])
        pts = []
        for pt in lidardata.T:
            if pt[-1] in labels_allowed:
                pts.append(pt)

        arr = np.array(kClosest(np.array(pts), target, 1))
    
        # CAM_BACK
        translation = np.array([-0.02832603, -0.00345137, -1.57910346])
        rotation = np.array([[ 0.00242171,  0.99998907, -0.00400023], [-0.01675361, -0.00395911, -0.99985181], [-0.99985672,  0.00248837,  0.01674384]])
        camera_intrinsic = np.array([[809.2209905677063, 0.0, 829.2196003259838], [0.0, 809.2209905677063, 481.77842384512485], [0.0, 0.0, 1.0]])

        cam_imgs = [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right]
        cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']
        calibration_data = {
            'CAM_FRONT_LEFT':{
                'translation': [-1.57525595, -0.50051938, -1.50696033],
                'rotation': [[ 0.82254604, -0.56868433, -0.00401771], [ 0.00647832,  0.01643407, -0.99984396], [ 0.56866162,  0.82239167,  0.01720189]],
                'camera_intrinsic': [[1257.8625342125129, 0.0, 827.2410631095686], [0.0, 1257.8625342125129, 450.915498205774], [0.0, 0.0, 1.0]]
            },
            'CAM_FRONT':{
                'translation': [-1.72200568, -0.00475453, -1.49491292],
                'rotation': [[ 0.01026021, -0.99987258, -0.01222952], [ 0.00843345,  0.01231626, -0.99988859], [ 0.9999118 ,  0.01015593,  0.00855874]],
                'camera_intrinsic': [[1252.8131021185304, 0.0, 826.588114781398], [0.0, 1252.8131021185304, 469.9846626224581], [0.0, 0.0, 1.0]]
            },
            'CAM_FRONT_RIGHT':{
                'translation': [-1.58082566,  0.49907871, -1.51749368],
                'rotation': [[-0.84397973, -0.53614138, -0.01583178], [ 0.01645551,  0.00362107, -0.99985804], [ 0.5361226 , -0.84412044,  0.00576637]],
                'camera_intrinsic': [[1256.7485116440405, 0.0, 817.7887570959712], [0.0, 1256.7485116440403, 451.9541780095127], [0.0, 0.0, 1.0]]
            },
            'CAM_BACK_LEFT':{
                'translation': [-1.035691  , -0.48479503, -1.59097015],
                'rotation': [[ 0.94776036,  0.31896113,  0.00375564], [ 0.00866572, -0.0139763 , -0.99986478], [-0.31886551,  0.94766474, -0.01601021]],
                'camera_intrinsic': [[1256.7414812095406, 0.0, 792.1125740759628], [0.0, 1256.7414812095406, 492.7757465151356], [0.0, 0.0, 1.0]]
            },
            'CAM_BACK':{
                'translation': [-0.02832603, -0.00345137, -1.57910346],
                'rotation': [[ 0.00242171,  0.99998907, -0.00400023], [-0.01675361, -0.00395911, -0.99985181], [-0.99985672,  0.00248837,  0.01674384]],
                'camera_intrinsic': [[809.2209905677063, 0.0, 829.2196003259838], [0.0, 809.2209905677063, 481.77842384512485], [0.0, 0.0, 1.0]]
            },
            'CAM_BACK_RIGHT':{
                'translation': [-1.0148781 ,  0.48056822, -1.56239545],
                'rotation': [[-0.93477554,  0.35507456, -0.01080503], [ 0.01587584,  0.0113705 , -0.99980932], [-0.35488399, -0.93476883, -0.01626597]],
                'camera_intrinsic': [[1259.5137405846733, 0.0, 807.2529053838625], [0.0, 1259.5137405846733, 501.19579884916527], [0.0, 0.0, 1.0]]
            }
        }

        min_dist = 1.0
        flag = 0
        cam_img = None
        for cam_ind, cam_key in enumerate(cam_keys):
            barr = np.copy(arr)
            ppc = LidarPointCloud(barr.T);
            ppc.translate(np.array(calibration_data[cam_key]['translation']))
            ppc.rotate(np.array(calibration_data[cam_key]['rotation']))
            ddepths_img = ppc.points[2, :]
            points = view_points(ppc.points[:3, :], np.array(calibration_data[cam_key]['camera_intrinsic']), normalize=True)
            mask = np.ones(ddepths_img.shape[0], dtype=bool)
            mask = np.logical_and(mask, ddepths_img > min_dist)
            mask = np.logical_and(mask, points[0, :] > 1)
            mask = np.logical_and(mask, points[0, :] < 1600 - 1)
            mask = np.logical_and(mask, points[1, :] > 1)
            mask = np.logical_and(mask, points[1, :] < 900 - 1)

            if mask.sum() > 0:
                # found
                flag = 1
                points = points[:, mask]
                ddepths_img = ddepths_img[mask]
                cam_img = cam_imgs[cam_ind]
                break

        if flag == 0:
            # no point able to back-project, just use front cam
            cam_img = cam_front

        predictor.set_image(cam_img)
        input_point = np.array(points.astype(np.int32).T[:, :2])
        input_label = np.array([1 for i in range(len(points[0]))])
        masks, scores, logits = predictor.predict(point_coords=input_point,point_labels=input_label,multimask_output=True)
        idxs = (masks[-1].astype(np.uint8) * 200) > 0
        img_copy = np.copy(cam_img)
        img_copy[~idxs] = 0
        img_copy = crop_around_bounding_box(img_copy)
        return img_copy

    bev_img.select(get_select_coords, [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, bev_img, lidar_data, tolerance], image)
    
    upload_button.click(upload_img, [image, text_input, chat_state], [image, text_input, upload_button, chat_state, img_list])
    
    text_input.submit(gradio_ask, [text_input, chatbot, chat_state], [text_input, chatbot, chat_state]).then(
        gradio_answer, [chatbot, chat_state, img_list, num_beams, temperature], [chatbot, chat_state, img_list]
    )
    clear.click(gradio_reset, [chat_state, img_list], [chatbot, image, text_input, upload_button, chat_state, img_list], queue=False)

demo.launch(share=True, enable_queue=True)

click2chat/click2chat_minigpt4.py
llmbev/talk2bev
from argparse import ArgumentParser
from PIL import Image
import torch
import os
import matplotlib as mpl
import torch.utils.data
import numpy as np
import json
from tqdm import tqdm
from nuscenes.nuscenes import NuScenes
from matplotlib import pyplot as plt
import cv2
from stp3.datas.reducedNuscenesData import FuturePredictionDataset
from stp3.trainer import TrainingModule
import argparse
import os

import numpy as np
import torch

from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from llava.conversation import conv_templates, SeparatorStyle
from llava.model.builder import load_pretrained_model
from llava.mm_utils import tokenizer_image_token, get_model_name_from_path, KeywordsStoppingCriteria

from transformers import TextStreamer
# Use a pipeline as a high-level helper
from transformers import pipeline

from PIL import Image

from minigpt4.common.config import Config
from minigpt4.common.dist_utils import get_rank
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Chat, CONV_VISION

# imports modules for registration
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *

from nuscenes.utils.data_classes import LidarPointCloud
from nuscenes.utils.geometry_utils import view_points

# ========================================
#             SAM Initialization
# ========================================
def init_sam(model_type = "vit_h", checkpoint="/home/t1/vikrant.dewangan/llm-bev/MiniGPT-4/sam_vit_h_4b8939.pth", device="cuda:0"):
    from segment_anything import sam_model_registry, SamPredictor


    sam = sam_model_registry[model_type](checkpoint="/home/t1/vikrant.dewangan/llm-bev/MiniGPT-4/sam_vit_h_4b8939.pth")
    sam.to(device=device)
    predictor = SamPredictor(sam)
    return predictor
# ========================================

def find_bounding_box(image):
    # Find the coordinates of non-zero (foreground) pixels along each channel
    foreground_pixels = np.array(np.where(np.any(image != 0, axis=2)))

    # Calculate the bounding box coordinates
    min_y, min_x = np.min(foreground_pixels, axis=1)
    max_y, max_x = np.max(foreground_pixels, axis=1)

    return min_y, min_x, max_y, max_x

def find_bounding_box_masks(image):
    # Find the coordinates of non-zero (foreground) pixels along each channel
    foreground_pixels = np.array(np.where(image != 0))

    # Calculate the bounding box coordinates
    min_y, min_x = np.min(foreground_pixels, axis=1)
    max_y, max_x = np.max(foreground_pixels, axis=1)

    return min_y, min_x, max_y, max_x

def crop_around_bounding_box(image, masks, get_black=False):
    if get_black:
        min_y, min_x, max_y, max_x = find_bounding_box(masks)
    else:
        min_y, min_x, max_y, max_x = find_bounding_box_masks(masks)

    # Crop the image using the bounding box coordinates
    cropped_image = image[min_y:max_y+1, min_x:max_x+1, :]

    return cropped_image

def kClosest(points, target, K):
    # Convert the points and target lists to NumPy arrays
    points = np.array(points)
    target = np.array([target])
    # Calculate the distance between each point and the target using NumPy broadcasting
    distances = np.sqrt(np.sum((points[:, :2] - target[:, :2])**2, axis=1))
    # Get the indices of the K closest points using argsort
    closest_indices = np.argsort(distances)[:K]
    # Get the K closest points
    closest_points = points[closest_indices]
    return closest_points.tolist()

def get_image_projection(cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right, arr):
    cam_imgs = [cam_left, cam_front, cam_right, cam_rear_left, cam_rear, cam_rear_right]
    cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']
    calibration_data = json.load(open("calibration.json"))
    min_dist = 1.0
    flag = 0
    cam_img = None
    matched_cam = None
    matched_points = None
    for cam_ind, cam_key in enumerate(cam_keys):
        barr = np.copy(arr)
        ppc = LidarPointCloud(barr.T);
        ppc.translate(np.array(calibration_data[cam_key]['translation']))
        ppc.rotate(np.array(calibration_data[cam_key]['rotation']))
        ddepths_img = ppc.points[2, :]
        points = view_points(ppc.points[:3, :], np.array(calibration_data[cam_key]['camera_intrinsic']), normalize=True)
        
        mask = np.ones(ddepths_img.shape[0], dtype=bool)
        mask = np.logical_and(mask, ddepths_img > min_dist)
        mask = np.logical_and(mask, points[0, :] > 1)
        mask = np.logical_and(mask, points[0, :] < 1600 - 1)
        mask = np.logical_and(mask, points[1, :] > 1)
        mask = np.logical_and(mask, points[1, :] < 900 - 1)

        if mask.sum() > 0:
            # found
            flag = 1
            points = points[:, mask]
            ddepths_img = ddepths_img[mask]
            cam_img = cam_imgs[cam_ind]
            matched_points = points
            matched_cam = cam_key
            break

    if flag == 0:
        # no point able to back-project, just use front cam
        matched_points = [150, 100] # hard-coded
        cam_img = cam_front
    return cam_img, points, matched_points, matched_cam

def extract_mask(predictor, cam_img, points):
    """
        arr: points to be back-projected: (K closest points)
    """
    predictor.set_image(cam_img)
    input_point = np.array(points.astype(np.int32).T[:, :2])
    input_label = np.array([1 for i in range(len(points[0]))])
    masks, scores, logits = predictor.predict(point_coords=input_point,point_labels=input_label,multimask_output=True)
    idxs = (masks[-1].astype(np.uint8) * 200) > 0
    img_copy = np.copy(cam_img)
    # img_copy[~idxs] = 0
    img_copy = crop_around_bounding_box(img_copy, masks[-1])
    return img_copy, masks

device = torch.device('cuda:0')
device2 = torch.device('cuda:1')

predictor = init_sam(device=device2)
print("Initializaed SAM")
print('Initialization Finished')

to_visualize = True
compute_losses = False

cam_names = ["CAM_FRONT_LEFT", "CAM_FRONT", "CAM_FRONT_RIGHT", "CAM_BACK_LEFT", "CAM_BACK", "CAM_BACK_RIGHT"]

def eval(args):
    data_path = args.data_path
    save_path = args.save_path
    trainer = TrainingModule.load_from_checkpoint(args.checkpoint, strict=True)
    trainer.cfg.INSTANCE_FLOW.ENABLED = True
    trainer.cfg.INSTANCE_SEG.ENABLED = True    
    print(f'Loaded weights from \n {args.checkpoint}')
    trainer.eval()

    trainer.to(device)
    model = trainer.model

    x = np.linspace(0, 19, 20)
    y = np.linspace(-3, 3, 7)
    xv, yv = np.meshgrid(x, y)
    grid = np.dstack((xv, yv))
    grid = np.concatenate((grid), axis=0)

    cfg = model.cfg
    cfg.GPUS = "[0]"
    cfg.BATCHSIZE = 1
    cfg.LIFT.GT_DEPTH = False
    cfg.DATASET.DATAROOT = data_path
    cfg.DATASET.MAP_FOLDER = data_path
    cfg.DEBUG = True

    nusc = NuScenes(version='v1.0-{}'.format("trainval"), dataroot=data_path, verbose=False)
    valdata = FuturePredictionDataset(nusc, 0, cfg)
    valdata.indices = valdata.indices
    valloader = torch.utils.data.DataLoader(
        valdata, batch_size=cfg.BATCHSIZE, shuffle=False, num_workers=0, pin_memory=True, drop_last=False
    )

    for index, batch in enumerate(tqdm(valloader)):
        cur_scene_token = batch['scene_token'][0]

        image = batch['image']
        intrinsics = batch['intrinsics']
        extrinsics = batch['extrinsics']
        future_egomotion = batch['future_egomotion']

        output = model(
            image.to(device), intrinsics.to(device), extrinsics.to(device), future_egomotion.to(device)
        )
        os.makedirs(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index)), exist_ok=True)
        for idx in range(6):
            img = batch['unnormalized_images'][0,2,idx].numpy().astype(np.uint8)
            Image.fromarray(img).save(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index), cam_names[idx] + ".png"))

        arr = np.zeros((200, 200, 3))
        whe = np.where(batch['hdmap'].squeeze()[2,1] > 0)
        arr[whe[0], whe[1]] = np.array([255,255,255])

        if args.bev == "pred":
            pred = torch.argmax(output['segmentation'], dim=2).squeeze()[2].cpu().numpy()
            whe = np.where(pred > 0)
        else:
            whe = np.where(batch['hdmap'].squeeze()[2,1] > 0)

        arr[whe[0], whe[1]] = np.array([0,0,255])
        barr = np.copy(arr)
        Image.fromarray(arr.astype(np.uint8)).save(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index), "pred_bev.png"))
        pts = []
        lidardata = batch['point_clouds'][2].squeeze()[2].numpy()
        for ptind, pt in enumerate(lidardata.T):
            pts.append(pt)

        bev_2d = np.logical_and(barr[:,:,2]>0,barr[:,:,0]==0)
        labels, pts_ = cv2.connectedComponents(bev_2d.astype(np.uint8))
        matched_imgs = []
        cam_imgs = []
        for cin in range(6):
            cam_imgs.append(batch['unnormalized_images'][0,2,cin].numpy())

        objects_json = []
        for idx in range(1, labels + 1):
            # Create a JSON object for each component
            x, y = np.where(pts_ == idx)
            bevy, bevx = np.where(pts_ == idx)
            bevy = 200-bevy

            obj = {
                "object_id": idx,
                "bev_centroid": [(np.mean(bevx).astype(np.int) - 100)/2, (np.mean(bevy).astype(np.int) - 100)/2],
                "matched_coords": [x.tolist(), y.tolist()],
                "bev_area": len(x)/4,
            }
            target_x, target_y = np.mean(x).astype(np.uint8), np.mean(y).astype(np.uint8)

            min_ann_dist = 1e11
            best_ann = {}
            for _, anns in enumerate(batch["categories_map"][2]):
                dist = np.linalg.norm(anns[1][0][:2] - np.array([(target_x - 100)/2, (target_y - 100)/2]))
                annotation = anns[0]
                if dist < min_ann_dist:
                    min_ann_dist = dist
                    best_ann = annotation
            if min_ann_dist > 5:
                continue
            keys = best_ann.keys()
            for key in keys:
                if type(best_ann[key]) == torch.Tensor:
                    best_ann[key] = best_ann[key].tolist()
                    continue
                for itemind, item in enumerate(best_ann[key]):
                    if type(item) == torch.Tensor:
                        best_ann[key][itemind] = item.tolist()
                    elif type(item) == list:
                        for listind in range(len(item)):
                            if type(item[listind]) == torch.Tensor:
                                best_ann[key][itemind][listind] = best_ann[key][itemind][listind].tolist()
            obj["annotation"] = best_ann
            pts = []
            lidardata = batch['point_clouds'][2].squeeze()[2].numpy()
            for ptind, pt in enumerate(lidardata.T):
                if batch['point_clouds_labels'][2].squeeze()[2][ptind] in labels_allowed:
                    pts.append(pt)

            dts = np.array(pts)
            bbox = batch["bottom_corners"][2][best_ann["token"][0]].squeeze().T.numpy()
            
            # filter out points within bbox                
            min_x = np.min(bbox[:, 0])
            max_x = np.max(bbox[:, 0])
            min_y = np.min(bbox[:, 1])
            max_y = np.max(bbox[:, 1])
            min_z = np.min(bbox[:, 2])
            max_z = np.max(bbox[:, 2])
            mask_x = (dts[:, 0] >= min_x) & (dts[:, 0] <= max_x)
            mask_y = (dts[:, 1] >= min_y) & (dts[:, 1] <= max_y)
            mask_z = (dts[:, 2] >= min_z) & (dts[:, 2] <= max_z)
            mask = mask_x & mask_y & mask_z
            indices = np.where(mask)
            dts = dts[indices]
            max_dist = max(abs(target_x - 100), abs(target_y - 100))
            elem = int(min(np.ceil(((max_dist + 50)/100) * (len(dts) - 1)), len(dts) - 1))
            z_filter = sorted(dts[:, 2])[elem]
            minind = np.argmin(np.linalg.norm(dts[:, :3] - np.array([[0.0, 0.0, z_filter]]), axis=1))
            arr = np.expand_dims(dts[minind], axis=1).T

            cam_img, points, matched_point, matched_cam = get_image_projection(batch['unnormalized_images'][0,2,0].numpy(), 
                                                                        batch['unnormalized_images'][0,2,1].numpy(), 
                                                                        batch['unnormalized_images'][0,2,2].numpy(), 
                                                                        batch['unnormalized_images'][0,2,3].numpy(), 
                                                                        batch['unnormalized_images'][0,2,4].numpy(), 
                                                                        batch['unnormalized_images'][0,2,5].numpy(), arr)
            if matched_cam == None:
                continue
            obj["matched_cam"] = matched_cam
            obj["matched_point"] = matched_point.tolist()
            img_cropped, masks = extract_mask(predictor, cam_img, points)
            matched_imgs.append(img_cropped)
            objects_json.append(obj)

        for matched_img_ind, matched_img in enumerate(matched_imgs):
            np.save(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index), "crops", f"{matched_img_ind + 1}_matched_img_pred.npy"), matched_img)
        for cimg_ind, cimg in enumerate(cam_imgs):
            np.save(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index), "cam_imgs", f"{cimg_ind + 1}_cimg.npy"), cimg)
        
        if args.bev == "pred":
            with open(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index), "answer_pred_both.json"), "w") as f:
                json.dump(objects_json, f, indent=4)
        else:
            with open(os.path.join(save_path, str(cur_scene_token[0]) + "_" + "{0:0=6d}".format(index), "answer_gt.json"), "w") as f:
                json.dump(objects_json, f, indent=4)

        print("DONE SAVED");print();print();print();print();

if __name__ == '__main__':
    parser = ArgumentParser(description='STP3 evaluation')
    parser.add_argument('--checkpoint', default='last.ckpt', type=str, help='path to checkpoint')
    parser.add_argument('--dataroot', default=None, type=str)
    parser.add_argument('--bev', default="pred", help='use pred or gt for bev type', type=str)
    parser.add_argument('--data_path', required=True, help='path to nuscenes v1.0 trainval dataset', type=str)
    parser.add_argument('--save_path', default="./", help='saving data directory', type=str)
    args = parser.parse_args()
    eval(args)

data/generate_base.py
llmbev/talk2bev
import os
import numpy as np
import json
from tqdm import tqdm
import argparse
from utils import init_minigp4, minigpt4_inference, init_blip2, init_instructblip2, instructblip2_inference

cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']

def run_minigpt4(index_start, index_end, gpu, model_path, save_path = "/raid/t1/scratch/vikrant.dewangan/datas", json_name="answer_gt.json"):
    json_list = os.listdir(save_path)
    num_to_json = {}
    for ind, objp in enumerate(json_list):
        if os.path.exists(os.path.join(save_path, objp, "answer_gt.json")):
            num_to_json[int(objp.split("_")[1])] = objp

    chat = init_minigp4()

    for ind in tqdm(range(index_start, index_end)):
        objs = json.load(open(os.path.join(save_path, num_to_json[ind], json_name)))
        imc = {}
        print(os.path.join(save_path, num_to_json[ind]))
        for j in range(6):
            img = np.load(os.path.join(save_path, num_to_json[ind], "crop_imgs", f"{j + 1}_cimg.npy"))
            user_message = "Describe this image."
            answer = minigpt4_inference(chat, img, user_message)
            imc[cam_keys[j]] = {}
            imc[cam_keys[j]]["description"] = answer

            user_message = "Is there anything unusual about this image?"
            answer = minigpt4_inference(chat, img, user_message)
            imc[cam_keys[j]]["unusual"] = answer

            user_message = "Describe the weather in this image. Is it day or night?"
            answer = minigpt4_inference(chat, img, user_message)
            imc[cam_keys[j]]["weather"] = answer
        
        print("Image Captioning done")
        for indd, obj in enumerate(objs):
            if "pred" in json_name:
                img = np.load(os.path.join(save_path, num_to_json[ind], "crops", f"{indd + 1}_matched_img_pred.npy"))
            else:
                img = np.load(os.path.join(save_path, num_to_json[ind], "crops", f"{indd + 1}_matched_img.npy"))
            user_message = "Describe the central object in this image."
            answer = minigpt4_inference(chat, img, user_message)
            objs[indd]["minigpt4_crop_brief"] = answer

            user_message = "Is the object's indicator on? If yes, what direction does it want to turn. Please answer carefully."
            answer = minigpt4_inference(chat, img, user_message)
            objs[indd]["minigpt4_crop_lights1"] = answer

            user_message = "In this image, are the object's rear lights closer  to the viewer or forward lights? Please answer in one word - forward/rear."
            answer = minigpt4_inference(chat, img, user_message)
            objs[indd]["minigpt4_crop_lights2"] = answer.lower()

            user_message = "Is there any text written on the object? Please look carefully, and describe it."
            answer = minigpt4_inference(chat, img, user_message)
            objs[indd]["minigpt4_crop_text"] = answer

            objs[indd]["minigpt4_bg_description"] = imc[objs[indd]["matched_cam"]]["description"]
            objs[indd]["minigpt4_bg_unusual"] = imc[objs[indd]["matched_cam"]]["unusual"]
            objs[indd]["minigpt4_bg_weather"] = imc[objs[indd]["matched_cam"]]["weather"]

        print("Object Captioning Done")
        save_path2 = "/raid/t1/scratch/vikrant.dewangan/datas5"
        os.makedirs(os.path.join(save_path2, num_to_json[ind]), exist_ok=True)
        with open(os.path.join(save_path2, num_to_json[ind], json_name), "w") as f:
            json.dump(objs, f, indent=4)
        print("DONE SAVED", os.path.join(save_path2, num_to_json[ind]))

def run_instructblip2(index_start, index_end, gpu, model_path, save_path = "/raid/t1/scratch/vikrant.dewangan/datas", json_name="answer_gt.json"):
    json_list = os.listdir(save_path)
    num_to_json = {}
    for ind, objp in enumerate(json_list):
        if os.path.exists(os.path.join(save_path, objp, "answer_gt.json")):
            num_to_json[int(objp.split("_")[1])] = objp

    model_instructblip, vis_processors = init_instructblip2(model_name = "blip2_vicuna_instruct", device="cuda:0")

    for ind in tqdm(range(index_start, index_end)):
        objs = json.load(open(os.path.join(save_path, num_to_json[ind], json_name)))
        imc = {}
        for j in range(6):
            img = np.load(os.path.join(save_path, num_to_json[ind], "crop_imgs", f"{j + 1}_cimg.npy"))
            user_message = "Describe this image."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            imc[cam_keys[j]] = {}
            imc[cam_keys[j]]["description"] = answer

            user_message = "Is there anything unusual about this image?"
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            imc[cam_keys[j]]["unusual"] = answer

            user_message = "Describe the weather in this image. Is it day or night?"
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            imc[cam_keys[j]]["weather"] = answer
        
        print("Image Captioning done")
        for indd, obj in enumerate(objs):
            if "pred" in json_name:
                img = np.load(os.path.join(save_path, num_to_json[ind], "crops", f"{indd + 1}_matched_img_pred.npy"))
            else:
                img = np.load(os.path.join(save_path, num_to_json[ind], "crops", f"{indd + 1}_matched_img.npy"))
            user_message = "Describe the central object in this image."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["instructblip2_crop_brief"] = answer

            user_message = "Is the object's indicator on? If yes, what direction does it want to turn. Please answer carefully."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["instructblip2_crop_lights1"] = answer

            user_message = "In this image, are the object's rear lights closer  to the viewer or forward lights? Please answer in one word - forward/rear."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["instructblip2_crop_lights2"] = answer.lower()

            user_message = "Is there any text written on the object? Please look carefully, and describe it."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["instructblip2_crop_text"] = answer

            objs[indd]["instructblip2_bg_description"] = imc[objs[indd]["matched_cam"]]["description"]
            objs[indd]["instructblip2_bg_unusual"] = imc[objs[indd]["matched_cam"]]["unusual"]
            objs[indd]["instructblip2_bg_weather"] = imc[objs[indd]["matched_cam"]]["weather"]

        print("Object Captioning Done")
        save_path2 = "/raid/t1/scratch/vikrant.dewangan/datas4"
        os.makedirs(os.path.join(save_path2, num_to_json[ind]), exist_ok=True)
        with open(os.path.join(save_path2, num_to_json[ind], json_name), "w") as f:
            json.dump(objs, f, indent=4)
        print("DONE SAVED", os.path.join(save_path2, num_to_json[ind]))

def run_blip2(index_start, index_end, gpu, model_path, save_path = "/raid/t1/scratch/vikrant.dewangan/datas", json_name="answer_gt.json"):
    json_list = os.listdir(save_path)
    num_to_json = {}
    for ind, objp in enumerate(json_list):
        if os.path.exists(os.path.join(save_path, objp, "answer_gt.json")):
            num_to_json[int(objp.split("_")[1])] = objp

    model_instructblip, vis_processors = init_blip2(model_name = "blip2_t5", device="cuda:0", model_type="pretrain_flant5xxl")

    for ind in tqdm(range(index_start, index_end)):
        objs = json.load(open(os.path.join(save_path, num_to_json[ind], json_name)))
        imc = {}
        for j in range(6):
            img = np.load(os.path.join(save_path, num_to_json[ind], "crop_imgs", f"{j + 1}_cimg.npy"))
            user_message = "Describe this image."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            imc[cam_keys[j]] = {}
            imc[cam_keys[j]]["description"] = answer

            user_message = "Is there anything unusual about this image?"
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            imc[cam_keys[j]]["unusual"] = answer

            user_message = "Describe the weather in this image. Is it day or night?"
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            imc[cam_keys[j]]["weather"] = answer
        
        print("Image Captioning done")
        for indd, obj in enumerate(objs):
            if "pred" in json_name:
                img = np.load(os.path.join(save_path, num_to_json[ind], "crops", f"{indd + 1}_matched_img_pred.npy"))
            else:
                img = np.load(os.path.join(save_path, num_to_json[ind], "crops", f"{indd + 1}_matched_img.npy"))
            user_message = "Describe the central object in this image."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["blip2_crop_brief"] = answer

            user_message = "Is the object's indicator on? If yes, what direction does it want to turn. Please answer carefully."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["blip2_crop_lights1"] = answer

            user_message = "In this image, are the object's rear lights closer  to the viewer or forward lights? Please answer in one word - forward/rear."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["blip2_crop_lights2"] = answer.lower()

            user_message = "Is there any text written on the object? Please look carefully, and describe it."
            answer = instructblip2_inference(model_instructblip, img, vis_processors, device="cuda:0", user_message=user_message)
            objs[indd]["blip2_crop_text"] = answer

            objs[indd]["blip2_bg_description"] = imc[objs[indd]["matched_cam"]]["description"]
            objs[indd]["blip2_bg_unusual"] = imc[objs[indd]["matched_cam"]]["unusual"]
            objs[indd]["blip2_bg_weather"] = imc[objs[indd]["matched_cam"]]["weather"]

        save_path2 = "/raid/t1/scratch/vikrant.dewangan/datas7"
        os.makedirs(os.path.join(save_path2, num_to_json[ind]), exist_ok=True)
        with open(os.path.join(save_path2, num_to_json[ind], json_name), "w") as f:
            json.dump(objs, f, indent=4)
        print("DONE SAVED", os.path.join(save_path2, num_to_json[ind]))

def main():
    parser = argparse.ArgumentParser(description='Generate captions.')    
    parser.add_argument('--model_path', type=str, default="/raid/t1/scratch/vikrant.dewangan/LLaVA/ckpt-old/", help='save path for jsons')
    parser.add_argument('--data_path', type=str, default="/raid/t1/scratch/vikrant.dewangan/datas", help='save path for jsons')
    parser.add_argument('--gpu', type=str, default="cuda:0", help='gpu number')
    parser.add_argument('--bev', type=str, default="pred", help='pred or gt')
    parser.add_argument('--start', type=int, default=0, help='start index')
    parser.add_argument('--end', type=int, default=100, help='end index')
    parser.add_argument('--model', type=str, required=True, help='model name, use minigpt4 or instructblip2 or blip2')

    args = parser.parse_args()
    args.save_path = args.data_path
    if args.bev == "pred":
        args.json_name = "answer_pred_both.json"
    if args.bev == "gt":
        args.json_name = "answer_gt.json"

    if args.model == "minigpt4":
        run_minigpt4(args.start, args.end, args.gpu, args.model_path, args.save_path, json_name=args.json_name)
    elif args.model == "instructblip2":
        run_instructblip2(args.start, args.end, args.gpu, args.model_path, args.save_path, json_name=args.json_name)
    elif args.model == "blip2":
        run_blip2(args.start, args.end, args.gpu, args.model_path, args.save_path, json_name=args.json_name)

if __name__ == "__main__":
    main()
data/generate_captions.py
llmbev/talk2bev
import argparse
from fvcore.common.config import CfgNode as _CfgNode


def convert_to_dict(cfg_node, key_list=[]):
    """Convert a config node to dictionary."""
    _VALID_TYPES = {tuple, list, str, int, float, bool}
    if not isinstance(cfg_node, _CfgNode):
        if type(cfg_node) not in _VALID_TYPES:
            print(
                'Key {} with value {} is not a valid type; valid types: {}'.format(
                    '.'.join(key_list), type(cfg_node), _VALID_TYPES
                ),
            )
        return cfg_node
    else:
        cfg_dict = dict(cfg_node)
        for k, v in cfg_dict.items():
            cfg_dict[k] = convert_to_dict(v, key_list + [k])
        return cfg_dict


class CfgNode(_CfgNode):
    """Remove once https://github.com/rbgirshick/yacs/issues/19 is merged."""

    def convert_to_dict(self):
        return convert_to_dict(self)


CN = CfgNode

_C = CN()
_C.LOG_DIR = 'tensorboard_logs'
_C.TAG = 'default'

_C.GPUS = [0, 1, 2]  # which gpus to use
_C.PRECISION = 32  # 16bit or 32bit
_C.BATCHSIZE = 1
_C.EPOCHS = 20
_C.DEBUG = True
_C.DENOISE = False # to denoise the GT or not
_C.LOGGER = 'wandb'
_C.COMPARE_WITH_ORIG = True # to compare with original or not

_C.N_WORKERS = 0
_C.VIS_INTERVAL = 5000
_C.LOGGING_INTERVAL = 500

_C.PRETRAINED = CN()
_C.PRETRAINED.LOAD_WEIGHTS = True
_C.PRETRAINED.PATH = 'last.ckpt'

_C.DATASET = CN()
_C.DATASET.DATAROOT = '/mnt/e/datasets/nuscenes'
_C.DATASET.VERSION = 'trainval'
_C.DATASET.NAME = 'nuscenes'
_C.DATASET.MAP_FOLDER = "/scratch/tushar.choudhary/forecasting/v1.0-trainval"
_C.DATASET.IGNORE_INDEX = 255  # Ignore index when creating flow/offset labels
_C.DATASET.FILTER_INVISIBLE_VEHICLES = True  # Filter vehicles that are not visible from the cameras
_C.DATASET.SAVE_DIR = 'datas'
_C.DATASET.THRESHOLD = 25

_C.TIME_RECEPTIVE_FIELD = 3  # how many frames of temporal context (1 for single timeframe)
_C.N_FUTURE_FRAMES = 6  # how many time steps into the future to predict

_C.IMAGE = CN()
_C.IMAGE.FINAL_DIM = (224, 480)
_C.IMAGE.RESIZE_SCALE = 0.3
_C.IMAGE.TOP_CROP = 46
_C.IMAGE.ORIGINAL_HEIGHT = 900  # Original input RGB camera height
_C.IMAGE.ORIGINAL_WIDTH = 1600  # Original input RGB camera width
_C.IMAGE.NAMES = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']

_C.LIFT = CN()  # image to BEV lifting
_C.LIFT.X_BOUND = [-50.0, 50.0, 0.5]  # Forward
_C.LIFT.Y_BOUND = [-50.0, 50.0, 0.5]  # Sides
_C.LIFT.Z_BOUND = [-10.0, 10.0, 20.0]  # Height
_C.LIFT.D_BOUND = [2.0, 50.0, 1.0]
_C.LIFT.GT_DEPTH = False
_C.LIFT.DISCOUNT = 0.5

_C.EGO = CN()
_C.EGO.WIDTH = 1.85
_C.EGO.HEIGHT = 4.084

_C.MODEL = CN()

_C.MODEL.ENCODER = CN()
_C.MODEL.ENCODER.DOWNSAMPLE = 8
_C.MODEL.ENCODER.NAME = 'efficientnet-b4'
_C.MODEL.ENCODER.OUT_CHANNELS = 64
_C.MODEL.ENCODER.USE_DEPTH_DISTRIBUTION = True

_C.MODEL.TEMPORAL_MODEL = CN()
_C.MODEL.TEMPORAL_MODEL.NAME = 'temporal_block'  # type of temporal model
_C.MODEL.TEMPORAL_MODEL.START_OUT_CHANNELS = 64
_C.MODEL.TEMPORAL_MODEL.EXTRA_IN_CHANNELS = 0
_C.MODEL.TEMPORAL_MODEL.INBETWEEN_LAYERS = 0
_C.MODEL.TEMPORAL_MODEL.PYRAMID_POOLING = True
_C.MODEL.TEMPORAL_MODEL.INPUT_EGOPOSE = True

_C.MODEL.DISTRIBUTION = CN()
_C.MODEL.DISTRIBUTION.LATENT_DIM = 32
_C.MODEL.DISTRIBUTION.MIN_LOG_SIGMA = -5.0
_C.MODEL.DISTRIBUTION.MAX_LOG_SIGMA = 5.0

_C.MODEL.FUTURE_PRED = CN()
_C.MODEL.FUTURE_PRED.N_GRU_BLOCKS = 2
_C.MODEL.FUTURE_PRED.N_RES_LAYERS = 1
_C.MODEL.FUTURE_PRED.MIXTURE = True

_C.MODEL.DECODER = CN()

_C.MODEL.BN_MOMENTUM = 0.1

_C.SEMANTIC_SEG = CN()

_C.SEMANTIC_SEG.VEHICLE = CN()
_C.SEMANTIC_SEG.VEHICLE.WEIGHTS = [1.0, 2.0]
_C.SEMANTIC_SEG.VEHICLE.USE_TOP_K = True  # backprop only top-k hardest pixels
_C.SEMANTIC_SEG.VEHICLE.TOP_K_RATIO = 0.25

_C.SEMANTIC_SEG.PEDESTRIAN = CN()
_C.SEMANTIC_SEG.PEDESTRIAN.ENABLED = True
_C.SEMANTIC_SEG.PEDESTRIAN.WEIGHTS = [1.0, 10.0]
_C.SEMANTIC_SEG.PEDESTRIAN.USE_TOP_K = True
_C.SEMANTIC_SEG.PEDESTRIAN.TOP_K_RATIO = 0.25

_C.SEMANTIC_SEG.HDMAP = CN()
_C.SEMANTIC_SEG.HDMAP.ENABLED = True
_C.SEMANTIC_SEG.HDMAP.ELEMENTS = ['lane_divider', 'drivable_area']
_C.SEMANTIC_SEG.HDMAP.WEIGHTS = [[1.0, 5.0], [1.0, 1.0]]
_C.SEMANTIC_SEG.HDMAP.TRAIN_WEIGHT = [1, 1]
_C.SEMANTIC_SEG.HDMAP.USE_TOP_K = [True, False]
_C.SEMANTIC_SEG.HDMAP.TOP_K_RATIO = [0.25, 0.25]

_C.INSTANCE_SEG = CN()
_C.INSTANCE_SEG.ENABLED = False

_C.INSTANCE_FLOW = CN()
_C.INSTANCE_FLOW.ENABLED = False

_C.PROBABILISTIC = CN()
_C.PROBABILISTIC.ENABLED = True  # learn a distribution over futures
_C.PROBABILISTIC.METHOD = 'GAUSSIAN' # [BERNOULLI, GAUSSIAN, MIXGAUSSIAN]

_C.PLANNING = CN()
_C.PLANNING.ENABLED = True
_C.PLANNING.GRU_ENABLED = False
_C.PLANNING.GRU_STATE_SIZE = 256
_C.PLANNING.SAMPLE_NUM = 150
_C.PLANNING.COMMAND = ['LEFT', 'FORWARD', 'RIGHT']

# dense configs
_C.PLANNING.DENSE = CN()
_C.PLANNING.DENSE.ENABLED = True # to consider dense centerlines
_C.PLANNING.DENSE.OBS = False # avoid obstacles
_C.PLANNING.DENSE.NUM_OBS = 5 # avoid obstacles
_C.PLANNING.DENSE.PATH = '../' # centerlines path

_C.FUTURE_DISCOUNT = 0.95

_C.OPTIMIZER = CN()
_C.OPTIMIZER.LR = 1e-3
_C.OPTIMIZER.WEIGHT_DECAY = 1e-7
_C.GRAD_NORM_CLIP = 5

_C.COST_FUNCTION = CN()
_C.COST_FUNCTION.SAFETY = 0.1
_C.COST_FUNCTION.LAMBDA = 1.
_C.COST_FUNCTION.HEADWAY = 1.
_C.COST_FUNCTION.LRDIVIDER = 10. 
_C.COST_FUNCTION.COMFORT = 0.1
_C.COST_FUNCTION.PROGRESS = 0.5
_C.COST_FUNCTION.VOLUME = 100.

def get_parser():
    parser = argparse.ArgumentParser(description='Fiery training')
    parser.add_argument('--config-file', default='', metavar='FILE', help='path to config file')
    parser.add_argument(
        'opts', help='Modify config options using the command-line', default=None, nargs=argparse.REMAINDER,
    )
    return parser


def get_cfg(args=None, cfg_dict=None):
    """ First get default config. Then merge cfg_dict. Then merge according to args. """

    cfg = _C.clone()

    if cfg_dict is not None:
        tmp = CfgNode(cfg_dict)
        for i in tmp.COST_FUNCTION:
            tmp.COST_FUNCTION.update({i: float(tmp.COST_FUNCTION.get(i))})
        cfg.merge_from_other_cfg(tmp)

    if args is not None:
        if args.config_file:
            cfg.merge_from_file(args.config_file)
        cfg.merge_from_list(args.opts)
        # cfg.freeze()
    return cfg

data/stp3/config.py
llmbev/talk2bev
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from stp3.utils.tools import gen_dx_bx
from stp3.utils.geometry import calculate_birds_eye_view_parameters
from skimage.draw import polygon


class Cost_Function(nn.Module):
    def __init__(self, cfg):
        super(Cost_Function, self).__init__()

        self.safetycost = SafetyCost(cfg)
        self.headwaycost = HeadwayCost(cfg)
        self.lrdividercost = LR_divider(cfg)
        self.comfortcost = Comfort(cfg)
        self.progresscost = Progress(cfg)
        self.rulecost = Rule(cfg)
        self.costvolume = Cost_Volume(cfg)

        self.n_future = cfg.N_FUTURE_FRAMES


    def forward(self, cost_volume, trajs, semantic_pred, lane_divider, drivable_area, target_point):
        '''
        args:
            cost_volume: torch.Tensor<float> (B, n_future, 200, 200)
            trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
            semantic_pred: torch.Tensor<float> (B, n_future, 200, 200)
            drivable_area: torch.Tensor<float> (B, 1/2, 200, 200)
            lane_divider: torch.Tensor<float> (B, 1/2, 200, 200)
            target_points: torch.Tensor<float> (B, 2)
        returns:
            cost_fo: B, N, n_future
            cost_fc: B, N            
        '''
        trajs = trajs * torch.tensor([-1, 1], device=trajs.device)
        safetycost = torch.clamp(self.safetycost(trajs, semantic_pred), 0, 100)
        headwaycost = torch.clamp(self.headwaycost(trajs, semantic_pred, drivable_area), 0, 100)
        lrdividercost = torch.clamp(self.lrdividercost(trajs, lane_divider), 0, 100)
        comfortcost = torch.clamp(self.comfortcost(trajs), 0, 100)
        progresscost = torch.clamp(self.progresscost(trajs, target_point), -100, 100)
        rulecost = torch.clamp(self.rulecost(trajs, drivable_area), 0, 100)
        costvolume = torch.clamp(self.costvolume(trajs, cost_volume), 0, 100)

        cost_fo = safetycost + headwaycost + lrdividercost + rulecost + costvolume
        cost_fc = comfortcost + progresscost
        return cost_fc, cost_fo



class BaseCost(nn.Module):
    def __init__(self, cfg):
        super(BaseCost, self).__init__()
        self.cfg = cfg

        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2], bx[:2]
        self.dx = nn.Parameter(dx,requires_grad=False)
        self.bx = nn.Parameter(bx,requires_grad=False)

        _,_, self.bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )

        self.W = cfg.EGO.WIDTH
        self.H = cfg.EGO.HEIGHT


    def get_origin_points(self, lambda_=0):
        W = self.W
        H = self.H
        pts = np.array([
            [-H / 2. + 0.5 - lambda_, W / 2. + lambda_],
            [H / 2. + 0.5 + lambda_, W / 2. + lambda_],
            [H / 2. + 0.5 + lambda_, -W / 2. - lambda_],
            [-H / 2. + 0.5 - lambda_, -W / 2. - lambda_],
        ])
        pts = (pts - self.bx.cpu().numpy()) / (self.dx.cpu().numpy())
        pts[:, [0, 1]] = pts[:, [1, 0]]
        rr , cc = polygon(pts[:,1], pts[:,0])
        rc = np.concatenate([rr[:,None], cc[:,None]], axis=-1)
        return torch.from_numpy(rc).to(device=self.bx.device) # (27,2)

    def get_points(self, trajs, lambda_=0):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)
        return:
        List[ torch.Tensor<int> (B, N, n_future), torch.Tensor<int> (B, N, n_future)]
        '''
        rc = self.get_origin_points(lambda_)
        B, N, n_future, _ = trajs.shape

        trajs = trajs.view(B, N, n_future, 1, 2) / self.dx
        trajs[:,:,:,:,[0,1]] = trajs[:,:,:,:,[1,0]]
        trajs = trajs + rc

        rr = trajs[:,:,:,:,0].long()
        rr = torch.clamp(rr, 0, self.bev_dimension[0] - 1)

        cc = trajs[:,:,:,:,1].long()
        cc = torch.clamp(cc, 0, self.bev_dimension[1] - 1)

        return rr, cc

    def compute_area(self, semantic_pred, trajs, ego_velocity=None, _lambda=0):
        '''
        semantic: torch.Tensor<float> (B, n_future, 200, 200)
        trajs: torch.Tensor<float> (B, N, n_future, 2)
        ego_velocity: torch.Tensor<float> (B, N, n_future)
        '''
        _lambda = int(_lambda / self.dx[0])
        rr, cc = self.get_points(trajs, _lambda)
        B, N, n_future,_ = trajs.shape

        if ego_velocity is None:
            ego_velocity = torch.ones((B,N,n_future), device=semantic_pred.device)

        ii = torch.arange(B)
        kk = torch.arange(n_future)
        subcost = semantic_pred[ii[:, None, None, None], kk[None, None, :, None], rr, cc].sum(dim=-1)
        subcost = subcost * ego_velocity


        return subcost

    def discretize(self, trajs):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
        '''
        B, N, n_future, _ = trajs.shape

        xx, yy = trajs[:,:,:,0], trajs[:,:,:,1]

        # discretize
        yi = ((yy - self.bx[0]) / self.dx[0]).long()
        yi = torch.clamp(yi,0, self.bev_dimension[0]-1)

        xi = ((xx - self.bx[1]) / self.dx[1]).long()
        xi = torch.clamp(xi, 0, self.bev_dimension[1]-1)

        

        return yi, xi

    def evaluate(self, trajs, C):
        '''
        args:
            trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
            C: torch.Tensor<float> (B, n_future, 200, 200)
        returns:
            CS: torch.Tensor<float> (B, N, n_future)
        '''
        B, N, n_future, _ = trajs.shape

        ii = torch.arange(B)
        ti = torch.arange(n_future)

        Syi, Sxi = self.discretize(trajs)

        CS = C[ii[:, None, None], ti[None, None, :], Syi, Sxi]
        return CS

class Cost_Volume(BaseCost):
    def __init__(self, cfg):
        super(Cost_Volume, self).__init__(cfg)
        """
            cfg.COST_FUNCTION.VOLUME: 100
        """
        self.factor = cfg.COST_FUNCTION.VOLUME

    def forward(self, trajs, cost_volume):
        '''
        cost_volume: torch.Tensor<float> (B, n_future, 200, 200)
        trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
        '''

        cost_volume = torch.clamp(cost_volume, 0, 1000)

        return self.evaluate(trajs, cost_volume) * self.factor

class Rule(BaseCost):
    def __init__(self, cfg):
        super(Rule, self).__init__(cfg)

        self.factor = 5

    def forward(self, trajs, drivable_area):
        '''
            trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
            drivable_area: torch.Tensor<float> (B, 1/2, 200, 200)
        '''
        assert drivable_area.ndim == 4, 'drivable area ndim should be 4'
        if drivable_area.shape[1] == 2:
            drivable_area = torch.softmax(drivable_area, dim=1)[:, 1]
            mask = drivable_area < 0.5
            drivable_area[mask] = 0
        else:
            drivable_area = drivable_area[:, 0]

        B, _, n_future, _ = trajs.shape
        _,H,W = drivable_area.shape

        dangerous_area = torch.logical_not(drivable_area).float().view(B, 1, W, H).expand(B, n_future, W, H)
        subcost = self.compute_area(dangerous_area, trajs)

        return subcost * self.factor


class SafetyCost(BaseCost):
    def __init__(self, cfg):
        super(SafetyCost, self).__init__(cfg)
        self.w = nn.Parameter(torch.tensor([1.,1.]),requires_grad=False)

        self._lambda = cfg.COST_FUNCTION.LAMBDA
        self.factor = cfg.COST_FUNCTION.SAFETY

    def forward(self, trajs, semantic_pred):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
        semantic_pred: torch.Tensor<float> (B, n_future, 200, 200)
        ego_velocity: torch.Tensor<float> (B, N, n_future)
        '''
        B, N, n_future, _ = trajs.shape
        ego_velocity = torch.zeros((B,N,n_future), device=semantic_pred.device)
        for i in range(n_future):
            if i == 0:
                ego_velocity[:,:,i] = torch.sqrt((trajs[:,:,i] ** 2).sum(axis=-1)) / 0.5
            else:
                ego_velocity[:,:,i] = torch.sqrt(((trajs[:,:,i] - trajs[:,:,i-1]) ** 2).sum(dim=-1)) / 0.5

        # o_c(tau, t, 0)
        subcost1 = self.compute_area(semantic_pred, trajs)
        # o_c(tau, t, lambda) x v(tau, t)
        subcost2 = self.compute_area(semantic_pred, trajs, ego_velocity, self._lambda)

        subcost = subcost1 * self.w[0] + subcost2 * self.w[1]


        return subcost * self.factor


class HeadwayCost(BaseCost):
    def __init__(self, cfg):
        super(HeadwayCost, self).__init__(cfg)
        self.L = 10  # Longitudinal distance keep 10m
        self.factor = cfg.COST_FUNCTION.HEADWAY

    def forward(self, trajs, semantic_pred, drivable_area):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
        semantic_pred: torch.Tensor<float> (B, n_future, 200, 200)
        drivable_area: torch.Tensor<float> (B, 1/2, 200, 200)
        '''
        assert drivable_area.ndim == 4, 'drivable area ndim should be 4'
        if drivable_area.shape[1] == 2:
            drivable_area = torch.softmax(drivable_area, dim=1)[:, 1]
            mask = drivable_area < 0.5
            drivable_area[mask] = 0
        else:
            drivable_area = drivable_area[:, 0]
        B, N, n_future, _ = trajs.shape
        _, W, H = drivable_area.shape
        drivable_mask = drivable_area.view(B, 1, W, H).expand(B, n_future, W, H)
        semantic_pred_ = semantic_pred * drivable_mask
        tmp_trajs = trajs.clone()
        tmp_trajs[:,:,:,1] = tmp_trajs[:,:,:,1]+self.L

        subcost = self.compute_area(semantic_pred_, tmp_trajs)

        return subcost * self.factor

class LR_divider(BaseCost):
    def __init__(self, cfg):
        super(LR_divider, self).__init__(cfg)
        self.L = 1 # Keep a distance of 2m from the lane line
        self.factor = cfg.COST_FUNCTION.LRDIVIDER

    def forward(self, trajs, lane_divider):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)   N: sample number
        lane_divider: torch.Tensor<float> (B, 1/2, 200, 200)
        '''
        assert lane_divider.ndim == 4, 'lane_divider ndim should be 4'
        if lane_divider.shape[1] == 2:
            lane_divider = torch.softmax(lane_divider, dim=1)[:,1]
            mask = lane_divider <= 0.5
            lane_divider[mask] = 0
        else:
            lane_divider = lane_divider[:,0]

        B, N, n_future, _ = trajs.shape

        yy, xx = self.discretize(trajs)
        yx = torch.stack([yy,xx],dim=-1) # (B, N, n_future, 2)

        # lane divider
        res1 = []
        for i in range(B):
            index = torch.nonzero(lane_divider[i]) # (n, 2)
            if len(index) != 0:
                yx_batch = yx[i].view(N, n_future, 1, 2)
                distance = torch.sqrt((((yx_batch - index) * reversed(self.dx))**2).sum(dim=-1)) # (N, n_future, n)
                distance,_ = distance.min(dim=-1) # (N, n_future)
                index = distance > self.L
                distance = (self.L - distance) ** 2
                distance[index] = 0
            else:
                distance = torch.zeros((N, n_future),device=trajs.device)
            res1.append(distance)
        res1 = torch.stack(res1, dim=0)

        return res1 * self.factor


class Comfort(BaseCost):
    def __init__(self, cfg):
        super(Comfort, self).__init__(cfg)

        self.c_lat_acc = 3 # m/s2
        self.c_lon_acc = 3 # m/s2
        self.c_jerk = 1 # m/s3

        self.factor = cfg.COST_FUNCTION.COMFORT

    def forward(self, trajs, t_fin=3):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)
        '''
        B, N, n_future, _ = trajs.shape
        #  3s in n_future time steps
        frame_rate = t_fin / (n_future)
        lateral_velocity = torch.zeros((B,N,n_future), device=trajs.device)
        longitudinal_velocity = torch.zeros((B, N, n_future), device=trajs.device)
        lateral_acc = torch.zeros((B,N,n_future), device=trajs.device)
        longitudinal_acc = torch.zeros((B,N,n_future), device=trajs.device)
        for i in range(n_future):
            if i == 0:
                lateral_velocity[:,:,i] = trajs[:,:,i,0] / frame_rate
                longitudinal_velocity[:,:,i] = trajs[:,:,i,1] / frame_rate
            else:
                lateral_velocity[:,:,i] = (trajs[:,:,i,0] - trajs[:,:,i-1,0]) / frame_rate
                longitudinal_velocity[:,:,i] = (trajs[:,:,i,1] - trajs[:,:,i-1,1]) / frame_rate
        for i in range(1, n_future):
            lateral_acc[:,:,i] = (lateral_velocity[:,:,i] - lateral_velocity[:,:,i-1]) / frame_rate
            longitudinal_acc[:,:,i] = (longitudinal_velocity[:,:,i] - longitudinal_velocity[:,:,i-1]) / frame_rate
        lateral_acc, _ = torch.abs(lateral_acc).max(dim=-1)
        longitudinal_acc, _ = torch.abs(longitudinal_acc).max(dim=-1)

        # jerk
        ego_velocity = torch.zeros((B, N, n_future), device=trajs.device)
        ego_acc = torch.zeros((B,N,n_future), device=trajs.device)
        ego_jerk = torch.zeros((B,N,n_future), device=trajs.device)
        for i in range(n_future):
            if i == 0:
                ego_velocity[:, :, i] = torch.sqrt((trajs[:, :, i] ** 2).sum(dim=-1)) / frame_rate
            else:
                ego_velocity[:, :, i] = torch.sqrt(((trajs[:, :, i] - trajs[:, :, i - 1]) ** 2).sum(dim=-1)) / frame_rate
        for i in range(1, n_future):
            ego_acc[:,:,i] = (ego_velocity[:,:,i] - ego_velocity[:,:,i-1]) / frame_rate
        for i in range(2, n_future):
            ego_jerk[:,:,i] = (ego_acc[:,:,i] - ego_acc[:,:,i-1]) / frame_rate
        ego_jerk,_ = torch.abs(ego_jerk).max(dim=-1)

        subcost = torch.zeros((B, N), device=trajs.device)

        # 30 was the clamp for 0.5 frame_rate 
        lateral_clamp = 30
        lateral_acc = torch.clamp(torch.abs(lateral_acc) - self.c_lat_acc, 0,lateral_clamp)
        subcost += lateral_acc ** 2

        # 30 was the clamp for 0.5 frame_rate
        longitudinal_clamp = 30
        longitudinal_acc = torch.clamp(torch.abs(longitudinal_acc) - self.c_lon_acc, 0, longitudinal_clamp)
        subcost += longitudinal_acc ** 2

        # 20 was the clamp for 0.5 frame rate
        jerk_clamp = 20
        ego_jerk = torch.clamp(torch.abs(ego_jerk) - self.c_jerk, 0, jerk_clamp)
        subcost += ego_jerk ** 2

        return subcost * self.factor

class Progress(BaseCost):
    def __init__(self, cfg):
        super(Progress, self).__init__(cfg)
        self.factor = cfg.COST_FUNCTION.PROGRESS

    def forward(self, trajs, target_points):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)
        target_points: torch.Tensor<float> (B, 2)
        '''
        B, N, n_future, _ = trajs.shape
        subcost1, _ = trajs[:,:,:,1].max(dim=-1)

        if target_points.sum() < 0.5:
            subcost2 = 0
        else:
            trajs = trajs[:,:,-1]
            target_points = target_points.unsqueeze(1)
            subcost2 = ((trajs - target_points) ** 2).sum(dim=-1)

        return (subcost2 - subcost1) * self.factor


class ComfortMetric():
    def __init__(self):
        super(ComfortMetric, self).__init__()
        """
        This was the value for 0.5 frame rate
            self.c_lat_acc = 3 # m/s2
            self.c_lon_acc = 3 # m/s2
            self.c_jerk = 1 # m/s3
        """
        self.c_lat_acc = 3 # m/s2
        self.c_lon_acc = 3 # m/s2
        self.c_jerk = 1 # m/s3

    def compute_cost(self, trajs, t_fin=3):
        '''
        trajs: torch.Tensor<float> (B, N, n_future, 2)
        '''
        #  3s in n_future time steps
        frame_rate = self.frame_rate
        lateral_velocity = torch.zeros((B,N,n_future), device=trajs.device)
        longitudinal_velocity = torch.zeros((B, N, n_future), device=trajs.device)
        lateral_acc = torch.zeros((B,N,n_future), device=trajs.device)
        longitudinal_acc = torch.zeros((B,N,n_future), device=trajs.device)
        for i in range(n_future):
            if i == 0:
                lateral_velocity[:,:,i] = trajs[:,:,i,0] / frame_rate
                longitudinal_velocity[:,:,i] = trajs[:,:,i,1] / frame_rate
            else:
                lateral_velocity[:,:,i] = (trajs[:,:,i,0] - trajs[:,:,i-1,0]) / frame_rate
                longitudinal_velocity[:,:,i] = (trajs[:,:,i,1] - trajs[:,:,i-1,1]) / frame_rate
        for i in range(1, n_future):
            lateral_acc[:,:,i] = (lateral_velocity[:,:,i] - lateral_velocity[:,:,i-1]) / frame_rate
            longitudinal_acc[:,:,i] = (longitudinal_velocity[:,:,i] - longitudinal_velocity[:,:,i-1]) / frame_rate
        lateral_acc, _ = torch.abs(lateral_acc).max(dim=-1)
        longitudinal_acc, _ = torch.abs(longitudinal_acc).max(dim=-1)

        # jerk
        ego_velocity = torch.zeros((B, N, n_future), device=trajs.device)
        ego_acc = torch.zeros((B,N,n_future), device=trajs.device)
        ego_jerk = torch.zeros((B,N,n_future), device=trajs.device)
        for i in range(n_future):
            if i == 0:
                ego_velocity[:, :, i] = torch.sqrt((trajs[:, :, i] ** 2).sum(dim=-1)) / frame_rate
            else:
                ego_velocity[:, :, i] = torch.sqrt(((trajs[:, :, i] - trajs[:, :, i - 1]) ** 2).sum(dim=-1)) / frame_rate
        for i in range(1, n_future):
            ego_acc[:,:,i] = (ego_velocity[:,:,i] - ego_velocity[:,:,i-1]) / frame_rate
        for i in range(2, n_future):
            ego_jerk[:,:,i] = (ego_acc[:,:,i] - ego_acc[:,:,i-1]) / frame_rate
        ego_jerk,_ = torch.abs(ego_jerk).max(dim=-1)

        subcost = torch.zeros((B, N), device=trajs.device)

        # 30 was the clamp for 0.5 frame_rate 
        lateral_clamp = 30
        lateral_acc = torch.clamp(torch.abs(lateral_acc) - self.c_lat_acc, 0,lateral_clamp)
        subcost += lateral_acc ** 2

        # 30 was the clamp for 0.5 frame_rate
        longitudinal_clamp = 30
        longitudinal_acc = torch.clamp(torch.abs(longitudinal_acc) - self.c_lon_acc, 0, longitudinal_clamp)
        subcost += longitudinal_acc ** 2

        # 20 was the clamp for 0.5 frame rate
        jerk_clamp = 20
        ego_jerk = torch.clamp(torch.abs(ego_jerk) - self.c_jerk, 0, jerk_clamp)
        subcost += ego_jerk ** 2

        return subcost

data/stp3/cost.py
llmbev/talk2bev
import os
import json
import numpy as np
import cv2
import torch
import torch.utils.data
import torchvision
from PIL import Image
import PIL
from pyquaternion import Quaternion

from stp3.utils.geometry import (
    update_intrinsics,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
)
import stp3.utils.sampler as trajectory_sampler


class CarlaDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5  # SECOND
    def __init__(self, root_dir, is_train, cfg):
        super(CarlaDataset, self).__init__()
        self.root_dir = root_dir
        self.is_train = is_train
        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD
        self.cfg = cfg
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
             ]
        )

        self.front = []
        self.left = []
        self.right = []
        self.rear = []
        self.front_depth = []
        self.left_depth = []
        self.right_depth = []
        self.rear_depth = []
        self.topdown = []
        self.x = []
        self.y = []
        self.x_command = []
        self.y_command = []
        self.theta = []
        self.steer = []
        self.throttle = []
        self.brake = []
        self.command = []
        self.velocity = []
        self.hdmap = []

        self.get_train_val()

    def get_train_val(self):
        train_towns = ['Town01', 'Town02', 'Town03', 'Town04', 'Town06', 'Town10']
        val_towns = ['Town05']
        train_data, val_data = [], []
        for town in train_towns:
            train_data.append(os.path.join(self.root_dir, town+'_tiny'))
            train_data.append(os.path.join(self.root_dir, town+'_short'))
        for town in val_towns:
            val_data.append(os.path.join(self.root_dir, town+'_short'))

        require_data = train_data if self.is_train else val_data

        for subroot in require_data:
            preload_file = os.path.join(subroot, 'cam_mea_topdown_'+str(self.receptive_field)+'_'+str(self.sequence_length)+'.npy')

            if not os.path.exists(preload_file):
                preload_front = []
                preload_left = []
                preload_right = []
                preload_rear = []
                preload_front_depth = []
                preload_left_depth = []
                preload_right_depth = []
                preload_rear_depth = []
                preload_topdown = []
                preload_x = []
                preload_y = []
                preload_x_command = []
                preload_y_command = []
                preload_theta = []
                preload_steer = []
                preload_throttle = []
                preload_brake = []
                preload_command = []
                preload_velocity = []
                preload_hdmap = []

                root_files = os.listdir(subroot)
                routes = [folder for folder in root_files if not os.path.isfile(os.path.join(subroot, folder))]
                for route in routes:
                    route_dir = os.path.join(subroot, route)
                    num_seq = len(os.listdir(route_dir + "/rgb_front/")) - self.sequence_length
                    for seq in range(num_seq):
                        fronts, lefts, rights, rears = [], [], [], []
                        fr_depths, le_depths, ri_depths, re_depths = [], [], [], []
                        xs, ys, thetas = [], [], []
                        top_down, hdmap = [], []
                        for i in range(self.receptive_field):
                            filename = f"{str(seq+1+i).zfill(4)}.png"
                            fronts.append(route_dir + "/rgb_front/" + filename)
                            lefts.append(route_dir + "/rgb_left/" + filename)
                            rights.append(route_dir + "/rgb_right/" + filename)
                            rears.append(route_dir + "/rgb_rear/" + filename)
                            fr_depths.append(route_dir + "/depth_front/" + filename)
                            le_depths.append(route_dir + "/depth_left/" + filename)
                            ri_depths.append(route_dir + "/depth_right/" + filename)
                            re_depths.append(route_dir + "/depth_rear/" + filename)
                            top_down.append(route_dir + "/topdown/" + filename)
                            hdmap.append(route_dir + "/hdmap/" + filename)
                            # position
                            with open(route_dir + f"/measurements/{str(seq+1+i).zfill(4)}.json","r") as read_file:
                                data = json.load(read_file)
                            xs.append(data['x'])
                            ys.append(data['y'])
                            thetas.append(data['theta'])

                        preload_x_command.append(data['x_command'])
                        preload_y_command.append(data['y_command'])
                        preload_steer.append(data['steer'])
                        preload_throttle.append(data['throttle'])
                        preload_brake.append(data['brake'])
                        preload_command.append(data['command'])
                        preload_velocity.append(data['speed'])

                        for i in range(self.receptive_field, self.sequence_length):
                            filename = f"{str(seq + 1 + i).zfill(4)}.png"
                            top_down.append(route_dir + "/topdown/" + filename)

                            with open(route_dir + f"/measurements/{str(seq+1+i).zfill(4)}.json","r") as read_file:
                                data = json.load(read_file)
                            xs.append(data['x'])
                            ys.append(data['y'])
                            if np.isnan(data['theta']):
                                thetas.append(0)
                            else:
                                thetas.append(data['theta'])

                        preload_front.append(fronts)
                        preload_left.append(lefts)
                        preload_right.append(rights)
                        preload_rear.append(rears)
                        preload_front_depth.append(fr_depths)
                        preload_left_depth.append(le_depths)
                        preload_right_depth.append(ri_depths)
                        preload_rear_depth.append(re_depths)
                        preload_x.append(xs)
                        preload_y.append(ys)
                        preload_theta.append(thetas)
                        preload_topdown.append(top_down)
                        preload_hdmap.append(hdmap)

                preload_dict = {}
                preload_dict['front'] = preload_front
                preload_dict['left'] = preload_left
                preload_dict['right'] = preload_right
                preload_dict['rear'] = preload_rear
                preload_dict['front_depth'] = preload_front_depth
                preload_dict['left_depth'] = preload_left_depth
                preload_dict['right_depth'] = preload_right_depth
                preload_dict['rear_depth'] = preload_rear_depth
                preload_dict['topdown'] = preload_topdown
                preload_dict['x'] = preload_x
                preload_dict['y'] = preload_y
                preload_dict['x_command'] = preload_x_command
                preload_dict['y_command'] = preload_y_command
                preload_dict['theta'] = preload_theta
                preload_dict['steer'] = preload_steer
                preload_dict['throttle'] = preload_throttle
                preload_dict['brake'] = preload_brake
                preload_dict['command'] = preload_command
                preload_dict['velocity'] = preload_velocity
                preload_dict['hdmap'] = preload_hdmap
                np.save(preload_file, preload_dict)

            # load from npy if available
            preload_dict = np.load(preload_file, allow_pickle=True)
            self.front += preload_dict.item()['front']
            self.left += preload_dict.item()['left']
            self.right += preload_dict.item()['right']
            self.rear += preload_dict.item()['rear']
            self.front_depth += preload_dict.item()['front_depth']
            self.left_depth += preload_dict.item()['left_depth']
            self.right_depth += preload_dict.item()['right_depth']
            self.rear_depth += preload_dict.item()['rear_depth']
            self.topdown += preload_dict.item()['topdown']
            self.x += preload_dict.item()['x']
            self.y += preload_dict.item()['y']
            self.x_command += preload_dict.item()['x_command']
            self.y_command += preload_dict.item()['y_command']
            self.theta += preload_dict.item()['theta']
            self.steer += preload_dict.item()['steer']
            self.throttle += preload_dict.item()['throttle']
            self.brake += preload_dict.item()['brake']
            self.command += preload_dict.item()['command']
            self.velocity += preload_dict.item()['velocity']
            self.hdmap += preload_dict.item()['hdmap']
            print("Preloading " + str(len(preload_dict.item()['front'])) + " sequences from " + preload_file)

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

    def get_future_egomotion(self, seq_x, seq_y, seq_theta):
        future_egomotions = []

        def convert_to_matrix_numpy(x, y, theta):
            matrix = np.zeros((4,4), dtype=np.float32)
            matrix[:2, :2] = np.array([
                [np.cos(theta), -np.sin(theta)],
                [np.sin(theta), np.cos(theta)]
            ])
            matrix[2,2] = 1
            matrix[0,3] = x
            matrix[1,3] = y
            matrix[3,3] = 1
            return matrix

        for i in range(len(seq_x)-1):
            egopose_t0 = convert_to_matrix_numpy(seq_x[i], seq_y[i], seq_theta[i])
            egopose_t1 = convert_to_matrix_numpy(seq_x[i+1], seq_y[i+1], seq_theta[i+1])

            future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
            future_egomotion[3, :3] = 0.0
            future_egomotion[3, 3] = 1.0

            future_egomotion = torch.Tensor(future_egomotion).float()
            future_egomotion = mat2pose_vec(future_egomotion)
            future_egomotions.append(future_egomotion.unsqueeze(0))

        return torch.cat(future_egomotions, dim=0)

    def get_hdmap(self, path, scale, crop):
        image = Image.open(path)
        (width, height) = (int(image.width // scale), int(image.height // scale))
        im_resized = image.resize((width, height))
        image = np.asarray(im_resized)
        start_x = height // 2 - crop // 2
        start_y = width // 2 - crop // 2
        cropped_image = image[start_x:start_x + crop, start_y:start_y + crop]
        h, w, c = cropped_image.shape
        lane_index = np.all((cropped_image == [255, 0, 255]), axis=2)
        lane = np.zeros((h, w))
        lane[lane_index] = 1
        drivable_index = np.all((cropped_image == [54, 52, 46]), axis=2)
        drivable = np.zeros((h, w))
        drivable[drivable_index] = 1
        drivable = np.logical_or(drivable, lane)
        # down, right is the positive
        lane = lane[::-1,::-1]
        drivable = drivable[::-1,::-1]
        hdmap = np.concatenate([lane[None], drivable[None]], axis=0)
        return hdmap

    def get_labels(self, path, scale, crop):
        image = Image.open(path)
        (width, height) = (int(image.width // scale), int(image.height // scale))
        im_resized = image.resize((width, height), resample=PIL.Image.NEAREST)
        image = np.asarray(im_resized)
        start_x = height // 2 - crop // 2
        start_y = width // 2 - crop // 2
        cropped_image = image[start_x:start_x + crop, start_y:start_y + crop]
        h, w = cropped_image.shape
        vehicle_index = cropped_image == 10
        vehicle = np.zeros((h, w))
        vehicle[vehicle_index] = 1
        vehicle[89:112,96:105] = 0
        pedestrian_index = cropped_image == 4
        pedestrian = np.zeros((h, w))
        pedestrian[pedestrian_index] = 1
        vehicle = vehicle[::-1,::-1]
        pedestrian = pedestrian[::-1,::-1]
        return vehicle.copy(), pedestrian.copy()

    def get_trajectory_sampling(self, v0, steering):

        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories

    def get_cam_para(self):
        def get_cam_to_ego(dof):
            yaw = dof[5]
            rotation = Quaternion(scalar=np.cos(yaw/2), vector=[0, 0, np.sin(yaw/2)])
            translation = np.array(dof[:3])[:, None]
            cam_to_ego = np.vstack([
                np.hstack((rotation.rotation_matrix,translation)),
                np.array([0,0,0,1])
            ])
            return cam_to_ego

        cam_front = [1.3, 0.0, 2.3, 0.0, 0.0, 0.0] # x,y,z,roll,pitch, yaw
        cam_left = [1.3, 0.0, 2.3, 0.0, 0.0, -60.0]
        cam_right = [1.3, 0.0, 2.3, 0.0, 0.0, 60.0]
        cam_rear = [-1.3, 0.0, 2.3, 0.0, 0.0, 180.0]
        front_to_ego = torch.from_numpy(get_cam_to_ego(cam_front)).float().unsqueeze(0)
        left_to_ego = torch.from_numpy(get_cam_to_ego(cam_left)).float().unsqueeze(0)
        right_to_ego = torch.from_numpy(get_cam_to_ego(cam_right)).float().unsqueeze(0)
        rear_to_ego = torch.from_numpy(get_cam_to_ego(cam_rear)).float().unsqueeze(0)
        extrinsic = torch.cat([front_to_ego, left_to_ego, right_to_ego, rear_to_ego], dim=0)

        sensor_data = {
            'width': 400,
            'height': 300,
            'fov': 100
        }
        w = sensor_data['width']
        h = sensor_data['height']
        fov = sensor_data['fov']
        f = w / (2 * np.tan(fov * np.pi/ 360))
        Cu = w / 2
        Cv = h / 2
        intrinsic = torch.Tensor([
            [f, 0, Cu],
            [0, f, Cv],
            [0, 0, 1]
        ])
        intrinsic = update_intrinsics(
            intrinsic, (h-256)/2, (w-256)/2,
            scale_width=1,
            scale_height=1
        )
        intrinsic = intrinsic.unsqueeze(0).expand(4,3,3)

        return extrinsic, intrinsic

    def get_depth(self, data):
        """
        Computes the normalized depth
        """
        data = data.astype(np.float32)

        normalized = np.dot(data, [65536.0, 256.0, 1.0])
        normalized /= (256 * 256 * 256 - 1)
        return torch.from_numpy(normalized * 1000)

    def __getitem__(self, index):
        data = {}
        keys = ['image', 'depths', 'segmentation', 'pedestrian', 'extrinsics', 'intrinsics', 'hdmap', 'gt_trajectory']
        for key in keys:
            data[key] = []

        seq_fronts = self.front[index]
        seq_lefts = self.left[index]
        seq_rights = self.right[index]
        seq_rears = self.rear[index]
        seq_front_depths = self.front_depth[index]
        seq_left_depths = self.left_depth[index]
        seq_right_depths = self.right_depth[index]
        seq_rear_depths = self.rear_depth[index]
        seq_hdmaps = self.hdmap[index]
        seq_topdowns = self.topdown[index]
        seq_x = self.x[index]
        seq_y = self.y[index]
        seq_theta = self.theta[index]

        for i in range(self.receptive_field):
            images = []
            images.append(self.normalise_image(np.array(
                scale_and_crop_image(Image.open(seq_fronts[i]), scale=1., crop=256))).unsqueeze(0))
            images.append(self.normalise_image(np.array(
                scale_and_crop_image(Image.open(seq_lefts[i]), scale=1., crop=256))).unsqueeze(0))
            images.append(self.normalise_image(np.array(
                scale_and_crop_image(Image.open(seq_rights[i]), scale=1., crop=256))).unsqueeze(0))
            images.append(self.normalise_image(np.array(
                scale_and_crop_image(Image.open(seq_rears[i]), scale=1.,crop=256))).unsqueeze(0))
            images = torch.cat(images, dim=0)
            data['image'].append(images.unsqueeze(0))
            depths = []
            depths.append(self.get_depth(np.array(
                scale_and_crop_image(Image.open(seq_front_depths[i]), scale=1., crop=256))).unsqueeze(0))
            depths.append(self.get_depth(np.array(
                scale_and_crop_image(Image.open(seq_left_depths[i]), scale=1., crop=256))).unsqueeze(0))
            depths.append(self.get_depth(np.array(
                scale_and_crop_image(Image.open(seq_right_depths[i]), scale=1., crop=256))).unsqueeze(0))
            depths.append(self.get_depth(np.array(
                scale_and_crop_image(Image.open(seq_rear_depths[i]), scale=1., crop=256))).unsqueeze(0))
            depths = torch.cat(depths, dim=0)
            data['depths'].append(depths.unsqueeze(0))
            extrinsics, intrinsics = self.get_cam_para()
            data['extrinsics'].append(extrinsics.unsqueeze(0))
            data['intrinsics'].append(intrinsics.unsqueeze(0))
            data['hdmap'].append(torch.from_numpy(self.get_hdmap(seq_hdmaps[i], 1., 200)).unsqueeze(0))


            # fix for theta=nan in some measurements
            if np.isnan(seq_theta[i]):
                seq_theta[i] = 0.
        ego_x = seq_x[self.receptive_field-1]
        ego_y = seq_y[self.receptive_field-1]
        ego_theta = seq_theta[self.receptive_field-1]

        for i in range(self.sequence_length):
            if i >= self.receptive_field-1:
                local_waypoint = transform_2d_points(np.zeros((1, 3)),
                                                     np.pi / 2 - seq_theta[i], -seq_x[i], -seq_y[i],
                                                     np.pi / 2 - ego_theta, -ego_x, -ego_y)
                local_waypoint = local_waypoint * [1.0, -1.0, 1.0]
                data['gt_trajectory'].append(torch.from_numpy(local_waypoint))
            segmentation, pedestrian = self.get_labels(seq_topdowns[i], 1.1, 200)
            data['segmentation'].append(torch.from_numpy(segmentation).unsqueeze(0).unsqueeze(0))
            data['pedestrian'].append(torch.from_numpy(pedestrian).unsqueeze(0).unsqueeze(0))

        R = np.array([
            [np.cos(np.pi / 2 + ego_theta), -np.sin(np.pi / 2 + ego_theta)],
            [np.sin(np.pi / 2 + ego_theta), np.cos(np.pi / 2 + ego_theta)]
        ])
        local_command_point = np.array([self.x_command[index] - ego_x, self.y_command[index] - ego_y])
        local_command_point = R.T.dot(local_command_point)
        local_command_point = local_command_point * [1.0, -1.0]
        data['target_point'] = torch.from_numpy(local_command_point)

        if self.command[index] == 1:
            data['command'] = 'LEFT'
        elif self.command[index] == 2:
            data['command'] = 'RIGHT'
        elif self.command[index] == 3:
            data['command'] = 'FORWARD'
        else:
            data['command'] = 'LANE'
        data['steer'] = self.steer[index]
        data['throttle'] = self.throttle[index]
        data['brake'] = self.brake[index]
        data['velocity'] = self.velocity[index]
        data['future_egomotion'] = self.get_future_egomotion(seq_x, seq_y, seq_theta)
        data['sample_trajectory'] = torch.from_numpy(self.get_trajectory_sampling(self.velocity[index], self.steer[index])).float()

        for key, value in data.items():
            if key in keys:
                data[key] = torch.cat(value, dim=0)

        return data



def scale_and_crop_image(image, scale=1., crop=256):
    """
    Scale and crop a PIL image, returning a channels-first numpy array.
    """
    (width, height) = (int(image.width // scale), int(image.height // scale))
    im_resized = image.resize((width, height))
    image = np.asarray(im_resized)
    start_x = height//2 - crop//2
    start_y = width//2 - crop//2
    cropped_image = image[start_x:start_x+crop, start_y:start_y+crop]
    return cropped_image


def transform_2d_points(xyz, r1, t1_x, t1_y, r2, t2_x, t2_y):
    """
    Build a rotation matrix and take the dot product.
    """
    # z value to 1 for rotation
    xy1 = xyz.copy()
    xy1[:, 2] = 1

    c, s = np.cos(r1), np.sin(r1)
    r1_to_world = np.matrix([[c, s, t1_x], [-s, c, t1_y], [0, 0, 1]])

    # np.dot converts to a matrix, so we explicitly change it back to an array
    world = np.asarray(r1_to_world @ xy1.T)

    c, s = np.cos(r2), np.sin(r2)
    r2_to_world = np.matrix([[c, s, t2_x], [-s, c, t2_y], [0, 0, 1]])
    world_to_r2 = np.linalg.inv(r2_to_world)

    out = np.asarray(world_to_r2 @ world).T

    # reset z-coordinate
    out[:, 2] = xyz[:, 2]

    return out

data/stp3/datas/CarlaData.py
llmbev/talk2bev
import os
from PIL import Image

import numpy as np
import cv2
import torch
import torch.utils.data
import torch.nn.functional as F
import torchvision

from shapely.geometry import LineString

from pyquaternion import Quaternion
from shapely import affinity
from nuscenes.nuscenes import NuScenes, NuScenesExplorer
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.splits import create_splits_scenes
from nuscenes.utils.data_classes import Box
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.eval.common.utils import quaternion_yaw
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from stp3.utils.tools import ( gen_dx_bx, get_nusc_maps)
import matplotlib.pyplot as plt 
from math import cos, sin

from stp3.utils.geometry import (
    resize_and_crop_image,
    update_intrinsics,
    calculate_birds_eye_view_parameters,
    convert_egopose_to_matrix_numpy,
    pose_vec2mat,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
    get_global_pose,
    get_patch_coord,
    dfs,
    dfs_incoming,
    remove_overlapping_lane_seq,
    extend_both_sides,
    project_to_frenet_frame,
    project_to_cartesian_frame
)
from stp3.utils.instance import convert_instance_mask_to_center_and_offset_label
import stp3.utils.sampler as trajectory_sampler
from stp3.utils.spline import Spline2D
from stp3.utils.Optnode import *

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

resolution_meters = 0.1
t_fin = 3.5
num = 61
nvar = 11
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
debug = False
brdebug = False
problem = OPTNode(t_fin=t_fin, num=num)

class FuturePredictionDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5 #SECOND
    def __init__(self, nusc, is_train, cfg):
        self.nusc = nusc
        self.dataroot = self.nusc.dataroot
        self.nusc_exp = NuScenesExplorer(nusc)
        self.nusc_can = NuScenesCanBus(dataroot=self.dataroot)
        self.is_train = is_train
        self.cfg = cfg
        self.problem = OPTNode(t_fin=3.5, num=self.cfg.N_FUTURE_FRAMES + 1)

        if self.is_train == 0:
            self.mode = 'train'
        elif self.is_train == 1:
            self.mode = 'val'
        elif self.is_train == 2:
            self.mode = 'test'
        else:
            raise NotImplementedError

        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD

        self.scenes = self.get_scenes()
        self.ixes = self.prepro()
        self.indices = self.get_indices()

        # Image resizing and cropping
        self.augmentation_parameters = self.get_resizing_and_cropping_parameters()

        # Normalising input images
        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # Bird's-eye view parameters
        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_resolution, self.bev_start_position, self.bev_dimension = (
            bev_resolution.numpy(), bev_start_position.numpy(), bev_dimension.numpy()
        )

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # The number of sampled trajectories
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        # HD-map feature extractor
        self.nusc_maps = get_nusc_maps(self.cfg.DATASET.MAP_FOLDER)
        self.scene2map = {}
        for sce in self.nusc.scene:
            log = self.nusc.get('log', sce['log_token'])
            self.scene2map[sce['name']] = log['location']
        self.save_dir = cfg.DATASET.SAVE_DIR

    def get_scenes(self):
        # filter by scene split
        split = {'v1.0-trainval': {0: 'train', 1: 'val', 2: 'test'},
                 'v1.0-mini': {0: 'mini_train', 1: 'mini_val'},}[
            self.nusc.version
        ][self.is_train]

        blacklist = [419] + self.nusc_can.can_blacklist  # # scene-0419 does not have vehicle monitor data
        blacklist = ['scene-' + str(scene_no).zfill(4) for scene_no in blacklist]

        scenes = create_splits_scenes()[split][:]
        for scene_no in blacklist:
            if scene_no in scenes:
                scenes.remove(scene_no)

        return scenes

    def prepro(self):
        samples = [samp for samp in self.nusc.sample]

        # remove samples that aren't in this split
        samples = [samp for samp in samples if self.nusc.get('scene', samp['scene_token'])['name'] in self.scenes]

        # sort by scene, timestamp (only to make chronological viz easier)
        samples.sort(key=lambda x: (x['scene_token'], x['timestamp']))

        return samples

    def get_indices(self):
        indices = []
        for index in range(len(self.ixes)):
            is_valid_data = True
            previous_rec = None
            current_indices = []
            for t in range(self.sequence_length):
                index_t = index + t
                # Going over the dataset size limit.
                if index_t >= len(self.ixes):
                    is_valid_data = False
                    break
                rec = self.ixes[index_t]
                # Check if scene is the same
                if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
                    is_valid_data = False
                    break

                current_indices.append(index_t)
                previous_rec = rec

            if is_valid_data:
                indices.append(current_indices)

        return np.asarray(indices)

    def get_resizing_and_cropping_parameters(self):
        original_height, original_width = self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH
        final_height, final_width = self.cfg.IMAGE.FINAL_DIM

        resize_scale = self.cfg.IMAGE.RESIZE_SCALE
        resize_dims = (int(original_width * resize_scale), int(original_height * resize_scale))
        resized_width, resized_height = resize_dims

        crop_h = self.cfg.IMAGE.TOP_CROP
        crop_w = int(max(0, (resized_width - final_width) / 2))
        # Left, top, right, bottom crops.
        crop = (crop_w, crop_h, crop_w + final_width, crop_h + final_height)

        if resized_width != final_width:
            print('Zero padding left and right parts of the image.')
        if crop_h + final_height != resized_height:
            print('Zero padding bottom part of the image.')

        return {'scale_width': resize_scale,
                'scale_height': resize_scale,
                'resize_dims': resize_dims,
                'crop': crop,
                }

    def get_input_data(self, rec):
        """
        Parameters
        ----------
            rec: nuscenes identifier for a given timestamp

        Returns
        -------
            images: torch.Tensor<float> (N, 3, H, W)
            intrinsics: torch.Tensor<float> (3, 3)
            extrinsics: torch.Tensor(N, 4, 4)
        """
        images = []
        intrinsics = []
        extrinsics = []
        depths = []
        cameras = self.cfg.IMAGE.NAMES

        # The extrinsics we want are from the camera sensor to "flat egopose" as defined
        # https://github.com/nutonomy/nuscenes-devkit/blob/9b492f76df22943daf1dc991358d3d606314af27/python-sdk/nuscenes/nuscenes.py#L279
        # which corresponds to the position of the lidar.
        # This is because the labels are generated by projecting the 3D bounding box in this lidar's reference frame.

        # From lidar egopose to world.
        lidar_sample = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        lidar_pose = self.nusc.get('ego_pose', lidar_sample['ego_pose_token'])
        yaw = Quaternion(lidar_pose['rotation']).yaw_pitch_roll[0]
        lidar_rotation = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        lidar_translation = np.array(lidar_pose['translation'])[:, None]
        lidar_to_world = np.vstack([
            np.hstack((lidar_rotation.rotation_matrix, lidar_translation)),
            np.array([0, 0, 0, 1])
        ])

        for cam in cameras:
            camera_sample = self.nusc.get('sample_data', rec['data'][cam])

            # Transformation from world to egopose
            car_egopose = self.nusc.get('ego_pose', camera_sample['ego_pose_token'])
            egopose_rotation = Quaternion(car_egopose['rotation']).inverse
            egopose_translation = -np.array(car_egopose['translation'])[:, None]
            world_to_car_egopose = np.vstack([
                np.hstack((egopose_rotation.rotation_matrix, egopose_rotation.rotation_matrix @ egopose_translation)),
                np.array([0, 0, 0, 1])
            ])

            # From egopose to sensor
            sensor_sample = self.nusc.get('calibrated_sensor', camera_sample['calibrated_sensor_token'])
            intrinsic = torch.Tensor(sensor_sample['camera_intrinsic'])
            sensor_rotation = Quaternion(sensor_sample['rotation'])
            sensor_translation = np.array(sensor_sample['translation'])[:, None]
            car_egopose_to_sensor = np.vstack([
                np.hstack((sensor_rotation.rotation_matrix, sensor_translation)),
                np.array([0, 0, 0, 1])
            ])
            car_egopose_to_sensor = np.linalg.inv(car_egopose_to_sensor)

            # Combine all the transformation.
            # From sensor to lidar.
            lidar_to_sensor = car_egopose_to_sensor @ world_to_car_egopose @ lidar_to_world
            sensor_to_lidar = torch.from_numpy(np.linalg.inv(lidar_to_sensor)).float()

            # Load image
            image_filename = os.path.join(self.dataroot, camera_sample['filename'])
            img = Image.open(image_filename)
            # Resize and crop
            img = resize_and_crop_image(
                img, resize_dims=self.augmentation_parameters['resize_dims'], crop=self.augmentation_parameters['crop']
            )
            # Normalise image
            normalised_img = self.normalise_image(img)

            # Combine resize/cropping in the intrinsics
            top_crop = self.augmentation_parameters['crop'][1]
            left_crop = self.augmentation_parameters['crop'][0]
            intrinsic = update_intrinsics(
                intrinsic, top_crop, left_crop,
                scale_width=self.augmentation_parameters['scale_width'],
                scale_height=self.augmentation_parameters['scale_height']
            )

            # Get Depth
            # Depth data should under the dataroot path 
            if self.cfg.LIFT.GT_DEPTH:
                base_root = os.path.join(self.dataroot, 'depths') 
                filename = os.path.basename(camera_sample['filename']).split('.')[0] + '.npy'
                depth_file_name = os.path.join(base_root, cam, 'npy', filename)
                depth = torch.from_numpy(np.load(depth_file_name)).unsqueeze(0).unsqueeze(0)
                depth = F.interpolate(depth, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear')
                depth = depth.squeeze()
                crop = self.augmentation_parameters['crop']
                depth = depth[crop[1]:crop[3], crop[0]:crop[2]]
                depth = torch.round(depth)
                depths.append(depth.unsqueeze(0).unsqueeze(0))

            images.append(normalised_img.unsqueeze(0).unsqueeze(0))
            intrinsics.append(intrinsic.unsqueeze(0).unsqueeze(0))
            extrinsics.append(sensor_to_lidar.unsqueeze(0).unsqueeze(0))

        images, intrinsics, extrinsics = (torch.cat(images, dim=1),
                                          torch.cat(intrinsics, dim=1),
                                          torch.cat(extrinsics, dim=1)
                                          )
        if len(depths) > 0:
            depths = torch.cat(depths, dim=1)

        return images, intrinsics, extrinsics, depths

    def _get_top_lidar_pose(self, rec):
        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        trans = -np.array(egopose['translation'])
        yaw = Quaternion(egopose['rotation']).yaw_pitch_roll[0]
        rot = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)]).inverse
        return trans, rot

    def get_depth_from_lidar(self, lidar_sample, cam_sample):
        points, coloring, im = self.nusc_exp.map_pointcloud_to_image(lidar_sample, cam_sample)
        tmp_cam = np.zeros((self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH))
        points = points.astype(np.int)
        tmp_cam[points[1, :], points[0,:]] = coloring
        tmp_cam = torch.from_numpy(tmp_cam).unsqueeze(0).unsqueeze(0)
        tmp_cam = F.interpolate(tmp_cam, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear', align_corners=False, recompute_scale_factor=True)
        tmp_cam = tmp_cam.squeeze()
        crop = self.augmentation_parameters['crop']
        tmp_cam = tmp_cam[crop[1]:crop[3], crop[0]:crop[2]]
        tmp_cam = torch.round(tmp_cam)
        return tmp_cam


    def get_birds_eye_view_label(self, rec, instance_map, in_pred):
        translation, rotation = self._get_top_lidar_pose(rec)
        segmentation = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        pedestrian = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        # Background is ID 0
        instance = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))

        for annotation_token in rec['anns']:
            # Filter out all non vehicle instances
            annotation = self.nusc.get('sample_annotation', annotation_token)

            if self.cfg.DATASET.FILTER_INVISIBLE_VEHICLES and int(annotation['visibility_token']) == 1 and in_pred is False:
                continue
            if in_pred is True and annotation['instance_token'] not in instance_map:
                continue

            # NuScenes filter
            if 'vehicle' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                instance_id = instance_map[annotation['instance_token']]
                poly_region, z = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(instance, [poly_region], instance_id)
                cv2.fillPoly(segmentation, [poly_region], 1.0)
            elif 'human' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                poly_region, z = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(pedestrian, [poly_region], 1.0)


        return segmentation, instance, pedestrian, instance_map

    def _get_poly_region_in_image(self, instance_annotation, ego_translation, ego_rotation):
        box = Box(
            instance_annotation['translation'], instance_annotation['size'], Quaternion(instance_annotation['rotation'])
        )
        box.translate(ego_translation)
        box.rotate(ego_rotation)

        pts = box.bottom_corners()[:2].T
        pts = np.round((pts - self.bev_start_position[:2] + self.bev_resolution[:2] / 2.0) / self.bev_resolution[:2]).astype(np.int32)
        pts[:, [1, 0]] = pts[:, [0, 1]]

        z = box.bottom_corners()[2, 0]
        return pts, z

    def get_label(self, rec, instance_map, in_pred):
        segmentation_np, instance_np, pedestrian_np, instance_map = \
            self.get_birds_eye_view_label(rec, instance_map, in_pred)
        segmentation = torch.from_numpy(segmentation_np).long().unsqueeze(0).unsqueeze(0)
        instance = torch.from_numpy(instance_np).long().unsqueeze(0)
        pedestrian = torch.from_numpy(pedestrian_np).long().unsqueeze(0).unsqueeze(0)

        return segmentation, instance, pedestrian, instance_map

    def get_future_egomotion(self, rec, index):
        rec_t0 = rec

        # Identity
        future_egomotion = np.eye(4, dtype=np.float32)

        if index < len(self.ixes) - 1:
            rec_t1 = self.ixes[index + 1]

            if rec_t0['scene_token'] == rec_t1['scene_token']:
                egopose_t0 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t0['data']['LIDAR_TOP'])['ego_pose_token']
                )
                egopose_t1 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t1['data']['LIDAR_TOP'])['ego_pose_token']
                )

                egopose_t0 = convert_egopose_to_matrix_numpy(egopose_t0)
                egopose_t1 = convert_egopose_to_matrix_numpy(egopose_t1)

                future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
                future_egomotion[3, :3] = 0.0
                future_egomotion[3, 3] = 1.0

        future_egomotion = torch.Tensor(future_egomotion).float()

        # Convert to 6DoF vector
        future_egomotion = mat2pose_vec(future_egomotion)
        return future_egomotion.unsqueeze(0)

    def get_trajectory_sampling(self, rec=None, sample_indice=None):
        if rec is None and sample_indice is None:
            raise ValueError("No valid input rec or token")
        if rec is None and sample_indice is not None:
            rec = self.ixes[sample_indice]

        ref_scene = self.nusc.get("scene", rec['scene_token'])

        # vm_msgs = self.nusc_can.get_messages(ref_scene['name'], 'vehicle_monitor')
        # vm_uts = [msg['utime'] for msg in vm_msgs]
        pose_msgs = self.nusc_can.get_messages(ref_scene['name'],'pose')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_msgs = self.nusc_can.get_messages(ref_scene['name'], 'steeranglefeedback')
        steer_uts = [msg['utime'] for msg in steer_msgs]

        ref_utime = rec['timestamp']
        # vm_index = locate_message(vm_uts, ref_utime)
        # vm_data = vm_msgs[vm_index]
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]

        # initial speed
        # v0 = vm_data["vehicle_speed"] / 3.6  # km/h to m/s
        v0 = pose_data["vel"][0]  # [0] means longitudinal velocity  m/s

        # curvature (positive: turn left)
        # steering = np.deg2rad(vm_data["steering"])
        steering = steer_data["value"]

        location = self.scene2map[ref_scene['name']]
        # flip x axis if in left-hand traffic (singapore)
        flip_flag = True if location.startswith('singapore') else False
        if flip_flag:
            steering *= -1
        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories

    def generate_centerlines(self, rec, map_name):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        x, y = egopose["translation"][0], egopose["translation"][1]
        nusc_map = self.nusc_maps[map_name]
        radius_init = 2
        closest_lane = None
        while closest_lane == None or len(closest_lane) < 5:
            # keep increasing search space until lane is found
            closest_lane = nusc_map.get_closest_lane(x, y, radius=radius_init)
            radius_init *= 2
        # print(closest_lane, "AAAAAAAAAAAAAAA")
        candidates_future = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        candidates_past = dfs_incoming(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        lane_ids = []
        for past_lane_seq in candidates_past:
            for future_lane_seq in candidates_future:
                lane_ids.append(past_lane_seq + future_lane_seq[1:])
        lane_ids = remove_overlapping_lane_seq(lane_ids)
        centerlines = []
        for lane_id in lane_ids:
            centerline = []
            for lane in lane_id:
                lane_record = nusc_map.get_arcline_path(lane)
                cx = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                centerline.extend(cx)
            centerlines.append(np.array(centerline))
        lanes = self.nusc_maps[map_name].get_records_in_radius(x, y, 2000, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        discrete_points = self.nusc_maps[map_name].discretize_lanes(lanes, resolution_meters=resolution_meters)
        lanes = []
        for lane_id, points in discrete_points.items():
            lanes.append(np.array(points)[:, :2])
        
        patch_box = box_coords
        patch_angle = rot * 180 / np.pi

        patch_x = patch_box[0]
        patch_y = patch_box[1]

        patch = get_patch_coord(patch_box, patch_angle)

        line_list = []
        for lane in centerlines:
            line = LineString([(arr[0], arr[1]) for arr in lane[:, :2]])
            c = lane[:, :2]
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.rotate(new_line, -patch_angle, origin=(patch_x, patch_y), use_radians=False)
                new_line = affinity.affine_transform(new_line,
                                                        [1.0, 0.0, 0.0, 1.0, -patch_x, -patch_y])
                line_list.append(new_line)
        
        local_box = (0.0, 0.0, patch_box[2], patch_box[3])
        patch_x, patch_y, patch_h, patch_w = local_box

        patch = get_patch_coord(local_box)

        canvas_h = canvas_size[0]
        canvas_w = canvas_size[1]
        scale_height = canvas_h/patch_h
        scale_width = canvas_w/patch_w

        trans_x = -patch_x + patch_w / 2.0
        trans_y = -patch_y + patch_h / 2.0

        overall_lines = []
        for line in line_list:
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.affine_transform(new_line,
                                                     [1.0, 0.0, 0.0, 1.0, trans_x, trans_y])
                new_line = affinity.scale(new_line, xfact=scale_width, yfact=scale_height, origin=(0, 0))
                overall_lines.append([list(new_line.coords)])

        return overall_lines

    def voxelize_hd_map(self, rec):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        elements = self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS
        hd_features = self.nusc_maps[map_name].get_map_mask(box_coords, rot * 180 / np.pi , elements, canvas_size=canvas_size)
        #traffic = self.hd_traffic_light(map_name, center, stretch, dx, bx, canvas_size)
        #return torch.from_numpy(np.concatenate((hd_features, traffic), axis=0)[None]).float()
        hd_features = torch.from_numpy(hd_features[None]).float()
        hd_features = torch.transpose(hd_features,-2,-1) # (y,x) replace horizontal and vertical coordinates
        return hd_features

    def hd_traffic_light(self, map_name, center, stretch, dx, bx, canvas_size):

        roads = np.zeros(canvas_size)
        my_patch = (
            center[0] - stretch[0],
            center[1] - stretch[1],
            center[0] + stretch[0],
            center[1] + stretch[1],
        )
        tl_token = self.nusc_maps[map_name].get_records_in_patch(my_patch, ['traffic_light'], mode='intersect')['traffic_light']
        polys = []
        for token in tl_token:
            road_token =self.nusc_maps[map_name].get('traffic_light', token)['from_road_block_token']
            pt = self.nusc_maps[map_name].get('road_block', road_token)['polygon_token']
            polygon = self.nusc_maps[map_name].extract_polygon(pt)
            polys.append(np.array(polygon.exterior.xy).T)

        def get_rot(h):
            return torch.Tensor([
                [np.cos(h), np.sin(h)],
                [-np.sin(h), np.cos(h)],
            ])
        # convert to local coordinates in place
        rot = get_rot(np.arctan2(center[3], center[2])).T
        for rowi in range(len(polys)):
            polys[rowi] -= center[:2]
            polys[rowi] = np.dot(polys[rowi], rot)

        for la in polys:
            pts = (la - bx) / dx
            pts = np.int32(np.around(pts))
            cv2.fillPoly(roads, [pts], 1)

        return roads[None]

    def get_gt_trajectory(self, rec, ref_index):
        n_output = self.cfg.N_FUTURE_FRAMES
        gt_trajectory = np.zeros((n_output+1+2, 3), np.float64)

        egopose_cur = get_global_pose(rec, self.nusc, inverse=True)

        for i in range(2):
            index = ref_index - 2 + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]

                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)

                egopose_future = egopose_cur.dot(egopose_future)
                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i, :] = [origin[0], origin[1], theta]

        for i in range(n_output+1):
            index = ref_index + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]

                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)

                egopose_future = egopose_cur.dot(egopose_future)
                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i + 2, :] = [origin[0], origin[1], theta]

        if gt_trajectory[-1][0] >= 2:
            command = 'RIGHT'
        elif gt_trajectory[-1][0] <= -2:
            command = 'LEFT'
        else:
            command = 'FORWARD'

        return gt_trajectory, command

    def get_routed_map(self, gt_points):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        canvas_size = (
            int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
            int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        roads = np.zeros(canvas_size)
        W = 1.85
        pts = np.array([
            [-4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, -W / 2.],
            [-4.084 / 2. + 0.5, -W / 2.],
        ])
        pts = (pts - bx) / dx
        pts[:, [0, 1]] = pts[:, [1, 0]]

        pts = np.int32(np.around(pts))
        cv2.fillPoly(roads, [pts], 1)

        gt_points = gt_points[:-1].numpy()
        # 坐标原点在左上角
        target = pts.copy()
        target[:,0] = pts[:,0] + gt_points[0] / dx[0]
        target[:,1] = pts[:,1] - gt_points[1] / dx[1]
        target = np.int32(np.around(target))
        cv2.fillPoly(roads, [target], 1)
        return roads

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

    def __getitem__(self, index):
        """
        Returns
        -------
            data: dict with the following keys:
                image: torch.Tensor<float> (T, N, 3, H, W)
                    normalised cameras images with T the sequence length, and N the number of cameras.
                intrinsics: torch.Tensor<float> (T, N, 3, 3)
                    intrinsics containing resizing and cropping parameters.
                extrinsics: torch.Tensor<float> (T, N, 4, 4)
                    6 DoF pose from world coordinates to camera coordinates.
                segmentation: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                    (H_bev, W_bev) are the pixel dimensions in bird's-eye view.
                instance: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                centerness: torch.Tensor<float> (T, 1, H_bev, W_bev)
                offset: torch.Tensor<float> (T, 2, H_bev, W_bev)
                flow: torch.Tensor<float> (T, 2, H_bev, W_bev)
                future_egomotion: torch.Tensor<float> (T, 6)
                    6 DoF egomotion t -> t+1

        """
        data = {}
        keys = ['centerlines', 'gt_trajectory']
        for key in keys:
            data[key] = []

        instance_map = {}
        # Loop over all the frames in the sequence.
        total_centerlines = []
        for i, index_t in enumerate(self.indices[index]):
            if i >= self.receptive_field:
                in_pred = True
            else:
                in_pred = False
            rec = self.ixes[index_t]
            map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]
            """
                get lanes in 100m radius
            """
            if i == self.cfg.TIME_RECEPTIVE_FIELD-1:
                gt_trajectory, command = self.get_gt_trajectory(rec, index_t)
                gt_trajectory_prev = gt_trajectory
                traj=torch.tensor(gt_trajectory_prev)[:, :2].reshape(1,9,2)
                gt_trajectory = gt_trajectory[2:]
                data['gt_trajectory'] = torch.from_numpy(gt_trajectory).float()
                trajs = self.get_trajectory_sampling(rec)

                """
                    CREATE GRID
                    grid: (M, 2) -> sample grid
                """
                sample_trajectories = []

            if self.cfg.PLANNING.DENSE and i == self.cfg.TIME_RECEPTIVE_FIELD-1:
                """
                    ITERATE OVER CENTERLINES
                """
                centerlines = self.generate_centerlines(rec, map_name)
                data['centerlines'] = centerlines


        for key, value in data.items():
            if key in ['image', 'intrinsics', 'extrinsics', 'depths', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian']:
                if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
                    continue
                data[key] = torch.cat(value, dim=0)

        data['target_point'] = torch.tensor([0., 0.])
        instance_centerness, instance_offset, instance_flow = convert_instance_mask_to_center_and_offset_label(
            data['instance'], data['future_egomotion'],
            num_instances=len(instance_map), ignore_index=self.cfg.DATASET.IGNORE_INDEX, subtract_egomotion=True,
            spatial_extent=self.spatial_extent,
        )
        data['centerness'] = instance_centerness
        data['offset'] = instance_offset
        data['flow'] = instance_flow

        return data
data/stp3/datas/CenterlinesNuscenesData.py
llmbev/talk2bev
import os
from PIL import Image

import numpy as np
import cv2
import torch
import torch.utils.data
import torch.nn.functional as F
import torchvision
import os.path as osp

from shapely.geometry import LineString

from pyquaternion import Quaternion
from shapely import affinity
from nuscenes.nuscenes import NuScenes, NuScenesExplorer
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.splits import create_splits_scenes
from nuscenes.utils.data_classes import Box
from nuscenes.lidarseg.lidarseg_utils import colormap_to_colors, plt_to_cv2, get_stats, get_labels_in_coloring, create_lidarseg_legend, paint_points_label
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.eval.common.utils import quaternion_yaw
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from nuscenes.utils.data_classes import PointCloud, LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix
from nuscenes.utils.data_io import load_bin_file
from stp3.utils.tools import ( gen_dx_bx, get_nusc_maps)
import matplotlib.pyplot as plt 
from math import cos, sin

from stp3.utils.geometry import (
    resize_and_crop_image,
    update_intrinsics,
    calculate_birds_eye_view_parameters,
    convert_egopose_to_matrix_numpy,
    pose_vec2mat,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
    get_global_pose,
    get_patch_coord,
    dfs,
    dfs_incoming,
    remove_overlapping_lane_seq,
    extend_both_sides,
    project_to_frenet_frame,
    project_to_cartesian_frame,
    interp_arc,
    interp_arc_np,
    denoise,
    interp_polyline_by_fixed_waypt_interval
)
from stp3.utils.instance import convert_instance_mask_to_center_and_offset_label
import stp3.utils.sampler as trajectory_sampler
from stp3.utils.spline import Spline2D
from stp3.utils.Optnode import *

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

resolution_meters = 0.1
t_fin = 3.5
num = 61
nvar = 11
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
debug = False
brdebug = False
problem = OPTNode(t_fin=t_fin, num=num)

def distance(x1, y1, z1, x2, y2, z2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + 1 * (z1 - z2) ** 2) ** ( 1 / 2)
 
# Function to calculate K closest points
def kClosest(points, target, K):
    pts = []
    n = len(points)
    d = []

    for i in range(n):
        d.append({
            "first": distance(points[i][0], points[i][1], points[i][2], target[0], target[1], target[2]),
            "second": i
        })
     
    d = sorted(d, key=lambda l:l["first"])
 
    for i in range(K):
        pt = []
        pt.append(points[d[i]["second"]][0])
        pt.append(points[d[i]["second"]][1])
        pt.append(points[d[i]["second"]][2])
        pt.append(points[d[i]["second"]][3])
        pts.append(pt)
 
    return pts

class FuturePredictionDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5 #SECOND
    def __init__(self, nusc, is_train, cfg):
        self.nusc = nusc
        self.dataroot = self.nusc.dataroot
        self.nusc_exp = NuScenesExplorer(nusc)
        self.nusc_can = NuScenesCanBus(dataroot=self.dataroot)
        self.is_train = is_train
        self.cfg = cfg
        self.problem = OPTNode(t_fin=3.5, num=self.cfg.N_FUTURE_FRAMES + 1)

        if self.is_train == 0:
            self.mode = 'train'
        elif self.is_train == 1:
            self.mode = 'val'
        elif self.is_train == 2:
            self.mode = 'test'
        else:
            raise NotImplementedError


        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD

        self.emergency_ixes = []
        self.scenes = self.get_scenes()
        self.ixes = self.prepro()
        self.indices = self.get_indices()

        # import pdb; pdb.set_trace()

        # Image resizing and cropping
        self.augmentation_parameters = self.get_resizing_and_cropping_parameters()

        # Normalising input images
        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # Bird's-eye view parameters
        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_resolution, self.bev_start_position, self.bev_dimension = (
            bev_resolution.numpy(), bev_start_position.numpy(), bev_dimension.numpy()
        )

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # The number of sampled trajectories
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        # HD-map feature extractor
        self.nusc_maps = get_nusc_maps(self.cfg.DATASET.MAP_FOLDER)
        self.scene2map = {}
        for sce in self.nusc.scene:
            log = self.nusc.get('log', sce['log_token'])
            self.scene2map[sce['name']] = log['location']
        self.save_dir = cfg.DATASET.SAVE_DIR

        self.curbatchnum1 = None
        self.curbatchnum2 = None
        if self.cfg.PLANNING.DENSE.ENABLED:
            self.centerlines_path = self.cfg.PLANNING.DENSE.PATH

    def get_scenes(self):
        # filter by scene split
        split = {'v1.0-trainval': {0: 'train', 1: 'val', 2: 'test'},
                 'v1.0-mini': {0: 'mini_train', 1: 'mini_val'},}[
            self.nusc.version
        ][self.is_train]

        blacklist = [419] + self.nusc_can.can_blacklist  # # scene-0419 does not have vehicle monitor data
        blacklist = ['scene-' + str(scene_no).zfill(4) for scene_no in blacklist]

        scenes = create_splits_scenes()[split][:]
        for scene_no in blacklist:
            if scene_no in scenes:
                scenes.remove(scene_no)

        return scenes

    def prepro(self):
        samples = [samp for samp in self.nusc.sample]

        # remove samples that aren't in this split
        samples = [samp for samp in samples if self.nusc.get('scene', samp['scene_token'])['name'] in self.scenes]
        
        # sort by scene, timestamp (only to make chronological viz easier)
        samples.sort(key=lambda x: (x['scene_token'], x['timestamp']))

        for sample in samples:
            for ann in sample['anns']:
                category = self.nusc.get("sample_annotation", ann)["category_name"]
                if "police" in category:
                    self.emergency_ixes.append(sample)
                    break
        self.ixes = self.emergency_ixes

        return samples

    def get_indices(self):
        indices = []
        for index in range(len(self.ixes)):
            # print(index, end=' ')
            is_valid_data = True
            previous_rec = None
            current_indices = []
            # if self.ixes[index] not in self.emergency_ixes:
            #     continue
            for t in range(self.sequence_length):
                index_t = index + t
                # Going over the dataset size limit.
                if index_t >= len(self.ixes):
                    is_valid_data = False
                    break
                rec = self.ixes[index_t]
                # Check if scene is the same
                if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
                    is_valid_data = False
                    break

                current_indices.append(index_t)
                previous_rec = rec

            if is_valid_data:
                indices.append(current_indices)

        return np.asarray(indices)

    # def get_indices(self):
    #     indices = []
    #     scene_token_mapping = {}
    #     ixes_new = []
    #     for index in range(len(self.ixes)):
    #         if scene_token_mapping.get(self.ixes[index]['scene_token']) == None:
    #             scene_token_mapping[self.ixes[index]['scene_token']] = 1
    #             continue
    #         scene_token_mapping[self.ixes[index]['scene_token']] = scene_token_mapping[self.ixes[index]['scene_token']] + 1
    #         is_valid_data = True
    #         previous_rec = None
    #         current_indices = []
    #         for t in range(self.sequence_length):
    #             index_t = index + t
    #             # Going over the dataset size limit.
    #             if index_t >= len(self.ixes):
    #                 is_valid_data = False
    #                 break
    #             rec = self.ixes[index_t]
    #             ixes_new.append(self.ixes[index_t])
    #             # Check if scene is the same
    #             if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
    #                 is_valid_data = False
    #                 break

    #             current_indices.append(len(ixes_new) - 1)
    #             previous_rec = rec

    #         if is_valid_data:
    #             indices.append(current_indices)

    #     self.ixes = ixes_new
    #     return np.asarray(indices)

    def get_resizing_and_cropping_parameters(self):
        original_height, original_width = self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH
        final_height, final_width = self.cfg.IMAGE.FINAL_DIM

        resize_scale = self.cfg.IMAGE.RESIZE_SCALE
        resize_dims = (int(original_width * resize_scale), int(original_height * resize_scale))
        resized_width, resized_height = resize_dims

        crop_h = self.cfg.IMAGE.TOP_CROP
        crop_w = int(max(0, (resized_width - final_width) / 2))
        # Left, top, right, bottom crops.
        crop = (crop_w, crop_h, crop_w + final_width, crop_h + final_height)

        if resized_width != final_width:
            print('Zero padding left and right parts of the image.')
        if crop_h + final_height != resized_height:
            print('Zero padding bottom part of the image.')

        return {'scale_width': resize_scale,
                'scale_height': resize_scale,
                'resize_dims': resize_dims,
                'crop': crop,
                }

    def get_input_data(self, rec):
        """
        Parameters
        ----------
            rec: nuscenes identifier for a given timestamp

        Returns
        -------
            images: torch.Tensor<float> (N, 3, H, W)
            intrinsics: torch.Tensor<float> (3, 3)
            extrinsics: torch.Tensor(N, 4, 4)
        """
        images = []
        unnormalized_images = []
        intrinsics = []
        extrinsics = []
        depths = []
        translations = []
        rotations = []
        camera_intrinsics = []
        point_clouds = []
        point_clouds_labels = []
        point_clouds_panoptic_labels = []
        panoptic_mappings_list = []
        extra_rotations = []
        extra_translations = []
        cameras = self.cfg.IMAGE.NAMES

        # # The extrinsics we want are from the camera sensor to "flat egopose" as defined
        # # https://github.com/nutonomy/nuscenes-devkit/blob/9b492f76df22943daf1dc991358d3d606314af27/python-sdk/nuscenes/nuscenes.py#L279
        # # which corresponds to the position of the lidar.
        # # This is because the labels are generated by projecting the 3D bounding box in this lidar's reference frame.

        # # From lidar egopose to world.
        lidar_sample = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        lidar_pose = self.nusc.get('ego_pose', lidar_sample['ego_pose_token'])
        yaw = Quaternion(lidar_pose['rotation']).yaw_pitch_roll[0]
        lidar_rotation = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        lidar_translation = np.array(lidar_pose['translation'])[:, None]
        lidar_to_world = np.vstack([
            np.hstack((lidar_rotation.rotation_matrix, lidar_translation)),
            np.array([0, 0, 0, 1])
        ])

        camera_sample_ = self.nusc.get('sample_data', rec['data']["CAM_FRONT"])
        my_sample = self.nusc.get('sample', camera_sample_['sample_token'])
        my_scene_token = my_sample['scene_token']

        # # Here we just grab the front camera and the point sensor.
        pointsensor = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        pcl_path = osp.join(self.nusc.dataroot, pointsensor['filename'])
        # panoptic_mappings = [samp.token for samp in self.nusc.get_sample_data(rec['data']['LIDAR_TOP'])[1]]
        panoptic_mappings = []
        for ok in rec['anns']:
            panoptic_mappings.append(self.nusc.get('sample_annotation', ok)['token'])
        panoptic_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get('panoptic', rec['data']['LIDAR_TOP'])['filename'])
        panoptic_labels = load_bin_file(panoptic_labels_filename, type='panoptic').astype(np.int32)
        point_clouds_panoptic_labels.append(panoptic_labels)
        for cam in cameras:
            camera_sample = self.nusc.get('sample_data', rec['data'][cam])
            sample_token = camera_sample['token']

            camm = self.nusc.get('sample_data', rec['data'][cam])
            lidarseg_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get("lidarseg", rec['data']['LIDAR_TOP'])['filename'])
            points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)


            cs_record = self.nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token'])
            pc = LidarPointCloud.from_file(pcl_path)
            ### LIDAR TO EGO_{TL} TRANSLATION -> CONSTANT
            pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)
            pc.translate(np.array(cs_record['translation']))

            point_clouds.append(torch.tensor(np.array(pc.points)).unsqueeze(0).unsqueeze(0)) # 1, 1, N, 4
            point_clouds_labels.append(torch.tensor(np.array(points_label)).unsqueeze(0).unsqueeze(0)) # 1, 1, N, 

            poserecord = self.nusc.get('ego_pose', pointsensor['ego_pose_token'])
            extra_rotations.append(Quaternion(poserecord['rotation']).rotation_matrix)
            extra_translations.append(np.array(poserecord['translation']))
            pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)
            pc.translate(np.array(poserecord['translation']))

            # Third step: transform from global into the ego vehicle frame for the timestamp of the image.
            poserecord = self.nusc.get('ego_pose', camm['ego_pose_token'])
            extra_rotations.append(Quaternion(poserecord['rotation']).rotation_matrix.T)
            extra_translations.append(np.array(poserecord['translation']))            
            pc.translate(-np.array(poserecord['translation']))
            pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix.T)

            # Transformation from world to egopose
            car_egopose = self.nusc.get('ego_pose', camera_sample['ego_pose_token'])
            egopose_rotation = Quaternion(car_egopose['rotation']).inverse
            egopose_translation = -np.array(car_egopose['translation'])[:, None]
            world_to_car_egopose = np.vstack([
                np.hstack((egopose_rotation.rotation_matrix, egopose_rotation.rotation_matrix @ egopose_translation)),
                np.array([0, 0, 0, 1])
            ])

            # From egopose to sensor
            sensor_sample = self.nusc.get('calibrated_sensor', camera_sample['calibrated_sensor_token'])
            intrinsic = torch.Tensor(sensor_sample['camera_intrinsic'])
            sensor_rotation = Quaternion(sensor_sample['rotation'])
            sensor_translation = np.array(sensor_sample['translation'])[:, None]
            car_egopose_to_sensor = np.vstack([
                np.hstack((sensor_rotation.rotation_matrix, sensor_translation)),
                np.array([0, 0, 0, 1])
            ])
            car_egopose_to_sensor = np.linalg.inv(car_egopose_to_sensor)

            # Combine all the transformation.
            # From sensor to lidar.
            lidar_to_sensor = car_egopose_to_sensor @ world_to_car_egopose @ lidar_to_world
            sensor_to_lidar = torch.from_numpy(np.linalg.inv(lidar_to_sensor)).float()

            # Load image
            image_filename = os.path.join(self.dataroot, camera_sample['filename'])
            img = Image.open(image_filename)
            unnormalized_image = torch.tensor(np.array(img))
            # Resize and crop
            img = resize_and_crop_image(
                img, resize_dims=self.augmentation_parameters['resize_dims'], crop=self.augmentation_parameters['crop']
            )
            # Normalise image
            normalised_img = self.normalise_image(img)

            # # Combine resize/cropping in the intrinsics
            top_crop = self.augmentation_parameters['crop'][1]
            left_crop = self.augmentation_parameters['crop'][0]
            intrinsic = update_intrinsics(
                intrinsic, top_crop, left_crop,
                scale_width=self.augmentation_parameters['scale_width'],
                scale_height=self.augmentation_parameters['scale_height']
            )

            # Get Depth
            # Depth data should under the dataroot path 
            if self.cfg.LIFT.GT_DEPTH:
                base_root = os.path.join(self.dataroot, 'depths') 
                filename = os.path.basename(camera_sample['filename']).split('.')[0] + '.npy'
                depth_file_name = os.path.join(base_root, cam, 'npy', filename)
                depth = torch.from_numpy(np.load(depth_file_name)).unsqueeze(0).unsqueeze(0)
                depth = F.interpolate(depth, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear')
                depth = depth.squeeze()
                crop = self.augmentation_parameters['crop']
                depth = depth[crop[1]:crop[3], crop[0]:crop[2]]
                depth = torch.round(depth)
                depths.append(depth.unsqueeze(0).unsqueeze(0))
            
            camm = self.nusc.get('sample_data', rec['data'][cam])
            cs_record = self.nusc.get('calibrated_sensor', camm['calibrated_sensor_token'])
            images.append(normalised_img.unsqueeze(0).unsqueeze(0))
            intrinsics.append(intrinsic.unsqueeze(0).unsqueeze(0))
            extrinsics.append(sensor_to_lidar.unsqueeze(0).unsqueeze(0))
            unnormalized_images.append(unnormalized_image.unsqueeze(0).unsqueeze(0))
            rotations.append(torch.tensor(Quaternion(cs_record['rotation']).rotation_matrix.T).unsqueeze(0).unsqueeze(0))
            translations.append(torch.tensor(-np.array(cs_record['translation'])).unsqueeze(0).unsqueeze(0))
            camera_intrinsics.append(torch.tensor(np.array(cs_record['camera_intrinsic'])).unsqueeze(0).unsqueeze(0))

        images, intrinsics, extrinsics, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, panoptic_mappings_list = (torch.cat(images, dim=1),
                                            torch.cat(intrinsics, dim=1),
                                            torch.cat(extrinsics, dim=1),
                                            torch.cat(unnormalized_images, dim=1),
                                            torch.cat(rotations, dim=1),
                                            torch.cat(translations, dim=1),
                                            torch.cat(camera_intrinsics, dim=1),
                                            torch.cat(point_clouds, dim=1),
                                            torch.cat(point_clouds_labels, dim=1),
                                            point_clouds_panoptic_labels,
                                            panoptic_mappings
                                          )
        if len(depths) > 0:
            depths = torch.cat(depths, dim=1)

        return images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, my_scene_token, panoptic_mappings_list

    def _get_top_lidar_pose(self, rec):
        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        trans = -np.array(egopose['translation'])
        yaw = Quaternion(egopose['rotation']).yaw_pitch_roll[0]
        # rot = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        rot = Quaternion(egopose['rotation']).inverse
        return trans, rot

    def get_depth_from_lidar(self, lidar_sample, cam_sample):
        points, coloring, im = self.nusc_exp.map_pointcloud_to_image(lidar_sample, cam_sample)
        tmp_cam = np.zeros((self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH))
        points = points.astype(np.int)
        tmp_cam[points[1, :], points[0,:]] = coloring
        tmp_cam = torch.from_numpy(tmp_cam).unsqueeze(0).unsqueeze(0)
        tmp_cam = F.interpolate(tmp_cam, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear', align_corners=False, recompute_scale_factor=True)
        tmp_cam = tmp_cam.squeeze()
        crop = self.augmentation_parameters['crop']
        tmp_cam = tmp_cam[crop[1]:crop[3], crop[0]:crop[2]]
        tmp_cam = torch.round(tmp_cam)
        return tmp_cam


    def get_birds_eye_view_label(self, rec, instance_map, in_pred):
        translation, rotation = self._get_top_lidar_pose(rec)
        segmentation = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        pedestrian = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        # Background is ID 0
        instance = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        categories_map = []
        bottom_corners = {}

        for annotation_token in rec['anns']:
            # Filter out all non vehicle instances
            annotation = self.nusc.get('sample_annotation', annotation_token)

            if self.cfg.DATASET.FILTER_INVISIBLE_VEHICLES and int(annotation['visibility_token']) == 1 and in_pred is False:
                continue
            if in_pred is True and annotation['instance_token'] not in instance_map:
                continue

            # NuScenes filter
            if 'vehicle' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                instance_id = instance_map[annotation['instance_token']]
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(instance, [poly_region], instance_id)
                cv2.fillPoly(segmentation, [poly_region], 1.0)
                categories_map.append([annotation, np.mean(corners, axis=1)])
                bottom_corners[annotation['token']] = corners
            elif 'human' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(pedestrian, [poly_region], 1.0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def _get_poly_region_in_image(self, instance_annotation, ego_translation, ego_rotation):
        box = Box(
            instance_annotation['translation'], instance_annotation['size'], Quaternion(instance_annotation['rotation'])
        )
        # import pdb; pdb.set_trace()
        box.translate(ego_translation)
        box.rotate(ego_rotation)
        corners = box.corners()

        pts = box.bottom_corners()[:2].T
        pts = np.round((pts - self.bev_start_position[:2] + self.bev_resolution[:2] / 2.0) / self.bev_resolution[:2]).astype(np.int32)
        pts[:, [1, 0]] = pts[:, [0, 1]]

        z = box.bottom_corners()[2, 0]
        return pts, z, corners

    def get_label(self, rec, instance_map, in_pred):
        segmentation_np, instance_np, pedestrian_np, instance_map, categories_map, bottom_corners = \
            self.get_birds_eye_view_label(rec, instance_map, in_pred)
        segmentation = torch.from_numpy(segmentation_np).long().unsqueeze(0).unsqueeze(0)
        instance = torch.from_numpy(instance_np).long().unsqueeze(0)
        pedestrian = torch.from_numpy(pedestrian_np).long().unsqueeze(0).unsqueeze(0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def get_future_egomotion(self, rec, index):
        rec_t0 = rec

        # Identity
        future_egomotion = np.eye(4, dtype=np.float32)

        if index < len(self.ixes) - 1:
            rec_t1 = self.ixes[index + 1]

            if rec_t0['scene_token'] == rec_t1['scene_token']:
                egopose_t0 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t0['data']['LIDAR_TOP'])['ego_pose_token']
                )
                egopose_t1 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t1['data']['LIDAR_TOP'])['ego_pose_token']
                )

                egopose_t0 = convert_egopose_to_matrix_numpy(egopose_t0)
                egopose_t1 = convert_egopose_to_matrix_numpy(egopose_t1)

                future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
                future_egomotion[3, :3] = 0.0
                future_egomotion[3, 3] = 1.0

        future_egomotion = torch.Tensor(future_egomotion).float()

        # Convert to 6DoF vector
        future_egomotion = mat2pose_vec(future_egomotion)
        return future_egomotion.unsqueeze(0)

    def get_trajectory_sampling(self, rec=None, sample_indice=None):
        if rec is None and sample_indice is None:
            raise ValueError("No valid input rec or token")
        if rec is None and sample_indice is not None:
            rec = self.ixes[sample_indice]

        ref_scene = self.nusc.get("scene", rec['scene_token'])

        # vm_msgs = self.nusc_can.get_messages(ref_scene['name'], 'vehicle_monitor')
        # vm_uts = [msg['utime'] for msg in vm_msgs]
        pose_msgs = self.nusc_can.get_messages(ref_scene['name'],'pose')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_msgs = self.nusc_can.get_messages(ref_scene['name'], 'steeranglefeedback')
        steer_uts = [msg['utime'] for msg in steer_msgs]

        ref_utime = rec['timestamp']
        # vm_index = locate_message(vm_uts, ref_utime)
        # vm_data = vm_msgs[vm_index]
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]

        # initial speed
        # v0 = vm_data["vehicle_speed"] / 3.6  # km/h to m/s
        v0 = pose_data["vel"][0]  # [0] means longitudinal velocity  m/s

        # curvature (positive: turn left)
        # steering = np.deg2rad(vm_data["steering"])
        steering = steer_data["value"]

        location = self.scene2map[ref_scene['name']]
        # flip x axis if in left-hand traffic (singapore)
        flip_flag = True if location.startswith('singapore') else False
        if flip_flag:
            steering *= -1
        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories


    def generate_centerlines(self, rec, map_name):


        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        x, y = egopose["translation"][0], egopose["translation"][1]
        nusc_map = self.nusc_maps[map_name]
        radius_init = 2
        closest_lane = None
        while closest_lane == None or len(closest_lane) < 5:
            # keep increasing search space until lane is found
            closest_lane = nusc_map.get_closest_lane(x, y, radius=radius_init)
            radius_init *= 2
        candidates_future = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        candidates_past = dfs_incoming(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        lane_ids = []
        for past_lane_seq in candidates_past:
            for future_lane_seq in candidates_future:
                lane_ids.append(past_lane_seq + future_lane_seq[1:])
        lane_ids = remove_overlapping_lane_seq(lane_ids)
        centerlines = []
        for lane_id in lane_ids:
            centerline = []
            for lane in lane_id:
                lane_record = nusc_map.get_arcline_path(lane)
                cx = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                centerline.extend(cx)
            centerlines.append(np.array(centerline))
        lanes = self.nusc_maps[map_name].get_records_in_radius(x, y, 2000, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        discrete_points = self.nusc_maps[map_name].discretize_lanes(lanes, resolution_meters=resolution_meters)
        lanes = []
        for lane_id, points in discrete_points.items():
            lanes.append(np.array(points)[:, :2])
     
        patch_box = box_coords
        patch_angle = rot * 180 / np.pi

        patch_x = patch_box[0]
        patch_y = patch_box[1]

        patch = get_patch_coord(patch_box, patch_angle)

        line_list = []
        for lane in centerlines:
            line = LineString([(arr[0], arr[1]) for arr in lane[:, :2]])
            c = lane[:, :2]
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.rotate(new_line, -patch_angle, origin=(patch_x, patch_y), use_radians=False)
                new_line = affinity.affine_transform(new_line,
                                                        [1.0, 0.0, 0.0, 1.0, -patch_x, -patch_y])
                line_list.append(new_line)
        
        local_box = (0.0, 0.0, patch_box[2], patch_box[3])
        patch_x, patch_y, patch_h, patch_w = local_box

        patch = get_patch_coord(local_box)

        canvas_h = canvas_size[0]
        canvas_w = canvas_size[1]
        scale_height = canvas_h/patch_h
        scale_width = canvas_w/patch_w

        trans_x = -patch_x + patch_w / 2.0
        trans_y = -patch_y + patch_h / 2.0

        overall_lines = []
        for line in line_list:
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.affine_transform(new_line,
                                                     [1.0, 0.0, 0.0, 1.0, trans_x, trans_y])
                new_line = affinity.scale(new_line, xfact=scale_width, yfact=scale_height, origin=(0, 0))
                # overall_lines.append([list(new_line.coords)])

                gg = np.array(list(new_line.coords))
                # interp_arc(t: int, px: np.ndarray, py: np.ndarray) -> np.ndarray:
                gg = interp_arc(t=2000, px=gg[:,0], py=gg[:,1])
                overall_lines.append(gg)

        return overall_lines

    def voxelize_hd_map(self, rec):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        elements = self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS
        hd_features = self.nusc_maps[map_name].get_map_mask(box_coords, rot * 180 / np.pi , elements, canvas_size=canvas_size)
        #traffic = self.hd_traffic_light(map_name, center, stretch, dx, bx, canvas_size)
        #return torch.from_numpy(np.concatenate((hd_features, traffic), axis=0)[None]).float()
        hd_features = torch.from_numpy(hd_features[None]).float()
        hd_features = torch.transpose(hd_features,-2,-1) # (y,x) replace horizontal and vertical coordinates
        return hd_features

    def hd_traffic_light(self, map_name, center, stretch, dx, bx, canvas_size):

        roads = np.zeros(canvas_size)
        my_patch = (
            center[0] - stretch[0],
            center[1] - stretch[1],
            center[0] + stretch[0],
            center[1] + stretch[1],
        )
        tl_token = self.nusc_maps[map_name].get_records_in_patch(my_patch, ['traffic_light'], mode='intersect')['traffic_light']
        polys = []
        for token in tl_token:
            road_token =self.nusc_maps[map_name].get('traffic_light', token)['from_road_block_token']
            pt = self.nusc_maps[map_name].get('road_block', road_token)['polygon_token']
            polygon = self.nusc_maps[map_name].extract_polygon(pt)
            polys.append(np.array(polygon.exterior.xy).T)

        def get_rot(h):
            return torch.Tensor([
                [np.cos(h), np.sin(h)],
                [-np.sin(h), np.cos(h)],
            ])
        # convert to local coordinates in place
        rot = get_rot(np.arctan2(center[3], center[2])).T
        for rowi in range(len(polys)):
            polys[rowi] -= center[:2]
            polys[rowi] = np.dot(polys[rowi], rot)

        for la in polys:
            pts = (la - bx) / dx
            pts = np.int32(np.around(pts))
            cv2.fillPoly(roads, [pts], 1)

        return roads[None]

    def get_gt_trajectory(self, rec, ref_index):
        n_output = self.cfg.N_FUTURE_FRAMES
        n_past = self.cfg.TIME_RECEPTIVE_FIELD
        gt_trajectory = np.zeros((n_output+1+(n_past-1), 3), np.float64)

        """
            egopos_cur -> gives you sensor to global.
                Suppose a point P in sensor-frame.
                    egopos_cur.P gives you the point's global coordinates
        """
        egopose_cur = get_global_pose(rec, self.nusc, inverse=True)

        for i in range(n_past - 1):
            index = ref_index - (n_past - 1) + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i, :] = [origin[0], origin[1], theta]

        for i in range(n_output+1):
            index = ref_index + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                """
                    egopos_future -> gives you global to future.
                        Suppose a point P in ego-frame.
                """
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i + (n_past - 1), :] = [origin[0], origin[1], theta]

        """
            ?
        """
        if gt_trajectory[-1][0] >= 2:
            command = 'RIGHT'
        elif gt_trajectory[-1][0] <= -2:
            command = 'LEFT'
        else:
            command = 'FORWARD'

        return gt_trajectory, command

    def get_routed_map(self, gt_points):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        canvas_size = (
            int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
            int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        roads = np.zeros(canvas_size)
        W = 1.85
        pts = np.array([
            [-4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, -W / 2.],
            [-4.084 / 2. + 0.5, -W / 2.],
        ])
        pts = (pts - bx) / dx
        pts[:, [0, 1]] = pts[:, [1, 0]]

        pts = np.int32(np.around(pts))
        cv2.fillPoly(roads, [pts], 1)

        gt_points = gt_points[:-1].numpy()
        # 坐标原点在左上角
        target = pts.copy()
        target[:,0] = pts[:,0] + gt_points[0] / dx[0]
        target[:,1] = pts[:,1] - gt_points[1] / dx[1]
        target = np.int32(np.around(target))
        cv2.fillPoly(roads, [target], 1)
        return roads

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

    def __getitem__(self, index):
        """
        Returns
        -------
            data: dict with the following keys:
                image: torch.Tensor<float> (T, N, 3, H, W)
                    normalised cameras images with T the sequence length, and N the number of cameras.
                intrinsics: torch.Tensor<float> (T, N, 3, 3)
                    intrinsics containing resizing and cropping parameters.
                extrinsics: torch.Tensor<float> (T, N, 4, 4)
                    6 DoF pose from world coordinates to camera coordinates.
                segmentation: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                    (H_bev, W_bev) are the pixel dimensions in bird's-eye view.
                instance: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                centerness: torch.Tensor<float> (T, 1, H_bev, W_bev)
                offset: torch.Tensor<float> (T, 2, H_bev, W_bev)
                flow: torch.Tensor<float> (T, 2, H_bev, W_bev)
                future_egomotion: torch.Tensor<float> (T, 6)
                    6 DoF egomotion t -> t+1

        """
        data = {}
        keys = ['image', 'intrinsics', 'extrinsics', 'depths',
                'segmentation', 'instance', 'centerness', 'offset', 'flow', 'pedestrian',
                'future_egomotion', 'hdmap', 'gt_trajectory', 'indices', 'poses', 'unnormalized_images',\
                'rotations', 'translations', 'camera_intrinsics', 'point_clouds', 'point_clouds_labels', \
                'point_clouds_panoptic_labels', 'scene_token', 'categories_map', 'panoptic_mappings_list', \
                'bottom_corners'
                ]
        for key in keys:
            data[key] = []

        instance_map = {}
        # Loop over all the frames in the sequence.
        total_centerlines = []
        for i, index_t in enumerate(self.indices[index]):
            if i >= self.receptive_field:
                in_pred = True
            else:
                in_pred = False
            rec = self.ixes[index_t]
            if i == 0:
                data['rec_first'] = rec            
            map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

            if i < self.receptive_field:
                images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, scene_token, panoptic_mappings_list = self.get_input_data(rec)
                data['image'].append(images)
                data['unnormalized_images'].append(unnormalized_images)
                data['intrinsics'].append(intrinsics)
                data['extrinsics'].append(extrinsics)
                data['depths'].append(depths)
                data['translations'].append(translations)
                data['rotations'].append(rotations)
                data['camera_intrinsics'].append(camera_intrinsics)
                data['point_clouds'].append(point_clouds)
                data['point_clouds_labels'].append(point_clouds_labels)
                data['scene_token'].append(scene_token)

            segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners = self.get_label(rec, instance_map, in_pred)

            future_egomotion = self.get_future_egomotion(rec, index_t)
            hd_map_feature = self.voxelize_hd_map(rec)

            data['segmentation'].append(segmentation)
            data['instance'].append(instance)
            data['pedestrian'].append(pedestrian)
            data['future_egomotion'].append(future_egomotion)
            data['hdmap'].append(hd_map_feature)
            data['indices'].append(index_t)
            data['categories_map'].append(categories_map)
            # import pdb; pdb.set_trace()
            data['point_clouds_panoptic_labels'].append(point_clouds_panoptic_labels)
            data['panoptic_mappings_list'].append(panoptic_mappings_list)
            data['bottom_corners'].append(bottom_corners)

        for key, value in data.items():
            if key in ['image', 'intrinsics', 'extrinsics', 'depths', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian', 'camera_intrinsics', 'rotations', 'translations', 'unnormalized_images']:
                if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
                    continue
                data[key] = torch.cat(value, dim=0)

        # images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds

        # for key, value in data.items():
        #     if key in ['image', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian', 'unnormalized_images']:
        #         if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
        #             continue
        #         data[key] = torch.cat(value, dim=0)

        data['target_point'] = torch.tensor([0., 0.])
        instance_centerness, instance_offset, instance_flow = convert_instance_mask_to_center_and_offset_label(
            data['instance'], data['future_egomotion'],
            num_instances=len(instance_map), ignore_index=self.cfg.DATASET.IGNORE_INDEX, subtract_egomotion=True,
            spatial_extent=self.spatial_extent,
        )
        data['centerness'] = instance_centerness
        data['offset'] = instance_offset
        data['flow'] = instance_flow

        return data
data/stp3/datas/NuscenesData.py
llmbev/talk2bev
import os
from PIL import Image

import numpy as np
import cv2
import torch
import torch.utils.data
import torch.nn.functional as F
import torchvision

from shapely.geometry import LineString

from pyquaternion import Quaternion
from shapely import affinity
from nuscenes.nuscenes import NuScenes, NuScenesExplorer
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.splits import create_splits_scenes
from nuscenes.utils.data_classes import Box
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.eval.common.utils import quaternion_yaw
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from stp3.utils.tools import ( gen_dx_bx, get_nusc_maps)
import matplotlib.pyplot as plt 
from math import cos, sin

from stp3.utils.geometry import (
    resize_and_crop_image,
    update_intrinsics,
    calculate_birds_eye_view_parameters,
    convert_egopose_to_matrix_numpy,
    pose_vec2mat,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
    get_global_pose,
    get_patch_coord,
    dfs,
    dfs_incoming,
    remove_overlapping_lane_seq,
    extend_both_sides,
    project_to_frenet_frame,
    project_to_cartesian_frame
)
from stp3.utils.instance import convert_instance_mask_to_center_and_offset_label
import stp3.utils.sampler as trajectory_sampler
from stp3.utils.spline import Spline2D
from stp3.utils.Optnode import *

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

resolution_meters = 0.1
t_fin = 3.5
num = 61
nvar = 11
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
debug = True
brdebug = False
problem = OPTNode(t_fin=t_fin, num=num)

class FuturePredictionDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5 #SECOND
    def __init__(self, nusc, is_train, cfg):
        self.nusc = nusc
        self.dataroot = self.nusc.dataroot
        self.nusc_exp = NuScenesExplorer(nusc)
        self.nusc_can = NuScenesCanBus(dataroot=self.dataroot)
        self.is_train = is_train
        self.cfg = cfg
        self.problem = OPTNode(t_fin=3.5, num=self.cfg.N_FUTURE_FRAMES + 1)

        if self.is_train == 0:
            self.mode = 'train'
        elif self.is_train == 1:
            self.mode = 'val'
        elif self.is_train == 2:
            self.mode = 'test'
        else:
            raise NotImplementedError

        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD

        self.scenes = self.get_scenes()
        self.ixes = self.prepro()
        self.indices = self.get_indices()

        # Image resizing and cropping
        self.augmentation_parameters = self.get_resizing_and_cropping_parameters()

        # Normalising input images
        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # Bird's-eye view parameters
        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_resolution, self.bev_start_position, self.bev_dimension = (
            bev_resolution.numpy(), bev_start_position.numpy(), bev_dimension.numpy()
        )

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # The number of sampled trajectories
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        # HD-map feature extractor
        self.nusc_maps = get_nusc_maps(self.cfg.DATASET.MAP_FOLDER)
        self.scene2map = {}
        for sce in self.nusc.scene:
            log = self.nusc.get('log', sce['log_token'])
            self.scene2map[sce['name']] = log['location']
        self.save_dir = cfg.DATASET.SAVE_DIR

    def get_scenes(self):
        # filter by scene split
        split = {'v1.0-trainval': {0: 'train', 1: 'val', 2: 'test'},
                 'v1.0-mini': {0: 'mini_train', 1: 'mini_val'},}[
            self.nusc.version
        ][self.is_train]

        blacklist = [419] + self.nusc_can.can_blacklist  # # scene-0419 does not have vehicle monitor data
        blacklist = ['scene-' + str(scene_no).zfill(4) for scene_no in blacklist]

        scenes = create_splits_scenes()[split][:]
        for scene_no in blacklist:
            if scene_no in scenes:
                scenes.remove(scene_no)

        return scenes

    def prepro(self):
        samples = [samp for samp in self.nusc.sample]

        # remove samples that aren't in this split
        samples = [samp for samp in samples if self.nusc.get('scene', samp['scene_token'])['name'] in self.scenes]

        # sort by scene, timestamp (only to make chronological viz easier)
        samples.sort(key=lambda x: (x['scene_token'], x['timestamp']))

        return samples

    def get_indices(self):
        indices = []
        for index in range(len(self.ixes)):
            is_valid_data = True
            previous_rec = None
            current_indices = []
            for t in range(self.sequence_length):
                index_t = index + t
                # Going over the dataset size limit.
                if index_t >= len(self.ixes):
                    is_valid_data = False
                    break
                rec = self.ixes[index_t]
                # Check if scene is the same
                if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
                    is_valid_data = False
                    break

                current_indices.append(index_t)
                previous_rec = rec

            if is_valid_data:
                indices.append(current_indices)

        return np.asarray(indices)

    def get_resizing_and_cropping_parameters(self):
        original_height, original_width = self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH
        final_height, final_width = self.cfg.IMAGE.FINAL_DIM

        resize_scale = self.cfg.IMAGE.RESIZE_SCALE
        resize_dims = (int(original_width * resize_scale), int(original_height * resize_scale))
        resized_width, resized_height = resize_dims

        crop_h = self.cfg.IMAGE.TOP_CROP
        crop_w = int(max(0, (resized_width - final_width) / 2))
        # Left, top, right, bottom crops.
        crop = (crop_w, crop_h, crop_w + final_width, crop_h + final_height)

        if resized_width != final_width:
            print('Zero padding left and right parts of the image.')
        if crop_h + final_height != resized_height:
            print('Zero padding bottom part of the image.')

        return {'scale_width': resize_scale,
                'scale_height': resize_scale,
                'resize_dims': resize_dims,
                'crop': crop,
                }

    def get_input_data(self, rec):
        """
        Parameters
        ----------
            rec: nuscenes identifier for a given timestamp

        Returns
        -------
            images: torch.Tensor<float> (N, 3, H, W)
            intrinsics: torch.Tensor<float> (3, 3)
            extrinsics: torch.Tensor(N, 4, 4)
        """
        images = []
        intrinsics = []
        extrinsics = []
        depths = []
        cameras = self.cfg.IMAGE.NAMES

        # The extrinsics we want are from the camera sensor to "flat egopose" as defined
        # https://github.com/nutonomy/nuscenes-devkit/blob/9b492f76df22943daf1dc991358d3d606314af27/python-sdk/nuscenes/nuscenes.py#L279
        # which corresponds to the position of the lidar.
        # This is because the labels are generated by projecting the 3D bounding box in this lidar's reference frame.

        # From lidar egopose to world.
        lidar_sample = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        lidar_pose = self.nusc.get('ego_pose', lidar_sample['ego_pose_token'])
        yaw = Quaternion(lidar_pose['rotation']).yaw_pitch_roll[0]
        lidar_rotation = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        lidar_translation = np.array(lidar_pose['translation'])[:, None]
        lidar_to_world = np.vstack([
            np.hstack((lidar_rotation.rotation_matrix, lidar_translation)),
            np.array([0, 0, 0, 1])
        ])

        for cam in cameras:
            camera_sample = self.nusc.get('sample_data', rec['data'][cam])

            # Transformation from world to egopose
            car_egopose = self.nusc.get('ego_pose', camera_sample['ego_pose_token'])
            egopose_rotation = Quaternion(car_egopose['rotation']).inverse
            egopose_translation = -np.array(car_egopose['translation'])[:, None]
            world_to_car_egopose = np.vstack([
                np.hstack((egopose_rotation.rotation_matrix, egopose_rotation.rotation_matrix @ egopose_translation)),
                np.array([0, 0, 0, 1])
            ])

            # From egopose to sensor
            sensor_sample = self.nusc.get('calibrated_sensor', camera_sample['calibrated_sensor_token'])
            intrinsic = torch.Tensor(sensor_sample['camera_intrinsic'])
            sensor_rotation = Quaternion(sensor_sample['rotation'])
            sensor_translation = np.array(sensor_sample['translation'])[:, None]
            car_egopose_to_sensor = np.vstack([
                np.hstack((sensor_rotation.rotation_matrix, sensor_translation)),
                np.array([0, 0, 0, 1])
            ])
            car_egopose_to_sensor = np.linalg.inv(car_egopose_to_sensor)

            # Combine all the transformation.
            # From sensor to lidar.
            lidar_to_sensor = car_egopose_to_sensor @ world_to_car_egopose @ lidar_to_world
            sensor_to_lidar = torch.from_numpy(np.linalg.inv(lidar_to_sensor)).float()

            # Load image
            image_filename = os.path.join(self.dataroot, camera_sample['filename'])
            img = Image.open(image_filename)
            # Resize and crop
            img = resize_and_crop_image(
                img, resize_dims=self.augmentation_parameters['resize_dims'], crop=self.augmentation_parameters['crop']
            )
            # Normalise image
            normalised_img = self.normalise_image(img)

            # Combine resize/cropping in the intrinsics
            top_crop = self.augmentation_parameters['crop'][1]
            left_crop = self.augmentation_parameters['crop'][0]
            intrinsic = update_intrinsics(
                intrinsic, top_crop, left_crop,
                scale_width=self.augmentation_parameters['scale_width'],
                scale_height=self.augmentation_parameters['scale_height']
            )

            # Get Depth
            # Depth data should under the dataroot path 
            if self.cfg.LIFT.GT_DEPTH:
                base_root = os.path.join(self.dataroot, 'depths') 
                filename = os.path.basename(camera_sample['filename']).split('.')[0] + '.npy'
                depth_file_name = os.path.join(base_root, cam, 'npy', filename)
                depth = torch.from_numpy(np.load(depth_file_name)).unsqueeze(0).unsqueeze(0)
                depth = F.interpolate(depth, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear')
                depth = depth.squeeze()
                crop = self.augmentation_parameters['crop']
                depth = depth[crop[1]:crop[3], crop[0]:crop[2]]
                depth = torch.round(depth)
                depths.append(depth.unsqueeze(0).unsqueeze(0))

            images.append(normalised_img.unsqueeze(0).unsqueeze(0))
            intrinsics.append(intrinsic.unsqueeze(0).unsqueeze(0))
            extrinsics.append(sensor_to_lidar.unsqueeze(0).unsqueeze(0))

        images, intrinsics, extrinsics = (torch.cat(images, dim=1),
                                          torch.cat(intrinsics, dim=1),
                                          torch.cat(extrinsics, dim=1)
                                          )
        if len(depths) > 0:
            depths = torch.cat(depths, dim=1)

        return images, intrinsics, extrinsics, depths

    def _get_top_lidar_pose(self, rec):
        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        trans = -np.array(egopose['translation'])
        yaw = Quaternion(egopose['rotation']).yaw_pitch_roll[0]
        rot = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)]).inverse
        return trans, rot

    def get_depth_from_lidar(self, lidar_sample, cam_sample):
        points, coloring, im = self.nusc_exp.map_pointcloud_to_image(lidar_sample, cam_sample)
        tmp_cam = np.zeros((self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH))
        points = points.astype(np.int)
        tmp_cam[points[1, :], points[0,:]] = coloring
        tmp_cam = torch.from_numpy(tmp_cam).unsqueeze(0).unsqueeze(0)
        tmp_cam = F.interpolate(tmp_cam, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear', align_corners=False, recompute_scale_factor=True)
        tmp_cam = tmp_cam.squeeze()
        crop = self.augmentation_parameters['crop']
        tmp_cam = tmp_cam[crop[1]:crop[3], crop[0]:crop[2]]
        tmp_cam = torch.round(tmp_cam)
        return tmp_cam


    def get_birds_eye_view_label(self, rec, instance_map, in_pred):
        translation, rotation = self._get_top_lidar_pose(rec)
        segmentation = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        pedestrian = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        # Background is ID 0
        instance = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))

        for annotation_token in rec['anns']:
            # Filter out all non vehicle instances
            annotation = self.nusc.get('sample_annotation', annotation_token)

            if self.cfg.DATASET.FILTER_INVISIBLE_VEHICLES and int(annotation['visibility_token']) == 1 and in_pred is False:
                continue
            if in_pred is True and annotation['instance_token'] not in instance_map:
                continue

            # NuScenes filter
            if 'vehicle' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                instance_id = instance_map[annotation['instance_token']]
                poly_region, z = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(instance, [poly_region], instance_id)
                cv2.fillPoly(segmentation, [poly_region], 1.0)
            elif 'human' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                poly_region, z = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(pedestrian, [poly_region], 1.0)


        return segmentation, instance, pedestrian, instance_map

    def _get_poly_region_in_image(self, instance_annotation, ego_translation, ego_rotation):
        box = Box(
            instance_annotation['translation'], instance_annotation['size'], Quaternion(instance_annotation['rotation'])
        )
        box.translate(ego_translation)
        box.rotate(ego_rotation)

        pts = box.bottom_corners()[:2].T
        pts = np.round((pts - self.bev_start_position[:2] + self.bev_resolution[:2] / 2.0) / self.bev_resolution[:2]).astype(np.int32)
        pts[:, [1, 0]] = pts[:, [0, 1]]

        z = box.bottom_corners()[2, 0]
        return pts, z

    def get_label(self, rec, instance_map, in_pred):
        segmentation_np, instance_np, pedestrian_np, instance_map = \
            self.get_birds_eye_view_label(rec, instance_map, in_pred)
        segmentation = torch.from_numpy(segmentation_np).long().unsqueeze(0).unsqueeze(0)
        instance = torch.from_numpy(instance_np).long().unsqueeze(0)
        pedestrian = torch.from_numpy(pedestrian_np).long().unsqueeze(0).unsqueeze(0)

        return segmentation, instance, pedestrian, instance_map

    def get_future_egomotion(self, rec, index):
        rec_t0 = rec

        # Identity
        future_egomotion = np.eye(4, dtype=np.float32)

        if index < len(self.ixes) - 1:
            rec_t1 = self.ixes[index + 1]

            if rec_t0['scene_token'] == rec_t1['scene_token']:
                egopose_t0 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t0['data']['LIDAR_TOP'])['ego_pose_token']
                )
                egopose_t1 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t1['data']['LIDAR_TOP'])['ego_pose_token']
                )

                egopose_t0 = convert_egopose_to_matrix_numpy(egopose_t0)
                egopose_t1 = convert_egopose_to_matrix_numpy(egopose_t1)

                future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
                future_egomotion[3, :3] = 0.0
                future_egomotion[3, 3] = 1.0

        future_egomotion = torch.Tensor(future_egomotion).float()

        # Convert to 6DoF vector
        future_egomotion = mat2pose_vec(future_egomotion)
        return future_egomotion.unsqueeze(0)

    def get_trajectory_sampling(self, rec=None, sample_indice=None):
        if rec is None and sample_indice is None:
            raise ValueError("No valid input rec or token")
        if rec is None and sample_indice is not None:
            rec = self.ixes[sample_indice]

        ref_scene = self.nusc.get("scene", rec['scene_token'])

        # vm_msgs = self.nusc_can.get_messages(ref_scene['name'], 'vehicle_monitor')
        # vm_uts = [msg['utime'] for msg in vm_msgs]
        pose_msgs = self.nusc_can.get_messages(ref_scene['name'],'pose')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_msgs = self.nusc_can.get_messages(ref_scene['name'], 'steeranglefeedback')
        steer_uts = [msg['utime'] for msg in steer_msgs]

        ref_utime = rec['timestamp']
        # vm_index = locate_message(vm_uts, ref_utime)
        # vm_data = vm_msgs[vm_index]
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]

        # initial speed
        # v0 = vm_data["vehicle_speed"] / 3.6  # km/h to m/s
        v0 = pose_data["vel"][0]  # [0] means longitudinal velocity  m/s

        # curvature (positive: turn left)
        # steering = np.deg2rad(vm_data["steering"])
        steering = steer_data["value"]

        location = self.scene2map[ref_scene['name']]
        # flip x axis if in left-hand traffic (singapore)
        flip_flag = True if location.startswith('singapore') else False
        if flip_flag:
            steering *= -1
        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories

    def generate_centerlines(self, rec, map_name):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        x, y = egopose["translation"][0], egopose["translation"][1]
        nusc_map = self.nusc_maps[map_name]
        closest_lane = nusc_map.get_closest_lane(x, y, radius=2)
        candidates_future = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        candidates_past = dfs_incoming(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        plt.savefig("output_vis/debug_dfs.png")
        # lane_ids = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD)
        lane_ids = []
        for past_lane_seq in candidates_past:
            for future_lane_seq in candidates_future:
                lane_ids.append(past_lane_seq + future_lane_seq[1:])
        lane_ids = remove_overlapping_lane_seq(lane_ids)
        centerlines = []
        for lane_id in lane_ids:
            centerline = []
            for lane in lane_id:
                lane_record = nusc_map.get_arcline_path(lane)
                cx = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                centerline.extend(cx)
            centerlines.append(np.array(centerline))
        lanes = self.nusc_maps[map_name].get_records_in_radius(x, y, 2000, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        discrete_points = self.nusc_maps[map_name].discretize_lanes(lanes, resolution_meters=resolution_meters)
        lanes = []
        for lane_id, points in discrete_points.items():
            lanes.append(np.array(points)[:, :2])
        
        patch_box = box_coords
        patch_angle = rot * 180 / np.pi

        patch_x = patch_box[0]
        patch_y = patch_box[1]

        patch = get_patch_coord(patch_box, patch_angle)

        line_list = []
        for lane in centerlines:
            line = LineString([(arr[0], arr[1]) for arr in lane[:, :2]])
            c = lane[:, :2]
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.rotate(new_line, -patch_angle, origin=(patch_x, patch_y), use_radians=False)
                new_line = affinity.affine_transform(new_line,
                                                        [1.0, 0.0, 0.0, 1.0, -patch_x, -patch_y])
                line_list.append(new_line)
        
        local_box = (0.0, 0.0, patch_box[2], patch_box[3])
        patch_x, patch_y, patch_h, patch_w = local_box

        patch = get_patch_coord(local_box)

        canvas_h = canvas_size[0]
        canvas_w = canvas_size[1]
        scale_height = canvas_h/patch_h
        scale_width = canvas_w/patch_w

        trans_x = -patch_x + patch_w / 2.0
        trans_y = -patch_y + patch_h / 2.0

        overall_lines = []
        for line in line_list:
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.affine_transform(new_line,
                                                     [1.0, 0.0, 0.0, 1.0, trans_x, trans_y])
                new_line = affinity.scale(new_line, xfact=scale_width, yfact=scale_height, origin=(0, 0))
                overall_lines.append([list(new_line.coords)])

        return overall_lines

    def voxelize_hd_map(self, rec):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        elements = self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS
        hd_features = self.nusc_maps[map_name].get_map_mask(box_coords, rot * 180 / np.pi , elements, canvas_size=canvas_size)
        #traffic = self.hd_traffic_light(map_name, center, stretch, dx, bx, canvas_size)
        #return torch.from_numpy(np.concatenate((hd_features, traffic), axis=0)[None]).float()
        hd_features = torch.from_numpy(hd_features[None]).float()
        hd_features = torch.transpose(hd_features,-2,-1) # (y,x) replace horizontal and vertical coordinates
        return hd_features

    def hd_traffic_light(self, map_name, center, stretch, dx, bx, canvas_size):

        roads = np.zeros(canvas_size)
        my_patch = (
            center[0] - stretch[0],
            center[1] - stretch[1],
            center[0] + stretch[0],
            center[1] + stretch[1],
        )
        tl_token = self.nusc_maps[map_name].get_records_in_patch(my_patch, ['traffic_light'], mode='intersect')['traffic_light']
        polys = []
        for token in tl_token:
            road_token =self.nusc_maps[map_name].get('traffic_light', token)['from_road_block_token']
            pt = self.nusc_maps[map_name].get('road_block', road_token)['polygon_token']
            polygon = self.nusc_maps[map_name].extract_polygon(pt)
            polys.append(np.array(polygon.exterior.xy).T)

        def get_rot(h):
            return torch.Tensor([
                [np.cos(h), np.sin(h)],
                [-np.sin(h), np.cos(h)],
            ])
        # convert to local coordinates in place
        rot = get_rot(np.arctan2(center[3], center[2])).T
        for rowi in range(len(polys)):
            polys[rowi] -= center[:2]
            polys[rowi] = np.dot(polys[rowi], rot)

        for la in polys:
            pts = (la - bx) / dx
            pts = np.int32(np.around(pts))
            cv2.fillPoly(roads, [pts], 1)

        return roads[None]

    def get_gt_trajectory(self, rec, ref_index):
        n_output = self.cfg.N_FUTURE_FRAMES
        gt_trajectory = np.zeros((n_output+1+2, 3), np.float64)

        egopose_cur = get_global_pose(rec, self.nusc, inverse=True)

        for i in range(2):
            index = ref_index - 2 + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]

                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)

                egopose_future = egopose_cur.dot(egopose_future)
                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i, :] = [origin[0], origin[1], theta]

        for i in range(n_output+1):
            index = ref_index + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]

                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)

                egopose_future = egopose_cur.dot(egopose_future)
                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i + 2, :] = [origin[0], origin[1], theta]

        if gt_trajectory[-1][0] >= 2:
            command = 'RIGHT'
        elif gt_trajectory[-1][0] <= -2:
            command = 'LEFT'
        else:
            command = 'FORWARD'

        return gt_trajectory, command

    def get_routed_map(self, gt_points):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        canvas_size = (
            int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
            int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        roads = np.zeros(canvas_size)
        W = 1.85
        pts = np.array([
            [-4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, -W / 2.],
            [-4.084 / 2. + 0.5, -W / 2.],
        ])
        pts = (pts - bx) / dx
        pts[:, [0, 1]] = pts[:, [1, 0]]

        pts = np.int32(np.around(pts))
        cv2.fillPoly(roads, [pts], 1)

        gt_points = gt_points[:-1].numpy()
        # 坐标原点在左上角
        target = pts.copy()
        target[:,0] = pts[:,0] + gt_points[0] / dx[0]
        target[:,1] = pts[:,1] - gt_points[1] / dx[1]
        target = np.int32(np.around(target))
        cv2.fillPoly(roads, [target], 1)
        return roads

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

    def __getitem__(self, index):
        """
        Returns
        -------
            data: dict with the following keys:
                image: torch.Tensor<float> (T, N, 3, H, W)
                    normalised cameras images with T the sequence length, and N the number of cameras.
                intrinsics: torch.Tensor<float> (T, N, 3, 3)
                    intrinsics containing resizing and cropping parameters.
                extrinsics: torch.Tensor<float> (T, N, 4, 4)
                    6 DoF pose from world coordinates to camera coordinates.
                segmentation: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                    (H_bev, W_bev) are the pixel dimensions in bird's-eye view.
                instance: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                centerness: torch.Tensor<float> (T, 1, H_bev, W_bev)
                offset: torch.Tensor<float> (T, 2, H_bev, W_bev)
                flow: torch.Tensor<float> (T, 2, H_bev, W_bev)
                future_egomotion: torch.Tensor<float> (T, 6)
                    6 DoF egomotion t -> t+1

        """
        data = {}
        keys = ['image', 'intrinsics', 'extrinsics', 'depths',
                'segmentation', 'instance', 'centerness', 'offset', 'flow', 'pedestrian',
                'future_egomotion', 'hdmap', 'gt_trajectory', 'indices', 'poses'
                ]
        for key in keys:
            data[key] = []

        instance_map = {}
        # Loop over all the frames in the sequence.
        total_centerlines = []
        for i, index_t in enumerate(self.indices[index]):
            if i >= self.receptive_field:
                in_pred = True
            else:
                in_pred = False
            rec = self.ixes[index_t]
            map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

            if i < self.receptive_field:
                images, intrinsics, extrinsics, depths = self.get_input_data(rec)
                data['image'].append(images)
                data['intrinsics'].append(intrinsics)
                data['extrinsics'].append(extrinsics)
                data['depths'].append(depths)
            segmentation, instance, pedestrian, instance_map = self.get_label(rec, instance_map, in_pred)

            future_egomotion = self.get_future_egomotion(rec, index_t)
            hd_map_feature = self.voxelize_hd_map(rec)

            data['segmentation'].append(segmentation)
            data['instance'].append(instance)
            data['pedestrian'].append(pedestrian)
            data['future_egomotion'].append(future_egomotion)
            data['hdmap'].append(hd_map_feature)
            data['indices'].append(index_t)

            """
                get lanes in 100m radius
            """
            centerlines = self.generate_centerlines(rec, map_name)
            total_centerlines.append(centerlines)

            if i == self.cfg.TIME_RECEPTIVE_FIELD-1:
                gt_trajectory, command = self.get_gt_trajectory(rec, index_t)
                gt_trajectory_prev = gt_trajectory
                traj=torch.tensor(gt_trajectory_prev)[:, :2].reshape(1,9,2)
                gt_trajectory = gt_trajectory[2:]
                data['gt_trajectory'] = torch.from_numpy(gt_trajectory).float()
                data['command'] = command
                trajs = self.get_trajectory_sampling(rec)
                data['sample_trajectory'] = torch.from_numpy(trajs).float()


                """
                    DISPLAY DEBUG
                """
                display = np.ones((200, 200, 3))
                whe = np.where(hd_map_feature[0,1] >= 1)
                display[whe[0], whe[1]] = [0, 0, 0]
                whe = np.where(segmentation[0,0] >= 1)
                display[whe[0], whe[1]] = [0, 0, 1]

                """
                    CREATE GRID
                    grid: (M, 2) -> sample grid
                """
                sample_trajectories = []
                if debug and brdebug:breakpoint()
                if debug:
                    plt.figure(figsize=(10, 10))

                """
                    ITERATE OVER CENTERLINES
                """
                total_trajs = []
                total_eps = []
                data_centerlines = []
                for centerline in total_centerlines[-1]:
                    """
                        INITIALIZE GRID
                    """
                    x = np.linspace(0, 19, 20)
                    y = np.linspace(-3, 3, 7)
                    xv, yv = np.meshgrid(x, y)
                    grid = np.dstack((xv, yv))
                    grid = np.concatenate((grid), axis=0)

                    """
                        CONVERT CENTERLINE TO BOX FRAME
                    """
                    cx = np.array(centerline)[0]
                    cx = extend_both_sides(cx, resolution=resolution_meters, extension=30)
                    cx_o = np.zeros_like(cx) + cx
                    cx = cx * ddx
                    cx += bbx
                    cx[:, [1, 0]] = cx[:, [0, 1]]
                    cx[:, :1] = cx[:, :1] * -1
                    if debug:
                        gtt = gt_trajectory_prev[:, :2]
                        plt.plot(gtt[:, 0], gtt[:, 1], color="orange")
                        plt.plot(cx[:, 0], cx[:, 1], color="green")
                        plt.savefig(f"output_vis/{index}debug_cartesian.png")
                        plt.clf()
                    #################################################################################


                    """
                        CONVERT CENTERLINE TO SPLINE
                            spline: represents the spline
                            gt_trajectory_prev: (1, 9, 2) representing the GT
                            ssd_list: frenet frame of gt_trajectory_prev
                            s_offset: s coordinate at T=2
                            d_offset: d coordinate at T=2
                    """
                    trj = torch.tensor(gt_trajectory_prev)[:, :2].reshape(1, 9, 2)
                    oracle_centerline = torch.tensor(cx).reshape(1, cx.shape[0], 2)
                    spline = Spline2D(x=oracle_centerline[0,:, 0].numpy(), y=oracle_centerline[0,:, 1].numpy(), resolution=resolution_meters, extension=15)
                    ssd_list = []
                    for xy in trj[0]:
                        s,d = spline.calc_frenet_position(xy[0].numpy(), xy[1].numpy())
                        ssd_list.append([s, d])
                    ssd = torch.tensor(ssd_list).reshape(1, trj.shape[1], 2)
                    s_offset = ssd[0,2,0].detach()
                    d_offset = ssd[0,2,1].detach()
                    vel = ssd[0,2] - ssd[0,1]
                    theta = np.arctan2(vel[1], vel[0]) # to make angle 0 wrt ego-agent
                    vel_prev = ssd[0,1] - ssd[0, 0]
                    theta_prev = np.arctan2(vel_prev[1], vel_prev[0])
                    psidot_init = (theta - theta_prev)/0.5
                    psi_init = theta
                    v_init = np.linalg.norm(vel)/0.5
                    y_init = 0
                    x_init = 0
                    rot_inv = np.array([[cos(-theta), -sin(-theta)], [sin(-theta), cos(-theta)]])
                    rot = np.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])
                    if debug and brdebug:breakpoint()                    
                    #################################################################################

                    """
                        CONVERT CENTERLINE ITSELF TO CENTERLINE FRAME
                    """
                    sdd_list = []
                    trj = oracle_centerline[:, 10:-10]
                    for xy in trj[0]:
                        s,d = spline.calc_frenet_position(xy[0].numpy(), xy[1].numpy())
                        sdd_list.append([s, d])
                    sdd = torch.tensor(sdd_list).reshape(1, trj.shape[1], 2)
                    if debug:
                        plt.plot(sdd[0, :, 0].cpu().numpy(), sdd[0, :, 1].cpu().numpy(), color="green")
                        plt.scatter(grid[:, 0], grid[ :, 1], color="red", s=1)
                        plt.plot(ssd[0, :, 0], ssd[0, :, 1], color="blue")
                        plt.axis('equal')
                        plt.savefig(f"output_vis/{index}debug_sdd_vefore.png")
                        plt.clf()

                    if debug and brdebug:breakpoint()
                    ssd_frame = ssd -  torch.tensor([s_offset, d_offset])
                    sdd_frame = sdd -  torch.tensor([s_offset, d_offset])
                    eps_frame = grid - np.array([0, d_offset])
                    if debug:
                        plt.plot(sdd_frame[0, :, 0].cpu().numpy(), sdd_frame[0, :, 1].cpu().numpy(), color="green")
                        plt.scatter(eps_frame[:, 0], eps_frame[ :, 1], color="red", s=1)
                        plt.plot(ssd_frame[0, :, 0], ssd_frame[0, :, 1], color="blue")
                        plt.axis('equal')
                        plt.savefig(f"output_vis/{index}debug_sdd.png")
                        plt.clf()
                        plt.plot(ssd_frame[:, 0], ssd_frame[:, 1], color="blue");plt.savefig("output_vis/debug_aa.png");plt.clf()
                    #################################################################################

                    """
                        CONVERT GRID TO AGENT FRAME
                    """
                    grid = grid - np.array([0, d_offset])
                    if debug and brdebug:breakpoint()
                    #################################################################################

                    """
                        ROTATE EVERYTHING TO AGENT FRAME
                    """
                    sdd_frame[0] = torch.tensor(np.dot(rot_inv, sdd_frame[0].numpy().T).T)
                    ssd_frame = np.dot(rot_inv, ssd_frame[0].T).T
                    eps_frame = np.dot(rot_inv, grid.T).T

                    if debug:
                        plt.plot(ssd_frame[:, 0], ssd_frame[:, 1], color="blue")
                        plt.plot(sdd_frame[0, :, 0].cpu().numpy(), sdd_frame[0, :, 1].cpu().numpy(), color="green")
                        plt.scatter(eps_frame[:, 0], eps_frame[ :, 1], color="red", s=1)
                        plt.axis('equal')
                        plt.savefig(f"output_vis/{index}debug_sdd_rotated.png")
                        plt.clf()
                    #################################################################################

                    """
                        SOLVE THE PROBLEM
                    """                        
                    fixed_params = torch.zeros((1, 5)).expand(eps_frame.shape[0], 5)
                    fixed_params[0] = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init])
                    variable_params = torch.zeros((eps_frame.shape[0], 3))
                    variable_params[:, 0] = torch.tensor(eps_frame[:, 0])
                    variable_params[:, 1] = torch.tensor(eps_frame[:, 1])
                    variable_params[:, 2] = torch.tensor(-theta)
                    P=problem.P
                    cc, _ = problem.solve(fixed_params.double(), variable_params.double())
                    x_sol = np.dot(P, cc.T[:nvar]).T
                    y_sol = np.dot(P, cc.T[nvar:2*nvar]).T
                    min_val = 1e11
                    min_ind = 0
                    #################################################################################

                    """
                        COMPUTE MIN_TRAJ
                    """
                    for i in range(len(x_sol)):
                        traj = np.dstack((x_sol[i][::10], y_sol[i][::10]))[0]
                        dist = np.linalg.norm(traj - ssd_frame[2:])
                        if dist < min_val:
                            min_val = dist
                            min_ind = i
                    if debug:
                        for i in range(len(x_sol)):
                            plt.plot(x_sol[i], y_sol[i],linewidth=0.5, color="blue", alpha=0.5)
                            plt.scatter(x_sol[i][::10], y_sol[i][::10],s=1, color="blue", alpha=0.5)
                        plt.plot(ssd_frame[:, 0], ssd_frame[:, 1], color="red", zorder=100, label="GT")
                        plt.scatter(ssd_frame[:, 0], ssd_frame[:, 1], color="red", zorder=100, s=1)
                        plt.scatter(eps_frame[:, 0], eps_frame[:, 1], color="yellow", zorder=100, s=1)
                        plt.plot(sdd_frame[0, 180:-50, 0].cpu().numpy(), sdd_frame[0, 180:-50, 1].cpu().numpy(), color="green")
                        plt.plot(x_sol[min_ind][::10], y_sol[min_ind][::10], color="orange", zorder=100, label="predicted")
                        plt.scatter(x_sol[min_ind][::10], y_sol[min_ind][::10], color="orange", zorder=100, s=3)
                        plt.legend()
                        plt.axis('equal')
                        plt.savefig(f"output_vis/{index}debug.png")
                        plt.clf()
                    #################################################################################

                    """
                       PROJECT TO CARTESIAN FRAME 
                    """
                    trajs = np.dstack((x_sol[:, ::10], y_sol[:,::10])) # N, 7, 2
                    for i in range(len(trajs)):
                        trajs[i] = np.dot(rot, trajs[i].T).T
                    trajs = trajs + np.array([s_offset, d_offset])
                    sdd_frame[0] = torch.tensor(np.dot(rot, sdd_frame[0].numpy().T).T)
                    ssd_back = np.dot(rot, ssd_frame.T).T
                    eps_back = np.dot(rot, eps_frame.T).T
                    ssd_back = ssd_back + np.array([s_offset, d_offset])
                    eps_back = eps_back + np.array([s_offset, d_offset])
                    sdd_back = sdd_frame[0].numpy() + np.array([s_offset, d_offset])
                    if debug:
                        done = False
                        plt.plot(sdd_back[:, 0], sdd_back[:, 1], color="green")
                        plt.plot(ssd_back[:, 0], ssd_back[:, 1], color="red")
                        plt.scatter(ssd_back[:, 0], ssd_back[:, 1], color="red", s=3)
                        plt.scatter(eps_back[:, 0], eps_back[:, 1], color="yellow", s=2)
                        for j in range(len(trajs)):
                            if not done:
                                done = True
                                plt.plot(trajs[min_ind, :7, 0], trajs[min_ind, :7, 1], color="orange", zorder=2000)
                                plt.scatter(trajs[min_ind, :7, 0], trajs[min_ind, :7, 1], color="orange", s=2,zorder=2000)
                            plt.plot(trajs[j,:7, 0], trajs[j,:7, 1], color="blue", alpha=0.5)
                    #################################################################################

                    """
                        ADD TO OVERALL TRAJECTORY SET
                    """
                    for j in range(len(trajs)):
                        xys = []
                        for sds in trajs[j]:
                            x, y = spline.calc_global_position_offline(sds[0], sds[1])
                            xys.append([x, y])
                        trajs[j] = torch.tensor(xys)
                        total_trajs.append(trajs[j, :7])
                    cxx = oracle_centerline[0, :, 0].detach()
                    cyy = oracle_centerline[0, :, 1].detach()
                    data_centerlines.append(np.dstack((cxx, cyy)))
                    if debug:
                        plt.axis('equal')
                        plt.savefig(f"output_vis/{index}__debug.png")
                        plt.clf()
                        orc = []
                        for sds in sdd_back:
                            x, y = spline.calc_global_position_offline(sds[0], sds[1])
                            orc.append([x, y])
                        gd = []
                        for sds in eps_back:
                            x, y = spline.calc_global_position_offline(sds[0], sds[1])
                            gd.append([x, y])
                            total_eps.append([x, y])
                        gd = np.array(gd)
                        orc = np.array(orc)
                        for i in range(len(trajs)):
                            pass
                            plt.plot(trajs[i, :7, 0], trajs[i, :7, 1], color="blue", alpha=0.5)
                        plt.plot(trajs[min_ind, :7, 0], trajs[min_ind, :7, 1], color="orange", zorder=2000)
                        plt.scatter(trajs[min_ind, :7, 0], trajs[min_ind, :7, 1], color="orange", s=3,zorder=2000)
                        plt.plot(gt_trajectory_prev[ :, 0], gt_trajectory_prev[ : ,1], zorder=2000)
                        plt.plot(oracle_centerline[0, :, 0], oracle_centerline[0, : ,1])
                        for traj in data['sample_trajectory']:
                            pass
                            # plt.plot(traj[:, 0], traj[:, 1], "green", alpha=0.05)
                        plt.scatter(gd[:, 0], gd[:, 1], color="yellow", s=3, zorder=2000)
                        plt.scatter(orc[:, 0], orc[:, 1], color="red", s=3, zorder=2000)
                        plt.axis('equal');
                        plt.savefig(f"output_vis/{index}_debug.png")
                        plt.clf()

                trajs = torch.tensor(total_trajs).float()
                zero = torch.zeros((trajs.shape[0], trajs.shape[1], 1))
                trajs = torch.cat((trajs, zero), dim=2)
                if self.cfg.DENSE:
                    data['sample_trajectory'] = trajs
                gt_trajectory[:, :1] = gt_trajectory[:, :1] * -1
                gt_trajectory = (gt_trajectory[:, :2] - bbx) / ddx
                data["centerlines"] = data_centerlines


        for key, value in data.items():
            if key in ['image', 'intrinsics', 'extrinsics', 'depths', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian']:
                if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
                    continue
                data[key] = torch.cat(value, dim=0)

        data['target_point'] = torch.tensor([0., 0.])
        instance_centerness, instance_offset, instance_flow = convert_instance_mask_to_center_and_offset_label(
            data['instance'], data['future_egomotion'],
            num_instances=len(instance_map), ignore_index=self.cfg.DATASET.IGNORE_INDEX, subtract_egomotion=True,
            spatial_extent=self.spatial_extent,
        )
        data['centerness'] = instance_centerness
        data['offset'] = instance_offset
        data['flow'] = instance_flow

        return data
data/stp3/datas/__NuscenesData.py
llmbev/talk2bev
import torch
import torch.utils.data
from nuscenes.nuscenes import NuScenes
from stp3.datas.trainNuscenesData import FuturePredictionDataset
from stp3.datas.CarlaData import CarlaDataset

def prepare_dataloaders(cfg, return_dataset=False):
    if cfg.DATASET.NAME == 'nuscenes':
        # 28130 train and 6019 val
        dataroot = cfg.DATASET.DATAROOT
        nusc = NuScenes(version='v1.0-{}'.format(cfg.DATASET.VERSION), dataroot=dataroot, verbose=False)
        traindata = FuturePredictionDataset(nusc, 0, cfg)
        valdata = FuturePredictionDataset(nusc, 1, cfg)

        if cfg.DATASET.VERSION == 'mini':
            pass
            # traindata.indices = traindata.indices[:10]
            # valdata.indices = valdata.indices[:10]

        nworkers = cfg.N_WORKERS
        trainloader = torch.utils.data.DataLoader(
            traindata, batch_size=cfg.BATCHSIZE, shuffle=False, num_workers=10, pin_memory=True, drop_last=True
        )
        valloader = torch.utils.data.DataLoader(
            valdata, batch_size=cfg.BATCHSIZE, shuffle=True, num_workers=10, pin_memory=True, drop_last=False
        )
    elif cfg.DATASET.NAME == 'carla':
        dataroot = cfg.DATASET.DATAROOT
        traindata = CarlaDataset(dataroot, True, cfg)
        valdata = CarlaDataset(dataroot, False, cfg)
        nworkers = cfg.N_WORKERS
        trainloader = torch.utils.data.DataLoader(
            traindata, batch_size=cfg.BATCHSIZE, shuffle=True, num_workers=nworkers, pin_memory=True, drop_last=True
        )
        valloader = torch.utils.data.DataLoader(
            valdata, batch_size=cfg.BATCHSIZE, shuffle=False, num_workers=nworkers, pin_memory=True, drop_last=False)
    else:
        raise NotImplementedError

    if return_dataset:
        return trainloader, valloader, traindata, valdata
    else:
        return trainloader, valloader

data/stp3/datas/dataloaders.py
llmbev/talk2bev
import os
from PIL import Image

import numpy as np
import cv2
import torch
import torch.utils.data
import torch.nn.functional as F
import torchvision
import os.path as osp

from shapely.geometry import LineString

from pyquaternion import Quaternion
from shapely import affinity
from nuscenes.nuscenes import NuScenes, NuScenesExplorer
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.splits import create_splits_scenes
from nuscenes.utils.data_classes import Box
from nuscenes.lidarseg.lidarseg_utils import colormap_to_colors, plt_to_cv2, get_stats, get_labels_in_coloring, create_lidarseg_legend, paint_points_label
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.eval.common.utils import quaternion_yaw
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from nuscenes.utils.data_classes import PointCloud, LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix
from nuscenes.utils.data_io import load_bin_file
from stp3.utils.tools import ( gen_dx_bx, get_nusc_maps)
import matplotlib.pyplot as plt 
from math import cos, sin

from stp3.utils.geometry import (
    resize_and_crop_image,
    update_intrinsics,
    calculate_birds_eye_view_parameters,
    convert_egopose_to_matrix_numpy,
    pose_vec2mat,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
    get_global_pose,
    get_patch_coord,
    dfs,
    dfs_incoming,
    remove_overlapping_lane_seq,
    extend_both_sides,
    project_to_frenet_frame,
    project_to_cartesian_frame,
    interp_arc,
    interp_arc_np,
    denoise,
    interp_polyline_by_fixed_waypt_interval
)
from stp3.utils.instance import convert_instance_mask_to_center_and_offset_label
import stp3.utils.sampler as trajectory_sampler
from stp3.utils.spline import Spline2D
from stp3.utils.Optnode import *

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

resolution_meters = 0.1
t_fin = 3.5
num = 61
nvar = 11
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
debug = False
brdebug = False
problem = OPTNode(t_fin=t_fin, num=num)

def distance(x1, y1, z1, x2, y2, z2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + 1 * (z1 - z2) ** 2) ** ( 1 / 2)
 
# Function to calculate K closest points
def kClosest(points, target, K):
    pts = []
    n = len(points)
    d = []

    for i in range(n):
        d.append({
            "first": distance(points[i][0], points[i][1], points[i][2], target[0], target[1], target[2]),
            "second": i
        })
     
    d = sorted(d, key=lambda l:l["first"])
 
    for i in range(K):
        pt = []
        pt.append(points[d[i]["second"]][0])
        pt.append(points[d[i]["second"]][1])
        pt.append(points[d[i]["second"]][2])
        pt.append(points[d[i]["second"]][3])
        pts.append(pt)
 
    return pts

class FuturePredictionDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5 #SECOND
    def __init__(self, nusc, is_train, cfg):
        self.nusc = nusc
        self.dataroot = self.nusc.dataroot
        self.nusc_exp = NuScenesExplorer(nusc)
        self.nusc_can = NuScenesCanBus(dataroot=self.dataroot)
        self.is_train = is_train
        self.cfg = cfg
        self.problem = OPTNode(t_fin=3.5, num=self.cfg.N_FUTURE_FRAMES + 1)

        if self.is_train == 0:
            self.mode = 'train'
        elif self.is_train == 1:
            self.mode = 'val'
        elif self.is_train == 2:
            self.mode = 'test'
        else:
            raise NotImplementedError

        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD

        self.scenes = self.get_scenes()
        self.ixes = self.prepro()
        self.indices = self.get_indices()

        # Image resizing and cropping
        self.augmentation_parameters = self.get_resizing_and_cropping_parameters()

        # Normalising input images
        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # Bird's-eye view parameters
        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_resolution, self.bev_start_position, self.bev_dimension = (
            bev_resolution.numpy(), bev_start_position.numpy(), bev_dimension.numpy()
        )

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # The number of sampled trajectories
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        # HD-map feature extractor
        self.nusc_maps = get_nusc_maps(self.cfg.DATASET.MAP_FOLDER)
        self.scene2map = {}
        for sce in self.nusc.scene:
            log = self.nusc.get('log', sce['log_token'])
            self.scene2map[sce['name']] = log['location']
        self.save_dir = cfg.DATASET.SAVE_DIR

        self.curbatchnum1 = None
        self.curbatchnum2 = None
        if self.cfg.PLANNING.DENSE.ENABLED:
            self.centerlines_path = self.cfg.PLANNING.DENSE.PATH

    def get_scenes(self):
        # filter by scene split
        split = {'v1.0-trainval': {0: 'train', 1: 'val', 2: 'test'},
                 'v1.0-mini': {0: 'mini_train', 1: 'mini_val'},}[
            self.nusc.version
        ][self.is_train]

        blacklist = [419] + self.nusc_can.can_blacklist  # # scene-0419 does not have vehicle monitor data
        blacklist = ['scene-' + str(scene_no).zfill(4) for scene_no in blacklist]

        scenes = create_splits_scenes()[split][:]
        for scene_no in blacklist:
            if scene_no in scenes:
                scenes.remove(scene_no)

        return scenes

    def prepro(self):
        samples = [samp for samp in self.nusc.sample]

        # remove samples that aren't in this split
        samples = [samp for samp in samples if self.nusc.get('scene', samp['scene_token'])['name'] in self.scenes]
        
        # sort by scene, timestamp (only to make chronological viz easier)
        samples.sort(key=lambda x: (x['scene_token'], x['timestamp']))

        return samples

    def get_indices(self):
        indices = []
        for index in range(len(self.ixes)):
            is_valid_data = True
            previous_rec = None
            current_indices = []
            for t in range(self.sequence_length):
                index_t = index + t
                # Going over the dataset size limit.
                if index_t >= len(self.ixes):
                    is_valid_data = False
                    break
                rec = self.ixes[index_t]
                # Check if scene is the same
                if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
                    is_valid_data = False
                    break

                current_indices.append(index_t)
                previous_rec = rec

            if is_valid_data:
                indices.append(current_indices)

        return np.asarray(indices)

    # def get_indices(self):
    #     indices = []
    #     scene_token_mapping = {}
    #     ixes_new = []
    #     for index in range(len(self.ixes)):
    #         if scene_token_mapping.get(self.ixes[index]['scene_token']) == None:
    #             scene_token_mapping[self.ixes[index]['scene_token']] = 1
    #             continue
    #         scene_token_mapping[self.ixes[index]['scene_token']] = scene_token_mapping[self.ixes[index]['scene_token']] + 1
    #         is_valid_data = True
    #         previous_rec = None
    #         current_indices = []
    #         for t in range(self.sequence_length):
    #             index_t = index + t
    #             # Going over the dataset size limit.
    #             if index_t >= len(self.ixes):
    #                 is_valid_data = False
    #                 break
    #             rec = self.ixes[index_t]
    #             ixes_new.append(self.ixes[index_t])
    #             # Check if scene is the same
    #             if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
    #                 is_valid_data = False
    #                 break

    #             current_indices.append(len(ixes_new) - 1)
    #             previous_rec = rec

    #         if is_valid_data:
    #             indices.append(current_indices)

    #     self.ixes = ixes_new
    #     return np.asarray(indices)

    def get_resizing_and_cropping_parameters(self):
        original_height, original_width = self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH
        final_height, final_width = self.cfg.IMAGE.FINAL_DIM

        resize_scale = self.cfg.IMAGE.RESIZE_SCALE
        resize_dims = (int(original_width * resize_scale), int(original_height * resize_scale))
        resized_width, resized_height = resize_dims

        crop_h = self.cfg.IMAGE.TOP_CROP
        crop_w = int(max(0, (resized_width - final_width) / 2))
        # Left, top, right, bottom crops.
        crop = (crop_w, crop_h, crop_w + final_width, crop_h + final_height)

        if resized_width != final_width:
            print('Zero padding left and right parts of the image.')
        if crop_h + final_height != resized_height:
            print('Zero padding bottom part of the image.')

        return {'scale_width': resize_scale,
                'scale_height': resize_scale,
                'resize_dims': resize_dims,
                'crop': crop,
                }

    def get_input_data(self, rec):
        """
        Parameters
        ----------
            rec: nuscenes identifier for a given timestamp

        Returns
        -------
            images: torch.Tensor<float> (N, 3, H, W)
            intrinsics: torch.Tensor<float> (3, 3)
            extrinsics: torch.Tensor(N, 4, 4)
        """
        images = []
        unnormalized_images = []
        intrinsics = []
        extrinsics = []
        depths = []
        translations = []
        rotations = []
        camera_intrinsics = []
        point_clouds = []
        point_clouds_labels = []
        point_clouds_panoptic_labels = []
        panoptic_mappings_list = []
        extra_rotations = []
        extra_translations = []
        cameras = self.cfg.IMAGE.NAMES

        # # The extrinsics we want are from the camera sensor to "flat egopose" as defined
        # # https://github.com/nutonomy/nuscenes-devkit/blob/9b492f76df22943daf1dc991358d3d606314af27/python-sdk/nuscenes/nuscenes.py#L279
        # # which corresponds to the position of the lidar.
        # # This is because the labels are generated by projecting the 3D bounding box in this lidar's reference frame.

        # # From lidar egopose to world.
        lidar_sample = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        lidar_pose = self.nusc.get('ego_pose', lidar_sample['ego_pose_token'])
        yaw = Quaternion(lidar_pose['rotation']).yaw_pitch_roll[0]
        lidar_rotation = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        lidar_translation = np.array(lidar_pose['translation'])[:, None]
        lidar_to_world = np.vstack([
            np.hstack((lidar_rotation.rotation_matrix, lidar_translation)),
            np.array([0, 0, 0, 1])
        ])

        camera_sample_ = self.nusc.get('sample_data', rec['data']["CAM_FRONT"])
        my_sample = self.nusc.get('sample', camera_sample_['sample_token'])
        my_scene_token = my_sample['scene_token']

        # # Here we just grab the front camera and the point sensor.
        pointsensor = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        pcl_path = osp.join(self.nusc.dataroot, pointsensor['filename'])
        # panoptic_mappings = [samp.token for samp in self.nusc.get_sample_data(rec['data']['LIDAR_TOP'])[1]]
        panoptic_mappings = []
        for ok in rec['anns']:
            panoptic_mappings.append(self.nusc.get('sample_annotation', ok)['token'])
        panoptic_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get('panoptic', rec['data']['LIDAR_TOP'])['filename'])
        panoptic_labels = load_bin_file(panoptic_labels_filename, type='panoptic').astype(np.int32)
        point_clouds_panoptic_labels.append(panoptic_labels)
        for cam in cameras:
            camera_sample = self.nusc.get('sample_data', rec['data'][cam])
            sample_token = camera_sample['token']

            camm = self.nusc.get('sample_data', rec['data'][cam])
            lidarseg_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get("lidarseg", rec['data']['LIDAR_TOP'])['filename'])
            points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)


            cs_record = self.nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token'])
            pc = LidarPointCloud.from_file(pcl_path)
            ### LIDAR TO EGO_{TL} TRANSLATION -> CONSTANT
            pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)
            pc.translate(np.array(cs_record['translation']))

            point_clouds.append(torch.tensor(np.array(pc.points)).unsqueeze(0).unsqueeze(0)) # 1, 1, N, 4
            point_clouds_labels.append(torch.tensor(np.array(points_label)).unsqueeze(0).unsqueeze(0)) # 1, 1, N, 

            poserecord = self.nusc.get('ego_pose', pointsensor['ego_pose_token'])
            extra_rotations.append(Quaternion(poserecord['rotation']).rotation_matrix)
            extra_translations.append(np.array(poserecord['translation']))
            pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)
            pc.translate(np.array(poserecord['translation']))

            # Third step: transform from global into the ego vehicle frame for the timestamp of the image.
            poserecord = self.nusc.get('ego_pose', camm['ego_pose_token'])
            extra_rotations.append(Quaternion(poserecord['rotation']).rotation_matrix.T)
            extra_translations.append(np.array(poserecord['translation']))            
            pc.translate(-np.array(poserecord['translation']))
            pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix.T)

            # Transformation from world to egopose
            car_egopose = self.nusc.get('ego_pose', camera_sample['ego_pose_token'])
            egopose_rotation = Quaternion(car_egopose['rotation']).inverse
            egopose_translation = -np.array(car_egopose['translation'])[:, None]
            world_to_car_egopose = np.vstack([
                np.hstack((egopose_rotation.rotation_matrix, egopose_rotation.rotation_matrix @ egopose_translation)),
                np.array([0, 0, 0, 1])
            ])

            # From egopose to sensor
            sensor_sample = self.nusc.get('calibrated_sensor', camera_sample['calibrated_sensor_token'])
            intrinsic = torch.Tensor(sensor_sample['camera_intrinsic'])
            sensor_rotation = Quaternion(sensor_sample['rotation'])
            sensor_translation = np.array(sensor_sample['translation'])[:, None]
            car_egopose_to_sensor = np.vstack([
                np.hstack((sensor_rotation.rotation_matrix, sensor_translation)),
                np.array([0, 0, 0, 1])
            ])
            car_egopose_to_sensor = np.linalg.inv(car_egopose_to_sensor)

            # Combine all the transformation.
            # From sensor to lidar.
            lidar_to_sensor = car_egopose_to_sensor @ world_to_car_egopose @ lidar_to_world
            sensor_to_lidar = torch.from_numpy(np.linalg.inv(lidar_to_sensor)).float()

            # Load image
            image_filename = os.path.join(self.dataroot, camera_sample['filename'])
            img = Image.open(image_filename)
            unnormalized_image = torch.tensor(np.array(img))
            # Resize and crop
            img = resize_and_crop_image(
                img, resize_dims=self.augmentation_parameters['resize_dims'], crop=self.augmentation_parameters['crop']
            )
            # Normalise image
            normalised_img = self.normalise_image(img)

            # # Combine resize/cropping in the intrinsics
            top_crop = self.augmentation_parameters['crop'][1]
            left_crop = self.augmentation_parameters['crop'][0]
            intrinsic = update_intrinsics(
                intrinsic, top_crop, left_crop,
                scale_width=self.augmentation_parameters['scale_width'],
                scale_height=self.augmentation_parameters['scale_height']
            )

            # Get Depth
            # Depth data should under the dataroot path 
            if self.cfg.LIFT.GT_DEPTH:
                base_root = os.path.join(self.dataroot, 'depths') 
                filename = os.path.basename(camera_sample['filename']).split('.')[0] + '.npy'
                depth_file_name = os.path.join(base_root, cam, 'npy', filename)
                depth = torch.from_numpy(np.load(depth_file_name)).unsqueeze(0).unsqueeze(0)
                depth = F.interpolate(depth, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear')
                depth = depth.squeeze()
                crop = self.augmentation_parameters['crop']
                depth = depth[crop[1]:crop[3], crop[0]:crop[2]]
                depth = torch.round(depth)
                depths.append(depth.unsqueeze(0).unsqueeze(0))
            
            camm = self.nusc.get('sample_data', rec['data'][cam])
            cs_record = self.nusc.get('calibrated_sensor', camm['calibrated_sensor_token'])
            images.append(normalised_img.unsqueeze(0).unsqueeze(0))
            intrinsics.append(intrinsic.unsqueeze(0).unsqueeze(0))
            extrinsics.append(sensor_to_lidar.unsqueeze(0).unsqueeze(0))
            unnormalized_images.append(unnormalized_image.unsqueeze(0).unsqueeze(0))
            rotations.append(torch.tensor(Quaternion(cs_record['rotation']).rotation_matrix.T).unsqueeze(0).unsqueeze(0))
            translations.append(torch.tensor(-np.array(cs_record['translation'])).unsqueeze(0).unsqueeze(0))
            camera_intrinsics.append(torch.tensor(np.array(cs_record['camera_intrinsic'])).unsqueeze(0).unsqueeze(0))

        images, intrinsics, extrinsics, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, panoptic_mappings_list = (torch.cat(images, dim=1),
                                            torch.cat(intrinsics, dim=1),
                                            torch.cat(extrinsics, dim=1),
                                            torch.cat(unnormalized_images, dim=1),
                                            torch.cat(rotations, dim=1),
                                            torch.cat(translations, dim=1),
                                            torch.cat(camera_intrinsics, dim=1),
                                            torch.cat(point_clouds, dim=1),
                                            torch.cat(point_clouds_labels, dim=1),
                                            point_clouds_panoptic_labels,
                                            panoptic_mappings
                                          )
        if len(depths) > 0:
            depths = torch.cat(depths, dim=1)

        return images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, my_scene_token, panoptic_mappings_list

    def _get_top_lidar_pose(self, rec):
        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        trans = -np.array(egopose['translation'])
        yaw = Quaternion(egopose['rotation']).yaw_pitch_roll[0]
        # rot = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        rot = Quaternion(egopose['rotation']).inverse
        return trans, rot

    def get_depth_from_lidar(self, lidar_sample, cam_sample):
        points, coloring, im = self.nusc_exp.map_pointcloud_to_image(lidar_sample, cam_sample)
        tmp_cam = np.zeros((self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH))
        points = points.astype(np.int)
        tmp_cam[points[1, :], points[0,:]] = coloring
        tmp_cam = torch.from_numpy(tmp_cam).unsqueeze(0).unsqueeze(0)
        tmp_cam = F.interpolate(tmp_cam, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear', align_corners=False, recompute_scale_factor=True)
        tmp_cam = tmp_cam.squeeze()
        crop = self.augmentation_parameters['crop']
        tmp_cam = tmp_cam[crop[1]:crop[3], crop[0]:crop[2]]
        tmp_cam = torch.round(tmp_cam)
        return tmp_cam


    def get_birds_eye_view_label(self, rec, instance_map, in_pred):
        translation, rotation = self._get_top_lidar_pose(rec)
        segmentation = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        pedestrian = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        # Background is ID 0
        instance = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        categories_map = []
        bottom_corners = {}

        for annotation_token in rec['anns']:
            # Filter out all non vehicle instances
            annotation = self.nusc.get('sample_annotation', annotation_token)

            if self.cfg.DATASET.FILTER_INVISIBLE_VEHICLES and int(annotation['visibility_token']) == 1 and in_pred is False:
                continue
            if in_pred is True and annotation['instance_token'] not in instance_map:
                continue

            # NuScenes filter
            if 'vehicle' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                instance_id = instance_map[annotation['instance_token']]
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(instance, [poly_region], instance_id)
                cv2.fillPoly(segmentation, [poly_region], 1.0)
                categories_map.append([annotation, np.mean(corners, axis=1)])
                bottom_corners[annotation['token']] = corners
            elif 'human' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(pedestrian, [poly_region], 1.0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def _get_poly_region_in_image(self, instance_annotation, ego_translation, ego_rotation):
        box = Box(
            instance_annotation['translation'], instance_annotation['size'], Quaternion(instance_annotation['rotation'])
        )
        # import pdb; pdb.set_trace()
        box.translate(ego_translation)
        box.rotate(ego_rotation)
        corners = box.corners()

        pts = box.bottom_corners()[:2].T
        pts = np.round((pts - self.bev_start_position[:2] + self.bev_resolution[:2] / 2.0) / self.bev_resolution[:2]).astype(np.int32)
        pts[:, [1, 0]] = pts[:, [0, 1]]

        z = box.bottom_corners()[2, 0]
        return pts, z, corners

    def get_label(self, rec, instance_map, in_pred):
        segmentation_np, instance_np, pedestrian_np, instance_map, categories_map, bottom_corners = \
            self.get_birds_eye_view_label(rec, instance_map, in_pred)
        segmentation = torch.from_numpy(segmentation_np).long().unsqueeze(0).unsqueeze(0)
        instance = torch.from_numpy(instance_np).long().unsqueeze(0)
        pedestrian = torch.from_numpy(pedestrian_np).long().unsqueeze(0).unsqueeze(0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def get_future_egomotion(self, rec, index):
        rec_t0 = rec

        # Identity
        future_egomotion = np.eye(4, dtype=np.float32)

        if index < len(self.ixes) - 1:
            rec_t1 = self.ixes[index + 1]

            if rec_t0['scene_token'] == rec_t1['scene_token']:
                egopose_t0 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t0['data']['LIDAR_TOP'])['ego_pose_token']
                )
                egopose_t1 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t1['data']['LIDAR_TOP'])['ego_pose_token']
                )

                egopose_t0 = convert_egopose_to_matrix_numpy(egopose_t0)
                egopose_t1 = convert_egopose_to_matrix_numpy(egopose_t1)

                future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
                future_egomotion[3, :3] = 0.0
                future_egomotion[3, 3] = 1.0

        future_egomotion = torch.Tensor(future_egomotion).float()

        # Convert to 6DoF vector
        future_egomotion = mat2pose_vec(future_egomotion)
        return future_egomotion.unsqueeze(0)

    def get_trajectory_sampling(self, rec=None, sample_indice=None):
        if rec is None and sample_indice is None:
            raise ValueError("No valid input rec or token")
        if rec is None and sample_indice is not None:
            rec = self.ixes[sample_indice]

        ref_scene = self.nusc.get("scene", rec['scene_token'])

        # vm_msgs = self.nusc_can.get_messages(ref_scene['name'], 'vehicle_monitor')
        # vm_uts = [msg['utime'] for msg in vm_msgs]
        pose_msgs = self.nusc_can.get_messages(ref_scene['name'],'pose')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_msgs = self.nusc_can.get_messages(ref_scene['name'], 'steeranglefeedback')
        steer_uts = [msg['utime'] for msg in steer_msgs]

        ref_utime = rec['timestamp']
        # vm_index = locate_message(vm_uts, ref_utime)
        # vm_data = vm_msgs[vm_index]
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]

        # initial speed
        # v0 = vm_data["vehicle_speed"] / 3.6  # km/h to m/s
        v0 = pose_data["vel"][0]  # [0] means longitudinal velocity  m/s

        # curvature (positive: turn left)
        # steering = np.deg2rad(vm_data["steering"])
        steering = steer_data["value"]

        location = self.scene2map[ref_scene['name']]
        # flip x axis if in left-hand traffic (singapore)
        flip_flag = True if location.startswith('singapore') else False
        if flip_flag:
            steering *= -1
        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories


    def generate_centerlines(self, rec, map_name):


        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        x, y = egopose["translation"][0], egopose["translation"][1]
        nusc_map = self.nusc_maps[map_name]
        radius_init = 2
        closest_lane = None
        while closest_lane == None or len(closest_lane) < 5:
            # keep increasing search space until lane is found
            closest_lane = nusc_map.get_closest_lane(x, y, radius=radius_init)
            radius_init *= 2
        candidates_future = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        candidates_past = dfs_incoming(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        lane_ids = []
        for past_lane_seq in candidates_past:
            for future_lane_seq in candidates_future:
                lane_ids.append(past_lane_seq + future_lane_seq[1:])
        lane_ids = remove_overlapping_lane_seq(lane_ids)
        centerlines = []
        for lane_id in lane_ids:
            centerline = []
            for lane in lane_id:
                lane_record = nusc_map.get_arcline_path(lane)
                cx = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                centerline.extend(cx)
            centerlines.append(np.array(centerline))
        lanes = self.nusc_maps[map_name].get_records_in_radius(x, y, 2000, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        discrete_points = self.nusc_maps[map_name].discretize_lanes(lanes, resolution_meters=resolution_meters)
        lanes = []
        for lane_id, points in discrete_points.items():
            lanes.append(np.array(points)[:, :2])
     
        patch_box = box_coords
        patch_angle = rot * 180 / np.pi

        patch_x = patch_box[0]
        patch_y = patch_box[1]

        patch = get_patch_coord(patch_box, patch_angle)

        line_list = []
        for lane in centerlines:
            line = LineString([(arr[0], arr[1]) for arr in lane[:, :2]])
            c = lane[:, :2]
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.rotate(new_line, -patch_angle, origin=(patch_x, patch_y), use_radians=False)
                new_line = affinity.affine_transform(new_line,
                                                        [1.0, 0.0, 0.0, 1.0, -patch_x, -patch_y])
                line_list.append(new_line)
        
        local_box = (0.0, 0.0, patch_box[2], patch_box[3])
        patch_x, patch_y, patch_h, patch_w = local_box

        patch = get_patch_coord(local_box)

        canvas_h = canvas_size[0]
        canvas_w = canvas_size[1]
        scale_height = canvas_h/patch_h
        scale_width = canvas_w/patch_w

        trans_x = -patch_x + patch_w / 2.0
        trans_y = -patch_y + patch_h / 2.0

        overall_lines = []
        for line in line_list:
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.affine_transform(new_line,
                                                     [1.0, 0.0, 0.0, 1.0, trans_x, trans_y])
                new_line = affinity.scale(new_line, xfact=scale_width, yfact=scale_height, origin=(0, 0))
                # overall_lines.append([list(new_line.coords)])

                gg = np.array(list(new_line.coords))
                # interp_arc(t: int, px: np.ndarray, py: np.ndarray) -> np.ndarray:
                gg = interp_arc(t=2000, px=gg[:,0], py=gg[:,1])
                overall_lines.append(gg)

        return overall_lines

    def voxelize_hd_map(self, rec):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        elements = self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS
        hd_features = self.nusc_maps[map_name].get_map_mask(box_coords, rot * 180 / np.pi , elements, canvas_size=canvas_size)
        #traffic = self.hd_traffic_light(map_name, center, stretch, dx, bx, canvas_size)
        #return torch.from_numpy(np.concatenate((hd_features, traffic), axis=0)[None]).float()
        hd_features = torch.from_numpy(hd_features[None]).float()
        hd_features = torch.transpose(hd_features,-2,-1) # (y,x) replace horizontal and vertical coordinates
        return hd_features

    def hd_traffic_light(self, map_name, center, stretch, dx, bx, canvas_size):

        roads = np.zeros(canvas_size)
        my_patch = (
            center[0] - stretch[0],
            center[1] - stretch[1],
            center[0] + stretch[0],
            center[1] + stretch[1],
        )
        tl_token = self.nusc_maps[map_name].get_records_in_patch(my_patch, ['traffic_light'], mode='intersect')['traffic_light']
        polys = []
        for token in tl_token:
            road_token =self.nusc_maps[map_name].get('traffic_light', token)['from_road_block_token']
            pt = self.nusc_maps[map_name].get('road_block', road_token)['polygon_token']
            polygon = self.nusc_maps[map_name].extract_polygon(pt)
            polys.append(np.array(polygon.exterior.xy).T)

        def get_rot(h):
            return torch.Tensor([
                [np.cos(h), np.sin(h)],
                [-np.sin(h), np.cos(h)],
            ])
        # convert to local coordinates in place
        rot = get_rot(np.arctan2(center[3], center[2])).T
        for rowi in range(len(polys)):
            polys[rowi] -= center[:2]
            polys[rowi] = np.dot(polys[rowi], rot)

        for la in polys:
            pts = (la - bx) / dx
            pts = np.int32(np.around(pts))
            cv2.fillPoly(roads, [pts], 1)

        return roads[None]

    def get_gt_trajectory(self, rec, ref_index):
        n_output = self.cfg.N_FUTURE_FRAMES
        n_past = self.cfg.TIME_RECEPTIVE_FIELD
        gt_trajectory = np.zeros((n_output+1+(n_past-1), 3), np.float64)

        """
            egopos_cur -> gives you sensor to global.
                Suppose a point P in sensor-frame.
                    egopos_cur.P gives you the point's global coordinates
        """
        egopose_cur = get_global_pose(rec, self.nusc, inverse=True)

        for i in range(n_past - 1):
            index = ref_index - (n_past - 1) + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i, :] = [origin[0], origin[1], theta]

        for i in range(n_output+1):
            index = ref_index + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                """
                    egopos_future -> gives you global to future.
                        Suppose a point P in ego-frame.
                """
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i + (n_past - 1), :] = [origin[0], origin[1], theta]

        """
            ?
        """
        if gt_trajectory[-1][0] >= 2:
            command = 'RIGHT'
        elif gt_trajectory[-1][0] <= -2:
            command = 'LEFT'
        else:
            command = 'FORWARD'

        return gt_trajectory, command

    def get_routed_map(self, gt_points):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        canvas_size = (
            int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
            int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        roads = np.zeros(canvas_size)
        W = 1.85
        pts = np.array([
            [-4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, -W / 2.],
            [-4.084 / 2. + 0.5, -W / 2.],
        ])
        pts = (pts - bx) / dx
        pts[:, [0, 1]] = pts[:, [1, 0]]

        pts = np.int32(np.around(pts))
        cv2.fillPoly(roads, [pts], 1)

        gt_points = gt_points[:-1].numpy()
        # 坐标原点在左上角
        target = pts.copy()
        target[:,0] = pts[:,0] + gt_points[0] / dx[0]
        target[:,1] = pts[:,1] - gt_points[1] / dx[1]
        target = np.int32(np.around(target))
        cv2.fillPoly(roads, [target], 1)
        return roads

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

    def __getitem__(self, index):
        """
        Returns
        -------
            data: dict with the following keys:
                image: torch.Tensor<float> (T, N, 3, H, W)
                    normalised cameras images with T the sequence length, and N the number of cameras.
                intrinsics: torch.Tensor<float> (T, N, 3, 3)
                    intrinsics containing resizing and cropping parameters.
                extrinsics: torch.Tensor<float> (T, N, 4, 4)
                    6 DoF pose from world coordinates to camera coordinates.
                segmentation: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                    (H_bev, W_bev) are the pixel dimensions in bird's-eye view.
                instance: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                centerness: torch.Tensor<float> (T, 1, H_bev, W_bev)
                offset: torch.Tensor<float> (T, 2, H_bev, W_bev)
                flow: torch.Tensor<float> (T, 2, H_bev, W_bev)
                future_egomotion: torch.Tensor<float> (T, 6)
                    6 DoF egomotion t -> t+1

        """
        data = {}
        keys = ['image', 'intrinsics', 'extrinsics', 'depths',
                'segmentation', 'instance', 'centerness', 'offset', 'flow', 'pedestrian',
                'future_egomotion', 'hdmap', 'gt_trajectory', 'indices', 'poses', 'unnormalized_images',\
                'rotations', 'translations', 'camera_intrinsics', 'point_clouds', 'point_clouds_labels', \
                'point_clouds_panoptic_labels', 'scene_token', 'categories_map', 'panoptic_mappings_list', \
                'bottom_corners'
                ]
        for key in keys:
            data[key] = []

        instance_map = {}
        # Loop over all the frames in the sequence.
        total_centerlines = []
        for i, index_t in enumerate(self.indices[index]):
            if i >= self.receptive_field:
                in_pred = True
            else:
                in_pred = False
            rec = self.ixes[index_t]
            if i == 0:
                data['rec_first'] = rec            
            map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

            if i < self.receptive_field:
                images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, scene_token, panoptic_mappings_list = self.get_input_data(rec)
                data['image'].append(images)
                data['unnormalized_images'].append(unnormalized_images)
                data['intrinsics'].append(intrinsics)
                data['extrinsics'].append(extrinsics)
                data['depths'].append(depths)
                data['translations'].append(translations)
                data['rotations'].append(rotations)
                data['camera_intrinsics'].append(camera_intrinsics)
                data['point_clouds'].append(point_clouds)
                data['point_clouds_labels'].append(point_clouds_labels)
                data['scene_token'].append(scene_token)

            segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners = self.get_label(rec, instance_map, in_pred)

            future_egomotion = self.get_future_egomotion(rec, index_t)
            hd_map_feature = self.voxelize_hd_map(rec)

            data['segmentation'].append(segmentation)
            data['instance'].append(instance)
            data['pedestrian'].append(pedestrian)
            data['future_egomotion'].append(future_egomotion)
            data['hdmap'].append(hd_map_feature)
            data['indices'].append(index_t)
            data['categories_map'].append(categories_map)
            # import pdb; pdb.set_trace()
            data['point_clouds_panoptic_labels'].append(point_clouds_panoptic_labels)
            data['panoptic_mappings_list'].append(panoptic_mappings_list)
            data['bottom_corners'].append(bottom_corners)

        for key, value in data.items():
            if key in ['image', 'intrinsics', 'extrinsics', 'depths', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian', 'camera_intrinsics', 'rotations', 'translations', 'unnormalized_images']:
                if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
                    continue
                data[key] = torch.cat(value, dim=0)

        # images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds

        # for key, value in data.items():
        #     if key in ['image', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian', 'unnormalized_images']:
        #         if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
        #             continue
        #         data[key] = torch.cat(value, dim=0)

        data['target_point'] = torch.tensor([0., 0.])
        instance_centerness, instance_offset, instance_flow = convert_instance_mask_to_center_and_offset_label(
            data['instance'], data['future_egomotion'],
            num_instances=len(instance_map), ignore_index=self.cfg.DATASET.IGNORE_INDEX, subtract_egomotion=True,
            spatial_extent=self.spatial_extent,
        )
        data['centerness'] = instance_centerness
        data['offset'] = instance_offset
        data['flow'] = instance_flow

        return data
data/stp3/datas/fullNuscenesData.py
llmbev/talk2bev
import os
from PIL import Image

import numpy as np
import cv2
import torch
import torch.utils.data
import torch.nn.functional as F
import torchvision
import os.path as osp

from shapely.geometry import LineString

from pyquaternion import Quaternion
from shapely import affinity
from nuscenes.nuscenes import NuScenes, NuScenesExplorer
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.splits import create_splits_scenes
from nuscenes.utils.data_classes import Box
from nuscenes.lidarseg.lidarseg_utils import colormap_to_colors, plt_to_cv2, get_stats, get_labels_in_coloring, create_lidarseg_legend, paint_points_label
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.eval.common.utils import quaternion_yaw
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from nuscenes.utils.data_classes import PointCloud, LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix
from nuscenes.utils.data_io import load_bin_file
from stp3.utils.tools import ( gen_dx_bx, get_nusc_maps)
import matplotlib.pyplot as plt 
from math import cos, sin

from stp3.utils.geometry import (
    resize_and_crop_image,
    update_intrinsics,
    calculate_birds_eye_view_parameters,
    convert_egopose_to_matrix_numpy,
    pose_vec2mat,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
    get_global_pose,
    get_patch_coord,
    dfs,
    dfs_incoming,
    remove_overlapping_lane_seq,
    extend_both_sides,
    project_to_frenet_frame,
    project_to_cartesian_frame,
    interp_arc,
    interp_arc_np,
    denoise,
    interp_polyline_by_fixed_waypt_interval
)
from stp3.utils.instance import convert_instance_mask_to_center_and_offset_label
import stp3.utils.sampler as trajectory_sampler
from stp3.utils.spline import Spline2D
from stp3.utils.Optnode import *

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

resolution_meters = 0.1
t_fin = 3.5
num = 61
nvar = 11
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
debug = False
brdebug = False
problem = OPTNode(t_fin=t_fin, num=num)

def distance(x1, y1, z1, x2, y2, z2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + 1 * (z1 - z2) ** 2) ** ( 1 / 2)
 
# Function to calculate K closest points
def kClosest(points, target, K):
    pts = []
    n = len(points)
    d = []

    for i in range(n):
        d.append({
            "first": distance(points[i][0], points[i][1], points[i][2], target[0], target[1], target[2]),
            "second": i
        })
     
    d = sorted(d, key=lambda l:l["first"])
 
    for i in range(K):
        pt = []
        pt.append(points[d[i]["second"]][0])
        pt.append(points[d[i]["second"]][1])
        pt.append(points[d[i]["second"]][2])
        pt.append(points[d[i]["second"]][3])
        pts.append(pt)
 
    return pts

class FuturePredictionDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5 #SECOND
    def __init__(self, nusc, is_train, cfg):
        self.nusc = nusc
        self.dataroot = self.nusc.dataroot
        self.nusc_exp = NuScenesExplorer(nusc)
        self.nusc_can = NuScenesCanBus(dataroot=self.dataroot)
        self.is_train = is_train
        self.cfg = cfg
        self.problem = OPTNode(t_fin=3.5, num=self.cfg.N_FUTURE_FRAMES + 1)

        if self.is_train == 0:
            self.mode = 'train'
        elif self.is_train == 1:
            self.mode = 'val'
        elif self.is_train == 2:
            self.mode = 'test'
        else:
            raise NotImplementedError

        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD

        self.scenes = self.get_scenes()
        self.ixes = self.prepro()
        self.indices = self.get_indices()

        # Image resizing and cropping
        self.augmentation_parameters = self.get_resizing_and_cropping_parameters()

        # Normalising input images
        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # Bird's-eye view parameters
        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_resolution, self.bev_start_position, self.bev_dimension = (
            bev_resolution.numpy(), bev_start_position.numpy(), bev_dimension.numpy()
        )

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # The number of sampled trajectories
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        # HD-map feature extractor
        self.nusc_maps = get_nusc_maps(self.cfg.DATASET.MAP_FOLDER)
        self.scene2map = {}
        for sce in self.nusc.scene:
            log = self.nusc.get('log', sce['log_token'])
            self.scene2map[sce['name']] = log['location']
        self.save_dir = cfg.DATASET.SAVE_DIR

        self.curbatchnum1 = None
        self.curbatchnum2 = None
        if self.cfg.PLANNING.DENSE.ENABLED:
            self.centerlines_path = self.cfg.PLANNING.DENSE.PATH

    def get_scenes(self):
        # filter by scene split
        split = {'v1.0-trainval': {0: 'train', 1: 'val', 2: 'test'},
                 'v1.0-mini': {0: 'mini_train', 1: 'mini_val'},}[
            self.nusc.version
        ][self.is_train]

        blacklist = [419] + self.nusc_can.can_blacklist  # # scene-0419 does not have vehicle monitor data
        blacklist = ['scene-' + str(scene_no).zfill(4) for scene_no in blacklist]

        scenes = create_splits_scenes()[split][:]
        for scene_no in blacklist:
            if scene_no in scenes:
                scenes.remove(scene_no)

        return scenes

    def prepro(self):
        samples = [samp for samp in self.nusc.sample]

        # remove samples that aren't in this split
        samples = [samp for samp in samples if self.nusc.get('scene', samp['scene_token'])['name'] in self.scenes]
        
        # sort by scene, timestamp (only to make chronological viz easier)
        samples.sort(key=lambda x: (self.nusc.get("scene", x['scene_token'])['name'], x['timestamp']))

        return samples

    # def get_indices(self):
    #     indices = []
    #     for index in range(len(self.ixes)):
    #         is_valid_data = True
    #         previous_rec = None
    #         current_indices = []
    #         for t in range(self.sequence_length):
    #             index_t = index + t
    #             # Going over the dataset size limit.
    #             if index_t >= len(self.ixes):
    #                 is_valid_data = False
    #                 break
    #             rec = self.ixes[index_t]
    #             # Check if scene is the same
    #             if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
    #                 is_valid_data = False
    #                 break

    #             current_indices.append(index_t)
    #             previous_rec = rec

    #         if is_valid_data:
    #             indices.append(current_indices)

    #     return np.asarray(indices)

    def get_indices(self):
        indices = []
        scene_token_mapping = {}
        ixes_new = []
        selected_indices = [7, 15, 11]
        for index in range(len(self.ixes)):
            if scene_token_mapping.get(self.ixes[index]['scene_token']) == None:
                scene_token_mapping[self.ixes[index]['scene_token']] = 1
                continue
            scene_token_mapping[self.ixes[index]['scene_token']] = scene_token_mapping[self.ixes[index]['scene_token']] + 1
            if scene_token_mapping[self.ixes[index]['scene_token']] not in selected_indices:
                continue
            is_valid_data = True
            previous_rec = None
            current_indices = []
            for t in range(self.sequence_length):
                index_t = index + t
                # Going over the dataset size limit.
                if index_t >= len(self.ixes):
                    is_valid_data = False
                    break
                rec = self.ixes[index_t]
                ixes_new.append(self.ixes[index_t])
                # Check if scene is the same
                if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
                    is_valid_data = False
                    break

                current_indices.append(len(ixes_new) - 1)
                previous_rec = rec

            if is_valid_data:
                indices.append(current_indices)

        self.ixes = ixes_new
        return np.asarray(indices)

    def get_resizing_and_cropping_parameters(self):
        original_height, original_width = self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH
        final_height, final_width = self.cfg.IMAGE.FINAL_DIM

        resize_scale = self.cfg.IMAGE.RESIZE_SCALE
        resize_dims = (int(original_width * resize_scale), int(original_height * resize_scale))
        resized_width, resized_height = resize_dims

        crop_h = self.cfg.IMAGE.TOP_CROP
        crop_w = int(max(0, (resized_width - final_width) / 2))
        # Left, top, right, bottom crops.
        crop = (crop_w, crop_h, crop_w + final_width, crop_h + final_height)

        if resized_width != final_width:
            print('Zero padding left and right parts of the image.')
        if crop_h + final_height != resized_height:
            print('Zero padding bottom part of the image.')

        return {'scale_width': resize_scale,
                'scale_height': resize_scale,
                'resize_dims': resize_dims,
                'crop': crop,
                }

    def get_input_data(self, rec):
        """
        Parameters
        ----------
            rec: nuscenes identifier for a given timestamp

        Returns
        -------
            images: torch.Tensor<float> (N, 3, H, W)
            intrinsics: torch.Tensor<float> (3, 3)
            extrinsics: torch.Tensor(N, 4, 4)
        """
        images = []
        unnormalized_images = []
        intrinsics = []
        extrinsics = []
        depths = []
        translations = []
        rotations = []
        camera_intrinsics = []
        point_clouds = []
        point_clouds_labels = []
        point_clouds_panoptic_labels = []
        panoptic_mappings_list = []
        extra_rotations = []
        extra_translations = []
        cameras = self.cfg.IMAGE.NAMES

        # # The extrinsics we want are from the camera sensor to "flat egopose" as defined
        # # https://github.com/nutonomy/nuscenes-devkit/blob/9b492f76df22943daf1dc991358d3d606314af27/python-sdk/nuscenes/nuscenes.py#L279
        # # which corresponds to the position of the lidar.
        # # This is because the labels are generated by projecting the 3D bounding box in this lidar's reference frame.

        # # From lidar egopose to world.
        lidar_sample = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        lidar_pose = self.nusc.get('ego_pose', lidar_sample['ego_pose_token'])
        yaw = Quaternion(lidar_pose['rotation']).yaw_pitch_roll[0]
        lidar_rotation = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        lidar_translation = np.array(lidar_pose['translation'])[:, None]
        lidar_to_world = np.vstack([
            np.hstack((lidar_rotation.rotation_matrix, lidar_translation)),
            np.array([0, 0, 0, 1])
        ])

        camera_sample_ = self.nusc.get('sample_data', rec['data']["CAM_FRONT"])
        my_sample = self.nusc.get('sample', camera_sample_['sample_token'])
        my_scene_token = my_sample['scene_token']

        # # Here we just grab the front camera and the point sensor.
        pointsensor = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        pcl_path = osp.join(self.nusc.dataroot, pointsensor['filename'])
        # panoptic_mappings = [samp.token for samp in self.nusc.get_sample_data(rec['data']['LIDAR_TOP'])[1]]
        panoptic_mappings = []
        for ok in rec['anns']:
            panoptic_mappings.append(self.nusc.get('sample_annotation', ok)['token'])
        panoptic_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get('panoptic', rec['data']['LIDAR_TOP'])['filename'])
        panoptic_labels = load_bin_file(panoptic_labels_filename, type='panoptic').astype(np.int32)
        point_clouds_panoptic_labels.append(panoptic_labels)
        for cam in cameras:
            camera_sample = self.nusc.get('sample_data', rec['data'][cam])
            sample_token = camera_sample['token']

            camm = self.nusc.get('sample_data', rec['data'][cam])
            lidarseg_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get("lidarseg", rec['data']['LIDAR_TOP'])['filename'])
            points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)


            cs_record = self.nusc.get('calibrated_sensor', pointsensor['calibrated_sensor_token'])
            pc = LidarPointCloud.from_file(pcl_path)
            ### LIDAR TO EGO_{TL} TRANSLATION -> CONSTANT
            pc.rotate(Quaternion(cs_record['rotation']).rotation_matrix)
            pc.translate(np.array(cs_record['translation']))

            point_clouds.append(torch.tensor(np.array(pc.points)).unsqueeze(0).unsqueeze(0)) # 1, 1, N, 4
            point_clouds_labels.append(torch.tensor(np.array(points_label)).unsqueeze(0).unsqueeze(0)) # 1, 1, N, 

            poserecord = self.nusc.get('ego_pose', pointsensor['ego_pose_token'])
            extra_rotations.append(Quaternion(poserecord['rotation']).rotation_matrix)
            extra_translations.append(np.array(poserecord['translation']))
            pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix)
            pc.translate(np.array(poserecord['translation']))

            # Third step: transform from global into the ego vehicle frame for the timestamp of the image.
            poserecord = self.nusc.get('ego_pose', camm['ego_pose_token'])
            extra_rotations.append(Quaternion(poserecord['rotation']).rotation_matrix.T)
            extra_translations.append(np.array(poserecord['translation']))            
            pc.translate(-np.array(poserecord['translation']))
            pc.rotate(Quaternion(poserecord['rotation']).rotation_matrix.T)

            # Transformation from world to egopose
            car_egopose = self.nusc.get('ego_pose', camera_sample['ego_pose_token'])
            egopose_rotation = Quaternion(car_egopose['rotation']).inverse
            egopose_translation = -np.array(car_egopose['translation'])[:, None]
            world_to_car_egopose = np.vstack([
                np.hstack((egopose_rotation.rotation_matrix, egopose_rotation.rotation_matrix @ egopose_translation)),
                np.array([0, 0, 0, 1])
            ])

            # From egopose to sensor
            sensor_sample = self.nusc.get('calibrated_sensor', camera_sample['calibrated_sensor_token'])
            intrinsic = torch.Tensor(sensor_sample['camera_intrinsic'])
            sensor_rotation = Quaternion(sensor_sample['rotation'])
            sensor_translation = np.array(sensor_sample['translation'])[:, None]
            car_egopose_to_sensor = np.vstack([
                np.hstack((sensor_rotation.rotation_matrix, sensor_translation)),
                np.array([0, 0, 0, 1])
            ])
            car_egopose_to_sensor = np.linalg.inv(car_egopose_to_sensor)

            # Combine all the transformation.
            # From sensor to lidar.
            lidar_to_sensor = car_egopose_to_sensor @ world_to_car_egopose @ lidar_to_world
            sensor_to_lidar = torch.from_numpy(np.linalg.inv(lidar_to_sensor)).float()

            # Load image
            image_filename = os.path.join(self.dataroot, camera_sample['filename'])
            img = Image.open(image_filename)
            unnormalized_image = torch.tensor(np.array(img))
            # Resize and crop
            img = resize_and_crop_image(
                img, resize_dims=self.augmentation_parameters['resize_dims'], crop=self.augmentation_parameters['crop']
            )
            # Normalise image
            normalised_img = self.normalise_image(img)

            # # Combine resize/cropping in the intrinsics
            top_crop = self.augmentation_parameters['crop'][1]
            left_crop = self.augmentation_parameters['crop'][0]
            intrinsic = update_intrinsics(
                intrinsic, top_crop, left_crop,
                scale_width=self.augmentation_parameters['scale_width'],
                scale_height=self.augmentation_parameters['scale_height']
            )

            # Get Depth
            # Depth data should under the dataroot path 
            if self.cfg.LIFT.GT_DEPTH:
                base_root = os.path.join(self.dataroot, 'depths') 
                filename = os.path.basename(camera_sample['filename']).split('.')[0] + '.npy'
                depth_file_name = os.path.join(base_root, cam, 'npy', filename)
                depth = torch.from_numpy(np.load(depth_file_name)).unsqueeze(0).unsqueeze(0)
                depth = F.interpolate(depth, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear')
                depth = depth.squeeze()
                crop = self.augmentation_parameters['crop']
                depth = depth[crop[1]:crop[3], crop[0]:crop[2]]
                depth = torch.round(depth)
                depths.append(depth.unsqueeze(0).unsqueeze(0))
            
            camm = self.nusc.get('sample_data', rec['data'][cam])
            cs_record = self.nusc.get('calibrated_sensor', camm['calibrated_sensor_token'])
            images.append(normalised_img.unsqueeze(0).unsqueeze(0))
            intrinsics.append(intrinsic.unsqueeze(0).unsqueeze(0))
            extrinsics.append(sensor_to_lidar.unsqueeze(0).unsqueeze(0))
            unnormalized_images.append(unnormalized_image.unsqueeze(0).unsqueeze(0))
            rotations.append(torch.tensor(Quaternion(cs_record['rotation']).rotation_matrix.T).unsqueeze(0).unsqueeze(0))
            translations.append(torch.tensor(-np.array(cs_record['translation'])).unsqueeze(0).unsqueeze(0))
            camera_intrinsics.append(torch.tensor(np.array(cs_record['camera_intrinsic'])).unsqueeze(0).unsqueeze(0))

        images, intrinsics, extrinsics, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, panoptic_mappings_list = (torch.cat(images, dim=1),
                                            torch.cat(intrinsics, dim=1),
                                            torch.cat(extrinsics, dim=1),
                                            torch.cat(unnormalized_images, dim=1),
                                            torch.cat(rotations, dim=1),
                                            torch.cat(translations, dim=1),
                                            torch.cat(camera_intrinsics, dim=1),
                                            torch.cat(point_clouds, dim=1),
                                            torch.cat(point_clouds_labels, dim=1),
                                            point_clouds_panoptic_labels,
                                            panoptic_mappings
                                          )
        if len(depths) > 0:
            depths = torch.cat(depths, dim=1)

        return images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, my_scene_token, panoptic_mappings_list

    def _get_top_lidar_pose(self, rec):
        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        trans = -np.array(egopose['translation'])
        yaw = Quaternion(egopose['rotation']).yaw_pitch_roll[0]
        # rot = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        rot = Quaternion(egopose['rotation']).inverse
        return trans, rot

    def get_depth_from_lidar(self, lidar_sample, cam_sample):
        points, coloring, im = self.nusc_exp.map_pointcloud_to_image(lidar_sample, cam_sample)
        tmp_cam = np.zeros((self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH))
        points = points.astype(np.int)
        tmp_cam[points[1, :], points[0,:]] = coloring
        tmp_cam = torch.from_numpy(tmp_cam).unsqueeze(0).unsqueeze(0)
        tmp_cam = F.interpolate(tmp_cam, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear', align_corners=False, recompute_scale_factor=True)
        tmp_cam = tmp_cam.squeeze()
        crop = self.augmentation_parameters['crop']
        tmp_cam = tmp_cam[crop[1]:crop[3], crop[0]:crop[2]]
        tmp_cam = torch.round(tmp_cam)
        return tmp_cam


    def get_birds_eye_view_label(self, rec, instance_map, in_pred):
        translation, rotation = self._get_top_lidar_pose(rec)
        segmentation = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        pedestrian = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        # Background is ID 0
        instance = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        categories_map = []
        bottom_corners = {}

        for annotation_token in rec['anns']:
            # Filter out all non vehicle instances
            annotation = self.nusc.get('sample_annotation', annotation_token)

            if self.cfg.DATASET.FILTER_INVISIBLE_VEHICLES and int(annotation['visibility_token']) == 1 and in_pred is False:
                continue
            if in_pred is True and annotation['instance_token'] not in instance_map:
                continue

            # NuScenes filter
            if 'vehicle' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                instance_id = instance_map[annotation['instance_token']]
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(instance, [poly_region], instance_id)
                cv2.fillPoly(segmentation, [poly_region], 1.0)
                categories_map.append([annotation, np.mean(corners, axis=1)])
                bottom_corners[annotation['token']] = corners
            elif 'human' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(pedestrian, [poly_region], 1.0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def _get_poly_region_in_image(self, instance_annotation, ego_translation, ego_rotation):
        box = Box(
            instance_annotation['translation'], instance_annotation['size'], Quaternion(instance_annotation['rotation'])
        )
        # import pdb; pdb.set_trace()
        box.translate(ego_translation)
        box.rotate(ego_rotation)
        corners = box.corners()

        pts = box.bottom_corners()[:2].T
        pts = np.round((pts - self.bev_start_position[:2] + self.bev_resolution[:2] / 2.0) / self.bev_resolution[:2]).astype(np.int32)
        pts[:, [1, 0]] = pts[:, [0, 1]]

        z = box.bottom_corners()[2, 0]
        return pts, z, corners

    def get_label(self, rec, instance_map, in_pred):
        segmentation_np, instance_np, pedestrian_np, instance_map, categories_map, bottom_corners = \
            self.get_birds_eye_view_label(rec, instance_map, in_pred)
        segmentation = torch.from_numpy(segmentation_np).long().unsqueeze(0).unsqueeze(0)
        instance = torch.from_numpy(instance_np).long().unsqueeze(0)
        pedestrian = torch.from_numpy(pedestrian_np).long().unsqueeze(0).unsqueeze(0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def get_future_egomotion(self, rec, index):
        rec_t0 = rec

        # Identity
        future_egomotion = np.eye(4, dtype=np.float32)

        if index < len(self.ixes) - 1:
            rec_t1 = self.ixes[index + 1]

            if rec_t0['scene_token'] == rec_t1['scene_token']:
                egopose_t0 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t0['data']['LIDAR_TOP'])['ego_pose_token']
                )
                egopose_t1 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t1['data']['LIDAR_TOP'])['ego_pose_token']
                )

                egopose_t0 = convert_egopose_to_matrix_numpy(egopose_t0)
                egopose_t1 = convert_egopose_to_matrix_numpy(egopose_t1)

                future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
                future_egomotion[3, :3] = 0.0
                future_egomotion[3, 3] = 1.0

        future_egomotion = torch.Tensor(future_egomotion).float()

        # Convert to 6DoF vector
        future_egomotion = mat2pose_vec(future_egomotion)
        return future_egomotion.unsqueeze(0)

    def get_trajectory_sampling(self, rec=None, sample_indice=None):
        if rec is None and sample_indice is None:
            raise ValueError("No valid input rec or token")
        if rec is None and sample_indice is not None:
            rec = self.ixes[sample_indice]

        ref_scene = self.nusc.get("scene", rec['scene_token'])

        # vm_msgs = self.nusc_can.get_messages(ref_scene['name'], 'vehicle_monitor')
        # vm_uts = [msg['utime'] for msg in vm_msgs]
        pose_msgs = self.nusc_can.get_messages(ref_scene['name'],'pose')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_msgs = self.nusc_can.get_messages(ref_scene['name'], 'steeranglefeedback')
        steer_uts = [msg['utime'] for msg in steer_msgs]

        ref_utime = rec['timestamp']
        # vm_index = locate_message(vm_uts, ref_utime)
        # vm_data = vm_msgs[vm_index]
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]

        # initial speed
        # v0 = vm_data["vehicle_speed"] / 3.6  # km/h to m/s
        v0 = pose_data["vel"][0]  # [0] means longitudinal velocity  m/s

        # curvature (positive: turn left)
        # steering = np.deg2rad(vm_data["steering"])
        steering = steer_data["value"]

        location = self.scene2map[ref_scene['name']]
        # flip x axis if in left-hand traffic (singapore)
        flip_flag = True if location.startswith('singapore') else False
        if flip_flag:
            steering *= -1
        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories


    def generate_centerlines(self, rec, map_name):


        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        x, y = egopose["translation"][0], egopose["translation"][1]
        nusc_map = self.nusc_maps[map_name]
        radius_init = 2
        closest_lane = None
        while closest_lane == None or len(closest_lane) < 5:
            # keep increasing search space until lane is found
            closest_lane = nusc_map.get_closest_lane(x, y, radius=radius_init)
            radius_init *= 2
        candidates_future = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        candidates_past = dfs_incoming(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        lane_ids = []
        for past_lane_seq in candidates_past:
            for future_lane_seq in candidates_future:
                lane_ids.append(past_lane_seq + future_lane_seq[1:])
        lane_ids = remove_overlapping_lane_seq(lane_ids)
        centerlines = []
        for lane_id in lane_ids:
            centerline = []
            for lane in lane_id:
                lane_record = nusc_map.get_arcline_path(lane)
                cx = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                centerline.extend(cx)
            centerlines.append(np.array(centerline))
        lanes = self.nusc_maps[map_name].get_records_in_radius(x, y, 2000, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        discrete_points = self.nusc_maps[map_name].discretize_lanes(lanes, resolution_meters=resolution_meters)
        lanes = []
        for lane_id, points in discrete_points.items():
            lanes.append(np.array(points)[:, :2])
     
        patch_box = box_coords
        patch_angle = rot * 180 / np.pi

        patch_x = patch_box[0]
        patch_y = patch_box[1]

        patch = get_patch_coord(patch_box, patch_angle)

        line_list = []
        for lane in centerlines:
            line = LineString([(arr[0], arr[1]) for arr in lane[:, :2]])
            c = lane[:, :2]
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.rotate(new_line, -patch_angle, origin=(patch_x, patch_y), use_radians=False)
                new_line = affinity.affine_transform(new_line,
                                                        [1.0, 0.0, 0.0, 1.0, -patch_x, -patch_y])
                line_list.append(new_line)
        
        local_box = (0.0, 0.0, patch_box[2], patch_box[3])
        patch_x, patch_y, patch_h, patch_w = local_box

        patch = get_patch_coord(local_box)

        canvas_h = canvas_size[0]
        canvas_w = canvas_size[1]
        scale_height = canvas_h/patch_h
        scale_width = canvas_w/patch_w

        trans_x = -patch_x + patch_w / 2.0
        trans_y = -patch_y + patch_h / 2.0

        overall_lines = []
        for line in line_list:
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.affine_transform(new_line,
                                                     [1.0, 0.0, 0.0, 1.0, trans_x, trans_y])
                new_line = affinity.scale(new_line, xfact=scale_width, yfact=scale_height, origin=(0, 0))
                # overall_lines.append([list(new_line.coords)])

                gg = np.array(list(new_line.coords))
                # interp_arc(t: int, px: np.ndarray, py: np.ndarray) -> np.ndarray:
                gg = interp_arc(t=2000, px=gg[:,0], py=gg[:,1])
                overall_lines.append(gg)

        return overall_lines

    def voxelize_hd_map(self, rec):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        elements = self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS
        hd_features = self.nusc_maps[map_name].get_map_mask(box_coords, rot * 180 / np.pi , elements, canvas_size=canvas_size)
        #traffic = self.hd_traffic_light(map_name, center, stretch, dx, bx, canvas_size)
        #return torch.from_numpy(np.concatenate((hd_features, traffic), axis=0)[None]).float()
        hd_features = torch.from_numpy(hd_features[None]).float()
        hd_features = torch.transpose(hd_features,-2,-1) # (y,x) replace horizontal and vertical coordinates
        return hd_features

    def hd_traffic_light(self, map_name, center, stretch, dx, bx, canvas_size):

        roads = np.zeros(canvas_size)
        my_patch = (
            center[0] - stretch[0],
            center[1] - stretch[1],
            center[0] + stretch[0],
            center[1] + stretch[1],
        )
        tl_token = self.nusc_maps[map_name].get_records_in_patch(my_patch, ['traffic_light'], mode='intersect')['traffic_light']
        polys = []
        for token in tl_token:
            road_token =self.nusc_maps[map_name].get('traffic_light', token)['from_road_block_token']
            pt = self.nusc_maps[map_name].get('road_block', road_token)['polygon_token']
            polygon = self.nusc_maps[map_name].extract_polygon(pt)
            polys.append(np.array(polygon.exterior.xy).T)

        def get_rot(h):
            return torch.Tensor([
                [np.cos(h), np.sin(h)],
                [-np.sin(h), np.cos(h)],
            ])
        # convert to local coordinates in place
        rot = get_rot(np.arctan2(center[3], center[2])).T
        for rowi in range(len(polys)):
            polys[rowi] -= center[:2]
            polys[rowi] = np.dot(polys[rowi], rot)

        for la in polys:
            pts = (la - bx) / dx
            pts = np.int32(np.around(pts))
            cv2.fillPoly(roads, [pts], 1)

        return roads[None]

    def get_gt_trajectory(self, rec, ref_index):
        n_output = self.cfg.N_FUTURE_FRAMES
        n_past = self.cfg.TIME_RECEPTIVE_FIELD
        gt_trajectory = np.zeros((n_output+1+(n_past-1), 3), np.float64)

        """
            egopos_cur -> gives you sensor to global.
                Suppose a point P in sensor-frame.
                    egopos_cur.P gives you the point's global coordinates
        """
        egopose_cur = get_global_pose(rec, self.nusc, inverse=True)

        for i in range(n_past - 1):
            index = ref_index - (n_past - 1) + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i, :] = [origin[0], origin[1], theta]

        for i in range(n_output+1):
            index = ref_index + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                """
                    egopos_future -> gives you global to future.
                        Suppose a point P in ego-frame.
                """
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i + (n_past - 1), :] = [origin[0], origin[1], theta]

        """
            ?
        """
        if gt_trajectory[-1][0] >= 2:
            command = 'RIGHT'
        elif gt_trajectory[-1][0] <= -2:
            command = 'LEFT'
        else:
            command = 'FORWARD'

        return gt_trajectory, command

    def get_routed_map(self, gt_points):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        canvas_size = (
            int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
            int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        roads = np.zeros(canvas_size)
        W = 1.85
        pts = np.array([
            [-4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, -W / 2.],
            [-4.084 / 2. + 0.5, -W / 2.],
        ])
        pts = (pts - bx) / dx
        pts[:, [0, 1]] = pts[:, [1, 0]]

        pts = np.int32(np.around(pts))
        cv2.fillPoly(roads, [pts], 1)

        gt_points = gt_points[:-1].numpy()
        # 坐标原点在左上角
        target = pts.copy()
        target[:,0] = pts[:,0] + gt_points[0] / dx[0]
        target[:,1] = pts[:,1] - gt_points[1] / dx[1]
        target = np.int32(np.around(target))
        cv2.fillPoly(roads, [target], 1)
        return roads

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

    def __getitem__(self, index):
        """
        Returns
        -------
            data: dict with the following keys:
                image: torch.Tensor<float> (T, N, 3, H, W)
                    normalised cameras images with T the sequence length, and N the number of cameras.
                intrinsics: torch.Tensor<float> (T, N, 3, 3)
                    intrinsics containing resizing and cropping parameters.
                extrinsics: torch.Tensor<float> (T, N, 4, 4)
                    6 DoF pose from world coordinates to camera coordinates.
                segmentation: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                    (H_bev, W_bev) are the pixel dimensions in bird's-eye view.
                instance: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                centerness: torch.Tensor<float> (T, 1, H_bev, W_bev)
                offset: torch.Tensor<float> (T, 2, H_bev, W_bev)
                flow: torch.Tensor<float> (T, 2, H_bev, W_bev)
                future_egomotion: torch.Tensor<float> (T, 6)
                    6 DoF egomotion t -> t+1

        """
        data = {}
        keys = ['image', 'intrinsics', 'extrinsics', 'depths',
                'segmentation', 'instance', 'centerness', 'offset', 'flow', 'pedestrian',
                'future_egomotion', 'hdmap', 'gt_trajectory', 'indices', 'poses', 'unnormalized_images',\
                'rotations', 'translations', 'camera_intrinsics', 'point_clouds', 'point_clouds_labels', \
                'point_clouds_panoptic_labels', 'scene_token', 'categories_map', 'panoptic_mappings_list', \
                'bottom_corners'
                ]
        for key in keys:
            data[key] = []

        instance_map = {}
        # Loop over all the frames in the sequence.
        total_centerlines = []
        for i, index_t in enumerate(self.indices[index]):
            if i >= self.receptive_field:
                in_pred = True
            else:
                in_pred = False
            rec = self.ixes[index_t]
            if i == 0:
                data['rec_first'] = rec            
            map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

            if i < self.receptive_field:
                images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, scene_token, panoptic_mappings_list = self.get_input_data(rec)
                data['image'].append(images)
                data['unnormalized_images'].append(unnormalized_images)
                data['intrinsics'].append(intrinsics)
                data['extrinsics'].append(extrinsics)
                data['depths'].append(depths)
                data['translations'].append(translations)
                data['rotations'].append(rotations)
                data['camera_intrinsics'].append(camera_intrinsics)
                data['point_clouds'].append(point_clouds)
                data['point_clouds_labels'].append(point_clouds_labels)
                data['scene_token'].append(scene_token)

            segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners = self.get_label(rec, instance_map, in_pred)

            future_egomotion = self.get_future_egomotion(rec, index_t)
            hd_map_feature = self.voxelize_hd_map(rec)

            data['segmentation'].append(segmentation)
            data['instance'].append(instance)
            data['pedestrian'].append(pedestrian)
            data['future_egomotion'].append(future_egomotion)
            data['hdmap'].append(hd_map_feature)
            data['indices'].append(index_t)
            data['categories_map'].append(categories_map)
            # import pdb; pdb.set_trace()
            data['point_clouds_panoptic_labels'].append(point_clouds_panoptic_labels)
            data['panoptic_mappings_list'].append(panoptic_mappings_list)
            data['bottom_corners'].append(bottom_corners)

        for key, value in data.items():
            if key in ['image', 'intrinsics', 'extrinsics', 'depths', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian', 'camera_intrinsics', 'rotations', 'translations', 'unnormalized_images']:
                if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
                    continue
                data[key] = torch.cat(value, dim=0)

        # images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds

        # for key, value in data.items():
        #     if key in ['image', 'segmentation', 'instance', 'future_egomotion', 'hdmap', 'pedestrian', 'unnormalized_images']:
        #         if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
        #             continue
        #         data[key] = torch.cat(value, dim=0)

        data['target_point'] = torch.tensor([0., 0.])
        instance_centerness, instance_offset, instance_flow = convert_instance_mask_to_center_and_offset_label(
            data['instance'], data['future_egomotion'],
            num_instances=len(instance_map), ignore_index=self.cfg.DATASET.IGNORE_INDEX, subtract_egomotion=True,
            spatial_extent=self.spatial_extent,
        )
        data['centerness'] = instance_centerness
        data['offset'] = instance_offset
        data['flow'] = instance_flow

        return data
data/stp3/datas/reducedNuscenesData.py
llmbev/talk2bev
import os
from PIL import Image

import numpy as np
import cv2
import torch
import torch.utils.data
import torch.nn.functional as F
import torchvision
import os.path as osp

from shapely.geometry import LineString

from pyquaternion import Quaternion
from shapely import affinity
from nuscenes.nuscenes import NuScenes, NuScenesExplorer
from nuscenes.can_bus.can_bus_api import NuScenesCanBus
from nuscenes.utils.splits import create_splits_scenes
from nuscenes.utils.data_classes import Box
from nuscenes.lidarseg.lidarseg_utils import colormap_to_colors, plt_to_cv2, get_stats, get_labels_in_coloring, create_lidarseg_legend, paint_points_label
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.eval.common.utils import quaternion_yaw
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from nuscenes.utils.data_classes import PointCloud, LidarPointCloud, RadarPointCloud, Box
from nuscenes.utils.geometry_utils import view_points, box_in_image, BoxVisibility, transform_matrix
from nuscenes.utils.data_io import load_bin_file
from stp3.utils.tools import ( gen_dx_bx, get_nusc_maps)
import matplotlib.pyplot as plt 
from math import cos, sin

from stp3.utils.geometry import (
    resize_and_crop_image,
    update_intrinsics,
    calculate_birds_eye_view_parameters,
    convert_egopose_to_matrix_numpy,
    pose_vec2mat,
    mat2pose_vec,
    invert_matrix_egopose_numpy,
    get_global_pose,
    get_patch_coord,
    dfs,
    dfs_incoming,
    remove_overlapping_lane_seq,
    extend_both_sides,
    project_to_frenet_frame,
    project_to_cartesian_frame,
    interp_arc,
    interp_arc_np,
    denoise,
    interp_polyline_by_fixed_waypt_interval
)
from stp3.utils.instance import convert_instance_mask_to_center_and_offset_label
import stp3.utils.sampler as trajectory_sampler
from stp3.utils.spline import Spline2D
from stp3.utils.Optnode import *

def locate_message(utimes, utime):
    i = np.searchsorted(utimes, utime)
    if i == len(utimes) or (i > 0 and utime - utimes[i-1] < utimes[i] - utime):
        i -= 1
    return i

resolution_meters = 0.1
t_fin = 3.5
num = 61
nvar = 11
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
debug = False
brdebug = False
problem = OPTNode(t_fin=t_fin, num=num)

def distance(x1, y1, z1, x2, y2, z2):
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2 + 1 * (z1 - z2) ** 2) ** ( 1 / 2)
 
# Function to calculate K closest points
def kClosest(points, target, K):
    pts = []
    n = len(points)
    d = []

    for i in range(n):
        d.append({
            "first": distance(points[i][0], points[i][1], points[i][2], target[0], target[1], target[2]),
            "second": i
        })
     
    d = sorted(d, key=lambda l:l["first"])
 
    for i in range(K):
        pt = []
        pt.append(points[d[i]["second"]][0])
        pt.append(points[d[i]["second"]][1])
        pt.append(points[d[i]["second"]][2])
        pt.append(points[d[i]["second"]][3])
        pts.append(pt)
 
    return pts

class FuturePredictionDataset(torch.utils.data.Dataset):
    SAMPLE_INTERVAL = 0.5 #SECOND
    def __init__(self, nusc, is_train, cfg):
        self.nusc = nusc
        self.dataroot = self.nusc.dataroot
        self.nusc_exp = NuScenesExplorer(nusc)
        self.nusc_can = NuScenesCanBus(dataroot=self.dataroot)
        self.is_train = is_train
        self.cfg = cfg
        self.problem = OPTNode(t_fin=3.5, num=self.cfg.N_FUTURE_FRAMES + 1)

        if self.is_train == 0:
            self.mode = 'train'
        elif self.is_train == 1:
            self.mode = 'val'
        elif self.is_train == 2:
            self.mode = 'test'
        else:
            raise NotImplementedError

        self.sequence_length = cfg.TIME_RECEPTIVE_FIELD + cfg.N_FUTURE_FRAMES
        self.receptive_field = cfg.TIME_RECEPTIVE_FIELD

        self.scenes = self.get_scenes()
        self.ixes = self.prepro()
        self.indices = self.get_indices()

        # Image resizing and cropping
        self.augmentation_parameters = self.get_resizing_and_cropping_parameters()

        # Normalising input images
        self.normalise_image = torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )

        # Bird's-eye view parameters
        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_resolution, self.bev_start_position, self.bev_dimension = (
            bev_resolution.numpy(), bev_start_position.numpy(), bev_dimension.numpy()
        )

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # The number of sampled trajectories
        self.n_samples = self.cfg.PLANNING.SAMPLE_NUM

        # HD-map feature extractor
        self.nusc_maps = get_nusc_maps(self.cfg.DATASET.MAP_FOLDER)
        self.scene2map = {}
        for sce in self.nusc.scene:
            log = self.nusc.get('log', sce['log_token'])
            self.scene2map[sce['name']] = log['location']
        self.save_dir = cfg.DATASET.SAVE_DIR

        self.curbatchnum1 = None
        self.curbatchnum2 = None
        if self.cfg.PLANNING.DENSE.ENABLED:
            self.centerlines_path = self.cfg.PLANNING.DENSE.PATH

    def get_scenes(self):
        # filter by scene split
        split = {'v1.0-trainval': {0: 'train', 1: 'val', 2: 'test'},
                 'v1.0-mini': {0: 'mini_train', 1: 'mini_val'},}[
            self.nusc.version
        ][self.is_train]

        blacklist = [419] + self.nusc_can.can_blacklist  # # scene-0419 does not have vehicle monitor data
        blacklist = ['scene-' + str(scene_no).zfill(4) for scene_no in blacklist]

        scenes = create_splits_scenes()[split][:]
        for scene_no in blacklist:
            if scene_no in scenes:
                scenes.remove(scene_no)

        return scenes

    def prepro(self):
        samples = [samp for samp in self.nusc.sample]

        # remove samples that aren't in this split
        samples = [samp for samp in samples if self.nusc.get('scene', samp['scene_token'])['name'] in self.scenes]
        
        # sort by scene, timestamp (only to make chronological viz easier)
        samples.sort(key=lambda x: (x['scene_token'], x['timestamp']))

        return samples

    def get_indices(self):
        indices = []
        for index in range(len(self.ixes)):
            is_valid_data = True
            previous_rec = None
            current_indices = []
            for t in range(self.sequence_length):
                index_t = index + t
                # Going over the dataset size limit.
                if index_t >= len(self.ixes):
                    is_valid_data = False
                    break
                rec = self.ixes[index_t]
                # Check if scene is the same
                if (previous_rec is not None) and (rec['scene_token'] != previous_rec['scene_token']):
                    is_valid_data = False
                    break

                current_indices.append(index_t)
                previous_rec = rec

            if is_valid_data:
                indices.append(current_indices)

        return np.asarray(indices)

    def get_resizing_and_cropping_parameters(self):
        original_height, original_width = self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH
        final_height, final_width = self.cfg.IMAGE.FINAL_DIM

        resize_scale = self.cfg.IMAGE.RESIZE_SCALE
        resize_dims = (int(original_width * resize_scale), int(original_height * resize_scale))
        resized_width, resized_height = resize_dims

        crop_h = self.cfg.IMAGE.TOP_CROP
        crop_w = int(max(0, (resized_width - final_width) / 2))
        # Left, top, right, bottom crops.
        crop = (crop_w, crop_h, crop_w + final_width, crop_h + final_height)

        if resized_width != final_width:
            print('Zero padding left and right parts of the image.')
        if crop_h + final_height != resized_height:
            print('Zero padding bottom part of the image.')

        return {'scale_width': resize_scale,
                'scale_height': resize_scale,
                'resize_dims': resize_dims,
                'crop': crop,
                }

    def get_input_data(self, rec):
        """
        Parameters
        ----------
            rec: nuscenes identifier for a given timestamp

        Returns
        -------
            images: torch.Tensor<float> (N, 3, H, W)
            intrinsics: torch.Tensor<float> (3, 3)
            extrinsics: torch.Tensor(N, 4, 4)
        """
        images = []
        unnormalized_images = []
        intrinsics = []
        extrinsics = []
        depths = []
        translations = []
        rotations = []
        camera_intrinsics = []
        point_clouds = []
        point_clouds_labels = []
        point_clouds_panoptic_labels = []
        panoptic_mappings_list = []
        extra_rotations = []
        extra_translations = []
        cameras = self.cfg.IMAGE.NAMES

        # # The extrinsics we want are from the camera sensor to "flat egopose" as defined
        # # https://github.com/nutonomy/nuscenes-devkit/blob/9b492f76df22943daf1dc991358d3d606314af27/python-sdk/nuscenes/nuscenes.py#L279
        # # which corresponds to the position of the lidar.
        # # This is because the labels are generated by projecting the 3D bounding box in this lidar's reference frame.

        # # From lidar egopose to world.
        lidar_sample = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        lidar_pose = self.nusc.get('ego_pose', lidar_sample['ego_pose_token'])
        yaw = Quaternion(lidar_pose['rotation']).yaw_pitch_roll[0]
        lidar_rotation = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        lidar_translation = np.array(lidar_pose['translation'])[:, None]
        lidar_to_world = np.vstack([
            np.hstack((lidar_rotation.rotation_matrix, lidar_translation)),
            np.array([0, 0, 0, 1])
        ])

        camera_sample_ = self.nusc.get('sample_data', rec['data']["CAM_FRONT"])
        my_sample = self.nusc.get('sample', camera_sample_['sample_token'])
        my_scene_token = my_sample['scene_token']

        # # Here we just grab the front camera and the point sensor.
        pointsensor = self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])
        pcl_path = osp.join(self.nusc.dataroot, pointsensor['filename'])
        # panoptic_mappings = [samp.token for samp in self.nusc.get_sample_data(rec['data']['LIDAR_TOP'])[1]]
        # panoptic_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get('panoptic', rec['data']['LIDAR_TOP'])['filename'])
        # panoptic_labels = load_bin_file(panoptic_labels_filename, type='panoptic').astype(np.int32)
        # point_clouds_panoptic_labels.append(panoptic_labels)
        # lidarseg_labels_filename = osp.join(self.nusc.dataroot, self.nusc.get("lidarseg", rec['data']['LIDAR_TOP'])['filename'])
        # points_label = np.fromfile(lidarseg_labels_filename, dtype=np.uint8)
        for cam in cameras:
            camera_sample = self.nusc.get('sample_data', rec['data'][cam])
            sample_token = camera_sample['token']

            camm = self.nusc.get('sample_data', rec['data'][cam])

            # Transformation from world to egopose
            car_egopose = self.nusc.get('ego_pose', camera_sample['ego_pose_token'])
            egopose_rotation = Quaternion(car_egopose['rotation']).inverse
            egopose_translation = -np.array(car_egopose['translation'])[:, None]
            world_to_car_egopose = np.vstack([
                np.hstack((egopose_rotation.rotation_matrix, egopose_rotation.rotation_matrix @ egopose_translation)),
                np.array([0, 0, 0, 1])
            ])

            # From egopose to sensor
            sensor_sample = self.nusc.get('calibrated_sensor', camera_sample['calibrated_sensor_token'])
            intrinsic = torch.Tensor(sensor_sample['camera_intrinsic'])
            sensor_rotation = Quaternion(sensor_sample['rotation'])
            sensor_translation = np.array(sensor_sample['translation'])[:, None]
            car_egopose_to_sensor = np.vstack([
                np.hstack((sensor_rotation.rotation_matrix, sensor_translation)),
                np.array([0, 0, 0, 1])
            ])
            car_egopose_to_sensor = np.linalg.inv(car_egopose_to_sensor)

            # Combine all the transformation.
            # From sensor to lidar.
            lidar_to_sensor = car_egopose_to_sensor @ world_to_car_egopose @ lidar_to_world
            sensor_to_lidar = torch.from_numpy(np.linalg.inv(lidar_to_sensor)).float()

            # Load image
            image_filename = os.path.join(self.dataroot, camera_sample['filename'])
            img = Image.open(image_filename)
            unnormalized_image = torch.tensor(np.array(img))
            # Resize and crop
            img = resize_and_crop_image(
                img, resize_dims=self.augmentation_parameters['resize_dims'], crop=self.augmentation_parameters['crop']
            )
            # Normalise image
            normalised_img = self.normalise_image(img)

            # # Combine resize/cropping in the intrinsics
            top_crop = self.augmentation_parameters['crop'][1]
            left_crop = self.augmentation_parameters['crop'][0]
            intrinsic = update_intrinsics(
                intrinsic, top_crop, left_crop,
                scale_width=self.augmentation_parameters['scale_width'],
                scale_height=self.augmentation_parameters['scale_height']
            )
            
            camm = self.nusc.get('sample_data', rec['data'][cam])
            cs_record = self.nusc.get('calibrated_sensor', camm['calibrated_sensor_token'])
            images.append(normalised_img.unsqueeze(0).unsqueeze(0))
            intrinsics.append(intrinsic.unsqueeze(0).unsqueeze(0))
            extrinsics.append(sensor_to_lidar.unsqueeze(0).unsqueeze(0))
            unnormalized_images.append(unnormalized_image.unsqueeze(0).unsqueeze(0))
            rotations.append(torch.tensor(Quaternion(cs_record['rotation']).rotation_matrix.T).unsqueeze(0).unsqueeze(0))
            translations.append(torch.tensor(-np.array(cs_record['translation'])).unsqueeze(0).unsqueeze(0))
            camera_intrinsics.append(torch.tensor(np.array(cs_record['camera_intrinsic'])).unsqueeze(0).unsqueeze(0))

        images, intrinsics, extrinsics, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels = (torch.cat(images, dim=1),
                                            torch.cat(intrinsics, dim=1),
                                            torch.cat(extrinsics, dim=1),
                                            torch.cat(unnormalized_images, dim=1),
                                            torch.cat(rotations, dim=1),
                                            torch.cat(translations, dim=1),
                                            torch.cat(camera_intrinsics, dim=1),
                                            torch.cat(rotations, dim=1),
                                            torch.cat(rotations, dim=1),
                                            point_clouds_panoptic_labels
                                          )
        if len(depths) > 0:
            depths = torch.cat(depths, dim=1)

        return images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, rotations, rotations, rotations, rotations, rotations

    def _get_top_lidar_pose(self, rec):
        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        trans = -np.array(egopose['translation'])
        yaw = Quaternion(egopose['rotation']).yaw_pitch_roll[0]
        # rot = Quaternion(scalar=np.cos(yaw / 2), vector=[0, 0, np.sin(yaw / 2)])
        rot = Quaternion(egopose['rotation']).inverse
        return trans, rot

    def get_depth_from_lidar(self, lidar_sample, cam_sample):
        points, coloring, im = self.nusc_exp.map_pointcloud_to_image(lidar_sample, cam_sample)
        tmp_cam = np.zeros((self.cfg.IMAGE.ORIGINAL_HEIGHT, self.cfg.IMAGE.ORIGINAL_WIDTH))
        points = points.astype(np.int)
        tmp_cam[points[1, :], points[0,:]] = coloring
        tmp_cam = torch.from_numpy(tmp_cam).unsqueeze(0).unsqueeze(0)
        tmp_cam = F.interpolate(tmp_cam, scale_factor=self.cfg.IMAGE.RESIZE_SCALE, mode='bilinear', align_corners=False, recompute_scale_factor=True)
        tmp_cam = tmp_cam.squeeze()
        crop = self.augmentation_parameters['crop']
        tmp_cam = tmp_cam[crop[1]:crop[3], crop[0]:crop[2]]
        tmp_cam = torch.round(tmp_cam)
        return tmp_cam


    def get_birds_eye_view_label(self, rec, instance_map, in_pred):
        translation, rotation = self._get_top_lidar_pose(rec)
        segmentation = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        pedestrian = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        # Background is ID 0
        instance = np.zeros((self.bev_dimension[0], self.bev_dimension[1]))
        categories_map = []
        bottom_corners = {}

        for annotation_token in rec['anns']:
            # Filter out all non vehicle instances
            annotation = self.nusc.get('sample_annotation', annotation_token)

            if self.cfg.DATASET.FILTER_INVISIBLE_VEHICLES and int(annotation['visibility_token']) == 1 and in_pred is False:
                continue
            if in_pred is True and annotation['instance_token'] not in instance_map:
                continue

            # NuScenes filter
            if 'vehicle' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                instance_id = instance_map[annotation['instance_token']]
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(instance, [poly_region], instance_id)
                cv2.fillPoly(segmentation, [poly_region], 1.0)
                categories_map.append([annotation, np.mean(corners, axis=1)])
                bottom_corners[annotation['token']] = corners
            elif 'human' in annotation['category_name']:
                if annotation['instance_token'] not in instance_map:
                    instance_map[annotation['instance_token']] = len(instance_map) + 1
                poly_region, z, corners = self._get_poly_region_in_image(annotation, translation, rotation)
                cv2.fillPoly(pedestrian, [poly_region], 1.0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def _get_poly_region_in_image(self, instance_annotation, ego_translation, ego_rotation):
        box = Box(
            instance_annotation['translation'], instance_annotation['size'], Quaternion(instance_annotation['rotation'])
        )
        # import pdb; pdb.set_trace()
        box.translate(ego_translation)
        box.rotate(ego_rotation)
        corners = box.corners()

        pts = box.bottom_corners()[:2].T
        pts = np.round((pts - self.bev_start_position[:2] + self.bev_resolution[:2] / 2.0) / self.bev_resolution[:2]).astype(np.int32)
        pts[:, [1, 0]] = pts[:, [0, 1]]

        z = box.bottom_corners()[2, 0]
        return pts, z, corners

    def get_label(self, rec, instance_map, in_pred):
        segmentation_np, instance_np, pedestrian_np, instance_map, categories_map, bottom_corners = \
            self.get_birds_eye_view_label(rec, instance_map, in_pred)
        segmentation = torch.from_numpy(segmentation_np).long().unsqueeze(0).unsqueeze(0)
        instance = torch.from_numpy(instance_np).long().unsqueeze(0)
        pedestrian = torch.from_numpy(pedestrian_np).long().unsqueeze(0).unsqueeze(0)

        return segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners

    def get_future_egomotion(self, rec, index):
        rec_t0 = rec

        # Identity
        future_egomotion = np.eye(4, dtype=np.float32)

        if index < len(self.ixes) - 1:
            rec_t1 = self.ixes[index + 1]

            if rec_t0['scene_token'] == rec_t1['scene_token']:
                egopose_t0 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t0['data']['LIDAR_TOP'])['ego_pose_token']
                )
                egopose_t1 = self.nusc.get(
                    'ego_pose', self.nusc.get('sample_data', rec_t1['data']['LIDAR_TOP'])['ego_pose_token']
                )

                egopose_t0 = convert_egopose_to_matrix_numpy(egopose_t0)
                egopose_t1 = convert_egopose_to_matrix_numpy(egopose_t1)

                future_egomotion = invert_matrix_egopose_numpy(egopose_t1).dot(egopose_t0)
                future_egomotion[3, :3] = 0.0
                future_egomotion[3, 3] = 1.0

        future_egomotion = torch.Tensor(future_egomotion).float()

        # Convert to 6DoF vector
        future_egomotion = mat2pose_vec(future_egomotion)
        return future_egomotion.unsqueeze(0)

    def get_trajectory_sampling(self, rec=None, sample_indice=None):
        if rec is None and sample_indice is None:
            raise ValueError("No valid input rec or token")
        if rec is None and sample_indice is not None:
            rec = self.ixes[sample_indice]

        ref_scene = self.nusc.get("scene", rec['scene_token'])

        # vm_msgs = self.nusc_can.get_messages(ref_scene['name'], 'vehicle_monitor')
        # vm_uts = [msg['utime'] for msg in vm_msgs]
        pose_msgs = self.nusc_can.get_messages(ref_scene['name'],'pose')
        pose_uts = [msg['utime'] for msg in pose_msgs]
        steer_msgs = self.nusc_can.get_messages(ref_scene['name'], 'steeranglefeedback')
        steer_uts = [msg['utime'] for msg in steer_msgs]

        ref_utime = rec['timestamp']
        # vm_index = locate_message(vm_uts, ref_utime)
        # vm_data = vm_msgs[vm_index]
        pose_index = locate_message(pose_uts, ref_utime)
        pose_data = pose_msgs[pose_index]
        steer_index = locate_message(steer_uts, ref_utime)
        steer_data = steer_msgs[steer_index]

        # initial speed
        # v0 = vm_data["vehicle_speed"] / 3.6  # km/h to m/s
        v0 = pose_data["vel"][0]  # [0] means longitudinal velocity  m/s

        # curvature (positive: turn left)
        # steering = np.deg2rad(vm_data["steering"])
        steering = steer_data["value"]

        location = self.scene2map[ref_scene['name']]
        # flip x axis if in left-hand traffic (singapore)
        flip_flag = True if location.startswith('singapore') else False
        if flip_flag:
            steering *= -1
        Kappa = 2 * steering / 2.588

        # initial state
        T0 = np.array([0.0, 1.0])  # define front
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])  # define side

        t_start = 0  # second
        t_end = self.cfg.N_FUTURE_FRAMES * self.SAMPLE_INTERVAL  # second
        t_interval = self.SAMPLE_INTERVAL / 10
        tt = np.arange(t_start, t_end + t_interval, t_interval)
        sampled_trajectories_fine = trajectory_sampler.sample(v0, Kappa, T0, N0, tt, self.n_samples)
        sampled_trajectories = sampled_trajectories_fine[:, ::10]
        return sampled_trajectories


    def generate_centerlines(self, rec, map_name):


        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        x, y = egopose["translation"][0], egopose["translation"][1]
        nusc_map = self.nusc_maps[map_name]
        radius_init = 2
        closest_lane = None
        while closest_lane == None or len(closest_lane) < 5:
            # keep increasing search space until lane is found
            closest_lane = nusc_map.get_closest_lane(x, y, radius=radius_init)
            radius_init *= 2
        candidates_future = dfs(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        candidates_past = dfs_incoming(nusc_map, closest_lane, dist=0, threshold=self.cfg.DATASET.THRESHOLD, resolution_meters=resolution_meters)
        lane_ids = []
        for past_lane_seq in candidates_past:
            for future_lane_seq in candidates_future:
                lane_ids.append(past_lane_seq + future_lane_seq[1:])
        lane_ids = remove_overlapping_lane_seq(lane_ids)
        centerlines = []
        for lane_id in lane_ids:
            centerline = []
            for lane in lane_id:
                lane_record = nusc_map.get_arcline_path(lane)
                cx = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                centerline.extend(cx)
            centerlines.append(np.array(centerline))
        lanes = self.nusc_maps[map_name].get_records_in_radius(x, y, 2000, ['lane', 'lane_connector'])
        lanes = lanes['lane'] + lanes['lane_connector']
        discrete_points = self.nusc_maps[map_name].discretize_lanes(lanes, resolution_meters=resolution_meters)
        lanes = []
        for lane_id, points in discrete_points.items():
            lanes.append(np.array(points)[:, :2])
     
        patch_box = box_coords
        patch_angle = rot * 180 / np.pi

        patch_x = patch_box[0]
        patch_y = patch_box[1]

        patch = get_patch_coord(patch_box, patch_angle)

        line_list = []
        for lane in centerlines:
            line = LineString([(arr[0], arr[1]) for arr in lane[:, :2]])
            c = lane[:, :2]
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.rotate(new_line, -patch_angle, origin=(patch_x, patch_y), use_radians=False)
                new_line = affinity.affine_transform(new_line,
                                                        [1.0, 0.0, 0.0, 1.0, -patch_x, -patch_y])
                line_list.append(new_line)
        
        local_box = (0.0, 0.0, patch_box[2], patch_box[3])
        patch_x, patch_y, patch_h, patch_w = local_box

        patch = get_patch_coord(local_box)

        canvas_h = canvas_size[0]
        canvas_w = canvas_size[1]
        scale_height = canvas_h/patch_h
        scale_width = canvas_w/patch_w

        trans_x = -patch_x + patch_w / 2.0
        trans_y = -patch_y + patch_h / 2.0

        overall_lines = []
        for line in line_list:
            if not line.intersection(patch).is_empty:
                new_line = line
                new_line = affinity.affine_transform(new_line,
                                                     [1.0, 0.0, 0.0, 1.0, trans_x, trans_y])
                new_line = affinity.scale(new_line, xfact=scale_width, yfact=scale_height, origin=(0, 0))
                # overall_lines.append([list(new_line.coords)])

                gg = np.array(list(new_line.coords))
                # interp_arc(t: int, px: np.ndarray, py: np.ndarray) -> np.ndarray:
                gg = interp_arc(t=2000, px=gg[:,0], py=gg[:,1])
                overall_lines.append(gg)

        return overall_lines

    def voxelize_hd_map(self, rec):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        stretch = [self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1]]
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        egopose = self.nusc.get('ego_pose', self.nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
        map_name = self.scene2map[self.nusc.get('scene', rec['scene_token'])['name']]

        rot = Quaternion(egopose['rotation']).rotation_matrix
        rot = np.arctan2(rot[1,0], rot[0,0]) # in radian
        center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

        box_coords = (
            center[0],
            center[1],
            stretch[0]*2,
            stretch[1]*2
        ) # (x_center, y_center, width, height)
        canvas_size = (
                int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
                int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        elements = self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS
        hd_features = self.nusc_maps[map_name].get_map_mask(box_coords, rot * 180 / np.pi , elements, canvas_size=canvas_size)
        #traffic = self.hd_traffic_light(map_name, center, stretch, dx, bx, canvas_size)
        #return torch.from_numpy(np.concatenate((hd_features, traffic), axis=0)[None]).float()
        hd_features = torch.from_numpy(hd_features[None]).float()
        hd_features = torch.transpose(hd_features,-2,-1) # (y,x) replace horizontal and vertical coordinates
        return hd_features

    def hd_traffic_light(self, map_name, center, stretch, dx, bx, canvas_size):

        roads = np.zeros(canvas_size)
        my_patch = (
            center[0] - stretch[0],
            center[1] - stretch[1],
            center[0] + stretch[0],
            center[1] + stretch[1],
        )
        tl_token = self.nusc_maps[map_name].get_records_in_patch(my_patch, ['traffic_light'], mode='intersect')['traffic_light']
        polys = []
        for token in tl_token:
            road_token =self.nusc_maps[map_name].get('traffic_light', token)['from_road_block_token']
            pt = self.nusc_maps[map_name].get('road_block', road_token)['polygon_token']
            polygon = self.nusc_maps[map_name].extract_polygon(pt)
            polys.append(np.array(polygon.exterior.xy).T)

        def get_rot(h):
            return torch.Tensor([
                [np.cos(h), np.sin(h)],
                [-np.sin(h), np.cos(h)],
            ])
        # convert to local coordinates in place
        rot = get_rot(np.arctan2(center[3], center[2])).T
        for rowi in range(len(polys)):
            polys[rowi] -= center[:2]
            polys[rowi] = np.dot(polys[rowi], rot)

        for la in polys:
            pts = (la - bx) / dx
            pts = np.int32(np.around(pts))
            cv2.fillPoly(roads, [pts], 1)

        return roads[None]

    def get_gt_trajectory(self, rec, ref_index):
        n_output = self.cfg.N_FUTURE_FRAMES
        n_past = self.cfg.TIME_RECEPTIVE_FIELD
        gt_trajectory = np.zeros((n_output+1+(n_past-1), 3), np.float64)

        """
            egopos_cur -> gives you sensor to global.
                Suppose a point P in sensor-frame.
                    egopos_cur.P gives you the point's global coordinates
        """
        egopose_cur = get_global_pose(rec, self.nusc, inverse=True)

        for i in range(n_past - 1):
            index = ref_index - (n_past - 1) + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i, :] = [origin[0], origin[1], theta]

        for i in range(n_output+1):
            index = ref_index + i
            if index < len(self.ixes):
                rec_future = self.ixes[index]
                """
                    egopos_future -> gives you global to future.
                        Suppose a point P in ego-frame.
                """
                egopose_future = get_global_pose(rec_future, self.nusc, inverse=False)
                egopose_future = egopose_cur.dot(egopose_future)

                theta = quaternion_yaw(Quaternion(matrix=egopose_future))

                origin = np.array(egopose_future[:3, 3])

                gt_trajectory[i + (n_past - 1), :] = [origin[0], origin[1], theta]

        """
            ?
        """
        if gt_trajectory[-1][0] >= 2:
            command = 'RIGHT'
        elif gt_trajectory[-1][0] <= -2:
            command = 'LEFT'
        else:
            command = 'FORWARD'

        return gt_trajectory, command

    def get_routed_map(self, gt_points):
        dx, bx, _ = gen_dx_bx(self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2].numpy(), bx[:2].numpy()

        canvas_size = (
            int(self.cfg.LIFT.X_BOUND[1] * 2 / self.cfg.LIFT.X_BOUND[2]),
            int(self.cfg.LIFT.Y_BOUND[1] * 2 / self.cfg.LIFT.Y_BOUND[2])
        )

        roads = np.zeros(canvas_size)
        W = 1.85
        pts = np.array([
            [-4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, W / 2.],
            [4.084 / 2. + 0.5, -W / 2.],
            [-4.084 / 2. + 0.5, -W / 2.],
        ])
        pts = (pts - bx) / dx
        pts[:, [0, 1]] = pts[:, [1, 0]]

        pts = np.int32(np.around(pts))
        cv2.fillPoly(roads, [pts], 1)

        gt_points = gt_points[:-1].numpy()
        # 坐标原点在左上角
        target = pts.copy()
        target[:,0] = pts[:,0] + gt_points[0] / dx[0]
        target[:,1] = pts[:,1] - gt_points[1] / dx[1]
        target = np.int32(np.around(target))
        cv2.fillPoly(roads, [target], 1)
        return roads

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

    def __getitem__(self, index):
        """
        Returns
        -------
            data: dict with the following keys:
                image: torch.Tensor<float> (T, N, 3, H, W)
                    normalised cameras images with T the sequence length, and N the number of cameras.
                intrinsics: torch.Tensor<float> (T, N, 3, 3)
                    intrinsics containing resizing and cropping parameters.
                extrinsics: torch.Tensor<float> (T, N, 4, 4)
                    6 DoF pose from world coordinates to camera coordinates.
                segmentation: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                    (H_bev, W_bev) are the pixel dimensions in bird's-eye view.
                instance: torch.Tensor<int64> (T, 1, H_bev, W_bev)
                centerness: torch.Tensor<float> (T, 1, H_bev, W_bev)
                offset: torch.Tensor<float> (T, 2, H_bev, W_bev)
                flow: torch.Tensor<float> (T, 2, H_bev, W_bev)
                future_egomotion: torch.Tensor<float> (T, 6)
                    6 DoF egomotion t -> t+1

        """
        data = {}
        keys = ['image', 'intrinsics', 'extrinsics', 'depths',
                'segmentation', 'instance', 'centerness', 'offset', 'flow', 'pedestrian',
                'future_egomotion', 'hdmap', 'gt_trajectory', 'indices', 'poses', 'unnormalized_images',\
                'rotations', 'translations', 'camera_intrinsics', 'point_clouds', 'point_clouds_labels', \
                'point_clouds_panoptic_labels', 'scene_token', 'categories_map', 'panoptic_mappings_list', \
                'bottom_corners'
                ]
        for key in keys:
            data[key] = []

        instance_map = {}
        # Loop over all the frames in the sequence.
        total_centerlines = []
        for i, index_t in enumerate(self.indices[index]):
            if i >= self.receptive_field:
                in_pred = True
            else:
                in_pred = False
            rec = self.ixes[index_t]

            if i < self.receptive_field:
                images, intrinsics, extrinsics, depths, unnormalized_images, rotations, translations, camera_intrinsics, point_clouds, point_clouds_labels, point_clouds_panoptic_labels, scene_token, panoptic_mappings_list = self.get_input_data(rec)
                data['image'].append(images)
                data['unnormalized_images'].append(unnormalized_images)
                data['intrinsics'].append(intrinsics)
                data['extrinsics'].append(extrinsics)
                data['translations'].append(translations)
                data['rotations'].append(rotations)
                data['camera_intrinsics'].append(camera_intrinsics)

            segmentation, instance, pedestrian, instance_map, categories_map, bottom_corners = self.get_label(rec, instance_map, in_pred)

            future_egomotion = self.get_future_egomotion(rec, index_t)
            hd_map_feature = self.voxelize_hd_map(rec)

            data['segmentation'].append(segmentation)
            data['future_egomotion'].append(future_egomotion)
            data['hdmap'].append(hd_map_feature)
        for key, value in data.items():
            if key in ['image', 'intrinsics', 'extrinsics', 'depths', 'segmentation','future_egomotion', 'hdmap']:
                if key == 'depths' and self.cfg.LIFT.GT_DEPTH is False:
                    continue
                data[key] = torch.cat(value, dim=0)

        return data
data/stp3/datas/trainNuscenesData.py
llmbev/talk2bev
from collections import OrderedDict
from functools import partial

import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import DropPath


class ConvBlock(nn.Module):
    """2D convolution followed by
         - an optional normalisation (batch norm or instance norm)
         - an optional activation (ReLU, LeakyReLU, or tanh)
    """

    def __init__(
        self,
        in_channels,
        out_channels=None,
        kernel_size=3,
        stride=1,
        norm='bn',
        activation='relu',
        bias=False,
        transpose=False,
    ):
        super().__init__()
        out_channels = out_channels or in_channels
        padding = int((kernel_size - 1) / 2)
        self.conv = nn.Conv2d if not transpose else partial(nn.ConvTranspose2d, output_padding=1)
        self.conv = self.conv(in_channels, out_channels, kernel_size, stride, padding=padding, bias=bias)

        if norm == 'bn':
            self.norm = nn.BatchNorm2d(out_channels)
        elif norm == 'in':
            self.norm = nn.InstanceNorm2d(out_channels)
        elif norm == 'none':
            self.norm = None
        else:
            raise ValueError('Invalid norm {}'.format(norm))

        if activation == 'relu':
            self.activation = nn.ReLU(inplace=True)
        elif activation == 'lrelu':
            self.activation = nn.LeakyReLU(0.1, inplace=True)
        elif activation == 'elu':
            self.activation = nn.ELU(inplace=True)
        elif activation == 'tanh':
            self.activation = nn.Tanh()
        elif activation == 'none':
            self.activation = None
        else:
            raise ValueError('Invalid activation {}'.format(activation))

    def forward(self, x):
        x = self.conv(x)

        if self.norm:
            x = self.norm(x)
        if self.activation:
            x = self.activation(x)
        return x


class Bottleneck(nn.Module):
    """
    Defines a bottleneck module with a residual connection
    """

    def __init__(
        self,
        in_channels,
        out_channels=None,
        kernel_size=3,
        dilation=1,
        groups=1,
        upsample=False,
        downsample=False,
        dropout=0.0,
    ):
        super().__init__()
        self._downsample = downsample
        bottleneck_channels = int(in_channels / 2)
        out_channels = out_channels or in_channels
        padding_size = ((kernel_size - 1) * dilation + 1) // 2

        # Define the main conv operation
        assert dilation == 1
        if upsample:
            assert not downsample, 'downsample and upsample not possible simultaneously.'
            bottleneck_conv = nn.ConvTranspose2d(
                bottleneck_channels,
                bottleneck_channels,
                kernel_size=kernel_size,
                bias=False,
                dilation=1,
                stride=2,
                output_padding=padding_size,
                padding=padding_size,
                groups=groups,
            )
        elif downsample:
            bottleneck_conv = nn.Conv2d(
                bottleneck_channels,
                bottleneck_channels,
                kernel_size=kernel_size,
                bias=False,
                dilation=dilation,
                stride=2,
                padding=padding_size,
                groups=groups,
            )
        else:
            bottleneck_conv = nn.Conv2d(
                bottleneck_channels,
                bottleneck_channels,
                kernel_size=kernel_size,
                bias=False,
                dilation=dilation,
                padding=padding_size,
                groups=groups,
            )

        self.layers = nn.Sequential(
            OrderedDict(
                [
                    # First projection with 1x1 kernel
                    ('conv_down_project', nn.Conv2d(in_channels, bottleneck_channels, kernel_size=1, bias=False)),
                    ('abn_down_project', nn.Sequential(nn.BatchNorm2d(bottleneck_channels),
                                                       nn.ReLU(inplace=True))),
                    # Second conv block
                    ('conv', bottleneck_conv),
                    ('abn', nn.Sequential(nn.BatchNorm2d(bottleneck_channels), nn.ReLU(inplace=True))),
                    # Final projection with 1x1 kernel
                    ('conv_up_project', nn.Conv2d(bottleneck_channels, out_channels, kernel_size=1, bias=False)),
                    ('abn_up_project', nn.Sequential(nn.BatchNorm2d(out_channels),
                                                     nn.ReLU(inplace=True))),
                    # Regulariser
                    ('dropout', nn.Dropout2d(p=dropout)),
                ]
            )
        )

        if out_channels == in_channels and not downsample and not upsample:
            self.projection = None
        else:
            projection = OrderedDict()
            if upsample:
                projection.update({'upsample_skip_proj': Interpolate(scale_factor=2)})
            elif downsample:
                projection.update({'upsample_skip_proj': nn.MaxPool2d(kernel_size=2, stride=2)})
            projection.update(
                {
                    'conv_skip_proj': nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
                    'bn_skip_proj': nn.BatchNorm2d(out_channels),
                }
            )
            self.projection = nn.Sequential(projection)

    # pylint: disable=arguments-differ
    def forward(self, *args):
        (x,) = args
        x_residual = self.layers(x)
        if self.projection is not None:
            if self._downsample:
                # pad h/w dimensions if they are odd to prevent shape mismatch with residual layer
                x = nn.functional.pad(x, (0, x.shape[-1] % 2, 0, x.shape[-2] % 2), value=0)
            return x_residual + self.projection(x)
        return x_residual + x


class Interpolate(nn.Module):
    def __init__(self, scale_factor: int = 2):
        super().__init__()
        self._interpolate = nn.functional.interpolate
        self._scale_factor = scale_factor

    # pylint: disable=arguments-differ
    def forward(self, x):
        return self._interpolate(x, scale_factor=self._scale_factor, mode='bilinear', align_corners=False)


class UpsamplingConcat(nn.Module):
    def __init__(self, in_channels, out_channels, scale_factor=2):
        super().__init__()

        self.upsample = nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=False)

        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, x_to_upsample, x):
        x_to_upsample = self.upsample(x_to_upsample)
        x_to_upsample = torch.cat([x, x_to_upsample], dim=1)
        return self.conv(x_to_upsample)


class UpsamplingAdd(nn.Module):
    def __init__(self, in_channels, out_channels, scale_factor=2):
        super().__init__()
        self.upsample_layer = nn.Sequential(
            nn.Upsample(scale_factor=scale_factor, mode='bilinear', align_corners=False),
            nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
        )

    def forward(self, x, x_skip):
        x = self.upsample_layer(x)
        return x + x_skip

class ASPPConv(nn.Sequential):
    def __init__(self, in_channels, out_channels, dilation):
        modules = [
            nn.Conv2d(in_channels, out_channels, 3, padding=dilation, dilation=dilation, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        ]
        super(ASPPConv, self).__init__(*modules)


class ASPPPooling(nn.Sequential):
    def __init__(self, in_channels, out_channels):
        super(ASPPPooling, self).__init__(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU())

    def forward(self, x):
        size = x.shape[-2:]
        for mod in self:
            x = mod(x)
        return F.interpolate(x, size=size, mode='bilinear', align_corners=False)


class ASPP(nn.Module):
    def __init__(self, in_channels, atrous_rates, out_channels=256):
        super(ASPP, self).__init__()
        modules = []
        modules.append(nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()))

        rates = tuple(atrous_rates)
        for rate in rates:
            modules.append(ASPPConv(in_channels, out_channels, rate))

        modules.append(ASPPPooling(in_channels, out_channels))

        self.convs = nn.ModuleList(modules)

        self.project = nn.Sequential(
            nn.Conv2d(len(self.convs) * out_channels, out_channels, 1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(),
            nn.Dropout(0.5))

    def forward(self, x):
        res = []
        for conv in self.convs:
            res.append(conv(x))
        res = torch.cat(res, dim=1)
        return self.project(res)

class DeepLabHead(nn.Sequential):
    def __init__(self, in_channels, num_classes, hidden_channel=256):
        super(DeepLabHead, self).__init__(
            ASPP(in_channels, [12, 24, 36], hidden_channel),
            nn.Conv2d(hidden_channel, hidden_channel, 3, padding=1, bias=False),
            nn.BatchNorm2d(hidden_channel),
            nn.ReLU(),
            nn.Conv2d(hidden_channel, num_classes, 1)
        )


class LayerNorm(nn.Module):
    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
    with shape (batch_size, channels, height, width).
    """

    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError
        self.normalized_shape = (normalized_shape,)

    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x

class Block(nn.Module):
    r""" ConvNeXt Block. There are two equivalent implementations:
    (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
    (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
    We use (2) as we find it slightly faster in PyTorch

    Args:
        dim (int): Number of input channels.
        drop_path (float): Stochastic depth rate. Default: 0.0
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
    """

    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x

class Bottleblock(nn.Module):
    def __init__(self, in_channels, out_channels=None):
        super(Bottleblock, self).__init__()

        bottleneck_channels = int(in_channels / 2)
        out_channels = out_channels or in_channels

        self.layers = nn.Sequential(
            nn.Conv2d(in_channels, bottleneck_channels, kernel_size=7, bias=False, padding=3),
            LayerNorm(bottleneck_channels, eps=1e-6, data_format='channels_first'),
            nn.GELU(),
            nn.Conv2d(bottleneck_channels, bottleneck_channels, kernel_size=1, bias=False),
            LayerNorm(bottleneck_channels, eps=1e-6, data_format='channels_first'),
            nn.GELU(),
            nn.Conv2d(bottleneck_channels, out_channels, kernel_size=3, bias=False, padding=1),
            LayerNorm(out_channels, eps=1e-6, data_format='channels_first'),
            nn.GELU()
        )

        if out_channels == in_channels:
            self.projection = None
        else:
            self.projection = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
                nn.GELU()
            )

    def forward(self, *args):
        (x,) = args
        x_residual = self.layers(x)
        if self.projection is not None:
            return x_residual + self.projection(x)
        return x_residual + x
data/stp3/layers/convolutions.py
llmbev/talk2bev
from collections import OrderedDict

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

from stp3.layers.convolutions import Bottleneck, Block, Bottleblock
from stp3.utils.geometry import warp_features


class SpatialGRU(nn.Module):
    """A GRU cell that takes an input tensor [BxTxCxHxW] and an optional previous state and passes a
    convolutional gated recurrent unit over the data"""

    def __init__(self, input_size, hidden_size, gru_bias_init=0.0):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.gru_bias_init = gru_bias_init

        self.conv_update = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_reset = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_state_tilde = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_decoder = nn.Conv2d(hidden_size, input_size, kernel_size=1, bias=False)

    def forward(self, x, state=None):
        # pylint: disable=unused-argument, arguments-differ
        # Check size
        assert len(x.size()) == 5, 'Input tensor must be BxTxCxHxW.'

        # recurrent layers
        rnn_output = []
        b, timesteps, c, h, w = x.size()
        rnn_state = torch.zeros(b, self.hidden_size, h, w, device=x.device) if state is None else state
        for t in range(timesteps):
            x_t = x[:, t]

            rnn_state = self.gru_cell(x_t, rnn_state)
            rnn_output.append(self.conv_decoder(rnn_state))

        # reshape rnn output to batch tensor
        return torch.stack(rnn_output, dim=1)

    def gru_cell(self, x, state):
        # Compute gates
        x_and_state = torch.cat([x, state], dim=1)
        update_gate = self.conv_update(x_and_state)
        reset_gate = self.conv_reset(x_and_state)
        # Add bias to initialise gate as close to identity function
        update_gate = torch.sigmoid(update_gate + self.gru_bias_init)
        reset_gate = torch.sigmoid(reset_gate + self.gru_bias_init)

        # Compute proposal state, activation is defined in norm_act_config (can be tanh, ReLU etc)
        state_tilde = self.conv_state_tilde(torch.cat([x, (1.0 - reset_gate) * state], dim=1))

        output = (1.0 - update_gate) * state + update_gate * state_tilde
        return output

class Dual_GRU(nn.Module):
    def __init__(self, in_channels, latent_dim, n_future, mixture=True, gru_bias_init = 0.0):
        super(Dual_GRU, self).__init__()

        input_size = in_channels
        hidden_size = latent_dim

        self.n_future = n_future
        self.mixture = mixture

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.gru_bias_init = gru_bias_init

        self.conv_update_1 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_reset_1 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_state_tilde_1 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)

        self.conv_update_2 = nn.Conv2d(hidden_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_reset_2 = nn.Conv2d(hidden_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_state_tilde_2 = nn.Conv2d(hidden_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_decoder_2 = nn.Conv2d(hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)

        self.trusting_gate = nn.Sequential(
            Bottleblock(hidden_size+hidden_size, hidden_size),
            nn.Conv2d(hidden_size, 2, kernel_size=1, bias=False)
        )

    def forward(self, x, state):
        '''
        x: torch.Tensor [b, 1, input_size, h, w]
        state: torch.Tensor [b, n_present, hidden_size, h, w]
        '''
        b, s, c, h, w = x.shape
        assert c == self.input_size, f'feature sizes must match, got input {c} for layer with size {self.input_size}'
        n_present = state.shape[1]

        h = state[:, 0]
        pred_state = []

        # warmup
        for t in range(n_present - 1):
            cur_state = state[:, t]
            h = self.gru_cell_2(cur_state, h)

        # recurrent layers
        rnn_state1 = state[:, -1]
        rnn_state2 = state[:, -1]
        x = x[:, 0]

        for t in range(self.n_future):
            # propagate gru v1
            rnn_state1 = self.gru_cell_1(x, rnn_state1)
            # propagate gru v2
            h = self.gru_cell_2(rnn_state2, h)
            rnn_state2 = self.conv_decoder_2(h)

            # mix the two distribution
            mix_state = torch.cat([rnn_state1, rnn_state2], dim=1)
            trust_gate = self.trusting_gate(mix_state)
            trust_gate = torch.softmax(trust_gate, dim=1)
            cur_state = rnn_state2 * trust_gate[:,0:1] + rnn_state1 * trust_gate[:,1:]

            pred_state.append(cur_state)

            if self.mixture:
                rnn_state1 = cur_state
                rnn_state2 = cur_state

        x = torch.stack(pred_state, dim=1)

        return x

    def gru_cell_1(self, x, state):
        # Compute gates
        x_and_state = torch.cat([x, state], dim=1)
        update_gate = self.conv_update_1(x_and_state)
        reset_gate = self.conv_reset_1(x_and_state)
        # Add bias to initialise gate as close to identity function
        update_gate = torch.sigmoid(update_gate + self.gru_bias_init)
        reset_gate = torch.sigmoid(reset_gate + self.gru_bias_init)

        # Compute proposal state, activation is defined in norm_act_config (can be tanh, ReLU etc)
        state_tilde = self.conv_state_tilde_1(torch.cat([x, (1.0 - reset_gate) * state], dim=1))

        output = (1.0 - update_gate) * state + update_gate * state_tilde
        return output

    def gru_cell_2(self, x, state):
        # Compute gates
        x_and_state = torch.cat([x, state], dim=1)
        update_gate = self.conv_update_2(x_and_state)
        reset_gate = self.conv_reset_2(x_and_state)
        # Add bias to initialise gate as close to identity function
        update_gate = torch.sigmoid(update_gate + self.gru_bias_init)
        reset_gate = torch.sigmoid(reset_gate + self.gru_bias_init)

        # Compute proposal state, activation is defined in norm_act_config (can be tanh, ReLU etc)
        state_tilde = self.conv_state_tilde_2(torch.cat([x, (1.0 - reset_gate) * state], dim=1))

        output = (1.0 - update_gate) * state + update_gate * state_tilde
        return output

class BiGRU(nn.Module):
    def __init__(self, in_channels, gru_bias_init = 0.0):
        super(BiGRU, self).__init__()
        input_size = in_channels
        hidden_size = in_channels

        self.input_size = in_channels
        self.hidden_size = in_channels
        self.gru_bias_init = gru_bias_init

        self.conv_update_1 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_reset_1 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_state_tilde_1 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_decoder_1 = Bottleblock(hidden_size, input_size)

        self.conv_update_2 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_reset_2 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_state_tilde_2 = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size=3, bias=True, padding=1)
        self.conv_decoder_2 = Bottleblock(hidden_size, input_size)

        self.res_blocks = nn.Sequential(
            Bottleblock(in_channels + in_channels, in_channels),
            Block(in_channels, in_channels),
            Block(in_channels, in_channels)
        )


    def forward(self, x):
        '''
        x: torch.Tensor [b, s, input_size, h, w]
        '''
        b, s, c, h, w = x.shape

        rnn_state1 = x[:, 0]
        rnn_state2 = x[:, -1]

        states_1 = []
        states_2 = []

        for t in range(s):
            # propagate gru v1
            rnn_state1 = self.gru_cell_1(x[:, t], rnn_state1)
            # propagate gru v2
            rnn_state2 = self.gru_cell_2(x[:, s-t-1], rnn_state2)

            states_1.append(self.conv_decoder_1(rnn_state1))
            states_2.append(self.conv_decoder_2(rnn_state2))

        states_1 = torch.stack(states_1, dim=1)
        states_2 = torch.stack(states_2[::-1], dim=1)

        states = torch.cat([states_1, states_2], dim=2)

        b, s, c, h, w = states.shape
        x = self.res_blocks(states.view(b * s, c, h, w))
        x = x.view(b, s, *x.shape[1:])
        # reshape rnn output to batch tensor
        return x

    def gru_cell_1(self, x, state):
        # Compute gates
        x_and_state = torch.cat([x, state], dim=1)
        update_gate = self.conv_update_1(x_and_state)
        reset_gate = self.conv_reset_1(x_and_state)
        # Add bias to initialise gate as close to identity function
        update_gate = torch.sigmoid(update_gate + self.gru_bias_init)
        reset_gate = torch.sigmoid(reset_gate + self.gru_bias_init)

        # Compute proposal state, activation is defined in norm_act_config (can be tanh, ReLU etc)
        state_tilde = self.conv_state_tilde_1(torch.cat([x, (1.0 - reset_gate) * state], dim=1))

        output = (1.0 - update_gate) * state + update_gate * state_tilde
        return output

    def gru_cell_2(self, x, state):
        # Compute gates
        x_and_state = torch.cat([x, state], dim=1)
        update_gate = self.conv_update_2(x_and_state)
        reset_gate = self.conv_reset_2(x_and_state)
        # Add bias to initialise gate as close to identity function
        update_gate = torch.sigmoid(update_gate + self.gru_bias_init)
        reset_gate = torch.sigmoid(reset_gate + self.gru_bias_init)

        # Compute proposal state, activation is defined in norm_act_config (can be tanh, ReLU etc)
        state_tilde = self.conv_state_tilde_2(torch.cat([x, (1.0 - reset_gate) * state], dim=1))

        output = (1.0 - update_gate) * state + update_gate * state_tilde
        return output


class CausalConv3d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=(2, 3, 3), dilation=(1, 1, 1), bias=False):
        super().__init__()
        assert len(kernel_size) == 3, 'kernel_size must be a 3-tuple.'
        time_pad = (kernel_size[0] - 1) * dilation[0]
        height_pad = ((kernel_size[1] - 1) * dilation[1]) // 2
        width_pad = ((kernel_size[2] - 1) * dilation[2]) // 2

        # Pad temporally on the left
        # padding (padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back)
        self.pad = nn.ConstantPad3d(padding=(width_pad, width_pad, height_pad, height_pad, time_pad, 0), value=0)
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, dilation=dilation, stride=1, padding=0, bias=bias)
        self.norm = nn.BatchNorm3d(out_channels)
        self.activation = nn.ReLU(inplace=True)

    def forward(self, *inputs):
        (x,) = inputs
        x = self.pad(x)
        x = self.conv(x)
        x = self.norm(x)
        x = self.activation(x)
        return x

class TemporalConv3d(nn.Module):
    def __init__(self, in_channels, out_channels, n_present, n_future, kernel_size=(2, 3, 3), dilation=(1, 1, 1), bias=False):
        super(TemporalConv3d, self).__init__()
        assert len(kernel_size) == 3, 'kernel_size must be a 3-tuple.'

        time_pad = n_future - n_present + dilation[0] * (kernel_size[0] - 1)
        height_pad = ((kernel_size[1] - 1) * dilation[1]) // 2
        width_pad = ((kernel_size[2] - 1) * dilation[2]) // 2

        self.pad = nn.ConstantPad3d(padding=(width_pad, width_pad, height_pad, height_pad, time_pad // 2, time_pad - time_pad // 2), value=0)
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, dilation=dilation, stride=1, padding=0, bias=bias)
        self.norm = nn.BatchNorm3d(out_channels)
        self.activation = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.pad(x)
        x = self.conv(x)
        x = self.norm(x)
        x = self.activation(x)
        return x

class CausalMaxPool3d(nn.Module):
    def __init__(self, kernel_size=(2, 3, 3)):
        super().__init__()
        assert len(kernel_size) == 3, 'kernel_size must be a 3-tuple.'
        time_pad = kernel_size[0] - 1
        height_pad = (kernel_size[1] - 1) // 2
        width_pad = (kernel_size[2] - 1) // 2

        # Pad temporally on the left
        self.pad = nn.ConstantPad3d(padding=(width_pad, width_pad, height_pad, height_pad, time_pad, 0), value=0)
        self.max_pool = nn.MaxPool3d(kernel_size, stride=1)

    def forward(self, *inputs):
        (x,) = inputs
        x = self.pad(x)
        x = self.max_pool(x)
        return x


def conv_1x1x1_norm_activated(in_channels, out_channels):
    """1x1x1 3D convolution, normalization and activation layer."""
    return nn.Sequential(
        OrderedDict(
            [
                ('conv', nn.Conv3d(in_channels, out_channels, kernel_size=1, bias=False)),
                ('norm', nn.BatchNorm3d(out_channels)),
                ('activation', nn.ReLU(inplace=True)),
            ]
        )
    )


class Bottleneck3D(nn.Module):
    """
    Defines a bottleneck module with a residual connection
    """

    def __init__(self, in_channels, out_channels=None, kernel_size=(2, 3, 3), dilation=(1, 1, 1)):
        super().__init__()
        bottleneck_channels = in_channels // 2
        out_channels = out_channels or in_channels

        self.layers = nn.Sequential(
            OrderedDict(
                [
                    # First projection with 1x1 kernel
                    ('conv_down_project', conv_1x1x1_norm_activated(in_channels, bottleneck_channels)),
                    # Second conv block
                    (
                        'conv',
                        CausalConv3d(
                            bottleneck_channels,
                            bottleneck_channels,
                            kernel_size=kernel_size,
                            dilation=dilation,
                            bias=False,
                        ),
                    ),
                    # Final projection with 1x1 kernel
                    ('conv_up_project', conv_1x1x1_norm_activated(bottleneck_channels, out_channels)),
                ]
            )
        )

        if out_channels != in_channels:
            self.projection = nn.Sequential(
                nn.Conv3d(in_channels, out_channels, kernel_size=1, bias=False),
                nn.BatchNorm3d(out_channels),
            )
        else:
            self.projection = None

    def forward(self, *args):
        (x,) = args
        x_residual = self.layers(x)
        x_features = self.projection(x) if self.projection is not None else x
        return x_residual + x_features


class PyramidSpatioTemporalPooling(nn.Module):
    """ Spatio-temporal pyramid pooling.
        Performs 3D average pooling followed by 1x1x1 convolution to reduce the number of channels and upsampling.
        Setting contains a list of kernel_size: usually it is [(2, h, w), (2, h//2, w//2), (2, h//4, w//4)]
    """

    def __init__(self, in_channels, reduction_channels, pool_sizes):
        super().__init__()
        self.features = []
        for pool_size in pool_sizes:
            assert pool_size[0] == 2, (
                "Time kernel should be 2 as PyTorch raises an error when" "padding with more than half the kernel size"
            )
            stride = (1, *pool_size[1:])
            padding = (pool_size[0] - 1, 0, 0)
            self.features.append(
                nn.Sequential(
                    OrderedDict(
                        [
                            # Pad the input tensor but do not take into account zero padding into the average.
                            (
                                'avgpool',
                                torch.nn.AvgPool3d(
                                    kernel_size=pool_size, stride=stride, padding=padding, count_include_pad=False
                                ),
                            ),
                            ('conv_bn_relu', conv_1x1x1_norm_activated(in_channels, reduction_channels)),
                        ]
                    )
                )
            )
        self.features = nn.ModuleList(self.features)

    def forward(self, *inputs):
        (x,) = inputs
        b, _, t, h, w = x.shape
        # Do not include current tensor when concatenating
        out = []
        for f in self.features:
            # Remove unnecessary padded values (time dimension) on the right
            x_pool = f(x)[:, :, :-1].contiguous()
            c = x_pool.shape[1]
            x_pool = nn.functional.interpolate(
                x_pool.view(b * t, c, *x_pool.shape[-2:]), (h, w), mode='bilinear', align_corners=False
            )
            x_pool = x_pool.view(b, c, t, h, w)
            out.append(x_pool)
        out = torch.cat(out, 1)
        return out


class TemporalBlock(nn.Module):
    """ Temporal block with the following layers:
        - 2x3x3, 1x3x3, spatio-temporal pyramid pooling
        - dropout
        - skip connection.
    """

    def __init__(self, in_channels, out_channels=None, use_pyramid_pooling=False, pool_sizes=None):
        super().__init__()
        self.in_channels = in_channels
        self.half_channels = in_channels // 2
        self.out_channels = out_channels or self.in_channels
        self.kernels = [(2, 3, 3), (1, 3, 3)]

        # Flag for spatio-temporal pyramid pooling
        self.use_pyramid_pooling = use_pyramid_pooling

        # 3 convolution paths: 2x3x3, 1x3x3, 1x1x1
        self.convolution_paths = []
        for kernel_size in self.kernels:
            self.convolution_paths.append(
                nn.Sequential(
                    conv_1x1x1_norm_activated(self.in_channels, self.half_channels),
                    CausalConv3d(self.half_channels, self.half_channels, kernel_size=kernel_size),
                )
            )
        self.convolution_paths.append(conv_1x1x1_norm_activated(self.in_channels, self.half_channels))
        self.convolution_paths = nn.ModuleList(self.convolution_paths)

        agg_in_channels = len(self.convolution_paths) * self.half_channels

        if self.use_pyramid_pooling:
            assert pool_sizes is not None, "setting must contain the list of kernel_size, but is None."
            reduction_channels = self.in_channels // 3
            self.pyramid_pooling = PyramidSpatioTemporalPooling(self.in_channels, reduction_channels, pool_sizes)
            agg_in_channels += len(pool_sizes) * reduction_channels

        # Feature aggregation
        self.aggregation = nn.Sequential(
            conv_1x1x1_norm_activated(agg_in_channels, self.out_channels),)

        if self.out_channels != self.in_channels:
            self.projection = nn.Sequential(
                nn.Conv3d(self.in_channels, self.out_channels, kernel_size=1, bias=False),
                nn.BatchNorm3d(self.out_channels),
            )
        else:
            self.projection = None

    def forward(self, *inputs):
        (x,) = inputs
        x_paths = []
        for conv in self.convolution_paths:
            x_paths.append(conv(x))
        x_residual = torch.cat(x_paths, dim=1)
        if self.use_pyramid_pooling:
            x_pool = self.pyramid_pooling(x)
            x_residual = torch.cat([x_residual, x_pool], dim=1)
        x_residual = self.aggregation(x_residual)

        if self.out_channels != self.in_channels:
            x = self.projection(x)
        x = x + x_residual
        return x

data/stp3/layers/temporal.py
llmbev/talk2bev
import torch
import torch.nn as nn
import torch.nn.functional as F
DEBUG = False

class SpatialRegressionLoss(nn.Module):
    def __init__(self, norm, ignore_index=255, future_discount=1.0):
        super(SpatialRegressionLoss, self).__init__()
        self.norm = norm
        self.ignore_index = ignore_index
        self.future_discount = future_discount

        if norm == 1:
            self.loss_fn = F.l1_loss
        elif norm == 2:
            self.loss_fn = F.mse_loss
        else:
            raise ValueError(f'Expected norm 1 or 2, but got norm={norm}')

    def forward(self, prediction, target, n_present=3):
        assert len(prediction.shape) == 5, 'Must be a 5D tensor'
        # ignore_index is the same across all channels
        mask = target[:, :, :1] != self.ignore_index
        if mask.sum() == 0:
            return prediction.new_zeros(1)[0].float()
        loss = self.loss_fn(prediction, target, reduction='none')
        # Sum channel dimension
        loss = torch.sum(loss, dim=-3, keepdim=True)
        seq_len = loss.shape[1]
        assert seq_len >= n_present
        future_len = seq_len - n_present
        future_discounts = self.future_discount ** torch.arange(1, future_len+1, device=loss.device, dtype=loss.dtype)
        discounts = torch.cat([torch.ones(n_present, device=loss.device, dtype=loss.dtype), future_discounts], dim=0)
        discounts = discounts.view(1, seq_len, 1, 1, 1)
        loss = loss * discounts

        return loss[mask].mean()

class SegmentationLoss(nn.Module):
    def __init__(self, class_weights, ignore_index=255, use_top_k=False, top_k_ratio=1.0, future_discount=1.0):
        super().__init__()
        self.class_weights = class_weights
        self.ignore_index = ignore_index
        self.use_top_k = use_top_k
        self.top_k_ratio = top_k_ratio
        self.future_discount = future_discount

    def forward(self, prediction, target, n_present=3):
        if target.shape[-3] != 1:
            raise ValueError('segmentation label must be an index-label with channel dimension = 1.')
        b, s, c, h, w = prediction.shape
        prediction = prediction.view(b * s, c, h, w)
        target = target.view(b * s, h, w)
        loss = F.cross_entropy(
            prediction,
            target,
            ignore_index=self.ignore_index,
            reduction='none',
            weight=self.class_weights.to(target.device),
        )

        loss = loss.view(b, s, h, w)

        assert s >= n_present
        future_len = s - n_present
        future_discounts = self.future_discount ** torch.arange(1, future_len+1, device=loss.device, dtype=loss.dtype)
        discounts = torch.cat([torch.ones(n_present, device=loss.device, dtype=loss.dtype), future_discounts], dim=0)
        discounts = discounts.view(1, s, 1, 1)
        loss = loss * discounts

        loss = loss.view(b, s, -1)
        if self.use_top_k:
            # Penalises the top-k hardest pixels

            k = int(self.top_k_ratio * loss.shape[2])

            loss, _ = torch.sort(loss, dim=2, descending=True)

            loss = loss[:, :, :k]


        return torch.mean(loss)

class HDmapLoss(nn.Module):
    def __init__(self, class_weights, training_weights, use_top_k, top_k_ratio, ignore_index=255):
        super(HDmapLoss, self).__init__()
        self.class_weights = class_weights
        self.training_weights = training_weights
        self.ignore_index = ignore_index
        self.use_top_k = use_top_k
        self.top_k_ratio = top_k_ratio

    def forward(self, prediction, target):
        loss = 0
        for i in range(target.shape[-3]):
            cur_target = target[:, 0, i]
            b, h, w = cur_target.shape
            cur_prediction = prediction[:, 2*i:2*(i+1)]
            cur_loss = F.cross_entropy(
                cur_prediction,
                cur_target,
                ignore_index=self.ignore_index,
                reduction='none',
                weight=self.class_weights[i].to(target.device),
            )
            cur_loss = cur_loss.view(b, -1)
            if self.use_top_k[i]:
                k = int(self.top_k_ratio[i] * cur_loss.shape[1])
                cur_loss, _ = torch.sort(cur_loss, dim=1, descending=True)
                cur_loss = cur_loss[:, :k]
            loss += torch.mean(cur_loss) * self.training_weights[i]
        return loss

class DepthLoss(nn.Module):
    def __init__(self, class_weights=None, ignore_index=255):
        super(DepthLoss, self).__init__()
        self.class_weights = class_weights
        self.ignore_index = ignore_index

    def forward(self, prediction, target):
        b, s, n, d, h, w = prediction.shape

        prediction = prediction.view(b*s*n, d, h, w)
        target = target.view(b*s*n, h, w)
        loss = F.cross_entropy(
            prediction,
            target,
            ignore_index=self.ignore_index,
            reduction='none',
            weight=self.class_weights
        )
        return torch.mean(loss)


class ProbabilisticLoss(nn.Module):
    def __init__(self, method):
        super(ProbabilisticLoss, self).__init__()
        self.method = method

    def kl_div(self, present_mu, present_log_sigma, future_mu, future_log_sigma):
        var_future = torch.exp(2 * future_log_sigma)
        var_present = torch.exp(2 * present_log_sigma)
        kl_div = (
                present_log_sigma - future_log_sigma - 0.5 + (var_future + (future_mu - present_mu) ** 2) / (
                2 * var_present)
        )

        kl_loss = torch.mean(torch.sum(kl_div, dim=-1))
        return kl_loss

    def forward(self, output):
        if self.method == 'GAUSSIAN':

            present_mu = output['present_mu']

            present_log_sigma = output['present_log_sigma']

            future_mu = output['future_mu']

            future_log_sigma = output['future_log_sigma']


            kl_loss = self.kl_div(present_mu, present_log_sigma, future_mu, future_log_sigma)

        elif self.method == 'MIXGAUSSIAN':
            present_mu = output['present_mu']
            present_log_sigma = output['present_log_sigma']
            future_mu = output['future_mu']
            future_log_sigma = output['future_log_sigma']

            kl_loss = 0
            for i in range(len(present_mu)):
                kl_loss += self.kl_div(present_mu[i], present_log_sigma[i], future_mu[i], future_log_sigma[i])
        elif self.method == 'BERNOULLI':
            present_log_prob = output['present_log_prob']
            future_log_prob = output['future_log_prob']

            kl_loss = F.kl_div(present_log_prob, future_log_prob, reduction='batchmean', log_target=True)
        else:
            raise NotImplementedError


        return kl_loss
data/stp3/losses.py
llmbev/talk2bev
from typing import Optional

import torch
import torch.nn as nn
import numpy as np
from pytorch_lightning.metrics.metric import Metric
from pytorch_lightning.metrics.functional.classification import stat_scores_multiple_classes
from pytorch_lightning.metrics.functional.reduction import reduce
from skimage.draw import polygon

from stp3.utils.tools import gen_dx_bx
from stp3.utils.geometry import calculate_birds_eye_view_parameters


class IntersectionOverUnion(Metric):
    """Computes intersection-over-union."""
    def __init__(
        self,
        n_classes: int,
        ignore_index: Optional[int] = None,
        absent_score: float = 0.0,
        reduction: str = 'none',
        compute_on_step: bool = False,
    ):
        super().__init__(compute_on_step=compute_on_step)

        self.n_classes = n_classes
        self.ignore_index = ignore_index
        self.absent_score = absent_score
        self.reduction = reduction

        self.add_state('true_positive', default=torch.zeros(n_classes), dist_reduce_fx='sum')
        self.add_state('false_positive', default=torch.zeros(n_classes), dist_reduce_fx='sum')
        self.add_state('false_negative', default=torch.zeros(n_classes), dist_reduce_fx='sum')
        self.add_state('support', default=torch.zeros(n_classes), dist_reduce_fx='sum')

    def update(self, prediction: torch.Tensor, target: torch.Tensor):
        tps, fps, _, fns, sups = stat_scores_multiple_classes(prediction, target, self.n_classes)

        self.true_positive += tps
        self.false_positive += fps
        self.false_negative += fns
        self.support += sups

    def compute(self):
        scores = torch.zeros(self.n_classes, device=self.true_positive.device, dtype=torch.float32)

        for class_idx in range(self.n_classes):
            if class_idx == self.ignore_index:
                continue

            tp = self.true_positive[class_idx]
            fp = self.false_positive[class_idx]
            fn = self.false_negative[class_idx]
            sup = self.support[class_idx]

            # If this class is absent in the target (no support) AND absent in the pred (no true or false
            # positives), then use the absent_score for this class.
            if sup + tp + fp == 0:
                scores[class_idx] = self.absent_score
                continue

            denominator = tp + fp + fn
            score = tp.to(torch.float) / denominator
            scores[class_idx] = score

        # Remove the ignored class index from the scores.
        if (self.ignore_index is not None) and (0 <= self.ignore_index < self.n_classes):
            scores = torch.cat([scores[:self.ignore_index], scores[self.ignore_index+1:]])

        return reduce(scores, reduction=self.reduction)


class PanopticMetric(Metric):
    def __init__(
        self,
        n_classes: int,
        temporally_consistent: bool = True,
        vehicles_id: int = 1,
        compute_on_step: bool = False,
    ):
        super().__init__(compute_on_step=compute_on_step)

        self.n_classes = n_classes
        self.temporally_consistent = temporally_consistent
        self.vehicles_id = vehicles_id
        self.keys = ['iou', 'true_positive', 'false_positive', 'false_negative']

        self.add_state('iou', default=torch.zeros(n_classes), dist_reduce_fx='sum')
        self.add_state('true_positive', default=torch.zeros(n_classes), dist_reduce_fx='sum')
        self.add_state('false_positive', default=torch.zeros(n_classes), dist_reduce_fx='sum')
        self.add_state('false_negative', default=torch.zeros(n_classes), dist_reduce_fx='sum')

    def update(self, pred_instance, gt_instance):
        """
        Update state with predictions and targets.

        Parameters
        ----------
            pred_instance: (b, s, h, w)
                Temporally consistent instance segmentation prediction.
            gt_instance: (b, s, h, w)
                Ground truth instance segmentation.
        """
        batch_size, sequence_length = gt_instance.shape[:2]
        # Process labels
        assert gt_instance.min() == 0, 'ID 0 of gt_instance must be background'
        pred_segmentation = (pred_instance > 0).long()
        gt_segmentation = (gt_instance > 0).long()

        for b in range(batch_size):
            unique_id_mapping = {}
            for t in range(sequence_length):
                result = self.panoptic_metrics(
                    pred_segmentation[b, t].detach(),
                    pred_instance[b, t].detach(),
                    gt_segmentation[b, t],
                    gt_instance[b, t],
                    unique_id_mapping,
                )

                self.iou += result['iou']
                self.true_positive += result['true_positive']
                self.false_positive += result['false_positive']
                self.false_negative += result['false_negative']

    def compute(self):
        denominator = torch.maximum(
            (self.true_positive + self.false_positive / 2 + self.false_negative / 2),
            torch.ones_like(self.true_positive)
        )
        pq = self.iou / denominator
        sq = self.iou / torch.maximum(self.true_positive, torch.ones_like(self.true_positive))
        rq = self.true_positive / denominator

        return {'pq': pq,
                'sq': sq,
                'rq': rq,
                # If 0, it means there wasn't any detection.
                # 'denominator': (self.true_positive + self.false_positive / 2 + self.false_negative / 2),
                }

    def panoptic_metrics(self, pred_segmentation, pred_instance, gt_segmentation, gt_instance, unique_id_mapping):
        """
        Computes panoptic quality metric components.

        Parameters
        ----------
            pred_segmentation: [H, W] range {0, ..., n_classes-1} (>= n_classes is void)
            pred_instance: [H, W] range {0, ..., n_instances} (zero means background)
            gt_segmentation: [H, W] range {0, ..., n_classes-1} (>= n_classes is void)
            gt_instance: [H, W] range {0, ..., n_instances} (zero means background)
            unique_id_mapping: instance id mapping to check consistency
        """
        n_classes = self.n_classes

        result = {key: torch.zeros(n_classes, dtype=torch.float32, device=gt_instance.device) for key in self.keys}

        assert pred_segmentation.dim() == 2
        assert pred_segmentation.shape == pred_instance.shape == gt_segmentation.shape == gt_instance.shape

        n_instances = int(torch.cat([pred_instance, gt_instance]).max().item())
        n_all_things = n_instances + n_classes  # Classes + instances.
        n_things_and_void = n_all_things + 1

        # Now 1 is background; 0 is void (not used). 2 is vehicle semantic class but since it overlaps with
        # instances, it is not present.
        # and the rest are instance ids starting from 3
        prediction, pred_to_cls = self.combine_mask(pred_segmentation, pred_instance, n_classes, n_all_things)
        target, target_to_cls = self.combine_mask(gt_segmentation, gt_instance, n_classes, n_all_things)

        # Compute ious between all stuff and things
        # hack for bincounting 2 arrays together
        x = prediction + n_things_and_void * target
        bincount_2d = torch.bincount(x.long(), minlength=n_things_and_void ** 2)
        if bincount_2d.shape[0] != n_things_and_void ** 2:
            raise ValueError('Incorrect bincount size.')
        conf = bincount_2d.reshape((n_things_and_void, n_things_and_void))
        # Drop void class
        conf = conf[1:, 1:]

        # Confusion matrix contains intersections between all combinations of classes
        union = conf.sum(0).unsqueeze(0) + conf.sum(1).unsqueeze(1) - conf
        iou = torch.where(union > 0, (conf.float() + 1e-9) / (union.float() + 1e-9), torch.zeros_like(union).float())

        # In the iou matrix, first dimension is target idx, second dimension is pred idx.
        # Mapping will contain a tuple that maps prediction idx to target idx for segments matched by iou.
        mapping = (iou > 0.5).nonzero(as_tuple=False)

        # Check that classes match.
        is_matching = pred_to_cls[mapping[:, 1]] == target_to_cls[mapping[:, 0]]
        mapping = mapping[is_matching]
        tp_mask = torch.zeros_like(conf, dtype=torch.bool)
        tp_mask[mapping[:, 0], mapping[:, 1]] = True

        # First ids correspond to "stuff" i.e. semantic seg.
        # Instance ids are offset accordingly
        for target_id, pred_id in mapping:
            cls_id = pred_to_cls[pred_id]

            if self.temporally_consistent and cls_id == self.vehicles_id:
                if target_id.item() in unique_id_mapping and unique_id_mapping[target_id.item()] != pred_id.item():
                    # Not temporally consistent
                    result['false_negative'][target_to_cls[target_id]] += 1
                    result['false_positive'][pred_to_cls[pred_id]] += 1
                    unique_id_mapping[target_id.item()] = pred_id.item()
                    continue

            result['true_positive'][cls_id] += 1
            result['iou'][cls_id] += iou[target_id][pred_id]
            unique_id_mapping[target_id.item()] = pred_id.item()

        for target_id in range(n_classes, n_all_things):
            # If this is a true positive do nothing.
            if tp_mask[target_id, n_classes:].any():
                continue
            # If this target instance didn't match with any predictions and was present set it as false negative.
            if target_to_cls[target_id] != -1:
                result['false_negative'][target_to_cls[target_id]] += 1

        for pred_id in range(n_classes, n_all_things):
            # If this is a true positive do nothing.
            if tp_mask[n_classes:, pred_id].any():
                continue
            # If this predicted instance didn't match with any prediction, set that predictions as false positive.
            if pred_to_cls[pred_id] != -1 and (conf[:, pred_id] > 0).any():
                result['false_positive'][pred_to_cls[pred_id]] += 1

        return result

    def combine_mask(self, segmentation: torch.Tensor, instance: torch.Tensor, n_classes: int, n_all_things: int):
        """Shifts all things ids by num_classes and combines things and stuff into a single mask

        Returns a combined mask + a mapping from id to segmentation class.
        """
        instance = instance.view(-1)
        instance_mask = instance > 0
        instance = instance - 1 + n_classes

        segmentation = segmentation.clone().view(-1)
        segmentation_mask = segmentation < n_classes  # Remove void pixels.

        # Build an index from instance id to class id.
        instance_id_to_class_tuples = torch.cat(
            (
                instance[instance_mask & segmentation_mask].unsqueeze(1),
                segmentation[instance_mask & segmentation_mask].unsqueeze(1),
            ),
            dim=1,
        )
        instance_id_to_class = -instance_id_to_class_tuples.new_ones((n_all_things,))
        instance_id_to_class[instance_id_to_class_tuples[:, 0]] = instance_id_to_class_tuples[:, 1]
        instance_id_to_class[torch.arange(n_classes, device=segmentation.device)] = torch.arange(
            n_classes, device=segmentation.device
        )

        segmentation[instance_mask] = instance[instance_mask]
        segmentation += 1  # Shift all legit classes by 1.
        segmentation[~segmentation_mask] = 0  # Shift void class to zero.

        return segmentation, instance_id_to_class

class PlanningMetric(Metric):
    def __init__(
        self,
        cfg,
        n_future=4,
        compute_on_step: bool = False,
    ):
        super().__init__(compute_on_step=compute_on_step)
        dx, bx, _ = gen_dx_bx(cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND)
        dx, bx = dx[:2], bx[:2]
        self.dx = nn.Parameter(dx, requires_grad=False)
        self.bx = nn.Parameter(bx, requires_grad=False)

        _, _, self.bev_dimension = calculate_birds_eye_view_parameters(
            cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND
        )
        self.bev_dimension = self.bev_dimension.numpy()

        self.W = cfg.EGO.WIDTH
        self.H = cfg.EGO.HEIGHT

        self.n_future = n_future

        self.add_state("obj_col", default=torch.zeros(self.n_future), dist_reduce_fx="sum")
        self.add_state("obj_box_col", default=torch.zeros(self.n_future), dist_reduce_fx="sum")
        self.add_state("L2", default=torch.zeros(self.n_future),dist_reduce_fx="sum")
        self.add_state("total", default=torch.tensor(0), dist_reduce_fx="sum")


    def evaluate_single_coll(self, traj, segmentation):
        '''
        gt_segmentation
        traj: torch.Tensor (n_future, 2)
        segmentation: torch.Tensor (n_future, 200, 200)
        '''
        pts = np.array([
            [-self.H / 2. + 0.5, self.W / 2.],
            [self.H / 2. + 0.5, self.W / 2.],
            [self.H / 2. + 0.5, -self.W / 2.],
            [-self.H / 2. + 0.5, -self.W / 2.],
        ])
        pts = (pts - self.bx.cpu().numpy()) / (self.dx.cpu().numpy())
        pts[:, [0, 1]] = pts[:, [1, 0]]
        rr, cc = polygon(pts[:,1], pts[:,0])
        rc = np.concatenate([rr[:,None], cc[:,None]], axis=-1)

        n_future, _ = traj.shape
        trajs = traj.view(n_future, 1, 2)
        trajs[:,:,[0,1]] = trajs[:,:,[1,0]] # can also change original tensor
        trajs = trajs / self.dx
        trajs = trajs.cpu().numpy() + rc # (n_future, 32, 2)

        r = trajs[:,:,0].astype(np.int32)
        r = np.clip(r, 0, self.bev_dimension[0] - 1)

        c = trajs[:,:,1].astype(np.int32)
        c = np.clip(c, 0, self.bev_dimension[1] - 1)

        collision = np.full(n_future, False)
        for t in range(n_future):
            rr = r[t]
            cc = c[t]
            I = np.logical_and(
                np.logical_and(rr >= 0, rr < self.bev_dimension[0]),
                np.logical_and(cc >= 0, cc < self.bev_dimension[1]),
            )
            collision[t] = np.any(segmentation[t, rr[I], cc[I]].cpu().numpy())

        return torch.from_numpy(collision).to(device=traj.device)

    def evaluate_coll(self, trajs, gt_trajs, segmentation):
        '''
        trajs: torch.Tensor (B, n_future, 2)
        gt_trajs: torch.Tensor (B, n_future, 2)
        segmentation: torch.Tensor (B, n_future, 200, 200)
        '''
        B, n_future, _ = trajs.shape
        trajs = trajs * torch.tensor([-1, 1], device=trajs.device)
        gt_trajs = gt_trajs * torch.tensor([-1, 1], device=gt_trajs.device)

        obj_coll_sum = torch.zeros(n_future, device=segmentation.device)
        obj_box_coll_sum = torch.zeros(n_future, device=segmentation.device)

        for i in range(B):
            gt_box_coll = self.evaluate_single_coll(gt_trajs[i], segmentation[i])

            xx, yy = trajs[i,:,0], trajs[i, :, 1]
            yi = ((yy - self.bx[0]) / self.dx[0]).long()
            xi = ((xx - self.bx[1]) / self.dx[1]).long()

            m1 = torch.logical_and(
                torch.logical_and(yi >= 0, yi < self.bev_dimension[0]),
                torch.logical_and(xi >= 0, xi < self.bev_dimension[1]),
            )
            m1 = torch.logical_and(m1, torch.logical_not(gt_box_coll))

            ti = torch.arange(n_future)
            obj_coll_sum[ti[m1]] += segmentation[i, ti[m1], yi[m1], xi[m1]].long()

            m2 = torch.logical_not(gt_box_coll)
            box_coll = self.evaluate_single_coll(trajs[i], segmentation[i])
            obj_box_coll_sum[ti[m2]] += (box_coll[ti[m2]]).long()

        return obj_coll_sum, obj_box_coll_sum

    def compute_L2(self, trajs, gt_trajs):
        '''
        trajs: torch.Tensor (B, n_future, 3)
        gt_trajs: torch.Tensor (B, n_future, 3)
        '''

        return torch.sqrt(((trajs[:, :, :2] - gt_trajs[:, :, :2]) ** 2).sum(dim=-1))

    def update(self, trajs, gt_trajs, segmentation):
        '''
        trajs: torch.Tensor (B, n_future, 3)
        gt_trajs: torch.Tensor (B, n_future, 3)
        segmentation: torch.Tensor (B, n_future, 200, 200)
        '''
        assert trajs.shape == gt_trajs.shape
        L2 = self.compute_L2(trajs, gt_trajs)
        obj_coll_sum, obj_box_coll_sum = self.evaluate_coll(trajs[:,:,:2], gt_trajs[:,:,:2], segmentation)

        self.obj_col += obj_coll_sum
        self.obj_box_col += obj_box_coll_sum
        self.L2 += L2.sum(dim=0)
        self.total +=len(trajs)

    def compute(self):
        return {
            'obj_col': self.obj_col / self.total,
            'obj_box_col': self.obj_box_col / self.total,
            'L2' : self.L2 / self.total
        }
data/stp3/metrics.py
llmbev/talk2bev
import torch
import torch.nn as nn
from torchvision.models.resnet import resnet18

from stp3.layers.convolutions import UpsamplingAdd, DeepLabHead


class Decoder(nn.Module):
    def __init__(self, in_channels, n_classes, n_present, n_hdmap, predict_gate):
        super().__init__()
        self.perceive_hdmap = predict_gate['perceive_hdmap']
        self.predict_pedestrian = predict_gate['predict_pedestrian']
        self.predict_instance = predict_gate['predict_instance']
        self.predict_future_flow = predict_gate['predict_future_flow']
        self.planning = predict_gate['planning']

        self.n_classes = n_classes
        self.n_present = n_present
        if self.predict_instance is False and self.predict_future_flow is True:
            raise ValueError('flow cannot be True when not predicting instance')

        backbone = resnet18(pretrained=False, zero_init_residual=True)

        self.first_conv = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = backbone.bn1
        self.relu = backbone.relu

        self.layer1 = backbone.layer1
        self.layer2 = backbone.layer2
        self.layer3 = backbone.layer3

        shared_out_channels = in_channels
        self.up3_skip = UpsamplingAdd(256, 128, scale_factor=2)
        self.up2_skip = UpsamplingAdd(128, 64, scale_factor=2)
        self.up1_skip = UpsamplingAdd(64, shared_out_channels, scale_factor=2)

        self.segmentation_head = nn.Sequential(
            nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(shared_out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(shared_out_channels, self.n_classes, kernel_size=1, padding=0),
        )

        if self.predict_pedestrian:
            self.pedestrian_head = nn.Sequential(
                nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(shared_out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(shared_out_channels, self.n_classes, kernel_size=1, padding=0),
            )

        if self.perceive_hdmap:
            self.hdmap_head = nn.Sequential(
                nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(shared_out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(shared_out_channels, 2 * n_hdmap, kernel_size=1, padding=0),
            )

        if self.predict_instance:
            self.instance_offset_head = nn.Sequential(
                nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(shared_out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(shared_out_channels, 2, kernel_size=1, padding=0),
            )
            self.instance_center_head = nn.Sequential(
                nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(shared_out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(shared_out_channels, 1, kernel_size=1, padding=0),
                nn.Sigmoid(),
            )

        if self.predict_future_flow:
            self.instance_future_head = nn.Sequential(
                nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(shared_out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(shared_out_channels, 2, kernel_size=1, padding=0),
            )

        if self.planning:
            self.costvolume_head = nn.Sequential(
                nn.Conv2d(shared_out_channels, shared_out_channels, kernel_size=3, padding=1, bias=False),
                nn.BatchNorm2d(shared_out_channels),
                nn.ReLU(inplace=True),
                nn.Conv2d(shared_out_channels, 1, kernel_size=1, padding=0),
            )

    def forward(self, x):
        b, s, c, h, w = x.shape
        x = x.view(b * s, c, h, w)
        # (H, W)
        skip_x = {'1': x}

        # (H/2, W/2)
        x = self.first_conv(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.layer1(x)
        skip_x['2'] = x

        # (H/4 , W/4)
        x = self.layer2(x)
        skip_x['3'] = x

        # (H/8, W/8)
        x = self.layer3(x)  # (b*s, 256, 25, 25)

        # First upsample to (H/4, W/4)
        x = self.up3_skip(x, skip_x['3'])

        # Second upsample to (H/2, W/2)
        x = self.up2_skip(x, skip_x['2'])

        # Third upsample to (H, W)
        x = self.up1_skip(x, skip_x['1'])

        segmentation_output = self.segmentation_head(x)
        pedestrian_output = self.pedestrian_head(x) if self.predict_pedestrian else None
        hdmap_output = self.hdmap_head(x.view(b, s, *x.shape[1:])[:,self.n_present-1]) if self.perceive_hdmap else None
        instance_center_output = self.instance_center_head(x) if self.predict_instance else None
        instance_offset_output = self.instance_offset_head(x) if self.predict_instance else None
        instance_future_output = self.instance_future_head(x) if self.predict_future_flow else None
        costvolume = self.costvolume_head(x).squeeze(1) if self.planning else None
        return {
            'segmentation': segmentation_output.view(b, s, *segmentation_output.shape[1:]),
            'pedestrian': pedestrian_output.view(b, s, *pedestrian_output.shape[1:])
            if pedestrian_output is not None else None,
            'hdmap' : hdmap_output,
            'instance_center': instance_center_output.view(b, s, *instance_center_output.shape[1:])
            if instance_center_output is not None else None,
            'instance_offset': instance_offset_output.view(b, s, *instance_offset_output.shape[1:])
            if instance_offset_output is not None else None,
            'instance_flow': instance_future_output.view(b, s, *instance_future_output.shape[1:])
            if instance_future_output is not None else None,
            'costvolume': costvolume.view(b, s, *costvolume.shape[1:])
            if costvolume is not None else None,
        }

data/stp3/models/decoder.py
llmbev/talk2bev
import torch
import torch.nn as nn

from stp3.layers.convolutions import Bottleneck


class DistributionModule(nn.Module):
    """
    A convolutional net that parametrises a diagonal Gaussian distribution.
    """

    def __init__(
        self, in_channels, latent_dim, method="GAUSSIAN"):
        super().__init__()
        self.compress_dim = in_channels // 2
        self.latent_dim = latent_dim
        self.method = method

        if method == 'GAUSSIAN':
            self.encoder = DistributionEncoder(in_channels, self.compress_dim)
            self.decoder = nn.Sequential(
                nn.AdaptiveAvgPool2d(1), nn.Conv2d(self.compress_dim, out_channels=2 * self.latent_dim, kernel_size=1)
            )
        elif method == 'MIXGAUSSIAN':
            self.encoder = DistributionEncoder(in_channels, self.compress_dim)
            self.decoder = nn.Sequential(
                nn.AdaptiveAvgPool2d(1), nn.Conv2d(self.compress_dim, out_channels=6 * self.latent_dim + 3, kernel_size=1)
            )
        elif method == 'BERNOULLI':
            self.encoder = nn.Sequential(
                Bottleneck(in_channels, self.latent_dim)
            )
            self.decoder = nn.LogSigmoid()
        else:
            raise NotImplementedError

    def forward(self, s_t):
        b, s = s_t.shape[:2]
        assert s == 1
        encoding = self.encoder(s_t[:, 0])

        if self.method == 'GAUSSIAN':
            decoder = self.decoder(encoding).view(b, 1, 2 * self.latent_dim)
        elif self.method == 'MIXGAUSSIAN':
            decoder = self.decoder(encoding).view(b, 1, 6 * self.latent_dim + 3)
        elif self.method == 'BERNOULLI':
            decoder = self.decoder(encoding)
        else:
            raise NotImplementedError

        return decoder


class DistributionEncoder(nn.Module):
    """Encodes s_t or (s_t, y_{t+1}, ..., y_{t+H}).
    """
    def __init__(self, in_channels, out_channels):
        super().__init__()

        self.model = nn.Sequential(
            Bottleneck(in_channels, out_channels=out_channels, downsample=True),
            Bottleneck(out_channels, out_channels=out_channels, downsample=True),
            Bottleneck(out_channels, out_channels=out_channels, downsample=True),
            Bottleneck(out_channels, out_channels=out_channels, downsample=True),
        )

    def forward(self, s_t):
        return self.model(s_t)

data/stp3/models/distributions.py
llmbev/talk2bev
import torch.nn as nn
import numpy as np
from efficientnet_pytorch import EfficientNet

from stp3.layers.convolutions import UpsamplingConcat, DeepLabHead



class Encoder(nn.Module):
    def __init__(self, cfg, D):
        super().__init__()
        self.D = D
        self.C = cfg.OUT_CHANNELS
        self.use_depth_distribution = cfg.USE_DEPTH_DISTRIBUTION
        self.downsample = cfg.DOWNSAMPLE
        self.version = cfg.NAME.split('-')[1]

        self.backbone = EfficientNet.from_pretrained(cfg.NAME)
        self.delete_unused_layers()
        if self.version == 'b4':
            self.reduction_channel = [0, 24, 32, 56, 160, 448]
        elif self.version == 'b0':
            self.reduction_channel = [0, 16, 24, 40, 112, 320]
        else:
            raise NotImplementedError
        self.upsampling_out_channel = [0, 48, 64, 128, 512]

        # index = np.log2(self.downsample).astype(np.int)
        index = np.log2(self.downsample).astype(int)


        if self.use_depth_distribution:
            self.depth_layer_1 = DeepLabHead(self.reduction_channel[index+1], self.reduction_channel[index+1], hidden_channel=64)
            self.depth_layer_2 = UpsamplingConcat(self.reduction_channel[index+1] + self.reduction_channel[index], self.D)

        self.feature_layer_1 = DeepLabHead(self.reduction_channel[index+1], self.reduction_channel[index+1], hidden_channel=64)
        self.feature_layer_2 = UpsamplingConcat(self.reduction_channel[index+1] + self.reduction_channel[index], self.C)



    def delete_unused_layers(self):
        indices_to_delete = []
        for idx in range(len(self.backbone._blocks)):
            if self.downsample == 8:
                if self.version == 'b0' and idx > 10:
                    indices_to_delete.append(idx)
                if self.version == 'b4' and idx > 21:
                    indices_to_delete.append(idx)

        for idx in reversed(indices_to_delete):
            del self.backbone._blocks[idx]

        del self.backbone._conv_head
        del self.backbone._bn1
        del self.backbone._avg_pooling
        del self.backbone._dropout
        del self.backbone._fc

    def get_features_depth(self, x):
        # Adapted from https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py#L231
        endpoints = dict()

        # Stem
        x = self.backbone._swish(self.backbone._bn0(self.backbone._conv_stem(x)))
        prev_x = x

        # Blocks
        for idx, block in enumerate(self.backbone._blocks):
            drop_connect_rate = self.backbone._global_params.drop_connect_rate
            if drop_connect_rate:
                drop_connect_rate *= float(idx) / len(self.backbone._blocks)
            x = block(x, drop_connect_rate=drop_connect_rate)
            if prev_x.size(2) > x.size(2):
                endpoints['reduction_{}'.format(len(endpoints) + 1)] = prev_x
            prev_x = x

            if self.downsample == 8:
                if self.version == 'b0' and idx == 10:
                    break
                if self.version == 'b4' and idx == 21:
                    break

        # Head
        endpoints['reduction_{}'.format(len(endpoints) + 1)] = x

        index = np.log2(self.downsample).astype(np.int)
        input_1 = endpoints['reduction_{}'.format(index + 1)]
        input_2 = endpoints['reduction_{}'.format(index)]

        feature = self.feature_layer_1(input_1)
        feature = self.feature_layer_2(feature, input_2)

        if self.use_depth_distribution:
            depth = self.depth_layer_1(input_1)
            depth = self.depth_layer_2(depth, input_2)
        else:
            depth = None

        return feature, depth

    def forward(self, x):
        feature, depth = self.get_features_depth(x)  # get feature vector

        # if self.use_depth_distribution:
        #     depth_prob = depth.softmax(dim=1)
        #     feature = depth_prob.unsqueeze(1) * feature.unsqueeze(2)  # outer product depth and features
        # else:
        #     feature = feature.unsqueeze(2).repeat(1, 1, self.D, 1, 1)

        return feature, depth

data/stp3/models/encoder.py
llmbev/talk2bev
import torch
import torch.nn as nn

from stp3.layers.convolutions import Bottleneck, Block, DeepLabHead
from stp3.layers.temporal import SpatialGRU, Dual_GRU, BiGRU

class FuturePrediction(nn.Module):
    def __init__(self, in_channels, latent_dim, n_future, mixture=True, n_gru_blocks=2, n_res_layers=1):
        super(FuturePrediction, self).__init__()
        self.n_spatial_gru = n_gru_blocks

        gru_in_channels = latent_dim
        self.dual_grus = Dual_GRU(gru_in_channels, in_channels, n_future=n_future, mixture=mixture)
        self.res_blocks1 = nn.Sequential(*[Block(in_channels) for _ in range(n_res_layers)])

        self.spatial_grus = []
        self.res_blocks = []
        for i in range(self.n_spatial_gru):
            self.spatial_grus.append(SpatialGRU(in_channels, in_channels))
            if i < self.n_spatial_gru - 1:
                self.res_blocks.append(nn.Sequential(*[Block(in_channels) for _ in range(n_res_layers)]))
            else:
                self.res_blocks.append(DeepLabHead(in_channels, in_channels, 128))

        self.spatial_grus = torch.nn.ModuleList(self.spatial_grus)
        self.res_blocks = torch.nn.ModuleList(self.res_blocks)


    def forward(self, x, state):
        # x has shape (b, 1, c, h, w), state: torch.Tensor [b, n_present, hidden_size, h, w]
        x = self.dual_grus(x, state)

        b, n_future, c, h, w = x.shape
        x = self.res_blocks1(x.view(b * n_future, c, h, w))
        x = x.view(b, n_future, c, h, w)

        x = torch.cat([state, x], dim=1)

        hidden_state = x[:, 0]
        for i in range(self.n_spatial_gru):
            x = self.spatial_grus[i](x, hidden_state)

            b, s, c, h, w = x.shape
            x = self.res_blocks[i](x.view(b*s, c, h, w))
            x = x.view(b, s, c, h, w)

        return x
data/stp3/models/future_prediction.py
llmbev/talk2bev
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from src.layers.convolutions import Bottleneck
from src.layers.temporal import SpatialGRU, Dual_GRU, BiGRU
from src.cost import Cost_Function

class Planning(nn.Module):
    def __init__(self, cfg, feature_channel, gru_input_size=6, gru_state_size=256):
        super(Planning, self).__init__()
        self.cost_function = Cost_Function(cfg)

        self.sample_num = cfg.PLANNING.SAMPLE_NUM
        self.commands = cfg.PLANNING.COMMAND
        assert self.sample_num % 3 == 0
        self.num = int(self.sample_num / 3)

        self.reduce_channel = nn.Sequential(
            Bottleneck(feature_channel, feature_channel, downsample=True),
            Bottleneck(feature_channel, int(feature_channel/2), downsample=True),
            Bottleneck(int(feature_channel/2), int(feature_channel/2), downsample=True),
            Bottleneck(int(feature_channel/2), int(feature_channel/8))
        )

        self.GRU = nn.GRUCell(gru_input_size, gru_state_size)
        self.decoder = nn.Sequential(
            nn.Linear(gru_state_size, gru_state_size),
            nn.ReLU(inplace=True),
            nn.Linear(gru_state_size, 2)
        )


    def compute_L2(self, trajs, gt_traj):
        '''
        trajs: torch.Tensor (B, N, n_future, 3)
        gt_traj: torch.Tensor (B,1, n_future, 3)
        '''
        if trajs.ndim == 4 and gt_traj.ndim == 4:
            return ((trajs[:,:,:,:2] - gt_traj[:,:,:,:2]) ** 2).sum(dim=-1)
        if trajs.ndim == 3 and gt_traj.ndim == 3:
            return ((trajs[:, :, :2] - gt_traj[:, :, :2]) ** 2).sum(dim=-1)

        raise ValueError('trajs ndim != gt_traj ndim')

    def select(self, trajs, cost_volume, semantic_pred, lane_divider, drivable_area, target_points, k=1):
        '''
        trajs: torch.Tensor (B, N, n_future, 3)
        cost_volume: torch.Tensor (B, n_future, 200, 200)
        semantic_pred: torch.Tensor(B, n_future, 200, 200)
        lane_divider: torch.Tensor(B, 1/2, 200, 200)
        drivable_area: torch.Tensor(B, 1/2, 200, 200)
        target_points: torch.Tensor<float> (B, 2)
        '''
        sm_cost_fc, sm_cost_fo = self.cost_function(cost_volume, trajs[:,:,:,:2], semantic_pred, lane_divider, drivable_area, target_points)

        CS = sm_cost_fc + sm_cost_fo.sum(dim=-1)
        CC, KK = torch.topk(CS, k, dim=-1, largest=False)

        ii = torch.arange(len(trajs))
        select_traj = trajs[ii[:,None], KK].squeeze(1) # (B, n_future, 3)

        return select_traj, CS

    def loss(self, trajs, gt_trajs, cost_volume, semantic_pred, lane_divider, drivable_area, target_points):
        '''
        trajs: torch.Tensor (B, N, n_future, 3)
        gt_trajs: torch.Tensor (B, n_future, 3)
        cost_volume: torch.Tensor (B, n_future, 200, 200)
        semantic_pred: torch.Tensor(B, n_future, 200, 200)
        lane_divider: torch.Tensor(B, 1/2, 200, 200)
        drivable_area: torch.Tensor(B, 1/2, 200, 200)
        target_points: torch.Tensor<float> (B, 2)
        '''
        sm_cost_fc, sm_cost_fo = self.cost_function(cost_volume, trajs[:, :, :, :2], semantic_pred, lane_divider, drivable_area, target_points)
        print(sm_cost_fc.shape, sm_cost_fo.shape, "sm_costs")

        if gt_trajs.ndim == 3:
            gt_trajs = gt_trajs[:, None]

        gt_cost_fc, gt_cost_fo = self.cost_function(cost_volume, gt_trajs[:, :, :, :2], semantic_pred, lane_divider, drivable_area, target_points)

        L, _ = F.relu(
            F.relu(gt_cost_fo - sm_cost_fo).sum(-1) + (gt_cost_fc - sm_cost_fc) + self.compute_L2(trajs, gt_trajs).mean(
                dim=-1)).max(dim=-1)

        return torch.mean(L), gt_cost_fc, gt_cost_fo, sm_cost_fc, sm_cost_fo

    def forward(self,cam_front, trajs, gt_trajs, cost_volume, semantic_pred, hd_map, commands, target_points):
        '''
        cam_front: torch.Tensor (B, 64, 60, 28)
        trajs: torch.Tensor (B, N, n_future, 3)
        gt_trajs: torch.Tensor (B, n_future, 3)
        cost_volume: torch.Tensor (B, n_future, 200, 200)
        semantic_pred: torch.Tensor(B, n_future, 200, 200)
        hd_map: torch.Tensor (B, 2/4, 200, 200)
        commands: List (B)
        target_points: (B, 2)
        '''

        cur_trajs = []
        for i in range(len(commands)):
            command = commands[i]
            traj = trajs[i]
            if command == 'LEFT':
                cur_trajs.append(traj[:self.num].repeat(3, 1, 1))
            elif command == 'FORWARD':
                cur_trajs.append(traj[self.num:self.num * 2].repeat(3, 1, 1))
            elif command == 'RIGHT':
                cur_trajs.append(traj[self.num * 2:].repeat(3, 1, 1))
            else:
                cur_trajs.append(traj)
        cur_trajs = torch.stack(cur_trajs)
        print(cur_trajs.shape, "cur_trajs", trajs.shape, "trajs")
        cur_trajs = trajs


        if hd_map.shape[1] == 2:
            lane_divider = hd_map[:, 0:1]
            drivable_area = hd_map[:, 1:2]
        elif hd_map.shape[1] == 4:
            lane_divider = hd_map[:, 0:2]
            drivable_area = hd_map[:, 2:4]
        else:
            raise NotImplementedError

        if self.training:
            loss, gt_cost_fc, gt_cost_fo, sm_cost_fc, sm_cost_fo = self.loss(cur_trajs, gt_trajs, cost_volume, semantic_pred, lane_divider, drivable_area, target_points)
        else:
            loss, gt_cost_fc, gt_cost_fo, sm_cost_fc, sm_cost_fo = self.loss(cur_trajs, gt_trajs, cost_volume, semantic_pred, lane_divider, drivable_area, target_points)
            loss = 0


        cam_front = self.reduce_channel(cam_front)
        h0 = cam_front.flatten(start_dim=1) # (B, 256/128)
        final_traj, CS = self.select(cur_trajs, cost_volume, semantic_pred, lane_divider, drivable_area, target_points) # (B, n_future, 3)
        target_points = target_points.to(dtype=h0.dtype)
        b, s, _ = final_traj.shape
        x = torch.zeros((b, 2), device=h0.device)
        output_traj = []
        for i in range(s):
            x = torch.cat([x, final_traj[:,i,:2], target_points], dim=-1) # (B, 6)
            h0 = self.GRU(x, h0)
            x = self.decoder(h0) # (B, 2)
            output_traj.append(x)
        output_traj = torch.stack(output_traj, dim=1) # (B, 4, 2)

        output_traj = torch.cat(
            [output_traj, torch.zeros((*output_traj.shape[:-1],1), device=output_traj.device)], dim=-1
        )

        if self.training:
            loss = loss*0.5 + (F.smooth_l1_loss(final[:,:,:2], gt_trajs[:,:,:2], reduction='none')*torch.tensor([10., 1.], device=loss.device)).mean()


        return loss, output_traj, CS, cur_trajs, gt_cost_fc, gt_cost_fo, final_traj, sm_cost_fc, sm_cost_fo

data/stp3/models/planning_model.py
llmbev/talk2bev
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from stp3.models.encoder import Encoder
from stp3.models.temporal_model import TemporalModelIdentity, TemporalModel
from stp3.models.distributions import DistributionModule
from stp3.models.future_prediction import FuturePrediction
from stp3.models.decoder import Decoder
from stp3.models.planning_model import Planning
from stp3.utils.network import pack_sequence_dim, unpack_sequence_dim, set_bn_momentum
from stp3.utils.geometry import calculate_birds_eye_view_parameters, VoxelsSumming, pose_vec2mat

class STP3(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg

        bev_resolution, bev_start_position, bev_dimension = calculate_birds_eye_view_parameters(
            self.cfg.LIFT.X_BOUND, self.cfg.LIFT.Y_BOUND, self.cfg.LIFT.Z_BOUND
        )
        self.bev_resolution = nn.Parameter(bev_resolution, requires_grad=False)
        self.bev_start_position = nn.Parameter(bev_start_position, requires_grad=False)
        self.bev_dimension = nn.Parameter(bev_dimension, requires_grad=False)

        self.encoder_downsample = self.cfg.MODEL.ENCODER.DOWNSAMPLE
        self.encoder_out_channels = self.cfg.MODEL.ENCODER.OUT_CHANNELS

        self.frustum = self.create_frustum()
        self.depth_channels, _, _, _ = self.frustum.shape
        self.discount = self.cfg.LIFT.DISCOUNT

        if self.cfg.TIME_RECEPTIVE_FIELD == 1:
            assert self.cfg.MODEL.TEMPORAL_MODEL.NAME == 'identity'

        # temporal block
        self.receptive_field = self.cfg.TIME_RECEPTIVE_FIELD
        self.n_future = self.cfg.N_FUTURE_FRAMES
        self.latent_dim = self.cfg.MODEL.DISTRIBUTION.LATENT_DIM

        # Spatial extent in bird's-eye view, in meters
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])
        self.bev_size = (self.bev_dimension[0].item(), self.bev_dimension[1].item())

        # Encoder
        self.encoder = Encoder(cfg=self.cfg.MODEL.ENCODER, D=self.depth_channels)

        # Temporal model
        temporal_in_channels = self.encoder_out_channels
        if self.cfg.MODEL.TEMPORAL_MODEL.INPUT_EGOPOSE:
            temporal_in_channels += 6
        if self.cfg.MODEL.TEMPORAL_MODEL.NAME == 'identity':
            self.temporal_model = TemporalModelIdentity(temporal_in_channels, self.receptive_field)
        elif cfg.MODEL.TEMPORAL_MODEL.NAME == 'temporal_block':
            self.temporal_model = TemporalModel(
                temporal_in_channels,
                self.receptive_field,
                input_shape=self.bev_size,
                start_out_channels=self.cfg.MODEL.TEMPORAL_MODEL.START_OUT_CHANNELS,
                extra_in_channels=self.cfg.MODEL.TEMPORAL_MODEL.EXTRA_IN_CHANNELS,
                n_spatial_layers_between_temporal_layers=self.cfg.MODEL.TEMPORAL_MODEL.INBETWEEN_LAYERS,
                use_pyramid_pooling=self.cfg.MODEL.TEMPORAL_MODEL.PYRAMID_POOLING,
            )
        else:
            raise NotImplementedError(f'Temporal module {self.cfg.MODEL.TEMPORAL_MODEL.NAME}.')

        self.future_pred_in_channels = self.temporal_model.out_channels
        if self.n_future > 0:
            # probabilistic sampling
            if self.cfg.PROBABILISTIC.ENABLED:
                # Distribution networks
                self.present_distribution = DistributionModule(
                    self.future_pred_in_channels,
                    self.latent_dim,
                    method=self.cfg.PROBABILISTIC.METHOD
                )

            # Future prediction
            self.future_prediction = FuturePrediction(
                in_channels=self.future_pred_in_channels,
                latent_dim=self.latent_dim,
                n_future=self.n_future,
                mixture=self.cfg.MODEL.FUTURE_PRED.MIXTURE,
                n_gru_blocks=self.cfg.MODEL.FUTURE_PRED.N_GRU_BLOCKS,
                n_res_layers=self.cfg.MODEL.FUTURE_PRED.N_RES_LAYERS,
            )

        # Decoder
        self.decoder = Decoder(
            in_channels=self.future_pred_in_channels,
            n_classes=len(self.cfg.SEMANTIC_SEG.VEHICLE.WEIGHTS),
            n_present=self.receptive_field,
            n_hdmap=len(self.cfg.SEMANTIC_SEG.HDMAP.ELEMENTS),
            predict_gate = {
                'perceive_hdmap': self.cfg.SEMANTIC_SEG.HDMAP.ENABLED,
                'predict_pedestrian': self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED,
                'predict_instance': self.cfg.INSTANCE_SEG.ENABLED,
                'predict_future_flow': self.cfg.INSTANCE_FLOW.ENABLED,
                'planning': self.cfg.PLANNING.ENABLED,
            }
        )

        # Cost function
        # Carla 128, Nuscenes 256
        if self.cfg.PLANNING.ENABLED:
            self.planning = Planning(cfg, self.encoder_out_channels, 6, gru_state_size=self.cfg.PLANNING.GRU_STATE_SIZE)

        set_bn_momentum(self, self.cfg.MODEL.BN_MOMENTUM)

    def create_frustum(self):
        # Create grid in image plane
        h, w = self.cfg.IMAGE.FINAL_DIM
        downsampled_h, downsampled_w = h // self.encoder_downsample, w // self.encoder_downsample

        # Depth grid
        depth_grid = torch.arange(*self.cfg.LIFT.D_BOUND, dtype=torch.float)
        depth_grid = depth_grid.view(-1, 1, 1).expand(-1, downsampled_h, downsampled_w)
        n_depth_slices = depth_grid.shape[0]

        # x and y grids
        x_grid = torch.linspace(0, w - 1, downsampled_w, dtype=torch.float)
        x_grid = x_grid.view(1, 1, downsampled_w).expand(n_depth_slices, downsampled_h, downsampled_w)
        y_grid = torch.linspace(0, h - 1, downsampled_h, dtype=torch.float)
        y_grid = y_grid.view(1, downsampled_h, 1).expand(n_depth_slices, downsampled_h, downsampled_w)

        # Dimension (n_depth_slices, downsampled_h, downsampled_w, 3)
        # containing data points in the image: left-right, top-bottom, depth
        frustum = torch.stack((x_grid, y_grid, depth_grid), -1)
        return nn.Parameter(frustum, requires_grad=False)

    def forward(self, image, intrinsics, extrinsics, future_egomotion):
        output = {}

        # Only process features from the past and present
        image = image[:, :self.receptive_field].contiguous()
        intrinsics = intrinsics[:, :self.receptive_field].contiguous()
        extrinsics = extrinsics[:, :self.receptive_field].contiguous()
        future_egomotion = future_egomotion[:, :self.receptive_field].contiguous()

        # Lifting features and project to bird's-eye view
        x, depth, cam_front = self.calculate_birds_eye_view_features(image, intrinsics, extrinsics, future_egomotion) # (3,3,64,200,200)
        output = {**output, 'depth_prediction': depth, 'cam_front':cam_front}

        if self.cfg.MODEL.TEMPORAL_MODEL.INPUT_EGOPOSE:
            b, s, c = future_egomotion.shape
            h, w = x.shape[-2:]
            future_egomotions_spatial = future_egomotion.view(b, s, c, 1, 1).expand(b, s, c, h, w)
            # at time 0, no egomotion so feed zero vector
            future_egomotions_spatial = torch.cat([torch.zeros_like(future_egomotions_spatial[:, :1]),
                                                   future_egomotions_spatial[:, :(self.receptive_field-1)]], dim=1)
            x = torch.cat([x, future_egomotions_spatial], dim=-3)

        #  Temporal model
        states = self.temporal_model(x)

        if self.n_future > 0:
            present_state = states[:, -1:].contiguous()

            b, _, c, h, w = present_state.shape

            if self.cfg.PROBABILISTIC.ENABLED:
                sample = self.distribution_forward(
                    present_state,
                    min_log_sigma=self.cfg.MODEL.DISTRIBUTION.MIN_LOG_SIGMA,
                    max_log_sigma=self.cfg.MODEL.DISTRIBUTION.MAX_LOG_SIGMA,
                )
                future_prediction_input = sample
            else:
                future_prediction_input = present_state.new_zeros(b, 1, self.latent_dim, h, w)

            # predict the future
            states = self.future_prediction(future_prediction_input, states)

            # predict BEV outputs
            bev_output = self.decoder(states)

        else:
            # Perceive BEV outputs
            bev_output = self.decoder(states)

        output = {**output, **bev_output}

        return output

    def get_geometry(self, intrinsics, extrinsics):
        """Calculate the (x, y, z) 3D position of the features.
        """
        rotation, translation = extrinsics[..., :3, :3], extrinsics[..., :3, 3]
        B, N, _ = translation.shape
        # Add batch, camera dimension, and a dummy dimension at the end
        points = self.frustum.unsqueeze(0).unsqueeze(0).unsqueeze(-1).to(intrinsics.device)

        # Camera to ego reference frame
        points = torch.cat((points[:, :, :, :, :, :2] * points[:, :, :, :, :, 2:3], points[:, :, :, :, :, 2:3]), 5)
        combined_transformation = rotation.matmul(torch.inverse(intrinsics))
        points = combined_transformation.view(B, N, 1, 1, 1, 3, 3).matmul(points).squeeze(-1)
        points += translation.view(B, N, 1, 1, 1, 3)

        # The 3 dimensions in the ego reference frame are: (forward, sides, height)
        return points

    def encoder_forward(self, x, cam_front_index=1):
        # batch, n_cameras, channels, height, width
        b, n, c, h, w = x.shape

        x = x.view(b * n, c, h, w) # (9 * 6, 3, 224, 480)
        x, depth = self.encoder(x) # (9 * 6, 64, 28, 60)
        if self.cfg.PLANNING.ENABLED:
            cam_front = x.view(b, n, *x.shape[1:])[:, cam_front_index]
        else:
            cam_front = None

        if self.cfg.MODEL.ENCODER.USE_DEPTH_DISTRIBUTION:
            depth_prob = depth.softmax(dim=1)
            x = depth_prob.unsqueeze(1) * x.unsqueeze(2)  # outer product depth and features
        else:
            x = x.unsqueeze(2).repeat(1, 1, self.depth_channels, 1, 1)

        x = x.view(b, n, *x.shape[1:])
        x = x.permute(0, 1, 3, 4, 5, 2) # channel dimension
        depth = depth.view(b, n, *depth.shape[1:])

        return x, depth, cam_front

    def projection_to_birds_eye_view(self, x, geometry, future_egomotion):
        """ Adapted from https://github.com/nv-tlabs/lift-splat-shoot/blob/master/src/models.py#L200"""
        # batch, s, n_cameras, depth, height, width, channels
        batch, s, n, d, h, w, c = x.shape
        output = torch.zeros(
            (batch, s, c, self.bev_dimension[0], self.bev_dimension[1]), dtype=torch.float, device=x.device
        )

        future_egomotion_mat = pose_vec2mat(future_egomotion)  # (3,3,4,4)
        rotation, translation = future_egomotion_mat[..., :3, :3], future_egomotion_mat[..., :3, 3]

        def voxel_to_pixel(geometry_b, x_b):
            # Mask out points that are outside the considered spatial extent.
            mask = (
                    (geometry_b[:, 0] >= 0)
                    & (geometry_b[:, 0] < self.bev_dimension[0])
                    & (geometry_b[:, 1] >= 0)
                    & (geometry_b[:, 1] < self.bev_dimension[1])
                    & (geometry_b[:, 2] >= 0)
                    & (geometry_b[:, 2] < self.bev_dimension[2])
            )
            x_b = x_b[mask]
            geometry_b = geometry_b[mask]

            # Sort tensors so that those within the same voxel are consecutives.
            ranks = (
                    geometry_b[:, 0] * (self.bev_dimension[1] * self.bev_dimension[2])
                    + geometry_b[:, 1] * (self.bev_dimension[2])
                    + geometry_b[:, 2]
            )
            ranks_indices = ranks.argsort()
            x_b, geometry_b, ranks = x_b[ranks_indices], geometry_b[ranks_indices], ranks[ranks_indices]

            # Project to bird's-eye view by summing voxels.
            x_b, geometry_b = VoxelsSumming.apply(x_b, geometry_b, ranks)
            return geometry_b, x_b

        # Number of 3D points
        N = n * d * h * w
        for b in range(batch):
            flow_b = x[b]
            flow_geo = geometry[b]

            #####  transform the 3D voxel to current frame  #####
            for t in range(s):
                if t != s - 1:
                    flow_geo_tmp = flow_geo[:t + 1]
                    rotation_b = rotation[b, t].view(1, 1, 1, 1, 1, 3, 3)
                    translation_b = translation[b, t].view(1, 1, 1, 1, 1, 3)
                    flow_geo_tmp = rotation_b.matmul(flow_geo_tmp.unsqueeze(-1)).squeeze(-1)
                    flow_geo_tmp += translation_b
                    flow_geo[:t + 1] = flow_geo_tmp

            bev_feature = torch.zeros((self.bev_dimension[2], self.bev_dimension[0], self.bev_dimension[1], c),
                                      device=flow_b.device)

            for t in range(s):
                # flatten x
                x_b = flow_b[t].reshape(N, c)

                # Convert positions to integer indices
                geometry_b = (
                        (flow_geo[t] - (self.bev_start_position - self.bev_resolution / 2.0)) / self.bev_resolution)
                geometry_b = geometry_b.view(N, 3).long()
                geometry_b, x_b = voxel_to_pixel(geometry_b, x_b)

                tmp_bev_feature = torch.zeros((self.bev_dimension[2], self.bev_dimension[0], self.bev_dimension[1], c),
                                              device=flow_b.device)
                tmp_bev_feature[geometry_b[:, 2], geometry_b[:, 0], geometry_b[:, 1]] = x_b

                bev_feature = bev_feature * self.discount + tmp_bev_feature
                tmp_bev_feature = bev_feature.permute((0, 3, 1, 2))
                tmp_bev_feature = tmp_bev_feature.squeeze(0)
                output[b, t] = tmp_bev_feature

        return output

    def calculate_birds_eye_view_features(self, x, intrinsics, extrinsics, future_egomotion):
        print(x.shape)
        b, s, n, c, h, w = x.shape
        # Reshape
        x = pack_sequence_dim(x)
        intrinsics = pack_sequence_dim(intrinsics)
        extrinsics = pack_sequence_dim(extrinsics)

        geometry = self.get_geometry(intrinsics, extrinsics)
        x, depth, cam_front = self.encoder_forward(x)
        x = unpack_sequence_dim(x, b, s)
        geometry = unpack_sequence_dim(geometry, b, s)
        depth = unpack_sequence_dim(depth, b, s)
        cam_front = unpack_sequence_dim(cam_front, b, s)[:,-1] if cam_front is not None else None

        x = self.projection_to_birds_eye_view(x, geometry, future_egomotion)
        return x, depth, cam_front

    def distribution_forward(self, present_features, min_log_sigma, max_log_sigma):
        """
        Parameters
        ----------
            present_features: 5-D output from dynamics module with shape (b, 1, c, h, w)

        Returns
        -------
            sample: sample taken from present/future distribution, broadcast to shape (b, s, latent_dim, h, w)
        """
        b, s, _, h, w = present_features.size()
        assert s == 1

        def get_mu_sigma(mu_log_sigma):
            mu = mu_log_sigma[:, :, :self.latent_dim]
            log_sigma = mu_log_sigma[:, :, self.latent_dim:2*self.latent_dim]
            log_sigma = torch.clamp(log_sigma, min_log_sigma, max_log_sigma)
            sigma = torch.exp(log_sigma)
            if self.training:
                gaussian_noise = torch.randn((b, s, self.latent_dim), device=present_features.device)
            else:
                gaussian_noise = torch.zeros((b, s, self.latent_dim), device=present_features.device)
            sample = mu + sigma * gaussian_noise
            return mu, log_sigma, sample


        if self.cfg.PROBABILISTIC.METHOD == 'GAUSSIAN':
            mu_log_sigma = self.present_distribution(present_features)
            present_mu, present_log_sigma, present_sample = get_mu_sigma(mu_log_sigma)
            sample = present_sample

            # Spatially broadcast sample to the dimensions of present_features
            sample = sample.view(b, s, self.latent_dim, 1, 1).expand(b, s, self.latent_dim, h, w)

        elif self.cfg.PROBABILISTIC.METHOD == "BERNOULLI":
            present_log_prob = self.present_distribution(present_features)
            if self.training:
                bernoulli_noise = torch.randn((b, self.latent_dim, h, w), device=present_features.device)
            else:
                bernoulli_noise = torch.zeros((b, self.latent_dim, h, w), device=present_features.device)
            sample = torch.exp(present_log_prob) + bernoulli_noise

            sample = sample.view(b, s, self.latent_dim, h, w)


        elif self.cfg.PROBABILISTIC.METHOD == 'MIXGAUSSIAN':
            mu_log_sigma = self.present_distribution(present_features)
            present_mu1, present_log_sigma1, present_sample1 = get_mu_sigma(mu_log_sigma[:, :, :2*self.latent_dim])
            present_mu2, present_log_sigma2, present_sample2 = get_mu_sigma(mu_log_sigma[:, :, 2 * self.latent_dim : 4 * self.latent_dim])
            present_mu3, present_log_sigma3, present_sample3 = get_mu_sigma(mu_log_sigma[:, :, 4 * self.latent_dim : 6 * self.latent_dim])
            coefficient = mu_log_sigma[:, :, 6 * self.latent_dim:]
            coefficient = torch.softmax(coefficient, dim=-1)
            sample = present_sample1 * coefficient[:,:,0:1] + \
                     present_sample2 * coefficient[:,:,1:2] + \
                     present_sample3 * coefficient[:,:,2:3]

            # Spatially broadcast sample to the dimensions of present_features
            sample = sample.view(b, s, self.latent_dim, 1, 1).expand(b, s, self.latent_dim, h, w)

        else:
            raise NotImplementedError

        return sample

    def select_best_traj(self, trajs, cost_volume, lane_divider, semantic_pred, k=1):
        '''
        trajs: torch.Tensor (B, N, n_future, 3)
        fuser_feature: torch.Tensor (B, n_present, fuser_channel, 200, 200)
        hd_map: torch.Tensor(B, 5, 200, 200)
        semantic_pred: torch.Tensor(B, n_future, 200, 200)
        '''
        sm_cost_fc, sm_cost_fo = self.cost_function(cost_volume, trajs[:,:,:,:2], semantic_pred, lane_divider)

        CS = sm_cost_fc + sm_cost_fo.sum(dim=-1)
        CC, KK = torch.topk(CS, k, dim=-1, largest=False)

        ii = torch.arange(len(trajs))
        select_traj = trajs[ii[:,None], KK].squeeze(1) # (B, n_future, 3)

        return select_traj
data/stp3/models/stp3.py
llmbev/talk2bev
import torch
import torch.nn as nn

from stp3.layers.temporal import Bottleneck3D, TemporalBlock
from stp3.layers.convolutions import ConvBlock, Bottleneck, DeepLabHead

class TemporalModel(nn.Module):
    def __init__(
            self, in_channels, receptive_field, input_shape, start_out_channels=64, extra_in_channels=0,
            n_spatial_layers_between_temporal_layers=0, use_pyramid_pooling=True):
        super().__init__()
        self.receptive_field = receptive_field
        n_temporal_layers = receptive_field - 1

        h, w = input_shape
        modules = []

        block_in_channels = in_channels
        block_out_channels = start_out_channels

        for _ in range(n_temporal_layers):
            if use_pyramid_pooling:
                use_pyramid_pooling = True
                pool_sizes = [(2, h, w)]
            else:
                use_pyramid_pooling = False
                pool_sizes = None
            temporal = TemporalBlock(
                block_in_channels,
                block_out_channels,
                use_pyramid_pooling=use_pyramid_pooling,
                pool_sizes=pool_sizes,
            )
            spatial = [
                Bottleneck3D(block_out_channels, block_out_channels, kernel_size=(1, 3, 3))
                for _ in range(n_spatial_layers_between_temporal_layers)
            ]
            temporal_spatial_layers = nn.Sequential(temporal, *spatial)
            modules.extend(temporal_spatial_layers)

            block_in_channels = block_out_channels
            block_out_channels += extra_in_channels

        self.out_channels = block_in_channels

        self.final_conv = DeepLabHead(block_out_channels, block_out_channels, hidden_channel=128)

        self.model = nn.Sequential(*modules)

    def forward(self, x):
        # Reshape input tensor to (batch, C, time, H, W)
        x = x.permute(0, 2, 1, 3, 4)
        x = self.model(x)
        x = x.permute(0, 2, 1, 3, 4).contiguous()

        b, s, c, h, w = x.shape
        x = x.view(b * s, c, h, w)
        x = self.final_conv(x)
        x = x.view(b, s, c, h, w)
        return x


class TemporalModelIdentity(nn.Module):
    def __init__(self, in_channels, receptive_field):
        super().__init__()
        self.receptive_field = receptive_field
        self.out_channels = in_channels

    def forward(self, x):
        return x
data/stp3/models/temporal_model.py
llmbev/talk2bev
import torch
import torch.nn as nn
import pytorch_lightning as pl

from stp3.config import get_cfg
from stp3.models.stp3 import STP3
from stp3.losses import SpatialRegressionLoss, SegmentationLoss, HDmapLoss, DepthLoss
from stp3.metrics import IntersectionOverUnion, PanopticMetric, PlanningMetric
from stp3.utils.geometry import cumulative_warp_features_reverse, cumulative_warp_features
from stp3.utils.instance import predict_instance_segmentation_and_trajectories
from stp3.utils.visualisation import visualise_output


class TrainingModule(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()

        # see config.py for details
        self.hparams = hparams
        # pytorch lightning does not support saving YACS CfgNone
        cfg = get_cfg(cfg_dict=self.hparams)
        self.cfg = cfg
        self.n_classes = len(self.cfg.SEMANTIC_SEG.VEHICLE.WEIGHTS)
        self.hdmap_class = cfg.SEMANTIC_SEG.HDMAP.ELEMENTS

        # Bird's-eye view extent in meters
        assert self.cfg.LIFT.X_BOUND[1] > 0 and self.cfg.LIFT.Y_BOUND[1] > 0
        self.spatial_extent = (self.cfg.LIFT.X_BOUND[1], self.cfg.LIFT.Y_BOUND[1])

        # Model
        self.model = STP3(cfg)

        self.losses_fn = nn.ModuleDict()

        # Semantic segmentation
        self.losses_fn['segmentation'] = SegmentationLoss(
            class_weights=torch.Tensor(self.cfg.SEMANTIC_SEG.VEHICLE.WEIGHTS),
            use_top_k=self.cfg.SEMANTIC_SEG.VEHICLE.USE_TOP_K,
            top_k_ratio=self.cfg.SEMANTIC_SEG.VEHICLE.TOP_K_RATIO,
            future_discount=self.cfg.FUTURE_DISCOUNT,
        )
        self.model.segmentation_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)
        self.metric_vehicle_val = IntersectionOverUnion(self.n_classes)

        # Pedestrian segmentation
        if self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
            self.losses_fn['pedestrian'] = SegmentationLoss(
                class_weights=torch.Tensor(self.cfg.SEMANTIC_SEG.PEDESTRIAN.WEIGHTS),
                use_top_k=self.cfg.SEMANTIC_SEG.PEDESTRIAN.USE_TOP_K,
                top_k_ratio=self.cfg.SEMANTIC_SEG.PEDESTRIAN.TOP_K_RATIO,
                future_discount=self.cfg.FUTURE_DISCOUNT,
            )
            self.model.pedestrian_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)
            self.metric_pedestrian_val = IntersectionOverUnion(self.n_classes)

        # HD map
        if self.cfg.SEMANTIC_SEG.HDMAP.ENABLED:
            self.losses_fn['hdmap'] = HDmapLoss(
                class_weights=torch.Tensor(self.cfg.SEMANTIC_SEG.HDMAP.WEIGHTS),
                training_weights=self.cfg.SEMANTIC_SEG.HDMAP.TRAIN_WEIGHT,
                use_top_k=self.cfg.SEMANTIC_SEG.HDMAP.USE_TOP_K,
                top_k_ratio=self.cfg.SEMANTIC_SEG.HDMAP.TOP_K_RATIO,
            )
            self.metric_hdmap_val = []
            for i in range(len(self.hdmap_class)):
                self.metric_hdmap_val.append(IntersectionOverUnion(2, absent_score=1))
            self.model.hdmap_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)
            self.metric_hdmap_val = nn.ModuleList(self.metric_hdmap_val)

        # Depth
        if self.cfg.LIFT.GT_DEPTH:
            self.losses_fn['depths'] = DepthLoss()
            self.model.depths_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)

        # Instance segmentation
        if self.cfg.INSTANCE_SEG.ENABLED:
            self.losses_fn['instance_center'] = SpatialRegressionLoss(
                norm=2, future_discount=self.cfg.FUTURE_DISCOUNT
            )
            self.losses_fn['instance_offset'] = SpatialRegressionLoss(
                norm=1, future_discount=self.cfg.FUTURE_DISCOUNT, ignore_index=self.cfg.DATASET.IGNORE_INDEX
            )
            self.model.centerness_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)
            self.model.offset_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)
            self.metric_panoptic_val = PanopticMetric(n_classes=self.n_classes)

        # Instance flow
        if self.cfg.INSTANCE_FLOW.ENABLED:
            self.losses_fn['instance_flow'] = SpatialRegressionLoss(
                norm=1, future_discount=self.cfg.FUTURE_DISCOUNT, ignore_index=self.cfg.DATASET.IGNORE_INDEX
            )
            self.model.flow_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)

        # Planning
        if self.cfg.PLANNING.ENABLED:
            self.metric_planning_val = PlanningMetric(self.cfg, self.cfg.N_FUTURE_FRAMES)
            self.model.planning_weight = nn.Parameter(torch.tensor(0.0), requires_grad=True)

        self.training_step_count = 0

    def shared_step(self, batch, is_train):
        image = batch['image']
        intrinsics = batch['intrinsics']
        extrinsics = batch['extrinsics']
        future_egomotion = batch['future_egomotion']
        command = batch['command']
        trajs = batch['sample_trajectory']
        target_points = batch['target_point']
        B = len(image)

        # Warp labels
        labels = self.prepare_future_labels(batch)

        # Forward pass
        output = self.model(
            image, intrinsics, extrinsics, future_egomotion,
        )

        #####
        # Loss computation
        #####
        loss = {}
        if is_train:
            # segmentation
            segmentation_factor = 1 / (2 * torch.exp(self.model.segmentation_weight))
            loss['segmentation'] = segmentation_factor * self.losses_fn['segmentation'](
                output['segmentation'], labels['segmentation'], self.model.receptive_field
            )
            loss['segmentation_uncertainty'] = 0.5 * self.model.segmentation_weight

            # Pedestrian
            if self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
                pedestrian_factor = 1 / (2 * torch.exp(self.model.pedestrian_weight))
                loss['pedestrian'] = pedestrian_factor * self.losses_fn['pedestrian'](
                    output['pedestrian'], labels['pedestrian'], self.model.receptive_field
                )
                loss['pedestrian_uncertainty'] = 0.5 * self.model.pedestrian_weight

            # hdmap loss
            if self.cfg.SEMANTIC_SEG.HDMAP.ENABLED:
                hdmap_factor = 1 / (2 * torch.exp(self.model.hdmap_weight))
                loss['hdmap'] = hdmap_factor * self.losses_fn['hdmap'](output['hdmap'], labels['hdmap'])
                loss['hdmap_uncertainty'] = 0.5 * self.model.hdmap_weight

            if self.cfg.INSTANCE_SEG.ENABLED:
                # instance center
                centerness_factor = 1 / (2 * torch.exp(self.model.centerness_weight))
                loss['instance_center'] = centerness_factor * self.losses_fn['instance_center'](
                    output['instance_center'], labels['centerness'], self.model.receptive_field
                )
                loss['centerness_uncertainty'] = 0.5 * self.model.centerness_weight

                # instance offset
                offset_factor = 1 / (2 * torch.exp(self.model.offset_weight))
                loss['instance_offset'] = offset_factor * self.losses_fn['instance_offset'](
                    output['instance_offset'], labels['offset'], self.model.receptive_field
                )
                loss['offset_uncertainty'] = 0.5 * self.model.offset_weight

            # depth loss
            if self.cfg.LIFT.GT_DEPTH:
                depths_factor = 1 / (2 * torch.exp(self.model.depths_weight))
                loss['depths'] = depths_factor * self.losses_fn['depths'](output['depth_prediction'], labels['depths'])
                loss['depths_uncertainty'] = 0.5 * self.model.depths_weight

            # instance flow
            if self.cfg.INSTANCE_FLOW.ENABLED:
                flow_factor = 1 / (2 * torch.exp(self.model.flow_weight))
                loss['instance_flow'] = flow_factor * self.losses_fn['instance_flow'](
                    output['instance_flow'], labels['flow'], self.model.receptive_field
                )
                loss['flow_uncertainty'] = 0.5 * self.model.flow_weight

            # Planning
            if self.cfg.PLANNING.ENABLED:
                receptive_field = self.model.receptive_field
                planning_factor = 1 / (2 * torch.exp(self.model.planning_weight))
                occupancy = torch.logical_or(labels['segmentation'][:, receptive_field:].squeeze(2),
                                             labels['pedestrian'][:, receptive_field:].squeeze(2))
                pl_loss, final_traj, CS, cur_trajs, gt_cost_fc, gt_cost_fo, final_traj, sm_cost_fc, sm_cost_fo = self.model.planning(
                    cam_front=output['cam_front'].detach(),
                    trajs=trajs[:, :, 1:],
                    gt_trajs=labels['gt_trajectory'][:, 1:],
                    cost_volume=output['costvolume'][:, receptive_field:],
                    semantic_pred=occupancy,
                    hd_map=labels['hdmap'],
                    commands=command,
                    target_points=target_points
                )
                loss['planning'] = planning_factor * pl_loss
                loss['planning_uncertainty'] = 0.5 * self.model.planning_weight
                output = {**output, 'selected_traj': torch.cat(
                    [torch.zeros((B, 1, 3), device=final_traj.device), final_traj], dim=1)}
            else:
                output = {**output, 'selected_traj': labels['gt_trajectory']}

        # Metrics
        else:
            n_present = self.model.receptive_field

            # semantic segmentation metric
            seg_prediction = output['segmentation'].detach()
            seg_prediction = torch.argmax(seg_prediction, dim=2, keepdim=True)
            self.metric_vehicle_val(seg_prediction[:, n_present - 1:], labels['segmentation'][:, n_present - 1:])

            # pedestrian segmentation metric
            if self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
                pedestrian_prediction = output['pedestrian'].detach()
                pedestrian_prediction = torch.argmax(pedestrian_prediction, dim=2, keepdim=True)
                self.metric_pedestrian_val(pedestrian_prediction[:, n_present - 1:],
                                           labels['pedestrian'][:, n_present - 1:])
            else:
                pedestrian_prediction = torch.zeros_like(seg_prediction)

            # hdmap metric
            if self.cfg.SEMANTIC_SEG.HDMAP.ENABLED:
                for i in range(len(self.hdmap_class)):
                    hdmap_prediction = output['hdmap'][:, 2 * i:2 * (i + 1)].detach()
                    hdmap_prediction = torch.argmax(hdmap_prediction, dim=1, keepdim=True)
                    self.metric_hdmap_val[i](hdmap_prediction, labels['hdmap'][:, i:i + 1])

            # instance segmentation metric
            if self.cfg.INSTANCE_SEG.ENABLED:
                pred_consistent_instance_seg = predict_instance_segmentation_and_trajectories(
                    output, compute_matched_centers=False
                )
                self.metric_panoptic_val(pred_consistent_instance_seg[:, n_present - 1:],
                                         labels['instance'][:, n_present - 1:])

            # planning metric
            if self.cfg.PLANNING.ENABLED:
                occupancy = torch.logical_or(seg_prediction, pedestrian_prediction)
                _, final_traj, CS, cur_trajs, gt_cost_fc, gt_cost_fo, final_traj, sm_cost_fc, sm_cost_fo = self.model.planning(
                    cam_front=output['cam_front'].detach(),
                    trajs=trajs[:, :, 1:],
                    gt_trajs=labels['gt_trajectory'][:, 1:],
                    cost_volume=output['costvolume'][:, n_present:].detach(),
                    semantic_pred=occupancy[:, n_present:].squeeze(2),
                    hd_map=output['hdmap'].detach(),
                    commands=command,
                    target_points=target_points
                )
                occupancy = torch.logical_or(labels['segmentation'][:, n_present:].squeeze(2),
                                             labels['pedestrian'][:, n_present:].squeeze(2))
                self.metric_planning_val(final_traj, labels['gt_trajectory'][:, 1:], occupancy)
                output = {**output,
                          'selected_traj': torch.cat([torch.zeros((B, 1, 3), device=final_traj.device), final_traj],
                                                     dim=1)}
            else:
                output = {**output, 'selected_traj': labels['gt_trajectory']}

        return output, labels, loss

    def prepare_future_labels(self, batch):
        labels = {}

        segmentation_labels = batch['segmentation']
        hdmap_labels = batch['hdmap']
        future_egomotion = batch['future_egomotion']
        gt_trajectory = batch['gt_trajectory']

        # present frame hd map gt
        labels['hdmap'] = hdmap_labels[:, self.model.receptive_field - 1].long().contiguous()

        # gt trajectory
        labels['gt_trajectory'] = gt_trajectory
        # first rec token
        labels['rec_first'] = batch['rec_first']

        # gt trajectory
        labels['gt_trajectory_prev'] = batch['gt_trajectory_prev']


        # Past frames gt depth
        if self.cfg.LIFT.GT_DEPTH:
            depths = batch['depths']
            depth_labels = depths[:, :self.model.receptive_field, :, ::self.model.encoder_downsample,
                           ::self.model.encoder_downsample]
            depth_labels = torch.clamp(depth_labels, self.cfg.LIFT.D_BOUND[0], self.cfg.LIFT.D_BOUND[1] - 1) - \
                           self.cfg.LIFT.D_BOUND[0]
            depth_labels = depth_labels.long().contiguous()
            labels['depths'] = depth_labels

        # Warp labels to present's reference frame
        segmentation_labels_past = cumulative_warp_features(
            segmentation_labels[:, :self.model.receptive_field].float(),
            future_egomotion[:, :self.model.receptive_field],
            mode='nearest', spatial_extent=self.spatial_extent,
        ).long().contiguous()[:, :-1]
        segmentation_labels = cumulative_warp_features_reverse(
            segmentation_labels[:, (self.model.receptive_field - 1):].float(),
            future_egomotion[:, (self.model.receptive_field - 1):],
            mode='nearest', spatial_extent=self.spatial_extent,
        ).long().contiguous()
        labels['segmentation'] = torch.cat([segmentation_labels_past, segmentation_labels], dim=1)

        if self.cfg.SEMANTIC_SEG.HDMAP.ENABLED:
            hdmap_labels = batch['hdmap'][:, :, 1:2]
            hdmap_labels_past = cumulative_warp_features(
                hdmap_labels[:, :self.model.receptive_field].float(),
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()[:, :-1]
            hdmap_labels = cumulative_warp_features_reverse(
                hdmap_labels[:, (self.model.receptive_field - 1):].float(),
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()
            labels['hdmap_warped_road'] = torch.cat([hdmap_labels_past, hdmap_labels], dim=1)
            hdmap_labels = batch['hdmap'][:, :, 0:1]
            hdmap_labels_past = cumulative_warp_features(
                hdmap_labels[:, :self.model.receptive_field].float(),
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()[:, :-1]
            hdmap_labels = cumulative_warp_features_reverse(
                hdmap_labels[:, (self.model.receptive_field - 1):].float(),
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()
            labels['hdmap_warped_lane'] = torch.cat([hdmap_labels_past, hdmap_labels], dim=1)            
            pass
        if self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
            pedestrian_labels = batch['pedestrian']
            pedestrian_labels_past = cumulative_warp_features(
                pedestrian_labels[:, :self.model.receptive_field].float(),
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()[:, :-1]
            pedestrian_labels = cumulative_warp_features_reverse(
                pedestrian_labels[:, (self.model.receptive_field - 1):].float(),
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()
            labels['pedestrian'] = torch.cat([pedestrian_labels_past, pedestrian_labels], dim=1)

        # Warp instance labels to present's reference frame
        if self.cfg.INSTANCE_SEG.ENABLED:
            gt_instance = batch['instance']
            instance_center_labels = batch['centerness']
            instance_offset_labels = batch['offset']
            gt_instance_past = cumulative_warp_features(
                gt_instance[:, :self.model.receptive_field].float().unsqueeze(2),
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()[:, :-1, 0]
            gt_instance = cumulative_warp_features_reverse(
                gt_instance[:, (self.model.receptive_field - 1):].float().unsqueeze(2),
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).long().contiguous()[:, :, 0]
            labels['instance'] = torch.cat([gt_instance_past, gt_instance], dim=1)

            instance_center_labels_past = cumulative_warp_features(
                instance_center_labels[:, :self.model.receptive_field],
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).contiguous()[:, :-1]
            instance_center_labels = cumulative_warp_features_reverse(
                instance_center_labels[:, (self.model.receptive_field - 1):],
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).contiguous()
            labels['centerness'] = torch.cat([instance_center_labels_past, instance_center_labels], dim=1)

            instance_offset_labels_past = cumulative_warp_features(
                instance_offset_labels[:, :self.model.receptive_field],
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).contiguous()[:, :-1]
            instance_offset_labels = cumulative_warp_features_reverse(
                instance_offset_labels[:, (self.model.receptive_field - 1):],
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).contiguous()
            labels['offset'] = torch.cat([instance_offset_labels_past, instance_offset_labels], dim=1)

        if self.cfg.INSTANCE_FLOW.ENABLED:
            instance_flow_labels = batch['flow']
            instance_flow_labels_past = cumulative_warp_features(
                instance_flow_labels[:, :self.model.receptive_field],
                future_egomotion[:, :self.model.receptive_field],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).contiguous()[:, :-1]
            instance_flow_labels = cumulative_warp_features_reverse(
                instance_flow_labels[:, (self.model.receptive_field - 1):],
                future_egomotion[:, (self.model.receptive_field - 1):],
                mode='nearest', spatial_extent=self.spatial_extent,
            ).contiguous()
            labels['flow'] = torch.cat([instance_flow_labels_past, instance_flow_labels], dim=1)

        return labels

    def visualise(self, labels, output, batch_idx, prefix='train'):
        visualisation_video = visualise_output(labels, output, self.cfg)
        name = f'{prefix}_outputs'
        if prefix == 'val':
            name = name + f'_{batch_idx}'
        self.logger.experiment.add_video(name, visualisation_video, global_step=self.training_step_count, fps=2)

    def training_step(self, batch, batch_idx):
        output, labels, loss = self.shared_step(batch, True)
        self.training_step_count += 1
        for key, value in loss.items():
            self.logger.experiment.add_scalar('step_train_loss_' + key, value, global_step=self.training_step_count)
        if self.training_step_count % self.cfg.VIS_INTERVAL == 0:
            self.visualise(labels, output, batch_idx, prefix='train')
        return sum(loss.values())

    def validation_step(self, batch, batch_idx):
        output, labels, loss = self.shared_step(batch, False)
        scores = self.metric_vehicle_val.compute()
        self.log('step_val_seg_iou_dynamic', scores[1])
        self.log('step_predicted_traj_x', output['selected_traj'][0, -1, 0])
        self.log('step_target_traj_x', labels['gt_trajectory'][0, -1, 0])
        self.log('step_predicted_traj_y', output['selected_traj'][0, -1, 1])
        self.log('step_target_traj_y', labels['gt_trajectory'][0, -1, 1])

        if batch_idx == 0:
            self.visualise(labels, output, batch_idx, prefix='val')

    def shared_epoch_end(self, step_outputs, is_train):
        if not is_train:
            scores = self.metric_vehicle_val.compute()
            self.logger.experiment.add_scalar('epoch_val_all_seg_iou_dynamic', scores[1],
                                              global_step=self.training_step_count)
            self.metric_vehicle_val.reset()

            if self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
                scores = self.metric_pedestrian_val.compute()
                self.logger.experiment.add_scalar('epoch_val_all_seg_iou_pedestrian', scores[1],
                                                  global_step=self.training_step_count)
                self.metric_pedestrian_val.reset()

            if self.cfg.SEMANTIC_SEG.HDMAP.ENABLED:
                for i, name in enumerate(self.hdmap_class):
                    scores = self.metric_hdmap_val[i].compute()
                    self.logger.experiment.add_scalar('epoch_val_hdmap_iou_' + name, scores[1],
                                                      global_step=self.training_step_count)
                    self.metric_hdmap_val[i].reset()

            if self.cfg.INSTANCE_SEG.ENABLED:
                scores = self.metric_panoptic_val.compute()
                for key, value in scores.items():
                    self.logger.experiment.add_scalar(f'epoch_val_all_ins_{key}_vehicle', value[1].item(),
                                                      global_step=self.training_step_count)
                self.metric_panoptic_val.reset()

            if self.cfg.PLANNING.ENABLED:
                scores = self.metric_planning_val.compute()
                for key, value in scores.items():
                    self.logger.experiment.add_scalar('epoch_val_plan_' + key, value.mean(),
                                                      global_step=self.training_step_count)
                self.metric_planning_val.reset()

        self.logger.experiment.add_scalar('epoch_segmentation_weight',
                                          1 / (2 * torch.exp(self.model.segmentation_weight)),
                                          global_step=self.training_step_count)
        if self.cfg.LIFT.GT_DEPTH:
            self.logger.experiment.add_scalar('epoch_depths_weight', 1 / (2 * torch.exp(self.model.depths_weight)),
                                              global_step=self.training_step_count)
        if self.cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
            self.logger.experiment.add_scalar('epoch_pedestrian_weight',
                                              1 / (2 * torch.exp(self.model.pedestrian_weight)),
                                              global_step=self.training_step_count)
        if self.cfg.SEMANTIC_SEG.HDMAP.ENABLED:
            self.logger.experiment.add_scalar('epoch_hdmap_weight', 1 / (2 * torch.exp(self.model.hdmap_weight)),
                                              global_step=self.training_step_count)
        if self.cfg.INSTANCE_SEG.ENABLED:
            self.logger.experiment.add_scalar('epoch_centerness_weight',
                                              1 / (2 * torch.exp(self.model.centerness_weight)),
                                              global_step=self.training_step_count)
            self.logger.experiment.add_scalar('epoch_offset_weight', 1 / (2 * torch.exp(self.model.offset_weight)),
                                              global_step=self.training_step_count)
        if self.cfg.INSTANCE_FLOW.ENABLED:
            self.logger.experiment.add_scalar('epoch_flow_weight', 1 / (2 * torch.exp(self.model.flow_weight)),
                                              global_step=self.training_step_count)
        if self.cfg.PLANNING.ENABLED:
            self.logger.experiment.add_scalar('epoch_planning_weight', 1 / (2 * torch.exp(self.model.planning_weight)),
                                              global_step=self.training_step_count)

    def training_epoch_end(self, step_outputs):
        self.shared_epoch_end(step_outputs, True)

    def validation_epoch_end(self, step_outputs):
        self.shared_epoch_end(step_outputs, False)

    def configure_optimizers(self):
        params = self.model.parameters()
        optimizer = torch.optim.Adam(
            params, lr=self.cfg.OPTIMIZER.LR, weight_decay=self.cfg.OPTIMIZER.WEIGHT_DECAY
        )

        return optimizer

data/stp3/trainer.py
llmbev/talk2bev
import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader

import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader
#from bernstein import bernstesin_coeff_order10_new

def bernstein_coeff_order10_new(n, tmin, tmax, t_actual):
    l = tmax - tmin
    t = (t_actual - tmin) / l

    P0 = scipy.special.binom(n, 0) * ((1 - t) ** (n - 0)) * t ** 0
    P1 = scipy.special.binom(n, 1) * ((1 - t) ** (n - 1)) * t ** 1
    P2 = scipy.special.binom(n, 2) * ((1 - t) ** (n - 2)) * t ** 2
    P3 = scipy.special.binom(n, 3) * ((1 - t) ** (n - 3)) * t ** 3
    P4 = scipy.special.binom(n, 4) * ((1 - t) ** (n - 4)) * t ** 4
    P5 = scipy.special.binom(n, 5) * ((1 - t) ** (n - 5)) * t ** 5
    P6 = scipy.special.binom(n, 6) * ((1 - t) ** (n - 6)) * t ** 6
    P7 = scipy.special.binom(n, 7) * ((1 - t) ** (n - 7)) * t ** 7
    P8 = scipy.special.binom(n, 8) * ((1 - t) ** (n - 8)) * t ** 8
    P9 = scipy.special.binom(n, 9) * ((1 - t) ** (n - 9)) * t ** 9
    P10 = scipy.special.binom(n, 10) * ((1 - t) ** (n - 10)) * t ** 10

    P0dot = -10.0 * (-t + 1) ** 9
    P1dot = -90.0 * t * (-t + 1) ** 8 + 10.0 * (-t + 1) ** 9
    P2dot = -360.0 * t ** 2 * (-t + 1) ** 7 + 90.0 * t * (-t + 1) ** 8
    P3dot = -840.0 * t ** 3 * (-t + 1) ** 6 + 360.0 * t ** 2 * (-t + 1) ** 7
    P4dot = -1260.0 * t ** 4 * (-t + 1) ** 5 + 840.0 * t ** 3 * (-t + 1) ** 6
    P5dot = -1260.0 * t ** 5 * (-t + 1) ** 4 + 1260.0 * t ** 4 * (-t + 1) ** 5
    P6dot = -840.0 * t ** 6 * (-t + 1) ** 3 + 1260.0 * t ** 5 * (-t + 1) ** 4
    P7dot = -360.0 * t ** 7 * (-t + 1) ** 2 + 840.0 * t ** 6 * (-t + 1) ** 3
    P8dot = 45.0 * t ** 8 * (2 * t - 2) + 360.0 * t ** 7 * (-t + 1) ** 2
    P9dot = -10.0 * t ** 9 + 9 * t ** 8 * (-10.0 * t + 10.0)
    P10dot = 10.0 * t ** 9

    P0ddot = 90.0 * (-t + 1) ** 8
    P1ddot = 720.0 * t * (-t + 1) ** 7 - 180.0 * (-t + 1) ** 8
    P2ddot = 2520.0 * t ** 2 * (-t + 1) ** 6 - 1440.0 * t * (-t + 1) ** 7 + 90.0 * (-t + 1) ** 8
    P3ddot = 5040.0 * t ** 3 * (-t + 1) ** 5 - 5040.0 * t ** 2 * (-t + 1) ** 6 + 720.0 * t * (-t + 1) ** 7
    P4ddot = 6300.0 * t ** 4 * (-t + 1) ** 4 - 10080.0 * t ** 3 * (-t + 1) ** 5 + 2520.0 * t ** 2 * (-t + 1) ** 6
    P5ddot = 5040.0 * t ** 5 * (-t + 1) ** 3 - 12600.0 * t ** 4 * (-t + 1) ** 4 + 5040.0 * t ** 3 * (-t + 1) ** 5
    P6ddot = 2520.0 * t ** 6 * (-t + 1) ** 2 - 10080.0 * t ** 5 * (-t + 1) ** 3 + 6300.0 * t ** 4 * (-t + 1) ** 4
    P7ddot = -360.0 * t ** 7 * (2 * t - 2) - 5040.0 * t ** 6 * (-t + 1) ** 2 + 5040.0 * t ** 5 * (-t + 1) ** 3
    P8ddot = 90.0 * t ** 8 + 720.0 * t ** 7 * (2 * t - 2) + 2520.0 * t ** 6 * (-t + 1) ** 2
    P9ddot = -180.0 * t ** 8 + 72 * t ** 7 * (-10.0 * t + 10.0)
    P10ddot = 90.0 * t ** 8
    90.0 * t ** 8

    P = np.hstack((P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10))
    Pdot = np.hstack((P0dot, P1dot, P2dot, P3dot, P4dot, P5dot, P6dot, P7dot, P8dot, P9dot, P10dot)) / l
    Pddot = np.hstack((P0ddot, P1ddot, P2ddot, P3ddot, P4ddot, P5ddot, P6ddot, P7ddot, P8ddot, P9ddot, P10ddot)) / (l ** 2)
    return P, Pdot, Pddot


class OPTNode():
    def __init__(self, rho_eq=1.0, rho_goal=1.0, rho_nonhol=1.0, rho_psi=1.0, maxiter=5000, weight_smoothness=1.0, weight_smoothness_psi=1.0, t_fin=2.0, num=30, bernstein_coeff_order10_new=bernstein_coeff_order10_new, device="cpu"):
        super().__init__()
        self.rho_eq = rho_eq
        self.rho_goal = rho_goal
        self.rho_nonhol = rho_nonhol
        self.rho_psi = rho_psi
        self.maxiter = maxiter
        self.weight_smoothness = weight_smoothness
        self.weight_smoothness_psi = weight_smoothness_psi
        self.bernstein_coeff_order10_new = bernstein_coeff_order10_new
        
        self.device = device
        
        self.t_fin = t_fin
        self.num = num
        self.t = self.t_fin / self.num

        #self.num_batch = 10
        
        tot_time = np.linspace(0.0, self.t_fin, self.num)
        tot_time_copy = tot_time.reshape(self.num, 1)
        self.P, self.Pdot, self.Pddot = bernstein_coeff_order10_new(10, tot_time_copy[0], tot_time_copy[-1], tot_time_copy)
        self.nvar = np.shape(self.P)[1]
        
        self.cost_smoothness = self.weight_smoothness * np.dot(self.Pddot.T, self.Pddot)
        self.cost_smoothness_psi = self.weight_smoothness_psi * np.dot(self.Pddot.T, self.Pddot)
        self.lincost_smoothness_psi = np.zeros(self.nvar)

        self.A_eq = np.vstack((self.P[0], self.P[-1]))
        self.A_eq_psi = np.vstack((self.P[0], self.Pdot[0], self.P[-1]))
        
        self.P = torch.tensor(self.P, dtype=torch.double).to(device)
        self.Pdot = torch.tensor(self.Pdot, dtype=torch.double).to(device)
        self.Pddot = torch.tensor(self.Pddot, dtype=torch.double).to(device)
        self.A_eq = torch.tensor(self.A_eq, dtype=torch.double).to(device)        
        self.A_eq_psi = torch.tensor(self.A_eq_psi, dtype=torch.double).to(device)
        self.cost_smoothness = torch.tensor(self.cost_smoothness, dtype=torch.double).to(device)
        self.cost_smoothness_psi = torch.tensor(self.cost_smoothness_psi, dtype=torch.double).to(device)
        self.lincost_smoothness_psi = torch.tensor(self.lincost_smoothness_psi, dtype=torch.double).to(device)
    
        self.A_nonhol = self.Pdot
        self.A_psi = self.P
        
        self.lamda_x = None
        self.lamda_y = None
        self.lamda_psi = None
        
    def compute_x(self, v, psi, b_eq_x, b_eq_y):
        b_nonhol_x = v * torch.cos(psi)
        b_nonhol_y = v * torch.sin(psi)
    
        cost = self.cost_smoothness + self.rho_nonhol * torch.matmul(self.A_nonhol.T, self.A_nonhol) + self.rho_eq * torch.matmul(self.A_eq.T, self.A_eq)
        lincost_x = -self.lamda_x - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_x.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_x.T).T
        lincost_y = -self.lamda_y - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_y.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_y.T).T

        cost_inv = torch.linalg.inv(cost)

        sol_x = torch.matmul(-cost_inv, lincost_x.T).T
        sol_y = torch.matmul(-cost_inv, lincost_y.T).T

        x = torch.matmul(self.P, sol_x.T).T
        xdot = torch.matmul(self.Pdot, sol_x.T).T

        y = torch.matmul(self.P, sol_y.T).T
        ydot = torch.matmul(self.Pdot, sol_y.T).T
         
        return sol_x, sol_y, x, xdot, y, ydot
    
    def compute_psi(self, psi, lamda_psi, psi_temp, b_eq_psi):
        cost = self.cost_smoothness_psi + self.rho_psi * torch.matmul(self.A_psi.T, self.A_psi) + self.rho_eq * torch.matmul(self.A_eq_psi.T, self.A_eq_psi)
        lincost_psi = -self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, psi_temp.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, b_eq_psi.T).T

        cost_inv = torch.linalg.inv(cost)

        sol_psi = torch.matmul(-cost_inv, lincost_psi.T).T

        psi = torch.matmul(self.P, sol_psi.T).T

        res_psi = torch.matmul(self.A_psi, sol_psi.T).T - psi_temp
        res_eq_psi = torch.matmul(self.A_eq_psi, sol_psi.T).T - b_eq_psi

        self.lamda_psi = self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, res_psi.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, res_eq_psi.T).T

        return sol_psi, torch.linalg.norm(res_psi), torch.linalg.norm(res_eq_psi), psi

    
    def solve(self, fixed_params, variable_params):
        batch_size, _ = fixed_params.size()
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin = torch.chunk(variable_params, 3, dim=1)

        b_eq_x = torch.cat((x_init, x_fin), dim=1)
        b_eq_y = torch.cat((y_init, y_fin), dim=1)
        b_eq_psi = torch.cat((psi_init, psidot_init, psi_fin), dim=1)
        
        v = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * v_init
        psi = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * psi_init
        xdot = v * torch.cos(psi)
        ydot = v * torch.sin(psi)
        
        self.lamda_x = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_y = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_psi = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        
        res_psi_arr = []
        res_eq_psi_arr = []
        res_eq_arr = []
        res_nonhol_arr = []
        for i in range(0, self.maxiter):
            psi_temp = torch.atan2(ydot, xdot)
            c_psi, res_psi, res_eq_psi, psi = self.compute_psi(psi, self.lamda_psi, psi_temp, b_eq_psi)
            c_x, c_y, x, xdot, y, ydot = self.compute_x(v, psi, b_eq_x, b_eq_y)
            
            res_eq_psi_arr.append(res_eq_psi)
            res_psi_arr.append(res_psi)
            v = torch.sqrt(xdot ** 2 + ydot ** 2)
            #v[:, 0] = v_init[:, 0]

            res_eq_x = torch.matmul(self.A_eq, c_x.T).T - b_eq_x
            res_nonhol_x = xdot - v * torch.cos(psi)

            res_eq_y = torch.matmul(self.A_eq, c_y.T).T - b_eq_y
            res_nonhol_y = ydot - v * torch.sin(psi)

            res_eq_arr.append(torch.linalg.norm(torch.sqrt(res_eq_x**2 + res_eq_y**2)))
            res_nonhol_arr.append(torch.linalg.norm(torch.sqrt(res_nonhol_x**2 + res_nonhol_y**2)))
            
            self.lamda_x = self.lamda_x - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_x.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_x.T).T
            self.lamda_y = self.lamda_y - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_y.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_y.T).T
        
        primal_sol = torch.hstack((c_x, c_y, c_psi, v))
        return primal_sol, None
    
    def objective(self, fixed_params, variable_params, y):
        c_x = y[:, :self.nvar]
        c_y = y[:, self.nvar:2*self.nvar]
        c_psi = y[:, 2*self.nvar:3*self.nvar]
        v = y[:, 3*self.nvar:]
        
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin = torch.chunk(variable_params, 3, dim=1)
        
        x = torch.matmul(self.P, c_x.T).T
        y = torch.matmul(self.P, c_y.T).T
        psi = torch.matmul(self.P, c_psi.T).T
        xdot = torch.matmul(self.Pdot, c_x.T).T
        ydot = torch.matmul(self.Pdot, c_y.T).T
        psidot = torch.matmul(self.Pdot, c_psi.T).T
        xddot = torch.matmul(self.Pddot, c_x.T).T
        yddot = torch.matmul(self.Pddot, c_y.T).T
        psiddot = torch.matmul(self.Pddot, c_psi.T).T
        
        cost_nonhol = 0.5*self.rho_nonhol*torch.sum((xdot - v*torch.cos(psi)) ** 2, 1) + 0.5*self.rho_nonhol*torch.sum((ydot - v*torch.sin(psi)) ** 2, 1)
        cost_pos = 0.5*self.rho_eq*(torch.sum((x[:, -1] - x_fin) ** 2, 1) + torch.sum((y[:, -1] - y_fin) ** 2, 1) + torch.sum((x[:, 0] - x_init) ** 2, 1) + torch.sum((y[:, 0] - y_init) ** 2, 1))
        cost_psi = 0.5*self.rho_eq*(torch.sum((psi[:, -1] - psi_fin) ** 2, 1) + torch.sum((psi[:, 0] - psi_init) ** 2, 1)
                                    + torch.sum((psidot[:, 0] - psidot_init) ** 2, 1))
        #cost_v = 0.5*self.rho_eq*torch.sum((v[:, 0] - v_init) ** 2, 1)
        cost_cancel = torch.diagonal(torch.matmul(-self.lamda_x, c_x.T) + torch.matmul(-self.lamda_y, c_y.T) + torch.matmul(-self.lamda_psi, c_psi.T))
        
        cost_smoothness = 0.5*self.weight_smoothness*(torch.sum(xddot**2, 1) + torch.sum(yddot**2, 1)) + 0.5*self.weight_smoothness_psi*torch.sum(psiddot**2, 1)
        return cost_nonhol + cost_pos + cost_psi + cost_smoothness + cost_cancel #+ cost_v 

data/stp3/utils/Optnode.py
llmbev/talk2bev
import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader

import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader
#from bernstein import bernstesin_coeff_order10_new

def bernstein_coeff_order10_new(n, tmin, tmax, t_actual):
    l = tmax - tmin
    t = (t_actual - tmin) / l

    P0 = scipy.special.binom(n, 0) * ((1 - t) ** (n - 0)) * t ** 0
    P1 = scipy.special.binom(n, 1) * ((1 - t) ** (n - 1)) * t ** 1
    P2 = scipy.special.binom(n, 2) * ((1 - t) ** (n - 2)) * t ** 2
    P3 = scipy.special.binom(n, 3) * ((1 - t) ** (n - 3)) * t ** 3
    P4 = scipy.special.binom(n, 4) * ((1 - t) ** (n - 4)) * t ** 4
    P5 = scipy.special.binom(n, 5) * ((1 - t) ** (n - 5)) * t ** 5
    P6 = scipy.special.binom(n, 6) * ((1 - t) ** (n - 6)) * t ** 6
    P7 = scipy.special.binom(n, 7) * ((1 - t) ** (n - 7)) * t ** 7
    P8 = scipy.special.binom(n, 8) * ((1 - t) ** (n - 8)) * t ** 8
    P9 = scipy.special.binom(n, 9) * ((1 - t) ** (n - 9)) * t ** 9
    P10 = scipy.special.binom(n, 10) * ((1 - t) ** (n - 10)) * t ** 10

    P0dot = -10.0 * (-t + 1) ** 9
    P1dot = -90.0 * t * (-t + 1) ** 8 + 10.0 * (-t + 1) ** 9
    P2dot = -360.0 * t ** 2 * (-t + 1) ** 7 + 90.0 * t * (-t + 1) ** 8
    P3dot = -840.0 * t ** 3 * (-t + 1) ** 6 + 360.0 * t ** 2 * (-t + 1) ** 7
    P4dot = -1260.0 * t ** 4 * (-t + 1) ** 5 + 840.0 * t ** 3 * (-t + 1) ** 6
    P5dot = -1260.0 * t ** 5 * (-t + 1) ** 4 + 1260.0 * t ** 4 * (-t + 1) ** 5
    P6dot = -840.0 * t ** 6 * (-t + 1) ** 3 + 1260.0 * t ** 5 * (-t + 1) ** 4
    P7dot = -360.0 * t ** 7 * (-t + 1) ** 2 + 840.0 * t ** 6 * (-t + 1) ** 3
    P8dot = 45.0 * t ** 8 * (2 * t - 2) + 360.0 * t ** 7 * (-t + 1) ** 2
    P9dot = -10.0 * t ** 9 + 9 * t ** 8 * (-10.0 * t + 10.0)
    P10dot = 10.0 * t ** 9

    P0ddot = 90.0 * (-t + 1) ** 8
    P1ddot = 720.0 * t * (-t + 1) ** 7 - 180.0 * (-t + 1) ** 8
    P2ddot = 2520.0 * t ** 2 * (-t + 1) ** 6 - 1440.0 * t * (-t + 1) ** 7 + 90.0 * (-t + 1) ** 8
    P3ddot = 5040.0 * t ** 3 * (-t + 1) ** 5 - 5040.0 * t ** 2 * (-t + 1) ** 6 + 720.0 * t * (-t + 1) ** 7
    P4ddot = 6300.0 * t ** 4 * (-t + 1) ** 4 - 10080.0 * t ** 3 * (-t + 1) ** 5 + 2520.0 * t ** 2 * (-t + 1) ** 6
    P5ddot = 5040.0 * t ** 5 * (-t + 1) ** 3 - 12600.0 * t ** 4 * (-t + 1) ** 4 + 5040.0 * t ** 3 * (-t + 1) ** 5
    P6ddot = 2520.0 * t ** 6 * (-t + 1) ** 2 - 10080.0 * t ** 5 * (-t + 1) ** 3 + 6300.0 * t ** 4 * (-t + 1) ** 4
    P7ddot = -360.0 * t ** 7 * (2 * t - 2) - 5040.0 * t ** 6 * (-t + 1) ** 2 + 5040.0 * t ** 5 * (-t + 1) ** 3
    P8ddot = 90.0 * t ** 8 + 720.0 * t ** 7 * (2 * t - 2) + 2520.0 * t ** 6 * (-t + 1) ** 2
    P9ddot = -180.0 * t ** 8 + 72 * t ** 7 * (-10.0 * t + 10.0)
    P10ddot = 90.0 * t ** 8
    90.0 * t ** 8

    P = np.hstack((P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10))
    Pdot = np.hstack((P0dot, P1dot, P2dot, P3dot, P4dot, P5dot, P6dot, P7dot, P8dot, P9dot, P10dot)) / l
    Pddot = np.hstack((P0ddot, P1ddot, P2ddot, P3ddot, P4ddot, P5ddot, P6ddot, P7ddot, P8ddot, P9ddot, P10ddot)) / (l ** 2)
    return P, Pdot, Pddot


#####################################################
#####################################################
##################### NEW OPTIMIZER #################
#######################################################
######################################################
######################################################################################################################
######################################################################################################################
######################################################################################################################

class OPTNode_batched():
    def __init__(self, P, Pddot, A_eq, A_obs, Q_smoothness, x_obs, y_obs, num=12, num_obs=4, nvar=11, a_obs=1.0, b_obs=1.0, rho_obs=0.3, rho_eq=10.0, weight_smoothness=10, maxiter=300, eps=1e-7, num_tot=48, batch_size=30):
        super().__init__()
        self.device = 'cpu'
        device = self.device
        self.P = torch.tensor(P, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        self.Pddot = torch.tensor(Pddot, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        self.A_eq = torch.tensor(A_eq, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        self.A_obs = torch.tensor(A_obs, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        self.Q_smoothness = torch.tensor(Q_smoothness, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        self.x_obs = torch.tensor(x_obs, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        self.y_obs = torch.tensor(y_obs, dtype=torch.double).to(device).expand(batch_size, -1, -1)
        
        
        self.num = num
        self.num_obs = num_obs
        self.eps = eps
        self.nvar = nvar        
        self.a_obs = a_obs
        self.b_obs = b_obs        
        self.rho_eq = rho_eq
        self.num_obs = num_obs
        self.maxiter = maxiter
        self.num_tot = num_tot
        self.rho_obs = rho_obs
        self.weight_smoothness = weight_smoothness
        self.batch_size = batch_size        

    def optimize2(self, b, lamda_x, lamda_y):
        device = self.device
        # print(b.shape)
        bx_eq_tensor, by_eq_tensor = torch.split(b, 6, dim=0)
        
        d_obs = torch.ones(self.batch_size, self.num_obs, self.num, dtype=torch.double).to(device)
        alpha_obs = torch.zeros(self.batch_size, self.num_obs, self.num, dtype=torch.double).to(device)
        ones_tensor = torch.ones((self.batch_size, self.num_obs, self.num), dtype=torch.double).to(device)
        # print(self.Pddot.shape)
        # trans = self.Pddot.permute(0, 2, 1)
        # print(trans.shape, " AA")
        # import pdb;pdb.set_trace()
        cost_smoothness = self.weight_smoothness * torch.bmm(self.Pddot.permute(0, 2, 1), self.Pddot)
        cost = cost_smoothness + self.rho_obs * torch.bmm(self.A_obs.permute(0, 2, 1), self.A_obs) + self.rho_eq * torch.bmm(self.A_eq.permute(0, 2, 1), self.A_eq)
        # import pdb;pdb.set_trace()
        for i in range(self.maxiter):
            temp_x_obs = d_obs * torch.cos(alpha_obs) * self.a_obs
            temp_y_obs = d_obs * torch.sin(alpha_obs) * self.b_obs

            # import pdb;pdb.set_trace()
            b_obs_x = self.x_obs.view(self.batch_size, self.num * self.num_obs) + temp_x_obs.view(self.batch_size, self.num * self.num_obs)
            b_obs_y = self.y_obs.view(self.batch_size, self.num * self.num_obs) + temp_y_obs.view(self.batch_size, self.num * self.num_obs)
            # import pdb;pdb.set_trace()
#             print(self.A_eq.permute(0, 2, 1).shape, bx_eq_tensor.unsqueeze(2).shape, " POO", lamda_x.shape)
#             print(self.A_obs.shape, b_obs_x.unsqueeze(2).shape, torch.bmm(self.A_obs.permute(0, 2, 1), b_obs_x.unsqueeze(2)).shape, " obs")
#             print(self.A_eq.shape, bx_eq_tensor.unsqueeze(2).shape, torch.bmm(self.A_eq.permute(0, 2, 1), bx_eq_tensor.unsqueeze(2).permute(1, 0, 2)).shape, " eq")
            
            # print(lamda_x.shape, " lamda_x")
            # import pdb;pdb.set_trace()
            lincost_x = -lamda_x - self.rho_obs * torch.bmm(self.A_obs.permute(0, 2, 1), b_obs_x.unsqueeze(2)) - self.rho_eq * torch.bmm(self.A_eq.permute(0, 2, 1), bx_eq_tensor.unsqueeze(2).permute(1, 0, 2))
            lincost_y = -lamda_y - self.rho_obs * torch.bmm(self.A_obs.permute(0, 2, 1), b_obs_y.unsqueeze(2)) - self.rho_eq * torch.bmm(self.A_eq.permute(0, 2, 1), by_eq_tensor.unsqueeze(2).permute(1, 0, 2))
            # import pdb;pdb.set_trace()
            # lincost_x = lincost_x.view(-1, 1)
            # lincost_y = lincost_y.view(-1, 1)
            # print(cost.shape, " AAA")
            # print(cost_inv.shape)
            # print(cost_inv.shape, lincost_x.shape, lincost_x.T.shape, " lincost")
            cost_inv = torch.zeros_like(cost)
            for j in range(self.batch_size):
                cost_inv[j] = torch.linalg.inv(cost[j])

            sol_x = torch.bmm(-cost_inv, lincost_x)
            sol_y = torch.bmm(-cost_inv, lincost_y)
            # import pdb;pdb.set_trace()
            
            # print(torch.linalg.lstsq(lincost_x, -cost).solution)
            # sol_x, _ = torch.linalg.lstsq(lincost_x, -cost)
            # sol_x = torch.linalg.lstsq(lincost_x, -cost).solution
            # sol_y, _ = torch.linalg.lstsq(lincost_y, -cost)
            # sol_y = torch.linalg.lstsq(lincost_y, -cost).solution
            # print(sol_x.shape, self.P.shape, " sol")

            # sol_x = sol_x.view(-1)
            # sol_y = sol_y.view(-1)

            x = torch.bmm(self.P, sol_x)
            y = torch.bmm(self.P, sol_y)
            # import pdb;pdb.set_trace()

            # print(x.shape, y.shape, " xy")
            # print(self.x_obs.shape, " x_obs")
            wc_alpha = x.permute(0, 2, 1) - self.x_obs
            ws_alpha = y.permute(0, 2, 1) - self.y_obs
            alpha_obs = torch.atan2(ws_alpha * self.a_obs, wc_alpha * self.b_obs)
            # import pdb;pdb.set_trace()

            c1_d = self.rho_obs * (self.a_obs ** 2 * torch.cos(alpha_obs) ** 2 + self.b_obs ** 2 * torch.sin(alpha_obs) ** 2)
            c2_d = self.rho_obs * (self.a_obs * wc_alpha * torch.cos(alpha_obs) + self.b_obs * ws_alpha * torch.sin(alpha_obs))
            d_temp = c2_d / c1_d
            d_obs = torch.max(d_temp, ones_tensor)
            # import pdb;pdb.set_trace()

            # print(alpha_obs.shape, wc_alpha.shape, " alpga_obs")
            res_x_obs_vec = wc_alpha - self.a_obs * d_obs * torch.cos(alpha_obs)
            res_y_obs_vec = ws_alpha - self.b_obs * d_obs * torch.sin(alpha_obs)
            # import pdb;pdb.set_trace()

            # print(torch.bmm(self.A_eq, sol_x).shape, " res_eq_x_vec")
            res_eq_x_vec = torch.bmm(self.A_eq, sol_x) - bx_eq_tensor.unsqueeze(2).permute(1, 0, 2)
            res_eq_y_vec = torch.bmm(self.A_eq, sol_y) - by_eq_tensor.unsqueeze(2).permute(1, 0, 2)
            # import pdb;pdb.set_trace()

            # print(res_x_obs_vec.shape, "res_x_obs_vec")
            # print(res_eq_x_vec.shape, "res_eq_x_vec")
            # print(self.A_obs.permute(0, 2, 1).shape, res_x_obs_vec.view(self.batch_size, -1).unsqueeze(2).shape)
            # print(self.A_eq.permute(0, 2, 1).shape, res_eq_x_vec.shape)
            # print(torch.bmm(self.A_obs.permute(0, 2, 1), res_x_obs_vec.view(self.batch_size, -1).unsqueeze(2)).shape)
            # print(torch.bmm(self.A_eq.permute(0, 2, 1), res_eq_x_vec).shape)
            # print(lamda_x.shape)
            lamda_x -= self.rho_obs * torch.bmm(self.A_obs.permute(0, 2, 1), res_x_obs_vec.view(self.batch_size, -1).unsqueeze(2)) + self.rho_eq * torch.bmm(self.A_eq.permute(0, 2, 1), res_eq_x_vec)
            lamda_y -= self.rho_obs * torch.bmm(self.A_obs.permute(0, 2, 1), res_y_obs_vec.view(self.batch_size, -1).unsqueeze(2)) + self.rho_eq * torch.bmm(self.A_eq.permute(0, 2, 1), res_eq_y_vec)
            # import pdb;pdb.set_trace()

        sol = torch.cat([sol_x, sol_y], dim=1)
        return sol
    
    
    def solve(self, b, lamda_x, lamda_y):
        device = self.device
        batch_size, _ = b.size()
        b = b.transpose(0, 1)
        lamda_x = lamda_x.unsqueeze(2)
        lamda_y = lamda_y.unsqueeze(2)
        sol = self.optimize2(b, lamda_x, lamda_y)
        return sol.squeeze(), None
data/stp3/utils/Optnode_naive.py
llmbev/talk2bev
import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader

import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader
#from bernstein import bernstesin_coeff_order10_new

def bernstein_coeff_order10_new(n, tmin, tmax, t_actual):
    l = tmax - tmin
    t = (t_actual - tmin) / l

    P0 = scipy.special.binom(n, 0) * ((1 - t) ** (n - 0)) * t ** 0
    P1 = scipy.special.binom(n, 1) * ((1 - t) ** (n - 1)) * t ** 1
    P2 = scipy.special.binom(n, 2) * ((1 - t) ** (n - 2)) * t ** 2
    P3 = scipy.special.binom(n, 3) * ((1 - t) ** (n - 3)) * t ** 3
    P4 = scipy.special.binom(n, 4) * ((1 - t) ** (n - 4)) * t ** 4
    P5 = scipy.special.binom(n, 5) * ((1 - t) ** (n - 5)) * t ** 5
    P6 = scipy.special.binom(n, 6) * ((1 - t) ** (n - 6)) * t ** 6
    P7 = scipy.special.binom(n, 7) * ((1 - t) ** (n - 7)) * t ** 7
    P8 = scipy.special.binom(n, 8) * ((1 - t) ** (n - 8)) * t ** 8
    P9 = scipy.special.binom(n, 9) * ((1 - t) ** (n - 9)) * t ** 9
    P10 = scipy.special.binom(n, 10) * ((1 - t) ** (n - 10)) * t ** 10

    P0dot = -10.0 * (-t + 1) ** 9
    P1dot = -90.0 * t * (-t + 1) ** 8 + 10.0 * (-t + 1) ** 9
    P2dot = -360.0 * t ** 2 * (-t + 1) ** 7 + 90.0 * t * (-t + 1) ** 8
    P3dot = -840.0 * t ** 3 * (-t + 1) ** 6 + 360.0 * t ** 2 * (-t + 1) ** 7
    P4dot = -1260.0 * t ** 4 * (-t + 1) ** 5 + 840.0 * t ** 3 * (-t + 1) ** 6
    P5dot = -1260.0 * t ** 5 * (-t + 1) ** 4 + 1260.0 * t ** 4 * (-t + 1) ** 5
    P6dot = -840.0 * t ** 6 * (-t + 1) ** 3 + 1260.0 * t ** 5 * (-t + 1) ** 4
    P7dot = -360.0 * t ** 7 * (-t + 1) ** 2 + 840.0 * t ** 6 * (-t + 1) ** 3
    P8dot = 45.0 * t ** 8 * (2 * t - 2) + 360.0 * t ** 7 * (-t + 1) ** 2
    P9dot = -10.0 * t ** 9 + 9 * t ** 8 * (-10.0 * t + 10.0)
    P10dot = 10.0 * t ** 9

    P0ddot = 90.0 * (-t + 1) ** 8
    P1ddot = 720.0 * t * (-t + 1) ** 7 - 180.0 * (-t + 1) ** 8
    P2ddot = 2520.0 * t ** 2 * (-t + 1) ** 6 - 1440.0 * t * (-t + 1) ** 7 + 90.0 * (-t + 1) ** 8
    P3ddot = 5040.0 * t ** 3 * (-t + 1) ** 5 - 5040.0 * t ** 2 * (-t + 1) ** 6 + 720.0 * t * (-t + 1) ** 7
    P4ddot = 6300.0 * t ** 4 * (-t + 1) ** 4 - 10080.0 * t ** 3 * (-t + 1) ** 5 + 2520.0 * t ** 2 * (-t + 1) ** 6
    P5ddot = 5040.0 * t ** 5 * (-t + 1) ** 3 - 12600.0 * t ** 4 * (-t + 1) ** 4 + 5040.0 * t ** 3 * (-t + 1) ** 5
    P6ddot = 2520.0 * t ** 6 * (-t + 1) ** 2 - 10080.0 * t ** 5 * (-t + 1) ** 3 + 6300.0 * t ** 4 * (-t + 1) ** 4
    P7ddot = -360.0 * t ** 7 * (2 * t - 2) - 5040.0 * t ** 6 * (-t + 1) ** 2 + 5040.0 * t ** 5 * (-t + 1) ** 3
    P8ddot = 90.0 * t ** 8 + 720.0 * t ** 7 * (2 * t - 2) + 2520.0 * t ** 6 * (-t + 1) ** 2
    P9ddot = -180.0 * t ** 8 + 72 * t ** 7 * (-10.0 * t + 10.0)
    P10ddot = 90.0 * t ** 8
    90.0 * t ** 8

    P = np.hstack((P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10))
    Pdot = np.hstack((P0dot, P1dot, P2dot, P3dot, P4dot, P5dot, P6dot, P7dot, P8dot, P9dot, P10dot)) / l
    Pddot = np.hstack((P0ddot, P1ddot, P2ddot, P3ddot, P4ddot, P5ddot, P6ddot, P7ddot, P8ddot, P9ddot, P10ddot)) / (l ** 2)
    return P, Pdot, Pddot


#####################################################
#####################################################
##################### NEW OPTIMIZER #################
#######################################################
######################################################
class OPTNode_obs():
    def __init__(self, rho_eq=1.0, rho_goal=1.0, a_obs=1.0, b_obs=1.0, x_obs=[], y_obs=[], rho_obs=3.0, num_obs=0, rho_nonhol=1.0, rho_psi=1.0, maxiter=5000, weight_smoothness=1.0, weight_smoothness_psi=1.0, t_fin=2.0, num=30, bernstein_coeff_order10_new=None, device="cpu"):
        super().__init__()
        self.rho_eq = rho_eq
        self.rho_goal = rho_goal
        self.rho_nonhol = rho_nonhol
        self.rho_psi = rho_psi
        self.rho_obs = rho_obs
        self.maxiter = maxiter
        self.a_obs = a_obs
        self.b_obs = b_obs
        self.x_obs = torch.tensor(x_obs, dtype=torch.double).to(device)
        self.y_obs = torch.tensor(y_obs, dtype=torch.double).to(device)
        self.num_obs = num_obs
        self.weight_smoothness = weight_smoothness
        self.weight_smoothness_psi = weight_smoothness_psi
        self.bernstein_coeff_order10_new = bernstein_coeff_order10_new
        
        self.device = device
        
        self.t_fin = t_fin
        self.num = num
        self.t = self.t_fin / self.num

        #self.num_batch = 10
        
        tot_time = np.linspace(0.0, self.t_fin, self.num)
        tot_time_copy = tot_time.reshape(self.num, 1)
        self.P, self.Pdot, self.Pddot = bernstein_coeff_order10_new(10, tot_time_copy[0], tot_time_copy[-1], tot_time_copy)
        self.nvar = np.shape(self.P)[1]
        
        self.cost_smoothness = self.weight_smoothness * np.dot(self.Pddot.T, self.Pddot)
        self.cost_smoothness_psi = self.weight_smoothness_psi * np.dot(self.Pddot.T, self.Pddot)
        self.lincost_smoothness_psi = np.zeros(self.nvar)

        self.A_eq = np.vstack((self.P[0], self.P[-1]))
        self.A_obs = np.tile(self.P, (num_obs, 1))
        self.A_eq_psi = np.vstack((self.P[0], self.Pdot[0], self.P[-1]))
        
        self.P = torch.tensor(self.P, dtype=torch.double).to(device)
        self.Pdot = torch.tensor(self.Pdot, dtype=torch.double).to(device)
        self.Pddot = torch.tensor(self.Pddot, dtype=torch.double).to(device)
        self.A_obs = torch.tensor(self.A_obs, dtype=torch.double).to(device)
        self.A_eq = torch.tensor(self.A_eq, dtype=torch.double).to(device)        
        self.A_eq_psi = torch.tensor(self.A_eq_psi, dtype=torch.double).to(device)
        self.cost_smoothness = torch.tensor(self.cost_smoothness, dtype=torch.double).to(device)
        self.cost_smoothness_psi = torch.tensor(self.cost_smoothness_psi, dtype=torch.double).to(device)
        self.lincost_smoothness_psi = torch.tensor(self.lincost_smoothness_psi, dtype=torch.double).to(device)
    
        self.A_nonhol = self.Pdot
        self.A_psi = self.P
        
        self.lamda_x = None
        self.lamda_y = None
        self.lamda_psi = None
        
    def compute_x(self, v, psi, b_eq_x, b_eq_y):
        b_nonhol_x = v * torch.cos(psi)
        b_nonhol_y = v * torch.sin(psi)
        batch_size = b_eq_x.shape[0]
        b_obs_x = self.x_obs.reshape(self.num, self.num_obs).expand(batch_size, self.num, self.num_obs).view(batch_size, self.num * self.num_obs, 1) + self.temp_x_obs.reshape(batch_size, self.num * self.num_obs, 1)
        b_obs_y = self.y_obs.reshape(self.num, self.num_obs).expand(batch_size, self.num, self.num_obs).view(batch_size, self.num * self.num_obs, 1) + self.temp_y_obs.reshape(batch_size, self.num * self.num_obs, 1)
    
        cost = self.cost_smoothness + self.rho_nonhol * torch.matmul(self.A_nonhol.T, self.A_nonhol) + self.rho_eq * torch.matmul(self.A_eq.T, self.A_eq) + self.rho_obs * torch.matmul(self.A_obs.T, self.A_obs)
        lincost_x = -self.lamda_x - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_x.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_x.T).T - self.rho_obs * torch.matmul(self.A_obs.T.expand(batch_size, self.nvar, self.num * self.num_obs), b_obs_x).squeeze()
        lincost_y = -self.lamda_y - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_y.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_y.T).T - self.rho_obs * torch.matmul(self.A_obs.T.expand(batch_size, self.nvar, self.num * self.num_obs), b_obs_y).squeeze()

        cost_inv = torch.linalg.inv(cost)

        sol_x = torch.matmul(-cost_inv, lincost_x.T).T
        sol_y = torch.matmul(-cost_inv, lincost_y.T).T

        x = torch.matmul(self.P, sol_x.T).T
        xdot = torch.matmul(self.Pdot, sol_x.T).T

        y = torch.matmul(self.P, sol_y.T).T
        ydot = torch.matmul(self.Pdot, sol_y.T).T
         
        return sol_x, sol_y, x, xdot, y, ydot
    
    def compute_psi(self, psi, lamda_psi, psi_temp, b_eq_psi):
        cost = self.cost_smoothness_psi + self.rho_psi * torch.matmul(self.A_psi.T, self.A_psi) + self.rho_eq * torch.matmul(self.A_eq_psi.T, self.A_eq_psi)
        lincost_psi = -self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, psi_temp.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, b_eq_psi.T).T

        cost_inv = torch.linalg.inv(cost)

        sol_psi = torch.matmul(-cost_inv, lincost_psi.T).T

        psi = torch.matmul(self.P, sol_psi.T).T

        res_psi = torch.matmul(self.A_psi, sol_psi.T).T - psi_temp
        res_eq_psi = torch.matmul(self.A_eq_psi, sol_psi.T).T - b_eq_psi

        self.lamda_psi = self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, res_psi.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, res_eq_psi.T).T

        return sol_psi, torch.linalg.norm(res_psi), torch.linalg.norm(res_eq_psi), psi

    
    def solve(self, fixed_params, variable_params):
        device = self.device
        batch_size, _ = fixed_params.size()
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin = torch.chunk(variable_params, 3, dim=1)
        
        b_eq_x = torch.cat((x_init, x_fin), dim=1)
        b_eq_y = torch.cat((y_init, y_fin), dim=1)
        b_eq_psi = torch.cat((psi_init, psidot_init, psi_fin), dim=1)

        d_obs = torch.ones(batch_size, self.num_obs, self.num, dtype=torch.double).to(device)
        alpha_obs = torch.zeros(batch_size, self.num_obs, self.num, dtype=torch.double).to(device)
        ones_tensor = torch.ones((batch_size, self.num_obs, self.num), dtype=torch.double).to(device)        
        
        v = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * v_init
        psi = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * psi_init
        xdot = v * torch.cos(psi)
        ydot = v * torch.sin(psi)
        
        self.lamda_x = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_y = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_psi = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        
        res_psi_arr = []
        res_eq_psi_arr = []
        res_eq_arr = []
        res_nonhol_arr = []
        for i in range(0, self.maxiter):
            self.temp_x_obs = d_obs * torch.cos(alpha_obs) * self.a_obs
            self.temp_y_obs = d_obs * torch.sin(alpha_obs) * self.b_obs

            psi_temp = torch.atan2(ydot, xdot)
            c_psi, res_psi, res_eq_psi, psi = self.compute_psi(psi, self.lamda_psi, psi_temp, b_eq_psi)
            c_x, c_y, x, xdot, y, ydot = self.compute_x(v, psi, b_eq_x, b_eq_y)

            # x -> (150, 6), x_obs -> (20, 6)
            wc_alpha = x.unsqueeze(1) - self.x_obs.expand(batch_size, self.num_obs, self.num) # (150, 1, 6) - (150, 20, 6) -> (150, 20, 6)
            ws_alpha = y.unsqueeze(1) - self.y_obs.expand(batch_size, self.num_obs, self.num) # (150, 1, 6) - (150, 20, 6) -> (150, 20, 6)
            alpha_obs = torch.atan2(ws_alpha * self.a_obs, wc_alpha * self.b_obs)

            c1_d = self.rho_obs * (self.a_obs ** 2 * torch.cos(alpha_obs) ** 2 + self.b_obs ** 2 * torch.sin(alpha_obs) ** 2)
            c2_d = self.rho_obs * (self.a_obs * wc_alpha * torch.cos(alpha_obs) + self.b_obs * ws_alpha * torch.sin(alpha_obs))
            d_temp = c2_d / c1_d
            d_obs = torch.max(d_temp, ones_tensor)
            
            res_x_obs_vec = wc_alpha - self.a_obs * d_obs * torch.cos(alpha_obs)
            res_y_obs_vec = ws_alpha - self.b_obs * d_obs * torch.sin(alpha_obs)            
            
            res_eq_psi_arr.append(res_eq_psi)
            res_psi_arr.append(res_psi)
            v = torch.sqrt(xdot ** 2 + ydot ** 2)

            res_eq_x = torch.matmul(self.A_eq, c_x.T).T - b_eq_x
            res_nonhol_x = xdot - v * torch.cos(psi)

            res_eq_y = torch.matmul(self.A_eq, c_y.T).T - b_eq_y
            res_nonhol_y = ydot - v * torch.sin(psi)

            res_eq_arr.append(torch.linalg.norm(torch.sqrt(res_eq_x**2 + res_eq_y**2)))
            res_nonhol_arr.append(torch.linalg.norm(torch.sqrt(res_nonhol_x**2 + res_nonhol_y**2)))
            
            self.lamda_x = self.lamda_x - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_x.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_x.T).T - self.rho_obs * torch.matmul(self.A_obs.T.expand(batch_size, self.nvar, self.num * self.num_obs), res_x_obs_vec.reshape(batch_size, self.num * self.num_obs, 1)).squeeze()
            self.lamda_y = self.lamda_y - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_y.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_y.T).T - self.rho_obs * torch.matmul(self.A_obs.T.expand(batch_size, self.nvar, self.num * self.num_obs), res_y_obs_vec.reshape(batch_size, self.num * self.num_obs, 1)).squeeze()

        print(c_x.shape)
        print(c_y.shape)
        print(c_psi.shape)
        print(v.shape)
        print(alpha_obs.shape)
        print(d_obs.shape)
        print("c_x, c_y, c_psi, v, alpha_obs, d_obs")
        """
            c_x -> (150, nvar)
            c_y -> (150, nvar)
            c_psi -> (150, nvar)
            v -> (150, num)
            alpha_obs -> (150, num_obs, num)
            d_obs -> (150, num_obs, num)
        """
        primal_sol = torch.hstack((c_x, c_y, c_psi, v, alpha_obs.reshape(batch_size, self.num_obs * self.num), d_obs.reshape(batch_size, self.num_obs * self.num)))
        return primal_sol, None

data/stp3/utils/Optnode_obs_batched.py
llmbev/talk2bev
import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader

import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader
#from bernstein import bernstesin_coeff_order10_new

def bernstein_coeff_order10_new(n, tmin, tmax, t_actual):
    l = tmax - tmin
    t = (t_actual - tmin) / l

    P0 = scipy.special.binom(n, 0) * ((1 - t) ** (n - 0)) * t ** 0
    P1 = scipy.special.binom(n, 1) * ((1 - t) ** (n - 1)) * t ** 1
    P2 = scipy.special.binom(n, 2) * ((1 - t) ** (n - 2)) * t ** 2
    P3 = scipy.special.binom(n, 3) * ((1 - t) ** (n - 3)) * t ** 3
    P4 = scipy.special.binom(n, 4) * ((1 - t) ** (n - 4)) * t ** 4
    P5 = scipy.special.binom(n, 5) * ((1 - t) ** (n - 5)) * t ** 5
    P6 = scipy.special.binom(n, 6) * ((1 - t) ** (n - 6)) * t ** 6
    P7 = scipy.special.binom(n, 7) * ((1 - t) ** (n - 7)) * t ** 7
    P8 = scipy.special.binom(n, 8) * ((1 - t) ** (n - 8)) * t ** 8
    P9 = scipy.special.binom(n, 9) * ((1 - t) ** (n - 9)) * t ** 9
    P10 = scipy.special.binom(n, 10) * ((1 - t) ** (n - 10)) * t ** 10

    P0dot = -10.0 * (-t + 1) ** 9
    P1dot = -90.0 * t * (-t + 1) ** 8 + 10.0 * (-t + 1) ** 9
    P2dot = -360.0 * t ** 2 * (-t + 1) ** 7 + 90.0 * t * (-t + 1) ** 8
    P3dot = -840.0 * t ** 3 * (-t + 1) ** 6 + 360.0 * t ** 2 * (-t + 1) ** 7
    P4dot = -1260.0 * t ** 4 * (-t + 1) ** 5 + 840.0 * t ** 3 * (-t + 1) ** 6
    P5dot = -1260.0 * t ** 5 * (-t + 1) ** 4 + 1260.0 * t ** 4 * (-t + 1) ** 5
    P6dot = -840.0 * t ** 6 * (-t + 1) ** 3 + 1260.0 * t ** 5 * (-t + 1) ** 4
    P7dot = -360.0 * t ** 7 * (-t + 1) ** 2 + 840.0 * t ** 6 * (-t + 1) ** 3
    P8dot = 45.0 * t ** 8 * (2 * t - 2) + 360.0 * t ** 7 * (-t + 1) ** 2
    P9dot = -10.0 * t ** 9 + 9 * t ** 8 * (-10.0 * t + 10.0)
    P10dot = 10.0 * t ** 9

    P0ddot = 90.0 * (-t + 1) ** 8
    P1ddot = 720.0 * t * (-t + 1) ** 7 - 180.0 * (-t + 1) ** 8
    P2ddot = 2520.0 * t ** 2 * (-t + 1) ** 6 - 1440.0 * t * (-t + 1) ** 7 + 90.0 * (-t + 1) ** 8
    P3ddot = 5040.0 * t ** 3 * (-t + 1) ** 5 - 5040.0 * t ** 2 * (-t + 1) ** 6 + 720.0 * t * (-t + 1) ** 7
    P4ddot = 6300.0 * t ** 4 * (-t + 1) ** 4 - 10080.0 * t ** 3 * (-t + 1) ** 5 + 2520.0 * t ** 2 * (-t + 1) ** 6
    P5ddot = 5040.0 * t ** 5 * (-t + 1) ** 3 - 12600.0 * t ** 4 * (-t + 1) ** 4 + 5040.0 * t ** 3 * (-t + 1) ** 5
    P6ddot = 2520.0 * t ** 6 * (-t + 1) ** 2 - 10080.0 * t ** 5 * (-t + 1) ** 3 + 6300.0 * t ** 4 * (-t + 1) ** 4
    P7ddot = -360.0 * t ** 7 * (2 * t - 2) - 5040.0 * t ** 6 * (-t + 1) ** 2 + 5040.0 * t ** 5 * (-t + 1) ** 3
    P8ddot = 90.0 * t ** 8 + 720.0 * t ** 7 * (2 * t - 2) + 2520.0 * t ** 6 * (-t + 1) ** 2
    P9ddot = -180.0 * t ** 8 + 72 * t ** 7 * (-10.0 * t + 10.0)
    P10ddot = 90.0 * t ** 8
    90.0 * t ** 8

    P = np.hstack((P0, P1, P2, P3, P4, P5, P6, P7, P8, P9, P10))
    Pdot = np.hstack((P0dot, P1dot, P2dot, P3dot, P4dot, P5dot, P6dot, P7dot, P8dot, P9dot, P10dot)) / l
    Pddot = np.hstack((P0ddot, P1ddot, P2ddot, P3ddot, P4ddot, P5ddot, P6ddot, P7ddot, P8ddot, P9ddot, P10ddot)) / (l ** 2)
    return P, Pdot, Pddot


#####################################################
#####################################################
##################### NEW OPTIMIZER #################
#######################################################
######################################################
class OPTNode_obs():
    def __init__(self, rho_eq=1.0, rho_goal=1.0, a_obs=1.0, b_obs=1.0, x_obs=[], y_obs=[], rho_obs=3.0, num_obs=0, rho_nonhol=1.0, rho_psi=1.0, maxiter=5000, weight_smoothness=1.0, weight_smoothness_psi=1.0, t_fin=2.0, num=30, bernstein_coeff_order10_new=None, device="cpu"):
        super().__init__()
        self.rho_eq = rho_eq
        self.rho_goal = rho_goal
        self.rho_nonhol = rho_nonhol
        self.rho_psi = rho_psi
        self.rho_obs = rho_obs
        self.maxiter = maxiter
        self.a_obs = a_obs
        self.b_obs = b_obs
        self.x_obs = torch.tensor(x_obs, dtype=torch.double).to(device)
        self.y_obs = torch.tensor(y_obs, dtype=torch.double).to(device)
        self.num_obs = num_obs
        self.weight_smoothness = weight_smoothness
        self.weight_smoothness_psi = weight_smoothness_psi
        self.bernstein_coeff_order10_new = bernstein_coeff_order10_new
        
        self.device = device
        
        self.t_fin = t_fin
        self.num = num
        self.t = self.t_fin / self.num

        #self.num_batch = 10
        
        tot_time = np.linspace(0.0, self.t_fin, self.num)
        tot_time_copy = tot_time.reshape(self.num, 1)
        self.P, self.Pdot, self.Pddot = bernstein_coeff_order10_new(10, tot_time_copy[0], tot_time_copy[-1], tot_time_copy)
        self.nvar = np.shape(self.P)[1]
        
        self.cost_smoothness = self.weight_smoothness * np.dot(self.Pddot.T, self.Pddot)
        self.cost_smoothness_psi = self.weight_smoothness_psi * np.dot(self.Pddot.T, self.Pddot)
        self.lincost_smoothness_psi = np.zeros(self.nvar)

        self.A_eq = np.vstack((self.P[0], self.P[-1]))
        self.A_obs = np.tile(self.P, (num_obs, 1))
        self.A_eq_psi = np.vstack((self.P[0], self.Pdot[0], self.P[-1]))
        
        self.P = torch.tensor(self.P, dtype=torch.double).to(device)
        self.Pdot = torch.tensor(self.Pdot, dtype=torch.double).to(device)
        self.Pddot = torch.tensor(self.Pddot, dtype=torch.double).to(device)
        self.A_obs = torch.tensor(self.A_obs, dtype=torch.double).to(device)
        self.A_eq = torch.tensor(self.A_eq, dtype=torch.double).to(device)        
        self.A_eq_psi = torch.tensor(self.A_eq_psi, dtype=torch.double).to(device)
        self.cost_smoothness = torch.tensor(self.cost_smoothness, dtype=torch.double).to(device)
        self.cost_smoothness_psi = torch.tensor(self.cost_smoothness_psi, dtype=torch.double).to(device)
        self.lincost_smoothness_psi = torch.tensor(self.lincost_smoothness_psi, dtype=torch.double).to(device)
    
        self.A_nonhol = self.Pdot
        self.A_psi = self.P
        
        self.lamda_x = None
        self.lamda_y = None
        self.lamda_psi = None
        
    def compute_x(self, v, psi, b_eq_x, b_eq_y):
        b_nonhol_x = v * torch.cos(psi)
        b_nonhol_y = v * torch.sin(psi)
    
        b_obs_x = self.x_obs.view(self.num * self.num_obs) + self.temp_x_obs.view(self.num * self.num_obs)
        b_obs_y = self.y_obs.view(self.num * self.num_obs) + self.temp_y_obs.view(self.num * self.num_obs)
    
        cost = self.cost_smoothness + self.rho_nonhol * torch.matmul(self.A_nonhol.T, self.A_nonhol) + self.rho_eq * torch.matmul(self.A_eq.T, self.A_eq) + self.rho_obs * torch.matmul(self.A_obs.T, self.A_obs)
        lincost_x = -self.lamda_x - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_x.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_x.T).T - self.rho_obs * torch.matmul(self.A_obs.T, b_obs_x)
        lincost_y = -self.lamda_y - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_y.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_y.T).T - self.rho_obs * torch.matmul(self.A_obs.T, b_obs_y)

        cost_inv = torch.linalg.inv(cost)

        sol_x = torch.matmul(-cost_inv, lincost_x.T).T
        sol_y = torch.matmul(-cost_inv, lincost_y.T).T

        x = torch.matmul(self.P, sol_x.T).T
        xdot = torch.matmul(self.Pdot, sol_x.T).T

        y = torch.matmul(self.P, sol_y.T).T
        ydot = torch.matmul(self.Pdot, sol_y.T).T
         
        return sol_x, sol_y, x, xdot, y, ydot
    
    def compute_psi(self, psi, lamda_psi, psi_temp, b_eq_psi):
        cost = self.cost_smoothness_psi + self.rho_psi * torch.matmul(self.A_psi.T, self.A_psi) + self.rho_eq * torch.matmul(self.A_eq_psi.T, self.A_eq_psi)
        lincost_psi = -self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, psi_temp.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, b_eq_psi.T).T

        cost_inv = torch.linalg.inv(cost)

        sol_psi = torch.matmul(-cost_inv, lincost_psi.T).T

        psi = torch.matmul(self.P, sol_psi.T).T

        res_psi = torch.matmul(self.A_psi, sol_psi.T).T - psi_temp
        res_eq_psi = torch.matmul(self.A_eq_psi, sol_psi.T).T - b_eq_psi

        self.lamda_psi = self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, res_psi.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, res_eq_psi.T).T

        return sol_psi, torch.linalg.norm(res_psi), torch.linalg.norm(res_eq_psi), psi

    
    def solve(self, fixed_params, variable_params):
        device = self.device
        batch_size, _ = fixed_params.size()
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin = torch.chunk(variable_params, 3, dim=1)
        
        b_eq_x = torch.cat((x_init, x_fin), dim=1)
        b_eq_y = torch.cat((y_init, y_fin), dim=1)
        b_eq_psi = torch.cat((psi_init, psidot_init, psi_fin), dim=1)

        d_obs = torch.ones(self.batch_size, self.num_obs, self.num, dtype=torch.double).to(device)
        alpha_obs = torch.zeros(self.batch_size, self.num_obs, self.num, dtype=torch.double).to(device)
        ones_tensor = torch.ones((self.batch_size, self.num_obs, self.num), dtype=torch.double).to(device)
        

        v = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * v_init
        psi = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * psi_init
        xdot = v * torch.cos(psi)
        ydot = v * torch.sin(psi)
        
        self.lamda_x = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_y = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_psi = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        
        res_psi_arr = []
        res_eq_psi_arr = []
        res_eq_arr = []
        res_nonhol_arr = []
        for i in range(0, self.maxiter):
            self.temp_x_obs = d_obs * torch.cos(alpha_obs) * self.a_obs
            self.temp_y_obs = d_obs * torch.sin(alpha_obs) * self.b_obs

            psi_temp = torch.atan2(ydot, xdot)
            c_psi, res_psi, res_eq_psi, psi = self.compute_psi(psi, self.lamda_psi, psi_temp, b_eq_psi)
            c_x, c_y, x, xdot, y, ydot = self.compute_x(v, psi, b_eq_x, b_eq_y)
            
            wc_alpha = x - self.x_obs
            ws_alpha = y - self.y_obs
            alpha_obs = torch.atan2(ws_alpha * self.a_obs, wc_alpha * self.b_obs)

            c1_d = self.rho_obs * (self.a_obs ** 2 * torch.cos(alpha_obs) ** 2 + self.b_obs ** 2 * torch.sin(alpha_obs) ** 2)
            c2_d = self.rho_obs * (self.a_obs * wc_alpha * torch.cos(alpha_obs) + self.b_obs * ws_alpha * torch.sin(alpha_obs))
            d_temp = c2_d / c1_d
            d_obs = torch.max(d_temp, ones_tensor)
            
            res_x_obs_vec = wc_alpha - self.a_obs * d_obs * torch.cos(alpha_obs)
            res_y_obs_vec = ws_alpha - self.b_obs * d_obs * torch.sin(alpha_obs)            
            
            res_eq_psi_arr.append(res_eq_psi)
            res_psi_arr.append(res_psi)
            v = torch.sqrt(xdot ** 2 + ydot ** 2)
            #v[:, 0] = v_init[:, 0]

            res_eq_x = torch.matmul(self.A_eq, c_x.T).T - b_eq_x
            res_nonhol_x = xdot - v * torch.cos(psi)

            res_eq_y = torch.matmul(self.A_eq, c_y.T).T - b_eq_y
            res_nonhol_y = ydot - v * torch.sin(psi)

            res_eq_arr.append(torch.linalg.norm(torch.sqrt(res_eq_x**2 + res_eq_y**2)))
            res_nonhol_arr.append(torch.linalg.norm(torch.sqrt(res_nonhol_x**2 + res_nonhol_y**2)))
            
            self.lamda_x = self.lamda_x - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_x.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_x.T).T - self.rho_obs * torch.matmul(self.A_obs.T, res_x_obs_vec.view(-1))
            self.lamda_y = self.lamda_y - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_y.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_y.T).T - self.rho_obs * torch.matmul(self.A_obs.T, res_y_obs_vec.view(-1))
        
        primal_sol = torch.hstack((c_x, c_y, c_psi, v))
        return primal_sol, None
    
    def objective(self, fixed_params, variable_params, y):
        c_x = y[:, :self.nvar]
        c_y = y[:, self.nvar:2*self.nvar]
        c_psi = y[:, 2*self.nvar:3*self.nvar]
        v = y[:, 3*self.nvar:]        
        
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin = torch.chunk(variable_params, 3, dim=1)
        
        x = torch.matmul(self.P, c_x.T).T
        y = torch.matmul(self.P, c_y.T).T
        psi = torch.matmul(self.P, c_psi.T).T
        xdot = torch.matmul(self.Pdot, c_x.T).T
        ydot = torch.matmul(self.Pdot, c_y.T).T
        psidot = torch.matmul(self.Pdot, c_psi.T).T
        xddot = torch.matmul(self.Pddot, c_x.T).T
        yddot = torch.matmul(self.Pddot, c_y.T).T
        psiddot = torch.matmul(self.Pddot, c_psi.T).T

        cost_nonhol = 0.5*self.rho_nonhol*torch.sum((xdot - v*torch.cos(psi)) ** 2, 1) + 0.5*self.rho_nonhol*torch.sum((ydot - v*torch.sin(psi)) ** 2, 1)
        cost_pos = 0.5*self.rho_eq*(torch.sum((x[:, -1] - x_fin) ** 2, 1) + torch.sum((y[:, -1] - y_fin) ** 2, 1) + torch.sum((x[:, 0] - x_init) ** 2, 1) + torch.sum((y[:, 0] - y_init) ** 2, 1))
        cost_psi = 0.5*self.rho_eq*(torch.sum((psi[:, -1] - psi_fin) ** 2, 1) + torch.sum((psi[:, 0] - psi_init) ** 2, 1)
                                    + torch.sum((psidot[:, 0] - psidot_init) ** 2, 1))
        #cost_v = 0.5*self.rho_eq*torch.sum((v[:, 0] - v_init) ** 2, 1)
        cost_cancel = torch.diagonal(torch.matmul(-self.lamda_x, c_x.T) + torch.matmul(-self.lamda_y, c_y.T) + torch.matmul(-self.lamda_psi, c_psi.T))
        
        cost_smoothness = 0.5*self.weight_smoothness*(torch.sum(xddot**2, 1) + torch.sum(yddot**2, 1)) + 0.5*self.weight_smoothness_psi*torch.sum(psiddot**2, 1)
        return cost_nonhol + cost_pos + cost_psi + cost_smoothness + cost_cancel #+ cost_v 

data/stp3/utils/Optnode_obs_unbatched.py
llmbev/talk2bev
import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader

import os
import sys
import warnings
warnings.filterwarnings('ignore')

import torch
import numpy as np
import scipy.special
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt

from scipy.linalg import block_diag
from torch.utils.data import Dataset, DataLoader
#from bernstein import bernstesin_coeff_order10_new

class OPTNode_waypoint():
    def __init__(self, rho_eq=1.0, rho_goal=1.0, rho_nonhol=1.0, rho_psi=1.0, maxiter=5000, weight_smoothness=1.0, weight_smoothness_psi=1.0, t_fin=2.0, num=30, bernstein_coeff_order10_new=None, device="cpu"):
        super().__init__()
        self.rho_eq = rho_eq
        self.rho_goal = rho_goal
        self.rho_nonhol = rho_nonhol
        self.rho_psi = rho_psi
        self.maxiter = maxiter
        self.weight_smoothness = weight_smoothness
        self.weight_smoothness_psi = weight_smoothness_psi

        self.device = device        
        
        self.t_fin = t_fin
        self.num = num
        self.t = self.t_fin / self.num

        #self.num_batch = 10
        
        tot_time = np.linspace(0.0, self.t_fin, self.num)
        tot_time_copy = tot_time.reshape(self.num, 1)
        self.P, self.Pdot, self.Pddot = bernstein_coeff_order10_new(10, tot_time_copy[0], tot_time_copy[-1], tot_time_copy)
        self.nvar = np.shape(self.P)[1]
        
        self.cost_smoothness = self.weight_smoothness * np.dot(self.Pddot.T, self.Pddot)
        self.cost_smoothness_psi = self.weight_smoothness_psi * np.dot(self.Pddot.T, self.Pddot)
        self.lincost_smoothness_psi = np.zeros(self.nvar)

        # self.A_eq = np.vstack((self.P[0], self.P[9], self.P[19], self.P[-1]))
        # self.A_eq_psi = np.vstack((self.P[0], self.Pdot[0], self.P[-1]))
        self.A_eq = np.vstack((self.P[0], self.P[5], self.P[10], self.P[15], self.P[20], self.P[25], self.P[30]))
        # self.A_eq = np.vstack((self.P[0], self.P[5], self.P[10], self.P[15], self.P[20], self.P[25], self.P[30]))s
        self.A_eq_psi = np.vstack((self.P[0], self.Pdot[0], self.P[-1]))
        
        
        self.P = torch.tensor(self.P, dtype=torch.double).to(device)
        self.Pdot = torch.tensor(self.Pdot, dtype=torch.double).to(device)
        self.Pddot = torch.tensor(self.Pddot, dtype=torch.double).to(device)
        self.A_eq = torch.tensor(self.A_eq, dtype=torch.double).to(device)        
        self.A_eq_psi = torch.tensor(self.A_eq_psi, dtype=torch.double).to(device)
        self.cost_smoothness = torch.tensor(self.cost_smoothness, dtype=torch.double).to(device)
        self.cost_smoothness_psi = torch.tensor(self.cost_smoothness_psi, dtype=torch.double).to(device)
        self.lincost_smoothness_psi = torch.tensor(self.lincost_smoothness_psi, dtype=torch.double).to(device)
        
        self.A_nonhol = self.Pdot
        self.A_psi = self.P
        
        self.lamda_x = None
        self.lamda_y = None
        self.lamda_psi = None
        
    def compute_x(self, v, psi, b_eq_x, b_eq_y):
        b_nonhol_x = v * torch.cos(psi)
        b_nonhol_y = v * torch.sin(psi)
    
        cost = self.cost_smoothness + self.rho_nonhol * torch.matmul(self.A_nonhol.T, self.A_nonhol) + self.rho_eq * torch.matmul(self.A_eq.T, self.A_eq)
        lincost_x = -self.lamda_x - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_x.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_x.T).T
        lincost_y = -self.lamda_y - self.rho_nonhol * torch.matmul(self.A_nonhol.T, b_nonhol_y.T).T - self.rho_eq * torch.matmul(self.A_eq.T, b_eq_y.T).T

        cost_inv = torch.linalg.inv(cost)

        sol_x = torch.matmul(-cost_inv, lincost_x.T).T
        sol_y = torch.matmul(-cost_inv, lincost_y.T).T

        x = torch.matmul(self.P, sol_x.T).T
        xdot = torch.matmul(self.Pdot, sol_x.T).T

        y = torch.matmul(self.P, sol_y.T).T
        ydot = torch.matmul(self.Pdot, sol_y.T).T
         
        return sol_x, sol_y, x, xdot, y, ydot
    
    def compute_psi(self, psi, lamda_psi, psi_temp, b_eq_psi):
        cost = self.cost_smoothness_psi + self.rho_psi * torch.matmul(self.A_psi.T, self.A_psi) + self.rho_eq * torch.matmul(self.A_eq_psi.T, self.A_eq_psi)
        lincost_psi = -self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, psi_temp.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, b_eq_psi.T).T

        cost_inv = torch.linalg.inv(cost)

        sol_psi = torch.matmul(-cost_inv, lincost_psi.T).T

        psi = torch.matmul(self.P, sol_psi.T).T

        res_psi = torch.matmul(self.A_psi, sol_psi.T).T - psi_temp
        res_eq_psi = torch.matmul(self.A_eq_psi, sol_psi.T).T - b_eq_psi

        self.lamda_psi = self.lamda_psi - self.rho_psi * torch.matmul(self.A_psi.T, res_psi.T).T - self.rho_eq * torch.matmul(self.A_eq_psi.T, res_eq_psi.T).T

        return sol_psi, np.linalg.norm(res_psi), np.linalg.norm(res_eq_psi), psi

    
    def solve(self, fixed_params, variable_params):
        batch_size, _ = fixed_params.size()
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin, x_mid1, y_mid1, x_mid2, y_mid2, x_mid3, y_mid3, x_mid4, y_mid4, x_mid5, y_mid5 = torch.chunk(variable_params, 13, dim=1)

        b_eq_x = torch.cat((x_init, x_mid1, x_mid2, x_mid3, x_mid4, x_mid5, x_fin), dim=1)
        b_eq_y = torch.cat((y_init, y_mid1, y_mid2, y_mid3, y_mid4, y_mid5, y_fin), dim=1)
        
        b_eq_psi = torch.cat((psi_init, psidot_init, psi_fin), dim=1)
        
        v = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * v_init
        psi = torch.ones(batch_size, self.num, dtype=torch.double).to(self.device) * psi_init
        xdot = v * torch.cos(psi)
        ydot = v * torch.sin(psi)
        
        self.lamda_x = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_y = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        self.lamda_psi = torch.zeros(batch_size, self.nvar, dtype=torch.double).to(self.device)
        
        res_psi_arr = []
        res_eq_psi_arr = []
        res_eq_arr = []
        res_nonhol_arr = []
        for i in range(0, self.maxiter):
            psi_temp = torch.atan2(ydot, xdot)
            c_psi, res_psi, res_eq_psi, psi = self.compute_psi(psi, self.lamda_psi, psi_temp, b_eq_psi)
            c_x, c_y, x, xdot, y, ydot = self.compute_x(v, psi, b_eq_x, b_eq_y)
            
            res_eq_psi_arr.append(res_eq_psi)
            res_psi_arr.append(res_psi)
            v = torch.sqrt(xdot ** 2 + ydot ** 2)
            #v[:, 0] = v_init[:, 0]

            res_eq_x = torch.matmul(self.A_eq, c_x.T).T - b_eq_x
            res_nonhol_x = xdot - v * torch.cos(psi)

            res_eq_y = torch.matmul(self.A_eq, c_y.T).T - b_eq_y
            res_nonhol_y = ydot - v * torch.sin(psi)

            res_eq_arr.append(np.linalg.norm(np.sqrt(res_eq_x**2 + res_eq_y**2)))
            res_nonhol_arr.append(np.linalg.norm(np.sqrt(res_nonhol_x**2 + res_nonhol_y**2)))

            self.lamda_x = self.lamda_x - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_x.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_x.T).T
            self.lamda_y = self.lamda_y - self.rho_eq * torch.matmul(self.A_eq.T, res_eq_y.T).T - self.rho_nonhol * torch.matmul(self.A_nonhol.T, res_nonhol_y.T).T
        
        primal_sol = torch.hstack((c_x, c_y, c_psi, v))
        return primal_sol, None
    
    def objective(self, fixed_params, variable_params, y):
        c_x = y[:, :self.nvar]
        c_y = y[:, self.nvar:2*self.nvar]
        c_psi = y[:, 2*self.nvar:3*self.nvar]
        v = y[:, 3*self.nvar:]
        
        x_init, y_init, v_init, psi_init, psidot_init = torch.chunk(fixed_params, 5, dim=1)
        x_fin, y_fin, psi_fin, x_mid1, y_mid1, x_mid2, y_mid2, x_mid3, y_mid3, x_mid4, y_mid4, x_mid5, y_mid5 = torch.chunk(variable_params, 13, dim=1)
        
        x = torch.matmul(self.P, c_x.T).T
        y = torch.matmul(self.P, c_y.T).T
        psi = torch.matmul(self.P, c_psi.T).T
        xdot = torch.matmul(self.Pdot, c_x.T).T
        ydot = torch.matmul(self.Pdot, c_y.T).T
        psidot = torch.matmul(self.Pdot, c_psi.T).T
        xddot = torch.matmul(self.Pddot, c_x.T).T
        yddot = torch.matmul(self.Pddot, c_y.T).T
        psiddot = torch.matmul(self.Pddot, c_psi.T).T
        
        cost_nonhol = 0.5*self.rho_nonhol*torch.sum((xdot - v*torch.cos(psi)) ** 2, 1) + 0.5*self.rho_nonhol*torch.sum((ydot - v*torch.sin(psi)) ** 2, 1)
        cost_pos = 0.5*self.rho_eq*(torch.sum((x[:, -1] - x_fin) ** 2, 1) + torch.sum((y[:, -1] - y_fin) ** 2, 1) + torch.sum((x[:, 0] - x_init) ** 2, 1) + torch.sum((y[:, 0] - y_init) ** 2, 1))
        cost_pos += 0.5*self.rho_eq*(torch.sum((x[:, 5] - x_mid1)**2, 1) 
                                     + torch.sum((y[:, 5] - y_mid1)**2, 1) 
                                     + torch.sum((x[:, 10] - x_mid2)**2, 1) 
                                     + torch.sum((y[:, 10] - y_mid2)**2, 1)
                                     + torch.sum((x[:, 15] - x_mid3)**2, 1)
                                     + torch.sum((y[:, 15] - y_mid3)**2, 1)
                                     + torch.sum((x[:, 20] - x_mid4)**2, 1)
                                     + torch.sum((y[:, 20] - y_mid4)**2, 1) 
                                     + torch.sum((x[:, 25] - x_mid5)**2, 1) 
                                     + torch.sum((y[:, 25] - y_mid5)**2, 1)
                                    )
        cost_psi = 0.5*self.rho_eq*(torch.sum((psi[:, -1] - psi_fin) ** 2, 1) + torch.sum((psi[:, 0] - psi_init) ** 2, 1)
                                    + torch.sum((psidot[:, 0] - psidot_init) ** 2, 1))
        #cost_v = 0.5*self.rho_eq*torch.sum((v[:, 0] - v_init) ** 2, 1)
        cost_cancel = torch.diagonal(torch.matmul(-self.lamda_x, c_x.T) + torch.matmul(-self.lamda_y, c_y.T) + torch.matmul(-self.lamda_psi, c_psi.T))
        
        cost_smoothness = 0.5*self.weight_smoothness*(torch.sum(xddot**2, 1) + torch.sum(yddot**2, 1)) + 0.5*self.weight_smoothness_psi*torch.sum(psiddot**2, 1)
        return cost_nonhol + cost_pos + cost_psi + cost_smoothness + cost_cancel #+ cost_v 

data/stp3/utils/Optnode_waypoint.py
llmbev/talk2bev
import torch
import torch.nn as nn
import pytorch_lightning as pl
import matplotlib.pyplot as plt
import numpy as np
import io
import copy
from io import BytesIO
import matplotlib
import matplotlib as mpl
import PIL
from PIL import Image

from stp3.utils.geometry import cumulative_warp_features_reverse, cumulative_warp_features, extract_trajs, extract_obs_from_centerness, generate_instance_colours

def prepare_future_labels(batch, model):
    labels = {}

    cfg = model.cfg

    segmentation_labels = batch['segmentation']
    hdmap_labels = batch['hdmap']
    future_egomotion = batch['future_egomotion']
    gt_trajectory = batch['gt_trajectory']
    labels['sample_trajectory'] = batch['sample_trajectory']

    # present frame hd map gt
    labels['hdmap'] = hdmap_labels[:, model.receptive_field - 1].long().contiguous()

    # first rec token
    if cfg.DATASET.VERSION == 'nuscenes':
        labels['rec_first'] = batch['rec_first']
        labels['gt_trajectory_prev'] = batch['gt_trajectory_prev']

    # gt trajectory
    labels['gt_trajectory'] = gt_trajectory
    spatial_extent = (cfg.LIFT.X_BOUND[1], cfg.LIFT.Y_BOUND[1])

    # Past frames gt depth
    if cfg.LIFT.GT_DEPTH:
        depths = batch['depths']
        depth_labels = depths[:, :model.receptive_field, :, ::model.encoder_downsample,
                        ::model.encoder_downsample]
        depth_labels = torch.clamp(depth_labels, cfg.LIFT.D_BOUND[0], cfg.LIFT.D_BOUND[1] - 1) - \
                        cfg.LIFT.D_BOUND[0]
        depth_labels = depth_labels.long().contiguous()
        labels['depths'] = depth_labels

    # Warp labels to present's reference frame
    segmentation_labels_past = cumulative_warp_features(
        segmentation_labels[:, :model.receptive_field].float(),
        future_egomotion[:, :model.receptive_field],
        mode='nearest', spatial_extent=spatial_extent,
    ).long().contiguous()[:, :-1]
    segmentation_labels = cumulative_warp_features_reverse(
        segmentation_labels[:, (model.receptive_field - 1):].float(),
        future_egomotion[:, (model.receptive_field - 1):],
        mode='nearest', spatial_extent=spatial_extent,
    ).long().contiguous()
    labels['segmentation'] = torch.cat([segmentation_labels_past, segmentation_labels], dim=1)

    if cfg.SEMANTIC_SEG.HDMAP.ENABLED:
        hdmap_labels = batch['hdmap'][:, :, 1:2]
        hdmap_labels_past = cumulative_warp_features(
            hdmap_labels[:, :model.receptive_field].float(),
            future_egomotion[:, :model.receptive_field],
            mode='nearest', spatial_extent=spatial_extent,
        ).long().contiguous()[:, :-1]
        hdmap_labels = cumulative_warp_features_reverse(
            hdmap_labels[:, (model.receptive_field - 1):].float(),
            future_egomotion[:, (model.receptive_field - 1):],
            mode='nearest', spatial_extent=spatial_extent,
        ).long().contiguous()
        labels['hdmap_warped_road'] = torch.cat([hdmap_labels_past, hdmap_labels], dim=1)
        hdmap_labels = batch['hdmap'][:, :, 0:1]
        hdmap_labels_past = cumulative_warp_features(
            hdmap_labels[:, :model.receptive_field].float(),
            future_egomotion[:, :model.receptive_field],
            mode='nearest', spatial_extent=spatial_extent,
        ).long().contiguous()[:, :-1]
        hdmap_labels = cumulative_warp_features_reverse(
            hdmap_labels[:, (model.receptive_field - 1):].float(),
            future_egomotion[:, (model.receptive_field - 1):],
            mode='nearest', spatial_extent=spatial_extent,
        ).long().contiguous()
        labels['hdmap_warped_lane'] = torch.cat([hdmap_labels_past, hdmap_labels], dim=1)            

    return labels

data/stp3/utils/data.py
llmbev/talk2bev
import math
import PIL
import numpy as np
import torch
import torch.nn.functional as F
from typing import Tuple
import matplotlib.pyplot as plt
import pdb

from pyquaternion import Quaternion
from shapely import affinity
from shapely.geometry import Polygon, MultiPolygon, LineString, Point, box
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from shapely.geometry import LineString
from skimage.draw import polygon
from math import cos, sin
from stp3.utils.Optnode_obs_unbatched import OPTNode_obs
from stp3.utils.Optnode import OPTNode, bernstein_coeff_order10_new
from stp3.utils.Optnode_naive import OPTNode_batched
from stp3.utils.Optnode_waypoint import OPTNode_waypoint

import pdb

from typing import Iterable, List, Sequence, Set, Tuple

n_present = 3

W = 1.85
H = 4.084
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
bx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
dx = np.array([0.5, 0.5])

def get_rot(theta):
    rot = np.array([[math.cos(theta), -math.sin(theta)], [math.sin(theta), math.cos(theta)]])
    return rot

def get_origin_points(lambda_=0, theta=0):
    pts = np.array([
        [-H / 2. + 0.5 - lambda_, W / 2. + lambda_],
        [H / 2. + 0.5 + lambda_, W / 2. + lambda_],
        [H / 2. + 0.5 + lambda_, -W / 2. - lambda_],
        [-H / 2. + 0.5 - lambda_, -W / 2. - lambda_],
    ])
    rot = get_rot(theta)
    pts = np.dot(rot, pts.T).T
    pts = (pts - bx) / (dx)    
    pts[:, [0, 1]] = pts[:, [1, 0]]
    rr , cc = polygon(pts[:,1], pts[:,0])
    rc = np.concatenate([rr[:,None], cc[:,None]], axis=-1)
    return rc # (27,2)

def get_points(self, pts, lambda_=0, bev_dimension=None):
    '''
    pts: np array (n_future, 2)
    return:
    List[ torch.Tensor<int> (B, N, n_future), torch.Tensor<int> (B, N, n_future)]
    '''
    n_future, _ = pts.shape
    rc = get_origin_points(lambda_)
    if len(pts) > 1:
        theta = pts[1:] - pts[:-1]
    else:
        theta = 0
    pts = pts + rc
    pts = torch.tensor(pts)

    rr = pts[:,0].long()
    rr = torch.clamp(rr, 0, bev_dimension[0] - 1)

    cc = pts[:,1].long()
    cc = torch.clamp(cc, 0, bev_dimension[1] - 1)

    return rr, cc

def extract_trajs(centerlines, trj, seg_prediction=None, hdmap_prediction=None, nvar=11, viz=False, debugg=False, device='cpu', labels=None, griddim=None, problem=None, avoid_obs=False, obs=None, a_obs=1, b_obs=1, ind="0", bev_dimension=None):
    save_device = device
    device = 'cpu'
    save_obs = obs
    bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
    ddx = np.array([0.5, 0.5])
    if griddim == None:
        X_MAX = 50
        X_MIN = 0
        NUM_X = 30
        Y_MAX = 5
        Y_MIN = -5
        NUM_Y = 5
    else:
        X_MIN, X_MAX, NUM_X, Y_MIN, Y_MAX, NUM_Y = griddim


    # get trajectories
    x = torch.linspace(X_MIN, X_MAX, NUM_X).to(device)
    y = torch.linspace(Y_MIN, Y_MAX, NUM_Y).to(device)
    xv, yv = torch.meshgrid(x, y)
    print(device)
    grid = torch.dstack((xv, yv)).to(device)
    ccc = centerlines.detach().cpu().numpy()
    cx, _ = interp_polyline_by_fixed_waypt_interval_np(ccc, 0.1)
    cx = extend_both_sides(cx, resolution=0.1, extension=500)
    cx = cx * ddx
    cx += bbx
    cx[:, [1, 0]] = cx[:, [0, 1]]
    cx[:, :1] = cx[:, :1] * -1
    cx = torch.tensor(cx).to(device)

    trr = trj.squeeze().detach().cpu()
    ref_line_ = torch.tensor(cx.reshape(1, cx.shape[0], 2))
    ref_line = torch.zeros((1, ref_line_.shape[1] - 1, 3))
    ref_line[:, :, :2] = ref_line_[:, 1:] 
    off = ref_line_[:, 1:] - ref_line_[:,:-1]
    ref_line[:, :, 2] = torch.atan2(off[:, :, 1], off[:, :, 0])
    ref_line = ref_line.float().to(device)
    trj = trj.float().to(device)

    test_ref_line=ref_line[:, 100:-100]
    sdd = project_to_frenet_frame(test_ref_line,  ref_line)

    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)


    cx = interp_arc_np(2000, cx.detach().cpu().numpy())
    cx = extend_both_sides(cx, resolution=0.1, extension=50)
    cx = torch.tensor(cx).to(device)
    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)
    s_offset = ssd[:,2,0]
    d_offset = ssd[0,2,1]
    vel = ssd[0,2] - ssd[0,1]
    vel_prev = ssd[0,1] - ssd[0, 0]
    theta = torch.atan2(vel[1], vel[0]) # to make angle 0 wrt ego-agent
    theta_prev = torch.atan2(vel_prev[1], vel_prev[0])
    psidot_init = (theta - theta_prev)/0.5
    psi_init = theta
    v_init = torch.linalg.norm(vel)/0.5
    y_init = 0
    x_init = 0
    rot_inv = torch.tensor([[cos(-theta), -sin(-theta)], [sin(-theta), cos(-theta)]]).to(device)
    rot = torch.tensor([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]]).to(device)


    # optimized smapling
    grid_ref = grid.reshape(1, NUM_X * NUM_Y, 2)
    s_offset = ssd[:,2,0]
    grid_ref = grid_ref + torch.tensor([s_offset, 0])
    grid_ref_cartesian = project_to_cartesian_frame(grid_ref, ref_line) # in cartesian frame
    grid_ref_cartesian = (grid_ref_cartesian[0, :, :2])/ ddx
    grid_ref_cartesian[ :, :1] = grid_ref_cartesian[:, :1] * -1
    grid_ref_cartesian[:, [0, 1]] = grid_ref_cartesian[:, [1, 0]]
    # pdb.set_trace()

    rectangle = get_origin_points()
    grid_ref_cartesian = grid_ref_cartesian.unsqueeze(1) + rectangle
    gc = grid_ref_cartesian.long()
    hdmap = hdmap_prediction[0, 1]
    seg = seg_prediction[0, -1, 0]
    map_vals = hdmap[gc[:,:,0], gc[:,:,1]].sum(-1)
    seg_vals = seg[gc[:,:,0], gc[:,:,1]].sum(-1)

    cutoff = 0
    take_map = (map_vals >= (len(rectangle) - 1 - cutoff)) # do not take any point if any portion outside hdmap prediction

    seg_map = (seg_vals == 0) # do not take any point if any portion intersects with segmentation prediction

    final_grid = []
    grid_ref = grid.reshape(1, NUM_X * NUM_Y, 2)
    for i in range(len(grid_ref[0])):
        if take_map[i] and seg_map[i]:
            print(grid_ref[0][i])
            final_grid.append(grid_ref[0][i].detach().cpu().numpy())
    final_grid = np.array(final_grid)
    grid = final_grid


    # translate
    sdd_frame = sdd -  torch.tensor([s_offset, d_offset]).to(device)
    ssd_frame = ssd -  torch.tensor([s_offset, d_offset]).to(device)
    eps_frame = torch.tensor(grid).to(device) - torch.tensor([0, d_offset]).to(device) # (20, 7, 2)

    # rotate
    sdd_frame[0] = torch.matmul(rot_inv, sdd_frame[0].T).T # (1, M, 2)
    ssd_frame = torch.matmul(rot_inv, ssd_frame[0].T).T # ()
    # eps_frame = eps_frame.permute(1, 2, 0) # (7, 2, 20)
    eps_frame = torch.matmul(rot_inv, eps_frame.T).T # (2, 2) x (M, 2) -> (M, 2)

    """
        save params of the optimizer before calling solve
    """
    # obs -> (M, num, 2) : M obstacles for num timesteps
    # avoid obstacles
    num_obs = obs.shape[0]
    num = obs.shape[1]
    # obs = obs.reshape(1, num_obs * num, 2) # (1, M * num, 2)

    # first bring to agent frame
    obs = obs * ddx
    obs += bbx
    obs[:, :, [1, 0]] = obs[:, :, [0, 1]]
    obs[:, :, :1] = obs[:, :, :1] * -1

    ssd_obs = torch.zeros((1, num_obs, num, 2))
    for n in range(num_obs):
        print(obs[n].unsqueeze(0).shape, "OK")
        print(project_to_frenet_frame(obs[n].unsqueeze(0).double().to(device), ref_line.double().to(device)).shape)
        ssd_obs[:, n] = project_to_frenet_frame(obs[n].unsqueeze(0).double().to(device), ref_line.double().to(device)) # (1, M * num, 2)

    # translate
    ssd_obs_frame = ssd_obs[0] - torch.tensor([s_offset, d_offset]).to(device) # (M, num, 2)
    # print(ssd_obs.shape, ssd_obs_frame.shape, " OBS")
    for i in range(num_obs):
        ssd_obs_frame[i] = torch.matmul(rot_inv, ssd_obs_frame[i].float().T).T
    # ssd_obs_frame = ssd_obs_frame[0].permute(1, 2, 0) # (num, 2, M)
    # ssd_obs_frame = torch.matmul(rot_inv, ssd_obs_frame.float()) # (2, 2) x (num, 2, M) -> (num, 2, M)
    # ssd_obs_frame = ssd_obs_frame.reshape(num_obs, num, 2)  # (M, num, 2)
    x_obs = ssd_obs_frame[:, : ,0]
    y_obs = ssd_obs_frame[:, : ,1]
    # x_obs, y_obs = problem.x_obs, problem.y_obs
    obs_frame = torch.dstack((x_obs, y_obs))

    x_eps = eps_frame[:,0].flatten()
    y_eps = eps_frame[:,1].flatten()

    fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
    fixed_params = fixed_params.expand(len(x_eps), 5)
    variable_params = torch.zeros((len(x_eps), 3)).to(device)
    variable_params[:, 0] = torch.tensor(x_eps)
    variable_params[:, 1] = torch.tensor(y_eps)
    variable_params[:, 2] = torch.tensor(-theta)

    """
        prepare x_obs and y_obs
    """
    num_obs = len(obs_frame)
    num_local_obs = 61
    x_obs = np.ones((num_obs, num_local_obs))
    y_obs = np.ones((num_obs, num_local_obs))
    b_inp = np.ones((len(fixed_params), 12))
    for i, ob in enumerate(obs_frame.detach().cpu().numpy()):
        # extrapolating obstacle points to (61, 2)
        x_obs[i][0] = ob[0, 0]
        x_obs[i][-1] = ob[-1, 0]
        y_obs[i][0] = ob[0, 1]
        y_obs[i][-1] = ob[-1, 1]
        pts = interp_arc_np(t=num_local_obs - 2 + 1, points=ob[1:-1])
        pts = pts[1:] # nan at start for some reason
        x_obs[i][1:-1] = pts[:, 0]
        y_obs[i][1:-1] = pts[:, 1]
        if viz:
            plt.plot(ob[:, 0], ob[:, 1])
            plt.scatter(ob[:, 0], ob[:, 1])
    print(x_obs.shape, y_obs.shape)
    for i in range(len(fixed_params)):
        x_init, y_init = fixed_params[i,0], fixed_params[i, 1]
        x_fin, y_fin = variable_params[i, 0], variable_params[i, 1]
        b_inp[i] = np.array([x_init, 0, 0, x_fin, 0, 0, y_init, 0, 0, y_fin, 0, 0])

    """
        prepare Optnode, refine collision bound trajectories
    """
    t_fin = 2
    a_obs = 2.0
    b_obs = 2.0
    # rho_obs = 0.3
    rho_obs = 1.2
    rho_eq = 10.0
    weight_smoothness = 10
    tot_time = np.linspace(0.0, t_fin, num=num_local_obs)
    tot_time_copy = tot_time.reshape(num_local_obs, 1)
    P, Pdot, Pddot = bernstein_coeff_order10_new(10, tot_time_copy[0], tot_time_copy[-1], tot_time_copy)
    A_obs = np.tile(P, (num_obs, 1))
    A_eq = np.vstack((P[0], Pdot[0], Pddot[0], P[-1], Pdot[-1], Pddot[-1]))
    Q_smoothness = np.dot(Pddot.T, Pddot)
    a_obs = 2.0
    b_obs = 2.0
    # rho_obs = 0.3
    rho_obs = 1.2
    rho_eq = 10.0
    weight_smoothness = 10
    batch_size = len(fixed_params)
    x_obs = np.ones((batch_size, num_obs, num_local_obs)) * x_obs
    y_obs = np.ones((batch_size, num_obs, num_local_obs)) * y_obs
    opt_node = OPTNode_batched(P, Pddot, A_eq, A_obs, Q_smoothness, x_obs, y_obs,num=num_local_obs, num_obs=num_obs, nvar=nvar, a_obs=a_obs, b_obs=b_obs, rho_obs=2, rho_eq=10,weight_smoothness=weight_smoothness, maxiter=3000, eps=1e-7, num_tot=num_local_obs*num_obs,batch_size=batch_size)
    lamda_x = torch.zeros(batch_size, nvar, dtype=torch.double)
    lamda_y = torch.zeros(batch_size, nvar, dtype=torch.double)
    sol = opt_node.solve(torch.tensor(b_inp), lamda_x, lamda_y)[0]
    x_pred = torch.matmul(torch.tensor(P), sol[:, :nvar].transpose(0, 1)) # (num_local_obs, batch_size)
    y_pred = torch.matmul(torch.tensor(P), sol[:, nvar:2*nvar].transpose(0, 1))# (num_local_obs, batch_size)
    x_pred = x_pred.T
    y_pred = y_pred.T

    not_to_refine_inds = []
    refine_inds = []
    for i in range(len(x_pred)):
        thetas = torch.atan2(y_pred[1:, i] - y_pred[:-1, i], x_pred[1:, i] - x_pred[:-1, i])
        unique_thetas = len(torch.unique(torch.clamp(thetas - thetas[0], 1e-5, 1e11)))
        if unique_thetas == 1:
            # straight line
            not_to_refine_inds.append(i)
        else:
            refine_inds.append(i)

    x_pred_waypt = x_pred[refine_inds]
    y_pred_waypt = y_pred[refine_inds]
    eps_frame_nonholo = eps_frame[not_to_refine_inds]
    eps_frame_waypt = eps_frame[refine_inds]
    torch.save(x_obs, f"obs/x_obs_{ind}.pt")
    torch.save(y_obs, f"obs/y_obs_{ind}.pt")
    torch.save(obs_frame, f"obs/obs_{ind}.pt")
    torch.save(fixed_params, f"obs/fixed_params_{ind}.pt")
    torch.save(variable_params, f"obs/variable_params_{ind}.pt")

    problem = OPTNode(t_fin=problem.t_fin, num=problem.num, device='cpu')
    x_eps = eps_frame_nonholo[:,0].flatten()
    y_eps = eps_frame_nonholo[:,1].flatten()
    fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
    fixed_params = fixed_params.expand(len(eps_frame_nonholo), 5)
    variable_params = torch.zeros((len(eps_frame_nonholo), 3)).to(device)
    variable_params[:, 0] = torch.tensor(x_eps)
    variable_params[:, 1] = torch.tensor(y_eps)
    variable_params[:, 2] = torch.tensor(-theta)
    P = problem.P
    cc, _ = problem.solve(fixed_params.double(), variable_params.double())
    x_sol = torch.matmul(P, cc.T[:nvar]).T
    y_sol = torch.matmul(P, cc.T[nvar:2*nvar]).T    
    trajs = torch.dstack((x_sol, y_sol))



    """
        use waypoint optimizer to refine obstacle avoided trajectories
    """
    if len(eps_frame_waypt) > 1:
        x_eps = eps_frame_waypt[:,0].flatten()
        y_eps = eps_frame_waypt[:,1].flatten()
        fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
        fixed_params = fixed_params.expand(len(eps_frame_waypt), 5)
        variable_params = torch.zeros((len(eps_frame_waypt), 13)).to(device)
        variable_params[:, 0] = torch.tensor(x_eps)
        variable_params[:, 1] = torch.tensor(y_eps)
        variable_params[:, 2] = torch.tensor(-theta)
        variable_params[:, 3] = torch.tensor(x_pred_waypt[:, 10])
        variable_params[:, 4] = torch.tensor(y_pred_waypt[:, 10])
        variable_params[:, 5] = torch.tensor(x_pred_waypt[:, 20])
        variable_params[:, 6] = torch.tensor(y_pred_waypt[:, 20])
        variable_params[:, 7] = torch.tensor(x_pred_waypt[:, 30])
        variable_params[:, 8] = torch.tensor(y_pred_waypt[:, 30])
        variable_params[:, 9] = torch.tensor(x_pred_waypt[:, 40])
        variable_params[:, 10] = torch.tensor(y_pred_waypt[:, 40])
        variable_params[:, 11] = torch.tensor(x_pred_waypt[:, 50])
        variable_params[:, 12] = torch.tensor(y_pred_waypt[:, 50])
        problem_waypt = OPTNode_waypoint(rho_eq=1.1, num=61, t_fin=t_fin, bernstein_coeff_order10_new=bernstein_coeff_order10_new)
        cc, _  = problem_waypt.solve(fixed_params.double(), variable_params.double())
        x_sol = torch.matmul(P, cc.T[:nvar]).T
        y_sol = torch.matmul(P, cc.T[nvar:2*nvar]).T    
        trajs_waypt = torch.dstack((x_sol, y_sol))
        trajs_waypt = trajs_waypt[:, ::10] # (B, 7, 2)
        trajs = torch.cat((trajs, trajs_waypt), dim=0)

    trajs = trajs.permute(0, 2, 1).float()
    trajs = torch.matmul(rot, trajs)
    trajs = trajs.permute(0, 2, 1)
    trajs = trajs + torch.tensor([s_offset, d_offset]).to(device)

    min_val = 1e11
    min_ind = 0

    dense_trajs = project_to_cartesian_frame(trajs.reshape(1, trajs.shape[0] * trajs.shape[1], 2), ref_line).reshape(1, trajs.shape[0], trajs.shape[1], 2)#(B, N, n_future, 2)  
    return dense_trajs.to(save_device)

def extrapolate_points(points, num_points, num=6):
    """
        num_points 
            -> [1, 2, -1, -1, 5, 6] -> (-1) indicates tracking unavaiable for that timestep.
        points 
            -> [(), (), (), ()] 
    """
    total_points = []
    prev_timestep = 0
    count_unavailable = (np.array(num_points) == -1).sum() # number of points where it is unavailable.
    if not count_unavailable:
        for i in range(num):
            # ? is this possible
            total_points.append([-1e11, -1e11])
    elif count_unavailable == 1:
        for i in range(num):
            total_points.append(points[0])
    else:
        whe = np.where(points != -1)[0]

    dist = np.linalg.norm(points[1:] - points[:-1], axis=1).mean()
    
    slopes = np.zeros_like(num_points)
    cnt = 0
    for i in range(len(num_points) - 1):
        if num_points[i] != -1 and num_points[i + 1] != -1:
            # overwrite slope at these points
            slope[i] = num_points[i + 1] - num_points[i]

    return np.array(total_points)

def extract_obs_from_centerness(labels, n_present=3, num_obs=20, num=6, traj=None):
    """
        instance: (B, 9, 200, 200)
        centerness: (B, 9, 1, 200, 200)
        offset: (B, 9, 2, 200, 200)
        flow: (B, 9, 2, 200, 200)
        traj: (B, 3, 2)

        returns: 
            obs: (num_obs, num, 2) in bev space
    """
    bx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
    dx = np.array([0.5, 0.5])

    instance_seg = labels['instance'].detach().cpu().numpy()
    centerness = labels['centerness'].detach().cpu().numpy()
    offset = labels['offset'].detach().cpu().numpy()
    flow = labels['flow'].detach().cpu().numpy()
    traj = traj.detach().cpu().numpy() # (1, 3, 2) for B=1
    vel = np.linalg.norm(traj[:, -1] - traj[:, -2])
    traj = extend_both_sides(traj[0], resolution=vel, extension=num,reverse=False) # extend forward by num timesteps

    unique_ids = np.unique(labels['instance'].detach().cpu().numpy())
    instance_map = dict(zip(unique_ids, unique_ids))
    instance_colors = generate_instance_colours(instance_map)

    """
        instance: (6, 200, 200)
        centerness: (6, 1, 200, 200)
        offset: (6, 2, 200, 200)
        flow: (6, 2, 200, 200)
        we use centerness for now
    """
    instance = instance_seg[0, n_present - 1:]
    centerness = centerness[0, n_present - 1:, 0]
    offset = centerness[0, n_present - 1:, 0]
    flow = centerness[0, n_present - 1:, 0]

    instance_trajs = {}
    top_num_obs_ids = []
    for instance_id in unique_ids:
        if instance_id == 0:
            # background
            continue
        # compute across time
        points = []
        time_dict = {}
        num_points = []
        for i in range(0, 6):
            whe = np.where(instance[i] == instance_id)
            if len(whe[0]) == 0:
                # instance not present in scene
                num_points.append(-1)
                continue
            max_num = np.argmax(centerness[i][whe[0], whe[1]])
            points.append([whe[0][max_num], whe[1][max_num]])
            num_points.append(i)
        if len(points) == 0:
            for i in range(6):
                points.append([-1e11, -1e11])
        points = np.array(points)
        if len(points) == 6:
            instance_trajs[instance_id] = points

    x_obs = np.ones((num_obs, num)) * -5e5
    y_obs = np.ones((num_obs, num)) * -5e5
    obs_ids = list(instance_trajs.keys())
    distance_dict = {}
    for obs in obs_ids:
        points = np.array(instance_trajs[obs])
        distance_dict[obs] = np.linalg.norm(traj[0, -num:] - points)
    distance_dict = {k: v for k, v in sorted(distance_dict.items(), key=lambda item: item[1])}    # sort points by distance
    obs_ids = list(distance_dict.keys())
    obs_cnt = 0
    for instance_ids in obs_ids:
        if instance_ids == 0: continue
        points = np.array(instance_trajs[instance_ids])
        try:
            x_obs[obs_cnt] = points[:, 0]
            y_obs[obs_cnt] = points[:, 1]
            obs_cnt = obs_cnt + 1
        except Exception as e:
            print("Exception", e, x_obs.shape, points.shape)
            pass
        if obs_cnt >= num_obs:
            break
    for i in range(len(x_obs)):
        plt.plot(x_obs[i], y_obs[i])
    plt.plot(traj[:, 0], traj[:, 1])
    plt.savefig("output_vis/debug_trj.png")
    plt.clf()
    x_obs = torch.tensor(x_obs)
    y_obs = torch.tensor(y_obs)
    return torch.dstack((x_obs, y_obs)) # (num_obs, num, 2)

def interp_arc(t: int, px: np.ndarray, py: np.ndarray) -> np.ndarray:
    # px, py = eliminate_duplicates_2d(px, py)
    eq_spaced_points = np.linspace(0, 1, t)
    n = px.size
    assert px.size == py.size
    pxy = np.array((px, py)).T  # 2d polyline
    chordlen = np.linalg.norm(np.diff(pxy, axis=0), axis=1)
    chordlen = chordlen / np.sum(chordlen)
    cumarc = np.append(0, np.cumsum(chordlen))
    tbins = np.digitize(eq_spaced_points, cumarc)
    tbins[np.where((tbins <= 0) | (eq_spaced_points <= 0))] = 1
    tbins[np.where((tbins >= n) | (eq_spaced_points >= 1))] = n - 1
    s = np.divide((eq_spaced_points - cumarc[tbins - 1]), chordlen[tbins - 1])
    pt = pxy[tbins - 1, :] + np.multiply((pxy[tbins, :] - pxy[tbins - 1, :]), (np.vstack([s] * 2)).T)
    return pt
    

def get_polyline_length(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = torch.diff(polyline, axis=0)
    return float(torch.linalg.norm(offsets, axis=1).sum())

def interp_polyline_by_fixed_waypt_interval(polyline, waypt_interval):
    """Resample waypoints of a polyline so that waypoints appear roughly at fixed intervals from the start.
    Args:
        polyline: array pf shape (N,2) or (N,3) representing a polyline.
        waypt_interval: space interval between waypoints, in meters.
    Returns:
        interp_polyline: array of shape (N,2) or (N,3) representing a resampled/interpolated polyline.
        num_waypts: number of computed waypoints.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")

    # get the total length in meters of the line segment
    len_m = get_polyline_length(polyline)

    # count number of waypoints to get the desired length
    # add one for the extra endpoint
    num_waypts = torch.floor(torch.tensor(len_m / waypt_interval)) + 1
    interp_polyline = interp_arc(t=int(num_waypts.detach().cpu().item()), points=polyline)
    return interp_polyline, num_waypts

def get_polyline_length_np(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = np.diff(polyline, axis=0)
    return float(np.linalg.norm(offsets, axis=1).sum())

def interp_arc_np(t, points):
    """Linearly interpolate equally-spaced points along a polyline, either in 2d or 3d.

    We use a chordal parameterization so that interpolated arc-lengths
    will approximate original polyline chord lengths.
        Ref: M. Floater and T. Surazhsky, Parameterization for curve
            interpolation. 2005.
            https://www.mathworks.com/matlabcentral/fileexchange/34874-interparc

    For the 2d case, we remove duplicate consecutive points, since these have zero
    distance and thus cause division by zero in chord length computation.

    Args:
        t: number of points that will be uniformly interpolated and returned
        points: Numpy array of shape (N,2) or (N,3), representing 2d or 3d-coordinates of the arc.

    Returns:
        Numpy array of shape (N,2)

    Raises:
        ValueError: If `points` is not in R^2 or R^3.
    """
    if points.ndim != 2:
        raise ValueError("Input array must be (N,2) or (N,3) in shape.")

    # the number of points on the curve itself
    n, _ = points.shape

    # equally spaced in arclength -- the number of points that will be uniformly interpolated
    eq_spaced_points = np.linspace(0, 1, t)

    # Compute the chordal arclength of each segment.
    # Compute differences between each x coord, to get the dx's
    # Do the same to get dy's. Then the hypotenuse length is computed as a norm.
    chordlen= np.linalg.norm(np.diff(points, axis=0), axis=1)
    # Normalize the arclengths to a unit total
    chordlen = chordlen / np.sum(chordlen)
    # cumulative arclength

    cumarc = np.zeros(len(chordlen) + 1)
    cumarc[1:] = np.cumsum(chordlen)

    # which interval did each point fall in, in terms of eq_spaced_points? (bin index)
    tbins: NDArrayInt = np.digitize(eq_spaced_points, bins=cumarc).astype(int)

    # #catch any problems at the ends
    tbins[np.where((tbins <= 0) | (eq_spaced_points <= 0))] = 1
    tbins[np.where((tbins >= n) | (eq_spaced_points >= 1))] = n - 1

    s = np.divide((eq_spaced_points - cumarc[tbins - 1]), chordlen[tbins - 1])
    anchors = points[tbins - 1, :]
    # broadcast to scale each row of `points` by a different row of s
    offsets = (points[tbins, :] - points[tbins - 1, :]) * s.reshape(-1, 1)
    points_interp: NDArrayFloat = anchors + offsets

    return points_interp

def interp_polyline_by_fixed_waypt_interval_np(polyline, waypt_interval):
    """Resample waypoints of a polyline so that waypoints appear roughly at fixed intervals from the start.
    Args:
        polyline: array pf shape (N,2) or (N,3) representing a polyline.
        waypt_interval: space interval between waypoints, in meters.
    Returns:
        interp_polyline: array of shape (N,2) or (N,3) representing a resampled/interpolated polyline.
        num_waypts: number of computed waypoints.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")

    # get the total length in meters of the line segment
    len_m = get_polyline_length_np(polyline)

    # count number of waypoints to get the desired length
    # add one for the extra endpoint
    num_waypts = math.floor(len_m / waypt_interval) + 1
    interp_polyline = interp_arc_np(t=num_waypts, points=polyline)
    return interp_polyline, num_waypts

def get_polyline_length_np(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = np.diff(polyline, axis=0)
    return float(np.linalg.norm(offsets, axis=1).sum())

def extend_both_sides(oracle_centerline, resolution=1, extension=15, forward=True, reverse=True):
    """
        oracle_centerline: (N, 2)
        resolution (m)
        extension (m)
    """
    x, y = oracle_centerline[:, 0], oracle_centerline[:, 1]
    angle_start = np.arctan2(y[1] - y[0], x[1] - x[0])
    angle_end = np.arctan2(y[-1] - y[-2], x[-1] - x[-2])
    
    x_start = x[0] + np.cos(angle_start) * np.arange(- extension, 0, resolution)
    x_end = x[-1] + np.cos(angle_end) * np.arange(0, 0 + extension, resolution)

    y_start = y[0] + np.sin(angle_start) * np.arange(- extension, 0, resolution)
    y_end = y[-1] + np.sin(angle_end) * np.arange(0, 0 + extension, resolution)

    if forward and reverse:
        x = np.concatenate((x_start, x.reshape(x.shape[0]), x_end))
        y = np.concatenate((y_start, y.reshape(y.shape[0]), y_end))
    elif forward: 
        x = np.concatenate((x.reshape(x.shape[0]), x_end))
        y = np.concatenate((y.reshape(y.shape[0]), y_end))
    elif reverse:
        x = np.concatenate((x_start, x.reshape(x.shape[0])))
        y = np.concatenate((y_start, y.reshape(y.shape[0])))

    return np.dstack((x,y))[0]

def project_to_frenet_frame(traj, ref_line):
    distance_to_ref = torch.cdist(traj[:, :, :2], ref_line[:, :, :2])
    k = torch.argmin(distance_to_ref, dim=-1).view(-1, traj.shape[1], 1).expand(-1, -1, 3)
    ref_points = torch.gather(ref_line, 1, k)
    x_r, y_r, theta_r = ref_points[:, :, 0], ref_points[:, :, 1], ref_points[:, :, 2]
    x, y = traj[:, :, 0], traj[:, :, 1]
    s = 0.1 * (k[:, :, 0])
    l = torch.sign((y-y_r)*torch.cos(theta_r)-(x-x_r)*torch.sin(theta_r)) * torch.sqrt(torch.square(x-x_r)+torch.square(y-y_r))
    sl = torch.stack([s, l], dim=-1)

    return sl

def project_to_cartesian_frame(traj, ref_line):
    k = (10 * traj[:, :, 0]).long()
    k = torch.clip(k, 0, len(ref_line[0, :]) - 1)
    ref_points = torch.gather(ref_line, 1, k.view(-1, traj.shape[1], 1).expand(-1, -1, 3))
    x_r, y_r, theta_r = ref_points[:, :, 0], ref_points[:, :, 1], ref_points[:, :, 2]
    x = x_r - traj[:, :, 1] * torch.sin(theta_r)
    y = y_r + traj[:, :, 1] * torch.cos(theta_r)
    xy = torch.stack([x, y], dim=-1)
    return xy

def denoise(arr, w = 7):
    gt_x, gt_y = arr[:, 0], arr[:, 1]
    # denoising
    w = w
    gt_x_t = []
    gt_y_t = []
    for iq in range(len(gt_x)):
        if iq >= w and iq + w <= len(gt_x):
            gt_x_t.append(np.mean(gt_x[iq: iq + w]))
            gt_y_t.append(np.mean(gt_y[iq: iq + w]))
        elif iq < w:
            okx = np.mean(gt_x[w: w + w])
            gt_x_t.append(gt_x[0] + (okx - gt_x[0]) * (iq) / w)
            oky = np.mean(gt_y[w: w + w])
            gt_y_t.append(gt_y[0] + (oky - gt_y[0]) * (iq) / w)
        else:
            okx = np.mean(gt_x[len(gt_x) - w:len(gt_x) - w  + w])
            oky = np.mean(gt_y[len(gt_x) - w: len(gt_x) - w + w])
            gt_x_t.append(okx + (gt_x[-1] - okx) * (w - (len(gt_x) - iq)) / w)
            gt_y_t.append(oky + (gt_y[-1] - oky) * (w - (len(gt_y) - iq)) / w)                   

    gt_x = gt_x_t
    gt_y = gt_y_t
    print(gt_x, gt_y)
    return np.dstack((gt_x, gt_y))[0]

def generate_instance_colours(instance_map):
    # Most distinct 22 colors (kelly colors from https://stackoverflow.com/questions/470690/how-to-automatically-generate
    # -n-distinct-colors)
    # plus some colours from AD40k
    INSTANCE_COLOURS = np.asarray([
        [0, 0, 0],
        [255, 179, 0],
        [128, 62, 117],
        [255, 104, 0],
        [166, 189, 215],
        [193, 0, 32],
        [206, 162, 98],
        [129, 112, 102],
        [0, 125, 52],
        [246, 118, 142],
        [0, 83, 138],
        [255, 122, 92],
        [83, 55, 122],
        [255, 142, 0],
        [179, 40, 81],
        [244, 200, 0],
        [127, 24, 13],
        [147, 170, 0],
        [89, 51, 21],
        [241, 58, 19],
        [35, 44, 22],
        [112, 224, 255],
        [70, 184, 160],
        [153, 0, 255],
        [71, 255, 0],
        [255, 0, 163],
        [255, 204, 0],
        [0, 255, 235],
        [255, 0, 235],
        [255, 0, 122],
        [255, 245, 0],
        [10, 190, 212],
        [214, 255, 0],
        [0, 204, 255],
        [20, 0, 255],
        [255, 255, 0],
        [0, 153, 255],
        [0, 255, 204],
        [41, 255, 0],
        [173, 0, 255],
        [0, 245, 255],
        [71, 0, 255],
        [0, 255, 184],
        [0, 92, 255],
        [184, 255, 0],
        [255, 214, 0],
        [25, 194, 194],
        [92, 0, 255],
        [220, 220, 220],
        [255, 9, 92],
        [112, 9, 255],
        [8, 255, 214],
        [255, 184, 6],
        [10, 255, 71],
        [255, 41, 10],
        [7, 255, 255],
        [224, 255, 8],
        [102, 8, 255],
        [255, 61, 6],
        [255, 194, 7],
        [0, 255, 20],
        [255, 8, 41],
        [255, 5, 153],
        [6, 51, 255],
        [235, 12, 255],
        [160, 150, 20],
        [0, 163, 255],
        [140, 140, 140],
        [250, 10, 15],
        [20, 255, 0],
    ])

    return {instance_id: INSTANCE_COLOURS[global_instance_id % len(INSTANCE_COLOURS)] for
            instance_id, global_instance_id in instance_map.items()
            }


def is_overlapping_lane_seq(lane_seq1: Sequence[int], lane_seq2: Sequence[int]) -> bool:
    """
    Check if the 2 lane sequences are overlapping.
    Overlapping is defined as::
        s1------s2-----------------e1--------e2
    Here lane2 starts somewhere on lane 1 and ends after it, OR::
        s1------s2-----------------e2--------e1
    Here lane2 starts somewhere on lane 1 and ends before it
    Args:
        lane_seq1: list of lane ids
        lane_seq2: list of lane ids
    Returns:
        bool, True if the lane sequences overlap
    """

    if lane_seq2[0] in lane_seq1[1:] and lane_seq1[-1] in lane_seq2[:-1]:
        return True
    elif set(lane_seq2) <= set(lane_seq1):
        return True
    return False

def remove_overlapping_lane_seq(lane_seqs: List[List[int]]) -> List[List[int]]:
    """
    Remove lane sequences which are overlapping to some extent
    Args:
        lane_seqs (list of list of integers): List of list of lane ids (Eg. [[12345, 12346, 12347], [12345, 12348]])
    Returns:
        List of sequence of lane ids (e.g. ``[[12345, 12346, 12347], [12345, 12348]]``)
    """
    redundant_lane_idx: Set[int] = set()
    for i in range(len(lane_seqs)):
        for j in range(len(lane_seqs)):
            if i in redundant_lane_idx or i == j:
                continue
            if is_overlapping_lane_seq(lane_seqs[i], lane_seqs[j]):
                redundant_lane_idx.add(j)

    unique_lane_seqs = [lane_seqs[i] for i in range(len(lane_seqs)) if i not in redundant_lane_idx]
    return unique_lane_seqs

def dfs_incoming(nusc_map, lane_id, dist, threshold, resolution_meters=1):
    if dist > threshold:
        return [[lane_id]]
    else:
        traversed_lanes = []
        parent_lanes = (
                nusc_map.get_incoming_lane_ids(lane_id)
        )            
        if parent_lanes is not None:
            for parent in parent_lanes:
                # centerline = nusc_map.get_outgoing_lane_ids(child)
                try:
                    lane_record = nusc_map.get_arcline_path(parent)
                    centerline = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                    cl_length = LineString(centerline).length
                    curr_lane_ids = dfs_incoming(
                        nusc_map,
                        parent,
                        dist + cl_length,
                        threshold
                    )
                    traversed_lanes.extend(curr_lane_ids)
                except:
                    pass
        if len(traversed_lanes) == 0:
            return [[lane_id]]
        lanes_to_return = []
        for lane_seq in traversed_lanes:
            lanes_to_return.append(lane_seq + [lane_id])
        return lanes_to_return


def dfs(nusc_map, lane_id, dist, threshold, resolution_meters=1):
    if dist > threshold:
        return [[lane_id]]
    else:
        traversed_lanes = []
        child_lanes = (
                nusc_map.get_outgoing_lane_ids(lane_id)
        )            
        if child_lanes is not None:
            for child in child_lanes:
                # centerline = nusc_map.get_outgoing_lane_ids(child)
                try:
                    lane_record = nusc_map.get_arcline_path(child)
                    centerline = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                    cl_length = LineString(centerline).length
                    curr_lane_ids = dfs(
                        nusc_map,
                        child,
                        dist + cl_length,
                        threshold
                    )
                    traversed_lanes.extend(curr_lane_ids)
                except:
                    pass
        if len(traversed_lanes) == 0:
            return [[lane_id]]
        lanes_to_return = []
        for lane_seq in traversed_lanes:
            lanes_to_return.append([lane_id] + lane_seq)
        return lanes_to_return


def resize_and_crop_image(img, resize_dims, crop):
    # Bilinear resizing followed by cropping
    img = img.resize(resize_dims, resample=PIL.Image.BILINEAR)
    img = img.crop(crop)
    return img

def get_patch_coord(patch_box: Tuple[float, float, float, float],
                    patch_angle: float = 0.0) -> Polygon:
    """
    Convert patch_box to shapely Polygon coordinates.
    :param patch_box: Patch box defined as [x_center, y_center, height, width].
    :param patch_angle: Patch orientation in degrees.
    :return: Box Polygon for patch_box.
    """
    patch_x, patch_y, patch_h, patch_w = patch_box

    x_min = patch_x - patch_w / 2.0
    y_min = patch_y - patch_h / 2.0
    x_max = patch_x + patch_w / 2.0
    y_max = patch_y + patch_h / 2.0

    patch = box(x_min, y_min, x_max, y_max)
    patch = affinity.rotate(patch, patch_angle, origin=(patch_x, patch_y), use_radians=False)

    return patch


def update_intrinsics(intrinsics, top_crop=0.0, left_crop=0.0, scale_width=1.0, scale_height=1.0):
    """
    Parameters
    ----------
        intrinsics: torch.Tensor (3, 3)
        top_crop: float
        left_crop: float
        scale_width: float
        scale_height: float
    """
    updated_intrinsics = intrinsics.clone()
    # Adjust intrinsics scale due to resizing
    updated_intrinsics[0, 0] *= scale_width
    updated_intrinsics[0, 2] *= scale_width
    updated_intrinsics[1, 1] *= scale_height
    updated_intrinsics[1, 2] *= scale_height

    # Adjust principal point due to cropping
    updated_intrinsics[0, 2] -= left_crop
    updated_intrinsics[1, 2] -= top_crop

    return updated_intrinsics


def calculate_birds_eye_view_parameters(x_bounds, y_bounds, z_bounds):
    """
    Parameters
    ----------
        x_bounds: Forward direction in the ego-car.
        y_bounds: Sides
        z_bounds: Height

    Returns
    -------
        bev_resolution: Bird's-eye view bev_resolution
        bev_start_position Bird's-eye view first element
        bev_dimension Bird's-eye view tensor spatial dimension
    """
    bev_resolution = torch.tensor([row[2] for row in [x_bounds, y_bounds, z_bounds]])
    bev_start_position = torch.tensor([row[0] + row[2] / 2.0 for row in [x_bounds, y_bounds, z_bounds]])
    bev_dimension = torch.tensor([(row[1] - row[0]) / row[2] for row in [x_bounds, y_bounds, z_bounds]],
                                 dtype=torch.long)

    return bev_resolution, bev_start_position, bev_dimension


def convert_egopose_to_matrix_numpy(egopose):
    transformation_matrix = np.zeros((4, 4), dtype=np.float32)
    rotation = Quaternion(egopose['rotation']).rotation_matrix
    translation = np.array(egopose['translation'])
    transformation_matrix[:3, :3] = rotation
    transformation_matrix[:3, 3] = translation
    transformation_matrix[3, 3] = 1.0
    return transformation_matrix

def get_global_pose(rec, nusc, inverse=False):
    lidar_sample_data = nusc.get('sample_data', rec['data']['LIDAR_TOP'])

    sd_ep = nusc.get("ego_pose", lidar_sample_data["ego_pose_token"])
    sd_cs = nusc.get("calibrated_sensor", lidar_sample_data["calibrated_sensor_token"])
    if inverse is False:
        global_from_ego = transform_matrix(sd_ep["translation"], Quaternion(sd_ep["rotation"]), inverse=False)
        ego_from_sensor = transform_matrix(sd_cs["translation"], Quaternion(sd_cs["rotation"]), inverse=False)
        pose = global_from_ego.dot(ego_from_sensor)
    else:
        sensor_from_ego = transform_matrix(sd_cs["translation"], Quaternion(sd_cs["rotation"]), inverse=True)
        ego_from_global = transform_matrix(sd_ep["translation"], Quaternion(sd_ep["rotation"]), inverse=True)
        pose = sensor_from_ego.dot(ego_from_global)
    return pose

def invert_matrix_egopose_numpy(egopose):
    """ Compute the inverse transformation of a 4x4 egopose numpy matrix."""
    inverse_matrix = np.zeros((4, 4), dtype=np.float32)
    rotation = egopose[:3, :3]
    translation = egopose[:3, 3]
    inverse_matrix[:3, :3] = rotation.T
    inverse_matrix[:3, 3] = -np.dot(rotation.T, translation)
    inverse_matrix[3, 3] = 1.0
    return inverse_matrix


def mat2pose_vec(matrix: torch.Tensor):
    """
    Converts a 4x4 pose matrix into a 6-dof pose vector
    Args:
        matrix (ndarray): 4x4 pose matrix
    Returns:
        vector (ndarray): 6-dof pose vector comprising translation components (tx, ty, tz) and
        rotation components (rx, ry, rz)
    """

    # M[1, 2] = -sinx*cosy, M[2, 2] = +cosx*cosy
    rotx = torch.atan2(-matrix[..., 1, 2], matrix[..., 2, 2])

    # M[0, 2] = +siny, M[1, 2] = -sinx*cosy, M[2, 2] = +cosx*cosy
    cosy = torch.sqrt(matrix[..., 1, 2] ** 2 + matrix[..., 2, 2] ** 2)
    roty = torch.atan2(matrix[..., 0, 2], cosy)

    # M[0, 0] = +cosy*cosz, M[0, 1] = -cosy*sinz
    rotz = torch.atan2(-matrix[..., 0, 1], matrix[..., 0, 0])

    rotation = torch.stack((rotx, roty, rotz), dim=-1)

    # Extract translation params
    translation = matrix[..., :3, 3]
    return torch.cat((translation, rotation), dim=-1)


def euler2mat(angle: torch.Tensor):
    """Convert euler angles to rotation matrix.
    Reference: https://github.com/pulkitag/pycaffe-utils/blob/master/rot_utils.py#L174
    Args:
        angle: rotation angle along 3 axis (in radians) [Bx3]
    Returns:
        Rotation matrix corresponding to the euler angles [Bx3x3]
    """
    shape = angle.shape
    angle = angle.view(-1, 3)
    x, y, z = angle[:, 0], angle[:, 1], angle[:, 2]

    cosz = torch.cos(z)
    sinz = torch.sin(z)

    zeros = torch.zeros_like(z)
    ones = torch.ones_like(z)
    zmat = torch.stack([cosz, -sinz, zeros, sinz, cosz, zeros, zeros, zeros, ones], dim=1).view(-1, 3, 3)

    cosy = torch.cos(y)
    siny = torch.sin(y)

    ymat = torch.stack([cosy, zeros, siny, zeros, ones, zeros, -siny, zeros, cosy], dim=1).view(-1, 3, 3)

    cosx = torch.cos(x)
    sinx = torch.sin(x)

    xmat = torch.stack([ones, zeros, zeros, zeros, cosx, -sinx, zeros, sinx, cosx], dim=1).view(-1, 3, 3)

    rot_mat = xmat.bmm(ymat).bmm(zmat)
    rot_mat = rot_mat.view(*shape[:-1], 3, 3)
    return rot_mat


def pose_vec2mat(vec: torch.Tensor):
    """
    Convert 6DoF parameters to transformation matrix.
    Args:
        vec: 6DoF parameters in the order of tx, ty, tz, rx, ry, rz [B,6]
    Returns:
        A transformation matrix [B,4,4]
    """
    translation = vec[..., :3].unsqueeze(-1)  # [...x3x1]
    rot = vec[..., 3:].contiguous()  # [...x3]
    rot_mat = euler2mat(rot)  # [...,3,3]
    transform_mat = torch.cat([rot_mat, translation], dim=-1)  # [...,3,4]
    transform_mat = torch.nn.functional.pad(transform_mat, [0, 0, 0, 1], value=0)  # [...,4,4]
    transform_mat[..., 3, 3] = 1.0
    return transform_mat


def invert_pose_matrix(x):
    """
    Parameters
    ----------
        x: [B, 4, 4] batch of pose matrices

    Returns
    -------
        out: [B, 4, 4] batch of inverse pose matrices
    """
    assert len(x.shape) == 3 and x.shape[1:] == (4, 4), 'Only works for batch of pose matrices.'

    transposed_rotation = torch.transpose(x[:, :3, :3], 1, 2)
    translation = x[:, :3, 3:]

    inverse_mat = torch.cat([transposed_rotation, -torch.bmm(transposed_rotation, translation)], dim=-1) # [B,3,4]
    inverse_mat = torch.nn.functional.pad(inverse_mat, [0, 0, 0, 1], value=0)  # [B,4,4]
    inverse_mat[..., 3, 3] = 1.0
    return inverse_mat


def warp_features(x, flow, mode='nearest', spatial_extent=None):
    """ Applies a rotation and translation to feature map x.
        Args:
            x: (b, c, h, w) feature map
            flow: (b, 6) 6DoF vector (only uses the xy poriton)
            mode: use 'nearest' when dealing with categorical inputs
        Returns:
            in plane transformed feature map
        """
    if flow is None:
        return x
    b, c, h, w = x.shape
    # z-rotation
    angle = flow[:, 5].clone()  # torch.atan2(flow[:, 1, 0], flow[:, 0, 0])
    # x-y translation
    translation = flow[:, :2].clone()  # flow[:, :2, 3]

    # Normalise translation. Need to divide by how many meters is half of the image.
    # because translation of 1.0 correspond to translation of half of the image.
    translation[:, 0] /= spatial_extent[0]
    translation[:, 1] /= spatial_extent[1]
    # forward axis is inverted
    translation[:, 0] *= -1

    cos_theta = torch.cos(angle)
    sin_theta = torch.sin(angle)

    # output = Rot.input + translation
    # tx and ty are inverted as is the case when going from real coordinates to numpy coordinates
    # translation_pos_0 -> positive value makes the image move to the left
    # translation_pos_1 -> positive value makes the image move to the top
    # Angle -> positive value in rad makes the image move in the trigonometric way
    transformation = torch.stack([cos_theta, -sin_theta, translation[:, 1],
                                  sin_theta, cos_theta, translation[:, 0]], dim=-1).view(b, 2, 3)

    # Note that a rotation will preserve distances only if height = width. Otherwise there's
    # resizing going on. e.g. rotation of pi/2 of a 100x200 image will make what's in the center of the image
    # elongated.
    grid = torch.nn.functional.affine_grid(transformation, size=x.shape, align_corners=False)
    grid = grid.to(dtype=x.dtype)
    warped_x = torch.nn.functional.grid_sample(x, grid, mode=mode, padding_mode='zeros', align_corners=False)

    return warped_x


def cumulative_warp_features(x, flow, mode='nearest', spatial_extent=None):
    """ Warps a sequence of feature maps by accumulating incremental 2d flow.

    x[:, -1] remains unchanged
    x[:, -2] is warped using flow[:, -2]
    x[:, -3] is warped using flow[:, -3] @ flow[:, -2]
    ...
    x[:, 0] is warped using flow[:, 0] @ ... @ flow[:, -3] @ flow[:, -2]

    Args:
        x: (b, t, c, h, w) sequence of feature maps
        flow: (b, t, 6) sequence of 6 DoF pose
            from t to t+1 (only uses the xy poriton)

    """
    sequence_length = x.shape[1]
    if sequence_length == 1:
        return x

    flow = pose_vec2mat(flow)

    out = [x[:, -1]]
    cum_flow = flow[:, -2]
    for t in reversed(range(sequence_length - 1)):
        out.append(warp_features(x[:, t], mat2pose_vec(cum_flow), mode=mode, spatial_extent=spatial_extent))
        # @ is the equivalent of torch.bmm
        cum_flow = flow[:, t - 1] @ cum_flow

    return torch.stack(out[::-1], 1)


def cumulative_warp_features_reverse(x, flow, mode='nearest', spatial_extent=None):
    """ Warps a sequence of feature maps by accumulating incremental 2d flow.

    x[:, 0] remains unchanged
    x[:, 1] is warped using flow[:, 0].inverse()
    x[:, 2] is warped using flow[:, 0].inverse() @ flow[:, 1].inverse()
    ...

    Args:
        x: (b, t, c, h, w) sequence of feature maps
        flow: (b, t, 6) sequence of 6 DoF pose
            from t to t+1 (only uses the xy poriton)

    """
    flow = pose_vec2mat(flow)

    out = [x[:,0]]
    
    for i in range(1, x.shape[1]):
        if i==1:
            cum_flow = invert_pose_matrix(flow[:, 0])
        else:
            cum_flow = cum_flow @ invert_pose_matrix(flow[:,i-1])
        out.append( warp_features(x[:,i], mat2pose_vec(cum_flow), mode, spatial_extent=spatial_extent))
    return torch.stack(out, 1)


class VoxelsSumming(torch.autograd.Function):
    """Adapted from https://github.com/nv-tlabs/lift-splat-shoot/blob/master/src/tools.py#L193"""
    @staticmethod
    def forward(ctx, x, geometry, ranks):
        """The features `x` and `geometry` are ranked by voxel positions."""
        # Cumulative sum of all features.
        x = x.cumsum(0)

        # Indicates the change of voxel.
        mask = torch.ones(x.shape[0], device=x.device, dtype=torch.bool)
        mask[:-1] = ranks[1:] != ranks[:-1]

        x, geometry = x[mask], geometry[mask]
        # Calculate sum of features within a voxel.
        x = torch.cat((x[:1], x[1:] - x[:-1]))

        ctx.save_for_backward(mask)
        ctx.mark_non_differentiable(geometry)

        return x, geometry

    @staticmethod
    def backward(ctx, grad_x, grad_geometry):
        (mask,) = ctx.saved_tensors
        # Since the operation is summing, we simply need to send gradient
        # to all elements that were part of the summation process.
        indices = torch.cumsum(mask, 0)
        indices[mask] -= 1

        output_grad = grad_x[indices]

        return output_grad, None, None

def ___extract_trajs(centerlines, trj, nvar=11, viz=False, debugg=False, device='cpu', labels=None, griddim=None, problem=None, avoid_obs=False, obs=None):
    """
        this is just saved to understand where to insert the plotting code.
    """
    save_device = device
    device = 'cpu'
    bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
    ddx = np.array([0.5, 0.5])
    if griddim == None:
        X_MAX = 30
        X_MIN = 0
        Y_MAX = 2
        Y_MIN = -2
        NUM_X = 30
        NUM_Y = 5
    else:
        X_MIN, X_MAX, NUM_X, Y_MIN, Y_MAX, NUM_Y = griddim

    # get trajectories
    x = torch.linspace(X_MIN, X_MAX, NUM_X).to(device)
    y = torch.linspace(Y_MIN, Y_MAX, NUM_Y).to(device)
    xv, yv = torch.meshgrid(x, y)
    print(device)
    grid = torch.dstack((xv, yv)).to(device)
    ccc = centerlines.detach().cpu().numpy()
    cx, _ = interp_polyline_by_fixed_waypt_interval_np(ccc, 0.1)
    cx = extend_both_sides(cx, resolution=0.1, extension=100)
    cx = cx * ddx
    cx += bbx
    cx[:, [1, 0]] = cx[:, [0, 1]]
    cx[:, :1] = cx[:, :1] * -1
    cx = torch.tensor(cx).to(device)
    if viz:
        trr = trj.squeeze().detach().cpu()
        plt.plot(cx[:, 0], cx[:, 1], color='grey')
        plt.scatter(cx[:, 0], cx[:, 1], color='grey')
        plt.plot(trr[:, 0], trr[:, 1], color='blue')
        plt.scatter(trr[:, 0], trr[:, 1], color='blue')
        # print(trr[:, 0], trr[:, 1])
        plt.xlim([-5,5])
        plt.ylim([-5,5])
        plt.savefig('output_vis/debug_actual.png'); plt.clf()

    trr = trj.squeeze().detach().cpu()
    ref_line_ = torch.tensor(cx.reshape(1, cx.shape[0], 2))
    ref_line = torch.zeros((1, ref_line_.shape[1] - 1, 3))
    ref_line[:, :, :2] = ref_line_[:, 1:] 
    off = ref_line_[:, 1:] - ref_line_[:,:-1]
    ref_line[:, :, 2] = torch.atan2(off[:, :, 1], off[:, :, 0])
    ref_line = ref_line.float().to(device)
    trj = trj.float().to(device)

    test_ref_line=ref_line[:, 100:-100]
    sdd = project_to_frenet_frame(test_ref_line,  ref_line)
    if debugg: pdb.set_trace()
    if viz:
        plt.scatter(sdd[0, :, 0].detach().cpu(), sdd[0, :, 1].detach().cpu(), color='grey')
        plt.show()

    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)

    if avoid_obs:
        # obs -> (M, num, 2) : M obstacles for num timesteps
        # avoid obstacles
        num_obs = 20
        a_obs = 1
        b_obs = 1
        x_obs = obs[:, :, 0]
        y_obs = obs[:, :, 1]

        ssd_obs = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)
        problem = OPTNode(
            x_obs=x_obs, 
            y_obs=y_obs, 
            a_obs=a_obs, 
            b_obs=b_obs, 
            num_obs=num_obs, 
            rho_obs=1, 
            rho_eq=problem.rho_eq, 
            t_fin=problem.t_fin, 
            num=problem.num, 
            bernstein_coeff_order10_new=problem.bernstein_coeff_order10_new, 
            device = centerlines.device
        )



    cx = interp_arc_np(2000, cx.detach().cpu().numpy())
    cx = extend_both_sides(cx, resolution=0.1, extension=50)
    cx = torch.tensor(cx).to(device)
    if debugg: pdb.set_trace()

    if viz:
        gt = labels['gt_trajectory_prev'][0]
        plt.scatter(gt[:2, 0].detach().cpu(), gt[:2, 1].detach().cpu(), color='red')
        plt.plot(gt[:, 0].detach().cpu(), gt[:, 1].detach().cpu(), color='red', label='observed')
        plt.scatter(gt[2:, 0].detach().cpu(), gt[2:, 1].detach().cpu(), color='orange')
        plt.plot(gt[2:, 0].detach().cpu(), gt[2:, 1].detach().cpu(), color='orange', label='future')
        # plt.axis('equal')
        # plt.xlim([-5, 5])
        # plt.ylim([-5, 5])
        plt.legend()
        plt.savefig("output_vis/debug_mid.png");plt.clf()

    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)
    s_offset = ssd[:,2,0]
    d_offset = ssd[0,2,1]
    vel = ssd[0,2] - ssd[0,1]
    vel_prev = ssd[0,1] - ssd[0, 0]
    theta = torch.atan2(vel[1], vel[0]) # to make angle 0 wrt ego-agent
    theta_prev = torch.atan2(vel_prev[1], vel_prev[0])
    psidot_init = (theta - theta_prev)/0.5
    psi_init = theta
    v_init = torch.linalg.norm(vel)/0.5
    y_init = 0
    x_init = 0
    rot_inv = torch.tensor([[cos(-theta), -sin(-theta)], [sin(-theta), cos(-theta)]]).to(device)
    rot = torch.tensor([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]]).to(device)
    if debugg: pdb.set_trace()

    sdd_frame = sdd -  torch.tensor([s_offset, d_offset]).to(device)
    ssd_frame = ssd -  torch.tensor([s_offset, d_offset]).to(device)
    eps_frame = torch.tensor(grid).to(device) - torch.tensor([0, d_offset]).to(device) # (20, 7, 2)


    if debugg: pdb.set_trace()
    if viz:
        for eps in eps_frame:
            ep = eps.detach().cpu().numpy()
            plt.scatter(ep[:, 0], ep[:, 1])
        ss = ssd_frame[0].detach().cpu().numpy()
        plt.scatter(ss[:2, 0], ss[:2, 1], color='red')
        plt.scatter(ss[2:, 0], ss[2:, 1], color='orange')
        sd = sdd_frame[0].detach().cpu().numpy()
        plt.scatter(sd[:, 0], sd[:, 1], color='grey')
        plt.plot(sd[:, 0], sd[:, 1], color='grey')
        plt.savefig("output_vis/debug_ss.png"); plt.clf()
        plt.clf()

    if debugg: pdb.set_trace()
    # rotate
    sdd_frame[0] = torch.matmul(rot_inv, sdd_frame[0].T).T # (1, M, 2)
    ssd_frame = torch.matmul(rot_inv, ssd_frame[0].T).T # ()
    eps_frame = eps_frame.permute(1, 2, 0) # (7, 2, 20)
    eps_frame = torch.matmul(rot_inv, eps_frame) # (2, 2) x (7, 2, 20) -> (7, 2, 20)
    if debugg: pdb.set_trace()

    if viz:
        ss = ssd_frame.detach().cpu().numpy()
        plt.scatter(ss[:2, 0], ss[:2, 1], color='red')
        plt.scatter(ss[2:, 0], ss[2:, 1], color='orange')
        sd = sdd_frame[0].detach().cpu().numpy()
        plt.scatter(sd[:, 0], sd[:, 1], color='grey', label='centerline')
        plt.plot(sd[:, 0], sd[:, 1], color='grey')
        # for ep in eps_frame:
        #     plt.scatter(ep[:, 0], ep[:, 1], color='yellow')
        plt.scatter(eps_frame[:, 0, :].flatten().detach().cpu().numpy(), eps_frame[:, 1, :].flatten().detach().cpu().numpy(), color='yellow')
        # plt.savefig("output_vis/debug_ss.png"); plt.clf()
        plt.xlim([-25, 25])
        plt.ylim([-25, 25])
        plt.legend()
        plt.savefig("output_vis/debug_ss_rotated.png"); plt.clf()
        plt.clf()

    if debugg: pdb.set_trace()
    x_eps = eps_frame[:,0,:].flatten()
    y_eps = eps_frame[:,1,:].flatten()
    fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
    fixed_params = fixed_params.expand(len(x_eps), 5)
    variable_params = torch.zeros((len(x_eps), 3)).to(device)
    variable_params[:, 0] = torch.tensor(x_eps)
    variable_params[:, 1] = torch.tensor(y_eps)
    variable_params[:, 2] = torch.tensor(-theta)
    P = problem.P
    cc, _ = problem.solve(fixed_params.double().detach().cpu(), variable_params.double().detach().cpu())
    x_sol = torch.matmul(P, cc.T[:nvar]).T
    y_sol = torch.matmul(P, cc.T[nvar:2*nvar]).T

    if debugg: pdb.set_trace()

    if viz:
        ss = ssd_frame.detach().cpu().numpy()
        plt.scatter(eps_frame[:, 0, :].flatten().detach().cpu().numpy(), eps_frame[:, 1, :].flatten().detach().cpu().numpy(), color='yellow')
        sd = sdd_frame[0].detach().cpu().numpy()
        plt.scatter(sd[:, 0], sd[:, 1], color='grey', label='centerline')
        plt.plot(sd[:, 0], sd[:, 1], color='grey')
        for i in range(len(x_sol)):
            plt.plot(x_sol[i].flatten().detach().cpu().numpy(), y_sol[i].flatten().detach().cpu().numpy())
        # plt.savefig("output_vis/debug_ss.png"); plt.clf()
        plt.scatter(ss[:2, 0], ss[:2, 1], color='red')
        plt.scatter(ss[2:, 0], ss[2:, 1], color='orange')
        plt.xlim([-5, 5])
        plt.ylim([-5, 5])
        plt.legend()
        plt.savefig("output_vis/debug_frenet.png")
        plt.clf()

    if debugg: pdb.set_trace()
    trajs = torch.dstack((x_sol, y_sol)).to(device) # N, 7, 2
    trajs = trajs.permute(0, 2, 1).float()
    trajs = torch.matmul(rot, trajs)
    trajs = trajs.permute(0, 2, 1)
    trajs = trajs + torch.tensor([s_offset, d_offset]).to(device)

    min_val = 1e11
    min_ind = 0

    dense_trajs = project_to_cartesian_frame(trajs.reshape(1, trajs.shape[0] * trajs.shape[1], 2), ref_line).reshape(1, trajs.shape[0], trajs.shape[1], 2)#(B, N, n_future, 2)  
    return dense_trajs.to(save_device)


data/stp3/utils/geometry.py
llmbev/talk2bev
import math
import PIL
import numpy as np
import torch
import torch.nn.functional as F
from typing import Tuple
import matplotlib.pyplot as plt
import pdb

from pyquaternion import Quaternion
from shapely import affinity
from shapely.geometry import Polygon, MultiPolygon, LineString, Point, box
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.map_expansion.arcline_path_utils import discretize_lane, ArcLinePath
from shapely.geometry import LineString
from skimage.draw import polygon
from math import cos, sin
from stp3.utils.Optnode_obs_unbatched import OPTNode_obs
from stp3.utils.Optnode import OPTNode, bernstein_coeff_order10_new
from stp3.utils.Optnode_naive import OPTNode_batched
from stp3.utils.Optnode_waypoint import OPTNode_waypoint

import pdb

from typing import Iterable, List, Sequence, Set, Tuple

n_present = 3

W = 1.85
H = 4.084
bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
ddx = np.array([0.5, 0.5])
bx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
dx = np.array([0.5, 0.5])

def get_rot(theta):
    rot = np.array([[math.cos(theta), -math.sin(theta)], [math.sin(theta), math.cos(theta)]])
    return rot

def get_origin_points(lambda_=0, theta=0):
    pts = np.array([
        [-H / 2. + 0.5 - lambda_, W / 2. + lambda_],
        [H / 2. + 0.5 + lambda_, W / 2. + lambda_],
        [H / 2. + 0.5 + lambda_, -W / 2. - lambda_],
        [-H / 2. + 0.5 - lambda_, -W / 2. - lambda_],
    ])
    rot = get_rot(theta)
    pts = np.dot(rot, pts.T).T
    pts = (pts - bx) / (dx)    
    pts[:, [0, 1]] = pts[:, [1, 0]]
    rr , cc = polygon(pts[:,1], pts[:,0])
    rc = np.concatenate([rr[:,None], cc[:,None]], axis=-1)
    return rc # (27,2)

def get_points(self, pts, lambda_=0, bev_dimension=None):
    '''
    pts: np array (n_future, 2)
    return:
    List[ torch.Tensor<int> (B, N, n_future), torch.Tensor<int> (B, N, n_future)]
    '''
    n_future, _ = pts.shape
    rc = get_origin_points(lambda_)
    if len(pts) > 1:
        theta = pts[1:] - pts[:-1]
    else:
        theta = 0
    pts = pts + rc
    pts = torch.tensor(pts)

    rr = pts[:,0].long()
    rr = torch.clamp(rr, 0, bev_dimension[0] - 1)

    cc = pts[:,1].long()
    cc = torch.clamp(cc, 0, bev_dimension[1] - 1)

    return rr, cc

def extract_trajs(centerlines, trj, seg_prediction=None, hdmap_prediction=None, nvar=11, viz=False, debugg=False, device='cpu', labels=None, griddim=None, problem=None, avoid_obs=False, obs=None, a_obs=1, b_obs=1, ind="0", bev_dimension=None):
    save_device = device
    device = 'cpu'
    save_obs = obs
    bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
    ddx = np.array([0.5, 0.5])
    if griddim == None:
        X_MAX = 50
        X_MIN = 0
        NUM_X = 30
        Y_MAX = 5
        Y_MIN = -5
        NUM_Y = 5
    else:
        X_MIN, X_MAX, NUM_X, Y_MIN, Y_MAX, NUM_Y = griddim


    # get trajectories
    x = torch.linspace(X_MIN, X_MAX, NUM_X).to(device)
    y = torch.linspace(Y_MIN, Y_MAX, NUM_Y).to(device)
    xv, yv = torch.meshgrid(x, y)
    print(device)
    grid = torch.dstack((xv, yv)).to(device)
    ccc = centerlines.detach().cpu().numpy()
    cx, _ = interp_polyline_by_fixed_waypt_interval_np(ccc, 0.1)
    cx = extend_both_sides(cx, resolution=0.1, extension=500)
    cx = cx * ddx
    cx += bbx
    cx[:, [1, 0]] = cx[:, [0, 1]]
    cx[:, :1] = cx[:, :1] * -1
    cx = torch.tensor(cx).to(device)

    trr = trj.squeeze().detach().cpu()
    ref_line_ = torch.tensor(cx.reshape(1, cx.shape[0], 2))
    ref_line = torch.zeros((1, ref_line_.shape[1] - 1, 3))
    ref_line[:, :, :2] = ref_line_[:, 1:] 
    off = ref_line_[:, 1:] - ref_line_[:,:-1]
    ref_line[:, :, 2] = torch.atan2(off[:, :, 1], off[:, :, 0])
    ref_line = ref_line.float().to(device)
    trj = trj.float().to(device)

    test_ref_line=ref_line[:, 100:-100]
    sdd = project_to_frenet_frame(test_ref_line,  ref_line)

    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)


    cx = interp_arc_np(2000, cx.detach().cpu().numpy())
    cx = extend_both_sides(cx, resolution=0.1, extension=50)
    cx = torch.tensor(cx).to(device)
    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)
    s_offset = ssd[:,2,0]
    d_offset = ssd[0,2,1]
    vel = ssd[0,2] - ssd[0,1]
    vel_prev = ssd[0,1] - ssd[0, 0]
    theta = torch.atan2(vel[1], vel[0]) # to make angle 0 wrt ego-agent
    theta_prev = torch.atan2(vel_prev[1], vel_prev[0])
    psidot_init = (theta - theta_prev)/0.5
    psi_init = theta
    v_init = torch.linalg.norm(vel)/0.5
    y_init = 0
    x_init = 0
    rot_inv = torch.tensor([[cos(-theta), -sin(-theta)], [sin(-theta), cos(-theta)]]).to(device)
    rot = torch.tensor([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]]).to(device)


    # optimized smapling
    grid_ref = grid.reshape(1, NUM_X * NUM_Y, 2)
    s_offset = ssd[:,2,0]
    grid_ref = grid_ref + torch.tensor([s_offset, 0])
    grid_ref_cartesian = project_to_cartesian_frame(grid_ref, ref_line) # in cartesian frame
    grid_ref_cartesian = (grid_ref_cartesian[0, :, :2])/ ddx
    grid_ref_cartesian[ :, :1] = grid_ref_cartesian[:, :1] * -1
    grid_ref_cartesian[:, [0, 1]] = grid_ref_cartesian[:, [1, 0]]
    # pdb.set_trace()

    rectangle = get_origin_points()
    grid_ref_cartesian = grid_ref_cartesian.unsqueeze(1) + rectangle
    gc = grid_ref_cartesian.long()
    hdmap = hdmap_prediction[0, 1]
    seg = seg_prediction[0, -1, 0]
    map_vals = hdmap[gc[:,:,0], gc[:,:,1]].sum(-1)
    seg_vals = seg[gc[:,:,0], gc[:,:,1]].sum(-1)

    cutoff = 0
    take_map = (map_vals >= (len(rectangle) - 1 - cutoff)) # do not take any point if any portion outside hdmap prediction

    seg_map = (seg_vals == 0) # do not take any point if any portion intersects with segmentation prediction

    final_grid = []
    grid_ref = grid.reshape(1, NUM_X * NUM_Y, 2)
    for i in range(len(grid_ref[0])):
        if take_map[i] and seg_map[i]:
            print(grid_ref[0][i])
            final_grid.append(grid_ref[0][i].detach().cpu().numpy())
    final_grid = np.array(final_grid)
    grid = final_grid


    # translate
    sdd_frame = sdd -  torch.tensor([s_offset, d_offset]).to(device)
    ssd_frame = ssd -  torch.tensor([s_offset, d_offset]).to(device)
    eps_frame = torch.tensor(grid).to(device) - torch.tensor([0, d_offset]).to(device) # (20, 7, 2)

    # rotate
    sdd_frame[0] = torch.matmul(rot_inv, sdd_frame[0].T).T # (1, M, 2)
    ssd_frame = torch.matmul(rot_inv, ssd_frame[0].T).T # ()
    # eps_frame = eps_frame.permute(1, 2, 0) # (7, 2, 20)
    eps_frame = torch.matmul(rot_inv, eps_frame.T).T # (2, 2) x (M, 2) -> (M, 2)

    """
        save params of the optimizer before calling solve
    """
    # obs -> (M, num, 2) : M obstacles for num timesteps
    # avoid obstacles
    num_obs = obs.shape[0]
    num = obs.shape[1]
    # obs = obs.reshape(1, num_obs * num, 2) # (1, M * num, 2)

    # first bring to agent frame
    obs = obs * ddx
    obs += bbx
    obs[:, :, [1, 0]] = obs[:, :, [0, 1]]
    obs[:, :, :1] = obs[:, :, :1] * -1

    ssd_obs = torch.zeros((1, num_obs, num, 2))
    for n in range(num_obs):
        print(obs[n].unsqueeze(0).shape, "OK")
        print(project_to_frenet_frame(obs[n].unsqueeze(0).double().to(device), ref_line.double().to(device)).shape)
        ssd_obs[:, n] = project_to_frenet_frame(obs[n].unsqueeze(0).double().to(device), ref_line.double().to(device)) # (1, M * num, 2)

    # translate
    ssd_obs_frame = ssd_obs[0] - torch.tensor([s_offset, d_offset]).to(device) # (M, num, 2)
    # print(ssd_obs.shape, ssd_obs_frame.shape, " OBS")
    for i in range(num_obs):
        ssd_obs_frame[i] = torch.matmul(rot_inv, ssd_obs_frame[i].float().T).T
    # ssd_obs_frame = ssd_obs_frame[0].permute(1, 2, 0) # (num, 2, M)
    # ssd_obs_frame = torch.matmul(rot_inv, ssd_obs_frame.float()) # (2, 2) x (num, 2, M) -> (num, 2, M)
    # ssd_obs_frame = ssd_obs_frame.reshape(num_obs, num, 2)  # (M, num, 2)
    x_obs = ssd_obs_frame[:, : ,0]
    y_obs = ssd_obs_frame[:, : ,1]
    # x_obs, y_obs = problem.x_obs, problem.y_obs
    obs_frame = torch.dstack((x_obs, y_obs))

    x_eps = eps_frame[:,0].flatten()
    y_eps = eps_frame[:,1].flatten()

    fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
    fixed_params = fixed_params.expand(len(x_eps), 5)
    variable_params = torch.zeros((len(x_eps), 3)).to(device)
    variable_params[:, 0] = torch.tensor(x_eps)
    variable_params[:, 1] = torch.tensor(y_eps)
    variable_params[:, 2] = torch.tensor(-theta)

    """
        prepare x_obs and y_obs
    """
    num_obs = len(obs_frame)
    num_local_obs = 61
    x_obs = np.ones((num_obs, num_local_obs))
    y_obs = np.ones((num_obs, num_local_obs))
    b_inp = np.ones((len(fixed_params), 12))
    for i, ob in enumerate(obs_frame.detach().cpu().numpy()):
        # extrapolating obstacle points to (61, 2)
        x_obs[i][0] = ob[0, 0]
        x_obs[i][-1] = ob[-1, 0]
        y_obs[i][0] = ob[0, 1]
        y_obs[i][-1] = ob[-1, 1]
        pts = interp_arc_np(t=num_local_obs - 2 + 1, points=ob[1:-1])
        pts = pts[1:] # nan at start for some reason
        x_obs[i][1:-1] = pts[:, 0]
        y_obs[i][1:-1] = pts[:, 1]
        if viz:
            plt.plot(ob[:, 0], ob[:, 1])
            plt.scatter(ob[:, 0], ob[:, 1])
    print(x_obs.shape, y_obs.shape)
    for i in range(len(fixed_params)):
        x_init, y_init = fixed_params[i,0], fixed_params[i, 1]
        x_fin, y_fin = variable_params[i, 0], variable_params[i, 1]
        b_inp[i] = np.array([x_init, 0, 0, x_fin, 0, 0, y_init, 0, 0, y_fin, 0, 0])

    import pdb; pdb.set_trace()
    """
        prepare Optnode, refine collision bound trajectories
    """
    t_fin = 2
    a_obs = 2.0
    b_obs = 2.0
    # rho_obs = 0.3
    rho_obs = 1.2
    rho_eq = 10.0
    weight_smoothness = 10
    tot_time = np.linspace(0.0, t_fin, num=num_local_obs)
    tot_time_copy = tot_time.reshape(num_local_obs, 1)
    P, Pdot, Pddot = bernstein_coeff_order10_new(10, tot_time_copy[0], tot_time_copy[-1], tot_time_copy)
    A_obs = np.tile(P, (num_obs, 1))
    A_eq = np.vstack((P[0], Pdot[0], Pddot[0], P[-1], Pdot[-1], Pddot[-1]))
    Q_smoothness = np.dot(Pddot.T, Pddot)
    a_obs = 2.0
    b_obs = 2.0
    # rho_obs = 0.3
    rho_obs = 1.2
    rho_eq = 10.0
    weight_smoothness = 10
    batch_size = len(fixed_params)
    x_obs = np.ones((batch_size, num_obs, num_local_obs)) * x_obs
    y_obs = np.ones((batch_size, num_obs, num_local_obs)) * y_obs
    opt_node = OPTNode_batched(P, Pddot, A_eq, A_obs, Q_smoothness, x_obs, y_obs,num=num_local_obs, num_obs=num_obs, nvar=nvar, a_obs=a_obs, b_obs=b_obs, rho_obs=2, rho_eq=10,weight_smoothness=weight_smoothness, maxiter=3000, eps=1e-7, num_tot=num_local_obs*num_obs,batch_size=batch_size)
    lamda_x = torch.zeros(batch_size, nvar, dtype=torch.double)
    lamda_y = torch.zeros(batch_size, nvar, dtype=torch.double)
    sol = opt_node.solve(torch.tensor(b_inp), lamda_x, lamda_y)[0]
    import pdb; pdb.set_trace()
    x_pred = torch.matmul(torch.tensor(P), sol[:, :nvar].transpose(0, 1)) # (num_local_obs, batch_size)
    y_pred = torch.matmul(torch.tensor(P), sol[:, nvar:2*nvar].transpose(0, 1))# (num_local_obs, batch_size)
    x_pred = x_pred.T
    y_pred = y_pred.T

    not_to_refine_inds = []
    refine_inds = []
    for i in range(len(x_pred)):
        torch.atan2(y_pred[1:, i] - y_pred[:-1, i], x_pred[1:, i] - x_pred[:-1, i])
        unique_thetas = len(torch.unique(torch.clamp(thetas - thetas[0], 1e-5, 1e11)))
        if len(unique_thetas) == 1:
            # straight line
            not_to_refine_inds.append(i)
        else:
            refine_inds.append(i)
    import pdb; pdb.set_trace()

    x_pred_waypt = x_pred[refine_inds]
    y_pred_waypt = y_pred[refine_inds]
    eps_frame_nonholo = eps_frame[not_to_refine_inds]
    eps_frame_waypt = eps_frame[refine_inds]
    torch.save(x_obs, f"obs/x_obs_{ind}.pt")
    torch.save(y_obs, f"obs/y_obs_{ind}.pt")
    torch.save(obs_frame, f"obs/obs_{ind}.pt")
    torch.save(fixed_params, f"obs/fixed_params_{ind}.pt")
    torch.save(variable_params, f"obs/variable_params_{ind}.pt")

    import pdb; pdb.set_trace()
    problem = OPTNode(t_fin=problem.t_fin, num=problem.num, device='cpu')
    x_eps = eps_frame_nonholo[:,0].flatten()
    y_eps = eps_frame_nonholo[:,1].flatten()
    fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
    fixed_params = fixed_params.expand(len(eps_frame_nonholo), 5)
    variable_params = torch.zeros((len(eps_frame_nonholo), 3)).to(device)
    variable_params[:, 0] = torch.tensor(x_eps)
    variable_params[:, 1] = torch.tensor(y_eps)
    variable_params[:, 2] = torch.tensor(-theta)
    P = problem.P
    cc, _ = problem.solve(fixed_params.double(), variable_params.double())
    x_sol = torch.matmul(P, cc.T[:nvar]).T
    y_sol = torch.matmul(P, cc.T[nvar:2*nvar]).T    
    trajs = torch.dstack((x_sol, y_sol))
    import pdb; pdb.set_trace()



    """
        use waypoint optimizer to refine obstacle avoided trajectories
    """
    if len(eps_frame_waypt) > 1:
        x_eps = eps_frame_waypt[:,0].flatten()
        y_eps = eps_frame_waypt[:,1].flatten()
        fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
        fixed_params = fixed_params.expand(len(eps_frame_waypt), 5)
        variable_params = torch.zeros((len(eps_frame_waypt), 12)).to(device)
        variable_params[:, 0] = torch.tensor(x_eps)
        variable_params[:, 1] = torch.tensor(y_eps)
        variable_params[:, 2] = torch.tensor(-theta)
        variable_params[:, 3] = torch.tensor(x_pred_waypt[:, 10])
        variable_params[:, 4] = torch.tensor(y_pred_waypt[:, 10])
        variable_params[:, 5] = torch.tensor(x_pred_waypt[:, 20])
        variable_params[:, 6] = torch.tensor(y_pred_waypt[:, 20])
        variable_params[:, 7] = torch.tensor(x_pred_waypt[:, 30])
        variable_params[:, 8] = torch.tensor(y_pred_waypt[:, 30])
        variable_params[:, 9] = torch.tensor(x_pred_waypt[:, 40])
        variable_params[:, 10] = torch.tensor(y_pred_waypt[:, 40])
        variable_params[:, 11] = torch.tensor(x_pred_waypt[:, 50])
        variable_params[:, 12] = torch.tensor(y_pred_waypt[:, 50])
        problem_waypt = OPTNode_waypoint(rho_eq=1.1, num=61, t_fin=t_fin, bernstein_coeff_order10_new=bernstein_coeff_order10_new)
        cc, _  = problem_waypt.solve(fixed_params.double(), variable_params.double())

    trajs = trajs.permute(0, 2, 1).float()
    trajs = torch.matmul(rot, trajs)
    trajs = trajs.permute(0, 2, 1)
    trajs = trajs + torch.tensor([s_offset, d_offset]).to(device)

    min_val = 1e11
    min_ind = 0

    dense_trajs = project_to_cartesian_frame(trajs.reshape(1, trajs.shape[0] * trajs.shape[1], 2), ref_line).reshape(1, trajs.shape[0], trajs.shape[1], 2)#(B, N, n_future, 2)  
    return dense_trajs.to(save_device)

def extrapolate_points(points, num_points, num=6):
    """
        num_points 
            -> [1, 2, -1, -1, 5, 6] -> (-1) indicates tracking unavaiable for that timestep.
        points 
            -> [(), (), (), ()] 
    """
    total_points = []
    prev_timestep = 0
    count_unavailable = (np.array(num_points) == -1).sum() # number of points where it is unavailable.
    if not count_unavailable:
        for i in range(num):
            # ? is this possible
            total_points.append([-1e11, -1e11])
    elif count_unavailable == 1:
        for i in range(num):
            total_points.append(points[0])
    else:
        whe = np.where(points != -1)[0]

    dist = np.linalg.norm(points[1:] - points[:-1], axis=1).mean()
    
    slopes = np.zeros_like(num_points)
    cnt = 0
    for i in range(len(num_points) - 1):
        if num_points[i] != -1 and num_points[i + 1] != -1:
            # overwrite slope at these points
            slope[i] = num_points[i + 1] - num_points[i]

    return np.array(total_points)

def extract_obs_from_centerness(labels, n_present=3, num_obs=20, num=6, traj=None):
    """
        instance: (B, 9, 200, 200)
        centerness: (B, 9, 1, 200, 200)
        offset: (B, 9, 2, 200, 200)
        flow: (B, 9, 2, 200, 200)
        traj: (B, 3, 2)

        returns: 
            obs: (num_obs, num, 2) in bev space
    """
    bx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
    dx = np.array([0.5, 0.5])

    instance_seg = labels['instance'].detach().cpu().numpy()
    centerness = labels['centerness'].detach().cpu().numpy()
    offset = labels['offset'].detach().cpu().numpy()
    flow = labels['flow'].detach().cpu().numpy()
    traj = traj.detach().cpu().numpy() # (1, 3, 2) for B=1
    vel = np.linalg.norm(traj[:, -1] - traj[:, -2])
    traj = extend_both_sides(traj[0], resolution=vel, extension=num,reverse=False) # extend forward by num timesteps

    unique_ids = np.unique(labels['instance'].detach().cpu().numpy())
    instance_map = dict(zip(unique_ids, unique_ids))
    instance_colors = generate_instance_colours(instance_map)

    """
        instance: (6, 200, 200)
        centerness: (6, 1, 200, 200)
        offset: (6, 2, 200, 200)
        flow: (6, 2, 200, 200)
        we use centerness for now
    """
    instance = instance_seg[0, n_present - 1:]
    centerness = centerness[0, n_present - 1:, 0]
    offset = centerness[0, n_present - 1:, 0]
    flow = centerness[0, n_present - 1:, 0]

    instance_trajs = {}
    top_num_obs_ids = []
    for instance_id in unique_ids:
        if instance_id == 0:
            # background
            continue
        # compute across time
        points = []
        time_dict = {}
        num_points = []
        for i in range(0, 6):
            whe = np.where(instance[i] == instance_id)
            if len(whe[0]) == 0:
                # instance not present in scene
                num_points.append(-1)
                continue
            max_num = np.argmax(centerness[i][whe[0], whe[1]])
            points.append([whe[0][max_num], whe[1][max_num]])
            num_points.append(i)
        if len(points) == 0:
            for i in range(6):
                points.append([-1e11, -1e11])
        points = np.array(points)
        if len(points) == 6:
            instance_trajs[instance_id] = points

    x_obs = np.ones((num_obs, num)) * -5e5
    y_obs = np.ones((num_obs, num)) * -5e5
    obs_ids = list(instance_trajs.keys())
    distance_dict = {}
    for obs in obs_ids:
        points = np.array(instance_trajs[obs])
        distance_dict[obs] = np.linalg.norm(traj[0, -num:] - points)
    distance_dict = {k: v for k, v in sorted(distance_dict.items(), key=lambda item: item[1])}    # sort points by distance
    obs_ids = list(distance_dict.keys())
    obs_cnt = 0
    for instance_ids in obs_ids:
        if instance_ids == 0: continue
        points = np.array(instance_trajs[instance_ids])
        try:
            x_obs[obs_cnt] = points[:, 0]
            y_obs[obs_cnt] = points[:, 1]
            obs_cnt = obs_cnt + 1
        except Exception as e:
            print("Exception", e, x_obs.shape, points.shape)
            pass
        if obs_cnt >= num_obs:
            break
    for i in range(len(x_obs)):
        plt.plot(x_obs[i], y_obs[i])
    plt.plot(traj[:, 0], traj[:, 1])
    plt.savefig("output_vis/debug_trj.png")
    plt.clf()
    x_obs = torch.tensor(x_obs)
    y_obs = torch.tensor(y_obs)
    return torch.dstack((x_obs, y_obs)) # (num_obs, num, 2)

def interp_arc(t: int, px: np.ndarray, py: np.ndarray) -> np.ndarray:
    # px, py = eliminate_duplicates_2d(px, py)
    eq_spaced_points = np.linspace(0, 1, t)
    n = px.size
    assert px.size == py.size
    pxy = np.array((px, py)).T  # 2d polyline
    chordlen = np.linalg.norm(np.diff(pxy, axis=0), axis=1)
    chordlen = chordlen / np.sum(chordlen)
    cumarc = np.append(0, np.cumsum(chordlen))
    tbins = np.digitize(eq_spaced_points, cumarc)
    tbins[np.where((tbins <= 0) | (eq_spaced_points <= 0))] = 1
    tbins[np.where((tbins >= n) | (eq_spaced_points >= 1))] = n - 1
    s = np.divide((eq_spaced_points - cumarc[tbins - 1]), chordlen[tbins - 1])
    pt = pxy[tbins - 1, :] + np.multiply((pxy[tbins, :] - pxy[tbins - 1, :]), (np.vstack([s] * 2)).T)
    return pt
    

def get_polyline_length(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = torch.diff(polyline, axis=0)
    return float(torch.linalg.norm(offsets, axis=1).sum())

def interp_polyline_by_fixed_waypt_interval(polyline, waypt_interval):
    """Resample waypoints of a polyline so that waypoints appear roughly at fixed intervals from the start.
    Args:
        polyline: array pf shape (N,2) or (N,3) representing a polyline.
        waypt_interval: space interval between waypoints, in meters.
    Returns:
        interp_polyline: array of shape (N,2) or (N,3) representing a resampled/interpolated polyline.
        num_waypts: number of computed waypoints.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")

    # get the total length in meters of the line segment
    len_m = get_polyline_length(polyline)

    # count number of waypoints to get the desired length
    # add one for the extra endpoint
    num_waypts = torch.floor(torch.tensor(len_m / waypt_interval)) + 1
    interp_polyline = interp_arc(t=int(num_waypts.detach().cpu().item()), points=polyline)
    return interp_polyline, num_waypts

def get_polyline_length_np(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = np.diff(polyline, axis=0)
    return float(np.linalg.norm(offsets, axis=1).sum())

def interp_arc_np(t, points):
    """Linearly interpolate equally-spaced points along a polyline, either in 2d or 3d.

    We use a chordal parameterization so that interpolated arc-lengths
    will approximate original polyline chord lengths.
        Ref: M. Floater and T. Surazhsky, Parameterization for curve
            interpolation. 2005.
            https://www.mathworks.com/matlabcentral/fileexchange/34874-interparc

    For the 2d case, we remove duplicate consecutive points, since these have zero
    distance and thus cause division by zero in chord length computation.

    Args:
        t: number of points that will be uniformly interpolated and returned
        points: Numpy array of shape (N,2) or (N,3), representing 2d or 3d-coordinates of the arc.

    Returns:
        Numpy array of shape (N,2)

    Raises:
        ValueError: If `points` is not in R^2 or R^3.
    """
    if points.ndim != 2:
        raise ValueError("Input array must be (N,2) or (N,3) in shape.")

    # the number of points on the curve itself
    n, _ = points.shape

    # equally spaced in arclength -- the number of points that will be uniformly interpolated
    eq_spaced_points = np.linspace(0, 1, t)

    # Compute the chordal arclength of each segment.
    # Compute differences between each x coord, to get the dx's
    # Do the same to get dy's. Then the hypotenuse length is computed as a norm.
    chordlen= np.linalg.norm(np.diff(points, axis=0), axis=1)
    # Normalize the arclengths to a unit total
    chordlen = chordlen / np.sum(chordlen)
    # cumulative arclength

    cumarc = np.zeros(len(chordlen) + 1)
    cumarc[1:] = np.cumsum(chordlen)

    # which interval did each point fall in, in terms of eq_spaced_points? (bin index)
    tbins: NDArrayInt = np.digitize(eq_spaced_points, bins=cumarc).astype(int)

    # #catch any problems at the ends
    tbins[np.where((tbins <= 0) | (eq_spaced_points <= 0))] = 1
    tbins[np.where((tbins >= n) | (eq_spaced_points >= 1))] = n - 1

    s = np.divide((eq_spaced_points - cumarc[tbins - 1]), chordlen[tbins - 1])
    anchors = points[tbins - 1, :]
    # broadcast to scale each row of `points` by a different row of s
    offsets = (points[tbins, :] - points[tbins - 1, :]) * s.reshape(-1, 1)
    points_interp: NDArrayFloat = anchors + offsets

    return points_interp

def interp_polyline_by_fixed_waypt_interval_np(polyline, waypt_interval):
    """Resample waypoints of a polyline so that waypoints appear roughly at fixed intervals from the start.
    Args:
        polyline: array pf shape (N,2) or (N,3) representing a polyline.
        waypt_interval: space interval between waypoints, in meters.
    Returns:
        interp_polyline: array of shape (N,2) or (N,3) representing a resampled/interpolated polyline.
        num_waypts: number of computed waypoints.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")

    # get the total length in meters of the line segment
    len_m = get_polyline_length_np(polyline)

    # count number of waypoints to get the desired length
    # add one for the extra endpoint
    num_waypts = math.floor(len_m / waypt_interval) + 1
    interp_polyline = interp_arc_np(t=num_waypts, points=polyline)
    return interp_polyline, num_waypts

def get_polyline_length_np(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = np.diff(polyline, axis=0)
    return float(np.linalg.norm(offsets, axis=1).sum())

def extend_both_sides(oracle_centerline, resolution=1, extension=15, forward=True, reverse=True):
    """
        oracle_centerline: (N, 2)
        resolution (m)
        extension (m)
    """
    x, y = oracle_centerline[:, 0], oracle_centerline[:, 1]
    angle_start = np.arctan2(y[1] - y[0], x[1] - x[0])
    angle_end = np.arctan2(y[-1] - y[-2], x[-1] - x[-2])
    
    x_start = x[0] + np.cos(angle_start) * np.arange(- extension, 0, resolution)
    x_end = x[-1] + np.cos(angle_end) * np.arange(0, 0 + extension, resolution)

    y_start = y[0] + np.sin(angle_start) * np.arange(- extension, 0, resolution)
    y_end = y[-1] + np.sin(angle_end) * np.arange(0, 0 + extension, resolution)

    if forward and reverse:
        x = np.concatenate((x_start, x.reshape(x.shape[0]), x_end))
        y = np.concatenate((y_start, y.reshape(y.shape[0]), y_end))
    elif forward: 
        x = np.concatenate((x.reshape(x.shape[0]), x_end))
        y = np.concatenate((y.reshape(y.shape[0]), y_end))
    elif reverse:
        x = np.concatenate((x_start, x.reshape(x.shape[0])))
        y = np.concatenate((y_start, y.reshape(y.shape[0])))

    return np.dstack((x,y))[0]

def project_to_frenet_frame(traj, ref_line):
    distance_to_ref = torch.cdist(traj[:, :, :2], ref_line[:, :, :2])
    k = torch.argmin(distance_to_ref, dim=-1).view(-1, traj.shape[1], 1).expand(-1, -1, 3)
    ref_points = torch.gather(ref_line, 1, k)
    x_r, y_r, theta_r = ref_points[:, :, 0], ref_points[:, :, 1], ref_points[:, :, 2]
    x, y = traj[:, :, 0], traj[:, :, 1]
    s = 0.1 * (k[:, :, 0])
    l = torch.sign((y-y_r)*torch.cos(theta_r)-(x-x_r)*torch.sin(theta_r)) * torch.sqrt(torch.square(x-x_r)+torch.square(y-y_r))
    sl = torch.stack([s, l], dim=-1)

    return sl

def project_to_cartesian_frame(traj, ref_line):
    k = (10 * traj[:, :, 0]).long()
    k = torch.clip(k, 0, len(ref_line[0, :]) - 1)
    ref_points = torch.gather(ref_line, 1, k.view(-1, traj.shape[1], 1).expand(-1, -1, 3))
    x_r, y_r, theta_r = ref_points[:, :, 0], ref_points[:, :, 1], ref_points[:, :, 2]
    x = x_r - traj[:, :, 1] * torch.sin(theta_r)
    y = y_r + traj[:, :, 1] * torch.cos(theta_r)
    xy = torch.stack([x, y], dim=-1)
    return xy

def denoise(arr, w = 7):
    gt_x, gt_y = arr[:, 0], arr[:, 1]
    # denoising
    w = w
    gt_x_t = []
    gt_y_t = []
    for iq in range(len(gt_x)):
        if iq >= w and iq + w <= len(gt_x):
            gt_x_t.append(np.mean(gt_x[iq: iq + w]))
            gt_y_t.append(np.mean(gt_y[iq: iq + w]))
        elif iq < w:
            okx = np.mean(gt_x[w: w + w])
            gt_x_t.append(gt_x[0] + (okx - gt_x[0]) * (iq) / w)
            oky = np.mean(gt_y[w: w + w])
            gt_y_t.append(gt_y[0] + (oky - gt_y[0]) * (iq) / w)
        else:
            okx = np.mean(gt_x[len(gt_x) - w:len(gt_x) - w  + w])
            oky = np.mean(gt_y[len(gt_x) - w: len(gt_x) - w + w])
            gt_x_t.append(okx + (gt_x[-1] - okx) * (w - (len(gt_x) - iq)) / w)
            gt_y_t.append(oky + (gt_y[-1] - oky) * (w - (len(gt_y) - iq)) / w)                   

    gt_x = gt_x_t
    gt_y = gt_y_t
    print(gt_x, gt_y)
    return np.dstack((gt_x, gt_y))[0]

def generate_instance_colours(instance_map):
    # Most distinct 22 colors (kelly colors from https://stackoverflow.com/questions/470690/how-to-automatically-generate
    # -n-distinct-colors)
    # plus some colours from AD40k
    INSTANCE_COLOURS = np.asarray([
        [0, 0, 0],
        [255, 179, 0],
        [128, 62, 117],
        [255, 104, 0],
        [166, 189, 215],
        [193, 0, 32],
        [206, 162, 98],
        [129, 112, 102],
        [0, 125, 52],
        [246, 118, 142],
        [0, 83, 138],
        [255, 122, 92],
        [83, 55, 122],
        [255, 142, 0],
        [179, 40, 81],
        [244, 200, 0],
        [127, 24, 13],
        [147, 170, 0],
        [89, 51, 21],
        [241, 58, 19],
        [35, 44, 22],
        [112, 224, 255],
        [70, 184, 160],
        [153, 0, 255],
        [71, 255, 0],
        [255, 0, 163],
        [255, 204, 0],
        [0, 255, 235],
        [255, 0, 235],
        [255, 0, 122],
        [255, 245, 0],
        [10, 190, 212],
        [214, 255, 0],
        [0, 204, 255],
        [20, 0, 255],
        [255, 255, 0],
        [0, 153, 255],
        [0, 255, 204],
        [41, 255, 0],
        [173, 0, 255],
        [0, 245, 255],
        [71, 0, 255],
        [0, 255, 184],
        [0, 92, 255],
        [184, 255, 0],
        [255, 214, 0],
        [25, 194, 194],
        [92, 0, 255],
        [220, 220, 220],
        [255, 9, 92],
        [112, 9, 255],
        [8, 255, 214],
        [255, 184, 6],
        [10, 255, 71],
        [255, 41, 10],
        [7, 255, 255],
        [224, 255, 8],
        [102, 8, 255],
        [255, 61, 6],
        [255, 194, 7],
        [0, 255, 20],
        [255, 8, 41],
        [255, 5, 153],
        [6, 51, 255],
        [235, 12, 255],
        [160, 150, 20],
        [0, 163, 255],
        [140, 140, 140],
        [250, 10, 15],
        [20, 255, 0],
    ])

    return {instance_id: INSTANCE_COLOURS[global_instance_id % len(INSTANCE_COLOURS)] for
            instance_id, global_instance_id in instance_map.items()
            }


def is_overlapping_lane_seq(lane_seq1: Sequence[int], lane_seq2: Sequence[int]) -> bool:
    """
    Check if the 2 lane sequences are overlapping.
    Overlapping is defined as::
        s1------s2-----------------e1--------e2
    Here lane2 starts somewhere on lane 1 and ends after it, OR::
        s1------s2-----------------e2--------e1
    Here lane2 starts somewhere on lane 1 and ends before it
    Args:
        lane_seq1: list of lane ids
        lane_seq2: list of lane ids
    Returns:
        bool, True if the lane sequences overlap
    """

    if lane_seq2[0] in lane_seq1[1:] and lane_seq1[-1] in lane_seq2[:-1]:
        return True
    elif set(lane_seq2) <= set(lane_seq1):
        return True
    return False

def remove_overlapping_lane_seq(lane_seqs: List[List[int]]) -> List[List[int]]:
    """
    Remove lane sequences which are overlapping to some extent
    Args:
        lane_seqs (list of list of integers): List of list of lane ids (Eg. [[12345, 12346, 12347], [12345, 12348]])
    Returns:
        List of sequence of lane ids (e.g. ``[[12345, 12346, 12347], [12345, 12348]]``)
    """
    redundant_lane_idx: Set[int] = set()
    for i in range(len(lane_seqs)):
        for j in range(len(lane_seqs)):
            if i in redundant_lane_idx or i == j:
                continue
            if is_overlapping_lane_seq(lane_seqs[i], lane_seqs[j]):
                redundant_lane_idx.add(j)

    unique_lane_seqs = [lane_seqs[i] for i in range(len(lane_seqs)) if i not in redundant_lane_idx]
    return unique_lane_seqs

def dfs_incoming(nusc_map, lane_id, dist, threshold, resolution_meters=1):
    if dist > threshold:
        return [[lane_id]]
    else:
        traversed_lanes = []
        parent_lanes = (
                nusc_map.get_incoming_lane_ids(lane_id)
        )            
        if parent_lanes is not None:
            for parent in parent_lanes:
                # centerline = nusc_map.get_outgoing_lane_ids(child)
                try:
                    lane_record = nusc_map.get_arcline_path(parent)
                    centerline = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                    cl_length = LineString(centerline).length
                    curr_lane_ids = dfs_incoming(
                        nusc_map,
                        parent,
                        dist + cl_length,
                        threshold
                    )
                    traversed_lanes.extend(curr_lane_ids)
                except:
                    pass
        if len(traversed_lanes) == 0:
            return [[lane_id]]
        lanes_to_return = []
        for lane_seq in traversed_lanes:
            lanes_to_return.append(lane_seq + [lane_id])
        return lanes_to_return


def dfs(nusc_map, lane_id, dist, threshold, resolution_meters=1):
    if dist > threshold:
        return [[lane_id]]
    else:
        traversed_lanes = []
        child_lanes = (
                nusc_map.get_outgoing_lane_ids(lane_id)
        )            
        if child_lanes is not None:
            for child in child_lanes:
                # centerline = nusc_map.get_outgoing_lane_ids(child)
                try:
                    lane_record = nusc_map.get_arcline_path(child)
                    centerline = np.array(discretize_lane(lane_record, resolution_meters=resolution_meters))
                    cl_length = LineString(centerline).length
                    curr_lane_ids = dfs(
                        nusc_map,
                        child,
                        dist + cl_length,
                        threshold
                    )
                    traversed_lanes.extend(curr_lane_ids)
                except:
                    pass
        if len(traversed_lanes) == 0:
            return [[lane_id]]
        lanes_to_return = []
        for lane_seq in traversed_lanes:
            lanes_to_return.append([lane_id] + lane_seq)
        return lanes_to_return


def resize_and_crop_image(img, resize_dims, crop):
    # Bilinear resizing followed by cropping
    img = img.resize(resize_dims, resample=PIL.Image.BILINEAR)
    img = img.crop(crop)
    return img

def get_patch_coord(patch_box: Tuple[float, float, float, float],
                    patch_angle: float = 0.0) -> Polygon:
    """
    Convert patch_box to shapely Polygon coordinates.
    :param patch_box: Patch box defined as [x_center, y_center, height, width].
    :param patch_angle: Patch orientation in degrees.
    :return: Box Polygon for patch_box.
    """
    patch_x, patch_y, patch_h, patch_w = patch_box

    x_min = patch_x - patch_w / 2.0
    y_min = patch_y - patch_h / 2.0
    x_max = patch_x + patch_w / 2.0
    y_max = patch_y + patch_h / 2.0

    patch = box(x_min, y_min, x_max, y_max)
    patch = affinity.rotate(patch, patch_angle, origin=(patch_x, patch_y), use_radians=False)

    return patch


def update_intrinsics(intrinsics, top_crop=0.0, left_crop=0.0, scale_width=1.0, scale_height=1.0):
    """
    Parameters
    ----------
        intrinsics: torch.Tensor (3, 3)
        top_crop: float
        left_crop: float
        scale_width: float
        scale_height: float
    """
    updated_intrinsics = intrinsics.clone()
    # Adjust intrinsics scale due to resizing
    updated_intrinsics[0, 0] *= scale_width
    updated_intrinsics[0, 2] *= scale_width
    updated_intrinsics[1, 1] *= scale_height
    updated_intrinsics[1, 2] *= scale_height

    # Adjust principal point due to cropping
    updated_intrinsics[0, 2] -= left_crop
    updated_intrinsics[1, 2] -= top_crop

    return updated_intrinsics


def calculate_birds_eye_view_parameters(x_bounds, y_bounds, z_bounds):
    """
    Parameters
    ----------
        x_bounds: Forward direction in the ego-car.
        y_bounds: Sides
        z_bounds: Height

    Returns
    -------
        bev_resolution: Bird's-eye view bev_resolution
        bev_start_position Bird's-eye view first element
        bev_dimension Bird's-eye view tensor spatial dimension
    """
    bev_resolution = torch.tensor([row[2] for row in [x_bounds, y_bounds, z_bounds]])
    bev_start_position = torch.tensor([row[0] + row[2] / 2.0 for row in [x_bounds, y_bounds, z_bounds]])
    bev_dimension = torch.tensor([(row[1] - row[0]) / row[2] for row in [x_bounds, y_bounds, z_bounds]],
                                 dtype=torch.long)

    return bev_resolution, bev_start_position, bev_dimension


def convert_egopose_to_matrix_numpy(egopose):
    transformation_matrix = np.zeros((4, 4), dtype=np.float32)
    rotation = Quaternion(egopose['rotation']).rotation_matrix
    translation = np.array(egopose['translation'])
    transformation_matrix[:3, :3] = rotation
    transformation_matrix[:3, 3] = translation
    transformation_matrix[3, 3] = 1.0
    return transformation_matrix

def get_global_pose(rec, nusc, inverse=False):
    lidar_sample_data = nusc.get('sample_data', rec['data']['LIDAR_TOP'])

    sd_ep = nusc.get("ego_pose", lidar_sample_data["ego_pose_token"])
    sd_cs = nusc.get("calibrated_sensor", lidar_sample_data["calibrated_sensor_token"])
    if inverse is False:
        global_from_ego = transform_matrix(sd_ep["translation"], Quaternion(sd_ep["rotation"]), inverse=False)
        ego_from_sensor = transform_matrix(sd_cs["translation"], Quaternion(sd_cs["rotation"]), inverse=False)
        pose = global_from_ego.dot(ego_from_sensor)
    else:
        sensor_from_ego = transform_matrix(sd_cs["translation"], Quaternion(sd_cs["rotation"]), inverse=True)
        ego_from_global = transform_matrix(sd_ep["translation"], Quaternion(sd_ep["rotation"]), inverse=True)
        pose = sensor_from_ego.dot(ego_from_global)
    return pose

def invert_matrix_egopose_numpy(egopose):
    """ Compute the inverse transformation of a 4x4 egopose numpy matrix."""
    inverse_matrix = np.zeros((4, 4), dtype=np.float32)
    rotation = egopose[:3, :3]
    translation = egopose[:3, 3]
    inverse_matrix[:3, :3] = rotation.T
    inverse_matrix[:3, 3] = -np.dot(rotation.T, translation)
    inverse_matrix[3, 3] = 1.0
    return inverse_matrix


def mat2pose_vec(matrix: torch.Tensor):
    """
    Converts a 4x4 pose matrix into a 6-dof pose vector
    Args:
        matrix (ndarray): 4x4 pose matrix
    Returns:
        vector (ndarray): 6-dof pose vector comprising translation components (tx, ty, tz) and
        rotation components (rx, ry, rz)
    """

    # M[1, 2] = -sinx*cosy, M[2, 2] = +cosx*cosy
    rotx = torch.atan2(-matrix[..., 1, 2], matrix[..., 2, 2])

    # M[0, 2] = +siny, M[1, 2] = -sinx*cosy, M[2, 2] = +cosx*cosy
    cosy = torch.sqrt(matrix[..., 1, 2] ** 2 + matrix[..., 2, 2] ** 2)
    roty = torch.atan2(matrix[..., 0, 2], cosy)

    # M[0, 0] = +cosy*cosz, M[0, 1] = -cosy*sinz
    rotz = torch.atan2(-matrix[..., 0, 1], matrix[..., 0, 0])

    rotation = torch.stack((rotx, roty, rotz), dim=-1)

    # Extract translation params
    translation = matrix[..., :3, 3]
    return torch.cat((translation, rotation), dim=-1)


def euler2mat(angle: torch.Tensor):
    """Convert euler angles to rotation matrix.
    Reference: https://github.com/pulkitag/pycaffe-utils/blob/master/rot_utils.py#L174
    Args:
        angle: rotation angle along 3 axis (in radians) [Bx3]
    Returns:
        Rotation matrix corresponding to the euler angles [Bx3x3]
    """
    shape = angle.shape
    angle = angle.view(-1, 3)
    x, y, z = angle[:, 0], angle[:, 1], angle[:, 2]

    cosz = torch.cos(z)
    sinz = torch.sin(z)

    zeros = torch.zeros_like(z)
    ones = torch.ones_like(z)
    zmat = torch.stack([cosz, -sinz, zeros, sinz, cosz, zeros, zeros, zeros, ones], dim=1).view(-1, 3, 3)

    cosy = torch.cos(y)
    siny = torch.sin(y)

    ymat = torch.stack([cosy, zeros, siny, zeros, ones, zeros, -siny, zeros, cosy], dim=1).view(-1, 3, 3)

    cosx = torch.cos(x)
    sinx = torch.sin(x)

    xmat = torch.stack([ones, zeros, zeros, zeros, cosx, -sinx, zeros, sinx, cosx], dim=1).view(-1, 3, 3)

    rot_mat = xmat.bmm(ymat).bmm(zmat)
    rot_mat = rot_mat.view(*shape[:-1], 3, 3)
    return rot_mat


def pose_vec2mat(vec: torch.Tensor):
    """
    Convert 6DoF parameters to transformation matrix.
    Args:
        vec: 6DoF parameters in the order of tx, ty, tz, rx, ry, rz [B,6]
    Returns:
        A transformation matrix [B,4,4]
    """
    translation = vec[..., :3].unsqueeze(-1)  # [...x3x1]
    rot = vec[..., 3:].contiguous()  # [...x3]
    rot_mat = euler2mat(rot)  # [...,3,3]
    transform_mat = torch.cat([rot_mat, translation], dim=-1)  # [...,3,4]
    transform_mat = torch.nn.functional.pad(transform_mat, [0, 0, 0, 1], value=0)  # [...,4,4]
    transform_mat[..., 3, 3] = 1.0
    return transform_mat


def invert_pose_matrix(x):
    """
    Parameters
    ----------
        x: [B, 4, 4] batch of pose matrices

    Returns
    -------
        out: [B, 4, 4] batch of inverse pose matrices
    """
    assert len(x.shape) == 3 and x.shape[1:] == (4, 4), 'Only works for batch of pose matrices.'

    transposed_rotation = torch.transpose(x[:, :3, :3], 1, 2)
    translation = x[:, :3, 3:]

    inverse_mat = torch.cat([transposed_rotation, -torch.bmm(transposed_rotation, translation)], dim=-1) # [B,3,4]
    inverse_mat = torch.nn.functional.pad(inverse_mat, [0, 0, 0, 1], value=0)  # [B,4,4]
    inverse_mat[..., 3, 3] = 1.0
    return inverse_mat


def warp_features(x, flow, mode='nearest', spatial_extent=None):
    """ Applies a rotation and translation to feature map x.
        Args:
            x: (b, c, h, w) feature map
            flow: (b, 6) 6DoF vector (only uses the xy poriton)
            mode: use 'nearest' when dealing with categorical inputs
        Returns:
            in plane transformed feature map
        """
    if flow is None:
        return x
    b, c, h, w = x.shape
    # z-rotation
    angle = flow[:, 5].clone()  # torch.atan2(flow[:, 1, 0], flow[:, 0, 0])
    # x-y translation
    translation = flow[:, :2].clone()  # flow[:, :2, 3]

    # Normalise translation. Need to divide by how many meters is half of the image.
    # because translation of 1.0 correspond to translation of half of the image.
    translation[:, 0] /= spatial_extent[0]
    translation[:, 1] /= spatial_extent[1]
    # forward axis is inverted
    translation[:, 0] *= -1

    cos_theta = torch.cos(angle)
    sin_theta = torch.sin(angle)

    # output = Rot.input + translation
    # tx and ty are inverted as is the case when going from real coordinates to numpy coordinates
    # translation_pos_0 -> positive value makes the image move to the left
    # translation_pos_1 -> positive value makes the image move to the top
    # Angle -> positive value in rad makes the image move in the trigonometric way
    transformation = torch.stack([cos_theta, -sin_theta, translation[:, 1],
                                  sin_theta, cos_theta, translation[:, 0]], dim=-1).view(b, 2, 3)

    # Note that a rotation will preserve distances only if height = width. Otherwise there's
    # resizing going on. e.g. rotation of pi/2 of a 100x200 image will make what's in the center of the image
    # elongated.
    grid = torch.nn.functional.affine_grid(transformation, size=x.shape, align_corners=False)
    grid = grid.to(dtype=x.dtype)
    warped_x = torch.nn.functional.grid_sample(x, grid, mode=mode, padding_mode='zeros', align_corners=False)

    return warped_x


def cumulative_warp_features(x, flow, mode='nearest', spatial_extent=None):
    """ Warps a sequence of feature maps by accumulating incremental 2d flow.

    x[:, -1] remains unchanged
    x[:, -2] is warped using flow[:, -2]
    x[:, -3] is warped using flow[:, -3] @ flow[:, -2]
    ...
    x[:, 0] is warped using flow[:, 0] @ ... @ flow[:, -3] @ flow[:, -2]

    Args:
        x: (b, t, c, h, w) sequence of feature maps
        flow: (b, t, 6) sequence of 6 DoF pose
            from t to t+1 (only uses the xy poriton)

    """
    sequence_length = x.shape[1]
    if sequence_length == 1:
        return x

    flow = pose_vec2mat(flow)

    out = [x[:, -1]]
    cum_flow = flow[:, -2]
    for t in reversed(range(sequence_length - 1)):
        out.append(warp_features(x[:, t], mat2pose_vec(cum_flow), mode=mode, spatial_extent=spatial_extent))
        # @ is the equivalent of torch.bmm
        cum_flow = flow[:, t - 1] @ cum_flow

    return torch.stack(out[::-1], 1)


def cumulative_warp_features_reverse(x, flow, mode='nearest', spatial_extent=None):
    """ Warps a sequence of feature maps by accumulating incremental 2d flow.

    x[:, 0] remains unchanged
    x[:, 1] is warped using flow[:, 0].inverse()
    x[:, 2] is warped using flow[:, 0].inverse() @ flow[:, 1].inverse()
    ...

    Args:
        x: (b, t, c, h, w) sequence of feature maps
        flow: (b, t, 6) sequence of 6 DoF pose
            from t to t+1 (only uses the xy poriton)

    """
    flow = pose_vec2mat(flow)

    out = [x[:,0]]
    
    for i in range(1, x.shape[1]):
        if i==1:
            cum_flow = invert_pose_matrix(flow[:, 0])
        else:
            cum_flow = cum_flow @ invert_pose_matrix(flow[:,i-1])
        out.append( warp_features(x[:,i], mat2pose_vec(cum_flow), mode, spatial_extent=spatial_extent))
    return torch.stack(out, 1)


class VoxelsSumming(torch.autograd.Function):
    """Adapted from https://github.com/nv-tlabs/lift-splat-shoot/blob/master/src/tools.py#L193"""
    @staticmethod
    def forward(ctx, x, geometry, ranks):
        """The features `x` and `geometry` are ranked by voxel positions."""
        # Cumulative sum of all features.
        x = x.cumsum(0)

        # Indicates the change of voxel.
        mask = torch.ones(x.shape[0], device=x.device, dtype=torch.bool)
        mask[:-1] = ranks[1:] != ranks[:-1]

        x, geometry = x[mask], geometry[mask]
        # Calculate sum of features within a voxel.
        x = torch.cat((x[:1], x[1:] - x[:-1]))

        ctx.save_for_backward(mask)
        ctx.mark_non_differentiable(geometry)

        return x, geometry

    @staticmethod
    def backward(ctx, grad_x, grad_geometry):
        (mask,) = ctx.saved_tensors
        # Since the operation is summing, we simply need to send gradient
        # to all elements that were part of the summation process.
        indices = torch.cumsum(mask, 0)
        indices[mask] -= 1

        output_grad = grad_x[indices]

        return output_grad, None, None

def ___extract_trajs(centerlines, trj, nvar=11, viz=False, debugg=False, device='cpu', labels=None, griddim=None, problem=None, avoid_obs=False, obs=None):
    """
        this is just saved to understand where to insert the plotting code.
    """
    save_device = device
    device = 'cpu'
    bbx = np.array([-50.0 + 0.5/2.0, -50.0 + 0.5/2.0])
    ddx = np.array([0.5, 0.5])
    if griddim == None:
        X_MAX = 30
        X_MIN = 0
        Y_MAX = 2
        Y_MIN = -2
        NUM_X = 30
        NUM_Y = 5
    else:
        X_MIN, X_MAX, NUM_X, Y_MIN, Y_MAX, NUM_Y = griddim

    # get trajectories
    x = torch.linspace(X_MIN, X_MAX, NUM_X).to(device)
    y = torch.linspace(Y_MIN, Y_MAX, NUM_Y).to(device)
    xv, yv = torch.meshgrid(x, y)
    print(device)
    grid = torch.dstack((xv, yv)).to(device)
    ccc = centerlines.detach().cpu().numpy()
    cx, _ = interp_polyline_by_fixed_waypt_interval_np(ccc, 0.1)
    cx = extend_both_sides(cx, resolution=0.1, extension=100)
    cx = cx * ddx
    cx += bbx
    cx[:, [1, 0]] = cx[:, [0, 1]]
    cx[:, :1] = cx[:, :1] * -1
    cx = torch.tensor(cx).to(device)
    if viz:
        trr = trj.squeeze().detach().cpu()
        plt.plot(cx[:, 0], cx[:, 1], color='grey')
        plt.scatter(cx[:, 0], cx[:, 1], color='grey')
        plt.plot(trr[:, 0], trr[:, 1], color='blue')
        plt.scatter(trr[:, 0], trr[:, 1], color='blue')
        # print(trr[:, 0], trr[:, 1])
        plt.xlim([-5,5])
        plt.ylim([-5,5])
        plt.savefig('output_vis/debug_actual.png'); plt.clf()

    trr = trj.squeeze().detach().cpu()
    ref_line_ = torch.tensor(cx.reshape(1, cx.shape[0], 2))
    ref_line = torch.zeros((1, ref_line_.shape[1] - 1, 3))
    ref_line[:, :, :2] = ref_line_[:, 1:] 
    off = ref_line_[:, 1:] - ref_line_[:,:-1]
    ref_line[:, :, 2] = torch.atan2(off[:, :, 1], off[:, :, 0])
    ref_line = ref_line.float().to(device)
    trj = trj.float().to(device)

    test_ref_line=ref_line[:, 100:-100]
    sdd = project_to_frenet_frame(test_ref_line,  ref_line)
    if debugg: pdb.set_trace()
    if viz:
        plt.scatter(sdd[0, :, 0].detach().cpu(), sdd[0, :, 1].detach().cpu(), color='grey')
        plt.show()

    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)

    if avoid_obs:
        # obs -> (M, num, 2) : M obstacles for num timesteps
        # avoid obstacles
        num_obs = 20
        a_obs = 1
        b_obs = 1
        x_obs = obs[:, :, 0]
        y_obs = obs[:, :, 1]

        ssd_obs = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)
        problem = OPTNode(
            x_obs=x_obs, 
            y_obs=y_obs, 
            a_obs=a_obs, 
            b_obs=b_obs, 
            num_obs=num_obs, 
            rho_obs=1, 
            rho_eq=problem.rho_eq, 
            t_fin=problem.t_fin, 
            num=problem.num, 
            bernstein_coeff_order10_new=problem.bernstein_coeff_order10_new, 
            device = centerlines.device
        )



    cx = interp_arc_np(2000, cx.detach().cpu().numpy())
    cx = extend_both_sides(cx, resolution=0.1, extension=50)
    cx = torch.tensor(cx).to(device)
    if debugg: pdb.set_trace()

    if viz:
        gt = labels['gt_trajectory_prev'][0]
        plt.scatter(gt[:2, 0].detach().cpu(), gt[:2, 1].detach().cpu(), color='red')
        plt.plot(gt[:, 0].detach().cpu(), gt[:, 1].detach().cpu(), color='red', label='observed')
        plt.scatter(gt[2:, 0].detach().cpu(), gt[2:, 1].detach().cpu(), color='orange')
        plt.plot(gt[2:, 0].detach().cpu(), gt[2:, 1].detach().cpu(), color='orange', label='future')
        # plt.axis('equal')
        # plt.xlim([-5, 5])
        # plt.ylim([-5, 5])
        plt.legend()
        plt.savefig("output_vis/debug_mid.png");plt.clf()

    ssd = project_to_frenet_frame(trj, ref_line) # (1, 7, 2)
    s_offset = ssd[:,2,0]
    d_offset = ssd[0,2,1]
    vel = ssd[0,2] - ssd[0,1]
    vel_prev = ssd[0,1] - ssd[0, 0]
    theta = torch.atan2(vel[1], vel[0]) # to make angle 0 wrt ego-agent
    theta_prev = torch.atan2(vel_prev[1], vel_prev[0])
    psidot_init = (theta - theta_prev)/0.5
    psi_init = theta
    v_init = torch.linalg.norm(vel)/0.5
    y_init = 0
    x_init = 0
    rot_inv = torch.tensor([[cos(-theta), -sin(-theta)], [sin(-theta), cos(-theta)]]).to(device)
    rot = torch.tensor([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]]).to(device)
    if debugg: pdb.set_trace()

    sdd_frame = sdd -  torch.tensor([s_offset, d_offset]).to(device)
    ssd_frame = ssd -  torch.tensor([s_offset, d_offset]).to(device)
    eps_frame = torch.tensor(grid).to(device) - torch.tensor([0, d_offset]).to(device) # (20, 7, 2)


    if debugg: pdb.set_trace()
    if viz:
        for eps in eps_frame:
            ep = eps.detach().cpu().numpy()
            plt.scatter(ep[:, 0], ep[:, 1])
        ss = ssd_frame[0].detach().cpu().numpy()
        plt.scatter(ss[:2, 0], ss[:2, 1], color='red')
        plt.scatter(ss[2:, 0], ss[2:, 1], color='orange')
        sd = sdd_frame[0].detach().cpu().numpy()
        plt.scatter(sd[:, 0], sd[:, 1], color='grey')
        plt.plot(sd[:, 0], sd[:, 1], color='grey')
        plt.savefig("output_vis/debug_ss.png"); plt.clf()
        plt.clf()

    if debugg: pdb.set_trace()
    # rotate
    sdd_frame[0] = torch.matmul(rot_inv, sdd_frame[0].T).T # (1, M, 2)
    ssd_frame = torch.matmul(rot_inv, ssd_frame[0].T).T # ()
    eps_frame = eps_frame.permute(1, 2, 0) # (7, 2, 20)
    eps_frame = torch.matmul(rot_inv, eps_frame) # (2, 2) x (7, 2, 20) -> (7, 2, 20)
    if debugg: pdb.set_trace()

    if viz:
        ss = ssd_frame.detach().cpu().numpy()
        plt.scatter(ss[:2, 0], ss[:2, 1], color='red')
        plt.scatter(ss[2:, 0], ss[2:, 1], color='orange')
        sd = sdd_frame[0].detach().cpu().numpy()
        plt.scatter(sd[:, 0], sd[:, 1], color='grey', label='centerline')
        plt.plot(sd[:, 0], sd[:, 1], color='grey')
        # for ep in eps_frame:
        #     plt.scatter(ep[:, 0], ep[:, 1], color='yellow')
        plt.scatter(eps_frame[:, 0, :].flatten().detach().cpu().numpy(), eps_frame[:, 1, :].flatten().detach().cpu().numpy(), color='yellow')
        # plt.savefig("output_vis/debug_ss.png"); plt.clf()
        plt.xlim([-25, 25])
        plt.ylim([-25, 25])
        plt.legend()
        plt.savefig("output_vis/debug_ss_rotated.png"); plt.clf()
        plt.clf()

    if debugg: pdb.set_trace()
    x_eps = eps_frame[:,0,:].flatten()
    y_eps = eps_frame[:,1,:].flatten()
    fixed_params = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init]).to(device)
    fixed_params = fixed_params.expand(len(x_eps), 5)
    variable_params = torch.zeros((len(x_eps), 3)).to(device)
    variable_params[:, 0] = torch.tensor(x_eps)
    variable_params[:, 1] = torch.tensor(y_eps)
    variable_params[:, 2] = torch.tensor(-theta)
    P = problem.P
    cc, _ = problem.solve(fixed_params.double().detach().cpu(), variable_params.double().detach().cpu())
    x_sol = torch.matmul(P, cc.T[:nvar]).T
    y_sol = torch.matmul(P, cc.T[nvar:2*nvar]).T

    if debugg: pdb.set_trace()

    if viz:
        ss = ssd_frame.detach().cpu().numpy()
        plt.scatter(eps_frame[:, 0, :].flatten().detach().cpu().numpy(), eps_frame[:, 1, :].flatten().detach().cpu().numpy(), color='yellow')
        sd = sdd_frame[0].detach().cpu().numpy()
        plt.scatter(sd[:, 0], sd[:, 1], color='grey', label='centerline')
        plt.plot(sd[:, 0], sd[:, 1], color='grey')
        for i in range(len(x_sol)):
            plt.plot(x_sol[i].flatten().detach().cpu().numpy(), y_sol[i].flatten().detach().cpu().numpy())
        # plt.savefig("output_vis/debug_ss.png"); plt.clf()
        plt.scatter(ss[:2, 0], ss[:2, 1], color='red')
        plt.scatter(ss[2:, 0], ss[2:, 1], color='orange')
        plt.xlim([-5, 5])
        plt.ylim([-5, 5])
        plt.legend()
        plt.savefig("output_vis/debug_frenet.png")
        plt.clf()

    if debugg: pdb.set_trace()
    trajs = torch.dstack((x_sol, y_sol)).to(device) # N, 7, 2
    trajs = trajs.permute(0, 2, 1).float()
    trajs = torch.matmul(rot, trajs)
    trajs = trajs.permute(0, 2, 1)
    trajs = trajs + torch.tensor([s_offset, d_offset]).to(device)

    min_val = 1e11
    min_ind = 0

    dense_trajs = project_to_cartesian_frame(trajs.reshape(1, trajs.shape[0] * trajs.shape[1], 2), ref_line).reshape(1, trajs.shape[0], trajs.shape[1], 2)#(B, N, n_future, 2)  
    return dense_trajs.to(save_device)


data/stp3/utils/go_save.py
llmbev/talk2bev
from typing import Tuple

import torch
import torch.nn.functional as F
import numpy as np
from scipy.optimize import linear_sum_assignment

from stp3.utils.geometry import mat2pose_vec, pose_vec2mat, warp_features


# set ignore index to 0 for vis
def convert_instance_mask_to_center_and_offset_label(instance_img, future_egomotion, num_instances, ignore_index=255,
                                                     subtract_egomotion=True, sigma=3, spatial_extent=None):
    seq_len, h, w = instance_img.shape
    center_label = torch.zeros(seq_len, 1, h, w)
    offset_label = ignore_index * torch.ones(seq_len, 2, h, w)
    future_displacement_label = ignore_index * torch.ones(seq_len, 2, h, w)
    # x is vertical displacement, y is horizontal displacement
    x, y = torch.meshgrid(torch.arange(h, dtype=torch.float), torch.arange(w, dtype=torch.float))

    if subtract_egomotion:
        future_egomotion_inv = mat2pose_vec(pose_vec2mat(future_egomotion).inverse())

    # Compute warped instance segmentation
    warped_instance_seg = {}
    for t in range(1, seq_len):
        warped_inst_t = warp_features(instance_img[t].unsqueeze(0).unsqueeze(1).float(),
                                      future_egomotion_inv[t - 1].unsqueeze(0), mode='nearest',
                                      spatial_extent=spatial_extent)
        warped_instance_seg[t] = warped_inst_t[0, 0]

    # Ignore id 0 which is the background
    for instance_id in range(1, num_instances+1):
        prev_xc = None
        prev_yc = None
        prev_mask = None
        for t in range(seq_len):
            instance_mask = (instance_img[t] == instance_id)
            if instance_mask.sum() == 0:
                # this instance is not in this frame
                prev_xc = None
                prev_yc = None
                prev_mask = None
                continue

            xc = x[instance_mask].mean().round().long()
            yc = y[instance_mask].mean().round().long()

            off_x = xc - x
            off_y = yc - y
            g = torch.exp(-(off_x ** 2 + off_y ** 2) / sigma ** 2)
            center_label[t, 0] = torch.maximum(center_label[t, 0], g)
            offset_label[t, 0, instance_mask] = off_x[instance_mask]
            offset_label[t, 1, instance_mask] = off_y[instance_mask]

            if prev_xc is not None:
                # old method
                # cur_pt = torch.stack((xc, yc)).unsqueeze(0).float()
                # if subtract_egomotion:
                #     cur_pt = warp_points(cur_pt, future_egomotion_inv[t - 1])
                # cur_pt = cur_pt.squeeze(0)

                warped_instance_mask = warped_instance_seg[t] == instance_id
                if warped_instance_mask.sum() > 0:
                    warped_xc = x[warped_instance_mask].mean().round()
                    warped_yc = y[warped_instance_mask].mean().round()

                    delta_x = warped_xc - prev_xc
                    delta_y = warped_yc - prev_yc
                    future_displacement_label[t - 1, 0, prev_mask] = delta_x
                    future_displacement_label[t - 1, 1, prev_mask] = delta_y

            prev_xc = xc
            prev_yc = yc
            prev_mask = instance_mask

    return center_label, offset_label, future_displacement_label


def find_instance_centers(center_prediction: torch.Tensor, conf_threshold: float = 0.1, nms_kernel_size: float = 3):
    assert len(center_prediction.shape) == 3
    center_prediction = F.threshold(center_prediction, threshold=conf_threshold, value=-1)

    nms_padding = (nms_kernel_size - 1) // 2
    maxpooled_center_prediction = F.max_pool2d(
        center_prediction, kernel_size=nms_kernel_size, stride=1, padding=nms_padding
    )

    # Filter all elements that are not the maximum (i.e. the center of the heatmap instance)
    center_prediction[center_prediction != maxpooled_center_prediction] = -1
    return torch.nonzero(center_prediction > 0)[:, 1:]


def group_pixels(centers: torch.Tensor, offset_predictions: torch.Tensor) -> torch.Tensor:
    width, height = offset_predictions.shape[-2:]
    x_grid = (
        torch.arange(width, dtype=offset_predictions.dtype, device=offset_predictions.device)
        .view(1, width, 1)
        .repeat(1, 1, height)
    )
    y_grid = (
        torch.arange(height, dtype=offset_predictions.dtype, device=offset_predictions.device)
        .view(1, 1, height)
        .repeat(1, width, 1)
    )
    pixel_grid = torch.cat((x_grid, y_grid), dim=0)
    center_locations = (pixel_grid + offset_predictions).view(2, width * height, 1).permute(2, 1, 0)
    centers = centers.view(-1, 1, 2)

    distances = torch.norm(centers - center_locations, dim=-1)

    instance_id = torch.argmin(distances, dim=0).reshape(1, width, height) + 1
    return instance_id


def get_instance_segmentation_and_centers(
    center_predictions: torch.Tensor,
    offset_predictions: torch.Tensor,
    foreground_mask: torch.Tensor,
    conf_threshold: float = 0.1,
    nms_kernel_size: float = 3,
    max_n_instance_centers: int = 100,
) -> Tuple[torch.Tensor, torch.Tensor]:
    width, height = center_predictions.shape[-2:]
    center_predictions = center_predictions.view(1, width, height)
    offset_predictions = offset_predictions.view(2, width, height)
    foreground_mask = foreground_mask.view(1, width, height)

    centers = find_instance_centers(center_predictions, conf_threshold=conf_threshold, nms_kernel_size=nms_kernel_size)
    if not len(centers):
        return torch.zeros(center_predictions.shape, dtype=torch.int64, device=center_predictions.device), \
               torch.zeros((0, 2), device=centers.device)

    if len(centers) > max_n_instance_centers:
        # print(f'There are a lot of detected instance centers: {centers.shape}')
        centers = centers[:max_n_instance_centers].clone()

    instance_ids = group_pixels(centers, offset_predictions)
    instance_seg = (instance_ids * foreground_mask.float()).long()

    # Make the indices of instance_seg consecutive
    instance_seg = make_instance_seg_consecutive(instance_seg)

    return instance_seg.long(), centers


def update_instance_ids(instance_seg, old_ids, new_ids):
    """
    Parameters
    ----------
        instance_seg: torch.Tensor arbitrary shape
        old_ids: 1D tensor containing the list of old ids, must be all present in instance_seg.
        new_ids: 1D tensor with the new ids, aligned with old_ids

    Returns
        new_instance_seg: torch.Tensor same shape as instance_seg with new ids
    """
    indices = torch.arange(old_ids.max() + 1, device=instance_seg.device)
    for old_id, new_id in zip(old_ids, new_ids):
        indices[old_id] = new_id

    return indices[instance_seg].long()


def make_instance_seg_consecutive(instance_seg):
    # Make the indices of instance_seg consecutive
    unique_ids = torch.unique(instance_seg)
    new_ids = torch.arange(len(unique_ids), device=instance_seg.device)
    instance_seg = update_instance_ids(instance_seg, unique_ids, new_ids)
    return instance_seg


def make_instance_id_temporally_consistent(pred_inst, future_flow, matching_threshold=3.0):
    """
    Parameters
    ----------
        pred_inst: torch.Tensor (1, seq_len, h, w)
        future_flow: torch.Tensor(1, seq_len, 2, h, w)
        matching_threshold: distance threshold for a match to be valid.

    Returns
    -------
    consistent_instance_seg: torch.Tensor(1, seq_len, h, w)

    1. time t. Loop over all detected instances. Use flow to compute new centers at time t+1.
    2. Store those centers
    3. time t+1. Re-identify instances by comparing position of actual centers, and flow-warped centers.
        Make the labels at t+1 consistent with the matching
    4. Repeat
    """
    assert pred_inst.shape[0] == 1, 'Assumes batch size = 1'

    # Initialise instance segmentations with prediction corresponding to the present
    consistent_instance_seg = [pred_inst[0, 0]]
    largest_instance_id = consistent_instance_seg[0].max().item()

    _, seq_len, h, w = pred_inst.shape
    device = pred_inst.device
    for t in range(seq_len - 1):
        # Compute predicted future instance means
        grid = torch.stack(torch.meshgrid(
            torch.arange(h, dtype=torch.float, device=device), torch.arange(w, dtype=torch.float, device=device)
        ))

        # Add future flow
        grid = grid + future_flow[0, t]
        warped_centers = []
        # Go through all ids, except the background
        t_instance_ids = torch.unique(consistent_instance_seg[-1])[1:].cpu().numpy()

        if len(t_instance_ids) == 0:
            # No instance so nothing to update
            consistent_instance_seg.append(pred_inst[0, t + 1])
            continue

        for instance_id in t_instance_ids:
            instance_mask = (consistent_instance_seg[-1] == instance_id)
            warped_centers.append(grid[:, instance_mask].mean(dim=1))
        warped_centers = torch.stack(warped_centers)

        # Compute actual future instance means
        centers = []
        grid = torch.stack(torch.meshgrid(
            torch.arange(h, dtype=torch.float, device=device), torch.arange(w, dtype=torch.float, device=device)
        ))
        n_instances = int(pred_inst[0, t + 1].max().item())

        if n_instances == 0:
            # No instance, so nothing to update.
            consistent_instance_seg.append(pred_inst[0, t + 1])
            continue

        for instance_id in range(1, n_instances + 1):
            instance_mask = (pred_inst[0, t + 1] == instance_id)
            centers.append(grid[:, instance_mask].mean(dim=1))
        centers = torch.stack(centers)

        # Compute distance matrix between warped centers and actual centers
        distances = torch.norm(centers.unsqueeze(0) - warped_centers.unsqueeze(1), dim=-1).cpu().numpy()
        # outputs (row, col) with row: index in frame t, col: index in frame t+1
        # the missing ids in col must be added (correspond to new instances)
        ids_t, ids_t_one = linear_sum_assignment(distances)
        matching_distances = distances[ids_t, ids_t_one]
        # Offset by one as id=0 is the background
        ids_t += 1
        ids_t_one += 1

        # swap ids_t with real ids. as those ids correspond to the position in the distance matrix.
        id_mapping = dict(zip(np.arange(1, len(t_instance_ids) + 1), t_instance_ids))
        ids_t = np.vectorize(id_mapping.__getitem__, otypes=[np.int64])(ids_t)

        # Filter low quality match
        ids_t = ids_t[matching_distances < matching_threshold]
        ids_t_one = ids_t_one[matching_distances < matching_threshold]

        # Elements that are in t+1, but weren't matched
        remaining_ids = set(torch.unique(pred_inst[0, t + 1]).cpu().numpy()).difference(set(ids_t_one))
        # remove background
        remaining_ids.remove(0)
        #  Set remaining_ids to a new unique id
        for remaining_id in list(remaining_ids):
            largest_instance_id += 1
            ids_t = np.append(ids_t, largest_instance_id)
            ids_t_one = np.append(ids_t_one, remaining_id)

        consistent_instance_seg.append(update_instance_ids(pred_inst[0, t + 1], old_ids=ids_t_one, new_ids=ids_t))

    consistent_instance_seg = torch.stack(consistent_instance_seg).unsqueeze(0)
    return consistent_instance_seg


def predict_instance_segmentation_and_trajectories(
        output, compute_matched_centers=False, make_consistent=True, vehicles_id=1,
):
    preds = output['segmentation'].detach()
    preds = torch.argmax(preds, dim=2, keepdim=True)
    foreground_masks = preds.squeeze(2) == vehicles_id

    batch_size, seq_len = preds.shape[:2]
    pred_inst = []
    for b in range(batch_size):
        pred_inst_batch = []
        for t in range(seq_len):
            pred_instance_t, _ = get_instance_segmentation_and_centers(
                output['instance_center'][b, t].detach(),
                output['instance_offset'][b, t].detach(),
                foreground_masks[b, t].detach()
            )
            pred_inst_batch.append(pred_instance_t)
        pred_inst.append(torch.stack(pred_inst_batch, dim=0))

    pred_inst = torch.stack(pred_inst).squeeze(2)

    if make_consistent:
        if output['instance_flow'] is None:
            print('Using zero flow because instance_future_output is None')
            output['instance_flow'] = torch.zeros_like(output['instance_offset'])
        consistent_instance_seg = []
        for b in range(batch_size):
            consistent_instance_seg.append(
                make_instance_id_temporally_consistent(pred_inst[b:b+1],
                                                       output['instance_flow'][b:b+1].detach())
            )
        consistent_instance_seg = torch.cat(consistent_instance_seg, dim=0)
    else:
        consistent_instance_seg = pred_inst

    if compute_matched_centers:
        assert batch_size == 1
        # Generate trajectories
        matched_centers = {}
        _, seq_len, h, w = consistent_instance_seg.shape
        grid = torch.stack(torch.meshgrid(
            torch.arange(h, dtype=torch.float, device=preds.device),
            torch.arange(w, dtype=torch.float, device=preds.device)
        ))

        for instance_id in torch.unique(consistent_instance_seg[0, 0])[1:].cpu().numpy():
            for t in range(seq_len):
                instance_mask = consistent_instance_seg[0, t] == instance_id
                if instance_mask.sum() > 0:
                    matched_centers[instance_id] = matched_centers.get(instance_id, []) + [
                        grid[:, instance_mask].mean(dim=-1)]

        for key, value in matched_centers.items():
            matched_centers[key] = torch.stack(value).cpu().numpy()[:, ::-1]

        return consistent_instance_seg, matched_centers

    return consistent_instance_seg



data/stp3/utils/instance.py
llmbev/talk2bev
import torch
import torch.nn as nn
import torchvision

def pack_sequence_dim(x):
    b, s = x.shape[:2]
    return x.view(b * s, *x.shape[2:])


def unpack_sequence_dim(x, b, s):
    return x.view(b, s, *x.shape[1:])


def preprocess_batch(batch, device, unsqueeze=False):
    for key, value in batch.items():
        if torch.is_tensor(value):
            batch[key] = value.to(device)
            if unsqueeze:
                batch[key] = batch[key].unsqueeze(0)


def set_module_grad(module, requires_grad=False):
    for p in module.parameters():
        p.requires_grad = requires_grad


def set_bn_momentum(model, momentum=0.1):
    for m in model.modules():
        if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
            m.momentum = momentum


class NormalizeInverse(torchvision.transforms.Normalize):
    #  https://discuss.pytorch.org/t/simple-way-to-inverse-transform-normalization/4821/8
    def __init__(self, mean, std):
        mean = torch.as_tensor(mean)
        std = torch.as_tensor(std)
        std_inv = 1 / (std + 1e-7)
        mean_inv = -mean * std_inv
        super().__init__(mean=mean_inv, std=std_inv)

    def __call__(self, tensor):
        return super().__call__(tensor.clone())

data/stp3/utils/network.py
llmbev/talk2bev
# sampler.py
# trajectory sampler
# including a mix of clothoids, straight lines, and circles
import numpy as np
import torch
from scipy.special import fresnel

# s_offset = ssd[0,2,0].detach()
# d_offset = ssd[0,2,1].detach()
# vel = ssd[0,2] - ssd[0,1]
# theta = np.arctan2(vel[1], vel[0]) # to make angle 0 wrt ego-agent
# vel_prev = ssd[0,1] - ssd[0, 0]
# theta_prev = np.arctan2(vel_prev[1], vel_prev[0])
# psidot_init = (theta - theta_prev)/0.5
# psi_init = theta
# v_init = np.linalg.norm(vel)/0.5
# y_init = 0
# x_init = 0
# rot_inv = np.array([[cos(-theta), -sin(-theta)], [sin(-theta), cos(-theta)]])
# rot = np.array([[cos(theta), -sin(theta)], [sin(theta), cos(theta)]])
# #################################################################################

# """
#     CONVERT CENTERLINE ITSELF TO CENTERLINE FRAME
# """
# sdd_list = []
# trj = oracle_centerline[:, 10:-10]
# for xy in trj[0]:
#     s,d = spline.calc_frenet_position(xy[0].numpy(), xy[1].numpy())
#     sdd_list.append([s, d])
# sdd = torch.tensor(sdd_list).reshape(1, trj.shape[1], 2)

# ssd_frame = ssd -  torch.tensor([s_offset, d_offset])
# sdd_frame = sdd -  torch.tensor([s_offset, d_offset])
# eps_frame = grid - np.array([0, d_offset])
# #################################################################################

# """
#     CONVERT GRID TO AGENT FRAME
# """
# grid = grid - np.array([0, d_offset])
# #################################################################################

# """
#     ROTATE EVERYTHING TO AGENT FRAME
# """
# sdd_frame[0] = torch.tensor(np.dot(rot_inv, sdd_frame[0].numpy().T).T)
# ssd_frame = np.dot(rot_inv, ssd_frame[0].T).T
# eps_frame = np.dot(rot_inv, grid.T).T

# #################################################################################

# """
#     SOLVE THE PROBLEM
# """                        
# fixed_params = torch.zeros((1, 5)).expand(eps_frame.shape[0], 5)
# fixed_params[0] = torch.tensor([x_init, y_init, v_init, psi_init, psidot_init])
# variable_params = torch.zeros((eps_frame.shape[0], 3))
# variable_params[:, 0] = torch.tensor(eps_frame[:, 0])
# variable_params[:, 1] = torch.tensor(eps_frame[:, 1])
# variable_params[:, 2] = torch.tensor(-theta)
# P = problem.P
# cc, _ = problem.solve(fixed_params.double(), variable_params.double())
# x_sol = np.dot(P, cc.T[:nvar]).T
# y_sol = np.dot(P, cc.T[nvar:2*nvar]).T
# min_val = 1e11
# min_ind = 0
# #################################################################################

# """
#     COMPUTE MIN_TRAJ
# """
# for i in range(len(x_sol)):
#     traj = np.dstack((x_sol[i][::10], y_sol[i][::10]))[0]
#     dist = np.linalg.norm(traj - ssd_frame[2:])
#     if dist < min_val:
#         min_val = dist
#         min_ind = i
# #################################################################################

# """
#     PROJECT TO CARTESIAN FRAME 
# """
# trajs = np.dstack((x_sol[:, ::10], y_sol[:,::10])) # N, 7, 2
# for i in range(len(trajs)):
#     trajs[i] = np.dot(rot, trajs[i].T).T
# trajs = trajs + np.array([s_offset, d_offset])
# sdd_frame[0] = torch.tensor(np.dot(rot, sdd_frame[0].numpy().T).T)
# ssd_back = np.dot(rot, ssd_frame.T).T
# eps_back = np.dot(rot, eps_frame.T).T
# ssd_back = ssd_back + np.array([s_offset, d_offset])
# eps_back = eps_back + np.array([s_offset, d_offset])
# sdd_back = sdd_frame[0].numpy() + np.array([s_offset, d_offset])
# #################################################################################

# """
#     ADD TO OVERALL TRAJECTORY SET
# """
# for j in range(len(trajs)):
#     xys = []
#     for sds in trajs[j]:
#         x, y = spline.calc_global_position_offline(sds[0], sds[1])
#         xys.append([x, y])
#     trajs[j] = torch.tensor(xys)
#     total_trajs.append(trajs[j, :7])
# cxx = oracle_centerline[0, :, 0].detach()
# cyy = oracle_centerline[0, :, 1].detach()
# data_centerlines.append(np.dstack((cxx, cyy)))


def sample(v0, Kappa, T0, N0, tt, M, possibility = None):
    '''
    :param v0: initial velocity
    :param Kappa: curvature
    :param T0: initial tangent vector
    :param N0: initial normal vector
    :param tt: time stamp
    :param M: the number of sample
    :param possibility: torch.Tensor [3]
    :param debug: whether in debug mode
    :return: the nparray of trajectory
    '''
    # sample accelerations
    if possibility is None:
        possibility = [0.4, 0.2, 0.4]

    straight_num = int(M * possibility[1])
    left_num = int(M * possibility[0])
    right_num = int(M * possibility[2])

    accelerations = 10*(np.random.rand(M)-0.5) + 2  # -3m/s^2 to 7m/s^2

    # sample velocities
    # randomly sample a velocity <=15m/s at 80% of time
    v_options = np.stack((np.full(M, v0), 15*np.random.rand(M)))
    v_selections = (np.random.rand(M) >= 0.2).astype(int)
    velocities = v_options[v_selections, np.arange(M)]

    # generate longitudinal distances
    L = velocities[:, None] * tt[None, :] + accelerations[:, None] * (tt[None, :]**2) / 2
    L_straight = L[:straight_num]
    L = L[straight_num:]
    # print("L:", L)

    # scaling factor which determine the Clothiod curve  6 ~ 80
    alphas = (80 - 6) * np.random.rand(left_num + right_num) + 6

    ############################################################################
    # sample M straight lines
    line_points = L_straight[:, :, None] * T0[None, None, :]
    line_thetas = np.zeros_like(L_straight)
    lines = np.concatenate((line_points, line_thetas[:, :, None]), axis=-1)

    ############################################################################
    # sample M circles
    Krappa = min(-0.01, Kappa) if Kappa <= 0 else max(0.01, Kappa)
    radius = np.abs(1 / Krappa)
    center = np.array([-1 / Krappa, 0])
    circle_phis = L / radius if Krappa >= 0 else np.pi - L/radius

    circle_points = np.dstack([
        center[0] + radius * np.cos(circle_phis),
        center[1] + radius * np.sin(circle_phis),
    ])

    # rotate thetas, wrap
    circle_thetas = L/radius if Krappa >= 0 else -L/radius
    circle_thetas = (circle_thetas + np.pi) % (2 * np.pi) - np.pi
    circles = np.concatenate((circle_points, circle_thetas[:, :, None]), axis=-1)

    ############################################################################
    # sample M clothoids
    # Xi0 = Kappa / np.pi
    # Xis = Xi0 + L
    Xi0 = np.abs(Kappa) / np.pi
    Xis = Xi0 + L

    #
    # Ss, Cs = fresnel((Xis - Xi0) / alphas[:, None])
    Ss, Cs = fresnel(Xis / alphas[:, None])

    clothoid_points = alphas[:, None, None] * (Cs[:, :, None]*T0[None, None, :] + Ss[:, :, None]*N0[None, None, :])

    #
    Xs = clothoid_points[:, :, 0] - clothoid_points[:, 0, 0, None]
    Ys = clothoid_points[:, :, 1] - clothoid_points[:, 0, 1, None]
    clothoid_theta0s = 0.5 * np.pi * ((Kappa / np.pi / alphas) ** 2)
    clothoid_theta0s = clothoid_theta0s[:, None]
    signed_clothoid_theta0s = clothoid_theta0s * np.sign(Kappa)
    # when kappa is positive, the clothoid curves left, theta is positive
    # we will rotate it clockwise by theta
    # when kappa is negative, the clothoid curves right, theta is negative
    # we will rotate it counterclockwise by theta
    clothoid_points[:, :, 0] = np.cos(signed_clothoid_theta0s) * Xs + np.sin(signed_clothoid_theta0s) * Ys
    clothoid_points[:, :, 1] = - np.sin(signed_clothoid_theta0s) * Xs + np.cos(signed_clothoid_theta0s) * Ys

    # tangent vector: http://mathworld.wolfram.com/CornuSpiral.html
    clothoid_thetas = 0.5 * np.pi * ((Xis / alphas[:, None])**2)
    clothoid_thetas = clothoid_thetas - clothoid_theta0s
    signed_clothoid_thetas = clothoid_thetas * np.sign(Kappa)
    # clothoid_thetas = clothoid_thetas if Krappa >= 0 else -clothoid_thetas
    # wrap
    # clothoid_thetas = (clothoid_thetas + np.pi) % (2 * np.pi) - np.pi
    wrapped_signed_clothoid_thetas = (signed_clothoid_thetas + np.pi) % (2 * np.pi) - np.pi
    # wrapped_signed_clothoid_thetas = (signed_clothoid_thetas) % (2 * np.pi)
    #
    clothoids = np.concatenate((clothoid_points, wrapped_signed_clothoid_thetas[:, :, None]), axis=-1)

    ############################################################################
    # pick M in total
    t_options = np.stack((circles, clothoids))
    t_selections = np.random.choice([0, 1], size=left_num + right_num, p=(0.2, 0.8))
    ############################################################################

    trajs = t_options[t_selections, np.arange(left_num + right_num)]

    # toss a coin for vertical flipping
    # left_possibility = possibility[1] / (possibility[1] + possibility[2])
    # if Kappa > 0:
    #     heads = (np.random.rand(M) <= left_possibility.item())
    # else:
    #     heads = (np.random.rand(M) <= (1- left_possibility).item())
    # tails = np.logical_not(heads)
    #
    # # NOTE theta means what here
    # conditions = [heads[:, None, None], tails[:, None, None]]
    # choices = [trajs, np.dstack((
    #     -trajs[:, :, 0], trajs[:, :, 1], -trajs[:, :, 2]
    # ))]
    #
    # trajectories = np.select(conditions, choices)
    if Kappa > 0:
        left_curve = trajs[: left_num]
        right_curve = trajs[left_num: left_num + right_num]
        right_curve = np.dstack((
            -right_curve[:, :, 0], right_curve[:, :, 1], -right_curve[:, :, 2]
        ))
    else:
        right_curve = trajs[: left_num]
        left_curve = trajs[left_num: left_num + right_num]
        left_curve = np.dstack((
            -left_curve[:, :, 0], left_curve[:, :, 1], -left_curve[:, :, 2]
        ))

    trajectories = np.concatenate([left_curve, lines, right_curve], axis=0)
    mask = np.argsort(trajectories[:, -1, 0])
    trajectories = trajectories[mask]

    return trajectories

if __name__ == "__main__":
    from nuscenes.nuscenes import NuScenes
    from nuscenes.can_bus.can_bus_api import NuScenesCanBus
    import matplotlib.pyplot as plt
    nusc = NuScenes("v1.0-mini", '/home/hsc/data/Nuscenes')
    nusc_can = NuScenesCanBus(dataroot='/home/hsc/data/Nuscenes')

    for scene in nusc.scene:
        scene_name = scene["name"]
        scene_id = int(scene_name[-4:])
        if scene_id in nusc_can.can_blacklist:
            print(f"skipping {scene_name}")
            continue
        pose = nusc_can.get_messages(scene_name, "pose") # The current pose of the ego vehicle, sampled at 50 HZ
        saf = nusc_can.get_messages(scene_name, "steeranglefeedback") # Steering angle feedback in radians at 100 HZ
        vm = nusc_can.get_messages(scene_name, "vehicle_monitor") # information most, but sample at 2 HZ
        # NOTE: I tried to verify if the relevant measurements are consistent
        # across multiple tables that contain redundant information
        # NOTE: verified pose's velocity matches vehicle monitor's
        # but the pose table offers at a much higher frequency
        # NOTE: same that steeranglefeedback's steering angle matches vehicle monitor's
        # but the steeranglefeedback table offers at a much higher frequency
        print(pose[23])
        print(saf[45])
        print(vm[0])

        # initial velocity (m/s)
        v0 = pose[23]["vel"][0]
        # curvature
        #Kappa = 2 * saf[45]["value"] / 2.588  # 2 x \phi / distance between front and rear
        Kappa = 0
        # T0: longitudinal axis  Tangent vector
        T0 = np.array([0.0, 1.0])
        # N0: normal directional vector  Normal vector
        N0 = np.array([1.0, 0.0]) if Kappa <= 0 else np.array([-1.0, 0.0])
        # tt: time stamps
        tt = np.arange(0.0, 3.01, 0.01)
        # M: number of samples
        M = 1800
        #
        debug = False
        #
        trajectories = sample(v0, Kappa, T0, N0, tt, M)

        trajectories = trajectories[:,::100]

        #
        for i in range(len(trajectories)):
            trajectory = trajectories[i]
            plt.plot(trajectory[:, 0], trajectory[:, 1])
        plt.grid(False)
        plt.axis("equal")

        plt.show()

        break

data/stp3/utils/sampler.py
llmbev/talk2bev
"""
Cubic spline planner
Author: Atsushi Sakai(@Atsushi_twi)
Source: https://github.com/AtsushiSakai/PythonRobotics/blob/master/PathPlanning/CubicSpline/cubic_spline_planner.py
"""
import math
import numpy as np
import bisect


class Spline:
    """
    Cubic Spline class
    """

    def __init__(self, x, y):
        self.b, self.c, self.d, self.w = [], [], [], []

        self.x = np.array(x)
        self.y = np.array(y)

        self.eps = np.finfo(float).eps

        self.nx = len(x)  # dimension of x
        h = np.diff(x)

        # calc coefficient c
        self.a = np.array([iy for iy in y])

        # calc coefficient c
        A = self.__calc_A(h)
        B = self.__calc_B(h)
        self.c = np.linalg.solve(A, B)
        #  print(self.c1)

        # calc spline coefficient b and d
        for i in range(self.nx - 1):
            self.d.append((self.c[i + 1] - self.c[i]) / (3.0 * h[i] + self.eps))
            tb = (self.a[i + 1] - self.a[i]) / (h[i] + self.eps) - h[i] * \
                (self.c[i + 1] + 2.0 * self.c[i]) / 3.0
            self.b.append(tb)
        self.b = np.array(self.b)
        self.d = np.array(self.d)

    def calc(self, t):
        """
        Calc position
        if t is outside of the input x, return None
        """
        t = np.asarray(t)
        mask = np.logical_and(t < self.x[0], t > self.x[-1])
        t[mask] = self.x[0]

        i = self.__search_index(t)
        dx = t - self.x[i.astype(int)]
        result = self.a[i] + self.b[i] * dx + \
            self.c[i] * dx ** 2.0 + self.d[i] * dx ** 3.0

        result = np.asarray(result)
        result[mask] = None
        return result

    def calcd(self, t):
        """
        Calc first derivative
        if t is outside of the input x, return None
        """
        t = np.asarray(t)
        mask = np.logical_and(t < self.x[0], t > self.x[-1])
        t[mask] = 0

        i = self.__search_index(t)
        dx = t - self.x[i]
        result = self.b[i] + 2.0 * self.c[i] * dx + 3.0 * self.d[i] * dx ** 2.0

        result = np.asarray(result)
        result[mask] = None
        return result

    def calcdd(self, t):
        """
        Calc second derivative
        """
        t = np.asarray(t)
        mask = np.logical_and(t < self.x[0], t > self.x[-1])
        t[mask] = 0

        i = self.__search_index(t)
        dx = t - self.x[i]
        result = 2.0 * self.c[i] + 6.0 * self.d[i] * dx

        result = np.asarray(result)
        result[mask] = None
        return result

    def __search_index(self, x):
        """
        search data segment index
        """
        indices = np.asarray(np.searchsorted(self.x, x, "left") - 1)
        indices[indices <= 0] = 0
        return indices

    def __calc_A(self, h):
        """
        calc matrix A for spline coefficient c
        """
        A = np.zeros((self.nx, self.nx))
        A[0, 0] = 1.0
        for i in range(self.nx - 1):
            if i != (self.nx - 2):
                A[i + 1, i + 1] = 2.0 * (h[i] + h[i + 1])
            A[i + 1, i] = h[i]
            A[i, i + 1] = h[i]

        A[0, 1] = 0.0
        A[self.nx - 1, self.nx - 2] = 0.0
        A[self.nx - 1, self.nx - 1] = 1.0
        #  print(A)
        return A

    def __calc_B(self, h):
        """
        calc matrix B for spline coefficient c
        """
        B = np.zeros(self.nx)
        for i in range(self.nx - 2):
            B[i + 1] = 3.0 * (self.a[i + 2] - self.a[i + 1]) / (h[i + 1] + self.eps) \
                       - 3.0 * (self.a[i + 1] - self.a[i]) / (h[i] + self.eps)
        return B


class Spline2D:
    """
    2D Cubic Spline class
    """

    def __init__(self, x, y, resolution=0.1, extension=300, debug=False):
        # resolution of 1m, extension of 300m on both sides
        # trail and append linear extensions
        # angle_start = np.arctan2(y[1] - y[0], x[1] - x[0])
        # angle_end = np.arctan2(y[-1] - y[-2], x[-1] - x[-2])
        
        # x_start = np.arange(x[0] - extension, x[0], 0.5)
        # x_end = np.arange(x[-1], x[-1] + extension, 0.5)
        
        # y_start = y[0] + np.tan(angle_start) * np.arange(- extension, 0, 0.5)
        # y_end = y[-1] + np.tan(angle_end) * np.arange(0, 0 + extension, 0.5)

        # x = np.concatenate((x_start, x.reshape(x.shape[0]), x_end))
        # y = np.concatenate((y_start, y.reshape(y.shape[0]), y_end))

        self.splinex = x
        self.spliney = y
        
        if debug:
            import matplotlib.pyplot as plt
            plt.plot(x,y)
            plt.title("Reference Spline")
            plt.show()
            plt.clf()
        
        self.s = self.__calc_s(x, y)
        self.sx = Spline(self.s, x)
        self.sy = Spline(self.s, y)

        self.s_fine = np.arange(0, self.s[-1], resolution)
        xy = np.array([self.calc_global_position_online(s_i) for s_i in self.s_fine])

        self.x_fine = xy[:, 0]
        self.y_fine = xy[:, 1]

    def __calc_s(self, x, y):
        dx = np.diff(x)
        dy = np.diff(y)
        self.ds = np.hypot(dx, dy)
        s = [0]
        s.extend(np.cumsum(self.ds))
        return s

    def calc_global_position_online(self, s):
        """
        calc global position of points on the line, s: float
        return: x: float; y: float; the global coordinate of given s on the spline
        """
        x = self.sx.calc(s)
        y = self.sy.calc(s)

        return x, y

    def calc_global_position_offline(self, s, d):
        """
        calc global position of points in the frenet coordinate w.r.t. the line.
        s: float, longitudinal; d: float, lateral;
        return: x, float; y, float;
        """
        s_x = self.sx.calc(s)
        s_y = self.sy.calc(s)

        theta = math.atan2(self.sy.calcd(s), self.sx.calcd(s))
        x = s_x - math.sin(theta) * d
        y = s_y + math.cos(theta) * d
        return x, y

    def calc_frenet_position(self, x, y):
        """
        cal the frenet position of given global coordinate (x, y)
        return s: the longitudinal; d: the lateral
        """
        # find nearst x, y
        diff = np.hypot(self.x_fine - x, self.y_fine - y)
        idx = np.argmin(diff)
        [x_s, y_s] = self.x_fine[idx], self.y_fine[idx]
        s = self.s_fine[idx]

        # compute theta
        theta = math.atan2(self.sy.calcd(s), self.sx.calcd(s))
        d_x, d_y = x - x_s, y - y_s
        cross_rd_nd = math.cos(theta) * d_y - math.sin(theta) * d_x
        d = math.copysign(np.hypot(d_x, d_y), cross_rd_nd)
        return s, d

    def calc_curvature(self, s):
        """
        calc curvature
        """
        dx = self.sx.calcd(s)
        ddx = self.sx.calcdd(s)
        dy = self.sy.calcd(s)
        ddy = self.sy.calcdd(s)
        k = (ddy * dx - ddx * dy) / ((dx ** 2 + dy ** 2)**(3 / 2))
        return k

    def calc_yaw(self, s):
        """
        calc yaw
        """
        dx = self.sx.calcd(s)
        dy = self.sy.calcd(s)
        yaw = np.arctan2(dy, dx)
        return yaw


def calc_spline_course(x, y, ds=0.1):
    sp = Spline2D(x, y)
    s = list(np.arange(0, sp.s[-1], ds))

    rx, ry, ryaw, rk = [], [], [], []
    for i_s in s:
        ix, iy = sp.calc_global_position_online(i_s)
        rx.append(ix)
        ry.append(iy)
        ryaw.append(sp.calc_yaw(i_s))
        rk.append(sp.calc_curvature(i_s))

    return rx, ry, ryaw, rk, s


def main():  # pragma: no cover
    print("Spline 2D test")
    import matplotlib.pyplot as plt
    # x = [-2.5, 0.0, 2.5, 5.0, 7.5, 3.0, -1.0]
    # y = [0.7, -6, 5, 6.5, 0.0, 5.0, -2.0]
    x = [float(i) for i in range(30)]
    y = [float(0.0) for i in range(30)]
    for i in range(len(y)):
        if i % 2 == 0.0:
            y[i] = float(i)

    ds = 0.1  # [m] distance of each intepolated points

    sp = Spline2D(x, y)
    s = np.arange(0, sp.s[-1], ds)

    rx, ry, ryaw, rk = [], [], [], []
    rx, ry = sp.calc_global_position_online(s)
    ryaw = sp.calc_yaw(s)
    rk = sp.calc_curvature(s)

    plt.subplots(1)
    plt.plot(x, y, "xb", label="input")
    plt.plot(rx, ry, "-r", label="spline")
    plt.grid(True)
    plt.axis("equal")
    plt.xlabel("x[m]")
    plt.ylabel("y[m]")
    plt.legend()

    plt.subplots(1)
    plt.plot(s, [np.rad2deg(iyaw) for iyaw in ryaw], "-r", label="yaw")
    plt.grid(True)
    plt.legend()
    plt.xlabel("line length[m]")
    plt.ylabel("yaw angle[deg]")

    plt.subplots(1)
    plt.plot(s, rk, "-r", label="curvature")
    plt.grid(True)
    plt.legend()
    plt.xlabel("line length[m]")
    plt.ylabel("curvature [1/m]")

    plt.show()


if __name__ == '__main__':
    main()

data/stp3/utils/spline.py
llmbev/talk2bev
"""
Copyright (C) 2020 NVIDIA Corporation.  All rights reserved.
Licensed under the NVIDIA Source Code License. See LICENSE at https://github.com/nv-tlabs/lift-splat-shoot.
Authors: Jonah Philion and Sanja Fidler
"""

import os
import numpy as np
import torch
import torch.nn as nn
import torchvision
import shutil
from tqdm import tqdm
from pyquaternion import Quaternion
from PIL import Image
from functools import reduce
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
from nuscenes.utils.data_classes import LidarPointCloud
from nuscenes.utils.geometry_utils import transform_matrix
from nuscenes.map_expansion.map_api import NuScenesMap


def get_lidar_data(nusc, sample_rec, nsweeps, min_distance):
    """
    Returns at most nsweeps of lidar in the ego frame.
    Returned tensor is 5(x, y, z, reflectance, dt) x N
    Adapted from https://github.com/nutonomy/nuscenes-devkit/blob/master/python-sdk/nuscenes/utils/data_classes.py#L56
    """
    points = np.zeros((5, 0))

    # Get reference pose and timestamp.
    ref_sd_token = sample_rec['data']['LIDAR_TOP']
    ref_sd_rec = nusc.get('sample_data', ref_sd_token)
    ref_pose_rec = nusc.get('ego_pose', ref_sd_rec['ego_pose_token'])
    ref_cs_rec = nusc.get('calibrated_sensor', ref_sd_rec['calibrated_sensor_token'])
    ref_time = 1e-6 * ref_sd_rec['timestamp']

    # Homogeneous transformation matrix from global to _current_ ego car frame.
    car_from_global = transform_matrix(ref_pose_rec['translation'], Quaternion(ref_pose_rec['rotation']),
                                        inverse=True)

    # Aggregate current and previous sweeps.
    sample_data_token = sample_rec['data']['LIDAR_TOP']
    current_sd_rec = nusc.get('sample_data', sample_data_token)
    for _ in range(nsweeps):
        # Load up the pointcloud and remove points close to the sensor.
        current_pc = LidarPointCloud.from_file(os.path.join(nusc.dataroot, current_sd_rec['filename']))
        current_pc.remove_close(min_distance)

        # Get past pose.
        current_pose_rec = nusc.get('ego_pose', current_sd_rec['ego_pose_token'])
        global_from_car = transform_matrix(current_pose_rec['translation'],
                                            Quaternion(current_pose_rec['rotation']), inverse=False)

        # Homogeneous transformation matrix from sensor coordinate frame to ego car frame.
        current_cs_rec = nusc.get('calibrated_sensor', current_sd_rec['calibrated_sensor_token'])
        car_from_current = transform_matrix(current_cs_rec['translation'], Quaternion(current_cs_rec['rotation']),
                                            inverse=False)

        # Fuse four transformation matrices into one and perform transform.
        trans_matrix = reduce(np.dot, [car_from_global, global_from_car, car_from_current])
        current_pc.transform(trans_matrix)

        # Add time vector which can be used as a temporal feature.
        time_lag = ref_time - 1e-6 * current_sd_rec['timestamp']
        times = time_lag * np.ones((1, current_pc.nbr_points()))

        new_points = np.concatenate((current_pc.points, times), 0)
        points = np.concatenate((points, new_points), 1)

        # Abort if there are no previous sweeps.
        if current_sd_rec['prev'] == '':
            break
        else:
            current_sd_rec = nusc.get('sample_data', current_sd_rec['prev'])

    return points


def ego_to_cam(points, rot, trans, intrins):
    """Transform points (3 x N) from ego frame into a pinhole camera
    """
    points = points - trans.unsqueeze(1)
    points = rot.permute(1, 0).matmul(points)

    points = intrins.matmul(points)
    points[:2] /= points[2:3]

    return points


def cam_to_ego(points, rot, trans, intrins):
    """Transform points (3 x N) from pinhole camera with depth
    to the ego frame
    """
    points = torch.cat((points[:2] * points[2:3], points[2:3]))
    points = intrins.inverse().matmul(points)

    points = rot.matmul(points)
    points += trans.unsqueeze(1)

    return points


def get_only_in_img_mask(pts, H, W):
    """pts should be 3 x N
    """
    return (pts[2] > 0) &\
        (pts[0] > 1) & (pts[0] < W - 1) &\
        (pts[1] > 1) & (pts[1] < H - 1)


def get_rot(h):
    return torch.Tensor([
        [np.cos(h), np.sin(h)],
        [-np.sin(h), np.cos(h)],
    ])


def img_transform(img, post_rot, post_tran,
                  resize, resize_dims, crop,
                  flip, rotate):
    # adjust image
    img = img.resize(resize_dims)
    img = img.crop(crop)
    if flip:
        img = img.transpose(method=Image.FLIP_LEFT_RIGHT)
    img = img.rotate(rotate)

    # post-homography transformation
    post_rot *= resize
    post_tran -= torch.Tensor(crop[:2])
    if flip:
        A = torch.Tensor([[-1, 0], [0, 1]])
        b = torch.Tensor([crop[2] - crop[0], 0])
        post_rot = A.matmul(post_rot)
        post_tran = A.matmul(post_tran) + b
    A = get_rot(rotate/180*np.pi)
    b = torch.Tensor([crop[2] - crop[0], crop[3] - crop[1]]) / 2
    b = A.matmul(-b) + b
    post_rot = A.matmul(post_rot)
    post_tran = A.matmul(post_tran) + b

    return img, post_rot, post_tran


class NormalizeInverse(torchvision.transforms.Normalize):
    #  https://discuss.pytorch.org/t/simple-way-to-inverse-transform-normalization/4821/8
    def __init__(self, mean, std):
        mean = torch.as_tensor(mean)
        std = torch.as_tensor(std)
        std_inv = 1 / (std + 1e-7)
        mean_inv = -mean * std_inv
        super().__init__(mean=mean_inv, std=std_inv)

    def __call__(self, tensor):
        return super().__call__(tensor.clone())


denormalize_img = torchvision.transforms.Compose((
            NormalizeInverse(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225]),
            torchvision.transforms.ToPILImage(),
        ))


normalize_img = torchvision.transforms.Compose((
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
))


def gen_dx_bx(xbound, ybound, zbound):
    dx = torch.Tensor([row[2] for row in [xbound, ybound, zbound]])
    bx = torch.Tensor([row[0] + row[2]/2.0 for row in [xbound, ybound, zbound]])
    nx = torch.LongTensor([(row[1] - row[0]) / row[2] for row in [xbound, ybound, zbound]])

    return dx, bx, nx


def cumsum_trick(x, geom_feats, ranks):
    x = x.cumsum(0)
    kept = torch.ones(x.shape[0], device=x.device, dtype=torch.bool)
    kept[:-1] = (ranks[1:] != ranks[:-1])

    x, geom_feats = x[kept], geom_feats[kept]
    x = torch.cat((x[:1], x[1:] - x[:-1]))

    return x, geom_feats


class QuickCumsum(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, geom_feats, ranks):
        x = x.cumsum(0)
        kept = torch.ones(x.shape[0], device=x.device, dtype=torch.bool)
        kept[:-1] = (ranks[1:] != ranks[:-1])

        x, geom_feats = x[kept], geom_feats[kept]
        x = torch.cat((x[:1], x[1:] - x[:-1]))

        # save kept for backward
        ctx.save_for_backward(kept)

        # no gradient for geom_feats
        ctx.mark_non_differentiable(geom_feats)

        return x, geom_feats

    @staticmethod
    def backward(ctx, gradx, gradgeom):
        kept, = ctx.saved_tensors
        back = torch.cumsum(kept, 0)
        back[kept] -= 1

        val = gradx[back]

        return val, None, None


class SimpleLoss(torch.nn.Module):
    def __init__(self, pos_weight):
        super(SimpleLoss, self).__init__()
        self.loss_fn = torch.nn.BCEWithLogitsLoss(pos_weight=torch.Tensor([pos_weight]))

    def forward(self, ypred, ytgt):
        loss = self.loss_fn(ypred, ytgt)
        return loss


def get_batch_iou(preds, binimgs):
    """Assumes preds has NOT been sigmoided yet
    """
    with torch.no_grad():
        pred = (preds > 0)
        tgt = binimgs.bool()
        intersect = (pred & tgt).sum().float().item()
        union = (pred | tgt).sum().float().item()
    return intersect, union, intersect / union if (union > 0) else 1.0


def get_val_info(model, valloader, loss_fn, device, use_tqdm=False):
    model.eval()
    total_loss = 0.0
    total_intersect = 0.0
    total_union = 0
    print('running eval...')
    loader = tqdm(valloader) if use_tqdm else valloader
    with torch.no_grad():
        for batch in loader:
            allimgs, rots, trans, intrins, post_rots, post_trans, binimgs = batch
            preds = model(allimgs.to(device), rots.to(device),
                          trans.to(device), intrins.to(device), post_rots.to(device),
                          post_trans.to(device))
            binimgs = binimgs.to(device)

            # loss
            total_loss += loss_fn(preds, binimgs).item() * preds.shape[0]

            # iou
            intersect, union, _ = get_batch_iou(preds, binimgs)
            total_intersect += intersect
            total_union += union

    model.train()
    return {
            'loss': total_loss / len(valloader.dataset),
            'iou': total_intersect / total_union,
            }


def add_ego(bx, dx):
    # approximate rear axel
    W = 1.85
    pts = np.array([
        [-4.084/2.+0.5, W/2.],
        [4.084/2.+0.5, W/2.],
        [4.084/2.+0.5, -W/2.],
        [-4.084/2.+0.5, -W/2.],
    ])
    pts = (pts - bx) / dx
    pts[:, [0,1]] = pts[:, [1,0]]
    plt.fill(pts[:, 0], pts[:, 1], '#76b900')


def get_nusc_maps(map_folder):
    nusc_maps = {map_name: NuScenesMap(dataroot=map_folder,
                map_name=map_name) for map_name in [
                    "singapore-hollandvillage",
                    "singapore-queenstown",
                    "boston-seaport",
                    "singapore-onenorth",
                ]}
    return nusc_maps


def plot_nusc_map(rec, nusc_maps, nusc, scene2map, dx, bx):
    egopose = nusc.get('ego_pose', nusc.get('sample_data', rec['data']['LIDAR_TOP'])['ego_pose_token'])
    map_name = scene2map[nusc.get('scene', rec['scene_token'])['name']]

    rot = Quaternion(egopose['rotation']).rotation_matrix
    rot = np.arctan2(rot[1, 0], rot[0, 0])
    center = np.array([egopose['translation'][0], egopose['translation'][1], np.cos(rot), np.sin(rot)])

    poly_names = ['road_segment', 'lane']
    line_names = ['road_divider', 'lane_divider']
    lmap = get_local_map(nusc_maps[map_name], center,
                         50.0, poly_names, line_names)
    for name in poly_names:
        for la in lmap[name]:
            pts = (la - bx) / dx
            plt.fill(pts[:, 1], pts[:, 0], c=(1.00, 0.50, 0.31), alpha=0.2)
    for la in lmap['road_divider']:
        pts = (la - bx) / dx
        plt.plot(pts[:, 1], pts[:, 0], c=(0.0, 0.0, 1.0), alpha=0.5)
    for la in lmap['lane_divider']:
        pts = (la - bx) / dx
        plt.plot(pts[:, 1], pts[:, 0], c=(159./255., 0.0, 1.0), alpha=0.5)


def get_local_map(nmap, center, stretch, layer_names, line_names):
    # need to get the map here...
    box_coords = (
        center[0] - stretch,
        center[1] - stretch,
        center[0] + stretch,
        center[1] + stretch,
    )

    polys = {}

    # polygons
    records_in_patch = nmap.get_records_in_patch(box_coords,
                                                 layer_names=layer_names,
                                                 mode='intersect')
    for layer_name in layer_names:
        polys[layer_name] = []
        for token in records_in_patch[layer_name]:
            poly_record = nmap.get(layer_name, token)
            if layer_name == 'drivable_area':
                polygon_tokens = poly_record['polygon_tokens']
            else:
                polygon_tokens = [poly_record['polygon_token']]

            for polygon_token in polygon_tokens:
                polygon = nmap.extract_polygon(polygon_token)
                polys[layer_name].append(np.array(polygon.exterior.xy).T)

    # lines
    for layer_name in line_names:
        polys[layer_name] = []
        for record in getattr(nmap, layer_name):
            token = record['token']

            line = nmap.extract_line(record['line_token'])
            if line.is_empty:  # Skip lines without nodes
                continue
            xs, ys = line.xy

            polys[layer_name].append(
                np.array([xs, ys]).T
                )

    # convert to local coordinates in place
    rot = get_rot(np.arctan2(center[3], center[2])).T
    for layer_name in polys:
        for rowi in range(len(polys[layer_name])):
            polys[layer_name][rowi] -= center[:2]
            polys[layer_name][rowi] = np.dot(polys[layer_name][rowi], rot)

    return polys

def create_exp_dir(path, scripts_to_save=None, prefix=' '):
    if not os.path.exists(path):
        os.makedirs(path)
    print('Experiment dir : {}'.format(path))

    if scripts_to_save is not None:
        os.mkdir(os.path.join(path, 'scripts'))
        for script in scripts_to_save:
            if script.startswith(prefix):
                continue
            dir_name = os.path.dirname(script)
            if not os.path.exists(os.path.join(path,'scripts', dir_name)):
                os.makedirs(os.path.join(path,'scripts', dir_name))
            dst_file = os.path.join(path, 'scripts', script)
            shutil.copyfile(script, dst_file)

class KalmanFilter(nn.Module):
    def __init__(self, F = None, B = None, H = None, Q = None, R = None, P = None, x0 = None):
        super(KalmanFilter, self).__init__()

        if (F is None or H is None):
            raise ValueError('Set F and H in kalman setting')

        c, h, w = F.shape
        self.c = c
        self.h = h
        self.w = w

        self.F = F
        self.H = H

        self.B = nn.Parameter(torch.zeros(1), requires_grad=False) if B is None else nn.Parameter(B)
        self.Q = nn.Parameter(torch.eye(h, w).view(1, h, w).expand(c, h, w), requires_grad=False) if Q is None \
            else nn.Parameter(Q)
        self.R = nn.Parameter(torch.eye(h, w).view(1, h, w).expand(c, h, w), requires_grad=False) if R is None \
            else nn.Parameter(R)
        self.P = nn.Parameter(torch.eye(h, w).view(1, h, w).expand(c, h, w), requires_grad=False) if P is None \
            else nn.Parameter(P)
        self.x = nn.Parameter(torch.zeros(c, h, w), requires_grad=False) if x0 is None \
            else nn.Parameter(x0)

    def predict(self, u = None):
        if u == None:
            u = torch.zeros(1, device=self.B.device)

        self.x.data = self.F @ self.x + self.B @ u
        F_T = torch.transpose(self.F, -1, -2)
        self.P.data = self.F @ self.P @ F_T + self.Q
        return self.x

    def update(self, z):
        y = z - self.H @ self.x
        H_T = torch.transpose(self.H, -1, -2)
        S = self.R + self.H @ self.P @ H_T
        K = self.P @ H_T @ torch.inverse(S)
        self.x.data = self.x + K @ y
        I = torch.eye(self.h, self.w).view(1, self.h, self.w).expand(self.c, self.h, self.w).to(device=z.device)
        tmp = (I - K @ self.H)
        tmp_T = torch.transpose(tmp, -1, -2)
        K_T = torch.transpose(K, -1, -2)
        self.P.data = tmp @ self.P @ tmp_T + K @ self.R @ K_T



data/stp3/utils/tools.py
llmbev/talk2bev
import numpy as np
import torch
import matplotlib.pylab
import matplotlib.pyplot as plt
from stp3.utils.tools import gen_dx_bx


from stp3.utils.instance import predict_instance_segmentation_and_trajectories

DEFAULT_COLORMAP = matplotlib.pylab.cm.jet


def flow_to_image(flow: np.ndarray, autoscale: bool = False) -> np.ndarray:
    """
    Applies colour map to flow which should be a 2 channel image tensor HxWx2. Returns a HxWx3 numpy image
    Code adapted from: https://github.com/liruoteng/FlowNet/blob/master/models/flownet/scripts/flowlib.py
    """
    u = flow[0, :, :]
    v = flow[1, :, :]

    # Convert to polar coordinates
    rad = np.sqrt(u ** 2 + v ** 2)
    maxrad = np.max(rad)

    # Normalise flow maps
    if autoscale:
        u /= maxrad + np.finfo(float).eps
        v /= maxrad + np.finfo(float).eps

    # visualise flow with cmap
    return np.uint8(compute_color(u, v) * 255)


def _normalise(image: np.ndarray) -> np.ndarray:
    lower = np.min(image)
    delta = np.max(image) - lower
    if delta == 0:
        delta = 1
    image = (image.astype(np.float32) - lower) / delta
    return image


def apply_colour_map(
    image: np.ndarray, cmap: matplotlib.colors.LinearSegmentedColormap = DEFAULT_COLORMAP, autoscale: bool = False
) -> np.ndarray:
    """
    Applies a colour map to the given 1 or 2 channel numpy image. if 2 channel, must be 2xHxW.
    Returns a HxWx3 numpy image
    """
    if image.ndim == 2 or (image.ndim == 3 and image.shape[0] == 1):
        if image.ndim == 3:
            image = image[0]
        # grayscale scalar image
        if autoscale:
            image = _normalise(image)
        return cmap(image)[:, :, :3]
    if image.shape[0] == 2:
        # 2 dimensional UV
        return flow_to_image(image, autoscale=autoscale)
    if image.shape[0] == 3:
        # normalise rgb channels
        if autoscale:
            image = _normalise(image)
        return np.transpose(image, axes=[1, 2, 0])
    raise Exception('Image must be 1, 2 or 3 channel to convert to colour_map (CxHxW)')


def heatmap_image(
    image: np.ndarray, cmap: matplotlib.colors.LinearSegmentedColormap = DEFAULT_COLORMAP, autoscale: bool = True
) -> np.ndarray:
    """Colorize an 1 or 2 channel image with a colourmap."""
    if not issubclass(image.dtype.type, np.floating):
        raise ValueError(f"Expected a ndarray of float type, but got dtype {image.dtype}")
    if not (image.ndim == 2 or (image.ndim == 3 and image.shape[0] in [1, 2])):
        raise ValueError(f"Expected a ndarray of shape [H, W] or [1, H, W] or [2, H, W], but got shape {image.shape}")
    heatmap_np = apply_colour_map(image, cmap=cmap, autoscale=autoscale)
    heatmap_np = np.uint8(heatmap_np * 255)
    return heatmap_np


def compute_color(u: np.ndarray, v: np.ndarray) -> np.ndarray:
    assert u.shape == v.shape
    [h, w] = u.shape
    img = np.zeros([h, w, 3])
    nan_mask = np.isnan(u) | np.isnan(v)
    u[nan_mask] = 0
    v[nan_mask] = 0

    colorwheel = make_color_wheel()
    ncols = np.size(colorwheel, 0)

    rad = np.sqrt(u ** 2 + v ** 2)
    a = np.arctan2(-v, -u) / np.pi
    f_k = (a + 1) / 2 * (ncols - 1) + 1
    k_0 = np.floor(f_k).astype(int)
    k_1 = k_0 + 1
    k_1[k_1 == ncols + 1] = 1
    f = f_k - k_0

    for i in range(0, np.size(colorwheel, 1)):
        tmp = colorwheel[:, i]
        col0 = tmp[k_0 - 1] / 255
        col1 = tmp[k_1 - 1] / 255
        col = (1 - f) * col0 + f * col1

        idx = rad <= 1
        col[idx] = 1 - rad[idx] * (1 - col[idx])
        notidx = np.logical_not(idx)

        col[notidx] *= 0.75
        img[:, :, i] = col * (1 - nan_mask)

    return img


def make_color_wheel() -> np.ndarray:
    """
    Create colour wheel.
    Code adapted from https://github.com/liruoteng/FlowNet/blob/master/models/flownet/scripts/flowlib.py
    """
    red_yellow = 15
    yellow_green = 6
    green_cyan = 4
    cyan_blue = 11
    blue_magenta = 13
    magenta_red = 6

    ncols = red_yellow + yellow_green + green_cyan + cyan_blue + blue_magenta + magenta_red
    colorwheel = np.zeros([ncols, 3])

    col = 0

    # red_yellow
    colorwheel[0:red_yellow, 0] = 255
    colorwheel[0:red_yellow, 1] = np.transpose(np.floor(255 * np.arange(0, red_yellow) / red_yellow))
    col += red_yellow

    # yellow_green
    colorwheel[col : col + yellow_green, 0] = 255 - np.transpose(
        np.floor(255 * np.arange(0, yellow_green) / yellow_green)
    )
    colorwheel[col : col + yellow_green, 1] = 255
    col += yellow_green

    # green_cyan
    colorwheel[col : col + green_cyan, 1] = 255
    colorwheel[col : col + green_cyan, 2] = np.transpose(np.floor(255 * np.arange(0, green_cyan) / green_cyan))
    col += green_cyan

    # cyan_blue
    colorwheel[col : col + cyan_blue, 1] = 255 - np.transpose(np.floor(255 * np.arange(0, cyan_blue) / cyan_blue))
    colorwheel[col : col + cyan_blue, 2] = 255
    col += cyan_blue

    # blue_magenta
    colorwheel[col : col + blue_magenta, 2] = 255
    colorwheel[col : col + blue_magenta, 0] = np.transpose(np.floor(255 * np.arange(0, blue_magenta) / blue_magenta))
    col += +blue_magenta

    # magenta_red
    colorwheel[col : col + magenta_red, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, magenta_red) / magenta_red))
    colorwheel[col : col + magenta_red, 0] = 255

    return colorwheel


def make_contour(img, colour=[0, 0, 0], double_line=False):
    h, w = img.shape[:2]
    out = img.copy()
    # Vertical lines
    out[np.arange(h), np.repeat(0, h)] = colour
    out[np.arange(h), np.repeat(w - 1, h)] = colour

    # Horizontal lines
    out[np.repeat(0, w), np.arange(w)] = colour
    out[np.repeat(h - 1, w), np.arange(w)] = colour

    if double_line:
        out[np.arange(h), np.repeat(1, h)] = colour
        out[np.arange(h), np.repeat(w - 2, h)] = colour

        # Horizontal lines
        out[np.repeat(1, w), np.arange(w)] = colour
        out[np.repeat(h - 2, w), np.arange(w)] = colour
    return out


def plot_instance_map(instance_image, instance_map, instance_colours=None, bg_image=None):
    if isinstance(instance_image, torch.Tensor):
        instance_image = instance_image.cpu().numpy()
    assert isinstance(instance_image, np.ndarray)
    if instance_colours is None:
        instance_colours = generate_instance_colours(instance_map)
    if len(instance_image.shape) > 2:
        instance_image = instance_image.reshape((instance_image.shape[-2], instance_image.shape[-1]))

    if bg_image is None:
        plot_image = 255 * np.ones((instance_image.shape[0], instance_image.shape[1], 3), dtype=np.uint8)
    else:
        plot_image = bg_image

    for key, value in instance_colours.items():
        plot_image[instance_image == key] = value

    return plot_image


def visualise_output(labels, output, cfg):
    semantic_colours = np.array([[255, 255, 255], [0, 0, 0]], dtype=np.uint8)

    if cfg.INSTANCE_SEG.ENABLED:
        consistent_instance_seg = predict_instance_segmentation_and_trajectories(
            output, compute_matched_centers=False
        )

    sequence_length = labels['segmentation'].shape[1]
    b = 0
    video = []
    for t in range(sequence_length):
        out_t = []
        # Ground truth
        semantic_seg = labels['segmentation'].squeeze(2).cpu().numpy()
        semantic_plot = semantic_colours[semantic_seg[b, t][::-1, ::-1]]
        semantic_plot = make_contour(semantic_plot)

        if cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
            pedestrian_seg = labels['pedestrian'].squeeze(2).cpu().numpy()
            pedestrian_plot = semantic_colours[pedestrian_seg[b, t][::-1, ::-1]]
            pedestrian_plot = make_contour(pedestrian_plot)
        else:
            pedestrian_plot = np.zeros_like(semantic_plot)

        if cfg.INSTANCE_SEG.ENABLED:
            unique_ids = torch.unique(labels['instance'][b, t]).cpu().numpy()[1:]
            instance_map = dict(zip(unique_ids, unique_ids))
            instance_plot = plot_instance_map(labels['instance'][b, t].cpu(), instance_map)[::-1, ::-1]
            instance_plot = make_contour(instance_plot)

            center_plot = heatmap_image(labels['centerness'][b, t, 0].cpu().numpy())[::-1, ::-1]
            center_plot = make_contour(center_plot)

            offset_plot = labels['offset'][b, t].cpu().numpy()
            offset_plot[:, semantic_seg[b, t] != 1] = 0
            offset_plot = flow_to_image(offset_plot)[::-1, ::-1]
            offset_plot = make_contour(offset_plot)
        else:
            instance_plot = np.zeros_like(semantic_plot)
            center_plot = np.zeros_like(semantic_plot)
            offset_plot = np.zeros_like(semantic_plot)

        if cfg.INSTANCE_FLOW.ENABLED:
            future_flow_plot = labels['flow'][b, t].cpu().numpy()
            future_flow_plot[:, semantic_seg[b, t] != 1] = 0
            future_flow_plot = flow_to_image(future_flow_plot)[::-1, ::-1]
            future_flow_plot = make_contour(future_flow_plot)
        else:
            future_flow_plot = np.zeros_like(semantic_plot)

        planning_plot = plot_planning(labels['hdmap'][b], labels['gt_trajectory'][b], cfg)
        planning_plot = make_contour(planning_plot)

        out_t.append(np.concatenate([instance_plot, future_flow_plot,
                                     semantic_plot, center_plot, offset_plot, pedestrian_plot, planning_plot], axis=0))

        # Predictions
        semantic_seg = output['segmentation'].argmax(dim=2).detach().cpu().numpy()
        semantic_plot = semantic_colours[semantic_seg[b, t][::-1, ::-1]]
        semantic_plot = make_contour(semantic_plot)

        if cfg.SEMANTIC_SEG.PEDESTRIAN.ENABLED:
            pedestrian_seg = output['pedestrian'].argmax(dim=2).detach().cpu().numpy()
            pedestrian_plot = semantic_colours[pedestrian_seg[b, t][::-1, ::-1]]
            pedestrian_plot = make_contour(pedestrian_plot)
        else:
            pedestrian_plot = np.zeros_like(semantic_plot)

        if cfg.INSTANCE_SEG.ENABLED:
            unique_ids = torch.unique(consistent_instance_seg[b, t]).cpu().numpy()[1:]
            instance_map = dict(zip(unique_ids, unique_ids))
            instance_plot = plot_instance_map(consistent_instance_seg[b, t].cpu(), instance_map)[::-1, ::-1]
            instance_plot = make_contour(instance_plot)

            center_plot = heatmap_image(output['instance_center'][b, t, 0].detach().cpu().numpy())[::-1, ::-1]
            center_plot = make_contour(center_plot)

            offset_plot = output['instance_offset'][b, t].detach().cpu().numpy()
            offset_plot[:, semantic_seg[b, t] != 1] = 0
            offset_plot = flow_to_image(offset_plot)[::-1, ::-1]
            offset_plot = make_contour(offset_plot)
        else:
            instance_plot = np.zeros_like(semantic_plot)
            center_plot = np.zeros_like(semantic_plot)
            offset_plot = np.zeros_like(semantic_plot)

        if cfg.INSTANCE_FLOW.ENABLED:
            future_flow_plot = output['instance_flow'][b, t].detach().cpu().numpy()
            future_flow_plot[:, semantic_seg[b, t] != 1] = 0
            future_flow_plot = flow_to_image(future_flow_plot)[::-1, ::-1]
            future_flow_plot = make_contour(future_flow_plot)
        else:
            future_flow_plot = np.zeros_like(semantic_plot)

        if cfg.SEMANTIC_SEG.HDMAP.ENABLED:
            hdmap1 = output['hdmap'][:, :2].argmax(dim=1, keepdims=True)
            hdmap2 = output['hdmap'][:, 2:].argmax(dim=1, keepdims=True)
            hdmap = torch.cat([hdmap1, hdmap2], dim=1)
            planning_plot = plot_planning(hdmap[b], output['selected_traj'][b], cfg)
            planning_plot = make_contour(planning_plot)
        else:
            planning_plot = np.zeros_like(semantic_plot)

        out_t.append(np.concatenate([instance_plot, future_flow_plot,
                                     semantic_plot, center_plot, offset_plot, pedestrian_plot, planning_plot], axis=0))
        out_t = np.concatenate(out_t, axis=1)
        # Shape (C, H, W)
        out_t = out_t.transpose((2, 0, 1))

        video.append(out_t)

    # Shape (B, T, C, H, W)
    video = np.stack(video)[None]
    return video


def convert_figure_numpy(figure):
    """ Convert figure to numpy image """
    figure_np = np.frombuffer(figure.canvas.tostring_rgb(), dtype=np.uint8)
    figure_np = figure_np.reshape(figure.canvas.get_width_height()[::-1] + (3,))
    return figure_np

def plot_planning(hd_map, traj, cfg):
    '''
    hd_map: torch.tensor (2, 200, 200)
    traj: torch.tensor (n_future, 2)
    '''
    if isinstance(hd_map, torch.Tensor):
        hd_map = hd_map.detach().cpu().numpy()
    if isinstance(traj, torch.Tensor):
        traj = traj.detach().cpu().numpy()

    h, w = hd_map.shape[-2:]
    fig = plt.figure(figsize=(w/100, h/100))

    dx, bx, _ = gen_dx_bx(cfg.LIFT.X_BOUND, cfg.LIFT.Y_BOUND, cfg.LIFT.Z_BOUND)
    dx, bx = dx[:2].numpy(), bx[:2].numpy()

    map = np.ones((h, w, 3))
    yx = np.nonzero(hd_map[0])
    c = np.array([1.00, 0.50, 0.31])
    map[yx[0], yx[1], :] = c
    yx = np.nonzero(hd_map[1])
    c = np.array([0.50, 0.24, 0.46])
    map[yx[0], yx[1], :] = c
    plt.imshow(map, alpha=0.2)
    plt.axis('off')

    plt.xlim((w, 0))
    plt.ylim((0, h))
    W = cfg.EGO.WIDTH
    H = cfg.EGO.HEIGHT
    pts = np.array([
        [-H / 2. + 0.5, W / 2.],
        [H / 2. + 0.5, W / 2.],
        [H / 2. + 0.5, -W / 2.],
        [-H / 2. + 0.5, -W / 2.],
    ])
    pts = (pts - bx) / dx
    pts[:, [0, 1]] = pts[:, [1, 0]]
    plt.fill(pts[:, 0], pts[:, 1], '#76b900')

    gt = (traj[:, :2] - bx) / dx
    plt.plot(gt[:, 0], gt[:, 1])

    plt.draw()
    figure_numpy = convert_figure_numpy(fig)
    plt.close()


    return figure_numpy

def generate_instance_colours(instance_map):
    # Most distinct 22 colors (kelly colors from https://stackoverflow.com/questions/470690/how-to-automatically-generate
    # -n-distinct-colors)
    # plus some colours from AD40k
    INSTANCE_COLOURS = np.asarray([
        [0, 0, 0],
        [255, 179, 0],
        [128, 62, 117],
        [255, 104, 0],
        [166, 189, 215],
        [193, 0, 32],
        [206, 162, 98],
        [129, 112, 102],
        [0, 125, 52],
        [246, 118, 142],
        [0, 83, 138],
        [255, 122, 92],
        [83, 55, 122],
        [255, 142, 0],
        [179, 40, 81],
        [244, 200, 0],
        [127, 24, 13],
        [147, 170, 0],
        [89, 51, 21],
        [241, 58, 19],
        [35, 44, 22],
        [112, 224, 255],
        [70, 184, 160],
        [153, 0, 255],
        [71, 255, 0],
        [255, 0, 163],
        [255, 204, 0],
        [0, 255, 235],
        [255, 0, 235],
        [255, 0, 122],
        [255, 245, 0],
        [10, 190, 212],
        [214, 255, 0],
        [0, 204, 255],
        [20, 0, 255],
        [255, 255, 0],
        [0, 153, 255],
        [0, 255, 204],
        [41, 255, 0],
        [173, 0, 255],
        [0, 245, 255],
        [71, 0, 255],
        [0, 255, 184],
        [0, 92, 255],
        [184, 255, 0],
        [255, 214, 0],
        [25, 194, 194],
        [92, 0, 255],
        [220, 220, 220],
        [255, 9, 92],
        [112, 9, 255],
        [8, 255, 214],
        [255, 184, 6],
        [10, 255, 71],
        [255, 41, 10],
        [7, 255, 255],
        [224, 255, 8],
        [102, 8, 255],
        [255, 61, 6],
        [255, 194, 7],
        [0, 255, 20],
        [255, 8, 41],
        [255, 5, 153],
        [6, 51, 255],
        [235, 12, 255],
        [160, 150, 20],
        [0, 163, 255],
        [140, 140, 140],
        [250, 10, 15],
        [20, 255, 0],
    ])

    return {instance_id: INSTANCE_COLOURS[global_instance_id % len(INSTANCE_COLOURS)] for
            instance_id, global_instance_id in instance_map.items()
            }

data/stp3/utils/visualisation.py
llmbev/talk2bev
from PIL import Image
import argparse
from PIL import Image

from minigpt4.common.config import Config
from minigpt4.common.registry import registry
from minigpt4.conversation.conversation import Chat, CONV_VISION

# imports modules for registration
from minigpt4.datasets.builders import *
from minigpt4.models import *
from minigpt4.processors import *
from minigpt4.runners import *
from minigpt4.tasks import *

from llava.conversation import conv_templates

from lavis.models import load_model_and_preprocess

from tqdm import tqdm
import argparse

cam_keys = ['CAM_FRONT_LEFT', 'CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_LEFT', 'CAM_BACK', 'CAM_BACK_RIGHT']

# ========================================
#             InsructBLIP-2 Model Initialization
# ========================================
def init_instructblip2(model_name = "blip2_vicuna_instruct", device="cuda:0"):
    model, vis_processors, _ = load_model_and_preprocess(
        name=model_name,
        model_type="vicuna13b",
        is_eval=True,
        device=device,
    )
    return model, vis_processors
# ========================================

# ========================================
#             BLIP-2 Model Initialization
# ========================================
def init_blip2(model_name = "blip2_vicuna_instruct", device="cuda:0", model_type="vicuna13b"):
    model, vis_processors, _ = load_model_and_preprocess(
        name=model_name,
        model_type=model_type,
        is_eval=True,
        device=device,
    )
    return model, vis_processors
# ========================================


# ========================================
#             MiniGPT4 Initialization
# ========================================
def init_minigp4():
    parser = argparse.ArgumentParser(description="Demo")
    parser.add_argument("--cfg-path", default="eval_configs/minigpt4_eval.yaml", help="path to configuration file.")
    parser.add_argument("--gpu-id", type=int, default=0, help="specify the gpu to load the model.")
    parser.add_argument("--sam-checkpoint", type=str, default="sam_vit_h_4b8939.pth", help="path to sam weights.")
    parser.add_argument('--model_path', type=str, default="/raid/t1/scratch/vikrant.dewangan/LLaVA/ckpt-old/", help='save path for jsons')
    parser.add_argument('--save_path', type=str, default="/raid/t1/scratch/vikrant.dewangan/datas", help='save path for jsons')
    parser.add_argument('--gpu', type=str, default="cuda:0", help='save path for jsons')
    parser.add_argument('--json_name', type=str, default="answer_pred_both.json", help='save path for jsons')
    parser.add_argument('--start', type=int, default=0, help='start index')
    parser.add_argument('--end', type=int, default=100, help='end index')

    parser.add_argument(
        "--options",
        nargs="+",
        help="override some settings in the used config, the key-value pair "
        "in xxx=yyy format will be merged into config file (deprecate), "
        "change to --cfg-options instead.",
    )
    args = parser.parse_args()
    print('Initializing Chat')
    cfg = Config(args)

    model_config = cfg.model_cfg
    model_config.device_8bit = args.gpu_id
    model_cls = registry.get_model_class(model_config.arch)
    model = model_cls.from_config(model_config).to('cuda:{}'.format(args.gpu_id))

    vis_processor_cfg = cfg.datasets_cfg.cc_sbu_align.vis_processor.train
    vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
    chat = Chat(model, vis_processor, device='cuda:{}'.format(args.gpu_id))
    return chat
# ========================================

def reset_conv(model_name = "llava"):
    if 'llama-2' in model_name.lower():
        conv_mode = "llava_llama_2"
    elif "v1" in model_name.lower():
        conv_mode = "llava_v1"
    elif "mpt" in model_name.lower():
        conv_mode = "mpt"
    else:
        conv_mode = "llava_v0"
    print('reset conv')

    conv = conv_templates[conv_mode].copy()
    if "mpt" in model_name.lower():
        roles = ('user', 'assistant')
    else:
        roles = conv.roles
    return conv

def minigpt4_inference(chat, img_cropped, user_message):
    img_list = []
    chat_state = CONV_VISION.copy()  # Reset chat state to default template
    llm_message = chat.upload_img(Image.fromarray(img_cropped), chat_state, img_list)

    print('Upload done')

    chat.ask(user_message, chat_state)
    llm_message = chat.answer(
            conv=chat_state,
            img_list=img_list,
            # num_beams=num_beams,
            num_beams=1,
            # temperature=temperature,
            temperature=0.7,
            max_new_tokens=300,
            max_length=2000
    )[0]
    return llm_message

def instructblip2_inference(model_instructblip, img_cropped, vis_processors, device="cuda:0", user_message="describe the central object in the scene."):
    image = vis_processors["eval"](Image.fromarray(img_cropped)).unsqueeze(0).to(device)

    samples = {
        "image": image,
        "prompt": user_message,
    }

    output_blip = model_instructblip.generate(
        samples,
        length_penalty=float(1),
        repetition_penalty=float(1),
        num_beams=5,
        max_length=256,
        min_length=1,
        top_p=0.2,
        use_nucleus_sampling=False,
    )

    return output_blip[0]
data/utils.py
llmbev/talk2bev
import openai
import os

# chat gpt credentials
ORGANIZATION = os.getenv("ORGANIZATION")
API_KEY = os.getenv("API_KEY")

# question types to ask
QUESTION_TYPES = ["SPATIAL_REASONING", "INSTANCE_ATTRIBUTE", "INSTANCE_COUNTING", "VISUAL_REASONING"]

# to generate per type
NUM_PER_TYPE = 10

JSONS_LIMIT = 1000

# example conversation to prompt GPT
EXAMPLE_MCQ = "Here is an example\
    What is our planet's name?\
        A. Earth\
        B. Sun\
        C. Mars\
        D. Jupyter\
        \n\n \
        Answer: A. Earth"

# scene data format
SCENE_FORMAT = {
	"bev_centroid": "2D coordinates of the center of the object",
	"bev_area": "2D area of the object",
  "brief_label": "A brief caption describing the object.",
  "bg_description": "A description of the background around the object.",
  "weather": "Weather description of the scene"
}

# question context
QUESTION_CONTEXT = {
  "SPATIAL_REASONING": "The questions should be about spatial relations between two objects. The question should be mainly based on the coordinates of the two objects. To answer the question, one should find the two mentioned objects, and find their relative spatial relation to answer the question.",
  "INSTANCE_ATTRIBUTE": "The questions should be about the attribute of a certain object, such as its color, shape or fine-grained type. Do not use the object ID in the question.",
  "INSTANCE_COUNTING": "The questions should involve the number of appearance of a certain object. Start with 'How many ....'. The choices of the question should be numbers. To answer the question, one should find and count all of the mentioned objects in the image. Make sure the options are far apart.",
  "VISUAL_REASONING": "Create complex questions beyond describing the scene. The questions should aid the user who is driving the ego-vehicle. To answer such question, one should first understanding the visual content, then based on the background knowledge or reasoning, either explain why the things are happening that way, or provide guides and help to user's request. Make the question challenging by not including the visual content details in the question so that the user needs to reason about that first. Do not use the object ID in the question. Do not mention the coordinates."
}

# question generation format
QUESTION_FORMAT = {
   "question_number": "Please give the question number of this question.",
   "question": "Please give the question.",
   "options": "this should be a list of options",
   "correct_option": "this should be the correct options",
}

# answer format - spatial operators
ANSWER_FORMAT = {
   "inferred_query": "your interpretation of the user query in a succinct form",
   "relevant_objects": "python list of all relevant object ids for the user query",
   "query_achievable": "whether or not the user-specified query is achievable using the objects and descriptions provided in the scene.",
   "spatial_reasoning_functions": "If the query needs calling one or more spatial reasoning functions, this field contains a list of function calls that conform to the API above. Else, this field contains an empty list.",
   "explanation": "A brief explanation of what the relevant objects, and how it addresses the task."
}

# spatial operators
SPATIAL_OPERATORS_LIST = [
	"filter_front(list_of_objects, object_id) : Within the the list of objects, it returns the list of objects to front of the object with id as object_id",
	"filter_left(list_of_objects, object_id) : Within the the list of objects, it returns the list of objects to left of the object with id as object_id",
	"filter_right(list_of_objects, object_id) : Within the the list of objects, it returns the list of objects to right of the object with id as object_id",
	"filter_rear(list_of_objects, object_id) : Within the the list of objects, it returns the list of objects to rear of the object with id as object_id",
	"find_distance(list_of_objects, object_id_1, object_id_2) : Within the the list of objects, it returns the distance within 2 objects object_id_1 and object_id_2",
	"find_objects_within_distance(list_of_objects, object_id, distance) : Within the the list of objects, it returns the list of objects within distance to the object with id object_id",
    "get_k_closest_objects(list_of_objects, object_id, k) : Within the the list of objects, it returns the list of k closest objects to the object with id object_id",
    "get_k_farthest_objects(list_of_objects, object_id, k) : Within the the list of objects, it gets the k farthest objects to the object with id object_id",
    "filter_objects_with_tag(list_of_objects, object_id, tagname, d) : Within the the list of objects, it finds the objects which have tag as tagname and are within d distance to the object with id object_id",
    "filter_color(list_of_objects, object_id, colorname, d): Within the the list of objects, itfinds the objects which have color as colorname and are within d distance to the object with id object_id",
    "filter_size(list_of_objects, object_id, distance, min_size, max_size): Within the the list of objects, itfinds the objects which have size between min_size and max_size and are within d distance to the object with id object_id"
]

def setup_question_generation(question_type):
    message = f"You will be given, as input a 2D road scene in Bird's Eye View, as a list. The ego-vehicle is at (100, 100) facing along the positive Y-axis. Each entry in the list describes one object in the scene, with the following five fields: \
    \n{str(SCENE_FORMAT)}\n\
    Once you have parsed the JSON and are ready to generate question about the scene, Create {NUM_PER_TYPE} distinct multi-choice question about the scene, and provide the choices and answer. Each question should have 4 options, out of which only 1 should be correct. Do not use the object ID in the question. Do not mention the coordinates."# + "You have to return a list of JSONs each containing an MCQ question."
    message += QUESTION_CONTEXT[question_type]
    return {
        "role": "system",
        "content": message  + "Please provide answer as well. \n" + str(EXAMPLE_MCQ) + "\n . NOTE: DO NOT ask simple questions like 'What is the central object in the scene?', or What is the color of the car?." + "\n\n. "
        # "content": message  + "Please provide answer as well. \n" + "\n . NOTE: DO NOT ask simple questions like 'What is the central object in the scene?'." + "\n\n. For each question, your JSON should contain - " + str(QUESTION_FORMAT) + "\n. Your final output should be a list of JSONs, and each string should be in double quotes. Please ensure your list can be parsed in a python program. Your string should be in double quotes. Please ensure this."
    }

def add_conversation_context(conversation_type):
    if conversation_type == "MCQ":
        return  "The user will then begin to ask Multiple Choice Questions, and the task is to answer various user queries about the scene. For each question, answer with just one correct option. You should only output the correct answer. Do not give any explanation"

    if conversation_type == "SPATIAL":
        return  f"The user will then begin to ask questions, and the task is to answer various user queries about the scene. These questions will involve spatial reasoning. To assist with such queries, we have the following available functions:.\
                    {str(SPATIAL_OPERATORS_LIST)}\
                    For each user question, respond with a JSON dictionary with the following fields:\
                    {str(ANSWER_FORMAT)}\
                    The Object ID of ego-vehicle is 0. Only output this JSON. Do not output any explanation."

def setup_conversation(conversation_type="MCQ"):
    message = f"You will be given, as input a 2D road scene in Bird's Eye View, as a list. The ego-vehicle is at (0,0) facing along the positive Y-axis. Each entry in the list describes one object in the scene, with the following five fields: \
                    \n\n {str(SCENE_FORMAT)}  \n\n  \
                Once you have parsed the JSON and are ready to answer questions about the scene, please wait for the user to input JSON."
    message += add_conversation_context(conversation_type)
    return {
        "role": "system",
        "content": message
    }

class ChatGPTInteface:
    def __init__(self, API_KEY, organization, model_name="gpt-4") -> None:
        openai.api_key = API_KEY
        openai.organization = organization
        self.model_name = "gpt-4"

    def generate_question(self, data, question_type="DEFAULT"):
        system_message = setup_question_generation(question_type)
        user_message = {
            "role": "user",
            "content": str(data)
        }
        response = openai.ChatCompletion.create(
            model=self.model_name,
            messages=[
                system_message, # prompt template
                user_message # data
            ],
            temperature=0,
            max_tokens=1024
        )
        return response

    def generate_conversation(self, data, question, conversation_type):
        system_message = setup_conversation(conversation_type)
        user_message = {
            "role": "user",
            "content": str(data) + ". The question is as follows \n" + str(question)
        }
        response = openai.ChatCompletion.create(
            model=self.model_name,
            messages=[
                system_message,
                user_message
            ],
            temperature=0,
            max_tokens=1024
        )
        return response

evaluation/chatgptinterface.py
llmbev/talk2bev
from chatgptinterface import ChatGPTInteface
import json
import numpy as np
import pandas as pd
import os

def calculate_jaccard_similarity(str1, str2):
    set1 = set(str1.lower().split())
    set2 = set(str2.lower().split())
    
    intersection = len(set1.intersection(set2))
    union = len(set1) + len(set2) - intersection
    
    jaccard_similarity = intersection / union
    return jaccard_similarity

def are_strings_similar(str1, str2, threshold=0.75):
    similarity = calculate_jaccard_similarity(str1, str2)
    return similarity >= threshold

ORGANIZATION = os.getenv("ORGANIZATION")
API_KEY = os.getenv("API_KEY")
save_path = "datas"

chatGPTInteface = ChatGPTInteface(API_KEY=API_KEY, organization=ORGANIZATION)
interface = ChatGPTInteface(API_KEY, ORGANIZATION, model_name="gpt-3.5-turbo")
json_list = sorted(os.listdir(save_path))

counts = []
logs = []
for json_item in json_list:
    gt_json = json.load(open(os.path.join(save_path, json_item, "scene", "answer_gt.json")))
    pred_json = json.load(open(os.path.join(save_path, json_item, "scene", "answer_pred_both.json")))
    for i in gt_json:
        del i["matched_coords"]
        del i["annotation"]

    for i in pred_json:
        del i["matched_coords"]
        del i["annotation"]

    question_types = ["SPATIAL_REASONING", "INSTANCE_ATTRIBUTE", "INSTANCE_COUNTING", "VISUAL_REASONING"]
    for qntype in question_types:
        try:
            response = interface.generate_question(gt_json, question_type=qntype)
            question_with_answer = response["choices"][0]["message"]["content"]
            print("response \n", question_with_answer)
            print("\n\n")
            answer_ind = response["choices"][0]["message"]["content"].lower().find("answer")
            generated_question = question_with_answer[:answer_ind]
            correct_answer = question_with_answer[answer_ind+8:]
            print("separated question \n", generated_question)
            print("\n\n")
            print("separated answer \n", correct_answer)
            response = interface.generate_conversation(pred_json, generated_question, conversation_type="MCQ")
            chatgpt_answer = response["choices"][0]["message"]["content"]
            print("\n\n\n")
            print("selected answer \n", chatgpt_answer)
            count = are_strings_similar(chatgpt_answer, correct_answer)
            print(count)
            counts.append(count)
            logs.append([
                generated_question,
                correct_answer,
                chatgpt_answer
            ])
        except Exception as e:
            print(e)
            pass

counts = np.array(counts)
print("MCQ Accuracy: \t", (counts.sum())/len(counts))
df = pd.DataFrame( data=logs, columns=["Question", "Correct Answer", "ChatGPT Answer"])
df.to_csv("logs.csv")

evaluation/eval_mcq.py
llmbev/talk2bev
import numpy as np
import re
import json
import decimal
import os

from spatial_operators import *

s1 = 0
s2 = 0
dir1 = 'spatial_gt/' # path to GT jsons
dir2 = 'spatial_pred/' # path to answer jsons

files = os.listdir(dir1)

for fll in files:
    key = fll[:-5]

    json_gt = json.load(open(dir1[8:]+key+'.json'))
    json_pd = json.load(open(dir2[8:]+key+'.json'))

    # Question 1
    d1 = get_object_dist(json_gt, json_gt[0]['token'][0])
    d2 = get_object_dist(json_pd, json_gt[0]['token'][0])
    s1 += dist_score(d1, d2)

    # Question 2
    d1 = get_object_dist(json_gt, json_gt[1]['token'][0])
    d2 = get_object_dist(json_pd, json_gt[1]['token'][0])
    s1 += dist_score(d1, d2)

    # Question 3
    d1 = distance_between_objects(json_gt, json_gt[0]['token'][0], json_gt[1]['token'][0])
    d2 = distance_between_objects(json_pd, json_gt[0]['token'][0], json_gt[1]['token'][0])
    s1 += dist_score(d1, d2)

    # Question 4
    jgt = get_k_closest_jsons(json_gt, 2)
    jpd = get_k_closest_jsons(json_pd, 2)
    d1 = distance_between_objects(json_gt, jgt[0]['token'][0], jgt[1]['token'][0])
    d2 = distance_between_objects(json_pd, jpd[0]['token'][0], jpd[1]['token'][0])
    s1 += dist_score(d1, d2)

    dd = 20

    # Question 5
    jgt = distance_filtering(json_gt, dd)
    jpd = distance_filtering(json_pd, dd)
    s2 += iou(jgt, jpd)

    # Question 6
    jgt = get_k_closest_jsons(json_gt, 3)
    jpd = get_k_closest_jsons(json_pd, 3)
    s2 += iou(jgt, jpd)

    # Question 7
    jgt = filter_front(json_gt)
    jgt = distance_filtering(jgt, dd)
    jpd = filter_front(json_pd)
    jpd = distance_filtering(jpd, dd)
    s2 += iou(jgt, jpd)

    # Question 8
    jgt = filter_rear(json_gt)
    jgt = get_k_closest_jsons(jgt, 3)
    jpd = filter_rear(json_pd)
    jpd = get_k_closest_jsons(jpd, 3)
    s2 += iou(jgt, jpd)

s1 = s1/4/7
s1 = decimal.Decimal(s1)
s1 = s1.quantize(decimal.Decimal('0.00'))

s2 = s2/4/7
s2 = decimal.Decimal(s2)
s2 = s2.quantize(decimal.Decimal('0.00'))

print("IoU:", s2)
print("Distance Score:", s1)

evaluation/eval_spops.py
llmbev/talk2bev
import math
import heapq
import numpy as np

main_coverage_angle = 90
diagonal_coverage_angle = 90

def calculate_angle(obj_x, obj_y):
    dx1, dy1 = 100, 0
    dx2, dy2 = obj_x - 100, obj_y - 100
    dot_product = dx1 * dx2 + dy1 * dy2
    magnitude_v1 = math.sqrt(dx1**2 + dy1**2)
    magnitude_v2 = math.sqrt(dx2**2 + dy2**2)
    cosine_angle = dot_product / (magnitude_v1 * magnitude_v2)
    angle_radians = math.acos(cosine_angle)
    angle_degrees = math.degrees(angle_radians)
    cross_product = dx1 * dy2 - dx2 * dy1
    if cross_product < 0:
        angle_degrees = -angle_degrees
    return angle_degrees
def filter_rear(json_list):
    rear_json = []
    for gg in json_list:
        obj_x, obj_y = gg['bev_centroid']
        angle = calculate_angle(obj_x, obj_y)
        if angle >= (180 - main_coverage_angle/2) or angle <= (-180 + main_coverage_angle/2):
            rear_json.append(gg)
    return rear_json
def filter_left(json_list):
    left_json = []
    for gg in json_list:
        obj_x, obj_y = gg['bev_centroid']
        angle = calculate_angle(obj_x, obj_y)
        if angle >= (90 - main_coverage_angle/2) and angle <= (90 + main_coverage_angle/2):
            left_json.append(gg)
    return left_json
def filter_front(json_list):
    front_json = []
    for gg in json_list:
        obj_x, obj_y = gg['bev_centroid']
        angle = calculate_angle(obj_x, obj_y)
        if angle >= -main_coverage_angle/2 and angle <= main_coverage_angle/2:
            front_json.append(gg)
    return front_json
def filter_right(json_list):
    right_json = []
    for gg in json_list:
        obj_x, obj_y = gg['bev_centroid']
        angle = calculate_angle(obj_x, obj_y)
        if angle >= (-90 - main_coverage_angle/2) and angle <= (-90 + main_coverage_angle/2):
            right_json.append(gg)
    return right_json
def distance_filtering(json_list, dist):
    dist_px = 2*dist
    filtered_json = []
    for gg in json_list:
        obj_x, obj_y = gg['bev_centroid']
        if math.sqrt((obj_x-100)**2 + (obj_y-100)**2) <= dist_px:
            filtered_json.append(gg)
    return filtered_json
def distance_between_objects(json_list, id1, id2):
    x1, y1 = -1, -1
    x2, y2 = -1, -1
    for gg in json_list:
        if gg['token'][0]==id1:
            x1, y1 = gg['bev_centroid']
        if gg['token'][0]==id2:
            x2, y2 = gg['bev_centroid']
    return math.sqrt((x1-x2)**2 + (y1-y2)**2)
def get_object_dist(json_list, object_id):
    x, y = -1, -1
    for gg in json_list:
        if gg['token'][0]==object_id:
            x, y = gg['bev_centroid']
            # print(object_id, x, y)
    return math.sqrt((x-100)**2 + (y-100)**2)
def get_k_closest_jsons(json_list, k):    
    filtered_json = []
    # print("---")
    for gg in json_list:
        x, y = gg['bev_centroid']
        distance = math.sqrt((x-100)**2 + (y-100)**2)
        # print(distance, gg['token'][0])
        heapq.heappush(filtered_json, (-distance, gg))
        if len(filtered_json) > k:
            heapq.heappop(filtered_json)
    return [json_data for (_, json_data) in filtered_json]
def get_k_farthest_jsons(json_list, k):    
    filtered_json = []
    for gg in json_list:
        x, y = gg['bev_centroid']
        distance = math.sqrt((x-100)**2 + (y-100)**2)
        heapq.heappush(filtered_json, (distance, gg))
        if len(filtered_json) > k:
            heapq.heappop(filtered_json)
    return [json_data for (_, json_data) in filtered_json]


##############################
######### Evaluation #########
##############################


def dist_score(d1, d2):
    max_difference = 200 * np.sqrt(2)
    difference = abs(d1 - d2)
    if difference >= max_difference:
        score = 0
    else:
        score = 1 - difference / max_difference
    return score

def iou(query_result, ground_truth):
    query_result_ids = [json_str["token"][0] for json_str in query_result]
    ground_truth_ids = [json_str["object_id"] for json_str in ground_truth]
    if len(query_result_ids) == 0 and len(ground_truth_ids) == 0:
        return 1
    else:
        intersection = len(set(query_result_ids) & set(ground_truth_ids))
        union = len(set(query_result_ids) | set(ground_truth_ids))
        iou = intersection / union
        return iou

def evaluate_spatial_query(query_result, ground_truth):
    
    query_result_ids = [json_str["object_id"] for json_str in query_result]
    ground_truth_ids = [json_str["object_id"] for json_str in ground_truth]

    if len(query_result_ids) == 0 and len(ground_truth_ids) == 0:
        return 1
    
    else:
        intersection = len(set(query_result_ids) & set(ground_truth_ids))
        union = len(set(query_result_ids) | set(ground_truth_ids))
        iou = intersection / union
        return iou

# Distance should be in meters
def calculate_score(d1, d2):
    max_difference = 100 * np.sqrt(2)
    difference = abs(d1 - d2)
    if difference >= max_difference:
        score = 0
    else:
        score = 1 - difference / max_difference
    return score

main_coverage_angle = 90

def calculate_angle(obj_x, obj_y, org_x, org_y):
    dx1, dy1 = 200 - org_x, 0
    dx2, dy2 = obj_x - org_x, obj_y - org_y
    dot_product = dx1 * dx2 + dy1 * dy2
    magnitude_v1 = math.sqrt(dx1**2 + dy1**2)
    magnitude_v2 = math.sqrt(dx2**2 + dy2**2)
    cosine_angle = dot_product / (magnitude_v1 * magnitude_v2)
    angle_radians = math.acos(cosine_angle)
    angle_degrees = math.degrees(angle_radians)
    cross_product = dx1 * dy2 - dx2 * dy1

    if cross_product < 0:
        angle_degrees = -angle_degrees

    return angle_degrees

def filter_rear(json_list, object_id):
    rear_json = []
    org_x, org_y = 100, 100
    
    if object_id != 0:      #passing object_id as 0 for ego vehicle
        for obj in json_list:
            if obj["object_id"] == object_id:
                org_x, org_y = obj['bev_centroid']

    for obj in json_list: 
        if obj['object_id'] != object_id:
            obj_x, obj_y = obj['bev_centroid']
            angle = calculate_angle(obj_x, obj_y, org_x, org_y)
            if angle >= (180 - main_coverage_angle/2) or angle <= (-180 + main_coverage_angle/2):
                rear_json.append(obj)
    return rear_json

def filter_left(json_list, object_id):
    left_json = []
    org_x, org_y = 100, 100
    
    if object_id != 0:      #passing object_id as 0 for ego vehicle
        for obj in json_list:
            if obj["object_id"] == object_id:
                org_x, org_y = obj['bev_centroid']

    for obj in json_list:
        if obj['object_id'] != object_id:
            obj_x, obj_y = obj['bev_centroid']
            angle = calculate_angle(obj_x, obj_y, org_x, org_y)
            if angle >= (90 - main_coverage_angle/2) and angle <= (90 + main_coverage_angle/2):
                left_json.append(obj)
    return left_json

def filter_front(json_list, object_id):
    front_json = []
    org_x, org_y = 100, 100
    
    if object_id != 0:      #passing object_id as 0 for ego vehicle
        for obj in json_list:
            if obj["object_id"] == object_id:
                org_x, org_y = obj['bev_centroid']

    for obj in json_list:
        if obj['object_id'] != object_id:
            obj_x, obj_y = obj['bev_centroid']
            angle = calculate_angle(obj_x, obj_y, org_x, org_y)
            if angle >= -main_coverage_angle/2 and angle <= main_coverage_angle/2:
                front_json.append(obj)
    return front_json

def filter_right(json_list, object_id):
    right_json = []
    org_x, org_y = 100, 100
    
    if object_id != 0:      #passing object_id as 0 for ego vehicle
        for obj in json_list:
            if obj["object_id"] == object_id:
                org_x, org_y = obj['bev_centroid']

    for obj in json_list:
        if obj['object_id'] != object_id:
            obj_x, obj_y = obj['bev_centroid']
            angle = calculate_angle(obj_x, obj_y, org_x, org_y)
            if angle >= (-90 - main_coverage_angle/2) and angle <= (-90 + main_coverage_angle/2):
                right_json.append(obj)
    return right_json

#assuming distance in meters
def find_distance(json_list, object_id1, object_id2):
    if object_id1 == 0:
        x1, y1 = 100, 100
    if object_id2 == 0:
        x2, y2 = 100, 100

    for obj in json_list:
        if obj['object_id'] == object_id1:
            x1, y1 = obj['bev_centroid']
        if obj['object_id'] == object_id2:
            x2, y2 = obj['bev_centroid']
    
    dist = (math.sqrt((x1-x2)**2 + (y1-y2)**2)) / 2
    return dist

def find_objects_within_distance(json_list, object_id, distance):
    nearby_json = []

    for obj in json_list:
        if obj['object_id'] != object_id:
            dist = find_distance(json_list, object_id, obj['object_id'])
            if dist <= distance:
                nearby_json.append(obj)
    return nearby_json

def get_k_closest_objects(json_list, object_id, k):
    closest_json = []
    closest_dist = []
    objects = []

    for obj in json_list:
        if obj['object_id'] != object_id:
            dist = find_distance(json_list, object_id, obj['object_id'])
            objects.append((obj, dist))
    closest_objects = sorted(objects, key=lambda obj: obj[1])

    for i in range(k):
        closest_json.append(closest_objects[i][0])
        closest_dist.append(closest_objects[i][1])

    return closest_json, closest_dist

def get_k_farthest_objects(json_list, object_id, k):
    farthest_json = []
    farthest_dist = []
    objects = []

    for obj in json_list:
        if obj['object_id'] != object_id:
            dist = find_distance(json_list, object_id, obj['object_id'])
            objects.append((obj, dist))
    farthest_objects = sorted(objects, key=lambda obj: obj[1], reverse=True)

    for i in range(k):
        farthest_json.append(farthest_objects[i][0])
        farthest_dist.append(farthest_objects[i][1])

    return farthest_json, farthest_dist

def filter_objects_with_tag(json_list, object_id, tagname, distance=1000):
    type_json = []

    nearby_json = find_objects_within_distance(json_list, object_id, distance)
    for obj in nearby_json:
        if obj['tag'] == tagname:
            type_json.append(obj)
    return type_json

def filter_color(json_list, object_id, color, distance=1000):
    type_json = []

    nearby_json = find_objects_within_distance(json_list, object_id, distance)
    for obj in nearby_json:
        if obj['color'] == color:
            type_json.append(obj)
    return type_json

def filter_size(json_list, object_id, distance, min_size, max_size):
    size_json = []

    nearby_json = find_objects_within_distance(json_list, object_id, distance)
    for obj in nearby_json:
        if obj['size'] >= min_size and obj['size'] <= max_size:
            size_json.append(obj)
    return size_json

#finds direction of object_id2 with respect to object_id1
def find_direction_of_object(json_list, object_id1, object_id2):
    org_x, org_y = 100, 100

    if object_id1 != 0:      #passing object_id as 0 for ego vehicle
        for obj in json_list:
            if obj["object_id"] == object_id1:
                org_x, org_y = obj['bev_centroid']

    for obj in json_list: 
        if obj['object_id'] == object_id2:
            obj_x, obj_y = obj['bev_centroid']
            angle = calculate_angle(obj_x, obj_y, org_x, org_y)
            if angle >= (180 - main_coverage_angle/2) or angle <= (-180 + main_coverage_angle/2):
                return 'REAR'
            elif angle >= (90 - main_coverage_angle/2) and angle <= (90 + main_coverage_angle/2):
                return 'LEFT'
            elif angle >= -main_coverage_angle/2 and angle <= main_coverage_angle/2:
                return 'FRONT'
            else:
                return 'RIGHT'



# frontFilter(objs)
# leftFilter(objs)
# rightFilter(objs)
# rearFilter(objs)
# distFilter(objs, X)
# objDistance(objs, id)
# kClosest(objs, k)
# kFarthest(objs, k)
# findDist(objs, id1, id2)
# objsInDist(objs, id, dist)
# kClosestToObj(objs, id, k)
# kFarthestToObj(objs, id, k)
evaluation/spatial_operators.py
llmbev/talk2bev