diff --git a/app.py b/app.py index 21fbb7e..71744f8 100644 --- a/app.py +++ b/app.py @@ -10,10 +10,10 @@ from fastapi.staticfiles import StaticFiles from cleaner import ( DEFAULT_MODEL, CleanerError, - _get_image_info, build_prompt, clean_garment, - encode_image_bytes, + get_image_info, + resolve_api_key, ) DEFAULT_PROMPT = build_prompt("beige") @@ -53,23 +53,25 @@ async def clean_endpoint( api_key: str | None = Form(default=None), restore_res: bool = Form(default=True), ): - effective_api_key = (api_key.strip() if api_key and api_key.strip() else None) or os.environ.get( - "OPENROUTER_API_KEY" - ) + effective_api_key = resolve_api_key(api_key) if not effective_api_key: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="API key is required. Provide it in the form or set OPENROUTER_API_KEY.", ) - file_bytes = await file.read() + try: + file_bytes = await file.read() + finally: + await file.close() + if not file_bytes: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Uploaded file is empty.", ) - orig_dims, _ = _get_image_info(file_bytes) + orig_dims, _ = get_image_info(file_bytes) if orig_dims is None: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -89,21 +91,7 @@ async def clean_endpoint( detail=str(exc), ) - return { - "image": encode_image_bytes(clean_result.image_bytes, mime=clean_result.media_type), - "cost": clean_result.cost, - "width": clean_result.width, - "height": clean_result.height, - "original_dimensions": list(clean_result.original_dimensions) - if clean_result.original_dimensions - else None, - "model_dimensions": list(clean_result.model_dimensions) - if clean_result.model_dimensions - else None, - "was_rescaled": clean_result.was_rescaled, - "media_type": clean_result.media_type, - } - + return clean_result.to_payload() if STATIC_DIR.is_dir(): app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") diff --git a/cleaner.py b/cleaner.py index b18ada3..8e92638 100644 --- a/cleaner.py +++ b/cleaner.py @@ -26,7 +26,14 @@ EXTENSIONS = { } -def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, str | None]: +def resolve_api_key(api_key: str | None = None) -> str | None: + """Resolve OpenRouter API key from parameter or OPENROUTER_API_KEY environment variable.""" + if api_key and api_key.strip(): + return api_key.strip() + return os.environ.get("OPENROUTER_API_KEY") + + +def get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, str | None]: """Extract (width, height) dimensions and mime type from bytes, file path, or data URI.""" try: from PIL import Image @@ -34,7 +41,6 @@ def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, if isinstance(data_or_path, bytes): raw_bytes = data_or_path elif isinstance(data_or_path, str) and data_or_path.startswith("data:"): - # Parse data URI: data:;base64, comma_idx = data_or_path.find(",") if comma_idx != -1: raw_bytes = base64.b64decode(data_or_path[comma_idx + 1:]) @@ -50,6 +56,9 @@ def _get_image_info(data_or_path: bytes | str) -> tuple[tuple[int, int] | None, pass return None, None + +_get_image_info = get_image_info + def build_prompt(bg_choice: str) -> str: bg_description = BACKGROUND_PRESETS.get(bg_choice.lower(), bg_choice) return ( @@ -132,6 +141,22 @@ class CleanResult: return self.model_dimensions != self.original_dimensions return False + def to_payload(self) -> dict[str, Any]: + """Serialize result to a dictionary suitable for API responses.""" + return { + "image": encode_image_bytes(self.image_bytes, mime=self.media_type), + "cost": self.cost, + "width": self.width, + "height": self.height, + "original_dimensions": list(self.original_dimensions) + if self.original_dimensions + else None, + "model_dimensions": list(self.model_dimensions) + if self.model_dimensions + else None, + "was_rescaled": self.was_rescaled, + "media_type": self.media_type, + } def call_openrouter_images( image_url_or_data_uri: str, diff --git a/tests/test_cleaner.py b/tests/test_cleaner.py index 460ae80..7baaa0e 100644 --- a/tests/test_cleaner.py +++ b/tests/test_cleaner.py @@ -285,3 +285,35 @@ def test_clean_garment_file(tmp_path, monkeypatch): assert res.width == 250 assert res.height == 350 assert res.cost == 0.03 + +def test_resolve_api_key(monkeypatch): + import cleaner + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + assert cleaner.resolve_api_key(None) is None + assert cleaner.resolve_api_key(" ") is None + assert cleaner.resolve_api_key("custom-key") == "custom-key" + monkeypatch.setenv("OPENROUTER_API_KEY", "env-key") + assert cleaner.resolve_api_key(None) == "env-key" + assert cleaner.resolve_api_key("custom-key") == "custom-key" + + +def test_clean_result_to_payload(): + raw = b"img-data" + res = CleanResult( + image_bytes=raw, + media_type="image/jpeg", + cost=0.015, + width=100, + height=200, + original_dimensions=(100, 200), + model_dimensions=(50, 100), + ) + payload = res.to_payload() + assert payload["cost"] == 0.015 + assert payload["width"] == 100 + assert payload["height"] == 200 + assert payload["original_dimensions"] == [100, 200] + assert payload["model_dimensions"] == [50, 100] + assert payload["was_rescaled"] is True + assert payload["media_type"] == "image/jpeg" + assert payload["image"].startswith("data:image/jpeg;base64,")