Huang Qi
c8fa5db7d3
Coverage: - test_model: SimpleNet forward (parametrized over batch sizes and both unsqueezed and flat input shapes), layer dimensions, differentiability, and ONNX round-trip - test_inference: load_model resolution order (bundled, cwd override, explicit path, missing path), and predict shape/dtype/determinism plus endpoint sanity across 8 cardinal/diagonal targets - test_train: _load_csv parsing, TrajectoryDataset indexing, full train() pipeline producing a single-file ONNX, plus a smoke test against the real data shipped under data/ - test_cli: --help for the three console scripts and a real run of mouse-visualize via both the entry point and python -m Wire up pytest via dependency-groups and tool.pytest.ini_options.
104 lines
3.2 KiB
Python
104 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
import pytest
|
|
import torch
|
|
|
|
from mouse_control.train import TrajectoryDataset, _load_csv, train
|
|
|
|
|
|
def _write_csv(path: Path, rows: list[list[tuple[int, int]]]) -> None:
|
|
"""Write rows in the project's quirky single-quoted-pair format."""
|
|
with open(path, "w", newline="") as f:
|
|
writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
|
|
for row in rows:
|
|
# First column is the target (last keypoint), then 10 keypoints.
|
|
target_x, target_y = row[-1]
|
|
cells = [f"{target_x},{target_y}"] + [f"{x},{y}" for x, y in row]
|
|
writer.writerow(cells)
|
|
|
|
|
|
@pytest.fixture
|
|
def synthetic_csv(tmp_path: Path) -> Path:
|
|
"""Tiny linear trajectories from origin -> target."""
|
|
path = tmp_path / "synth.csv"
|
|
rows = []
|
|
rng = np.random.default_rng(0)
|
|
for _ in range(20):
|
|
tx, ty = int(rng.integers(-150, 151)), int(rng.integers(-150, 151))
|
|
keypoints = [
|
|
(int(round(tx * i / 9)), int(round(ty * i / 9))) for i in range(10)
|
|
]
|
|
rows.append(keypoints)
|
|
_write_csv(path, rows)
|
|
return path
|
|
|
|
|
|
def test_load_csv_shapes(synthetic_csv: Path) -> None:
|
|
inputs, labels = _load_csv(synthetic_csv)
|
|
assert inputs.shape == (20, 1, 2)
|
|
assert labels.shape == (20, 10, 2)
|
|
assert inputs.dtype == torch.float32
|
|
assert labels.dtype == torch.float32
|
|
|
|
|
|
def test_load_csv_round_trip(synthetic_csv: Path) -> None:
|
|
"""The target column (input) must equal the last keypoint of the labels."""
|
|
inputs, labels = _load_csv(synthetic_csv)
|
|
last_keypoint = labels[:, -1, :]
|
|
np.testing.assert_array_equal(inputs.squeeze(1).numpy(), last_keypoint.numpy())
|
|
|
|
|
|
def test_dataset_indexing(synthetic_csv: Path) -> None:
|
|
inputs, labels = _load_csv(synthetic_csv)
|
|
ds = TrajectoryDataset(inputs, labels)
|
|
assert len(ds) == 20
|
|
x, y = ds[0]
|
|
assert x.shape == (1, 2)
|
|
assert y.shape == (10, 2)
|
|
|
|
|
|
def test_load_csv_missing_file_raises(tmp_path: Path) -> None:
|
|
with pytest.raises(FileNotFoundError):
|
|
_load_csv(tmp_path / "missing.csv")
|
|
|
|
|
|
def test_train_end_to_end_produces_single_file_onnx(
|
|
synthetic_csv: Path, tmp_path: Path,
|
|
) -> None:
|
|
"""A quick training run must export a valid, self-contained ONNX model."""
|
|
output = tmp_path / "trained.onnx"
|
|
train(
|
|
train_csv=synthetic_csv,
|
|
test_csv=synthetic_csv,
|
|
output=output,
|
|
epochs=3,
|
|
batch_size=8,
|
|
)
|
|
assert output.exists()
|
|
assert output.stat().st_size > 1000
|
|
# external_data=False -> no .data sidecar
|
|
assert not output.with_suffix(output.suffix + ".data").exists()
|
|
|
|
session = ort.InferenceSession(str(output))
|
|
inp = np.array([[[100.0, 50.0]]], dtype=np.float32)
|
|
out = session.run(None, {"input": inp})[0]
|
|
assert out.shape == (1, 10, 2)
|
|
|
|
|
|
def test_train_with_real_data(train_csv: Path, test_csv: Path, tmp_path: Path) -> None:
|
|
"""Smoke test against the project's actual data."""
|
|
output = tmp_path / "real.onnx"
|
|
train(
|
|
train_csv=train_csv,
|
|
test_csv=test_csv,
|
|
output=output,
|
|
epochs=2,
|
|
batch_size=32,
|
|
)
|
|
assert output.exists()
|