Spaces:
Running on Zero
Running on Zero
| """RGB / depth / instance-segmentation cameras on one pose, and decoding their frames.""" | |
| from __future__ import annotations | |
| import math | |
| import queue | |
| from dataclasses import dataclass | |
| import numpy as np | |
| from .config import HFOV, RGB_H, RGB_W, SENSOR_H, SENSOR_W, TAG2CWM, TAG_PEDESTRIAN, TAGS_SOLID | |
| RGB, DEPTH, SEG = "sensor.camera.rgb", "sensor.camera.depth", "sensor.camera.instance_segmentation" | |
| class SensorRig: | |
| def __init__(self, world, bp, initial_tf): | |
| self.sensors, self.queues = [], {} | |
| self.rgb = self._add(world, bp, initial_tf, RGB, RGB_W, RGB_H, motion_blur_intensity=0.0) | |
| self._add(world, bp, initial_tf, DEPTH, SENSOR_W, SENSOR_H) | |
| self._add(world, bp, initial_tf, SEG, SENSOR_W, SENSOR_H) | |
| def _add(self, world, bp, tf, kind, w, h, **attrs): | |
| b = bp.find(kind) | |
| b.set_attribute("image_size_x", str(w)) | |
| b.set_attribute("image_size_y", str(h)) | |
| b.set_attribute("fov", f"{HFOV:.4f}") | |
| for k, v in attrs.items(): | |
| if b.has_attribute(k): | |
| b.set_attribute(k, str(v)) | |
| s = world.spawn_actor(b, tf) | |
| q = queue.Queue() | |
| s.listen(q.put) | |
| self.sensors.append(s) | |
| self.queues[kind] = q | |
| return s | |
| def place(self, tf): | |
| for s in self.sensors: | |
| s.set_transform(tf) | |
| def grab(self, frame: int, timeout: float = 120.0) -> dict: | |
| got = {} | |
| for kind, q in self.queues.items(): | |
| while True: | |
| img = q.get(timeout=timeout) | |
| if img.frame == frame: | |
| got[kind] = img | |
| break | |
| return got | |
| def stop(self): | |
| for s in self.sensors: | |
| s.stop() | |
| class Frame: | |
| rgb: np.ndarray # [H,W,3] uint8 | |
| z: np.ndarray # [SENSOR_H,SENSOR_W] planar depth (m), sky -> 0, clipped at 0.3 | |
| ids: np.ndarray # CWM semantic ids at sensor resolution | |
| tag: np.ndarray # raw CARLA tags | |
| hero: np.ndarray # protagonist mask | |
| def decode(imgs: dict, hero_loc, rgb_sensor) -> Frame: | |
| bgra = np.frombuffer(imgs[RGB].raw_data, np.uint8).reshape(RGB_H, RGB_W, 4) | |
| d = np.frombuffer(imgs[DEPTH].raw_data, np.uint8).reshape(SENSOR_H, SENSOR_W, 4).astype(np.float64) | |
| z = (d[..., 2] + d[..., 1] * 256.0 + d[..., 0] * 65536.0) / (256.0 ** 3 - 1) * 1000.0 # 1000 m far | |
| s = np.frombuffer(imgs[SEG].raw_data, np.uint8).reshape(SENSOR_H, SENSOR_W, 4) | |
| tag = s[..., 2] | |
| inst = s[..., 1].astype(np.int32) + 256 * s[..., 0].astype(np.int32) | |
| ids = TAG2CWM[tag].copy() | |
| hero = hero_mask(tag, inst, hero_loc, np.array(rgb_sensor.get_transform().get_inverse_matrix())) | |
| ids[hero] = 8 | |
| sky = (ids == 1) | (z >= 999.0) | |
| ids[sky] = 1 | |
| z = np.where(sky, 0.0, np.clip(z, 0.3, None)).astype(np.float32) | |
| return Frame(np.ascontiguousarray(bgra[..., 2::-1]), z, ids, tag, hero) | |
| def hero_mask(tag, inst, hero_loc, w2c) -> np.ndarray: | |
| """The instance channel is an internal object id, not the actor id: project the protagonist's | |
| body into the image and take the pedestrian instance found there (majority vote).""" | |
| fpx = (SENSOR_W / 2) / math.tan(math.radians(HFOV) / 2) | |
| votes = {} | |
| for dz in (-0.4, 0.0, 0.4): | |
| pc = w2c @ np.array([hero_loc.x, hero_loc.y, hero_loc.z + dz, 1.0]) | |
| if pc[0] <= 0.3: | |
| continue | |
| u = int(SENSOR_W / 2 + fpx * pc[1] / pc[0]) | |
| v = int(SENSOR_H / 2 - fpx * pc[2] / pc[0]) | |
| y0, y1, x0, x1 = max(v - 6, 0), min(v + 7, SENSOR_H), max(u - 6, 0), min(u + 7, SENSOR_W) | |
| if y0 >= y1 or x0 >= x1: | |
| continue | |
| win_t, win_i = tag[y0:y1, x0:x1], inst[y0:y1, x0:x1] | |
| for key in win_i[win_t == TAG_PEDESTRIAN].ravel(): | |
| votes[int(key)] = votes.get(int(key), 0) + 1 | |
| if not votes: | |
| return np.zeros_like(tag, bool) | |
| return (tag == TAG_PEDESTRIAN) & (inst == max(votes, key=votes.get)) | |
| def leg_clip_fraction(f: Frame) -> float: | |
| """Fraction of the protagonist's leg box covered by solid scenery (building, wall, fence, | |
| vegetation, props) at its own depth. Ground classes are excluded -- they sit at foot depth.""" | |
| ys, xs = np.nonzero(f.hero) | |
| if len(ys) <= 80: | |
| return 0.0 | |
| y0, y1, x0, x1 = ys.min(), ys.max() + 1, xs.min(), xs.max() + 1 | |
| legs = slice(y0 + int(0.55 * (y1 - y0)), y1) | |
| hz = float(np.median(f.z[ys, xs])) | |
| solid = np.isin(f.tag[legs, x0:x1], TAGS_SOLID) & (np.abs(f.z[legs, x0:x1] - hz) < 0.7) | |
| return float(solid.mean()) | |