ffmpeg_denoise_images.py

#!/usr/bin/env python3
"""FFmpeg-based image denoiser.

MIT License

Copyright (c) 2025 Frederic Devernay

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

Denoise images from a directory or single file using FFmpeg denoising filters,
preserving metadata and original quality settings.

Available per-frame denoising filters:
- nlmeans: Non-local means denoiser (good for general noise reduction)
- bm3d: Block-Matching 3D denoiser in 2D mode (high quality, slower)
- dctdnoiz: 2D DCT-based denoiser (fast, good for JPEG artifacts)
- fftdnoiz: 3D FFT denoiser in 2D mode (spatial-only when prev=0, next=0)
- owdenoise: Wavelet denoiser (spatial processing)
- vaguedenoiser: Wavelet denoiser (spatial processing)

Unavailable filters (require temporal processing):
- atadenoise: Adaptive Temporal Averaging (requires multiple frames)
- hqdn3d: High Quality 3D Denoiser (has temporal component)

Usage:
    # Process directory of images
    ffmpeg_denoise_images.py input_dir output_dir [options]
    
    # Process single image
    ffmpeg_denoise_images.py input.jpg output.jpg [options]

Examples:
    # Basic usage with directory
    ffmpeg_denoise_images.py ./noisy_images ./clean_images
    
    # Process single file with bm3d filter
    ffmpeg_denoise_images.py noisy.jpg clean.jpg --filter bm3d --bm3d-sigma 3.0
    
    # Use custom filter string for advanced control
    ffmpeg_denoise_images.py input.png output.png --custom-filter "nlmeans=s=2.0:p=7:r=15"

Requirements:
    - FFmpeg (with denoising filters compiled)
    - exiftool (for metadata preservation)
    - Python 3.10+
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import List

import ffmpeg
from exiftool import ExifToolHelper


def validate_input(input_path: Path) -> List[Path]:
    """Validate input path and return list of image files.
    
    Args:
        input_path: Path to input directory or file
        
    Returns:
        List of image file paths
        
    Raises:
        SystemExit: If input is invalid
    """
    if not input_path.exists():
        print(f"Error: Input path '{input_path}' does not exist")
        raise SystemExit(1)
    
    # Handle single file input
    if input_path.is_file():
        return [input_path]
    
    # Handle directory input
    if not input_path.is_dir():
        print(f"Error: '{input_path}' is neither a file nor a directory")
        raise SystemExit(1)
    
    files = []
    for item in input_path.iterdir():
        if item.is_dir():
            print(f"Error: Directory contains subdirectory '{item.name}'. Only files are allowed.")
            raise SystemExit(1)
        elif item.is_file():
            files.append(item)
    
    if not files:
        print(f"Error: No files found in directory '{input_path}'")
        raise SystemExit(1)
        
    return files


def check_ffmpeg_support(file_path: Path) -> bool:
    """Check if FFmpeg can read the image format.
    
    Args:
        file_path: Path to image file
        
    Returns:
        True if supported, False otherwise
    """
    try:
        ffmpeg.probe(str(file_path))
        return True
    except ffmpeg.Error:
        return False


def get_image_quality_settings(file_path: Path) -> dict[str, str]:
    """Get quality settings for an image file.
    
    Args:
        file_path: Path to image file
        
    Returns:
        Dictionary of FFmpeg output parameters to preserve quality
    """
    try:
        probe = ffmpeg.probe(str(file_path))
        format_name = probe['format']['format_name'].lower()
        
        # JPEG quality detection
        if 'jpeg' in format_name or 'jpg' in format_name:
            # Try to detect JPEG quality from file
            try:
                with ExifToolHelper() as et:
                    metadata = et.get_metadata(str(file_path))[0]
                    # Look for quality indicators in metadata
                    if 'JPEG:JPEGQuality' in metadata:
                        quality = metadata['JPEG:JPEGQuality']
                        return {'q:v': str(quality)}
            except Exception:
                pass
            # Default high quality for JPEG
            return {'q:v': '2'}  # FFmpeg scale: 2-31, lower is better
        
        # PNG - lossless, no quality setting needed
        elif 'png' in format_name:
            return {}
        
        # WEBP quality detection
        elif 'webp' in format_name:
            return {'q:v': '95'}  # High quality default
        
        # TIFF - typically lossless
        elif 'tiff' in format_name or 'tif' in format_name:
            return {}
        
        # Default high quality for other formats
        else:
            return {'q:v': '2'}
            
    except Exception:
        # Fallback to high quality
        return {'q:v': '2'}


def copy_metadata(source_path: Path, target_path: Path) -> bool:
    """Copy metadata from source to target image using exiftool.
    
    Args:
        source_path: Path to source image
        target_path: Path to target image
        
    Returns:
        True if metadata was successfully copied, False otherwise
    """
    try:
        with ExifToolHelper() as et:
            # Copy all metadata from source to target
            et.execute("-TagsFromFile", str(source_path), "-all:all", str(target_path))
        return True
    except Exception:
        return False
    """Copy metadata from source to target image using exiftool.
    
    Args:
        source_path: Path to source image
        target_path: Path to target image
        
    Returns:
        True if metadata was successfully copied, False otherwise
    """
    try:
        with ExifToolHelper() as et:
            # Copy all metadata from source to target
            et.execute("-TagsFromFile", str(source_path), "-all:all", str(target_path))
        return True
    except Exception:
        return False


def get_output_path(input_path: Path, output_path: Path, file_path: Path) -> Path:
    """Get output path for a file based on input/output configuration.
    
    Args:
        input_path: Original input path (file or directory)
        output_path: Specified output path (file or directory)
        file_path: Current file being processed
        
    Returns:
        Output path for the file
    """
    # If input is a single file, output should be the specified path
    if input_path.is_file():
        return output_path
    
    # If input is a directory, output file goes in output directory
    return output_path / file_path.name


def get_filter_with_params(filter_name: str, args: argparse.Namespace) -> str:
    """Get filter string with parameters based on filter type and arguments.
    
    Args:
        filter_name: Name of the denoising filter
        args: Parsed command line arguments
        
    Returns:
        Filter string with parameters
    """
    if args.custom_filter:
        return args.custom_filter
    
    if filter_name == "nlmeans":
        return f"nlmeans=s={args.nlmeans_strength}"
    elif filter_name == "bm3d":
        # Use 2D mode (group=1) for per-frame processing
        return f"bm3d=sigma={args.bm3d_sigma}:group=1"
    elif filter_name == "dctdnoiz":
        return f"dctdnoiz=sigma={args.dctdnoiz_sigma}"
    elif filter_name == "fftdnoiz":
        # Use spatial-only mode (prev=0, next=0)
        return f"fftdnoiz=sigma={args.fftdnoiz_sigma}:prev=0:next=0"
    elif filter_name == "owdenoise":
        return f"owdenoise=ls={args.owdenoise_strength}"
    elif filter_name == "vaguedenoiser":
        return f"vaguedenoiser=threshold={args.vaguedenoiser_threshold}"
    else:
        return filter_name
    """Get filter string with parameters based on filter type and arguments.
    
    Args:
        filter_name: Name of the denoising filter
        args: Parsed command line arguments
        
    Returns:
        Filter string with parameters
    """
    if args.custom_filter:
        return args.custom_filter
    
    if filter_name == "nlmeans":
        return f"nlmeans=s={args.nlmeans_strength}"
    elif filter_name == "bm3d":
        # Use 2D mode (group=1) for per-frame processing
        return f"bm3d=sigma={args.bm3d_sigma}:group=1"
    elif filter_name == "dctdnoiz":
        return f"dctdnoiz=sigma={args.dctdnoiz_sigma}"
    elif filter_name == "fftdnoiz":
        # Use spatial-only mode (prev=0, next=0)
        return f"fftdnoiz=sigma={args.fftdnoiz_sigma}:prev=0:next=0"
    elif filter_name == "owdenoise":
        return f"owdenoise=ls={args.owdenoise_strength}"
    elif filter_name == "vaguedenoiser":
        return f"vaguedenoiser=threshold={args.vaguedenoiser_threshold}"
    else:
        return filter_name


def denoise_image(input_path: Path, output_path: Path, filter_string: str) -> bool:
    """Denoise a single image using FFmpeg.
    
    Args:
        input_path: Path to input image
        output_path: Path to output image
        filter_string: Complete filter string with parameters
        
    Returns:
        True if successful, False otherwise
    """
    try:
        # Get quality settings for the input image
        quality_settings = get_image_quality_settings(input_path)
        
        stream = ffmpeg.input(str(input_path))
        # Apply the filter using vf (video filter)
        stream = ffmpeg.output(stream, str(output_path), vf=filter_string, **quality_settings)
        ffmpeg.run(stream, overwrite_output=True, quiet=True)
        return True
    except ffmpeg.Error as e:
        print(f"Error processing {input_path.name}: {e}")
        return False


def main() -> None:
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    
    parser.add_argument(
        "input_path",
        type=Path,
        help="Input directory containing images or single image file to denoise"
    )
    
    parser.add_argument(
        "output_path", 
        type=Path,
        help="Output directory for denoised images or single output image file"
    )
    
    parser.add_argument(
        "--filter",
        choices=["nlmeans", "bm3d", "dctdnoiz", "fftdnoiz", "owdenoise", "vaguedenoiser"],
        default="nlmeans",
        help="Denoising filter to use (default: nlmeans)"
    )
    
    # Filter-specific parameters
    parser.add_argument(
        "--nlmeans-strength", type=float, default=1.0,
        help="nlmeans denoising strength (default: 1.0)"
    )
    parser.add_argument(
        "--bm3d-sigma", type=float, default=1.0,
        help="bm3d denoising strength (default: 1.0)"
    )
    parser.add_argument(
        "--dctdnoiz-sigma", type=float, default=0.0,
        help="dctdnoiz noise sigma (default: 0.0 for auto)"
    )
    parser.add_argument(
        "--fftdnoiz-sigma", type=float, default=1.0,
        help="fftdnoiz denoising strength (default: 1.0)"
    )
    parser.add_argument(
        "--owdenoise-strength", type=float, default=1.0,
        help="owdenoise luma strength (default: 1.0)"
    )
    parser.add_argument(
        "--vaguedenoiser-threshold", type=float, default=2.0,
        help="vaguedenoiser filtering strength (default: 2.0)"
    )
    parser.add_argument(
        "--custom-filter", type=str,
        help="Custom filter string (overrides filter selection and parameters)"
    )
    
    args = parser.parse_args()
    
    # Validate input path
    files = validate_input(args.input_path)
    
    # Create output directory if input is directory, or output parent if single file
    if args.input_path.is_dir():
        args.output_path.mkdir(parents=True, exist_ok=True)
    else:
        args.output_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Track metadata warnings
    metadata_failed = []
    
    # Process each file
    for file_path in files:
        print(f"Processing {file_path.name}...")
        
        # Check FFmpeg support
        if not check_ffmpeg_support(file_path):
            print(f"Error: FFmpeg cannot read format of '{file_path.name}'")
            raise SystemExit(1)
        
        # Get filter string with parameters
        filter_string = get_filter_with_params(args.filter, args)
        
        # Get output path for this file
        output_file_path = get_output_path(args.input_path, args.output_path, file_path)
        
        # Denoise image
        if not denoise_image(file_path, output_file_path, filter_string):
            raise SystemExit(1)
        
        # Copy metadata
        if not copy_metadata(file_path, output_file_path):
            metadata_failed.append(file_path.name)
    
    print(f"Successfully processed {len(files)} {'image' if len(files) == 1 else 'images'} using {args.filter} filter")
    
    # Show metadata warnings
    if metadata_failed:
        print(f"\nWarning: Could not preserve metadata for {len(metadata_failed)} files:")
        for filename in metadata_failed:
            print(f"  - {filename}")
        print("These files were denoised but lost their metadata.")


if __name__ == "__main__":
    main()