Spaces:
Running
Running
| # 📁 File: tools/image_enhancer/tools.py | |
| import os | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ImageEnhance, ImageFilter | |
| from utils import validate_path, ensure_path_from_img, temp_output | |
| # === TOOL IMPLEMENTATIONS === | |
| def apply_color_correction(input_path: str, | |
| brightness: float = 1.0, | |
| contrast: float = 1.0, | |
| saturation: float = 1.0, | |
| output_path: str = None): | |
| """ | |
| Applies color correction to an image including brightness, contrast and saturation adjustments. | |
| Args: | |
| input_path (str): Path to the input image file | |
| output_path (str): Path where the corrected image will be saved | |
| brightness (float): Brightness factor (1.0 is original, <1.0 darker, >1.0 brighter) | |
| contrast (float): Contrast factor (1.0 is original, <1.0 less contrast, >1.0 more contrast) | |
| saturation (float): Saturation factor (1.0 is original, <1.0 less saturated, >1.0 more saturated) | |
| Returns: | |
| None: The function saves the output file to the specified path | |
| Raises: | |
| FileNotFoundError: If the input file does not exist | |
| ValueError: If the input is not a valid image | |
| OSError: If there's an error creating the output directory or writing the file | |
| """ | |
| try: | |
| if output_path is None: | |
| output_path = temp_output() | |
| img = ensure_path_from_img(input_path) | |
| # Open the image | |
| img = Image.open(img) | |
| # Apply brightness adjustment | |
| if brightness != 1.0: | |
| enhancer = ImageEnhance.Brightness(img) | |
| img = enhancer.enhance(brightness) | |
| # Apply contrast adjustment | |
| if contrast != 1.0: | |
| enhancer = ImageEnhance.Contrast(img) | |
| img = enhancer.enhance(contrast) | |
| # Apply saturation adjustment | |
| if saturation != 1.0: | |
| enhancer = ImageEnhance.Color(img) | |
| img = enhancer.enhance(saturation) | |
| # Ensure the output directory exists | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Save the result | |
| img.save(output_path) | |
| except Exception as e: | |
| # Re-raise with more context | |
| raise type(e)(f"Error in color correction: {str(e)}") | |
| return str(output_path) | |
| def apply_sharpening(input_path: str, strength: float = 1.0, output_path: str = None)-> str: | |
| """ | |
| Applies sharpening to an image to enhance edges and details. | |
| Args: | |
| input_path (str): Path to the input image file | |
| output_path (str): Path where the sharpened image will be saved | |
| strength (float): Sharpening strength (1.0 is standard, >1.0 is stronger) | |
| Returns: | |
| output_path (str): Path where the sharpened image will be saved | |
| Raises: | |
| FileNotFoundError: If the input file does not exist | |
| ValueError: If the input is not a valid image | |
| OSError: If there's an error creating the output directory or writing the file | |
| """ | |
| try: | |
| if output_path is None: | |
| output_path = temp_output() | |
| # Validate input file exists | |
| if not os.path.exists(input_path): | |
| raise FileNotFoundError(f"Input file not found: {input_path}") | |
| img = ensure_path_from_img(input_path) | |
| # Open the image | |
| img = Image.open(img) | |
| # Apply sharpening | |
| if strength <= 1.0: | |
| # Use PIL's built-in SHARPEN filter for moderate sharpening | |
| sharpened = img.filter(ImageFilter.SHARPEN) | |
| else: | |
| # For stronger sharpening, use UnsharpMask with adjusted parameters | |
| radius = 2.0 | |
| percent = int(150 * strength) | |
| threshold = 3 | |
| sharpened = img.filter(ImageFilter.UnsharpMask(radius, percent, threshold)) | |
| # Ensure the output directory exists | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Save the result | |
| sharpened.save(output_path) | |
| except Exception as e: | |
| # Re-raise with more context | |
| raise type(e)(f"Error in sharpening: {str(e)}") | |
| return str(output_path) | |
| # --- Constants for Denoising Parameters --- | |
| # These values control the filtering strength. | |
| # h: Luminance component filter strength. Larger h means more noise removal. | |
| # hColor: Color component filter strength. | |
| # templateWindowSize: Should be odd. Recommended: 7. | |
| # searchWindowSize: Should be odd. Recommended: 21. | |
| DENOISE_PARAMS = { | |
| "low": { | |
| "h": 5, | |
| "hColor": 5, | |
| "templateWindowSize": 7, | |
| "searchWindowSize": 21, | |
| }, | |
| "medium": { | |
| "h": 10, | |
| "hColor": 10, | |
| "templateWindowSize": 7, | |
| "searchWindowSize": 21, | |
| }, | |
| "high": { | |
| "h": 15, | |
| "hColor": 15, | |
| "templateWindowSize": 7, | |
| "searchWindowSize": 21, | |
| } | |
| } | |
| def reduce_noise(input_path: str, strength: str = "medium", output_path: str = None) -> str: | |
| """ | |
| Reduces noise in an image using cv2.fastNlMeansDenoisingColored. | |
| Args: | |
| input_path (str): Path to the input image file. | |
| output_path (str): Path where the denoised image will be saved. | |
| strength (str): Noise reduction strength ("low", "medium", or "high"). | |
| Returns: | |
| output_path (str): Path where the sharpened image will be saved | |
| Raises: | |
| FileNotFoundError: If the input file does not exist. | |
| ValueError: If the strength parameter is invalid or the image cannot be loaded. | |
| OSError: If there's an error creating the output directory or writing the file. | |
| """ | |
| if output_path is None: | |
| output_path = temp_output() | |
| if not os.path.exists(input_path): | |
| raise FileNotFoundError(f"Input file not found: {input_path}") | |
| if strength not in DENOISE_PARAMS: | |
| raise ValueError(f"Invalid strength value: {strength}. Must be one of {list(DENOISE_PARAMS.keys())}") | |
| try: | |
| img = cv2.imread(input_path) | |
| if img is None: | |
| # This can happen if the file exists but is corrupted or not a valid image format. | |
| raise ValueError(f"Failed to load image from path: {input_path}. The file may be corrupted or in an unsupported format.") | |
| # Get parameters for the chosen strength | |
| params = DENOISE_PARAMS[strength] | |
| # Apply noise reduction | |
| denoised = cv2.fastNlMeansDenoisingColored( | |
| img, | |
| None, | |
| params["h"], | |
| params["hColor"], | |
| params["templateWindowSize"], | |
| params["searchWindowSize"] | |
| ) | |
| # Ensure the output directory exists | |
| output_dir = os.path.dirname(output_path) | |
| if output_dir: | |
| os.makedirs(output_dir, exist_ok=True) | |
| # Save the result | |
| success = cv2.imwrite(output_path, denoised) | |
| if not success: | |
| raise OSError(f"Failed to save the denoised image to {output_path}") | |
| except (cv2.error, OSError) as e: | |
| # Catch specific errors from OpenCV or file system and re-raise | |
| # them to provide context without losing the original error. | |
| raise OSError(f"An error occurred during image processing or file saving for {input_path}") from e | |
| return str(output_path) | |
| def upscale_resolution(input_path: str, scale_factor: int = 2, method: str = "bicubic", output_path: str = None) -> str: | |
| """ | |
| Increases the resolution of an image using various upscaling methods. | |
| Args: | |
| input_path (str): Path to the input image file | |
| scale_factor (int): How much to increase resolution (2 = 2x, 3 = 3x, 4 = 4x) | |
| method (str): Upscaling method ("nearest", "bilinear", "bicubic", or "lanczos") | |
| output_path (str): Path where the upscaled image will be saved | |
| Returns: | |
| output_path (str): Path where the sharpened image will be saved | |
| Raises: | |
| FileNotFoundError: If the input file does not exist | |
| ValueError: If the scale_factor or method parameters are invalid | |
| OSError: If there's an error creating the output directory or writing the file | |
| """ | |
| try: | |
| if output_path is None: | |
| output_path = temp_output() | |
| # Validate input file exists | |
| if not os.path.exists(input_path): | |
| raise FileNotFoundError(f"Input file not found: {input_path}") | |
| # Validate scale factor | |
| if scale_factor not in [2, 3, 4]: | |
| raise ValueError(f"Invalid scale factor: {scale_factor}. Must be 2, 3, or 4") | |
| # Map method names to PIL resampling filters | |
| method_map = { | |
| "nearest": Image.NEAREST, | |
| "bilinear": Image.BILINEAR, | |
| "bicubic": Image.BICUBIC, | |
| "lanczos": Image.LANCZOS | |
| } | |
| # Validate method name | |
| if method.lower() not in method_map: | |
| raise ValueError(f"Invalid method: {method}. Must be one of {list(method_map.keys())}") | |
| # Use bicubic as default if method is not recognized | |
| resampling_filter = method_map.get(method.lower(), Image.BICUBIC) | |
| # Open the image | |
| img = Image.open(input_path) | |
| # Get original dimensions | |
| width, height = img.size | |
| # Calculate new dimensions | |
| new_width = width * scale_factor | |
| new_height = height * scale_factor | |
| # Resize the image | |
| upscaled = img.resize((new_width, new_height), resampling_filter) | |
| # Ensure the output directory exists | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Save the result | |
| upscaled.save(output_path) | |
| except Exception as e: | |
| # Re-raise with more context | |
| raise type(e)(f"Error in upscaling: {str(e)}") | |
| return str(output_path) | |
| def auto_enhance(input_path: str, output_path: str = None) -> str: | |
| """ | |
| Automatically enhances an image by applying balanced corrections to | |
| brightness, contrast, color, and sharpness. | |
| Args: | |
| input_path (str): Path to the input image file | |
| output_path (str): Path where the enhanced image will be saved | |
| Returns: | |
| output_path (str): Path where the sharpened image will be saved | |
| Raises: | |
| FileNotFoundError: If the input file does not exist | |
| ValueError: If the input is not a valid image | |
| OSError: If there's an error creating the output directory or writing the file | |
| """ | |
| try: | |
| if output_path is None: | |
| output_path = temp_output() | |
| # Validate input file exists | |
| if not os.path.exists(input_path): | |
| raise FileNotFoundError(f"Input file not found: {input_path}") | |
| # Open the image | |
| img = Image.open(input_path) | |
| # Auto-enhance: apply a series of balanced enhancements | |
| # Step 1: Auto-contrast | |
| enhancer = ImageEnhance.Contrast(img) | |
| img = enhancer.enhance(1.2) | |
| # Step 2: Moderate brightness adjustment | |
| enhancer = ImageEnhance.Brightness(img) | |
| img = enhancer.enhance(1.1) | |
| # Step 3: Moderate color enhancement | |
| enhancer = ImageEnhance.Color(img) | |
| img = enhancer.enhance(1.1) | |
| # Step 4: Moderate sharpening | |
| enhancer = ImageEnhance.Sharpness(img) | |
| img = enhancer.enhance(1.3) | |
| # Ensure the output directory exists | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| # Save the result | |
| img.save(output_path) | |
| except Exception as e: | |
| # Re-raise with more context | |
| raise type(e)(f"Error in auto enhancement: {str(e)}") | |
| return str(output_path) |