Enhances generated figures using AI image enhancement services. Supports providers: - bianxie: Uses BianXie AI API (OpenAI-compatible) - openrouter: Uses OpenRouter API (requires modalities parameter) - gemini: Uses Google Gemini native API
| 246 | |
| 247 | |
| 248 | class ImageEnhancer: |
| 249 | """ |
| 250 | Enhances generated figures using AI image enhancement services. |
| 251 | |
| 252 | Supports providers: |
| 253 | - bianxie: Uses BianXie AI API (OpenAI-compatible) |
| 254 | - openrouter: Uses OpenRouter API (requires modalities parameter) |
| 255 | - gemini: Uses Google Gemini native API |
| 256 | """ |
| 257 | |
| 258 | def __init__(self, config: "Config"): |
| 259 | """ |
| 260 | Initialize the image enhancer. |
| 261 | |
| 262 | Args: |
| 263 | config: AutoFigure SDK configuration |
| 264 | """ |
| 265 | self.config = config |
| 266 | |
| 267 | def enhance( |
| 268 | self, |
| 269 | input_path: str, |
| 270 | output_path: Optional[str] = None, |
| 271 | enhancement_input: str = "", |
| 272 | style: Optional[str] = None, |
| 273 | input_type: str = "code2prompt", |
| 274 | ) -> Optional[str]: |
| 275 | """ |
| 276 | Enhance an image file using the configured provider. |
| 277 | |
| 278 | Args: |
| 279 | input_path: Path to the input image (PNG) |
| 280 | output_path: Path for the enhanced output (optional) |
| 281 | enhancement_input: Enhancement prompt or code |
| 282 | style: Art style for enhancement |
| 283 | input_type: 'none' | 'code' | 'code2prompt' |
| 284 | - 'none': Direct visual enhancement without code reference |
| 285 | - 'code': Use code as reference (svg_code mode) |
| 286 | - 'code2prompt': Use LLM-generated detailed prompt |
| 287 | |
| 288 | Returns: |
| 289 | Path to the enhanced image, or None on failure |
| 290 | """ |
| 291 | api_key = self.config.enhancement_api_key |
| 292 | if not api_key: |
| 293 | print("[ImageEnhancer] No enhancement API key configured, skipping enhancement") |
| 294 | return None |
| 295 | |
| 296 | input_path = Path(input_path) |
| 297 | if not input_path.exists(): |
| 298 | print(f"[ImageEnhancer] Input file not found: {input_path}") |
| 299 | return None |
| 300 | |
| 301 | # Default output path |
| 302 | if output_path is None: |
| 303 | output_path = str(input_path.parent / f"{input_path.stem}_enhanced.png") |
| 304 | |
| 305 | provider = self.config.enhancement_provider |
no outgoing calls
no test coverage detected