diff --git a/.ultralytics/Ultralytics/persistent_cache.json b/.ultralytics/Ultralytics/persistent_cache.json new file mode 100644 index 0000000..cf4bc9c --- /dev/null +++ b/.ultralytics/Ultralytics/persistent_cache.json @@ -0,0 +1,3 @@ +{ + "cpu_info": "AMD Ryzen 7 5800H with Radeon Graphics" +} \ No newline at end of file diff --git a/.ultralytics/Ultralytics/settings.json b/.ultralytics/Ultralytics/settings.json new file mode 100644 index 0000000..8644039 --- /dev/null +++ b/.ultralytics/Ultralytics/settings.json @@ -0,0 +1,21 @@ +{ + "settings_version": "0.0.6", + "datasets_dir": "D:\\\u684c\u9762\\\u76ae\u80a4\u5206\u6790\u4eea\\py-skin-analysis-worker-main-e1281e11629a703fa96b8fd5d434b32749e64b53\\datasets", + "weights_dir": "weights", + "runs_dir": "runs", + "uuid": "857db6b65a8087ce7afdf4a38b0cf1a8cfda560ed281365e348293d9b06bc403", + "sync": true, + "api_key": "", + "openai_api_key": "", + "clearml": true, + "comet": true, + "dvc": true, + "hub": true, + "mlflow": true, + "neptune": true, + "raytune": true, + "tensorboard": false, + "wandb": false, + "vscode_msg": true, + "openvino_msg": true +} \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..1536057 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,22 @@ +FROM skin-py-cpu-base:v1.0.2 + +WORKDIR /app + +ENV PYTHONUNBUFFERED=1 +ENV PIP_NO_CACHE_DIR=1 + +COPY wanzheng.py . +COPY api_server.py . +COPY task_queue.py . +COPY resident_worker.py . +COPY config ./config + +EXPOSE 8071 + +CMD ["python", "api_server.py"] + +# 服务器执行 nvidia-smi 查看CUDA版本 比如 CUDA 12.1 就是/whl/cu121 +# RUN pip install \ +# torch==2.4.1 \ +# torchvision==0.19.1 \ +# --index-url https://download.pytorch.org/whl/cu121 diff --git a/TASK_QUEUE.md b/TASK_QUEUE.md new file mode 100644 index 0000000..f688396 --- /dev/null +++ b/TASK_QUEUE.md @@ -0,0 +1,60 @@ +# 任务队列配置 + +任务提交后先写入 SQLite，再由固定数量的工作线程分发给常驻分析子进程。 +服务启动时每个子进程加载一次 YOLO，并使用 640×640 空白图片预热；全部就绪后 API 才开放。 +后续任务复用同一模型。YOLO 推理通过模型独占锁串行执行，其他图片分析仍可并行。 +原有提交、查询和列表接口路径不变，重复提交等待中或运行中的任务不会重复执行。 + +| 环境变量 | 默认值 | 说明 | +| --- | --- | --- | +| `MAX_CONCURRENT_TASKS` | `1` | 同时运行的分析子进程数，必须大于 0 | +| `MAX_PENDING_TASKS` | `100` | 等待队列容量，不含运行中的任务，必须大于 0 | +| `TASK_DB_PATH` | 项目下 `data/tasks.sqlite3` | SQLite 文件路径 | +| `MODEL_PATH` | 项目下 `best.pt` | 模型路径，Compose 使用 `/app/best.pt` | +| `IMAGE_WORKERS` | `4` | 每个任务的图片处理线程数 | + +每增加一个任务并发槽位都会常驻一份模型，占用额外内存或显存。 +启动预热超时为 180 秒，每个任务超时为 3600 秒。超时、崩溃或停止会回收对应进程； +下一个任务重新创建并预热该槽位，不自动重试已经失败的任务。 + +Compose 将数据库放在 `/data/analysis/.worker/tasks.sqlite3`，由现有的 +`/data/analysis:/data/analysis` 挂载持久保存。目录必须可写，数据库应放在本机磁盘， +不要放到网络共享文件系统。修改配置后重启服务生效。 + +队列满时，新任务返回 HTTP 429 和 `Retry-After: 5`；已有活动任务仍正常返回状态。 +`/health` 增加 `tasks.pending` 和 `queue` 配置字段。 + +## 重启与停止 + +- 已完成和失败的任务保留，可以继续查询。 +- `pending` 任务重启后继续执行；输入目录在提交时解析并保存，跨天不会重新选日期。 +- 异常退出留下的 `running` 任务启动时标记为 `failed`，错误提示重新提交。 +- 正常停止会终止、回收正在执行的分析子进程，并将对应任务标记失败；等待任务保留。 +- 不自动重试中断任务，避免重复生成结果或重复回调。 + +使用单个 API worker 和单个服务实例。独占锁会拒绝第二个进程同时打开同一队列， +避免并发上限被成倍放大。不同数据库之间不共享限制。 +如果只强制杀死 API 父进程，操作系统可能留下分析子进程；再次启动前需要确认这些进程已结束。 +容器部署应停止/重启整个容器。 + +## 回调与范围 + +回调仅记录已知运行任务的结果，任务仍保持 `running`，直到常驻进程返回该任务的结果。 +失败回调不会被成功结果覆盖；未知任务返回 404；等待中和已结束任务不接受状态覆盖。 +外部回调载荷仍为 `task_id` 和 `status`。 + +图片算法、输出目录规则和报告格式保持原样。空目录、部分分析失败、报告写入失败会明确返回失败， +避免把不完整处理当成测速成功。回调仍同步执行，没有加入独立通知队列。 +现有输出目录可能发生覆盖，调整并发数前应先处理任务输出隔离。 +历史内存字典无法在重启后迁移；持久化从此版本接收的任务开始。 + +## 验证 + +安装服务依赖和 `httpx` 后，在项目根目录执行： + +```text +python -B -m unittest discover -s tests -v +``` + +队列测试只使用 Python 标准库。接口测试使用 FastAPI TestClient，子进程测试运行短 Python +程序，不加载 YOLO、不发送后端回调。缺少接口测试依赖时会明确显示跳过。 diff --git a/api_server.py b/api_server.py new file mode 100644 index 0000000..6836efb --- /dev/null +++ b/api_server.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +皮肤分析 API 接口服务 +""" + +import sys +import os +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from fastapi import FastAPI, HTTPException, Query +from contextlib import asynccontextmanager +import asyncio +from pydantic import BaseModel, Field +import uvicorn +import logging +from datetime import datetime +from typing import Optional +from pathlib import Path +from config.settings import (API_HOST, API_CALLBACK, TASK_DB_PATH, + MAX_CONCURRENT_TASKS, MAX_PENDING_TASKS, MODEL_PATH, IMAGE_WORKERS) +from task_queue import TaskQueue, QueueFull +from resident_worker import ResidentPool + +# ==================== 日志配置 ==================== +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s | %(levelname)s | %(message)s' +) +logger = logging.getLogger(__name__) + +# ==================== 服务生命周期 ==================== +@asynccontextmanager +async def lifespan(app): + await asyncio.to_thread(task_queue.start) + try: + yield + finally: + await asyncio.to_thread(task_queue.stop) + + +app = FastAPI( + title="皮肤分析任务接口", + description="接收文件夹名称，触发皮肤分析任务", + version="1.0.0", + lifespan=lifespan, +) + +# ==================== 数据模型 ==================== +class TaskRequest(BaseModel): + folder_name: str = Field(..., description="文件夹全路径") + task_id: str = Field(..., description="目标文件夹名称（任务 ID）") + +class TaskResponse(BaseModel): + code: int = Field(..., description="状态码") + message: str = Field(..., description="消息") + task_id: str = Field(..., description="任务 ID") + status: str = Field(..., description="任务状态") + +class TaskStatus(BaseModel): + task_id: str = Field(..., description="任务 ID") + status: str = Field(..., description="状态") + created_at: str = Field(..., description="创建时间") + completed_at: Optional[str] = Field(default=None, description="完成时间") + error: Optional[str] = Field(default=None, description="错误信息") + +class CallbackRequest(BaseModel): + task_id: str = Field(..., description="任务 ID") + status: int = Field(..., description="处理状态，1=成功，0=失败") + +# 每个任务槽位对应一个常驻模型进程；启动预热后才开放 API。 +analysis_pool = ResidentPool(MODEL_PATH, MAX_CONCURRENT_TASKS, IMAGE_WORKERS, + API_HOST + API_CALLBACK) +task_queue = TaskQueue(TASK_DB_PATH, MAX_CONCURRENT_TASKS, + MAX_PENDING_TASKS, analysis_pool) + +# ==================== API 接口 ==================== +@app.get("/") +async def root(): + return { + "service": "皮肤分析任务接口", + "version": "1.0.0", + "status": "running" + } + +@app.get("/health") +def health_check(): + return { + "status": "healthy", + "timestamp": datetime.now().isoformat(), + "tasks": task_queue.counts(), + "queue": {"max_concurrent": task_queue.concurrency, + "max_pending": task_queue.max_pending} + } + + +@app.post("/api/v1/analysis/submit", response_model=TaskResponse) +def submit_task(request: TaskRequest): + task_id = request.task_id if request.task_id else Path(request.folder_name).name + if not task_id or len(task_id) < 3: + raise HTTPException(status_code=400, detail="任务 ID 格式错误，至少 3 个字符") + # Resolve the date at submission, not after a queued task starts tomorrow. + input_dir = Path(request.folder_name) + if not input_dir.is_absolute(): + input_dir = Path("/data/analysis") / datetime.now().strftime("%Y%m%d") / task_id + try: + task, created = task_queue.submit( + task_id, {"folder_name": request.folder_name, "task_id": task_id}, + str(input_dir.resolve())) + except QueueFull as exc: + raise HTTPException(status_code=429, detail=str(exc), + headers={"Retry-After": "5"}) from exc + return TaskResponse(code=200, + message="任务已加入队列" if created else "任务已存在，正在等待或运行中", + task_id=task_id, status=task["status"]) + +@app.post("/api/v1/analysis/callback") +def analysis_callback(request: CallbackRequest): + if request.status not in (0, 1): + raise HTTPException(status_code=400, detail="状态只允许 0 或 1") + status = task_queue.callback(request.task_id, request.status) + if status is None: + raise HTTPException(status_code=404, detail="任务不存在") + return {"code": 200, "message": "回调接收成功", "task_id": request.task_id, "status": status} + +@app.get("/api/v1/analysis/status/{task_id}", response_model=TaskStatus) +def get_task_status(task_id: str): + task = task_queue.get(task_id) + if task is None: + raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}") + + return TaskStatus( + task_id=task["task_id"], + status=task["status"], + created_at=task["created_at"], + completed_at=task.get("completed_at"), + error=task.get("error") + ) + +@app.get("/api/v1/analysis/list") +def list_tasks( + status: Optional[str] = Query(default=None), + limit: int = Query(default=20, ge=1, le=100) +): + tasks = task_queue.list(status, limit) + return {"total": len(tasks), "tasks": tasks} + +# ==================== 启动服务 ==================== +if __name__ == "__main__": + print("="*60) + print("皮肤分析 API 服务") + print("="*60) + print("监听地址：http://0.0.0.0:8071") + print("提交接口：POST /api/v1/analysis/submit") + print("状态查询：GET /api/v1/analysis/status/{task_id}") + print("健康检查：GET /health") + print("="*60) + print() + + uvicorn.run(app, host="0.0.0.0", port=8071) diff --git a/best.pt b/best.pt new file mode 100644 index 0000000..1456d9a Binary files /dev/null and b/best.pt differ diff --git a/config/__init__.py b/config/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/config/__pycache__/__init__.cpython-310.pyc b/config/__pycache__/__init__.cpython-310.pyc new file mode 100644 index 0000000..1de2f04 Binary files /dev/null and b/config/__pycache__/__init__.cpython-310.pyc differ diff --git a/config/__pycache__/settings.cpython-310.pyc b/config/__pycache__/settings.cpython-310.pyc new file mode 100644 index 0000000..58591f0 Binary files /dev/null and b/config/__pycache__/settings.cpython-310.pyc differ diff --git a/config/settings.py b/config/settings.py new file mode 100644 index 0000000..86c01b3 --- /dev/null +++ b/config/settings.py @@ -0,0 +1,21 @@ +import os +from pathlib import Path + +API_HOST = os.getenv( + "API_HOST", + "http://skin-test-api.ailuowan.com" +) + +API_CALLBACK = os.getenv( + "API_CALLBACK", + "/api/v1/analysis/callback" +) + +# SQLite must reside on persistent local storage; Compose uses /data/analysis. +TASK_DB_PATH = os.getenv( + "TASK_DB_PATH", str(Path(__file__).resolve().parent.parent / "data" / "tasks.sqlite3") +) +MAX_CONCURRENT_TASKS = int(os.getenv("MAX_CONCURRENT_TASKS", "1")) +MAX_PENDING_TASKS = int(os.getenv("MAX_PENDING_TASKS", "100")) +MODEL_PATH = os.getenv("MODEL_PATH", str(Path(__file__).resolve().parent.parent / "best.pt")) +IMAGE_WORKERS = int(os.getenv("IMAGE_WORKERS", "4")) diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..3e11eb4 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,53 @@ +services: + # 哪个项目-什么语言-做什么 + skin-py-analysis: + build: /data/www/py-skin-analysis-worker + image: analysis-py:latest + working_dir: /app + expose: + - "8071" + volumes: + - /data/www/py-skin-analysis-worker:/app + - /data/analysis:/data/analysis + environment: + - PYTHONUNBUFFERED=1 + - TASK_DB_PATH=/data/analysis/.worker/tasks.sqlite3 + - MAX_CONCURRENT_TASKS=1 + - MAX_PENDING_TASKS=100 + - MODEL_PATH=/app/best.pt + - IMAGE_WORKERS=4 + command: python api_server.py + restart: unless-stopped + networks: + - default + - common-network +networks: + default: + external: true + name: skin-network + + common-network: + external: true + +# services: +# # 皮肤分析 API 服务 +# skin-api: +# build: . +# image: skin-analysis:latest +# container_name: skin-analysis-api +# ports: +# - "8071:8071" +# volumes: +# # 挂载数据目录 +# - D:/data/analysis:/data/analysis +# environment: +# - PYTHONUNBUFFERED=1 +# command: python api_server.py +# restart: unless-stopped +# deploy: +# resources: +# reservations: +# devices: +# - driver: nvidia +# count: 1 +# capabilities: [gpu] diff --git a/resident_worker.py b/resident_worker.py new file mode 100644 index 0000000..b700ce6 --- /dev/null +++ b/resident_worker.py @@ -0,0 +1,135 @@ +"""Spawned model processes, reused across tasks and replaceable after a crash.""" +import logging +import multiprocessing +import queue +import time + +logger = logging.getLogger(__name__) + + +def analysis_process(connection, model_path, workers): + try: + from wanzheng import SafeYOLO, analyze_directory + model = SafeYOLO(model_path) + connection.send({'ready': True}) + while True: + task = connection.recv() + if task is None: + return + result = analyze_directory(task['input_dir'], task['task_id'], model, + workers, task.get('api_url')) + connection.send(result) + except EOFError: + pass + except BaseException as exc: + logger.exception('常驻分析进程异常') + try: + connection.send({'error': str(exc) or type(exc).__name__}) + except (OSError, EOFError): + pass + finally: + connection.close() + + +class ResidentPool: + def __init__(self, model_path, concurrency=1, workers=4, api_url=None, + startup_timeout=180, task_timeout=3600, target=analysis_process): + if min(concurrency, workers, startup_timeout, task_timeout) <= 0: + raise ValueError('并发、线程数和超时必须大于 0') + self.model_path = str(model_path) + self.concurrency = concurrency + self.workers = workers + self.api_url = api_url + self.startup_timeout = startup_timeout + self.task_timeout = task_timeout + self.target = target + self.context = multiprocessing.get_context('spawn') + self.slots = queue.Queue() + + def receive(self, slot, timeout, stop_event=None): + connection, process = slot + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if stop_event is not None and stop_event.is_set(): + raise RuntimeError('服务停止，任务中断，请重新提交') + if connection.poll(0.1): + try: + return connection.recv() + except (EOFError, OSError) as exc: + raise RuntimeError('常驻分析进程连接中断') from exc + if not process.is_alive(): + raise RuntimeError(f'常驻分析进程退出：{process.exitcode}') + raise RuntimeError('常驻分析进程等待超时') + + @staticmethod + def dispose(slot): + if slot is None: + return + connection, process = slot + connection.close() + if process.is_alive(): + process.terminate() + process.join(timeout=5) + if process.is_alive(): + process.kill() + process.join() + process.close() + + def spawn(self, stop_event=None): + parent, child = self.context.Pipe() + process = self.context.Process(target=self.target, + args=(child, self.model_path, self.workers)) + try: + process.start() + except BaseException: + parent.close() + child.close() + process.close() + raise + child.close() + slot = parent, process + try: + message = self.receive(slot, self.startup_timeout, stop_event) + if not message.get('ready'): + raise RuntimeError(message.get('error', '模型预热失败')) + return slot + except BaseException: + self.dispose(slot) + raise + + def start(self): + try: + for _ in range(self.concurrency): + self.slots.put(self.spawn()) + except BaseException: + self.stop() + raise + + def stop(self): + # TaskQueue joins its worker threads before closing idle slots. + while True: + try: + self.dispose(self.slots.get_nowait()) + except queue.Empty: + break + + def __call__(self, task, stop_event): + slot = self.slots.get() + try: + if stop_event.is_set(): + raise RuntimeError('服务停止，任务中断，请重新提交') + if slot is None or not slot[1].is_alive(): + self.dispose(slot) + slot = None + slot = self.spawn(stop_event) + slot[0].send(dict(task, api_url=self.api_url)) + result = self.receive(slot, self.task_timeout, stop_event) + except BaseException: + self.dispose(slot) + slot = None + raise + finally: + self.slots.put(slot) + if result.get('error'): + raise RuntimeError(result['error']) + return result diff --git a/task_queue.py b/task_queue.py new file mode 100644 index 0000000..1cc1268 --- /dev/null +++ b/task_queue.py @@ -0,0 +1,187 @@ +"""SQLite task queue for one service process with bounded worker threads.""" +import json +import logging +import sqlite3 +import threading +from contextlib import contextmanager +from datetime import datetime +from pathlib import Path + +logger = logging.getLogger(__name__) + + +class QueueFull(Exception): + pass + + +class TaskQueue: + def __init__(self, path, concurrency, max_pending, runner): + if concurrency < 1 or max_pending < 1: + raise ValueError("任务并发数和队列容量必须大于 0") + self.path = Path(path).resolve() + self.concurrency = concurrency + self.max_pending = max_pending + self.runner = runner + self.stop_event = threading.Event() + self.threads = [] + self.lock_file = None + + @contextmanager + def connect(self): + conn = sqlite3.connect(str(self.path), timeout=30) + conn.row_factory = sqlite3.Row + try: + with conn: + yield conn + finally: + conn.close() + + def start(self): + self.path.parent.mkdir(parents=True, exist_ok=True) + # Hold a kernel lock for the full process lifecycle. A second API worker + # must not recover this process's running tasks or multiply the limit. + lock_file = open(str(self.path) + ".lock", "a+b") + try: + import os + if os.name == "nt": + import msvcrt + lock_file.seek(0, 2) + if lock_file.tell() == 0: + lock_file.write(b"0") + lock_file.flush() + lock_file.seek(0) + msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) + else: + import fcntl + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as exc: + lock_file.close() + raise RuntimeError("任务数据库已被其他服务使用，请使用单个 API worker") from exc + self.lock_file = lock_file + try: + with self.connect() as db: + db.execute("PRAGMA journal_mode=WAL") + db.execute("""CREATE TABLE IF NOT EXISTS tasks ( + task_id TEXT PRIMARY KEY, status TEXT NOT NULL, + created_at TEXT NOT NULL, completed_at TEXT, error TEXT, + request TEXT, input_dir TEXT, callback_status INTEGER + )""") + db.execute("CREATE INDEX IF NOT EXISTS tasks_status_created ON tasks(status, created_at)") + db.execute("""UPDATE tasks SET status='failed', completed_at=?, + error='服务中断，任务未完成，请重新提交' WHERE status='running'""", + (datetime.now().isoformat(),)) + self.stop_event.clear() + if hasattr(self.runner, 'start'): + self.runner.start() + for number in range(self.concurrency): + thread = threading.Thread(target=self.work, name=f"analysis-{number}") + self.threads.append(thread) + thread.start() + except BaseException: + self.stop() + raise + + def stop(self): + self.stop_event.set() + for thread in self.threads: + if thread.ident is not None: + thread.join() + self.threads.clear() + if hasattr(self.runner, 'stop'): + self.runner.stop() + if self.lock_file: + self.lock_file.close() + self.lock_file = None + + @staticmethod + def decode(row): + if row is None: + return None + task = dict(row) + task['request'] = json.loads(task['request']) if task['request'] else None + return task + + def submit(self, task_id, request, input_dir): + with self.connect() as db: + db.execute("BEGIN IMMEDIATE") + row = db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone() + if row and row['status'] in ('pending', 'running'): + return self.decode(row), False + pending = db.execute("SELECT COUNT(*) FROM tasks WHERE status='pending'").fetchone()[0] + if pending >= self.max_pending: + raise QueueFull("等待队列已满，请稍后重试") + db.execute("""INSERT INTO tasks VALUES (?, 'pending', ?, NULL, NULL, ?, ?, NULL) + ON CONFLICT(task_id) DO UPDATE SET status='pending', + created_at=excluded.created_at, completed_at=NULL, error=NULL, + request=excluded.request, input_dir=excluded.input_dir, callback_status=NULL""", + (task_id, datetime.now().isoformat(), json.dumps(request, ensure_ascii=False), input_dir)) + return self.decode(db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()), True + + def claim(self): + with self.connect() as db: + db.execute("BEGIN IMMEDIATE") + row = db.execute("SELECT * FROM tasks WHERE status='pending' ORDER BY created_at, rowid LIMIT 1").fetchone() + if row is None: + return None + db.execute("UPDATE tasks SET status='running' WHERE task_id=?", (row['task_id'],)) + task = self.decode(row) + task['status'] = 'running' + return task + + def finish(self, task_id, error=None): + with self.connect() as db: + # A callback cannot release a slot before the subprocess exits. + # Preserve a failure callback even when the script exits with zero. + db.execute("""UPDATE tasks SET + status=CASE WHEN ? IS NOT NULL OR callback_status=0 THEN 'failed' ELSE 'completed' END, + error=CASE WHEN ? IS NOT NULL THEN ? WHEN callback_status=0 THEN '分析任务回调失败' ELSE NULL END, + completed_at=? WHERE task_id=? AND status='running'""", + (error, error, error, datetime.now().isoformat(), task_id)) + + def work(self): + while not self.stop_event.is_set(): + try: + task = self.claim() + if task is None: + self.stop_event.wait(0.25) + continue + error = None + try: + self.runner(task, self.stop_event) + except Exception as exc: + error = str(exc) or type(exc).__name__ + logger.exception("任务失败：%s", task['task_id']) + self.finish(task['task_id'], error) + except Exception: + logger.exception("任务队列异常") + self.stop_event.wait(1) + + def get(self, task_id): + with self.connect() as db: + return self.decode(db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone()) + + def list(self, status=None, limit=20): + with self.connect() as db: + if status: + rows = db.execute("SELECT * FROM tasks WHERE status=? ORDER BY created_at DESC LIMIT ?", (status, limit)).fetchall() + else: + rows = db.execute("SELECT * FROM tasks ORDER BY created_at DESC LIMIT ?", (limit,)).fetchall() + return [self.decode(row) for row in rows] + + def counts(self): + counts = dict.fromkeys(('pending', 'running', 'completed', 'failed'), 0) + with self.connect() as db: + counts.update(dict(db.execute("SELECT status, COUNT(*) FROM tasks GROUP BY status").fetchall())) + counts['total'] = sum(counts.values()) + return counts + + def callback(self, task_id, status): + with self.connect() as db: + db.execute("BEGIN IMMEDIATE") + row = db.execute("SELECT * FROM tasks WHERE task_id=?", (task_id,)).fetchone() + if row is None: + return None + if row['status'] == 'running': + db.execute("UPDATE tasks SET callback_status=? WHERE task_id=?", (status, task_id)) + # Late callbacks must not rewrite terminal state or consume pending work. + return row['status'] diff --git a/wanzheng.py b/wanzheng.py new file mode 100644 index 0000000..4eb2d19 --- /dev/null +++ b/wanzheng.py @@ -0,0 +1,653 @@ +#!/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 +import threading +import time + +from concurrent.futures import ThreadPoolExecutor, as_completed +from tqdm import tqdm +import logging +from typing import Dict, List, Optional, Tuple +import requests + +from config.settings import API_HOST, API_CALLBACK + +# ==================== 日志配置 ==================== +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: + # 直接全图检测 + 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 + +class SafeYOLO: + """One initialized model; serialize prediction while OpenCV runs in parallel.""" + def __init__(self, model_path): + self._lock = threading.Lock() + started = time.perf_counter() + self._model = YOLO(str(model_path)) + self.names = self._model.names + # Complete predictor creation, layer fusion and device setup before threads. + self(np.zeros((640, 640, 3), dtype=np.uint8), conf=0.2, verbose=False) + logger.info("模型加载及预热完成：%.3f 秒，设备：%s", + time.perf_counter() - started, self._model.predictor.device) + + def __call__(self, *args, **kwargs): + with self._lock: + return self._model(*args, **kwargs) + + +def analyze_directory(input_dir, task_id, model, workers=4, api_url=None): + """Reuse a warmed model. Return structured status to the resident parent.""" + started = time.perf_counter() + input_dir = Path(input_dir) + results = [] + error = None + try: + if not input_dir.is_dir(): + raise RuntimeError(f"输入目录不存在：{input_dir}") + image_files = sorted(p for p in input_dir.iterdir() if p.is_file() + and p.suffix.lower() in {'.jpg', '.jpeg', '.png', '.bmp', '.webp'}) + if not image_files: + raise RuntimeError("输入目录没有可处理的图片") + output_base = input_dir.parent / 'output' + for subdir in ['yolo', 'redzone', 'face_red', 'sensitivity', 'pigment', 'texts']: + (output_base / subdir).mkdir(parents=True, exist_ok=True) + logger.info("任务 %s：%d 张图片，%d 个图片线程，复用已预热模型", task_id, len(image_files), workers) + with ThreadPoolExecutor(max_workers=workers) as executor: + futures = [executor.submit(process_single_image, (p, model, output_base, False)) + for p in image_files] + for future in as_completed(futures): + results.append(future.result()) + if not save_text_report(results, output_base / 'texts'): + raise RuntimeError("文本报告保存失败") + failed = [r for r in results if r['status'] != 'success'] + if failed: + error = "; ".join(f"{r['filename']}: {r.get('errors') or r.get('error') or r['status']}" + for r in failed) + except Exception as exc: + logger.exception("分析任务失败：%s", task_id) + error = str(exc) or type(exc).__name__ + analysis_seconds = time.perf_counter() - started + logger.info("任务 %s 分析结束：%.3f 秒，成功 %d，未完整成功 %d", task_id, analysis_seconds, + sum(r['status'] == 'success' for r in results), + sum(r['status'] != 'success' for r in results)) + callback_ok = None + if api_url: + callback_ok = callback_api(task_id, 0 if error else 1, api_url) + return {'task_id': task_id, 'error': error, 'images': results, + 'analysis_seconds': analysis_seconds, + 'total_seconds': time.perf_counter() - started, 'callback_ok': callback_ok} + + +def main(): + parser = argparse.ArgumentParser(description='皮肤分析工具') + parser.add_argument('input_dir') + parser.add_argument('--task-id') + parser.add_argument('--model', '-m', default='best.pt') + parser.add_argument('--workers', '-w', type=int, default=4) + parser.add_argument('--verbose', '-v', action='store_true') + parser.add_argument('--api-url', default=API_HOST + API_CALLBACK) + args = parser.parse_args() + setup_logging(args.verbose) + if args.workers < 1: + parser.error('--workers 必须大于 0') + try: + model = SafeYOLO(args.model) + except Exception: + logger.exception("模型加载或预热失败") + callback_api(args.task_id, 0, args.api_url) + return 1 + result = analyze_directory(args.input_dir, args.task_id, model, args.workers, args.api_url) + return 1 if result['error'] else 0 + + +if __name__ == "__main__": + raise SystemExit(main())