diff --git a/tests/data/test_video_utils.py b/tests/data/test_video_utils.py new file mode 100644 index 00000000..f9073396 --- /dev/null +++ b/tests/data/test_video_utils.py @@ -0,0 +1,316 @@ +"""Unit tests for the video / audio preview utilities.""" +import numpy as np +import pytest + +from PIL import Image + +from weightslab.data import video_utils as vu + + +def _clip(frames=8, height=32, width=32, channels=3, dtype=np.uint8): + rng = np.random.default_rng(0) + arr = rng.integers(0, 255, (frames, height, width, channels)) + return arr.astype(dtype) + + +# --------------------------------------------------------------------------- +# Task-type predicates +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("task,expected", [ + ("video_generation", True), + ("image_generation", True), + ("audio_generation", True), + (" Video_Generation ", True), + ("classification", False), + ("segmentation", False), + ("detection_pointcloud", False), + ("", False), + (None, False), +]) +def test_is_generation_task(task, expected): + assert vu.is_generation_task(task) is expected + + +def test_is_video_task_is_narrower_than_is_generation_task(): + assert vu.is_video_task("video_generation") + assert not vu.is_video_task("image_generation") + assert not vu.is_video_task("audio_generation") + + +def test_generation_tasks_do_not_collide_with_point_cloud_routing(): + """The point-cloud router matches substrings "3d"/"point"/"lidar". + + A generative task name containing any of them would silently be sent to the + BEV renderer, so guard the chosen names against that regression. + """ + from weightslab.data.point_cloud_utils import is_point_cloud_task + + for task in (vu.VIDEO_GENERATION_TASK, vu.IMAGE_GENERATION_TASK, + vu.AUDIO_GENERATION_TASK): + assert not is_point_cloud_task(task) + + +# --------------------------------------------------------------------------- +# Shape heuristics / routing +# --------------------------------------------------------------------------- +def test_looks_like_video_accepts_clip_shape(): + assert vu.looks_like_video(_clip()) + + +@pytest.mark.parametrize("arr", [ + None, + np.zeros((32, 32, 3), dtype=np.uint8), # single image + np.zeros((8, 32, 32, 7), dtype=np.uint8), # implausible channel count + np.zeros((1, 32, 32, 3), dtype=np.uint8), # single frame + np.zeros((100, 4), dtype=np.float32), # point cloud +]) +def test_looks_like_video_rejects_non_clips(arr): + assert not vu.looks_like_video(arr) + + +def test_video_routing_requires_explicit_opt_in(): + """A bare 4-D array must NOT be treated as video. + + Volumetric scans have the same rank; without this guard every medical + volume would start rendering as a clip. + """ + class Dataset: + pass + + assert not vu.is_video_sample(Dataset(), _clip()) + + +def test_video_routing_via_task_type_and_media_kind(): + class Bare: + pass + + class Declared: + media_kind = "video" + + class Tasked: + task_type = "video_generation" + + assert vu.is_video_sample(Bare(), _clip(), "video_generation") + assert vu.is_video_sample(Declared(), _clip()) + assert vu.is_video_sample(Tasked(), _clip()) + + +def test_volumetric_data_is_not_hijacked_by_media_kind_image(): + class Volume: + media_kind = "image" + + assert not vu.is_video_sample(Volume(), _clip()) + + +# --------------------------------------------------------------------------- +# Frames / poster +# --------------------------------------------------------------------------- +def test_poster_uses_middle_frame_by_default(): + clip = np.zeros((5, 8, 8, 3), dtype=np.uint8) + clip[2] = 200 # middle frame is the distinctive one + + poster = vu.video_poster_frame(object(), clip) + assert np.asarray(poster).max() == 200 + + +def test_poster_honours_dataset_hook(): + class Dataset: + def render_video_poster(self, frames): + return Image.new("RGB", (11, 7), (1, 2, 3)) + + poster = vu.video_poster_frame(Dataset(), _clip()) + assert poster.size == (11, 7) + + +def test_poster_falls_back_when_dataset_hook_raises(): + class Dataset: + def render_video_poster(self, frames): + raise RuntimeError("boom") + + # A broken user hook must degrade to the built-in poster, never propagate. + poster = vu.video_poster_frame(Dataset(), _clip()) + assert poster.size == (32, 32) + + +@pytest.mark.parametrize("scale,offset,label", [ + (255.0, 0.0, "raw 0-255"), + (1.0, 0.0, "normalized 0-1"), + (2.0, -1.0, "generative -1..1"), +]) +def test_float_frames_are_rescaled_to_visible_range(scale, offset, label): + # Every frame spans the full 0..1 range, so the poster (the middle frame) + # is representative regardless of which frame gets picked. + ramp = np.linspace(0, 1, 8 * 8 * 3).reshape(1, 8, 8, 3) + frames = (np.tile(ramp, (4, 1, 1, 1)) * scale + offset).astype(np.float32) + + poster = np.asarray(vu.video_poster_frame(object(), frames)) + assert poster.dtype == np.uint8 + # A washed-out or all-black poster means the range detection misfired. + assert poster.max() > 200, label + + +def test_low_contrast_frames_are_not_stretched_to_full_range(): + """Rescaling maps by convention, it does not auto-level. + + A dim clip must stay dim: min/max normalizing here would misrepresent a + genuinely low-contrast generation as a healthy one. + """ + frames = np.full((4, 8, 8, 3), 0.5, dtype=np.float32) # mid-grey, 0..1 + poster = np.asarray(vu.video_poster_frame(object(), frames)) + assert 120 <= poster.min() <= 135 + assert 120 <= poster.max() <= 135 + + +def test_video_frame_at_clamps_out_of_range_index(): + clip = _clip(frames=4) + assert vu.video_frame_at(clip, -5).size == (32, 32) + assert vu.video_frame_at(clip, 99).size == (32, 32) + + +def test_video_frame_at_selects_the_requested_frame(): + clip = np.zeros((4, 8, 8, 3), dtype=np.uint8) + clip[3] = 123 + assert np.asarray(vu.video_frame_at(clip, 3)).max() == 123 + + +def test_grayscale_and_rgba_clips_are_previewable(): + assert vu.video_poster_frame(object(), _clip(channels=1)).mode == "L" + assert vu.video_poster_frame(object(), _clip(channels=4)).mode == "RGBA" + + +# --------------------------------------------------------------------------- +# Audio +# --------------------------------------------------------------------------- +def test_encode_audio_wav_roundtrip(): + import io + import wave + + sample_rate = 16_000 + tone = np.sin(2 * np.pi * 440 * np.arange(sample_rate) / sample_rate) + + data = vu.encode_audio_wav(tone, sample_rate) + with wave.open(io.BytesIO(data), "rb") as handle: + assert handle.getframerate() == sample_rate + assert handle.getnchannels() == 1 + assert handle.getsampwidth() == 2 + assert handle.getnframes() == sample_rate + + +def test_encode_audio_wav_handles_stereo_and_channel_first(): + sample_rate = 8_000 + stereo = np.zeros((2, sample_rate), dtype=np.float32) # [C, S] + data = vu.encode_audio_wav(stereo, sample_rate) + + import io + import wave + with wave.open(io.BytesIO(data), "rb") as handle: + # [C, S] must be transposed to [S, C], not read as a 2-sample clip. + assert handle.getnchannels() == 2 + assert handle.getnframes() == sample_rate + + +def test_encode_audio_wav_empty_is_safe(): + assert vu.encode_audio_wav(np.zeros(0), 16_000) == b"" + assert vu.encode_audio_wav(np.zeros(10), 0) == b"" + + +def test_load_sample_audio_without_hook(): + assert vu.load_sample_audio(object(), 0) == (None, 0) + + +def test_load_sample_audio_survives_failing_hook(): + class Dataset: + def get_audio(self, index): + raise RuntimeError("no audio") + + assert vu.load_sample_audio(Dataset(), 0) == (None, 0) + + +# --------------------------------------------------------------------------- +# Descriptors / fps +# --------------------------------------------------------------------------- +def test_get_fps_prefers_dataset_attribute(): + class Dataset: + fps = 24 + + assert vu.get_fps(Dataset()) == 24.0 + + +def test_get_fps_ignores_invalid_values(monkeypatch): + class Dataset: + fps = "not-a-number" + + monkeypatch.delenv("WL_VIDEO_FPS", raising=False) + assert vu.get_fps(Dataset()) == 8.0 + + +def test_describe_clip_reports_duration(): + class Dataset: + num_frames = 24 + fps = 12 + + def get_audio(self, index): + return None + + info = vu.describe_clip(Dataset()) + assert info["frame_count"] == 24 + assert info["fps"] == 12.0 + assert info["has_audio"] is True + assert info["duration_seconds"] == pytest.approx(2.0) + + +def test_describe_clip_omits_unknown_frame_count(): + info = vu.describe_clip(object()) + assert "frame_count" not in info + assert info["has_audio"] is False + + +# --------------------------------------------------------------------------- +# Muxing +# --------------------------------------------------------------------------- +@pytest.mark.skipif(not vu.ffmpeg_available(), reason="ffmpeg not installed") +def test_encode_video_mp4_produces_a_real_container(): + data = vu.encode_video_mp4(_clip(frames=6), fps=8) + assert data[4:8] == b"ftyp" + + +@pytest.mark.skipif(not vu.ffmpeg_available(), reason="ffmpeg not installed") +def test_encode_video_mp4_pads_odd_dimensions(): + """H.264 rejects odd width/height; the encoder must pad, not fail.""" + data = vu.encode_video_mp4(_clip(frames=4, height=31, width=33), fps=8) + assert data[4:8] == b"ftyp" + + +@pytest.mark.skipif(not vu.ffmpeg_available(), reason="ffmpeg not installed") +def test_encode_clip_muxes_audio(): + sample_rate = 16_000 + tone = np.sin(2 * np.pi * 440 * np.arange(sample_rate * 2) / sample_rate) + + data, mime, has_audio = vu.encode_clip( + _clip(frames=16), fps=8, audio=tone, sample_rate=sample_rate) + assert mime == "video/mp4" + assert has_audio is True + # A muxed AAC track makes the file substantially larger than the video-only + # encode of the same frames. + video_only, _, _ = vu.encode_clip(_clip(frames=16), fps=8) + assert len(data) > len(video_only) + + +def test_encode_clip_falls_back_to_gif_without_ffmpeg(monkeypatch): + monkeypatch.setattr(vu, "_ffmpeg_binary", lambda: "") + + data, mime, has_audio = vu.encode_clip(_clip(frames=4), fps=8) + assert mime == "image/gif" + assert has_audio is False + assert data[:6] in (b"GIF87a", b"GIF89a") + + +def test_encode_video_mp4_returns_empty_without_ffmpeg(monkeypatch): + monkeypatch.setattr(vu, "_ffmpeg_binary", lambda: "") + assert vu.encode_video_mp4(_clip(), fps=8) == b"" + + +def test_encode_empty_clip_is_safe(): + empty = np.zeros((0, 8, 8, 3), dtype=np.uint8) + assert vu.encode_video_mp4(empty, fps=8) == b"" + assert vu.encode_video_gif(empty, fps=8) == b"" diff --git a/tests/gRPC/test_get_media.py b/tests/gRPC/test_get_media.py new file mode 100644 index 00000000..ee95e07d --- /dev/null +++ b/tests/gRPC/test_get_media.py @@ -0,0 +1,402 @@ +"""Servicer-level tests for the GetMedia streaming RPC.""" +import numpy as np +import pytest + +import weightslab.proto.experiment_service_pb2 as pb2 + +from weightslab.data import video_utils as vu +from weightslab.trainer.services.data_service import ( + DataService, + _build_media_stats, + _DEFAULT_MEDIA_CHUNK_BYTES, + _media_chunk_bytes, +) + + +SAMPLE_ID = 7 +SAMPLE_RATE = 16_000 + + +class _FakeVideoDataset: + task_type = "video_generation" + fps = 8 + num_frames = 12 + + def __init__(self, frames=12, with_audio=True, height=32, width=32): + rng = np.random.default_rng(0) + self._clip = rng.integers( + 0, 255, (frames, height, width, 3)).astype(np.uint8) + self._with_audio = with_audio + self.num_frames = frames + + def get_index_from_sample_id(self, sample_id): + if int(sample_id) != SAMPLE_ID: + raise KeyError(sample_id) + return 0 + + def __getitem__(self, idx): + return self._clip, f"uid_{SAMPLE_ID:06d}", None, None + + def get_items(self, idx, include_metadata=False, include_labels=False, + include_images=False): + return self._clip, f"uid_{SAMPLE_ID:06d}", None, None + + def get_audio(self, index): + if not self._with_audio: + return None + duration = self._clip.shape[0] / float(self.fps) + t = np.arange(int(SAMPLE_RATE * duration)) / SAMPLE_RATE + return np.sin(2 * np.pi * 440 * t).astype(np.float32), SAMPLE_RATE + + +class _FakeImageDataset: + """A non-video dataset, to prove GetMedia refuses it cleanly.""" + task_type = "classification" + + def get_index_from_sample_id(self, sample_id): + return 0 + + def __getitem__(self, idx): + return np.zeros((32, 32, 3), dtype=np.uint8), "uid", None, None + + def get_items(self, idx, **kwargs): + return np.zeros((32, 32, 3), dtype=np.uint8), "uid", None, None + + +class _StubService: + """Minimal stand-in exposing only what GetMedia touches.""" + _MEDIA_CHUNK_BYTES = DataService._MEDIA_CHUNK_BYTES + _MEDIA_CACHE_ENTRIES = DataService._MEDIA_CACHE_ENTRIES + GetMedia = DataService.GetMedia + _stream_media = DataService._stream_media + _media_cache_put = DataService._media_cache_put + _locate_sample = DataService._locate_sample + + def __init__(self, dataset): + self._dataset = dataset + self._media_cache = {} + + def _get_dataset(self, origin): + return self._dataset if origin == "train_loader" else None + + +def _collect(stub, **kwargs): + kwargs.setdefault("sample_id", str(SAMPLE_ID)) + kwargs.setdefault("origin", "train_loader") + return list(stub.GetMedia(pb2.MediaRequest(**kwargs), context=None)) + + +ffmpeg_required = pytest.mark.skipif( + not vu.ffmpeg_available(), reason="ffmpeg not installed") + + +# --------------------------------------------------------------------------- +# Happy paths +# --------------------------------------------------------------------------- +@ffmpeg_required +def test_get_media_streams_a_playable_clip(): + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub) + + assert all(c.success for c in chunks) + first = chunks[0] + assert first.mime_type == "video/mp4" + assert first.frame_count == 12 + assert first.fps == pytest.approx(8.0) + assert first.has_audio is True + assert (first.width, first.height) == (32, 32) + assert first.duration_seconds == pytest.approx(1.5) + assert first.total_chunks == len(chunks) + + payload = b"".join(c.data for c in chunks) + assert len(payload) == first.total_bytes + assert payload[4:8] == b"ftyp" + + +@ffmpeg_required +def test_only_the_first_chunk_carries_the_header(): + stub = _StubService(_FakeVideoDataset(frames=32, height=128, width=128)) + # Force many chunks so there is a meaningful tail to check. + stub._MEDIA_CHUNK_BYTES = 1024 + chunks = _collect(stub) + + assert len(chunks) > 1 + for index, chunk in enumerate(chunks[1:], start=1): + assert chunk.mime_type == "" + assert chunk.frame_count == 0 + assert chunk.total_bytes == 0 + assert chunk.chunk_index == index + + +def test_get_media_audio_kind_returns_wav(): + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub, kind="audio") + + assert all(c.success for c in chunks) + first = chunks[0] + assert first.mime_type == "audio/wav" + assert first.sample_rate == SAMPLE_RATE + assert first.has_audio is True + assert first.duration_seconds == pytest.approx(1.5, abs=0.01) + + payload = b"".join(c.data for c in chunks) + assert payload[:4] == b"RIFF" + assert payload[8:12] == b"WAVE" + + +@ffmpeg_required +def test_max_frames_subsamples_across_the_whole_clip(): + """Capping must subsample, not truncate — the tail matters for video.""" + stub = _StubService(_FakeVideoDataset(frames=40)) + chunks = _collect(stub, max_frames=10) + + assert chunks[0].frame_count == 10 + # 40 source frames at 8 fps played back as 10 frames at 8 fps. + assert chunks[0].duration_seconds == pytest.approx(10 / 8.0) + + +# --------------------------------------------------------------------------- +# Caching +# --------------------------------------------------------------------------- +@ffmpeg_required +def test_repeated_requests_reuse_the_encoded_clip(): + stub = _StubService(_FakeVideoDataset()) + first = _collect(stub) + assert len(stub._media_cache) == 1 + + # Re-encoding would be observable as a different byte payload only by luck, + # so assert on the cache itself plus payload identity. + second = _collect(stub) + assert len(stub._media_cache) == 1 + assert b"".join(c.data for c in first) == b"".join(c.data for c in second) + + +@ffmpeg_required +def test_media_cache_evicts_oldest_entries(): + stub = _StubService(_FakeVideoDataset()) + for max_frames in range(1, DataService._MEDIA_CACHE_ENTRIES + 4): + _collect(stub, max_frames=max_frames + 1) + + assert len(stub._media_cache) <= DataService._MEDIA_CACHE_ENTRIES + + +@ffmpeg_required +def test_video_and_audio_are_cached_separately(): + stub = _StubService(_FakeVideoDataset()) + _collect(stub) + _collect(stub, kind="audio") + assert len(stub._media_cache) == 2 + + +# --------------------------------------------------------------------------- +# Failure modes — every one must yield a clean chunk, never raise +# --------------------------------------------------------------------------- +def test_unknown_sample_fails_gracefully(): + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub, sample_id="404") + + assert len(chunks) == 1 + assert chunks[0].success is False + assert "not found" in chunks[0].message + + +def test_unknown_origin_fails_gracefully(): + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub, origin="nope_loader") + + assert len(chunks) == 1 + assert chunks[0].success is False + + +def test_non_video_sample_fails_gracefully(): + stub = _StubService(_FakeImageDataset()) + chunks = _collect(stub) + + assert len(chunks) == 1 + assert chunks[0].success is False + assert "not a video" in chunks[0].message + + +def test_audio_request_without_audio_fails_gracefully(): + stub = _StubService(_FakeVideoDataset(with_audio=False)) + chunks = _collect(stub, kind="audio") + + assert len(chunks) == 1 + assert chunks[0].success is False + assert "no audio" in chunks[0].message + + +def test_encode_failure_reports_actionable_message(monkeypatch): + monkeypatch.setattr(vu, "_ffmpeg_binary", lambda: "") + monkeypatch.setattr(vu, "encode_video_gif", lambda *a, **k: b"") + + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub) + + assert chunks[0].success is False + assert "ffmpeg" in chunks[0].message + + +def test_clip_without_ffmpeg_still_streams_as_gif(monkeypatch): + monkeypatch.setattr(vu, "_ffmpeg_binary", lambda: "") + + stub = _StubService(_FakeVideoDataset(frames=4)) + chunks = _collect(stub) + + assert all(c.success for c in chunks) + assert chunks[0].mime_type == "image/gif" + assert chunks[0].has_audio is False + payload = b"".join(c.data for c in chunks) + assert payload[:6] in (b"GIF87a", b"GIF89a") + + +# --------------------------------------------------------------------------- +# Media attached to metadata fields (wl.save_media) — the generated-output path +# --------------------------------------------------------------------------- +def _store_clip(field="pred_video", sample_id=str(SAMPLE_ID), size=9000, **meta): + from weightslab.data import media_store + + payload = { + "frame_count": 16, "fps": 8.0, "has_audio": True, + "width": 64, "height": 64, "duration_seconds": 2.0, + } + payload.update(meta) + media_store.put(field, sample_id, b"\x00" * size, "video/mp4", "video", + poster=b"poster", meta=payload) + return payload + + +def test_field_media_streams_from_the_store(): + from weightslab.data import media_store + media_store.clear() + _store_clip() + + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub, field="pred_video") + + assert all(c.success for c in chunks) + first = chunks[0] + assert first.mime_type == "video/mp4" + assert first.frame_count == 16 + assert first.has_audio is True + assert first.duration_seconds == pytest.approx(2.0) + assert sum(len(c.data) for c in chunks) == 9000 + + +def test_field_media_does_not_touch_the_dataset(): + """A generated clip lives in the store, not in the dataset. + + The field path must therefore work even when the sample is not a video + sample at all — otherwise attaching media to a classification dataset + would be impossible. + """ + from weightslab.data import media_store + media_store.clear() + _store_clip() + + stub = _StubService(_FakeImageDataset()) + chunks = _collect(stub, field="pred_video") + + assert all(c.success for c in chunks) + assert chunks[0].mime_type == "video/mp4" + + +def test_missing_field_media_fails_gracefully(): + from weightslab.data import media_store + media_store.clear() + + stub = _StubService(_FakeVideoDataset()) + chunks = _collect(stub, field="pred_video") + + assert len(chunks) == 1 + assert chunks[0].success is False + assert "pred_video" in chunks[0].message + + +def test_field_media_is_not_confused_with_the_input_clip(): + """field="" and field="pred_video" must resolve to different payloads.""" + from weightslab.data import media_store + media_store.clear() + _store_clip(size=1234) + + stub = _StubService(_FakeVideoDataset()) + attached = b"".join(c.data for c in _collect(stub, field="pred_video")) + assert len(attached) == 1234 + + if vu.ffmpeg_available(): + own = b"".join(c.data for c in _collect(stub)) + assert own != attached + assert own[4:8] == b"ftyp" + + +def test_media_stats_carry_descriptor_and_poster(): + import pandas as pd + from weightslab.data import media_store + media_store.clear() + _store_clip() + + row = pd.Series({ + "sample_id": SAMPLE_ID, + "media:pred_video": '{"kind":"video","fps":8}', + "loss": 0.5, + }) + stats = _build_media_stats(row) + + assert len(stats) == 1 + assert stats[0].name == "media:pred_video" + assert stats[0].type == "media" + assert stats[0].value_string == '{"kind":"video","fps":8}' + assert stats[0].thumbnail == b"poster" + + +def test_media_stats_survive_an_evicted_payload(): + """The column must keep its shape even when the bytes are gone.""" + import pandas as pd + from weightslab.data import media_store + media_store.clear() + + row = pd.Series({ + "sample_id": SAMPLE_ID, + "media:pred_video": '{"kind":"video"}', + }) + stats = _build_media_stats(row) + + assert len(stats) == 1 + assert stats[0].thumbnail == b"" # no still, but the descriptor is intact + + +def test_rows_without_media_emit_no_media_stats(): + import pandas as pd + + row = pd.Series({"sample_id": SAMPLE_ID, "loss": 0.5, "label": "cat"}) + stats = _build_media_stats(row) + assert stats == [] + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- +def test_media_chunk_bytes_defaults(monkeypatch): + monkeypatch.delenv("WL_MEDIA_CHUNK_BYTES", raising=False) + assert _media_chunk_bytes() == _DEFAULT_MEDIA_CHUNK_BYTES + + +@pytest.mark.parametrize("raw", ["0", "-5", "not-an-int", ""]) +def test_media_chunk_bytes_rejects_bad_values(monkeypatch, raw): + monkeypatch.setenv("WL_MEDIA_CHUNK_BYTES", raw) + assert _media_chunk_bytes() == _DEFAULT_MEDIA_CHUNK_BYTES + + +def test_media_chunk_bytes_honours_override(monkeypatch): + monkeypatch.setenv("WL_MEDIA_CHUNK_BYTES", "4096") + assert _media_chunk_bytes() == 4096 + + +@ffmpeg_required +def test_chunk_sizes_respect_the_configured_bound(): + stub = _StubService(_FakeVideoDataset(frames=24, height=96, width=96)) + stub._MEDIA_CHUNK_BYTES = 2048 + chunks = _collect(stub) + + assert all(len(c.data) <= 2048 for c in chunks) + assert sum(len(c.data) for c in chunks) == chunks[0].total_bytes diff --git a/tests/gRPC/test_get_point_cloud.py b/tests/gRPC/test_get_point_cloud.py index e74f6629..49532df4 100644 --- a/tests/gRPC/test_get_point_cloud.py +++ b/tests/gRPC/test_get_point_cloud.py @@ -42,6 +42,9 @@ class _StubService: """Minimal stand-in exposing only what GetPointCloud touches.""" _POINT_CLOUD_CHUNK_BYTES = DataService._POINT_CLOUD_CHUNK_BYTES GetPointCloud = DataService.GetPointCloud + _stream_point_cloud = DataService._stream_point_cloud + # Sample lookup is shared with GetMedia. + _locate_sample = DataService._locate_sample def __init__(self, dataset): self._dataset = dataset @@ -88,6 +91,86 @@ def test_get_point_cloud_unknown_sample_fails_gracefully(): assert "not found" in chunks[0].message +# --------------------------------------------------------------------------- +# Point cloud attached to a metadata field (wl.save_media) — the +# generated/predicted-cloud path, streamed with no dataset lookup at all. +# --------------------------------------------------------------------------- +def _store_cloud(field="pred_cloud", sample_id="7", n_points=200): + from weightslab.data import media_store + from weightslab.data.point_cloud_utils import pack_point_cloud + + rng = np.random.default_rng(1) + points = np.stack([ + rng.uniform(0, 10, n_points), rng.uniform(0, 10, n_points), + rng.uniform(0, 2, n_points), rng.uniform(0, 1, n_points), + ], axis=1).astype(np.float32) + data, num_points, num_features = pack_point_cloud(points) + media_store.put(field, sample_id, data, "application/octet-stream", + "pointcloud", poster=b"poster", + meta={"num_points": num_points, "num_features": num_features, + "pc_range": [0, 0, 0, 10, 10, 2], + "feature_names": ["x", "y", "z", "intensity"]}) + return data, num_points, num_features + + +def test_field_point_cloud_streams_from_the_store(): + from weightslab.data import media_store + media_store.clear() + data, num_points, num_features = _store_cloud() + + stub = _StubService(_FakeLidarDataset()) + chunks = _collect(stub, pb2.PointCloudRequest(sample_id="7", field="pred_cloud")) + + assert all(c.success for c in chunks) + first = chunks[0] + assert first.num_points == num_points + assert first.num_features == num_features + assert list(first.pc_range) == [0, 0, 0, 10, 10, 2] + assert list(first.feature_names) == ["x", "y", "z", "intensity"] + assert b"".join(c.data for c in chunks) == data + + +def test_field_point_cloud_does_not_touch_the_dataset(): + """A generated cloud lives in the store; the sample need not be a cloud at all.""" + from weightslab.data import media_store + media_store.clear() + _store_cloud() + + class _FakeNonCloudDataset: + task_type = "classification" + def get_index_from_sample_id(self, sample_id): return 0 + def __getitem__(self, idx): return np.zeros((8, 8, 3)), "uid", None, None + def get_items(self, idx, **kwargs): return np.zeros((8, 8, 3)), "uid", None, None + + stub = _StubService(_FakeNonCloudDataset()) + chunks = _collect(stub, pb2.PointCloudRequest(sample_id="7", field="pred_cloud")) + assert all(c.success for c in chunks) + + +def test_missing_field_point_cloud_fails_gracefully(): + from weightslab.data import media_store + media_store.clear() + + stub = _StubService(_FakeLidarDataset()) + chunks = _collect(stub, pb2.PointCloudRequest(sample_id="7", field="pred_cloud")) + assert len(chunks) == 1 + assert chunks[0].success is False + assert "pred_cloud" in chunks[0].message + + +def test_field_point_cloud_is_not_confused_with_the_sample_own_cloud(): + from weightslab.data import media_store + media_store.clear() + _store_cloud(n_points=200) + + stub = _StubService(_FakeLidarDataset(n_points=50_000)) + own = _collect(stub, pb2.PointCloudRequest(sample_id="7", origin="train_loader")) + attached = _collect(stub, pb2.PointCloudRequest(sample_id="7", field="pred_cloud")) + + assert own[0].num_points == 50_000 + assert attached[0].num_points == 200 + + def test_point_cloud_chunk_bytes_default(monkeypatch): monkeypatch.delenv("WL_POINT_CLOUD_CHUNK_BYTES", raising=False) assert _point_cloud_chunk_bytes() == _DEFAULT_POINT_CLOUD_CHUNK_BYTES == (1 << 20) diff --git a/weightslab/__init__.py b/weightslab/__init__.py index 8802753e..03ea4b38 100644 --- a/weightslab/__init__.py +++ b/weightslab/__init__.py @@ -36,7 +36,7 @@ # Everything re-exported straight from .src (attribute name == export name). for _name in ( "watch_or_edit", "start_training", "serve", "keep_serving", "save_signals", - "save_instance_signals", "save_group_signals", "tag_samples", + "save_instance_signals", "save_group_signals", "save_media", "tag_samples", "register_categorical_tag", "set_categorical_tag", "discard_samples", "get_samples_by_tag", "get_discarded_samples", "signal", "eval_fn", "compute_signals", "SignalContext", "BatchSignalContext", "StaleSignalError", @@ -191,6 +191,7 @@ def _clean(v: str) -> str: "save_signals", "save_instance_signals", "save_group_signals", + "save_media", "signal", "compute_signals", "set_log_directory", diff --git a/weightslab/data/data_samples_with_ops.py b/weightslab/data/data_samples_with_ops.py index d5ef1f5b..aa2df7ce 100644 --- a/weightslab/data/data_samples_with_ops.py +++ b/weightslab/data/data_samples_with_ops.py @@ -82,6 +82,66 @@ def ALL(cls): return list(map(lambda c: c.value, cls)) +def _partition_metadata_for_auto_render(safe_meta: dict): + """Split one sample's metadata dict into (raw_values, media_values). + + ``media_values`` are the ones auto_render_metadata will attach via the + media store: anything media_store.classify_kind recognizes (an array or a + string), plus any "_audio" key -- always consumed as another + field's audio companion, never rendered as its own column. Everything + else (plain scalars) stays a raw DataFrame value, exactly like today -- + including task_type: a dataset combining several main-sample modalities + behind one loader (see main.py's CombinedDataset) stamps it as a control + string, not content, and it must keep landing in its own reserved column + (SampleStatsEx.TASK_TYPE), never treated as a "text" media field. + """ + from weightslab.data import media_store as _ms + + media_keys = { + k for k, v in safe_meta.items() + if k != SampleStatsEx.TASK_TYPE.value + and (k.endswith("_audio") or _ms.classify_kind(v) is not None) + } + if not media_keys: + return safe_meta, {} + raw = {k: v for k, v in safe_meta.items() if k not in media_keys} + media = {k: v for k, v in safe_meta.items() if k in media_keys} + return raw, media + + +def _auto_render_pending_metadata(origin: str, pending: list, dataset=None) -> None: + """Attach every classified metadata value the same way wl.save_media() + would, once per sample. Run right after register_split so every row + already exists for save_media's DataFrame update to find. + """ + from weightslab.data import media_store as _ms + from weightslab.data.video_utils import get_fps + from weightslab.src import save_media + + fps = get_fps(dataset) + for sid, media_meta in pending: + for field, value in media_meta.items(): + if field.endswith("_audio"): + continue # only ever consumed below, alongside its video field + kind = _ms.classify_kind(value) + if kind is None: + continue + audio, sample_rate = None, 0 + if kind == _ms.KIND_VIDEO: + companion = media_meta.get(f"{field}_audio") + if isinstance(companion, (tuple, list)) and len(companion) == 2: + audio, sample_rate = [companion[0]], int(companion[1] or 0) + try: + save_media( + field, batch_ids=[sid], media=[value], kind=kind, + fps=fps, audio=audio, sample_rate=sample_rate, + origin=origin, dataset=dataset, + ) + except Exception as exc: + logger.warning("auto_render_metadata: %s failed for sample %s: %r", + field, sid, exc) + + class DataSampleTrackingWrapper(Dataset): """Wrapper for PyTorch datasets that tracks per-sample statistics and supports tag-based labeling. @@ -104,6 +164,13 @@ class DataSampleTrackingWrapper(Dataset): preload_metadata: Whether to attempt preloading metadata into the stats dataframe defaults (can speed up access but may increase init time) preload_uids: Whether to attempt preloading unique IDs from metadata instead of generating them (requires metadata to have unique sample_id) keep_leakages: Whether to keep cross-loader duplicates that may cause data leakage (not recommended, use for debugging only) + auto_render_metadata: Auto-classify every preloaded metadata value by shape/type + (array -> image/mask/video/pointcloud, str -> text) and attach it via the + same path as wl.save_media(), once per sample at registration -- no + explicit save_media() call needed. A value named "_audio" that is a + (samples, sample_rate) pair is muxed into ""'s video instead of + becoming its own field. Off by default: every other dataset's metadata + dict keeps landing as plain DataFrame values, unchanged. Examples: @@ -141,6 +208,7 @@ def __init__( preload_metadata: bool = True, preload_uids: bool = False, keep_leakages: bool = False, + auto_render_metadata: bool = False, **_, ): if len(wrapped_dataset) == 0: @@ -291,6 +359,7 @@ def __init__( default_data = [] expanded_uids = [] physical_uids = [] + _pending_auto_render = [] if auto_render_metadata else None # User information logger.info( @@ -382,6 +451,10 @@ def __init__( if preload_metadata and metadata: # Do not overwrite standard managed keys with raw un-casted metadata payload safe_meta = {k: v for k, v in metadata.items() if k not in {SampleStatsEx.GROUP_ID.value, SampleStatsEx.SAMPLE_ID.value, SampleStatsEx.ORIGIN.value}} + if auto_render_metadata: + safe_meta, media_meta = _partition_metadata_for_auto_render(safe_meta) + if media_meta: + _pending_auto_render.append((sid, media_meta)) row.update(safe_meta) for k, v in list(safe_meta.items()): if isinstance(v, dict): @@ -406,6 +479,13 @@ def __init__( use_cache=self.array_use_cache ) + # Auto-render every classified metadata value now that every row + # above actually exists in the ledger -- save_media's DataFrame + # update requires the row to already be registered, so this MUST + # run after register_split, not inside the loop that built it. + if auto_render_metadata and _pending_auto_render: + _auto_render_pending_metadata(self._dataset_split, _pending_auto_render, wrapped_dataset) + # Log tag-based labeling configuration if enabled if self._use_tags: with self._df_lock: diff --git a/weightslab/data/data_utils.py b/weightslab/data/data_utils.py index 3f5aac51..661d6c7b 100644 --- a/weightslab/data/data_utils.py +++ b/weightslab/data/data_utils.py @@ -386,6 +386,12 @@ def _get_image_array_and_metadata(wrapped, index, rank: int = 0) -> tuple: if hasattr(np_img, 'numpy'): np_img = np_img.numpy() + # A text-generation sample's "image" slot is the prompt itself (a plain + # str) — bail out before any array-shaped logic below, which assumes + # .ndim/.shape exist. + if isinstance(np_img, str): + return np_img, False, None + is_volumetric = np_img.ndim >= 4 # 3 is for RGB; while 4 is 3D # TODO (GP): Should be fix because this will not work with grayscale image wo. color channel # For 4D volumetric data, detect and transpose channel-first formats: @@ -538,7 +544,7 @@ def load_metadata(dataset, sample_id): return None -def load_raw_image(dataset, index, slice_idx: int = None, rank: int = 0) -> Image.Image: +def load_raw_image(dataset, index, slice_idx: int = None, rank: int = 0) -> Image.Image | None: """Load raw image from dataset at given index. For 4D volumetric data (Z, H, W, C) or (Z, H, W), extracts a single slice. @@ -551,7 +557,8 @@ def load_raw_image(dataset, index, slice_idx: int = None, rank: int = 0) -> Imag rank: Member rank for grouped data Returns: - PIL Image of the 2D slice + PIL Image of the 2D slice, or None for a text-generation sample + (its "image" slot is a plain string — there is no image to load). """ # Get dataset wrapper if exists wrapped = getattr(dataset, "wrapped_dataset", dataset) @@ -567,6 +574,14 @@ def load_raw_image(dataset, index, slice_idx: int = None, rank: int = 0) -> Imag elif hasattr(wrapped, '__getitem__') or hasattr(wrapped, "data") or hasattr(wrapped, "dataset"): np_img, is_volumetric, original_shape = _get_image_array_and_metadata(wrapped, index, rank=rank) + # A text-generation sample's "image" slot is a plain str, not pixels — + # there is nothing to build a PIL preview from. None is this + # function's existing "nothing to show" convention (its only caller, + # data_service.py's natural-sort default-signal computation, already + # treats a None return as "skip this sample" rather than an error). + if isinstance(np_img, str): + return None + # Handle 4D volumetric data if is_volumetric: np_img = _extract_slice_from_4d(np_img, slice_idx=slice_idx) @@ -619,6 +634,15 @@ def load_raw_image_array(dataset, index, rank: int = 0) -> tuple: if hasattr(wrapped, '__getitem__'): np_img, is_volumetric, original_shape = _get_image_array_and_metadata(wrapped, index, rank=rank) + # Text-generation samples (a prompt, an RLHF conversation, ...) have no + # pixels at all — the text itself is the payload. Every check below + # this point assumes an array (.ndim/.shape), so this must come first. + # The caller (load_raw_image_array's callers in data_service.py) + # detects this case via isinstance(..., str) and skips image encoding + # entirely in favor of a plain text stat. + if isinstance(np_img, str): + return np_img, False, original_shape, None + # Tabular samples (1-D feature vectors) have no spatial dims and cannot be # PIL-encoded as an image. Render a small heatmap for display continuity; # the caller transmits the actual values via build_tabular_raw_data_stat. @@ -649,6 +673,17 @@ def load_raw_image_array(dataset, index, rank: int = 0) -> tuple: thumb_arr = np.asarray(thumb_pil) return thumb_arr, False, tuple(thumb_arr.shape), thumb_pil + # Video samples are [T, H, W, C] and share their rank with volumetric + # scans, so the dataset must opt in (task_type or media_kind) before we + # treat them as clips. Preview them with a single poster frame; the + # playable clip is streamed separately by the GetMedia RPC, which keeps + # the grid and list cheap regardless of clip length. + from weightslab.data.video_utils import is_video_sample, video_poster_frame + if is_video_sample(dataset, np_img, _raw_task): + poster_pil = video_poster_frame(dataset, np_img) + poster_arr = np.asarray(poster_pil) + return poster_arr, False, tuple(poster_arr.shape), poster_pil + # Extract middle slice for thumbnail if is_volumetric: middle_slice = _extract_slice_from_4d(np_img, slice_idx=None) @@ -696,6 +731,44 @@ def load_raw_point_cloud(dataset, index, rank: int = 0): return np.asarray(arr, dtype=np.float32) +def load_raw_video(dataset, index, rank: int = 0): + """Load the raw clip array [T, H, W, C] for one sample (no poster render). + + Same plucking rules as ``load_raw_image_array`` but returns every frame so + the GetMedia RPC can mux and stream a playable clip. Returns None when the + sample is not a video. + """ + from weightslab.data.video_utils import is_video_sample + + wrapped = getattr(dataset, "wrapped_dataset", dataset) + if not hasattr(wrapped, '__getitem__'): + return None + arr, _, _ = _get_image_array_and_metadata(wrapped, index, rank=rank) + if arr is None: + return None + task_type = getattr(wrapped, "task_type", getattr(dataset, "task_type", None)) + if not task_type: + # Some datasets vary task_type PER SAMPLE via the metadata dict their + # own __getitem__ returns (the tracked loader's documented "more than + # two elements" contract) instead of one fixed dataset-level + # attribute -- e.g. a dataset combining several main-sample + # modalities behind one loader. Peek at it before falling back to the + # shape-only heuristic in is_video_sample, which cannot tell a video + # clip apart from volumetric image data on shape alone. + try: + raw_item = wrapped[index] + if isinstance(raw_item, tuple) and len(raw_item) > 3: + for m in raw_item[3:]: + if isinstance(m, dict) and m.get("task_type"): + task_type = m["task_type"] + break + except Exception: + pass + if not is_video_sample(dataset, arr, task_type): + return None + return np.asarray(arr) + + def load_uid(dataset, sample_id): """Load uid from dataset at given index. diff --git a/weightslab/data/media_store.py b/weightslab/data/media_store.py new file mode 100644 index 00000000..5afad7f4 --- /dev/null +++ b/weightslab/data/media_store.py @@ -0,0 +1,219 @@ +"""Process-local store for per-sample media attached to metadata columns. + +Any metadata field can carry media — an image, a mask, a video clip, or a point +cloud — not just the sample's own input. That is what makes a *generated* clip +viewable next to the clip it was trained against: the training script attaches +it with ``wl.save_media(field="pred_video", ...)`` and the studio then shows a +thumbnail in the list and opens the full viewer on click. + +Split of responsibilities: + + * the dataframe column ``media:`` holds only a small JSON **descriptor** + (kind, mime, geometry, fps, duration) — cheap to query, sort and ship in a + metadata response. + * the encoded **bytes** live here, keyed by ``(field, sample_id)``, and are + streamed on demand by the GetMedia RPC. + +Bytes are held in memory rather than on disk on purpose: attached media is +regenerated every epoch, so persisting it would mean writing (and garbage +collecting) thousands of files nobody asks for. The store is capped by total +bytes and evicts least-recently-used entries, so a long run cannot grow without +bound — attaching media to more samples than the cap holds simply means the +oldest thumbnails stop being openable, which degrades gracefully. +""" +import json +import logging +import os +import threading +from collections import OrderedDict + +import numpy as np + +logger = logging.getLogger(__name__) + +# Dataframe column prefix. Mirrors the existing "tag:" convention. +MEDIA_COLUMN_PREFIX = "media:" + +# Media kinds a descriptor may declare. +KIND_IMAGE = "image" +KIND_MASK = "mask" +KIND_VIDEO = "video" +KIND_AUDIO = "audio" +KIND_POINTCLOUD = "pointcloud" +KIND_TEXT = "text" + +_VALID_KINDS = (KIND_IMAGE, KIND_MASK, KIND_VIDEO, KIND_AUDIO, KIND_POINTCLOUD, KIND_TEXT) + +_DEFAULT_MAX_BYTES = 256 * 1024 * 1024 # 256 MiB + +_lock = threading.RLock() +# (field, sample_id) -> entry dict. Insertion-ordered => LRU. +_entries: "OrderedDict[tuple, dict]" = OrderedDict() +_total_bytes = 0 + + +def media_column(field: str) -> str: + """Dataframe column name for a media field.""" + return f"{MEDIA_COLUMN_PREFIX}{field}" + + +def field_from_column(column: str) -> str: + """Inverse of :func:`media_column` ("" when not a media column).""" + text = str(column or "") + if text.startswith(MEDIA_COLUMN_PREFIX): + return text[len(MEDIA_COLUMN_PREFIX):] + return "" + + +def is_media_column(column: str) -> bool: + return str(column or "").startswith(MEDIA_COLUMN_PREFIX) + + +def classify_kind(value): + """Best-effort auto-detection of a metadata value's media kind, from its + shape/type alone — no field-name convention, no explicit kind= needed. + + Returns one of KIND_TEXT / KIND_VIDEO / KIND_POINTCLOUD / KIND_MASK / + KIND_IMAGE, or None for a plain scalar (already fine as a raw DataFrame + value, no encoding needed). Used by DataSampleTrackingWrapper's + auto_render_metadata option, so a dataset's own __getitem__ metadata dict + can carry raw content directly, with no wl.save_media() call required. + """ + if value is None or isinstance(value, (bool, int, float)): + return None + if isinstance(value, str): + return KIND_TEXT + + if hasattr(value, "detach"): + value = value.detach().cpu() + if hasattr(value, "numpy"): + value = value.numpy() + try: + arr = np.asarray(value) + except Exception: + return None + if arr.ndim == 0 or arr.size == 0: + return None + + from weightslab.data.point_cloud_utils import looks_like_point_cloud + from weightslab.data.video_utils import looks_like_video + + if looks_like_video(arr): + return KIND_VIDEO + if looks_like_point_cloud(arr): + return KIND_POINTCLOUD + if arr.ndim == 2 and np.issubdtype(arr.dtype, np.integer): + span = int(arr.max()) - int(arr.min()) + if span < 64: + return KIND_MASK + return KIND_IMAGE + + +def max_bytes() -> int: + """Byte cap for the store (env WL_MEDIA_STORE_MAX_BYTES).""" + raw = os.environ.get("WL_MEDIA_STORE_MAX_BYTES") + if not raw: + return _DEFAULT_MAX_BYTES + try: + value = int(raw) + except (TypeError, ValueError): + logger.warning( + "WL_MEDIA_STORE_MAX_BYTES=%r is not an integer — using default %d", + raw, _DEFAULT_MAX_BYTES) + return _DEFAULT_MAX_BYTES + return value if value > 0 else _DEFAULT_MAX_BYTES + + +def put(field: str, sample_id, data: bytes, mime: str, kind: str, + poster: bytes = b"", meta: dict = None) -> dict: + """Store one sample's media and return its descriptor dict. + + ``poster`` is the still shown in the grid/list; for an image kind it may be + the image itself. ``meta`` carries kind-specific extras (fps, frame_count, + has_audio, num_points, ...) and is merged into the descriptor. + """ + global _total_bytes + + if kind not in _VALID_KINDS: + raise ValueError(f"kind must be one of {_VALID_KINDS}, got {kind!r}") + + entry = { + "field": str(field), + "sample_id": str(sample_id), + "data": data or b"", + "mime": str(mime or ""), + "kind": str(kind), + "poster": poster or b"", + "meta": dict(meta or {}), + } + key = (str(field), str(sample_id)) + size = len(entry["data"]) + len(entry["poster"]) + + with _lock: + previous = _entries.pop(key, None) + if previous is not None: + _total_bytes -= len(previous["data"]) + len(previous["poster"]) + _entries[key] = entry + _total_bytes += size + + cap = max_bytes() + while _total_bytes > cap and len(_entries) > 1: + _, evicted = _entries.popitem(last=False) + _total_bytes -= len(evicted["data"]) + len(evicted["poster"]) + + return descriptor(entry) + + +def descriptor(entry: dict) -> dict: + """The small JSON-able summary that goes into the dataframe column.""" + out = { + "kind": entry["kind"], + "mime": entry["mime"], + "bytes": len(entry.get("data") or b""), + } + out.update(entry.get("meta") or {}) + return out + + +def descriptor_json(entry: dict) -> str: + try: + return json.dumps(descriptor(entry)) + except (TypeError, ValueError): + return json.dumps({"kind": entry.get("kind", "")}) + + +def get(field: str, sample_id) -> dict: + """Fetch an entry, refreshing its LRU position. None when absent.""" + key = (str(field), str(sample_id)) + with _lock: + entry = _entries.pop(key, None) + if entry is None: + return None + _entries[key] = entry # refresh recency + return entry + + +def get_poster(field: str, sample_id) -> bytes: + """Poster bytes for a stored entry (b"" when absent).""" + entry = get(field, sample_id) + return (entry or {}).get("poster") or b"" + + +def fields() -> list: + """Every media field currently holding at least one entry.""" + with _lock: + return sorted({key[0] for key in _entries}) + + +def stats() -> dict: + """Store occupancy, for diagnostics and tests.""" + with _lock: + return {"entries": len(_entries), "bytes": _total_bytes, "cap": max_bytes()} + + +def clear() -> None: + """Drop everything (used between runs and by tests).""" + global _total_bytes + with _lock: + _entries.clear() + _total_bytes = 0 diff --git a/weightslab/data/video_utils.py b/weightslab/data/video_utils.py new file mode 100644 index 00000000..b17d87cb --- /dev/null +++ b/weightslab/data/video_utils.py @@ -0,0 +1,475 @@ +"""Video / audio preview utilities (generative media support). + +Video samples (task_type "video_generation", or any dataset declaring +``media_kind = "video"``) cannot be PIL-encoded as a single image, so the +studio pipeline previews them the same way point clouds are previewed — with a +server-rendered still — and ships the playable media over a separate stream: + + * thumbnails / preview cache / grid / list -> ``video_poster_frame`` + (one representative frame, encoded by the normal WebP path). The grid and + list therefore stay cheap: they never carry video bytes. + * the modal player fetches the muxed clip via the GetMedia RPC + (``encode_video_mp4`` does frame + audio muxing), and can page through + individual frames via ``video_frame_at``. + * audio-only samples ("audio_generation") stream a WAV/AAC track through the + same RPC with ``kind = "audio"``. + +Image-generation samples (task_type "image_generation") need none of this — +they are ordinary 2D arrays and flow through the existing image path. They are +covered here only by ``is_generation_task`` so the label/prediction gates in +data_service treat generated output as free-form rather than as a mask. + +Datasets can replace or supply any part of this by defining attributes: + + * ``media_kind`` -> "video" | "audio" | "image" (disambiguates 4-D arrays, + which would otherwise be indistinguishable from volumetric medical data) + * ``fps`` / ``frame_rate`` -> float, defaults to WL_VIDEO_FPS (8) + * ``get_audio(index) -> (np.ndarray [S] or [S, C], sample_rate)`` + * ``render_video_poster(frames) -> PIL.Image | np.ndarray`` + +The frame array layout accepted everywhere below is ``[T, H, W, C]`` uint8 or +float. ``_get_image_array_and_metadata`` in data_utils already normalizes the +common PyTorch ``[T, C, H, W]`` layout into that form. +""" +import io +import logging +import os +import shutil +import subprocess +import tempfile +import wave + +import numpy as np + +from PIL import Image + +logger = logging.getLogger(__name__) + +# Canonical task types for generative media. TEXT_GENERATION_TASK lives here +# too (not just video/audio/image) because this module is the shared registry +# data_service.py checks before treating a target/prediction as a mask — text +# generation needs that exclusion exactly like the others, even though its +# own preview path (a plain string, no pixels) is handled in data_utils.py. +VIDEO_GENERATION_TASK = "video_generation" +IMAGE_GENERATION_TASK = "image_generation" +AUDIO_GENERATION_TASK = "audio_generation" +TEXT_GENERATION_TASK = "text_generation" + +_GENERATION_TASKS = ( + VIDEO_GENERATION_TASK, + IMAGE_GENERATION_TASK, + AUDIO_GENERATION_TASK, + TEXT_GENERATION_TASK, +) + +# Dataset attribute names. +MEDIA_KIND_ATTRS = ("media_kind", "wl_media_kind") +FPS_ATTRS = ("fps", "frame_rate") +AUDIO_HOOK = "get_audio" +POSTER_HOOK = "render_video_poster" + +# Global decorator registry (set via wl.video_poster). +_REGISTERED_POSTER_FN = None + +# A clip must have at least this many frames before a 4-D array is even +# considered video. Below it, the array is far more likely to be a volumetric +# stack that the existing middle-slice path already handles well. +MIN_VIDEO_FRAMES = 2 +# Channel counts that make sense as displayable frames. +_VIDEO_CHANNELS = (1, 3, 4) + + +def register_poster_fn(fn): + """Register a global video poster renderer (see wl.video_poster).""" + global _REGISTERED_POSTER_FN + _REGISTERED_POSTER_FN = fn + return fn + + +# --------------------------------------------------------------------------- +# Task-type / media-kind predicates +# --------------------------------------------------------------------------- +def is_generation_task(task_type) -> bool: + """True for any generative task type (video, image, or audio generation).""" + return str(task_type or "").strip().lower() in _GENERATION_TASKS + + +def is_video_task(task_type) -> bool: + """True when the task type denotes video output.""" + return str(task_type or "").strip().lower() == VIDEO_GENERATION_TASK + + +def is_audio_task(task_type) -> bool: + """True when the task type denotes audio-only output.""" + return str(task_type or "").strip().lower() == AUDIO_GENERATION_TASK + + +def get_media_kind(dataset) -> str: + """Return the dataset's declared media kind ("" when it declares none).""" + wrapped = getattr(dataset, "wrapped_dataset", dataset) + for source in (wrapped, dataset): + for attr in MEDIA_KIND_ATTRS: + value = getattr(source, attr, None) + if value: + return str(value).strip().lower() + return "" + + +def get_fps(dataset, default: float = None) -> float: + """Frame rate for a video dataset (dataset attr -> WL_VIDEO_FPS -> 8).""" + wrapped = getattr(dataset, "wrapped_dataset", dataset) + for source in (wrapped, dataset): + for attr in FPS_ATTRS: + value = getattr(source, attr, None) + if value: + try: + parsed = float(value) + except (TypeError, ValueError): + continue + if parsed > 0: + return parsed + if default is not None: + return default + try: + return max(1.0, float(os.environ.get("WL_VIDEO_FPS", "8"))) + except (TypeError, ValueError): + return 8.0 + + +def looks_like_video(arr) -> bool: + """Shape heuristic: does this array look like a [T, H, W, C] clip? + + Deliberately conservative — 4-D arrays are ambiguous (volumetric scans have + the same rank), so callers must additionally confirm via the task type or + the dataset's ``media_kind``. This function only rules out shapes that + could not be a clip at all. + """ + if arr is None: + return False + shape = getattr(arr, "shape", None) + if shape is None or len(shape) != 4: + return False + frames, height, width, channels = shape + return ( + frames >= MIN_VIDEO_FRAMES + and height >= 2 + and width >= 2 + and channels in _VIDEO_CHANNELS + ) + + +def is_video_sample(dataset, arr, task_type=None) -> bool: + """Full routing decision for the preview pipeline. + + A sample is treated as video when it has a clip-shaped array AND the + dataset opts in — either through the task type or an explicit + ``media_kind`` attribute. Without the opt-in, volumetric data keeps its + existing middle-slice behaviour. + """ + if not looks_like_video(arr): + return False + if is_video_task(task_type): + return True + if get_media_kind(dataset) == "video": + return True + raw_task = getattr( + getattr(dataset, "wrapped_dataset", dataset), "task_type", + getattr(dataset, "task_type", None)) + return is_video_task(raw_task) + + +# --------------------------------------------------------------------------- +# Frame handling +# --------------------------------------------------------------------------- +def _frames_to_uint8(frames) -> np.ndarray: + """Normalize a [T, H, W, C] array to contiguous uint8 RGB/L frames.""" + arr = np.asarray(frames) + if arr.dtype != np.uint8: + arr = arr.astype(np.float32) + finite = arr[np.isfinite(arr)] + peak = float(finite.max()) if finite.size else 0.0 + low = float(finite.min()) if finite.size else 0.0 + # [-1, 1] generative output is the common case; then [0, 1]; then raw. + if low < -0.01: + arr = (arr + 1.0) * 127.5 + elif peak <= 1.0 + 1e-6: + arr = arr * 255.0 + arr = np.clip(arr, 0, 255) + arr = arr.astype(np.uint8) + return np.ascontiguousarray(arr) + + +def video_frame_at(frames, index: int) -> Image.Image: + """Return frame ``index`` of a clip as a PIL image (clamped to range).""" + arr = _frames_to_uint8(frames) + if arr.shape[0] == 0: + raise ValueError("Cannot extract a frame from an empty clip.") + idx = max(0, min(int(index), arr.shape[0] - 1)) + return _frame_to_pil(arr[idx]) + + +def _frame_to_pil(frame: np.ndarray) -> Image.Image: + """One [H, W, C] uint8 frame -> PIL image.""" + if frame.ndim == 2: + return Image.fromarray(frame, mode="L") + channels = frame.shape[-1] + if channels == 1: + return Image.fromarray(frame[..., 0], mode="L") + if channels == 4: + return Image.fromarray(frame, mode="RGBA") + if channels == 3: + return Image.fromarray(frame, mode="RGB") + return Image.fromarray(frame[..., 0], mode="L") + + +def video_poster_frame(dataset, frames) -> Image.Image: + """Render the still shown in the grid, list, and as the player poster. + + Resolution order mirrors the point-cloud thumbnail hook chain: + dataset ``render_video_poster`` -> globally registered fn -> middle frame. + The middle frame is used rather than the first because generative clips + frequently start from noise or a black fade-in. + """ + arr = _frames_to_uint8(frames) + + wrapped = getattr(dataset, "wrapped_dataset", dataset) + for source in (wrapped, dataset): + hook = getattr(source, POSTER_HOOK, None) + if callable(hook): + try: + rendered = hook(arr) + if rendered is not None: + return rendered if isinstance(rendered, Image.Image) \ + else _frame_to_pil(_frames_to_uint8(rendered[None])[0]) + except Exception as exc: + logger.warning("Dataset %s hook failed: %r", POSTER_HOOK, exc) + + if _REGISTERED_POSTER_FN is not None: + try: + rendered = _REGISTERED_POSTER_FN(arr) + if rendered is not None: + return rendered if isinstance(rendered, Image.Image) \ + else _frame_to_pil(_frames_to_uint8(rendered[None])[0]) + except Exception as exc: + logger.warning("Registered video poster fn failed: %r", exc) + + return _frame_to_pil(arr[arr.shape[0] // 2]) + + +# --------------------------------------------------------------------------- +# Audio handling +# --------------------------------------------------------------------------- +def load_sample_audio(dataset, index): + """Fetch ``(samples, sample_rate)`` for one sample, or ``(None, 0)``. + + The dataset opts in by defining ``get_audio(index)``. Returned samples may + be [S] mono or [S, C]; float in [-1, 1] or int16. + """ + wrapped = getattr(dataset, "wrapped_dataset", dataset) + for source in (wrapped, dataset): + hook = getattr(source, AUDIO_HOOK, None) + if callable(hook): + try: + result = hook(index) + except Exception as exc: + logger.warning("Dataset %s(%s) failed: %r", AUDIO_HOOK, index, exc) + return None, 0 + if result is None: + return None, 0 + samples, sample_rate = result + if samples is None: + return None, 0 + return np.asarray(samples), int(sample_rate or 0) + return None, 0 + + +def _audio_to_int16(samples: np.ndarray) -> np.ndarray: + """Normalize audio to interleaved int16, shape [S, C].""" + arr = np.asarray(samples) + if arr.ndim == 1: + arr = arr[:, None] + elif arr.ndim == 2 and arr.shape[0] < arr.shape[1]: + # [C, S] channel-first -> [S, C]; safe because real clips have far more + # samples than channels. + arr = arr.T + if arr.dtype != np.int16: + arr = np.asarray(arr, dtype=np.float32) + peak = float(np.abs(arr).max()) if arr.size else 0.0 + if peak > 1.0: + arr = arr / peak + arr = np.clip(arr, -1.0, 1.0) * 32767.0 + arr = arr.astype(np.int16) + return np.ascontiguousarray(arr) + + +def encode_audio_wav(samples, sample_rate: int) -> bytes: + """Encode audio to a WAV container (browser-playable, no ffmpeg needed).""" + pcm = _audio_to_int16(samples) + if pcm.size == 0 or sample_rate <= 0: + return b"" + buffer = io.BytesIO() + with wave.open(buffer, "wb") as handle: + handle.setnchannels(int(pcm.shape[1])) + handle.setsampwidth(2) + handle.setframerate(int(sample_rate)) + handle.writeframes(pcm.tobytes()) + return buffer.getvalue() + + +# --------------------------------------------------------------------------- +# Muxing +# --------------------------------------------------------------------------- +def _ffmpeg_binary() -> str: + """Locate an ffmpeg executable, preferring the pip-installable bundle. + + Returns "" when none is available; callers degrade to an animated preview + rather than failing the request. + """ + override = os.environ.get("WL_FFMPEG_BINARY", "").strip() + if override: + return override + try: + import imageio_ffmpeg + return imageio_ffmpeg.get_ffmpeg_exe() + except Exception: + pass + return shutil.which("ffmpeg") or "" + + +def ffmpeg_available() -> bool: + """True when video muxing to MP4 is possible in this process.""" + return bool(_ffmpeg_binary()) + + +def encode_video_mp4(frames, fps: float, audio=None, sample_rate: int = 0, + crf: int = None) -> bytes: + """Mux ``frames`` (+ optional ``audio``) into a browser-playable MP4. + + H.264 + AAC in an MP4 with the moov atom moved to the front, so the studio + can play it straight from a Blob URL without waiting for the whole file. + Returns b"" when ffmpeg is unavailable — callers fall back to ``encode_video_gif``. + """ + binary = _ffmpeg_binary() + if not binary: + return b"" + + arr = _frames_to_uint8(frames) + if arr.shape[0] == 0: + return b"" + # H.264 requires even dimensions; pad rather than crop so nothing is lost. + if arr.shape[1] % 2 or arr.shape[2] % 2: + arr = np.pad( + arr, + ((0, 0), (0, arr.shape[1] % 2), (0, arr.shape[2] % 2), (0, 0)), + mode="edge") + if arr.shape[-1] == 1: + arr = np.repeat(arr, 3, axis=-1) + elif arr.shape[-1] == 4: + arr = arr[..., :3] + + frame_count, height, width, _ = arr.shape + if crf is None: + try: + crf = int(os.environ.get("WL_VIDEO_CRF", "23")) + except (TypeError, ValueError): + crf = 23 + + wav_bytes = b"" + if audio is not None and sample_rate: + wav_bytes = encode_audio_wav(audio, sample_rate) + + tmp_dir = tempfile.mkdtemp(prefix="wl_video_") + out_path = os.path.join(tmp_dir, "clip.mp4") + audio_path = os.path.join(tmp_dir, "audio.wav") + try: + cmd = [ + binary, "-hide_banner", "-loglevel", "error", "-y", + "-f", "rawvideo", "-pix_fmt", "rgb24", + "-s", f"{width}x{height}", "-r", f"{max(1.0, float(fps))}", + "-i", "-", + ] + if wav_bytes: + with open(audio_path, "wb") as handle: + handle.write(wav_bytes) + cmd += ["-i", audio_path, "-c:a", "aac", "-b:a", "128k", "-shortest"] + cmd += [ + "-c:v", "libx264", "-preset", "veryfast", "-crf", str(crf), + # yuv420p + even dims is what browsers actually decode. + "-pix_fmt", "yuv420p", + "-movflags", "+faststart", + out_path, + ] + proc = subprocess.run( + cmd, input=arr.tobytes(), stdout=subprocess.PIPE, + stderr=subprocess.PIPE, check=False) + if proc.returncode != 0 or not os.path.exists(out_path): + logger.warning( + "ffmpeg failed to encode %d frames: %s", + frame_count, proc.stderr.decode("utf-8", "replace")[:400]) + return b"" + with open(out_path, "rb") as handle: + return handle.read() + except Exception as exc: + logger.warning("MP4 encoding failed: %r", exc) + return b"" + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + +def encode_video_gif(frames, fps: float) -> bytes: + """Animated-GIF fallback used when ffmpeg is not installed. + + Silent and much larger than H.264, but it keeps the modal player useful on + a bare ``pip install weightslab`` with no system ffmpeg. + """ + arr = _frames_to_uint8(frames) + if arr.shape[0] == 0: + return b"" + images = [_frame_to_pil(frame).convert("P", palette=Image.ADAPTIVE) + for frame in arr] + buffer = io.BytesIO() + images[0].save( + buffer, format="GIF", save_all=True, append_images=images[1:], + duration=max(1, int(1000.0 / max(1.0, float(fps)))), loop=0) + return buffer.getvalue() + + +def describe_clip(dataset) -> dict: + """Cheap clip descriptor for the grid, read from dataset attributes only. + + Called once per rendered row, so it must never decode a clip. Datasets that + do not advertise ``num_frames`` simply omit it and the studio falls back to + whatever the GetMedia stream reports when the player opens. + """ + wrapped = getattr(dataset, "wrapped_dataset", dataset) + info = {"kind": "video", "fps": get_fps(dataset)} + + for attr in ("num_frames", "clip_length", "frames_per_clip"): + for source in (wrapped, dataset): + value = getattr(source, attr, None) + if value: + try: + info["frame_count"] = int(value) + except (TypeError, ValueError): + continue + break + if "frame_count" in info: + break + + info["has_audio"] = any( + callable(getattr(source, AUDIO_HOOK, None)) for source in (wrapped, dataset)) + if info.get("frame_count"): + info["duration_seconds"] = round( + info["frame_count"] / max(1.0, info["fps"]), 3) + return info + + +def encode_clip(frames, fps: float, audio=None, sample_rate: int = 0): + """Encode a clip for transport, returning ``(bytes, mime, has_audio)``.""" + mp4 = encode_video_mp4(frames, fps, audio=audio, sample_rate=sample_rate) + if mp4: + return mp4, "video/mp4", bool(audio is not None and sample_rate) + gif = encode_video_gif(frames, fps) + if gif: + return gif, "image/gif", False + return b"", "", False diff --git a/weightslab/proto/experiment_service.proto b/weightslab/proto/experiment_service.proto index e926ab55..abedc1d4 100644 --- a/weightslab/proto/experiment_service.proto +++ b/weightslab/proto/experiment_service.proto @@ -33,6 +33,12 @@ service ExperimentService { // Raw point cloud of one sample (task_type "detection_pointcloud"), server-streamed // in binary chunks for the interactive 3D viewer. rpc GetPointCloud (PointCloudRequest) returns (stream PointCloudChunk); + // Playable media of one sample (task_type "video_generation" / + // "audio_generation"), server-streamed in binary chunks for the modal + // player. Grid and list previews never call this — they use the poster + // frame carried by the ordinary raw_data stat — so clip length does not + // affect browsing cost. + rpc GetMedia (MediaRequest) returns (stream MediaChunk); rpc EditDataSample (DataEditsRequest) returns (DataEditsResponse); rpc GetDataSplits (Empty) returns (DataSplitsResponse); rpc CheckAgentHealth (Empty) returns (AgentHealthResponse); @@ -514,6 +520,10 @@ message PointCloudRequest { string sample_id = 1; string origin = 2; // loader name (e.g. "train_loader") int32 max_points = 3; // 0 = use the server default cap + // Metadata field to stream, i.e. the "media:" column attached by + // wl.save_media (e.g. "pred_cloud"). Empty means the sample's own point + // cloud (task_type "detection_pointcloud"), which is the existing behavior. + string field = 4; } message PointCloudChunk { @@ -533,6 +543,39 @@ message PointCloudChunk { repeated string feature_names = 9; } +message MediaRequest { + string sample_id = 1; + string origin = 2; // loader name (e.g. "train_loader") + // "video" (default) streams the muxed clip; "audio" streams only the + // soundtrack, which the frame explorer uses to draw a waveform without + // re-downloading the video. + string kind = 3; + int32 max_frames = 4; // 0 = use the server default cap + // Metadata field to stream, i.e. the "media:" column attached by + // wl.save_media (e.g. "pred_video"). Empty means the sample's own input + // clip, which is what the poster in the grid shows. + string field = 5; +} + +message MediaChunk { + bool success = 1; + string message = 2; + // Set on the first chunk only: + string mime_type = 3; // "video/mp4" | "image/gif" | "audio/wav" + int32 frame_count = 4; // frames actually encoded (after any capping) + float fps = 5; + bool has_audio = 6; // true when the muxed clip carries a soundtrack + int32 width = 7; + int32 height = 8; + int32 total_bytes = 9; // full payload size, for progress reporting + float duration_seconds = 10; + int32 sample_rate = 11; // audio streams only + // Every chunk: + bytes data = 12; + int32 chunk_index = 13; + int32 total_chunks = 14; +} + enum SampleEditType { EDIT_OVERRIDE = 0; EDIT_ACCUMULATE = 1; diff --git a/weightslab/proto/experiment_service_pb2.py b/weightslab/proto/experiment_service_pb2.py index d299a943..9a49b699 100644 --- a/weightslab/proto/experiment_service_pb2.py +++ b/weightslab/proto/experiment_service_pb2.py @@ -24,7 +24,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n)weightslab/proto/experiment_service.proto\"\x89\x01\n\x1aGetLatestLoggerDataRequest\x12\x1c\n\x14request_full_history\x18\x01 \x01(\x08\x12\x12\n\nmax_points\x18\x02 \x01(\x05\x12\x17\n\x0f\x62reak_by_slices\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\"\x81\x02\n\x0fLoggerDataPoint\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x11\n\tmodel_age\x18\x02 \x01(\x05\x12\x14\n\x0cmetric_value\x18\x03 \x01(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x04 \x01(\t\x12\x11\n\ttimestamp\x18\x05 \x01(\x03\x12\x11\n\tsample_id\x18\x06 \x01(\t\x12\x1c\n\x14is_evaluation_marker\x18\x07 \x01(\x08\x12\x12\n\nsplit_name\x18\x08 \x01(\t\x12\x17\n\x0f\x65valuation_tags\x18\t \x03(\t\x12\x12\n\npoint_note\x18\n \x01(\t\x12\x12\n\naudit_mode\x18\x0b \x01(\x08\"[\n\x1bGetLatestLoggerDataResponse\x12 \n\x06points\x18\x01 \x03(\x0b\x32\x10.LoggerDataPoint\x12\x1a\n\x12weightslab_version\x18\x02 \x01(\t\"\x07\n\x05\x45mpty\"/\n\x08NeuronId\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tneuron_id\x18\x02 \x01(\x05\"\x91\x02\n\x0fWeightOperation\x12*\n\x07op_type\x18\x01 \x01(\x0e\x32\x14.WeightOperationTypeH\x00\x88\x01\x01\x12\x15\n\x08layer_id\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1d\n\nneuron_ids\x18\x03 \x03(\x0b\x32\t.NeuronId\x12\x16\n\x0eneurons_to_add\x18\t \x01(\x05\x12 \n\x18zerofy_from_incoming_ids\x18\x0b \x03(\x05\x12\x1c\n\x14zerofy_to_neuron_ids\x18\x0c \x03(\x05\x12+\n\x11zerofy_predicates\x18\r \x03(\x0e\x32\x10.ZerofyPredicateB\n\n\x08_op_typeB\x0b\n\t_layer_id\"_\n\x17WeightsOperationRequest\x12/\n\x10weight_operation\x18\x01 \x01(\x0b\x32\x10.WeightOperationH\x00\x88\x01\x01\x42\x13\n\x11_weight_operation\"<\n\x18WeightsOperationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xc1\x05\n\x0fHyperParameters\x12\x1c\n\x0f\x65xperiment_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12!\n\x14training_steps_to_do\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x17\n\nbatch_size\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12 \n\x13\x66ull_eval_frequency\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12 \n\x13\x63heckpont_frequency\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x18\n\x0bis_training\x18\x07 \x01(\x08H\x06\x88\x01\x01\x12\x15\n\x08nb_steps\x18\x08 \x01(\x05H\x07\x88\x01\x01\x12\x19\n\x0c\x61uditor_mode\x18\t \x01(\x08H\x08\x88\x01\x01\x12\x1d\n\x10train_batch_size\x18\n \x01(\x05H\t\x88\x01\x01\x12\x1b\n\x0eval_batch_size\x18\x0b \x01(\x05H\n\x88\x01\x01\x12\x1c\n\x0ftest_batch_size\x18\x0c \x01(\x05H\x0b\x88\x01\x01\x12\x1c\n\x0f\x65valuation_mode\x18\r \x01(\x08H\x0c\x88\x01\x01\x12\x1e\n\x11\x65valuation_config\x18\x0e \x01(\tH\r\x88\x01\x01\x42\x12\n\x10_experiment_nameB\x17\n\x15_training_steps_to_doB\x10\n\x0e_learning_rateB\r\n\x0b_batch_sizeB\x16\n\x14_full_eval_frequencyB\x16\n\x14_checkpont_frequencyB\x0e\n\x0c_is_trainingB\x0b\n\t_nb_stepsB\x0f\n\r_auditor_modeB\x13\n\x11_train_batch_sizeB\x11\n\x0f_val_batch_sizeB\x12\n\x10_test_batch_sizeB\x12\n\x10_evaluation_modeB\x14\n\x12_evaluation_config\",\n\rMetricsStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"~\n\rAnnotatStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12.\n\x08metadata\x18\x02 \x03(\x0b\x32\x1c.AnnotatStatus.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x90\x02\n\x10TrainingStatusEx\x12\x16\n\ttimestamp\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0f\x65xperiment_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tmodel_age\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12+\n\x0emetrics_status\x18\x04 \x01(\x0b\x32\x0e.MetricsStatusH\x03\x88\x01\x01\x12+\n\x0e\x61nnotat_status\x18\x05 \x01(\x0b\x32\x0e.AnnotatStatusH\x04\x88\x01\x01\x42\x0c\n\n_timestampB\x12\n\x10_experiment_nameB\x0c\n\n_model_ageB\x11\n\x0f_metrics_statusB\x11\n\x0f_annotat_status\"]\n\x15HyperParameterCommand\x12/\n\x10hyper_parameters\x18\x01 \x01(\x0b\x32\x10.HyperParametersH\x00\x88\x01\x01\x42\x13\n\x11_hyper_parameters\">\n\x14\x44\x65nySamplesOperation\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\"0\n\x17LoadCheckpointOperation\x12\x15\n\rcheckpoint_id\x18\x01 \x01(\x05\"b\n\x11PlotNoteOperation\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x0c\n\x04note\x18\x04 \x01(\t\"L\n\x17SaveCheckpointOperation\x12\x19\n\x11save_architecture\x18\x01 \x01(\x08\x12\x16\n\x0esave_optimizer\x18\x02 \x01(\x08\"\x1a\n\x18RestartInstanceOperation\"\x8d\x08\n\x0eTrainerCommand\x12\x1c\n\x14get_hyper_parameters\x18\x04 \x01(\x08\x12\x1e\n\x16get_interactive_layers\x18\x05 \x01(\x08\x12\x1d\n\x10get_data_records\x18\x06 \x01(\tH\x00\x88\x01\x01\x12%\n\x18get_single_layer_info_id\x18\x08 \x01(\x05H\x01\x88\x01\x01\x12;\n\x16hyper_parameter_change\x18\x01 \x01(\x0b\x32\x16.HyperParameterCommandH\x02\x88\x01\x01\x12:\n\x16\x64\x65ny_samples_operation\x18\x07 \x01(\x0b\x32\x15.DenySamplesOperationH\x03\x88\x01\x01\x12?\n\x1b\x64\x65ny_eval_samples_operation\x18\n \x01(\x0b\x32\x15.DenySamplesOperationH\x04\x88\x01\x01\x12@\n\x19load_checkpoint_operation\x18\t \x01(\x0b\x32\x18.LoadCheckpointOperationH\x05\x88\x01\x01\x12\x42\n\x1eremove_from_denylist_operation\x18\x0b \x01(\x0b\x32\x15.DenySamplesOperationH\x06\x88\x01\x01\x12G\n#remove_eval_from_denylist_operation\x18\x0c \x01(\x0b\x32\x15.DenySamplesOperationH\x07\x88\x01\x01\x12\x34\n\x13plot_note_operation\x18\r \x01(\x0b\x32\x12.PlotNoteOperationH\x08\x88\x01\x01\x12@\n\x19save_checkpoint_operation\x18\x0e \x01(\x0b\x32\x18.SaveCheckpointOperationH\t\x88\x01\x01\x12\x39\n\x11restart_operation\x18\x0f \x01(\x0b\x32\x19.RestartInstanceOperationH\n\x88\x01\x01\x42\x13\n\x11_get_data_recordsB\x1b\n\x19_get_single_layer_info_idB\x19\n\x17_hyper_parameter_changeB\x19\n\x17_deny_samples_operationB\x1e\n\x1c_deny_eval_samples_operationB\x1c\n\x1a_load_checkpoint_operationB!\n\x1f_remove_from_denylist_operationB&\n$_remove_eval_from_denylist_operationB\x16\n\x14_plot_note_operationB\x1c\n\x1a_save_checkpoint_operationB\x14\n\x12_restart_operation\"\x9d\x01\n\x12HyperParameterDesc\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x1c\n\x0fnumerical_value\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cstring_value\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_numerical_valueB\x0f\n\r_string_value\"\xf2\x02\n\x10NeuronStatistics\x12!\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronIdH\x00\x88\x01\x01\x12\x17\n\nneuron_age\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1f\n\x12train_trigger_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x1e\n\x11\x65val_trigger_rate\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x07 \x01(\x02H\x04\x88\x01\x01\x12\x36\n\x0bincoming_lr\x18\x08 \x03(\x0b\x32!.NeuronStatistics.IncomingLrEntry\x1a\x31\n\x0fIncomingLrEntry\x12\x0b\n\x03key\x18\x01 \x01(\x05\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\x42\x0c\n\n_neuron_idB\r\n\x0b_neuron_ageB\x15\n\x13_train_trigger_rateB\x14\n\x12_eval_trigger_rateB\x10\n\x0e_learning_rate\"\xf0\x02\n\x13LayerRepresentation\x12\x15\n\x08layer_id\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rneurons_count\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12#\n\x16incoming_neurons_count\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x13\n\x06stride\x18\x07 \x01(\x05H\x06\x88\x01\x01\x12-\n\x12neurons_statistics\x18\n \x03(\x0b\x32\x11.NeuronStatisticsB\x0b\n\t_layer_idB\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x10\n\x0e_neurons_countB\x19\n\x17_incoming_neurons_countB\x0e\n\x0c_kernel_sizeB\t\n\x07_stride\"H\n\x11\x41\x63tivationRequest\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tsample_id\x18\x02 \x01(\t\x12\x0e\n\x06origin\x18\x03 \x01(\t\"H\n\rActivationMap\x12\x11\n\tneuron_id\x18\x01 \x01(\x05\x12\x0e\n\x06values\x18\x02 \x03(\x02\x12\t\n\x01H\x18\x03 \x01(\x05\x12\t\n\x01W\x18\x04 \x01(\x05\"d\n\x12\x41\x63tivationResponse\x12\x12\n\nlayer_type\x18\x01 \x01(\t\x12\x15\n\rneurons_count\x18\x02 \x01(\x05\x12#\n\x0b\x61\x63tivations\x18\x03 \x03(\x0b\x32\x0e.ActivationMap\"\x93\x01\n\tTaskField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x15\n\x0b\x66loat_value\x18\x02 \x01(\x02H\x00\x12\x13\n\tint_value\x18\x03 \x01(\x05H\x00\x12\x16\n\x0cstring_value\x18\x04 \x01(\tH\x00\x12\x15\n\x0b\x62ytes_value\x18\x05 \x01(\x0cH\x00\x12\x14\n\nbool_value\x18\x06 \x01(\x08H\x00\x42\x07\n\x05value\"\x87\x03\n\x0eRecordMetadata\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x14\n\x0csample_label\x18\x02 \x03(\x05\x12\x19\n\x11sample_prediction\x18\x03 \x03(\x05\x12=\n\x10sample_last_loss\x18\x04 \x03(\x0b\x32#.RecordMetadata.SampleLastLossEntry\x12\x19\n\x11sample_encounters\x18\x05 \x01(\x05\x12\x18\n\x10sample_discarded\x18\x06 \x01(\x08\x12 \n\x0c\x65xtra_fields\x18\x07 \x03(\x0b\x32\n.TaskField\x12\x16\n\x0eprediction_raw\x18\t \x01(\x0c\x12\x11\n\ttask_type\x18\n \x01(\t\x12\x19\n\x11sample_label_text\x18\x0b \x03(\t\x12\x1e\n\x16sample_prediction_text\x18\x0c \x03(\t\x1a\x35\n\x13SampleLastLossEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x93\x01\n\x10SampleStatistics\x12\x13\n\x06origin\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0csample_count\x18\x07 \x01(\x05H\x01\x88\x01\x01\x12\x11\n\ttask_type\x18\t \x01(\t\x12 \n\x07records\x18\x08 \x03(\x0b\x32\x0f.RecordMetadataB\t\n\x07_originB\x0f\n\r_sample_count\"\xe6\x01\n\x0f\x43ommandResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x33\n\x16hyper_parameters_descs\x18\x03 \x03(\x0b\x32\x13.HyperParameterDesc\x12\x33\n\x15layer_representations\x18\x04 \x03(\x0b\x32\x14.LayerRepresentation\x12\x31\n\x11sample_statistics\x18\x05 \x01(\x0b\x32\x11.SampleStatisticsH\x00\x88\x01\x01\x42\x14\n\x12_sample_statistics\"U\n\rSampleRequest\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_origin\"\xad\x02\n\x15SampleRequestResponse\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05label\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x64\x61ta\x18\x04 \x01(\x0cH\x03\x88\x01\x01\x12\x1a\n\rerror_message\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x15\n\x08raw_data\x18\x06 \x01(\x0cH\x05\x88\x01\x01\x12\x11\n\x04mask\x18\x07 \x01(\x0cH\x06\x88\x01\x01\x12\x17\n\nprediction\x18\x08 \x01(\x0cH\x07\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_originB\x08\n\x06_labelB\x07\n\x05_dataB\x10\n\x0e_error_messageB\x0b\n\t_raw_dataB\x07\n\x05_maskB\r\n\x0b_prediction\"\x92\x01\n\x12\x42\x61tchSampleRequest\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x19\n\x0cresize_width\x18\x03 \x01(\x05H\x00\x88\x01\x01\x12\x1a\n\rresize_height\x18\x04 \x01(\x05H\x01\x88\x01\x01\x42\x0f\n\r_resize_widthB\x10\n\x0e_resize_height\">\n\x13\x42\x61tchSampleResponse\x12\'\n\x07samples\x18\x01 \x03(\x0b\x32\x16.SampleRequestResponse\".\n\x0eWeightsRequest\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\"\x9d\x02\n\x0fWeightsResponse\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x10\n\x08incoming\x18\x04 \x01(\x05\x12\x10\n\x08outgoing\x18\x05 \x01(\x05\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x0f\n\x07weights\x18\x07 \x03(\x02\x12\x0f\n\x07success\x18\x0b \x01(\x08\x12\x1a\n\rerror_message\x18\x0c \x01(\tH\x03\x88\x01\x01\x42\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x0e\n\x0c_kernel_sizeB\x10\n\x0e_error_message\"R\n\x10\x44\x61taQueryRequest\x12\r\n\x05query\x18\x01 \x01(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\x12\x1b\n\x13is_natural_language\x18\x03 \x01(\x08\"5\n\x11\x43\x61tegoricalTagDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\ncategories\x18\x02 \x03(\t\"\xa9\x02\n\x11\x44\x61taQueryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1d\n\x15number_of_all_samples\x18\x03 \x01(\x05\x12%\n\x1dnumber_of_samples_in_the_loop\x18\x04 \x01(\x05\x12#\n\x1bnumber_of_discarded_samples\x18\x05 \x01(\x05\x12\x13\n\x0bunique_tags\x18\x06 \x03(\t\x12+\n\x11\x61gent_intent_type\x18\x07 \x01(\x0e\x32\x10.AgentIntentType\x12\x17\n\x0f\x61nalysis_result\x18\x08 \x01(\t\x12,\n\x10\x63\x61tegorical_tags\x18\t \x03(\x0b\x32\x12.CategoricalTagDef\"\xc2\x01\n\x12\x44\x61taSamplesRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12 \n\x18include_transformed_data\x18\x03 \x01(\x08\x12\x18\n\x10include_raw_data\x18\x04 \x01(\x08\x12\x19\n\x11stats_to_retrieve\x18\x05 \x03(\t\x12\x14\n\x0cresize_width\x18\x06 \x01(\x05\x12\x15\n\rresize_height\x18\x07 \x01(\x05\"m\n\x08\x44\x61taStat\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\r\n\x05shape\x18\x03 \x03(\x05\x12\r\n\x05value\x18\x04 \x03(\x02\x12\x14\n\x0cvalue_string\x18\x05 \x01(\t\x12\x11\n\tthumbnail\x18\x06 \x01(\x0c\">\n\nDataRecord\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x1d\n\ndata_stats\x18\x02 \x03(\x0b\x32\t.DataStat\"Z\n\x13\x44\x61taSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12!\n\x0c\x64\x61ta_records\x18\x03 \x03(\x0b\x32\x0b.DataRecord\"C\n\x0fHistogramSubBar\x12\x0e\n\x06origin\x18\x01 \x01(\t\x12\x11\n\tdiscarded\x18\x02 \x01(\x08\x12\r\n\x05\x63ount\x18\x03 \x01(\x03\"h\n\x0cHistogramBin\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03\x61vg\x18\x03 \x01(\x01\x12\r\n\x05\x63ount\x18\x04 \x01(\x03\x12\"\n\x08sub_bars\x18\x05 \x03(\x0b\x32\x10.HistogramSubBar\"[\n\x17\x43\x61tegoricalHistogramBar\x12\r\n\x05label\x18\x01 \x01(\t\x12\r\n\x05\x63ount\x18\x02 \x01(\x03\x12\"\n\x08sub_bars\x18\x03 \x03(\x0b\x32\x10.HistogramSubBar\"4\n\x10HistogramRequest\x12\x0e\n\x06\x63olumn\x18\x01 \x01(\t\x12\x10\n\x08max_bins\x18\x02 \x01(\x05\"\xb2\x01\n\x11HistogramResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\ntotal_rows\x18\x03 \x01(\x03\x12\x1b\n\x04\x62ins\x18\x04 \x03(\x0b\x32\r.HistogramBin\x12\x16\n\x0eis_categorical\x18\x05 \x01(\x08\x12\x32\n\x10\x63\x61tegorical_bars\x18\x06 \x03(\x0b\x32\x18.CategoricalHistogramBar\"W\n\x12GetMetaDataRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12\x17\n\x0fmodal_sample_id\x18\x03 \x01(\t\"\x99\x01\n\x13GetMetaDataResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1a\n\x12\x61ll_metadata_names\x18\x03 \x03(\t\x12!\n\x0cgrid_records\x18\x04 \x03(\x0b\x32\x0b.DataRecord\x12!\n\x0cmodal_record\x18\x05 \x01(\x0b\x32\x0b.DataRecord\"Y\n\x1aGetSignalTrajectoryRequest\x12\x13\n\x0bsignal_name\x18\x01 \x01(\t\x12\x12\n\nsample_ids\x18\x02 \x03(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"4\n\x10SignalTrajectory\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x03(\x02\"}\n\x1bGetSignalTrajectoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x13\n\x0bsignal_name\x18\x03 \x01(\t\x12\'\n\x0ctrajectories\x18\x04 \x03(\x0b\x32\x11.SignalTrajectory\"J\n\x11PointCloudRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"\xbf\x01\n\x0fPointCloudChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nnum_points\x18\x03 \x01(\x05\x12\x14\n\x0cnum_features\x18\x04 \x01(\x05\x12\x10\n\x08pc_range\x18\x05 \x03(\x02\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\x07 \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x08 \x01(\x05\x12\x15\n\rfeature_names\x18\t \x03(\t\"\xdc\x01\n\x10\x44\x61taEditsRequest\x12\x11\n\tstat_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66loat_value\x18\x02 \x01(\x02\x12\x14\n\x0cstring_value\x18\x03 \x01(\t\x12\x12\n\nbool_value\x18\x04 \x01(\x08\x12\x1d\n\x04type\x18\x05 \x01(\x0e\x32\x0f.SampleEditType\x12\x13\n\x0bsamples_ids\x18\x06 \x03(\t\x12\x16\n\x0esample_origins\x18\x07 \x03(\t\x12\x16\n\x0eis_categorical\x18\x08 \x01(\x08\x12\x12\n\ncategories\x18\t \x03(\t\"5\n\x11\x44\x61taEditsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\":\n\x12\x44\x61taSplitsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x13\n\x0bsplit_names\x18\x02 \x03(\t\"9\n\x13\x41gentHealthResponse\x12\x11\n\tavailable\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"^\n\x16InitializeAgentRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12$\n\x08provider\x18\x02 \x01(\x0e\x32\x12.AgentProviderType\x12\r\n\x05model\x18\x03 \x01(\t\";\n\x17InitializeAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"(\n\x17\x43hangeAgentModelRequest\x12\r\n\x05model\x18\x01 \x01(\t\"<\n\x18\x43hangeAgentModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x17\n\x15GetAgentModelsRequest\"J\n\x16GetAgentModelsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0e\n\x06models\x18\x02 \x03(\t\x12\x0f\n\x07message\x18\x03 \x01(\t\"6\n\x12ResetAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"3\n\x18RestoreCheckpointRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\"=\n\x19RestoreCheckpointResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"R\n\x18TriggerEvaluationRequest\x12\x12\n\nsplit_name\x18\x01 \x01(\t\x12\x0c\n\x04tags\x18\x02 \x03(\t\x12\x14\n\x0cuse_full_set\x18\x03 \x01(\x08\"=\n\x19TriggerEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x1c\n\x1aGetEvaluationStatusRequest\"\x81\x01\n\x1bGetEvaluationStatusResponse\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x0f\n\x07\x63urrent\x18\x02 \x01(\x05\x12\r\n\x05total\x18\x03 \x01(\x05\x12\x0f\n\x07message\x18\x04 \x01(\t\x12\r\n\x05\x65rror\x18\x05 \x01(\t\x12\x12\n\nsplit_name\x18\x06 \x01(\t\")\n\x17\x43\x61ncelEvaluationRequest\x12\x0e\n\x06reason\x18\x01 \x01(\t\"<\n\x18\x43\x61ncelEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"7\n\x16RunNotebookCellRequest\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07\x63\x65ll_id\x18\x02 \x01(\t\"2\n\x10NotebookCellDone\x12\x12\n\nexec_count\x18\x01 \x01(\x05\x12\n\n\x02ok\x18\x02 \x01(\x08\"\x1e\n\x1cInterruptNotebookCellRequest\":\n\x1dInterruptNotebookCellResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\xbd\x01\n\x11NotebookCellChunk\x12\x0f\n\x07\x63\x65ll_id\x18\x01 \x01(\t\x12\x10\n\x06stdout\x18\x02 \x01(\tH\x00\x12\x10\n\x06stderr\x18\x03 \x01(\tH\x00\x12\x15\n\x0bresult_text\x18\x04 \x01(\tH\x00\x12\x13\n\timage_png\x18\x05 \x01(\x0cH\x00\x12\x19\n\x0f\x65rror_traceback\x18\x06 \x01(\tH\x00\x12!\n\x04\x64one\x18\x07 \x01(\x0b\x32\x11.NotebookCellDoneH\x00\x42\t\n\x07payload\"S\n\x10NotebookResponse\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0f\n\x07\x65xisted\x18\x02 \x01(\x08\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"7\n\x13SaveNotebookRequest\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\"M\n\x14SaveNotebookResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"C\n\x1bGenerateNotebookCodeRequest\x12\x0e\n\x06prompt\x18\x01 \x01(\t\x12\x14\n\x0c\x63ontext_code\x18\x02 \x01(\t\"\\\n\x1cGenerateNotebookCodeResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x13\n\x0b\x65xplanation\x18\x02 \x01(\t\x12\n\n\x02ok\x18\x03 \x01(\x08\x12\r\n\x05\x65rror\x18\x04 \x01(\t*d\n\x13WeightOperationType\x12\n\n\x06ZEROFY\x10\x00\x12\x10\n\x0cREINITIALIZE\x10\x01\x12\n\n\x06\x46REEZE\x10\x02\x12\x12\n\x0eREMOVE_NEURONS\x10\t\x12\x0f\n\x0b\x41\x44\x44_NEURONS\x10\n*o\n\x0fZerofyPredicate\x12\x19\n\x15ZEROFY_PREDICATE_NONE\x10\x00\x12 \n\x1cZEROFY_PREDICATE_WITH_FROZEN\x10\x01\x12\x1f\n\x1bZEROFY_PREDICATE_WITH_OLDER\x10\x02*M\n\x0f\x41gentIntentType\x12\x12\n\x0eINTENT_UNKNOWN\x10\x00\x12\x11\n\rINTENT_FILTER\x10\x01\x12\x13\n\x0fINTENT_ANALYSIS\x10\x02*I\n\x0eSampleEditType\x12\x11\n\rEDIT_OVERRIDE\x10\x00\x12\x13\n\x0f\x45\x44IT_ACCUMULATE\x10\x01\x12\x0f\n\x0b\x45\x44IT_REMOVE\x10\x02*,\n\x11\x41gentProviderType\x12\x17\n\x13PROVIDER_OPENROUTER\x10\x00\x32\x9d\x0e\n\x11\x45xperimentService\x12P\n\x13GetLatestLoggerData\x12\x1b.GetLatestLoggerDataRequest\x1a\x1c.GetLatestLoggerDataResponse\x12\x36\n\x11\x45xperimentCommand\x12\x0f.TrainerCommand\x1a\x10.CommandResponse\x12H\n\x11ManipulateWeights\x12\x18.WeightsOperationRequest\x1a\x19.WeightsOperationResponse\x12/\n\nGetWeights\x12\x0f.WeightsRequest\x1a\x10.WeightsResponse\x12\x39\n\x0eGetActivations\x12\x12.ActivationRequest\x1a\x13.ActivationResponse\x12\x37\n\nGetSamples\x12\x13.BatchSampleRequest\x1a\x14.BatchSampleResponse\x12\x37\n\x0e\x41pplyDataQuery\x12\x11.DataQueryRequest\x1a\x12.DataQueryResponse\x12;\n\x0eGetDataSamples\x12\x13.DataSamplesRequest\x1a\x14.DataSamplesResponse\x12\x35\n\x0cGetHistogram\x12\x11.HistogramRequest\x1a\x12.HistogramResponse\x12\x38\n\x0bGetMetaData\x12\x13.GetMetaDataRequest\x1a\x14.GetMetaDataResponse\x12P\n\x13GetSignalTrajectory\x12\x1b.GetSignalTrajectoryRequest\x1a\x1c.GetSignalTrajectoryResponse\x12\x37\n\rGetPointCloud\x12\x12.PointCloudRequest\x1a\x10.PointCloudChunk0\x01\x12\x37\n\x0e\x45\x64itDataSample\x12\x11.DataEditsRequest\x1a\x12.DataEditsResponse\x12,\n\rGetDataSplits\x12\x06.Empty\x1a\x13.DataSplitsResponse\x12\x30\n\x10\x43heckAgentHealth\x12\x06.Empty\x1a\x14.AgentHealthResponse\x12\x44\n\x0fInitializeAgent\x12\x17.InitializeAgentRequest\x1a\x18.InitializeAgentResponse\x12G\n\x10\x43hangeAgentModel\x12\x18.ChangeAgentModelRequest\x1a\x19.ChangeAgentModelResponse\x12\x41\n\x0eGetAgentModels\x12\x16.GetAgentModelsRequest\x1a\x17.GetAgentModelsResponse\x12)\n\nResetAgent\x12\x06.Empty\x1a\x13.ResetAgentResponse\x12@\n\x0fRunNotebookCell\x12\x17.RunNotebookCellRequest\x1a\x12.NotebookCellChunk0\x01\x12V\n\x15InterruptNotebookCell\x12\x1d.InterruptNotebookCellRequest\x1a\x1e.InterruptNotebookCellResponse\x12(\n\x0bGetNotebook\x12\x06.Empty\x1a\x11.NotebookResponse\x12;\n\x0cSaveNotebook\x12\x14.SaveNotebookRequest\x1a\x15.SaveNotebookResponse\x12S\n\x14GenerateNotebookCode\x12\x1c.GenerateNotebookCodeRequest\x1a\x1d.GenerateNotebookCodeResponse\x12J\n\x11RestoreCheckpoint\x12\x19.RestoreCheckpointRequest\x1a\x1a.RestoreCheckpointResponse\x12J\n\x11TriggerEvaluation\x12\x19.TriggerEvaluationRequest\x1a\x1a.TriggerEvaluationResponse\x12P\n\x13GetEvaluationStatus\x12\x1b.GetEvaluationStatusRequest\x1a\x1c.GetEvaluationStatusResponse\x12G\n\x10\x43\x61ncelEvaluation\x12\x18.CancelEvaluationRequest\x1a\x19.CancelEvaluationResponseb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n)weightslab/proto/experiment_service.proto\"\x89\x01\n\x1aGetLatestLoggerDataRequest\x12\x1c\n\x14request_full_history\x18\x01 \x01(\x08\x12\x12\n\nmax_points\x18\x02 \x01(\x05\x12\x17\n\x0f\x62reak_by_slices\x18\x03 \x01(\x08\x12\x0c\n\x04tags\x18\x04 \x03(\t\x12\x12\n\ngraph_name\x18\x05 \x01(\t\"\x81\x02\n\x0fLoggerDataPoint\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x11\n\tmodel_age\x18\x02 \x01(\x05\x12\x14\n\x0cmetric_value\x18\x03 \x01(\x02\x12\x17\n\x0f\x65xperiment_hash\x18\x04 \x01(\t\x12\x11\n\ttimestamp\x18\x05 \x01(\x03\x12\x11\n\tsample_id\x18\x06 \x01(\t\x12\x1c\n\x14is_evaluation_marker\x18\x07 \x01(\x08\x12\x12\n\nsplit_name\x18\x08 \x01(\t\x12\x17\n\x0f\x65valuation_tags\x18\t \x03(\t\x12\x12\n\npoint_note\x18\n \x01(\t\x12\x12\n\naudit_mode\x18\x0b \x01(\x08\"[\n\x1bGetLatestLoggerDataResponse\x12 \n\x06points\x18\x01 \x03(\x0b\x32\x10.LoggerDataPoint\x12\x1a\n\x12weightslab_version\x18\x02 \x01(\t\"\x07\n\x05\x45mpty\"/\n\x08NeuronId\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tneuron_id\x18\x02 \x01(\x05\"\x91\x02\n\x0fWeightOperation\x12*\n\x07op_type\x18\x01 \x01(\x0e\x32\x14.WeightOperationTypeH\x00\x88\x01\x01\x12\x15\n\x08layer_id\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1d\n\nneuron_ids\x18\x03 \x03(\x0b\x32\t.NeuronId\x12\x16\n\x0eneurons_to_add\x18\t \x01(\x05\x12 \n\x18zerofy_from_incoming_ids\x18\x0b \x03(\x05\x12\x1c\n\x14zerofy_to_neuron_ids\x18\x0c \x03(\x05\x12+\n\x11zerofy_predicates\x18\r \x03(\x0e\x32\x10.ZerofyPredicateB\n\n\x08_op_typeB\x0b\n\t_layer_id\"_\n\x17WeightsOperationRequest\x12/\n\x10weight_operation\x18\x01 \x01(\x0b\x32\x10.WeightOperationH\x00\x88\x01\x01\x42\x13\n\x11_weight_operation\"<\n\x18WeightsOperationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\xc1\x05\n\x0fHyperParameters\x12\x1c\n\x0f\x65xperiment_name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12!\n\x14training_steps_to_do\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x17\n\nbatch_size\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12 \n\x13\x66ull_eval_frequency\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12 \n\x13\x63heckpont_frequency\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x18\n\x0bis_training\x18\x07 \x01(\x08H\x06\x88\x01\x01\x12\x15\n\x08nb_steps\x18\x08 \x01(\x05H\x07\x88\x01\x01\x12\x19\n\x0c\x61uditor_mode\x18\t \x01(\x08H\x08\x88\x01\x01\x12\x1d\n\x10train_batch_size\x18\n \x01(\x05H\t\x88\x01\x01\x12\x1b\n\x0eval_batch_size\x18\x0b \x01(\x05H\n\x88\x01\x01\x12\x1c\n\x0ftest_batch_size\x18\x0c \x01(\x05H\x0b\x88\x01\x01\x12\x1c\n\x0f\x65valuation_mode\x18\r \x01(\x08H\x0c\x88\x01\x01\x12\x1e\n\x11\x65valuation_config\x18\x0e \x01(\tH\r\x88\x01\x01\x42\x12\n\x10_experiment_nameB\x17\n\x15_training_steps_to_doB\x10\n\x0e_learning_rateB\r\n\x0b_batch_sizeB\x16\n\x14_full_eval_frequencyB\x16\n\x14_checkpont_frequencyB\x0e\n\x0c_is_trainingB\x0b\n\t_nb_stepsB\x0f\n\r_auditor_modeB\x13\n\x11_train_batch_sizeB\x11\n\x0f_val_batch_sizeB\x12\n\x10_test_batch_sizeB\x12\n\x10_evaluation_modeB\x14\n\x12_evaluation_config\",\n\rMetricsStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02\"~\n\rAnnotatStatus\x12\x0c\n\x04name\x18\x01 \x01(\t\x12.\n\x08metadata\x18\x02 \x03(\x0b\x32\x1c.AnnotatStatus.MetadataEntry\x1a/\n\rMetadataEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x90\x02\n\x10TrainingStatusEx\x12\x16\n\ttimestamp\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x1c\n\x0f\x65xperiment_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x16\n\tmodel_age\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12+\n\x0emetrics_status\x18\x04 \x01(\x0b\x32\x0e.MetricsStatusH\x03\x88\x01\x01\x12+\n\x0e\x61nnotat_status\x18\x05 \x01(\x0b\x32\x0e.AnnotatStatusH\x04\x88\x01\x01\x42\x0c\n\n_timestampB\x12\n\x10_experiment_nameB\x0c\n\n_model_ageB\x11\n\x0f_metrics_statusB\x11\n\x0f_annotat_status\"]\n\x15HyperParameterCommand\x12/\n\x10hyper_parameters\x18\x01 \x01(\x0b\x32\x10.HyperParametersH\x00\x88\x01\x01\x42\x13\n\x11_hyper_parameters\">\n\x14\x44\x65nySamplesOperation\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\"0\n\x17LoadCheckpointOperation\x12\x15\n\rcheckpoint_id\x18\x01 \x01(\x05\"b\n\x11PlotNoteOperation\x12\x13\n\x0bmetric_name\x18\x01 \x01(\t\x12\x17\n\x0f\x65xperiment_hash\x18\x02 \x01(\t\x12\x11\n\tmodel_age\x18\x03 \x01(\x05\x12\x0c\n\x04note\x18\x04 \x01(\t\"L\n\x17SaveCheckpointOperation\x12\x19\n\x11save_architecture\x18\x01 \x01(\x08\x12\x16\n\x0esave_optimizer\x18\x02 \x01(\x08\"\x1a\n\x18RestartInstanceOperation\"\x8d\x08\n\x0eTrainerCommand\x12\x1c\n\x14get_hyper_parameters\x18\x04 \x01(\x08\x12\x1e\n\x16get_interactive_layers\x18\x05 \x01(\x08\x12\x1d\n\x10get_data_records\x18\x06 \x01(\tH\x00\x88\x01\x01\x12%\n\x18get_single_layer_info_id\x18\x08 \x01(\x05H\x01\x88\x01\x01\x12;\n\x16hyper_parameter_change\x18\x01 \x01(\x0b\x32\x16.HyperParameterCommandH\x02\x88\x01\x01\x12:\n\x16\x64\x65ny_samples_operation\x18\x07 \x01(\x0b\x32\x15.DenySamplesOperationH\x03\x88\x01\x01\x12?\n\x1b\x64\x65ny_eval_samples_operation\x18\n \x01(\x0b\x32\x15.DenySamplesOperationH\x04\x88\x01\x01\x12@\n\x19load_checkpoint_operation\x18\t \x01(\x0b\x32\x18.LoadCheckpointOperationH\x05\x88\x01\x01\x12\x42\n\x1eremove_from_denylist_operation\x18\x0b \x01(\x0b\x32\x15.DenySamplesOperationH\x06\x88\x01\x01\x12G\n#remove_eval_from_denylist_operation\x18\x0c \x01(\x0b\x32\x15.DenySamplesOperationH\x07\x88\x01\x01\x12\x34\n\x13plot_note_operation\x18\r \x01(\x0b\x32\x12.PlotNoteOperationH\x08\x88\x01\x01\x12@\n\x19save_checkpoint_operation\x18\x0e \x01(\x0b\x32\x18.SaveCheckpointOperationH\t\x88\x01\x01\x12\x39\n\x11restart_operation\x18\x0f \x01(\x0b\x32\x19.RestartInstanceOperationH\n\x88\x01\x01\x42\x13\n\x11_get_data_recordsB\x1b\n\x19_get_single_layer_info_idB\x19\n\x17_hyper_parameter_changeB\x19\n\x17_deny_samples_operationB\x1e\n\x1c_deny_eval_samples_operationB\x1c\n\x1a_load_checkpoint_operationB!\n\x1f_remove_from_denylist_operationB&\n$_remove_eval_from_denylist_operationB\x16\n\x14_plot_note_operationB\x1c\n\x1a_save_checkpoint_operationB\x14\n\x12_restart_operation\"\x9d\x01\n\x12HyperParameterDesc\x12\r\n\x05label\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\x12\x0c\n\x04type\x18\x03 \x01(\t\x12\x1c\n\x0fnumerical_value\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cstring_value\x18\x05 \x01(\tH\x01\x88\x01\x01\x42\x12\n\x10_numerical_valueB\x0f\n\r_string_value\"\xf2\x02\n\x10NeuronStatistics\x12!\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronIdH\x00\x88\x01\x01\x12\x17\n\nneuron_age\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12\x1f\n\x12train_trigger_rate\x18\x03 \x01(\x02H\x02\x88\x01\x01\x12\x1e\n\x11\x65val_trigger_rate\x18\x04 \x01(\x02H\x03\x88\x01\x01\x12\x1a\n\rlearning_rate\x18\x07 \x01(\x02H\x04\x88\x01\x01\x12\x36\n\x0bincoming_lr\x18\x08 \x03(\x0b\x32!.NeuronStatistics.IncomingLrEntry\x1a\x31\n\x0fIncomingLrEntry\x12\x0b\n\x03key\x18\x01 \x01(\x05\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\x42\x0c\n\n_neuron_idB\r\n\x0b_neuron_ageB\x15\n\x13_train_trigger_rateB\x14\n\x12_eval_trigger_rateB\x10\n\x0e_learning_rate\"\xf0\x02\n\x13LayerRepresentation\x12\x15\n\x08layer_id\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x02\x88\x01\x01\x12\x1a\n\rneurons_count\x18\x04 \x01(\x05H\x03\x88\x01\x01\x12#\n\x16incoming_neurons_count\x18\x05 \x01(\x05H\x04\x88\x01\x01\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x05\x88\x01\x01\x12\x13\n\x06stride\x18\x07 \x01(\x05H\x06\x88\x01\x01\x12-\n\x12neurons_statistics\x18\n \x03(\x0b\x32\x11.NeuronStatisticsB\x0b\n\t_layer_idB\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x10\n\x0e_neurons_countB\x19\n\x17_incoming_neurons_countB\x0e\n\x0c_kernel_sizeB\t\n\x07_stride\"H\n\x11\x41\x63tivationRequest\x12\x10\n\x08layer_id\x18\x01 \x01(\x05\x12\x11\n\tsample_id\x18\x02 \x01(\t\x12\x0e\n\x06origin\x18\x03 \x01(\t\"H\n\rActivationMap\x12\x11\n\tneuron_id\x18\x01 \x01(\x05\x12\x0e\n\x06values\x18\x02 \x03(\x02\x12\t\n\x01H\x18\x03 \x01(\x05\x12\t\n\x01W\x18\x04 \x01(\x05\"d\n\x12\x41\x63tivationResponse\x12\x12\n\nlayer_type\x18\x01 \x01(\t\x12\x15\n\rneurons_count\x18\x02 \x01(\x05\x12#\n\x0b\x61\x63tivations\x18\x03 \x03(\x0b\x32\x0e.ActivationMap\"\x93\x01\n\tTaskField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x15\n\x0b\x66loat_value\x18\x02 \x01(\x02H\x00\x12\x13\n\tint_value\x18\x03 \x01(\x05H\x00\x12\x16\n\x0cstring_value\x18\x04 \x01(\tH\x00\x12\x15\n\x0b\x62ytes_value\x18\x05 \x01(\x0cH\x00\x12\x14\n\nbool_value\x18\x06 \x01(\x08H\x00\x42\x07\n\x05value\"\x87\x03\n\x0eRecordMetadata\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x14\n\x0csample_label\x18\x02 \x03(\x05\x12\x19\n\x11sample_prediction\x18\x03 \x03(\x05\x12=\n\x10sample_last_loss\x18\x04 \x03(\x0b\x32#.RecordMetadata.SampleLastLossEntry\x12\x19\n\x11sample_encounters\x18\x05 \x01(\x05\x12\x18\n\x10sample_discarded\x18\x06 \x01(\x08\x12 \n\x0c\x65xtra_fields\x18\x07 \x03(\x0b\x32\n.TaskField\x12\x16\n\x0eprediction_raw\x18\t \x01(\x0c\x12\x11\n\ttask_type\x18\n \x01(\t\x12\x19\n\x11sample_label_text\x18\x0b \x03(\t\x12\x1e\n\x16sample_prediction_text\x18\x0c \x03(\t\x1a\x35\n\x13SampleLastLossEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x02:\x02\x38\x01\"\x93\x01\n\x10SampleStatistics\x12\x13\n\x06origin\x18\x06 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0csample_count\x18\x07 \x01(\x05H\x01\x88\x01\x01\x12\x11\n\ttask_type\x18\t \x01(\t\x12 \n\x07records\x18\x08 \x03(\x0b\x32\x0f.RecordMetadataB\t\n\x07_originB\x0f\n\r_sample_count\"\xe6\x01\n\x0f\x43ommandResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x33\n\x16hyper_parameters_descs\x18\x03 \x03(\x0b\x32\x13.HyperParameterDesc\x12\x33\n\x15layer_representations\x18\x04 \x03(\x0b\x32\x14.LayerRepresentation\x12\x31\n\x11sample_statistics\x18\x05 \x01(\x0b\x32\x11.SampleStatisticsH\x00\x88\x01\x01\x42\x14\n\x12_sample_statistics\"U\n\rSampleRequest\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_origin\"\xad\x02\n\x15SampleRequestResponse\x12\x16\n\tsample_id\x18\x01 \x01(\tH\x00\x88\x01\x01\x12\x13\n\x06origin\x18\x02 \x01(\tH\x01\x88\x01\x01\x12\x12\n\x05label\x18\x03 \x01(\x05H\x02\x88\x01\x01\x12\x11\n\x04\x64\x61ta\x18\x04 \x01(\x0cH\x03\x88\x01\x01\x12\x1a\n\rerror_message\x18\x05 \x01(\tH\x04\x88\x01\x01\x12\x15\n\x08raw_data\x18\x06 \x01(\x0cH\x05\x88\x01\x01\x12\x11\n\x04mask\x18\x07 \x01(\x0cH\x06\x88\x01\x01\x12\x17\n\nprediction\x18\x08 \x01(\x0cH\x07\x88\x01\x01\x42\x0c\n\n_sample_idB\t\n\x07_originB\x08\n\x06_labelB\x07\n\x05_dataB\x10\n\x0e_error_messageB\x0b\n\t_raw_dataB\x07\n\x05_maskB\r\n\x0b_prediction\"\x92\x01\n\x12\x42\x61tchSampleRequest\x12\x12\n\nsample_ids\x18\x01 \x03(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x19\n\x0cresize_width\x18\x03 \x01(\x05H\x00\x88\x01\x01\x12\x1a\n\rresize_height\x18\x04 \x01(\x05H\x01\x88\x01\x01\x42\x0f\n\r_resize_widthB\x10\n\x0e_resize_height\">\n\x13\x42\x61tchSampleResponse\x12\'\n\x07samples\x18\x01 \x03(\x0b\x32\x16.SampleRequestResponse\".\n\x0eWeightsRequest\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\"\x9d\x02\n\x0fWeightsResponse\x12\x1c\n\tneuron_id\x18\x01 \x01(\x0b\x32\t.NeuronId\x12\x17\n\nlayer_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x17\n\nlayer_type\x18\x03 \x01(\tH\x01\x88\x01\x01\x12\x10\n\x08incoming\x18\x04 \x01(\x05\x12\x10\n\x08outgoing\x18\x05 \x01(\x05\x12\x18\n\x0bkernel_size\x18\x06 \x01(\x05H\x02\x88\x01\x01\x12\x0f\n\x07weights\x18\x07 \x03(\x02\x12\x0f\n\x07success\x18\x0b \x01(\x08\x12\x1a\n\rerror_message\x18\x0c \x01(\tH\x03\x88\x01\x01\x42\r\n\x0b_layer_nameB\r\n\x0b_layer_typeB\x0e\n\x0c_kernel_sizeB\x10\n\x0e_error_message\"R\n\x10\x44\x61taQueryRequest\x12\r\n\x05query\x18\x01 \x01(\t\x12\x12\n\naccumulate\x18\x02 \x01(\x08\x12\x1b\n\x13is_natural_language\x18\x03 \x01(\x08\"5\n\x11\x43\x61tegoricalTagDef\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\ncategories\x18\x02 \x03(\t\"\xa9\x02\n\x11\x44\x61taQueryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1d\n\x15number_of_all_samples\x18\x03 \x01(\x05\x12%\n\x1dnumber_of_samples_in_the_loop\x18\x04 \x01(\x05\x12#\n\x1bnumber_of_discarded_samples\x18\x05 \x01(\x05\x12\x13\n\x0bunique_tags\x18\x06 \x03(\t\x12+\n\x11\x61gent_intent_type\x18\x07 \x01(\x0e\x32\x10.AgentIntentType\x12\x17\n\x0f\x61nalysis_result\x18\x08 \x01(\t\x12,\n\x10\x63\x61tegorical_tags\x18\t \x03(\x0b\x32\x12.CategoricalTagDef\"\xc2\x01\n\x12\x44\x61taSamplesRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12 \n\x18include_transformed_data\x18\x03 \x01(\x08\x12\x18\n\x10include_raw_data\x18\x04 \x01(\x08\x12\x19\n\x11stats_to_retrieve\x18\x05 \x03(\t\x12\x14\n\x0cresize_width\x18\x06 \x01(\x05\x12\x15\n\rresize_height\x18\x07 \x01(\x05\"m\n\x08\x44\x61taStat\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\r\n\x05shape\x18\x03 \x03(\x05\x12\r\n\x05value\x18\x04 \x03(\x02\x12\x14\n\x0cvalue_string\x18\x05 \x01(\t\x12\x11\n\tthumbnail\x18\x06 \x01(\x0c\">\n\nDataRecord\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x1d\n\ndata_stats\x18\x02 \x03(\x0b\x32\t.DataStat\"Z\n\x13\x44\x61taSamplesResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12!\n\x0c\x64\x61ta_records\x18\x03 \x03(\x0b\x32\x0b.DataRecord\"C\n\x0fHistogramSubBar\x12\x0e\n\x06origin\x18\x01 \x01(\t\x12\x11\n\tdiscarded\x18\x02 \x01(\x08\x12\r\n\x05\x63ount\x18\x03 \x01(\x03\"h\n\x0cHistogramBin\x12\x0b\n\x03min\x18\x01 \x01(\x01\x12\x0b\n\x03max\x18\x02 \x01(\x01\x12\x0b\n\x03\x61vg\x18\x03 \x01(\x01\x12\r\n\x05\x63ount\x18\x04 \x01(\x03\x12\"\n\x08sub_bars\x18\x05 \x03(\x0b\x32\x10.HistogramSubBar\"[\n\x17\x43\x61tegoricalHistogramBar\x12\r\n\x05label\x18\x01 \x01(\t\x12\r\n\x05\x63ount\x18\x02 \x01(\x03\x12\"\n\x08sub_bars\x18\x03 \x03(\x0b\x32\x10.HistogramSubBar\"4\n\x10HistogramRequest\x12\x0e\n\x06\x63olumn\x18\x01 \x01(\t\x12\x10\n\x08max_bins\x18\x02 \x01(\x05\"\xb2\x01\n\x11HistogramResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\ntotal_rows\x18\x03 \x01(\x03\x12\x1b\n\x04\x62ins\x18\x04 \x03(\x0b\x32\r.HistogramBin\x12\x16\n\x0eis_categorical\x18\x05 \x01(\x08\x12\x32\n\x10\x63\x61tegorical_bars\x18\x06 \x03(\x0b\x32\x18.CategoricalHistogramBar\"W\n\x12GetMetaDataRequest\x12\x13\n\x0bstart_index\x18\x01 \x01(\x05\x12\x13\n\x0brecords_cnt\x18\x02 \x01(\x05\x12\x17\n\x0fmodal_sample_id\x18\x03 \x01(\t\"\x99\x01\n\x13GetMetaDataResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x1a\n\x12\x61ll_metadata_names\x18\x03 \x03(\t\x12!\n\x0cgrid_records\x18\x04 \x03(\x0b\x32\x0b.DataRecord\x12!\n\x0cmodal_record\x18\x05 \x01(\x0b\x32\x0b.DataRecord\"Y\n\x1aGetSignalTrajectoryRequest\x12\x13\n\x0bsignal_name\x18\x01 \x01(\t\x12\x12\n\nsample_ids\x18\x02 \x03(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\"4\n\x10SignalTrajectory\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x03(\x02\"}\n\x1bGetSignalTrajectoryResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x13\n\x0bsignal_name\x18\x03 \x01(\t\x12\'\n\x0ctrajectories\x18\x04 \x03(\x0b\x32\x11.SignalTrajectory\"Y\n\x11PointCloudRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x12\n\nmax_points\x18\x03 \x01(\x05\x12\r\n\x05\x66ield\x18\x04 \x01(\t\"\xbf\x01\n\x0fPointCloudChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x12\n\nnum_points\x18\x03 \x01(\x05\x12\x14\n\x0cnum_features\x18\x04 \x01(\x05\x12\x10\n\x08pc_range\x18\x05 \x03(\x02\x12\x0c\n\x04\x64\x61ta\x18\x06 \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\x07 \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x08 \x01(\x05\x12\x15\n\rfeature_names\x18\t \x03(\t\"b\n\x0cMediaRequest\x12\x11\n\tsample_id\x18\x01 \x01(\t\x12\x0e\n\x06origin\x18\x02 \x01(\t\x12\x0c\n\x04kind\x18\x03 \x01(\t\x12\x12\n\nmax_frames\x18\x04 \x01(\x05\x12\r\n\x05\x66ield\x18\x05 \x01(\t\"\x92\x02\n\nMediaChunk\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\x12\x11\n\tmime_type\x18\x03 \x01(\t\x12\x13\n\x0b\x66rame_count\x18\x04 \x01(\x05\x12\x0b\n\x03\x66ps\x18\x05 \x01(\x02\x12\x11\n\thas_audio\x18\x06 \x01(\x08\x12\r\n\x05width\x18\x07 \x01(\x05\x12\x0e\n\x06height\x18\x08 \x01(\x05\x12\x13\n\x0btotal_bytes\x18\t \x01(\x05\x12\x18\n\x10\x64uration_seconds\x18\n \x01(\x02\x12\x13\n\x0bsample_rate\x18\x0b \x01(\x05\x12\x0c\n\x04\x64\x61ta\x18\x0c \x01(\x0c\x12\x13\n\x0b\x63hunk_index\x18\r \x01(\x05\x12\x14\n\x0ctotal_chunks\x18\x0e \x01(\x05\"\xdc\x01\n\x10\x44\x61taEditsRequest\x12\x11\n\tstat_name\x18\x01 \x01(\t\x12\x13\n\x0b\x66loat_value\x18\x02 \x01(\x02\x12\x14\n\x0cstring_value\x18\x03 \x01(\t\x12\x12\n\nbool_value\x18\x04 \x01(\x08\x12\x1d\n\x04type\x18\x05 \x01(\x0e\x32\x0f.SampleEditType\x12\x13\n\x0bsamples_ids\x18\x06 \x03(\t\x12\x16\n\x0esample_origins\x18\x07 \x03(\t\x12\x16\n\x0eis_categorical\x18\x08 \x01(\x08\x12\x12\n\ncategories\x18\t \x03(\t\"5\n\x11\x44\x61taEditsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\":\n\x12\x44\x61taSplitsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x13\n\x0bsplit_names\x18\x02 \x03(\t\"9\n\x13\x41gentHealthResponse\x12\x11\n\tavailable\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"^\n\x16InitializeAgentRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12$\n\x08provider\x18\x02 \x01(\x0e\x32\x12.AgentProviderType\x12\r\n\x05model\x18\x03 \x01(\t\";\n\x17InitializeAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"(\n\x17\x43hangeAgentModelRequest\x12\r\n\x05model\x18\x01 \x01(\t\"<\n\x18\x43hangeAgentModelResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x17\n\x15GetAgentModelsRequest\"J\n\x16GetAgentModelsResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0e\n\x06models\x18\x02 \x03(\t\x12\x0f\n\x07message\x18\x03 \x01(\t\"6\n\x12ResetAgentResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"3\n\x18RestoreCheckpointRequest\x12\x17\n\x0f\x65xperiment_hash\x18\x01 \x01(\t\"=\n\x19RestoreCheckpointResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"R\n\x18TriggerEvaluationRequest\x12\x12\n\nsplit_name\x18\x01 \x01(\t\x12\x0c\n\x04tags\x18\x02 \x03(\t\x12\x14\n\x0cuse_full_set\x18\x03 \x01(\x08\"=\n\x19TriggerEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"\x1c\n\x1aGetEvaluationStatusRequest\"\x81\x01\n\x1bGetEvaluationStatusResponse\x12\x0e\n\x06status\x18\x01 \x01(\t\x12\x0f\n\x07\x63urrent\x18\x02 \x01(\x05\x12\r\n\x05total\x18\x03 \x01(\x05\x12\x0f\n\x07message\x18\x04 \x01(\t\x12\r\n\x05\x65rror\x18\x05 \x01(\t\x12\x12\n\nsplit_name\x18\x06 \x01(\t\")\n\x17\x43\x61ncelEvaluationRequest\x12\x0e\n\x06reason\x18\x01 \x01(\t\"<\n\x18\x43\x61ncelEvaluationResponse\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0f\n\x07message\x18\x02 \x01(\t\"7\n\x16RunNotebookCellRequest\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x0f\n\x07\x63\x65ll_id\x18\x02 \x01(\t\"2\n\x10NotebookCellDone\x12\x12\n\nexec_count\x18\x01 \x01(\x05\x12\n\n\x02ok\x18\x02 \x01(\x08\"\x1e\n\x1cInterruptNotebookCellRequest\":\n\x1dInterruptNotebookCellResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\r\n\x05\x65rror\x18\x02 \x01(\t\"\xbd\x01\n\x11NotebookCellChunk\x12\x0f\n\x07\x63\x65ll_id\x18\x01 \x01(\t\x12\x10\n\x06stdout\x18\x02 \x01(\tH\x00\x12\x10\n\x06stderr\x18\x03 \x01(\tH\x00\x12\x15\n\x0bresult_text\x18\x04 \x01(\tH\x00\x12\x13\n\timage_png\x18\x05 \x01(\x0cH\x00\x12\x19\n\x0f\x65rror_traceback\x18\x06 \x01(\tH\x00\x12!\n\x04\x64one\x18\x07 \x01(\x0b\x32\x11.NotebookCellDoneH\x00\x42\t\n\x07payload\"S\n\x10NotebookResponse\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0f\n\x07\x65xisted\x18\x02 \x01(\x08\x12\x0c\n\x04path\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"7\n\x13SaveNotebookRequest\x12\x12\n\nipynb_json\x18\x01 \x01(\t\x12\x0c\n\x04name\x18\x02 \x01(\t\"M\n\x14SaveNotebookResponse\x12\n\n\x02ok\x18\x01 \x01(\x08\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\r\n\x05\x65rror\x18\x03 \x01(\t\x12\x0c\n\x04name\x18\x04 \x01(\t\"C\n\x1bGenerateNotebookCodeRequest\x12\x0e\n\x06prompt\x18\x01 \x01(\t\x12\x14\n\x0c\x63ontext_code\x18\x02 \x01(\t\"\\\n\x1cGenerateNotebookCodeResponse\x12\x0c\n\x04\x63ode\x18\x01 \x01(\t\x12\x13\n\x0b\x65xplanation\x18\x02 \x01(\t\x12\n\n\x02ok\x18\x03 \x01(\x08\x12\r\n\x05\x65rror\x18\x04 \x01(\t*d\n\x13WeightOperationType\x12\n\n\x06ZEROFY\x10\x00\x12\x10\n\x0cREINITIALIZE\x10\x01\x12\n\n\x06\x46REEZE\x10\x02\x12\x12\n\x0eREMOVE_NEURONS\x10\t\x12\x0f\n\x0b\x41\x44\x44_NEURONS\x10\n*o\n\x0fZerofyPredicate\x12\x19\n\x15ZEROFY_PREDICATE_NONE\x10\x00\x12 \n\x1cZEROFY_PREDICATE_WITH_FROZEN\x10\x01\x12\x1f\n\x1bZEROFY_PREDICATE_WITH_OLDER\x10\x02*M\n\x0f\x41gentIntentType\x12\x12\n\x0eINTENT_UNKNOWN\x10\x00\x12\x11\n\rINTENT_FILTER\x10\x01\x12\x13\n\x0fINTENT_ANALYSIS\x10\x02*I\n\x0eSampleEditType\x12\x11\n\rEDIT_OVERRIDE\x10\x00\x12\x13\n\x0f\x45\x44IT_ACCUMULATE\x10\x01\x12\x0f\n\x0b\x45\x44IT_REMOVE\x10\x02*,\n\x11\x41gentProviderType\x12\x17\n\x13PROVIDER_OPENROUTER\x10\x00\x32\xc7\x0e\n\x11\x45xperimentService\x12P\n\x13GetLatestLoggerData\x12\x1b.GetLatestLoggerDataRequest\x1a\x1c.GetLatestLoggerDataResponse\x12\x36\n\x11\x45xperimentCommand\x12\x0f.TrainerCommand\x1a\x10.CommandResponse\x12H\n\x11ManipulateWeights\x12\x18.WeightsOperationRequest\x1a\x19.WeightsOperationResponse\x12/\n\nGetWeights\x12\x0f.WeightsRequest\x1a\x10.WeightsResponse\x12\x39\n\x0eGetActivations\x12\x12.ActivationRequest\x1a\x13.ActivationResponse\x12\x37\n\nGetSamples\x12\x13.BatchSampleRequest\x1a\x14.BatchSampleResponse\x12\x37\n\x0e\x41pplyDataQuery\x12\x11.DataQueryRequest\x1a\x12.DataQueryResponse\x12;\n\x0eGetDataSamples\x12\x13.DataSamplesRequest\x1a\x14.DataSamplesResponse\x12\x35\n\x0cGetHistogram\x12\x11.HistogramRequest\x1a\x12.HistogramResponse\x12\x38\n\x0bGetMetaData\x12\x13.GetMetaDataRequest\x1a\x14.GetMetaDataResponse\x12P\n\x13GetSignalTrajectory\x12\x1b.GetSignalTrajectoryRequest\x1a\x1c.GetSignalTrajectoryResponse\x12\x37\n\rGetPointCloud\x12\x12.PointCloudRequest\x1a\x10.PointCloudChunk0\x01\x12(\n\x08GetMedia\x12\r.MediaRequest\x1a\x0b.MediaChunk0\x01\x12\x37\n\x0e\x45\x64itDataSample\x12\x11.DataEditsRequest\x1a\x12.DataEditsResponse\x12,\n\rGetDataSplits\x12\x06.Empty\x1a\x13.DataSplitsResponse\x12\x30\n\x10\x43heckAgentHealth\x12\x06.Empty\x1a\x14.AgentHealthResponse\x12\x44\n\x0fInitializeAgent\x12\x17.InitializeAgentRequest\x1a\x18.InitializeAgentResponse\x12G\n\x10\x43hangeAgentModel\x12\x18.ChangeAgentModelRequest\x1a\x19.ChangeAgentModelResponse\x12\x41\n\x0eGetAgentModels\x12\x16.GetAgentModelsRequest\x1a\x17.GetAgentModelsResponse\x12)\n\nResetAgent\x12\x06.Empty\x1a\x13.ResetAgentResponse\x12@\n\x0fRunNotebookCell\x12\x17.RunNotebookCellRequest\x1a\x12.NotebookCellChunk0\x01\x12V\n\x15InterruptNotebookCell\x12\x1d.InterruptNotebookCellRequest\x1a\x1e.InterruptNotebookCellResponse\x12(\n\x0bGetNotebook\x12\x06.Empty\x1a\x11.NotebookResponse\x12;\n\x0cSaveNotebook\x12\x14.SaveNotebookRequest\x1a\x15.SaveNotebookResponse\x12S\n\x14GenerateNotebookCode\x12\x1c.GenerateNotebookCodeRequest\x1a\x1d.GenerateNotebookCodeResponse\x12J\n\x11RestoreCheckpoint\x12\x19.RestoreCheckpointRequest\x1a\x1a.RestoreCheckpointResponse\x12J\n\x11TriggerEvaluation\x12\x19.TriggerEvaluationRequest\x1a\x1a.TriggerEvaluationResponse\x12P\n\x13GetEvaluationStatus\x12\x1b.GetEvaluationStatusRequest\x1a\x1c.GetEvaluationStatusResponse\x12G\n\x10\x43\x61ncelEvaluation\x12\x18.CancelEvaluationRequest\x1a\x19.CancelEvaluationResponseb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -37,16 +37,16 @@ _globals['_NEURONSTATISTICS_INCOMINGLRENTRY']._serialized_options = b'8\001' _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._loaded_options = None _globals['_RECORDMETADATA_SAMPLELASTLOSSENTRY']._serialized_options = b'8\001' - _globals['_WEIGHTOPERATIONTYPE']._serialized_start=10979 - _globals['_WEIGHTOPERATIONTYPE']._serialized_end=11079 - _globals['_ZEROFYPREDICATE']._serialized_start=11081 - _globals['_ZEROFYPREDICATE']._serialized_end=11192 - _globals['_AGENTINTENTTYPE']._serialized_start=11194 - _globals['_AGENTINTENTTYPE']._serialized_end=11271 - _globals['_SAMPLEEDITTYPE']._serialized_start=11273 - _globals['_SAMPLEEDITTYPE']._serialized_end=11346 - _globals['_AGENTPROVIDERTYPE']._serialized_start=11348 - _globals['_AGENTPROVIDERTYPE']._serialized_end=11392 + _globals['_WEIGHTOPERATIONTYPE']._serialized_start=11371 + _globals['_WEIGHTOPERATIONTYPE']._serialized_end=11471 + _globals['_ZEROFYPREDICATE']._serialized_start=11473 + _globals['_ZEROFYPREDICATE']._serialized_end=11584 + _globals['_AGENTINTENTTYPE']._serialized_start=11586 + _globals['_AGENTINTENTTYPE']._serialized_end=11663 + _globals['_SAMPLEEDITTYPE']._serialized_start=11665 + _globals['_SAMPLEEDITTYPE']._serialized_end=11738 + _globals['_AGENTPROVIDERTYPE']._serialized_start=11740 + _globals['_AGENTPROVIDERTYPE']._serialized_end=11784 _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_start=46 _globals['_GETLATESTLOGGERDATAREQUEST']._serialized_end=183 _globals['_LOGGERDATAPOINT']._serialized_start=186 @@ -158,67 +158,71 @@ _globals['_GETSIGNALTRAJECTORYRESPONSE']._serialized_start=8460 _globals['_GETSIGNALTRAJECTORYRESPONSE']._serialized_end=8585 _globals['_POINTCLOUDREQUEST']._serialized_start=8587 - _globals['_POINTCLOUDREQUEST']._serialized_end=8661 - _globals['_POINTCLOUDCHUNK']._serialized_start=8664 - _globals['_POINTCLOUDCHUNK']._serialized_end=8855 - _globals['_DATAEDITSREQUEST']._serialized_start=8858 - _globals['_DATAEDITSREQUEST']._serialized_end=9078 - _globals['_DATAEDITSRESPONSE']._serialized_start=9080 - _globals['_DATAEDITSRESPONSE']._serialized_end=9133 - _globals['_DATASPLITSRESPONSE']._serialized_start=9135 - _globals['_DATASPLITSRESPONSE']._serialized_end=9193 - _globals['_AGENTHEALTHRESPONSE']._serialized_start=9195 - _globals['_AGENTHEALTHRESPONSE']._serialized_end=9252 - _globals['_INITIALIZEAGENTREQUEST']._serialized_start=9254 - _globals['_INITIALIZEAGENTREQUEST']._serialized_end=9348 - _globals['_INITIALIZEAGENTRESPONSE']._serialized_start=9350 - _globals['_INITIALIZEAGENTRESPONSE']._serialized_end=9409 - _globals['_CHANGEAGENTMODELREQUEST']._serialized_start=9411 - _globals['_CHANGEAGENTMODELREQUEST']._serialized_end=9451 - _globals['_CHANGEAGENTMODELRESPONSE']._serialized_start=9453 - _globals['_CHANGEAGENTMODELRESPONSE']._serialized_end=9513 - _globals['_GETAGENTMODELSREQUEST']._serialized_start=9515 - _globals['_GETAGENTMODELSREQUEST']._serialized_end=9538 - _globals['_GETAGENTMODELSRESPONSE']._serialized_start=9540 - _globals['_GETAGENTMODELSRESPONSE']._serialized_end=9614 - _globals['_RESETAGENTRESPONSE']._serialized_start=9616 - _globals['_RESETAGENTRESPONSE']._serialized_end=9670 - _globals['_RESTORECHECKPOINTREQUEST']._serialized_start=9672 - _globals['_RESTORECHECKPOINTREQUEST']._serialized_end=9723 - _globals['_RESTORECHECKPOINTRESPONSE']._serialized_start=9725 - _globals['_RESTORECHECKPOINTRESPONSE']._serialized_end=9786 - _globals['_TRIGGEREVALUATIONREQUEST']._serialized_start=9788 - _globals['_TRIGGEREVALUATIONREQUEST']._serialized_end=9870 - _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_start=9872 - _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_end=9933 - _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_start=9935 - _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_end=9963 - _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_start=9966 - _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_end=10095 - _globals['_CANCELEVALUATIONREQUEST']._serialized_start=10097 - _globals['_CANCELEVALUATIONREQUEST']._serialized_end=10138 - _globals['_CANCELEVALUATIONRESPONSE']._serialized_start=10140 - _globals['_CANCELEVALUATIONRESPONSE']._serialized_end=10200 - _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_start=10202 - _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_end=10257 - _globals['_NOTEBOOKCELLDONE']._serialized_start=10259 - _globals['_NOTEBOOKCELLDONE']._serialized_end=10309 - _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_start=10311 - _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_end=10341 - _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_start=10343 - _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_end=10401 - _globals['_NOTEBOOKCELLCHUNK']._serialized_start=10404 - _globals['_NOTEBOOKCELLCHUNK']._serialized_end=10593 - _globals['_NOTEBOOKRESPONSE']._serialized_start=10595 - _globals['_NOTEBOOKRESPONSE']._serialized_end=10678 - _globals['_SAVENOTEBOOKREQUEST']._serialized_start=10680 - _globals['_SAVENOTEBOOKREQUEST']._serialized_end=10735 - _globals['_SAVENOTEBOOKRESPONSE']._serialized_start=10737 - _globals['_SAVENOTEBOOKRESPONSE']._serialized_end=10814 - _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_start=10816 - _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_end=10883 - _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_start=10885 - _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_end=10977 - _globals['_EXPERIMENTSERVICE']._serialized_start=11395 - _globals['_EXPERIMENTSERVICE']._serialized_end=13216 + _globals['_POINTCLOUDREQUEST']._serialized_end=8676 + _globals['_POINTCLOUDCHUNK']._serialized_start=8679 + _globals['_POINTCLOUDCHUNK']._serialized_end=8870 + _globals['_MEDIAREQUEST']._serialized_start=8872 + _globals['_MEDIAREQUEST']._serialized_end=8970 + _globals['_MEDIACHUNK']._serialized_start=8973 + _globals['_MEDIACHUNK']._serialized_end=9247 + _globals['_DATAEDITSREQUEST']._serialized_start=9250 + _globals['_DATAEDITSREQUEST']._serialized_end=9470 + _globals['_DATAEDITSRESPONSE']._serialized_start=9472 + _globals['_DATAEDITSRESPONSE']._serialized_end=9525 + _globals['_DATASPLITSRESPONSE']._serialized_start=9527 + _globals['_DATASPLITSRESPONSE']._serialized_end=9585 + _globals['_AGENTHEALTHRESPONSE']._serialized_start=9587 + _globals['_AGENTHEALTHRESPONSE']._serialized_end=9644 + _globals['_INITIALIZEAGENTREQUEST']._serialized_start=9646 + _globals['_INITIALIZEAGENTREQUEST']._serialized_end=9740 + _globals['_INITIALIZEAGENTRESPONSE']._serialized_start=9742 + _globals['_INITIALIZEAGENTRESPONSE']._serialized_end=9801 + _globals['_CHANGEAGENTMODELREQUEST']._serialized_start=9803 + _globals['_CHANGEAGENTMODELREQUEST']._serialized_end=9843 + _globals['_CHANGEAGENTMODELRESPONSE']._serialized_start=9845 + _globals['_CHANGEAGENTMODELRESPONSE']._serialized_end=9905 + _globals['_GETAGENTMODELSREQUEST']._serialized_start=9907 + _globals['_GETAGENTMODELSREQUEST']._serialized_end=9930 + _globals['_GETAGENTMODELSRESPONSE']._serialized_start=9932 + _globals['_GETAGENTMODELSRESPONSE']._serialized_end=10006 + _globals['_RESETAGENTRESPONSE']._serialized_start=10008 + _globals['_RESETAGENTRESPONSE']._serialized_end=10062 + _globals['_RESTORECHECKPOINTREQUEST']._serialized_start=10064 + _globals['_RESTORECHECKPOINTREQUEST']._serialized_end=10115 + _globals['_RESTORECHECKPOINTRESPONSE']._serialized_start=10117 + _globals['_RESTORECHECKPOINTRESPONSE']._serialized_end=10178 + _globals['_TRIGGEREVALUATIONREQUEST']._serialized_start=10180 + _globals['_TRIGGEREVALUATIONREQUEST']._serialized_end=10262 + _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_start=10264 + _globals['_TRIGGEREVALUATIONRESPONSE']._serialized_end=10325 + _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_start=10327 + _globals['_GETEVALUATIONSTATUSREQUEST']._serialized_end=10355 + _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_start=10358 + _globals['_GETEVALUATIONSTATUSRESPONSE']._serialized_end=10487 + _globals['_CANCELEVALUATIONREQUEST']._serialized_start=10489 + _globals['_CANCELEVALUATIONREQUEST']._serialized_end=10530 + _globals['_CANCELEVALUATIONRESPONSE']._serialized_start=10532 + _globals['_CANCELEVALUATIONRESPONSE']._serialized_end=10592 + _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_start=10594 + _globals['_RUNNOTEBOOKCELLREQUEST']._serialized_end=10649 + _globals['_NOTEBOOKCELLDONE']._serialized_start=10651 + _globals['_NOTEBOOKCELLDONE']._serialized_end=10701 + _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_start=10703 + _globals['_INTERRUPTNOTEBOOKCELLREQUEST']._serialized_end=10733 + _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_start=10735 + _globals['_INTERRUPTNOTEBOOKCELLRESPONSE']._serialized_end=10793 + _globals['_NOTEBOOKCELLCHUNK']._serialized_start=10796 + _globals['_NOTEBOOKCELLCHUNK']._serialized_end=10985 + _globals['_NOTEBOOKRESPONSE']._serialized_start=10987 + _globals['_NOTEBOOKRESPONSE']._serialized_end=11070 + _globals['_SAVENOTEBOOKREQUEST']._serialized_start=11072 + _globals['_SAVENOTEBOOKREQUEST']._serialized_end=11127 + _globals['_SAVENOTEBOOKRESPONSE']._serialized_start=11129 + _globals['_SAVENOTEBOOKRESPONSE']._serialized_end=11206 + _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_start=11208 + _globals['_GENERATENOTEBOOKCODEREQUEST']._serialized_end=11275 + _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_start=11277 + _globals['_GENERATENOTEBOOKCODERESPONSE']._serialized_end=11369 + _globals['_EXPERIMENTSERVICE']._serialized_start=11787 + _globals['_EXPERIMENTSERVICE']._serialized_end=13650 # @@protoc_insertion_point(module_scope) diff --git a/weightslab/proto/experiment_service_pb2_grpc.py b/weightslab/proto/experiment_service_pb2_grpc.py index 2f9127b2..7625d40a 100644 --- a/weightslab/proto/experiment_service_pb2_grpc.py +++ b/weightslab/proto/experiment_service_pb2_grpc.py @@ -5,7 +5,7 @@ from weightslab.proto import experiment_service_pb2 as weightslab_dot_proto_dot_experiment__service__pb2 -GRPC_GENERATED_VERSION = '1.76.0' +GRPC_GENERATED_VERSION = '1.74.0' GRPC_VERSION = grpc.__version__ _version_not_supported = False @@ -18,7 +18,7 @@ if _version_not_supported: raise RuntimeError( f'The grpc package installed is at version {GRPC_VERSION},' - + ' but the generated code in weightslab/proto/experiment_service_pb2_grpc.py depends on' + + f' but the generated code in weightslab/proto/experiment_service_pb2_grpc.py depends on' + f' grpcio>={GRPC_GENERATED_VERSION}.' + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}' + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.' @@ -94,6 +94,11 @@ def __init__(self, channel): request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudRequest.SerializeToString, response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudChunk.FromString, _registered_method=True) + self.GetMedia = channel.unary_stream( + '/ExperimentService/GetMedia', + request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.MediaRequest.SerializeToString, + response_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.MediaChunk.FromString, + _registered_method=True) self.EditDataSample = channel.unary_unary( '/ExperimentService/EditDataSample', request_serializer=weightslab_dot_proto_dot_experiment__service__pb2.DataEditsRequest.SerializeToString, @@ -264,6 +269,17 @@ def GetPointCloud(self, request, context): context.set_details('Method not implemented!') raise NotImplementedError('Method not implemented!') + def GetMedia(self, request, context): + """Playable media of one sample (task_type "video_generation" / + "audio_generation"), server-streamed in binary chunks for the modal + player. Grid and list previews never call this — they use the poster + frame carried by the ordinary raw_data stat — so clip length does not + affect browsing cost. + """ + context.set_code(grpc.StatusCode.UNIMPLEMENTED) + context.set_details('Method not implemented!') + raise NotImplementedError('Method not implemented!') + def EditDataSample(self, request, context): """Missing associated documentation comment in .proto file.""" context.set_code(grpc.StatusCode.UNIMPLEMENTED) @@ -433,6 +449,11 @@ def add_ExperimentServiceServicer_to_server(servicer, server): request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudRequest.FromString, response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.PointCloudChunk.SerializeToString, ), + 'GetMedia': grpc.unary_stream_rpc_method_handler( + servicer.GetMedia, + request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.MediaRequest.FromString, + response_serializer=weightslab_dot_proto_dot_experiment__service__pb2.MediaChunk.SerializeToString, + ), 'EditDataSample': grpc.unary_unary_rpc_method_handler( servicer.EditDataSample, request_deserializer=weightslab_dot_proto_dot_experiment__service__pb2.DataEditsRequest.FromString, @@ -848,6 +869,33 @@ def GetPointCloud(request, metadata, _registered_method=True) + @staticmethod + def GetMedia(request, + target, + options=(), + channel_credentials=None, + call_credentials=None, + insecure=False, + compression=None, + wait_for_ready=None, + timeout=None, + metadata=None): + return grpc.experimental.unary_stream( + request, + target, + '/ExperimentService/GetMedia', + weightslab_dot_proto_dot_experiment__service__pb2.MediaRequest.SerializeToString, + weightslab_dot_proto_dot_experiment__service__pb2.MediaChunk.FromString, + options, + channel_credentials, + insecure, + call_credentials, + compression, + wait_for_ready, + timeout, + metadata, + _registered_method=True) + @staticmethod def EditDataSample(request, target, diff --git a/weightslab/src.py b/weightslab/src.py index 62dd5290..af80ab5c 100644 --- a/weightslab/src.py +++ b/weightslab/src.py @@ -2782,6 +2782,206 @@ def save_group_signals( DATAFRAME_M.update_by_groups_bulk(origin=origin, group_ids=active_group_ids, updates_list=all_updates) +def _encode_one_media(item, kind: str, fps: float, audio, sample_rate: int, + dataset=None): + """Encode one media item for transport. Returns (data, mime, poster, meta). + + ``poster`` is the still the grid and list render; the full payload is only + fetched when the user opens the sample. + """ + import io as _io + + import numpy as _np + from PIL import Image as _Image + + from weightslab.data import media_store as _ms + from weightslab.data import video_utils as _vu + from weightslab.trainer.services.data_image_utils import encode_image_webp + + def _to_numpy(value): + if hasattr(value, "detach"): + value = value.detach().cpu() + if hasattr(value, "numpy"): + value = value.numpy() + return _np.asarray(value) + + if kind == _ms.KIND_TEXT: + # The item IS the text (a caption, a reference completion, ...) — no + # array conversion, no poster: the studio shows a "" tag, same + # as an unstillable audio field. + text = "" if item is None else str(item) + data = text.encode("utf-8") + return data, "text/plain", b"", {"length": len(text)} + + arr = _to_numpy(item) if item is not None else None + + if kind == _ms.KIND_AUDIO: + data = _vu.encode_audio_wav(arr, sample_rate) + duration = (len(arr) / float(sample_rate)) if sample_rate else 0.0 + # Audio has no still; the list shows a "♪" chip instead of a thumbnail. + return data, "audio/wav", b"", { + "sample_rate": int(sample_rate or 0), + "duration_seconds": round(float(duration), 3), + } + + if kind == _ms.KIND_POINTCLOUD: + from weightslab.data.point_cloud_utils import ( + get_pc_range, get_point_feature_names, pack_point_cloud, + render_thumbnail_2d_for_dataset, + ) + points = _np.asarray(arr, dtype=_np.float32) + data, num_points, num_features = pack_point_cloud(points) + poster = b"" + try: + poster = encode_image_webp( + render_thumbnail_2d_for_dataset(dataset, points), quality=70) + except Exception as exc: + logger.debug("point-cloud poster render failed: %r", exc) + # pc_range/feature_names are what the interactive 3D viewer needs to + # frame the camera and offer the right colour modes; without them here + # a field-streamed cloud (GetPointCloud with `field` set) would open + # with an unframed, unlabeled scene. + pc_range = get_pc_range(dataset, points) or () + feature_names = get_point_feature_names(dataset, num_features) + return data, "application/octet-stream", poster, { + "num_points": int(num_points), + "num_features": int(num_features), + "pc_range": [float(v) for v in pc_range], + "feature_names": list(feature_names), + } + + if kind == _ms.KIND_VIDEO: + # [T, C, H, W] is the common torch layout; normalize to [T, H, W, C]. + if arr.ndim == 4 and arr.shape[1] in (1, 3, 4) and arr.shape[1] < min(arr.shape[2], arr.shape[3]): + arr = _np.transpose(arr, (0, 2, 3, 1)) + data, mime, has_audio = _vu.encode_clip( + arr, fps, audio=audio, sample_rate=sample_rate) + poster = b"" + try: + poster = encode_image_webp(_vu.video_poster_frame(dataset, arr), quality=70) + except Exception as exc: + logger.debug("video poster render failed: %r", exc) + return data, mime, poster, { + "frame_count": int(arr.shape[0]), + "fps": float(fps), + "has_audio": bool(has_audio), + "width": int(arr.shape[2]), + "height": int(arr.shape[1]), + "duration_seconds": round(float(arr.shape[0]) / max(1.0, float(fps)), 3), + } + + # image / mask: a single still. [C, H, W] -> [H, W, C]. + if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[0] < min(arr.shape[1], arr.shape[2]): + arr = _np.transpose(arr, (1, 2, 0)) + frame = _vu._frames_to_uint8(arr[None])[0] + pil = _vu._frame_to_pil(frame) + buffer = _io.BytesIO() + pil.save(buffer, format="PNG") + data = buffer.getvalue() + poster = encode_image_webp(pil.convert("RGB"), quality=70) + return data, "image/png", poster, { + "width": int(pil.width), "height": int(pil.height), + } + + +def save_media( + field: str, + batch_ids, + media, + kind: str = "image", + fps: float = 8.0, + audio=None, + sample_rate: int = 0, + origin: str | None = None, + dataset=None, +): + """Attach media to a metadata field so the studio can visualize it. + + This is what makes a **generated** result viewable next to its input: any + field — not just the sample's own input — can carry an image, a mask, a + video clip, an audio track or a point cloud. In the studio the field becomes + a column showing a thumbnail per row (sized to the list's row height), and + clicking a thumbnail opens the full viewer: zoomable image, mask overlay, + video player with frame scrubbing and sound, or the interactive 3D cloud. + + Args: + field: metadata field name, e.g. "pred_video". Stored in the dataframe + as the column ``media:pred_video``. + batch_ids: the sample ids this media belongs to (same order as ``media``). + media: a sequence, one entry per id. Video entries are ``[T, H, W, C]`` + (or ``[T, C, H, W]``); images/masks are ``[H, W]``/``[H, W, C]``; + point clouds are ``[N, F]`` float; audio is a 1-D waveform; text + entries are plain strings. + kind: one of "image", "mask", "video", "audio", "pointcloud", "text". + fps: frame rate for video entries. + audio: optional per-sample waveforms muxed into video entries; pass a + single waveform to reuse it for every sample. + sample_rate: sample rate for ``audio`` (required for it to be muxed). + origin: loader name; defaults to the active one. + dataset: optional dataset, used only for its poster/thumbnail hooks. + + Example: + >>> generated = to_clip_space(sampled) # [B, T, H, W, C] uint8 + >>> wl.save_media("pred_video", batch_ids=uids, media=generated, + ... kind="video", fps=8, audio=waveforms, sample_rate=16000) + """ + from weightslab.data import media_store + + global DATAFRAME_M + if DATAFRAME_M is None: + DATAFRAME_M = get_dataframe() + if DATAFRAME_M is None: + logger.warning("save_media: no dataframe registered; skipping.") + return + + if batch_ids is None or media is None: + return + if isinstance(batch_ids, th.Tensor): + ids = [str(i) for i in batch_ids.detach().cpu().tolist()] + else: + ids = [str(i) for i in batch_ids] + if len(ids) == 0: + return + + origin = origin or get_active_origin() or "train_loader" + column = media_store.media_column(field) + + # A single waveform is broadcast to every sample; a sequence is per-sample. + audio_is_per_sample = ( + audio is not None + and hasattr(audio, "__len__") + and len(audio) == len(ids) + and hasattr(audio[0], "__len__") + ) + + for index, sample_id in enumerate(ids): + if index >= len(media): + break + try: + item_audio = audio[index] if audio_is_per_sample else audio + data, mime, poster, meta = _encode_one_media( + media[index], kind=kind, fps=fps, audio=item_audio, + sample_rate=sample_rate, dataset=dataset) + if not data: + continue + entry = { + "field": field, "sample_id": sample_id, "data": data, + "mime": mime, "kind": kind, "poster": poster, "meta": meta, + } + media_store.put(field, sample_id, data, mime, kind, + poster=poster, meta=meta) + # Only the small descriptor goes in the dataframe; the bytes stay in + # the media store and are streamed on demand by GetMedia. + DATAFRAME_M.update_values( + origin=origin, + sample_id=int(sample_id) if str(sample_id).isdigit() else sample_id, + updates={column: media_store.descriptor_json(entry)}, + ) + except Exception as exc: + logger.warning("save_media(%s) failed for sample %s: %r", + field, sample_id, exc) + + def clear_all(): """Clear all WeightsLab registries (models, dataloaders, etc.).""" ledgers.clear_all() diff --git a/weightslab/trainer/services/data_service.py b/weightslab/trainer/services/data_service.py index 64192cfd..2fa78b73 100755 --- a/weightslab/trainer/services/data_service.py +++ b/weightslab/trainer/services/data_service.py @@ -38,6 +38,10 @@ POINT_CLOUD_DETECTION_TASK, is_point_cloud_detection_task, ) +from weightslab.data.video_utils import ( + describe_clip, is_generation_task, is_video_task, +) +from weightslab.data import media_store from weightslab.trainer.trainer_tools import execute_df_operation, generate_overview, encode_image_to_raw_bytes from weightslab.data.data_utils import load_raw_image_array @@ -109,6 +113,72 @@ def _point_cloud_chunk_bytes() -> int: return val +# Streamed chunk size for GetMedia. Smaller than the point-cloud chunk because +# these messages cross grpc-web to a browser, where oversized frames are the +# usual cause of stalled streams. Override with WL_MEDIA_CHUNK_BYTES. +_DEFAULT_MEDIA_CHUNK_BYTES = 1 << 18 # 256 KiB + + +def _media_chunk_bytes() -> int: + """Read WL_MEDIA_CHUNK_BYTES; non-positive/invalid falls back to the default.""" + return _env_int("WL_MEDIA_CHUNK_BYTES", _DEFAULT_MEDIA_CHUNK_BYTES) + + +# Task types whose target/prediction is never a spatial mask or box set, and so +# must skip the RLE/detection interpretation entirely. For generative tasks the +# target is a caption or a reference clip: RLE-encoding a [T, H, W, C] clip +# would be both meaningless and ruinously expensive, so they route to the +# text/scalar branch alongside classification and tabular. +_NON_MASK_TASKS = ("classification", "tabular") + + +def _is_non_mask_task(task_type) -> bool: + """True when labels/predictions for this task must not be read as masks.""" + return task_type in _NON_MASK_TASKS or is_generation_task(task_type) + + +def _build_media_stats(row) -> list: + """DataStats for every ``media:`` column present on this row. + + Emits stat type "media" carrying the descriptor in ``value_string`` and the + poster still in ``thumbnail``. Rows whose media has been evicted from the + store still get the descriptor, so the column keeps its shape; the frontend + just has no still to draw. + + A module-level function, not a method: it needs no instance state, and + binding it to the service would force every white-box test stub of + _process_sample_row to grow another attribute. + """ + stats = [] + try: + sample_id = row.get(SampleStatsEx.SAMPLE_ID.value) if hasattr(row, "get") else None + if sample_id is None: + # sample_id is the dataframe INDEX, not a column, on real rows. + sample_id = getattr(row, "name", None) + if isinstance(sample_id, tuple): + sample_id = sample_id[0] + keys = row.keys() if hasattr(row, "keys") else [] + for column in keys: + if not media_store.is_media_column(column): + continue + value = row.get(column) + if value is None or (isinstance(value, float) and np.isnan(value)): + continue + field = media_store.field_from_column(column) + stats.append( + create_data_stat( + name=str(column), + stat_type="media", + shape=[1], + value_string=str(value), + thumbnail=media_store.get_poster(field, sample_id), + ) + ) + except Exception as exc: + logger.debug("media stat build failed: %r", exc) + return stats + + def peek_is_tabular_sample(dataset, sample_id) -> bool: """True when ``sample_id``'s raw model input is a 1-D feature vector. @@ -488,6 +558,11 @@ def __init__(self, ctx): else: self._preview_cache_ready.set() # No preload → mark immediately ready + # Encoded-media LRU for GetMedia (insertion-ordered dict). Muxing a clip + # costs an ffmpeg round-trip, so opening, closing and re-opening the + # same sample — or paging its frames — must not re-encode. + self._media_cache: dict = {} + logger.info("DataService initialized.") @staticmethod @@ -575,7 +650,12 @@ def _build_preview_cache(self) -> None: member_rank = 0 _pv_np_img, _, _, pil_img = load_raw_image_array(dataset, ds_idx, rank=member_rank) - if looks_like_tabular(_pv_np_img): + if isinstance(_pv_np_img, str): + # Text-generation sample: cache the text itself, no image. + stats.append(create_data_stat( + 'raw_data', 'text', shape=[1], value_string=_pv_np_img)) + pil_img = None + elif looks_like_tabular(_pv_np_img): # Tabular sample: cache the feature values (lossless) as a # 'vector' raw_data stat + heatmap thumbnail, not an image. stats.append(build_tabular_raw_data_stat(_pv_np_img)) @@ -1432,7 +1512,7 @@ def _process_sample_row(self, args): # ====== Step 7: Process labels ====== t_label_convert = 0.0 - if not skip_label_for_request and task_type not in ("classification", "tabular"): + if not skip_label_for_request and not _is_non_mask_task(task_type): t0_gt_conv = time.time() label_raw = row.get(SampleStatsEx.TARGET.value) if label is None else label if label_raw is None and dataset is not None: @@ -1706,7 +1786,7 @@ def _json_default(o): if skip_prediction_for_request: pred = None - if task_type not in ("classification", "tabular") and pred is not None: + if not _is_non_mask_task(task_type) and pred is not None: t0_pmask = time.time() # 3D (point cloud) detection predictions: same dual payload as @@ -1889,10 +1969,22 @@ def _json_default(o): else: np_img, is_volumetric, original_shape, middle_pil = None, False, [], None + # Text-generation input: the sample's payload IS the text (a + # prompt, a conversation turn, ...) — there is no image to + # encode at all, lossy or otherwise. + if isinstance(np_img, str): + data_stats.append( + create_data_stat( + name='raw_data', + stat_type='text', + shape=[1], + value_string=np_img, + ) + ) # Tabular input: the model input is a 1-D feature vector, not an # image. Transmit the actual values (lossless) as a 'vector' # raw_data stat instead of a lossy WebP image. - if looks_like_tabular(np_img): + elif looks_like_tabular(np_img): data_stats.append(build_tabular_raw_data_stat(np_img)) elif middle_pil is not None: original_size = middle_pil.size @@ -1993,12 +2085,35 @@ def _json_default(o): ) ) + # For video samples the bytes above are only the poster + # frame, so advertise the clip's shape here. This lets the + # grid badge the cell and the modal size its player without + # anyone paying for a GetMedia round-trip first. + if is_video_task(task_type): + media_info = describe_clip(dataset) + if media_info: + data_stats.append( + create_data_stat( + name='media_info', + stat_type='string', + shape=[1], + value_string=json.dumps(media_info), + ) + ) + # Eagerly release intermediate image data to reduce peak memory # across concurrent thread-pool workers. del raw_data_bytes, middle_pil if np_img is not None: del np_img + # ====== Step 9b: Media attached to metadata fields ====== + # Any column written by wl.save_media becomes its own viewable + # column: the poster travels here (so the list can draw a thumbnail + # per row) while the payload stays in the store until the user opens + # it. This is what lets a generated clip sit beside its target. + data_stats.extend(_build_media_stats(row)) + # ====== Step 10: Create DataRecord ====== total_time = time.time() - start_total @@ -3801,7 +3916,13 @@ def _build_metadata_only_response(self, df_slice: pd.DataFrame, requested_cols=N # response to 100s of MB and silently break the histogram fetch. if not requested_cols: _HEAVY_BLOB_COLS = {"prediction_raw"} - requested_cols = [c for c in df_slice.columns if c not in _HEAVY_BLOB_COLS] + # media: columns are emitted as their own "media" stat (with a + # poster) by _build_media_stats; shipping them here too would add a + # second column showing the raw descriptor JSON as text. + requested_cols = [ + c for c in df_slice.columns + if c not in _HEAVY_BLOB_COLS and not media_store.is_media_column(c) + ] metadata_cols = [ col for col in requested_cols @@ -3926,6 +4047,9 @@ def _get_all_metadata_column_names(self) -> list: name = str(col) if col in _HEAVY_BLOB_COLS or col in _INTERNAL_COLS: continue + # Media columns are surfaced as "media" stats, not text columns. + if media_store.is_media_column(col): + continue if name in seen: continue seen.add(name) @@ -4824,6 +4948,191 @@ def GetHistogram(self, request, context): # configurable via the WL_POINT_CLOUD_CHUNK_BYTES env var (default 1 MiB). _POINT_CLOUD_CHUNK_BYTES = _point_cloud_chunk_bytes() + # Same, for GetMedia (already-compressed MP4/WAV bytes per message). + _MEDIA_CHUNK_BYTES = _media_chunk_bytes() + # How many encoded clips to keep. Clips are megabytes each, so this only + # needs to cover "the sample the user is currently looking at". + _MEDIA_CACHE_ENTRIES = 3 + + def _locate_sample(self, sample_id: int, origin: str = ""): + """Resolve ``sample_id`` to ``(dataset, ds_index, member_rank)``. + + Shared by the media streaming RPCs. When ``origin`` is empty every + known loader is scanned, so a sample can still be fetched when the + caller does not know which loader owns it. + Returns ``(None, None, 0)`` when the sample cannot be located. + """ + origins = [origin] if origin else list(get_dataloaders().keys()) + for candidate_origin in origins: + candidate = self._get_dataset(candidate_origin) + if candidate is None: + continue + try: + if hasattr(candidate, "get_physical_location"): + ds_idx, member_rank = candidate.get_physical_location(sample_id) + else: + ds_idx = candidate.get_index_from_sample_id(sample_id) + member_rank = 0 + return candidate, ds_idx, member_rank + except (KeyError, ValueError, AttributeError, IndexError): + continue + return None, None, 0 + + def GetMedia(self, request, context): + """Stream one sample's playable media (clip or soundtrack) in chunks. + + Used by the studio's modal player for generative media samples + (task_type "video_generation" / "audio_generation"). The clip is muxed + to MP4 (H.264 + AAC) on demand, capped by request ``max_frames`` / + env WL_MAX_VIDEO_FRAMES, and memoized in a small LRU so scrubbing and + re-opening the same sample does not re-encode. The first chunk carries + the mime type, geometry, fps and duration so the client can size the + player before any bytes arrive. + """ + from weightslab.data.data_utils import load_raw_video + from weightslab.data import video_utils as vu + + try: + sample_id = int(str(request.sample_id)) + kind = (request.kind or "video").strip().lower() + server_cap = int(os.environ.get("WL_MAX_VIDEO_FRAMES", "512")) + max_frames = int(request.max_frames) if request.max_frames > 0 else server_cap + if server_cap > 0: + max_frames = min(max_frames, server_cap) + + # A field request serves media attached to a metadata column by + # wl.save_media (e.g. a generated clip); it is already encoded, so + # it streams straight out of the store with no re-encode. + field = (request.field or "").strip() + if field: + entry = media_store.get(field, sample_id) + if entry is None: + yield pb2.MediaChunk( + success=False, + message=(f"No media stored for field {field!r} on sample " + f"{request.sample_id}"), + ) + return + meta = entry.get("meta") or {} + cached = { + "data": entry["data"], "mime": entry["mime"], + "frame_count": int(meta.get("frame_count", 0) or 0), + "fps": float(meta.get("fps", 0.0) or 0.0), + "has_audio": bool(meta.get("has_audio", False)), + "width": int(meta.get("width", 0) or 0), + "height": int(meta.get("height", 0) or 0), + "sample_rate": int(meta.get("sample_rate", 0) or 0), + "duration": float(meta.get("duration_seconds", 0.0) or 0.0), + } + yield from self._stream_media(cached) + return + + cache_key = (sample_id, request.origin or "", kind, max_frames) + cached = self._media_cache.get(cache_key) + if cached is None: + dataset, ds_idx, member_rank = self._locate_sample( + sample_id, request.origin) + if dataset is None or ds_idx is None: + yield pb2.MediaChunk( + success=False, + message=(f"Sample {request.sample_id} not found " + f"(origin={request.origin or 'any'})"), + ) + return + + frames = load_raw_video(dataset, ds_idx, rank=member_rank) + audio, sample_rate = vu.load_sample_audio(dataset, ds_idx) + + if kind == "audio": + if audio is None or not sample_rate: + yield pb2.MediaChunk( + success=False, + message=f"Sample {request.sample_id} has no audio track", + ) + return + data = vu.encode_audio_wav(audio, sample_rate) + cached = { + "data": data, "mime": "audio/wav", "frame_count": 0, + "fps": 0.0, "has_audio": True, "width": 0, "height": 0, + "sample_rate": sample_rate, + "duration": (len(audio) / float(sample_rate)) if sample_rate else 0.0, + } + else: + if frames is None: + yield pb2.MediaChunk( + success=False, + message=f"Sample {request.sample_id} is not a video", + ) + return + # Cap long clips by uniform subsampling so the playback + # still spans the whole clip rather than truncating it. + if max_frames > 0 and frames.shape[0] > max_frames: + picks = np.linspace( + 0, frames.shape[0] - 1, max_frames).astype(int) + frames = frames[picks] + fps = vu.get_fps(dataset) + data, mime, has_audio = vu.encode_clip( + frames, fps, audio=audio, sample_rate=sample_rate) + if not data: + yield pb2.MediaChunk( + success=False, + message=("Failed to encode clip — install ffmpeg " + "(pip install imageio-ffmpeg) for MP4 output"), + ) + return + cached = { + "data": data, "mime": mime, + "frame_count": int(frames.shape[0]), "fps": float(fps), + "has_audio": has_audio, + "width": int(frames.shape[2]), "height": int(frames.shape[1]), + "sample_rate": int(sample_rate or 0), + "duration": float(frames.shape[0]) / max(1.0, float(fps)), + } + self._media_cache_put(cache_key, cached) + + yield from self._stream_media(cached) + except Exception as e: + logger.error("Error in GetMedia: %s", str(e), exc_info=True) + yield pb2.MediaChunk( + success=False, + message=f"Failed to retrieve media: {str(e)}", + ) + + def _stream_media(self, payload: dict): + """Chunk one encoded payload out, header on the first chunk only.""" + data = payload["data"] + chunk_size = self._MEDIA_CHUNK_BYTES + total_chunks = max(1, (len(data) + chunk_size - 1) // chunk_size) + for i in range(total_chunks): + chunk = pb2.MediaChunk( + success=True, + data=data[i * chunk_size:(i + 1) * chunk_size], + chunk_index=i, + total_chunks=total_chunks, + ) + if i == 0: + chunk.mime_type = payload["mime"] + chunk.frame_count = payload["frame_count"] + chunk.fps = payload["fps"] + chunk.has_audio = payload["has_audio"] + chunk.width = payload["width"] + chunk.height = payload["height"] + chunk.total_bytes = len(data) + chunk.duration_seconds = payload["duration"] + chunk.sample_rate = payload["sample_rate"] + yield chunk + + def _media_cache_put(self, key, value) -> None: + """Insert into the encoded-media LRU, evicting the oldest entries. + + Encoded clips are large, so this is deliberately tiny — it exists to + make re-opening and scrubbing the *same* sample free, not to hold a + working set. + """ + self._media_cache[key] = value + while len(self._media_cache) > self._MEDIA_CACHE_ENTRIES: + self._media_cache.pop(next(iter(self._media_cache))) + def GetPointCloud(self, request, context): """Stream one sample's raw point cloud as binary float32 chunks. @@ -4840,26 +5149,33 @@ def GetPointCloud(self, request, context): try: sample_id = int(str(request.sample_id)) - origins = [request.origin] if request.origin else [] - if not origins: - # No origin provided: scan known loaders for the sample. - origins = list(get_dataloaders().keys()) - - dataset, ds_idx, member_rank = None, None, 0 - for origin in origins: - candidate = self._get_dataset(origin) - if candidate is None: - continue - try: - if hasattr(candidate, "get_physical_location"): - ds_idx, member_rank = candidate.get_physical_location(sample_id) - else: - ds_idx = candidate.get_index_from_sample_id(sample_id) - member_rank = 0 - dataset = candidate - break - except (KeyError, ValueError, AttributeError, IndexError): - continue + + # A field request serves a cloud attached to a metadata column by + # wl.save_media (e.g. a generated/predicted cloud); it is already + # packed, so it streams straight out of the store with no dataset + # lookup at all — mirrors GetMedia's field branch. + field = (request.field or "").strip() + if field: + entry = media_store.get(field, sample_id) + if entry is None: + yield pb2.PointCloudChunk( + success=False, + message=(f"No point cloud stored for field {field!r} on " + f"sample {request.sample_id}"), + ) + return + meta = entry.get("meta") or {} + yield from self._stream_point_cloud( + entry["data"], + num_points=int(meta.get("num_points", 0) or 0), + num_features=int(meta.get("num_features", 0) or 0), + pc_range=meta.get("pc_range") or (), + feature_names=meta.get("feature_names") or [], + ) + return + + dataset, ds_idx, member_rank = self._locate_sample( + sample_id, request.origin) if dataset is None or ds_idx is None: yield pb2.PointCloudChunk( @@ -4886,21 +5202,9 @@ def GetPointCloud(self, request, context): pc_range = get_pc_range(ds, points) or () feature_names = get_point_feature_names(ds, num_features) - chunk_size = self._POINT_CLOUD_CHUNK_BYTES - total_chunks = max(1, (len(data) + chunk_size - 1) // chunk_size) - for i in range(total_chunks): - chunk = pb2.PointCloudChunk( - success=True, - data=data[i * chunk_size:(i + 1) * chunk_size], - chunk_index=i, - total_chunks=total_chunks, - ) - if i == 0: - chunk.num_points = num_points - chunk.num_features = num_features - chunk.pc_range.extend(float(v) for v in pc_range) - chunk.feature_names.extend(feature_names) - yield chunk + yield from self._stream_point_cloud( + data, num_points=num_points, num_features=num_features, + pc_range=pc_range, feature_names=feature_names) except Exception as e: logger.error("Error in GetPointCloud: %s", str(e), exc_info=True) yield pb2.PointCloudChunk( @@ -4908,6 +5212,25 @@ def GetPointCloud(self, request, context): message=f"Failed to retrieve point cloud: {str(e)}", ) + def _stream_point_cloud(self, data: bytes, num_points: int, num_features: int, + pc_range, feature_names): + """Chunk one packed cloud out, header on the first chunk only.""" + chunk_size = self._POINT_CLOUD_CHUNK_BYTES + total_chunks = max(1, (len(data) + chunk_size - 1) // chunk_size) + for i in range(total_chunks): + chunk = pb2.PointCloudChunk( + success=True, + data=data[i * chunk_size:(i + 1) * chunk_size], + chunk_index=i, + total_chunks=total_chunks, + ) + if i == 0: + chunk.num_points = num_points + chunk.num_features = num_features + chunk.pc_range.extend(float(v) for v in pc_range) + chunk.feature_names.extend(feature_names) + yield chunk + def _set_h5_persistence_enabled(self, enabled: bool) -> bool: """Toggle dataframe manager H5 persistence flag when available.""" if self._df_manager is None: diff --git a/weightslab/trainer/trainer_services.py b/weightslab/trainer/trainer_services.py index 6ea86831..f4baaf15 100644 --- a/weightslab/trainer/trainer_services.py +++ b/weightslab/trainer/trainer_services.py @@ -369,6 +369,11 @@ def GetPointCloud(self, request, context): # Server-streaming RPC: delegate the generator directly. return self._exp_service.data_service.GetPointCloud(request, context) + def GetMedia(self, request, context): + logger.debug(f"\nExperimentServiceServicer.GetMedia({request})") + # Server-streaming RPC: delegate the generator directly. + return self._exp_service.data_service.GetMedia(request, context) + def EditDataSample(self, request, context): logger.debug(f"\nExperimentServiceServicer.EditDataSample({request})") return self._exp_service.data_service.EditDataSample(request, context)