| | |
| |
|
| | import numpy as np |
| | import torch |
| |
|
| |
|
| | def chunk_sequence( |
| | data, |
| | indices, |
| | *, |
| | names=None, |
| | max_length=100, |
| | min_length=1, |
| | max_delay_s=None, |
| | max_inter_dist=None, |
| | max_total_dist=None, |
| | ): |
| | sort_array = data.get("capture_time", data.get("index")) |
| | if sort_array is None: |
| | sort_array = indices if names is None else names |
| | indices = sorted(indices, key=lambda i: sort_array[i].tolist()) |
| | centers = torch.stack([data["t_c2w"][i][:2] for i in indices]).numpy() |
| | dists = np.linalg.norm(np.diff(centers, axis=0), axis=-1) |
| | if "capture_time" in data: |
| | times = torch.stack([data["capture_time"][i] for i in indices]) |
| | times = times.double() / 1e3 |
| | delays = np.diff(times, axis=0) |
| | else: |
| | delays = np.zeros_like(dists) |
| | chunks = [[indices[0]]] |
| | dist_total = 0 |
| | for dist, delay, idx in zip(dists, delays, indices[1:]): |
| | dist_total += dist |
| | if ( |
| | (max_inter_dist is not None and dist > max_inter_dist) |
| | or (max_total_dist is not None and dist_total > max_total_dist) |
| | or (max_delay_s is not None and delay > max_delay_s) |
| | or len(chunks[-1]) >= max_length |
| | ): |
| | chunks.append([]) |
| | dist_total = 0 |
| | chunks[-1].append(idx) |
| | chunks = list(filter(lambda c: len(c) >= min_length, chunks)) |
| | chunks = sorted(chunks, key=len, reverse=True) |
| | return chunks |
