import random
from dataclasses import dataclass
from typing import Dict, List, Tuple
import numpy as np
import scipy.spatial.transform
import yaml
from kitcar_utils.camera.camera_specs import CameraSpecs
from kitcar_utils.geometry import Vector
[docs]@dataclass
class SampleGenerator:
camera_specs: CameraSpecs
"""Camera specification including calibration."""
ts_sizes: Dict[str, Tuple[float, float, float]]
"""Sizes of all traffic signs."""
min_distance: float
"""Minimal distance to generated sign."""
max_distance: float
"""Maximal distance to generated sign."""
def __init__(self, camera_specs_path, ts_sizes_path, min_distance, max_distance):
self.camera_specs = CameraSpecs.from_yaml(camera_specs_path)
with open(ts_sizes_path) as file:
self.ts_sizes = yaml.load(file, Loader=yaml.SafeLoader)
self.min_distance = min_distance
self.max_distance = max_distance
[docs] def sample_world_vec(self) -> Vector:
"""Sample vector to point in the fov."""
angle = 1 / 2 * self.camera_specs.horizontal_fov * random.uniform(-1, 1)
distance = random.uniform(self.min_distance, self.max_distance)
return Vector(r=distance, phi=angle)
[docs] def _H(self, height: float) -> np.matrix:
"""Calculate transform from vehicle coords to img at a given height."""
tt = np.array([0, 0, -height]).reshape((3, 1))
t = 0.001 * self.camera_specs.t - self.camera_specs.R @ tt
M = np.zeros((3, 3))
M[:, :2] = self.camera_specs.R[:, :2]
M[:, 2:3] = t
H = self.camera_specs.A @ M
return H
[docs] def vehicle_point_to_img(self, vec: Vector) -> np.ndarray:
"""Transform vector in vehicle coords to pixel."""
hom = self._H(vec.z) @ np.array([vec.x, vec.y, 1]).reshape((3, 1))
return np.array(hom[:, 0] / hom[2]).reshape(3)
[docs] def sign_coords(self, sign_name: str, angle_x: float, angle_z: float) -> List[Vector]:
"""Calculate rotated sign coordinates.
Args:
sign_name: Name of the sign.
angle_x: Rotation around x axis.
angle_z: Rotation around z axis.
"""
sign_width, sign_height, foot_height = self.ts_sizes[
sign_name if sign_name in self.ts_sizes else "default"
]
rot_mat = scipy.spatial.transform.Rotation.from_euler(
"xz", [angle_x, angle_z], degrees=True
)
tspoints = [
[0.0, +sign_width / 2.0, foot_height + sign_height],
[0.0, -sign_width / 2.0, foot_height + sign_height],
[0.0, +sign_width / 2.0, foot_height],
[0.0, -sign_width / 2.0, foot_height],
]
tspoints = [rot_mat.apply(p) for p in tspoints]
return [Vector(p) for p in tspoints]