import numpy as np
import rerun as rr
import rerun.blueprint as rrb
CATEGORY_COLORS = {
1: (100, 149, 237), # pedestrian
2: (65, 105, 225), # people
3: (255, 215, 0), # bicycle
4: (220, 20, 60), # car
5: (50, 205, 50), # van
6: (255, 140, 0), # truck
7: (138, 43, 226), # tricycle
8: (0, 206, 209), # awning-tricycle
9: (255, 69, 0), # bus
10: (255, 20, 147), # motor
11: (169, 169, 169), # others
}
rr.init('auraison/visdrone', spawn=False)
rr.save(RRD_PATH)
blueprint = rrb.Spatial2DView(
name=f'VisDrone VID -- {SEQ_NAME}',
origin='/drone',
)
rr.send_blueprint(blueprint)
print(f'Logging {len(frame_names)} frames ...')
for frame_path in frame_names:
fname = os.path.basename(frame_path)
frame_idx = int(fname.replace('.jpg', ''))
rr.set_time('frame', sequence=frame_idx)
img_bytes = zf.read(frame_path)
rr.log('/drone/image', rr.EncodedImage(contents=img_bytes, media_type='image/jpeg'))
boxes = ann_by_frame.get(frame_idx, [])
if boxes:
mins = np.array([[x, y] for x, y, w, h, _ in boxes], dtype=np.float32)
sizes = np.array([[w, h] for x, y, w, h, _ in boxes], dtype=np.float32)
colors = [CATEGORY_COLORS.get(cat, (200, 200, 200)) for *_, cat in boxes]
labels = [CATEGORY_NAMES[cat] for *_, cat in boxes]
rr.log('/drone/boxes', rr.Boxes2D(mins=mins, sizes=sizes, colors=colors, labels=labels))
else:
rr.log('/drone/boxes', rr.Clear(recursive=False))
if frame_idx % 10 == 0:
print(f' frame {frame_idx:5d} -- {len(boxes):2d} boxes')
print(f'\nRecording saved: {RRD_PATH}')
rrd_size_mb = os.path.getsize(RRD_PATH) / 1024 / 1024
print(f'File size: {rrd_size_mb:.1f} MiB')
zf.close()