Initial commit: Telegram Sticker Bot Manager
Features: - Create and manage sticker packs - Add existing sticker sets to database - Automatic image optimization (512px, WebP conversion) - GIF to animated video sticker conversion (WebM, VP9) - Background removal (local rembg + optional Remove.bg API) - Hybrid API with retry logic and usage tracking - Premium status detection - SQLite database for persistence - Docker containerization - Complete OPENCODE.md documentation Technical Stack: - python-telegram-bot 22.7 - Pillow for image processing - FFmpeg for video conversion - rembg for local background removal - aiosqlite for async database operations
This commit is contained in:
183
background_remover.py
Normal file
183
background_remover.py
Normal file
@@ -0,0 +1,183 @@
|
||||
"""Background removal module supporting local (rembg) and API (remove.bg) methods."""
|
||||
import io
|
||||
import logging
|
||||
import time
|
||||
from typing import Optional, Tuple, Union
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Remove.bg API endpoint
|
||||
REMOVE_BG_API_URL = "https://api.remove.bg/v1.0/removebg"
|
||||
MAX_API_CALLS_PER_MONTH = 50
|
||||
|
||||
# Retry configuration
|
||||
MAX_RETRIES = 3
|
||||
RETRY_DELAYS = [2, 5, 10] # seconds
|
||||
|
||||
|
||||
async def remove_background_local(image_bytes: bytes) -> Optional[bytes]:
|
||||
"""
|
||||
Remove background using local rembg library (U²Net model).
|
||||
|
||||
Args:
|
||||
image_bytes: Raw image bytes
|
||||
|
||||
Returns:
|
||||
PNG image bytes with transparent background, or None if failed
|
||||
"""
|
||||
try:
|
||||
from rembg import remove
|
||||
|
||||
# Open image
|
||||
input_image = Image.open(io.BytesIO(image_bytes))
|
||||
|
||||
# Remove background
|
||||
output_image = remove(input_image)
|
||||
|
||||
# Save to bytes
|
||||
output_buffer = io.BytesIO()
|
||||
output_image.save(output_buffer, format='PNG')
|
||||
output_buffer.seek(0)
|
||||
|
||||
logger.info("Background removed successfully (local)")
|
||||
return output_buffer.getvalue()
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error removing background locally: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def remove_background_api(
|
||||
image_bytes: bytes,
|
||||
api_key: str
|
||||
) -> Tuple[Union[bytes, str], str]:
|
||||
"""
|
||||
Remove background using remove.bg API with retry logic.
|
||||
|
||||
Args:
|
||||
image_bytes: Raw image bytes
|
||||
api_key: Remove.bg API key
|
||||
|
||||
Returns:
|
||||
Tuple of (result, message)
|
||||
- Success: (image_bytes, "success")
|
||||
- Rate limit: ("LIMIT_REACHED", error_message)
|
||||
- Other error: ("API_ERROR", error_message)
|
||||
"""
|
||||
headers = {
|
||||
"X-Api-Key": api_key
|
||||
}
|
||||
|
||||
last_error = None
|
||||
|
||||
# Retry loop
|
||||
for attempt in range(MAX_RETRIES):
|
||||
try:
|
||||
# Prepare multipart form data
|
||||
files = {
|
||||
'image_file': ('image.png', io.BytesIO(image_bytes), 'image/png')
|
||||
}
|
||||
|
||||
response = requests.post(
|
||||
REMOVE_BG_API_URL,
|
||||
headers=headers,
|
||||
files=files,
|
||||
timeout=30
|
||||
)
|
||||
|
||||
# Check response
|
||||
if response.status_code == 200:
|
||||
# Success - return image
|
||||
logger.info("Background removed successfully (API)")
|
||||
return response.content, "success"
|
||||
|
||||
elif response.status_code in (402, 429):
|
||||
# Rate limit / Payment required (limit reached)
|
||||
error_data = response.json()
|
||||
error_msg = error_data.get('errors', [{}])[0].get('title', 'Rate limit exceeded')
|
||||
logger.warning(f"Remove.bg API limit reached: {error_msg}")
|
||||
return "LIMIT_REACHED", error_msg
|
||||
|
||||
else:
|
||||
# Other API error
|
||||
error_data = response.json() if response.text else {}
|
||||
error_msg = error_data.get('errors', [{}])[0].get('title', f'HTTP {response.status_code}')
|
||||
logger.error(f"Remove.bg API error: {error_msg}")
|
||||
last_error = error_msg
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
last_error = "Timeout - API antwortet nicht"
|
||||
logger.warning(f"API timeout (attempt {attempt + 1}/{MAX_RETRIES})")
|
||||
|
||||
except requests.exceptions.ConnectionError:
|
||||
last_error = "Verbindungsfehler - API nicht erreichbar"
|
||||
logger.warning(f"API connection error (attempt {attempt + 1}/{MAX_RETRIES})")
|
||||
|
||||
except Exception as e:
|
||||
last_error = str(e)
|
||||
logger.error(f"API call error (attempt {attempt + 1}/{MAX_RETRIES}): {e}")
|
||||
|
||||
# Wait before retry (except for last attempt)
|
||||
if attempt < MAX_RETRIES - 1:
|
||||
delay = RETRY_DELAYS[attempt]
|
||||
logger.info(f"Retrying in {delay} seconds...")
|
||||
time.sleep(delay)
|
||||
|
||||
# All retries failed
|
||||
logger.error(f"All API retries failed: {last_error}")
|
||||
return "API_ERROR", f"API nicht erreichbar nach {MAX_RETRIES} Versuchen: {last_error}"
|
||||
|
||||
|
||||
async def should_show_api_limit_warning(api_key: str, current_usage: int) -> bool:
|
||||
"""
|
||||
Check if we should show warning on the API button.
|
||||
|
||||
Args:
|
||||
api_key: API key (if None, no warning)
|
||||
current_usage: Current month usage count
|
||||
|
||||
Returns:
|
||||
True if limit reached or very close
|
||||
"""
|
||||
if not api_key:
|
||||
return False
|
||||
|
||||
# Show warning when at or above limit
|
||||
return current_usage >= MAX_API_CALLS_PER_MONTH
|
||||
|
||||
|
||||
async def validate_image_for_bg_removal(image_bytes: bytes) -> Tuple[bool, str]:
|
||||
"""
|
||||
Validate if image is suitable for background removal.
|
||||
|
||||
Args:
|
||||
image_bytes: Raw image bytes
|
||||
|
||||
Returns:
|
||||
Tuple of (is_valid, error_message)
|
||||
"""
|
||||
try:
|
||||
image = Image.open(io.BytesIO(image_bytes))
|
||||
|
||||
# Check format
|
||||
if image.format not in ('PNG', 'JPEG', 'JPG', 'WEBP', 'BMP'):
|
||||
return False, f"Format '{image.format}' nicht unterstützt. Nutze PNG, JPG oder WebP."
|
||||
|
||||
# Check size
|
||||
width, height = image.size
|
||||
if width < 64 or height < 64:
|
||||
return False, "Bild zu klein (min. 64x64 Pixel)"
|
||||
|
||||
if width > 4096 or height > 4096:
|
||||
return False, "Bild zu groß (max. 4096x4096 Pixel)"
|
||||
|
||||
# Check file size
|
||||
if len(image_bytes) > 10 * 1024 * 1024: # 10MB
|
||||
return False, "Datei zu groß (max. 10MB)"
|
||||
|
||||
return True, ""
|
||||
|
||||
except Exception as e:
|
||||
return False, f"Ungültiges Bild: {str(e)}"
|
||||
Reference in New Issue
Block a user