395 lines
11 KiB
Python
395 lines
11 KiB
Python
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())
|