CoolFace
Apppublic

imran-decoder/filecrackhead1

sourceHugging Faceupdated 6mo agoView on Hugging Face
0likes
file_compressor.py278 linesDownload Raw Back to services
1"""2File compression service.3 4Provides lossy and lossless compression for:5  - Images  (Pillow)6  - PDFs    (PyMuPDF / fitz)7  - Generic (gzip)8"""9 10import gzip11import io12from pathlib import Path13from typing import Optional14 15import fitz  # PyMuPDF16from PIL import Image17 18from app.config import get_settings19from app.utils.logging_utils import get_logger20 21logger = get_logger(__name__)22settings = get_settings()23 24# ---------------------------------------------------------------------------25# Format categories26# ---------------------------------------------------------------------------27_IMAGE_EXTENSIONS = {28    "jpg", "jpeg", "png", "gif", "bmp", "tiff", "tif", "webp",29    "ico", "heic", "heif", "avif", "tga", "pcx", "ppm",30}31_PDF_EXTENSIONS = {"pdf"}32 33# Lossy formats that benefit from a quality knob34_LOSSY_FORMATS = {"jpg", "jpeg", "webp", "avif", "heic", "heif"}35 36# Formats that support lossless optimisation natively37_LOSSLESS_OPT_FORMATS = {"png", "webp", "tiff", "tif"}38 39 40# ---------------------------------------------------------------------------41# Image compression42# ---------------------------------------------------------------------------43 44def _compress_image(45    input_path: Path,46    output_path: Path,47    mode: str,48    quality: int,49) -> Path:50    """Compress an image using Pillow."""51    img = Image.open(input_path)52 53    # Strip EXIF / metadata efficiently (no pixel-level copy)54    clean = img.copy()55    img.close()56    clean.info = {}  # Drop all metadata (EXIF, ICC, XMP, etc.)57    if hasattr(clean, '_exif'):58        del clean._exif59 60    ext = output_path.suffix.lstrip(".").lower()61    save_kwargs: dict = {}62 63    if mode == "lossy":64        # For lossy-capable formats, use quality parameter65        if ext in {"jpg", "jpeg"}:66            if clean.mode in ("RGBA", "P", "LA"):67                clean = clean.convert("RGB")68            save_kwargs["quality"] = quality69            save_kwargs["optimize"] = True70        elif ext == "webp":71            save_kwargs["quality"] = quality72            save_kwargs["method"] = 6  # slowest / best compression73        elif ext in {"avif", "heic", "heif"}:74            save_kwargs["quality"] = quality75        elif ext == "png":76            # PNG is lossless; best we can do is optimize + reduce colors77            save_kwargs["optimize"] = True78            if quality < 50:79                clean = clean.quantize(colors=max(16, quality * 5))80        elif ext == "gif":81            save_kwargs["optimize"] = True82            if clean.mode not in ("P", "L"):83                clean = clean.convert("P", palette=Image.ADAPTIVE, colors=256)84        else:85            # BMP, TIFF, TGA, PCX, PPM – just save as-is86            pass87    else:88        # Lossless mode89        if ext in {"png"}:90            save_kwargs["optimize"] = True91        elif ext == "webp":92            save_kwargs["lossless"] = True93        elif ext in {"tiff", "tif"}:94            save_kwargs["compression"] = "tiff_lzw"95 96    try:97        clean.save(str(output_path), **save_kwargs)98    finally:99        clean.close()100    return output_path101 102 103# ---------------------------------------------------------------------------104# PDF compression105# ---------------------------------------------------------------------------106 107def _compress_pdf(108    input_path: Path,109    output_path: Path,110    mode: str,111    quality: int,112) -> Path:113    """Compress a PDF using PyMuPDF (fitz).114 115    Lossy:  recompress embedded images at reduced quality.116    Lossless: deflate streams and remove duplication only.117    """118    doc = fitz.open(str(input_path))119 120    if mode == "lossy":121        # Re-encode every embedded image at lower quality122        for page_num in range(len(doc)):123            page = doc[page_num]124            image_list = page.get_images(full=True)125            for img_info in image_list:126                xref = img_info[0]127                try:128                    base_image = doc.extract_image(xref)129                    if base_image is None:130                        continue131                    image_bytes = base_image["image"]132                    pil_img = Image.open(io.BytesIO(image_bytes))133                    if pil_img.mode in ("RGBA", "P", "LA"):134                        pil_img = pil_img.convert("RGB")135 136                    buf = io.BytesIO()137                    pil_img.save(buf, format="JPEG", quality=quality, optimize=True)138                    buf.seek(0)139 140                    # Replace the image in the PDF141                    page.replace_image(xref, stream=buf.getvalue())142                except Exception as exc:143                    logger.debug(144                        "Could not recompress image xref %d on page %d: %s",145                        xref, page_num, exc,146                    )147                    continue148 149    # Save with garbage collection and deflation150    doc.save(151        str(output_path),152        garbage=4,          # max garbage collection153        deflate=True,       # deflate streams154        clean=True,         # clean unused objects155    )156    doc.close()157    return output_path158 159 160# ---------------------------------------------------------------------------161# Generic gzip compression162# ---------------------------------------------------------------------------163 164def _compress_generic(input_path: Path, output_path: Path) -> Path:165    """Compress any file using gzip."""166    gz_path = output_path.with_suffix(output_path.suffix + ".gz")167    with open(input_path, "rb") as f_in:168        with gzip.open(str(gz_path), "wb", compresslevel=9) as f_out:169            while chunk := f_in.read(1024 * 1024):  # 1MB chunks170                f_out.write(chunk)171    return gz_path172 173 174# ---------------------------------------------------------------------------175# Target-size binary search (images & PDFs only)176# ---------------------------------------------------------------------------177 178def _compress_to_target_size(179    input_path: Path,180    output_path: Path,181    ext: str,182    target_size_kb: int,183    compress_fn,184    max_iterations: int = 10,185) -> Path:186    """Iteratively adjust quality to hit a target file size."""187    lo, hi = 1, 100188    best_path: Optional[Path] = None189    best_diff = float("inf")190    tolerance_kb = max(1, int(target_size_kb * 0.05))  # 5% tolerance191 192    for _ in range(max_iterations):193        mid = (lo + hi) // 2194        result = compress_fn(input_path, output_path, "lossy", mid)195        result_size_kb = result.stat().st_size / 1024196 197        diff = abs(result_size_kb - target_size_kb)198        if diff < best_diff:199            best_diff = diff200            best_path = result201 202        # Early exit if within 5% tolerance203        if diff <= tolerance_kb:204            break205 206        if result_size_kb > target_size_kb:207            hi = mid - 1208        elif result_size_kb < target_size_kb:209            lo = mid + 1210        else:211            break212 213        if lo > hi:214            break215 216    return best_path or output_path217 218 219# ---------------------------------------------------------------------------220# Public API221# ---------------------------------------------------------------------------222 223def get_file_category(ext: str) -> str:224    """Return 'image', 'pdf', or 'generic' for the given extension."""225    ext = ext.lower().lstrip(".")226    if ext in _IMAGE_EXTENSIONS:227        return "image"228    if ext in _PDF_EXTENSIONS:229        return "pdf"230    return "generic"231 232 233def compress_file(234    input_path: Path,235    output_path: Path,236    mode: str = "lossy",237    quality: int = 60,238    target_size_kb: Optional[int] = None,239) -> Path:240    """241    Compress a single file.242 243    Args:244        input_path:      Source file.245        output_path:     Desired output path.246        mode:            'lossy' or 'lossless'.247        quality:         1-100, used for lossy compression.248        target_size_kb:  Optional target size in KB (overrides quality).249 250    Returns:251        Path to the compressed file.252    """253    ext = input_path.suffix.lstrip(".").lower()254    category = get_file_category(ext)255 256    if category == "image":257        if target_size_kb is not None:258            return _compress_to_target_size(259                input_path, output_path, ext, target_size_kb, _compress_image,260            )261        return _compress_image(input_path, output_path, mode, quality)262 263    elif category == "pdf":264        if target_size_kb is not None:265            return _compress_to_target_size(266                input_path, output_path, ext, target_size_kb, _compress_pdf,267            )268        return _compress_pdf(input_path, output_path, mode, quality)269 270    else:271        # Generic files get gzip; target_size_kb not applicable272        if mode == "lossy":273            logger.info(274                "File '%s' does not support lossy compression; using gzip (lossless).",275                input_path.name,276            )277        return _compress_generic(input_path, output_path)278