407 lines
15 KiB
Python
407 lines
15 KiB
Python
import torch
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
class ImageQualityGate:
|
|
"""
|
|
Single-node image quality control for AI-generated images.
|
|
Analyzes Blur, Sharpness, Noise, Clipping, and Entropy.
|
|
Three-tier verdict: PASS (green), SO-SO (yellow), FAIL (red).
|
|
Designed for LoRA training dataset filtering workflows.
|
|
"""
|
|
|
|
# Shot type presets apply multipliers to base thresholds.
|
|
# "custom" uses raw slider values with no scaling.
|
|
SHOT_PROFILES = {
|
|
"custom": {
|
|
"blur": 1.0,
|
|
"sharpness": 1.0,
|
|
"noise": 1.0,
|
|
"clipping": 1.0,
|
|
"entropy": 1.0,
|
|
"block_size": 24,
|
|
},
|
|
"close-up": {
|
|
"blur": 0.6,
|
|
"sharpness": 1.5,
|
|
"noise": 0.8,
|
|
"clipping": 0.8,
|
|
"entropy": 1.1,
|
|
"block_size": 16,
|
|
},
|
|
"medium": {
|
|
"blur": 1.0,
|
|
"sharpness": 1.0,
|
|
"noise": 1.0,
|
|
"clipping": 1.0,
|
|
"entropy": 1.0,
|
|
"block_size": 24,
|
|
},
|
|
"wide / full-body": {
|
|
"blur": 1.6,
|
|
"sharpness": 0.6,
|
|
"noise": 1.3,
|
|
"clipping": 1.2,
|
|
"entropy": 0.85,
|
|
"block_size": 32,
|
|
},
|
|
}
|
|
|
|
# Margin zone multiplier for SO-SO detection.
|
|
# Upper-bound metrics (blur, noise, clip): score <= threshold*MF → marginal
|
|
# Lower-bound metrics (sharp, entropy): score >= threshold/MF → marginal
|
|
MARGIN_FACTOR = 1.4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"shot_type": (["custom", "close-up", "medium", "wide / full-body"],),
|
|
"blur_mode": (["full image", "center-weighted"],),
|
|
"blur_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.40,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"tooltip": "Max blur ratio. Presets scale this; 'custom' uses as-is.",
|
|
},
|
|
),
|
|
"blur_var_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 30.0,
|
|
"min": 5.0,
|
|
"max": 500.0,
|
|
"step": 5.0,
|
|
"tooltip": "Laplacian variance per block. AI-generated: 20-50, photos: 80-150.",
|
|
},
|
|
),
|
|
"sharpness_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 15.0,
|
|
"min": 0.0,
|
|
"max": 500.0,
|
|
"step": 0.5,
|
|
"tooltip": "Min sharpness. Presets scale this; 'custom' uses as-is.",
|
|
},
|
|
),
|
|
"noise_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 25.0,
|
|
"min": 0.0,
|
|
"max": 100.0,
|
|
"step": 0.5,
|
|
"tooltip": "Max noise level. Presets scale this; 'custom' uses as-is.",
|
|
},
|
|
),
|
|
"clipping_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.02,
|
|
"min": 0.0,
|
|
"max": 1.0,
|
|
"step": 0.001,
|
|
"tooltip": "Max clipping ratio. Presets scale this; 'custom' uses as-is.",
|
|
},
|
|
),
|
|
"entropy_threshold": (
|
|
"FLOAT",
|
|
{
|
|
"default": 5.0,
|
|
"min": 0.0,
|
|
"max": 8.0,
|
|
"step": 0.1,
|
|
"tooltip": "Min entropy in bits. Presets scale this; 'custom' uses as-is.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "INT", "STRING", "FLOAT", "FLOAT", "FLOAT", "FLOAT", "FLOAT")
|
|
RETURN_NAMES = ("badge", "verdict", "report", "blur_score", "sharpness_score", "noise_score", "clipping_score", "entropy_score")
|
|
FUNCTION = "analyze"
|
|
CATEGORY = "image/analysis"
|
|
|
|
def analyze(self, image, shot_type, blur_mode, blur_threshold, blur_var_threshold,
|
|
sharpness_threshold, noise_threshold, clipping_threshold, entropy_threshold):
|
|
|
|
profile = self.SHOT_PROFILES[shot_type]
|
|
|
|
# Effective thresholds (custom = 1.0 multiplier = raw slider value)
|
|
eff_blur_t = blur_threshold * profile["blur"]
|
|
eff_sharp_t = sharpness_threshold * profile["sharpness"]
|
|
eff_noise_t = noise_threshold * profile["noise"]
|
|
eff_clip_t = clipping_threshold * profile["clipping"]
|
|
eff_entropy_t = entropy_threshold * profile["entropy"]
|
|
eff_block = profile["block_size"]
|
|
|
|
# Tensor to numpy
|
|
img_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
|
|
|
|
# --- Run all 5 analyses ---
|
|
center_weighted = (blur_mode == "center-weighted")
|
|
blur_score = self._blur_detection(gray, eff_block, center_weighted, blur_var_threshold)
|
|
sharpness_score = self._sharpness_hybrid(gray)
|
|
noise_score = self._noise_estimation(gray)
|
|
clipping_score = self._clipping_analysis(gray)
|
|
entropy_score = self._entropy_analysis(gray)
|
|
|
|
# --- Three-tier evaluation per metric ---
|
|
mf = self.MARGIN_FACTOR
|
|
|
|
checks = [
|
|
self._evaluate_upper(f"Blur [{blur_mode}]", blur_score, eff_blur_t, mf),
|
|
self._evaluate_lower("Sharpness", sharpness_score, eff_sharp_t, mf),
|
|
self._evaluate_upper("Noise", noise_score, eff_noise_t, mf),
|
|
self._evaluate_upper("Clipping", clipping_score, eff_clip_t, mf),
|
|
self._evaluate_lower("Entropy", entropy_score, eff_entropy_t, mf),
|
|
]
|
|
|
|
# --- Overall verdict ---
|
|
# verdict int: 0=FAIL, 1=SO-SO, 2=PASS (for Switch node routing)
|
|
states = [c["state"] for c in checks]
|
|
if "fail" in states:
|
|
verdict_label = "FAIL"
|
|
verdict_int = 0
|
|
elif "marginal" in states:
|
|
verdict_label = "SO-SO"
|
|
verdict_int = 1
|
|
else:
|
|
verdict_label = "PASS"
|
|
verdict_int = 2
|
|
|
|
# --- Text report ---
|
|
report = self._build_report(verdict_label, shot_type, blur_mode, blur_var_threshold, checks)
|
|
|
|
# --- Visual badge ---
|
|
badge_tensor = self._render_badge(verdict_label, checks, shot_type, blur_mode, blur_var_threshold)
|
|
|
|
return (badge_tensor, verdict_int, report, blur_score, sharpness_score,
|
|
noise_score, clipping_score, entropy_score)
|
|
|
|
# ----------------------------------------------------------------
|
|
# Three-tier evaluation
|
|
# ----------------------------------------------------------------
|
|
|
|
def _evaluate_upper(self, name, score, threshold, margin_factor):
|
|
"""Evaluate metric where lower is better (blur, noise, clipping)."""
|
|
margin_limit = threshold * margin_factor
|
|
if score <= threshold:
|
|
state = "pass"
|
|
elif score <= margin_limit:
|
|
state = "marginal"
|
|
else:
|
|
state = "fail"
|
|
return {
|
|
"name": name,
|
|
"score": score,
|
|
"threshold": threshold,
|
|
"condition": f"< {threshold:.3f}",
|
|
"margin": f"< {margin_limit:.3f}",
|
|
"state": state,
|
|
}
|
|
|
|
def _evaluate_lower(self, name, score, threshold, margin_factor):
|
|
"""Evaluate metric where higher is better (sharpness, entropy)."""
|
|
margin_limit = threshold / margin_factor
|
|
if score >= threshold:
|
|
state = "pass"
|
|
elif score >= margin_limit:
|
|
state = "marginal"
|
|
else:
|
|
state = "fail"
|
|
return {
|
|
"name": name,
|
|
"score": score,
|
|
"threshold": threshold,
|
|
"condition": f"> {threshold:.1f}",
|
|
"margin": f"> {margin_limit:.1f}",
|
|
"state": state,
|
|
}
|
|
|
|
# ----------------------------------------------------------------
|
|
# Analysis methods
|
|
# ----------------------------------------------------------------
|
|
|
|
def _blur_detection(self, gray, block_size, center_weighted=False, var_threshold=30.0):
|
|
"""
|
|
Localized Laplacian variance blur detection.
|
|
|
|
In center-weighted mode, blocks are weighted by distance from center:
|
|
- Center 30%: weight 3.0x (subject area)
|
|
- Mid ring 30-60%: weight 1.0x
|
|
- Outer 60%+: weight 0.3x (background)
|
|
"""
|
|
h, w = gray.shape
|
|
cx, cy = w / 2.0, h / 2.0
|
|
max_dist = np.sqrt(cx ** 2 + cy ** 2)
|
|
|
|
weighted_blur = 0.0
|
|
total_weight = 0.0
|
|
|
|
for y in range(0, h - block_size + 1, block_size):
|
|
for x in range(0, w - block_size + 1, block_size):
|
|
block = gray[y:y + block_size, x:x + block_size]
|
|
lap_var = cv2.Laplacian(block, cv2.CV_64F).var()
|
|
is_blurry = 1.0 if lap_var < var_threshold else 0.0
|
|
|
|
if center_weighted:
|
|
bx = x + block_size / 2.0
|
|
by = y + block_size / 2.0
|
|
dist = np.sqrt((bx - cx) ** 2 + (by - cy) ** 2) / max_dist
|
|
if dist < 0.3:
|
|
weight = 3.0
|
|
elif dist < 0.6:
|
|
weight = 1.0
|
|
else:
|
|
weight = 0.3
|
|
else:
|
|
weight = 1.0
|
|
|
|
weighted_blur += is_blurry * weight
|
|
total_weight += weight
|
|
|
|
return weighted_blur / max(total_weight, 1.0)
|
|
|
|
def _sharpness_hybrid(self, gray):
|
|
"""Hybrid sharpness: Laplacian variance + Tenengrad."""
|
|
lap = cv2.Laplacian(gray, cv2.CV_64F)
|
|
lap_score = lap.var()
|
|
gx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
|
|
gy = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
|
|
ten_score = np.mean(gx ** 2 + gy ** 2)
|
|
return (lap_score + ten_score / 1000.0) / 2.0
|
|
|
|
def _noise_estimation(self, gray):
|
|
"""Gaussian difference noise estimation."""
|
|
blurred = cv2.GaussianBlur(gray.astype(np.float64), (5, 5), 0)
|
|
diff = np.abs(gray.astype(np.float64) - blurred)
|
|
return np.mean(diff)
|
|
|
|
def _clipping_analysis(self, gray, threshold=5):
|
|
"""Highlight/shadow clipping ratio."""
|
|
total = gray.size
|
|
shadows = np.sum(gray <= threshold)
|
|
highlights = np.sum(gray >= 255 - threshold)
|
|
return (shadows + highlights) / total
|
|
|
|
def _entropy_analysis(self, gray):
|
|
"""Shannon entropy in bits (0-8)."""
|
|
hist, _ = np.histogram(gray.flatten(), bins=256, range=(0, 256))
|
|
hist = hist / hist.sum()
|
|
hist = hist[hist > 0]
|
|
return -np.sum(hist * np.log2(hist))
|
|
|
|
# ----------------------------------------------------------------
|
|
# Report
|
|
# ----------------------------------------------------------------
|
|
|
|
def _build_report(self, verdict_label, shot_type, blur_mode, var_threshold, checks):
|
|
lines = []
|
|
if verdict_label == "PASS":
|
|
lines.append("PASS — Suitable for dataset")
|
|
elif verdict_label == "SO-SO":
|
|
lines.append("SO-SO — Marginal quality, review recommended")
|
|
else:
|
|
lines.append("FAIL — Rejected")
|
|
lines.append("-" * 44)
|
|
lines.append(f"Shot: {shot_type} | Blur: {blur_mode} | Var_t: {var_threshold:.0f}")
|
|
lines.append("")
|
|
for c in checks:
|
|
if c["state"] == "pass":
|
|
tag = "[OK] "
|
|
elif c["state"] == "marginal":
|
|
tag = "[~] "
|
|
else:
|
|
tag = "[FAIL]"
|
|
lines.append(f" {tag} {c['name']}: {c['score']:.4f} (need {c['condition']})")
|
|
return "\n".join(lines)
|
|
|
|
# ----------------------------------------------------------------
|
|
# Badge renderer
|
|
# ----------------------------------------------------------------
|
|
|
|
def _render_badge(self, verdict_label, checks, shot_type, blur_mode, var_threshold):
|
|
"""Renders a green/yellow/red visual badge with metric breakdown."""
|
|
width, height = 420, 310
|
|
if verdict_label == "PASS":
|
|
bg = (34, 120, 34)
|
|
elif verdict_label == "SO-SO":
|
|
bg = (180, 140, 20)
|
|
else:
|
|
bg = (170, 30, 30)
|
|
|
|
img = Image.new("RGB", (width, height), bg)
|
|
draw = ImageDraw.Draw(img)
|
|
|
|
try:
|
|
font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 32)
|
|
font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
|
|
font_meta = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 12)
|
|
except (OSError, IOError):
|
|
font_large = ImageFont.load_default()
|
|
font_small = ImageFont.load_default()
|
|
font_meta = ImageFont.load_default()
|
|
|
|
# Verdict header
|
|
icons = {"PASS": "\u2713", "SO-SO": "~", "FAIL": "\u2717"}
|
|
header = f"{icons.get(verdict_label, '')} {verdict_label}"
|
|
draw.text((width // 2 - 60, 10), header, fill="white", font=font_large)
|
|
|
|
# Shot info
|
|
meta = f"[{shot_type}]"
|
|
draw.text((width // 2 - 30, 50), meta, fill=(220, 220, 220), font=font_meta)
|
|
|
|
# Separator
|
|
draw.line([(15, 72), (width - 15, 72)], fill=(255, 255, 255, 120), width=1)
|
|
|
|
# Metric rows
|
|
y = 82
|
|
for c in checks:
|
|
state = c["state"]
|
|
if state == "pass":
|
|
icon = "\u2713"
|
|
color = (180, 255, 180)
|
|
icon_color = (80, 200, 80)
|
|
elif state == "marginal":
|
|
icon = "~"
|
|
color = (255, 240, 160)
|
|
icon_color = (220, 200, 60)
|
|
else:
|
|
icon = "\u2717"
|
|
color = (255, 170, 170)
|
|
icon_color = (255, 80, 80)
|
|
|
|
draw.text((20, y), icon, fill=icon_color, font=font_small)
|
|
draw.text((42, y), f"{c['name']}:", fill=color, font=font_small)
|
|
draw.text((210, y), f"{c['score']:.3f}", fill="white", font=font_small)
|
|
draw.text((290, y), f"({c['condition']})", fill=(170, 170, 170), font=font_small)
|
|
y += 40
|
|
|
|
# Footer
|
|
draw.text((width // 2 - 48, height - 20), "ImageQualityGate", fill=(200, 200, 200), font=font_meta)
|
|
|
|
img_np = np.array(img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(img_np).unsqueeze(0)
|
|
|
|
|
|
# ----------------------------------------------------------------
|
|
# ComfyUI Registration
|
|
# ----------------------------------------------------------------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"ImageQualityGate": ImageQualityGate
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"ImageQualityGate": "Image Quality Gate"
|
|
}
|