713 lines
28 KiB
Python
713 lines
28 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
皮肤分析工具 - 后端 API 版
|
||
输出 JSON 数据到文件夹，完成后回调后端接口
|
||
"""
|
||
|
||
import cv2
|
||
import numpy as np
|
||
from ultralytics import YOLO
|
||
from pathlib import Path
|
||
|
||
import os
|
||
import argparse
|
||
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from tqdm import tqdm
|
||
import logging
|
||
import threading
|
||
from typing import Dict, List, Optional, Tuple
|
||
import requests
|
||
|
||
from config.settings import API_HOST, API_CALLBACK
|
||
|
||
# YOLO 模型非线程安全（fuse/BN 会竞态报 AttributeError: bn），推理需串行
|
||
_yolo_infer_lock = threading.Lock()
|
||
|
||
# ==================== 日志配置 ====================
|
||
def setup_logging(verbose: bool = False):
|
||
level = logging.DEBUG if verbose else logging.INFO
|
||
logging.basicConfig(
|
||
level=level,
|
||
format='%(asctime)s | %(levelname)s | %(message)s',
|
||
datefmt='%H:%M:%S'
|
||
)
|
||
return logging.getLogger(__name__)
|
||
|
||
logger = setup_logging()
|
||
|
||
# ==================== 支持中文路径的读写函数 ====================
|
||
def cv2_imread_chinese(file_path: str) -> Optional[np.ndarray]:
|
||
if not os.path.exists(file_path):
|
||
logger.warning(f"文件不存在：{file_path}")
|
||
return None
|
||
try:
|
||
img = cv2.imdecode(np.fromfile(file_path, dtype=np.uint8), cv2.IMREAD_COLOR)
|
||
return img
|
||
except Exception as e:
|
||
logger.error(f"读取失败 {file_path}: {e}")
|
||
return None
|
||
|
||
def cv2_imwrite_chinese(file_path: str, img: np.ndarray) -> bool:
|
||
try:
|
||
ext = os.path.splitext(file_path)[1]
|
||
if ext.lower() not in ['.jpg', '.jpeg', '.png', '.bmp']:
|
||
ext = '.jpg'
|
||
cv2.imencode(ext, img)[1].tofile(file_path)
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"写入失败 {file_path}: {e}")
|
||
return False
|
||
|
||
# ==================== YOLO 痤疮检测 ====================
|
||
def detect_acne(img: np.ndarray, image_path: str, model: YOLO, output_path: str) -> Tuple[bool, Dict, str]:
|
||
try:
|
||
# 共享模型下串行推理，避免多线程 fuse/BN 竞态
|
||
with _yolo_infer_lock:
|
||
results = model(img, conf=0.2, verbose=False)
|
||
|
||
annotated_img = img.copy()
|
||
class_counts = {}
|
||
|
||
for result in results:
|
||
if result.boxes is not None:
|
||
for box in result.boxes:
|
||
cls_id = int(box.cls[0])
|
||
class_name = model.names[cls_id]
|
||
x1, y1, x2, y2 = map(int, box.xyxy[0])
|
||
conf = float(box.conf[0])
|
||
|
||
color = (0, 0, 255)
|
||
cv2.rectangle(annotated_img, (x1, y1), (x2, y2), color, 2)
|
||
label = f"{class_name} {conf:.2f}"
|
||
(text_width, text_height), baseline = cv2.getTextSize(
|
||
label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2
|
||
)
|
||
cv2.rectangle(annotated_img, (x1, y1 - text_height - baseline),
|
||
(x1 + text_width, y1), color, -1)
|
||
cv2.putText(annotated_img, label, (x1, y1 - baseline),
|
||
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
|
||
|
||
class_counts[class_name] = class_counts.get(class_name, 0) + 1
|
||
|
||
if cv2_imwrite_chinese(output_path, annotated_img):
|
||
logger.debug(f"YOLO 完成：{os.path.basename(output_path)}")
|
||
return True, class_counts, ""
|
||
else:
|
||
return False, {}, "写入失败"
|
||
|
||
except Exception as e:
|
||
logger.error(f"YOLO 失败：{e}")
|
||
return False, {}, str(e)
|
||
|
||
def build_lut(anchors: np.ndarray) -> np.ndarray:
|
||
"""从锚点构建自定义 256 级颜色映射 LUT"""
|
||
lut = np.zeros((256, 3), dtype=np.uint8)
|
||
for c in range(3):
|
||
lut[:, c] = np.interp(np.arange(256), anchors[:, 0], anchors[:, c + 1])
|
||
return lut.reshape(256, 1, 3)
|
||
|
||
|
||
def detect_skin_region(img: np.ndarray):
|
||
"""多色彩空间皮肤检测，返回 (skin_soft, skin_binary, a_float)"""
|
||
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
|
||
_, a_ch, _ = cv2.split(lab)
|
||
a_float = a_ch.astype(np.float32)
|
||
|
||
mask_lab = cv2.inRange(lab, (40, 120, 105), (235, 185, 175))
|
||
ycrcb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
|
||
mask_ycrcb = cv2.inRange(ycrcb, (40, 133, 77), (240, 175, 135))
|
||
skin_mask = cv2.bitwise_or(mask_lab, mask_ycrcb)
|
||
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
|
||
skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel, iterations=3)
|
||
skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel, iterations=1)
|
||
skin_mask = cv2.dilate(skin_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)), iterations=1)
|
||
|
||
contours, _ = cv2.findContours(skin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if contours:
|
||
largest = max(contours, key=cv2.contourArea)
|
||
skin_mask = np.zeros_like(skin_mask)
|
||
cv2.drawContours(skin_mask, [largest], -1, 255, -1)
|
||
hull = cv2.convexHull(largest)
|
||
cv2.drawContours(skin_mask, [hull], -1, 255, -1)
|
||
|
||
skin_soft = cv2.GaussianBlur(skin_mask.astype(np.float32), (21, 21), 0) / 255.0
|
||
return skin_soft, skin_mask, a_float
|
||
|
||
|
||
def compute_heatmap_idx(a_float: np.ndarray, skin_mask: np.ndarray,
|
||
baseline_percentile: float, sensitivity: float) -> np.ndarray:
|
||
"""计算红度热力图索引（0-255）"""
|
||
skin_pixels = a_float[skin_mask > 128]
|
||
if len(skin_pixels) > 100:
|
||
baseline = np.percentile(skin_pixels, baseline_percentile)
|
||
else:
|
||
baseline = np.median(a_float)
|
||
redness = np.maximum(0, a_float - baseline) * sensitivity
|
||
vmax = np.percentile(redness, 98)
|
||
if vmax > 0:
|
||
idx = np.clip(redness, 0, vmax) / vmax * 255.0
|
||
else:
|
||
idx = np.zeros_like(a_float)
|
||
return np.clip(idx, 0, 255).astype(np.uint8)
|
||
|
||
|
||
def add_dark_details(result: np.ndarray, img: np.ndarray,
|
||
intensity: float = 0.50) -> np.ndarray:
|
||
"""叠加原图暗部细节（眉毛/头发/眼睛），保持五官可见"""
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
dark = (255.0 - gray.astype(np.float32)) / 255.0
|
||
dark = np.clip(dark * intensity, 0, 1)
|
||
dark = cv2.GaussianBlur(dark, (5, 5), 0)
|
||
return (result.astype(np.float32) * (1 - np.dstack([dark] * 3))).astype(np.uint8)
|
||
|
||
|
||
# ==================== RBX 红区分析 ====================
|
||
def analyze_red_zone(img: np.ndarray, output_path: str) -> Tuple[bool, Dict, str]:
|
||
try:
|
||
# 直接全图皮肤检测
|
||
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
|
||
_, a_ch, _ = cv2.split(lab)
|
||
a_float = a_ch.astype(np.float32)
|
||
|
||
mask_lab = cv2.inRange(lab, (40, 120, 105), (235, 185, 175))
|
||
ycrcb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
|
||
mask_ycrcb = cv2.inRange(ycrcb, (40, 133, 77), (240, 175, 135))
|
||
skin_mask = cv2.bitwise_or(mask_lab, mask_ycrcb)
|
||
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
|
||
skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel, iterations=3)
|
||
skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel, iterations=1)
|
||
skin_mask = cv2.dilate(skin_mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)), iterations=1)
|
||
|
||
contours, _ = cv2.findContours(skin_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if contours:
|
||
largest = max(contours, key=cv2.contourArea)
|
||
skin_mask = np.zeros_like(skin_mask)
|
||
cv2.drawContours(skin_mask, [largest], -1, 255, -1)
|
||
hull = cv2.convexHull(largest)
|
||
cv2.drawContours(skin_mask, [hull], -1, 255, -1)
|
||
|
||
skin_soft = cv2.GaussianBlur(skin_mask.astype(np.float32), (31, 31), 0) / 255.0
|
||
skin_pixels = a_float[skin_mask > 128]
|
||
if len(skin_pixels) > 100:
|
||
baseline_a = np.percentile(skin_pixels, 15)
|
||
else:
|
||
baseline_a = np.median(a_float)
|
||
|
||
redness_map = np.maximum(0, a_float - baseline_a) * 1.8
|
||
vmax = np.percentile(redness_map, 98)
|
||
if vmax > 0:
|
||
redness_map = np.clip(redness_map, 0, vmax)
|
||
heatmap_idx = (redness_map / vmax) * 255.0
|
||
else:
|
||
heatmap_idx = np.zeros_like(a_float)
|
||
heatmap_idx = np.clip(heatmap_idx, 0, 255).astype(np.uint8)
|
||
heatmap_idx = cv2.bilateralFilter(heatmap_idx, 9, 75, 75)
|
||
heatmap_idx = cv2.GaussianBlur(heatmap_idx, (7, 7), 0)
|
||
|
||
# 皮肤健康度统计
|
||
skin_area = np.sum(skin_mask > 128)
|
||
if skin_area > 0:
|
||
within_skin = heatmap_idx[skin_mask > 128]
|
||
normal_pct = float(np.sum(within_skin < 55) / skin_area * 100)
|
||
mild_pct = float(np.sum((within_skin >= 55) & (within_skin < 110)) / skin_area * 100)
|
||
warning_pct = float(np.sum((within_skin >= 110) & (within_skin < 170)) / skin_area * 100)
|
||
severe_pct = float(np.sum(within_skin >= 170) / skin_area * 100)
|
||
health_score = float(normal_pct * 1.0 + mild_pct * 0.7 + warning_pct * 0.4 + severe_pct * 0.1)
|
||
else:
|
||
normal_pct = mild_pct = warning_pct = severe_pct = health_score = 0.0
|
||
health_stats = {
|
||
'health_score': health_score, 'normal_pct': normal_pct,
|
||
'mild_pct': mild_pct, 'warning_pct': warning_pct, 'severe_pct': severe_pct,
|
||
}
|
||
|
||
alpha_red = heatmap_idx.astype(np.float32) / 255.0
|
||
alpha_red = np.clip(alpha_red * 1.5, 0, 1)
|
||
alpha_red = cv2.GaussianBlur(alpha_red, (5, 5), 0)
|
||
|
||
anchors = np.array([
|
||
[0, 232, 228, 242],
|
||
[3, 205, 170, 220],
|
||
[6, 150, 90, 215],
|
||
[10, 70, 30, 210],
|
||
[15, 35, 15, 225],
|
||
[20, 18, 14, 242],
|
||
[255, 2, 2, 97 ],
|
||
], dtype=np.float32)
|
||
heatmap_color = cv2.applyColorMap(heatmap_idx, build_lut(anchors))
|
||
heatmap_color = cv2.GaussianBlur(heatmap_color, (3, 3), 0)
|
||
|
||
hsv = cv2.cvtColor(heatmap_color, cv2.COLOR_BGR2HSV).astype(np.float32)
|
||
hsv[:, :, 1] = np.clip(hsv[:, :, 1] * 1.5, 0, 255)
|
||
heatmap_color = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
|
||
heatmap_float = heatmap_color.astype(np.float32) / 255.0
|
||
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float32)
|
||
skin_texture = cv2.GaussianBlur(gray, (5, 5), 0)
|
||
skin_texture = np.clip(215 + skin_texture * 0.12, 215, 252)
|
||
skin_texture = np.dstack([skin_texture] * 3) / 255.0
|
||
|
||
alpha_3 = np.dstack([alpha_red] * 3)
|
||
result = skin_texture * (1 - alpha_3 * 0.8) + heatmap_float * (alpha_3 * 0.8)
|
||
|
||
skin_alpha = np.dstack([skin_soft] * 3)
|
||
|
||
highpass = gray - cv2.GaussianBlur(gray, (11, 11), 0)
|
||
highpass = np.clip(highpass * 0.08, -8, 8)
|
||
highpass = cv2.GaussianBlur(highpass, (3, 3), 0)
|
||
highpass_3 = np.dstack([highpass] * 3) / 255.0
|
||
result = np.clip(result + highpass_3 * skin_alpha * 0.5, 0, 1)
|
||
|
||
dark_mask = np.clip(1.0 - gray / 55.0, 0, 1)
|
||
dark_mask = cv2.GaussianBlur(dark_mask, (11, 11), 0)
|
||
dark_3 = np.dstack([dark_mask] * 3)
|
||
white = np.ones_like(heatmap_float)
|
||
result = result * (1 - dark_3) + white * dark_3
|
||
|
||
result = (result * 255).astype(np.uint8)
|
||
|
||
if cv2_imwrite_chinese(output_path, result):
|
||
logger.debug(f"红区完成：{os.path.basename(output_path)}")
|
||
return True, health_stats, ""
|
||
return False, {}, "写入失败"
|
||
except Exception as e:
|
||
logger.error(f"红区失败：{e}")
|
||
return False, {}, str(e)
|
||
|
||
# ==================== 人脸红区检测 ====================
|
||
def analyze_face_red(img: np.ndarray, output_path: str) -> Tuple[bool, str]:
|
||
"""交叉偏振人脸皮肤分析：红度 + 色素"""
|
||
try:
|
||
# 直接全图皮肤检测
|
||
lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
|
||
l_ch, a_ch, b_ch = cv2.split(lab)
|
||
a_float = a_ch.astype(np.float32)
|
||
l_float = l_ch.astype(np.float32)
|
||
|
||
mask_lab = cv2.inRange(lab, (40, 120, 105), (235, 185, 175))
|
||
ycrcb = cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
|
||
mask_ycrcb = cv2.inRange(ycrcb, (40, 133, 77), (240, 175, 135))
|
||
skin = cv2.bitwise_or(mask_lab, mask_ycrcb)
|
||
k = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
|
||
skin = cv2.morphologyEx(skin, cv2.MORPH_CLOSE, k, iterations=3)
|
||
|
||
contours, _ = cv2.findContours(skin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if contours:
|
||
skin[:] = 0
|
||
cv2.drawContours(skin, [max(contours, key=cv2.contourArea)], -1, 255, -1)
|
||
|
||
skin_soft = cv2.GaussianBlur(skin.astype(np.float32), (15, 15), 0) / 255.0
|
||
|
||
# 3. 红度分析
|
||
skin_pixels = a_float[skin > 128]
|
||
baseline = np.percentile(skin_pixels, 20) if len(skin_pixels) > 100 else np.median(a_float)
|
||
redness = np.maximum(0, a_float - baseline) * 2.0
|
||
vmax = np.percentile(redness[skin > 128], 97) if np.sum(skin > 128) > 100 else 0
|
||
red_map = np.clip(redness / vmax * 255.0, 0, 255).astype(np.uint8) if vmax > 0 else np.zeros_like(a_float, dtype=np.uint8)
|
||
red_map = cv2.bilateralFilter(red_map, 7, 50, 50)
|
||
|
||
# 4. 色素分析
|
||
l_blur = cv2.GaussianBlur(l_float, (31, 31), 0)
|
||
pigment = np.abs(l_float - l_blur)
|
||
p_max = np.percentile(pigment[skin > 128], 95) if np.sum(skin > 128) > 100 else 0
|
||
pig_map = np.clip(pigment / p_max * 255.0, 0, 255).astype(np.uint8) if p_max > 0 else np.zeros_like(l_float, dtype=np.uint8)
|
||
|
||
# 5. 可视化：原图 + 红度(红) + 色素(蓝)
|
||
img_f = img.astype(np.float32)
|
||
skin_3 = np.dstack([skin_soft] * 3)
|
||
|
||
red_alpha = np.clip(red_map.astype(np.float32) / 200.0, 0, 1)
|
||
red_alpha = cv2.GaussianBlur(red_alpha, (5, 5), 0)
|
||
red_ol = np.zeros_like(img)
|
||
red_ol[:, :, 2] = 180
|
||
red_3 = np.dstack([red_alpha] * 3)
|
||
result = img_f * (1 - red_3 * 0.5) + red_ol.astype(np.float32) * (red_3 * 0.5)
|
||
|
||
pig_alpha = np.clip(pig_map.astype(np.float32) / 200.0, 0, 1)
|
||
pig_alpha = cv2.GaussianBlur(pig_alpha, (5, 5), 0)
|
||
pig_ol = np.zeros_like(img)
|
||
pig_ol[:, :, 0] = 160
|
||
pig_3 = np.dstack([pig_alpha] * 3)
|
||
result = result * (1 - pig_3 * 0.4) + pig_ol.astype(np.float32) * (pig_3 * 0.4)
|
||
|
||
# 非皮肤压暗
|
||
dark = img_f * 0.3 + 50
|
||
result = result * skin_3 + dark * (1 - skin_3)
|
||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
if cv2_imwrite_chinese(output_path, result):
|
||
logger.debug(f"交叉偏振分析完成：{os.path.basename(output_path)}")
|
||
return True, ""
|
||
return False, "写入失败"
|
||
except Exception as e:
|
||
logger.error(f"交叉偏振分析失败：{e}")
|
||
return False, str(e)
|
||
|
||
# ==================== 热力敏感图分析 ====================
|
||
def analyze_heat_sensitivity(img: np.ndarray, output_path: str) -> Tuple[bool, str]:
|
||
try:
|
||
# 直接全图皮肤检测
|
||
skin_soft, skin_binary, a_float = detect_skin_region(img)
|
||
heatmap_idx = compute_heatmap_idx(a_float, skin_binary, 30, 1.8)
|
||
heatmap_idx = cv2.GaussianBlur(heatmap_idx, (7, 7), 0)
|
||
heatmap_idx = cv2.bilateralFilter(heatmap_idx, 7, 60, 60)
|
||
|
||
anchors = np.array([
|
||
[0, 192, 97, 16 ],
|
||
[1, 233, 241, 86 ],
|
||
[60, 8, 213, 204],
|
||
[100, 0, 250, 245],
|
||
[160, 0, 150, 255],
|
||
[245, 33, 38, 200],
|
||
[255, 4, 38, 152],
|
||
], dtype=np.float32)
|
||
heatmap_color = cv2.applyColorMap(heatmap_idx, build_lut(anchors))
|
||
heatmap_color = cv2.convertScaleAbs(heatmap_color, alpha=1.05, beta=0)
|
||
|
||
result = (heatmap_color.astype(np.float32) * 0.92 +
|
||
np.full_like(img, 255, dtype=np.float32) * 0.08)
|
||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||
|
||
result = add_dark_details(result, img)
|
||
|
||
if cv2_imwrite_chinese(output_path, result):
|
||
logger.debug(f"热力敏感图完成：{os.path.basename(output_path)}")
|
||
return True, ""
|
||
return False, "写入失败"
|
||
except Exception as e:
|
||
logger.error(f"热力敏感图失败：{e}")
|
||
return False, str(e)
|
||
|
||
# ==================== UV 荧光分析 ====================
|
||
def detect_skin_uv(img: np.ndarray):
|
||
"""UV 照片皮肤检测 — 亮度阈值 + 形态学"""
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
|
||
# UV 照片: 皮肤有微弱的 UV 反射/荧光, 背景几乎全黑
|
||
_, skin = cv2.threshold(gray, 18, 255, cv2.THRESH_BINARY)
|
||
|
||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (11, 11))
|
||
skin = cv2.morphologyEx(skin, cv2.MORPH_CLOSE, kernel, iterations=3)
|
||
skin = cv2.morphologyEx(skin, cv2.MORPH_OPEN, kernel, iterations=1)
|
||
skin = cv2.dilate(skin, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7)), iterations=1)
|
||
|
||
contours, _ = cv2.findContours(skin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
if contours:
|
||
largest = max(contours, key=cv2.contourArea)
|
||
skin = np.zeros_like(skin)
|
||
cv2.drawContours(skin, [largest], -1, 255, -1)
|
||
hull = cv2.convexHull(largest)
|
||
cv2.drawContours(skin, [hull], -1, 255, -1)
|
||
|
||
skin_soft = cv2.GaussianBlur(skin.astype(np.float32), (21, 21), 0) / 255.0
|
||
return skin_soft, skin
|
||
|
||
|
||
def analyze_pigment(img: np.ndarray, output_path: str) -> Tuple[bool, str]:
|
||
"""UV 荧光分析 — 黑底 + 皮肤灰度 + 荧光绿点"""
|
||
try:
|
||
skin_soft, skin_binary = detect_skin_uv(img)
|
||
|
||
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).astype(np.float32)
|
||
clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(16, 16))
|
||
skin_base = clahe.apply(gray.astype(np.uint8)).astype(np.float32)
|
||
|
||
skin_base = skin_base * skin_soft
|
||
skin_base = np.power(skin_base / 255.0, 1.2) * 255.0
|
||
|
||
skin_base_3 = np.dstack([skin_base / 255.0] * 3)
|
||
result_u8 = (np.clip(skin_base_3, 0, 1) * 255).astype(np.uint8)
|
||
|
||
if cv2_imwrite_chinese(output_path, result_u8):
|
||
logger.debug(f"UV 荧光完成：{os.path.basename(output_path)}")
|
||
return True, ""
|
||
else:
|
||
return False, "写入失败"
|
||
|
||
except Exception as e:
|
||
logger.error(f"UV 荧光失败：{e}")
|
||
return False, str(e)
|
||
|
||
# ==================== 单张图片处理 ====================
|
||
def process_single_image(args: Tuple) -> Dict:
|
||
img_path, model, output_base, verbose = args
|
||
|
||
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}
|
||
|
||
if img_path.suffix.lower() not in IMAGE_EXTENSIONS:
|
||
return {'status': 'skip', 'filename': img_path.name}
|
||
|
||
logger.info(f"处理：{img_path.name}")
|
||
|
||
img = cv2_imread_chinese(str(img_path))
|
||
if img is None:
|
||
return {'status': 'error', 'filename': img_path.name, 'error': '读取失败'}
|
||
|
||
img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
|
||
|
||
result = {
|
||
'status': 'success',
|
||
'filename': img_path.name,
|
||
'stem': img_path.stem,
|
||
'acne_counts': {},
|
||
'health_score': 0,
|
||
'normal_pct': 0,
|
||
'mild_pct': 0,
|
||
'warning_pct': 0,
|
||
'severe_pct': 0,
|
||
'errors': []
|
||
}
|
||
|
||
# 1. YOLO 痤疮检测
|
||
try:
|
||
yolo_out = output_base / "yolo" / f"yolo_{img_path.stem}.jpg"
|
||
success, acne_counts, error = detect_acne(img, str(img_path), model, str(yolo_out))
|
||
if success:
|
||
result['acne_counts'] = acne_counts
|
||
else:
|
||
result['errors'].append(f"YOLO: {error}")
|
||
except Exception as e:
|
||
result['errors'].append(f"YOLO: {e}")
|
||
|
||
# 3. RBX 红区分析
|
||
try:
|
||
red_out = output_base / "redzone" / f"red_{img_path.stem}.jpg"
|
||
success, health_stats, error = analyze_red_zone(img, str(red_out))
|
||
if success:
|
||
result.update(health_stats)
|
||
else:
|
||
result['errors'].append(f"红区：{error}")
|
||
except Exception as e:
|
||
result['errors'].append(f"红区：{e}")
|
||
|
||
# 4. 人脸红区检测
|
||
try:
|
||
face_red_out = output_base / "face_red" / f"face_red_{img_path.stem}.jpg"
|
||
success, error = analyze_face_red(img, str(face_red_out))
|
||
if not success:
|
||
result['errors'].append(f"人脸红区：{error}")
|
||
except Exception as e:
|
||
result['errors'].append(f"人脸红区：{e}")
|
||
|
||
# 6. 热力敏感图分析
|
||
try:
|
||
sensitivity_out = output_base / "sensitivity" / f"sensitivity_{img_path.stem}.jpg"
|
||
success, error = analyze_heat_sensitivity(img, str(sensitivity_out))
|
||
if not success:
|
||
result['errors'].append(f"热力敏感图：{error}")
|
||
except Exception as e:
|
||
result['errors'].append(f"热力敏感图：{e}")
|
||
|
||
# 7. 色素铅图分析
|
||
try:
|
||
pig_out = output_base / "pigment" / f"pigment_{img_path.stem}.jpg"
|
||
success, error = analyze_pigment(img, str(pig_out))
|
||
if not success:
|
||
result['errors'].append(f"色素：{error}")
|
||
except Exception as e:
|
||
result['errors'].append(f"色素：{e}")
|
||
|
||
if result['errors']:
|
||
result['status'] = 'partial'
|
||
logger.warning(f"{img_path.name} 部分失败：{result['errors']}")
|
||
else:
|
||
logger.info(f"{img_path.name} 完成")
|
||
|
||
return result
|
||
|
||
# ==================== 保存文本报告 ====================
|
||
def save_text_report(results: List[Dict], texts_dir: Path):
|
||
try:
|
||
# 确保texts目录存在
|
||
if not texts_dir.exists():
|
||
texts_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
text_path = texts_dir / "analysis_report.txt"
|
||
|
||
valid = [r for r in results if r.get('health_score', 0) > 0]
|
||
if valid:
|
||
avg_health = np.mean([r['health_score'] for r in valid])
|
||
avg_normal = np.mean([r['normal_pct'] for r in valid])
|
||
avg_mild = np.mean([r['mild_pct'] for r in valid])
|
||
avg_warning = np.mean([r['warning_pct'] for r in valid])
|
||
avg_severe = np.mean([r['severe_pct'] for r in valid])
|
||
else:
|
||
avg_health = avg_normal = avg_mild = avg_warning = avg_severe = 0.0
|
||
|
||
with open(text_path, 'w', encoding='utf-8') as f:
|
||
f.write(f"Health Score::{avg_health:.1f} / 100;;\n")
|
||
f.write(f"Normal area::{avg_normal:.1f}%;;\n")
|
||
f.write(f"Mild abnormality::{avg_mild:.1f}%;;\n")
|
||
f.write(f"Moderate abnormality::{avg_warning:.1f}%;;\n")
|
||
f.write(f"Severe abnormality::{avg_severe:.1f}%;;\n")
|
||
|
||
logger.info(f"文本报告已保存：{text_path}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"文本报告生成失败：{e}")
|
||
return False
|
||
|
||
# ==================== 回调后端接口 ====================
|
||
def callback_api(task_id: str, status: int, api_url: str = API_HOST+API_CALLBACK):
|
||
"""完成后回调后端接口"""
|
||
try:
|
||
payload = {
|
||
"task_id": task_id,
|
||
"status": status # 1=成功，0=失败
|
||
}
|
||
response = requests.post(api_url, json=payload, timeout=10)
|
||
if response.status_code == 200:
|
||
logger.info(f"回调成功：{task_id}")
|
||
return True
|
||
else:
|
||
logger.warning(f"回调失败：{response.status_code, api_url}")
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"回调异常：{e}")
|
||
return False
|
||
|
||
# ==================== 主函数 ====================
|
||
def main():
|
||
global logger
|
||
|
||
parser = argparse.ArgumentParser(description='皮肤分析工具 - 后端 API 版')
|
||
parser.add_argument('--task-id', type=str,
|
||
help='任务 ID')
|
||
parser.add_argument('input_dir', type=str,
|
||
help='输入目录完整路径（后端传入的 folder_name）')
|
||
parser.add_argument('--model', '-m', type=str, default='best.pt',
|
||
help='YOLO 痤疮模型路径 (默认：cuochuang.pt)')
|
||
parser.add_argument('--workers', '-w', type=int, default=3,
|
||
help='并发工作线程数 (默认：3)')
|
||
parser.add_argument('--verbose', '-v', action='store_true',
|
||
help='详细输出模式')
|
||
parser.add_argument('--api-url', type=str, default=API_HOST+API_CALLBACK,
|
||
help='回调 API 地址')
|
||
args = parser.parse_args()
|
||
|
||
# 初始化日志
|
||
logger = setup_logging(args.verbose)
|
||
|
||
# 输入输出路径
|
||
input_dir = Path(args.input_dir)
|
||
# task_id = input_dir.name # 从文件夹名称自动提取 task_id // 修改为不从目录获取任务ID
|
||
task_id = args.task_id
|
||
output_base = input_dir.parent / 'output'
|
||
|
||
logger.info("=" * 60)
|
||
logger.info("皮肤分析工具 - 后端 API 版")
|
||
logger.info("=" * 60)
|
||
logger.info(f"任务 ID: {task_id}")
|
||
logger.info(f"输入目录：{input_dir}")
|
||
logger.info(f"输出目录：{output_base}")
|
||
logger.info(f"模型路径：{args.model}")
|
||
logger.info(f"并发线程：{args.workers}")
|
||
logger.info(f"回调地址：{args.api_url}")
|
||
|
||
model_path = args.model
|
||
|
||
if not input_dir.exists():
|
||
logger.error(f"输入目录不存在：{input_dir}")
|
||
callback_api(task_id, status=0, api_url=args.api_url)
|
||
return
|
||
|
||
# 创建输出目录
|
||
texts_dir = output_base / 'texts'
|
||
|
||
try:
|
||
# 首先创建输出目录
|
||
output_base.mkdir(exist_ok=True, parents=True)
|
||
logger.info(f"输出目录创建成功：{output_base}")
|
||
|
||
# 然后创建子目录
|
||
for subdir in ['yolo', 'redzone', 'face_red', 'sensitivity', 'pigment', 'texts']:
|
||
subdir_path = output_base / subdir
|
||
subdir_path.mkdir(exist_ok=True)
|
||
logger.info(f"子目录创建成功：{subdir_path}")
|
||
except Exception as e:
|
||
logger.error(f"目录创建失败：{e}")
|
||
callback_api(task_id, status=0, api_url=args.api_url)
|
||
return
|
||
|
||
# 加载痤疮模型
|
||
logger.info("加载 YOLO 痤疮模型...")
|
||
try:
|
||
model = YOLO(model_path)
|
||
# 主线程预 fuse，避免首批多线程推理时争抢 BN
|
||
try:
|
||
if hasattr(model, "fuse"):
|
||
model.fuse()
|
||
elif hasattr(model, "model") and hasattr(model.model, "fuse"):
|
||
model.model.fuse()
|
||
except Exception as fuse_err:
|
||
logger.warning(f"YOLO 预 fuse 跳过：{fuse_err}")
|
||
logger.info("痤疮模型加载成功")
|
||
except Exception as e:
|
||
logger.error(f"痤疮模型加载失败：{e}")
|
||
callback_api(task_id, status=0, api_url=args.api_url)
|
||
return
|
||
|
||
# 获取图片列表
|
||
try:
|
||
image_files = list(input_dir.iterdir())
|
||
total_images = len(image_files)
|
||
logger.info(f"发现 {total_images} 个文件")
|
||
except Exception as e:
|
||
logger.error(f"获取图片列表失败：{e}")
|
||
callback_api(task_id, status=0, api_url=args.api_url)
|
||
return
|
||
|
||
# 并发处理
|
||
results = []
|
||
tasks = [(img_path, model, output_base, args.verbose)
|
||
for img_path in image_files]
|
||
|
||
logger.info("开始并发处理...")
|
||
try:
|
||
with ThreadPoolExecutor(max_workers=args.workers) as executor:
|
||
futures = {executor.submit(process_single_image, task): task
|
||
for task in tasks}
|
||
|
||
for future in tqdm(as_completed(futures), total=len(futures),
|
||
desc="处理进度", unit="图"):
|
||
result = future.result()
|
||
results.append(result)
|
||
except Exception as e:
|
||
logger.error(f"并发处理失败：{e}")
|
||
callback_api(task_id, status=0, api_url=args.api_url)
|
||
return
|
||
|
||
# 保存文本报告
|
||
logger.info("\n保存文本报告中...")
|
||
logger.info(f"文本目录：{texts_dir}")
|
||
try:
|
||
save_text_report(results, texts_dir)
|
||
logger.info("文本报告保存成功")
|
||
except Exception as e:
|
||
logger.error(f"保存数据失败：{e}")
|
||
callback_api(task_id, status=0, api_url=args.api_url)
|
||
return
|
||
|
||
# 统计总结
|
||
success = sum(1 for r in results if r['status'] == 'success')
|
||
partial = sum(1 for r in results if r['status'] == 'partial')
|
||
errors = sum(1 for r in results if r['status'] == 'error')
|
||
|
||
logger.info("\n" + "=" * 60)
|
||
logger.info("全部完成！")
|
||
logger.info(f"成功：{success} | 部分：{partial} | 失败：{errors}")
|
||
logger.info(f"输出目录：{output_base}")
|
||
logger.info(f"文本报告：{texts_dir / 'analysis_report.txt'}")
|
||
logger.info("=" * 60)
|
||
|
||
# 回调后端接口（状态 1=成功）
|
||
status = 1 if errors == 0 or (success + partial) > 0 else 0
|
||
callback_api(task_id, status=status, api_url=args.api_url)
|
||
|
||
if __name__ == "__main__":
|
||
main()
|