from __future__ import annotations import os import subprocess import threading from dataclasses import dataclass from pathlib import Path from typing import Callable, List, Optional from PIL import Image, ImageOps from PyQt6.QtCore import QThread, pyqtSignal from .media import ( FFmpegTools, MediaError, available_encoders, compact_error, find_ffmpeg_tools, image_output_extension, unique_output_path, video_output_extension, ) from .models import ImageOptions, MediaKind, VideoOptions ProgressCallback = Callable[[int], None] CancelCheck = Callable[[], bool] class CancelledError(Exception): pass def _check_cancel(cancelled: Optional[CancelCheck]) -> None: if cancelled and cancelled(): raise CancelledError() def _safe_remove(path: Path) -> None: try: path.unlink() except OSError: pass # --------------------------------------------------------------------------- # Images # --------------------------------------------------------------------------- def _target_image_size(width: int, height: int, options: ImageOptions): if options.resize_mode == "fit": max_width = max(1, options.max_width) max_height = max(1, options.max_height) ratio = min(1.0, max_width / width, max_height / height) target = (max(1, round(width * ratio)), max(1, round(height * ratio))) elif options.resize_mode == "percent": ratio = max(1, options.scale_percent) / 100.0 target = (max(1, round(width * ratio)), max(1, round(height * ratio))) else: target = (width, height) return target def _has_alpha(image: Image.Image) -> bool: return image.mode in ("RGBA", "LA", "PA") or ( image.mode == "P" and "transparency" in image.info ) def _prepare_image_mode(image: Image.Image, target_format: str) -> Image.Image: if target_format == "JPEG": if _has_alpha(image): rgba = image.convert("RGBA") background = Image.new("RGB", rgba.size, (255, 255, 255)) background.paste(rgba, mask=rgba.getchannel("A")) return background return image.convert("RGB") if image.mode != "RGB" else image if target_format == "WEBP": if _has_alpha(image): return image.convert("RGBA") return image.convert("RGB") if image.mode not in ("RGB",) else image if target_format == "PNG": if image.mode not in ("RGB", "RGBA", "L", "LA", "P", "I", "I;16"): return image.convert("RGBA") return image if target_format == "BMP": return image.convert("RGB") if image.mode != "RGB" else image return image def _resolve_image_format(source_format: str, options: ImageOptions) -> str: mapping = {"jpeg": "JPEG", "webp": "WEBP", "png": "PNG"} if options.output_format in mapping: return mapping[options.output_format] source_format = (source_format or "").upper() if source_format in {"JPEG", "PNG", "WEBP", "BMP", "TIFF"}: return source_format raise MediaError("This image format cannot be saved in place.") def compress_image( source: Path, destination: Path, options: ImageOptions, progress: Optional[ProgressCallback] = None, cancelled: Optional[CancelCheck] = None, ) -> Path: """Resize and re-encode one image. Saves atomically via a .part file.""" progress = progress or (lambda _value: None) _check_cancel(cancelled) destination.parent.mkdir(parents=True, exist_ok=True) temporary = destination.with_name(destination.name + ".part") try: with Image.open(source) as opened: if getattr(opened, "is_animated", False) and getattr(opened, "n_frames", 1) > 1: raise MediaError("Animated images are not supported yet.") source_format = opened.format or "" icc_profile = opened.info.get("icc_profile") if options.keep_metadata else None # Detach from the lazy file handle before leaving the context manager. image = ImageOps.exif_transpose(opened).copy() progress(8) _check_cancel(cancelled) target_format = _resolve_image_format(source_format, options) target_size = _target_image_size(image.width, image.height, options) if target_size != image.size: image = image.resize(target_size, Image.Resampling.LANCZOS) progress(42) _check_cancel(cancelled) image = _prepare_image_mode(image, target_format) progress(58) exif_bytes = None if options.keep_metadata: try: exif = image.getexif() exif.pop(0x0112, None) # Rotation has already been applied. if len(exif): exif_bytes = exif.tobytes() except Exception: exif_bytes = None save_options = {} if target_format == "JPEG": save_options.update( quality=max(1, min(100, options.quality)), optimize=True, progressive=True, subsampling=0 if options.quality >= 90 else 2, ) elif target_format == "WEBP": save_options.update( quality=max(1, min(100, options.quality)), method=6, ) elif target_format == "PNG": save_options.update(optimize=True, compress_level=9) elif target_format == "TIFF": save_options.update(compression="tiff_lzw") supports_metadata = target_format in {"JPEG", "WEBP", "PNG", "TIFF"} if supports_metadata and exif_bytes: save_options["exif"] = exif_bytes if supports_metadata and icc_profile: save_options["icc_profile"] = icc_profile _check_cancel(cancelled) progress(72) _safe_remove(temporary) image.save(temporary, format=target_format, **save_options) progress(96) _check_cancel(cancelled) os.replace(temporary, destination) progress(100) return destination except BaseException: _safe_remove(temporary) raise # --------------------------------------------------------------------------- # Video # --------------------------------------------------------------------------- _VIDEO_PRESETS = { "fast": "veryfast", "balanced": "medium", "best": "slow", } _VP9_SPEED = {"fast": "5", "balanced": "2", "best": "0"} _AV1_AOM_SPEED = {"fast": "8", "balanced": "6", "best": "4"} _AV1_SVT_SPEED = {"fast": "10", "balanced": "8", "best": "6"} def calculate_crf(codec: str, quality: int) -> int: """Map the simple 20-95 quality scale onto each codec's CRF range.""" quality = max(0, min(100, quality)) if codec in {"h264", "h265"}: return max(16, min(34, round(36 - 0.18 * quality))) if codec == "vp9": return max(18, min(46, round(46 - 0.22 * quality))) return max(18, min(50, round(52 - 0.28 * quality))) def _scale_filter(width: int, height: int, target_short_edge: int) -> Optional[str]: if not width or not height or not target_short_edge: return None if min(width, height) <= target_short_edge: return None if width >= height: return f"scale=-2:{int(target_short_edge)}:flags=lanczos,setsar=1" return f"scale={int(target_short_edge)}:-2:flags=lanczos,setsar=1" def build_video_command( source: Path, destination: Path, options: VideoOptions, width: int, height: int, tools: Optional[FFmpegTools] = None, ) -> List[str]: tools = tools or find_ffmpeg_tools() if tools.ffmpeg is None: raise MediaError("FFmpeg was not found. Add it to PATH to compress videos.") encoders = available_encoders(tools) codec = options.codec audio_bitrate = int(options.audio_bitrate) crf = str(calculate_crf(codec, options.quality)) speed = options.speed if options.speed in _VIDEO_PRESETS else "balanced" cmd = [ str(tools.ffmpeg), "-hide_banner", "-nostdin", "-loglevel", "error", "-y", "-i", str(source), "-map", "0:v:0", ] if audio_bitrate > 0: cmd += ["-map", "0:a?"] else: cmd += ["-an"] cmd += ["-sn", "-dn"] if not options.keep_metadata: cmd += ["-map_metadata", "-1", "-map_chapters", "-1"] scale = _scale_filter(width, height, options.resolution) if scale: cmd += ["-vf", scale] elif width and height and (width % 2 or height % 2): # yuv420p needs even dimensions, and we never want to distort the source. cmd += ["-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2:0:0"] if codec == "h264": if "libx264" not in encoders: raise MediaError("This FFmpeg build cannot encode H.264.") cmd += [ "-c:v", "libx264", "-preset", _VIDEO_PRESETS[speed], "-crf", crf, "-pix_fmt", "yuv420p", ] elif codec == "h265": if "libx265" not in encoders: raise MediaError("This FFmpeg build cannot encode H.265.") cmd += [ "-c:v", "libx265", "-preset", _VIDEO_PRESETS[speed], "-crf", crf, "-pix_fmt", "yuv420p", "-tag:v", "hvc1", "-x265-params", "log-level=error", ] elif codec == "vp9": if "libvpx-vp9" not in encoders: raise MediaError("This FFmpeg build cannot encode VP9.") cmd += [ "-c:v", "libvpx-vp9", "-b:v", "0", "-crf", crf, "-deadline", "good", "-cpu-used", _VP9_SPEED[speed], "-row-mt", "1", "-pix_fmt", "yuv420p", ] elif codec == "av1": if "libsvtav1" in encoders: cmd += [ "-c:v", "libsvtav1", "-crf", crf, "-preset", _AV1_SVT_SPEED[speed], "-pix_fmt", "yuv420p", ] elif "libaom-av1" in encoders: cmd += [ "-c:v", "libaom-av1", "-b:v", "0", "-crf", crf, "-cpu-used", _AV1_AOM_SPEED[speed], "-row-mt", "1", "-pix_fmt", "yuv420p", ] else: raise MediaError("This FFmpeg build cannot encode AV1.") else: raise MediaError("Unknown video codec.") if audio_bitrate > 0: audio_codec = "libopus" if destination.suffix.lower() == ".webm" else "aac" cmd += ["-c:a", audio_codec, "-b:a", f"{audio_bitrate}k"] if destination.suffix.lower() == ".mp4": cmd += ["-movflags", "+faststart"] cmd += ["-progress", "pipe:1", "-nostats", str(destination)] return cmd def _progress_line_ignored(line: str) -> bool: prefixes = ( "frame=", "fps=", "stream_", "bitrate=", "total_size=", "out_time_us=", "out_time_ms=", "out_time=", "dup_frames=", "drop_frames=", "speed=", "progress=", ) return line.startswith(prefixes) def compress_video( source: Path, destination: Path, options: VideoOptions, width: int, height: int, duration: float, progress: Optional[ProgressCallback] = None, cancelled: Optional[CancelCheck] = None, process_changed: Optional[Callable[[Optional[subprocess.Popen]], None]] = None, tools: Optional[FFmpegTools] = None, ) -> Path: progress = progress or (lambda _value: None) tools = tools or find_ffmpeg_tools() _check_cancel(cancelled) destination.parent.mkdir(parents=True, exist_ok=True) cmd = build_video_command(source, destination, options, width, height, tools) failure: Optional[BaseException] = None last_percent = -1 error_lines: List[str] = [] process: Optional[subprocess.Popen] = None try: process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, encoding="utf-8", errors="replace", bufsize=1, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0) if os.name == "nt" else 0, ) if process_changed: process_changed(process) assert process.stdout is not None for raw_line in process.stdout: if cancelled and cancelled(): try: process.kill() except OSError: pass line = raw_line.strip() if not line: continue if line.startswith("out_time_us=") or line.startswith("out_time_ms="): if duration > 0: try: microseconds = int(line.split("=", 1)[1]) percent = int(min(99, max(0, microseconds / 1_000_000.0 / duration * 100))) if percent > last_percent: last_percent = percent progress(percent) except (TypeError, ValueError): pass continue if line.startswith("progress=end"): if duration > 0: progress(100) continue if not _progress_line_ignored(line): error_lines.append(line) if len(error_lines) > 20: error_lines = error_lines[-20:] return_code = process.wait() if process_changed: process_changed(None) process = None if cancelled and cancelled(): raise CancelledError() if return_code != 0: message = compact_error("\n".join(error_lines)) or f"FFmpeg stopped with exit code {return_code}." raise MediaError(message) if not destination.exists() or destination.stat().st_size == 0: raise MediaError("FFmpeg did not produce an output file.") progress(100) return destination except BaseException as exc: failure = exc raise finally: if process is not None: try: process.kill() process.wait(timeout=3) except Exception: pass if process_changed: process_changed(None) if failure is not None: _safe_remove(destination) # --------------------------------------------------------------------------- # Worker # --------------------------------------------------------------------------- @dataclass(frozen=True) class CompressionJob: item_id: str source: Path kind: MediaKind width: int height: int duration: float = 0.0 class CompressionWorker(QThread): itemStarted = pyqtSignal(str) itemProgress = pyqtSignal(str, int) itemFinished = pyqtSignal(str, str, str, str, int) overallProgress = pyqtSignal(int) summary = pyqtSignal(object) def __init__( self, jobs: List[CompressionJob], image_options: ImageOptions, video_options: VideoOptions, output_dir: Optional[Path], suffix: str = "-compressed", parent=None, ): super().__init__(parent) self._jobs = list(jobs) self._image_options = ImageOptions(**vars(image_options)) self._video_options = VideoOptions(**vars(video_options)) self._output_dir = Path(output_dir) if output_dir else None self._suffix = suffix self._cancelled = threading.Event() self._process_lock = threading.Lock() self._process: Optional[subprocess.Popen] = None self._tools = find_ffmpeg_tools() @property def cancelled(self) -> bool: return self._cancelled.is_set() def cancel(self) -> None: self._cancelled.set() with self._process_lock: process = self._process if process is not None and process.poll() is None: try: process.kill() except OSError: pass def _set_process(self, process: Optional[subprocess.Popen]) -> None: with self._process_lock: self._process = process def _is_cancelled(self) -> bool: return self._cancelled.is_set() def run(self) -> None: # noqa: C901 - a small, explicit batch loop is clearest here total = len(self._jobs) completed = 0 errors = 0 cancelled_count = 0 output_bytes = 0 done_input_bytes = 0 if total == 0: self.summary.emit( { "total": 0, "done": 0, "errors": 0, "cancelled": 0, "output_bytes": 0, "input_bytes": 0, } ) return for index, job in enumerate(self._jobs): if self._cancelled.is_set(): break self.itemStarted.emit(job.item_id) destination: Optional[Path] = None def report(value: int, job_id: str = job.item_id, item_index: int = index) -> None: clamped = max(0, min(100, int(value))) self.itemProgress.emit(job_id, clamped) overall = int(((item_index + clamped / 100.0) / total) * 100) self.overallProgress.emit(min(99, overall)) try: extension = ( image_output_extension(job.source, self._image_options.output_format) if job.kind == MediaKind.IMAGE else video_output_extension(self._video_options.codec) ) destination = unique_output_path( job.source, self._output_dir, self._suffix, extension ) if not job.source.exists(): raise MediaError("The original file is no longer available.") if job.kind == MediaKind.IMAGE: compress_image( job.source, destination, self._image_options, progress=report, cancelled=self._is_cancelled, ) else: compress_video( job.source, destination, self._video_options, width=job.width, height=job.height, duration=job.duration, progress=report, cancelled=self._is_cancelled, process_changed=self._set_process, tools=self._tools, ) output_size = destination.stat().st_size output_bytes += output_size done_input_bytes += job.source.stat().st_size completed += 1 self.itemFinished.emit(job.item_id, "done", "", str(destination), output_size) except CancelledError: cancelled_count += 1 if destination is not None: _safe_remove(destination) self.itemFinished.emit(job.item_id, "cancelled", "", "", 0) except MediaError as exc: errors += 1 if destination is not None: _safe_remove(destination) self.itemFinished.emit(job.item_id, "error", compact_error(str(exc)), "", 0) except Exception as exc: errors += 1 if destination is not None: _safe_remove(destination) self.itemFinished.emit(job.item_id, "error", compact_error(str(exc)), "", 0) finally: self._set_process(None) self.overallProgress.emit(min(100, int(((index + 1) / total) * 100))) if self._cancelled.is_set(): break if self._cancelled.is_set(): # Anything that never started should not sit at "Working" forever. completed_or_handled = completed + errors + cancelled_count for job in self._jobs[completed_or_handled:]: self.itemFinished.emit(job.item_id, "cancelled", "", "", 0) cancelled_count += 1 self.summary.emit( { "total": total, "done": completed, "errors": errors, "cancelled": cancelled_count, "output_bytes": output_bytes, "input_bytes": done_input_bytes, } )