feat: Add AI Refine option for Balloon tool and implement refinement process with BEN2

tests/test_server.py

@@ -215,6 +215,7 @@ class ServerTest(unittest.TestCase):
             {
                 "download-model",
                 "download-required-models",
+                "refine-balloon",
                 "remove-background",
                 "extract-frames",
                 "generate-depth",
@@ -381,6 +382,92 @@ class ServerTest(unittest.TestCase):
 
         self.assertEqual(status["state"], "succeeded")
 
+    def test_balloon_lasso_input_crops_and_clears_outside_rgba(self):
+        import numpy as np
+        from PIL import Image
+
+        balloon = sys.modules["server.jobs.balloon"]
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            source_path = root / "source.png"
+            output_path = root / "lasso.png"
+            pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8)
+            pixels[1, 1, 3] = 64
+            Image.fromarray(pixels, mode="RGBA").save(source_path)
+
+            _path, _selection, origin = balloon.create_lasso_input(
+                source_path,
+                ((1, 1), (5, 1), (1, 5)),
+                output_path,
+            )
+            result = np.asarray(Image.open(output_path).convert("RGBA"))
+
+        self.assertEqual(origin, (1, 1))
+        self.assertEqual(result.shape, (4, 4, 4))
+        self.assertEqual(tuple(result[0, 0]), (20, 40, 60, 64))
+        self.assertEqual(tuple(result[3, 3]), (0, 0, 0, 0))
+
+    def test_balloon_job_sends_the_masked_crop_to_ben2(self):
+        import numpy as np
+        from PIL import Image
+
+        balloon = sys.modules["server.jobs.balloon"]
+        source_path = self.storage_root / "source.png"
+        pixels = np.full((6, 6, 4), (20, 40, 60, 255), dtype=np.uint8)
+        Image.fromarray(pixels, mode="RGBA").save(source_path)
+        job_directory = self.storage_root / "job"
+        context = SimpleNamespace(
+            directory=job_directory,
+            progress=lambda *_args: None,
+            check_cancelled=lambda: None,
+            resource=lambda _name: SimpleNamespace(
+                directory=lambda _model: self.storage_root / "model"
+            ),
+        )
+        received = []
+
+        def infer(input_path, *_args, **_kwargs):
+            image = Image.open(input_path).convert("RGBA")
+            received.append(np.asarray(image))
+            return image, 0.0, 1.0
+
+        with patch.object(balloon, "infer_foreground", side_effect=infer):
+            result = balloon.run(
+                context,
+                {
+                    "input": str(source_path),
+                    "polygon": ((1, 1), (5, 1), (1, 5)),
+                    "model": "BEN2_BASE",
+                    "device": "cpu",
+                },
+            )
+
+        self.assertEqual(result, {"contour": "contour.json"})
+        self.assertEqual(tuple(received[0][3, 3]), (0, 0, 0, 0))
+        self.assertTrue((job_directory / "lasso-input.png").is_file())
+        self.assertTrue((job_directory / "foreground.png").is_file())
+        self.assertTrue((job_directory / "contour.json").is_file())
+
+    def test_balloon_contour_uses_alpha_above_half_and_largest_component(self):
+        import numpy as np
+
+        balloon = sys.modules["server.jobs.balloon"]
+        alpha = np.zeros((6, 7), dtype=np.uint8)
+        alpha[0, 0] = 255
+        alpha[1:4, 2:5] = 128
+        alpha[1, 1] = 127
+        selection_mask = np.ones_like(alpha, dtype=bool)
+
+        contour = balloon.contour_from_alpha(
+            alpha,
+            selection_mask,
+            origin=(10, 20),
+        )
+
+        self.assertEqual(
+            set(map(tuple, contour)),
+            {(12, 21), (15, 21), (15, 24), (12, 24)},
+        )
 
 if __name__ == "__main__":
     unittest.main()