Files
CARLA_RGB_depth_semantic_co…/carla_road_collector_belka.py
2026-04-16 11:55:28 +03:00

706 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
CARLA Road-Following Map Collector
====================================
Сбор мультимодальных снимков вдоль дорог карты CARLA.
Три типа съёмки на каждой точке:
• satellite/ — вид сверху (pitch = -90°)
• drone/ — вид под углом с орбиты
• street/ — вид с уровня человеческого роста (~1.7 м)
Структура выходных файлов:
carla_road_dataset/Town10HD/
├── wp000000_r5_l-1/
│ ├── satellite/
│ │ ├── rgb.png
│ │ ├── semantic_label.png
│ │ ├── semantic_color.png
│ │ ├── depth_uint16.png
│ │ └── depth_color.png
│ ├── drone/
│ │ ├── yaw000_rgb.png
│ │ ├── yaw000_semantic_label.png
│ │ ├── yaw000_semantic_color.png
│ │ ├── yaw000_depth_uint16.png
│ │ ├── yaw000_depth_color.png
│ │ └── yaw090_… (и т.д.)
│ └── street/
│ ├── forward_rgb.png
│ ├── forward_semantic_label.png
│ ├── forward_semantic_color.png
│ ├── forward_depth_uint16.png
│ ├── forward_depth_color.png
│ ├── left_rgb.png
│ ├── right_rgb.png
│ └── backward_rgb.png
├── wp000001_r5_l-1/
│ └── …
└── metadata.json
Usage:
python carla_road_collector.py --host localhost --port 2000 \\
--output ./dataset --route dense --step 20
Requirements:
pip install carla numpy opencv-python tqdm
"""
import carla
import numpy as np
import cv2
import math
import argparse
import json
import threading
from pathlib import Path
from dataclasses import dataclass, field
from typing import Optional, List, Dict, Tuple
from tqdm import tqdm
import os
# ---------------------------------------------------------------------------
# Semantic palette
# ---------------------------------------------------------------------------
SEMANTIC_PALETTE: Dict[int, Tuple[int, int, int]] = {
0: (0, 0, 0), 1: (70, 70, 70), 2: (100, 40, 40),
3: (55, 90, 80), 4: (220, 20, 60), 5: (153, 153, 153),
6: (157, 234, 50), 7: (128, 64, 128), 8: (244, 35, 232),
9: (107, 142, 35), 10: (0, 0, 142), 11: (102, 102, 156),
12: (220, 220, 0), 13: (70, 130, 180), 14: (81, 0, 81),
15: (150, 100, 100), 16: (230, 150, 140), 17: (180, 165, 180),
18: (250, 170, 30), 19: (110, 190, 160), 20: (170, 120, 50),
21: (45, 60, 150), 22: (145, 170, 100),
}
_SEM_NAMES = {
0:"Unlabeled", 1:"Building", 2:"Fence", 3:"Other", 4:"Pedestrian",
5:"Pole", 6:"RoadLine", 7:"Road", 8:"SideWalk", 9:"Vegetation",
10:"Vehicle", 11:"Wall", 12:"TrafficSign", 13:"Sky", 14:"Ground",
15:"Bridge", 16:"RailTrack", 17:"GuardRail", 18:"TrafficLight",
19:"Static", 20:"Dynamic", 21:"Water", 22:"Terrain",
}
_PALETTE_LUT: Optional[np.ndarray] = None
def _get_lut() -> np.ndarray:
global _PALETTE_LUT
if _PALETTE_LUT is None:
lut = np.zeros((256, 3), dtype=np.uint8)
for tag, bgr in SEMANTIC_PALETTE.items():
lut[tag] = bgr
_PALETTE_LUT = lut
return _PALETTE_LUT
def semantic_to_color(label: np.ndarray) -> np.ndarray:
return _get_lut()[label]
# ---------------------------------------------------------------------------
# Depth helpers
# ---------------------------------------------------------------------------
def decode_depth_meters(img: carla.Image) -> np.ndarray:
arr = np.frombuffer(img.raw_data, dtype=np.uint8).reshape(
(img.height, img.width, 4))
R = arr[:, :, 2].astype(np.float32)
G = arr[:, :, 1].astype(np.float32)
B = arr[:, :, 0].astype(np.float32)
return (R + G * 256.0 + B * 65536.0) / (256.0 ** 3 - 1) * 1000.0
def depth_to_uint16(depth_m: np.ndarray, max_d: float) -> np.ndarray:
return (np.clip(depth_m, 0, max_d) / max_d * 65535).astype(np.uint16)
def depth_to_colormap(depth_m: np.ndarray, max_d: float) -> np.ndarray:
norm = (np.clip(depth_m, 0, max_d) / max_d * 255).astype(np.uint8)
return cv2.applyColorMap(norm, cv2.COLORMAP_TURBO)
# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
@dataclass
class RoadCollectorConfig:
# Маршрут
route_mode: str = "dense" # "topology" | "dense"
step_meters: float = 20.0
min_dist_between_captures: float = 10.0
# Спутник
sat_altitude: float = 60.0
sat_fov: float = 90.0
sat_image_w: int = 1024
sat_image_h: int = 1024
# Дрон
drone_altitude: float = 40.0
drone_pitch: float = -35.0
drone_yaw_offsets: List[float] = field(
default_factory=lambda: [0.0, 90.0, 180.0, 270.0])
drone_fov: float = 90.0
drone_image_w: int = 1024
drone_image_h: int = 1024
# Улица (вид с уровня человека)
street_height: float = 1.9 # высота камеры над землёй, метры
street_fov: float = 90.0
street_image_w: int = 1024
street_image_h: int = 1024
# Стороны съёмки: forward=вперёд по дороге, остальные — относительно него
street_yaw_offsets: List[float] = field(
default_factory=lambda: [0.0, 90.0, 180.0, 270.0])
street_yaw_labels: List[str] = field(
default_factory=lambda: ["forward", "right", "backward", "left"])
# Модальности
capture_rgb: bool = True
capture_semantic: bool = True
capture_depth: bool = True
depth_max_meters: float = 100.0
# Виды
capture_satellite: bool = True
capture_drone: bool = True
capture_street: bool = True
# Вывод
output_dir: str = "./carla_road_dataset"
save_metadata: bool = True
# Синхронизация
settle_ticks: int = 5
# ---------------------------------------------------------------------------
# Sensor infrastructure
# ---------------------------------------------------------------------------
class _SensorWrapper:
def __init__(self):
self._data = None
self._event = threading.Event()
def callback(self, data):
self._data = data
self._event.set()
def wait(self, timeout: float = 10.0):
if not self._event.wait(timeout) or self._data is None:
raise RuntimeError("Sensor timed out.")
return self._data
class MultiSensorCapture:
"""Спавним RGB + Semantic + Depth в одной точке, снимаем за один тик."""
_BP = {
"rgb": "sensor.camera.rgb",
"semantic": "sensor.camera.semantic_segmentation",
"depth": "sensor.camera.depth",
}
def __init__(self, world: carla.World, transform: carla.Transform,
w: int, h: int, fov: float, modalities: List[str]):
self._wrappers: Dict[str, _SensorWrapper] = {}
self._actors: List[carla.Actor] = []
bp_lib = world.get_blueprint_library()
for name in modalities:
bp = bp_lib.find(self._BP[name])
bp.set_attribute("image_size_x", str(w))
bp.set_attribute("image_size_y", str(h))
bp.set_attribute("fov", str(fov))
actor = world.spawn_actor(bp, transform)
wrapper = _SensorWrapper()
actor.listen(wrapper.callback)
self._wrappers[name] = wrapper
self._actors.append(actor)
def capture(self, tick_fn, settle_ticks: int = 5) -> Dict[str, carla.Image]:
for _ in range(settle_ticks):
tick_fn()
return {name: w.wait() for name, w in self._wrappers.items()}
def destroy(self):
for a in self._actors:
if a and a.is_alive:
a.destroy()
self._actors.clear()
# ---------------------------------------------------------------------------
# Сохранение одного набора изображений (все модальности) в папку
# ---------------------------------------------------------------------------
def save_shot(images: Dict[str, carla.Image],
folder: Path,
prefix: str,
cfg: RoadCollectorConfig) -> Dict:
"""
Сохраняет все активные модальности в `folder`.
Имена файлов: <prefix>_rgb.png, <prefix>_semantic_label.png и т.д.
Если prefix пустой — просто rgb.png, semantic_label.png и т.д.
Возвращает словарь с путями файлов и статистикой глубины.
"""
folder.mkdir(parents=True, exist_ok=True)
def fname(tag: str) -> Path:
name = f"{prefix}_{tag}.png" if prefix else f"{tag}.png"
return folder / name
saved: Dict = {}
# ── RGB ──────────────────────────────────────────────────────────────────
if "rgb" in images:
img = images["rgb"]
arr = np.frombuffer(img.raw_data, dtype=np.uint8).reshape(
(img.height, img.width, 4))
p = fname("rgb")
cv2.imwrite(str(p), arr[:, :, :3])
saved["rgb"] = p.name
# ── Semantic ─────────────────────────────────────────────────────────────
if "semantic" in images:
img = images["semantic"]
arr = np.frombuffer(img.raw_data, dtype=np.uint8).reshape(
(img.height, img.width, 4))
label = arr[:, :, 2] # тег хранится в R-канале
p_lbl = fname("semantic_label")
cv2.imwrite(str(p_lbl), label)
saved["semantic_label"] = p_lbl.name
p_col = fname("semantic_color")
cv2.imwrite(str(p_col), semantic_to_color(label))
saved["semantic_color"] = p_col.name
# ── Depth ─────────────────────────────────────────────────────────────────
if "depth" in images:
depth_m = decode_depth_meters(images["depth"])
p_u16 = fname("depth_uint16")
cv2.imwrite(str(p_u16), depth_to_uint16(depth_m, cfg.depth_max_meters))
saved["depth_uint16"] = p_u16.name
p_vis = fname("depth_color")
cv2.imwrite(str(p_vis), depth_to_colormap(depth_m, cfg.depth_max_meters))
saved["depth_color"] = p_vis.name
saved["depth_stats"] = {
"min_m": float(np.nanmin(depth_m)),
"max_m": float(np.nanmax(depth_m)),
"mean_m": float(np.nanmean(depth_m)),
}
return saved
def _active_modalities(cfg: RoadCollectorConfig) -> List[str]:
m = [k for k, v in [("rgb", cfg.capture_rgb),
("semantic", cfg.capture_semantic),
("depth", cfg.capture_depth)] if v]
if not m:
raise ValueError("Включите хотя бы одну модальность.")
return m
# ---------------------------------------------------------------------------
# Route builders
# ---------------------------------------------------------------------------
def build_route_topology(carla_map: carla.Map) -> List[carla.Waypoint]:
topology = carla_map.get_topology()
seen = set()
waypoints = []
for wp_s, wp_e in topology:
key = (wp_s.road_id, wp_s.lane_id)
if key in seen:
continue
seen.add(key)
ls, le = wp_s.transform.location, wp_e.transform.location
mid = carla.Location(x=(ls.x+le.x)/2, y=(ls.y+le.y)/2, z=(ls.z+le.z)/2)
wp = carla_map.get_waypoint(mid, project_to_road=True,
lane_type=carla.LaneType.Driving)
if wp:
waypoints.append(wp)
print(f"[Route] topology: {len(waypoints)} сегментов")
return waypoints
def build_route_dense(carla_map: carla.Map,
step: float, min_dist: float) -> List[carla.Waypoint]:
all_wps = [wp for wp in carla_map.generate_waypoints(step)
if wp.lane_type == carla.LaneType.Driving]
kept = []
for wp in all_wps:
loc = wp.transform.location
if not any(loc.distance(k.transform.location) < min_dist for k in kept):
kept.append(wp)
# ! reverse
kept = kept[::-1]
print(f"[Route] dense: {len(all_wps)}{len(kept)} точек "
f"(шаг={step}м, min_dist={min_dist}м)")
return kept
# ---------------------------------------------------------------------------
# Capture at one waypoint — satellite / drone / street
# ---------------------------------------------------------------------------
def capture_at_waypoint(world: carla.World,
wp: carla.Waypoint,
idx: int,
cfg: RoadCollectorConfig,
map_dir: Path) -> dict:
"""
Создаёт папку для одной точки маршрута и снимает все три вида.
Структура папки:
wp000042_r12_l-1/
├── satellite/
│ ├── rgb.png
│ ├── semantic_label.png
│ ├── semantic_color.png
│ ├── depth_uint16.png
│ └── depth_color.png
├── drone/
│ ├── yaw000_rgb.png … yaw000_depth_color.png
│ └── yaw090_rgb.png …
└── street/
├── forward_rgb.png … forward_depth_color.png
├── right_rgb.png …
├── backward_rgb.png …
└── left_rgb.png …
"""
loc = wp.transform.location
road_yaw = wp.transform.rotation.yaw
modalities = _active_modalities(cfg)
# Папка этой точки
cell_name = f"wp{idx:06d}_r{wp.road_id}_l{wp.lane_id}"
cell_dir = map_dir / cell_name
cell_dir.mkdir(parents=True, exist_ok=True)
meta: dict = {
"cell": cell_name,
"waypoint_idx": idx,
"road_id": wp.road_id,
"lane_id": wp.lane_id,
"road_yaw_deg": road_yaw,
"location": {"x": loc.x, "y": loc.y, "z": loc.z},
"satellite": None,
"drone": [],
"street": [],
}
# ── 1. SATELLITE ──────────────────────────────────────────────────────────
if cfg.capture_satellite:
cam_z = loc.z + cfg.sat_altitude
#! foward
sat_yaw = road_yaw+180
#! reverse
# sat_yaw = road_yaw
transform = carla.Transform(
carla.Location(x=loc.x, y=loc.y, z=cam_z),
carla.Rotation(pitch=-90.0, yaw=sat_yaw, roll=0.0)
)
sensor = MultiSensorCapture(world, transform,
cfg.sat_image_w, cfg.sat_image_h,
cfg.sat_fov, modalities)
try:
images = sensor.capture(world.tick, cfg.settle_ticks)
finally:
sensor.destroy()
sat_dir = cell_dir / "satellite"
saved = save_shot(images, sat_dir, prefix="", cfg=cfg)
meta["satellite"] = {
"camera": {"x": loc.x, "y": loc.y, "z": cam_z,
"pitch": -90.0, "yaw": sat_yaw}
}
# ── 2. DRONE ──────────────────────────────────────────────────────────────
if cfg.capture_drone:
cam_z = loc.z + cfg.drone_altitude
pitch_rad = math.radians(abs(cfg.drone_pitch))
orbit_r = (cfg.drone_altitude / math.tan(pitch_rad)
if pitch_rad > 1e-6 else 0.0)
drone_dir = cell_dir / "drone"
for yaw_off in cfg.drone_yaw_offsets:
abs_yaw = road_yaw + yaw_off # выравниваем по дороге
yaw_rad = math.radians(abs_yaw)
#! foward
cam_x = loc.x - (orbit_r * math.cos(yaw_rad))/4
cam_y = loc.y - (orbit_r * math.sin(yaw_rad))/4
look_yaw = abs_yaw
#! reverse
# cam_x = loc.x + (orbit_r * math.cos(yaw_rad))/4
# cam_y = loc.y + (orbit_r * math.sin(yaw_rad))/4
# look_yaw = abs_yaw + 180
transform = carla.Transform(
carla.Location(x=cam_x, y=cam_y, z=cam_z),
carla.Rotation(pitch=cfg.drone_pitch, yaw=look_yaw, roll=0.0)
)
sensor = MultiSensorCapture(world, transform,
cfg.drone_image_w, cfg.drone_image_h,
cfg.drone_fov, modalities)
try:
images = sensor.capture(world.tick, cfg.settle_ticks)
finally:
sensor.destroy()
prefix = f"yaw{int(yaw_off):03d}"
saved = save_shot(images, drone_dir, prefix=prefix, cfg=cfg)
meta["drone"].append({
"yaw_offset_from_road": yaw_off,
"camera": {"x": cam_x, "y": cam_y, "z": cam_z,
"pitch": cfg.drone_pitch, "yaw": look_yaw},
"orbit_radius": orbit_r,
"altitude": cfg.drone_altitude,
"fov": cfg.drone_fov,
"files": saved,
})
# ── 3. STREET ─────────────────────────────────────────────────────────────
if cfg.capture_street:
cam_z = loc.z + cfg.street_height
street_dir = cell_dir / "street"
labels = cfg.street_yaw_labels
for i, yaw_off in enumerate(cfg.street_yaw_offsets):
#! foward
abs_yaw = road_yaw + yaw_off # 0° = смотрим вперёд по дороге
#! reverse
# abs_yaw += 180
label = labels[i] if i < len(labels) else f"yaw{int(yaw_off):03d}"
cam_x = loc.x - orbit_r * math.cos(yaw_rad)
cam_y = loc.y - orbit_r * math.sin(yaw_rad)
transform = carla.Transform(
carla.Location(x=loc.x, y=loc.y, z=cam_z), # old v
carla.Rotation(pitch=0.0, yaw=abs_yaw, roll=0.0)
)
sensor = MultiSensorCapture(world, transform,
cfg.street_image_w, cfg.street_image_h,
cfg.street_fov, modalities)
try:
images = sensor.capture(world.tick, cfg.settle_ticks)
finally:
sensor.destroy()
saved = save_shot(images, street_dir, prefix=label, cfg=cfg)
meta["street"].append({
"direction": label,
"yaw_offset_from_road": yaw_off,
"camera": {"x": cam_x, "y": cam_y, "z": cam_z,
"pitch": 0.0, "yaw": abs_yaw},
"height": cfg.street_height,
"fov": cfg.street_fov,
"files": saved,
})
return meta
# ---------------------------------------------------------------------------
# Main collection loop
# ---------------------------------------------------------------------------
def collect_along_roads(cfg: RoadCollectorConfig,
host: str = "localhost",
port: int = 2000):
print(f"[CARLA] Подключение к {host}:{port}")
client = carla.Client(host, port)
client.set_timeout(30.0)
world = client.get_world()
settings = world.get_settings()
settings.synchronous_mode = True
settings.fixed_delta_seconds = 0.05
settings.no_rendering_mode = False
world.apply_settings(settings)
carla_map = world.get_map()
map_name = carla_map.name
modalities = _active_modalities(cfg)
views = [v for v, on in [("satellite", cfg.capture_satellite),
("drone", cfg.capture_drone),
("street", cfg.capture_street)] if on]
print(f"[CARLA] Карта: {map_name}")
print(f"[Config] Виды: {views} | Модальности: {modalities}")
map_dir = Path(cfg.output_dir) / map_name.replace("/", "_")
map_dir.mkdir(parents=True, exist_ok=True)
# Маршрут
if cfg.route_mode == "topology":
waypoints = build_route_topology(carla_map)
elif cfg.route_mode == "dense":
waypoints = build_route_dense(carla_map, cfg.step_meters,
cfg.min_dist_between_captures)
else:
raise ValueError(f"Неизвестный route_mode: {cfg.route_mode!r}")
print(f"[Route] Точек для съёмки: {len(waypoints)}")
all_meta = []
try:
for idx, wp in enumerate(tqdm(waypoints, desc="Capture", unit="wp")):
try:
meta = capture_at_waypoint(world, wp, idx, cfg, map_dir)
all_meta.append(meta)
except Exception as exc:
loc = wp.transform.location
print(f"\n [!] wp{idx:06d} ({loc.x:.1f}, {loc.y:.1f}) → {exc}")
finally:
settings.synchronous_mode = False
world.apply_settings(settings)
if cfg.save_metadata:
meta_path = map_dir / "metadata.json"
with open(meta_path, "w") as f:
json.dump({
"map": map_name,
"config": {
"route_mode": cfg.route_mode,
"step_meters": cfg.step_meters,
"views": views,
"modalities": modalities,
"sat_altitude": cfg.sat_altitude,
"drone_altitude": cfg.drone_altitude,
"drone_pitch": cfg.drone_pitch,
"drone_yaw_offsets": cfg.drone_yaw_offsets,
"street_height": cfg.street_height,
"street_yaw_offsets": cfg.street_yaw_offsets,
"street_yaw_labels": cfg.street_yaw_labels,
"depth_max_meters": cfg.depth_max_meters,
},
"semantic_palette": {
str(k): {"bgr": list(v), "name": _SEM_NAMES.get(k, "?")}
for k, v in SEMANTIC_PALETTE.items()
},
"total_cells": len(all_meta),
"cells": all_meta,
}, f, indent=2)
print(f"\n[Meta] Сохранено → {meta_path}")
total_imgs = sum(
len([f for f in (m.get("satellite") or {}).get("files", {})
if not f.endswith("_stats")])
+ sum(len(d["files"]) for d in m.get("drone", []))
+ sum(len(s["files"]) for s in m.get("street", []))
for m in all_meta
)
print(f"\n✓ Готово. {len(all_meta)} ячеек | ~{total_imgs} файлов → {map_dir}")
return all_meta
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def parse_args():
p = argparse.ArgumentParser(
description="CARLA road collector — satellite + drone + street",
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
p.add_argument("--host", default="localhost")
p.add_argument("--port", type=int, default=2000)
p.add_argument("--output", default="./carla_road_map11_EPIC_FULL")
# p.add_argument("--output", default="./carla_road_map11rev_EPIC_FULL")
g = p.add_argument_group("Маршрут")
g.add_argument("--route", default="dense",
choices=["topology", "dense"])
g.add_argument("--step", type=float, default=20.0, metavar="M", # Было 1000
help="Шаг между точками (dense mode, метры)")
g.add_argument("--min-dist", type=float, default=50.0, metavar="M")# Было 100
g2 = p.add_argument_group("Камеры")
g2.add_argument("--sat-altitude", type=float, default=60.0)
g2.add_argument("--sat-fov", type=float, default=90.0) # было 90
g2.add_argument("--drone-altitude", type=float, default=40.0) # было 40
g2.add_argument("--drone-pitch", type=float, default=-35.0)
g2.add_argument("--drone-yaws", nargs="+", type=float,
#default=[0.0, 90.0, 180.0, 270.0],
default=[180.0], metavar="YAW")
g2.add_argument("--street-height", type=float, default=1.7,
help="Высота уличной камеры над землёй (м)")
g2.add_argument("--street-fov", type=float, default=90.0)
g2.add_argument("--street-yaws", nargs="+", type=float,
#default=[0.0, 90.0, 180.0, 270.0]
default=[0.0], metavar="YAW",
help="Направления улицы: 0=вперёд по дороге")
g2.add_argument("--image-size", type=int, default=1024)
g2.add_argument("--depth-max", type=float, default=100.0)
g2.add_argument("--settle-ticks", type=int, default=5)
g3 = p.add_argument_group("Модальности")
g3.add_argument("--no-rgb", action="store_true")
g3.add_argument("--no-semantic", action="store_false")
g3.add_argument("--no-depth", action="store_false")
p.set_defaults(no_rgb=False, no_semantic=False, no_depth=False)
g4 = p.add_argument_group("Виды")
g4.add_argument("--no-satellite", action="store_true")
g4.add_argument("--no-drone", action="store_true")
g4.add_argument("--no-street", action="store_true")
return p.parse_args()
if __name__ == "__main__":
args = parse_args()
cfg = RoadCollectorConfig(
route_mode=args.route,
step_meters=args.step,
min_dist_between_captures=args.min_dist,
sat_altitude=args.sat_altitude,
sat_fov=args.sat_fov,
sat_image_w=args.image_size,
sat_image_h=args.image_size,
drone_altitude=args.drone_altitude,
drone_pitch=args.drone_pitch,
drone_yaw_offsets=args.drone_yaws,
drone_fov=90.0,
drone_image_w=args.image_size,
drone_image_h=args.image_size,
street_height=args.street_height,
street_fov=args.street_fov,
street_yaw_offsets=args.street_yaws,
street_yaw_labels=["forward", "right", "backward", "left"],
street_image_w=args.image_size,
street_image_h=args.image_size,
capture_rgb=not args.no_rgb,
capture_semantic=not args.no_semantic,
capture_depth=not args.no_depth,
depth_max_meters=args.depth_max,
capture_satellite=not args.no_satellite,
capture_drone=not args.no_drone,
capture_street=not args.no_street,
output_dir=args.output,
settle_ticks=args.settle_ticks,
)
collect_along_roads(cfg=cfg, host=args.host, port=args.port)