1733 lines
78 KiB
Python
1733 lines
78 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
face_checker.py — LoRA 训练素材自动审核 + 整理工具
|
||
====================================================
|
||
功能：
|
||
1. 扫描照片目录，自动判断每张图是否适合做 LoRA 训练素材
|
||
2. 检测项：分辨率 / 清晰度 / 人脸数 / 人脸占比 / 正侧脸角度 / 五官遮挡 / 重复图
|
||
3. 生成可视化 HTML 审核报告（图片墙 + 红黄绿标记 + 原因）
|
||
4. 一键把合格照片整理成训练集（重命名 img_001.jpg + EXIF 修正）
|
||
|
||
用法：
|
||
python face_checker.py check <照片目录> [--out 报告输出目录]
|
||
python face_checker.py prepare <照片目录> --out <训练集目录>
|
||
python face_checker.py batch <照片目录> --out <训练集目录>
|
||
|
||
依赖：opencv-python, pillow, numpy（musubi-tuner 虚拟环境已具备）
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import math
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import threading
|
||
import traceback
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
|
||
# 模块级编码兜底：无论 CLI 还是被 GUI import，stdout/stderr 都按 UTF-8 处理，避免 GBK 控制台报错
|
||
for _stream in (sys.stdout, sys.stderr):
|
||
try:
|
||
_stream.reconfigure(encoding="utf-8", errors="replace")
|
||
except Exception:
|
||
pass
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image, ImageOps, ImageDraw
|
||
|
||
# ---------- 常量 ----------
|
||
SCRIPT_DIR = Path(__file__).resolve().parent
|
||
# OpenCV DNN 的 ONNX 加载不支持非 ASCII 路径（C++ 层窄字符 IO），
|
||
# 因此模型优先从 ASCII 路径加载，找不到才回退到脚本目录
|
||
_ASCII_CANDIDATES = [
|
||
Path(r"D:\AI\sd\musubi-tuner\models_yunet\face_detection_yunet.onnx"),
|
||
Path(r"D:\AI\sd\musubi-tuner\face_detection_yunet.onnx"),
|
||
]
|
||
YUNET_PATH = next((p for p in _ASCII_CANDIDATES if p.exists()), SCRIPT_DIR / "models" / "face_detection_yunet.onnx")
|
||
|
||
MIN_RESOLUTION = 1024 # 合格最小短边分辨率
|
||
BLUR_THRESHOLD = 100.0 # Laplacian 方差低于此值判模糊（经验值，可调）
|
||
MIN_FACE_RATIO = 0.05 # 人脸框面积 / 图面积，低于此值判"人脸太小"
|
||
DUPLICATE_THRESHOLD = 5 # pHash 汉明距离低于此值判重复
|
||
SIDE_ANGLE_THRESHOLD = 0.35 # 侧脸判定阈值（landmark 几何比）
|
||
IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
|
||
|
||
STATUS = {"PASS": "合格", "WARN": "警告", "FAIL": "不合格"}
|
||
|
||
|
||
# ---------- 人脸检测 ----------
|
||
def _ensure_ascii_model_path(path):
|
||
"""OpenCV ONNX importer 在 Windows 上读不了非 ASCII 路径（中文目录）。
|
||
若路径含非 ASCII 字符，自动拷贝到 %TEMP% 的 ASCII 路径并返回该路径。"""
|
||
try:
|
||
path.encode("ascii")
|
||
return path # 纯 ASCII，直接用
|
||
except UnicodeEncodeError:
|
||
pass
|
||
import shutil
|
||
import tempfile
|
||
src = Path(path)
|
||
dst = Path(tempfile.gettempdir()) / f"yunet_{hash(str(src)) & 0xffffffff}.onnx"
|
||
if not dst.exists():
|
||
shutil.copy(src, dst)
|
||
return str(dst)
|
||
|
||
|
||
class FaceDetector:
|
||
"""基于 OpenCV YuNet 的人脸检测 + 5 点 landmark"""
|
||
|
||
def __init__(self, model_path=YUNET_PATH):
|
||
if not Path(model_path).exists():
|
||
raise FileNotFoundError(
|
||
f"人脸检测模型不存在: {model_path}\n请下载 face_detection_yunet.onnx 放到 tools/models/ 目录"
|
||
)
|
||
# Windows 上 OpenCV ONNX importer 读不了含中文/非 ASCII 的路径 → 自动拷贝到 %TEMP% ASCII 路径
|
||
model_path = _ensure_ascii_model_path(str(model_path))
|
||
self.detector = cv2.FaceDetectorYN_create(model_path, "", (320, 320), 0.6, 0.3, 5000)
|
||
|
||
def detect(self, img_bgr):
|
||
"""返回人脸列表，每项 = [x, y, w, h, landmarks(10 个浮点: 右眼x,y 左眼x,y 鼻x,y 右嘴x,y 左嘴x,y), score]"""
|
||
h, w = img_bgr.shape[:2]
|
||
self.detector.setInputSize((w, h))
|
||
_, faces = self.detector.detect(img_bgr)
|
||
if faces is None:
|
||
return []
|
||
return faces.tolist()
|
||
|
||
|
||
def load_image(path):
|
||
"""读取图片并修正 EXIF 方向，返回 (RGB ndarray, 原图信息)"""
|
||
img = Image.open(path)
|
||
img = ImageOps.exif_transpose(img)
|
||
return img
|
||
|
||
|
||
def exif_orientation(path):
|
||
try:
|
||
with Image.open(path) as im:
|
||
return im.getexif().get(274, 1)
|
||
except Exception:
|
||
return 1
|
||
|
||
|
||
# ---------- 检测项 ----------
|
||
def check_resolution(img):
|
||
w, h = img.size
|
||
short = min(w, h)
|
||
if short >= MIN_RESOLUTION:
|
||
return STATUS["PASS"], f"{w}×{h} ✓"
|
||
return STATUS["FAIL"], f"{w}×{h}，短边 {short}px < {MIN_RESOLUTION}px（微信压缩图常见）"
|
||
|
||
|
||
def check_sharpness(img, face_box=None):
|
||
"""清晰度检测：优先评估人脸区域（LoRA 选图关键），没人脸才看全图"""
|
||
gray = np.array(img.convert("L"))
|
||
if gray.size == 0:
|
||
return STATUS["FAIL"], "图片为空"
|
||
|
||
if face_box is not None:
|
||
x, y, w, h = [int(v) for v in face_box]
|
||
hh, ww = gray.shape
|
||
# 放大 2 倍评估区域，覆盖整张脸 + 周围
|
||
x0, y0 = max(0, x - w), max(0, y - h)
|
||
x1, y1 = min(ww, x + 2 * w), min(hh, y + 2 * h)
|
||
if (x1 - x0) > 40 and (y1 - y0) > 40:
|
||
gray = gray[y0:y1, x0:x1]
|
||
|
||
# 缩小后再算 Laplacian，避免大图噪声干扰
|
||
scale = 1024.0 / max(gray.shape)
|
||
if scale < 1.0:
|
||
gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
|
||
lap = cv2.Laplacian(gray, cv2.CV_64F).var()
|
||
if lap < BLUR_THRESHOLD * 0.5:
|
||
return STATUS["FAIL"], f"模糊（清晰度 {lap:.0f}，阈值 {BLUR_THRESHOLD}）"
|
||
if lap < BLUR_THRESHOLD:
|
||
return STATUS["WARN"], f"偏模糊（清晰度 {lap:.0f}，建议 ≥{BLUR_THRESHOLD}）"
|
||
return STATUS["PASS"], f"清晰（{lap:.0f}）"
|
||
|
||
|
||
def check_face(img, detector):
|
||
rgb = np.array(img.convert("RGB"))
|
||
bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
|
||
faces = detector.detect(bgr)
|
||
if not faces:
|
||
return STATUS["FAIL"], "未检测到人脸", None
|
||
|
||
if len(faces) > 1:
|
||
return STATUS["WARN"], f"检测到 {len(faces)} 张脸（LoRA 训练需要单人照）", faces
|
||
|
||
f = faces[0]
|
||
x, y, w, h, lms, score = f[0], f[1], f[2], f[3], f[4:14], f[14]
|
||
img_h, img_w = rgb.shape[:2]
|
||
face_area_ratio = (w * h) / (img_w * img_h)
|
||
|
||
reasons = []
|
||
# 人脸占比：按构图类型给差异化提示（半身/全身照脸小是正常的，不判警告；只作信息）
|
||
if face_area_ratio < 0.005:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.2f}%（过小，面部信息极少，仅适合作构图/全身补充）")
|
||
elif face_area_ratio < 0.02:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.1f}%（全身照属正常，脸部需清晰可辨）")
|
||
elif face_area_ratio < 0.05:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.1f}%（半身照，脸部偏小但可用）")
|
||
else:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.1f}%（特写，理想）")
|
||
status = STATUS["PASS"]
|
||
|
||
# 角度判断：双眼 x 距离 vs 双眼 y 偏差；鼻尖是否偏离双眼中心
|
||
re_x, re_y, le_x, le_y = lms[0], lms[1], lms[2], lms[3]
|
||
nose_x, nose_y = lms[4], lms[5]
|
||
eye_dx = abs(re_x - le_x)
|
||
eye_dy = abs(re_y - le_y)
|
||
eye_center_x = (re_x + le_x) / 2
|
||
nose_offset = abs(nose_x - eye_center_x) / max(eye_dx, 1e-6)
|
||
|
||
if eye_dx < 1e-3 or nose_offset > 1.8:
|
||
angle_note = "纯侧面/低头（特征不全）"
|
||
if status == STATUS["PASS"]:
|
||
status = STATUS["WARN"]
|
||
reasons.append(f"角度：{angle_note}")
|
||
elif nose_offset > SIDE_ANGLE_THRESHOLD:
|
||
angle_note = f"较大侧脸（鼻偏移 {nose_offset:.2f}）"
|
||
reasons.append(f"角度：{angle_note}（可用，但正脸更佳）")
|
||
else:
|
||
reasons.append(f"角度：正脸/微侧（鼻偏移 {nose_offset:.2f}）")
|
||
|
||
# 遮挡粗判：landmark 是否全部落在人脸框内
|
||
lm_pts = [(lms[0], lms[1]), (lms[2], lms[3]), (lms[4], lms[5]), (lms[6], lms[7]), (lms[8], lms[9])]
|
||
margin = 0.15 * w
|
||
out_count = sum(1 for px, py in lm_pts if not (x - margin <= px <= x + w + margin and y - margin <= py <= y + h + margin))
|
||
if out_count > 0:
|
||
status = STATUS["WARN"]
|
||
reasons.append(f"五官关键点 {out_count}/5 偏离（疑似口罩/墨镜/刘海遮挡）")
|
||
|
||
reasons.append(f"人脸置信度 {score:.2f}")
|
||
return status, "；".join(reasons), faces
|
||
|
||
|
||
def phash(img, hash_size=16):
|
||
"""感知哈希，用于重复检测"""
|
||
small = img.convert("L").resize((hash_size, hash_size), Image.LANCZOS)
|
||
arr = np.asarray(small, dtype=np.float32)
|
||
med = np.median(arr)
|
||
return (arr > med).astype(np.uint8).flatten()
|
||
|
||
|
||
def hamming(a, b):
|
||
return int(np.count_nonzero(a != b))
|
||
|
||
|
||
# ---------- 主流程 ----------
|
||
DETECT_MAX_SIDE = 1600 # 保留常量（后续可能用）；当前检测用原图
|
||
THUMB_DIR_NAME = ".thumbs"
|
||
_thread_local = threading.local()
|
||
|
||
|
||
def _get_thread_detector():
|
||
"""每个线程独立的 FaceDetector（YuNet 实例非线程安全，共享会误检/串行化）"""
|
||
if not hasattr(_thread_local, "det"):
|
||
_thread_local.det = FaceDetector()
|
||
return _thread_local.det
|
||
|
||
|
||
def _nms(faces, iou_threshold=0.5):
|
||
"""按 score 降序，合并 IoU 重叠的人脸框"""
|
||
if not faces:
|
||
return []
|
||
faces = sorted(faces, key=lambda f: -f[14])
|
||
keep = []
|
||
for f in faces:
|
||
x0, y0, w0, h0 = f[0], f[1], f[2], f[3]
|
||
dup = False
|
||
for k in keep:
|
||
x1, y1, w1, h1 = k[0], k[1], k[2], k[3]
|
||
ix = max(0, min(x0 + w0, x1 + w1) - max(x0, x1))
|
||
iy = max(0, min(y0 + h0, y1 + h1) - max(y0, y1))
|
||
inter = ix * iy
|
||
union = w0 * h0 + w1 * h1 - inter
|
||
if union > 0 and inter / union > iou_threshold:
|
||
dup = True
|
||
break
|
||
if not dup:
|
||
keep.append(f)
|
||
return keep
|
||
|
||
|
||
def _detect_faces_fast(img_rgb, detector, score_threshold=0.6):
|
||
"""
|
||
快速人脸检测（原图检测，保持质量；速度靠外层多线程并行）。
|
||
用 NMS 合并重叠框（合影时多脸不误合并）。
|
||
过滤：人脸中心 y 超过画面 85% 的大框基本是误检（腿/物体当脸）。
|
||
img_rgb: RGB ndarray；返回 faces（原图坐标）。
|
||
"""
|
||
h, w = img_rgb.shape[:2]
|
||
bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
|
||
faces = detector.detect(bgr)
|
||
if not faces:
|
||
return []
|
||
faces = _nms([f for f in faces if f[14] >= score_threshold])
|
||
# 位置合理性过滤：人脸中心 y 在画面最底部 15% 的大框视为误检
|
||
h_img = h
|
||
valid = [f for f in faces if (f[1] + f[3] / 2) / h_img <= 0.85]
|
||
return valid if valid else faces # 全部被过滤则返回原始（防误杀特殊构图）
|
||
|
||
|
||
def _analyze_one(p, detector=None):
|
||
"""
|
||
单张图快速分析：
|
||
- 完整解码（保持画质）+ LANCZOS 缩略到 ≤1600 再分析
|
||
- 缩略图检测人脸（NMS+阈值）
|
||
- 清晰度在人脸区域（缩略图尺度，阈值 60）
|
||
返回结果 dict（供 analyze_images 聚合）。
|
||
"""
|
||
det = _get_thread_detector() if detector is None else detector
|
||
img = load_image(p) # 完整解码 + EXIF 修正
|
||
w0, h0 = img.size
|
||
small = img.copy()
|
||
small.thumbnail((1600, 1600), Image.LANCZOS)
|
||
rgb = np.array(small.convert("RGB"))
|
||
|
||
# GUI 缩略图（供 gradio 快速显示，避免全尺寸加载）
|
||
try:
|
||
thumb_dir = Path(p).parent / THUMB_DIR_NAME
|
||
thumb_dir.mkdir(exist_ok=True)
|
||
thumb = small.copy()
|
||
thumb.thumbnail((400, 400), Image.LANCZOS)
|
||
thumb.save(thumb_dir / f"{Path(p).stem}.jpg", quality=85)
|
||
except Exception:
|
||
pass
|
||
|
||
# 分辨率检查（原始尺寸）
|
||
short = min(w0, h0)
|
||
if short >= MIN_RESOLUTION:
|
||
r_res = (STATUS["PASS"], f"{w0}×{h0} ✓")
|
||
else:
|
||
r_res = (STATUS["FAIL"], f"{w0}×{h0}，短边 {short}px < {MIN_RESOLUTION}px（微信压缩图常见）")
|
||
|
||
# 人脸检测（缩略图尺度），坐标映射回原图尺寸（含 landmarks）
|
||
faces = _detect_faces_fast(rgb, det)
|
||
if faces:
|
||
sx = w0 / rgb.shape[1]
|
||
sy = h0 / rgb.shape[0]
|
||
mapped = []
|
||
for f in faces:
|
||
nf = [f[0]*sx, f[1]*sy, f[2]*sx, f[3]*sy]
|
||
lms = list(f[4:14]) # 5 点 landmark (x,y)*5
|
||
for k in range(0, 10, 2):
|
||
lms[k] *= sx
|
||
lms[k+1] *= sy
|
||
mapped.append(nf + lms + [f[14]])
|
||
faces = mapped
|
||
face_box = faces[0][0:4] if faces else None
|
||
|
||
# 面部/整体清晰度（缩略图尺度）
|
||
lap_face = _lap_of(Image.fromarray(rgb), face_box) if face_box else 0.0
|
||
lap_all = _lap_of(Image.fromarray(rgb), None)
|
||
if face_box is not None and lap_face >= 60:
|
||
r_sharp = (STATUS["PASS"], f"清晰（{lap_face:.0f}）")
|
||
elif face_box is not None:
|
||
r_sharp = (STATUS["FAIL"], f"面部模糊（清晰度 {lap_face:.0f}，阈值 60）")
|
||
else:
|
||
r_sharp = (STATUS["WARN"], f"无人脸，全图清晰度 {lap_all:.0f}")
|
||
|
||
r_face, face_note, _ = check_face_from_faces(Image.fromarray(rgb), faces)
|
||
|
||
all_st = [r_res[0], r_sharp[0], r_face]
|
||
if STATUS["FAIL"] in all_st:
|
||
status = STATUS["FAIL"]
|
||
elif STATUS["WARN"] in all_st:
|
||
status = STATUS["WARN"]
|
||
else:
|
||
status = STATUS["PASS"]
|
||
|
||
return {
|
||
"file": p.name, "path": str(p), "size": [w0, h0],
|
||
"bytes": p.stat().st_size, "status": status,
|
||
"reasons": [f"尺寸：{r_res[1]}", f"清晰度：{r_sharp[1]}", f"人脸：{face_note}"],
|
||
"faces": faces, "phash": phash(Image.fromarray(rgb)).tolist(),
|
||
"exif_orientation": exif_orientation(p),
|
||
}
|
||
|
||
|
||
def analyze_images(img_dir, detector, verbose=True):
|
||
"""分析目录下所有图片（draft 快解码 + 缩略图检测），返回结果列表"""
|
||
img_dir = Path(img_dir)
|
||
images = sorted(
|
||
[p for p in img_dir.iterdir() if p.suffix.lower() in IMG_EXTS and not p.name.startswith(".")]
|
||
)
|
||
if not images:
|
||
print(f"[WARN] 目录 {img_dir} 中没有图片")
|
||
return []
|
||
|
||
results = []
|
||
for p in images:
|
||
r = _analyze_one(p, detector)
|
||
results.append(r)
|
||
if verbose:
|
||
flag = {STATUS["PASS"]: "[OK]", STATUS["WARN"]: "[WARN]", STATUS["FAIL"]: "[FAIL]"}[r["status"]]
|
||
print(f"{flag} {r['file']} [{r['size'][0]}x{r['size'][1]}] {r['status']}")
|
||
|
||
results.sort(key=lambda r: r["file"])
|
||
|
||
# 重复检测（两两 pHash 距离）
|
||
for i in range(len(results)):
|
||
for j in range(i + 1, len(results)):
|
||
a, b = results[i], results[j]
|
||
if a["phash"] is None or b["phash"] is None:
|
||
continue
|
||
d = hamming(np.array(a["phash"]), np.array(b["phash"]))
|
||
if d < DUPLICATE_THRESHOLD:
|
||
note = f"与 {b['file']} 高度相似（可能重复）"
|
||
if note not in a["reasons"]:
|
||
a["reasons"].append(note)
|
||
if a["status"] == STATUS["PASS"]:
|
||
a["status"] = STATUS["WARN"]
|
||
note = f"与 {a['file']} 高度相似（可能重复）"
|
||
if note not in b["reasons"]:
|
||
b["reasons"].append(note)
|
||
if b["status"] == STATUS["PASS"]:
|
||
b["status"] = STATUS["WARN"]
|
||
return results
|
||
|
||
|
||
def check_face_from_faces(img, faces):
|
||
"""
|
||
基于已检测的 faces 做判断（避免重复检测）。
|
||
返回 (status, note, faces)。供 _analyze_one / 其他模块复用。
|
||
"""
|
||
if not faces:
|
||
return STATUS["FAIL"], "未检测到人脸", None
|
||
if len(faces) > 1:
|
||
return STATUS["WARN"], f"检测到 {len(faces)} 张脸（LoRA 训练需要单人照）", faces
|
||
f = faces[0]
|
||
x, y, w, h, lms, score = f[0], f[1], f[2], f[3], f[4:14], f[14]
|
||
rgb = np.array(img.convert("RGB"))
|
||
img_h, img_w = rgb.shape[:2]
|
||
face_area_ratio = (w * h) / (img_w * img_h)
|
||
|
||
reasons = []
|
||
if face_area_ratio < 0.005:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.2f}%（过小，面部信息极少，仅适合作构图/全身补充）")
|
||
elif face_area_ratio < 0.02:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.1f}%（全身照属正常，脸部需清晰可辨）")
|
||
elif face_area_ratio < 0.05:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.1f}%（半身照，脸部偏小但可用）")
|
||
else:
|
||
reasons.append(f"人脸占比 {face_area_ratio*100:.1f}%（特写，理想）")
|
||
status = STATUS["PASS"]
|
||
|
||
re_x, re_y, le_x, le_y = lms[0], lms[1], lms[2], lms[3]
|
||
nose_x, nose_y = lms[4], lms[5]
|
||
eye_dx = abs(re_x - le_x)
|
||
eye_center_x = (re_x + le_x) / 2
|
||
nose_offset = abs(nose_x - eye_center_x) / max(eye_dx, 1e-6)
|
||
|
||
if eye_dx < 1e-3 or nose_offset > 1.8:
|
||
if status == STATUS["PASS"]:
|
||
status = STATUS["WARN"]
|
||
reasons.append(f"角度：纯侧面/低头（特征不全）")
|
||
elif nose_offset > SIDE_ANGLE_THRESHOLD:
|
||
reasons.append(f"角度：较大侧脸（鼻偏移 {nose_offset:.2f}）（可用，但正脸更佳）")
|
||
else:
|
||
reasons.append(f"角度：正脸/微侧（鼻偏移 {nose_offset:.2f}）")
|
||
|
||
lm_pts = [(lms[0], lms[1]), (lms[2], lms[3]), (lms[4], lms[5]), (lms[6], lms[7]), (lms[8], lms[9])]
|
||
margin = 0.15 * w
|
||
out_count = sum(1 for px, py in lm_pts if not (x - margin <= px <= x + w + margin and y - margin <= py <= y + h + margin))
|
||
if out_count > 0:
|
||
status = STATUS["WARN"]
|
||
reasons.append(f"五官关键点 {out_count}/5 偏离（疑似口罩/墨镜/刘海遮挡）")
|
||
|
||
reasons.append(f"人脸置信度 {score:.2f}")
|
||
return status, "；".join(reasons), faces
|
||
|
||
|
||
def draw_annotation(img, faces, path_out):
|
||
"""在人脸框 + landmark 上画标注，保存预览图"""
|
||
draw = ImageDraw.Draw(img)
|
||
for f in faces or []:
|
||
x, y, w, h = f[0], f[1], f[2], f[3]
|
||
lms = f[4:14]
|
||
draw.rectangle([x, y, x + w, y + h], outline=(0, 255, 0), width=3)
|
||
for k in range(0, 10, 2):
|
||
cx, cy = lms[k], lms[k + 1]
|
||
r = 4
|
||
draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=(255, 0, 0))
|
||
img.save(path_out)
|
||
|
||
|
||
# ---------- 智能选图 ----------
|
||
def face_meta(r):
|
||
"""提取单张图的构图/角度特征，用于多样性选图"""
|
||
meta = {"ratio": 0.0, "angle": 0.0, "n_faces": 0}
|
||
if r["faces"] and r["size"][0] > 0:
|
||
f = r["faces"][0]
|
||
w, h = r["size"]
|
||
meta["ratio"] = (f[2] * f[3]) / (w * h)
|
||
meta["n_faces"] = len(r["faces"])
|
||
lms = f[4:14]
|
||
re_x, re_y, le_x, le_y, nose_x = lms[0], lms[1], lms[2], lms[3], lms[4]
|
||
eye_center = (re_x + le_x) / 2
|
||
eye_dx = abs(re_x - le_x) + 1e-6
|
||
meta["angle"] = abs(nose_x - eye_center) / eye_dx
|
||
return meta
|
||
|
||
|
||
def compose_label(ratio):
|
||
"""按人脸框占比分类构图（YuNet 框的是脸部，特写照脸框占比约 8-20%）"""
|
||
if ratio >= 0.08:
|
||
return "特写"
|
||
if ratio >= 0.02:
|
||
return "半身"
|
||
if ratio >= 0.005:
|
||
return "全身"
|
||
return "人很小"
|
||
|
||
|
||
def pick_diverse(results, count=20):
|
||
"""
|
||
从审核结果中贪心挑选多样化的 count 张。
|
||
策略：farthest-point sampling —— 每步选"与已选集合在 视觉/构图/角度 上差异最大"的图。
|
||
效果：自动覆盖 多角度 + 多发型/妆容（视觉差异）+ 多构图（特写/半身/全身）。
|
||
"""
|
||
ok = [r for r in results if r["status"] != STATUS["FAIL"] and r["phash"] is not None]
|
||
if not ok:
|
||
return []
|
||
|
||
phs = [np.array(r["phash"], dtype=np.uint8) for r in ok]
|
||
metas = [face_meta(r) for r in ok]
|
||
quality = [0 if r["status"] == STATUS["PASS"] else 1 for r in ok]
|
||
|
||
def dist(i, j):
|
||
hd = hamming(phs[i], phs[j])
|
||
rd = abs(metas[i]["ratio"] - metas[j]["ratio"]) * 100 # 构图差异（特写 vs 全身）
|
||
ad = abs(metas[i]["angle"] - metas[j]["angle"]) * 6 # 角度差异
|
||
return hd + rd + ad
|
||
|
||
n = len(ok)
|
||
order = sorted(range(n), key=lambda i: (quality[i], -(metas[i]["ratio"] > 0.05)))
|
||
selected = [order[0]]
|
||
while len(selected) < min(count, n):
|
||
best_i, best_d = None, -1
|
||
for i in range(n):
|
||
if i in selected:
|
||
continue
|
||
d = min(dist(i, s) for s in selected)
|
||
d -= 5 if quality[i] else 0 # WARN 图轻微降权，合格图优先
|
||
if d > best_d:
|
||
best_d, best_i = d, i
|
||
if best_i is None:
|
||
break
|
||
selected.append(best_i)
|
||
|
||
picked = [ok[i] for i in selected]
|
||
# 标注构图标签
|
||
for r in picked:
|
||
r["compose"] = compose_label(face_meta(r)["ratio"])
|
||
return picked
|
||
|
||
|
||
def build_pick_html(picked, out_path, count):
|
||
cards = []
|
||
for r in picked:
|
||
compose = r.get("compose", "")
|
||
reasons_html = "<br>".join(f"· {x}" for x in r["reasons"][:2])
|
||
size = f"{r['size'][0]}×{r['size'][1]}" if r["size"][0] else "?"
|
||
cards.append(f"""
|
||
<div class="card">
|
||
<div class="thumb-wrap"><img class="thumb" src="file:///{r['path']}" loading="lazy" onerror="this.style.display='none'"></div>
|
||
<div class="info">
|
||
<div class="row"><span class="badge" style="background:#16a34a">推荐</span>
|
||
<span class="badge2">{compose}</span><span class="fname">{r['file']}</span></div>
|
||
<div class="meta">{size} · {r['bytes']/1024:.0f}KB</div>
|
||
<div class="reasons">{reasons_html}</div>
|
||
</div>
|
||
</div>""")
|
||
html = f"""<!DOCTYPE html>
|
||
<html lang="zh"><head><meta charset="utf-8">
|
||
<title>LoRA 智能选图推荐</title>
|
||
<style>
|
||
body {{ font-family:"Microsoft YaHei",sans-serif; background:#f5f5f5; margin:20px; }}
|
||
h1 {{ font-size:22px; }}
|
||
.summary {{ background:#fff; padding:12px 18px; border-radius:8px; margin-bottom:16px; box-shadow:0 1px 3px #ccc; }}
|
||
.grid {{ display:grid; grid-template-columns:repeat(auto-fill,minmax(280px,1fr)); gap:14px; }}
|
||
.card {{ background:#fff; border:2px solid #16a34a; border-radius:10px; overflow:hidden; box-shadow:0 1px 3px #ccc; }}
|
||
.thumb-wrap {{ height:220px; overflow:hidden; background:#eee; }}
|
||
.thumb {{ width:100%; height:100%; object-fit:cover; }}
|
||
.info {{ padding:10px 12px; }}
|
||
.badge {{ color:#fff; background:#16a34a; padding:2px 10px; border-radius:4px; font-weight:bold; font-size:13px; }}
|
||
.badge2 {{ color:#333; background:#eee; padding:2px 8px; border-radius:4px; font-size:12px; margin-left:6px; }}
|
||
.fname {{ font-weight:bold; margin-left:8px; word-break:break-all; }}
|
||
.meta {{ color:#888; font-size:12px; margin:4px 0; }}
|
||
.reasons {{ font-size:12px; color:#666; line-height:1.5; }}
|
||
</style></head><body>
|
||
<h1>LoRA 智能选图推荐（{len(picked)}/{count}）</h1>
|
||
<div class="summary">按视觉差异 + 构图 + 角度 自动挑出的多样化组合，覆盖多发型/多妆容/多角度/多年龄段。可直接用这些照片训练。</div>
|
||
<div class="grid">{''.join(cards)}</div>
|
||
</body></html>"""
|
||
Path(out_path).write_text(html, encoding="utf-8")
|
||
return out_path
|
||
|
||
|
||
def crop_single(img, faces, margin_scale=1.8, min_side=1024):
|
||
"""
|
||
合影处理：按最大人脸裁切单人区域（正方形，人脸居中放大 margin_scale 倍）。
|
||
返回 (裁切后的 PIL 图, 说明文本)。
|
||
"""
|
||
rgb = np.array(img.convert("RGB"))
|
||
h, w = rgb.shape[:2]
|
||
if not faces:
|
||
return None, "未检测到人脸，无法裁切"
|
||
# 取面积最大的人脸（假设是主体）
|
||
f = max(faces, key=lambda x: x[2] * x[3])
|
||
fx, fy, fw, fh = int(f[0]), int(f[1]), int(f[2]), int(f[3])
|
||
cx, cy = fx + fw / 2.0, fy + fh / 2.0
|
||
side = int(max(fw, fh) * margin_scale)
|
||
# 居中裁切，越界时对齐边界
|
||
x0 = max(0, int(cx - side / 2))
|
||
y0 = max(0, int(cy - side / 2))
|
||
x0 = min(x0, max(0, w - side))
|
||
y0 = min(y0, max(0, h - side))
|
||
x1, y1 = min(w, x0 + side), min(h, y0 + side)
|
||
crop = rgb[y0:y1, x0:x1]
|
||
out = Image.fromarray(crop)
|
||
short = min(out.size)
|
||
note = f"裁切 {out.size[0]}×{out.size[1]}（脸框 {fw}×{fh}）"
|
||
if short < min_side:
|
||
note += f"；⚠️ 短边 {short}px < {min_side}px，裁后分辨率不足，建议放弃"
|
||
else:
|
||
note += f"；短边 {short}px ✓"
|
||
return out, note
|
||
|
||
|
||
def crop_dir(img_dir, out_dir, detector, margin_scale=1.8, min_side=1024):
|
||
"""批量裁切：photos -> croped 目录"""
|
||
img_dir = Path(img_dir)
|
||
out_dir = Path(out_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
images = sorted([p for p in img_dir.iterdir() if p.suffix.lower() in IMG_EXTS and not p.name.startswith(".")])
|
||
done, skipped = 0, []
|
||
for p in images:
|
||
try:
|
||
img = load_image(p)
|
||
rgb = np.array(img.convert("RGB"))
|
||
bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
|
||
faces = detector.detect(bgr)
|
||
if len(faces) > 1:
|
||
print(f"[多人 {len(faces)}] {p.name} -> 裁切最大人脸")
|
||
crop, note = crop_single(img, faces, margin_scale, min_side)
|
||
if crop is None:
|
||
skipped.append(f"{p.name}（{note}）")
|
||
continue
|
||
crop.save(out_dir / p.name, quality=95)
|
||
done += 1
|
||
print(f" [OK] {p.name}: {note}")
|
||
except Exception as e:
|
||
skipped.append(f"{p.name}（{e}）")
|
||
print(f"\n裁切完成: {done} 张 -> {out_dir}")
|
||
if skipped:
|
||
print(f"跳过 {len(skipped)} 张: {skipped}")
|
||
return done
|
||
|
||
|
||
# ---------- 人体检测（OpenCV HOG，零下载） ----------
|
||
class PersonDetector:
|
||
"""基于 OpenCV HOG 的人体检测（用于全身照识别/裁剪）"""
|
||
|
||
def __init__(self):
|
||
self.hog = cv2.HOGDescriptor()
|
||
self.hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
|
||
|
||
def detect(self, img_bgr, max_side=1600):
|
||
"""返回人体框列表 [[x,y,w,h],...]；大图先缩小加速"""
|
||
h, w = img_bgr.shape[:2]
|
||
scale = 1.0
|
||
if max(h, w) > max_side:
|
||
scale = max_side / max(h, w)
|
||
img_bgr = cv2.resize(img_bgr, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA)
|
||
boxes, _ = self.hog.detectMultiScale(img_bgr, winStride=(8, 8), padding=(8, 8), scale=1.05)
|
||
if len(boxes) == 0:
|
||
return []
|
||
boxes = boxes.astype(float).tolist()
|
||
if scale != 1.0:
|
||
boxes = [[x / scale, y / scale, bw / scale, bh / scale] for x, y, bw, bh in boxes]
|
||
return boxes
|
||
|
||
|
||
def crop_centered_halfbody(img, face):
|
||
"""
|
||
半身水平居中（保留原图高度，按人脸水平中心裁掉背景多余侧）：
|
||
修复"人偏左/偏右，另一半全是背景"的构图。与 crop_fullbody 的水平居中思路一致，
|
||
但不改高度（半身不看脚部），只把人物移到画面水平中心。
|
||
"""
|
||
rgb = np.array(img.convert("RGB"))
|
||
h, w = rgb.shape[:2]
|
||
fx, fy, fw, fh = [int(v) for v in face[:4]]
|
||
face_cx = fx + fw / 2
|
||
cx_ratio = face_cx / w
|
||
# 只在明显偏位（中心偏离画面中心 >10%）时裁剪，否则原样返回
|
||
if abs(cx_ratio - 0.5) <= 0.10:
|
||
return img, f"半身保留 {w}×{h}（已居中，无需裁）"
|
||
# 目标宽度：以人物为中心取 高度×0.9（半身自然比例），不超原宽
|
||
body_w = min(w, int(h * 0.9))
|
||
x0 = max(0, int(face_cx - body_w / 2))
|
||
x1 = min(w, x0 + body_w)
|
||
x0 = max(0, x1 - body_w)
|
||
crop = rgb[0:h, x0:x1]
|
||
out = Image.fromarray(crop)
|
||
note = f"半身居中 {out.size[0]}×{out.size[1]}（人脸原在 {cx_ratio*100:.0f}% 处，裁背景侧）"
|
||
return out, note
|
||
|
||
|
||
def crop_fullbody(img, face):
|
||
"""
|
||
全身构图（保留完整人物，不裁上下）：
|
||
- 垂直：保留原图完整高度（手举过头顶/脚部都不被截）
|
||
- 水平：把人物居中（裁剪背景多余侧，人物移到画面中心）
|
||
- 竖构图（h*0.75 >= w）时若人脸明显偏位（>10%），同样收窄宽度居中（原逻辑此时不裁导致人偏一边）
|
||
不强行裁到固定宽高比（那是手臂/脚被截的根源）。
|
||
"""
|
||
rgb = np.array(img.convert("RGB"))
|
||
h, w = rgb.shape[:2]
|
||
fx, fy, fw, fh = [int(v) for v in face[:4]]
|
||
face_cx = fx + fw / 2
|
||
cx_ratio = face_cx / w
|
||
body_w = min(w, int(h * 0.75))
|
||
# 竖构图不触发裁剪（body_w==w）但人脸明显偏位 → 收窄到 h*0.6 以便居中
|
||
if body_w >= w and abs(cx_ratio - 0.5) > 0.10:
|
||
body_w = min(w, int(h * 0.6))
|
||
x0 = max(0, int(face_cx - body_w / 2))
|
||
x1 = min(w, x0 + body_w)
|
||
x0 = max(0, x1 - body_w)
|
||
crop = rgb[0:h, x0:x1] # 高度完整保留，只裁水平
|
||
out = Image.fromarray(crop)
|
||
note = f"全身构图 {out.size[0]}×{out.size[1]}（人物居中，高度完整）"
|
||
return out, note
|
||
|
||
|
||
def crop_face_portrait(img, face, margin=2.5):
|
||
"""
|
||
脸部特写裁剪：人脸居中，边长 = 脸短边 × margin。
|
||
margin 2.0 保留完整发型（发际线/刘海/鬓角不被裁掉）。
|
||
可从任意构图（半身/全身）裁出脸部特写素材。
|
||
"""
|
||
rgb = np.array(img.convert("RGB"))
|
||
fx, fy, fw, fh = [int(v) for v in face[:4]]
|
||
# 脸朝向：鼻尖相对脸框中心偏移 → 侧脸时后脑勺在鼻的反方向
|
||
cx = fx + fw / 2
|
||
nose_offset = (face[8] - cx) / fw if len(face) >= 9 else 0.0
|
||
if abs(nose_offset) > 0.30:
|
||
# 明显侧转：裁框中心向后脑方向平移 0.35×脸宽
|
||
back_dir = -1 if nose_offset > 0 else 1
|
||
cx = cx + back_dir * 0.35 * fw
|
||
side = int(max(fw, fh) * margin)
|
||
# 向上偏：顶部从 脸框上缘 - 1.2×fh（完整保留头顶/发饰/蝴蝶结）
|
||
top = max(0, int(fy - 1.2 * fh))
|
||
bottom = min(rgb.shape[0], top + side)
|
||
if bottom - top < side: # 到底部贴边了，向上补足
|
||
top = max(0, bottom - side)
|
||
x0 = max(0, int(cx - side / 2))
|
||
x1 = min(rgb.shape[1], x0 + side)
|
||
x0 = max(0, x1 - side)
|
||
crop = rgb[top:bottom, x0:x1]
|
||
out = Image.fromarray(crop)
|
||
note = f"脸部特写 {out.size[0]}×{out.size[1]}（脸短边 {min(fw,fh):.0f}px，保留发型）"
|
||
return out, note
|
||
|
||
|
||
def face_sharpness(img, face):
|
||
"""
|
||
人脸清晰度（Laplacian）：紧贴脸框核心 70% 区域，避免背景稀释。
|
||
真实照片的雀斑/眉毛清晰时应有高值。
|
||
"""
|
||
rgb = np.array(img.convert("RGB"))
|
||
h, w = rgb.shape[:2]
|
||
fx, fy, fw, fh = [int(v) for v in face[:4]]
|
||
cx, cy = fx + fw / 2, fy + fh / 2
|
||
cw, ch = fw * 0.7, fh * 0.7
|
||
x0 = max(0, int(cx - cw / 2)); y0 = max(0, int(cy - ch / 2))
|
||
x1 = min(w, int(cx + cw / 2)); y1 = min(h, int(cy + ch / 2))
|
||
if (x1 - x0) < 30 or (y1 - y0) < 30:
|
||
return 0.0
|
||
gray = np.array(Image.fromarray(rgb[y0:y1, x0:x1]).convert("L"))
|
||
return float(cv2.Laplacian(gray, cv2.CV_64F).var())
|
||
|
||
|
||
def face_angle(face):
|
||
"""
|
||
人脸角度：用 landmarks 判断正脸/半侧面/侧面（ComfyUI/A1111 社区标准措辞）。
|
||
nose_offset = 鼻尖偏离双眼中心的距离（相对眼距），实测标定。
|
||
返回 (中文标签, 英文标签)。
|
||
"""
|
||
lms = face[4:14]
|
||
if len(lms) < 10:
|
||
return "正脸", "front view"
|
||
re_x, re_y, le_x, le_y, nose_x = lms[0], lms[1], lms[2], lms[3], lms[4]
|
||
eye_dx = abs(re_x - le_x)
|
||
eye_center = (re_x + le_x) / 2
|
||
if eye_dx < 1e-3:
|
||
return "正脸", "front view"
|
||
nose_offset = abs(nose_x - eye_center) / eye_dx
|
||
if nose_offset > 0.6:
|
||
return "侧面", "side view"
|
||
if nose_offset > 0.15:
|
||
return "半侧面", "three-quarter view"
|
||
return "正脸", "front view"
|
||
|
||
|
||
# ---------- auto 自动流水线 ----------
|
||
# 标签模板：只标角度+构图，不含姿势/服装/背景硬编码（避免雷同；细节由 GUI 打标或 VLM 补充）。
|
||
# ⚠️ 特写模板不再写死 neutral expression——表情由 VLM 描述提供，否则会生成"neutral expression, 大笑"矛盾标签。
|
||
AUTO_TEMPLATES = {
|
||
"特写": "photorealistic portrait, soft natural lighting",
|
||
"半身": "half body shot, natural lighting",
|
||
"全身": "full body shot, natural lighting",
|
||
}
|
||
|
||
# 本地 Ollama 配置（qwen3-vl 视觉模型，本地稳定，不依赖小果磁盘）
|
||
OLLAMA_URL = "http://localhost:11434/api/chat"
|
||
OLLAMA_MODEL = "huihui_ai/qwen3-vl-abliterated:8b"
|
||
# 小果 oMLX 配置（串行调用！多个可用模型）
|
||
OMLX_URL = "http://192.168.1.122:18003/v1/chat/completions"
|
||
OMLX_MODEL = "Qwen3-VL-30B-A3B-Thinking-4bit"
|
||
OMLX_MODEL_32B = "qwen2.5-VL-32B-abliterated-MLX-Q8"
|
||
|
||
SENSENOVA_URL = "https://token.sensenova.cn/v1/chat/completions"
|
||
SENSENOVA_KEY = "sk-aRNj3UwKSLPsDfh15QNTPwbHxahblfaO"
|
||
SENSENOVA_MODEL = "sensenova-6.7-flash-lite"
|
||
|
||
# 后端模型映射：checkbox 选项名 → (来源, 模型)
|
||
BACKEND_MODELS = {
|
||
"ollama": ("ollama", OLLAMA_MODEL),
|
||
"omlx": ("omlx", OMLX_MODEL),
|
||
"omlx-32b": ("omlx", OMLX_MODEL_32B),
|
||
"sensenova": ("sensenova", SENSENOVA_MODEL),
|
||
}
|
||
# 打标主用优先级（勾选组合时选最高的）
|
||
CAPTION_PRIORITY = ["omlx-32b", "omlx", "ollama", "sensenova"]
|
||
|
||
|
||
def _norm_backend(b):
|
||
"""打标/描述用单个后端。若传入列表（勾选组合），按优先级取主用：omlx-32b > omlx > ollama > sensenova"""
|
||
if isinstance(b, (list, tuple)):
|
||
for pref in CAPTION_PRIORITY:
|
||
if pref in b:
|
||
return pref
|
||
return "ollama"
|
||
return b
|
||
|
||
|
||
def _resolve_model(backend):
|
||
"""返回 (base_url, model_name)。omlx-32b / omlx 都用小果 URL，只是模型不同"""
|
||
src, model = BACKEND_MODELS.get(backend, ("ollama", OLLAMA_MODEL))
|
||
if src == "omlx":
|
||
return OMLX_URL, model
|
||
if src == "sensenova":
|
||
return SENSENOVA_URL, model
|
||
return OLLAMA_URL, model
|
||
|
||
|
||
def describe_image_omlx(img_path, timeout=120, face_only=False, backend="ollama"):
|
||
"""
|
||
VLM 图像识别（看图描述人物）。backend：
|
||
"ollama" = 本地 qwen3-vl-abliterated:8b（默认，稳定自主）；
|
||
"omlx" = 小果 Qwen3-VL-30B-A3B-Thinking-4bit（细节更准，串行调用，依赖小果在线）。
|
||
face_only=True：只描述面部特征+发型（特写用，不提服装/背景/姿势）；
|
||
face_only=False：完整描述（姿势/服装/光线/背景/发型）。
|
||
返回结构化描述，用于生成准确 caption。
|
||
"""
|
||
import base64, json, urllib.request
|
||
backend = _norm_backend(backend)
|
||
with open(img_path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode()
|
||
if face_only:
|
||
prompt = (
|
||
"请只描述这张图片中人物的**表情、妆容和发型**，不要描述肤色、斑点、五官、脸型等面部特征，"
|
||
"不要描述服装、背景、姿势。\n"
|
||
"格式：表情（微笑/露齿笑/大笑/中性/严肃/惊讶等）, 妆容（淡妆/浓妆/口红/眼影等）, 发型（发色/长度/造型/发饰）。\n"
|
||
"要求：只要短语，不要解释，不要思考过程。\n"
|
||
"例：大笑, 淡妆口红, 黑色齐刘海短发"
|
||
)
|
||
else:
|
||
prompt = (
|
||
"请用简洁的中文短语描述这张图片中的人物，格式：姿势, 表情, 服装, 光线, 背景, 妆容, 发型。\n"
|
||
"绝对不要描述肤色、斑点、五官、脸型等面部特征。\n"
|
||
"要求：只要短语，不要解释，不要思考过程，不要分点。\n"
|
||
"例：站立, 微笑, 红色礼服, 自然光, 户外花园, 淡妆, 长发马尾"
|
||
)
|
||
try:
|
||
if backend == "omlx":
|
||
payload = {
|
||
"model": _resolve_model(backend)[1],
|
||
"messages": [{"role": "user", "content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
|
||
]}],
|
||
"max_tokens": 2048,
|
||
}
|
||
req = urllib.request.Request(OMLX_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
text = (result["choices"][0]["message"].get("content") or "").strip()
|
||
else:
|
||
payload = {
|
||
"model": OLLAMA_MODEL,
|
||
"messages": [{"role": "user", "content": prompt, "images": [b64]}],
|
||
"stream": False,
|
||
}
|
||
req = urllib.request.Request(OLLAMA_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
text = result["message"]["content"].strip()
|
||
_skip_prefix = ("用户", "我", "这", "图片", "根据", "请", "好", "以下", "要", "所以")
|
||
for line in text.splitlines():
|
||
line = line.strip().strip("*#")
|
||
# Thinking 模型 content 可能带"要简洁短语，所以整合：xxx"前言 → 取冒号后实质内容
|
||
if ("：" in line or ":" in line) and line.startswith(_skip_prefix):
|
||
line = line.replace("：", ":").split(":")[-1].strip()
|
||
if "," in line and len(line) < 120 and not line.startswith(_skip_prefix):
|
||
return line
|
||
# fallback：跳过思考/寒暄行，取第一条实质内容（修 30B content 混入 thinking 开头的问题）
|
||
for line in text.splitlines():
|
||
line = line.strip().strip("*#")
|
||
if line and not line.startswith(_skip_prefix):
|
||
return line[:120]
|
||
return ""
|
||
except Exception as e:
|
||
return f"（VLM[{backend}] 调用失败: {e}）"
|
||
# 默认筛选配额（总 20）
|
||
DEFAULT_QUOTA = {"特写": 10, "半身": 6, "全身": 4}
|
||
# 配额装满后补录线：同类未选中里质量 ≥ 此值的图直接晋级（高分不因配额被挤掉）
|
||
TOPUP_QUALITY = 90
|
||
|
||
# 同场景去重：pHash 距离 ≤ 阈值 或 HSV 直方图相关性 ≥ 阈值 视为同批次，每簇最多保留 MAX_PER_CLUSTER 张
|
||
# （pHash 对"同造型微调姿势"失效——实测双胞胎距离 69/256；直方图 0.971 精准锁定）
|
||
DUP_PHASH_THRESHOLD = 12
|
||
DUP_HIST_CORREL = 0.85
|
||
MAX_PER_CLUSTER = 1
|
||
|
||
|
||
def _hsv_hist(img):
|
||
"""HSV 颜色直方图（16×8×8），捕捉同背景同服装同造型（pHash 抓不住的批次特征）"""
|
||
rgb = np.array(img.convert("RGB").resize((256, 256)))
|
||
hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
|
||
h = cv2.calcHist([hsv], [0, 1, 2], None, [16, 8, 8], [0, 180, 0, 256, 0, 256])
|
||
cv2.normalize(h, h)
|
||
return h
|
||
|
||
|
||
def _same_cluster(a, b):
|
||
"""两张图是否同批次：pHash 近（同构图）或直方图近似（同背景同造型）"""
|
||
if hamming(np.array(a["phash"]), np.array(b["phash"])) <= DUP_PHASH_THRESHOLD:
|
||
return True
|
||
if a.get("hist") is not None and b.get("hist") is not None:
|
||
return cv2.compareHist(a["hist"], b["hist"], cv2.HISTCMP_CORREL) >= DUP_HIST_CORREL
|
||
return False
|
||
|
||
ANGLE_VLM_PROMPT = (
|
||
"判断这张图片中人物脸部的朝向和拍摄角度。只回答以下之一：\n"
|
||
"正面（脸部正对镜头，双眼基本对称，平视）\n"
|
||
"前侧（脸部转向一侧约30-60度，双眼可见但不对称）\n"
|
||
"侧面（脸部转向一侧约70-90度，只能看到一只眼睛）\n"
|
||
"俯拍（相机从上方俯视拍摄，能看到头顶/发旋，额头偏大，下巴偏小）\n"
|
||
"仰拍（相机从下方仰视拍摄，能看到下巴下侧/鼻孔，下巴偏大，额头偏小）\n"
|
||
"注意：判断的是拍摄机位高低，不是人物朝向。只回答'正面'、'前侧'、'侧面'、'俯拍'、'仰拍'一个词。"
|
||
)
|
||
_ANGLE_ZH2EN = {"正脸": "front view", "正面": "front view", "半侧面": "three-quarter view",
|
||
"前侧": "three-quarter view", "侧面": "side view",
|
||
"俯拍": "high angle view", "仰拍": "low angle view"}
|
||
|
||
|
||
def classify_angle_vlm(img_path, backend="ollama", timeout=120):
|
||
"""
|
||
VLM 语义判断人脸朝向（比 5-landmark 几何法可靠得多：大侧脸/微侧都能判准）。
|
||
返回 "front view" / "three-quarter view" / "side view"；失败返回 None（调用方退回几何法）。
|
||
"""
|
||
backend = _norm_backend(backend)
|
||
import base64, json, urllib.request
|
||
with open(img_path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode()
|
||
try:
|
||
if backend == "omlx":
|
||
payload = {
|
||
"model": _resolve_model(backend)[1],
|
||
"messages": [{"role": "user", "content": [
|
||
{"type": "text", "text": ANGLE_VLM_PROMPT},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
|
||
]}],
|
||
"max_tokens": 1024,
|
||
}
|
||
req = urllib.request.Request(OMLX_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
text = (result["choices"][0]["message"].get("content") or "").strip()
|
||
else:
|
||
payload = {
|
||
"model": OLLAMA_MODEL,
|
||
"messages": [{"role": "user", "content": ANGLE_VLM_PROMPT, "images": [b64]}],
|
||
"stream": False,
|
||
}
|
||
req = urllib.request.Request(OLLAMA_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
text = result["message"]["content"].strip()
|
||
for zh, en in _ANGLE_ZH2EN.items():
|
||
if zh in text:
|
||
return en
|
||
return None
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def classify_expression_vlm(img_path, backend="ollama", timeout=120):
|
||
"""
|
||
VLM 判表情（大笑/露齿笑/微笑/中性/严肃/惊讶/其他）——多表情覆盖是脸部 LoRA 泛化的关键。
|
||
返回表情中文词；失败返回 None。
|
||
"""
|
||
import base64, json, urllib.request
|
||
with open(img_path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode()
|
||
_prompt = (
|
||
"判断这张图片中人物的表情。只回答以下之一：微笑/露齿笑/大笑/中性/严肃/惊讶/其他。\n"
|
||
"微笑=嘴角微上扬不露齿；露齿笑=笑容露出牙齿；大笑=张嘴大笑/笑得很开。\n"
|
||
"只回答一个词。"
|
||
)
|
||
try:
|
||
if backend == "omlx":
|
||
payload = {"model": _resolve_model(backend)[1], "messages": [{"role": "user", "content": [
|
||
{"type": "text", "text": _prompt},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}]}],
|
||
"max_tokens": 512}
|
||
req = urllib.request.Request(OMLX_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
text = (result["choices"][0]["message"].get("content") or "").strip()
|
||
else:
|
||
payload = {"model": OLLAMA_MODEL, "messages": [{"role": "user", "content": _prompt, "images": [b64]}], "stream": False}
|
||
req = urllib.request.Request(OLLAMA_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
result = json.loads(resp.read())
|
||
text = result["message"]["content"].strip()
|
||
for kw in ("微笑", "露齿笑", "大笑", "中性", "严肃", "惊讶"):
|
||
if kw in text:
|
||
return kw
|
||
return "其他"
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
# ===== 集成投票分类（3 模型自由描述 + 多数投票，比单模型强制单选可靠得多）=====
|
||
FREE_DESCRIBE_PROMPT = (
|
||
"请观察这张人物照片，用自然语言描述：\n"
|
||
"1) 表情状态（嘴角/嘴型/眉毛的具体状态：是微笑上扬、露齿、大笑张嘴、还是平静/严肃/惊讶？）\n"
|
||
"2) 拍摄角度（相机位置：平视？从上方俯视能看到头顶？从下方仰视能看到下巴下侧？"
|
||
"人物脸正对镜头还是转向一侧、约多少度？）\n"
|
||
"只描述你实际看到的视觉特征，不要给分类标签。"
|
||
)
|
||
|
||
def _free_describe(img_path, backend, timeout=120):
|
||
"""各后端自由描述图片，返回自然语言描述（失败返回 None）"""
|
||
import base64, json, urllib.request
|
||
with open(img_path, "rb") as f:
|
||
b64 = base64.b64encode(f.read()).decode()
|
||
try:
|
||
if backend == "sensenova":
|
||
payload = {"model": SENSENOVA_MODEL, "messages": [{"role": "user", "content": [
|
||
{"type": "text", "text": FREE_DESCRIBE_PROMPT},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}]}],
|
||
"max_tokens": 512}
|
||
req = urllib.request.Request(SENSENOVA_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json", "Authorization": f"Bearer {SENSENOVA_KEY}"},
|
||
method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
return json.loads(resp.read())["choices"][0]["message"]["content"].strip()
|
||
elif backend in ("omlx", "omlx-32b"):
|
||
payload = {"model": _resolve_model(backend)[1], "messages": [{"role": "user", "content": [
|
||
{"type": "text", "text": FREE_DESCRIBE_PROMPT},
|
||
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}}]}],
|
||
"max_tokens": 512}
|
||
req = urllib.request.Request(OMLX_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
msg = json.loads(resp.read())["choices"][0]["message"]
|
||
return (msg.get("content") or msg.get("reasoning_content") or "").strip()
|
||
else: # ollama
|
||
payload = {"model": OLLAMA_MODEL, "messages": [{"role": "user", "content": FREE_DESCRIBE_PROMPT, "images": [b64]}], "stream": False}
|
||
req = urllib.request.Request(OLLAMA_URL, data=json.dumps(payload).encode(),
|
||
headers={"Content-Type": "application/json"}, method="POST")
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
return json.loads(resp.read())["message"]["content"].strip()
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _extract_angle(text):
|
||
"""从自由描述中提取角度。极度保守：只在描述明确无歧义时才判定，否则返回 None（不瞎标）。"""
|
||
if not text:
|
||
return None
|
||
# 俯拍：必须明确"头顶/发旋/上方俯视"
|
||
if any(k in text for k in ("能看到头顶", "看到头顶", "头顶上方", "从上方俯视", "明显俯视", "俯拍", "从上往下")):
|
||
return "high angle view"
|
||
# 仰拍：必须明确"下巴下侧/从下往上/仰视"
|
||
if any(k in text for k in ("能看到下巴下侧", "看到下巴下侧", "从下方仰视", "明显仰视", "仰拍", "从下往上", "鼻孔下方视角")):
|
||
return "low angle view"
|
||
# 侧面：必须明确"只能看到一只眼睛"或"约90度"或"正侧面"
|
||
if any(k in text for k in ("只能看到一只眼睛", "仅能看到一只眼睛", "正侧面", "完全侧", "侧脸90")):
|
||
return "side view"
|
||
# 半侧：明确"约30-60度"/"半侧面"/"明显偏转但不是90"
|
||
if any(k in text for k in ("约30", "约45", "30-60", "半侧面", "偏转约")):
|
||
return "three-quarter view"
|
||
# 正面：明确"正对镜头/直视/完全正面"
|
||
if any(k in text for k in ("正对镜头", "直视镜头", "完全正面", "正面朝向", "正面平视")):
|
||
return "front view"
|
||
# 其他情况不判定（避免"转向一侧""鼻孔"等模糊词误触发）
|
||
return None
|
||
|
||
|
||
def _extract_expr(text):
|
||
"""从自由描述中提取表情。保守：明确特征才判定，模糊不标。"""
|
||
if not text:
|
||
return None
|
||
if any(k in text for k in ("大笑", "张嘴大笑", "笑得很开", "开怀大笑", "哈哈大笑")):
|
||
return "大笑"
|
||
if any(k in text for k in ("露齿", "露出牙齿", "咧嘴笑", "露牙")):
|
||
return "露齿笑"
|
||
if any(k in text for k in ("惊讶", "瞪大眼睛", "眼睛睁大", "眉毛挑高", "挑眉")):
|
||
return "惊讶"
|
||
if any(k in text for k in ("严肃", "肃穆", "面无表情", "没有笑容")):
|
||
return "严肃"
|
||
if any(k in text for k in ("中性", "平静", "淡然", "自然表情")):
|
||
return "中性"
|
||
if any(k in text for k in ("微笑", "嘴角上扬", "浅浅的笑", "温和的笑", "微笑")):
|
||
return "微笑"
|
||
return None
|
||
|
||
|
||
def classify_ensemble_vlm(img_path, backends=("omlx", "ollama", "sensenova", "omlx-32b"), timeout=120):
|
||
"""
|
||
三投票：多模型自由描述 + 规则提取 + 多数投票判 角度/表情。
|
||
多后端并行调用（本地 ollama / 小果 omlx / 云端 sensenova 互相独立；每后端仅 1 路）。
|
||
同源去重：omlx 与 omlx-32b 是同一台 oMLX 服务器——并发 2 路会撑爆小果内存，只保留优先级更高的 omlx-32b。
|
||
返回 (angle_en, expr_zh, confidence)：
|
||
confidence = "一致"(全同) / "多数"(2/3同) / "不确定"(全分歧，需人工)
|
||
失败/无法返回 (None, None, "失败")。
|
||
"""
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
# 同源去重：同一服务器只保留一路（omlx-32b 比 omlx 强，优先）
|
||
backends = tuple(b for b in backends if not (b == "omlx" and "omlx-32b" in backends))
|
||
if not backends:
|
||
return None, None, "失败"
|
||
angles, exprs = [], []
|
||
# 整体硬性超时：单张图全后端判定最多等 timeout 秒，超时即收已完成的（其他后端后台继续，不阻塞）；
|
||
# 防止某个后端（如小果 oMLX）挂起时无限拖住整条流水线
|
||
ex = ThreadPoolExecutor(max_workers=len(backends))
|
||
try:
|
||
futs = {ex.submit(_free_describe, img_path, b, timeout=timeout): b for b in backends}
|
||
try:
|
||
for f in as_completed(futs, timeout=timeout):
|
||
try:
|
||
desc = f.result()
|
||
except Exception:
|
||
desc = None
|
||
if desc:
|
||
a = _extract_angle(desc)
|
||
e = _extract_expr(desc)
|
||
if a:
|
||
angles.append(a)
|
||
if e:
|
||
exprs.append(e)
|
||
except TimeoutError:
|
||
pass # 整体超时：收下已完成的后端结果，其余后台继续
|
||
finally:
|
||
ex.shutdown(wait=False) # 不等待挂起线程，立即返回，避免 with 退出时二次阻塞
|
||
if not angles and not exprs:
|
||
return None, None, "失败"
|
||
|
||
def vote(vals):
|
||
from collections import Counter
|
||
if not vals:
|
||
return None
|
||
c = Counter(vals)
|
||
top, n = c.most_common(1)[0]
|
||
if n >= 2:
|
||
return top, "多数" if n < len(vals) else "一致"
|
||
return top, "不确定" # 全分歧，取第一个（标记人工复核）
|
||
|
||
ang, ang_conf = vote(angles)
|
||
expr, expr_conf = vote(exprs)
|
||
conf = "失败"
|
||
if ang_conf == "一致" or expr_conf == "一致":
|
||
conf = "一致"
|
||
elif ang_conf == "多数" or expr_conf == "多数":
|
||
conf = "多数"
|
||
elif ang_conf == "不确定" or expr_conf == "不确定":
|
||
conf = "不确定"
|
||
return ang, expr, conf
|
||
|
||
|
||
def vlm_caption(img_path, kind, backend="ollama", trigger="lm_face_v1", angle_en=None):
|
||
"""
|
||
统一 caption 生成（auto 流水线 / 候选换图 / 打标重新生成 三处共用的唯一实现）：
|
||
VLM 判角度（可用 angle_en 参数跳过重复判定）+ VLM 描述（特写只描述面部+发型）+ 模板拼装。
|
||
kind: 特写/半身/全身/其他（其他无模板短语）。
|
||
返回 (caption, angle_en)；角度判定失败返回 (None, None)。
|
||
"""
|
||
if not angle_en:
|
||
angle_en = classify_angle_vlm(str(img_path), backend=backend)
|
||
if not angle_en:
|
||
return None, None
|
||
desc = describe_image_omlx(str(img_path), face_only=(kind == "特写"), backend=backend)
|
||
tpl = AUTO_TEMPLATES.get(kind, "")
|
||
if tpl:
|
||
caption = f"{trigger}, {angle_en} {tpl}, {desc}" if desc else f"{trigger}, {angle_en} {tpl}"
|
||
else:
|
||
caption = f"{trigger}, {angle_en}, {desc}" if desc else f"{trigger}, {angle_en}"
|
||
return caption, angle_en
|
||
|
||
|
||
def face_quality_score(img, face_box, lap_face=None):
|
||
"""
|
||
综合质量分 0-100：人脸像素 + 面部清晰度 + 分辨率。
|
||
权重设计：清晰度(35) 略低于像素(40)，但"脸小极清晰"仍能胜过"脸大偏糊"。
|
||
"""
|
||
score = 0
|
||
# 人脸绝对像素（决定 LoRA 训练有效性）
|
||
if face_box:
|
||
fw, fh = face_box[2], face_box[3]
|
||
face_px = min(fw, fh)
|
||
if face_px >= 300:
|
||
score += 40
|
||
elif face_px >= 200:
|
||
score += 32
|
||
elif face_px >= 150:
|
||
score += 22
|
||
elif face_px >= 100:
|
||
score += 12
|
||
else:
|
||
score += 4
|
||
# 面部清晰度（人脸区域 Laplacian；真实照片细腻纹理值域 20-100+，雀斑清晰图约 25-60）
|
||
if lap_face is not None:
|
||
if lap_face >= 100:
|
||
score += 35
|
||
elif lap_face >= 60:
|
||
score += 30
|
||
elif lap_face >= 35:
|
||
score += 24
|
||
elif lap_face >= 20:
|
||
score += 15
|
||
elif lap_face >= 10:
|
||
score += 6
|
||
# 分辨率
|
||
short = min(img.size)
|
||
if short >= 1500:
|
||
score += 25
|
||
elif short >= 1200:
|
||
score += 22
|
||
elif short >= 1024:
|
||
score += 18
|
||
elif short >= 800:
|
||
score += 10
|
||
else:
|
||
score += 3
|
||
return min(100, score)
|
||
|
||
|
||
def _lap_of(img, face_box=None):
|
||
"""计算（人脸区域/全图）Laplacian 清晰度值"""
|
||
gray = np.array(img.convert("L"))
|
||
if face_box is not None:
|
||
x, y, w, h = [int(v) for v in face_box]
|
||
hh, ww = gray.shape
|
||
x0, y0 = max(0, x - w), max(0, y - h)
|
||
x1, y1 = min(ww, x + 2 * w), min(hh, y + 2 * h)
|
||
if (x1 - x0) > 40 and (y1 - y0) > 40:
|
||
gray = gray[y0:y1, x0:x1]
|
||
scale = 1024.0 / max(gray.shape)
|
||
if scale < 1.0:
|
||
gray = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
|
||
return float(cv2.Laplacian(gray, cv2.CV_64F).var())
|
||
|
||
|
||
# 角度保底配额：非正面素材稀缺且对泛化至关重要，优先必录（候选有才占名额）
|
||
ANGLE_MIN_QUOTA = {"side view": 2, "three-quarter view": 3, "high angle view": 1, "low angle view": 1}
|
||
# 表情保底配额：多表情覆盖是脸部 LoRA 泛化的关键（露齿笑/大笑/惊讶等稀缺表情优先必录）
|
||
EXPR_MIN_QUOTA = {"大笑": 2, "露齿笑": 2, "惊讶": 1}
|
||
|
||
|
||
def smart_pick(kind_items, target, quality_key="quality"):
|
||
"""
|
||
每类内部筛选：角度保底 + 表情保底 + 同簇去重硬上限 + 质量 + 贪心多样性。
|
||
角度优先用 item["angle_en"]（VLM 判定），表情用 item["expr"]（VLM 判定）。
|
||
kind_items: [result_dict,...]；返回选中的列表，且给每项加 pick_rank。
|
||
"""
|
||
def _ang(r):
|
||
return r.get("angle_en") or face_angle(r["face"])[1]
|
||
|
||
def _expr(r):
|
||
return r.get("expr") or "其他"
|
||
|
||
items = sorted(kind_items, key=lambda r: -r[quality_key])
|
||
if not items:
|
||
return []
|
||
selected, rest = [], list(items)
|
||
# 1. 角度保底：side view / three-quarter view 各保前 N 张最高质量（候选不足则跳过）
|
||
for ang, min_n in ANGLE_MIN_QUOTA.items():
|
||
cands = [r for r in rest if _ang(r) == ang]
|
||
for best in cands[:min_n]:
|
||
if len(selected) >= target:
|
||
break
|
||
selected.append(best)
|
||
rest.remove(best)
|
||
# 2. 表情保底：大笑/张嘴/严肃 优先必录（稀缺表情对泛化价值高）
|
||
for expr, min_n in EXPR_MIN_QUOTA.items():
|
||
cands = [r for r in rest if _expr(r) == expr]
|
||
for best in cands[:min_n]:
|
||
if len(selected) >= target:
|
||
break
|
||
selected.append(best)
|
||
rest.remove(best)
|
||
# 3. 剩余名额：同簇去重硬上限 + 质量 + pHash 多样性贪心
|
||
while len(selected) < min(target, len(items)) and rest:
|
||
best_i, best_score = -1, -1e9
|
||
for i, it in enumerate(rest):
|
||
# 与已选集合的最小 phash 距离
|
||
d = min(hamming(np.array(it["phash"]), np.array(s["phash"])) for s in selected) if selected else 0
|
||
# 同簇硬上限：pHash 近或直方图近似（同批次）且簇已满 → 禁止入选
|
||
if sum(1 for s in selected if _same_cluster(it, s)) >= MAX_PER_CLUSTER:
|
||
continue
|
||
score = it[quality_key] * 0.6 + min(d, 40) * 0.4
|
||
if score > best_score:
|
||
best_score, best_i = score, i
|
||
if best_i < 0:
|
||
break # 剩余全部撞簇上限
|
||
selected.append(rest.pop(best_i))
|
||
for i, r in enumerate(selected):
|
||
r["pick_rank"] = i + 1
|
||
return selected
|
||
|
||
|
||
def auto_process(src_dir, out_dir, face_det, person_det, trigger="lm_face_v1",
|
||
face_margin=1.8, min_side=MIN_RESOLUTION, quota=None, limit=True,
|
||
use_omlx=False, backend="ollama"):
|
||
"""
|
||
全自动素材流水线（v2）：
|
||
1. 逐张：人脸+人体检测；面部清晰度/整体清晰度分开
|
||
2. 综合质量评分（脸像素+清晰度+分辨率）
|
||
3. 构图决策（占比+脸像素双指标）并裁剪：
|
||
- 占比≥8% 或 脸短边≥256px → 特写候选（裁脸）
|
||
- 占比≥2% 且 脸短边≥150px → 半身
|
||
- 有人体 → 全身候选（裁全身）
|
||
- 裁后短边<min_side → 降级/保留原样（不硬裁）
|
||
4. 超量筛选（limit=True）：按配额每类质量+多样性挑最优
|
||
5. 打标草稿 + 淘汰记录
|
||
输出：out/特写/ 半身/ 全身/ 未选中/ 淘汰/
|
||
"""
|
||
src_dir, out_dir = Path(src_dir), Path(out_dir)
|
||
subdirs = {k: out_dir / k for k in AUTO_TEMPLATES}
|
||
unused_dir = out_dir / "未选中"
|
||
reject_dir = out_dir / "淘汰"
|
||
for d in [*subdirs.values(), unused_dir, reject_dir]:
|
||
d.mkdir(parents=True, exist_ok=True)
|
||
# 清空上一轮的产出素材（特写/半身/全身），防止旧轮次文件残留被误当本轮结果
|
||
# （教训：v3b 残留 full_003/face_011 与本轮 full_002/face_010 字节相同，被误报"重复入选"）
|
||
for d in subdirs.values():
|
||
for old in d.glob("*"):
|
||
if old.is_file():
|
||
old.unlink()
|
||
|
||
images = sorted([p for p in src_dir.iterdir() if p.suffix.lower() in IMG_EXTS and not p.name.startswith(".")])
|
||
if not images:
|
||
print(f"[WARN] {src_dir} 没有图片")
|
||
return {}, []
|
||
|
||
all_items = {k: [] for k in AUTO_TEMPLATES}
|
||
rejected = []
|
||
for p in images:
|
||
try:
|
||
img = load_image(p)
|
||
w, h = img.size
|
||
rgb = np.array(img.convert("RGB"))
|
||
bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
|
||
# 统一用 _detect_faces_fast（缩略图 + NMS + score 过滤，防腿部/物体误检）
|
||
faces = _detect_faces_fast(rgb, face_det)
|
||
persons = person_det.detect(bgr)
|
||
|
||
# 硬淘汰：无人脸。分辨率判定看**人脸区域像素**（脸短边≥250px 就算可用特写，
|
||
# 不看整图短边——高清面部图哪怕原图小，脸本身像素也够训练）。
|
||
if not faces:
|
||
# 漏检兜底：俯拍/仰拍等特殊角度 YuNet 常漏检 → 降阈值重试
|
||
faces = _detect_faces_fast(rgb, face_det, score_threshold=0.3)
|
||
if not faces:
|
||
# 仍无：用 VLM 复核是否真有脸（避免把"漏检的俯拍脸"误淘汰）
|
||
_vlm_sees_face = False
|
||
if use_omlx:
|
||
_ang = classify_angle_vlm(str(p), backend=backend)
|
||
_vlm_sees_face = _ang is not None
|
||
if _vlm_sees_face:
|
||
# VLM 确认有脸但检测器漏检 → 归入未选中待人工（不硬淘汰）
|
||
print(f"[注意] {p.name}: 检测器漏检但 VLM 确认有脸（{_ang}），放入未选中待人工")
|
||
pending_dir = out_dir / "未选中"
|
||
pending_dir.mkdir(exist_ok=True)
|
||
dst = pending_dir / p.name
|
||
dst.write_bytes(p.read_bytes()) # 始终覆盖（防旧版残留）
|
||
continue
|
||
rejected.append((p.name, "未检测到人脸"))
|
||
print(f"[淘汰] {p.name}: 无人脸")
|
||
continue
|
||
|
||
# 多人合影 → 取最大人脸；其余脸视为"多余人物"
|
||
f = max(faces, key=lambda x: x[2] * x[3]) if len(faces) > 1 else faces[0]
|
||
fx, fy, fw, fh = int(f[0]), int(f[1]), int(f[2]), int(f[3])
|
||
ratio = (fw * fh) / (w * h)
|
||
face_px = min(fw, fh)
|
||
|
||
# 人脸像素判定：脸短边 <250px 才淘汰（高清面部图 250px+ 就够训练）
|
||
if face_px < 250:
|
||
rejected.append((p.name, f"人脸像素不足（脸短边 {face_px}px < 250px，训练效果差）"))
|
||
print(f"[淘汰] {p.name}: 人脸像素不足（{face_px}px）")
|
||
continue
|
||
|
||
# 面部清晰度：紧贴脸框核心（避免背景稀释）。
|
||
# 注意：真实照片的雀斑/眉毛等细腻纹理 Laplacian 值较低（20-40 正常），
|
||
# 阈值 15 能区分"清晰"（>20）和"真糊"（<15，对焦失败/强模糊）。
|
||
lap_face = face_sharpness(img, (fx, fy, fw, fh))
|
||
if lap_face < 15:
|
||
rejected.append((p.name, f"面部模糊（清晰度 {lap_face:.0f}，阈值 15）"))
|
||
print(f"[淘汰] {p.name}: 面部模糊")
|
||
continue
|
||
lap_all = _lap_of(img, None)
|
||
|
||
item = {
|
||
"file": p.name, "src": p, "img": img,
|
||
"w": w, "h": h, "ratio": ratio, "face_px": face_px,
|
||
"face": list(f), # 完整 face（含 landmarks，face_angle 需要）
|
||
"lap_face": lap_face, "lap_all": lap_all,
|
||
"faces_n": len(faces), "persons": persons,
|
||
"phash": phash(img).tolist(),
|
||
"hist": _hsv_hist(img),
|
||
}
|
||
item["quality"] = face_quality_score(img, (fx, fy, fw, fh), lap_face)
|
||
|
||
# 构图分类（主要看人脸占比，位置仅辅助全身判定）
|
||
# 特写：脸是画面主体（≥8%）；半身/胸像：上半身可见（1.5-8%）；全身：整个人可见（<1.5%）
|
||
face_cy = fy + fh / 2
|
||
face_top = face_cy < h * 0.5 # 人脸在上半部（全身照特征）
|
||
if ratio >= 0.08:
|
||
kind = "特写"
|
||
elif ratio >= 0.015:
|
||
kind = "半身"
|
||
else:
|
||
# 占比小：人脸在上部 → 全身（人物从上到下）；人脸在中下部 → 半身（大头照/脸部局部）
|
||
kind = "全身" if face_top else "半身"
|
||
item["kind"] = kind
|
||
all_items[kind].append(item)
|
||
|
||
print(f"[{kind}] {p.name} | 脸{face_px}px 占比{ratio*100:.1f}% 清晰{lap_face:.0f}/{lap_all:.0f} 质量{item['quality']}")
|
||
|
||
except Exception as e:
|
||
rejected.append((p.name, f"处理异常: {e}"))
|
||
print(f"[异常] {p.name}: {e}")
|
||
|
||
# VLM 角度分类（语义判断，远比 5-landmark 几何法准：大侧脸/微侧都能判对）+ 表情分类（多表情覆盖）。
|
||
# 串行调用（并发会撑爆小果统一内存）；不开 VLM 时退回几何法并在标签阶段标注。
|
||
if use_omlx:
|
||
all_cands = [r for rs in all_items.values() for r in rs]
|
||
print(f"\n[VLM 角度+表情分类] {len(all_cands)} 张候选，串行判定中（models={backend}）...")
|
||
if isinstance(backend, (list, tuple)) and len(backend) > 1:
|
||
_models = tuple(backend)
|
||
for r in all_cands:
|
||
ang, expr, conf = classify_ensemble_vlm(str(r["src"]), backends=_models)
|
||
r["angle_en"] = ang or face_angle(r["face"])[1]
|
||
r["expr"] = expr or "其他"
|
||
print(f" [分类] {r['file']}: 角度={r['angle_en']} 表情={r['expr']} 置信={conf}")
|
||
else:
|
||
_bk = _norm_backend(backend)
|
||
for r in all_cands:
|
||
ang = classify_angle_vlm(str(r["src"]), backend=_bk)
|
||
r["angle_en"] = ang or face_angle(r["face"])[1]
|
||
r["expr"] = classify_expression_vlm(str(r["src"]), backend=_bk) or "其他"
|
||
print(f" [分类] {r['file']}: 角度={r['angle_en']} 表情={r['expr']}")
|
||
if ang is None:
|
||
print(f" [WARN] {r['file']}: VLM 角度失败，退回几何法（{r['angle_en']}）")
|
||
else:
|
||
print("\n[角度] 未开 VLM（--omlx），用几何法粗判（大侧脸/微侧可能不准，建议开 VLM）")
|
||
|
||
# 超量筛选：每类按配额挑最优；配额装满后，同类未选中里质量 ≥ TOPUP_QUALITY 的直接补录（高分不因配额被挤掉）
|
||
picked_all, unused_all = [], []
|
||
quota = quota or DEFAULT_QUOTA
|
||
for kind, items in all_items.items():
|
||
if not items:
|
||
continue
|
||
target = quota.get(kind, 0) if limit else len(items)
|
||
picked = smart_pick(items, target)
|
||
if limit and len(items) > target:
|
||
# 补录：同类未选中里质量达标者全部晋级（不设上限，高分全收）
|
||
used_ids = {id(r) for r in picked}
|
||
topup = [r for r in items if id(r) not in used_ids and r.get("quality", 0) >= TOPUP_QUALITY]
|
||
topup.sort(key=lambda r: -r.get("quality", 0))
|
||
for r in topup:
|
||
r["pick_rank"] = len(picked) + 1
|
||
picked.append(r)
|
||
if topup:
|
||
print(f"[补录] {kind}: 质量≥{TOPUP_QUALITY} 的未选中图补录 {len(topup)} 张（{', '.join(r['file'] for r in topup[:5])}{'...' if len(topup)>5 else ''}）")
|
||
picked_all.extend(picked)
|
||
used_ids = {id(r) for r in picked}
|
||
unused_all.extend(r for r in items if id(r) not in used_ids)
|
||
|
||
# 写出选中素材：每图可产出 特写（裁脸） + 构图素材（半身/全身）
|
||
counts = {k: 0 for k in AUTO_TEMPLATES}
|
||
for r in sorted(picked_all, key=lambda x: ({"特写": 0, "半身": 1, "全身": 2}[x["kind"]], x["pick_rank"])):
|
||
kind = r["kind"]
|
||
img = r["img"]
|
||
face = r["face"]
|
||
|
||
# 1. 脸部特写：只要人脸像素足够（≥150）就额外裁一张特写（一张图可产多个素材）
|
||
# 标签：角度 + **只描述面部+发型**（特写只剩面部，不含被裁掉的服装/背景/姿势）
|
||
if r["face_px"] >= 150:
|
||
crop_p, note_p = crop_face_portrait(img, face, margin=2.5)
|
||
angle_en = r.get("angle_en") or face_angle(face)[1]
|
||
angle_zh = {"front view": "正脸", "three-quarter view": "半侧面", "side view": "侧面", "high angle view": "俯拍", "low angle view": "仰拍"}.get(angle_en, angle_en)
|
||
counts["特写"] += 1
|
||
fname_p = f"face_{counts['特写']:03d}.jpg"
|
||
crop_p.save(subdirs["特写"] / fname_p, quality=95)
|
||
# 特写 caption：face_only（只描述面部+发型，不提被裁掉的服装/背景/姿势）
|
||
if use_omlx:
|
||
caption_p, _ = vlm_caption(str(r["src"]), "特写", backend=backend, trigger=trigger, angle_en=angle_en)
|
||
caption_p = caption_p or f"{trigger}, {angle_en} {AUTO_TEMPLATES['特写']}"
|
||
else:
|
||
caption_p = f"{trigger}, {angle_en} {AUTO_TEMPLATES['特写']}"
|
||
(subdirs["特写"] / f"face_{counts['特写']:03d}.txt").write_text(caption_p, encoding="utf-8")
|
||
print(f" [√] 特写/{fname_p} | {note_p} | {angle_zh} | {caption_p[:60]}")
|
||
|
||
# 2. 构图素材：半身/全身（标签按角度 + oMLX 自动描述）
|
||
if kind == "特写":
|
||
# 本来就是特写占比：保留原图（脸部已占主体）
|
||
crop, note = img.copy(), f"特写原图 {r['w']}×{r['h']}"
|
||
kind = "特写"
|
||
# 特写不再重复输出（上面已产 face_xxx），跳过构图素材
|
||
if r["face_px"] >= 150:
|
||
continue
|
||
elif kind == "半身":
|
||
crop, note = crop_centered_halfbody(img, face)
|
||
else:
|
||
crop, note = crop_fullbody(img, face)
|
||
|
||
counts[kind] += 1
|
||
stem = {"特写": "face", "半身": "half", "全身": "full"}[kind]
|
||
fname = f"{stem}_{counts[kind]:03d}.jpg"
|
||
crop.save(subdirs[kind] / fname, quality=95)
|
||
# 构图素材 caption：角度 + VLM 自动描述（匹配图）
|
||
angle_en = r.get("angle_en") or face_angle(face)[1]
|
||
angle_zh = {"front view": "正脸", "three-quarter view": "半侧面", "side view": "侧面", "high angle view": "俯拍", "low angle view": "仰拍"}.get(angle_en, angle_en)
|
||
if use_omlx:
|
||
caption, _ = vlm_caption(str(r["src"]), kind, backend=backend, trigger=trigger, angle_en=angle_en)
|
||
caption = caption or f"{trigger}, {angle_en} {AUTO_TEMPLATES[kind]}"
|
||
else:
|
||
caption = f"{trigger}, {angle_en} {AUTO_TEMPLATES[kind]}"
|
||
(subdirs[kind] / f"{stem}_{counts[kind]:03d}.txt").write_text(caption, encoding="utf-8")
|
||
print(f" [√] {kind}/{fname} | {note} | {angle_zh} | {caption[:60]}")
|
||
|
||
# 未选中 + 淘汰（始终覆盖——修 bug：之前 exists 跳过导致旧版残留缓存，用户修正的图不被更新）
|
||
for r in unused_all:
|
||
dst = unused_dir / r["file"]
|
||
dst.write_bytes(r["src"].read_bytes())
|
||
if rejected:
|
||
(reject_dir / "淘汰原因.txt").write_text("\n".join(f"{n}: {r}" for n, r in rejected), encoding="utf-8")
|
||
for n, r in rejected:
|
||
src = src_dir / n
|
||
if src.exists():
|
||
dst = reject_dir / n
|
||
dst.write_bytes(src.read_bytes())
|
||
|
||
print(f"\n===== 流水线完成 =====")
|
||
for k in AUTO_TEMPLATES:
|
||
print(f" {k}: {counts[k]} 张 -> {subdirs[k]}")
|
||
print(f" 未选中: {len(unused_all)} 张（见 {unused_dir}）")
|
||
print(f" 淘汰: {len(rejected)} 张（原因见 {reject_dir}/淘汰原因.txt）")
|
||
return counts, rejected
|
||
|
||
|
||
# ---------- HTML 报告 ----------
|
||
def build_html(results, out_path):
|
||
cards = []
|
||
for r in results:
|
||
color = {STATUS["PASS"]: "#16a34a", STATUS["WARN"]: "#d97706", STATUS["FAIL"]: "#dc2626"}[r["status"]]
|
||
reasons_html = "<br>".join(f"· {x}" for x in r["reasons"])
|
||
size = f"{r['size'][0]}×{r['size'][1]}" if r["size"][0] else "?"
|
||
cards.append(f"""
|
||
<div class="card" style="border-color:{color}">
|
||
<div class="thumb-wrap"><img class="thumb" src="file:///{r['path']}" loading="lazy" onerror="this.style.display='none'"></div>
|
||
<div class="info">
|
||
<div class="row"><span class="badge" style="background:{color}">{r['status']}</span><span class="fname">{r['file']}</span></div>
|
||
<div class="meta">{size} · {r['bytes']/1024:.0f}KB</div>
|
||
<div class="reasons" style="color:{color}">{reasons_html}</div>
|
||
</div>
|
||
</div>""")
|
||
|
||
counts = {v: sum(1 for r in results if r["status"] == v) for v in STATUS.values()}
|
||
html = f"""<!DOCTYPE html>
|
||
<html lang="zh"><head><meta charset="utf-8">
|
||
<title>LoRA 素材审核报告</title>
|
||
<style>
|
||
body {{ font-family: "Microsoft YaHei", sans-serif; background:#f5f5f5; margin:20px; }}
|
||
h1 {{ font-size:22px; }}
|
||
.summary {{ background:#fff; padding:12px 18px; border-radius:8px; margin-bottom:16px; box-shadow:0 1px 3px #ccc; }}
|
||
.summary b {{ margin-right:20px; }}
|
||
.grid {{ display:grid; grid-template-columns:repeat(auto-fill,minmax(280px,1fr)); gap:14px; }}
|
||
.card {{ background:#fff; border:3px solid; border-radius:10px; overflow:hidden; box-shadow:0 1px 3px #ccc; }}
|
||
.thumb-wrap {{ height:220px; overflow:hidden; background:#eee; }}
|
||
.thumb {{ width:100%; height:100%; object-fit:cover; }}
|
||
.info {{ padding:10px 12px; }}
|
||
.badge {{ color:#fff; padding:2px 10px; border-radius:4px; font-weight:bold; font-size:13px; }}
|
||
.fname {{ font-weight:bold; margin-left:8px; word-break:break-all; }}
|
||
.meta {{ color:#888; font-size:12px; margin:4px 0; }}
|
||
.reasons {{ font-size:13px; line-height:1.6; }}
|
||
</style></head><body>
|
||
<h1>📸 LoRA 训练素材审核报告</h1>
|
||
<div class="summary">
|
||
<b>✅ 合格: {counts['合格']}</b>
|
||
<b>🟡 警告: {counts['警告']}</b>
|
||
<b>❌ 不合格: {counts['不合格']}</b>
|
||
<b>共 {len(results)} 张</b>
|
||
</div>
|
||
<div class="grid">{''.join(cards)}</div>
|
||
</body></html>"""
|
||
Path(out_path).write_text(html, encoding="utf-8")
|
||
return out_path
|
||
|
||
|
||
# ---------- 整理训练集 ----------
|
||
def prepare_dataset(img_dir, results, out_dir, min_status="WARN"):
|
||
"""把达到最低状态（默认 WARN，即排除 FAIL）的图片复制到训练集并重命名"""
|
||
out_dir = Path(out_dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
ok = [r for r in results if r["status"] in ("PASS", "WARN") and r["status"] != "FAIL"]
|
||
# 按状态排序：PASS 优先
|
||
ok.sort(key=lambda r: 0 if r["status"] == "PASS" else 1)
|
||
|
||
copied = 0
|
||
for idx, r in enumerate(ok, 1):
|
||
src = Path(r["path"])
|
||
dst = out_dir / f"img_{idx:03d}.jpg"
|
||
try:
|
||
img = load_image(src).convert("RGB")
|
||
img.save(dst, quality=95)
|
||
copied += 1
|
||
print(f" [OK] {r['file']} -> {dst.name} ({r['status']})")
|
||
except Exception as e:
|
||
print(f" [ERR] {r['file']} 复制失败: {e}")
|
||
print(f"\n共整理 {copied} 张到 {out_dir}")
|
||
return copied
|
||
|
||
|
||
# ---------- CLI ----------
|
||
def main():
|
||
# Windows 控制台 GBK 编码兜底，避免 emoji/特殊字符打印报错
|
||
for stream in (sys.stdout, sys.stderr):
|
||
try:
|
||
stream.reconfigure(encoding="utf-8", errors="replace")
|
||
except Exception:
|
||
pass
|
||
|
||
parser = argparse.ArgumentParser(description="LoRA 训练素材审核与整理工具")
|
||
sub = parser.add_subparsers(dest="cmd", required=True)
|
||
|
||
p_check = sub.add_parser("check", help="审核照片，生成报告")
|
||
p_check.add_argument("dir", help="照片目录")
|
||
p_check.add_argument("--out", default=None, help="报告输出目录（默认同照片目录）")
|
||
|
||
p_prep = sub.add_parser("prepare", help="整理合格照片到训练集")
|
||
p_prep.add_argument("dir", help="照片目录")
|
||
p_prep.add_argument("--out", required=True, help="训练集输出目录")
|
||
|
||
p_batch = sub.add_parser("batch", help="审核 + 整理一条龙")
|
||
p_batch.add_argument("dir", help="照片目录")
|
||
p_batch.add_argument("--out", required=True, help="训练集输出目录")
|
||
|
||
p_pick = sub.add_parser("pick", help="智能选图：自动挑出多样化的 N 张")
|
||
p_pick.add_argument("dir", help="照片目录")
|
||
p_pick.add_argument("--count", type=int, default=20, help="要选多少张（默认 20）")
|
||
p_pick.add_argument("--out", default=None, help="报告输出目录（默认同照片目录）")
|
||
|
||
p_crop = sub.add_parser("crop", help="合影处理：按最大人脸裁切单人（处理多人合影）")
|
||
p_crop.add_argument("dir", help="照片目录")
|
||
p_crop.add_argument("--out", required=True, help="裁切输出目录")
|
||
p_crop.add_argument("--scale", type=float, default=1.8, help="裁切放大倍数（默认 1.8）")
|
||
|
||
p_auto = sub.add_parser("auto", help="全自动流水线：构图分类+自动裁剪+打标草稿+超量筛选")
|
||
p_auto.add_argument("dir", help="照片目录")
|
||
p_auto.add_argument("--out", required=True, help="输出目录（生成 特写/半身/全身/未选中/淘汰）")
|
||
p_auto.add_argument("--trigger", default="lm_face_v1", help="触发词（默认 lm_face_v1）")
|
||
p_auto.add_argument("--limit", type=int, default=20,
|
||
help="目标筛选数量（默认 20：特写10/半身6/全身4；0=不筛选全部保留）")
|
||
p_auto.add_argument("--omlx", action="store_true",
|
||
help="用 VLM 自动描述图像（姿势/服装/光线/背景/发型，匹配图）")
|
||
p_auto.add_argument("--backend", choices=["ollama", "omlx"], default="ollama",
|
||
help="VLM 后端：ollama=本地8B（默认，稳定）；omlx=小果30B-Thinking（细节更准，串行）")
|
||
|
||
args = parser.parse_args()
|
||
|
||
detector = FaceDetector()
|
||
|
||
if args.cmd == "check":
|
||
results = analyze_images(args.dir, detector)
|
||
out_dir = Path(args.out) if args.out else Path(args.dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
report = out_dir / "素材审核报告.html"
|
||
build_html(results, report)
|
||
(out_dir / "素材审核结果.json").write_text(json.dumps(results, ensure_ascii=False, indent=1), encoding="utf-8")
|
||
print(f"\n报告已生成: {report}")
|
||
|
||
elif args.cmd == "prepare":
|
||
results = analyze_images(args.dir, detector, verbose=False)
|
||
prepare_dataset(args.dir, results, args.out)
|
||
|
||
elif args.cmd == "batch":
|
||
results = analyze_images(args.dir, detector)
|
||
out_dir = Path(args.dir)
|
||
report = out_dir / "素材审核报告.html"
|
||
build_html(results, report)
|
||
print(f"报告已生成: {report}")
|
||
prepare_dataset(args.dir, results, args.out)
|
||
|
||
elif args.cmd == "pick":
|
||
results = analyze_images(args.dir, detector)
|
||
picked = pick_diverse(results, count=args.count)
|
||
out_dir = Path(args.out) if args.out else Path(args.dir)
|
||
out_dir.mkdir(parents=True, exist_ok=True)
|
||
report = out_dir / f"智能选图推荐_{len(picked)}张.html"
|
||
build_pick_html(picked, report, args.count)
|
||
(out_dir / "智能选图结果.json").write_text(
|
||
json.dumps(picked, ensure_ascii=False, indent=1), encoding="utf-8"
|
||
)
|
||
print(f"\n推荐选图报告: {report}")
|
||
print(f"共推荐 {len(picked)} 张（目标 {args.count}）")
|
||
for r in picked:
|
||
print(f" [{r.get('compose','')}] {r['file']} ({r['status']})")
|
||
|
||
elif args.cmd == "crop":
|
||
crop_dir(args.dir, args.out, detector, margin_scale=args.scale)
|
||
|
||
elif args.cmd == "auto":
|
||
person_det = PersonDetector()
|
||
if args.limit > 0:
|
||
n1 = max(1, round(args.limit * 0.5))
|
||
n2 = max(1, round(args.limit * 0.3))
|
||
quota = {"特写": n1, "半身": n2, "全身": max(1, args.limit - n1 - n2)}
|
||
else:
|
||
quota = None
|
||
auto_process(args.dir, args.out, detector, person_det, trigger=args.trigger,
|
||
quota=quota, limit=args.limit > 0, use_omlx=args.omlx, backend=args.backend)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|