diff --git a/CHANGELOG.md b/CHANGELOG.md index 9ad920d..38f91fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ Contributors add user-facing entries under `[Unreleased]` in the same PR. Mainta ## [Unreleased] +### Changed + +- **`creative/bg_remover`**: Hardened the skill by reusing rembg sessions across executions, validating Base64 input and image files, rejecting empty, oversized (>25 MB), or invalid images, and automatically creating parent directories for `output_path`. Updated documentation and expanded bundle tests for the new validation behavior. + ## [0.4.6] - 2026-07-23 ### Added diff --git a/docs/skills/bg_remover.md b/docs/skills/bg_remover.md index 36823a3..d8e1960 100644 --- a/docs/skills/bg_remover.md +++ b/docs/skills/bg_remover.md @@ -123,7 +123,7 @@ Skill instructions: when to invoke, input/output conventions, URL and cloud pre- ### The Body (`skill.py`) -Lazy-imports `rembg` and Pillow, runs `new_session` + `remove`, and returns structured JSON with transparent PNG bytes. +Lazy-imports `rembg` and Pillow, reuses cached rembg sessions, runs `remove`, and returns structured JSON with transparent PNG bytes. ## Usage Examples diff --git a/skills/creative/bg_remover/instructions.md b/skills/creative/bg_remover/instructions.md index ee0e3c6..ac952c8 100644 --- a/skills/creative/bg_remover/instructions.md +++ b/skills/creative/bg_remover/instructions.md @@ -32,6 +32,8 @@ For S3, GCS, Azure Blob, or Cloudflare R2: download the object to a temp file On the **first run** in a fresh environment, processing may take a few minutes while rembg downloads the ONNX model (~176 MB for `isnet-general-use`). Later runs reuse the cache and are much faster. +Background removal sessions are also reused across executions for the same model, reducing repeated initialization overhead. + ## Input (one required: `image` OR `input_path`) | Scenario | Parameter | @@ -43,6 +45,14 @@ On the **first run** in a fresh environment, processing may take a few minutes w If **both** `image` and `input_path` are sent, **`image` wins** (do not double-submit). +### Input validation + +- Invalid Base64 payloads are rejected. +- `input_path` must reference an existing file (directories are rejected). +- Empty image files are rejected. +- Images larger than **25 MB** are rejected. +- Image integrity is validated before processing. + ### Example payloads Chat / attachment: @@ -91,6 +101,8 @@ Omit `model` for default `isnet-general-use`. Set `alpha_matting` only when edge | Chat or API only | Omit `output_path`; use `image_base64` from the result (always present on success). | | Save next to original | Same directory, new name (e.g. `1223_no_bg.png`). | +If the parent directory of `output_path` does not exist, it is created automatically before writing the PNG. + ## Interpreting a successful result When `success` is `true`: @@ -109,10 +121,12 @@ Runtime: `rembg`, `pillow`, `onnxruntime`. Install: `pip install "skillware[crea First `execute()` downloads the ONNX model to the rembg cache (`~/.u2net/` on Linux/macOS, `%USERPROFILE%\.u2net\` on Windows). +Subsequent executions reuse cached rembg sessions for the selected model, reducing repeated initialization overhead. + ## Errors | `error_code` | Response | | :--- | :--- | -| `INVALID_INPUT` | Ask for an upload, attachment, base64 payload, or local `input_path`. | +| `INVALID_INPUT` | Invalid Base64, missing input, directory path, empty image, oversized image (>25 MB), corrupt image, or unsupported input. | | `MISSING_DEPENDENCY` | Ask the user to run `pip install "skillware[creative_bg_remover]"` (or `pip install rembg pillow onnxruntime`), then retry. | | `PROCESSING_FAILED` | Surface the `error` string; input may be missing, corrupt, or unsupported. | diff --git a/skills/creative/bg_remover/manifest.yaml b/skills/creative/bg_remover/manifest.yaml index 3e59e81..0160641 100644 --- a/skills/creative/bg_remover/manifest.yaml +++ b/skills/creative/bg_remover/manifest.yaml @@ -1,5 +1,5 @@ name: "creative/bg_remover" -version: "0.1.0" +version: "0.2.0" description: "Remove image backgrounds locally using rembg." short_description: "Offline background removal using rembg with transparent PNG output." diff --git a/skills/creative/bg_remover/skill.py b/skills/creative/bg_remover/skill.py index ceac48e..0aa237f 100644 --- a/skills/creative/bg_remover/skill.py +++ b/skills/creative/bg_remover/skill.py @@ -5,15 +5,29 @@ from skillware.core.base_skill import BaseSkill +MAX_IMAGE_BYTES = 25 * 1024 * 1024 # 25 MB + class BackgroundRemover(BaseSkill): """Remove image backgrounds locally using rembg.""" + _sessions = {} + + @classmethod + def _get_session(cls, model: str): + """Load and reuse rembg sessions across executions.""" + if model not in cls._sessions: + from rembg import new_session + + cls._sessions[model] = new_session(model) + + return cls._sessions[model] + @property def manifest(self) -> Dict[str, Any]: return { "name": "creative/bg_remover", - "version": "0.1.0", + "version": "0.2.0", "description": ( "Remove image backgrounds locally using rembg " "and return a transparent PNG." @@ -26,7 +40,7 @@ def execute(self, params: Dict[str, Any]) -> Dict[str, Any]: try: try: from PIL import Image - from rembg import new_session, remove + from rembg import remove except ImportError: return { "success": False, @@ -53,11 +67,58 @@ def execute(self, params: Dict[str, Any]) -> Dict[str, Any]: # Read image bytes if image_b64: - image_bytes = base64.b64decode(image_b64) + try: + image_bytes = base64.b64decode(image_b64, validate=True) + except Exception: + return { + "success": False, + "error": "Invalid base64 image.", + "error_code": "INVALID_INPUT", + } else: - image_bytes = Path(input_path).read_bytes() + input_file = Path(input_path) + + if input_file.is_dir(): + return { + "success": False, + "error": "Input path must be a file, not a directory.", + "error_code": "INVALID_INPUT", + } + + if not input_file.exists(): + return { + "success": False, + "error": f"Input file '{input_path}' was not found.", + "error_code": "FILE_NOT_FOUND", + } + + image_bytes = input_file.read_bytes() + + if len(image_bytes) > MAX_IMAGE_BYTES: + return { + "success": False, + "error": f"Input image exceeds the maximum size of {MAX_IMAGE_BYTES // (1024 * 1024)} MB.", + "error_code": "INVALID_INPUT", + } + + if len(image_bytes) == 0: + return { + "success": False, + "error": "Input image is empty.", + "error_code": "INVALID_INPUT", + } + + try: + image = Image.open(io.BytesIO(image_bytes)) + image.verify() + except Exception: + return { + "success": False, + "error": "Input is not a valid image.", + "error_code": "INVALID_INPUT", + } - session = new_session(model) + session = self._get_session(model) output_bytes = remove( image_bytes, @@ -76,7 +137,9 @@ def execute(self, params: Dict[str, Any]) -> Dict[str, Any]: # Optional save if output_path: - Path(output_path).write_bytes(buffer.getvalue()) + output_file = Path(output_path) + output_file.parent.mkdir(parents=True, exist_ok=True) + output_file.write_bytes(buffer.getvalue()) return { "success": True, diff --git a/skills/creative/bg_remover/test_skill.py b/skills/creative/bg_remover/test_skill.py index 92f50fb..16695d5 100644 --- a/skills/creative/bg_remover/test_skill.py +++ b/skills/creative/bg_remover/test_skill.py @@ -117,7 +117,7 @@ def test_invalid_base64(skill): result = skill.execute({"image": "not_base64"}) assert result["success"] is False - assert result["error_code"] == "PROCESSING_FAILED" + assert result["error_code"] == "INVALID_INPUT" def test_missing_dependency(monkeypatch, skill): @@ -249,4 +249,80 @@ def test_missing_input_path(skill): result = skill.execute({"input_path": "/nonexistent/path/image.png"}) assert result["success"] is False - assert result["error_code"] == "PROCESSING_FAILED" + assert result["error_code"] == "FILE_NOT_FOUND" + + +def test_directory_input(skill, tmp_path): + result = skill.execute({"input_path": str(tmp_path)}) + + assert result["success"] is False + assert result["error_code"] == "INVALID_INPUT" + + +def test_empty_input_file(skill, tmp_path): + empty = tmp_path / "empty.png" + empty.write_bytes(b"") + + result = skill.execute({"input_path": str(empty)}) + + assert result["success"] is False + assert result["error_code"] == "INVALID_INPUT" + + +def test_invalid_image_file(skill, tmp_path): + bad = tmp_path / "bad.png" + bad.write_text("not an image") + + result = skill.execute({"input_path": str(bad)}) + + assert result["success"] is False + assert result["error_code"] == "INVALID_INPUT" + + +def test_output_directory_created(skill, tmp_path): + image = Image.new("RGB", (64, 64), "white") + + input_file = tmp_path / "input.png" + image.save(input_file) + + output_file = tmp_path / "nested" / "folder" / "output.png" + + result = skill.execute( + { + "input_path": str(input_file), + "output_path": str(output_file), + } + ) + + assert result["success"] is True + assert output_file.exists() + + +def test_session_reuse(skill, monkeypatch): + calls = {"count": 0} + + def fake_remove(image_bytes, *args, **kwargs): + img = Image.new("RGBA", (100, 100), (255, 0, 0, 0)) + buffer = io.BytesIO() + img.save(buffer, format="PNG") + return buffer.getvalue() + + def fake_new_session(model_name="u2net", *args, **kwargs): + calls["count"] += 1 + return object() + + fake_module = types.ModuleType("rembg") + fake_module.remove = fake_remove + fake_module.new_session = fake_new_session + fake_module.__spec__ = importlib.util.spec_from_loader("rembg", loader=None) + + monkeypatch.setitem(sys.modules, "rembg", fake_module) + + BackgroundRemover._sessions.clear() + + skill.execute({"image": create_image()}) + skill.execute({"image": create_image()}) + + assert calls["count"] == 1 + + BackgroundRemover._sessions.clear()