imran-decoder/filecrackhead1
0
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 