Spaces:
Running on Zero
Running on Zero
File size: 4,557 Bytes
1811aab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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()
@dataclass
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())
|