645 lines
21 KiB
Python
645 lines
21 KiB
Python
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,
|
|
}
|
|
)
|