335 lines
10 KiB
Python
335 lines
10 KiB
Python
import re
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
def _optional_cv2():
|
|
try:
|
|
import cv2
|
|
except Exception:
|
|
return None
|
|
return cv2
|
|
|
|
|
|
def _normalize_selector(selector):
|
|
text = re.sub(r"\s+", " ", str(selector or "").strip().lower())
|
|
if not text:
|
|
return "best"
|
|
if "best" in text:
|
|
return "best"
|
|
if any(term in text for term in ("no face", "faceless", "face out", "without face")):
|
|
return "no_face"
|
|
if any(term in text for term in ("blurry", "blurred", "blur", "out of focus", "soft focus")):
|
|
return "blurry"
|
|
if any(term in text for term in ("moving", "motion", "movement", "dynamic", "action")):
|
|
return "moving"
|
|
if any(term in text for term in ("static", "still", "steady", "posed", "pose")):
|
|
return "static"
|
|
if any(term in text for term in ("flicker", "lighting", "light fluctuation", "flash")):
|
|
return "lighting"
|
|
return "best"
|
|
|
|
|
|
def _selector_rank(selector):
|
|
text = re.sub(r"\s+", " ", str(selector or "").strip().lower())
|
|
patterns = (
|
|
r"(?:rank|option|variant|choice|pick)\s*#?\s*(\d+)",
|
|
r"#\s*(\d+)",
|
|
)
|
|
for pattern in patterns:
|
|
match = re.search(pattern, text)
|
|
if match:
|
|
return max(1, int(match.group(1)))
|
|
return 1
|
|
|
|
|
|
def _normalized(values):
|
|
values = np.asarray(values, dtype=np.float32)
|
|
if values.size == 0:
|
|
return values
|
|
lo = float(np.min(values))
|
|
hi = float(np.max(values))
|
|
if hi - lo < 1e-6:
|
|
return np.full_like(values, 0.5, dtype=np.float32)
|
|
return (values - lo) / (hi - lo)
|
|
|
|
|
|
def _smooth(values, radius=2):
|
|
values = np.asarray(values, dtype=np.float32)
|
|
if values.size <= 1:
|
|
return values
|
|
out = np.empty_like(values)
|
|
for idx in range(values.size):
|
|
lo = max(0, idx - radius)
|
|
hi = min(values.size, idx + radius + 1)
|
|
out[idx] = float(np.mean(values[lo:hi]))
|
|
return out
|
|
|
|
|
|
def _edge_weights(count):
|
|
weights = np.ones(count, dtype=np.float32)
|
|
if count < 25:
|
|
return weights
|
|
margin = max(1, int(round(count * 0.08)))
|
|
weights[:margin] = 0.65
|
|
weights[-margin:] = 0.65
|
|
return weights
|
|
|
|
|
|
def _center_tiebreak(count):
|
|
if count <= 1:
|
|
return np.zeros(count, dtype=np.float32)
|
|
midpoint = (count - 1) / 2.0
|
|
distances = np.abs(np.arange(count, dtype=np.float32) - midpoint)
|
|
return (1.0 - (distances / max(midpoint, 1.0))) * 1e-4
|
|
|
|
|
|
def _default_min_separation(frame_count, desired_count):
|
|
if desired_count <= 1 or frame_count <= 2:
|
|
return 0
|
|
ratio_gap = int(round(frame_count * 0.12))
|
|
spread_cap = frame_count // max(desired_count + 1, 1)
|
|
return max(1, min(spread_cap, max(8, ratio_gap)))
|
|
|
|
|
|
def _gray_frame(rgb_frame):
|
|
frame = rgb_frame.astype(np.float32)
|
|
return frame[..., 0] * 0.299 + frame[..., 1] * 0.587 + frame[..., 2] * 0.114
|
|
|
|
|
|
def _blur_scores(gray_frames):
|
|
cv2 = _optional_cv2()
|
|
scores = []
|
|
for gray in gray_frames:
|
|
if cv2 is not None:
|
|
scores.append(float(cv2.Laplacian(gray.astype(np.uint8), cv2.CV_64F).var()))
|
|
continue
|
|
gy, gx = np.gradient(gray.astype(np.float32))
|
|
scores.append(float(np.var(gx) + np.var(gy)))
|
|
return np.asarray(scores, dtype=np.float32)
|
|
|
|
|
|
def _motion_and_lighting(gray_frames):
|
|
count = len(gray_frames)
|
|
motion = np.zeros(count, dtype=np.float32)
|
|
lighting = np.zeros(count, dtype=np.float32)
|
|
means = np.asarray([float(np.mean(gray)) for gray in gray_frames], dtype=np.float32)
|
|
|
|
for idx, gray in enumerate(gray_frames):
|
|
diffs = []
|
|
light_diffs = []
|
|
if idx > 0:
|
|
diffs.append(float(np.mean(np.abs(gray - gray_frames[idx - 1]))))
|
|
light_diffs.append(abs(float(means[idx] - means[idx - 1])))
|
|
if idx + 1 < count:
|
|
diffs.append(float(np.mean(np.abs(gray - gray_frames[idx + 1]))))
|
|
light_diffs.append(abs(float(means[idx] - means[idx + 1])))
|
|
motion[idx] = float(np.mean(diffs)) if diffs else 0.0
|
|
lighting[idx] = float(np.mean(light_diffs)) if light_diffs else 0.0
|
|
|
|
return _smooth(motion), _smooth(lighting)
|
|
|
|
|
|
def _mediapipe_face_scores(rgb_frames):
|
|
try:
|
|
import mediapipe as mp
|
|
except Exception:
|
|
return None
|
|
|
|
scores = []
|
|
try:
|
|
detector = mp.solutions.face_detection.FaceDetection(
|
|
model_selection=1,
|
|
min_detection_confidence=0.25,
|
|
)
|
|
except Exception:
|
|
return None
|
|
|
|
try:
|
|
for frame in rgb_frames:
|
|
result = detector.process(frame)
|
|
best = 0.0
|
|
for detection in result.detections or []:
|
|
confidence = float(detection.score[0]) if detection.score else 0.0
|
|
box = detection.location_data.relative_bounding_box
|
|
area = max(0.0, float(box.width)) * max(0.0, float(box.height))
|
|
best = max(best, confidence * (1.0 + min(area, 0.25)))
|
|
scores.append(best)
|
|
except Exception:
|
|
return None
|
|
finally:
|
|
detector.close()
|
|
|
|
return np.asarray(scores, dtype=np.float32)
|
|
|
|
|
|
def _opencv_face_scores(rgb_frames):
|
|
cv2 = _optional_cv2()
|
|
if cv2 is None:
|
|
return np.zeros(len(rgb_frames), dtype=np.float32)
|
|
|
|
cascade_names = (
|
|
"haarcascade_frontalface_default.xml",
|
|
"haarcascade_frontalface_alt2.xml",
|
|
"haarcascade_profileface.xml",
|
|
)
|
|
cascades = []
|
|
for name in cascade_names:
|
|
path = cv2.data.haarcascades + name
|
|
cascade = cv2.CascadeClassifier(path)
|
|
if not cascade.empty():
|
|
cascades.append(cascade)
|
|
if not cascades:
|
|
return np.zeros(len(rgb_frames), dtype=np.float32)
|
|
|
|
scores = []
|
|
for frame in rgb_frames:
|
|
gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
|
|
height, width = gray.shape[:2]
|
|
min_size = (max(20, width // 12), max(20, height // 12))
|
|
best_area = 0.0
|
|
for cascade in cascades:
|
|
faces = cascade.detectMultiScale(
|
|
gray,
|
|
scaleFactor=1.1,
|
|
minNeighbors=4,
|
|
minSize=min_size,
|
|
)
|
|
for _, _, face_w, face_h in faces:
|
|
best_area = max(best_area, float(face_w * face_h) / float(width * height))
|
|
scores.append(min(1.0, best_area * 12.0))
|
|
return np.asarray(scores, dtype=np.float32)
|
|
|
|
|
|
def _score_metrics(metrics, mode):
|
|
count = len(metrics["blur"])
|
|
sharp = _normalized(metrics["blur"])
|
|
motion = _normalized(metrics["motion"])
|
|
lighting = _normalized(metrics["lighting"])
|
|
face = _normalized(metrics["face"])
|
|
|
|
stable = 1.0 - lighting
|
|
low_motion = 1.0 - motion
|
|
no_face = 1.0 - face
|
|
blurry = 1.0 - sharp
|
|
moderate_motion = np.clip(1.0 - np.abs(motion - 0.55) * 2.0, 0.0, 1.0)
|
|
|
|
if mode == "no_face":
|
|
score = no_face * 0.55 + sharp * 0.20 + stable * 0.15 + low_motion * 0.10
|
|
elif mode == "blurry":
|
|
score = blurry * 0.70 + stable * 0.15 + face * 0.10 + low_motion * 0.05
|
|
elif mode == "moving":
|
|
score = motion * 0.55 + sharp * 0.25 + stable * 0.15 + face * 0.05
|
|
elif mode == "static":
|
|
score = low_motion * 0.45 + sharp * 0.35 + stable * 0.15 + face * 0.05
|
|
elif mode == "lighting":
|
|
score = lighting * 0.65 + sharp * 0.20 + motion * 0.10 + face * 0.05
|
|
else:
|
|
score = sharp * 0.50 + stable * 0.20 + moderate_motion * 0.15 + face * 0.15
|
|
|
|
return score * _edge_weights(count) + _center_tiebreak(count)
|
|
|
|
|
|
def _ranked_indices(metrics, mode, count=1, min_separation=None):
|
|
frame_count = len(metrics["blur"])
|
|
if frame_count == 0:
|
|
return []
|
|
|
|
wanted = min(max(1, int(count)), frame_count)
|
|
gap = (
|
|
_default_min_separation(frame_count, wanted)
|
|
if min_separation is None
|
|
else max(0, int(min_separation))
|
|
)
|
|
scores = _score_metrics(metrics, mode)
|
|
order = [int(index) for index in np.argsort(scores)[::-1]]
|
|
chosen = []
|
|
relax_by = max(1, int(round(max(gap, 1) * 0.25)))
|
|
|
|
while len(chosen) < wanted:
|
|
added = False
|
|
for index in order:
|
|
if index in chosen:
|
|
continue
|
|
if all(abs(index - other) >= gap for other in chosen):
|
|
chosen.append(index)
|
|
added = True
|
|
if len(chosen) == wanted:
|
|
break
|
|
if len(chosen) == wanted or gap == 0:
|
|
break
|
|
if not added or len(chosen) < wanted:
|
|
gap = max(0, gap - relax_by)
|
|
|
|
return chosen
|
|
|
|
|
|
class AIA_SelectVideoFrame:
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "select"
|
|
CATEGORY = "AIA/video"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"selector": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
@staticmethod
|
|
def _detect_face_scores(rgb_frames):
|
|
scores = _mediapipe_face_scores(rgb_frames)
|
|
if scores is not None:
|
|
return scores
|
|
return _opencv_face_scores(rgb_frames)
|
|
|
|
@classmethod
|
|
def _metrics(cls, frames):
|
|
rgb_frames = np.clip(frames * 255.0, 0, 255).astype(np.uint8)
|
|
gray_frames = [_gray_frame(frame) for frame in rgb_frames]
|
|
blur = _blur_scores(gray_frames)
|
|
motion, lighting = _motion_and_lighting(gray_frames)
|
|
face = cls._detect_face_scores(rgb_frames)
|
|
return {
|
|
"blur": blur,
|
|
"motion": motion,
|
|
"lighting": lighting,
|
|
"face": np.asarray(face, dtype=np.float32),
|
|
}
|
|
|
|
@staticmethod
|
|
def _choose_index(metrics, mode):
|
|
return _ranked_indices(metrics, mode, count=1)[0]
|
|
|
|
@staticmethod
|
|
def _ranked_indices(metrics, mode, count=1, min_separation=None):
|
|
return _ranked_indices(metrics, mode, count=count, min_separation=min_separation)
|
|
|
|
def select(self, image, selector):
|
|
if image is None or image.shape[0] == 0:
|
|
raise ValueError("AIA_SelectVideoFrame received an empty image batch")
|
|
|
|
frames = image.detach().cpu().numpy()
|
|
mode = _normalize_selector(selector)
|
|
rank = _selector_rank(selector)
|
|
metrics = self._metrics(frames)
|
|
indices = self._ranked_indices(metrics, mode, count=rank)
|
|
index = indices[min(rank - 1, len(indices) - 1)]
|
|
print(
|
|
"AIA_SelectVideoFrame: "
|
|
f"selector={selector!r} mode={mode} rank={rank} index={index}/{image.shape[0]}"
|
|
)
|
|
return (image[index : index + 1].to(device=image.device, dtype=image.dtype),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AIA_SelectVideoFrame": AIA_SelectVideoFrame,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AIA_SelectVideoFrame": "AIA Select Video Frame",
|
|
}
|