Instructions to use nvidia/Flora-TempSeg with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- PEFT
How to use nvidia/Flora-TempSeg with PEFT:
from peft import PeftModel from transformers import AutoModelForCausalLM base_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-9B") model = PeftModel.from_pretrained(base_model, "nvidia/Flora-TempSeg") - Notebooks
- Google Colab
- Kaggle
Add Flora inference examples
#1
by OmarLaymounNvidia - opened
- SOURCE_MANIFEST.sha256 +1 -1
- batch_utils/CODE_MANIFEST.json +173 -0
- batch_utils/README.md +107 -0
- batch_utils/__init__.py +3 -0
- batch_utils/__main__.py +3 -0
- batch_utils/api.py +58 -0
- batch_utils/bootstrap.py +18 -0
- batch_utils/cli.py +81 -0
- batch_utils/config.py +12 -0
- batch_utils/media.py +106 -0
- batch_utils/models.json +18 -0
- batch_utils/models.py +189 -0
- batch_utils/pipeline.py +169 -0
- batch_utils/preprocessing.py +56 -0
- batch_utils/requirements.txt +10 -0
- batch_utils/vendor/MANIFEST.json +29 -0
- batch_utils/vendor/flora_unified/__init__.py +1 -0
- batch_utils/vendor/flora_unified/unified_input.py +89 -0
- batch_utils/vendor/production_adapter/TEMPSEG_PROMPTS.json +27 -0
- batch_utils/vendor/production_adapter/adapter.py +102 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/__init__.py +1 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/exact_frame_serving.py +47 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/schema.py +88 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/__init__.py +1 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/u2_contract.py +58 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/u2_worker.py +44 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/__init__.py +1 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/contracts.py +73 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_subtask.py +852 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_temporal.py +987 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_worker.py +23 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/interface.py +24 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/scope_scheduler.py +163 -0
- batch_utils/vendor/production_adapter/core/flora_runtime/worker_transport.py +18 -0
- batch_utils/vendor/production_gdn_fix.py +2 -0
- batch_utils/vendor/pts_input_adapter.py +36 -0
- batch_utils/vendor/robot_caption_geometry.py +60 -0
- batch_utils/vendor/support/canonical_rgb.py +25 -0
- batch_utils/vendor/support/composers.py +26 -0
- batch_utils/vendor/support/plans.py +38 -0
- batch_utils/vendor/support/u2_contract.py +49 -0
- batch_utils/vendor/tempseg_fp32_loader.py +51 -0
- batch_utils/vendor/verify_gdn.py +36 -0
- infer.py +5 -0
SOURCE_MANIFEST.sha256
CHANGED
|
@@ -1,5 +1,5 @@
|
|
| 1 |
837c0aa8748126d2193799e7bec51d6550ca4ba07767591e66dcb370c3115e44 MODEL_PROVENANCE.json
|
| 2 |
-
|
| 3 |
ac4d8abd2a14045b22151354a7493049d651ec10c0b019fbe4927805e45c4b81 adapter_config.json
|
| 4 |
1d80fe544e4cd1a7ff6b0ea4be71e988972408dab81d4c8af72e3a628739d6c3 adapter_model.safetensors
|
| 5 |
d28ca68768ac2ce0ad42eaaf54d5bac4f8a298f7d017f5a04f3866094c65a839 chat_template.jinja
|
|
|
|
| 1 |
837c0aa8748126d2193799e7bec51d6550ca4ba07767591e66dcb370c3115e44 MODEL_PROVENANCE.json
|
| 2 |
+
88cb6731fc521a0404faed7c5045c3db77185341849bc4e0b6e60be3443f7d6c README.md
|
| 3 |
ac4d8abd2a14045b22151354a7493049d651ec10c0b019fbe4927805e45c4b81 adapter_config.json
|
| 4 |
1d80fe544e4cd1a7ff6b0ea4be71e988972408dab81d4c8af72e3a628739d6c3 adapter_model.safetensors
|
| 5 |
d28ca68768ac2ce0ad42eaaf54d5bac4f8a298f7d017f5a04f3866094c65a839 chat_template.jinja
|
batch_utils/CODE_MANIFEST.json
ADDED
|
@@ -0,0 +1,173 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"files": {
|
| 3 |
+
"batch_utils/README.md": {
|
| 4 |
+
"bytes": 4125,
|
| 5 |
+
"sha256": "f9a9dd0a67b914bb1304a53674c03deec9f6a9aaa580781a13d88b365284a65c"
|
| 6 |
+
},
|
| 7 |
+
"batch_utils/__init__.py": {
|
| 8 |
+
"bytes": 91,
|
| 9 |
+
"sha256": "50a358089004c69e4c83c41a55eb44c52179d3fa20be4d4c5b92fabbc8271ac4"
|
| 10 |
+
},
|
| 11 |
+
"batch_utils/__main__.py": {
|
| 12 |
+
"bytes": 78,
|
| 13 |
+
"sha256": "28de248c91c00642578a541827361af9b0e5e049e4c61dd6ec1f9dc78580c24e"
|
| 14 |
+
},
|
| 15 |
+
"batch_utils/api.py": {
|
| 16 |
+
"bytes": 3600,
|
| 17 |
+
"sha256": "ec5e17b5b02c20a40e52beab7d5ebfd49e32129ef03eb918790df720b0d97177"
|
| 18 |
+
},
|
| 19 |
+
"batch_utils/bootstrap.py": {
|
| 20 |
+
"bytes": 679,
|
| 21 |
+
"sha256": "83e49883136af2088038c6070303d0e89071a318e98252b40ff988230e491817"
|
| 22 |
+
},
|
| 23 |
+
"batch_utils/cli.py": {
|
| 24 |
+
"bytes": 4761,
|
| 25 |
+
"sha256": "9dc0256b45252d7512b1c69eb7eba3b681725d7c6109287dc3892ec9a9d668e1"
|
| 26 |
+
},
|
| 27 |
+
"batch_utils/config.py": {
|
| 28 |
+
"bytes": 503,
|
| 29 |
+
"sha256": "0967244955e8f128e03cc6a85973a41d1571f78acf4a9bc6f91e6682781f9891"
|
| 30 |
+
},
|
| 31 |
+
"batch_utils/media.py": {
|
| 32 |
+
"bytes": 5135,
|
| 33 |
+
"sha256": "e0ee696b63f3171d5163aadb960322f422183981c3768efd76c409976f3fda5b"
|
| 34 |
+
},
|
| 35 |
+
"batch_utils/models.json": {
|
| 36 |
+
"bytes": 612,
|
| 37 |
+
"sha256": "210e56e9ed46f61c6e18a974b78fccfb6a6aee3a6b0cb6cda5b5f774f934f536"
|
| 38 |
+
},
|
| 39 |
+
"batch_utils/models.py": {
|
| 40 |
+
"bytes": 11436,
|
| 41 |
+
"sha256": "92b3ab4b376f9a4dd3c7f75568746467f17a969430c63330667f2c6d2760f1fe"
|
| 42 |
+
},
|
| 43 |
+
"batch_utils/pipeline.py": {
|
| 44 |
+
"bytes": 11853,
|
| 45 |
+
"sha256": "588d4105197706da8acf3b591ba9671239505b53030f1d30cba0088ddad5432e"
|
| 46 |
+
},
|
| 47 |
+
"batch_utils/preprocessing.py": {
|
| 48 |
+
"bytes": 3311,
|
| 49 |
+
"sha256": "6bc9a2c7f08f92b4933f68c0fc674f938339eceeebdec186eef1514294b167ee"
|
| 50 |
+
},
|
| 51 |
+
"batch_utils/requirements.txt": {
|
| 52 |
+
"bytes": 163,
|
| 53 |
+
"sha256": "66d607ebfbd6525e005db045812453c4356d7df4e42e7b40b4994bd0e5e0cd51"
|
| 54 |
+
},
|
| 55 |
+
"batch_utils/vendor/MANIFEST.json": {
|
| 56 |
+
"bytes": 3225,
|
| 57 |
+
"sha256": "91ce5d8a5f13a583d48384c5de7f46696f5cd896ad2c1f3ce0eead9a1acae6d8"
|
| 58 |
+
},
|
| 59 |
+
"batch_utils/vendor/flora_unified/__init__.py": {
|
| 60 |
+
"bytes": 5,
|
| 61 |
+
"sha256": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2"
|
| 62 |
+
},
|
| 63 |
+
"batch_utils/vendor/flora_unified/unified_input.py": {
|
| 64 |
+
"bytes": 5978,
|
| 65 |
+
"sha256": "2aa2b521191b68913b03d7582b7a3bcfad49d90c014a1a1bb7571b5e4f08d7ef"
|
| 66 |
+
},
|
| 67 |
+
"batch_utils/vendor/production_adapter/TEMPSEG_PROMPTS.json": {
|
| 68 |
+
"bytes": 6654,
|
| 69 |
+
"sha256": "66edd6fb43b7fa55cc6257beee5c42a1a4555ab9790320c31114d77bd07a05f1"
|
| 70 |
+
},
|
| 71 |
+
"batch_utils/vendor/production_adapter/adapter.py": {
|
| 72 |
+
"bytes": 7526,
|
| 73 |
+
"sha256": "083c6788b561d726a85d9f3f8f516593a9f6f76a43f315581c821871e3f82312"
|
| 74 |
+
},
|
| 75 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/__init__.py": {
|
| 76 |
+
"bytes": 5,
|
| 77 |
+
"sha256": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2"
|
| 78 |
+
},
|
| 79 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/exact_frame_serving.py": {
|
| 80 |
+
"bytes": 2120,
|
| 81 |
+
"sha256": "f921e279b5d52f723b3e77e25441763683d05c6e9cd2f8ff204466a50255b346"
|
| 82 |
+
},
|
| 83 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/schema.py": {
|
| 84 |
+
"bytes": 3489,
|
| 85 |
+
"sha256": "42c523e4db1a4da3fe0d0d578f3cf6e94aed7666c06e3f18081c209159de04b9"
|
| 86 |
+
},
|
| 87 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/__init__.py": {
|
| 88 |
+
"bytes": 5,
|
| 89 |
+
"sha256": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2"
|
| 90 |
+
},
|
| 91 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/u2_contract.py": {
|
| 92 |
+
"bytes": 3511,
|
| 93 |
+
"sha256": "873476e78a9e3e222ae9e47fcb7a04c8c18ef518b2847c401a2e1cfb2292fe90"
|
| 94 |
+
},
|
| 95 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/u2_worker.py": {
|
| 96 |
+
"bytes": 2526,
|
| 97 |
+
"sha256": "870641228b0be442613d520232d451047742ca2ea41fd0cd51dbf0a58ba567d2"
|
| 98 |
+
},
|
| 99 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/__init__.py": {
|
| 100 |
+
"bytes": 5,
|
| 101 |
+
"sha256": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2"
|
| 102 |
+
},
|
| 103 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/contracts.py": {
|
| 104 |
+
"bytes": 3361,
|
| 105 |
+
"sha256": "8591d849b0b65fda6b21511557d4ee5eabd47f8332e2b6831f40d5bb472d978c"
|
| 106 |
+
},
|
| 107 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_subtask.py": {
|
| 108 |
+
"bytes": 77158,
|
| 109 |
+
"sha256": "6fbf8337bab8491148a78a822187d74397b7f0552700252451a18cd7c78c8e04"
|
| 110 |
+
},
|
| 111 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_temporal.py": {
|
| 112 |
+
"bytes": 66254,
|
| 113 |
+
"sha256": "32f2b99c97176f5e10934450da638577268c2e5d9977c2ea138ebf1b7eda58fa"
|
| 114 |
+
},
|
| 115 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_worker.py": {
|
| 116 |
+
"bytes": 1090,
|
| 117 |
+
"sha256": "8a31a77180c8d406113d4b494644153564a6f46c8a9b68ad7babb3cc1695edec"
|
| 118 |
+
},
|
| 119 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/interface.py": {
|
| 120 |
+
"bytes": 996,
|
| 121 |
+
"sha256": "a5424b2cc3a4c879552dd49bae207a58fb9e751359f441c03af167a2e1879359"
|
| 122 |
+
},
|
| 123 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/scope_scheduler.py": {
|
| 124 |
+
"bytes": 12654,
|
| 125 |
+
"sha256": "be95411b39124d83811fe74dcac2d911169ecf475a81b7460c64b19fbbc7e5db"
|
| 126 |
+
},
|
| 127 |
+
"batch_utils/vendor/production_adapter/core/flora_runtime/worker_transport.py": {
|
| 128 |
+
"bytes": 953,
|
| 129 |
+
"sha256": "30fb64869eb6a27ab96bf394c105f19fcb42590be4424c66b76b89e17631693a"
|
| 130 |
+
},
|
| 131 |
+
"batch_utils/vendor/production_gdn_fix.py": {
|
| 132 |
+
"bytes": 202,
|
| 133 |
+
"sha256": "a3f52604bfa2ec58901df63215df5e47d37e1d7974aa56d36a1e4efe8cad3fc6"
|
| 134 |
+
},
|
| 135 |
+
"batch_utils/vendor/pts_input_adapter.py": {
|
| 136 |
+
"bytes": 2340,
|
| 137 |
+
"sha256": "d784d6b77005eb6945dee334e4c9136fbededfa2049d5945244f198517984767"
|
| 138 |
+
},
|
| 139 |
+
"batch_utils/vendor/robot_caption_geometry.py": {
|
| 140 |
+
"bytes": 2739,
|
| 141 |
+
"sha256": "1c874d3c63359663276d31a3e9c3112a5982ad4000f44c808541b2471f5f7d2d"
|
| 142 |
+
},
|
| 143 |
+
"batch_utils/vendor/support/canonical_rgb.py": {
|
| 144 |
+
"bytes": 1010,
|
| 145 |
+
"sha256": "1e3230c60b45d8e017060d8a5dd2b5f48e2199192f09dcff4d04c05b4349e35a"
|
| 146 |
+
},
|
| 147 |
+
"batch_utils/vendor/support/composers.py": {
|
| 148 |
+
"bytes": 1263,
|
| 149 |
+
"sha256": "180b97a6a634979c1b0dbe720ed7a1abd818dc7ee31fdc4eeca7a8277e43d1b6"
|
| 150 |
+
},
|
| 151 |
+
"batch_utils/vendor/support/plans.py": {
|
| 152 |
+
"bytes": 3762,
|
| 153 |
+
"sha256": "d75b918804c9ec10f8642d362a54a74c3eb875abf22e8a60228a0d6b2cbb570e"
|
| 154 |
+
},
|
| 155 |
+
"batch_utils/vendor/support/u2_contract.py": {
|
| 156 |
+
"bytes": 3321,
|
| 157 |
+
"sha256": "50190914b719bbc9b8f2b53e52a6bd2b142a452e2fb1c5c897ea4c31123ad4dd"
|
| 158 |
+
},
|
| 159 |
+
"batch_utils/vendor/tempseg_fp32_loader.py": {
|
| 160 |
+
"bytes": 3419,
|
| 161 |
+
"sha256": "7ee9889f5aa5909abf3ecd1c34e3ac2f81a32a2098d94ae3ba9f33831e7a07c0"
|
| 162 |
+
},
|
| 163 |
+
"batch_utils/vendor/verify_gdn.py": {
|
| 164 |
+
"bytes": 2149,
|
| 165 |
+
"sha256": "96596ff2db4425f4dea71798aeabb581aa83e1ad8a878cfcb9efb8dbac80cda6"
|
| 166 |
+
},
|
| 167 |
+
"infer.py": {
|
| 168 |
+
"bytes": 171,
|
| 169 |
+
"sha256": "8cddbb43ba5b9d884e8f771b9c9e219ce2e0a3c1d71a94a75d02b1d3b730f93d"
|
| 170 |
+
}
|
| 171 |
+
},
|
| 172 |
+
"schema_version": "flora_inference_code_manifest_v1"
|
| 173 |
+
}
|
batch_utils/README.md
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# FLORA Data Engine inference
|
| 2 |
+
|
| 3 |
+
This directory contains the complete **TempSeg → Captioner** workflow, identical
|
| 4 |
+
in both model repositories. Use either copy. TempSeg predicts action intervals;
|
| 5 |
+
the Captioner describes those predicted intervals. The native windowing,
|
| 6 |
+
adaptive retries, stitching, prompts, source timestamps, and camera layouts are
|
| 7 |
+
included; no internal source checkout or running demo is required.
|
| 8 |
+
|
| 9 |
+
## Setup
|
| 10 |
+
|
| 11 |
+
Python 3.12 and a CUDA GPU with 80 GB VRAM are the validated configuration.
|
| 12 |
+
The two models run sequentially. Download/cache space of at least 150 GB is
|
| 13 |
+
recommended for the released FP32 Captioner, base model, and adapter.
|
| 14 |
+
|
| 15 |
+
From a downloaded copy of either repository:
|
| 16 |
+
|
| 17 |
+
```bash
|
| 18 |
+
python -m pip install torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cu128
|
| 19 |
+
python -m pip install -r batch_utils/requirements.txt
|
| 20 |
+
hf auth login
|
| 21 |
+
```
|
| 22 |
+
|
| 23 |
+
Both private model repositories must be accessible to the logged-in account.
|
| 24 |
+
`HF_TOKEN` is also supported; do not put a token in source code or command-line
|
| 25 |
+
arguments. Weights are fetched automatically at the revisions in `models.json`.
|
| 26 |
+
Set `FLORA_MODEL_CACHE` to choose the local model directory.
|
| 27 |
+
|
| 28 |
+
The NVIDIA repositories are selected by default. To override them:
|
| 29 |
+
|
| 30 |
+
```bash
|
| 31 |
+
export FLORA_TEMPSEG_REPO="nvidia/Flora-TempSeg"
|
| 32 |
+
export FLORA_CAPTIONER_REPO="nvidia/Flora-Captioner"
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
## Single video
|
| 36 |
+
|
| 37 |
+
```bash
|
| 38 |
+
python infer.py --video videos/human.mp4 --profile human --output human.json
|
| 39 |
+
```
|
| 40 |
+
|
| 41 |
+
Three synchronized robot cameras:
|
| 42 |
+
|
| 43 |
+
```bash
|
| 44 |
+
python infer.py --video videos/top.mp4 \
|
| 45 |
+
--left-wrist videos/left.mp4 --right-wrist videos/right.mp4 \
|
| 46 |
+
--profile yam --output yam.json
|
| 47 |
+
```
|
| 48 |
+
|
| 49 |
+
Use `--profile dex` for Dex or `--profile tianji` for Tianji inputs.
|
| 50 |
+
YAM defaults to native 16-second temporal windows; use `--yam32` for the native
|
| 51 |
+
32-second YAM profile. Human uses one video. Robot profiles accept one view or
|
| 52 |
+
all three; `--yam32` requires all three. Wrist files must share the main video's
|
| 53 |
+
relative frame timestamps, frame count, and duration. Inputs are not resampled.
|
| 54 |
+
|
| 55 |
+
The validated input limits are 120 seconds, 256 MiB per file, 4096 pixels per
|
| 56 |
+
side, 7200 frames, and at least 4 source FPS. Variable-rate video is supported
|
| 57 |
+
by the single-view Human path. `--instruction` optionally supplies an episode
|
| 58 |
+
task; leave it blank for the default workflow.
|
| 59 |
+
|
| 60 |
+
CPU-only media validation performs no inference or model downloads:
|
| 61 |
+
|
| 62 |
+
```bash
|
| 63 |
+
python -m batch_utils --video videos/human.mp4 --profile human --validate-only
|
| 64 |
+
```
|
| 65 |
+
|
| 66 |
+
## JSONL jobs
|
| 67 |
+
|
| 68 |
+
Each row specifies one job. Camera paths are relative to the JSONL file:
|
| 69 |
+
|
| 70 |
+
```json
|
| 71 |
+
{"id":"human01","video":"videos/human.mp4","profile":"human"}
|
| 72 |
+
{"id":"robot01","video":"videos/top.mp4","left_wrist":"videos/left.mp4","right_wrist":"videos/right.mp4","profile":"yam","yam32":true}
|
| 73 |
+
```
|
| 74 |
+
|
| 75 |
+
```bash
|
| 76 |
+
python infer.py --input-jsonl jobs.jsonl --output-dir results
|
| 77 |
+
```
|
| 78 |
+
|
| 79 |
+
Jobs and predicted segments run sequentially; JSONL is not GPU batch scheduling.
|
| 80 |
+
Existing outputs are preserved unless `--overwrite` is explicit. An error stops
|
| 81 |
+
the command; successfully written earlier jobs remain available.
|
| 82 |
+
|
| 83 |
+
## Python usage
|
| 84 |
+
|
| 85 |
+
```python
|
| 86 |
+
from batch_utils import DataEnginePipeline
|
| 87 |
+
|
| 88 |
+
engine = DataEnginePipeline()
|
| 89 |
+
result = engine.annotate("videos/human.mp4", profile="human")
|
| 90 |
+
for segment in result["segments"]:
|
| 91 |
+
print(segment["start_s"], segment["end_s"], segment["caption"])
|
| 92 |
+
```
|
| 93 |
+
|
| 94 |
+
Outputs contain clip-relative seconds, caption completion status, camera roles,
|
| 95 |
+
input SHA-256, and model revisions. Native overlapping intervals are retained.
|
| 96 |
+
`include_audit=True` / `--include-audit` adds temporal and caption-input records.
|
| 97 |
+
The code preserves FP32 LoRA/GDN values and validates checkpoint manifests.
|
| 98 |
+
Full manifest verification happens on first use in the process; reuse requires
|
| 99 |
+
unchanged manifest bytes and all file identities/timestamps. FP32 parameter
|
| 100 |
+
values are checked on every load. First-time downloads and verification are
|
| 101 |
+
substantial; later jobs reuse the verified local files.
|
| 102 |
+
|
| 103 |
+
Example output (illustrative):
|
| 104 |
+
|
| 105 |
+
```json
|
| 106 |
+
{"segments":[{"start_s":0.8,"end_s":3.2,"caption":"Pick up the cup with the right hand.","status":"complete"}]}
|
| 107 |
+
```
|
batch_utils/__init__.py
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from .api import DataEnginePipeline
|
| 2 |
+
__all__ = ['DataEnginePipeline']
|
| 3 |
+
__version__ = '1.0.0'
|
batch_utils/__main__.py
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from .cli import main
|
| 2 |
+
if __name__ == '__main__':
|
| 3 |
+
raise SystemExit(main())
|
batch_utils/api.py
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pathlib import Path
|
| 2 |
+
import math
|
| 3 |
+
import threading
|
| 4 |
+
from .config import MAX_UPLOAD_BYTES, MODELS
|
| 5 |
+
_REQUEST_LOCK = threading.RLock()
|
| 6 |
+
_PROFILES = {'human': 'Human', 'yam': 'YAM', 'dex': 'Dex', 'tianji': 'Tianji'}
|
| 7 |
+
|
| 8 |
+
class DataEnginePipeline:
|
| 9 |
+
|
| 10 |
+
@staticmethod
|
| 11 |
+
def _inputs(video, profile, left_wrist, right_wrist, yam32, instruction):
|
| 12 |
+
from .media import InputError, probe, check_synchronized
|
| 13 |
+
if not isinstance(profile, str) or profile.lower() not in _PROFILES:
|
| 14 |
+
raise InputError('profile must be human, yam, dex, or tianji')
|
| 15 |
+
canonical = _PROFILES[profile.lower()]
|
| 16 |
+
if bool(left_wrist) != bool(right_wrist):
|
| 17 |
+
raise InputError('Provide both left and right wrist videos together.')
|
| 18 |
+
if canonical == 'Human' and left_wrist:
|
| 19 |
+
raise InputError('The human profile accepts one video.')
|
| 20 |
+
if not isinstance(yam32, bool):
|
| 21 |
+
raise InputError('yam32 must be true or false.')
|
| 22 |
+
if yam32 and (canonical != 'YAM' or not left_wrist):
|
| 23 |
+
raise InputError('yam32 requires YAM and all three camera views.')
|
| 24 |
+
if not isinstance(instruction, str) or len(instruction.strip()) > 2000:
|
| 25 |
+
raise InputError('instruction must be a string of at most 2000 characters.')
|
| 26 |
+
paths = {'main': video}
|
| 27 |
+
if left_wrist:
|
| 28 |
+
paths.update(left_wrist=left_wrist, right_wrist=right_wrist)
|
| 29 |
+
videos = {}
|
| 30 |
+
for role, value in paths.items():
|
| 31 |
+
if not isinstance(value, (str, Path)) or not str(value).strip():
|
| 32 |
+
raise InputError('Each camera path must be a nonempty local path.')
|
| 33 |
+
path = Path(value).expanduser().resolve(strict=True)
|
| 34 |
+
if not path.is_file() or not 0 < path.stat().st_size <= MAX_UPLOAD_BYTES:
|
| 35 |
+
raise InputError('Each input must be a nonempty local video of at most 256 MiB.')
|
| 36 |
+
videos[role] = probe(path)
|
| 37 |
+
check_synchronized(videos)
|
| 38 |
+
return (videos, canonical)
|
| 39 |
+
|
| 40 |
+
def inspect_inputs(self, video, *, profile='human', left_wrist=None, right_wrist=None, yam32=False, instruction=''):
|
| 41 |
+
videos, canonical = self._inputs(video, profile, left_wrist, right_wrist, yam32, instruction)
|
| 42 |
+
return {'status': 'inputs_validated', 'profile': canonical, 'yam32': yam32, 'duration_s': videos['main'].duration_s, 'model_revisions': MODELS, 'new_model_calls': 0, 'cameras': {role: {'sha256': v.sha256, 'source_fps': str(v.rate), 'frames': len(v.pts_us), 'size': list(v.size), 'constant_frame_rate': v.cfr, 'duration_s': v.duration_s} for role, v in videos.items()}}
|
| 43 |
+
|
| 44 |
+
def annotate(self, video, *, profile='human', left_wrist=None, right_wrist=None, yam32=False, instruction='', progress=None, include_audit=False):
|
| 45 |
+
import torch
|
| 46 |
+
if not torch.cuda.is_available():
|
| 47 |
+
raise RuntimeError('CUDA is required for inference. Use --validate-only to check inputs on CPU.')
|
| 48 |
+
from . import pipeline
|
| 49 |
+
with _REQUEST_LOCK:
|
| 50 |
+
videos, canonical = self._inputs(video, profile, left_wrist, right_wrist, yam32, instruction)
|
| 51 |
+
result = pipeline.annotate(videos, canonical, yam32, instruction, progress or (lambda *a, **k: None))
|
| 52 |
+
duration = result['duration_s']
|
| 53 |
+
for segment in result['segments']:
|
| 54 |
+
if not (math.isfinite(segment['start_s']) and math.isfinite(segment['end_s']) and (0 <= segment['start_s'] < segment['end_s'] <= duration + 0.001)):
|
| 55 |
+
raise RuntimeError('The pipeline returned an invalid caption interval.')
|
| 56 |
+
if not include_audit:
|
| 57 |
+
result.pop('pipeline', None)
|
| 58 |
+
return result
|
batch_utils/bootstrap.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import hashlib
|
| 2 |
+
import json
|
| 3 |
+
import sys
|
| 4 |
+
from .config import VENDOR
|
| 5 |
+
|
| 6 |
+
def initialize():
|
| 7 |
+
manifest = json.loads((VENDOR / 'MANIFEST.json').read_text())
|
| 8 |
+
for relative, expected in manifest.items():
|
| 9 |
+
path = VENDOR / relative
|
| 10 |
+
if not path.resolve().is_relative_to(VENDOR.resolve()):
|
| 11 |
+
raise RuntimeError('Invalid runtime manifest path')
|
| 12 |
+
if hashlib.sha256(path.read_bytes()).hexdigest() != expected:
|
| 13 |
+
raise RuntimeError(f'Production code changed: {relative}')
|
| 14 |
+
for path in [VENDOR, VENDOR / 'support', VENDOR / 'production_adapter']:
|
| 15 |
+
sys.path.insert(0, str(path))
|
| 16 |
+
import adapter
|
| 17 |
+
return adapter
|
| 18 |
+
production = initialize()
|
batch_utils/cli.py
ADDED
|
@@ -0,0 +1,81 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pathlib import Path
|
| 2 |
+
import argparse
|
| 3 |
+
import json
|
| 4 |
+
import re
|
| 5 |
+
import sys
|
| 6 |
+
from .api import DataEnginePipeline
|
| 7 |
+
|
| 8 |
+
def parser():
|
| 9 |
+
p = argparse.ArgumentParser(description='FLORA: TempSeg -> Captioner with native video preprocessing.')
|
| 10 |
+
source = p.add_mutually_exclusive_group(required=True)
|
| 11 |
+
source.add_argument('--video', type=Path, help='Main/egocentric local video.')
|
| 12 |
+
source.add_argument('--input-jsonl', type=Path, help='One video job per line; paths relative to this JSONL.')
|
| 13 |
+
p.add_argument('--profile', choices=['human', 'yam', 'dex', 'tianji'], default='human')
|
| 14 |
+
p.add_argument('--left-wrist', type=Path)
|
| 15 |
+
p.add_argument('--right-wrist', type=Path)
|
| 16 |
+
p.add_argument('--yam32', action='store_true', help='Use the native 32-second three-view YAM temporal profile.')
|
| 17 |
+
p.add_argument('--instruction', default='', help='Optional episode instruction; blank by default.')
|
| 18 |
+
p.add_argument('--output', type=Path, help='Single-video result; default annotations.json.')
|
| 19 |
+
p.add_argument('--output-dir', type=Path, help='Required with --input-jsonl for inference.')
|
| 20 |
+
p.add_argument('--validate-only', action='store_true', help='Check media/timestamps on CPU; no weight downloads or generation.')
|
| 21 |
+
p.add_argument('--include-audit', action='store_true', help='Include temporal/model-input audit records in the result.')
|
| 22 |
+
p.add_argument('--overwrite', action='store_true', help='Explicitly allow replacing existing output files.')
|
| 23 |
+
return p
|
| 24 |
+
|
| 25 |
+
def jobs(args, p):
|
| 26 |
+
if args.video:
|
| 27 |
+
if args.output_dir:
|
| 28 |
+
p.error('--output-dir is for JSONL jobs; use --output for a single video')
|
| 29 |
+
return [(dict(video=args.video, profile=args.profile, left_wrist=args.left_wrist, right_wrist=args.right_wrist, yam32=args.yam32, instruction=args.instruction), args.output or Path('annotations.json'))]
|
| 30 |
+
if args.output or args.left_wrist or args.right_wrist or args.yam32 or args.instruction or (args.profile != 'human'):
|
| 31 |
+
p.error('Put per-video profile/cameras/instruction in JSONL rows.')
|
| 32 |
+
if not args.output_dir and (not args.validate_only):
|
| 33 |
+
p.error('--output-dir is required with --input-jsonl')
|
| 34 |
+
source = args.input_jsonl.expanduser().resolve(strict=True)
|
| 35 |
+
items, ids = ([], set())
|
| 36 |
+
allowed = {'id', 'video', 'profile', 'left_wrist', 'right_wrist', 'yam32', 'instruction'}
|
| 37 |
+
for line_number, line in enumerate(source.read_text().splitlines(), 1):
|
| 38 |
+
if not line.strip():
|
| 39 |
+
continue
|
| 40 |
+
row = json.loads(line)
|
| 41 |
+
if not isinstance(row, dict) or set(row) - allowed or 'video' not in row:
|
| 42 |
+
raise ValueError(f'Invalid JSONL job on line {line_number}.')
|
| 43 |
+
name = row.pop('id', f'video_{line_number:06d}')
|
| 44 |
+
if not isinstance(name, str) or not re.fullmatch('[A-Za-z0-9][A-Za-z0-9_.-]{0,95}', name) or name in ids:
|
| 45 |
+
raise ValueError(f'Invalid or duplicate job id on line {line_number}.')
|
| 46 |
+
ids.add(name)
|
| 47 |
+
for key in ['video', 'left_wrist', 'right_wrist']:
|
| 48 |
+
if row.get(key) is not None:
|
| 49 |
+
if not isinstance(row[key], str) or not row[key].strip():
|
| 50 |
+
raise ValueError(f'Invalid camera path on JSONL line {line_number}.')
|
| 51 |
+
value = Path(row[key]).expanduser()
|
| 52 |
+
row[key] = value if value.is_absolute() else source.parent / value
|
| 53 |
+
items.append((row, (args.output_dir or Path('.')) / (name + '.json')))
|
| 54 |
+
if not items:
|
| 55 |
+
raise ValueError('The JSONL contains no video jobs.')
|
| 56 |
+
return items
|
| 57 |
+
|
| 58 |
+
def main(argv=None):
|
| 59 |
+
p = parser()
|
| 60 |
+
args = p.parse_args(argv)
|
| 61 |
+
try:
|
| 62 |
+
items = jobs(args, p)
|
| 63 |
+
if not args.validate_only and (not args.overwrite):
|
| 64 |
+
for _, output in items:
|
| 65 |
+
if output.exists():
|
| 66 |
+
raise FileExistsError(f'Output already exists: {output}. Use --overwrite to replace it.')
|
| 67 |
+
engine = DataEnginePipeline()
|
| 68 |
+
for kwargs, output in items:
|
| 69 |
+
if args.validate_only:
|
| 70 |
+
print(json.dumps(engine.inspect_inputs(**kwargs), allow_nan=False))
|
| 71 |
+
continue
|
| 72 |
+
result = engine.annotate(**kwargs, include_audit=args.include_audit, progress=lambda fraction, desc='': print(f'{fraction:.0%} {desc}', file=sys.stderr, flush=True))
|
| 73 |
+
payload = json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False) + '\n'
|
| 74 |
+
output.parent.mkdir(parents=True, exist_ok=True)
|
| 75 |
+
with output.open('w' if args.overwrite else 'x') as f:
|
| 76 |
+
f.write(payload)
|
| 77 |
+
print(json.dumps({'output': str(output), 'segments': len(result['segments'])}))
|
| 78 |
+
return 0
|
| 79 |
+
except (OSError, ValueError, RuntimeError) as error:
|
| 80 |
+
print(f'FLORA inference failed: {error}', file=sys.stderr)
|
| 81 |
+
return 1
|
batch_utils/config.py
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from pathlib import Path
|
| 2 |
+
import json
|
| 3 |
+
import os
|
| 4 |
+
ROOT = Path(__file__).resolve().parent
|
| 5 |
+
VENDOR = ROOT / 'vendor'
|
| 6 |
+
MODELS = json.loads((ROOT / 'models.json').read_text())
|
| 7 |
+
MODEL_CACHE = Path(os.environ.get('FLORA_MODEL_CACHE', Path.home() / '.cache' / 'flora' / 'models'))
|
| 8 |
+
MAX_DURATION_S = 120
|
| 9 |
+
MAX_UPLOAD_BYTES = 256 * 1024 * 1024
|
| 10 |
+
MAX_DIMENSION = 4096
|
| 11 |
+
MAX_FRAME_COUNT = 7200
|
| 12 |
+
PROFILE_POOLS = {'Human': 'human', 'Human actions': 'human', 'YAM': 'yam', 'Dex': 'dex', 'Tianji': 'tianji', 'Robot actions': 'yam'}
|
batch_utils/media.py
ADDED
|
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from dataclasses import dataclass
|
| 2 |
+
from fractions import Fraction
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import hashlib
|
| 5 |
+
import av
|
| 6 |
+
from .config import MAX_DURATION_S, MAX_DIMENSION, MAX_FRAME_COUNT
|
| 7 |
+
|
| 8 |
+
class InputError(ValueError):
|
| 9 |
+
pass
|
| 10 |
+
|
| 11 |
+
@dataclass(frozen=True)
|
| 12 |
+
class Video:
|
| 13 |
+
path: Path
|
| 14 |
+
sha256: str
|
| 15 |
+
rate: Fraction
|
| 16 |
+
pts_us: tuple
|
| 17 |
+
duration_us: int
|
| 18 |
+
size: tuple
|
| 19 |
+
cfr: bool
|
| 20 |
+
|
| 21 |
+
@property
|
| 22 |
+
def duration_s(self):
|
| 23 |
+
return self.duration_us / 1000000
|
| 24 |
+
|
| 25 |
+
def probe(path):
|
| 26 |
+
with av.open(str(path), options={'err_detect': 'explode'}) as container:
|
| 27 |
+
if len(container.streams.video) != 1:
|
| 28 |
+
raise InputError('Use a file containing one video stream.')
|
| 29 |
+
stream = container.streams.video[0]
|
| 30 |
+
stream.codec_context.thread_count = 1
|
| 31 |
+
rate = stream.average_rate
|
| 32 |
+
if not rate or rate < 4:
|
| 33 |
+
raise InputError('The input needs at least 4 source frames per second.')
|
| 34 |
+
size = (stream.codec_context.width, stream.codec_context.height)
|
| 35 |
+
if min(size) <= 0 or max(size) > MAX_DIMENSION:
|
| 36 |
+
raise InputError('Video dimensions must be at most 4096 pixels per side.')
|
| 37 |
+
times, zero, last_length = ([], None, None)
|
| 38 |
+
for frame in container.decode(stream):
|
| 39 |
+
if frame.pts is None:
|
| 40 |
+
raise InputError('Every source frame needs a presentation timestamp.')
|
| 41 |
+
zero = frame.pts if zero is None else zero
|
| 42 |
+
t = Fraction(frame.pts - zero) * stream.time_base * 1000000
|
| 43 |
+
if times and t <= times[-1]:
|
| 44 |
+
raise InputError('Source timestamps must be strictly increasing.')
|
| 45 |
+
if t >= MAX_DURATION_S * 1000000 or len(times) >= MAX_FRAME_COUNT:
|
| 46 |
+
raise InputError(f'Video exceeds {MAX_DURATION_S:g} seconds or {MAX_FRAME_COUNT} frames.')
|
| 47 |
+
times.append(t)
|
| 48 |
+
last_length = Fraction(frame.duration) * stream.time_base * 1000000 if frame.duration else None
|
| 49 |
+
if len(times) < 2:
|
| 50 |
+
raise InputError('The video contains too few source frames.')
|
| 51 |
+
tail = last_length or Fraction(1000000, rate)
|
| 52 |
+
duration = round(times[-1] + tail)
|
| 53 |
+
if duration > MAX_DURATION_S * 1000000:
|
| 54 |
+
raise InputError(f'Video exceeds {MAX_DURATION_S:g} seconds.')
|
| 55 |
+
cfr = all((abs(t - Fraction(i * 1000000, rate)) <= 1 for i, t in enumerate(times)))
|
| 56 |
+
with path.open('rb') as f:
|
| 57 |
+
digest = hashlib.file_digest(f, 'sha256').hexdigest()
|
| 58 |
+
return Video(path, digest, rate, tuple(times), duration, size, cfr)
|
| 59 |
+
|
| 60 |
+
def check_synchronized(videos):
|
| 61 |
+
main = videos['main']
|
| 62 |
+
if len(videos) == 3:
|
| 63 |
+
for role, video in videos.items():
|
| 64 |
+
if video.rate != main.rate or video.pts_us != main.pts_us or video.duration_us != main.duration_us:
|
| 65 |
+
raise InputError('Three camera files must have identical relative frame timestamps, frame counts and duration.')
|
| 66 |
+
if not main.cfr:
|
| 67 |
+
raise InputError('The three-view profile requires constant-rate synchronized source videos.')
|
| 68 |
+
|
| 69 |
+
def decode_selected(video, wanted, canonical=False):
|
| 70 |
+
wanted = set(wanted)
|
| 71 |
+
result = {}
|
| 72 |
+
if not wanted:
|
| 73 |
+
return result
|
| 74 |
+
if min(wanted) < 0 or max(wanted) >= len(video.pts_us):
|
| 75 |
+
raise InputError('A requested source frame is outside the video.')
|
| 76 |
+
with av.open(str(video.path), options={'err_detect': 'explode'}) as container:
|
| 77 |
+
stream = container.streams.video[0]
|
| 78 |
+
stream.codec_context.thread_count = 1
|
| 79 |
+
zero = None
|
| 80 |
+
for i, frame in enumerate(container.decode(stream)):
|
| 81 |
+
zero = frame.pts if zero is None else zero
|
| 82 |
+
if i in wanted:
|
| 83 |
+
if Fraction(frame.pts - zero) * stream.time_base * 1000000 != video.pts_us[i]:
|
| 84 |
+
raise InputError('Source frame timestamps changed while decoding.')
|
| 85 |
+
use_canonical = canonical and frame.format.name == 'yuv420p' and (frame.width % 16 == 0) and (frame.height % 2 == 0) and (int(frame.colorspace) in (2, 5, 6)) and (int(frame.color_range) in (0, 1))
|
| 86 |
+
if use_canonical:
|
| 87 |
+
from canonical_rgb import canonical_rgb
|
| 88 |
+
result[i] = canonical_rgb(frame)
|
| 89 |
+
else:
|
| 90 |
+
result[i] = frame.to_image().convert('RGB')
|
| 91 |
+
if i >= max(wanted):
|
| 92 |
+
break
|
| 93 |
+
if set(result) != wanted:
|
| 94 |
+
raise InputError('A requested source frame could not be decoded.')
|
| 95 |
+
return result
|
| 96 |
+
|
| 97 |
+
def frames_for(videos, wanted, mosaic='YAM'):
|
| 98 |
+
frames = {role: decode_selected(video, wanted, canonical=len(videos) == 3) for role, video in videos.items()}
|
| 99 |
+
if len(videos) == 1:
|
| 100 |
+
return frames['main']
|
| 101 |
+
from composers import compose_top_dominant_mosaic, compose_frame
|
| 102 |
+
if mosaic == 'YAM':
|
| 103 |
+
return {i: compose_top_dominant_mosaic(frames['main'][i], frames['left_wrist'][i], frames['right_wrist'][i]) for i in wanted}
|
| 104 |
+
if mosaic == 'Dex':
|
| 105 |
+
return {i: compose_frame(dict(ego=frames['main'][i], left_wrist=frames['left_wrist'][i], right_wrist=frames['right_wrist'][i]), ('ego', 'left_wrist', 'right_wrist')) for i in wanted}
|
| 106 |
+
raise InputError('Choose a supported camera layout.')
|
batch_utils/models.json
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"tempseg": {
|
| 3 |
+
"repo": "nvidia/Flora-TempSeg",
|
| 4 |
+
"revision": "114a1076f70eb53f4a58ac045505db3087e79347",
|
| 5 |
+
"sha256": "1d80fe544e4cd1a7ff6b0ea4be71e988972408dab81d4c8af72e3a628739d6c3",
|
| 6 |
+
"repo_env": "FLORA_TEMPSEG_REPO"
|
| 7 |
+
},
|
| 8 |
+
"captioner": {
|
| 9 |
+
"repo": "nvidia/Flora-Captioner",
|
| 10 |
+
"revision": "fef0bbfe11829ca7f3539d58da7a6a2c2f109dc6",
|
| 11 |
+
"source_manifest_identity": "2ab6cafd286df7f446bfb84fbc4af7c54773d5e890dd277ac5a0453d0d4fb778",
|
| 12 |
+
"repo_env": "FLORA_CAPTIONER_REPO"
|
| 13 |
+
},
|
| 14 |
+
"tempseg_base": {
|
| 15 |
+
"repo": "Qwen/Qwen3.5-9B",
|
| 16 |
+
"revision": "c202236235762e1c871ad0ccb60c8ee5ba337b9a"
|
| 17 |
+
}
|
| 18 |
+
}
|
batch_utils/models.py
ADDED
|
@@ -0,0 +1,189 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from collections import Counter
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
import gc
|
| 4 |
+
import hashlib
|
| 5 |
+
import json
|
| 6 |
+
import os
|
| 7 |
+
import threading
|
| 8 |
+
import time
|
| 9 |
+
from datetime import datetime, timezone
|
| 10 |
+
from .config import MODELS, VENDOR, MODEL_CACHE
|
| 11 |
+
_LOCK = threading.Lock()
|
| 12 |
+
_MANIFEST_LOCK = threading.Lock()
|
| 13 |
+
_VERIFIED_MANIFESTS = {}
|
| 14 |
+
|
| 15 |
+
class NotReady(RuntimeError):
|
| 16 |
+
pass
|
| 17 |
+
|
| 18 |
+
def verify_cuda_residency(model):
|
| 19 |
+
devices = Counter((str(p.device) for p in model.parameters()))
|
| 20 |
+
if not devices or any((not d.startswith('cuda:') for d in devices)):
|
| 21 |
+
raise NotReady('GPU memory is insufficient for this native Captioner configuration.')
|
| 22 |
+
return dict(devices)
|
| 23 |
+
|
| 24 |
+
def runtime_event(event, **values):
|
| 25 |
+
print('FLORA_RUNTIME ' + json.dumps(dict(event=event, utc=datetime.now(timezone.utc).isoformat(), **values), sort_keys=True), flush=True)
|
| 26 |
+
|
| 27 |
+
def ready():
|
| 28 |
+
if os.environ.get('FLORA_ENABLE_INFERENCE', '0') != '1':
|
| 29 |
+
return (False, 'Preview ready · live annotation is paused.')
|
| 30 |
+
import torch
|
| 31 |
+
if not torch.cuda.is_available():
|
| 32 |
+
return (False, 'Preview ready · GPU hardware is needed for annotation.')
|
| 33 |
+
if sum((torch.cuda.get_device_properties(i).total_memory for i in range(torch.cuda.device_count()))) < 75 * 1024 ** 3:
|
| 34 |
+
return (False, 'Annotation needs an 80 GB GPU or a validated multi-GPU configuration.')
|
| 35 |
+
if not os.environ.get('HF_TOKEN'):
|
| 36 |
+
return (False, 'Annotation is waiting for access to the private weights.')
|
| 37 |
+
return (True, 'Ready to annotate · one video at a time.')
|
| 38 |
+
|
| 39 |
+
def snapshot(kind):
|
| 40 |
+
from huggingface_hub import snapshot_download
|
| 41 |
+
spec = dict(MODELS[kind])
|
| 42 |
+
spec['repo'] = os.environ.get(spec.get('repo_env', ''), spec['repo'])
|
| 43 |
+
if not spec['repo']:
|
| 44 |
+
raise ValueError('Set ' + spec['repo_env'] + ' to the released checkpoint repository.')
|
| 45 |
+
patterns = ['adapter_config.json', 'adapter_model.safetensors', 'config.json', 'generation_config.json', 'model.safetensors.index.json', 'model-*.safetensors', 'processor_config.json', 'tokenizer.json', 'tokenizer_config.json', 'chat_template.jinja', 'README.md', 'MODEL_PROVENANCE.json', 'SOURCE_MANIFEST.sha256']
|
| 46 |
+
return Path(snapshot_download(spec['repo'], revision=spec['revision'], token=os.environ.get('HF_TOKEN'), allow_patterns=patterns, local_dir=MODEL_CACHE / kind))
|
| 47 |
+
|
| 48 |
+
def processor(kind, path):
|
| 49 |
+
from transformers import AutoProcessor
|
| 50 |
+
from transformers.models.qwen3_vl.video_processing_qwen3_vl import Qwen3VLVideoProcessor
|
| 51 |
+
proc = AutoProcessor.from_pretrained(path, local_files_only=True, use_fast=True, min_pixels=1024, max_pixels=589824, video_processor=Qwen3VLVideoProcessor.from_pretrained(path, local_files_only=True, fps=4 if kind == 'tempseg' else 8, max_frames=128 if kind == 'tempseg' else 40))
|
| 52 |
+
proc.chat_template = (path / 'chat_template.jinja').read_text()
|
| 53 |
+
return proc
|
| 54 |
+
|
| 55 |
+
def validate_manifest(path):
|
| 56 |
+
with _MANIFEST_LOCK:
|
| 57 |
+
_validate_manifest(path)
|
| 58 |
+
|
| 59 |
+
def _validate_manifest(path):
|
| 60 |
+
started = time.monotonic()
|
| 61 |
+
allowed_root = path.parent.parent if path.parent.name == 'snapshots' else path
|
| 62 |
+
manifest = (path / 'SOURCE_MANIFEST.sha256').read_bytes()
|
| 63 |
+
key = str(path.resolve())
|
| 64 |
+
manifest_sha = hashlib.sha256(manifest).hexdigest()
|
| 65 |
+
files = []
|
| 66 |
+
|
| 67 |
+
def fingerprint(source):
|
| 68 |
+
resolved = source.resolve()
|
| 69 |
+
if not resolved.is_relative_to(allowed_root.resolve()):
|
| 70 |
+
raise RuntimeError('Invalid model manifest')
|
| 71 |
+
stat = resolved.stat()
|
| 72 |
+
return (str(resolved), stat.st_dev, stat.st_ino, stat.st_size, stat.st_mtime_ns, stat.st_ctime_ns)
|
| 73 |
+
for line in manifest.decode().splitlines():
|
| 74 |
+
if not line.strip():
|
| 75 |
+
continue
|
| 76 |
+
expected, filename = line.split(maxsplit=1)
|
| 77 |
+
relative = Path(filename.lstrip('*'))
|
| 78 |
+
if relative.is_absolute() or '..' in relative.parts:
|
| 79 |
+
raise RuntimeError('Invalid model manifest')
|
| 80 |
+
source = path / relative
|
| 81 |
+
if not source.resolve().is_relative_to(allowed_root.resolve()):
|
| 82 |
+
runtime_event('manifest_path_rejected', file=filename, snapshot=str(path), target=str(source.resolve()), allowed_root=str(allowed_root.resolve()))
|
| 83 |
+
raise RuntimeError('Invalid model manifest')
|
| 84 |
+
files.append((source, expected, fingerprint(source)))
|
| 85 |
+
before = tuple(((str(source), expected, stamp) for source, expected, stamp in files))
|
| 86 |
+
if _VERIFIED_MANIFESTS.get(key) == (manifest_sha, before):
|
| 87 |
+
runtime_event('manifest_verified_reuse', kind=path.name, files=len(files), payload_bytes_read=0, elapsed_s=round(time.monotonic() - started, 3))
|
| 88 |
+
return
|
| 89 |
+
_VERIFIED_MANIFESTS.pop(key, None)
|
| 90 |
+
for source, expected, _ in files:
|
| 91 |
+
with source.open('rb') as f:
|
| 92 |
+
if hashlib.file_digest(f, 'sha256').hexdigest() != expected:
|
| 93 |
+
raise RuntimeError('A model asset failed its integrity check')
|
| 94 |
+
after = tuple(((str(source), expected, fingerprint(source)) for source, expected, _ in files))
|
| 95 |
+
if before != after or manifest != (path / 'SOURCE_MANIFEST.sha256').read_bytes():
|
| 96 |
+
raise RuntimeError('A model asset changed during its integrity check')
|
| 97 |
+
_VERIFIED_MANIFESTS[key] = (manifest_sha, after)
|
| 98 |
+
runtime_event('manifest_full_verification', kind=path.name, files=len(files), payload_bytes_read=sum((stamp[3] for _, _, stamp in files)), elapsed_s=round(time.monotonic() - started, 3))
|
| 99 |
+
|
| 100 |
+
def preflight_adapter_download():
|
| 101 |
+
path = snapshot('tempseg')
|
| 102 |
+
validate_manifest(path)
|
| 103 |
+
runtime_event('adapter_download_preflight_pass', revision=MODELS['tempseg']['revision'], adapter_sha256=MODELS['tempseg']['sha256'], directory=str(path))
|
| 104 |
+
|
| 105 |
+
class Session:
|
| 106 |
+
|
| 107 |
+
def __init__(self):
|
| 108 |
+
self.model = self.proc = None
|
| 109 |
+
self.audit = {}
|
| 110 |
+
|
| 111 |
+
def load(self, kind):
|
| 112 |
+
import torch
|
| 113 |
+
previous_threads = torch.get_num_threads()
|
| 114 |
+
torch.set_num_threads(4)
|
| 115 |
+
load_started = time.monotonic()
|
| 116 |
+
from transformers import AutoModelForImageTextToText
|
| 117 |
+
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5PreTrainedModel
|
| 118 |
+
from production_gdn_fix import _keep_native_gdn_fp32
|
| 119 |
+
from verify_gdn import verify
|
| 120 |
+
self.close()
|
| 121 |
+
runtime_event('model_load_start', kind=kind, cpu_threads=torch.get_num_threads(), previous_cpu_threads=previous_threads)
|
| 122 |
+
_keep_native_gdn_fp32(Qwen3_5PreTrainedModel)
|
| 123 |
+
path = snapshot(kind)
|
| 124 |
+
runtime_event('model_assets_ready', kind=kind, elapsed_s=round(time.monotonic() - load_started, 3))
|
| 125 |
+
validate_manifest(path)
|
| 126 |
+
runtime_event('model_assets_verified', kind=kind, elapsed_s=round(time.monotonic() - load_started, 3))
|
| 127 |
+
if kind == 'tempseg':
|
| 128 |
+
base_spec = MODELS['tempseg_base']
|
| 129 |
+
from huggingface_hub import snapshot_download
|
| 130 |
+
base = snapshot_download(base_spec['repo'], revision=base_spec['revision'], token=os.environ.get('HF_TOKEN'), allow_patterns=['*.safetensors', '*.json', '*.jinja'])
|
| 131 |
+
from tempseg_fp32_loader import load
|
| 132 |
+
records = {}
|
| 133 |
+
self.model = load({'base': base, 'export': str(path), 'adapter_sha256': MODELS[kind]['sha256']}, lambda name, value: records.__setitem__(name, value))
|
| 134 |
+
self.audit[kind] = records
|
| 135 |
+
else:
|
| 136 |
+
index = json.loads((path / 'model.safetensors.index.json').read_text())
|
| 137 |
+
extra = {k for k in index['weight_map'] if '_extra_state' in k}
|
| 138 |
+
memory = {i: f'{max(1, int(torch.cuda.get_device_properties(i).total_memory / 1024 ** 3) - 12)}GiB' for i in range(torch.cuda.device_count())}
|
| 139 |
+
self.model, info = AutoModelForImageTextToText.from_pretrained(path, local_files_only=True, dtype=torch.bfloat16, attn_implementation={'text_config': 'sdpa', 'vision_config': 'sdpa'}, device_map='balanced', max_memory=memory, output_loading_info=True)
|
| 140 |
+
if any((info.get(k) for k in ['missing_keys', 'mismatched_keys', 'error_msgs'])) or not set(info.get('unexpected_keys', [])) <= extra:
|
| 141 |
+
raise RuntimeError('The released Captioner did not load completely')
|
| 142 |
+
devices = verify_cuda_residency(self.model)
|
| 143 |
+
self.model.eval()
|
| 144 |
+
self.audit[kind] = {'gdn': verify(self.model, path), 'loading_info': info, 'parameter_devices': devices, 'parameter_dtypes': dict(Counter((str(p.dtype) for p in self.model.parameters())))}
|
| 145 |
+
runtime_event('model_parameters_verified', kind=kind, elapsed_s=round(time.monotonic() - load_started, 3))
|
| 146 |
+
self.proc = processor(kind, path)
|
| 147 |
+
if kind == 'tempseg':
|
| 148 |
+
adapter = self.audit[kind]['ADAPTER_VALUE_VERIFICATION.json']
|
| 149 |
+
gdn = self.audit[kind]['GDN_VALUE_VERIFICATION.json']
|
| 150 |
+
verification = dict(adapter_tensors=adapter['tensors'], adapter_values_exact=adapter['all_saved_values_exact_before_and_after_placement'], adapter_fp32=adapter['holders_fp32_before_assignment'], gdn_values_exact=gdn['all_values_equal_saved_checkpoint'], gdn_fp32=gdn['all_runtime_fp32'])
|
| 151 |
+
else:
|
| 152 |
+
gdn = self.audit[kind]['gdn']
|
| 153 |
+
verification = dict(gdn_values_exact=gdn['all_values_equal_saved_checkpoint'], gdn_fp32=gdn['all_runtime_fp32'], parameter_dtypes=self.audit[kind]['parameter_dtypes'])
|
| 154 |
+
runtime_event('model_ready', kind=kind, revision=MODELS[kind]['revision'], torch_version=torch.__version__, gpu=torch.cuda.get_device_name(0), allocated_gib=round(torch.cuda.memory_allocated() / 1024 ** 3, 3), **verification)
|
| 155 |
+
return (self.model, self.proc)
|
| 156 |
+
|
| 157 |
+
def close(self):
|
| 158 |
+
self.model = self.proc = None
|
| 159 |
+
gc.collect()
|
| 160 |
+
import torch
|
| 161 |
+
if torch.cuda.is_available():
|
| 162 |
+
torch.cuda.empty_cache()
|
| 163 |
+
|
| 164 |
+
def generate(model, proc, inputs, budget):
|
| 165 |
+
import torch
|
| 166 |
+
inputs = dict(inputs)
|
| 167 |
+
inputs.pop('video_metadata', None)
|
| 168 |
+
n = inputs['input_ids'].shape[1]
|
| 169 |
+
if n + budget > 32768:
|
| 170 |
+
raise InputErrorForModel('The selected video exceeds the native model context budget.')
|
| 171 |
+
input_tensors = {k: {'dtype': str(v.dtype), 'shape': list(v.shape), 'sha256': hashlib.sha256(v.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest()} for k, v in inputs.items() if torch.is_tensor(v)}
|
| 172 |
+
input_tensor_digest = hashlib.sha256(json.dumps(input_tensors, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
|
| 173 |
+
device = next(model.parameters()).device
|
| 174 |
+
args = {k: v.to(device) if torch.is_tensor(v) else v for k, v in inputs.items()}
|
| 175 |
+
eos = proc.tokenizer.eos_token_id
|
| 176 |
+
started = time.monotonic()
|
| 177 |
+
runtime_event('generation_start', input_tokens=n, output_limit=budget, input_tensor_digest=input_tensor_digest, input_tensors=input_tensors)
|
| 178 |
+
with torch.inference_mode():
|
| 179 |
+
output = model.generate(**args, do_sample=False, num_beams=1, max_new_tokens=budget, eos_token_id=eos, pad_token_id=proc.tokenizer.pad_token_id, use_cache=True, max_time=300)
|
| 180 |
+
ids = output[0, n:].cpu().tolist()
|
| 181 |
+
stop = bool(ids) and ids[-1] in ([eos] if isinstance(eos, int) else eos)
|
| 182 |
+
text = proc.tokenizer.decode(ids, skip_special_tokens=True).strip()
|
| 183 |
+
runtime_event('generation_end', input_tokens=n, output_limit=budget, generated_tokens=len(ids), complete=stop and bool(text), elapsed_s=round(time.monotonic() - started, 3), peak_allocated_gib=round(torch.cuda.max_memory_allocated() / 1024 ** 3, 3))
|
| 184 |
+
if not stop or not text:
|
| 185 |
+
raise InputErrorForModel('The model did not return a complete output for this input.')
|
| 186 |
+
return text
|
| 187 |
+
|
| 188 |
+
class InputErrorForModel(RuntimeError):
|
| 189 |
+
pass
|
batch_utils/pipeline.py
ADDED
|
@@ -0,0 +1,169 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import asyncio
|
| 2 |
+
from fractions import Fraction
|
| 3 |
+
import hashlib
|
| 4 |
+
import json
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
import tempfile
|
| 7 |
+
import time
|
| 8 |
+
from .bootstrap import production as a
|
| 9 |
+
from .config import MODELS, PROFILE_POOLS
|
| 10 |
+
from .media import frames_for, check_synchronized, InputError
|
| 11 |
+
from .models import Session, generate, runtime_event
|
| 12 |
+
from .preprocessing import temporal_inputs, human_transport, instruction_suffix, three_view_caption_inputs
|
| 13 |
+
|
| 14 |
+
class Predictor:
|
| 15 |
+
|
| 16 |
+
def __init__(self, videos, source, pool, bound, model, proc, instruction, progress):
|
| 17 |
+
self.videos, self.source, self.pool, self.bound = (videos, source, pool, bound)
|
| 18 |
+
self.model, self.proc, self.instruction, self.progress = (model, proc, instruction, progress)
|
| 19 |
+
self.calls = []
|
| 20 |
+
|
| 21 |
+
def predict(self, request):
|
| 22 |
+
from flora_runtime.video_segmentation.fine_subtask import FINE_RESPONSE_SCHEMA
|
| 23 |
+
if request['source']['source_id'] != self.source.source_id or request['model_binding'] != self.bound:
|
| 24 |
+
raise RuntimeError('Production request binding changed')
|
| 25 |
+
predictions = []
|
| 26 |
+
for payload in request['windows']:
|
| 27 |
+
window = a.window_from_payload(payload, self.source)
|
| 28 |
+
plan = a.temporal_frame_plan(window, list(self.videos['main'].pts_us))
|
| 29 |
+
mosaic = 'Dex' if self.pool in ('dex', 'tianji') else 'YAM'
|
| 30 |
+
images = frames_for(self.videos, set(plan['source_frame_indices']), mosaic=mosaic)
|
| 31 |
+
if self.pool == 'human':
|
| 32 |
+
images = human_transport(images)
|
| 33 |
+
frames = [images[i] for i in plan['source_frame_indices']]
|
| 34 |
+
inputs, prompt = temporal_inputs(frames, window, self.pool, self.proc, list(self.videos), self.instruction)
|
| 35 |
+
self.progress(0.25, desc=f'Finding actions · window {len(self.calls) + 1}')
|
| 36 |
+
text = generate(self.model, self.proc, inputs, 1024)
|
| 37 |
+
predictions.append(dict(scope_id=window.scope_id, prediction_json=text))
|
| 38 |
+
self.calls.append(dict(window=window.worker_payload(), frame_plan=plan, prompt_sha256=hashlib.sha256(prompt.encode()).hexdigest(), output=text))
|
| 39 |
+
return dict(schema_version=FINE_RESPONSE_SCHEMA, source_id=self.source.source_id, model_io_contract=self.bound['model_io_contract'], model_binding=self.bound, predictions=predictions)
|
| 40 |
+
|
| 41 |
+
def segment(videos, pool, session, instruction, progress):
|
| 42 |
+
model, proc = session.load('tempseg')
|
| 43 |
+
main = videos['main']
|
| 44 |
+
source = a.AnnotationSource(source_id='upload_' + main.sha256[:24], location=str(main.path), duration_s=main.duration_s, fps=float(main.rate))
|
| 45 |
+
bound = dict(candidate_id='flora_final', checkpoint_sha256=MODELS['tempseg']['sha256'], model_io_contract=a.fine_window_profile(a.PROFILES[pool]).model_io_contract)
|
| 46 |
+
predictor = Predictor(videos, source, pool, bound, model, proc, instruction, progress)
|
| 47 |
+
if len(videos) == 3:
|
| 48 |
+
fine = a.FineSubtaskSegmenter(predictor=predictor, model_binding=bound, profile_id='robot_atomic_intent_v1', window_profile_id=a.PROFILES[pool], fail_on_unresolved=True, batch_concurrency=1, unresolved_stitch_tolerance_ticks=2, quality_policy=a.FineSegmentationQualityPolicy(drop_guard_only_segments=True))
|
| 49 |
+
engine = a.HierarchicalSubtaskSegmenter(coarse_segmenter=a.NoCoarse(), fine_segmenter=fine, fine_prior_mode='none', fine_failure_policy='adaptive_fine_then_conservative_completion', fine_output_coverage_policy='allow_gaps', adaptive_fine_retry_max_windows=32)
|
| 50 |
+
else:
|
| 51 |
+
engine = a.production_segmenter(predictor, bound, pool)
|
| 52 |
+
segments = asyncio.run(engine.segment_async(source))
|
| 53 |
+
return ([dict(start_s=s.span.start_s, end_s=s.span.end_s, metadata=dict(s.metadata)) for s in segments], predictor.calls)
|
| 54 |
+
|
| 55 |
+
def single_plans(videos, segments, profile):
|
| 56 |
+
if profile in ('Human', 'Human actions'):
|
| 57 |
+
from pts_input_adapter import native_blocks, microseconds
|
| 58 |
+
from flora_runtime.subtask_captioning import u2_contract
|
| 59 |
+
main = videos['main']
|
| 60 |
+
pts = [float(t / 1000000) for t in main.pts_us]
|
| 61 |
+
plans, dispositions = ([], [])
|
| 62 |
+
for i, s in enumerate(segments):
|
| 63 |
+
record = dict(event_id=f'event_{i:04d}', event_index=i, start_s=s['start_s'], end_s=s['end_s'])
|
| 64 |
+
try:
|
| 65 |
+
blocks = native_blocks(pts, microseconds(s['start_s']), microseconds(s['end_s']), main.duration_us, main.rate, u2_contract)
|
| 66 |
+
except ValueError:
|
| 67 |
+
dispositions.append(dict(record, status='unsupported_target_frames'))
|
| 68 |
+
continue
|
| 69 |
+
plans.append(dict(record, blocks=blocks, source_fps_num=main.rate.numerator, source_fps_den=main.rate.denominator))
|
| 70 |
+
dispositions.append(dict(record, status='caption_input'))
|
| 71 |
+
return (plans, dispositions)
|
| 72 |
+
main = videos['main']
|
| 73 |
+
if not main.cfr:
|
| 74 |
+
raise InputError('Use the Human actions profile for videos with variable frame timestamps.')
|
| 75 |
+
return a.caption_plans(segments, dict(fps_num=main.rate.numerator, fps_den=main.rate.denominator, frame_count=len(main.pts_us)))
|
| 76 |
+
|
| 77 |
+
def caption_single(videos, segments, profile, session, instruction, progress):
|
| 78 |
+
from flora_unified.unified_input import processor_inputs, prompt as native_prompt
|
| 79 |
+
from transformers.video_utils import VideoMetadata
|
| 80 |
+
plans, dispositions = single_plans(videos, segments, profile)
|
| 81 |
+
if not plans:
|
| 82 |
+
return ([dict(d, caption='', status=d['status']) for d in dispositions], [])
|
| 83 |
+
model, proc = session.load('captioner')
|
| 84 |
+
rows, audit = ({}, [])
|
| 85 |
+
with tempfile.TemporaryDirectory(prefix='flora-frames-') as temporary:
|
| 86 |
+
root = Path(temporary)
|
| 87 |
+
wanted = {i for plan in plans for block in plan['blocks'] for i in block['frame_indices']}
|
| 88 |
+
images = frames_for(videos, wanted)
|
| 89 |
+
receipts = {}
|
| 90 |
+
for i, image in images.items():
|
| 91 |
+
path = root / f'{i:012d}.png'
|
| 92 |
+
image.save(path, format='PNG')
|
| 93 |
+
receipts[i] = dict(path=path.name, png_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), size=list(image.size))
|
| 94 |
+
for n, plan in enumerate(plans):
|
| 95 |
+
progress(0.55 + 0.4 * n / len(plans), desc=f'Writing captions · {n + 1}/{len(plans)}')
|
| 96 |
+
sample = a.caption_sample(plan, receipts, ['main'])
|
| 97 |
+
if instruction.strip():
|
| 98 |
+
from flora_unified.unified_input import load_frames
|
| 99 |
+
frames, _ = load_frames(sample, root)
|
| 100 |
+
text = proc.apply_chat_template([dict(role='user', content=native_prompt(sample, proc.video_token) + instruction_suffix(instruction))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
|
| 101 |
+
metadata = [VideoMetadata(total_num_frames=len(f), fps=float(videos['main'].rate), frames_indices=list(b['block_local_timestamp_indices'])) for f, b in zip(frames, sample['blocks'], strict=True)]
|
| 102 |
+
inputs = proc(text=[text], videos=frames, video_metadata=metadata, do_sample_frames=False, return_metadata=True, return_tensors='pt')
|
| 103 |
+
else:
|
| 104 |
+
inputs, _, text = processor_inputs(sample, root, proc)
|
| 105 |
+
caption = generate(model, proc, inputs, 512)
|
| 106 |
+
rows[plan['event_id']] = dict(start_s=plan['start_s'], end_s=plan['end_s'], caption=caption, status='complete')
|
| 107 |
+
audit.append(dict(event_id=plan['event_id'], blocks=plan['blocks'], prompt_sha256=hashlib.sha256(text.encode()).hexdigest()))
|
| 108 |
+
return ([rows.get(d['event_id'], dict(start_s=d['start_s'], end_s=d['end_s'], caption='', status=d['status'])) for d in dispositions], audit)
|
| 109 |
+
|
| 110 |
+
def caption_three(videos, segments, session, instruction, progress, profile='YAM'):
|
| 111 |
+
from plans import plans_for_episode
|
| 112 |
+
from robot_caption_geometry import normalized, source_frames
|
| 113 |
+
main = videos['main']
|
| 114 |
+
from math import lcm
|
| 115 |
+
scale = lcm(*(t.denominator for t in main.pts_us)) * 1000000
|
| 116 |
+
ticks = [int(t * scale / 1000000) for t in main.pts_us]
|
| 117 |
+
clock = dict(ticks=ticks, time_base=str(Fraction(1, scale)), fps=str(main.rate), duration_s=main.duration_s, first_pts=0)
|
| 118 |
+
group = 'Dex' if profile in ('Dex', 'Tianji') else 'YAM'
|
| 119 |
+
episode = dict(episode_id='upload', source_id=main.sha256, group=group, units=[])
|
| 120 |
+
prepared, rows, audit = ([], [], [])
|
| 121 |
+
for i, s in enumerate(segments):
|
| 122 |
+
unit = dict(target_id=f'event_{i:04d}', unit_id=f'prediction:{i}', native_scope_s=[s['start_s'], s['end_s']])
|
| 123 |
+
try:
|
| 124 |
+
plan = plans_for_episode(dict(episode, units=[unit]), {role: clock for role in videos})[0]
|
| 125 |
+
prepared.append((i, plan))
|
| 126 |
+
except AssertionError:
|
| 127 |
+
rows.append((i, dict(start_s=s['start_s'], end_s=s['end_s'], caption='', status='unsupported_target_frames')))
|
| 128 |
+
if prepared:
|
| 129 |
+
model, proc = session.load('captioner')
|
| 130 |
+
for n, (i, plan) in enumerate(prepared):
|
| 131 |
+
progress(0.55 + 0.4 * n / len(prepared), desc=f'Writing captions · {n + 1}/{len(prepared)}')
|
| 132 |
+
wanted = {j for block in plan['blocks'] for j in block['indices']}
|
| 133 |
+
images = frames_for(videos, wanted, mosaic=group)
|
| 134 |
+
geometry = normalized(plan, images)
|
| 135 |
+
inputs, text = three_view_caption_inputs(plan, source_frames(geometry), proc, instruction)
|
| 136 |
+
caption = generate(model, proc, inputs, 512)
|
| 137 |
+
rows.append((i, dict(start_s=plan['native_scope_s'][0], end_s=plan['native_scope_s'][1], caption=caption, status='complete')))
|
| 138 |
+
audit.append(dict(event_id=plan['target_id'], blocks=plan['blocks'], common_geometry=geometry['common_geometry'], camera_layout=group, prompt_sha256=hashlib.sha256(text.encode()).hexdigest()))
|
| 139 |
+
return ([row for _, row in sorted(rows)], audit)
|
| 140 |
+
|
| 141 |
+
def annotate(videos, profile, long_windows, instruction, progress):
|
| 142 |
+
check_synchronized(videos)
|
| 143 |
+
instruction_suffix(instruction)
|
| 144 |
+
if profile not in PROFILE_POOLS:
|
| 145 |
+
raise InputError('Choose a supported action profile.')
|
| 146 |
+
if len(videos) == 3 and profile in ('Human', 'Human actions'):
|
| 147 |
+
raise InputError('Choose YAM, Dex or Tianji when uploading wrist cameras.')
|
| 148 |
+
if long_windows and len(videos) != 3:
|
| 149 |
+
raise InputError('The 32-second profile requires three synchronized camera views.')
|
| 150 |
+
if long_windows and profile not in ('YAM', 'Robot actions'):
|
| 151 |
+
raise InputError('The 32-second profile is available for YAM videos.')
|
| 152 |
+
pool = 'yam32' if long_windows else PROFILE_POOLS[profile]
|
| 153 |
+
start = time.monotonic()
|
| 154 |
+
session = Session()
|
| 155 |
+
try:
|
| 156 |
+
progress(0.1, desc='Loading the action segmenter')
|
| 157 |
+
segments, temporal_audit = segment(videos, pool, session, instruction, progress)
|
| 158 |
+
runtime_event('segmentation_complete', windows=len(temporal_audit), segments=len(segments))
|
| 159 |
+
progress(0.5, desc='Stitching finished · loading the captioner')
|
| 160 |
+
if len(videos) == 3:
|
| 161 |
+
captions, caption_audit = caption_three(videos, segments, session, instruction, progress, profile=profile)
|
| 162 |
+
else:
|
| 163 |
+
captions, caption_audit = caption_single(videos, segments, profile, session, instruction, progress)
|
| 164 |
+
for row, segment_row in zip(captions, segments, strict=True):
|
| 165 |
+
row['segmentation_provenance'] = segment_row['metadata'].get('segmentation_provenance', 'fine')
|
| 166 |
+
runtime_event('annotation_complete', segments=len(captions), complete_captions=sum((r['status'] == 'complete' for r in captions)), elapsed_s=round(time.monotonic() - start, 3))
|
| 167 |
+
return dict(schema_version='flora_demo_annotation_v1', duration_s=videos['main'].duration_s, segments=captions, elapsed_s=round(time.monotonic() - start, 2), source_clock='seconds relative to the first native source presentation timestamp', camera_roles=list(videos), profile=profile, model_revisions=MODELS, input_sha256={role: video.sha256 for role, video in videos.items()}, pipeline=dict(temporal_windows=temporal_audit, caption_inputs=caption_audit))
|
| 168 |
+
finally:
|
| 169 |
+
session.close()
|
batch_utils/preprocessing.py
ADDED
|
@@ -0,0 +1,56 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from fractions import Fraction
|
| 2 |
+
import io
|
| 3 |
+
import json
|
| 4 |
+
from PIL import Image
|
| 5 |
+
from .bootstrap import production as a
|
| 6 |
+
|
| 7 |
+
def instruction_suffix(instruction):
|
| 8 |
+
value = instruction.strip()
|
| 9 |
+
if not value:
|
| 10 |
+
return ''
|
| 11 |
+
if len(value) > 2000:
|
| 12 |
+
raise ValueError('Keep the optional task instruction below 2000 characters.')
|
| 13 |
+
return '\n' + json.dumps({'episode_instruction': value}, ensure_ascii=False, sort_keys=True, separators=(',', ':'), allow_nan=False)
|
| 14 |
+
|
| 15 |
+
def human_transport(frames):
|
| 16 |
+
from flora_runtime.exact_frame_serving import encode_regularized_image
|
| 17 |
+
output = {}
|
| 18 |
+
for index, image in frames.items():
|
| 19 |
+
payload = encode_regularized_image(image, encoding='jpeg_q98_444')['payload']
|
| 20 |
+
with Image.open(io.BytesIO(payload)) as decoded:
|
| 21 |
+
output[index] = decoded.convert('RGB')
|
| 22 |
+
return output
|
| 23 |
+
|
| 24 |
+
def temporal_inputs(frames, window, pool, proc, roles, instruction=''):
|
| 25 |
+
from transformers.video_utils import VideoMetadata
|
| 26 |
+
from flora_runtime.exact_frame_serving import regularized_resize_dims
|
| 27 |
+
profile = window.window_profile
|
| 28 |
+
if len(frames) != profile.slot_count:
|
| 29 |
+
raise ValueError('Incorrect native temporal slot count')
|
| 30 |
+
regularized = []
|
| 31 |
+
for image in frames:
|
| 32 |
+
image = image.convert('RGB')
|
| 33 |
+
height, width = regularized_resize_dims(height=image.height, width=image.width, min_pixels=256, max_pixels=40000 if len(roles) == 1 else 160000)
|
| 34 |
+
if image.size != (width, height):
|
| 35 |
+
image = image.resize((width, height), Image.Resampling.BICUBIC)
|
| 36 |
+
regularized.append(image)
|
| 37 |
+
prompt = a.temporal_prompt(window, pool) + instruction_suffix(instruction)
|
| 38 |
+
text = proc.apply_chat_template([dict(role='user', content=prompt.replace('<video>', proc.video_token))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
|
| 39 |
+
inputs = proc(text=[text], videos=[regularized], video_metadata=[VideoMetadata(total_num_frames=profile.slot_count, fps=profile.slot_fps, frames_indices=list(range(profile.slot_count)))], do_sample_frames=False, return_metadata=True, return_tensors='pt')
|
| 40 |
+
inputs.pop('video_metadata', None)
|
| 41 |
+
return (inputs, text)
|
| 42 |
+
|
| 43 |
+
def three_view_caption_inputs(plan, frames, proc, instruction=''):
|
| 44 |
+
from transformers.video_utils import VideoMetadata
|
| 45 |
+
from plans import HEADING, INSTRUCTION
|
| 46 |
+
time_base = Fraction(plan['time_base'])
|
| 47 |
+
content, videos, metadata = ([HEADING], [], [])
|
| 48 |
+
for block in plan['blocks']:
|
| 49 |
+
times = ', '.join((f'{float(t * time_base):.9f}' for t in block['pts_ticks']))
|
| 50 |
+
content.append(block['role'] + ':\nFrames are in this order; native source timestamps in seconds: [' + times + '].\n' + proc.video_token)
|
| 51 |
+
videos.append([frames[i] for i in block['indices']])
|
| 52 |
+
metadata.append(VideoMetadata(total_num_frames=len(block['indices']), fps=float(Fraction(plan['clock_fps'])), frames_indices=[i - block['source_start_frame'] for i in block['indices']]))
|
| 53 |
+
content.append(INSTRUCTION + instruction_suffix(instruction))
|
| 54 |
+
text = proc.apply_chat_template([dict(role='user', content=''.join(content))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
|
| 55 |
+
result = proc(text=[text], videos=videos, video_metadata=metadata, do_sample_frames=False, return_tensors='pt', padding=True)
|
| 56 |
+
return (result, text)
|
batch_utils/requirements.txt
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
torch==2.11.0
|
| 2 |
+
torchvision==0.26.0
|
| 3 |
+
huggingface_hub==1.33.0
|
| 4 |
+
transformers==5.8.1
|
| 5 |
+
peft==0.18.1
|
| 6 |
+
accelerate==1.12.0
|
| 7 |
+
safetensors==0.7.0
|
| 8 |
+
av==16.0.1
|
| 9 |
+
numpy<3
|
| 10 |
+
Pillow>=11,<13
|
batch_utils/vendor/MANIFEST.json
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"flora_unified/__init__.py": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2",
|
| 3 |
+
"flora_unified/unified_input.py": "2aa2b521191b68913b03d7582b7a3bcfad49d90c014a1a1bb7571b5e4f08d7ef",
|
| 4 |
+
"production_adapter/TEMPSEG_PROMPTS.json": "66edd6fb43b7fa55cc6257beee5c42a1a4555ab9790320c31114d77bd07a05f1",
|
| 5 |
+
"production_adapter/adapter.py": "083c6788b561d726a85d9f3f8f516593a9f6f76a43f315581c821871e3f82312",
|
| 6 |
+
"production_adapter/core/flora_runtime/__init__.py": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2",
|
| 7 |
+
"production_adapter/core/flora_runtime/exact_frame_serving.py": "f921e279b5d52f723b3e77e25441763683d05c6e9cd2f8ff204466a50255b346",
|
| 8 |
+
"production_adapter/core/flora_runtime/schema.py": "42c523e4db1a4da3fe0d0d578f3cf6e94aed7666c06e3f18081c209159de04b9",
|
| 9 |
+
"production_adapter/core/flora_runtime/subtask_captioning/__init__.py": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2",
|
| 10 |
+
"production_adapter/core/flora_runtime/subtask_captioning/u2_contract.py": "873476e78a9e3e222ae9e47fcb7a04c8c18ef518b2847c401a2e1cfb2292fe90",
|
| 11 |
+
"production_adapter/core/flora_runtime/subtask_captioning/u2_worker.py": "870641228b0be442613d520232d451047742ca2ea41fd0cd51dbf0a58ba567d2",
|
| 12 |
+
"production_adapter/core/flora_runtime/video_segmentation/__init__.py": "9f56e761d79bfdb34304a012586cb04d16b435ef6130091a97702e559260a2f2",
|
| 13 |
+
"production_adapter/core/flora_runtime/video_segmentation/contracts.py": "8591d849b0b65fda6b21511557d4ee5eabd47f8332e2b6831f40d5bb472d978c",
|
| 14 |
+
"production_adapter/core/flora_runtime/video_segmentation/fine_subtask.py": "6fbf8337bab8491148a78a822187d74397b7f0552700252451a18cd7c78c8e04",
|
| 15 |
+
"production_adapter/core/flora_runtime/video_segmentation/fine_temporal.py": "32f2b99c97176f5e10934450da638577268c2e5d9977c2ea138ebf1b7eda58fa",
|
| 16 |
+
"production_adapter/core/flora_runtime/video_segmentation/fine_worker.py": "8a31a77180c8d406113d4b494644153564a6f46c8a9b68ad7babb3cc1695edec",
|
| 17 |
+
"production_adapter/core/flora_runtime/video_segmentation/interface.py": "a5424b2cc3a4c879552dd49bae207a58fb9e751359f441c03af167a2e1879359",
|
| 18 |
+
"production_adapter/core/flora_runtime/video_segmentation/scope_scheduler.py": "be95411b39124d83811fe74dcac2d911169ecf475a81b7460c64b19fbbc7e5db",
|
| 19 |
+
"production_adapter/core/flora_runtime/worker_transport.py": "30fb64869eb6a27ab96bf394c105f19fcb42590be4424c66b76b89e17631693a",
|
| 20 |
+
"production_gdn_fix.py": "a3f52604bfa2ec58901df63215df5e47d37e1d7974aa56d36a1e4efe8cad3fc6",
|
| 21 |
+
"pts_input_adapter.py": "d784d6b77005eb6945dee334e4c9136fbededfa2049d5945244f198517984767",
|
| 22 |
+
"robot_caption_geometry.py": "1c874d3c63359663276d31a3e9c3112a5982ad4000f44c808541b2471f5f7d2d",
|
| 23 |
+
"support/canonical_rgb.py": "1e3230c60b45d8e017060d8a5dd2b5f48e2199192f09dcff4d04c05b4349e35a",
|
| 24 |
+
"support/composers.py": "180b97a6a634979c1b0dbe720ed7a1abd818dc7ee31fdc4eeca7a8277e43d1b6",
|
| 25 |
+
"support/plans.py": "d75b918804c9ec10f8642d362a54a74c3eb875abf22e8a60228a0d6b2cbb570e",
|
| 26 |
+
"support/u2_contract.py": "50190914b719bbc9b8f2b53e52a6bd2b142a452e2fb1c5c897ea4c31123ad4dd",
|
| 27 |
+
"tempseg_fp32_loader.py": "7ee9889f5aa5909abf3ecd1c34e3ac2f81a32a2098d94ae3ba9f33831e7a07c0",
|
| 28 |
+
"verify_gdn.py": "96596ff2db4425f4dea71798aeabb581aa83e1ad8a878cfcb9efb8dbac80cda6"
|
| 29 |
+
}
|
batch_utils/vendor/flora_unified/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pass
|
batch_utils/vendor/flora_unified/unified_input.py
ADDED
|
@@ -0,0 +1,89 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from fractions import Fraction
|
| 2 |
+
import hashlib
|
| 3 |
+
import math
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
from PIL import Image
|
| 6 |
+
INTENT_INSTRUCTION = 'Write one concise, self-contained English instruction for the local operation being attempted in TARGET.\nUse a natural imperative, concrete objects and the working hand; include meaningful support by the other hand and necessary source or destination. Say left hand, right hand or both hands. For a handover, name the object and both the giving and receiving hands.\nInfer the local goal from the video sequence. Use BEFORE and AFTER only to disambiguate TARGET, not to add operations occurring outside it. A failed or repeated attempt retains its established instruction; do not replace the goal with an accident, a success claim or a description of mechanical adjustments. Do not invent an object, hand role or hidden goal.\nOutput only the instruction, on one line, without timestamps, execution-quality commentary or the block names.'
|
| 7 |
+
|
| 8 |
+
def validate(sample):
|
| 9 |
+
blocks = sample['blocks']
|
| 10 |
+
roles = [b['role'] for b in blocks]
|
| 11 |
+
if roles not in [['TARGET'], ['BEFORE', 'TARGET'], ['TARGET', 'AFTER'], ['BEFORE', 'TARGET', 'AFTER']]:
|
| 12 |
+
raise ValueError('Missing TARGET, duplicate or reordered temporal role')
|
| 13 |
+
if sample['camera_count'] not in (1, 3):
|
| 14 |
+
raise ValueError('Unimplemented camera renderer')
|
| 15 |
+
for b in blocks:
|
| 16 |
+
start, end = (b['source_start_frame'], b['source_end_frame_exclusive'])
|
| 17 |
+
indices = b['block_local_timestamp_indices']
|
| 18 |
+
n = len(b['paths'])
|
| 19 |
+
fps = Fraction(b['clock_fps_num'], b['clock_fps_den'])
|
| 20 |
+
if fps <= 0 or end <= start or (not n) or (n != len(indices)) or (n != len(b['frames'])):
|
| 21 |
+
raise ValueError('Invalid exact frame plan')
|
| 22 |
+
if any((i < 0 or i >= end - start for i in indices)):
|
| 23 |
+
raise ValueError('Frame outside block clock')
|
| 24 |
+
if any((a > b for a, b in zip(indices, indices[1:]))):
|
| 25 |
+
raise ValueError('Non-monotone timestamps')
|
| 26 |
+
original = b.get('original_descriptor', {})
|
| 27 |
+
if 'frame_indices' in original and original['frame_indices'] != [start + i for i in indices]:
|
| 28 |
+
raise ValueError('Source frame ordinals disagree with block-local clock')
|
| 29 |
+
if len(set(indices)) != n and (not any(original.get('frame_cache_padding_mask', []))):
|
| 30 |
+
raise ValueError('Repeated frame requires explicit padding evidence')
|
| 31 |
+
if Fraction(b['effective_fps_num'], b['effective_fps_den']) != Fraction(n) * fps / (end - start):
|
| 32 |
+
raise ValueError('Effective sampling rate does not match block duration')
|
| 33 |
+
if n % 2:
|
| 34 |
+
raise ValueError('No implicit temporal pair padding')
|
| 35 |
+
expected = sample['request']['target']['native_span']
|
| 36 |
+
target = next((b for b in blocks if b['role'] == 'TARGET'))
|
| 37 |
+
if expected['start_frame'] is not None and (target['source_start_frame'], target['source_end_frame_exclusive']) != (expected['start_frame'], expected['end_frame_exclusive']):
|
| 38 |
+
raise ValueError('Accepted target span changed')
|
| 39 |
+
text = sample['request']['target']['text']
|
| 40 |
+
if hashlib.sha256(text.encode()).hexdigest() != sample['request']['target']['sha256']:
|
| 41 |
+
raise ValueError('Accepted target changed')
|
| 42 |
+
|
| 43 |
+
def prompt(sample, video_token):
|
| 44 |
+
validate(sample)
|
| 45 |
+
heading = ''
|
| 46 |
+
if sample['camera_count'] == 3:
|
| 47 |
+
if sample['camera_roles'] != ['main', 'left_wrist', 'right_wrist']:
|
| 48 |
+
raise ValueError('Camera roles do not match the mosaic renderer')
|
| 49 |
+
heading = 'Each frame is a synchronized camera mosaic: main camera above, left wrist camera below left, and right wrist camera below right. All panels show the same physical time. Describe each physical operation once across the views.\n\n'
|
| 50 |
+
return heading + '\n\n'.join((b['role'] + ':\n' + video_token for b in sample['blocks'])) + '\n\n' + INTENT_INSTRUCTION
|
| 51 |
+
|
| 52 |
+
def load_frames(sample, root, preserve_canvas=True):
|
| 53 |
+
validate(sample)
|
| 54 |
+
videos = []
|
| 55 |
+
sizes = []
|
| 56 |
+
max_pixels = 524288 if preserve_canvas and sample['camera_count'] == 3 else 262144
|
| 57 |
+
for b in sample['blocks']:
|
| 58 |
+
frames = []
|
| 59 |
+
shapes = []
|
| 60 |
+
for path, receipt in zip(b['paths'], b['frames'], strict=True):
|
| 61 |
+
p = Path(root) / path
|
| 62 |
+
if not p.resolve().is_relative_to(Path(root).resolve()):
|
| 63 |
+
raise ValueError('Unowned frame path')
|
| 64 |
+
raw = p.read_bytes()
|
| 65 |
+
if hashlib.sha256(raw).hexdigest() != receipt['png_sha256']:
|
| 66 |
+
raise ValueError('PNG changed')
|
| 67 |
+
with Image.open(p) as im:
|
| 68 |
+
im.load()
|
| 69 |
+
frame = im.convert('RGB')
|
| 70 |
+
if list(frame.size) != receipt['size']:
|
| 71 |
+
raise ValueError('Frame dimensions changed')
|
| 72 |
+
original = list(frame.size)
|
| 73 |
+
area = frame.width * frame.height
|
| 74 |
+
if area > max_pixels:
|
| 75 |
+
scale = math.sqrt(max_pixels / area)
|
| 76 |
+
frame = frame.resize((int(frame.width * scale), int(frame.height * scale)))
|
| 77 |
+
frames.append(frame)
|
| 78 |
+
shapes.append(dict(original=original, regularized=list(frame.size)))
|
| 79 |
+
videos.append(frames)
|
| 80 |
+
sizes.append(shapes)
|
| 81 |
+
return (videos, sizes)
|
| 82 |
+
|
| 83 |
+
def processor_inputs(sample, root, processor, preserve_canvas=True):
|
| 84 |
+
from transformers.video_utils import VideoMetadata
|
| 85 |
+
videos, sizes = load_frames(sample, root, preserve_canvas)
|
| 86 |
+
metadata = [VideoMetadata(total_num_frames=len(frames), fps=float(Fraction(b['clock_fps_num'], b['clock_fps_den'])), frames_indices=list(b['block_local_timestamp_indices'])) for frames, b in zip(videos, sample['blocks'], strict=True)]
|
| 87 |
+
text = processor.apply_chat_template([dict(role='user', content=prompt(sample, processor.video_token))], tokenize=False, add_generation_prompt=True, enable_thinking=False)
|
| 88 |
+
inputs = processor(text=[text], videos=videos, video_metadata=metadata, do_sample_frames=False, return_metadata=True, return_tensors='pt')
|
| 89 |
+
return (inputs, sizes, text)
|
batch_utils/vendor/production_adapter/TEMPSEG_PROMPTS.json
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"pool": "dex",
|
| 4 |
+
"prompt": "<video>\nThe video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\nThe emit core is 12 seconds with a 2-second guard on each side.\nSegment meaningful local operations. Acquisition and delivery are distinct when observable; retain distinct directional handovers. Ordinary contacts, adjustments and low-motion pauses within the same ongoing goal do not create new events. Keep clear failure/retry or execution-quality transitions as separate intervals even when their intended operation is unchanged. Represent genuine pure rest as an interval, not a gap.\nEdge padding repeats the nearest real sampled frame; ignore slots outside the valid tick envelope.\nTimes use 20 Hz local ticks in the half-open range [0,320).\nCanonical window envelope:\n{\"core_end_tick\":280,\"core_start_tick\":40,\"valid_end_tick\":320,\"valid_start_tick\":0}\nEmit only intervals with positive half-open intersection with [core_start_tick,core_end_tick). Clip intervals at valid video limits; lo/ro indicate continuation beyond the visible limit. lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\nReturn exactly one JSON object and no prose or Markdown, with only top-level key segments. Each row has keys exactly in order s,e,lo,ro. Sort rows by (s,e,lo,ro).",
|
| 5 |
+
"prompt_sha256": "a50644d790524e1d4c80a5ad154e940d9bb512b0816eb9d2703cf5721a2f1943"
|
| 6 |
+
},
|
| 7 |
+
{
|
| 8 |
+
"pool": "human",
|
| 9 |
+
"prompt": "<video>\nThe video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\nThe emit core is 12 seconds with a 2-second guard on each side.\nEdge padding repeats the nearest real frame; ignore slots outside the valid tick envelope.\nTimes use 20 Hz local ticks in the half-open range [0,320).\nCanonical window envelope:\n{\"valid_start_tick\":0,\"valid_end_tick\":320,\"core_start_tick\":40,\"core_end_tick\":280}\nEmit only subtask intervals with positive half-open intersection with [core_start_tick,core_end_tick).\nA quantized boundary alias may instead end exactly at core_start_tick or start exactly at core_end_tick.\nlo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\nReturn exactly one JSON object and no prose or Markdown.\nUse only the top-level key segments; each row's keys must appear exactly as s,e,lo,ro.\nPreserve concurrent or duplicate intervals as separate rows and sort rows by (s,e,lo,ro).\nAn empty list is allowed only when no interval has emit-core support.\n",
|
| 10 |
+
"prompt_sha256": "21521d5e3b854488f022f402cbb3c31137e2a1a8a114837b16e353c631bc350a"
|
| 11 |
+
},
|
| 12 |
+
{
|
| 13 |
+
"pool": "tianji",
|
| 14 |
+
"prompt": "<video>\nThe video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\nThe emit core is 12 seconds with a 2-second guard on each side.\nEdge padding repeats the nearest real frame; ignore slots outside the valid tick envelope.\nTimes use 20 Hz local ticks in the half-open range [0,320).\nCanonical window envelope:\n{\"core_end_tick\":280,\"core_start_tick\":40,\"valid_end_tick\":320,\"valid_start_tick\":40}\nEmit only atomic intervals with positive half-open intersection with [core_start_tick,core_end_tick).\nA quantized boundary alias may instead end exactly at core_start_tick or start exactly at core_end_tick.\nlo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\nReturn exactly one JSON object and no prose or Markdown.\nUse only the top-level key segments; each row's keys must appear exactly as s,e,lo,ro.\nPreserve concurrent or duplicate atomic intervals as separate rows and sort rows by (s,e,lo,ro).\nAn empty list is allowed only when no atomic interval has emit-core support.\n",
|
| 15 |
+
"prompt_sha256": "497e9adbd8ae4ffc76f54ee4f42bd20d597a44667709cee7fccf447c3f504243"
|
| 16 |
+
},
|
| 17 |
+
{
|
| 18 |
+
"pool": "yam32",
|
| 19 |
+
"prompt": "<video>\nThe video contains 128 ordered temporal slots sampled at 4 Hz for one 32-second local view.\nThe emit core is 24 seconds with a 4-second guard on each side.\nSegment meaningful local operations. Acquisition and delivery are distinct when observable; retain distinct directional handovers. Ordinary contacts, adjustments and low-motion pauses within the same ongoing goal do not create new events. Keep clear failure/retry or execution-quality transitions as separate intervals even when their intended operation is unchanged. Represent genuine pure rest as an interval, not a gap.\nEdge padding repeats the nearest real sampled frame; ignore slots outside the valid tick envelope.\nTimes use 20 Hz local ticks in the half-open range [0,640).\nCanonical window envelope:\n{\"core_end_tick\":560,\"core_start_tick\":80,\"valid_end_tick\":640,\"valid_start_tick\":0}\nEmit only intervals with positive half-open intersection with [core_start_tick,core_end_tick). Clip intervals at valid video limits; lo/ro indicate continuation beyond the visible limit. lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\nReturn exactly one JSON object and no prose or Markdown, with only top-level key segments. Each row has keys exactly in order s,e,lo,ro. Sort rows by (s,e,lo,ro).",
|
| 20 |
+
"prompt_sha256": "7c6913c78b3fddca0507d6f25534bb89797160ac6342b28f0117b17063ba116f"
|
| 21 |
+
},
|
| 22 |
+
{
|
| 23 |
+
"pool": "yam",
|
| 24 |
+
"prompt": "<video>\nThe video contains 64 ordered temporal slots sampled at 4 Hz for one 16-second local view.\nThe emit core is 12 seconds with a 2-second guard on each side.\nSegment meaningful local operations. Acquisition and delivery are distinct when observable; retain distinct directional handovers. Ordinary contacts, adjustments and low-motion pauses within the same ongoing goal do not create new events. Keep clear failure/retry or execution-quality transitions as separate intervals even when their intended operation is unchanged. Represent genuine pure rest as an interval, not a gap.\nEdge padding repeats the nearest real sampled frame; ignore slots outside the valid tick envelope.\nTimes use 20 Hz local ticks in the half-open range [0,320).\nCanonical window envelope:\n{\"core_end_tick\":280,\"core_start_tick\":40,\"valid_end_tick\":320,\"valid_start_tick\":0}\nEmit only intervals with positive half-open intersection with [core_start_tick,core_end_tick). Clip intervals at valid video limits; lo/ro indicate continuation beyond the visible limit. lo=true requires s==valid_start_tick; ro=true requires e==valid_end_tick.\nReturn exactly one JSON object and no prose or Markdown, with only top-level key segments. Each row has keys exactly in order s,e,lo,ro. Sort rows by (s,e,lo,ro).",
|
| 25 |
+
"prompt_sha256": "a50644d790524e1d4c80a5ad154e940d9bb512b0816eb9d2703cf5721a2f1943"
|
| 26 |
+
}
|
| 27 |
+
]
|
batch_utils/vendor/production_adapter/adapter.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from dataclasses import asdict
|
| 2 |
+
from fractions import Fraction
|
| 3 |
+
import hashlib
|
| 4 |
+
import json
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
import re
|
| 7 |
+
import sys
|
| 8 |
+
sys.path.insert(0, str(Path(__file__).resolve().parent / 'core'))
|
| 9 |
+
from flora_runtime.schema import AnnotationSource, TemporalSegment, TimeSpan
|
| 10 |
+
from flora_runtime.video_segmentation.fine_subtask import FineSubtaskSegmenter, HierarchicalSubtaskSegmenter, FineSegmentationQualityPolicy
|
| 11 |
+
from flora_runtime.video_segmentation.fine_temporal import FineWindow, fine_window_profile, TICK_US, HUMAN_FINE_WINDOW_PROFILE_ID, YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID, YAM_INTENT_32S_FINE_WINDOW_PROFILE_ID
|
| 12 |
+
from flora_runtime.video_segmentation.fine_worker import _nearest_index
|
| 13 |
+
from flora_runtime.subtask_captioning.u2_worker import _plan_blocks, U2WorkerError
|
| 14 |
+
HERE = Path(__file__).resolve().parent
|
| 15 |
+
TEMPLATES = {x['pool']: x for x in json.loads((HERE / 'TEMPSEG_PROMPTS.json').read_text())}
|
| 16 |
+
PROFILES = {'human': HUMAN_FINE_WINDOW_PROFILE_ID, 'tianji': YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID, 'dex': YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID, 'yam': YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID, 'yam32': YAM_INTENT_32S_FINE_WINDOW_PROFILE_ID}
|
| 17 |
+
|
| 18 |
+
def temporal_prompt(window, pool):
|
| 19 |
+
assert window.window_profile_id == PROFILES[pool]
|
| 20 |
+
row = TEMPLATES[pool]
|
| 21 |
+
template = row['prompt']
|
| 22 |
+
assert hashlib.sha256(template.encode()).hexdigest() == row['prompt_sha256']
|
| 23 |
+
pattern = '(Canonical window envelope:\\s*)(\\{[^{}]+\\})'
|
| 24 |
+
matches = list(re.finditer(pattern, template))
|
| 25 |
+
assert len(matches) == 1
|
| 26 |
+
match = matches[0]
|
| 27 |
+
keys = list(json.loads(match[2]))
|
| 28 |
+
assert set(keys) == {'valid_start_tick', 'valid_end_tick', 'core_start_tick', 'core_end_tick'}
|
| 29 |
+
envelope = json.dumps({k: getattr(window, k) for k in keys}, separators=(',', ':'))
|
| 30 |
+
return template[:match.start(2)] + envelope + template[match.end(2):]
|
| 31 |
+
|
| 32 |
+
def window_from_payload(payload, source):
|
| 33 |
+
start = payload['requested_view_start_us']
|
| 34 |
+
fields = {k: payload[k] for k in ('window_index', 'requested_view_start_us', 'requested_view_end_us', 'valid_start_tick', 'valid_end_tick', 'core_start_tick', 'core_end_tick', 'scope_id', 'window_profile_id', 'coarse_spans')}
|
| 35 |
+
duration = round(source.duration_s * 1000000)
|
| 36 |
+
fields.update(source_id=source.source_id, video_end_us=duration, valid_start_us=max(0, start), valid_end_us=min(duration, payload['requested_view_end_us']), core_start_us=max(0, start + payload['core_start_tick'] * TICK_US), core_end_us=min(duration, start + payload['core_end_tick'] * TICK_US))
|
| 37 |
+
return FineWindow(**fields)
|
| 38 |
+
|
| 39 |
+
class NoCoarse:
|
| 40 |
+
|
| 41 |
+
def segment(self, *_args, **_kwargs):
|
| 42 |
+
raise AssertionError('Unexpected coarse segmentation request')
|
| 43 |
+
|
| 44 |
+
def production_segmenter(predictor, binding, pool):
|
| 45 |
+
assert binding['model_io_contract'] == fine_window_profile(PROFILES[pool]).model_io_contract
|
| 46 |
+
fine = FineSubtaskSegmenter(predictor=predictor, model_binding=binding, profile_id='human_ego_atomic_v0_2', window_profile_id=PROFILES[pool], fail_on_unresolved=True, batch_concurrency=1, unresolved_stitch_tolerance_ticks=2, quality_policy=FineSegmentationQualityPolicy.from_value({'drop_guard_only_segments': True, 'hard_max_overlap_depth': 3, 'hard_max_segments_per_window': 23, 'max_adjacent_matching_ambiguities': 0, 'require_event_when_coarse_present': True, 'schema_version': 'subtask_annotation_atomic_fine_quality_policy_v1', 'soft_max_segments_per_window': 12}))
|
| 47 |
+
return HierarchicalSubtaskSegmenter(coarse_segmenter=NoCoarse(), fine_segmenter=fine, fine_prior_mode='none', fine_failure_policy='adaptive_fine_then_conservative_completion', fine_output_coverage_policy='allow_gaps', adaptive_fine_retry_max_windows=32)
|
| 48 |
+
|
| 49 |
+
def temporal_frame_plan(window, pts_us):
|
| 50 |
+
assert pts_us and pts_us[0] == 0 and all((a < b for a, b in zip(pts_us, pts_us[1:])))
|
| 51 |
+
assert pts_us[-1] < window.video_end_us
|
| 52 |
+
profile = window.window_profile
|
| 53 |
+
requested = [window.requested_view_start_us + i * profile.slot_spacing_us for i in range(profile.slot_count)]
|
| 54 |
+
pad = [not 0 <= t < window.video_end_us for t in requested]
|
| 55 |
+
indices = [None if padding else _nearest_index(t, pts_us, max_selection_error_us=profile.slot_spacing_us // 2) for t, padding in zip(requested, pad, strict=True)]
|
| 56 |
+
real = [i for i, v in enumerate(indices) if v is not None]
|
| 57 |
+
assert real
|
| 58 |
+
first, last = (real[0], real[-1])
|
| 59 |
+
indices[:first] = [indices[first]] * first
|
| 60 |
+
indices[last + 1:] = [indices[last]] * (len(indices) - last - 1)
|
| 61 |
+
assert all((v is not None for v in indices))
|
| 62 |
+
observed = [v for v, p in zip(indices, pad, strict=True) if not p]
|
| 63 |
+
assert len(set(observed)) == len(observed), 'Real temporal slots must not duplicate evidence'
|
| 64 |
+
return {'source_frame_indices': indices, 'padding_mask': pad, 'requested_us': requested, 'actual_pts_us': [float(pts_us[i]) for i in indices], 'model_clock_hz': profile.slot_fps, 'tick_us': TICK_US, 'max_real_error_us': float(max((abs(pts_us[i] - t) for i, t, p in zip(indices, requested, pad, strict=True) if not p)))}
|
| 65 |
+
|
| 66 |
+
def caption_plans(segments, media):
|
| 67 |
+
sources = [TemporalSegment(TimeSpan(s['start_s'], s['end_s']), dict(s.get('metadata', {}))) for s in segments]
|
| 68 |
+
sources = HierarchicalSubtaskSegmenter._with_refreshed_u2_context(tuple(sources))
|
| 69 |
+
plans, dispositions = ([], [])
|
| 70 |
+
for i, s in enumerate(sources):
|
| 71 |
+
eid = s.metadata.get('fine_event_id') or f'event_{i:06d}'
|
| 72 |
+
record = {'event_id': eid, 'start_s': s.span.start_s, 'end_s': s.span.end_s, 'event_index': i}
|
| 73 |
+
context = s.metadata['u2_context']
|
| 74 |
+
try:
|
| 75 |
+
blocks = _plan_blocks(context, media)
|
| 76 |
+
except U2WorkerError as error:
|
| 77 |
+
if 'TARGET span has fewer than' not in str(error):
|
| 78 |
+
raise
|
| 79 |
+
dispositions.append(dict(record, status='unsupported_target_frames', reason=str(error)))
|
| 80 |
+
continue
|
| 81 |
+
serialized = [asdict(b) for b in blocks]
|
| 82 |
+
for b in serialized:
|
| 83 |
+
b['frame_indices'] = list(b['frame_indices'])
|
| 84 |
+
plans.append(dict(record, blocks=serialized, source_fps_num=media['fps_num'], source_fps_den=media['fps_den'], metadata=dict(s.metadata)))
|
| 85 |
+
dispositions.append(dict(record, status='caption_input'))
|
| 86 |
+
assert len({x['event_id'] for x in dispositions}) == len(dispositions)
|
| 87 |
+
return (plans, dispositions)
|
| 88 |
+
|
| 89 |
+
def caption_sample(plan, frame_receipts, camera_roles):
|
| 90 |
+
blocks = []
|
| 91 |
+
clock = Fraction(plan['source_fps_num'], plan['source_fps_den'])
|
| 92 |
+
for b in plan['blocks']:
|
| 93 |
+
ids = b['frame_indices']
|
| 94 |
+
n = len(ids)
|
| 95 |
+
span = b['end_frame_exclusive'] - b['start_frame']
|
| 96 |
+
rate = Fraction(n) * clock / span
|
| 97 |
+
receipts = [frame_receipts[i] for i in ids]
|
| 98 |
+
blocks.append(dict(role=b['role'], paths=[r['path'] for r in receipts], frames=receipts, source_start_frame=b['start_frame'], source_end_frame_exclusive=b['end_frame_exclusive'], clock_fps_num=clock.numerator, clock_fps_den=clock.denominator, block_local_timestamp_indices=[i - b['start_frame'] for i in ids], effective_fps_num=rate.numerator, effective_fps_den=rate.denominator, original_descriptor={'frame_indices': ids}))
|
| 99 |
+
target = next((b for b in blocks if b['role'] == 'TARGET'))
|
| 100 |
+
text = 'REFERENCE_NOT_ADMITTED'
|
| 101 |
+
request = {'target': {'native_span': {'start_frame': target['source_start_frame'], 'end_frame_exclusive': target['source_end_frame_exclusive']}, 'text': text, 'sha256': hashlib.sha256(text.encode()).hexdigest()}}
|
| 102 |
+
return {'request': request, 'camera_count': len(camera_roles), 'camera_roles': camera_roles, 'blocks': blocks}
|
batch_utils/vendor/production_adapter/core/flora_runtime/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pass
|
batch_utils/vendor/production_adapter/core/flora_runtime/exact_frame_serving.py
ADDED
|
@@ -0,0 +1,47 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
import hashlib
|
| 5 |
+
import io
|
| 6 |
+
import math
|
| 7 |
+
from typing import Any
|
| 8 |
+
PIXEL_HASH_SCHEMA = 'rgb_u8_width_height_v1'
|
| 9 |
+
|
| 10 |
+
class ExactFrameServingError(ValueError):
|
| 11 |
+
pass
|
| 12 |
+
|
| 13 |
+
def pixel_sha256(image: Any) -> str:
|
| 14 |
+
rgb = image.convert('RGB')
|
| 15 |
+
header = f'{PIXEL_HASH_SCHEMA}\x00RGB\x00{rgb.width}\x00{rgb.height}\x00'.encode('ascii')
|
| 16 |
+
return hashlib.sha256(header + rgb.tobytes()).hexdigest()
|
| 17 |
+
|
| 18 |
+
def regularized_resize_dims(*, height: int, width: int, min_pixels: int=16 * 16, max_pixels: int=512 * 512) -> tuple[int, int]:
|
| 19 |
+
if min(height, width, min_pixels, max_pixels) <= 0 or min_pixels > max_pixels:
|
| 20 |
+
raise ExactFrameServingError('regularized resize arguments are invalid')
|
| 21 |
+
area = height * width
|
| 22 |
+
if area > max_pixels:
|
| 23 |
+
scale = math.sqrt(max_pixels / area)
|
| 24 |
+
return (int(height * scale), int(width * scale))
|
| 25 |
+
if area < min_pixels:
|
| 26 |
+
scale = math.sqrt(min_pixels / area)
|
| 27 |
+
return (int(height * scale), int(width * scale))
|
| 28 |
+
return (height, width)
|
| 29 |
+
|
| 30 |
+
def regularize_image(image: Any) -> Any:
|
| 31 |
+
from PIL import Image
|
| 32 |
+
source = image.convert('RGB')
|
| 33 |
+
height, width = regularized_resize_dims(height=source.height, width=source.width)
|
| 34 |
+
if source.size == (width, height):
|
| 35 |
+
return source
|
| 36 |
+
return source.resize((width, height), resample=Image.Resampling.BICUBIC)
|
| 37 |
+
|
| 38 |
+
def encode_regularized_image(image: Any, *, encoding: str) -> dict[str, Any]:
|
| 39 |
+
regularized = regularize_image(image)
|
| 40 |
+
stream = io.BytesIO()
|
| 41 |
+
if encoding == 'png_level6':
|
| 42 |
+
regularized.save(stream, format='PNG', optimize=False, compress_level=6)
|
| 43 |
+
elif encoding == 'jpeg_q98_444':
|
| 44 |
+
regularized.save(stream, format='JPEG', quality=98, subsampling=0, optimize=False)
|
| 45 |
+
else:
|
| 46 |
+
raise ExactFrameServingError(f'unsupported exact-frame cache encoding: {encoding}')
|
| 47 |
+
return {'payload': stream.getvalue(), 'width': regularized.width, 'height': regularized.height, 'rgb_pixel_sha256': pixel_sha256(regularized)}
|
batch_utils/vendor/production_adapter/core/flora_runtime/schema.py
ADDED
|
@@ -0,0 +1,88 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from collections.abc import Mapping
|
| 5 |
+
from dataclasses import dataclass, field
|
| 6 |
+
import math
|
| 7 |
+
import os
|
| 8 |
+
from typing import Any
|
| 9 |
+
|
| 10 |
+
def _require_nonempty(value: str, name: str) -> None:
|
| 11 |
+
if not value.strip():
|
| 12 |
+
raise ValueError(f'{name} must not be empty')
|
| 13 |
+
|
| 14 |
+
def _require_finite(value: float, name: str) -> None:
|
| 15 |
+
if not math.isfinite(value):
|
| 16 |
+
raise ValueError(f'{name} must be finite, got {value}')
|
| 17 |
+
|
| 18 |
+
@dataclass(frozen=True)
|
| 19 |
+
class VideoStream:
|
| 20 |
+
camera: str
|
| 21 |
+
location: str | os.PathLike[str]
|
| 22 |
+
fps: float | None = None
|
| 23 |
+
width: int | None = None
|
| 24 |
+
height: int | None = None
|
| 25 |
+
metadata: Mapping[str, Any] = field(default_factory=dict)
|
| 26 |
+
|
| 27 |
+
def __post_init__(self) -> None:
|
| 28 |
+
_require_nonempty(self.camera, 'camera')
|
| 29 |
+
location = os.fspath(self.location)
|
| 30 |
+
_require_nonempty(location, 'location')
|
| 31 |
+
object.__setattr__(self, 'location', location)
|
| 32 |
+
object.__setattr__(self, 'metadata', dict(self.metadata))
|
| 33 |
+
if self.fps is not None:
|
| 34 |
+
_require_finite(self.fps, 'fps')
|
| 35 |
+
if self.fps <= 0:
|
| 36 |
+
raise ValueError(f'fps must be positive, got {self.fps}')
|
| 37 |
+
for value, name in ((self.width, 'width'), (self.height, 'height')):
|
| 38 |
+
if value is not None and value <= 0:
|
| 39 |
+
raise ValueError(f'{name} must be positive, got {value}')
|
| 40 |
+
|
| 41 |
+
@dataclass(frozen=True)
|
| 42 |
+
class AnnotationSource:
|
| 43 |
+
source_id: str
|
| 44 |
+
location: str | os.PathLike[str]
|
| 45 |
+
video_streams: tuple[VideoStream, ...] = ()
|
| 46 |
+
duration_s: float | None = None
|
| 47 |
+
fps: float | None = None
|
| 48 |
+
metadata: Mapping[str, Any] = field(default_factory=dict)
|
| 49 |
+
|
| 50 |
+
def __post_init__(self) -> None:
|
| 51 |
+
_require_nonempty(self.source_id, 'source_id')
|
| 52 |
+
location = os.fspath(self.location)
|
| 53 |
+
_require_nonempty(location, 'location')
|
| 54 |
+
object.__setattr__(self, 'location', location)
|
| 55 |
+
object.__setattr__(self, 'video_streams', tuple(self.video_streams))
|
| 56 |
+
object.__setattr__(self, 'metadata', dict(self.metadata))
|
| 57 |
+
cameras = [stream.camera for stream in self.video_streams]
|
| 58 |
+
if len(cameras) != len(set(cameras)):
|
| 59 |
+
raise ValueError(f'video stream camera names must be unique, got {cameras}')
|
| 60 |
+
if self.duration_s is not None:
|
| 61 |
+
_require_finite(self.duration_s, 'duration_s')
|
| 62 |
+
if self.duration_s <= 0:
|
| 63 |
+
raise ValueError(f'duration_s must be positive, got {self.duration_s}')
|
| 64 |
+
if self.fps is not None:
|
| 65 |
+
_require_finite(self.fps, 'fps')
|
| 66 |
+
if self.fps <= 0:
|
| 67 |
+
raise ValueError(f'fps must be positive, got {self.fps}')
|
| 68 |
+
|
| 69 |
+
@dataclass(frozen=True, order=True)
|
| 70 |
+
class TimeSpan:
|
| 71 |
+
start_s: float
|
| 72 |
+
end_s: float
|
| 73 |
+
|
| 74 |
+
def __post_init__(self) -> None:
|
| 75 |
+
_require_finite(self.start_s, 'start_s')
|
| 76 |
+
_require_finite(self.end_s, 'end_s')
|
| 77 |
+
if self.start_s < 0:
|
| 78 |
+
raise ValueError(f'start_s must be non-negative, got {self.start_s}')
|
| 79 |
+
if self.end_s <= self.start_s:
|
| 80 |
+
raise ValueError(f'end_s must be greater than start_s, got [{self.start_s}, {self.end_s})')
|
| 81 |
+
|
| 82 |
+
@dataclass(frozen=True)
|
| 83 |
+
class TemporalSegment:
|
| 84 |
+
span: TimeSpan
|
| 85 |
+
metadata: Mapping[str, Any] = field(default_factory=dict)
|
| 86 |
+
|
| 87 |
+
def __post_init__(self) -> None:
|
| 88 |
+
object.__setattr__(self, 'metadata', dict(self.metadata))
|
batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pass
|
batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/u2_contract.py
ADDED
|
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from fractions import Fraction
|
| 5 |
+
TARGET_SAMPLE_FPS = 8
|
| 6 |
+
TARGET_MIN_FRAMES = 8
|
| 7 |
+
TARGET_MAX_FRAMES = 40
|
| 8 |
+
TARGET_FRAME_FACTOR = 2
|
| 9 |
+
CONTEXT_SAMPLE_FPS = 8
|
| 10 |
+
CONTEXT_MIN_FRAMES = 4
|
| 11 |
+
CONTEXT_MAX_FRAMES = 12
|
| 12 |
+
CONTEXT_FRAME_FACTOR = 4
|
| 13 |
+
MAX_TOTAL_FRAMES = 64
|
| 14 |
+
SAMPLER_VERSION = 'duration_adaptive_uniform_bin_centers_nested_context_v2'
|
| 15 |
+
|
| 16 |
+
class U2ContractError(ValueError):
|
| 17 |
+
pass
|
| 18 |
+
|
| 19 |
+
def uniform_bin_centers(start: int, end_exclusive: int, count: int) -> list[int]:
|
| 20 |
+
length = end_exclusive - start
|
| 21 |
+
if count <= 0:
|
| 22 |
+
return []
|
| 23 |
+
if length < count:
|
| 24 |
+
raise U2ContractError(f'span [{start}, {end_exclusive}) has fewer than {count} frames')
|
| 25 |
+
result = [start + (2 * index + 1) * length // (2 * count) for index in range(count)]
|
| 26 |
+
if len(result) != len(set(result)) or result[0] < start or result[-1] >= end_exclusive:
|
| 27 |
+
raise AssertionError('uniform frame sampler violated its contract')
|
| 28 |
+
return result
|
| 29 |
+
|
| 30 |
+
def duration_adaptive_sample(start: int, end_exclusive: int, *, source_fps_num: int, source_fps_den: int, sample_fps: int, min_frames: int, max_frames: int, frame_factor: int) -> tuple[list[int], dict[str, int | bool | str]]:
|
| 31 |
+
available = end_exclusive - start
|
| 32 |
+
if source_fps_num <= 0 or source_fps_den <= 0 or available <= 0:
|
| 33 |
+
raise U2ContractError('source FPS and frame span must be positive')
|
| 34 |
+
if min_frames % frame_factor or max_frames % frame_factor:
|
| 35 |
+
raise U2ContractError('adaptive frame bounds must align to frame factor')
|
| 36 |
+
duration = Fraction(available * source_fps_den, source_fps_num)
|
| 37 |
+
requested = min(Fraction(max_frames), max(Fraction(min_frames), duration * sample_fps))
|
| 38 |
+
desired_count = requested.numerator // (requested.denominator * frame_factor) * frame_factor
|
| 39 |
+
actual_count = min(desired_count, available // frame_factor * frame_factor)
|
| 40 |
+
if actual_count < frame_factor:
|
| 41 |
+
raise U2ContractError(f'span [{start}, {end_exclusive}) has fewer than {frame_factor} frames')
|
| 42 |
+
indices = uniform_bin_centers(start, end_exclusive, actual_count)
|
| 43 |
+
return (indices, {'policy': 'duration_fps_bounded_unique_frames_v1', 'sampler_version': SAMPLER_VERSION, 'source_fps_num': source_fps_num, 'source_fps_den': source_fps_den, 'sample_fps_num': sample_fps, 'sample_fps_den': 1, 'min_frames': min_frames, 'max_frames': max_frames, 'frame_factor': frame_factor, 'available_unique_source_frames': available, 'desired_frame_count': desired_count, 'actual_frame_count': actual_count, 'source_limited': actual_count < desired_count, 'selection': 'uniform_bin_centers'})
|
| 44 |
+
|
| 45 |
+
def _ceil_fraction(value: Fraction) -> int:
|
| 46 |
+
return -(-value.numerator // value.denominator)
|
| 47 |
+
|
| 48 |
+
def seconds_to_frame_span(start_sec: float | str, end_sec: float | str, *, source_fps_num: int, source_fps_den: int, source_frame_count: int) -> tuple[int, int]:
|
| 49 |
+
start = Fraction(str(start_sec))
|
| 50 |
+
end = Fraction(str(end_sec))
|
| 51 |
+
if start < 0 or end <= start:
|
| 52 |
+
raise U2ContractError(f'invalid seconds span [{start_sec}, {end_sec})')
|
| 53 |
+
fps = Fraction(source_fps_num, source_fps_den)
|
| 54 |
+
start_frame = max(0, min(source_frame_count, _ceil_fraction(start * fps)))
|
| 55 |
+
end_frame = max(0, min(source_frame_count, _ceil_fraction(end * fps)))
|
| 56 |
+
if end_frame <= start_frame:
|
| 57 |
+
raise U2ContractError(f'seconds span [{start_sec}, {end_sec}) maps to no source frames')
|
| 58 |
+
return (start_frame, end_frame)
|
batch_utils/vendor/production_adapter/core/flora_runtime/subtask_captioning/u2_worker.py
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from dataclasses import dataclass
|
| 5 |
+
from typing import Any
|
| 6 |
+
from flora_runtime.subtask_captioning.u2_contract import CONTEXT_FRAME_FACTOR, CONTEXT_MAX_FRAMES, CONTEXT_MIN_FRAMES, CONTEXT_SAMPLE_FPS, MAX_TOTAL_FRAMES, TARGET_FRAME_FACTOR, TARGET_MAX_FRAMES, TARGET_MIN_FRAMES, TARGET_SAMPLE_FPS, duration_adaptive_sample, seconds_to_frame_span
|
| 7 |
+
_ROLES = ('BEFORE', 'TARGET', 'AFTER')
|
| 8 |
+
|
| 9 |
+
class U2WorkerError(RuntimeError):
|
| 10 |
+
pass
|
| 11 |
+
|
| 12 |
+
@dataclass(frozen=True)
|
| 13 |
+
class _BlockPlan:
|
| 14 |
+
role: str
|
| 15 |
+
start_frame: int
|
| 16 |
+
end_frame_exclusive: int
|
| 17 |
+
frame_indices: tuple[int, ...]
|
| 18 |
+
sampling_receipt: dict[str, Any]
|
| 19 |
+
|
| 20 |
+
def _plan_blocks(u2_context: dict[str, Any], media: dict[str, Any]) -> tuple[_BlockPlan, ...]:
|
| 21 |
+
plans: list[_BlockPlan] = []
|
| 22 |
+
for role, key in zip(_ROLES, ('before', 'target', 'after'), strict=True):
|
| 23 |
+
span = u2_context.get(key)
|
| 24 |
+
if span is None:
|
| 25 |
+
if role == 'TARGET':
|
| 26 |
+
raise U2WorkerError('U2 TARGET context is absent')
|
| 27 |
+
continue
|
| 28 |
+
if not isinstance(span, dict):
|
| 29 |
+
raise U2WorkerError(f'U2 {key} span must be an object')
|
| 30 |
+
start, end = seconds_to_frame_span(span.get('start_s'), span.get('end_s'), source_fps_num=media['fps_num'], source_fps_den=media['fps_den'], source_frame_count=media['frame_count'])
|
| 31 |
+
target = role == 'TARGET'
|
| 32 |
+
frame_factor = TARGET_FRAME_FACTOR if target else CONTEXT_FRAME_FACTOR
|
| 33 |
+
if end - start < frame_factor:
|
| 34 |
+
if not target:
|
| 35 |
+
continue
|
| 36 |
+
raise U2WorkerError(f'U2 {role} span has fewer than the required {frame_factor}-frame group')
|
| 37 |
+
indices, receipt = duration_adaptive_sample(start, end, source_fps_num=media['fps_num'], source_fps_den=media['fps_den'], sample_fps=TARGET_SAMPLE_FPS if target else CONTEXT_SAMPLE_FPS, min_frames=TARGET_MIN_FRAMES if target else CONTEXT_MIN_FRAMES, max_frames=TARGET_MAX_FRAMES if target else CONTEXT_MAX_FRAMES, frame_factor=frame_factor)
|
| 38 |
+
plans.append(_BlockPlan(role=role, start_frame=start, end_frame_exclusive=end, frame_indices=tuple(indices), sampling_receipt=dict(receipt)))
|
| 39 |
+
total = sum((len(plan.frame_indices) for plan in plans))
|
| 40 |
+
if total > MAX_TOTAL_FRAMES:
|
| 41 |
+
raise U2WorkerError(f'U2 frame budget {total} exceeds {MAX_TOTAL_FRAMES}')
|
| 42 |
+
return tuple(plans)
|
| 43 |
+
if __name__ == '__main__':
|
| 44 |
+
raise SystemExit(main())
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
pass
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/contracts.py
ADDED
|
@@ -0,0 +1,73 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from collections.abc import Mapping
|
| 5 |
+
from dataclasses import dataclass, field
|
| 6 |
+
import math
|
| 7 |
+
from typing import Any, Literal
|
| 8 |
+
from flora_runtime.schema import TimeSpan
|
| 9 |
+
EventCompleteness = Literal['complete', 'left_censored', 'right_censored', 'both_censored', 'uncertain']
|
| 10 |
+
|
| 11 |
+
def _require_nonempty(value: str, label: str) -> None:
|
| 12 |
+
if not isinstance(value, str) or not value.strip():
|
| 13 |
+
raise ValueError(f'{label} must be non-empty')
|
| 14 |
+
|
| 15 |
+
def _validate_confidence(value: float | None, label: str) -> None:
|
| 16 |
+
if value is None:
|
| 17 |
+
return
|
| 18 |
+
if isinstance(value, bool) or not isinstance(value, int | float):
|
| 19 |
+
raise TypeError(f'{label} must be a number in [0, 1]')
|
| 20 |
+
if not math.isfinite(float(value)) or not 0.0 <= float(value) <= 1.0:
|
| 21 |
+
raise ValueError(f'{label} must be finite and in [0, 1]')
|
| 22 |
+
|
| 23 |
+
def span_contains(outer: TimeSpan, inner: TimeSpan, *, tolerance_s: float=1e-09) -> bool:
|
| 24 |
+
if tolerance_s < 0 or not math.isfinite(tolerance_s):
|
| 25 |
+
raise ValueError('tolerance_s must be finite and non-negative')
|
| 26 |
+
return outer.start_s <= inner.start_s + tolerance_s and outer.end_s + tolerance_s >= inner.end_s
|
| 27 |
+
|
| 28 |
+
@dataclass(frozen=True)
|
| 29 |
+
class ProcessingScope:
|
| 30 |
+
scope_id: str
|
| 31 |
+
source_id: str
|
| 32 |
+
core_span: TimeSpan
|
| 33 |
+
context_span: TimeSpan
|
| 34 |
+
proposal_confidence: float | None = None
|
| 35 |
+
metadata: Mapping[str, Any] = field(default_factory=dict)
|
| 36 |
+
|
| 37 |
+
def __post_init__(self) -> None:
|
| 38 |
+
_require_nonempty(self.scope_id, 'scope_id')
|
| 39 |
+
_require_nonempty(self.source_id, 'source_id')
|
| 40 |
+
if not span_contains(self.context_span, self.core_span):
|
| 41 |
+
raise ValueError('context_span must contain core_span')
|
| 42 |
+
_validate_confidence(self.proposal_confidence, 'proposal_confidence')
|
| 43 |
+
object.__setattr__(self, 'metadata', dict(self.metadata))
|
| 44 |
+
|
| 45 |
+
@dataclass(frozen=True)
|
| 46 |
+
class EventInterval:
|
| 47 |
+
event_id: str
|
| 48 |
+
source_id: str
|
| 49 |
+
span: TimeSpan
|
| 50 |
+
completeness: EventCompleteness = 'complete'
|
| 51 |
+
start_confidence: float | None = None
|
| 52 |
+
end_confidence: float | None = None
|
| 53 |
+
claim: str | None = None
|
| 54 |
+
source_scope_ids: tuple[str, ...] = ()
|
| 55 |
+
metadata: Mapping[str, Any] = field(default_factory=dict)
|
| 56 |
+
|
| 57 |
+
def __post_init__(self) -> None:
|
| 58 |
+
_require_nonempty(self.event_id, 'event_id')
|
| 59 |
+
_require_nonempty(self.source_id, 'source_id')
|
| 60 |
+
if self.completeness not in {'complete', 'left_censored', 'right_censored', 'both_censored', 'uncertain'}:
|
| 61 |
+
raise ValueError(f'unsupported completeness: {self.completeness!r}')
|
| 62 |
+
_validate_confidence(self.start_confidence, 'start_confidence')
|
| 63 |
+
_validate_confidence(self.end_confidence, 'end_confidence')
|
| 64 |
+
if self.claim is not None:
|
| 65 |
+
_require_nonempty(self.claim, 'claim')
|
| 66 |
+
object.__setattr__(self, 'claim', self.claim.strip())
|
| 67 |
+
scope_ids = tuple(self.source_scope_ids)
|
| 68 |
+
if len(scope_ids) != len(set(scope_ids)):
|
| 69 |
+
raise ValueError('source_scope_ids must be unique')
|
| 70 |
+
for scope_id in scope_ids:
|
| 71 |
+
_require_nonempty(scope_id, 'source_scope_id')
|
| 72 |
+
object.__setattr__(self, 'source_scope_ids', scope_ids)
|
| 73 |
+
object.__setattr__(self, 'metadata', dict(self.metadata))
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_subtask.py
ADDED
|
@@ -0,0 +1,852 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
import asyncio
|
| 5 |
+
from collections import Counter
|
| 6 |
+
from collections.abc import Mapping, Sequence
|
| 7 |
+
from dataclasses import asdict, dataclass, field, replace
|
| 8 |
+
import hashlib
|
| 9 |
+
import json
|
| 10 |
+
import math
|
| 11 |
+
from typing import Any
|
| 12 |
+
from flora_runtime.schema import AnnotationSource, TemporalSegment, TimeSpan
|
| 13 |
+
from flora_runtime.video_segmentation.contracts import EventInterval, ProcessingScope
|
| 14 |
+
from flora_runtime.video_segmentation.fine_temporal import DEFAULT_FINE_WINDOW_PROFILE_ID, MAX_BOUNDED_ANCHOR_TICKS, STITCH_TOLERANCE_TICKS, TICK_US, FineFragment, FineSegmentationResult, FineTemporalError, FineWindow, StitchingPolicy, UnresolvedFineEvent, analyze_stitching_seams, ceil_div, fine_window_profile, parse_prediction_with_diagnostics, stitch_predictions, window_from_scope
|
| 15 |
+
from flora_runtime.video_segmentation.interface import SegmentTopology, VideoSegmenter
|
| 16 |
+
from flora_runtime.video_segmentation.scope_scheduler import FINE_DIRECTIONAL_RETRY_SCOPE_NAMESPACE, schedule_fine_scopes, schedule_recentered_fine_retry_scopes, seconds_to_us
|
| 17 |
+
from flora_runtime.worker_transport import JsonPredictor, SubtaskAnnotationAdapterError, source_payload
|
| 18 |
+
FINE_REQUEST_SCHEMA = 'subtask_annotation_atomic_fine_batch_request_v1'
|
| 19 |
+
FINE_RESPONSE_SCHEMA = 'subtask_annotation_atomic_fine_batch_response_v1'
|
| 20 |
+
FINE_QUALITY_POLICY_SCHEMA = 'subtask_annotation_atomic_fine_quality_policy_v1'
|
| 21 |
+
PROVIDER_FREE_FINE_COMPLETION_POLICIES = {'adaptive_fine_then_conservative_completion'}
|
| 22 |
+
ADAPTIVE_FINE_REPAIR_SCHEMA = 'subtask_annotation_adaptive_fine_repair_v1'
|
| 23 |
+
ADAPTIVE_FINE_RETRY_MAX_WINDOWS = 4
|
| 24 |
+
ADAPTIVE_FINE_RETRY_HARD_MAX_WINDOWS = 64
|
| 25 |
+
CONSERVATIVE_FINE_COMPLETION_SCHEMA = 'subtask_annotation_conservative_fine_completion_v2'
|
| 26 |
+
UNRESOLVED_LOCAL_STITCH_SCHEMA = 'subtask_annotation_unresolved_local_stitch_v1'
|
| 27 |
+
SEGMENT_LOCAL_REPAIR_ENDPOINT_TOLERANCE_S = 0.5
|
| 28 |
+
|
| 29 |
+
class FineSegmentationOutputError(SubtaskAnnotationAdapterError):
|
| 30 |
+
|
| 31 |
+
def __init__(self, message: str, *, diagnostics: Mapping[str, object], partial_result: FineSegmentationResult | None=None) -> None:
|
| 32 |
+
super().__init__(message)
|
| 33 |
+
self.diagnostics = dict(diagnostics)
|
| 34 |
+
self.partial_result = partial_result
|
| 35 |
+
|
| 36 |
+
@dataclass(frozen=True)
|
| 37 |
+
class FineSegmentationQualityPolicy:
|
| 38 |
+
soft_max_segments_per_window: int = 32
|
| 39 |
+
hard_max_segments_per_window: int = 32
|
| 40 |
+
hard_max_overlap_depth: int = 32
|
| 41 |
+
max_adjacent_matching_ambiguities: int = 0
|
| 42 |
+
drop_guard_only_segments: bool = False
|
| 43 |
+
require_event_when_coarse_present: bool = False
|
| 44 |
+
|
| 45 |
+
def __post_init__(self) -> None:
|
| 46 |
+
integer_fields = ('soft_max_segments_per_window', 'hard_max_segments_per_window', 'hard_max_overlap_depth', 'max_adjacent_matching_ambiguities')
|
| 47 |
+
for name in integer_fields:
|
| 48 |
+
value = getattr(self, name)
|
| 49 |
+
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
| 50 |
+
raise ValueError(f'{name} must be a non-negative integer')
|
| 51 |
+
if self.soft_max_segments_per_window > self.hard_max_segments_per_window:
|
| 52 |
+
raise ValueError('soft_max_segments_per_window must not exceed the hard maximum')
|
| 53 |
+
if type(self.drop_guard_only_segments) is not bool:
|
| 54 |
+
raise TypeError('drop_guard_only_segments must be a boolean')
|
| 55 |
+
if type(self.require_event_when_coarse_present) is not bool:
|
| 56 |
+
raise TypeError('require_event_when_coarse_present must be a boolean')
|
| 57 |
+
|
| 58 |
+
@classmethod
|
| 59 |
+
def from_value(cls, value: object) -> FineSegmentationQualityPolicy:
|
| 60 |
+
if not isinstance(value, dict):
|
| 61 |
+
raise TypeError('fine quality policy must be an object')
|
| 62 |
+
expected = {'schema_version', 'soft_max_segments_per_window', 'hard_max_segments_per_window', 'hard_max_overlap_depth', 'max_adjacent_matching_ambiguities', 'drop_guard_only_segments', 'require_event_when_coarse_present'}
|
| 63 |
+
missing, unknown = (expected - set(value), set(value) - expected)
|
| 64 |
+
if missing or unknown:
|
| 65 |
+
raise ValueError(f'fine quality policy key mismatch; missing={sorted(missing)}, unknown={sorted(unknown)}')
|
| 66 |
+
if value['schema_version'] != FINE_QUALITY_POLICY_SCHEMA:
|
| 67 |
+
raise ValueError(f'fine quality policy schema must be {FINE_QUALITY_POLICY_SCHEMA!r}')
|
| 68 |
+
return cls(**{key: value[key] for key in expected - {'schema_version'}})
|
| 69 |
+
|
| 70 |
+
def to_dict(self) -> dict[str, object]:
|
| 71 |
+
return {'schema_version': FINE_QUALITY_POLICY_SCHEMA, 'soft_max_segments_per_window': self.soft_max_segments_per_window, 'hard_max_segments_per_window': self.hard_max_segments_per_window, 'hard_max_overlap_depth': self.hard_max_overlap_depth, 'max_adjacent_matching_ambiguities': self.max_adjacent_matching_ambiguities, 'drop_guard_only_segments': self.drop_guard_only_segments, 'require_event_when_coarse_present': self.require_event_when_coarse_present}
|
| 72 |
+
|
| 73 |
+
def _span_payload(span: TimeSpan | None) -> dict[str, float] | None:
|
| 74 |
+
if span is None:
|
| 75 |
+
return None
|
| 76 |
+
return {'start_s': span.start_s, 'end_s': span.end_s}
|
| 77 |
+
|
| 78 |
+
def _nearest_nonoverlapping_context(ordered: Sequence[EventInterval | TemporalSegment], index: int) -> tuple[TimeSpan | None, TimeSpan | None]:
|
| 79 |
+
target = ordered[index].span
|
| 80 |
+
before_candidates = [item.span for item in ordered if item.span.end_s <= target.start_s]
|
| 81 |
+
after_candidates = [item.span for item in ordered if item.span.start_s >= target.end_s]
|
| 82 |
+
before = max(before_candidates, key=lambda span: (span.end_s, span.start_s), default=None)
|
| 83 |
+
after = min(after_candidates, key=lambda span: (span.start_s, span.end_s), default=None)
|
| 84 |
+
return (before, after)
|
| 85 |
+
|
| 86 |
+
def events_to_temporal_segments(events: Sequence[EventInterval], *, diagnostics: Mapping[str, object] | None=None, resolution: str='fine') -> tuple[TemporalSegment, ...]:
|
| 87 |
+
ordered = tuple(sorted(events, key=lambda item: (item.span, item.event_id)))
|
| 88 |
+
result: list[TemporalSegment] = []
|
| 89 |
+
for index, event in enumerate(ordered):
|
| 90 |
+
before, after = _nearest_nonoverlapping_context(ordered, index)
|
| 91 |
+
metadata = dict(event.metadata)
|
| 92 |
+
metadata.setdefault('segmentation_provenance', 'fine')
|
| 93 |
+
metadata.setdefault('segmentation_resolution', resolution)
|
| 94 |
+
event_identity = {'fine_event_id': event.event_id} if metadata['segmentation_resolution'] == 'fine' else {'segment_event_id': event.event_id}
|
| 95 |
+
metadata.update({**event_identity, 'source_scope_ids': list(event.source_scope_ids), 'segment_topology': SegmentTopology.CONCURRENT.value, 'u2_context': {'before': _span_payload(before), 'target': _span_payload(event.span), 'after': _span_payload(after)}})
|
| 96 |
+
if diagnostics is not None:
|
| 97 |
+
metadata['fine_segmentation_diagnostics'] = dict(diagnostics)
|
| 98 |
+
result.append(TemporalSegment(span=event.span, metadata=metadata))
|
| 99 |
+
return tuple(result)
|
| 100 |
+
|
| 101 |
+
def _overlap_statistics(events: Sequence[EventInterval]) -> tuple[int, int]:
|
| 102 |
+
overlap_pairs = 0
|
| 103 |
+
for left_index, left in enumerate(events):
|
| 104 |
+
overlap_pairs += sum((min(left.span.end_s, right.span.end_s) > max(left.span.start_s, right.span.start_s) for right in events[left_index + 1:]))
|
| 105 |
+
points = [(item.span.start_s, 1) for item in events] + [(item.span.end_s, -1) for item in events]
|
| 106 |
+
depth = current = 0
|
| 107 |
+
for _, delta in sorted(points, key=lambda item: (item[0], item[1])):
|
| 108 |
+
current += delta
|
| 109 |
+
depth = max(depth, current)
|
| 110 |
+
return (overlap_pairs, depth)
|
| 111 |
+
|
| 112 |
+
def _exact_duplicate_intervals(events: Sequence[EventInterval]) -> list[dict[str, object]]:
|
| 113 |
+
counts = Counter(((item.span.start_s, item.span.end_s) for item in events))
|
| 114 |
+
return [{'start_s': start_s, 'end_s': end_s, 'count': count} for (start_s, end_s), count in sorted(counts.items()) if count > 1]
|
| 115 |
+
|
| 116 |
+
def _diagnostic_sha256(value: Mapping[str, object]) -> str:
|
| 117 |
+
payload = (json.dumps(value, sort_keys=True, ensure_ascii=False, allow_nan=False, separators=(',', ':')) + '\n').encode()
|
| 118 |
+
return hashlib.sha256(payload).hexdigest()
|
| 119 |
+
|
| 120 |
+
def _inference_context(source: AnnotationSource, coarse: Sequence[TemporalSegment], *, fine_prior_mode: str) -> dict[str, object]:
|
| 121 |
+
timestamp_prior = {'schema_version': 'subtask_annotation_coarse_timestamp_prior_v1', 'source_id': source.source_id, 'spans': [{'start_s': segment.span.start_s, 'end_s': segment.span.end_s} for segment in coarse]}
|
| 122 |
+
return {'fine_prior_mode': fine_prior_mode, 'coarse_prior_count': len(coarse), 'coarse_prior_sha256': _diagnostic_sha256(timestamp_prior)}
|
| 123 |
+
|
| 124 |
+
def _with_inference_context(diagnostics: Mapping[str, object], context: Mapping[str, object]) -> dict[str, object]:
|
| 125 |
+
result = dict(diagnostics)
|
| 126 |
+
result.update(context)
|
| 127 |
+
result['diagnostics_sha256'] = _diagnostic_sha256({key: value for key, value in result.items() if key != 'diagnostics_sha256'})
|
| 128 |
+
return result
|
| 129 |
+
|
| 130 |
+
@dataclass(frozen=True)
|
| 131 |
+
class _FineWindowBatchPrediction:
|
| 132 |
+
fragments_by_window: Mapping[int, tuple[FineFragment, ...]]
|
| 133 |
+
raw_segment_counts: tuple[int, ...]
|
| 134 |
+
emitted_segment_counts: tuple[int, ...]
|
| 135 |
+
guard_only_fragments: tuple[tuple[FineFragment, ...], ...]
|
| 136 |
+
|
| 137 |
+
@dataclass(frozen=True)
|
| 138 |
+
class _FineRetryPrediction:
|
| 139 |
+
candidates_by_window: Mapping[int, tuple[TemporalSegment, ...]]
|
| 140 |
+
diagnostics: Mapping[str, object]
|
| 141 |
+
|
| 142 |
+
@dataclass(frozen=True)
|
| 143 |
+
class FineSubtaskSegmenter:
|
| 144 |
+
predictor: JsonPredictor
|
| 145 |
+
model_binding: Mapping[str, Any]
|
| 146 |
+
profile_id: str = 'human_ego_atomic_v0_2'
|
| 147 |
+
window_profile_id: str = DEFAULT_FINE_WINDOW_PROFILE_ID
|
| 148 |
+
fail_on_unresolved: bool = True
|
| 149 |
+
batch_concurrency: int = 1
|
| 150 |
+
unresolved_stitch_tolerance_ticks: int = STITCH_TOLERANCE_TICKS
|
| 151 |
+
quality_policy: FineSegmentationQualityPolicy = field(default_factory=FineSegmentationQualityPolicy)
|
| 152 |
+
segment_topology = SegmentTopology.CONCURRENT
|
| 153 |
+
|
| 154 |
+
def __post_init__(self) -> None:
|
| 155 |
+
binding = dict(self.model_binding)
|
| 156 |
+
if not binding:
|
| 157 |
+
raise ValueError('model_binding must not be empty')
|
| 158 |
+
window_profile = fine_window_profile(self.window_profile_id)
|
| 159 |
+
declared_model_io_contract = binding.get('model_io_contract')
|
| 160 |
+
if declared_model_io_contract is not None and declared_model_io_contract != window_profile.model_io_contract:
|
| 161 |
+
raise ValueError('model_binding.model_io_contract must match the selected fine window profile')
|
| 162 |
+
if isinstance(self.batch_concurrency, bool) or not isinstance(self.batch_concurrency, int) or self.batch_concurrency <= 0:
|
| 163 |
+
raise ValueError('batch_concurrency must be a positive integer')
|
| 164 |
+
if isinstance(self.unresolved_stitch_tolerance_ticks, bool) or not isinstance(self.unresolved_stitch_tolerance_ticks, int) or (not STITCH_TOLERANCE_TICKS <= self.unresolved_stitch_tolerance_ticks <= MAX_BOUNDED_ANCHOR_TICKS):
|
| 165 |
+
raise ValueError(f'unresolved_stitch_tolerance_ticks must be an integer in [{STITCH_TOLERANCE_TICKS}, {MAX_BOUNDED_ANCHOR_TICKS}]')
|
| 166 |
+
object.__setattr__(self, 'model_binding', binding)
|
| 167 |
+
|
| 168 |
+
def _stitch_with_unresolved_local_candidates(self, windows: Sequence[FineWindow], fragments_by_window: Mapping[int, Sequence[FineFragment]]) -> FineSegmentationResult:
|
| 169 |
+
strict = stitch_predictions(windows, fragments_by_window)
|
| 170 |
+
tolerance_ticks = self.unresolved_stitch_tolerance_ticks
|
| 171 |
+
if not strict.unresolved or tolerance_ticks == STITCH_TOLERANCE_TICKS:
|
| 172 |
+
return strict
|
| 173 |
+
policy = StitchingPolicy.component_local_bounded_anchor(endpoint_tolerance_ticks=tolerance_ticks)
|
| 174 |
+
seam_diagnostics = analyze_stitching_seams(windows, fragments_by_window, policy=policy)
|
| 175 |
+
candidate_rows: list[dict[str, object]] = []
|
| 176 |
+
allowlist: set[tuple[int, int, int]] = set()
|
| 177 |
+
seams = seam_diagnostics.get('seams')
|
| 178 |
+
if not isinstance(seams, list):
|
| 179 |
+
raise FineTemporalError('component-local seam diagnostics are missing seam rows')
|
| 180 |
+
for seam_index, seam in enumerate(seams):
|
| 181 |
+
if not isinstance(seam, dict):
|
| 182 |
+
raise FineTemporalError('component-local seam diagnostic row must be an object')
|
| 183 |
+
local = seam.get('component_local_identity')
|
| 184 |
+
if not isinstance(local, dict):
|
| 185 |
+
raise FineTemporalError('component-local seam evidence is missing')
|
| 186 |
+
accepted_pairs = local.get('accepted_pairs')
|
| 187 |
+
if not isinstance(accepted_pairs, list):
|
| 188 |
+
raise FineTemporalError('component-local accepted-pair evidence is invalid')
|
| 189 |
+
for pair in accepted_pairs:
|
| 190 |
+
if not isinstance(pair, dict):
|
| 191 |
+
raise FineTemporalError('component-local accepted pair must be an object')
|
| 192 |
+
left_id = pair.get('left_observation_id')
|
| 193 |
+
right_id = pair.get('right_observation_id')
|
| 194 |
+
if isinstance(left_id, bool) or not isinstance(left_id, int) or isinstance(right_id, bool) or (not isinstance(right_id, int)):
|
| 195 |
+
raise FineTemporalError('component-local accepted pair identity is invalid')
|
| 196 |
+
allowlist.add((seam_index, left_id, right_id))
|
| 197 |
+
candidate_rows.append({'seam_index': seam_index, **pair})
|
| 198 |
+
if not allowlist:
|
| 199 |
+
return strict
|
| 200 |
+
relaxed = stitch_predictions(windows, fragments_by_window, policy=policy, component_pair_allowlist=allowlist)
|
| 201 |
+
if len(relaxed.unresolved) >= len(strict.unresolved):
|
| 202 |
+
raise FineTemporalError('component-local unresolved stitching did not reduce unresolved components')
|
| 203 |
+
unresolved_by_observation = {observation_id: component_index for component_index, component in enumerate(strict.unresolved) for observation_id in component.observation_ids}
|
| 204 |
+
resolved_by_observation = {observation_id: event for event in strict.events for observation_id in event.metadata.get('observation_ids', ()) if isinstance(observation_id, int)}
|
| 205 |
+
candidates: list[dict[str, object]] = []
|
| 206 |
+
seen_components: set[int] = set()
|
| 207 |
+
for row in candidate_rows:
|
| 208 |
+
partial_root = row.get('partial_component_root')
|
| 209 |
+
resolved_root = row.get('resolved_component_root')
|
| 210 |
+
if isinstance(partial_root, bool) or not isinstance(partial_root, int) or isinstance(resolved_root, bool) or (not isinstance(resolved_root, int)):
|
| 211 |
+
raise FineTemporalError('component-local candidate roots are invalid')
|
| 212 |
+
component_index = unresolved_by_observation.get(partial_root)
|
| 213 |
+
resolved_event = resolved_by_observation.get(resolved_root)
|
| 214 |
+
if component_index is None or resolved_event is None:
|
| 215 |
+
raise FineTemporalError('component-local candidate does not map to one strict unresolved and one strict resolved component')
|
| 216 |
+
if component_index in seen_components:
|
| 217 |
+
raise FineTemporalError('component-local candidate discovery produced multiple targets for one unresolved component')
|
| 218 |
+
seen_components.add(component_index)
|
| 219 |
+
candidates.append({'partial_component_index': component_index, 'resolved_event_id': resolved_event.event_id, 'resolved_event_span': _span_payload(resolved_event.span), 'pair': row})
|
| 220 |
+
local_stitch_diagnostics = {'schema_version': UNRESOLVED_LOCAL_STITCH_SCHEMA, 'mode': 'candidate_only_requires_adaptive_fine_confirmation_v1', 'strict_tolerance_ticks': STITCH_TOLERANCE_TICKS, 'candidate_tolerance_ticks': tolerance_ticks, 'tick_us': TICK_US, 'strict_unresolved_component_count': len(strict.unresolved), 'candidate_component_count': len(candidates), 'potential_remaining_unresolved_component_count': len(relaxed.unresolved), 'potential_absorbed_unresolved_component_count': len(strict.unresolved) - len(relaxed.unresolved), 'accepted_pair_count': len(allowlist), 'resolved_geometry_changed': False, 'semantic_quality_established': False, 'clean_acceptance_requires_adaptive_fine_confirmation': True, 'candidates': candidates}
|
| 221 |
+
return replace(strict, diagnostics={**dict(strict.diagnostics), 'unresolved_local_stitch': local_stitch_diagnostics})
|
| 222 |
+
|
| 223 |
+
def _assess_result(self, result: FineSegmentationResult, *, inference_arm: str, raw_segment_counts: Sequence[int], emitted_segment_counts: Sequence[int], guard_only_fragments: Sequence[Sequence[object]]) -> FineSegmentationResult:
|
| 224 |
+
overlap_pairs, overlap_depth = _overlap_statistics(result.events)
|
| 225 |
+
exact_duplicates = _exact_duplicate_intervals(result.events)
|
| 226 |
+
policy = self.quality_policy
|
| 227 |
+
window_profile = fine_window_profile(self.window_profile_id)
|
| 228 |
+
violations: list[dict[str, object]] = []
|
| 229 |
+
for index, count in enumerate(raw_segment_counts):
|
| 230 |
+
if count > policy.hard_max_segments_per_window:
|
| 231 |
+
violations.append({'code': 'hard_segment_count', 'window_index': index, 'observed': count, 'maximum': policy.hard_max_segments_per_window, 'failure_class': 'delivery_budget_exceeded'})
|
| 232 |
+
if overlap_depth > policy.hard_max_overlap_depth:
|
| 233 |
+
violations.append({'code': 'overlap_depth', 'observed': overlap_depth, 'maximum': policy.hard_max_overlap_depth, 'failure_class': 'delivery_budget_exceeded'})
|
| 234 |
+
if exact_duplicates:
|
| 235 |
+
violations.append({'code': 'exact_duplicate_intervals', 'observed': exact_duplicates, 'failure_class': 'structural_invalid'})
|
| 236 |
+
if result.adjacent_matching_ambiguities > policy.max_adjacent_matching_ambiguities:
|
| 237 |
+
violations.append({'code': 'stitch_ambiguity', 'observed': result.adjacent_matching_ambiguities, 'maximum': policy.max_adjacent_matching_ambiguities, 'failure_class': 'structural_invalid'})
|
| 238 |
+
if self.fail_on_unresolved and result.unresolved:
|
| 239 |
+
violations.append({'code': 'unresolved_components', 'observed': len(result.unresolved), 'maximum': 0, 'failure_class': 'structural_invalid'})
|
| 240 |
+
diagnostics: dict[str, object] = {'schema_version': 'subtask_annotation_atomic_fine_diagnostics_v1', 'inference_arm': inference_arm, 'window_profile_id': window_profile.profile_id, 'window_profile': window_profile.to_dict(), 'model_io_contract': window_profile.model_io_contract, 'target_granularity': window_profile.target_granularity, 'policy': policy.to_dict(), 'raw_segment_counts': list(raw_segment_counts), 'emitted_segment_counts': list(emitted_segment_counts), 'guard_only_dropped': [[asdict(fragment) for fragment in fragments] for fragments in guard_only_fragments], 'soft_limit_exceeded_windows': [index for index, count in enumerate(raw_segment_counts) if count > policy.soft_max_segments_per_window], 'resolved_event_count': len(result.events), 'unresolved_components': [asdict(item) for item in result.unresolved], 'adjacent_matching_ambiguities': result.adjacent_matching_ambiguities, 'overlap_pair_count': overlap_pairs, 'max_overlap_depth': overlap_depth, 'exact_duplicate_intervals': exact_duplicates, 'acceptance_failure_classes': sorted({str(item['failure_class']) for item in violations}), 'violations': violations}
|
| 241 |
+
if result.diagnostics:
|
| 242 |
+
diagnostics['stitch_diagnostics'] = dict(result.diagnostics)
|
| 243 |
+
diagnostics['diagnostics_sha256'] = _diagnostic_sha256(diagnostics)
|
| 244 |
+
assessed = replace(result, diagnostics=diagnostics)
|
| 245 |
+
if violations:
|
| 246 |
+
codes = ','.join((str(item['code']) for item in violations))
|
| 247 |
+
raise FineSegmentationOutputError(f'fine segmentation failed output acceptance gate: {codes}', diagnostics=diagnostics, partial_result=assessed)
|
| 248 |
+
return assessed
|
| 249 |
+
|
| 250 |
+
def _predict_scheduled_windows(self, source: AnnotationSource, windows: Sequence[FineWindow]) -> _FineWindowBatchPrediction:
|
| 251 |
+
return asyncio.run(self._predict_scheduled_windows_async(source, windows))
|
| 252 |
+
|
| 253 |
+
async def _predict_scheduled_windows_async(self, source: AnnotationSource, windows: Sequence[FineWindow]) -> _FineWindowBatchPrediction:
|
| 254 |
+
request = {'schema_version': FINE_REQUEST_SCHEMA, 'profile_id': self.profile_id, 'model_binding': dict(self.model_binding), 'source': source_payload(source), 'windows': [window.worker_payload() for window in windows], 'max_concurrency': min(self.batch_concurrency, len(windows))}
|
| 255 |
+
apredict = getattr(self.predictor, 'apredict', None)
|
| 256 |
+
if callable(apredict):
|
| 257 |
+
response = await apredict(request)
|
| 258 |
+
else:
|
| 259 |
+
response = await asyncio.to_thread(self.predictor.predict, request)
|
| 260 |
+
required_keys = {'schema_version', 'source_id', 'model_io_contract', 'model_binding', 'predictions'}
|
| 261 |
+
if set(response) != required_keys:
|
| 262 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation response must contain exactly schema_version, source_id, model_io_contract, model_binding, predictions')
|
| 263 |
+
if response['schema_version'] != FINE_RESPONSE_SCHEMA:
|
| 264 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation response schema mismatch')
|
| 265 |
+
if response['source_id'] != source.source_id:
|
| 266 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation response identity mismatch')
|
| 267 |
+
if response['model_io_contract'] != fine_window_profile(self.window_profile_id).model_io_contract:
|
| 268 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation model I/O contract mismatch')
|
| 269 |
+
if response['model_binding'] != self.model_binding:
|
| 270 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation worker model_binding mismatch')
|
| 271 |
+
prediction_rows = response['predictions']
|
| 272 |
+
if not isinstance(prediction_rows, list) or len(prediction_rows) != len(windows):
|
| 273 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation predictions must exactly cover every requested window')
|
| 274 |
+
expected_scope_ids = [window.scope_id for window in windows]
|
| 275 |
+
observed_scope_ids: list[str] = []
|
| 276 |
+
fragments_by_window = {}
|
| 277 |
+
raw_segment_counts: list[int] = []
|
| 278 |
+
emitted_segment_counts: list[int] = []
|
| 279 |
+
guard_only_fragments: list[Sequence[object]] = []
|
| 280 |
+
for window, prediction in zip(windows, prediction_rows, strict=True):
|
| 281 |
+
if not isinstance(prediction, dict) or set(prediction) != {'scope_id', 'prediction_json'}:
|
| 282 |
+
raise SubtaskAnnotationAdapterError('each fine-segmentation prediction must contain exactly scope_id and prediction_json')
|
| 283 |
+
observed_scope_ids.append(prediction['scope_id'])
|
| 284 |
+
if prediction['scope_id'] != window.scope_id:
|
| 285 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation prediction order or scope differs')
|
| 286 |
+
prediction_json = prediction['prediction_json']
|
| 287 |
+
if not isinstance(prediction_json, str):
|
| 288 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation prediction_json must be a string')
|
| 289 |
+
try:
|
| 290 |
+
parsed = parse_prediction_with_diagnostics(prediction_json, window, drop_guard_only=self.quality_policy.drop_guard_only_segments)
|
| 291 |
+
except FineTemporalError as exc:
|
| 292 |
+
raise FineSegmentationOutputError(f'fine-segmentation window {window.window_index} returned an invalid prediction: {exc}', diagnostics={'schema_version': 'subtask_annotation_atomic_fine_diagnostics_v1', 'stage': 'parse_prediction', 'window_index': window.window_index, 'error': str(exc)}) from exc
|
| 293 |
+
fragments_by_window[window.window_index] = parsed.fragments
|
| 294 |
+
raw_segment_counts.append(len(parsed.fragments) + len(parsed.guard_only_fragments))
|
| 295 |
+
emitted_segment_counts.append(len(parsed.fragments))
|
| 296 |
+
guard_only_fragments.append(parsed.guard_only_fragments)
|
| 297 |
+
if observed_scope_ids != expected_scope_ids or len(set(observed_scope_ids)) != len(observed_scope_ids):
|
| 298 |
+
raise SubtaskAnnotationAdapterError('fine-segmentation response scope coverage differs')
|
| 299 |
+
return _FineWindowBatchPrediction(fragments_by_window={index: tuple(fragments) for index, fragments in fragments_by_window.items()}, raw_segment_counts=tuple(raw_segment_counts), emitted_segment_counts=tuple(emitted_segment_counts), guard_only_fragments=tuple((tuple(items) for items in guard_only_fragments)))
|
| 300 |
+
|
| 301 |
+
def predict_repair_candidates(self, source: AnnotationSource, windows: Sequence[FineWindow], *, arm: str='adaptive_fine_repair') -> _FineRetryPrediction:
|
| 302 |
+
return asyncio.run(self.predict_repair_candidates_async(source, windows, arm=arm))
|
| 303 |
+
|
| 304 |
+
async def predict_repair_candidates_async(self, source: AnnotationSource, windows: Sequence[FineWindow], *, arm: str='adaptive_fine_repair') -> _FineRetryPrediction:
|
| 305 |
+
batch = await self._predict_scheduled_windows_async(source, windows)
|
| 306 |
+
hard_limit_violations = [{'code': 'hard_segment_count', 'window_index': window.window_index, 'observed': count, 'maximum': self.quality_policy.hard_max_segments_per_window, 'failure_class': 'delivery_budget_exceeded'} for window, count in zip(windows, batch.raw_segment_counts, strict=True) if count > self.quality_policy.hard_max_segments_per_window]
|
| 307 |
+
diagnostics: dict[str, object] = {'schema_version': ADAPTIVE_FINE_REPAIR_SCHEMA, 'stage': 'fine_retry_prediction', 'inference_arm': arm, 'worker_batch_call_count': 1, 'retry_window_count': len(windows), 'raw_segment_counts': list(batch.raw_segment_counts), 'emitted_segment_counts': list(batch.emitted_segment_counts), 'guard_only_dropped': [[asdict(fragment) for fragment in fragments] for fragments in batch.guard_only_fragments], 'violations': hard_limit_violations}
|
| 308 |
+
diagnostics['diagnostics_sha256'] = _diagnostic_sha256(diagnostics)
|
| 309 |
+
if hard_limit_violations:
|
| 310 |
+
raise FineSegmentationOutputError('adaptive fine retry exceeded the per-window segment budget', diagnostics=diagnostics)
|
| 311 |
+
candidates_by_window: dict[int, tuple[TemporalSegment, ...]] = {}
|
| 312 |
+
for window in windows:
|
| 313 |
+
candidates: list[TemporalSegment] = []
|
| 314 |
+
for fragment in batch.fragments_by_window[window.window_index]:
|
| 315 |
+
start_us = (window.global_tick_offset + fragment.s) * TICK_US
|
| 316 |
+
end_us = (window.global_tick_offset + fragment.e) * TICK_US
|
| 317 |
+
physical_start_us = max(window.valid_start_us, start_us)
|
| 318 |
+
physical_end_us = min(window.valid_end_us, end_us)
|
| 319 |
+
if physical_end_us <= physical_start_us:
|
| 320 |
+
raise FineSegmentationOutputError('adaptive fine retry produced an empty physical source interval', diagnostics={**diagnostics, 'window_index': window.window_index, 'fragment': asdict(fragment), 'model_tick_span_us': {'start': start_us, 'end': end_us}, 'physical_valid_span_us': {'start': window.valid_start_us, 'end': window.valid_end_us}})
|
| 321 |
+
physical_clip = None
|
| 322 |
+
if physical_start_us != start_us or physical_end_us != end_us:
|
| 323 |
+
physical_clip = {'contract': 'model_tick_span_intersect_physical_valid_span_v1', 'model_tick_span_us': {'start': start_us, 'end': end_us}, 'physical_valid_span_us': {'start': window.valid_start_us, 'end': window.valid_end_us}, 'clipped_start': physical_start_us != start_us, 'clipped_end': physical_end_us != end_us}
|
| 324 |
+
candidates.append(TemporalSegment(TimeSpan(physical_start_us / 1000000, physical_end_us / 1000000), {'scope_clipped_start': fragment.lo, 'scope_clipped_end': fragment.ro, 'fine_retry_scope_id': window.scope_id, 'fine_retry_window_index': window.window_index, **({'fine_retry_physical_clip': physical_clip} if physical_clip is not None else {})}))
|
| 325 |
+
candidates_by_window[window.window_index] = tuple(candidates)
|
| 326 |
+
return _FineRetryPrediction(candidates_by_window=candidates_by_window, diagnostics=diagnostics)
|
| 327 |
+
|
| 328 |
+
def segment_with_priors(self, source: AnnotationSource, coarse_proposals: Sequence[ProcessingScope], *, arm: str='no_prior_global_scan') -> FineSegmentationResult:
|
| 329 |
+
return asyncio.run(self.segment_with_priors_async(source, coarse_proposals, arm=arm))
|
| 330 |
+
|
| 331 |
+
async def segment_with_priors_async(self, source: AnnotationSource, coarse_proposals: Sequence[ProcessingScope], *, arm: str='no_prior_global_scan') -> FineSegmentationResult:
|
| 332 |
+
if source.duration_s is None:
|
| 333 |
+
raise SubtaskAnnotationAdapterError('fine segmentation requires exact source duration_s')
|
| 334 |
+
scopes = schedule_fine_scopes(source_id=source.source_id, duration_s=source.duration_s, coarse_proposals=coarse_proposals, arm=arm, window_profile_id=self.window_profile_id)
|
| 335 |
+
windows = tuple((window_from_scope(scope) for scope in scopes))
|
| 336 |
+
batch = await self._predict_scheduled_windows_async(source, windows)
|
| 337 |
+
result = self._stitch_with_unresolved_local_candidates(windows, batch.fragments_by_window)
|
| 338 |
+
return self._assess_result(result, inference_arm=arm, raw_segment_counts=batch.raw_segment_counts, emitted_segment_counts=batch.emitted_segment_counts, guard_only_fragments=batch.guard_only_fragments)
|
| 339 |
+
|
| 340 |
+
def segment(self, source: AnnotationSource) -> tuple[TemporalSegment, ...]:
|
| 341 |
+
return asyncio.run(self.segment_async(source))
|
| 342 |
+
|
| 343 |
+
async def segment_async(self, source: AnnotationSource) -> tuple[TemporalSegment, ...]:
|
| 344 |
+
result = await self.segment_with_priors_async(source, (), arm='no_prior_global_scan')
|
| 345 |
+
return events_to_temporal_segments(result.events, diagnostics=result.diagnostics)
|
| 346 |
+
|
| 347 |
+
@dataclass(frozen=True)
|
| 348 |
+
class _AcceptedFineRetry:
|
| 349 |
+
component_index: int
|
| 350 |
+
window: FineWindow
|
| 351 |
+
candidate: TemporalSegment
|
| 352 |
+
repaired_span: TimeSpan
|
| 353 |
+
repair_mode: str = 'boundary_completion'
|
| 354 |
+
confirmation_target_event_id: str | None = None
|
| 355 |
+
local_stitch_candidate: Mapping[str, object] | None = None
|
| 356 |
+
|
| 357 |
+
@dataclass(frozen=True)
|
| 358 |
+
class _MatchedFineRetryCandidate:
|
| 359 |
+
candidate: TemporalSegment
|
| 360 |
+
repaired_span: TimeSpan
|
| 361 |
+
repair_mode: str
|
| 362 |
+
local_stitch_candidate: Mapping[str, object] | None = None
|
| 363 |
+
|
| 364 |
+
def _tick_seconds(value: int) -> float:
|
| 365 |
+
return value * TICK_US / 1000000
|
| 366 |
+
|
| 367 |
+
def _same_event_geometry(left: TimeSpan, right: TimeSpan, *, tolerance_s: float=SEGMENT_LOCAL_REPAIR_ENDPOINT_TOLERANCE_S) -> bool:
|
| 368 |
+
return abs(left.start_s - right.start_s) <= tolerance_s and abs(left.end_s - right.end_s) <= tolerance_s
|
| 369 |
+
|
| 370 |
+
def _unresolved_local_stitch_candidates(diagnostics: Mapping[str, object]) -> dict[int, dict[str, object]]:
|
| 371 |
+
stitch_diagnostics = diagnostics.get('stitch_diagnostics')
|
| 372 |
+
if not isinstance(stitch_diagnostics, dict):
|
| 373 |
+
return {}
|
| 374 |
+
local = stitch_diagnostics.get('unresolved_local_stitch')
|
| 375 |
+
if not isinstance(local, dict):
|
| 376 |
+
return {}
|
| 377 |
+
candidates = local.get('candidates')
|
| 378 |
+
if not isinstance(candidates, list):
|
| 379 |
+
raise SubtaskAnnotationAdapterError('unresolved-local stitch diagnostics contain an invalid candidate ledger')
|
| 380 |
+
by_component: dict[int, dict[str, object]] = {}
|
| 381 |
+
for candidate in candidates:
|
| 382 |
+
if not isinstance(candidate, dict):
|
| 383 |
+
raise SubtaskAnnotationAdapterError('unresolved-local stitch candidate must be an object')
|
| 384 |
+
component_index = candidate.get('partial_component_index')
|
| 385 |
+
event_id = candidate.get('resolved_event_id')
|
| 386 |
+
span = candidate.get('resolved_event_span')
|
| 387 |
+
if isinstance(component_index, bool) or not isinstance(component_index, int) or component_index < 0 or (not isinstance(event_id, str)) or (not event_id) or (not isinstance(span, dict)) or (set(span) != {'start_s', 'end_s'}):
|
| 388 |
+
raise SubtaskAnnotationAdapterError('unresolved-local stitch candidate binding is invalid')
|
| 389 |
+
if component_index in by_component:
|
| 390 |
+
raise SubtaskAnnotationAdapterError('unresolved-local stitch candidate ledger repeats a component')
|
| 391 |
+
by_component[component_index] = dict(candidate)
|
| 392 |
+
return by_component
|
| 393 |
+
|
| 394 |
+
def _local_stitch_target_span(candidate: Mapping[str, object]) -> TimeSpan:
|
| 395 |
+
payload = candidate.get('resolved_event_span')
|
| 396 |
+
if not isinstance(payload, dict):
|
| 397 |
+
raise SubtaskAnnotationAdapterError('unresolved-local stitch candidate target span is missing')
|
| 398 |
+
start_s, end_s = (payload.get('start_s'), payload.get('end_s'))
|
| 399 |
+
if isinstance(start_s, bool) or not isinstance(start_s, int | float) or isinstance(end_s, bool) or (not isinstance(end_s, int | float)):
|
| 400 |
+
raise SubtaskAnnotationAdapterError('unresolved-local stitch candidate target span is invalid')
|
| 401 |
+
return TimeSpan(float(start_s), float(end_s))
|
| 402 |
+
|
| 403 |
+
def _retry_confirms_local_stitch(retry_candidate: TemporalSegment, local_candidate: Mapping[str, object]) -> bool:
|
| 404 |
+
return _same_event_geometry(retry_candidate.span, _local_stitch_target_span(local_candidate), tolerance_s=STITCH_TOLERANCE_TICKS * TICK_US / 1000000)
|
| 405 |
+
|
| 406 |
+
def _component_observation_span(component: UnresolvedFineEvent) -> TimeSpan:
|
| 407 |
+
return TimeSpan(_tick_seconds(min((item[0] for item in component.observation_spans_ticks))), _tick_seconds(max((item[1] for item in component.observation_spans_ticks))))
|
| 408 |
+
|
| 409 |
+
def _endpoint_consensus_ticks(values: Sequence[int]) -> int:
|
| 410 |
+
ordered = sorted(values)
|
| 411 |
+
if not ordered:
|
| 412 |
+
raise SubtaskAnnotationAdapterError('cannot form endpoint consensus without evidence')
|
| 413 |
+
middle = len(ordered) // 2
|
| 414 |
+
if len(ordered) % 2:
|
| 415 |
+
return ordered[middle]
|
| 416 |
+
return (ordered[middle - 1] + ordered[middle] + 1) // 2
|
| 417 |
+
|
| 418 |
+
def _conservative_component_completion_span(source: AnnotationSource, component: UnresolvedFineEvent) -> tuple[TimeSpan, str]:
|
| 419 |
+
if source.duration_s is None:
|
| 420 |
+
raise SubtaskAnnotationAdapterError('conservative fine completion requires exact source duration_s')
|
| 421 |
+
if not component.observation_spans_ticks:
|
| 422 |
+
raise SubtaskAnnotationAdapterError('conservative fine completion requires at least one observed fragment')
|
| 423 |
+
observed_start_tick = min((item[0] for item in component.observation_spans_ticks))
|
| 424 |
+
observed_end_tick = max((item[1] for item in component.observation_spans_ticks))
|
| 425 |
+
has_start = bool(component.closed_start_ticks)
|
| 426 |
+
has_end = bool(component.closed_end_ticks)
|
| 427 |
+
start_tick = _endpoint_consensus_ticks(component.closed_start_ticks) if has_start else observed_start_tick
|
| 428 |
+
end_tick = _endpoint_consensus_ticks(component.closed_end_ticks) if has_end else observed_end_tick
|
| 429 |
+
strategy = {(True, False): 'closed_start_plus_observed_censored_end_v1', (False, True): 'observed_censored_start_plus_closed_end_v1', (True, True): 'closed_endpoint_consensus_v1', (False, False): 'observed_fragment_envelope_v1'}[has_start, has_end]
|
| 430 |
+
start_s = max(0.0, _tick_seconds(start_tick))
|
| 431 |
+
end_s = min(source.duration_s, _tick_seconds(end_tick))
|
| 432 |
+
if end_s <= start_s:
|
| 433 |
+
start_s = max(0.0, _tick_seconds(observed_start_tick))
|
| 434 |
+
end_s = min(source.duration_s, _tick_seconds(observed_end_tick))
|
| 435 |
+
strategy = 'observed_fragment_envelope_after_endpoint_disagreement_v1'
|
| 436 |
+
if end_s <= start_s:
|
| 437 |
+
raise SubtaskAnnotationAdapterError('conservative fine completion produced an empty physical interval')
|
| 438 |
+
return (TimeSpan(start_s, end_s), strategy)
|
| 439 |
+
|
| 440 |
+
def _component_boundary_completion_span(component: UnresolvedFineEvent, segment: TemporalSegment) -> TimeSpan | None:
|
| 441 |
+
has_start = bool(component.closed_start_ticks)
|
| 442 |
+
has_end = bool(component.closed_end_ticks)
|
| 443 |
+
if has_start == has_end:
|
| 444 |
+
return None
|
| 445 |
+
observation = _component_observation_span(component)
|
| 446 |
+
if has_start:
|
| 447 |
+
if segment.metadata.get('scope_clipped_end') is True:
|
| 448 |
+
return None
|
| 449 |
+
if not segment.span.start_s < observation.end_s <= segment.span.end_s + 1e-09:
|
| 450 |
+
return None
|
| 451 |
+
start_s = _tick_seconds(_endpoint_consensus_ticks(component.closed_start_ticks))
|
| 452 |
+
end_s = segment.span.end_s
|
| 453 |
+
else:
|
| 454 |
+
if segment.metadata.get('scope_clipped_start') is True:
|
| 455 |
+
return None
|
| 456 |
+
if not segment.span.start_s <= observation.start_s + 1e-09 < segment.span.end_s:
|
| 457 |
+
return None
|
| 458 |
+
start_s = segment.span.start_s
|
| 459 |
+
end_s = _tick_seconds(_endpoint_consensus_ticks(component.closed_end_ticks))
|
| 460 |
+
if end_s <= start_s:
|
| 461 |
+
return None
|
| 462 |
+
return TimeSpan(start_s, end_s)
|
| 463 |
+
|
| 464 |
+
def _match_fine_retry_candidates(component_index: int, component: UnresolvedFineEvent, candidates: Sequence[TemporalSegment], local_stitch_candidates: Mapping[int, Mapping[str, object]]) -> tuple[tuple[_MatchedFineRetryCandidate, ...], tuple[dict[str, object], ...]]:
|
| 465 |
+
local_candidate = local_stitch_candidates.get(component_index)
|
| 466 |
+
matching: list[_MatchedFineRetryCandidate] = []
|
| 467 |
+
ignored: list[dict[str, object]] = []
|
| 468 |
+
for candidate in candidates:
|
| 469 |
+
if local_candidate is not None and _retry_confirms_local_stitch(candidate, local_candidate):
|
| 470 |
+
matching.append(_MatchedFineRetryCandidate(candidate=candidate, repaired_span=_local_stitch_target_span(local_candidate), repair_mode='unresolved_local_stitch_confirmed', local_stitch_candidate=local_candidate))
|
| 471 |
+
continue
|
| 472 |
+
repaired_span = _component_boundary_completion_span(component, candidate)
|
| 473 |
+
if repaired_span is not None:
|
| 474 |
+
matching.append(_MatchedFineRetryCandidate(candidate=candidate, repaired_span=repaired_span, repair_mode='boundary_completion'))
|
| 475 |
+
continue
|
| 476 |
+
ignored.append({'span': _span_payload(candidate.span), 'scope_clipped_start': candidate.metadata.get('scope_clipped_start'), 'scope_clipped_end': candidate.metadata.get('scope_clipped_end')})
|
| 477 |
+
return (tuple(matching), tuple(ignored))
|
| 478 |
+
|
| 479 |
+
def _covering_span(left: TimeSpan, right: TimeSpan) -> TimeSpan:
|
| 480 |
+
return TimeSpan(min(left.start_s, right.start_s), max(left.end_s, right.end_s))
|
| 481 |
+
|
| 482 |
+
def _unsupported_unresolved_component_indices(source: AnnotationSource, unresolved: Sequence[UnresolvedFineEvent]) -> list[int]:
|
| 483 |
+
return [index for index, component in enumerate(unresolved) if component.matching_ambiguous or (component.closed_start_ticks and component.closed_end_ticks) or (not component.closed_start_ticks and (not component.closed_end_ticks)) or (not component.closed_start_ticks and min((item[0] for item in component.observation_spans_ticks)) <= 0) or (not component.closed_end_ticks and source.duration_s is not None and (_tick_seconds(max((item[1] for item in component.observation_spans_ticks))) >= source.duration_s - 1e-09))]
|
| 484 |
+
|
| 485 |
+
def _adaptive_fine_failure(exc: FineSegmentationOutputError, *, code: str, detail: Mapping[str, object]) -> FineSegmentationOutputError:
|
| 486 |
+
diagnostics = dict(exc.diagnostics)
|
| 487 |
+
failures = list(diagnostics.get('adaptive_fine_repair_failures', []))
|
| 488 |
+
failures.append({'code': code, **dict(detail)})
|
| 489 |
+
diagnostics['adaptive_fine_repair_failures'] = failures
|
| 490 |
+
diagnostics['diagnostics_sha256'] = _diagnostic_sha256({key: value for key, value in diagnostics.items() if key != 'diagnostics_sha256'})
|
| 491 |
+
return FineSegmentationOutputError(f'fine segmentation could not be repaired by bounded fine retry: {code}', diagnostics=diagnostics, partial_result=exc.partial_result)
|
| 492 |
+
|
| 493 |
+
def _remaining_adaptive_fine_error(exc: FineSegmentationOutputError, result: FineSegmentationResult, *, repaired_events: tuple[EventInterval, ...], accepted_component_indices: Sequence[int], local_stitch_candidates: Mapping[int, Mapping[str, object]], deferred_failures: Sequence[Mapping[str, object]], repair_payload: Mapping[str, object], inference_context: Mapping[str, object]) -> FineSegmentationOutputError:
|
| 494 |
+
accepted = set(accepted_component_indices)
|
| 495 |
+
remaining_indices = [index for index in range(len(result.unresolved)) if index not in accepted]
|
| 496 |
+
if not remaining_indices:
|
| 497 |
+
raise SubtaskAnnotationAdapterError('remaining adaptive Fine error requires at least one unresolved component')
|
| 498 |
+
original_rows = exc.diagnostics.get('unresolved_components')
|
| 499 |
+
rows_are_aligned = isinstance(original_rows, list) and len(original_rows) == len(result.unresolved)
|
| 500 |
+
remaining_rows = [original_rows[index] for index in remaining_indices] if rows_are_aligned else []
|
| 501 |
+
repaired_rows = [original_rows[index] for index in sorted(accepted)] if rows_are_aligned else []
|
| 502 |
+
reduced_index = {old_index: new_index for new_index, old_index in enumerate(remaining_indices)}
|
| 503 |
+
remaining_local_candidates = []
|
| 504 |
+
for old_index in remaining_indices:
|
| 505 |
+
candidate = local_stitch_candidates.get(old_index)
|
| 506 |
+
if candidate is None:
|
| 507 |
+
continue
|
| 508 |
+
rebound = dict(candidate)
|
| 509 |
+
rebound['partial_component_index'] = reduced_index[old_index]
|
| 510 |
+
rebound['original_partial_component_index'] = old_index
|
| 511 |
+
remaining_local_candidates.append(rebound)
|
| 512 |
+
diagnostics = dict(exc.diagnostics)
|
| 513 |
+
failures = list(diagnostics.get('adaptive_fine_repair_failures', []))
|
| 514 |
+
failures.extend((dict(item) for item in deferred_failures))
|
| 515 |
+
diagnostics['adaptive_fine_repair_failures'] = failures
|
| 516 |
+
diagnostics['pre_repair_violations'] = list(diagnostics.get('violations', []))
|
| 517 |
+
diagnostics['repaired_unresolved_components'] = repaired_rows
|
| 518 |
+
diagnostics['unresolved_components'] = remaining_rows
|
| 519 |
+
violations = [dict(item) for item in diagnostics.get('violations', []) if isinstance(item, dict)]
|
| 520 |
+
for violation in violations:
|
| 521 |
+
if violation.get('code') == 'unresolved_components':
|
| 522 |
+
violation['observed'] = len(remaining_indices)
|
| 523 |
+
diagnostics['violations'] = violations
|
| 524 |
+
diagnostics['resolved_event_count'] = len(repaired_events)
|
| 525 |
+
diagnostics['overlap_pair_count'] = repair_payload['post_repair_overlap_pair_count']
|
| 526 |
+
diagnostics['max_overlap_depth'] = repair_payload['post_repair_max_overlap_depth']
|
| 527 |
+
diagnostics['exact_duplicate_intervals'] = repair_payload['post_repair_exact_duplicate_intervals']
|
| 528 |
+
diagnostics['adaptive_fine_repair'] = dict(repair_payload)
|
| 529 |
+
diagnostics['adaptive_fine_repair_sha256'] = _diagnostic_sha256(repair_payload)
|
| 530 |
+
stitch_diagnostics = dict(diagnostics.get('stitch_diagnostics', {}))
|
| 531 |
+
local_diagnostics = dict(stitch_diagnostics.get('unresolved_local_stitch', {}))
|
| 532 |
+
if local_diagnostics:
|
| 533 |
+
local_diagnostics.update({'accepted_pair_count': len(remaining_local_candidates), 'candidate_component_count': len(remaining_local_candidates), 'candidates': remaining_local_candidates, 'potential_absorbed_unresolved_component_count': len(remaining_local_candidates), 'potential_remaining_unresolved_component_count': len(remaining_indices) - len(remaining_local_candidates), 'strict_unresolved_component_count': len(remaining_indices)})
|
| 534 |
+
stitch_diagnostics['unresolved_local_stitch'] = local_diagnostics
|
| 535 |
+
diagnostics['stitch_diagnostics'] = stitch_diagnostics
|
| 536 |
+
diagnostics = _with_inference_context(diagnostics, inference_context)
|
| 537 |
+
partial_result = FineSegmentationResult(events=repaired_events, fragments=result.fragments, unresolved=tuple((result.unresolved[index] for index in remaining_indices)), adjacent_matching_ambiguities=result.adjacent_matching_ambiguities, diagnostics=diagnostics)
|
| 538 |
+
return FineSegmentationOutputError('bounded fine retry partially repaired the source; remaining components require conservative completion', diagnostics=diagnostics, partial_result=partial_result)
|
| 539 |
+
|
| 540 |
+
@dataclass(frozen=True)
|
| 541 |
+
class HierarchicalSubtaskSegmenter:
|
| 542 |
+
coarse_segmenter: VideoSegmenter
|
| 543 |
+
fine_segmenter: FineSubtaskSegmenter
|
| 544 |
+
fine_prior_mode: str = 'none'
|
| 545 |
+
fine_failure_policy: str = 'adaptive_fine_then_conservative_completion'
|
| 546 |
+
fine_output_coverage_policy: str = 'allow_gaps'
|
| 547 |
+
adaptive_fine_retry_max_windows: int = ADAPTIVE_FINE_RETRY_MAX_WINDOWS
|
| 548 |
+
segment_topology = SegmentTopology.CONCURRENT
|
| 549 |
+
|
| 550 |
+
def __post_init__(self):
|
| 551 |
+
if self.fine_prior_mode != 'none' or self.fine_failure_policy != 'adaptive_fine_then_conservative_completion' or self.fine_output_coverage_policy != 'allow_gaps':
|
| 552 |
+
raise ValueError('Use the released inference workflow configuration.')
|
| 553 |
+
if isinstance(self.adaptive_fine_retry_max_windows, bool) or not isinstance(self.adaptive_fine_retry_max_windows, int) or (not 1 <= self.adaptive_fine_retry_max_windows <= ADAPTIVE_FINE_RETRY_HARD_MAX_WINDOWS):
|
| 554 |
+
raise ValueError('Invalid adaptive inference retry budget')
|
| 555 |
+
|
| 556 |
+
def adaptive_fine_repair(self, source: AnnotationSource, exc: FineSegmentationOutputError, *, inference_context: Mapping[str, object] | None=None) -> tuple[TemporalSegment, ...]:
|
| 557 |
+
return asyncio.run(self.adaptive_fine_repair_async(source, exc, inference_context=inference_context))
|
| 558 |
+
|
| 559 |
+
async def adaptive_fine_repair_async(self, source: AnnotationSource, exc: FineSegmentationOutputError, *, inference_context: Mapping[str, object] | None=None) -> tuple[TemporalSegment, ...]:
|
| 560 |
+
if inference_context is None:
|
| 561 |
+
inference_context = _inference_context(source, (), fine_prior_mode='none')
|
| 562 |
+
result = exc.partial_result
|
| 563 |
+
violation_codes = [str(item.get('code')) for item in exc.diagnostics.get('violations', []) if isinstance(item, dict)]
|
| 564 |
+
if result is None or not result.unresolved or violation_codes != ['unresolved_components']:
|
| 565 |
+
raise _adaptive_fine_failure(exc, code='ineligible_fine_failure', detail={'violation_codes': violation_codes}) from exc
|
| 566 |
+
if source.duration_s is None:
|
| 567 |
+
raise _adaptive_fine_failure(exc, code='duration_unavailable', detail={}) from exc
|
| 568 |
+
progressive_provider_free = self.fine_failure_policy in PROVIDER_FREE_FINE_COMPLETION_POLICIES
|
| 569 |
+
window_profile = fine_window_profile(self.fine_segmenter.window_profile_id)
|
| 570 |
+
local_stitch_candidates = _unresolved_local_stitch_candidates(exc.diagnostics)
|
| 571 |
+
unsupported_components = _unsupported_unresolved_component_indices(source, result.unresolved)
|
| 572 |
+
if unsupported_components and (not progressive_provider_free):
|
| 573 |
+
raise _adaptive_fine_failure(exc, code='unsupported_unresolved_geometry', detail={'component_indices': unsupported_components}) from exc
|
| 574 |
+
if len(result.unresolved) > self.adaptive_fine_retry_max_windows and (not progressive_provider_free):
|
| 575 |
+
raise _adaptive_fine_failure(exc, code='retry_budget_exceeded', detail={'unresolved_component_count': len(result.unresolved), 'max_retry_windows': self.adaptive_fine_retry_max_windows}) from exc
|
| 576 |
+
deferred_failures: list[dict[str, object]] = []
|
| 577 |
+
unsupported_set = set(unsupported_components)
|
| 578 |
+
if unsupported_components:
|
| 579 |
+
deferred_failures.append({'code': 'unsupported_unresolved_geometry', 'component_indices': unsupported_components})
|
| 580 |
+
eligible_component_indices = [index for index in range(len(result.unresolved)) if index not in unsupported_set]
|
| 581 |
+
eligible_component_indices.sort(key=lambda index: (index not in local_stitch_candidates, index))
|
| 582 |
+
selected_component_indices = eligible_component_indices[:self.adaptive_fine_retry_max_windows]
|
| 583 |
+
budget_deferred_indices = eligible_component_indices[self.adaptive_fine_retry_max_windows:]
|
| 584 |
+
if budget_deferred_indices:
|
| 585 |
+
deferred_failures.append({'code': 'retry_budget_exceeded', 'unresolved_component_count': len(result.unresolved), 'max_retry_windows': self.adaptive_fine_retry_max_windows, 'deferred_component_indices': budget_deferred_indices, 'candidate_priority': 'unresolved_local_stitch_first_then_source_order_v1'})
|
| 586 |
+
if not selected_component_indices:
|
| 587 |
+
diagnostics = dict(exc.diagnostics)
|
| 588 |
+
failures = list(diagnostics.get('adaptive_fine_repair_failures', []))
|
| 589 |
+
failures.extend(deferred_failures)
|
| 590 |
+
diagnostics['adaptive_fine_repair_failures'] = failures
|
| 591 |
+
diagnostics['diagnostics_sha256'] = _diagnostic_sha256({key: value for key, value in diagnostics.items() if key != 'diagnostics_sha256'})
|
| 592 |
+
raise FineSegmentationOutputError('fine segmentation had no eligible bounded fine retry components', diagnostics=diagnostics, partial_result=result) from exc
|
| 593 |
+
component_edge_ticks = tuple(((component_index, max((item[1] for item in component.observation_spans_ticks)) if component.closed_start_ticks else min((item[0] for item in component.observation_spans_ticks))) for component_index in selected_component_indices for component in (result.unresolved[component_index],)))
|
| 594 |
+
try:
|
| 595 |
+
scopes = schedule_recentered_fine_retry_scopes(source_id=source.source_id, duration_s=source.duration_s, component_edge_ticks=component_edge_ticks, window_profile_id=window_profile.profile_id)
|
| 596 |
+
windows = tuple((window_from_scope(scope) for scope in scopes))
|
| 597 |
+
except Exception as geometry_error:
|
| 598 |
+
raise _adaptive_fine_failure(exc, code='invalid_retry_geometry', detail={'error_type': type(geometry_error).__name__, 'error': str(geometry_error)}) from geometry_error
|
| 599 |
+
try:
|
| 600 |
+
prediction = await self.fine_segmenter.predict_repair_candidates_async(source, windows)
|
| 601 |
+
except FineSegmentationOutputError as retry_error:
|
| 602 |
+
raise _adaptive_fine_failure(exc, code='retry_output_rejected', detail={'error': str(retry_error), 'retry_diagnostics': retry_error.diagnostics}) from retry_error
|
| 603 |
+
accepted_by_component: dict[int, _AcceptedFineRetry] = {}
|
| 604 |
+
retry_diagnostics_by_component: dict[int, dict[str, object]] = {}
|
| 605 |
+
incomplete_primary: list[tuple[int, int, dict[str, object]]] = []
|
| 606 |
+
for (component_index, edge_tick), window in zip(component_edge_ticks, windows, strict=True):
|
| 607 |
+
component = result.unresolved[component_index]
|
| 608 |
+
candidates = prediction.candidates_by_window[window.window_index]
|
| 609 |
+
matching_candidates, ignored_candidates = _match_fine_retry_candidates(component_index, component, candidates, local_stitch_candidates)
|
| 610 |
+
failure_detail = {'component_index': component_index, 'retry_scope_id': window.scope_id, 'retry_edge_tick': edge_tick, 'candidate_count': len(candidates), 'matching_candidate_count': len(matching_candidates), 'ignored_candidates': ignored_candidates, 'retry_prediction_diagnostics': dict(prediction.diagnostics)}
|
| 611 |
+
if len(matching_candidates) > 1:
|
| 612 |
+
if not progressive_provider_free:
|
| 613 |
+
raise _adaptive_fine_failure(exc, code='ambiguous_component_coverage', detail=failure_detail)
|
| 614 |
+
deferred_failures.append({'code': 'ambiguous_component_coverage', **failure_detail})
|
| 615 |
+
continue
|
| 616 |
+
if not matching_candidates:
|
| 617 |
+
if not candidates:
|
| 618 |
+
if not progressive_provider_free:
|
| 619 |
+
raise _adaptive_fine_failure(exc, code='incomplete_component_coverage', detail=failure_detail)
|
| 620 |
+
deferred_failures.append({'code': 'incomplete_component_coverage', **failure_detail})
|
| 621 |
+
continue
|
| 622 |
+
incomplete_primary.append((component_index, edge_tick, failure_detail))
|
| 623 |
+
continue
|
| 624 |
+
match = matching_candidates[0]
|
| 625 |
+
accepted_by_component[component_index] = _AcceptedFineRetry(component_index=component_index, window=window, candidate=match.candidate, repaired_span=match.repaired_span, repair_mode=match.repair_mode, confirmation_target_event_id=str(match.local_stitch_candidate['resolved_event_id']) if match.local_stitch_candidate is not None else None, local_stitch_candidate=match.local_stitch_candidate)
|
| 626 |
+
retry_diagnostics_by_component[component_index] = {'component_index': component_index, 'retry_pass': 'centered', 'repair_mode': match.repair_mode, 'retry_scope_id': window.scope_id, 'retry_edge_tick': edge_tick, 'requested_view_start_s': window.requested_view_start_us / 1000000, 'requested_view_end_s': window.requested_view_end_us / 1000000, 'candidate_count': len(candidates), 'candidate_evidence_span': _span_payload(match.candidate.span), 'repaired_span': _span_payload(match.repaired_span), 'ignored_candidates': list(ignored_candidates), **({'unresolved_local_stitch_confirmation': dict(match.local_stitch_candidate)} if match.local_stitch_candidate is not None else {})}
|
| 627 |
+
secondary_windows: tuple[FineWindow, ...] = ()
|
| 628 |
+
secondary_prediction: _FineRetryPrediction | None = None
|
| 629 |
+
if incomplete_primary:
|
| 630 |
+
secondary_edges = tuple(((component_index, edge_tick) for component_index, edge_tick, _ in incomplete_primary))
|
| 631 |
+
secondary_offsets = {component_index: window_profile.left_guard_us // TICK_US if result.unresolved[component_index].closed_start_ticks else window_profile.core_ticks - window_profile.right_guard_us // TICK_US for component_index, _, _ in incomplete_primary}
|
| 632 |
+
try:
|
| 633 |
+
secondary_scopes = schedule_recentered_fine_retry_scopes(source_id=source.source_id, duration_s=source.duration_s, component_edge_ticks=secondary_edges, arm='adaptive_fine_directional_retry', edge_core_offset_ticks=secondary_offsets, scope_namespace=FINE_DIRECTIONAL_RETRY_SCOPE_NAMESPACE, window_profile_id=window_profile.profile_id)
|
| 634 |
+
secondary_windows = tuple((window_from_scope(scope) for scope in secondary_scopes))
|
| 635 |
+
except Exception as geometry_error:
|
| 636 |
+
raise _adaptive_fine_failure(exc, code='invalid_directional_retry_geometry', detail={'error_type': type(geometry_error).__name__, 'error': str(geometry_error), 'component_indices': [item[0] for item in incomplete_primary]}) from geometry_error
|
| 637 |
+
try:
|
| 638 |
+
secondary_prediction = await self.fine_segmenter.predict_repair_candidates_async(source, secondary_windows, arm='adaptive_fine_directional_retry')
|
| 639 |
+
except FineSegmentationOutputError as retry_error:
|
| 640 |
+
raise _adaptive_fine_failure(exc, code='directional_retry_output_rejected', detail={'error': str(retry_error), 'retry_diagnostics': retry_error.diagnostics}) from retry_error
|
| 641 |
+
for (component_index, edge_tick, primary_failure), window in zip(incomplete_primary, secondary_windows, strict=True):
|
| 642 |
+
component = result.unresolved[component_index]
|
| 643 |
+
candidates = secondary_prediction.candidates_by_window[window.window_index]
|
| 644 |
+
matching_candidates, ignored_candidates = _match_fine_retry_candidates(component_index, component, candidates, local_stitch_candidates)
|
| 645 |
+
if len(matching_candidates) != 1:
|
| 646 |
+
failure = {'code': 'incomplete_component_coverage' if not matching_candidates else 'ambiguous_component_coverage', 'component_index': component_index, 'retry_pass': 'directional_missing_endpoint_context', 'retry_scope_id': window.scope_id, 'retry_edge_tick': edge_tick, 'candidate_count': len(candidates), 'matching_candidate_count': len(matching_candidates), 'ignored_candidates': list(ignored_candidates), 'primary_retry_failure': primary_failure, 'retry_prediction_diagnostics': dict(secondary_prediction.diagnostics)}
|
| 647 |
+
if not progressive_provider_free:
|
| 648 |
+
raise _adaptive_fine_failure(exc, code=str(failure.pop('code')), detail=failure)
|
| 649 |
+
deferred_failures.append(failure)
|
| 650 |
+
continue
|
| 651 |
+
match = matching_candidates[0]
|
| 652 |
+
accepted_by_component[component_index] = _AcceptedFineRetry(component_index=component_index, window=window, candidate=match.candidate, repaired_span=match.repaired_span, repair_mode=match.repair_mode, confirmation_target_event_id=str(match.local_stitch_candidate['resolved_event_id']) if match.local_stitch_candidate is not None else None, local_stitch_candidate=match.local_stitch_candidate)
|
| 653 |
+
retry_diagnostics_by_component[component_index] = {'component_index': component_index, 'retry_pass': 'directional_missing_endpoint_context', 'repair_mode': match.repair_mode, 'retry_scope_id': window.scope_id, 'retry_edge_tick': edge_tick, 'retry_edge_core_offset_tick': secondary_offsets[component_index], 'requested_view_start_s': window.requested_view_start_us / 1000000, 'requested_view_end_s': window.requested_view_end_us / 1000000, 'candidate_count': len(candidates), 'candidate_evidence_span': _span_payload(match.candidate.span), 'repaired_span': _span_payload(match.repaired_span), 'ignored_candidates': list(ignored_candidates), 'primary_retry_failure': primary_failure, **({'unresolved_local_stitch_confirmation': dict(match.local_stitch_candidate)} if match.local_stitch_candidate is not None else {})}
|
| 654 |
+
accepted = [accepted_by_component[index] for index in sorted(accepted_by_component)]
|
| 655 |
+
retry_diagnostics = [retry_diagnostics_by_component[index] for index in sorted(retry_diagnostics_by_component)]
|
| 656 |
+
combined_events = [replace(event, metadata={'segmentation_provenance': 'fine', **dict(event.metadata)}) for event in result.events]
|
| 657 |
+
retry_event_count_added = 0
|
| 658 |
+
suppressed_duplicate_count = 0
|
| 659 |
+
local_stitch_confirmation_count = 0
|
| 660 |
+
for repair in accepted:
|
| 661 |
+
if repair.repair_mode == 'unresolved_local_stitch_confirmed':
|
| 662 |
+
geometry_matches = [(0.0, index) for index, event in enumerate(combined_events) if event.event_id == repair.confirmation_target_event_id]
|
| 663 |
+
if len(geometry_matches) != 1 or repair.local_stitch_candidate is None:
|
| 664 |
+
raise _adaptive_fine_failure(exc, code='invalid_local_stitch_confirmation_target', detail={'component_index': repair.component_index, 'target_event_id': repair.confirmation_target_event_id, 'target_match_count': len(geometry_matches)})
|
| 665 |
+
else:
|
| 666 |
+
geometry_matches = [(abs(event.span.start_s - repair.repaired_span.start_s) + abs(event.span.end_s - repair.repaired_span.end_s), index) for index, event in enumerate(combined_events) if _same_event_geometry(event.span, repair.repaired_span)]
|
| 667 |
+
if geometry_matches:
|
| 668 |
+
_, event_index = min(geometry_matches)
|
| 669 |
+
matched = combined_events[event_index]
|
| 670 |
+
metadata = dict(matched.metadata)
|
| 671 |
+
lineage = metadata.get('adaptive_fine_repair')
|
| 672 |
+
prior_lineage = dict(lineage) if isinstance(lineage, dict) else {}
|
| 673 |
+
component_indices = set(prior_lineage.get('component_indices', []))
|
| 674 |
+
component_indices.add(repair.component_index)
|
| 675 |
+
retry_scope_ids = set(prior_lineage.get('retry_scope_ids', []))
|
| 676 |
+
retry_scope_ids.add(repair.window.scope_id)
|
| 677 |
+
metadata.update({'segmentation_provenance': 'fine_repaired_by_fine_retry', 'adaptive_fine_repair': {'component_indices': sorted(component_indices), 'retry_scope_ids': sorted(retry_scope_ids), 'candidate_evidence_span': _span_payload(repair.candidate.span), 'duplicate_suppressed': True, 'repair_mode': repair.repair_mode}, **({'unresolved_local_stitch_confirmation': {'schema_version': UNRESOLVED_LOCAL_STITCH_SCHEMA, 'status': 'confirmed_by_adaptive_fine_retry', 'strict_confirmation_tolerance_ticks': STITCH_TOLERANCE_TICKS, 'retry_candidate_span': _span_payload(repair.candidate.span), 'candidate': dict(repair.local_stitch_candidate)}} if repair.local_stitch_candidate is not None else {})})
|
| 678 |
+
combined_events[event_index] = replace(matched, source_scope_ids=tuple(sorted({*matched.source_scope_ids, repair.window.scope_id})), metadata=metadata)
|
| 679 |
+
suppressed_duplicate_count += 1
|
| 680 |
+
local_stitch_confirmation_count += int(repair.repair_mode == 'unresolved_local_stitch_confirmed')
|
| 681 |
+
continue
|
| 682 |
+
candidate_metadata = {key: value for key, value in repair.candidate.metadata.items() if key not in {'scope_clipped_start', 'scope_clipped_end'}}
|
| 683 |
+
combined_events.append(EventInterval(event_id=f'{source.source_id}:adaptive-fine-repair:{retry_event_count_added:06d}', source_id=source.source_id, span=repair.repaired_span, source_scope_ids=(*sorted({*result.unresolved[repair.component_index].source_scope_ids, repair.window.scope_id}),), metadata={**candidate_metadata, 'segmentation_provenance': 'fine_repaired_by_fine_retry', 'adaptive_fine_repair': {'component_indices': [repair.component_index], 'retry_scope_ids': [repair.window.scope_id], 'candidate_evidence_span': _span_payload(repair.candidate.span), 'duplicate_suppressed': False, 'repair_mode': repair.repair_mode}}))
|
| 684 |
+
retry_event_count_added += 1
|
| 685 |
+
ordered_events = tuple(sorted(combined_events, key=lambda item: (item.span, item.event_id)))
|
| 686 |
+
overlap_pairs, overlap_depth = _overlap_statistics(ordered_events)
|
| 687 |
+
exact_duplicates = _exact_duplicate_intervals(ordered_events)
|
| 688 |
+
post_repair_violations: list[dict[str, object]] = []
|
| 689 |
+
if overlap_depth > self.fine_segmenter.quality_policy.hard_max_overlap_depth:
|
| 690 |
+
post_repair_violations.append({'code': 'overlap_depth', 'observed': overlap_depth, 'maximum': self.fine_segmenter.quality_policy.hard_max_overlap_depth, 'failure_class': 'delivery_budget_exceeded'})
|
| 691 |
+
if exact_duplicates:
|
| 692 |
+
post_repair_violations.append({'code': 'exact_duplicate_intervals', 'observed': exact_duplicates, 'failure_class': 'structural_invalid'})
|
| 693 |
+
if post_repair_violations:
|
| 694 |
+
raise _adaptive_fine_failure(exc, code='invalid_combined_geometry', detail={'combined_event_count': len(ordered_events), 'violations': post_repair_violations})
|
| 695 |
+
repair_payload: dict[str, object] = {'schema_version': ADAPTIVE_FINE_REPAIR_SCHEMA, 'scheduler_contract': 'adaptive_recentered_then_directional_' + window_profile.retry_scheduler_contract(directional=True).removeprefix('adaptive_directional_') if secondary_windows else window_profile.retry_scheduler_contract(directional=False), 'model_io_contract': window_profile.model_io_contract, 'window_profile_id': window_profile.profile_id, 'max_retry_windows': self.adaptive_fine_retry_max_windows, 'unresolved_component_count': len(result.unresolved), 'selected_component_count': len(selected_component_indices), 'selected_component_indices': selected_component_indices, 'selection_priority': 'unresolved_local_stitch_first_then_source_order_v1', 'budget_deferred_component_count': len(budget_deferred_indices), 'unsupported_deferred_component_count': len(unsupported_components), 'repaired_component_count': len(accepted), 'worker_batch_call_count': 1 + int(bool(secondary_windows)), 'retry_window_count': len(windows) + len(secondary_windows), 'centered_retry_window_count': len(windows), 'directional_retry_window_count': len(secondary_windows), 'fine_event_count_before': len(result.events), 'retry_event_count_added': retry_event_count_added, 'retry_duplicate_suppressed': suppressed_duplicate_count, 'unresolved_local_stitch_candidate_count': len(local_stitch_candidates), 'unresolved_local_stitch_confirmation_count': local_stitch_confirmation_count, 'combined_event_count': len(ordered_events), 'post_repair_overlap_pair_count': overlap_pairs, 'post_repair_max_overlap_depth': overlap_depth, 'post_repair_exact_duplicate_intervals': exact_duplicates, 'retry_prediction_diagnostics': dict(prediction.diagnostics), **({'directional_retry_prediction_diagnostics': dict(secondary_prediction.diagnostics)} if secondary_prediction is not None else {}), 'components': retry_diagnostics, 'component_failures': deferred_failures}
|
| 696 |
+
remaining_component_indices = [index for index in range(len(result.unresolved)) if index not in accepted_by_component]
|
| 697 |
+
repair_payload['remaining_component_count'] = len(remaining_component_indices)
|
| 698 |
+
repair_payload['partial_success'] = bool(accepted) and bool(remaining_component_indices)
|
| 699 |
+
if remaining_component_indices:
|
| 700 |
+
raise _remaining_adaptive_fine_error(exc, result, repaired_events=ordered_events, accepted_component_indices=tuple(accepted_by_component), local_stitch_candidates=local_stitch_candidates, deferred_failures=deferred_failures, repair_payload=repair_payload, inference_context=inference_context) from exc
|
| 701 |
+
diagnostics = dict(exc.diagnostics)
|
| 702 |
+
diagnostics['pre_repair_violations'] = list(diagnostics.get('violations', []))
|
| 703 |
+
diagnostics['repaired_unresolved_components'] = list(diagnostics.get('unresolved_components', []))
|
| 704 |
+
diagnostics['violations'] = []
|
| 705 |
+
diagnostics['unresolved_components'] = []
|
| 706 |
+
diagnostics['acceptance_failure_classes'] = []
|
| 707 |
+
diagnostics['resolved_event_count'] = len(ordered_events)
|
| 708 |
+
diagnostics['overlap_pair_count'] = overlap_pairs
|
| 709 |
+
diagnostics['max_overlap_depth'] = overlap_depth
|
| 710 |
+
diagnostics['exact_duplicate_intervals'] = exact_duplicates
|
| 711 |
+
diagnostics['adaptive_fine_repair'] = repair_payload
|
| 712 |
+
diagnostics['adaptive_fine_repair_sha256'] = _diagnostic_sha256(repair_payload)
|
| 713 |
+
diagnostics = _with_inference_context(diagnostics, inference_context)
|
| 714 |
+
return events_to_temporal_segments(ordered_events, diagnostics=diagnostics)
|
| 715 |
+
|
| 716 |
+
def _conservative_fine_completion(self, source: AnnotationSource, exc: FineSegmentationOutputError, *, inference_context: Mapping[str, object]) -> tuple[TemporalSegment, ...]:
|
| 717 |
+
result = exc.partial_result
|
| 718 |
+
violations = [dict(item) for item in exc.diagnostics.get('violations', []) if isinstance(item, dict)]
|
| 719 |
+
violation_codes = {str(item.get('code')) for item in violations}
|
| 720 |
+
if result is None or not result.unresolved or (not violation_codes) or (not violation_codes <= {'unresolved_components', 'stitch_ambiguity'}) or (source.duration_s is None):
|
| 721 |
+
raise exc
|
| 722 |
+
combined_events = [replace(event, metadata={'segmentation_provenance': 'fine', 'segmentation_resolution': 'fine', **dict(event.metadata)}) for event in result.events]
|
| 723 |
+
completion_records: list[dict[str, object]] = []
|
| 724 |
+
added_count = 0
|
| 725 |
+
suppressed_duplicate_count = 0
|
| 726 |
+
local_stitch_degraded_count = 0
|
| 727 |
+
local_stitch_candidates = _unresolved_local_stitch_candidates(exc.diagnostics)
|
| 728 |
+
for component_index, component in enumerate(result.unresolved):
|
| 729 |
+
local_candidate = local_stitch_candidates.get(component_index)
|
| 730 |
+
if local_candidate is not None:
|
| 731 |
+
target_event_id = str(local_candidate['resolved_event_id'])
|
| 732 |
+
target_matches = [index for index, event in enumerate(combined_events) if event.event_id == target_event_id]
|
| 733 |
+
if len(target_matches) != 1:
|
| 734 |
+
raise SubtaskAnnotationAdapterError('unresolved-local degraded completion target is missing or ambiguous')
|
| 735 |
+
target_index = target_matches[0]
|
| 736 |
+
matched = combined_events[target_index]
|
| 737 |
+
strict_target_span = _local_stitch_target_span(local_candidate)
|
| 738 |
+
component_span, component_strategy = _conservative_component_completion_span(source, component)
|
| 739 |
+
completed_span = _covering_span(matched.span, component_span)
|
| 740 |
+
record = {'component_index': component_index, 'strategy': 'unconfirmed_unresolved_local_stitch_coverage_union_v2', 'observation_span': _span_payload(_component_observation_span(component)), 'component_completion_span': _span_payload(component_span), 'component_completion_strategy': component_strategy, 'strict_target_span': _span_payload(strict_target_span), 'pre_completion_target_span': _span_payload(matched.span), 'completed_span': _span_payload(completed_span), 'closed_start_ticks': list(component.closed_start_ticks), 'closed_end_ticks': list(component.closed_end_ticks), 'matching_ambiguous': component.matching_ambiguous, 'boundary_claim': 'observation_supported_union_of_strict_target_and_unresolved_component_v2', 'strict_target_geometry_expanded': completed_span != strict_target_span, 'candidate': dict(local_candidate), 'adaptive_fine_confirmation_failures': list(exc.diagnostics.get('adaptive_fine_repair_failures', [])), 'duplicate_suppressed': True}
|
| 741 |
+
metadata = dict(matched.metadata)
|
| 742 |
+
prior_evidence = metadata.get('unresolved_local_stitch_evidence', [])
|
| 743 |
+
evidence = list(prior_evidence) if isinstance(prior_evidence, list) else []
|
| 744 |
+
evidence.append(record)
|
| 745 |
+
metadata.update({'segmentation_provenance': 'fine_with_unconfirmed_unresolved_local_stitch', 'segmentation_resolution': 'fine_degraded', 'review_required': True, 'semantic_quality_established': False, 'unresolved_local_stitch_evidence': evidence})
|
| 746 |
+
combined_events[target_index] = replace(matched, span=completed_span, source_scope_ids=tuple(sorted({*matched.source_scope_ids, *component.source_scope_ids})), metadata=metadata)
|
| 747 |
+
suppressed_duplicate_count += 1
|
| 748 |
+
local_stitch_degraded_count += 1
|
| 749 |
+
completion_records.append(record)
|
| 750 |
+
continue
|
| 751 |
+
span, strategy = _conservative_component_completion_span(source, component)
|
| 752 |
+
record: dict[str, object] = {'component_index': component_index, 'strategy': strategy, 'observation_span': _span_payload(_component_observation_span(component)), 'completed_span': _span_payload(span), 'closed_start_ticks': list(component.closed_start_ticks), 'closed_end_ticks': list(component.closed_end_ticks), 'matching_ambiguous': component.matching_ambiguous, 'boundary_claim': 'observation_supported_conservative_extent_v1'}
|
| 753 |
+
exact_match_index = next((index for index, event in enumerate(combined_events) if event.span == span), None)
|
| 754 |
+
if exact_match_index is not None:
|
| 755 |
+
matched = combined_events[exact_match_index]
|
| 756 |
+
metadata = dict(matched.metadata)
|
| 757 |
+
prior_evidence = metadata.get('conservative_fine_completion_evidence', [])
|
| 758 |
+
evidence = list(prior_evidence) if isinstance(prior_evidence, list) else []
|
| 759 |
+
evidence.append(record)
|
| 760 |
+
metadata.update({'segmentation_provenance': 'fine_with_conservative_completion_evidence', 'segmentation_resolution': 'fine_degraded', 'review_required': True, 'semantic_quality_established': False, 'conservative_fine_completion_evidence': evidence})
|
| 761 |
+
combined_events[exact_match_index] = replace(matched, source_scope_ids=tuple(sorted({*matched.source_scope_ids, *component.source_scope_ids})), metadata=metadata)
|
| 762 |
+
record['duplicate_suppressed'] = True
|
| 763 |
+
suppressed_duplicate_count += 1
|
| 764 |
+
else:
|
| 765 |
+
record['duplicate_suppressed'] = False
|
| 766 |
+
combined_events.append(EventInterval(event_id=f'{source.source_id}:conservative-fine-completion:{component_index:06d}', source_id=source.source_id, span=span, source_scope_ids=tuple(sorted(component.source_scope_ids)), metadata={'segmentation_provenance': 'fine_conservative_completion', 'segmentation_resolution': 'fine_degraded', 'review_required': True, 'semantic_quality_established': False, 'boundary_claim': 'observation_supported_conservative_extent_v1', 'conservative_fine_completion': record}))
|
| 767 |
+
added_count += 1
|
| 768 |
+
completion_records.append(record)
|
| 769 |
+
ordered_events = tuple(sorted(combined_events, key=lambda item: (item.span, item.event_id)))
|
| 770 |
+
if not ordered_events:
|
| 771 |
+
raise exc
|
| 772 |
+
overlap_pairs, overlap_depth = _overlap_statistics(ordered_events)
|
| 773 |
+
exact_duplicates = _exact_duplicate_intervals(ordered_events)
|
| 774 |
+
if exact_duplicates:
|
| 775 |
+
raise SubtaskAnnotationAdapterError('conservative fine completion failed to suppress exact duplicate intervals')
|
| 776 |
+
completion_payload: dict[str, object] = {'schema_version': CONSERVATIVE_FINE_COMPLETION_SCHEMA, 'model_io_contract': fine_window_profile(self.fine_segmenter.window_profile_id).model_io_contract, 'window_profile_id': self.fine_segmenter.window_profile_id, 'terminal_contract': 'parsed_fine_evidence_only_no_coarse_provider_v1', 'boundary_claim': 'observation_supported_conservative_extent_v1', 'coarse_provider_called': False, 'additional_model_call_count': 0, 'max_retry_windows': self.adaptive_fine_retry_max_windows, 'fine_event_count_before': len(result.events), 'unresolved_component_count': len(result.unresolved), 'completion_event_count_added': added_count, 'completion_duplicate_suppressed': suppressed_duplicate_count, 'unresolved_local_stitch_candidate_count': len(local_stitch_candidates), 'unresolved_local_stitch_degraded_count': local_stitch_degraded_count, 'combined_event_count': len(ordered_events), 'post_completion_overlap_pair_count': overlap_pairs, 'post_completion_max_overlap_depth': overlap_depth, 'post_completion_exact_duplicate_intervals': exact_duplicates, 'components': completion_records}
|
| 777 |
+
diagnostics = dict(exc.diagnostics)
|
| 778 |
+
diagnostics['pre_completion_violations'] = violations
|
| 779 |
+
diagnostics['pre_completion_unresolved_components'] = list(diagnostics.get('unresolved_components', []))
|
| 780 |
+
diagnostics['violations'] = []
|
| 781 |
+
diagnostics['unresolved_components'] = []
|
| 782 |
+
diagnostics['acceptance_failure_classes'] = []
|
| 783 |
+
diagnostics['resolved_event_count'] = len(ordered_events)
|
| 784 |
+
diagnostics['overlap_pair_count'] = overlap_pairs
|
| 785 |
+
diagnostics['max_overlap_depth'] = overlap_depth
|
| 786 |
+
diagnostics['exact_duplicate_intervals'] = exact_duplicates
|
| 787 |
+
diagnostics['output_degraded'] = True
|
| 788 |
+
diagnostics['semantic_quality_established'] = False
|
| 789 |
+
diagnostics['conservative_fine_completion'] = completion_payload
|
| 790 |
+
diagnostics['conservative_fine_completion_sha256'] = _diagnostic_sha256(completion_payload)
|
| 791 |
+
diagnostics = _with_inference_context(diagnostics, inference_context)
|
| 792 |
+
return events_to_temporal_segments(ordered_events, diagnostics=diagnostics)
|
| 793 |
+
|
| 794 |
+
@staticmethod
|
| 795 |
+
def _bind_physical_source_envelope(source: AnnotationSource, segments: Sequence[TemporalSegment]) -> tuple[TemporalSegment, ...]:
|
| 796 |
+
if source.duration_s is None:
|
| 797 |
+
raise SubtaskAnnotationAdapterError('hierarchical fine segmentation requires exact source duration_s')
|
| 798 |
+
duration_us = seconds_to_us(source.duration_s, label='duration_s')
|
| 799 |
+
terminal_tick_us = ceil_div(duration_us, TICK_US) * TICK_US
|
| 800 |
+
bounded: list[TemporalSegment] = []
|
| 801 |
+
for segment in segments:
|
| 802 |
+
if segment.span.end_s <= source.duration_s + 1e-06:
|
| 803 |
+
bounded.append(segment)
|
| 804 |
+
continue
|
| 805 |
+
model_end_us = round(segment.span.end_s * 1000000)
|
| 806 |
+
if not math.isclose(segment.span.end_s, model_end_us / 1000000, abs_tol=1e-09) or model_end_us != terminal_tick_us or segment.span.start_s >= source.duration_s:
|
| 807 |
+
raise SubtaskAnnotationAdapterError(f'hierarchical segment for {source.source_id} ends at {segment.span.end_s}s beyond the physical source duration {source.duration_s}s and is not the terminal model tick')
|
| 808 |
+
metadata = dict(segment.metadata)
|
| 809 |
+
metadata['terminal_source_clip'] = {'contract': 'terminal_20hz_tick_to_physical_source_duration_v1', 'model_end_us': model_end_us, 'physical_end_us': duration_us}
|
| 810 |
+
bounded.append(TemporalSegment(TimeSpan(segment.span.start_s, source.duration_s), metadata))
|
| 811 |
+
ordered = tuple(sorted(bounded, key=lambda item: (item.span, repr(item.metadata))))
|
| 812 |
+
output: list[TemporalSegment] = []
|
| 813 |
+
for index, segment in enumerate(ordered):
|
| 814 |
+
before, after = _nearest_nonoverlapping_context(ordered, index)
|
| 815 |
+
metadata = dict(segment.metadata)
|
| 816 |
+
metadata['u2_context'] = {'before': _span_payload(before), 'target': _span_payload(segment.span), 'after': _span_payload(after)}
|
| 817 |
+
output.append(TemporalSegment(segment.span, metadata))
|
| 818 |
+
return tuple(output)
|
| 819 |
+
|
| 820 |
+
@staticmethod
|
| 821 |
+
def _with_refreshed_u2_context(segments: Sequence[TemporalSegment]) -> tuple[TemporalSegment, ...]:
|
| 822 |
+
ordered = tuple(sorted(segments, key=lambda item: (item.span, repr(item.metadata))))
|
| 823 |
+
output: list[TemporalSegment] = []
|
| 824 |
+
for index, segment in enumerate(ordered):
|
| 825 |
+
before, after = _nearest_nonoverlapping_context(ordered, index)
|
| 826 |
+
metadata = dict(segment.metadata)
|
| 827 |
+
metadata['u2_context'] = {'before': _span_payload(before), 'target': _span_payload(segment.span), 'after': _span_payload(after)}
|
| 828 |
+
output.append(TemporalSegment(segment.span, metadata))
|
| 829 |
+
return tuple(output)
|
| 830 |
+
|
| 831 |
+
def segment(self, source: AnnotationSource) -> tuple[TemporalSegment, ...]:
|
| 832 |
+
return asyncio.run(self.segment_async(source))
|
| 833 |
+
|
| 834 |
+
async def segment_async(self, source):
|
| 835 |
+
return self._bind_physical_source_envelope(source, await self._segment_impl(source))
|
| 836 |
+
|
| 837 |
+
async def _segment_impl(self, source):
|
| 838 |
+
inference_context = _inference_context(source, (), fine_prior_mode='none')
|
| 839 |
+
segment_with_priors_async = getattr(self.fine_segmenter, 'segment_with_priors_async', None)
|
| 840 |
+
try:
|
| 841 |
+
if callable(segment_with_priors_async):
|
| 842 |
+
result = await segment_with_priors_async(source, (), arm='no_prior_global_scan')
|
| 843 |
+
else:
|
| 844 |
+
result = await asyncio.to_thread(self.fine_segmenter.segment_with_priors, source, (), arm='no_prior_global_scan')
|
| 845 |
+
except FineSegmentationOutputError as exc:
|
| 846 |
+
contextual_error = FineSegmentationOutputError(str(exc), diagnostics=_with_inference_context(exc.diagnostics, inference_context), partial_result=exc.partial_result)
|
| 847 |
+
try:
|
| 848 |
+
return await self.adaptive_fine_repair_async(source, contextual_error, inference_context=inference_context)
|
| 849 |
+
except FineSegmentationOutputError as adaptive_error:
|
| 850 |
+
return self._conservative_fine_completion(source, adaptive_error, inference_context=inference_context)
|
| 851 |
+
diagnostics = _with_inference_context(result.diagnostics, inference_context)
|
| 852 |
+
return events_to_temporal_segments(result.events, diagnostics=diagnostics)
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_temporal.py
ADDED
|
@@ -0,0 +1,987 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from collections.abc import Collection, Iterable, Mapping, Sequence
|
| 5 |
+
from dataclasses import dataclass, field
|
| 6 |
+
from fractions import Fraction
|
| 7 |
+
import json
|
| 8 |
+
import math
|
| 9 |
+
from flora_runtime.schema import TimeSpan
|
| 10 |
+
from flora_runtime.video_segmentation.contracts import EventInterval, ProcessingScope
|
| 11 |
+
TICK_US = 50000
|
| 12 |
+
CORE_US = 6000000
|
| 13 |
+
LEFT_GUARD_US = 1000000
|
| 14 |
+
RIGHT_GUARD_US = 1000000
|
| 15 |
+
SLOT_COUNT = 64
|
| 16 |
+
SLOT_FPS = 8
|
| 17 |
+
LEGACY_FINE_WINDOW_PROFILE_ID = 'atomic_1plus6plus1_64slots_8hz'
|
| 18 |
+
HUMAN_FINE_WINDOW_PROFILE_ID = 'human_16s_4hz'
|
| 19 |
+
YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID = 'yam_atomic_2plus12plus2_64slots_4hz'
|
| 20 |
+
YAM_INTENT_32S_FINE_WINDOW_PROFILE_ID = 'yam_intent_4plus24plus4_128slots_4hz'
|
| 21 |
+
DEFAULT_FINE_WINDOW_PROFILE_ID = LEGACY_FINE_WINDOW_PROFILE_ID
|
| 22 |
+
MAX_SEGMENTS = 32
|
| 23 |
+
STITCH_TOLERANCE_TICKS = 2
|
| 24 |
+
MAX_BOUNDED_ANCHOR_TICKS = 4
|
| 25 |
+
STRICT_TRACE_STITCHING_CONTRACT = 'strict_shared_trace_endpoint_v1'
|
| 26 |
+
SHARED_TRACE_IOU_STITCHING_CONTRACT = 'shared_trace_iou_endpoint_v1'
|
| 27 |
+
COMPONENT_LOCAL_IDENTITY_STITCHING_CONTRACT = 'component_local_identity_anchor_v1'
|
| 28 |
+
COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT = 'component_local_bounded_anchor_certificate_v1'
|
| 29 |
+
COMPONENT_LOCAL_STITCHING_CONTRACTS = frozenset({COMPONENT_LOCAL_IDENTITY_STITCHING_CONTRACT, COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT})
|
| 30 |
+
STITCHING_CONTRACTS = {STRICT_TRACE_STITCHING_CONTRACT, SHARED_TRACE_IOU_STITCHING_CONTRACT, *COMPONENT_LOCAL_STITCHING_CONTRACTS}
|
| 31 |
+
STITCHING_SEAM_DIAGNOSTICS_SCHEMA = 'data_engine_fine_stitching_seam_diagnostics_v1'
|
| 32 |
+
WINDOW_SCHEMA = 'data_engine_atomic_fine_window_v1'
|
| 33 |
+
MODEL_IO_CONTRACT = 'flora_inference_atomic_temporal_segmentation_v1'
|
| 34 |
+
HUMAN_MODEL_IO_CONTRACT = 'flora_human_temporal_segmentation_v1'
|
| 35 |
+
YAM_ATOMIC_LONG_MODEL_IO_CONTRACT = 'data_engine_yam_atomic_temporal_segmentation_16s4hz_v1'
|
| 36 |
+
YAM_INTENT_32S_MODEL_IO_CONTRACT = 'data_engine_yam_intent_temporal_segmentation_32s4hz_v1'
|
| 37 |
+
ComponentPairKey = tuple[int, int, int]
|
| 38 |
+
|
| 39 |
+
class FineTemporalError(ValueError):
|
| 40 |
+
pass
|
| 41 |
+
|
| 42 |
+
class FinePredictionParseError(FineTemporalError):
|
| 43 |
+
pass
|
| 44 |
+
|
| 45 |
+
@dataclass(frozen=True)
|
| 46 |
+
class FineWindowProfile:
|
| 47 |
+
profile_id: str
|
| 48 |
+
left_guard_us: int
|
| 49 |
+
core_us: int
|
| 50 |
+
right_guard_us: int
|
| 51 |
+
slot_count: int
|
| 52 |
+
slot_fps: int
|
| 53 |
+
model_io_contract: str
|
| 54 |
+
target_granularity: str
|
| 55 |
+
|
| 56 |
+
def __post_init__(self) -> None:
|
| 57 |
+
if not self.profile_id or not self.model_io_contract:
|
| 58 |
+
raise FineTemporalError('fine window profile identity and model contract are required')
|
| 59 |
+
if self.target_granularity not in {'atomic_event', 'coarse_subtask_run'}:
|
| 60 |
+
raise FineTemporalError('fine window target granularity is unsupported')
|
| 61 |
+
values = (self.left_guard_us, self.core_us, self.right_guard_us, self.slot_count, self.slot_fps)
|
| 62 |
+
if any((isinstance(value, bool) or not isinstance(value, int) or value <= 0 for value in values)):
|
| 63 |
+
raise FineTemporalError('fine window geometry and sampling values must be positive integers')
|
| 64 |
+
if self.view_us % TICK_US or self.core_us % TICK_US:
|
| 65 |
+
raise FineTemporalError('fine window geometry must align to the 20 Hz model clock')
|
| 66 |
+
if self.left_guard_us != self.right_guard_us:
|
| 67 |
+
raise FineTemporalError('fine window profiles currently require symmetric guards')
|
| 68 |
+
if 1000000 % self.slot_fps:
|
| 69 |
+
raise FineTemporalError('fine window sample rate must divide one second exactly')
|
| 70 |
+
if self.view_us * self.slot_fps != self.slot_count * 1000000:
|
| 71 |
+
raise FineTemporalError('fine window slots, duration, and sample rate disagree')
|
| 72 |
+
|
| 73 |
+
@property
|
| 74 |
+
def view_us(self) -> int:
|
| 75 |
+
return self.left_guard_us + self.core_us + self.right_guard_us
|
| 76 |
+
|
| 77 |
+
@property
|
| 78 |
+
def view_ticks(self) -> int:
|
| 79 |
+
return self.view_us // TICK_US
|
| 80 |
+
|
| 81 |
+
@property
|
| 82 |
+
def core_ticks(self) -> int:
|
| 83 |
+
return self.core_us // TICK_US
|
| 84 |
+
|
| 85 |
+
@property
|
| 86 |
+
def slot_spacing_us(self) -> int:
|
| 87 |
+
return 1000000 // self.slot_fps
|
| 88 |
+
|
| 89 |
+
@property
|
| 90 |
+
def scheduler_contract(self) -> str:
|
| 91 |
+
core = self._duration_token(self.core_us)
|
| 92 |
+
guard = self._duration_token(self.left_guard_us)
|
| 93 |
+
return f'fixed_{core}_core_{guard}_guards_global_grid_v1'
|
| 94 |
+
|
| 95 |
+
def retry_scheduler_contract(self, *, directional: bool) -> str:
|
| 96 |
+
prefix = 'adaptive_directional' if directional else 'adaptive_recentered'
|
| 97 |
+
core = self._duration_token(self.core_us)
|
| 98 |
+
guard = self._duration_token(self.left_guard_us)
|
| 99 |
+
geometry = f'{core}_core_{guard}_guards'
|
| 100 |
+
return f'{prefix}_{geometry}_v1'
|
| 101 |
+
|
| 102 |
+
@staticmethod
|
| 103 |
+
def _duration_token(value_us: int) -> str:
|
| 104 |
+
return f'{value_us // 1000000}s' if value_us % 1000000 == 0 else f'{value_us}us'
|
| 105 |
+
|
| 106 |
+
def to_dict(self) -> dict[str, int | str]:
|
| 107 |
+
return {'profile_id': self.profile_id, 'left_guard_us': self.left_guard_us, 'core_us': self.core_us, 'right_guard_us': self.right_guard_us, 'view_us': self.view_us, 'slot_count': self.slot_count, 'slot_fps': self.slot_fps, 'model_io_contract': self.model_io_contract, 'target_granularity': self.target_granularity}
|
| 108 |
+
FINE_WINDOW_PROFILES: Mapping[str, FineWindowProfile] = {LEGACY_FINE_WINDOW_PROFILE_ID: FineWindowProfile(profile_id=LEGACY_FINE_WINDOW_PROFILE_ID, left_guard_us=LEFT_GUARD_US, core_us=CORE_US, right_guard_us=RIGHT_GUARD_US, slot_count=SLOT_COUNT, slot_fps=SLOT_FPS, model_io_contract=MODEL_IO_CONTRACT, target_granularity='atomic_event'), HUMAN_FINE_WINDOW_PROFILE_ID: FineWindowProfile(profile_id=HUMAN_FINE_WINDOW_PROFILE_ID, left_guard_us=2000000, core_us=12000000, right_guard_us=2000000, slot_count=64, slot_fps=4, model_io_contract=HUMAN_MODEL_IO_CONTRACT, target_granularity='coarse_subtask_run'), YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID: FineWindowProfile(profile_id=YAM_ATOMIC_LONG_FINE_WINDOW_PROFILE_ID, left_guard_us=2000000, core_us=12000000, right_guard_us=2000000, slot_count=64, slot_fps=4, model_io_contract=YAM_ATOMIC_LONG_MODEL_IO_CONTRACT, target_granularity='atomic_event'), YAM_INTENT_32S_FINE_WINDOW_PROFILE_ID: FineWindowProfile(profile_id=YAM_INTENT_32S_FINE_WINDOW_PROFILE_ID, left_guard_us=4000000, core_us=24000000, right_guard_us=4000000, slot_count=128, slot_fps=4, model_io_contract=YAM_INTENT_32S_MODEL_IO_CONTRACT, target_granularity='atomic_event')}
|
| 109 |
+
|
| 110 |
+
def fine_window_profile(profile_id: str=DEFAULT_FINE_WINDOW_PROFILE_ID) -> FineWindowProfile:
|
| 111 |
+
try:
|
| 112 |
+
return FINE_WINDOW_PROFILES[profile_id]
|
| 113 |
+
except (KeyError, TypeError) as exc:
|
| 114 |
+
raise FineTemporalError(f'unknown fine window profile: {profile_id!r}') from exc
|
| 115 |
+
|
| 116 |
+
@dataclass(frozen=True)
|
| 117 |
+
class StitchingPolicy:
|
| 118 |
+
contract: str = STRICT_TRACE_STITCHING_CONTRACT
|
| 119 |
+
endpoint_tolerance_ticks: int = STITCH_TOLERANCE_TICKS
|
| 120 |
+
min_shared_trace_iou_numerator: int = 1
|
| 121 |
+
min_shared_trace_iou_denominator: int = 2
|
| 122 |
+
|
| 123 |
+
def __post_init__(self) -> None:
|
| 124 |
+
if self.contract not in STITCHING_CONTRACTS:
|
| 125 |
+
raise FineTemporalError(f'unknown stitching contract: {self.contract}')
|
| 126 |
+
integer_fields = ('endpoint_tolerance_ticks', 'min_shared_trace_iou_numerator', 'min_shared_trace_iou_denominator')
|
| 127 |
+
for field_name in integer_fields:
|
| 128 |
+
value = getattr(self, field_name)
|
| 129 |
+
if isinstance(value, bool) or not isinstance(value, int):
|
| 130 |
+
raise FineTemporalError(f'{field_name} must be an integer')
|
| 131 |
+
if self.endpoint_tolerance_ticks < 0:
|
| 132 |
+
raise FineTemporalError('endpoint_tolerance_ticks must be non-negative')
|
| 133 |
+
if self.contract == COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT and (not STITCH_TOLERANCE_TICKS < self.endpoint_tolerance_ticks <= MAX_BOUNDED_ANCHOR_TICKS):
|
| 134 |
+
raise FineTemporalError('bounded-anchor tolerance must be above strict and at most four ticks')
|
| 135 |
+
if not 0 < self.min_shared_trace_iou_numerator <= self.min_shared_trace_iou_denominator:
|
| 136 |
+
raise FineTemporalError('shared-trace IoU threshold must lie in (0,1]')
|
| 137 |
+
|
| 138 |
+
@classmethod
|
| 139 |
+
def strict(cls, tolerance_ticks: int=STITCH_TOLERANCE_TICKS) -> StitchingPolicy:
|
| 140 |
+
return cls(contract=STRICT_TRACE_STITCHING_CONTRACT, endpoint_tolerance_ticks=tolerance_ticks)
|
| 141 |
+
|
| 142 |
+
@classmethod
|
| 143 |
+
def shared_trace_iou(cls, *, endpoint_tolerance_ticks: int, numerator: int=1, denominator: int=2) -> StitchingPolicy:
|
| 144 |
+
return cls(contract=SHARED_TRACE_IOU_STITCHING_CONTRACT, endpoint_tolerance_ticks=endpoint_tolerance_ticks, min_shared_trace_iou_numerator=numerator, min_shared_trace_iou_denominator=denominator)
|
| 145 |
+
|
| 146 |
+
@classmethod
|
| 147 |
+
def component_local_identity(cls, *, endpoint_tolerance_ticks: int=4, numerator: int=3, denominator: int=4) -> StitchingPolicy:
|
| 148 |
+
if endpoint_tolerance_ticks <= STITCH_TOLERANCE_TICKS:
|
| 149 |
+
raise FineTemporalError('component-local candidate tolerance must exceed the strict tolerance')
|
| 150 |
+
return cls(contract=COMPONENT_LOCAL_IDENTITY_STITCHING_CONTRACT, endpoint_tolerance_ticks=endpoint_tolerance_ticks, min_shared_trace_iou_numerator=numerator, min_shared_trace_iou_denominator=denominator)
|
| 151 |
+
|
| 152 |
+
@classmethod
|
| 153 |
+
def component_local_bounded_anchor(cls, *, endpoint_tolerance_ticks: int=4, numerator: int=3, denominator: int=4) -> StitchingPolicy:
|
| 154 |
+
if not STITCH_TOLERANCE_TICKS < endpoint_tolerance_ticks <= MAX_BOUNDED_ANCHOR_TICKS:
|
| 155 |
+
raise FineTemporalError('bounded-anchor candidate tolerance must be above strict and at most four ticks')
|
| 156 |
+
return cls(contract=COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT, endpoint_tolerance_ticks=endpoint_tolerance_ticks, min_shared_trace_iou_numerator=numerator, min_shared_trace_iou_denominator=denominator)
|
| 157 |
+
|
| 158 |
+
def to_dict(self) -> dict[str, int | str]:
|
| 159 |
+
return {'contract': self.contract, 'endpoint_tolerance_ticks': self.endpoint_tolerance_ticks, 'min_shared_trace_iou_numerator': self.min_shared_trace_iou_numerator, 'min_shared_trace_iou_denominator': self.min_shared_trace_iou_denominator}
|
| 160 |
+
|
| 161 |
+
def floor_div(value: int, divisor: int) -> int:
|
| 162 |
+
if divisor <= 0:
|
| 163 |
+
raise FineTemporalError('divisor must be positive')
|
| 164 |
+
return value // divisor
|
| 165 |
+
|
| 166 |
+
def ceil_div(value: int, divisor: int) -> int:
|
| 167 |
+
if divisor <= 0:
|
| 168 |
+
raise FineTemporalError('divisor must be positive')
|
| 169 |
+
return -(-value // divisor)
|
| 170 |
+
|
| 171 |
+
def round_half_up_div(value: int, divisor: int) -> int:
|
| 172 |
+
if divisor <= 0 or value < 0:
|
| 173 |
+
raise FineTemporalError('round_half_up_div requires non-negative input and positive divisor')
|
| 174 |
+
quotient, remainder = divmod(value, divisor)
|
| 175 |
+
return quotient + int(remainder * 2 >= divisor)
|
| 176 |
+
|
| 177 |
+
@dataclass(frozen=True)
|
| 178 |
+
class FineWindow:
|
| 179 |
+
source_id: str
|
| 180 |
+
window_index: int
|
| 181 |
+
video_end_us: int
|
| 182 |
+
requested_view_start_us: int
|
| 183 |
+
requested_view_end_us: int
|
| 184 |
+
valid_start_us: int
|
| 185 |
+
valid_end_us: int
|
| 186 |
+
core_start_us: int
|
| 187 |
+
core_end_us: int
|
| 188 |
+
valid_start_tick: int
|
| 189 |
+
valid_end_tick: int
|
| 190 |
+
core_start_tick: int
|
| 191 |
+
core_end_tick: int
|
| 192 |
+
scope_id: str
|
| 193 |
+
window_profile_id: str = DEFAULT_FINE_WINDOW_PROFILE_ID
|
| 194 |
+
coarse_spans: tuple[Mapping[str, object], ...] = ()
|
| 195 |
+
|
| 196 |
+
def __post_init__(self) -> None:
|
| 197 |
+
if not self.source_id or not self.scope_id or self.window_index < 0 or (self.video_end_us <= 0):
|
| 198 |
+
raise FineTemporalError('window identity, index, and duration are required')
|
| 199 |
+
profile = fine_window_profile(self.window_profile_id)
|
| 200 |
+
if self.requested_view_end_us - self.requested_view_start_us != profile.view_us:
|
| 201 |
+
raise FineTemporalError(f'requested view must match profile {profile.profile_id!r}')
|
| 202 |
+
if not self.requested_view_start_us <= self.valid_start_us < self.valid_end_us <= self.requested_view_end_us:
|
| 203 |
+
raise FineTemporalError('invalid visible-video envelope')
|
| 204 |
+
if not 0 <= self.valid_start_us < self.valid_end_us <= self.video_end_us:
|
| 205 |
+
raise FineTemporalError('valid envelope must lie inside the video')
|
| 206 |
+
if not self.valid_start_us <= self.core_start_us < self.core_end_us <= self.valid_end_us:
|
| 207 |
+
raise FineTemporalError('emit core must lie inside the valid envelope')
|
| 208 |
+
if not 0 <= self.valid_start_tick < self.valid_end_tick <= profile.view_ticks:
|
| 209 |
+
raise FineTemporalError('invalid visible tick envelope')
|
| 210 |
+
if not self.valid_start_tick <= self.core_start_tick < self.core_end_tick <= self.valid_end_tick:
|
| 211 |
+
raise FineTemporalError('invalid core tick envelope')
|
| 212 |
+
coarse_spans = tuple((dict(item) for item in self.coarse_spans))
|
| 213 |
+
if len(coarse_spans) > 64:
|
| 214 |
+
raise FineTemporalError('one view may contain at most 64 coarse spans')
|
| 215 |
+
previous: tuple[int, int, bool, bool] | None = None
|
| 216 |
+
for index, row in enumerate(coarse_spans):
|
| 217 |
+
if set(row) != {'s', 'e', 'lo', 'ro'}:
|
| 218 |
+
raise FineTemporalError(f'coarse span {index} must contain exactly s,e,lo,ro')
|
| 219 |
+
fragment = FineFragment(row['s'], row['e'], row['lo'], row['ro'])
|
| 220 |
+
if isinstance(fragment.s, bool) or isinstance(fragment.e, bool):
|
| 221 |
+
raise FineTemporalError(f'coarse span {index} ticks must be integers')
|
| 222 |
+
if not isinstance(fragment.s, int) or not isinstance(fragment.e, int):
|
| 223 |
+
raise FineTemporalError(f'coarse span {index} ticks must be integers')
|
| 224 |
+
if type(fragment.lo) is not bool or type(fragment.ro) is not bool:
|
| 225 |
+
raise FineTemporalError(f'coarse span {index} open flags must be booleans')
|
| 226 |
+
if not self.valid_start_tick <= fragment.s < fragment.e <= self.valid_end_tick:
|
| 227 |
+
raise FineTemporalError(f'coarse span {index} lies outside the valid envelope')
|
| 228 |
+
if fragment.lo and fragment.s != self.valid_start_tick:
|
| 229 |
+
raise FineTemporalError(f'coarse span {index} lo=true is not clipped to valid start')
|
| 230 |
+
if fragment.ro and fragment.e != self.valid_end_tick:
|
| 231 |
+
raise FineTemporalError(f'coarse span {index} ro=true is not clipped to valid end')
|
| 232 |
+
if previous is not None and fragment.sort_key <= previous:
|
| 233 |
+
raise FineTemporalError('coarse spans must be unique and in canonical order')
|
| 234 |
+
previous = fragment.sort_key
|
| 235 |
+
object.__setattr__(self, 'coarse_spans', coarse_spans)
|
| 236 |
+
|
| 237 |
+
@property
|
| 238 |
+
def global_tick_offset(self) -> int:
|
| 239 |
+
if self.requested_view_start_us % TICK_US:
|
| 240 |
+
raise FineTemporalError('requested view is not aligned to the 20 Hz model clock')
|
| 241 |
+
return self.requested_view_start_us // TICK_US
|
| 242 |
+
|
| 243 |
+
@property
|
| 244 |
+
def window_profile(self) -> FineWindowProfile:
|
| 245 |
+
return fine_window_profile(self.window_profile_id)
|
| 246 |
+
|
| 247 |
+
def worker_payload(self) -> dict[str, object]:
|
| 248 |
+
return {'schema_version': WINDOW_SCHEMA, 'model_io_contract': self.window_profile.model_io_contract, 'window_profile_id': self.window_profile_id, 'scope_id': self.scope_id, 'window_index': self.window_index, 'requested_view_start_us': self.requested_view_start_us, 'requested_view_end_us': self.requested_view_end_us, 'valid_start_tick': self.valid_start_tick, 'valid_end_tick': self.valid_end_tick, 'core_start_tick': self.core_start_tick, 'core_end_tick': self.core_end_tick, 'coarse_spans': [dict(item) for item in self.coarse_spans]}
|
| 249 |
+
|
| 250 |
+
@dataclass(frozen=True)
|
| 251 |
+
class FineFragment:
|
| 252 |
+
s: int
|
| 253 |
+
e: int
|
| 254 |
+
lo: bool
|
| 255 |
+
ro: bool
|
| 256 |
+
|
| 257 |
+
@property
|
| 258 |
+
def sort_key(self) -> tuple[int, int, bool, bool]:
|
| 259 |
+
return (self.s, self.e, self.lo, self.ro)
|
| 260 |
+
|
| 261 |
+
def _reject_constant(value: str) -> None:
|
| 262 |
+
raise FinePredictionParseError(f'non-finite JSON constant is forbidden: {value}')
|
| 263 |
+
|
| 264 |
+
def _object_without_duplicate_keys(pairs: list[tuple[str, object]]) -> dict[str, object]:
|
| 265 |
+
result: dict[str, object] = {}
|
| 266 |
+
for key, value in pairs:
|
| 267 |
+
if key in result:
|
| 268 |
+
raise FinePredictionParseError(f'duplicate JSON key is forbidden: {key}')
|
| 269 |
+
result[key] = value
|
| 270 |
+
return result
|
| 271 |
+
|
| 272 |
+
def _validate_fragment_envelope(fragment: FineFragment, window: FineWindow) -> None:
|
| 273 |
+
if isinstance(fragment.s, bool) or isinstance(fragment.e, bool):
|
| 274 |
+
raise FinePredictionParseError('segment ticks must be integers, not booleans')
|
| 275 |
+
if not isinstance(fragment.s, int) or not isinstance(fragment.e, int):
|
| 276 |
+
raise FinePredictionParseError('segment ticks must be integers')
|
| 277 |
+
if type(fragment.lo) is not bool or type(fragment.ro) is not bool:
|
| 278 |
+
raise FinePredictionParseError('open flags must be JSON booleans')
|
| 279 |
+
if not window.valid_start_tick <= fragment.s < fragment.e <= window.valid_end_tick <= window.window_profile.view_ticks:
|
| 280 |
+
raise FinePredictionParseError('segment lies outside the valid tick envelope')
|
| 281 |
+
if fragment.lo and fragment.s != window.valid_start_tick:
|
| 282 |
+
raise FinePredictionParseError('lo=true requires clipping to valid_start_tick')
|
| 283 |
+
if fragment.ro and fragment.e != window.valid_end_tick:
|
| 284 |
+
raise FinePredictionParseError('ro=true requires clipping to valid_end_tick')
|
| 285 |
+
|
| 286 |
+
def _has_emit_core_support(fragment: FineFragment, window: FineWindow) -> bool:
|
| 287 |
+
core_intersects = fragment.s < window.core_end_tick and fragment.e > window.core_start_tick
|
| 288 |
+
core_touches = fragment.e == window.core_start_tick or fragment.s == window.core_end_tick
|
| 289 |
+
return core_intersects or core_touches
|
| 290 |
+
|
| 291 |
+
def _validate_fragment(fragment: FineFragment, window: FineWindow) -> None:
|
| 292 |
+
_validate_fragment_envelope(fragment, window)
|
| 293 |
+
if not _has_emit_core_support(fragment, window):
|
| 294 |
+
raise FinePredictionParseError('segment has no emit-core support')
|
| 295 |
+
|
| 296 |
+
@dataclass(frozen=True)
|
| 297 |
+
class ParsedFinePrediction:
|
| 298 |
+
fragments: tuple[FineFragment, ...]
|
| 299 |
+
guard_only_fragments: tuple[FineFragment, ...]
|
| 300 |
+
|
| 301 |
+
def parse_prediction_with_diagnostics(payload: str, window: FineWindow, *, drop_guard_only: bool) -> ParsedFinePrediction:
|
| 302 |
+
if not isinstance(payload, str) or not payload:
|
| 303 |
+
raise FinePredictionParseError('prediction_json must be a non-empty string')
|
| 304 |
+
try:
|
| 305 |
+
value = json.loads(payload, object_pairs_hook=_object_without_duplicate_keys, parse_constant=_reject_constant)
|
| 306 |
+
except FinePredictionParseError:
|
| 307 |
+
raise
|
| 308 |
+
except (json.JSONDecodeError, TypeError, ValueError) as exc:
|
| 309 |
+
raise FinePredictionParseError('prediction is not one complete JSON object') from exc
|
| 310 |
+
if not isinstance(value, dict) or list(value) != ['segments']:
|
| 311 |
+
raise FinePredictionParseError('prediction must have exactly one top-level key: segments')
|
| 312 |
+
rows = value['segments']
|
| 313 |
+
if not isinstance(rows, list):
|
| 314 |
+
raise FinePredictionParseError('segments must be a list')
|
| 315 |
+
if len(rows) > MAX_SEGMENTS:
|
| 316 |
+
raise FinePredictionParseError(f'prediction contains more than {MAX_SEGMENTS} segments')
|
| 317 |
+
result: list[FineFragment] = []
|
| 318 |
+
guard_only: list[FineFragment] = []
|
| 319 |
+
previous: tuple[int, int, bool, bool] | None = None
|
| 320 |
+
for index, row in enumerate(rows):
|
| 321 |
+
if not isinstance(row, dict) or list(row) != ['s', 'e', 'lo', 'ro']:
|
| 322 |
+
raise FinePredictionParseError(f'segment {index}: keys must appear exactly as s,e,lo,ro')
|
| 323 |
+
fragment = FineFragment(row['s'], row['e'], row['lo'], row['ro'])
|
| 324 |
+
try:
|
| 325 |
+
_validate_fragment_envelope(fragment, window)
|
| 326 |
+
except FinePredictionParseError as exc:
|
| 327 |
+
raise FinePredictionParseError(f'segment {index}: {exc}') from exc
|
| 328 |
+
if previous is not None and fragment.sort_key < previous:
|
| 329 |
+
raise FinePredictionParseError('segments are not in canonical order')
|
| 330 |
+
previous = fragment.sort_key
|
| 331 |
+
if not _has_emit_core_support(fragment, window):
|
| 332 |
+
if not drop_guard_only:
|
| 333 |
+
raise FinePredictionParseError(f'segment {index}: segment has no emit-core support')
|
| 334 |
+
guard_only.append(fragment)
|
| 335 |
+
continue
|
| 336 |
+
result.append(fragment)
|
| 337 |
+
return ParsedFinePrediction(tuple(result), tuple(guard_only))
|
| 338 |
+
|
| 339 |
+
@dataclass(frozen=True)
|
| 340 |
+
class _Observation:
|
| 341 |
+
observation_id: int
|
| 342 |
+
window_index: int
|
| 343 |
+
row_ordinal: int
|
| 344 |
+
scope_id: str
|
| 345 |
+
start_tick: int
|
| 346 |
+
end_tick: int
|
| 347 |
+
lo: bool
|
| 348 |
+
ro: bool
|
| 349 |
+
valid_start_tick: int
|
| 350 |
+
valid_end_tick: int
|
| 351 |
+
|
| 352 |
+
@property
|
| 353 |
+
def has_closed_start(self) -> bool:
|
| 354 |
+
return not self.lo
|
| 355 |
+
|
| 356 |
+
@property
|
| 357 |
+
def has_closed_end(self) -> bool:
|
| 358 |
+
return not self.ro
|
| 359 |
+
|
| 360 |
+
@dataclass(frozen=True)
|
| 361 |
+
class UnresolvedFineEvent:
|
| 362 |
+
observation_ids: tuple[int, ...]
|
| 363 |
+
source_scope_ids: tuple[str, ...]
|
| 364 |
+
observation_spans_ticks: tuple[tuple[int, int], ...]
|
| 365 |
+
closed_start_ticks: tuple[int, ...]
|
| 366 |
+
closed_end_ticks: tuple[int, ...]
|
| 367 |
+
support_start_tick: int
|
| 368 |
+
support_end_tick: int
|
| 369 |
+
matching_ambiguous: bool
|
| 370 |
+
|
| 371 |
+
@dataclass(frozen=True)
|
| 372 |
+
class FineSegmentationResult:
|
| 373 |
+
events: tuple[EventInterval, ...]
|
| 374 |
+
fragments: tuple[EventInterval, ...]
|
| 375 |
+
unresolved: tuple[UnresolvedFineEvent, ...]
|
| 376 |
+
adjacent_matching_ambiguities: int
|
| 377 |
+
diagnostics: Mapping[str, object] = field(default_factory=dict)
|
| 378 |
+
|
| 379 |
+
def __post_init__(self) -> None:
|
| 380 |
+
object.__setattr__(self, 'diagnostics', dict(self.diagnostics))
|
| 381 |
+
|
| 382 |
+
def _observations(window: FineWindow, fragments: Sequence[FineFragment], *, first_id: int) -> tuple[_Observation, ...]:
|
| 383 |
+
result: list[_Observation] = []
|
| 384 |
+
for ordinal, fragment in enumerate(fragments):
|
| 385 |
+
_validate_fragment(fragment, window)
|
| 386 |
+
offset = window.global_tick_offset
|
| 387 |
+
result.append(_Observation(observation_id=first_id + ordinal, window_index=window.window_index, row_ordinal=ordinal, scope_id=window.scope_id, start_tick=offset + fragment.s, end_tick=offset + fragment.e, lo=fragment.lo, ro=fragment.ro, valid_start_tick=offset + window.valid_start_tick, valid_end_tick=offset + window.valid_end_tick))
|
| 388 |
+
return tuple(result)
|
| 389 |
+
|
| 390 |
+
def _shared_trace(item: _Observation, start: int, end: int) -> tuple[int, int] | None:
|
| 391 |
+
trace = (max(item.start_tick, start), min(item.end_tick, end))
|
| 392 |
+
return trace if trace[0] < trace[1] else None
|
| 393 |
+
|
| 394 |
+
@dataclass(frozen=True)
|
| 395 |
+
class _CandidateEvaluation:
|
| 396 |
+
cost: tuple[int, int] | None
|
| 397 |
+
rejection: str | None
|
| 398 |
+
|
| 399 |
+
def _candidate_evaluation(left: _Observation, right: _Observation, overlap_start: int, overlap_end: int, policy: StitchingPolicy) -> _CandidateEvaluation:
|
| 400 |
+
left_trace = _shared_trace(left, overlap_start, overlap_end)
|
| 401 |
+
right_trace = _shared_trace(right, overlap_start, overlap_end)
|
| 402 |
+
if left_trace is None or right_trace is None:
|
| 403 |
+
return _CandidateEvaluation(None, 'no_shared_trace')
|
| 404 |
+
deltas = (abs(left_trace[0] - right_trace[0]), abs(left_trace[1] - right_trace[1]))
|
| 405 |
+
if policy.contract in {STRICT_TRACE_STITCHING_CONTRACT, *COMPONENT_LOCAL_STITCHING_CONTRACTS}:
|
| 406 |
+
if max(deltas) > policy.endpoint_tolerance_ticks:
|
| 407 |
+
return _CandidateEvaluation(None, 'shared_trace_endpoint_delta')
|
| 408 |
+
trace_cost = sum(deltas)
|
| 409 |
+
else:
|
| 410 |
+
intersection = max(0, min(left_trace[1], right_trace[1]) - max(left_trace[0], right_trace[0]))
|
| 411 |
+
union = max(left_trace[1], right_trace[1]) - min(left_trace[0], right_trace[0])
|
| 412 |
+
if union <= 0 or intersection * policy.min_shared_trace_iou_denominator < union * policy.min_shared_trace_iou_numerator:
|
| 413 |
+
return _CandidateEvaluation(None, 'shared_trace_iou')
|
| 414 |
+
trace_cost = (union - intersection) * 1000000 // union
|
| 415 |
+
closed_cost = 0
|
| 416 |
+
if left.has_closed_start and right.has_closed_start:
|
| 417 |
+
delta = abs(left.start_tick - right.start_tick)
|
| 418 |
+
if delta > policy.endpoint_tolerance_ticks:
|
| 419 |
+
return _CandidateEvaluation(None, 'closed_start_delta')
|
| 420 |
+
closed_cost += delta
|
| 421 |
+
elif left.has_closed_start and left.start_tick > right.valid_start_tick + policy.endpoint_tolerance_ticks or (right.has_closed_start and right.start_tick > left.valid_start_tick + policy.endpoint_tolerance_ticks):
|
| 422 |
+
return _CandidateEvaluation(None, 'closed_start_visible_to_censored_peer')
|
| 423 |
+
if left.has_closed_end and right.has_closed_end:
|
| 424 |
+
delta = abs(left.end_tick - right.end_tick)
|
| 425 |
+
if delta > policy.endpoint_tolerance_ticks:
|
| 426 |
+
return _CandidateEvaluation(None, 'closed_end_delta')
|
| 427 |
+
closed_cost += delta
|
| 428 |
+
elif left.has_closed_end and left.end_tick < right.valid_end_tick - policy.endpoint_tolerance_ticks or (right.has_closed_end and right.end_tick < left.valid_end_tick - policy.endpoint_tolerance_ticks):
|
| 429 |
+
return _CandidateEvaluation(None, 'closed_end_visible_to_censored_peer')
|
| 430 |
+
return _CandidateEvaluation((trace_cost, closed_cost), None)
|
| 431 |
+
|
| 432 |
+
@dataclass(frozen=True)
|
| 433 |
+
class _MatchState:
|
| 434 |
+
count: int = 0
|
| 435 |
+
trace_cost: int = 0
|
| 436 |
+
closed_cost: int = 0
|
| 437 |
+
pairings: tuple[tuple[tuple[int, int], ...], ...] = ((),)
|
| 438 |
+
|
| 439 |
+
@property
|
| 440 |
+
def key(self) -> tuple[int, int, int]:
|
| 441 |
+
return (-self.count, self.trace_cost, self.closed_cost)
|
| 442 |
+
|
| 443 |
+
def _choose(states: Sequence[_MatchState]) -> _MatchState:
|
| 444 |
+
best_key = min((state.key for state in states))
|
| 445 |
+
optimal = [state for state in states if state.key == best_key]
|
| 446 |
+
pairings = tuple(sorted({pairing for state in optimal for pairing in state.pairings})[:2])
|
| 447 |
+
return _MatchState(-best_key[0], best_key[1], best_key[2], pairings)
|
| 448 |
+
|
| 449 |
+
def _match_adjacent(left: Sequence[_Observation], right: Sequence[_Observation], left_window: FineWindow, right_window: FineWindow, policy: StitchingPolicy) -> tuple[tuple[tuple[int, int], ...], bool, tuple[tuple[tuple[int, int], ...], ...]]:
|
| 450 |
+
if right_window.window_index != left_window.window_index + 1:
|
| 451 |
+
raise FineTemporalError('stitching requires the complete ordered global window grid')
|
| 452 |
+
overlap_start = max(left_window.global_tick_offset + left_window.valid_start_tick, right_window.global_tick_offset + right_window.valid_start_tick)
|
| 453 |
+
overlap_end = min(left_window.global_tick_offset + left_window.valid_end_tick, right_window.global_tick_offset + right_window.valid_end_tick)
|
| 454 |
+
if overlap_start >= overlap_end:
|
| 455 |
+
raise FineTemporalError('adjacent windows have no shared valid-visible guard')
|
| 456 |
+
n, m = (len(left), len(right))
|
| 457 |
+
dp = [[_MatchState() for _ in range(m + 1)] for _ in range(n + 1)]
|
| 458 |
+
for i in range(1, n + 1):
|
| 459 |
+
dp[i][0] = dp[i - 1][0]
|
| 460 |
+
for j in range(1, m + 1):
|
| 461 |
+
dp[0][j] = dp[0][j - 1]
|
| 462 |
+
for i in range(1, n + 1):
|
| 463 |
+
for j in range(1, m + 1):
|
| 464 |
+
candidates = [dp[i - 1][j], dp[i][j - 1]]
|
| 465 |
+
evaluation = _candidate_evaluation(left[i - 1], right[j - 1], overlap_start, overlap_end, policy)
|
| 466 |
+
if evaluation.cost is not None:
|
| 467 |
+
prior = dp[i - 1][j - 1]
|
| 468 |
+
pair = (left[i - 1].observation_id, right[j - 1].observation_id)
|
| 469 |
+
candidates.append(_MatchState(prior.count + 1, prior.trace_cost + evaluation.cost[0], prior.closed_cost + evaluation.cost[1], tuple(((*pairing, pair) for pairing in prior.pairings))))
|
| 470 |
+
dp[i][j] = _choose(candidates)
|
| 471 |
+
result = dp[n][m]
|
| 472 |
+
return (result.pairings[0], len(result.pairings) > 1, result.pairings)
|
| 473 |
+
|
| 474 |
+
class _UnionFind:
|
| 475 |
+
|
| 476 |
+
def __init__(self, values: Iterable[int]) -> None:
|
| 477 |
+
self.parent = {value: value for value in values}
|
| 478 |
+
|
| 479 |
+
def find(self, value: int) -> int:
|
| 480 |
+
parent = self.parent[value]
|
| 481 |
+
if parent != value:
|
| 482 |
+
self.parent[value] = self.find(parent)
|
| 483 |
+
return self.parent[value]
|
| 484 |
+
|
| 485 |
+
def union(self, left: int, right: int) -> None:
|
| 486 |
+
left_root, right_root = (self.find(left), self.find(right))
|
| 487 |
+
if left_root != right_root:
|
| 488 |
+
self.parent[max(left_root, right_root)] = min(left_root, right_root)
|
| 489 |
+
|
| 490 |
+
def _median(values: Sequence[int]) -> int:
|
| 491 |
+
ordered = sorted(values)
|
| 492 |
+
if not ordered:
|
| 493 |
+
raise FineTemporalError('cannot form endpoint consensus from no values')
|
| 494 |
+
middle = len(ordered) // 2
|
| 495 |
+
if len(ordered) % 2:
|
| 496 |
+
return ordered[middle]
|
| 497 |
+
value = Fraction(ordered[middle - 1] + ordered[middle], 2)
|
| 498 |
+
return math.floor(value + Fraction(1, 2))
|
| 499 |
+
|
| 500 |
+
@dataclass(frozen=True)
|
| 501 |
+
class _ComponentState:
|
| 502 |
+
root: int
|
| 503 |
+
observation_ids: tuple[int, ...]
|
| 504 |
+
starts: tuple[int, ...]
|
| 505 |
+
ends: tuple[int, ...]
|
| 506 |
+
|
| 507 |
+
@property
|
| 508 |
+
def resolved(self) -> bool:
|
| 509 |
+
return bool(self.starts) and bool(self.ends) and (self.starts[-1] - self.starts[0] <= STITCH_TOLERANCE_TICKS) and (self.ends[-1] - self.ends[0] <= STITCH_TOLERANCE_TICKS) and (_median(self.ends) > _median(self.starts))
|
| 510 |
+
|
| 511 |
+
@property
|
| 512 |
+
def partial_endpoint(self) -> tuple[str, tuple[int, ...]] | None:
|
| 513 |
+
if bool(self.starts) == bool(self.ends):
|
| 514 |
+
return None
|
| 515 |
+
name, values = ('start', self.starts) if self.starts else ('end', self.ends)
|
| 516 |
+
if values[-1] - values[0] > STITCH_TOLERANCE_TICKS:
|
| 517 |
+
return None
|
| 518 |
+
return (name, values)
|
| 519 |
+
|
| 520 |
+
@dataclass(frozen=True)
|
| 521 |
+
class _LocalCandidate:
|
| 522 |
+
seam_index: int
|
| 523 |
+
left_observation_id: int
|
| 524 |
+
right_observation_id: int
|
| 525 |
+
resolved_root: int
|
| 526 |
+
partial_root: int
|
| 527 |
+
anchor_endpoint: str
|
| 528 |
+
anchor_tick: int
|
| 529 |
+
anchor_delta_ticks: int
|
| 530 |
+
trace_intersection_ticks: int
|
| 531 |
+
trace_union_ticks: int
|
| 532 |
+
|
| 533 |
+
@property
|
| 534 |
+
def pair(self) -> tuple[int, int]:
|
| 535 |
+
return (self.left_observation_id, self.right_observation_id)
|
| 536 |
+
|
| 537 |
+
def to_dict(self) -> dict[str, object]:
|
| 538 |
+
result = {'left_observation_id': self.left_observation_id, 'right_observation_id': self.right_observation_id, 'resolved_component_root': self.resolved_root, 'partial_component_root': self.partial_root, 'anchor_endpoint': self.anchor_endpoint, 'anchor_tick': self.anchor_tick, 'trace_intersection_ticks': self.trace_intersection_ticks, 'trace_union_ticks': self.trace_union_ticks}
|
| 539 |
+
if self.anchor_delta_ticks:
|
| 540 |
+
result['anchor_delta_ticks'] = self.anchor_delta_ticks
|
| 541 |
+
return result
|
| 542 |
+
|
| 543 |
+
@dataclass(frozen=True)
|
| 544 |
+
class _AssociationPlan:
|
| 545 |
+
base_pairs_by_seam: tuple[tuple[tuple[int, int], ...], ...]
|
| 546 |
+
pairs_by_seam: tuple[tuple[tuple[int, int], ...], ...]
|
| 547 |
+
ambiguous_by_seam: tuple[bool, ...]
|
| 548 |
+
competing_by_seam: tuple[tuple[tuple[tuple[int, int], ...], ...], ...]
|
| 549 |
+
component_consistency_tolerance_ticks: int
|
| 550 |
+
component_local_by_seam: tuple[Mapping[str, object], ...]
|
| 551 |
+
component_local_summary: Mapping[str, object] | None
|
| 552 |
+
component_local_geometry_by_seam: tuple[Mapping[tuple[int, int], tuple[int, int]], ...]
|
| 553 |
+
|
| 554 |
+
def _component_states(observations: Sequence[_Observation], union_find: _UnionFind) -> dict[int, _ComponentState]:
|
| 555 |
+
grouped: dict[int, list[_Observation]] = {}
|
| 556 |
+
for observation in observations:
|
| 557 |
+
grouped.setdefault(union_find.find(observation.observation_id), []).append(observation)
|
| 558 |
+
return {root: _ComponentState(root=root, observation_ids=tuple(sorted((observation.observation_id for observation in component))), starts=tuple(sorted((observation.start_tick for observation in component if not observation.lo))), ends=tuple(sorted((observation.end_tick for observation in component if not observation.ro)))) for root, component in grouped.items()}
|
| 559 |
+
|
| 560 |
+
def _resolved_component_geometry(states: Mapping[int, _ComponentState]) -> tuple[tuple[int, int], ...]:
|
| 561 |
+
return tuple(sorted(((_median(state.starts), _median(state.ends)) for state in states.values() if state.resolved)))
|
| 562 |
+
|
| 563 |
+
def _trace_iou_counts(left: _Observation, right: _Observation, *, overlap_start: int, overlap_end: int) -> tuple[int, int]:
|
| 564 |
+
left_trace = _shared_trace(left, overlap_start, overlap_end)
|
| 565 |
+
right_trace = _shared_trace(right, overlap_start, overlap_end)
|
| 566 |
+
if left_trace is None or right_trace is None:
|
| 567 |
+
return (0, 0)
|
| 568 |
+
intersection = max(0, min(left_trace[1], right_trace[1]) - max(left_trace[0], right_trace[0]))
|
| 569 |
+
union = max(left_trace[1], right_trace[1]) - min(left_trace[0], right_trace[0])
|
| 570 |
+
return (intersection, union)
|
| 571 |
+
|
| 572 |
+
def _identity_anchor(left: _ComponentState, right: _ComponentState, *, maximum_delta_ticks: int) -> tuple[tuple[int, int, str, int, int] | None, str | None]:
|
| 573 |
+
if left.resolved == right.resolved:
|
| 574 |
+
return (None, 'requires_one_resolved_one_partial')
|
| 575 |
+
resolved, partial = (left, right) if left.resolved else (right, left)
|
| 576 |
+
partial_endpoint = partial.partial_endpoint
|
| 577 |
+
if partial_endpoint is None:
|
| 578 |
+
return (None, 'partial_component_is_not_one_sided')
|
| 579 |
+
endpoint_name, values = partial_endpoint
|
| 580 |
+
anchor_values = resolved.starts if endpoint_name == 'start' else resolved.ends
|
| 581 |
+
anchor_tick = _median(anchor_values)
|
| 582 |
+
maximum_delta = max((abs(value - anchor_tick) for value in values))
|
| 583 |
+
if maximum_delta > maximum_delta_ticks:
|
| 584 |
+
return (None, 'partial_endpoint_is_not_exactly_anchored' if maximum_delta_ticks == 0 else 'partial_endpoint_exceeds_anchor_tolerance')
|
| 585 |
+
if maximum_delta == 0 and _median((*anchor_values, *values)) != anchor_tick:
|
| 586 |
+
return (None, 'partial_endpoint_would_move_resolved_geometry')
|
| 587 |
+
return ((resolved.root, partial.root, endpoint_name, anchor_tick, maximum_delta), None)
|
| 588 |
+
|
| 589 |
+
def _monotone_with_pairs(pair: tuple[int, int], *, base_pairs: Sequence[tuple[int, int]], left: Sequence[_Observation], right: Sequence[_Observation]) -> bool:
|
| 590 |
+
left_position = {item.observation_id: index for index, item in enumerate(left)}
|
| 591 |
+
right_position = {item.observation_id: index for index, item in enumerate(right)}
|
| 592 |
+
left_index, right_index = (left_position[pair[0]], right_position[pair[1]])
|
| 593 |
+
return all((not (left_index < left_position[base_left] and right_index > right_position[base_right] or (left_index > left_position[base_left] and right_index < right_position[base_right])) for base_left, base_right in base_pairs))
|
| 594 |
+
|
| 595 |
+
def _component_local_absorptions(ordered: Sequence[FineWindow], by_window: Mapping[int, Sequence[_Observation]], observations: Sequence[_Observation], base_pairs_by_seam: Sequence[Sequence[tuple[int, int]]], base_union: _UnionFind, policy: StitchingPolicy) -> tuple[tuple[tuple[tuple[int, int], ...], ...], tuple[Mapping[str, object], ...], Mapping[str, object], tuple[Mapping[tuple[int, int], tuple[int, int]], ...]]:
|
| 596 |
+
states = _component_states(observations, base_union)
|
| 597 |
+
strict_policy = StitchingPolicy.strict()
|
| 598 |
+
candidates: list[_LocalCandidate] = []
|
| 599 |
+
rejection_counts: dict[str, int] = {}
|
| 600 |
+
seam_rejections: list[dict[str, int]] = [{} for _ in range(max(0, len(ordered) - 1))]
|
| 601 |
+
funnel_counts: dict[str, int] = {}
|
| 602 |
+
seam_funnel_counts: list[dict[str, int]] = [{} for _ in range(max(0, len(ordered) - 1))]
|
| 603 |
+
relaxed_rejection_counts: dict[str, int] = {}
|
| 604 |
+
identity_trace_rejection_buckets: dict[str, int] = {}
|
| 605 |
+
anchor_delta_ticks_histogram: dict[str, int] = {}
|
| 606 |
+
|
| 607 |
+
def reject(seam_index: int, reason: str) -> None:
|
| 608 |
+
rejection_counts[reason] = rejection_counts.get(reason, 0) + 1
|
| 609 |
+
current = seam_rejections[seam_index]
|
| 610 |
+
current[reason] = current.get(reason, 0) + 1
|
| 611 |
+
|
| 612 |
+
def advance(seam_index: int, stage: str) -> None:
|
| 613 |
+
funnel_counts[stage] = funnel_counts.get(stage, 0) + 1
|
| 614 |
+
current = seam_funnel_counts[seam_index]
|
| 615 |
+
current[stage] = current.get(stage, 0) + 1
|
| 616 |
+
|
| 617 |
+
def increment(values: dict[str, int], name: str) -> None:
|
| 618 |
+
values[name] = values.get(name, 0) + 1
|
| 619 |
+
for seam_index in range(len(ordered) - 1):
|
| 620 |
+
left_window, right_window = (ordered[seam_index], ordered[seam_index + 1])
|
| 621 |
+
left, right = (by_window[seam_index], by_window[seam_index + 1])
|
| 622 |
+
overlap_start = max(left_window.global_tick_offset + left_window.valid_start_tick, right_window.global_tick_offset + right_window.valid_start_tick)
|
| 623 |
+
overlap_end = min(left_window.global_tick_offset + left_window.valid_end_tick, right_window.global_tick_offset + right_window.valid_end_tick)
|
| 624 |
+
base_pairs = tuple(base_pairs_by_seam[seam_index])
|
| 625 |
+
matched_left = {pair[0] for pair in base_pairs}
|
| 626 |
+
matched_right = {pair[1] for pair in base_pairs}
|
| 627 |
+
for left_item in left:
|
| 628 |
+
if left_item.observation_id in matched_left:
|
| 629 |
+
continue
|
| 630 |
+
for right_item in right:
|
| 631 |
+
if right_item.observation_id in matched_right:
|
| 632 |
+
continue
|
| 633 |
+
advance(seam_index, 'unmatched_observation_pairs')
|
| 634 |
+
left_root = base_union.find(left_item.observation_id)
|
| 635 |
+
right_root = base_union.find(right_item.observation_id)
|
| 636 |
+
if left_root == right_root:
|
| 637 |
+
advance(seam_index, 'same_component_pairs')
|
| 638 |
+
continue
|
| 639 |
+
advance(seam_index, 'cross_component_pairs')
|
| 640 |
+
relaxed = _candidate_evaluation(left_item, right_item, overlap_start, overlap_end, policy)
|
| 641 |
+
if relaxed.cost is None:
|
| 642 |
+
advance(seam_index, 'relaxed_infeasible_pairs')
|
| 643 |
+
increment(relaxed_rejection_counts, str(relaxed.rejection or 'unknown'))
|
| 644 |
+
continue
|
| 645 |
+
advance(seam_index, 'relaxed_feasible_pairs')
|
| 646 |
+
strict = _candidate_evaluation(left_item, right_item, overlap_start, overlap_end, strict_policy)
|
| 647 |
+
if strict.cost is not None:
|
| 648 |
+
advance(seam_index, 'strict_feasible_pairs')
|
| 649 |
+
reject(seam_index, 'strict_feasible_not_selected')
|
| 650 |
+
continue
|
| 651 |
+
advance(seam_index, 'strict_infeasible_pairs')
|
| 652 |
+
intersection, union = _trace_iou_counts(left_item, right_item, overlap_start=overlap_start, overlap_end=overlap_end)
|
| 653 |
+
if union <= 0 or intersection * policy.min_shared_trace_iou_denominator < union * policy.min_shared_trace_iou_numerator:
|
| 654 |
+
if union <= 0 or intersection <= 0:
|
| 655 |
+
trace_bucket = 'zero'
|
| 656 |
+
elif intersection * 4 < union:
|
| 657 |
+
trace_bucket = 'gt_0_lt_0.25'
|
| 658 |
+
elif intersection * 2 < union:
|
| 659 |
+
trace_bucket = '0.25_to_lt_0.5'
|
| 660 |
+
else:
|
| 661 |
+
trace_bucket = '0.5_to_lt_0.75'
|
| 662 |
+
increment(identity_trace_rejection_buckets, trace_bucket)
|
| 663 |
+
reject(seam_index, 'identity_trace_iou')
|
| 664 |
+
continue
|
| 665 |
+
advance(seam_index, 'identity_trace_pass_pairs')
|
| 666 |
+
maximum_anchor_delta_ticks = policy.endpoint_tolerance_ticks if policy.contract == COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT else 0
|
| 667 |
+
anchor, reason = _identity_anchor(states[left_root], states[right_root], maximum_delta_ticks=maximum_anchor_delta_ticks)
|
| 668 |
+
if anchor is None:
|
| 669 |
+
if reason in {'partial_endpoint_is_not_exactly_anchored', 'partial_endpoint_exceeds_anchor_tolerance'}:
|
| 670 |
+
left_state, right_state = (states[left_root], states[right_root])
|
| 671 |
+
resolved, partial = (left_state, right_state) if left_state.resolved else (right_state, left_state)
|
| 672 |
+
partial_endpoint = partial.partial_endpoint
|
| 673 |
+
if partial_endpoint is None:
|
| 674 |
+
raise FineTemporalError('exact-anchor rejection lacks a partial endpoint')
|
| 675 |
+
endpoint_name, values = partial_endpoint
|
| 676 |
+
anchor_values = resolved.starts if endpoint_name == 'start' else resolved.ends
|
| 677 |
+
anchor_tick = _median(anchor_values)
|
| 678 |
+
maximum_delta = max((abs(value - anchor_tick) for value in values))
|
| 679 |
+
increment(anchor_delta_ticks_histogram, str(maximum_delta))
|
| 680 |
+
reject(seam_index, str(reason))
|
| 681 |
+
continue
|
| 682 |
+
advance(seam_index, 'identity_anchor_pass_pairs')
|
| 683 |
+
pair = (left_item.observation_id, right_item.observation_id)
|
| 684 |
+
if not _monotone_with_pairs(pair, base_pairs=base_pairs, left=left, right=right):
|
| 685 |
+
reject(seam_index, 'crosses_strict_pairing')
|
| 686 |
+
continue
|
| 687 |
+
advance(seam_index, 'monotone_pair_candidates')
|
| 688 |
+
resolved_root, partial_root, endpoint_name, anchor_tick, anchor_delta_ticks = anchor
|
| 689 |
+
candidates.append(_LocalCandidate(seam_index=seam_index, left_observation_id=pair[0], right_observation_id=pair[1], resolved_root=resolved_root, partial_root=partial_root, anchor_endpoint=endpoint_name, anchor_tick=anchor_tick, anchor_delta_ticks=anchor_delta_ticks, trace_intersection_ticks=intersection, trace_union_ticks=union))
|
| 690 |
+
by_component_pair: dict[tuple[int, int], list[_LocalCandidate]] = {}
|
| 691 |
+
for candidate in candidates:
|
| 692 |
+
key = (candidate.resolved_root, candidate.partial_root)
|
| 693 |
+
by_component_pair.setdefault(key, []).append(candidate)
|
| 694 |
+
unique_edges: list[_LocalCandidate] = []
|
| 695 |
+
for values in by_component_pair.values():
|
| 696 |
+
if len(values) != 1:
|
| 697 |
+
for candidate in values:
|
| 698 |
+
reject(candidate.seam_index, 'multiple_observation_edges')
|
| 699 |
+
continue
|
| 700 |
+
unique_edges.append(values[0])
|
| 701 |
+
advance(values[0].seam_index, 'unique_component_edge_pairs')
|
| 702 |
+
degrees: dict[int, int] = {}
|
| 703 |
+
for candidate in unique_edges:
|
| 704 |
+
degrees[candidate.resolved_root] = degrees.get(candidate.resolved_root, 0) + 1
|
| 705 |
+
degrees[candidate.partial_root] = degrees.get(candidate.partial_root, 0) + 1
|
| 706 |
+
reciprocal_unique: list[_LocalCandidate] = []
|
| 707 |
+
for candidate in unique_edges:
|
| 708 |
+
if degrees[candidate.resolved_root] != 1 or degrees[candidate.partial_root] != 1:
|
| 709 |
+
reject(candidate.seam_index, 'non_unique_component_identity')
|
| 710 |
+
continue
|
| 711 |
+
reciprocal_unique.append(candidate)
|
| 712 |
+
advance(candidate.seam_index, 'reciprocal_unique_pairs')
|
| 713 |
+
crossing: set[_LocalCandidate] = set()
|
| 714 |
+
by_seam_candidates: dict[int, list[_LocalCandidate]] = {}
|
| 715 |
+
for candidate in reciprocal_unique:
|
| 716 |
+
by_seam_candidates.setdefault(candidate.seam_index, []).append(candidate)
|
| 717 |
+
for seam_index, values in by_seam_candidates.items():
|
| 718 |
+
left, right = (by_window[seam_index], by_window[seam_index + 1])
|
| 719 |
+
for candidate in values:
|
| 720 |
+
peers = [peer.pair for peer in values if peer is not candidate]
|
| 721 |
+
if not _monotone_with_pairs(candidate.pair, base_pairs=peers, left=left, right=right):
|
| 722 |
+
crossing.add(candidate)
|
| 723 |
+
accepted = []
|
| 724 |
+
for candidate in reciprocal_unique:
|
| 725 |
+
if candidate in crossing:
|
| 726 |
+
reject(candidate.seam_index, 'crossing_local_candidates')
|
| 727 |
+
continue
|
| 728 |
+
accepted.append(candidate)
|
| 729 |
+
advance(candidate.seam_index, 'geometry_accepted_pairs')
|
| 730 |
+
accepted_by_seam: list[tuple[tuple[int, int], ...]] = []
|
| 731 |
+
geometry_by_seam: list[Mapping[tuple[int, int], tuple[int, int]]] = []
|
| 732 |
+
diagnostics_by_seam: list[Mapping[str, object]] = []
|
| 733 |
+
for seam_index in range(len(ordered) - 1):
|
| 734 |
+
seam_candidates = tuple((candidate for candidate in candidates if candidate.seam_index == seam_index))
|
| 735 |
+
seam_accepted = tuple((candidate for candidate in accepted if candidate.seam_index == seam_index))
|
| 736 |
+
accepted_by_seam.append(tuple(sorted((candidate.pair for candidate in seam_accepted))))
|
| 737 |
+
geometry_by_seam.append({candidate.pair: (_median(states[candidate.resolved_root].starts), _median(states[candidate.resolved_root].ends)) for candidate in seam_accepted})
|
| 738 |
+
diagnostics_by_seam.append({'candidate_pairs': [candidate.to_dict() for candidate in seam_candidates], 'accepted_pairs': [candidate.to_dict() for candidate in seam_accepted], 'rejection_counts': dict(sorted(seam_rejections[seam_index].items())), 'candidate_funnel': dict(sorted(seam_funnel_counts[seam_index].items()))})
|
| 739 |
+
summary = {'contract': policy.contract, 'strict_base_tolerance_ticks': STITCH_TOLERANCE_TICKS, 'candidate_tolerance_ticks': policy.endpoint_tolerance_ticks, 'minimum_identity_trace_iou': {'numerator': policy.min_shared_trace_iou_numerator, 'denominator': policy.min_shared_trace_iou_denominator}, 'candidate_pair_count': len(candidates), 'accepted_pair_count': len(accepted), 'absorbed_partial_component_count': len(accepted), 'resolved_event_count_change_by_contract': 0, 'resolved_event_geometry_change_by_contract': False, 'new_model_calls': 0, 'rejection_counts': dict(sorted(rejection_counts.items())), 'candidate_funnel': dict(sorted(funnel_counts.items())), 'candidate_funnel_details': {'relaxed_infeasible_rejection_counts': dict(sorted(relaxed_rejection_counts.items())), 'identity_trace_iou_rejection_buckets': dict(sorted(identity_trace_rejection_buckets.items())), 'anchor_maximum_delta_ticks_histogram': dict(sorted(anchor_delta_ticks_histogram.items(), key=lambda item: int(item[0])))}}
|
| 740 |
+
if policy.contract == COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT:
|
| 741 |
+
summary['anchor_contract'] = {'mode': 'bounded_certificate_only', 'maximum_delta_ticks': policy.endpoint_tolerance_ticks, 'resolved_geometry_authority': 'strict_resolved_component', 'explicit_allowlist_required_for_runtime': True}
|
| 742 |
+
return (tuple(accepted_by_seam), tuple(diagnostics_by_seam), summary, tuple(geometry_by_seam))
|
| 743 |
+
|
| 744 |
+
def _association_plan(ordered: Sequence[FineWindow], by_window: Mapping[int, Sequence[_Observation]], observations: Sequence[_Observation], policy: StitchingPolicy, component_pair_allowlist: Collection[ComponentPairKey] | None=None) -> _AssociationPlan:
|
| 745 |
+
match_policy = StitchingPolicy.strict() if policy.contract in COMPONENT_LOCAL_STITCHING_CONTRACTS else policy
|
| 746 |
+
base_pairs: list[tuple[tuple[int, int], ...]] = []
|
| 747 |
+
ambiguous: list[bool] = []
|
| 748 |
+
competing: list[tuple[tuple[tuple[int, int], ...], ...]] = []
|
| 749 |
+
for index in range(len(ordered) - 1):
|
| 750 |
+
pairs, is_ambiguous, alternatives = _match_adjacent(by_window[index], by_window[index + 1], ordered[index], ordered[index + 1], match_policy)
|
| 751 |
+
base_pairs.append(pairs)
|
| 752 |
+
ambiguous.append(is_ambiguous)
|
| 753 |
+
competing.append(alternatives)
|
| 754 |
+
if policy.contract not in COMPONENT_LOCAL_STITCHING_CONTRACTS:
|
| 755 |
+
return _AssociationPlan(base_pairs_by_seam=tuple(base_pairs), pairs_by_seam=tuple(base_pairs), ambiguous_by_seam=tuple(ambiguous), competing_by_seam=tuple(competing), component_consistency_tolerance_ticks=policy.endpoint_tolerance_ticks, component_local_by_seam=tuple(({} for _ in base_pairs)), component_local_summary=None, component_local_geometry_by_seam=tuple(({} for _ in base_pairs)))
|
| 756 |
+
base_union = _UnionFind((item.observation_id for item in observations))
|
| 757 |
+
for pairs in base_pairs:
|
| 758 |
+
for left_id, right_id in pairs:
|
| 759 |
+
base_union.union(left_id, right_id)
|
| 760 |
+
accepted, diagnostics, summary, geometry_by_seam = _component_local_absorptions(ordered, by_window, observations, base_pairs, base_union, policy)
|
| 761 |
+
if component_pair_allowlist is not None:
|
| 762 |
+
available = {(seam_index, *pair) for seam_index, seam_pairs in enumerate(accepted) for pair in seam_pairs}
|
| 763 |
+
requested = set(component_pair_allowlist)
|
| 764 |
+
unknown = requested - available
|
| 765 |
+
if unknown:
|
| 766 |
+
raise FineTemporalError(f'component pair allowlist contains non-candidates: {sorted(unknown)}')
|
| 767 |
+
filtered = tuple((tuple((pair for pair in seam_pairs if (seam_index, *pair) in requested)) for seam_index, seam_pairs in enumerate(accepted)))
|
| 768 |
+
rejected_count = sum(map(len, accepted)) - sum(map(len, filtered))
|
| 769 |
+
filtered_diagnostics = []
|
| 770 |
+
for seam_index, row in enumerate(diagnostics):
|
| 771 |
+
current = dict(row)
|
| 772 |
+
current['accepted_pairs'] = [pair for pair in row['accepted_pairs'] if (seam_index, int(pair['left_observation_id']), int(pair['right_observation_id'])) in requested]
|
| 773 |
+
rejection_counts = dict(row['rejection_counts'])
|
| 774 |
+
seam_rejected = len(accepted[seam_index]) - len(filtered[seam_index])
|
| 775 |
+
if seam_rejected:
|
| 776 |
+
rejection_counts['caption_certificate'] = seam_rejected
|
| 777 |
+
current['rejection_counts'] = dict(sorted(rejection_counts.items()))
|
| 778 |
+
filtered_diagnostics.append(current)
|
| 779 |
+
accepted = filtered
|
| 780 |
+
geometry_by_seam = tuple(({pair: geometry for pair, geometry in seam_geometry.items() if (seam_index, *pair) in requested} for seam_index, seam_geometry in enumerate(geometry_by_seam)))
|
| 781 |
+
diagnostics = tuple(filtered_diagnostics)
|
| 782 |
+
summary = dict(summary)
|
| 783 |
+
summary['accepted_pair_count'] = sum(map(len, accepted))
|
| 784 |
+
summary['absorbed_partial_component_count'] = sum(map(len, accepted))
|
| 785 |
+
rejection_counts = dict(summary['rejection_counts'])
|
| 786 |
+
if rejected_count:
|
| 787 |
+
rejection_counts['caption_certificate'] = rejected_count
|
| 788 |
+
summary['rejection_counts'] = dict(sorted(rejection_counts.items()))
|
| 789 |
+
combined = tuple((tuple(sorted((*strict_pairs, *local_pairs))) for strict_pairs, local_pairs in zip(base_pairs, accepted, strict=True)))
|
| 790 |
+
return _AssociationPlan(base_pairs_by_seam=tuple(base_pairs), pairs_by_seam=combined, ambiguous_by_seam=tuple(ambiguous), competing_by_seam=tuple(competing), component_consistency_tolerance_ticks=STITCH_TOLERANCE_TICKS, component_local_by_seam=diagnostics, component_local_summary=summary, component_local_geometry_by_seam=geometry_by_seam)
|
| 791 |
+
|
| 792 |
+
def _fragment_event(source_id: str, window: FineWindow, fragment: FineFragment, observation_id: int, row_ordinal: int) -> EventInterval:
|
| 793 |
+
start_us = max(0, (window.global_tick_offset + fragment.s) * TICK_US)
|
| 794 |
+
end_us = min(window.video_end_us, (window.global_tick_offset + fragment.e) * TICK_US)
|
| 795 |
+
if end_us <= start_us:
|
| 796 |
+
raise FineTemporalError('quantized local fragment is degenerate in source time')
|
| 797 |
+
completeness = {(False, False): 'complete', (True, False): 'left_censored', (False, True): 'right_censored', (True, True): 'both_censored'}[fragment.lo, fragment.ro]
|
| 798 |
+
return EventInterval(event_id=f'{source_id}:fragment:{observation_id:06d}', source_id=source_id, span=TimeSpan(start_us / 1000000, end_us / 1000000), completeness=completeness, source_scope_ids=(window.scope_id,), metadata={'window_index': window.window_index, 'row_ordinal': row_ordinal, 'observation_id': observation_id})
|
| 799 |
+
|
| 800 |
+
def _observation_payload(item: _Observation) -> dict[str, int | bool | str]:
|
| 801 |
+
return {'observation_id': item.observation_id, 'window_index': item.window_index, 'row_ordinal': item.row_ordinal, 'scope_id': item.scope_id, 'start_tick': item.start_tick, 'end_tick': item.end_tick, 'lo': item.lo, 'ro': item.ro, 'valid_start_tick': item.valid_start_tick, 'valid_end_tick': item.valid_end_tick}
|
| 802 |
+
|
| 803 |
+
def analyze_stitching_seams(windows: Sequence[FineWindow], fragments_by_window: Mapping[int, Sequence[FineFragment]], *, policy: StitchingPolicy | None=None, tolerance_ticks: int=STITCH_TOLERANCE_TICKS, component_pair_allowlist: Collection[ComponentPairKey] | None=None) -> dict[str, object]:
|
| 804 |
+
selected = policy or StitchingPolicy.strict(tolerance_ticks)
|
| 805 |
+
if policy is not None and tolerance_ticks != STITCH_TOLERANCE_TICKS:
|
| 806 |
+
raise FineTemporalError('pass either policy or a non-default tolerance_ticks, not both')
|
| 807 |
+
if component_pair_allowlist is not None and (policy is None or policy.contract not in COMPONENT_LOCAL_STITCHING_CONTRACTS):
|
| 808 |
+
raise FineTemporalError('component pair allowlist requires a component-local policy')
|
| 809 |
+
ordered = tuple(sorted(windows, key=lambda item: item.window_index))
|
| 810 |
+
if not ordered or [item.window_index for item in ordered] != list(range(len(ordered))):
|
| 811 |
+
raise FineTemporalError('windows must be one complete zero-based sequence')
|
| 812 |
+
if len({item.source_id for item in ordered}) != 1 or len({item.video_end_us for item in ordered}) != 1 or len({item.window_profile_id for item in ordered}) != 1:
|
| 813 |
+
raise FineTemporalError('all windows must belong to one source, duration, and profile')
|
| 814 |
+
unknown = set(fragments_by_window) - {item.window_index for item in ordered}
|
| 815 |
+
if unknown:
|
| 816 |
+
raise FineTemporalError(f'fragments supplied for unknown windows: {sorted(unknown)}')
|
| 817 |
+
by_window: dict[int, tuple[_Observation, ...]] = {}
|
| 818 |
+
observations: list[_Observation] = []
|
| 819 |
+
next_id = 0
|
| 820 |
+
for window in ordered:
|
| 821 |
+
current = _observations(window, tuple(fragments_by_window.get(window.window_index, ())), first_id=next_id)
|
| 822 |
+
next_id += len(current)
|
| 823 |
+
observations.extend(current)
|
| 824 |
+
by_window[window.window_index] = current
|
| 825 |
+
plan = _association_plan(ordered, by_window, observations, selected, component_pair_allowlist)
|
| 826 |
+
seams: list[dict[str, object]] = []
|
| 827 |
+
rejection_totals: dict[str, int] = {}
|
| 828 |
+
total_feasible = total_matched = total_ambiguous = 0
|
| 829 |
+
for index in range(len(ordered) - 1):
|
| 830 |
+
left_window, right_window = (ordered[index], ordered[index + 1])
|
| 831 |
+
left, right = (by_window[index], by_window[index + 1])
|
| 832 |
+
overlap_start = max(left_window.global_tick_offset + left_window.valid_start_tick, right_window.global_tick_offset + right_window.valid_start_tick)
|
| 833 |
+
overlap_end = min(left_window.global_tick_offset + left_window.valid_end_tick, right_window.global_tick_offset + right_window.valid_end_tick)
|
| 834 |
+
rejection_counts: dict[str, int] = {}
|
| 835 |
+
feasible_pairs: list[dict[str, object]] = []
|
| 836 |
+
for left_item in left:
|
| 837 |
+
for right_item in right:
|
| 838 |
+
evaluation = _candidate_evaluation(left_item, right_item, overlap_start, overlap_end, selected)
|
| 839 |
+
if evaluation.cost is None:
|
| 840 |
+
reason = str(evaluation.rejection)
|
| 841 |
+
rejection_counts[reason] = rejection_counts.get(reason, 0) + 1
|
| 842 |
+
rejection_totals[reason] = rejection_totals.get(reason, 0) + 1
|
| 843 |
+
continue
|
| 844 |
+
feasible_pairs.append({'left_observation_id': left_item.observation_id, 'right_observation_id': right_item.observation_id, 'trace_cost': evaluation.cost[0], 'closed_endpoint_cost': evaluation.cost[1]})
|
| 845 |
+
pairs = plan.pairs_by_seam[index]
|
| 846 |
+
ambiguous = plan.ambiguous_by_seam[index]
|
| 847 |
+
competing = plan.competing_by_seam[index]
|
| 848 |
+
matched_left = {item[0] for item in pairs}
|
| 849 |
+
matched_right = {item[1] for item in pairs}
|
| 850 |
+
total_feasible += len(feasible_pairs)
|
| 851 |
+
total_matched += len(pairs)
|
| 852 |
+
total_ambiguous += int(ambiguous)
|
| 853 |
+
seams.append({'left_window_index': left_window.window_index, 'right_window_index': right_window.window_index, 'overlap_start_tick': overlap_start, 'overlap_end_tick': overlap_end, 'left_observations': [_observation_payload(item) for item in left], 'right_observations': [_observation_payload(item) for item in right], 'feasible_pairs': feasible_pairs, 'matched_pairs': [list(item) for item in pairs], 'competing_optimal_pairings': [[list(item) for item in pairing] for pairing in competing], 'matching_ambiguous': ambiguous, 'unmatched_left_observation_ids': [item.observation_id for item in left if item.observation_id not in matched_left], 'unmatched_right_observation_ids': [item.observation_id for item in right if item.observation_id not in matched_right], 'rejection_counts': dict(sorted(rejection_counts.items())), 'component_local_identity': dict(plan.component_local_by_seam[index])})
|
| 854 |
+
return {'schema_version': STITCHING_SEAM_DIAGNOSTICS_SCHEMA, 'policy': selected.to_dict(), 'window_count': len(ordered), 'seam_count': len(seams), 'observation_count': next_id, 'feasible_pair_count': total_feasible, 'matched_pair_count': total_matched, 'ambiguous_seam_count': total_ambiguous, 'rejection_counts': dict(sorted(rejection_totals.items())), 'component_local_identity': dict(plan.component_local_summary) if plan.component_local_summary is not None else None, 'seams': seams}
|
| 855 |
+
|
| 856 |
+
def stitch_predictions(windows: Sequence[FineWindow], fragments_by_window: Mapping[int, Sequence[FineFragment]], *, tolerance_ticks: int=STITCH_TOLERANCE_TICKS, policy: StitchingPolicy | None=None, component_pair_allowlist: Collection[ComponentPairKey] | None=None) -> FineSegmentationResult:
|
| 857 |
+
if policy is not None and tolerance_ticks != STITCH_TOLERANCE_TICKS:
|
| 858 |
+
raise FineTemporalError('pass either policy or a non-default tolerance_ticks, not both')
|
| 859 |
+
if component_pair_allowlist is not None and (policy is None or policy.contract not in COMPONENT_LOCAL_STITCHING_CONTRACTS):
|
| 860 |
+
raise FineTemporalError('component pair allowlist requires a component-local policy')
|
| 861 |
+
selected = policy or StitchingPolicy.strict(tolerance_ticks)
|
| 862 |
+
if selected.contract == COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT and component_pair_allowlist is None:
|
| 863 |
+
raise FineTemporalError('bounded-anchor runtime requires an explicit component pair allowlist')
|
| 864 |
+
ordered = tuple(sorted(windows, key=lambda item: item.window_index))
|
| 865 |
+
if not ordered or [item.window_index for item in ordered] != list(range(len(ordered))):
|
| 866 |
+
raise FineTemporalError('windows must be one complete zero-based sequence')
|
| 867 |
+
if len({item.source_id for item in ordered}) != 1 or len({item.video_end_us for item in ordered}) != 1 or len({item.window_profile_id for item in ordered}) != 1:
|
| 868 |
+
raise FineTemporalError('all windows must belong to one source, duration, and profile')
|
| 869 |
+
unknown = set(fragments_by_window) - {item.window_index for item in ordered}
|
| 870 |
+
if unknown:
|
| 871 |
+
raise FineTemporalError(f'fragments supplied for unknown windows: {sorted(unknown)}')
|
| 872 |
+
observations: list[_Observation] = []
|
| 873 |
+
by_window: dict[int, tuple[_Observation, ...]] = {}
|
| 874 |
+
fragments: list[EventInterval] = []
|
| 875 |
+
next_id = 0
|
| 876 |
+
for window in ordered:
|
| 877 |
+
rows = tuple(fragments_by_window.get(window.window_index, ()))
|
| 878 |
+
current = _observations(window, rows, first_id=next_id)
|
| 879 |
+
for offset, fragment in enumerate(rows):
|
| 880 |
+
fragments.append(_fragment_event(window.source_id, window, fragment, next_id + offset, offset))
|
| 881 |
+
next_id += len(current)
|
| 882 |
+
observations.extend(current)
|
| 883 |
+
by_window[window.window_index] = current
|
| 884 |
+
plan = _association_plan(ordered, by_window, observations, selected, component_pair_allowlist)
|
| 885 |
+
baseline_geometry: tuple[tuple[int, int], ...] | None = None
|
| 886 |
+
baseline_unresolved_count: int | None = None
|
| 887 |
+
if selected.contract in COMPONENT_LOCAL_STITCHING_CONTRACTS:
|
| 888 |
+
baseline_union = _UnionFind((item.observation_id for item in observations))
|
| 889 |
+
for pairs in plan.base_pairs_by_seam:
|
| 890 |
+
for left_id, right_id in pairs:
|
| 891 |
+
baseline_union.union(left_id, right_id)
|
| 892 |
+
baseline_states = _component_states(observations, baseline_union)
|
| 893 |
+
baseline_geometry = _resolved_component_geometry(baseline_states)
|
| 894 |
+
baseline_unresolved_count = sum((not state.resolved for state in baseline_states.values()))
|
| 895 |
+
union_find = _UnionFind((item.observation_id for item in observations))
|
| 896 |
+
ambiguity_ids: set[int] = set()
|
| 897 |
+
ambiguity_count = 0
|
| 898 |
+
for pairs, ambiguous, competing in zip(plan.pairs_by_seam, plan.ambiguous_by_seam, plan.competing_by_seam, strict=True):
|
| 899 |
+
if ambiguous:
|
| 900 |
+
ambiguity_count += 1
|
| 901 |
+
ambiguity_ids.update((value for pairing in competing for pair in pairing for value in pair))
|
| 902 |
+
for left_id, right_id in pairs:
|
| 903 |
+
union_find.union(left_id, right_id)
|
| 904 |
+
authoritative_geometry_by_root: dict[int, tuple[int, int]] = {}
|
| 905 |
+
for geometry_by_pair in plan.component_local_geometry_by_seam:
|
| 906 |
+
for (left_id, right_id), geometry in geometry_by_pair.items():
|
| 907 |
+
root = union_find.find(left_id)
|
| 908 |
+
if root != union_find.find(right_id):
|
| 909 |
+
raise FineTemporalError('component-local geometry pair was not associated')
|
| 910 |
+
prior = authoritative_geometry_by_root.get(root)
|
| 911 |
+
if prior is not None and prior != geometry:
|
| 912 |
+
raise FineTemporalError('component-local candidates disagree on authoritative geometry')
|
| 913 |
+
authoritative_geometry_by_root[root] = geometry
|
| 914 |
+
if baseline_geometry is not None and baseline_unresolved_count is not None:
|
| 915 |
+
final_states = _component_states(observations, union_find)
|
| 916 |
+
accepted_count = sum((len(pairs) - len(base_pairs) for pairs, base_pairs in zip(plan.pairs_by_seam, plan.base_pairs_by_seam, strict=True)))
|
| 917 |
+
effective_geometry = tuple(sorted((authoritative_geometry_by_root.get(root, (_median(state.starts), _median(state.ends))) for root, state in final_states.items() if root in authoritative_geometry_by_root or state.resolved)))
|
| 918 |
+
if effective_geometry != baseline_geometry:
|
| 919 |
+
raise FineTemporalError('component-local identity relaxation changed resolved event geometry')
|
| 920 |
+
if sum((root not in authoritative_geometry_by_root and (not state.resolved) for root, state in final_states.items())) != baseline_unresolved_count - accepted_count:
|
| 921 |
+
raise FineTemporalError('component-local identity relaxation changed unsupported component structure')
|
| 922 |
+
grouped: dict[int, list[_Observation]] = {}
|
| 923 |
+
for observation in observations:
|
| 924 |
+
grouped.setdefault(union_find.find(observation.observation_id), []).append(observation)
|
| 925 |
+
ambiguous_roots = {union_find.find(value) for value in ambiguity_ids}
|
| 926 |
+
resolved: list[tuple[TimeSpan, tuple[str, ...], dict[str, object]]] = []
|
| 927 |
+
unresolved: list[UnresolvedFineEvent] = []
|
| 928 |
+
video_end_us = ordered[0].video_end_us
|
| 929 |
+
windows_by_index = {item.window_index: item for item in ordered}
|
| 930 |
+
for root, component in sorted(grouped.items()):
|
| 931 |
+
starts = tuple(sorted((item.start_tick for item in component if item.has_closed_start)))
|
| 932 |
+
ends = tuple(sorted((item.end_tick for item in component if item.has_closed_end)))
|
| 933 |
+
observation_spans = tuple(((item.start_tick, item.end_tick) for item in sorted(component, key=lambda item: item.observation_id)))
|
| 934 |
+
scope_ids = tuple(dict.fromkeys((item.scope_id for item in sorted(component, key=lambda item: item.window_index))))
|
| 935 |
+
raw_consistent = bool(starts) and bool(ends) and (starts[-1] - starts[0] <= plan.component_consistency_tolerance_ticks) and (ends[-1] - ends[0] <= plan.component_consistency_tolerance_ticks)
|
| 936 |
+
authoritative_geometry = authoritative_geometry_by_root.get(root)
|
| 937 |
+
consistent = authoritative_geometry is not None or raw_consistent
|
| 938 |
+
if authoritative_geometry is not None:
|
| 939 |
+
start_tick, end_tick = authoritative_geometry
|
| 940 |
+
start_us = max(0, start_tick * TICK_US)
|
| 941 |
+
end_us = min(video_end_us, end_tick * TICK_US)
|
| 942 |
+
else:
|
| 943 |
+
start_us = max(0, _median(starts) * TICK_US) if starts else None
|
| 944 |
+
end_us = min(video_end_us, _median(ends) * TICK_US) if ends else None
|
| 945 |
+
if not consistent or start_us is None or end_us is None or (end_us <= start_us):
|
| 946 |
+
support_window_indices = {item.window_index for item in component}
|
| 947 |
+
if not starts:
|
| 948 |
+
support_window_indices.update((item.window_index - 1 for item in component if item.window_index > 0))
|
| 949 |
+
if not ends:
|
| 950 |
+
support_window_indices.update((item.window_index + 1 for item in component if item.window_index + 1 < len(ordered)))
|
| 951 |
+
if starts and ends and (not consistent):
|
| 952 |
+
support_window_indices.update((item.window_index - 1 for item in component if item.window_index > 0))
|
| 953 |
+
support_window_indices.update((item.window_index + 1 for item in component if item.window_index + 1 < len(ordered)))
|
| 954 |
+
support_windows = tuple((windows_by_index[index] for index in sorted(support_window_indices)))
|
| 955 |
+
support_start_tick = min((item.global_tick_offset + item.valid_start_tick for item in support_windows))
|
| 956 |
+
support_end_tick = max((item.global_tick_offset + item.valid_end_tick for item in support_windows))
|
| 957 |
+
unresolved.append(UnresolvedFineEvent(tuple(sorted((item.observation_id for item in component))), scope_ids, observation_spans, starts, ends, support_start_tick, support_end_tick, root in ambiguous_roots))
|
| 958 |
+
continue
|
| 959 |
+
resolved.append((TimeSpan(start_us / 1000000, end_us / 1000000), scope_ids, {'observation_ids': tuple(sorted((item.observation_id for item in component))), 'matching_ambiguous': root in ambiguous_roots, **({'component_local_geometry_authority': 'strict_resolved_component', 'raw_closed_endpoints_consistent': raw_consistent} if authoritative_geometry is not None and selected.contract == COMPONENT_LOCAL_BOUNDED_ANCHOR_STITCHING_CONTRACT else {})}))
|
| 960 |
+
resolved.sort(key=lambda item: (item[0].start_s, item[0].end_s, item[1]))
|
| 961 |
+
identified = tuple((EventInterval(event_id=f'{ordered[0].source_id}:fine:{index:06d}', source_id=ordered[0].source_id, span=span, source_scope_ids=scope_ids, metadata=metadata) for index, (span, scope_ids, metadata) in enumerate(resolved)))
|
| 962 |
+
return FineSegmentationResult(events=identified, fragments=tuple(fragments), unresolved=tuple(unresolved), adjacent_matching_ambiguities=ambiguity_count, diagnostics={'component_local_identity': dict(plan.component_local_summary)} if plan.component_local_summary is not None else {})
|
| 963 |
+
|
| 964 |
+
def window_from_scope(scope: ProcessingScope) -> FineWindow:
|
| 965 |
+
metadata = dict(scope.metadata)
|
| 966 |
+
if metadata.get('schema_version') != WINDOW_SCHEMA:
|
| 967 |
+
raise FineTemporalError('processing scope is not a fine-window scope')
|
| 968 |
+
profile_id = metadata.get('window_profile_id', DEFAULT_FINE_WINDOW_PROFILE_ID)
|
| 969 |
+
profile = fine_window_profile(profile_id)
|
| 970 |
+
if metadata.get('model_io_contract') != profile.model_io_contract:
|
| 971 |
+
raise FineTemporalError('processing scope model I/O contract mismatch')
|
| 972 |
+
declared_profile = metadata.get('window_profile')
|
| 973 |
+
if declared_profile is not None and declared_profile != profile.to_dict():
|
| 974 |
+
raise FineTemporalError('processing scope fine window profile binding mismatch')
|
| 975 |
+
required = ('window_index', 'video_end_us', 'requested_view_start_us', 'requested_view_end_us', 'valid_start_us', 'valid_end_us', 'core_start_us', 'core_end_us', 'valid_start_tick', 'valid_end_tick', 'core_start_tick', 'core_end_tick')
|
| 976 |
+
values: dict[str, int] = {}
|
| 977 |
+
for key in required:
|
| 978 |
+
value = metadata.get(key)
|
| 979 |
+
if isinstance(value, bool) or not isinstance(value, int):
|
| 980 |
+
raise FineTemporalError(f'scope metadata {key} must be an integer')
|
| 981 |
+
values[key] = value
|
| 982 |
+
window = FineWindow(source_id=scope.source_id, scope_id=scope.scope_id, window_profile_id=profile.profile_id, coarse_spans=tuple(metadata.get('coarse_spans') or ()), **values)
|
| 983 |
+
expected_core = TimeSpan(window.core_start_us / 1000000, window.core_end_us / 1000000)
|
| 984 |
+
expected_context = TimeSpan(window.valid_start_us / 1000000, window.valid_end_us / 1000000)
|
| 985 |
+
if scope.core_span != expected_core or scope.context_span != expected_context:
|
| 986 |
+
raise FineTemporalError('float scope spans disagree with bound exact-us geometry')
|
| 987 |
+
return window
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/fine_worker.py
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from bisect import bisect_left
|
| 5 |
+
from collections.abc import Sequence
|
| 6 |
+
from fractions import Fraction
|
| 7 |
+
SLOT_SPACING_US = 125000
|
| 8 |
+
|
| 9 |
+
class FineWorkerError(RuntimeError):
|
| 10 |
+
pass
|
| 11 |
+
|
| 12 |
+
def _nearest_index(requested_us: int, times: Sequence[Fraction], *, max_selection_error_us: int=SLOT_SPACING_US // 2) -> int:
|
| 13 |
+
requested = Fraction(requested_us, 1)
|
| 14 |
+
position = bisect_left(times, requested)
|
| 15 |
+
candidates = [index for index in (position - 1, position) if 0 <= index < len(times)]
|
| 16 |
+
if not candidates:
|
| 17 |
+
raise FineWorkerError('no decoded frame is available for an in-video slot')
|
| 18 |
+
selected = min(candidates, key=lambda index: (abs(times[index] - requested), times[index], index))
|
| 19 |
+
if abs(times[selected] - requested) > max_selection_error_us:
|
| 20 |
+
raise FineWorkerError('nearest decoded frame is farther than half a profile slot')
|
| 21 |
+
return selected
|
| 22 |
+
if __name__ == '__main__':
|
| 23 |
+
raise SystemExit(main())
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/interface.py
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from collections.abc import Awaitable, Sequence
|
| 5 |
+
from enum import Enum
|
| 6 |
+
from typing import Protocol, runtime_checkable
|
| 7 |
+
from flora_runtime.schema import AnnotationSource, TemporalSegment, TimeSpan
|
| 8 |
+
|
| 9 |
+
class SegmentTopology(str, Enum):
|
| 10 |
+
EXCLUSIVE = 'exclusive_intervals'
|
| 11 |
+
CONCURRENT = 'concurrent_intervals'
|
| 12 |
+
|
| 13 |
+
def segment_topology(segmenter: object) -> SegmentTopology:
|
| 14 |
+
value = getattr(segmenter, 'segment_topology', SegmentTopology.EXCLUSIVE)
|
| 15 |
+
try:
|
| 16 |
+
return SegmentTopology(value)
|
| 17 |
+
except ValueError as exc:
|
| 18 |
+
raise ValueError(f'segmenter declares unsupported segment_topology: {value!r}') from exc
|
| 19 |
+
|
| 20 |
+
@runtime_checkable
|
| 21 |
+
class VideoSegmenter(Protocol):
|
| 22 |
+
|
| 23 |
+
def segment(self, source: AnnotationSource) -> Sequence[TemporalSegment | TimeSpan] | Awaitable[Sequence[TemporalSegment | TimeSpan]]:
|
| 24 |
+
...
|
batch_utils/vendor/production_adapter/core/flora_runtime/video_segmentation/scope_scheduler.py
ADDED
|
@@ -0,0 +1,163 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from collections.abc import Mapping, Sequence
|
| 5 |
+
from decimal import Decimal, InvalidOperation
|
| 6 |
+
from itertools import pairwise
|
| 7 |
+
import math
|
| 8 |
+
from flora_runtime.schema import TimeSpan
|
| 9 |
+
from flora_runtime.video_segmentation.contracts import ProcessingScope
|
| 10 |
+
from flora_runtime.video_segmentation.fine_temporal import DEFAULT_FINE_WINDOW_PROFILE_ID, TICK_US, WINDOW_SCHEMA, ceil_div, fine_window_profile, floor_div, round_half_up_div
|
| 11 |
+
|
| 12 |
+
class ScopeSchedulerError(ValueError):
|
| 13 |
+
pass
|
| 14 |
+
FINE_WINDOW_SCOPE_NAMESPACE = 'fine-window'
|
| 15 |
+
FINE_RETRY_SCOPE_NAMESPACE = 'fine-retry'
|
| 16 |
+
FINE_DIRECTIONAL_RETRY_SCOPE_NAMESPACE = 'fine-retry-directional'
|
| 17 |
+
|
| 18 |
+
def seconds_to_us(value: float, *, label: str) -> int:
|
| 19 |
+
if isinstance(value, bool) or not isinstance(value, int | float):
|
| 20 |
+
raise ScopeSchedulerError(f'{label} must be numeric')
|
| 21 |
+
try:
|
| 22 |
+
parsed = Decimal(str(value)) * 1000000
|
| 23 |
+
except (InvalidOperation, ValueError) as exc:
|
| 24 |
+
raise ScopeSchedulerError(f'{label} is not valid decimal seconds') from exc
|
| 25 |
+
if not parsed.is_finite() or parsed != parsed.to_integral_value():
|
| 26 |
+
raise ScopeSchedulerError(f'{label} must resolve to integral microseconds')
|
| 27 |
+
result = int(parsed)
|
| 28 |
+
if result < 0:
|
| 29 |
+
raise ScopeSchedulerError(f'{label} must be non-negative')
|
| 30 |
+
return result
|
| 31 |
+
|
| 32 |
+
def _quantized_prior(proposal: ProcessingScope, *, requested_start_us: int, valid_start_us: int, valid_end_us: int, valid_start_tick: int, valid_end_tick: int) -> dict[str, object] | None:
|
| 33 |
+
start_us = seconds_to_us(proposal.core_span.start_s, label=f'{proposal.scope_id}.start_s')
|
| 34 |
+
end_us = seconds_to_us(proposal.core_span.end_s, label=f'{proposal.scope_id}.end_s')
|
| 35 |
+
if start_us >= valid_end_us or end_us <= valid_start_us:
|
| 36 |
+
return None
|
| 37 |
+
visible_start = max(start_us, valid_start_us)
|
| 38 |
+
visible_end = min(end_us, valid_end_us)
|
| 39 |
+
lo, ro = (start_us < valid_start_us, end_us > valid_end_us)
|
| 40 |
+
start_delta, end_delta = (visible_start - requested_start_us, visible_end - requested_start_us)
|
| 41 |
+
s = valid_start_tick if lo else round_half_up_div(start_delta, TICK_US)
|
| 42 |
+
e = valid_end_tick if ro else round_half_up_div(end_delta, TICK_US)
|
| 43 |
+
if s >= e:
|
| 44 |
+
s, e = (floor_div(start_delta, TICK_US), ceil_div(end_delta, TICK_US))
|
| 45 |
+
if not valid_start_tick <= s < e <= valid_end_tick:
|
| 46 |
+
raise ScopeSchedulerError(f'proposal {proposal.scope_id} cannot be represented at 20 Hz')
|
| 47 |
+
return {'s': s, 'e': e, 'lo': lo, 'ro': ro}
|
| 48 |
+
|
| 49 |
+
def schedule_fine_scopes(*, source_id: str, duration_s: float, coarse_proposals: Sequence[ProcessingScope]=(), arm: str='no_prior_global_scan', window_profile_id: str=DEFAULT_FINE_WINDOW_PROFILE_ID) -> tuple[ProcessingScope, ...]:
|
| 50 |
+
if not isinstance(source_id, str) or not source_id.strip():
|
| 51 |
+
raise ScopeSchedulerError('source_id must be non-empty')
|
| 52 |
+
if not isinstance(arm, str) or not arm.strip():
|
| 53 |
+
raise ScopeSchedulerError('arm must be non-empty')
|
| 54 |
+
try:
|
| 55 |
+
profile = fine_window_profile(window_profile_id)
|
| 56 |
+
except ValueError as exc:
|
| 57 |
+
raise ScopeSchedulerError(str(exc)) from exc
|
| 58 |
+
video_end_us = seconds_to_us(duration_s, label='duration_s')
|
| 59 |
+
if video_end_us <= 0:
|
| 60 |
+
raise ScopeSchedulerError('duration_s must be positive')
|
| 61 |
+
proposals = tuple(coarse_proposals)
|
| 62 |
+
proposal_ids = [proposal.scope_id for proposal in proposals]
|
| 63 |
+
if len(proposal_ids) != len(set(proposal_ids)):
|
| 64 |
+
raise ScopeSchedulerError('coarse proposal scope_ids must be unique')
|
| 65 |
+
if any((proposal.source_id != source_id for proposal in proposals)):
|
| 66 |
+
raise ScopeSchedulerError('every coarse proposal must belong to source_id')
|
| 67 |
+
for proposal in proposals:
|
| 68 |
+
proposal_end_us = seconds_to_us(proposal.core_span.end_s, label=f'{proposal.scope_id}.end_s')
|
| 69 |
+
if proposal_end_us > video_end_us:
|
| 70 |
+
raise ScopeSchedulerError(f'coarse proposal {proposal.scope_id} extends past source')
|
| 71 |
+
proposals = tuple(sorted(proposals, key=lambda item: (item.core_span, item.scope_id)))
|
| 72 |
+
scopes: list[ProcessingScope] = []
|
| 73 |
+
window_index = 0
|
| 74 |
+
while True:
|
| 75 |
+
core_start_us = window_index * profile.core_us
|
| 76 |
+
if core_start_us >= video_end_us:
|
| 77 |
+
break
|
| 78 |
+
core_end_us = min(core_start_us + profile.core_us, video_end_us)
|
| 79 |
+
requested_start_us = core_start_us - profile.left_guard_us
|
| 80 |
+
requested_end_us = requested_start_us + profile.view_us
|
| 81 |
+
valid_start_us = max(0, requested_start_us)
|
| 82 |
+
valid_end_us = min(video_end_us, requested_end_us)
|
| 83 |
+
valid_start_tick = floor_div(valid_start_us - requested_start_us, TICK_US)
|
| 84 |
+
valid_end_tick = ceil_div(valid_end_us - requested_start_us, TICK_US)
|
| 85 |
+
core_start_tick = floor_div(core_start_us - requested_start_us, TICK_US)
|
| 86 |
+
core_end_tick = ceil_div(core_end_us - requested_start_us, TICK_US)
|
| 87 |
+
prior_groups: dict[tuple[int, int, bool, bool], list[str]] = {}
|
| 88 |
+
for proposal in proposals:
|
| 89 |
+
row = _quantized_prior(proposal, requested_start_us=requested_start_us, valid_start_us=valid_start_us, valid_end_us=valid_end_us, valid_start_tick=valid_start_tick, valid_end_tick=valid_end_tick)
|
| 90 |
+
if row is not None:
|
| 91 |
+
key = (int(row['s']), int(row['e']), bool(row['lo']), bool(row['ro']))
|
| 92 |
+
prior_groups.setdefault(key, []).append(proposal.scope_id)
|
| 93 |
+
ordered_priors = sorted(prior_groups.items())
|
| 94 |
+
coarse_rows = [{'s': key[0], 'e': key[1], 'lo': key[2], 'ro': key[3]} for key, _ in ordered_priors]
|
| 95 |
+
visible_ids = [scope_id for _, scope_ids in ordered_priors for scope_id in scope_ids]
|
| 96 |
+
coarse_span_lineage = [{'row_index': index, 'scope_ids': scope_ids} for index, (_, scope_ids) in enumerate(ordered_priors)]
|
| 97 |
+
if len(coarse_rows) > 64:
|
| 98 |
+
raise ScopeSchedulerError('one fine view contains more than 64 coarse priors')
|
| 99 |
+
scope_id = f'{source_id}:{FINE_WINDOW_SCOPE_NAMESPACE}:{window_index:06d}'
|
| 100 |
+
metadata = {'schema_version': WINDOW_SCHEMA, 'model_io_contract': profile.model_io_contract, 'window_profile_id': profile.profile_id, 'window_profile': profile.to_dict(), 'scheduler_contract': profile.scheduler_contract, 'arm': arm, 'guidance': 'coarse_prior' if coarse_rows else 'no_prior_global_scan', 'coarse_spans': coarse_rows, 'coarse_scope_ids': visible_ids, 'coarse_span_lineage': coarse_span_lineage, 'window_index': window_index, 'video_end_us': video_end_us, 'requested_view_start_us': requested_start_us, 'requested_view_end_us': requested_end_us, 'valid_start_us': valid_start_us, 'valid_end_us': valid_end_us, 'core_start_us': core_start_us, 'core_end_us': core_end_us, 'valid_start_tick': valid_start_tick, 'valid_end_tick': valid_end_tick, 'core_start_tick': core_start_tick, 'core_end_tick': core_end_tick}
|
| 101 |
+
scopes.append(ProcessingScope(scope_id=scope_id, source_id=source_id, core_span=TimeSpan(core_start_us / 1000000, core_end_us / 1000000), context_span=TimeSpan(valid_start_us / 1000000, valid_end_us / 1000000), metadata=metadata))
|
| 102 |
+
window_index += 1
|
| 103 |
+
if not scopes:
|
| 104 |
+
raise ScopeSchedulerError('scope scheduling produced no windows')
|
| 105 |
+
if scopes[0].core_span.start_s != 0 or not math.isclose(scopes[-1].core_span.end_s, duration_s):
|
| 106 |
+
raise ScopeSchedulerError('fine emit cores do not cover the complete source')
|
| 107 |
+
for left, right in pairwise(scopes):
|
| 108 |
+
if left.core_span.end_s != right.core_span.start_s:
|
| 109 |
+
raise ScopeSchedulerError('fine emit cores are not an exact partition')
|
| 110 |
+
return tuple(scopes)
|
| 111 |
+
|
| 112 |
+
def schedule_recentered_fine_retry_scopes(*, source_id: str, duration_s: float, component_edge_ticks: Sequence[tuple[int, int]], arm: str='adaptive_fine_repair', edge_core_offset_ticks: Mapping[int, int] | None=None, scope_namespace: str=FINE_RETRY_SCOPE_NAMESPACE, window_profile_id: str=DEFAULT_FINE_WINDOW_PROFILE_ID) -> tuple[ProcessingScope, ...]:
|
| 113 |
+
if not isinstance(source_id, str) or not source_id.strip():
|
| 114 |
+
raise ScopeSchedulerError('source_id must be non-empty')
|
| 115 |
+
if not isinstance(arm, str) or not arm.strip():
|
| 116 |
+
raise ScopeSchedulerError('arm must be non-empty')
|
| 117 |
+
try:
|
| 118 |
+
profile = fine_window_profile(window_profile_id)
|
| 119 |
+
except ValueError as exc:
|
| 120 |
+
raise ScopeSchedulerError(str(exc)) from exc
|
| 121 |
+
if not isinstance(scope_namespace, str) or not scope_namespace or any((not (character.isalnum() or character in '_-') for character in scope_namespace)):
|
| 122 |
+
raise ScopeSchedulerError('adaptive fine retry scope_namespace is invalid')
|
| 123 |
+
video_end_us = seconds_to_us(duration_s, label='duration_s')
|
| 124 |
+
if video_end_us <= 0:
|
| 125 |
+
raise ScopeSchedulerError('duration_s must be positive')
|
| 126 |
+
targets = tuple(component_edge_ticks)
|
| 127 |
+
if not targets:
|
| 128 |
+
raise ScopeSchedulerError('adaptive fine retry requires at least one component edge')
|
| 129 |
+
component_indices = [item[0] for item in targets]
|
| 130 |
+
if len(component_indices) != len(set(component_indices)):
|
| 131 |
+
raise ScopeSchedulerError('adaptive fine retry component indices must be unique')
|
| 132 |
+
offsets = dict(edge_core_offset_ticks or {})
|
| 133 |
+
if not set(offsets) <= set(component_indices):
|
| 134 |
+
raise ScopeSchedulerError('adaptive fine retry offsets contain an unknown component')
|
| 135 |
+
scopes: list[ProcessingScope] = []
|
| 136 |
+
for window_index, (component_index, edge_tick) in enumerate(targets):
|
| 137 |
+
if isinstance(component_index, bool) or not isinstance(component_index, int) or component_index < 0:
|
| 138 |
+
raise ScopeSchedulerError('adaptive fine retry component index must be non-negative')
|
| 139 |
+
if isinstance(edge_tick, bool) or not isinstance(edge_tick, int):
|
| 140 |
+
raise ScopeSchedulerError('adaptive fine retry edge tick must be an integer')
|
| 141 |
+
edge_core_offset_tick = offsets.get(component_index, floor_div(profile.core_us // 2, TICK_US))
|
| 142 |
+
if isinstance(edge_core_offset_tick, bool) or not isinstance(edge_core_offset_tick, int) or (not 0 <= edge_core_offset_tick <= profile.core_ticks):
|
| 143 |
+
raise ScopeSchedulerError('adaptive fine retry edge core offset must lie inside the emit core')
|
| 144 |
+
edge_us = edge_tick * TICK_US
|
| 145 |
+
if not 0 < edge_us < video_end_us:
|
| 146 |
+
raise ScopeSchedulerError('adaptive fine retry edge must be inside the source')
|
| 147 |
+
latest_full_core_start_us = floor_div(max(0, video_end_us - profile.core_us), TICK_US) * TICK_US
|
| 148 |
+
core_start_us = min(max(0, edge_us - edge_core_offset_tick * TICK_US), latest_full_core_start_us)
|
| 149 |
+
core_end_us = min(core_start_us + profile.core_us, video_end_us)
|
| 150 |
+
requested_start_us = core_start_us - profile.left_guard_us
|
| 151 |
+
requested_end_us = requested_start_us + profile.view_us
|
| 152 |
+
valid_start_us = max(0, requested_start_us)
|
| 153 |
+
valid_end_us = min(video_end_us, requested_end_us)
|
| 154 |
+
valid_start_tick = floor_div(valid_start_us - requested_start_us, TICK_US)
|
| 155 |
+
valid_end_tick = ceil_div(valid_end_us - requested_start_us, TICK_US)
|
| 156 |
+
core_start_tick = floor_div(core_start_us - requested_start_us, TICK_US)
|
| 157 |
+
core_end_tick = ceil_div(core_end_us - requested_start_us, TICK_US)
|
| 158 |
+
if not core_start_us <= edge_us <= core_end_us:
|
| 159 |
+
raise ScopeSchedulerError('adaptive fine retry edge does not lie in the emit core')
|
| 160 |
+
scope_id = f'{source_id}:{scope_namespace}:{window_index:06d}'
|
| 161 |
+
metadata = {'schema_version': WINDOW_SCHEMA, 'model_io_contract': profile.model_io_contract, 'window_profile_id': profile.profile_id, 'window_profile': profile.to_dict(), 'scheduler_contract': profile.retry_scheduler_contract(directional=component_index in offsets), 'arm': arm, 'guidance': 'no_prior_adaptive_retry', 'coarse_spans': [], 'coarse_scope_ids': [], 'coarse_span_lineage': [], 'window_index': window_index, 'retry_component_index': component_index, 'retry_edge_tick': edge_tick, 'retry_edge_core_offset_tick': edge_core_offset_tick, 'video_end_us': video_end_us, 'requested_view_start_us': requested_start_us, 'requested_view_end_us': requested_end_us, 'valid_start_us': valid_start_us, 'valid_end_us': valid_end_us, 'core_start_us': core_start_us, 'core_end_us': core_end_us, 'valid_start_tick': valid_start_tick, 'valid_end_tick': valid_end_tick, 'core_start_tick': core_start_tick, 'core_end_tick': core_end_tick}
|
| 162 |
+
scopes.append(ProcessingScope(scope_id=scope_id, source_id=source_id, core_span=TimeSpan(core_start_us / 1000000, core_end_us / 1000000), context_span=TimeSpan(valid_start_us / 1000000, valid_end_us / 1000000), metadata=metadata))
|
| 163 |
+
return tuple(scopes)
|
batch_utils/vendor/production_adapter/core/flora_runtime/worker_transport.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from collections.abc import Mapping
|
| 5 |
+
from typing import Any, Protocol, runtime_checkable
|
| 6 |
+
from flora_runtime.schema import AnnotationSource
|
| 7 |
+
|
| 8 |
+
class SubtaskAnnotationAdapterError(RuntimeError):
|
| 9 |
+
pass
|
| 10 |
+
|
| 11 |
+
@runtime_checkable
|
| 12 |
+
class JsonPredictor(Protocol):
|
| 13 |
+
|
| 14 |
+
def predict(self, payload: Mapping[str, Any]) -> Mapping[str, Any]:
|
| 15 |
+
...
|
| 16 |
+
|
| 17 |
+
def source_payload(source: AnnotationSource) -> dict[str, Any]:
|
| 18 |
+
return {'source_id': source.source_id, 'location': str(source.location), 'duration_s': source.duration_s, 'fps': source.fps, 'metadata': dict(source.metadata), 'video_streams': [{'camera': stream.camera, 'location': str(stream.location), 'fps': stream.fps, 'width': stream.width, 'height': stream.height, 'metadata': dict(stream.metadata)} for stream in source.video_streams]}
|
batch_utils/vendor/production_gdn_fix.py
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
def _keep_native_gdn_fp32(model_class: type) -> None:
|
| 2 |
+
model_class._keep_in_fp32_modules_strict = sorted(set(getattr(model_class, '_keep_in_fp32_modules_strict', None) or ()) | {'A_log', 'dt_bias'})
|
batch_utils/vendor/pts_input_adapter.py
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import bisect
|
| 2 |
+
from decimal import Decimal, ROUND_HALF_EVEN
|
| 3 |
+
|
| 4 |
+
def microseconds(value):
|
| 5 |
+
return int((Decimal(str(value)) * 1000000).to_integral_value(rounding=ROUND_HALF_EVEN))
|
| 6 |
+
|
| 7 |
+
def frame_span(pts_s, start_us, end_us):
|
| 8 |
+
if type(start_us) is not int or type(end_us) is not int or start_us < 0 or (end_us <= start_us):
|
| 9 |
+
raise ValueError('invalid integer-microsecond source interval')
|
| 10 |
+
times = [microseconds(t) for t in pts_s]
|
| 11 |
+
if not times or any((a >= b for a, b in zip(times, times[1:]))):
|
| 12 |
+
raise ValueError('PTS not strictly increasing after microsecond serialization')
|
| 13 |
+
return (bisect.bisect_left(times, start_us), bisect.bisect_left(times, end_us))
|
| 14 |
+
|
| 15 |
+
def native_blocks(pts_s, start_us, end_us, duration_us, rate, contract):
|
| 16 |
+
if end_us > duration_us:
|
| 17 |
+
raise ValueError('target beyond source')
|
| 18 |
+
windows = [('BEFORE', max(0, start_us - 2000000), start_us), ('TARGET', start_us, end_us), ('AFTER', end_us, min(duration_us, end_us + 2000000))]
|
| 19 |
+
blocks = []
|
| 20 |
+
for role, a, b in windows:
|
| 21 |
+
if a == b:
|
| 22 |
+
continue
|
| 23 |
+
lo, hi = frame_span(pts_s, a, b)
|
| 24 |
+
target = role == 'TARGET'
|
| 25 |
+
factor = contract.TARGET_FRAME_FACTOR if target else contract.CONTEXT_FRAME_FACTOR
|
| 26 |
+
if hi - lo < factor:
|
| 27 |
+
if target:
|
| 28 |
+
raise ValueError('target_has_fewer_than_native_frame_factor')
|
| 29 |
+
continue
|
| 30 |
+
ids, receipt = contract.duration_adaptive_sample(lo, hi, source_fps_num=rate.numerator, source_fps_den=rate.denominator, sample_fps=contract.TARGET_SAMPLE_FPS if target else contract.CONTEXT_SAMPLE_FPS, min_frames=contract.TARGET_MIN_FRAMES if target else contract.CONTEXT_MIN_FRAMES, max_frames=contract.TARGET_MAX_FRAMES if target else contract.CONTEXT_MAX_FRAMES, frame_factor=factor)
|
| 31 |
+
assert len(set(ids)) == len(ids) and all((lo <= i < hi for i in ids))
|
| 32 |
+
times = [microseconds(pts_s[i]) for i in ids]
|
| 33 |
+
assert all((a <= t < b for t in times))
|
| 34 |
+
blocks.append({'role': role, 'interval_us': [a, b], 'start_frame': lo, 'end_frame_exclusive': hi, 'frame_indices': ids, 'actual_pts_us': times, 'sampling_receipt': receipt, 'model_timestamp_clock': 'native ordinal/average FPS; exact source PTS retained separately'})
|
| 35 |
+
assert sum((len(b['frame_indices']) for b in blocks)) <= contract.MAX_TOTAL_FRAMES
|
| 36 |
+
return blocks
|
batch_utils/vendor/robot_caption_geometry.py
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import hashlib, io
|
| 2 |
+
from PIL import Image
|
| 3 |
+
PACKED_IMAGE_LIMIT_BYTES = 21_866_496
|
| 4 |
+
SCALES = [1.0, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2]
|
| 5 |
+
|
| 6 |
+
def sha(b):
|
| 7 |
+
return hashlib.sha256(b).hexdigest()
|
| 8 |
+
|
| 9 |
+
def png(im):
|
| 10 |
+
b = io.BytesIO()
|
| 11 |
+
im.save(b, format='PNG', compress_level=1, optimize=False)
|
| 12 |
+
return b.getvalue()
|
| 13 |
+
|
| 14 |
+
def normalized(plan, frames, expected=None):
|
| 15 |
+
factors = [expected['common_geometry']['scale']] if expected else SCALES
|
| 16 |
+
for factor in factors:
|
| 17 |
+
blocks = []
|
| 18 |
+
count = 0
|
| 19 |
+
for block in plan['blocks']:
|
| 20 |
+
assets = []
|
| 21 |
+
for start in range(0, len(block['indices']), 2):
|
| 22 |
+
indices = block['indices'][start:start + 2]
|
| 23 |
+
originals = [frames[i] for i in indices]
|
| 24 |
+
w, h = originals[0].size
|
| 25 |
+
nw = round(w * factor)
|
| 26 |
+
nh = round(h * factor)
|
| 27 |
+
packed = Image.new('RGB', (nw * len(indices), nh))
|
| 28 |
+
rgb = []
|
| 29 |
+
for pos, image in enumerate(originals):
|
| 30 |
+
if factor != 1:
|
| 31 |
+
image = image.resize((nw, nh), Image.Resampling.LANCZOS)
|
| 32 |
+
packed.paste(image, (pos * nw, 0))
|
| 33 |
+
rgb.append(sha(image.tobytes()))
|
| 34 |
+
blob = png(packed)
|
| 35 |
+
count += len(blob) + (-len(blob) % 3)
|
| 36 |
+
assets.append(dict(indices=indices, size=list(packed.size), sha256=sha(blob), rgb_sha256=rgb, bytes=len(blob), _blob=blob))
|
| 37 |
+
blocks.append(dict(role=block['role'], images=assets))
|
| 38 |
+
if count <= PACKED_IMAGE_LIMIT_BYTES:
|
| 39 |
+
break
|
| 40 |
+
else:
|
| 41 |
+
raise RuntimeError('No common caption input geometry fits')
|
| 42 |
+
if expected:
|
| 43 |
+
for old, new in zip(expected['frame_blocks'], blocks):
|
| 44 |
+
assert old['role'] == new['role'] and len(old['images']) == len(new['images'])
|
| 45 |
+
for a, b in zip(old['images'], new['images']):
|
| 46 |
+
assert a['indices'] == b['indices'] and a['rgb_sha256'] == b['rgb_sha256'] and (a['size'] == b['size'])
|
| 47 |
+
return dict(plan, frame_blocks=blocks, common_geometry=dict(scale=factor, source_pixel_unchanged=factor == 1, aligned_image_bytes=count, image_size_limit_bytes=PACKED_IMAGE_LIMIT_BYTES, sampling_frames_removed=0))
|
| 48 |
+
|
| 49 |
+
def source_frames(normalized):
|
| 50 |
+
out = {}
|
| 51 |
+
for block in normalized['frame_blocks']:
|
| 52 |
+
for a in block['images']:
|
| 53 |
+
image = Image.open(io.BytesIO(a['_blob'])).convert('RGB')
|
| 54 |
+
n = len(a['indices'])
|
| 55 |
+
w = image.width // n
|
| 56 |
+
for pos, i in enumerate(a['indices']):
|
| 57 |
+
frame = image.crop((pos * w, 0, (pos + 1) * w, image.height))
|
| 58 |
+
assert sha(frame.tobytes()) == a['rgb_sha256'][pos]
|
| 59 |
+
out[i] = frame
|
| 60 |
+
return out
|
batch_utils/vendor/support/canonical_rgb.py
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
from PIL import Image
|
| 3 |
+
|
| 4 |
+
def canonical_rgb(frame):
|
| 5 |
+
assert frame.format.name == 'yuv420p'
|
| 6 |
+
assert frame.width % 16 == 0 and frame.height % 2 == 0
|
| 7 |
+
assert int(frame.colorspace) in (2, 5, 6), ('unsupported colorspace', frame.colorspace)
|
| 8 |
+
assert int(frame.color_range) in (0, 1), ('unsupported range', frame.color_range)
|
| 9 |
+
planes = []
|
| 10 |
+
for p in frame.planes:
|
| 11 |
+
a = np.frombuffer(p, dtype=np.uint8).reshape(p.height, p.line_size)[:, :p.width].astype(np.int32)
|
| 12 |
+
planes.append(a)
|
| 13 |
+
y, u, v = planes
|
| 14 |
+
u = u.repeat(2, axis=0).repeat(2, axis=1)
|
| 15 |
+
v = v.repeat(2, axis=0).repeat(2, axis=1)
|
| 16 |
+
|
| 17 |
+
def q13(c):
|
| 18 |
+
return c * 8192 + 32768 >> 16
|
| 19 |
+
y = (y * 8 - 128) * q13(65536 * 255 // 219) >> 16
|
| 20 |
+
u = u * 8 - 1024
|
| 21 |
+
v = v * 8 - 1024
|
| 22 |
+
r = y + (v * q13(104597) >> 16)
|
| 23 |
+
g = y + (u * q13(-25675) >> 16) + (v * q13(-53279) >> 16)
|
| 24 |
+
b = y + (u * q13(132201) >> 16)
|
| 25 |
+
return Image.fromarray(np.clip(np.stack([r, g, b], axis=-1), 0, 255).astype(np.uint8))
|
batch_utils/vendor/support/composers.py
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from PIL import Image, ImageOps
|
| 2 |
+
from typing import Any
|
| 3 |
+
VIEWS = ('ego', 'left_wrist', 'right_wrist')
|
| 4 |
+
CAPTION_TOP_SIZE = (768, 432)
|
| 5 |
+
CAPTION_WRIST_SIZE = (384, 216)
|
| 6 |
+
CAPTION_MOSAIC_WIDTH = 768
|
| 7 |
+
CAPTION_MOSAIC_HEIGHT = 648
|
| 8 |
+
|
| 9 |
+
def compose_frame(frames, views):
|
| 10 |
+
canvas = Image.new('RGB', (640, 720))
|
| 11 |
+
for view, (x, y, w, h) in zip(VIEWS, [(0, 0, 640, 480), (0, 480, 320, 240), (320, 480, 320, 240)]):
|
| 12 |
+
if view in views:
|
| 13 |
+
panel = ImageOps.contain(frames[view].convert('RGB'), (w, h), Image.Resampling.BICUBIC)
|
| 14 |
+
canvas.paste(panel, (x + (w - panel.width) // 2, y + (h - panel.height) // 2))
|
| 15 |
+
return canvas
|
| 16 |
+
|
| 17 |
+
def compose_top_dominant_mosaic(top: Any, left: Any, right: Any) -> Any:
|
| 18 |
+
from PIL import Image
|
| 19 |
+
top_panel = top.convert('RGB').resize(CAPTION_TOP_SIZE, Image.Resampling.BICUBIC)
|
| 20 |
+
left_panel = left.convert('RGB').resize(CAPTION_WRIST_SIZE, Image.Resampling.BICUBIC)
|
| 21 |
+
right_panel = right.convert('RGB').resize(CAPTION_WRIST_SIZE, Image.Resampling.BICUBIC)
|
| 22 |
+
output = Image.new('RGB', (CAPTION_MOSAIC_WIDTH, CAPTION_MOSAIC_HEIGHT))
|
| 23 |
+
output.paste(top_panel, (0, 0))
|
| 24 |
+
output.paste(left_panel, (0, CAPTION_TOP_SIZE[1]))
|
| 25 |
+
output.paste(right_panel, (CAPTION_WRIST_SIZE[0], CAPTION_TOP_SIZE[1]))
|
| 26 |
+
return output
|
batch_utils/vendor/support/plans.py
ADDED
|
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import bisect, hashlib
|
| 2 |
+
from fractions import Fraction
|
| 3 |
+
import u2_contract as U2
|
| 4 |
+
ROLES = ('main', 'left_wrist', 'right_wrist')
|
| 5 |
+
HEADING = 'Each frame is a synchronized camera mosaic: main camera above, left wrist camera below left, and right wrist camera below right. All panels show the same physical time. Describe each physical operation once across the views.\n\n'
|
| 6 |
+
INSTRUCTION = 'Write one concise, self-contained English instruction for the local operation being attempted in TARGET.\nUse a natural imperative, concrete objects and the working hand; include meaningful support by the other hand and necessary source or destination. Say left hand, right hand or both hands. For a handover, name the object and both the giving and receiving hands.\nInfer the local goal from the video sequence. Use BEFORE and AFTER only to disambiguate TARGET, not to add operations occurring outside it. A failed or repeated attempt retains its established instruction; do not replace the goal with an accident, a success claim or a description of mechanical adjustments. Do not invent an object, hand role or hidden goal.\nOutput only the instruction, on one line, without timestamps, execution-quality commentary or the block names.'
|
| 7 |
+
|
| 8 |
+
def plans_for_episode(ep, clocks):
|
| 9 |
+
clock = clocks['main']
|
| 10 |
+
pts = [float(t * Fraction(clock['time_base'])) for t in clock['ticks']]
|
| 11 |
+
fps = Fraction(clock['fps'])
|
| 12 |
+
plans = []
|
| 13 |
+
for unit in ep['units']:
|
| 14 |
+
start, end = unit['native_scope_s']
|
| 15 |
+
assert 0 <= start < end <= clock['duration_s'] + 1e-05
|
| 16 |
+
blocks = []
|
| 17 |
+
spans = [('BEFORE', max(0.0, start - 2.0), start), ('TARGET', start, end), ('AFTER', end, min(clock['duration_s'], end + 2.0))]
|
| 18 |
+
for role, a, b in spans:
|
| 19 |
+
lo, hi = (bisect.bisect_left(pts, a - 1e-09), bisect.bisect_left(pts, b - 1e-09))
|
| 20 |
+
hi = min(hi, len(pts))
|
| 21 |
+
factor = U2.TARGET_FRAME_FACTOR if role == 'TARGET' else U2.CONTEXT_FRAME_FACTOR
|
| 22 |
+
if hi - lo < factor:
|
| 23 |
+
if role != 'TARGET':
|
| 24 |
+
continue
|
| 25 |
+
assert hi - lo == 1, (unit['target_id'], 'empty target')
|
| 26 |
+
indices = [lo, lo]
|
| 27 |
+
receipt = U2.repeated_single_target_receipt(lo, source_fps_num=fps.numerator, source_fps_den=fps.denominator)
|
| 28 |
+
padding = [False, True]
|
| 29 |
+
else:
|
| 30 |
+
target = role == 'TARGET'
|
| 31 |
+
indices, receipt = U2.duration_adaptive_sample(lo, hi, source_fps_num=fps.numerator, source_fps_den=fps.denominator, sample_fps=U2.TARGET_SAMPLE_FPS if target else U2.CONTEXT_SAMPLE_FPS, min_frames=U2.TARGET_MIN_FRAMES if target else U2.CONTEXT_MIN_FRAMES, max_frames=U2.TARGET_MAX_FRAMES if target else U2.CONTEXT_MAX_FRAMES, frame_factor=factor)
|
| 32 |
+
padding = [False] * len(indices)
|
| 33 |
+
blocks.append(dict(role=role, scope_s=[a, b], source_start_frame=lo, source_end_frame_exclusive=hi, indices=indices, pts_ticks=[clock['ticks'][i] for i in indices], padding_mask=padding, sampling_receipt=receipt))
|
| 34 |
+
assert any((b['role'] == 'TARGET' for b in blocks))
|
| 35 |
+
assert sum((len(b['indices']) for b in blocks)) <= U2.MAX_TOTAL_FRAMES
|
| 36 |
+
prompt = HEADING + '\n\n'.join((b['role'] + ':\n<video>' for b in blocks)) + '\n\n' + INSTRUCTION
|
| 37 |
+
plans.append(dict(target_id=unit['target_id'], unit_id=unit['unit_id'], episode_id=ep['episode_id'], source_id=ep['source_id'], group=ep['group'], native_scope_s=unit['native_scope_s'], camera_roles=list(ROLES), clock_fps=str(fps), time_base=clock['time_base'], blocks=blocks, prompt=prompt, prompt_sha256=hashlib.sha256(prompt.encode()).hexdigest(), context_policy='common_real_up_to_2s_each_side', source_clock_origin=clock['first_pts'], do_sample_frames=False))
|
| 38 |
+
return plans
|
batch_utils/vendor/support/u2_contract.py
ADDED
|
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
| 2 |
+
# SPDX-License-Identifier: Apache-2.0
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
from fractions import Fraction
|
| 5 |
+
TARGET_SAMPLE_FPS = 8
|
| 6 |
+
TARGET_MIN_FRAMES = 8
|
| 7 |
+
TARGET_MAX_FRAMES = 40
|
| 8 |
+
TARGET_FRAME_FACTOR = 2
|
| 9 |
+
CONTEXT_SAMPLE_FPS = 8
|
| 10 |
+
CONTEXT_MIN_FRAMES = 4
|
| 11 |
+
CONTEXT_MAX_FRAMES = 12
|
| 12 |
+
CONTEXT_FRAME_FACTOR = 4
|
| 13 |
+
MAX_TOTAL_FRAMES = 64
|
| 14 |
+
SAMPLER_VERSION = 'duration_adaptive_uniform_bin_centers_nested_context_v2'
|
| 15 |
+
SHORT_TARGET_POLICY = 'repeat_single_source_frame_to_temporal_pair_v1'
|
| 16 |
+
|
| 17 |
+
class U2ContractError(ValueError):
|
| 18 |
+
pass
|
| 19 |
+
|
| 20 |
+
def uniform_bin_centers(start: int, end_exclusive: int, count: int) -> list[int]:
|
| 21 |
+
length = end_exclusive - start
|
| 22 |
+
if count <= 0:
|
| 23 |
+
return []
|
| 24 |
+
if length < count:
|
| 25 |
+
raise U2ContractError(f'span [{start}, {end_exclusive}) has fewer than {count} frames')
|
| 26 |
+
result = [start + (2 * index + 1) * length // (2 * count) for index in range(count)]
|
| 27 |
+
if len(result) != len(set(result)) or result[0] < start or result[-1] >= end_exclusive:
|
| 28 |
+
raise AssertionError('uniform frame sampler violated its contract')
|
| 29 |
+
return result
|
| 30 |
+
|
| 31 |
+
def duration_adaptive_sample(start: int, end_exclusive: int, *, source_fps_num: int, source_fps_den: int, sample_fps: int, min_frames: int, max_frames: int, frame_factor: int) -> tuple[list[int], dict[str, int | bool | str]]:
|
| 32 |
+
available = end_exclusive - start
|
| 33 |
+
if source_fps_num <= 0 or source_fps_den <= 0 or available <= 0:
|
| 34 |
+
raise U2ContractError('source FPS and frame span must be positive')
|
| 35 |
+
if min_frames % frame_factor or max_frames % frame_factor:
|
| 36 |
+
raise U2ContractError('adaptive frame bounds must align to frame factor')
|
| 37 |
+
duration = Fraction(available * source_fps_den, source_fps_num)
|
| 38 |
+
requested = min(Fraction(max_frames), max(Fraction(min_frames), duration * sample_fps))
|
| 39 |
+
desired_count = requested.numerator // (requested.denominator * frame_factor) * frame_factor
|
| 40 |
+
actual_count = min(desired_count, available // frame_factor * frame_factor)
|
| 41 |
+
if actual_count < frame_factor:
|
| 42 |
+
raise U2ContractError(f'span [{start}, {end_exclusive}) has fewer than {frame_factor} frames')
|
| 43 |
+
indices = uniform_bin_centers(start, end_exclusive, actual_count)
|
| 44 |
+
return (indices, {'policy': 'duration_fps_bounded_unique_frames_v1', 'sampler_version': SAMPLER_VERSION, 'source_fps_num': source_fps_num, 'source_fps_den': source_fps_den, 'sample_fps_num': sample_fps, 'sample_fps_den': 1, 'min_frames': min_frames, 'max_frames': max_frames, 'frame_factor': frame_factor, 'available_unique_source_frames': available, 'desired_frame_count': desired_count, 'actual_frame_count': actual_count, 'source_limited': actual_count < desired_count, 'selection': 'uniform_bin_centers'})
|
| 45 |
+
|
| 46 |
+
def repeated_single_target_receipt(frame_index: int, *, source_fps_num: int, source_fps_den: int) -> dict[str, int | str]:
|
| 47 |
+
if frame_index < 0 or source_fps_num <= 0 or source_fps_den <= 0:
|
| 48 |
+
raise U2ContractError('single-target fallback received invalid source metadata')
|
| 49 |
+
return {'policy': SHORT_TARGET_POLICY, 'source_fps_num': source_fps_num, 'source_fps_den': source_fps_den, 'source_frame_index': frame_index, 'available_unique_source_frames': 1, 'actual_frame_count': TARGET_FRAME_FACTOR, 'repeated_frame_count': TARGET_FRAME_FACTOR}
|
batch_utils/vendor/tempseg_fp32_loader.py
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from collections import Counter
|
| 2 |
+
import hashlib
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import types
|
| 5 |
+
|
| 6 |
+
def digest(path):
|
| 7 |
+
with Path(path).open('rb') as stream:
|
| 8 |
+
return hashlib.file_digest(stream, 'sha256').hexdigest()
|
| 9 |
+
|
| 10 |
+
def load(spec, save):
|
| 11 |
+
import torch
|
| 12 |
+
from transformers import AutoModelForImageTextToText
|
| 13 |
+
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5PreTrainedModel
|
| 14 |
+
from safetensors.torch import load_file
|
| 15 |
+
from peft import PeftConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
|
| 16 |
+
from production_gdn_fix import _keep_native_gdn_fp32
|
| 17 |
+
from verify_gdn import verify
|
| 18 |
+
_keep_native_gdn_fp32(Qwen3_5PreTrainedModel)
|
| 19 |
+
base, info = AutoModelForImageTextToText.from_pretrained(spec['base'], local_files_only=True, dtype=torch.bfloat16, attn_implementation='sdpa', output_loading_info=True)
|
| 20 |
+
assert not any((info.get(k) for k in ['missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs'])), info
|
| 21 |
+
path = Path(spec['export']) / 'adapter_model.safetensors'
|
| 22 |
+
assert digest(path) == spec['adapter_sha256']
|
| 23 |
+
source = load_file(str(path), device='cpu')
|
| 24 |
+
assert len(source) == 496 and sum((v.numel() for v in source.values())) == 86556672
|
| 25 |
+
assert all((v.dtype == torch.float32 for v in source.values())), 'Expected corrected FP32 checkpoint'
|
| 26 |
+
model = get_peft_model(base, PeftConfig.from_pretrained(spec['export'], local_files_only=True), adapter_name='candidate', autocast_adapter_dtype=True)
|
| 27 |
+
initial = get_peft_model_state_dict(model, adapter_name='candidate')
|
| 28 |
+
assert set(initial) == set(source)
|
| 29 |
+
assert all((v.dtype == torch.float32 for v in initial.values())), 'PEFT must create FP32 adapter holders'
|
| 30 |
+
del initial
|
| 31 |
+
load_info = set_peft_model_state_dict(model, source, adapter_name='candidate')
|
| 32 |
+
assert not load_info.unexpected_keys, load_info
|
| 33 |
+
actual = get_peft_model_state_dict(model, adapter_name='candidate')
|
| 34 |
+
assert set(actual) == set(source)
|
| 35 |
+
assert all((actual[k].dtype == v.dtype and torch.equal(actual[k].detach().cpu(), v) for k, v in source.items())), 'Saved FP32 LoRA values must be loaded exactly'
|
| 36 |
+
del actual
|
| 37 |
+
model.to('cuda:0').eval()
|
| 38 |
+
actual = get_peft_model_state_dict(model, adapter_name='candidate')
|
| 39 |
+
assert all((actual[k].dtype == torch.float32 and torch.equal(actual[k].detach().cpu(), v) for k, v in source.items())), 'Final device placement must preserve saved FP32 values'
|
| 40 |
+
save('ADAPTER_VALUE_VERIFICATION.json', dict(tensors=496, elements=86556672, source_dtypes=dict(Counter((str(v.dtype) for v in source.values()))), holders_fp32_before_assignment=True, all_saved_values_exact_before_and_after_placement=True, adapter_sha256=spec['adapter_sha256'], merged=False))
|
| 41 |
+
del source, actual
|
| 42 |
+
gdn = verify(model, spec['base'], peft_prefix=True)
|
| 43 |
+
save('GDN_VALUE_VERIFICATION.json', gdn)
|
| 44 |
+
original_generate = model.generate
|
| 45 |
+
|
| 46 |
+
def generate_bf16(self, *args, **kwargs):
|
| 47 |
+
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
|
| 48 |
+
return original_generate(*args, **kwargs)
|
| 49 |
+
model.generate = types.MethodType(generate_bf16, model)
|
| 50 |
+
save('MODEL_LOADED.json', dict(adapter_persistent_dtype='torch.float32', compute_autocast='torch.bfloat16', gdn_values_exact=True, base_loading_info=info, parameter_dtypes=dict(Counter((str(p.dtype) for p in model.parameters()))), input_output_generation_policy_changed=False))
|
| 51 |
+
return model
|
batch_utils/vendor/verify_gdn.py
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from collections import Counter, defaultdict
|
| 2 |
+
import hashlib, json
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
|
| 5 |
+
def verify(model, path, *, peft_prefix=False):
|
| 6 |
+
import torch
|
| 7 |
+
from safetensors import safe_open
|
| 8 |
+
path = Path(path)
|
| 9 |
+
mapping = json.loads((path / 'model.safetensors.index.json').read_text())['weight_map']
|
| 10 |
+
wanted = {name for name in mapping if name.endswith(('.A_log', '.dt_bias'))}
|
| 11 |
+
assert wanted
|
| 12 |
+
actual = {}
|
| 13 |
+
for name, p in model.named_parameters():
|
| 14 |
+
if not name.endswith(('.A_log', '.dt_bias')):
|
| 15 |
+
continue
|
| 16 |
+
if peft_prefix:
|
| 17 |
+
assert name.startswith('base_model.model.')
|
| 18 |
+
name = name[len('base_model.model.'):]
|
| 19 |
+
assert name not in actual
|
| 20 |
+
actual[name] = p
|
| 21 |
+
assert set(actual) == wanted, {'missing': sorted(wanted - set(actual)), 'extra': sorted(set(actual) - wanted)}
|
| 22 |
+
groups = defaultdict(list)
|
| 23 |
+
for name in wanted:
|
| 24 |
+
groups[mapping[name]].append(name)
|
| 25 |
+
evidence = []
|
| 26 |
+
for filename, names in sorted(groups.items()):
|
| 27 |
+
with safe_open(str(path / filename), framework='pt', device='cpu') as source:
|
| 28 |
+
for name in sorted(names):
|
| 29 |
+
saved = source.get_tensor(name)
|
| 30 |
+
p = actual[name]
|
| 31 |
+
assert p.dtype == torch.float32, (name, str(p.dtype), 'must preserve FP32 before loading')
|
| 32 |
+
assert p.device.type == 'cuda', (name, 'must verify final device placement')
|
| 33 |
+
value = p.detach().cpu()
|
| 34 |
+
assert torch.equal(value, saved.float()), (name, 'checkpoint bits changed; upcasting rounded BF16 is insufficient')
|
| 35 |
+
evidence.append(dict(name=name, elements=p.numel(), source_dtype=str(saved.dtype), runtime_dtype=str(p.dtype), device=str(p.device), fp32_value_sha256=hashlib.sha256(value.contiguous().view(torch.uint8).numpy().tobytes()).hexdigest(), saved_value_exact=True))
|
| 36 |
+
return dict(status='PASS_NATIVE_GDN_FP32_CHECKPOINT_VALUES', tensors=len(evidence), source_dtypes=dict(Counter((r['source_dtype'] for r in evidence))), all_runtime_fp32=True, all_values_equal_saved_checkpoint=True, checked_after_final_device_placement=True, values=evidence)
|
infer.py
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run the complete paired Data Engine workflow from either model repository."""
|
| 2 |
+
from batch_utils.cli import main
|
| 3 |
+
|
| 4 |
+
if __name__ == "__main__":
|
| 5 |
+
raise SystemExit(main())
|