Buckets:
| import json | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| from types import SimpleNamespace | |
| from unittest.mock import patch | |
| import numpy as np | |
| import torch | |
| class SegmentationDataTests(unittest.TestCase): | |
| def test_exact_canonical_mapping_and_ignore_round_trip(self): | |
| from training.segmentation_data import CANONICAL_IDS, canonical_to_head, head_to_canonical | |
| self.assertEqual(list(CANONICAL_IDS), [1, 2, 12, 13, 14, 15, 16]) | |
| mask = np.array([[1, 2, 12, 13], [14, 15, 16, 255]], np.uint8) | |
| mapped = canonical_to_head(mask) | |
| np.testing.assert_array_equal(mapped, [[0, 1, 2, 3], [4, 5, 6, 255]]) | |
| np.testing.assert_array_equal(head_to_canonical(mapped), mask) | |
| for invalid in [np.array([[0]], np.uint8), np.array([[3]], np.uint8), np.array([[1.5]])]: | |
| with self.assertRaises(ValueError): | |
| canonical_to_head(invalid) | |
| def test_validity_and_categorical_transforms(self): | |
| from training.segmentation_data import convert_segmentation_sample, transform_example, canonical_to_head | |
| image = np.full((4, 4, 3), 100, np.uint8) | |
| mask = np.array([[1, 2, 12, 13], [14, 15, 16, 255]] * 2, np.uint8) | |
| valid = mask != 255 | |
| valid[0, 0] = False | |
| e = convert_segmentation_sample({"image": image, "mask": mask, "valid_mask": valid, | |
| "metadata": {"id": "synthetic", "task": "segmentation"}}) | |
| self.assertEqual(e.mask[0, 0], 255) | |
| transformed = transform_example(e, horizontal=True, vertical=True, rotations=1, size=(8, 8)) | |
| expected = np.repeat(np.repeat(np.rot90(e.mask[::-1, ::-1]), 2, 0), 2, 1) | |
| np.testing.assert_array_equal(transformed.mask, expected) | |
| np.testing.assert_array_equal(transformed.head_mask, canonical_to_head(expected)) | |
| np.testing.assert_array_equal(transformed.valid_mask, expected != 255) | |
| self.assertEqual(transformed.metadata, e.metadata) | |
| class SegmentationLossTests(unittest.TestCase): | |
| def test_ignore_pixels_do_not_change_segmentation_loss_or_iou(self): | |
| from training.segmentation_loss import SourceAwareSemanticLoss, segmentation_counts | |
| torch.set_num_threads(4) | |
| model = SimpleNamespace(model=[SimpleNamespace(nc=7)], args=SimpleNamespace(data=None), | |
| parameters=lambda: iter([torch.nn.Parameter(torch.zeros(1))])) | |
| criterion = SourceAwareSemanticLoss(model) | |
| logits = torch.randn(1, 7, 4, 4, requires_grad=True) | |
| target = torch.zeros(1, 4, 4, dtype=torch.long) | |
| valid = torch.ones(1, 4, 4, dtype=torch.bool) | |
| valid[:, :2] = False | |
| batch = {"semantic_mask": target, "valid_mask": valid} | |
| loss, _ = criterion((logits, logits.clone()), batch) | |
| changed = logits.detach().clone() | |
| changed[:, :, :2] += torch.randn_like(changed[:, :, :2]) * 100 | |
| other, _ = criterion((changed, changed.clone()), batch) | |
| torch.testing.assert_close(loss, other) | |
| loss.backward() | |
| self.assertFalse(logits.grad[:, :, :2].any()) | |
| truth = np.ones((4, 4), np.uint8) | |
| truth[:2] = 255 | |
| prediction = np.ones_like(truth) | |
| before = segmentation_counts(prediction, truth) | |
| prediction[:2] = 16 | |
| self.assertEqual(before, segmentation_counts(prediction, truth)) | |
| batch["valid_mask"].fill_(False) | |
| empty, _ = criterion((logits, logits.clone()), batch) | |
| self.assertEqual(empty.detach().item(), 0.) | |
| def test_actual_main_and_auxiliary_head_adaptation(self): | |
| from training.segmentation_loss import build_adapted_segmenter | |
| from training.model_factory import ROOT, load_pretrained_segmenter | |
| source = load_pretrained_segmenter(ROOT / "training/cache/checkpoints") | |
| before = source.model.model[-1].classifier[-1].weight.detach().clone() | |
| model, report = build_adapted_segmenter(source, seed=42) | |
| self.assertEqual(model.model[-1].nc, 7) | |
| self.assertEqual(report["fresh_canonical_ids"], [13, 14, 15, 16]) | |
| for branch in ("classifier", "aux_head"): | |
| old, new = getattr(source.model.model[-1], branch)[-1], getattr(model.model[-1], branch)[-1] | |
| for target, name in enumerate(["building", "road", "water"]): | |
| original = next(i for i, n in source.names.items() if n == name) | |
| torch.testing.assert_close(new.weight[target], old.weight[original]) | |
| torch.testing.assert_close(new.bias[target], old.bias[original]) | |
| torch.testing.assert_close(source.model.model[-1].classifier[-1].weight, before) | |
| model.train() | |
| predictions = model(torch.zeros(1, 3, 64, 64)) | |
| batch = {"semantic_mask": torch.zeros(1, 64, 64, dtype=torch.long), | |
| "valid_mask": torch.ones(1, 64, 64, dtype=torch.bool)} | |
| loss, _ = model.criterion(predictions, batch) | |
| self.assertTrue(torch.isfinite(loss)) | |
| loss.backward() | |
| self.assertTrue(model.model[-1].aux_head[-1].weight.grad.abs().sum() > 0) | |
| class SegmentationTrainerTests(unittest.TestCase): | |
| def test_synthetic_native_fit_and_reload(self): | |
| from copy import deepcopy | |
| from training.segmentation_trainer import SegmenterTrainConfig, make_segmentation_trainer | |
| from training.model_factory import ROOT | |
| from ultralytics import YOLO | |
| manifest = json.loads((ROOT / "evaluation/quick_eval_manifest.json").read_text()) | |
| selected = next(r["sample_id"] for r in manifest["selection_validation"] if r["task"] == "segmentation") | |
| mask = np.ones((1024, 1024), np.uint8) | |
| mask[:, 512:] = 2 | |
| mask[:16] = 255 | |
| generated = {"image": np.full((1024, 1024, 3), 100, np.uint8), "mask": mask, | |
| "valid_mask": mask != 255, "metadata": {"id": "synthetic", "task": "segmentation", "split": "train"}} | |
| def load(root, row, hashes): | |
| result = deepcopy(generated) | |
| result["metadata"].update(id=row["id"], split=row["split"]) | |
| return result | |
| with tempfile.TemporaryDirectory(dir=ROOT / "training/cache") as directory: | |
| config = SegmenterTrainConfig(output=Path(directory) / "fit", epochs=1, record_initial_metrics=True, | |
| freeze_bn_statistics=True) | |
| with patch("training.segmentation_trainer.dataset", side_effect=[ | |
| SimpleNamespace(samples=[{"id": "synthetic", "split": "train"}], expected_hashes={}), | |
| SimpleNamespace(samples=[{"id": selected, "split": "val"}], expected_hashes={})]), \ | |
| patch("training.segmentation_trainer.load_sample", side_effect=load): | |
| result = make_segmentation_trainer(config).fit(["synthetic"], [selected], max_seconds=120) | |
| self.assertEqual(result["status"], "COMPLETED") | |
| self.assertEqual(len(result["history"]), 2) | |
| self.assertTrue(result["history"][1]["bn_buffers_unchanged"]) | |
| self.assertEqual(result["history"][0]["evaluation_ids"], [selected]) | |
| self.assertEqual(len(result["best_metrics"]["per_class"]), 7) | |
| model = YOLO(result["best_checkpoint"]) | |
| self.assertEqual(model.model.args["imgsz"], 1024) | |
| self.assertTrue(model.ckpt["config"]["freeze_bn_statistics"]) | |
| self.assertEqual(model.ckpt["augmentation"]["scale_jitter"], False) | |
| with torch.inference_mode(): | |
| prediction = model.predict(generated["image"], imgsz=1024, device="cpu", save=False, verbose=False)[0] | |
| self.assertIsNotNone(prediction.semantic_mask) | |
| self.assertEqual(len(model.names), 7) | |
| self.assertEqual(model.task, "semantic") | |
| def test_budget_split_config_and_result_contract(self): | |
| from training.segmentation_trainer import SegmenterTrainConfig, make_segmentation_trainer | |
| from training.model_factory import ROOT | |
| manifest = json.loads((ROOT / "evaluation/quick_eval_manifest.json").read_text()) | |
| selected = next(r["sample_id"] for r in manifest["selection_validation"] if r["task"] == "segmentation") | |
| trainer = make_segmentation_trainer(SegmenterTrainConfig()) | |
| with self.assertRaises(ValueError): | |
| trainer.fit([], [], max_seconds=3601) | |
| with self.assertRaises(ValueError): | |
| SegmenterTrainConfig(epochs=11) | |
| with patch("training.segmentation_trainer.dataset", return_value=SimpleNamespace(samples=[{"id": "forbidden"}])): | |
| with self.assertRaises(ValueError): | |
| trainer.fit(["forbidden"], ["forbidden"], max_seconds=1) | |
| with tempfile.TemporaryDirectory(dir=ROOT / "training/cache") as directory: | |
| config = SegmenterTrainConfig.from_json(ROOT / "training/configs/segmenter.json", output=Path(directory) / "fit") | |
| with patch("training.segmentation_trainer.dataset", side_effect=[SimpleNamespace(samples=[{"id": "train"}]), | |
| SimpleNamespace(samples=[{"id": selected}])]), \ | |
| patch("training.segmentation_trainer.time.monotonic", side_effect=[0, 2, 3]), \ | |
| patch("training.segmentation_trainer.load_pretrained_segmenter") as load: | |
| result = make_segmentation_trainer(config).fit(["train"], [selected], max_seconds=1) | |
| load.assert_not_called() | |
| self.assertEqual(set(result), {"status", "best_checkpoint", "history", "best_metrics", "elapsed_seconds"}) | |
| self.assertEqual(result["status"], "TIME_LIMIT") | |
Xet Storage Details
- Size:
- 9.62 kB
- Xet hash:
- 54316c5d3098dc212570d3b5a9ce6ac594f5a1dad697a2ac496afc04e020502e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.