Hugging Face's logo

Spaces:
SmaugC137
/
toaster
Sleeping

App Files Community
toaster / tests /test_controller.py
SmaugC137's picture
SmaugC137
feat(api,web): delta label protocol — labelling is now O(edit), not O(cloud)
28ee467
Raw History Blame Contribute Delete
10 kB
"""Controller tests with a fake viewer — exercises the workflow without Qt/VTK."""
from __future__ import annotations
import numpy as np
from toaster.core import Grouping, Session
from toaster.interaction import InteractionController
class FakeViewer:
"""Records what the controller asks the renderer to do."""
def __init__(self):
self.last_highlight = None
self.recolored = []
self.point_cb = None
self.box_cb = None
self.cloud_colors = None
self.visible_mask = None
def set_cloud(self, xyz, colors):
self.cloud_colors = colors
def update_colors(self, indices, colors):
self.recolored.append((np.asarray(indices), np.asarray(colors)))
def highlight(self, indices):
self.last_highlight = np.asarray(indices)
def clear_highlight(self):
self.last_highlight = np.empty(0, dtype=np.int64)
def set_point_pick_callback(self, cb):
self.point_cb = cb
def set_box_pick_callback(self, cb):
self.box_cb = cb
def set_pick_mode(self, mode):
self.pick_mode = mode
def set_visible_mask(self, mask):
self.visible_mask = None if mask is None else np.asarray(mask)
def reset_camera(self):
pass
def render(self):
pass
def _session(two_clusters, schema):
two_clusters.ensure_labels(schema.unlabeled_id)
return Session(two_clusters, schema)
def test_pick_without_grouping_selects_single_point(two_clusters, schema):
session = _session(two_clusters, schema)
viewer = FakeViewer()
ctl = InteractionController(session, viewer)
ctl.on_pick(7)
assert session.selection.indices.tolist() == [7]
assert viewer.last_highlight.tolist() == [7]
def test_pick_with_active_grouping_selects_whole_group(two_clusters, schema):
session = _session(two_clusters, schema)
# Group 0 = first 50 points, group 1 = the rest.
gid = np.where(np.arange(two_clusters.n) < 50, 0, 1).astype(np.int32)
session.add_grouping(Grouping(gid))
ctl = InteractionController(session, FakeViewer())
ctl.on_pick(3) # a point in group 0
assert session.selection.count == 50
assert session.selection.indices.max() < 50
def test_shift_adds_ctrl_subtracts(two_clusters, schema):
session = _session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
ctl.on_pick(1)
ctl.on_pick(2, frozenset({"shift"}))
assert session.selection.indices.tolist() == [1, 2]
ctl.on_pick(1, frozenset({"ctrl"}))
assert session.selection.indices.tolist() == [2]
def test_assign_writes_labels_and_recolors(two_clusters, schema):
session = _session(two_clusters, schema)
viewer = FakeViewer()
ctl = InteractionController(session, viewer)
ctl.on_pick(5)
ctl.assign(2)
assert session.cloud.labels[5] == 2
assert viewer.recolored # the touched point was recoloured
# Selection is cleared after assigning.
assert session.selection.is_empty()
def _grouped_session(two_clusters, schema):
"""Session with a 2-group grouping active (group 0 = first 50, group 1 = rest)."""
two_clusters.ensure_labels(schema.unlabeled_id)
session = Session(two_clusters, schema)
gid = np.where(np.arange(two_clusters.n) < 50, 0, 1).astype(np.int32)
session.add_grouping(Grouping(gid, suggested_labels={1: 2}))
return session
def test_select_all_selects_every_point(two_clusters, schema):
session = _session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
ctl.select_all()
assert session.selection.count == two_clusters.n
assert session.selection.indices.tolist() == list(range(two_clusters.n))
def test_select_group_selects_whole_segment(two_clusters, schema):
session = _grouped_session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
ctl.select_group(0)
assert session.selection.count == 50
assert session.selection.indices.max() < 50
def test_assign_group_labels_the_segment(two_clusters, schema):
session = _grouped_session(two_clusters, schema)
viewer = FakeViewer()
ctl = InteractionController(session, viewer)
ctl.set_active_class(1)
touched = ctl.assign_group(0)
assert touched.size == 50
assert (session.cloud.labels[:50] == 1).all()
assert (session.cloud.labels[50:] == 0).all()
assert viewer.recolored # labelled points are repainted in any view
def test_apply_suggested_single_and_all(two_clusters, schema):
session = _grouped_session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
# Only group 1 carries a suggestion (-> class 2).
touched = ctl.apply_suggested(1)
assert touched.size == 50
assert (session.cloud.labels[50:] == 2).all()
# Group 0 has no suggestion -> no-op.
assert ctl.apply_suggested(0).size == 0
# "All suggested" applies every group that has one.
session2 = _grouped_session(two_clusters, schema)
ctl2 = InteractionController(session2, FakeViewer())
assert ctl2.apply_suggested(None).size == 50
def test_group_ops_noop_without_active_grouping(two_clusters, schema):
two_clusters.ensure_labels(schema.unlabeled_id)
session = Session(two_clusters, schema) # no grouping
ctl = InteractionController(session, FakeViewer())
assert ctl.assign_group(0).size == 0
assert ctl.apply_suggested().size == 0
def test_group_visibility_commands(two_clusters, schema):
session = _grouped_session(two_clusters, schema) # group 0 = first 50, group 1 = rest
viewer = FakeViewer()
ctl = InteractionController(session, viewer)
# Hide group 0 -> its points masked off, group 1 still visible.
ctl.set_group_visibility(0, False)
assert viewer.visible_mask is not None
assert viewer.visible_mask[:50].sum() == 0
assert viewer.visible_mask[50:].all()
assert next(s for s in ctl.snapshot().segments if s.id == 0).visible is False
# Hide all -> every group masked off.
ctl.hide_all_groups()
assert viewer.visible_mask is not None and viewer.visible_mask.sum() == 0
assert all(s.visible is False for s in ctl.snapshot().segments)
# Show all -> mask cleared.
ctl.show_all_groups()
assert viewer.visible_mask is None
assert all(s.visible for s in ctl.snapshot().segments)
def test_assign_visible_groups_labels_only_checked(two_clusters, schema):
session = _grouped_session(two_clusters, schema) # group 0 = first 50, group 1 = last 50
ctl = InteractionController(session, FakeViewer())
# Uncheck (hide) group 0, then assign -> only the visible group 1 is labelled.
ctl.set_group_visibility(0, False)
ctl.set_active_class(2)
touched = ctl.assign_visible_groups()
labels = session.cloud.labels
assert touched.size == 50
assert (labels[:50] == schema.unlabeled_id).all() # hidden group untouched
assert (labels[50:] == 2).all() # checked group labelled
# It is a single undoable batch.
ctl.undo()
assert (session.cloud.labels == schema.unlabeled_id).all()
def test_clear_grouping_discards_segmentation_keeps_labels(two_clusters, schema):
session = _grouped_session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
ctl.set_display_mode("grouping")
# Label a segment first — its labels must survive the grouping being dropped.
labelled = ctl.assign_group(0, 1)
assert labelled.size == 50
ctl.clear_grouping()
assert session.active_grouping is None
assert session.groupings == []
assert ctl.display_mode == "labels"
assert (session.cloud.labels[:50] == 1).all() # labels kept
assert ctl.snapshot().segments == []
# Clearing again is a harmless no-op.
ctl.clear_grouping()
assert session.active_grouping is None
def test_class_editing(two_clusters, schema):
from toaster.core import Selection
two_clusters.ensure_labels(schema.unlabeled_id)
session = Session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
# Add a class -> it becomes the active brush.
new_id = ctl.add_class("tree", (1, 2, 3))
assert session.active_class == new_id
assert session.schema.get(new_id).color == (1, 2, 3)
# Rename and recolour an existing class.
ctl.rename_class(1, "floor")
assert session.schema.get(1).name == "floor"
ctl.set_class_color(1, (9, 9, 9))
assert session.schema.get(1).color == (9, 9, 9)
# Label points class 1, then remove class 1 -> they fall back to unlabeled.
session.annotation.assign(Selection.from_indices([0, 1, 2], two_clusters.n), 1)
ctl.remove_class(1)
assert 1 not in [c.id for c in session.schema.classes]
assert (session.cloud.labels[:3] == schema.unlabeled_id).all()
def test_snapshot_is_a_flat_serializable_read_model(two_clusters, schema):
from dataclasses import asdict
session = _grouped_session(two_clusters, schema) # grouping with suggested {1: 2}
ctl = InteractionController(session, FakeViewer())
ctl.set_active_class(1)
snap = ctl.snapshot()
assert [c.id for c in snap.classes] == [0, 1, 2]
assert snap.active_class == 1
assert snap.class_name(1) == "a"
assert snap.active_grouping is not None and snap.active_grouping.n_groups == 2
assert len(snap.segments) == 2
seg1 = next(s for s in snap.segments if s.id == 1)
assert seg1.count == 50
assert seg1.suggested == 2
assert snap.has_suggestions is True
# No numpy / domain objects: a plain dict (i.e. wire-ready) round-trips.
asdict(snap)
def test_run_segmenter_sets_active_grouping(two_clusters, schema):
from toaster.segment import get_segmenter
session = _session(two_clusters, schema)
ctl = InteractionController(session, FakeViewer())
ctl.run_segmenter(get_segmenter("dbscan", eps=0.5, min_samples=5))
assert session.active_grouping is not None
assert session.active_grouping.n_groups == 2
assert ctl.display_mode == "grouping"