317 lines
12 KiB
Python
317 lines
12 KiB
Python
"""
|
|
Image Quality Analyzer
|
|
|
|
Description: ComfyUI custom node for PDF tools and media processing
|
|
Author: Eric Hiss (GitHub: EricRollei)
|
|
Contact: eric@historic.camera, eric@rollei.us
|
|
License: Dual License (Non-Commercial and Commercial Use)
|
|
Copyright (c) 2025 Eric Hiss. All rights reserved.
|
|
|
|
Dual License:
|
|
1. Non-Commercial Use: This software is licensed under the terms of the
|
|
Creative Commons Attribution-NonCommercial 4.0 International License.
|
|
To view a copy of this license, visit http://creativecommons.org/licenses/by-nc/4.0/
|
|
|
|
2. Commercial Use: For commercial use, a separate license is required.
|
|
Please contact Eric Hiss at eric@historic.camera or eric@rollei.us for licensing options.
|
|
|
|
Dependencies:
|
|
This code depends on several third-party libraries, each with its own license.
|
|
See CREDITS.md for a comprehensive list of dependencies and their licenses.
|
|
|
|
Third-party code:
|
|
- See CREDITS.md for complete list of dependencies
|
|
"""
|
|
|
|
"""
|
|
Image Quality Analyzer for Auto Enhancement Mode
|
|
|
|
Analyzes images to determine optimal enhancement parameters:
|
|
- Noise level → denoise_strength
|
|
- Edge sharpness → sharpen_strength
|
|
- Dynamic range → tone_map_strength
|
|
- Color saturation → color_enhance_strength
|
|
- JPEG artifacts → artifact_removal
|
|
"""
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Dict, Tuple
|
|
import warnings
|
|
warnings.filterwarnings('ignore')
|
|
|
|
class ImageQualityAnalyzer:
|
|
"""Analyzes image characteristics to determine optimal enhancement settings"""
|
|
|
|
def __init__(self, debug_mode: bool = False):
|
|
self.debug_mode = debug_mode
|
|
|
|
def analyze(self, image: Image.Image) -> Dict[str, float]:
|
|
"""
|
|
Analyze image and return optimal enhancement parameters
|
|
|
|
Returns:
|
|
dict with keys: denoise, sharpen, tone_map, color_enhance, artifact_removal
|
|
All values 0.0-1.0 except artifact_removal (bool)
|
|
"""
|
|
|
|
# Convert to numpy for analysis
|
|
if image.mode != 'RGB':
|
|
image = image.convert('RGB')
|
|
|
|
img_array = np.array(image).astype(np.float32) / 255.0
|
|
|
|
# Run all analyses
|
|
noise_level = self._analyze_noise(img_array)
|
|
sharpness = self._analyze_sharpness(img_array)
|
|
dynamic_range = self._analyze_dynamic_range(img_array)
|
|
color_saturation = self._analyze_color_saturation(img_array)
|
|
has_artifacts = self._detect_jpeg_artifacts(img_array, image)
|
|
|
|
# Convert to enhancement strengths
|
|
denoise = self._noise_to_denoise_strength(noise_level)
|
|
sharpen = self._sharpness_to_sharpen_strength(sharpness)
|
|
tone_map = self._dynamic_range_to_tonemap_strength(dynamic_range)
|
|
color_enhance = self._saturation_to_color_strength(color_saturation)
|
|
|
|
result = {
|
|
"denoise": denoise,
|
|
"sharpen": sharpen,
|
|
"tone_map": tone_map,
|
|
"color_enhance": color_enhance,
|
|
"artifact_removal": has_artifacts,
|
|
# Raw metrics for debugging
|
|
"noise_level": noise_level,
|
|
"sharpness_score": sharpness,
|
|
"dynamic_range": dynamic_range,
|
|
"saturation_score": color_saturation
|
|
}
|
|
|
|
if self.debug_mode:
|
|
print(f" 📊 Auto-Analysis Results:")
|
|
print(f" Noise: {noise_level:.3f} → Denoise: {denoise:.2f}")
|
|
print(f" Sharpness: {sharpness:.3f} → Sharpen: {sharpen:.2f}")
|
|
print(f" Dynamic Range: {dynamic_range:.3f} → Tone Map: {tone_map:.2f}")
|
|
print(f" Saturation: {color_saturation:.3f} → Color: {color_enhance:.2f}")
|
|
print(f" JPEG Artifacts: {'Yes' if has_artifacts else 'No'}")
|
|
|
|
return result
|
|
|
|
def _analyze_noise(self, img_array: np.ndarray) -> float:
|
|
"""
|
|
Estimate noise level using Laplacian variance method
|
|
Returns: 0.0 (no noise) to 1.0 (very noisy)
|
|
"""
|
|
try:
|
|
# Convert to grayscale for noise analysis
|
|
gray = np.dot(img_array[...,:3], [0.2989, 0.5870, 0.1140])
|
|
|
|
# Compute Laplacian
|
|
laplacian = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]])
|
|
|
|
# Convolve (simplified - using basic numpy)
|
|
h, w = gray.shape
|
|
result = np.zeros_like(gray)
|
|
|
|
for i in range(1, h-1):
|
|
for j in range(1, w-1):
|
|
result[i, j] = np.sum(gray[i-1:i+2, j-1:j+2] * laplacian)
|
|
|
|
# Variance of Laplacian as noise estimate
|
|
variance = np.var(result)
|
|
|
|
# Normalize: typical range is 0.0001 to 0.01
|
|
# High variance = high noise
|
|
noise_level = np.clip(variance * 100, 0, 1)
|
|
|
|
return float(noise_level)
|
|
|
|
except Exception as e:
|
|
if self.debug_mode:
|
|
print(f" ⚠️ Noise analysis failed: {e}")
|
|
return 0.1 # Default low noise
|
|
|
|
def _analyze_sharpness(self, img_array: np.ndarray) -> float:
|
|
"""
|
|
Estimate image sharpness using gradient magnitude
|
|
Returns: 0.0 (blurry) to 1.0 (sharp)
|
|
"""
|
|
try:
|
|
# Convert to grayscale
|
|
gray = np.dot(img_array[...,:3], [0.2989, 0.5870, 0.1140])
|
|
|
|
# Compute gradients (simple finite differences)
|
|
grad_x = np.diff(gray, axis=1)
|
|
grad_y = np.diff(gray, axis=0)
|
|
|
|
# Gradient magnitude (mean of absolute gradients)
|
|
mag_x = np.mean(np.abs(grad_x))
|
|
mag_y = np.mean(np.abs(grad_y))
|
|
sharpness = (mag_x + mag_y) / 2
|
|
|
|
# Normalize: typical range is 0.01 to 0.15
|
|
# Low gradient = blurry, high gradient = sharp
|
|
normalized = np.clip((sharpness - 0.02) / 0.10, 0, 1)
|
|
|
|
return float(normalized)
|
|
|
|
except Exception as e:
|
|
if self.debug_mode:
|
|
print(f" ⚠️ Sharpness analysis failed: {e}")
|
|
return 0.5 # Default medium sharpness
|
|
|
|
def _analyze_dynamic_range(self, img_array: np.ndarray) -> float:
|
|
"""
|
|
Analyze dynamic range (histogram spread)
|
|
Returns: 0.0 (full range used) to 1.0 (narrow range, needs expansion)
|
|
"""
|
|
try:
|
|
# Convert to grayscale
|
|
gray = np.dot(img_array[...,:3], [0.2989, 0.5870, 0.1140])
|
|
|
|
# Get percentiles (avoid extreme outliers)
|
|
p1 = np.percentile(gray, 1)
|
|
p99 = np.percentile(gray, 99)
|
|
|
|
# Dynamic range (0 to 1)
|
|
used_range = p99 - p1
|
|
|
|
# If range is narrow, needs tone mapping
|
|
# Full range (0.98) = 0.0 (no tone mapping needed)
|
|
# Narrow range (0.5) = 0.5 (moderate tone mapping)
|
|
# Very narrow (0.2) = 1.0 (strong tone mapping)
|
|
|
|
if used_range > 0.85:
|
|
need_tonemap = 0.0 # Good range
|
|
elif used_range > 0.6:
|
|
need_tonemap = (0.85 - used_range) / 0.25 # 0.0 to 1.0
|
|
else:
|
|
need_tonemap = 0.5 + (0.6 - used_range) / 0.8 # 0.5 to 1.0
|
|
|
|
return float(np.clip(need_tonemap, 0, 1))
|
|
|
|
except Exception as e:
|
|
if self.debug_mode:
|
|
print(f" ⚠️ Dynamic range analysis failed: {e}")
|
|
return 0.0 # Default no tone mapping
|
|
|
|
def _analyze_color_saturation(self, img_array: np.ndarray) -> float:
|
|
"""
|
|
Analyze color saturation
|
|
Returns: 0.0 (oversaturated) to 1.0 (desaturated, needs enhancement)
|
|
"""
|
|
try:
|
|
# Convert to HSV to get saturation channel
|
|
# Simple RGB to HSV conversion
|
|
r, g, b = img_array[:,:,0], img_array[:,:,1], img_array[:,:,2]
|
|
|
|
v = np.maximum(np.maximum(r, g), b)
|
|
s = v - np.minimum(np.minimum(r, g), b)
|
|
|
|
# Avoid division by zero
|
|
s = np.where(v > 0, s / (v + 1e-8), 0)
|
|
|
|
# Mean saturation
|
|
mean_sat = np.mean(s)
|
|
|
|
# If saturation is low, needs enhancement
|
|
# High sat (0.8+) = 0.0 (no enhancement needed)
|
|
# Medium sat (0.4-0.7) = 0.3 (light enhancement)
|
|
# Low sat (<0.3) = 0.8 (strong enhancement)
|
|
|
|
if mean_sat > 0.6:
|
|
need_color = 0.0
|
|
elif mean_sat > 0.3:
|
|
need_color = (0.6 - mean_sat) / 0.3 * 0.5 # 0.0 to 0.5
|
|
else:
|
|
need_color = 0.5 + (0.3 - mean_sat) / 0.3 * 0.5 # 0.5 to 1.0
|
|
|
|
return float(np.clip(need_color, 0, 1))
|
|
|
|
except Exception as e:
|
|
if self.debug_mode:
|
|
print(f" ⚠️ Color saturation analysis failed: {e}")
|
|
return 0.3 # Default light enhancement
|
|
|
|
def _detect_jpeg_artifacts(self, img_array: np.ndarray, image: Image.Image) -> bool:
|
|
"""
|
|
Detect JPEG compression artifacts
|
|
Returns: True if artifacts likely present
|
|
"""
|
|
try:
|
|
# Check image format/quality from metadata
|
|
format_info = image.format if hasattr(image, 'format') else None
|
|
|
|
if format_info == 'JPEG':
|
|
# Check quality if available
|
|
if hasattr(image, 'info') and 'quality' in image.info:
|
|
quality = image.info['quality']
|
|
if quality < 85:
|
|
return True
|
|
|
|
# Check for block artifacts (8x8 DCT blocks)
|
|
gray = np.dot(img_array[...,:3], [0.2989, 0.5870, 0.1140])
|
|
|
|
# Sample a region and check for 8-pixel periodicities
|
|
h, w = gray.shape
|
|
if h > 100 and w > 100:
|
|
sample = gray[50:150, 50:150]
|
|
|
|
# Check horizontal differences at 8-pixel intervals
|
|
diffs_8 = []
|
|
diffs_7 = []
|
|
|
|
for i in range(0, sample.shape[0]-8, 8):
|
|
diff_8 = np.abs(sample[i+7] - sample[i+8] if i+8 < sample.shape[0] else 0)
|
|
diff_7 = np.abs(sample[i+6] - sample[i+7] if i+7 < sample.shape[0] else 0)
|
|
diffs_8.append(np.mean(diff_8))
|
|
diffs_7.append(np.mean(diff_7))
|
|
|
|
if len(diffs_8) > 0 and len(diffs_7) > 0:
|
|
# If 8-pixel boundaries have stronger discontinuities
|
|
if np.mean(diffs_8) > np.mean(diffs_7) * 1.3:
|
|
return True
|
|
|
|
return False
|
|
|
|
except Exception as e:
|
|
if self.debug_mode:
|
|
print(f" ⚠️ Artifact detection failed: {e}")
|
|
return False
|
|
|
|
# Conversion functions: metric → enhancement strength
|
|
|
|
def _noise_to_denoise_strength(self, noise_level: float) -> float:
|
|
"""Convert noise level to denoise strength"""
|
|
# Low noise (0.0-0.2) = 0.0 (no denoising)
|
|
# Medium noise (0.2-0.5) = 0.3-0.5 (light/medium)
|
|
# High noise (0.5-1.0) = 0.5-0.8 (strong)
|
|
|
|
if noise_level < 0.15:
|
|
return 0.0
|
|
elif noise_level < 0.4:
|
|
return noise_level * 0.8 # 0.12 to 0.32
|
|
else:
|
|
return 0.3 + (noise_level - 0.4) * 0.5 # 0.3 to 0.6
|
|
|
|
def _sharpness_to_sharpen_strength(self, sharpness: float) -> float:
|
|
"""Convert sharpness score to sharpen strength (inverted)"""
|
|
# Sharp image (0.8-1.0) = 0.0-0.2 (minimal sharpening)
|
|
# Medium (0.4-0.7) = 0.3-0.5 (moderate)
|
|
# Blurry (0.0-0.3) = 0.6-0.8 (strong)
|
|
|
|
if sharpness > 0.7:
|
|
return 0.1
|
|
elif sharpness > 0.4:
|
|
return 0.3 + (0.7 - sharpness) * 0.7 # 0.3 to 0.5
|
|
else:
|
|
return 0.5 + (0.4 - sharpness) * 0.75 # 0.5 to 0.8
|
|
|
|
def _dynamic_range_to_tonemap_strength(self, need_tonemap: float) -> float:
|
|
"""Direct mapping - already computed correctly"""
|
|
return need_tonemap
|
|
|
|
def _saturation_to_color_strength(self, need_color: float) -> float:
|
|
"""Direct mapping - already computed correctly"""
|
|
return need_color
|