#!/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()