Add H3 scene image plan helpers

2026-08-12 13:57:05 +00:00
parent ba31c503b5
commit a2a9ff90f6
3 changed files with 328 additions and 0 deletions
@@ -1,4 +1,5 @@
 import importlib
+import json
 import os
 import sys
 import tempfile
@@ -318,6 +319,93 @@ class DumasImageNodeTests(unittest.TestCase):
        self.assertEqual(result["ui"]["images"], [])
        self.assertEqual(FakePILImage.saved_paths, [])

+    def test_h3_plan_scene_images_attach_and_extract(self):
+        attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
+        extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
+        plan = {"shots": [{"id": "intro"}, {"id": "middle"}]}
+        image_a = FakeTensorBatch()
+        image_b = FakeTensorBatch(width=10, height=10)
+
+        attached_plan, connected = attach_node.attach(
+            plan=plan,
+            scene_index=2,
+            image1=image_a,
+            image3=image_b,
+        )
+        extracted = extract_node.extract(attached_plan, 2)
+
+        self.assertEqual(connected, 2)
+        self.assertEqual(attached_plan["_dumas_scene_image_bindings"]["scene_counts"], {"2": 2})
+        self.assertIs(extracted[1], image_a)
+        self.assertIsNone(extracted[2])
+        self.assertIs(extracted[3], image_b)
+        self.assertEqual(extracted[-1], 2)
+        self.assertNotIn("_dumas_scene_image_bindings", plan)
+
+    def test_h3_plan_scene_images_support_multiple_scenes(self):
+        attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
+        extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
+        plan = {"shots": [{"id": "one"}, {"id": "two"}]}
+        image_1 = FakeTensorBatch()
+        image_2 = FakeTensorBatch(width=12, height=9)
+
+        plan_after_first, _connected = attach_node.attach(plan=plan, scene_index=1, image2=image_1)
+        plan_after_second, _connected = attach_node.attach(
+            plan=plan_after_first,
+            scene_index=2,
+            image6=image_2,
+        )
+
+        scene1 = extract_node.extract(plan_after_second, 1)
+        scene2 = extract_node.extract(plan_after_second, 2)
+
+        self.assertIs(scene1[2], image_1)
+        self.assertEqual(scene1[-1], 1)
+        self.assertIs(scene2[6], image_2)
+        self.assertEqual(scene2[-1], 1)
+        self.assertEqual(
+            plan_after_second["_dumas_scene_image_bindings"]["scene_counts"],
+            {"1": 1, "2": 1},
+        )
+
+    def test_h3_plan_scene_images_metadata_is_json_serializable(self):
+        attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
+        plan = {"shots": [{"id": "one"}]}
+
+        attached_plan, connected = attach_node.attach(
+            plan=plan,
+            scene_index=1,
+            image4=FakeTensorBatch(),
+        )
+
+        self.assertEqual(connected, 1)
+        json.dumps(attached_plan)
+
+    def test_h3_plan_scene_images_can_clear_a_scene_binding(self):
+        attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
+        extract_node = self.image_nodes.DumasH3PlanExtractSceneImagesNode()
+        plan = {"shots": [{"id": "one"}]}
+
+        attached_plan, connected = attach_node.attach(
+            plan=plan,
+            scene_index=1,
+            image1=FakeTensorBatch(),
+        )
+        cleared_plan, cleared = attach_node.attach(plan=attached_plan, scene_index=1)
+        extracted = extract_node.extract(cleared_plan, 1)
+
+        self.assertEqual(connected, 1)
+        self.assertEqual(cleared, 0)
+        self.assertNotIn("_dumas_scene_image_bindings", cleared_plan)
+        self.assertEqual(extracted[-1], 0)
+
+    def test_h3_plan_scene_images_reject_invalid_scene_index(self):
+        attach_node = self.image_nodes.DumasH3PlanAttachSceneImagesNode()
+        plan = {"shots": [{"id": "one"}]}
+
+        with self.assertRaises(ValueError):
+            attach_node.attach(plan=plan, scene_index=2, image1=FakeTensorBatch())
+

 if __name__ == "__main__":
    unittest.main()