Spaces:
Sleeping
Sleeping
| import io | |
| import subprocess | |
| from PIL import Image, ImageOps | |
| import cv2 | |
| import numpy as np | |
| from pathlib import Path | |
| import piexif | |
| class ImageCompressor: | |
| def __init__(self): | |
| self.supported_formats = {'.jpg', '.jpeg', '.png', '.webp', '.gif', '.bmp', '.tiff'} | |
| def compress(self, input_path, output_path, method="auto", **kwargs): | |
| """Main compression entry point""" | |
| img = Image.open(input_path) | |
| original_size = Path(input_path).stat().st_size | |
| # Get original format | |
| original_format = img.format | |
| if method == "auto": | |
| # Use AI analyzer to determine best method | |
| from .ai_analyzer import ContentAnalyzer | |
| analyzer = ContentAnalyzer() | |
| _, strategy = analyzer.analyze(input_path) | |
| method = strategy.get("method", "avif") | |
| kwargs = {**kwargs, **strategy} | |
| # Apply compression based on method | |
| if method == "avif": | |
| compressed = self._compress_avif(img, kwargs.get('quality', 75)) | |
| output_path = str(output_path).replace(Path(output_path).suffix, '.avif') | |
| elif method == "webp_lossless": | |
| compressed = self._compress_webp_lossless(img, kwargs.get('quality', 90)) | |
| output_path = str(output_path).replace(Path(output_path).suffix, '.webp') | |
| elif method == "png_optimized": | |
| compressed = self._compress_png_optimized(img, kwargs.get('colors', 256)) | |
| output_path = str(output_path).replace(Path(output_path).suffix, '.png') | |
| elif method == "png_quantized": | |
| compressed = self._compress_png_quantized(img, kwargs.get('colors', 128)) | |
| output_path = str(output_path).replace(Path(output_path).suffix, '.png') | |
| elif method == "jpeg_high_quality": | |
| compressed = self._compress_jpeg(img, kwargs.get('quality', 85)) | |
| output_path = str(output_path).replace(Path(output_path).suffix, '.jpg') | |
| else: | |
| # Default to WebP | |
| compressed = self._compress_webp(img, kwargs.get('quality', 80)) | |
| output_path = str(output_path).replace(Path(output_path).suffix, '.webp') | |
| # Save compressed image | |
| with open(output_path, 'wb') as f: | |
| f.write(compressed) | |
| compressed_size = Path(output_path).stat().st_size | |
| ratio = (1 - compressed_size / original_size) * 100 | |
| return output_path, ratio | |
| def _compress_avif(self, img, quality=75): | |
| """AVIF compression - best for photos""" | |
| output = io.BytesIO() | |
| # Convert RGBA to RGB if needed (AVIF doesn't support alpha well) | |
| if img.mode == 'RGBA': | |
| background = Image.new('RGB', img.size, (255, 255, 255)) | |
| background.paste(img, mask=img.split()[-1]) | |
| img = background | |
| # Save as AVIF (requires pillow-avif-plugin) | |
| img.save(output, format='AVIF', quality=quality, speed=4) | |
| return output.getvalue() | |
| def _compress_webp_lossless(self, img, quality=90): | |
| """Lossless WebP for text/screenshots""" | |
| output = io.BytesIO() | |
| img.save(output, format='WEBP', lossless=True, quality=quality, method=6) | |
| return output.getvalue() | |
| def _compress_webp(self, img, quality=80): | |
| """Standard WebP compression""" | |
| output = io.BytesIO() | |
| img.save(output, format='WEBP', quality=quality, method=4) | |
| return output.getvalue() | |
| def _compress_png_optimized(self, img, colors=256): | |
| """PNG with color palette reduction""" | |
| output = io.BytesIO() | |
| # Convert to palette mode | |
| if colors < 256: | |
| img = img.quantize(colors=colors, method=Image.MEDIANCUT) | |
| img.save(output, format='PNG', optimize=True) | |
| else: | |
| img.save(output, format='PNG', optimize=True) | |
| return output.getvalue() | |
| def _compress_png_quantized(self, img, colors=128): | |
| """Heavy PNG quantization for graphics""" | |
| output = io.BytesIO() | |
| # Reduce colors dramatically | |
| img_quantized = img.quantize(colors=colors, method=Image.FASTOCTREE) | |
| img_quantized.save(output, format='PNG', optimize=True) | |
| return output.getvalue() | |
| def _compress_jpeg(self, img, quality=85): | |
| """JPEG compression""" | |
| output = io.BytesIO() | |
| # Convert to RGB if needed | |
| if img.mode in ('RGBA', 'LA', 'P'): | |
| img = img.convert('RGB') | |
| img.save(output, format='JPEG', quality=quality, optimize=True, progressive=True) | |
| return output.getvalue() | |
| def smart_resize(self, img, target_dimension=2048): | |
| """Intelligently resize if image is too large""" | |
| width, height = img.size | |
| max_dim = max(width, height) | |
| if max_dim > target_dimension: | |
| scale = target_dimension / max_dim | |
| new_size = (int(width * scale), int(height * scale)) | |
| return img.resize(new_size, Image.Resampling.LANCZOS) | |
| return img | |
| def remove_metadata(self, image_path): | |
| """Strip all EXIF metadata""" | |
| piexif.remove(str(image_path)) |