Update to high speed encoder stream
This commit is contained in:
394
Python/clean_recordings.py
Normal file
394
Python/clean_recordings.py
Normal file
@@ -0,0 +1,394 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import statistics
|
||||
import sys
|
||||
from bisect import bisect_left
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
INPUT_COLUMNS = {"time_us", "turn", "pulse"}
|
||||
GP_COLUMNS = ("gp0_falling", "gp1_falling")
|
||||
OUTPUT_COLUMNS = ("time_us", "turn", "pulse", *GP_COLUMNS)
|
||||
DEFAULT_PPR = 256
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Events:
|
||||
turns: list[int]
|
||||
pulses: list[int]
|
||||
gp0_falling: list[int]
|
||||
gp1_falling: list[int]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CleanResult:
|
||||
turns: list[int]
|
||||
pulses: list[int]
|
||||
gp0_falling: list[int]
|
||||
gp1_falling: list[int]
|
||||
pulse_noise_removed: int
|
||||
pulses_inserted: int
|
||||
turn_noise_removed: int
|
||||
turn_phase: int | None
|
||||
|
||||
|
||||
def positive_intervals(times: list[int]) -> list[int]:
|
||||
return [b - a for a, b in zip(times, times[1:]) if b > a]
|
||||
|
||||
|
||||
def median_or_none(values: list[int] | list[float]) -> float | None:
|
||||
if not values:
|
||||
return None
|
||||
return float(statistics.median(values))
|
||||
|
||||
|
||||
def robust_median_interval(times: list[int]) -> float | None:
|
||||
intervals = positive_intervals(times)
|
||||
if not intervals:
|
||||
return None
|
||||
|
||||
intervals = sorted(intervals)
|
||||
if len(intervals) >= 20:
|
||||
lo = len(intervals) // 20
|
||||
hi = len(intervals) - lo
|
||||
intervals = intervals[lo:hi]
|
||||
|
||||
return float(statistics.median(intervals))
|
||||
|
||||
|
||||
def read_events(path: Path) -> Events:
|
||||
turns: list[int] = []
|
||||
pulses: list[int] = []
|
||||
gp0_falling: list[int] = []
|
||||
gp1_falling: list[int] = []
|
||||
|
||||
with path.open(newline="") as file:
|
||||
reader = csv.DictReader(file)
|
||||
fieldnames = set(reader.fieldnames or ())
|
||||
missing = INPUT_COLUMNS - fieldnames
|
||||
if missing:
|
||||
raise ValueError(f"missing columns: {', '.join(sorted(missing))}")
|
||||
|
||||
for row in reader:
|
||||
time_us = int(row["time_us"])
|
||||
if int(row["turn"]):
|
||||
turns.append(time_us)
|
||||
if int(row["pulse"]):
|
||||
pulses.append(time_us)
|
||||
if "gp0_falling" in fieldnames and int(row["gp0_falling"]):
|
||||
gp0_falling.append(time_us)
|
||||
if "gp1_falling" in fieldnames and int(row["gp1_falling"]):
|
||||
gp1_falling.append(time_us)
|
||||
|
||||
return Events(
|
||||
turns=sorted(turns),
|
||||
pulses=sorted(pulses),
|
||||
gp0_falling=sorted(gp0_falling),
|
||||
gp1_falling=sorted(gp1_falling),
|
||||
)
|
||||
|
||||
|
||||
def remove_close_pulses(pulses: list[int]) -> tuple[list[int], int]:
|
||||
if len(pulses) < 3:
|
||||
return sorted(set(pulses)), 0
|
||||
|
||||
median_interval = robust_median_interval(pulses)
|
||||
if median_interval is None:
|
||||
return sorted(set(pulses)), 0
|
||||
|
||||
min_interval = max(2, int(median_interval * 0.35))
|
||||
cleaned: list[int] = []
|
||||
removed = 0
|
||||
|
||||
for pulse in sorted(pulses):
|
||||
if cleaned and pulse <= cleaned[-1]:
|
||||
removed += 1
|
||||
continue
|
||||
if cleaned and pulse - cleaned[-1] < min_interval:
|
||||
removed += 1
|
||||
continue
|
||||
cleaned.append(pulse)
|
||||
|
||||
return cleaned, removed
|
||||
|
||||
|
||||
def local_interval(intervals: list[int], index: int, fallback: float) -> float:
|
||||
start = max(0, index - 16)
|
||||
end = min(len(intervals), index + 17)
|
||||
neighbors = intervals[start:index] + intervals[index + 1 : end]
|
||||
plausible = [dt for dt in neighbors if 0.25 * fallback <= dt <= 4.0 * fallback]
|
||||
return median_or_none(plausible) or fallback
|
||||
|
||||
|
||||
def interpolate_small_gaps(
|
||||
pulses: list[int],
|
||||
*,
|
||||
max_missing_pulses: int,
|
||||
gap_tolerance: float,
|
||||
) -> tuple[list[int], int]:
|
||||
if len(pulses) < 3:
|
||||
return pulses, 0
|
||||
|
||||
fallback = robust_median_interval(pulses)
|
||||
if fallback is None:
|
||||
return pulses, 0
|
||||
|
||||
intervals = positive_intervals(pulses)
|
||||
result = [pulses[0]]
|
||||
inserted = 0
|
||||
|
||||
for index, (start, end) in enumerate(zip(pulses, pulses[1:])):
|
||||
gap = end - start
|
||||
expected = local_interval(intervals, index, fallback)
|
||||
pulse_count = round(gap / expected)
|
||||
|
||||
if 2 <= pulse_count <= max_missing_pulses + 1:
|
||||
corrected_interval = gap / pulse_count
|
||||
error = abs(corrected_interval - expected) / expected
|
||||
if error <= gap_tolerance:
|
||||
for step in range(1, pulse_count):
|
||||
result.append(round(start + corrected_interval * step))
|
||||
inserted += 1
|
||||
|
||||
result.append(end)
|
||||
|
||||
return result, inserted
|
||||
|
||||
|
||||
def nearest_index(times: list[int], target: int) -> int | None:
|
||||
if not times:
|
||||
return None
|
||||
|
||||
index = bisect_left(times, target)
|
||||
candidates = []
|
||||
if index < len(times):
|
||||
candidates.append(index)
|
||||
if index > 0:
|
||||
candidates.append(index - 1)
|
||||
|
||||
return min(candidates, key=lambda candidate: abs(times[candidate] - target))
|
||||
|
||||
|
||||
def remove_close_turns(turns: list[int], min_separation_us: int) -> tuple[list[int], int]:
|
||||
cleaned: list[int] = []
|
||||
removed = 0
|
||||
|
||||
for turn in sorted(turns):
|
||||
if cleaned and turn - cleaned[-1] < min_separation_us:
|
||||
removed += 1
|
||||
continue
|
||||
cleaned.append(turn)
|
||||
|
||||
return cleaned, removed
|
||||
|
||||
|
||||
def choose_turn_phase(turns: list[int], pulses: list[int], ppr: int) -> int | None:
|
||||
if not turns or len(pulses) < ppr:
|
||||
return None
|
||||
|
||||
pulse_interval = robust_median_interval(pulses)
|
||||
if pulse_interval is None:
|
||||
return None
|
||||
|
||||
max_distance_us = max(1_000, int(pulse_interval * 10))
|
||||
scores = [0.0] * ppr
|
||||
|
||||
for turn in turns:
|
||||
index = nearest_index(pulses, turn)
|
||||
if index is None:
|
||||
continue
|
||||
distance = abs(pulses[index] - turn)
|
||||
if distance > max_distance_us:
|
||||
continue
|
||||
scores[index % ppr] += 1.0 - distance / max_distance_us
|
||||
|
||||
best_score = max(scores)
|
||||
if best_score <= 0:
|
||||
return None
|
||||
|
||||
return scores.index(best_score)
|
||||
|
||||
|
||||
def synthesize_turns(turns: list[int], pulses: list[int], ppr: int) -> tuple[list[int], int, int | None]:
|
||||
if len(pulses) < ppr:
|
||||
return turns, 0, None
|
||||
|
||||
pulse_interval = robust_median_interval(pulses) or 0.0
|
||||
expected_turn_interval = max(1, int(pulse_interval * ppr))
|
||||
debounced_turns, removed = remove_close_turns(
|
||||
turns,
|
||||
max(1, int(expected_turn_interval * 0.45)),
|
||||
)
|
||||
phase = choose_turn_phase(debounced_turns, pulses, ppr)
|
||||
if phase is None:
|
||||
phase = 0
|
||||
|
||||
match_window_us = max(1_000, int(expected_turn_interval * 0.15))
|
||||
synthesized: list[int] = []
|
||||
|
||||
for pulse_index in range(phase, len(pulses), ppr):
|
||||
pulse_time = pulses[pulse_index]
|
||||
turn_index = nearest_index(debounced_turns, pulse_time)
|
||||
if turn_index is not None and abs(debounced_turns[turn_index] - pulse_time) <= match_window_us:
|
||||
synthesized.append(debounced_turns[turn_index])
|
||||
else:
|
||||
synthesized.append(pulse_time)
|
||||
|
||||
return synthesized, removed, phase
|
||||
|
||||
|
||||
def clean_events(
|
||||
events: Events,
|
||||
*,
|
||||
ppr: int,
|
||||
max_missing_pulses: int,
|
||||
gap_tolerance: float,
|
||||
) -> CleanResult:
|
||||
pulses, pulse_noise_removed = remove_close_pulses(events.pulses)
|
||||
pulses, pulses_inserted = interpolate_small_gaps(
|
||||
pulses,
|
||||
max_missing_pulses=max_missing_pulses,
|
||||
gap_tolerance=gap_tolerance,
|
||||
)
|
||||
turns, turn_noise_removed, turn_phase = synthesize_turns(events.turns, pulses, ppr)
|
||||
|
||||
return CleanResult(
|
||||
turns=turns,
|
||||
pulses=pulses,
|
||||
gp0_falling=events.gp0_falling,
|
||||
gp1_falling=events.gp1_falling,
|
||||
pulse_noise_removed=pulse_noise_removed,
|
||||
pulses_inserted=pulses_inserted,
|
||||
turn_noise_removed=turn_noise_removed,
|
||||
turn_phase=turn_phase,
|
||||
)
|
||||
|
||||
|
||||
def write_events(path: Path, result: CleanResult) -> None:
|
||||
rows: dict[int, list[int]] = {}
|
||||
for turn in result.turns:
|
||||
rows.setdefault(turn, [0, 0, 0, 0])[0] = 1
|
||||
for pulse in result.pulses:
|
||||
rows.setdefault(pulse, [0, 0, 0, 0])[1] = 1
|
||||
for gp0 in result.gp0_falling:
|
||||
rows.setdefault(gp0, [0, 0, 0, 0])[2] = 1
|
||||
for gp1 in result.gp1_falling:
|
||||
rows.setdefault(gp1, [0, 0, 0, 0])[3] = 1
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with path.open("w", newline="") as file:
|
||||
writer = csv.writer(file)
|
||||
writer.writerow(OUTPUT_COLUMNS)
|
||||
for time_us in sorted(rows):
|
||||
writer.writerow((time_us, *rows[time_us]))
|
||||
|
||||
|
||||
def output_path_for(input_path: Path, input_root: Path, output_root: Path) -> Path:
|
||||
relative = input_path.relative_to(input_root)
|
||||
return output_root / relative.with_name(f"{relative.stem}_cleaned.csv")
|
||||
|
||||
|
||||
def process_file(
|
||||
input_path: Path,
|
||||
output_path: Path,
|
||||
*,
|
||||
ppr: int,
|
||||
max_missing_pulses: int,
|
||||
gap_tolerance: float,
|
||||
) -> CleanResult:
|
||||
events = read_events(input_path)
|
||||
result = clean_events(
|
||||
events,
|
||||
ppr=ppr,
|
||||
max_missing_pulses=max_missing_pulses,
|
||||
gap_tolerance=gap_tolerance,
|
||||
)
|
||||
write_events(output_path, result)
|
||||
return result
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Clean rotary turn and pulse event CSVs from recordings/."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-root",
|
||||
type=Path,
|
||||
default=Path("recordings"),
|
||||
help="Folder to search recursively for CSV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-root",
|
||||
type=Path,
|
||||
default=Path("recordings_cleaned"),
|
||||
help="Folder where cleaned CSV files are written.",
|
||||
)
|
||||
parser.add_argument("--ppr", type=int, default=DEFAULT_PPR, help="Pulse encoder pulses per turn.")
|
||||
parser.add_argument(
|
||||
"--max-missing-pulses",
|
||||
type=int,
|
||||
default=8,
|
||||
help="Largest pulse gap to interpolate. Larger gaps are treated as recording breaks.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gap-tolerance",
|
||||
type=float,
|
||||
default=0.45,
|
||||
help="Allowed fractional error when deciding whether a gap is missing pulses.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
input_root = args.input_root
|
||||
output_root = args.output_root
|
||||
|
||||
if not input_root.exists():
|
||||
print(f"Input folder does not exist: {input_root}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
csv_paths = sorted(
|
||||
path
|
||||
for path in input_root.rglob("*.csv")
|
||||
if not path.name.endswith("_cleaned.csv")
|
||||
)
|
||||
if not csv_paths:
|
||||
print(f"No CSV files found under {input_root}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
failures = 0
|
||||
for input_path in csv_paths:
|
||||
output_path = output_path_for(input_path, input_root, output_root)
|
||||
try:
|
||||
result = process_file(
|
||||
input_path,
|
||||
output_path,
|
||||
ppr=args.ppr,
|
||||
max_missing_pulses=args.max_missing_pulses,
|
||||
gap_tolerance=args.gap_tolerance,
|
||||
)
|
||||
except (OSError, ValueError) as exc:
|
||||
failures += 1
|
||||
print(f"Skipping {input_path}: {exc}", file=sys.stderr)
|
||||
continue
|
||||
|
||||
phase = "unknown" if result.turn_phase is None else str(result.turn_phase)
|
||||
print(
|
||||
f"{input_path} -> {output_path} "
|
||||
f"pulses={len(result.pulses)} "
|
||||
f"turns={len(result.turns)} "
|
||||
f"removed_pulses={result.pulse_noise_removed} "
|
||||
f"inserted_pulses={result.pulses_inserted} "
|
||||
f"removed_turns={result.turn_noise_removed} "
|
||||
f"phase={phase}"
|
||||
)
|
||||
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user