trap/trap/utils.py
2025-05-15 18:53:43 +02:00

201 lines
No EOL
7.6 KiB
Python

# lerp & inverse lerp from https://gist.github.com/laundmo/b224b1f4c8ef6ca5fe47e132c8deab56
import linecache
import math
import os
from pathlib import Path
import tracemalloc
from typing import Iterable
import cv2
import numpy as np
from trajectron.environment.map import GeometricMap
def lerp(a: float, b: float, t: float) -> float:
"""Linear interpolate on the scale given by a to b, using t as the point on that scale.
Examples
--------
50 == lerp(0, 100, 0.5)
4.2 == lerp(1, 5, 0.8)
"""
return (1 - t) * a + t * b
def inv_lerp(a: float, b: float, v: float) -> float:
"""Inverse Linar Interpolation, get the fraction between a and b on which v resides.
Examples
--------
0.5 == inv_lerp(0, 100, 50)
0.8 == inv_lerp(1, 5, 4.2)
"""
return (v - a) / (b - a)
def exponentialDecayRounded(a, b, decay, dt, abs_tolerance):
"""Exponential decay as alternative to Lerp
Introduced by Freya Holmér: https://www.youtube.com/watch?v=LSNQuFEDOyQ
"""
c = b + (a-b) * math.exp(-decay * dt)
if abs(b-c) < abs_tolerance:
return b
return c
def exponentialDecay(a, b, decay, dt):
"""Exponential decay as alternative to Lerp
Introduced by Freya Holmér: https://www.youtube.com/watch?v=LSNQuFEDOyQ
"""
return b + (a-b) * math.exp(-decay * dt)
def relativePointToPolar(origin, point) -> tuple[float, float]:
x, y = point[0] - origin[0], point[1] - origin[1]
return np.sqrt(x**2 + y**2), np.arctan2(y, x)
def relativePolarToPoint(origin, r, angle) -> tuple[float, float]:
return r * np.cos(angle) + origin[0], r * np.sin(angle) + origin[1]
# def line_intersection(line1, line2):
# xdiff = (line1[0][0] - line1[1][0], line2[0][0] - line2[1][0])
# ydiff = (line1[0][1] - line1[1][1], line2[0][1] - line2[1][1])
# def det(a, b):
# return a[0] * b[1] - a[1] * b[0]
# div = det(xdiff, ydiff)
# if div == 0:
# return None
# d = (det(*line1), det(*line2))
# x = det(d, xdiff) / div
# y = det(d, ydiff) / div
# return x, y
# def polyline_intersection(poly1, poly2):
# for i, p1_first_point in enumerate(poly1[:-1]):
# p1_second_point = poly1[i + 1]
# for j, p2_first_point in enumerate(poly2[:-1]):
# p2_second_point = poly2[j + 1]
# intersection = line_intersection((p1_first_point, p1_second_point), (p2_first_point, p2_second_point))
# if intersection:
# return intersection # returns x,y
# return None
def get_bins(bin_size: float):
return [[bin_size, 0], [bin_size, bin_size], [0, bin_size], [-bin_size, bin_size], [-bin_size, 0], [-bin_size, -bin_size], [0, -bin_size], [bin_size, -bin_size]]
def convert_world_space_to_img_space(H: cv2.Mat, scale=100):
"""Transform the given matrix so that it immediately converts
the points to img space"""
new_H = H.copy()
new_H[:2] = H[:2] * scale
return new_H
def convert_world_points_to_img_points(points: Iterable, scale=100):
"""Transform the given matrix so that it immediately converts
the points to img space"""
if isinstance(points, np.ndarray):
return np.array(points) * scale
return [[p[0]*scale, p[1]*scale] for p in points]
def display_top(snapshot: tracemalloc.Snapshot, key_type='lineno', limit=5):
snapshot = snapshot.filter_traces((
tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
tracemalloc.Filter(False, "<unknown>"),
))
top_stats = snapshot.statistics(key_type)
print("Top %s lines" % limit)
for index, stat in enumerate(top_stats[:limit], 1):
frame = stat.traceback[0]
# replace "/path/to/module/file.py" with "module/file.py"
filename = os.sep.join(frame.filename.split(os.sep)[-2:])
print("#%s: %s:%s: %.1f KiB"
% (index, filename, frame.lineno, stat.size / 1024))
line = linecache.getline(frame.filename, frame.lineno).strip()
if line:
print(' %s' % line)
other = top_stats[limit:]
if other:
size = sum(stat.size for stat in other)
print("%s other: %.1f KiB" % (len(other), size / 1024))
total = sum(stat.size for stat in top_stats)
print("Total allocated size: %.1f KiB" % (total / 1024))
class ImageMap(GeometricMap): # TODO Implement for image maps -> watch flipped coordinate system
def __init__(self, img: cv2.Mat, H_world_to_map: cv2.Mat, description=None):
# homography_matrix = np.loadtxt('H.txt')
# homography_matrix = H_img_to_world.copy()
# homography_matrix /= homography_matrix[2, 2] # normalise? https://github.com/StanfordASL/Trajectron-plus-plus/issues/14#issuecomment-637880857
# homography_matrix = np.linalg.inv(homography_matrix)
homography_matrix = H_world_to_map
# RGB png image has 3 layers
img = img.astype(np.uint8)
# img = cv2.resize(img, (img.shape[1]//10, img.shape[0]//10))
img_reverse = img[::-1,:,:] # origin to bottom left, instead of top-left
layers = np.transpose(img, (2, 1, 0)) # array order: layers, x, y
layers = layers.copy() # copy to apply negative stride
# layers =
#scale 255
#alternatively: morph image to world space with a scale, as in trajectron/experiments/nuscenes/process_data.py
super().__init__(layers, homography_matrix, description)
def to_map_points(self, scene_pts):
org_shape = None
if len(scene_pts.shape) > 2:
org_shape = scene_pts.shape
scene_pts = scene_pts.reshape((-1, 2))
N, dims = scene_pts.shape
points_with_one = np.ones((dims + 1, N))
points_with_one[:dims] = scene_pts.T
# map_points = np.fliplr((self.homography @ points_with_one).T[..., :dims]).astype(np.uint32)
# map_points = np.flipud((self.homography @ points_with_one).T[..., :dims]).astype(np.uint32)
map_points = (self.homography @ points_with_one).T[..., :dims].astype(np.uint32)
if org_shape is not None:
map_points = map_points.reshape(org_shape)
# print(scene_pts,'->', map_points)
# exit()
return map_points
# nuscener process_data.py
# type_map = dict()
# canvas_size = (np.round(3 * y_size).astype(int), np.round(3 * x_size).astype(int))
# homography = np.array([[3., 0., 0.], [0., 3., 0.], [0., 0., 3.]])
# layer_names = ['lane', 'road_segment', 'drivable_area', 'road_divider', 'lane_divider', 'stop_line',
# 'ped_crossing', 'stop_line', 'ped_crossing', 'walkway']
# map_mask = (nusc_map.get_map_mask(patch_box, patch_angle, layer_names, canvas_size) * 255.0).astype(
# np.uint8)
# map_mask = np.swapaxes(map_mask, 1, 2) # x axis comes first
# # PEDESTRIANS
# map_mask_pedestrian = np.stack((map_mask[9], map_mask[8], np.max(map_mask[:3], axis=0)), axis=0)
#
# type_map['PEDESTRIAN'] = GeometricMap(data=map_mask_pedestrian, homography=homography, description=', '.join(layer_names))
# Notes: map_mask is a list of masks
# map_mask = []
# _line_geom_to_mask
# def mask_for_lines(...):
# map_mask = np.zeros(canvas_size, np.uint8)
# if layer_name is 'traffic_light':
# return None
# for line in layer_geom:
# new_line = line.intersection(patch)
# if not new_line.is_empty:
# 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))
# map_mask = self.mask_for_lines(new_line, map_mask)