| import json |
| from types import SimpleNamespace |
|
|
| import numpy as np |
| import pytest |
| import torch |
| from PIL import Image |
|
|
| from image3d_student.data import CachedDataset, cache_sample, cell_indices, splat_targets |
| from image3d_student.engine import split_rows |
| from image3d_student.model import Student, distillation_loss |
| from image3d_student.splats import Splats, read_ply, write_ply |
| from image3d_student.teacher_idol import patch_attention |
|
|
|
|
| @pytest.fixture |
| def splats(): |
| return Splats(np.array([[-.2, .1, .0], [.2, -.1, .3]]), |
| np.array([[.9,.1,.2], [.1,.8,.4]]), np.array([[.01,.02,.03],[.04,.01,.02]]), |
| np.array([[1,0,0,0], [.5,.5,.5,.5]]), np.array([[.8],[.9]])) |
|
|
|
|
| def test_ply_roundtrip_preserves_frame_and_attributes(tmp_path, splats): |
| path = tmp_path / "splats.ply" |
| write_ply(path, splats) |
| reread = read_ply(path) |
| for name in splats.__dataclass_fields__: |
| np.testing.assert_allclose(getattr(reread, name), getattr(splats, name), atol=1e-6) |
|
|
|
|
| def test_non_gaussian_ply_rejected(tmp_path): |
| import trimesh |
| path = tmp_path / "mesh.ply" |
| trimesh.creation.box().export(path) |
| with pytest.raises(ValueError, match="Gaussian-splat"): |
| read_ply(path) |
|
|
|
|
| def test_coordinates_not_silently_clipped(): |
| with pytest.raises(ValueError, match="aligned"): |
| cell_indices(np.array([[1., 0, 0]]), 8) |
| assert cell_indices(np.array([[.5,.5,.5]]), 8)[0] == 511 |
|
|
|
|
| def test_voxel_assignment_selects_highest_opacity(splats): |
| splats.xyz[1] = splats.xyz[0] + .001 |
| target = splat_targets(splats, 8) |
| assert target["occ"].sum() == 1 |
| np.testing.assert_allclose(target["rgb"][target["occ"][:,0] > 0][0], splats.rgb[1]) |
|
|
|
|
| def test_cache_integrity_and_duplicate_protection(tmp_path, splats): |
| image = tmp_path / "input.png" |
| Image.new("RGB", (32,32), "red").save(image) |
| root = tmp_path / "dataset" |
| cache_sample(root, "a", image, splats, 8) |
| dataset = CachedDataset(root / "manifest.jsonl") |
| tensor, target = dataset[0] |
| assert tensor.shape == (3,128,128) |
| assert target["occ"].sum() == 2 |
| with pytest.raises(ValueError, match="already exists"): |
| cache_sample(root, "a", image, splats, 8) |
| (root / "images/a.png").write_bytes(b"changed") |
| with pytest.raises(ValueError, match="hash mismatch"): |
| CachedDataset(root / "manifest.jsonl") |
|
|
|
|
| def test_geometry_import_requires_frame_and_marks_teachers(tmp_path, splats): |
| import trimesh |
| image = tmp_path / "input.png" |
| Image.new("RGB", (32,32), "green").save(image) |
| mesh = tmp_path / "shape.obj" |
| trimesh.creation.icosphere(subdivisions=1, radius=.3).export(mesh) |
| with pytest.raises(ValueError, match="explicit"): |
| cache_sample(tmp_path / "bad", "a", image, splats, 8, coarse=mesh) |
| transform = tmp_path / "identity.json" |
| transform.write_text(json.dumps(np.eye(4).tolist())) |
| root = tmp_path / "good" |
| row = cache_sample(root, "a", image, splats, 8, coarse=mesh, coarse_transform=transform, |
| fine=mesh, fine_transform=transform) |
| assert row["teachers"] == {"appearance": "idol", "coarse": "InstantMesh", "fine": "Hi3D"} |
| with np.load(root / row["target"]) as target: |
| assert target["geometry_valid"].tolist() == [1,1] |
| assert target["geometry"].sum() > 0 |
|
|
|
|
| def test_missing_teachers_are_masked_and_quaternions_sign_invariant(splats): |
| torch.manual_seed(1) |
| model = Student(8,32) |
| target = {k: torch.from_numpy(v).float()[None] for k,v in splat_targets(splats,8).items() if k != "grid"} |
| prediction = model(torch.zeros(1,3,128,128)) |
| loss1, parts1 = distillation_loss(prediction, target, 8) |
| target["geometry"][:] = 1 |
| target["quat"] *= -1 |
| loss2, parts2 = distillation_loss(prediction, target, 8) |
| torch.testing.assert_close(loss1, loss2) |
| assert parts1["geometry"] == 0 |
| assert parts2["geometry"] == 0 |
| target["geometry_valid"][:] = 1 |
| loss3, parts3 = distillation_loss(prediction, target, 8) |
| assert parts3["geometry"] > 0 |
| assert loss3 > loss2 |
|
|
|
|
| def test_query_chunk_attention_matches_full_attention(): |
| from torch.nn import functional as F |
| torch.manual_seed(7) |
| layers = SimpleNamespace(RePo3DRotaryEmbedding=type("RoPE", (), {})) |
| chunked = patch_attention(layers, 3) |
| q = torch.randn(2,11,2,8) |
| k = torch.randn(2,17,2,8) |
| v = torch.randn(2,17,2,8) |
| expected = F.scaled_dot_product_attention(q.permute(0,2,1,3), k.permute(0,2,1,3), v.permute(0,2,1,3)).permute(0,2,1,3) |
| torch.testing.assert_close(chunked(q,k,v), expected, atol=1e-6, rtol=1e-5) |
| qkv = torch.randn(1,7,3,2,8) |
| q,k,v = qkv.unbind(2) |
| torch.testing.assert_close(chunked(qkv), chunked(q,k,v)) |
| assert layers.RePo3DRotaryEmbedding.real_valued |
|
|
|
|
| def test_split_keeps_object_views_together(): |
| rows = [{"id": str(i), "image_sha256": str(i), "metadata": {"object_id": str(i//2)}} for i in range(8)] |
| train, val = split_rows(rows, .25, 42) |
| assert not {r["metadata"]["object_id"] for r in train} & {r["metadata"]["object_id"] for r in val} |
| assert len(train) == 6 and len(val) == 2 |
|
|
|
|
| def test_actual_training_reduces_distillation_loss(splats): |
| torch.set_num_threads(2) |
| torch.manual_seed(42) |
| model = Student(8,32) |
| target = {k: torch.from_numpy(v).float()[None] for k,v in splat_targets(splats,8).items() if k != "grid"} |
| image = torch.rand(1,3,128,128) |
| optimizer = torch.optim.Adam(model.parameters(), lr=.003) |
| initial = float(distillation_loss(model(image), target, 8)[0].detach()) |
| for _ in range(25): |
| optimizer.zero_grad() |
| loss, _ = distillation_loss(model(image), target, 8) |
| loss.backward() |
| optimizer.step() |
| final = float(distillation_loss(model(image), target, 8)[0].detach()) |
| assert final < initial * .65, (initial, final) |
|
|
