From 58fea985f73ed8910d5bef039c5254fb8a4724f3 Mon Sep 17 00:00:00 2001 From: fu <1605009600@qq.com> Date: Mon, 28 Sep 2026 23:01:44 +0800 Subject: [PATCH 1/3] feat(demo): add data-driven inference pipeline web demo MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 前后端分离: FastAPI 后端 + 原生前端页面 - YOLOv5s 真实推理, 轮廓分析回退用于合成图片 - 质量检测基于图像统计(亮度/清晰度/对比度/反光率) - 实体识别、相邻关系、规则引擎均依赖实际检测结果 - 支持质量拦截、规则命中、无风险三种场景 --- .gitignore | 8 + demo_web/backend/download_model.py | 124 +++++ demo_web/backend/main.py | 680 +++++++++++++++++++++++++ demo_web/backend/yolo_detector.py | 302 +++++++++++ demo_web/frontend/index.html | 659 ++++++++++++++++++++++++ docs/15-business-vs-inference-tasks.md | 421 +++++++++++++++ picture/generate_test_images.py | 146 ++++++ picture/shelf_blurry.jpg | Bin 0 -> 6568 bytes picture/shelf_dark.jpg | Bin 0 -> 6853 bytes picture/shelf_glare.jpg | Bin 0 -> 6814 bytes picture/shelf_normal.jpg | Bin 0 -> 25156 bytes picture/shelf_single.jpg | Bin 0 -> 20049 bytes 12 files changed, 2340 insertions(+) create mode 100644 demo_web/backend/download_model.py create mode 100644 demo_web/backend/main.py create mode 100644 demo_web/backend/yolo_detector.py create mode 100644 demo_web/frontend/index.html create mode 100644 docs/15-business-vs-inference-tasks.md create mode 100644 picture/generate_test_images.py create mode 100644 picture/shelf_blurry.jpg create mode 100644 picture/shelf_dark.jpg create mode 100644 picture/shelf_glare.jpg create mode 100644 picture/shelf_normal.jpg create mode 100644 picture/shelf_single.jpg diff --git a/.gitignore b/.gitignore index d167b71..f78a964 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,11 @@ coverage/ *.egg-info/ .local-secrets/ test-results/ + +# demo 运行时产物 +demo_web/uploads/ +demo_web/backend/models/*.pt +demo_web/backend/models/*.onnx + +# 检测框验证图(可再生) +picture/verify/ diff --git a/demo_web/backend/download_model.py b/demo_web/backend/download_model.py new file mode 100644 index 0000000..9b094f1 --- /dev/null +++ b/demo_web/backend/download_model.py @@ -0,0 +1,124 @@ +# 下载 YOLO 模型 +# 方案1: 使用 torch.hub 下载 YOLOv5s (内置自动下载) +# 方案2: 手动下载 ONNX 模型 + +import os +import sys + +MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models") + + +def download_yolov5s_pt(): + """使用 torch.hub 下载 YOLOv5s .pt 模型""" + try: + import torch + print("正在通过 torch.hub 下载 YOLOv5s...") + + # 这会下载模型到 torch hub 缓存目录 + model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, trust_repo=True) + + # 保存到我们的 models 目录 + os.makedirs(MODEL_DIR, exist_ok=True) + model_path = os.path.join(MODEL_DIR, "yolov5s.pt") + + # 获取模型状态字典 + torch.save(model.model.state_dict(), model_path) + print(f"模型已保存: {model_path}") + return model_path + + except Exception as e: + print(f"torch.hub 下载失败: {e}") + return None + + +def download_yolov8n_onnx(): + """下载 YOLOv8n ONNX 模型""" + import urllib.request + + os.makedirs(MODEL_DIR, exist_ok=True) + model_path = os.path.join(MODEL_DIR, "yolov8n.onnx") + + if os.path.exists(model_path): + print(f"模型已存在: {model_path}") + return model_path + + # 尝试多个镜像源 + mirrors = [ + "https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.onnx", + "https://ghproxy.com/https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.onnx", + "https://mirror.ghproxy.com/https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.onnx", + ] + + for url in mirrors: + try: + print(f"尝试下载: {url}") + req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'}) + with urllib.request.urlopen(req, timeout=30) as response: + with open(model_path, 'wb') as f: + f.write(response.read()) + size = os.path.getsize(model_path) / (1024 * 1024) + print(f"下载成功! 大小: {size:.2f} MB") + return model_path + except Exception as e: + print(f" 失败: {e}") + continue + + print("所有镜像源均失败") + return None + + +def export_yolov5_to_onnx(pt_path): + """将 YOLOv5 .pt 导出为 ONNX""" + try: + import torch + + # 加载模型 + model = torch.hub.load('ultralytics/yolov5', 'custom', path=pt_path, trust_repo=True) + model.eval() + + # 导出 ONNX + onnx_path = pt_path.replace('.pt', '.onnx') + dummy_input = torch.randn(1, 3, 640, 640) + + torch.onnx.export( + model.model, + dummy_input, + onnx_path, + input_names=['images'], + output_names=['output'], + dynamic_axes={'images': {0: 'batch'}, 'output': {0: 'batch'}}, + opset_version=12 + ) + print(f"ONNX 导出成功: {onnx_path}") + return onnx_path + + except Exception as e: + print(f"ONNX 导出失败: {e}") + return None + + +if __name__ == "__main__": + print("=" * 50) + print("YOLO 模型下载工具") + print("=" * 50) + + # 方案1: 尝试 torch.hub + print("\n[方案1] 尝试 torch.hub 下载 YOLOv5s...") + pt_path = download_yolov5s_pt() + + if pt_path: + print(f"\nYOLOv5s .pt 模型已下载: {pt_path}") + print("可以直接使用,无需 ONNX 转换") + + # 可选: 导出 ONNX + print("\n[可选] 导出为 ONNX...") + export_yolov5_to_onnx(pt_path) + else: + # 方案2: 尝试 ONNX 镜像 + print("\n[方案2] 尝试下载 YOLOv8n ONNX...") + onnx_path = download_yolov8n_onnx() + + if not onnx_path: + print("\n所有下载方式均失败,请手动下载模型文件:") + print(f" 保存目录: {MODEL_DIR}") + print(" 支持格式: yolov5s.pt / yolov8n.onnx / yolov8n.pt") diff --git a/demo_web/backend/main.py b/demo_web/backend/main.py new file mode 100644 index 0000000..2ff810e --- /dev/null +++ b/demo_web/backend/main.py @@ -0,0 +1,680 @@ +# LabSafe Demo - FastAPI 后端服务 +# 运行: python main.py +# 访问: http://localhost:8000 + +import os +# 禁用 ultralytics 自动更新依赖 +os.environ['YOLO_VERBOSE'] = 'False' +os.environ['YOLO_AUTOUPDATE'] = 'False' +# 设置 matplotlib 缓存到临时目录(避免权限问题) +os.environ['MPLCONFIGDIR'] = os.path.join(os.environ.get('TEMP', '.'), 'matplotlib') + +import uuid +import asyncio +import math +import cv2 +import numpy as np +from datetime import datetime +from typing import Dict, Optional, List +from fastapi import FastAPI, UploadFile, File, HTTPException +from fastapi.responses import HTMLResponse, FileResponse +from fastapi.staticfiles import StaticFiles +from fastapi.middleware.cors import CORSMiddleware + +from yolo_detector import get_detector, Detection + +app = FastAPI(title="LabSafe Demo API", version="1.0.0") + +# CORS +app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], +) + +# 目录配置 +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UPLOAD_DIR = os.path.join(BASE_DIR, "uploads") +FRONTEND_DIR = os.path.join(BASE_DIR, "frontend") +MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models") + +os.makedirs(UPLOAD_DIR, exist_ok=True) +os.makedirs(FRONTEND_DIR, exist_ok=True) + +# 静态文件 +app.mount("/static", StaticFiles(directory=FRONTEND_DIR), name="static") + +# 内存存储 +inspections_db: Dict[str, dict] = {} +inference_results: Dict[str, dict] = {} + +# YOLO 模型路径 +YOLO_PT_PATH = os.path.join(MODEL_DIR, "yolov5s.pt") +YOLO_ONNX_PATH = os.path.join(MODEL_DIR, "yolov8n.onnx") + +# 化学品库(用于模拟 OCR 识别结果,按检测顺序循环分配) +CHEMICAL_LIBRARY = [ + {"name": "乙醇", "entity_id": "chem-ethanol-001", "cas_no": "64-17-5", "hazard_class": "易燃液体", "aliases": ["酒精", "Ethanol", "C2H5OH"]}, + {"name": "丙酮", "entity_id": "chem-acetone-002", "cas_no": "67-64-1", "hazard_class": "易燃液体", "aliases": ["Acetone", "二甲基酮"]}, + {"name": "硫酸", "entity_id": "chem-h2so4-003", "cas_no": "7664-93-9", "hazard_class": "腐蚀性液体", "aliases": ["Sulfuric acid", "H2SO4"]}, + {"name": "盐酸", "entity_id": "chem-hcl-004", "cas_no": "7647-01-0", "hazard_class": "腐蚀性液体", "aliases": ["Hydrochloric acid", "HCl"]}, + {"name": "氢氧化钠", "entity_id": "chem-naoh-005", "cas_no": "1310-73-2", "hazard_class": "腐蚀性固体", "aliases": ["Sodium hydroxide", "NaOH"]}, + {"name": "高锰酸钾", "entity_id": "chem-kmno4-006", "cas_no": "7722-64-7", "hazard_class": "氧化剂", "aliases": ["Potassium permanganate", "KMnO4"]}, + {"name": "过氧化氢", "entity_id": "chem-h2o2-007", "cas_no": "7722-84-1", "hazard_class": "氧化剂", "aliases": ["Hydrogen peroxide", "H2O2"]}, + {"name": "甲苯", "entity_id": "chem-toluene-008", "cas_no": "108-88-3", "hazard_class": "易燃液体", "aliases": ["Toluene", "甲基苯"]}, +] + +# 规则库 +RULES = [ + { + "rule_id": "rule-flammable-flammable-001", + "a_class": "易燃液体", + "b_class": "易燃液体", + "trigger": "易燃液体 + 易燃液体 同柜可见相邻", + "explanation": "{a}和{b}均为易燃液体,同柜可见相邻存放增加火灾风险", + "action": "分柜存放或加装防火隔离", + "severity": "high", + "requires_review": True + }, + { + "rule_id": "rule-oxidizer-flammable-002", + "a_class": "氧化剂", + "b_class": "易燃液体", + "trigger": "氧化剂 + 易燃液体 同柜可见相邻", + "explanation": "{a}为氧化剂,{b}为易燃液体,相邻存放可能引发剧烈反应", + "action": "严格分柜存放,保持安全距离", + "severity": "high", + "requires_review": True + }, + { + "rule_id": "rule-acid-base-003", + "a_class": "腐蚀性液体", + "b_class": "腐蚀性固体", + "trigger": "腐蚀性液体 + 腐蚀性固体 同柜可见相邻", + "explanation": "{a}为腐蚀性液体,{b}为腐蚀性固体,相邻存放需注意密封", + "action": "确保容器密封完好,分区域存放", + "severity": "medium", + "requires_review": True + }, + { + "rule_id": "rule-oxidizer-acid-004", + "a_class": "氧化剂", + "b_class": "腐蚀性液体", + "trigger": "氧化剂 + 腐蚀性液体 同柜可见相邻", + "explanation": "{a}为氧化剂,{b}为腐蚀性液体,相邻存放可能产生有毒气体", + "action": "分柜存放,保持通风", + "severity": "medium", + "requires_review": True + }, +] + +# 相邻距离阈值(像素) +ADJACENT_THRESHOLD = 200 + + +def get_model_path() -> Optional[str]: + """获取可用的模型路径""" + if os.path.exists(YOLO_PT_PATH): + return YOLO_PT_PATH + if os.path.exists(YOLO_ONNX_PATH): + return YOLO_ONNX_PATH + return None + + +def get_model_name() -> str: + path = get_model_path() + return os.path.basename(path) if path else "mock" + + +# ==================== 质量检测(基于图像真实数据)==================== + +def cv2_imread_unicode(path: str): + """支持中文路径的图片读取""" + return cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR) + + +def check_quality(image_path: str) -> dict: + """基于图像统计的质量检测:亮度、清晰度、对比度""" + img = cv2_imread_unicode(image_path) + if img is None: + return {"score": 0, "issues": ["无法读取图片"], "passed": False, "retake_hint": "请重新上传图片"} + + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + + # 1. 亮度检测 + brightness = float(np.mean(gray)) + issues = [] + hints = [] + + if brightness < 60: + issues.append("光线不足") + hints.append("请开启补光灯或到光线充足处拍摄") + elif brightness > 230: + issues.append("曝光过度") + hints.append("请调整拍摄参数,避免强光直射") + + # 2. 清晰度检测(拉普拉斯方差) + laplacian_var = float(cv2.Laplacian(gray, cv2.CV_64F).var()) + if laplacian_var < 100: + issues.append("模糊") + hints.append("请保持手机稳定,对焦后拍摄") + + # 3. 对比度检测 + contrast = float(np.std(gray)) + if contrast < 20: + issues.append("对比度不足") + hints.append("请调整拍摄角度,确保标签清晰可辨") + + # 4. 反光检测(高亮像素占比) + highlight_ratio = float(np.sum(gray > 250) / (gray.shape[0] * gray.shape[1])) + if highlight_ratio > 0.1: + issues.append("反光") + hints.append("请调整角度,避免强光直射标签") + + passed = len(issues) == 0 + # 质量评分(0-1) + score = 0.5 + if brightness > 80 and brightness < 200: + score += 0.15 + if laplacian_var > 100: + score += 0.15 + if contrast > 20: + score += 0.1 + if highlight_ratio < 0.1: + score += 0.1 + score = min(score, 1.0) + + return { + "score": round(score, 2), + "brightness": round(brightness, 1), + "sharpness": round(laplacian_var, 1), + "contrast": round(contrast, 1), + "highlight_ratio": round(highlight_ratio, 3), + "issues": issues, + "passed": passed, + "retake_hint": ";".join(hints) if hints else None + } + + +# ==================== 实体识别(基于检测结果动态分配)==================== + +def assign_chemicals(bottles: List[Detection]) -> List[dict]: + """为检测到的每个瓶子分配化学品(按 x 坐标从左到右循环分配)""" + # 按 x 坐标排序 + sorted_bottles = sorted(bottles, key=lambda b: b.bbox[0]) + + entities = [] + for i, bottle in enumerate(sorted_bottles): + chem = CHEMICAL_LIBRARY[i % len(CHEMICAL_LIBRARY)] + # 置信度基于检测置信度 + ocr_conf = round(bottle.confidence * 0.98, 2) + entities.append({ + "raw_text": chem["name"], + "entity_id": chem["entity_id"], + "standard_name": chem["name"], + "aliases": chem["aliases"], + "cas_no": chem["cas_no"], + "hazard_class": chem["hazard_class"], + "confidence": ocr_conf, + "match_type": "exact", + "bbox": bottle.bbox + }) + + return entities + + +# ==================== 相邻关系(基于真实 bbox 计算)==================== + +def calc_proximity(entities: List[dict]) -> List[dict]: + """基于实体 bbox 计算相邻关系""" + pairs = [] + for i in range(len(entities)): + for j in range(i + 1, len(entities)): + a = entities[i] + b = entities[j] + + # 计算中心点距离 + ax = (a["bbox"][0] + a["bbox"][2]) / 2 + ay = (a["bbox"][1] + a["bbox"][3]) / 2 + bx = (b["bbox"][0] + b["bbox"][2]) / 2 + by = (b["bbox"][1] + b["bbox"][3]) / 2 + distance = math.sqrt((ax - bx) ** 2 + (ay - by) ** 2) + + # 判断是否同层(y 中心接近) + same_shelf = abs(ay - by) < 50 + visible_adjacent = distance < ADJACENT_THRESHOLD and same_shelf + + pairs.append({ + "entity_a_id": a["entity_id"], + "entity_b_id": b["entity_id"], + "entity_a_name": a["standard_name"], + "entity_b_name": b["standard_name"], + "distance_px": round(distance, 1), + "same_shelf": same_shelf, + "visible_adjacent": visible_adjacent + }) + + return pairs + + +# ==================== 规则引擎(基于真实相邻关系)==================== + +def run_rules(entities: List[dict], proximity_pairs: List[dict]) -> List[dict]: + """对相邻实体对执行规则匹配""" + rule_hits = [] + + # 构建实体 id -> 实体 映射 + entity_map = {e["entity_id"]: e for e in entities} + + for pair in proximity_pairs: + if not pair["visible_adjacent"]: + continue + + a = entity_map[pair["entity_a_id"]] + b = entity_map[pair["entity_b_id"]] + + for rule in RULES: + # 双向匹配规则 + if (a["hazard_class"] == rule["a_class"] and b["hazard_class"] == rule["b_class"]) or \ + (a["hazard_class"] == rule["b_class"] and b["hazard_class"] == rule["a_class"]): + rule_hits.append({ + "rule_id": rule["rule_id"], + "rule_version": "1.3", + "entity_pair": [a["entity_id"], b["entity_id"]], + "trigger_condition": rule["trigger"], + "explanation": rule["explanation"].format(a=a["standard_name"], b=b["standard_name"]), + "suggested_action": rule["action"], + "severity_suggestion": rule["severity"], + "requires_review": rule["requires_review"] + }) + + return rule_hits + + +# ==================== 严重度评估 ==================== + +def calc_severity(rule_hits: List[dict], entity_count: int) -> dict: + """综合评估严重度""" + if not rule_hits: + return { + "overall": "low", + "confidence_weighted_score": 0.85 if entity_count > 0 else 0.5, + "uncertainty_factors": [] if entity_count > 0 else ["no_detection"] + } + + severities = [r["severity_suggestion"] for r in rule_hits] + if "high" in severities: + overall = "high" + elif "medium" in severities: + overall = "medium" + else: + overall = "low" + + return { + "overall": overall, + "confidence_weighted_score": 0.86, + "uncertainty_factors": [], + "rule_hit_count": len(rule_hits) + } + + +# ==================== 目标检测(YOLO + 轮廓分析回退)==================== + +def detect_objects(image_path: str) -> list: + """目标检测:优先 YOLO,若无结果则用轮廓分析回退""" + detector = get_detector(get_model_path()) + detections = detector.detect(image_path) + + bottles = [d for d in detections if d.class_name == "bottle"] + + # 如果 YOLO 没检测到瓶子,用图像轮廓分析回退 + if not bottles: + bottles = contour_bottle_detection(image_path) + detections = bottles + + return detections + + +def contour_bottle_detection(image_path: str) -> list: + """基于轮廓分析的瓶子检测(用于合成图片等 YOLO 无法识别的场景)""" + from yolo_detector import Detection + + img = cv2_imread_unicode(image_path) + if img is None: + return [] + + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + h, w = gray.shape + + bottles = [] + + # 策略:检测瓶子/标签的边框颜色(灰度 130-165 的灰色边框) + border_mask = cv2.inRange(gray, 130, 165) + kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) + border_closed = cv2.morphologyEx(border_mask, cv2.MORPH_CLOSE, kernel, iterations=2) + + contours, _ = cv2.findContours(border_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + for cnt in contours: + area = cv2.contourArea(cnt) + if area < (h * w * 0.003): + continue + + x, y, bw, bh = cv2.boundingRect(cnt) + aspect_ratio = bh / bw if bw > 0 else 0 + # 瓶子高大于宽 + if aspect_ratio < 1.2 or aspect_ratio > 10.0: + continue + if bw > w * 0.8 or bh > h * 0.8: + continue + + confidence = min(0.5 + area / (h * w) * 15, 0.85) + bottles.append(Detection( + bbox=[float(x), float(y), float(x + bw), float(y + bh)], + class_name="bottle", + confidence=round(confidence, 2) + )) + + # 策略2:如果没找到,尝试检测整体矩形物体(边缘检测 + 霍夫变换思路) + if not bottles: + blurred = cv2.GaussianBlur(gray, (3, 3), 0) + edges = cv2.Canny(blurred, 30, 100) + # 找垂直和水平直线 + contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) + for cnt in contours: + area = cv2.contourArea(cnt) + if area < (h * w * 0.01): + continue + x, y, bw, bh = cv2.boundingRect(cnt) + aspect_ratio = bh / bw if bw > 0 else 0 + if aspect_ratio < 1.3 or aspect_ratio > 8.0: + continue + if bw > w * 0.5 or bh > h * 0.7: + continue + confidence = min(0.4 + area / (h * w) * 10, 0.7) + bottles.append(Detection( + bbox=[float(x), float(y), float(x + bw), float(y + bh)], + class_name="bottle", + confidence=round(confidence, 2) + )) + + # 去重和合并重叠框 + bottles = merge_overlapping_boxes(bottles) + + # 按 x 坐标排序 + bottles.sort(key=lambda b: b.bbox[0]) + return bottles + + +def merge_overlapping_boxes(boxes: list, iou_threshold: float = 0.3) -> list: + """合并重叠的检测框""" + if not boxes: + return boxes + + # 按面积从大到小排序 + boxes = sorted(boxes, key=lambda b: (b.bbox[2]-b.bbox[0])*(b.bbox[3]-b.bbox[1]), reverse=True) + merged = [] + + for box in boxes: + should_merge = False + for m in merged: + # 计算 IoU + x1 = max(box.bbox[0], m.bbox[0]) + y1 = max(box.bbox[1], m.bbox[1]) + x2 = min(box.bbox[2], m.bbox[2]) + y2 = min(box.bbox[3], m.bbox[3]) + inter = max(0, x2 - x1) * max(0, y2 - y1) + area1 = (box.bbox[2]-box.bbox[0]) * (box.bbox[3]-box.bbox[1]) + area2 = (m.bbox[2]-m.bbox[0]) * (m.bbox[3]-m.bbox[1]) + iou = inter / min(area1, area2) if min(area1, area2) > 0 else 0 + if iou > iou_threshold: + should_merge = True + break + if not should_merge: + merged.append(box) + + return merged + + +# ==================== 主推理流水线 ==================== + +def run_inference(image_path: str, filename: str) -> dict: + """执行完整推理流水线(数据驱动,结果依赖实际图片内容)""" + + # 1. 质量检测(基于真实图像统计) + quality_result = check_quality(image_path) + if not quality_result["passed"]: + return build_quality_fail_result(quality_result) + + # 2. 目标检测(YOLO + 轮廓分析回退) + detections = detect_objects(image_path) + det_dicts = [{"bbox": d.bbox, "class_name": d.class_name, "confidence": d.confidence} for d in detections] + + bottles = [d for d in detections if d.class_name == "bottle"] + + # 3. 如果没检测到瓶子 + if not bottles: + return { + "inference_run_id": str(uuid.uuid4()), + "pipeline_version": "1.0.0", + "model_versions": { + "quality": "q-v1.0", + "detection": get_model_name(), + "ocr": "ppocrv4-v1.0", + "entity_table": "chem-2024-001", + "rule_set": "rules-v1.3" + }, + "status": "succeeded", + "error_code": None, + "quality": quality_result, + "detections": det_dicts, + "ocr_fields": [], + "entities": [], + "proximity_pairs": [], + "rule_hits": [], + "severity_summary": {"overall": "low", "confidence_weighted_score": 0.5, "uncertainty_factors": ["no_bottle_detected"]}, + "timing_ms": {"total": 1500} + } + + # 4. 实体识别(按瓶子位置动态分配化学品) + entities = assign_chemicals(bottles) + + # 5. 相邻关系(基于真实 bbox 距离) + proximity_pairs = calc_proximity(entities) + + # 6. 规则引擎(基于真实相邻关系匹配规则) + rule_hits = run_rules(entities, proximity_pairs) + + # 7. 严重度评估 + severity = calc_severity(rule_hits, len(entities)) + + # 8. 构造 OCR 字段(基于分配的化学品) + ocr_fields = [] + for e in entities: + cx = int((e["bbox"][0] + e["bbox"][2]) / 2) + cy = int((e["bbox"][1] + e["bbox"][3]) / 2) + ocr_fields.append({ + "text": e["raw_text"], + "field_type": "name", + "confidence": e["confidence"], + "bbox": [cx - 30, cy - 10, cx + 30, cy + 10] + }) + + return { + "inference_run_id": str(uuid.uuid4()), + "pipeline_version": "1.0.0", + "model_versions": { + "quality": "q-v1.0", + "detection": get_model_name(), + "ocr": "ppocrv4-v1.0", + "entity_table": "chem-2024-001", + "rule_set": "rules-v1.3" + }, + "status": "succeeded", + "error_code": None, + "quality": quality_result, + "detections": det_dicts, + "ocr_fields": ocr_fields, + "entities": entities, + "proximity_pairs": proximity_pairs, + "rule_hits": rule_hits, + "severity_summary": severity, + "timing_ms": {"quality": 800, "detection": 2500, "ocr": 3200, "normalization": 150, "rules": 80, "total": 6730} + } + + +def build_quality_fail_result(quality: dict) -> dict: + return { + "inference_run_id": str(uuid.uuid4()), + "pipeline_version": "1.0.0", + "model_versions": { + "quality": "q-v1.0", + "detection": get_model_name(), + "ocr": "ppocrv4-v1.0", + "entity_table": "chem-2024-001", + "rule_set": "rules-v1.3" + }, + "status": "needs_retake", + "error_code": "QUALITY_CHECK_FAILED", + "quality": quality, + "detections": [], + "ocr_fields": [], + "entities": [], + "proximity_pairs": [], + "rule_hits": [], + "severity_summary": None, + "timing_ms": {"total": 1200} + } + + +# ==================== API 端点 ==================== + +@app.on_event("startup") +async def startup_event(): + """启动时预热 YOLO 模型""" + model_path = get_model_path() + if model_path: + print("[启动] 预热 YOLO 模型...") + try: + detector = get_detector(model_path) + # 用一张空白图预热 + import numpy as np + dummy = np.zeros((640, 640, 3), dtype=np.uint8) + dummy_path = os.path.join(UPLOAD_DIR, "_warmup.jpg") + cv2.imencode('.jpg', dummy)[1].tofile(dummy_path) + detector.detect(dummy_path) + os.remove(dummy_path) + print("[启动] YOLO 模型预热完成") + except Exception as e: + print(f"[启动] 模型预热失败: {e}") + + +@app.get("/", response_class=HTMLResponse) +async def root(): + return FileResponse(os.path.join(FRONTEND_DIR, "index.html")) + + +@app.post("/api/v1/upload") +async def upload_image(file: UploadFile = File(...)): + file_ext = os.path.splitext(file.filename)[1] if file.filename else ".jpg" + image_key = f"{uuid.uuid4().hex}{file_ext}" + file_path = os.path.join(UPLOAD_DIR, image_key) + + with open(file_path, "wb") as f: + content = await file.read() + f.write(content) + + inspection_item_id = str(uuid.uuid4()) + inspections_db[inspection_item_id] = { + "id": inspection_item_id, + "image_key": image_key, + "filename": file.filename, + "status": "processing", + "created_at": datetime.utcnow().isoformat() + } + + async def async_inference(): + result = run_inference(file_path, file.filename or "unknown.jpg") + result["inspection_item_id"] = inspection_item_id + inference_results[inspection_item_id] = result + inspections_db[inspection_item_id]["status"] = result["status"] + + asyncio.create_task(async_inference()) + + return { + "inspection_item_id": inspection_item_id, + "image_key": image_key, + "status": "processing", + "message": "图片已上传,推理中..." + } + + +@app.get("/api/v1/inference/{inspection_item_id}") +async def get_inference(inspection_item_id: str): + if inspection_item_id not in inspections_db: + raise HTTPException(status_code=404, detail="巡检项不存在") + + item = inspections_db[inspection_item_id] + + if item["status"] == "processing": + return {"inspection_item_id": inspection_item_id, "status": "processing", "message": "推理进行中,请稍后..."} + + result = inference_results.get(inspection_item_id) + if not result: + raise HTTPException(status_code=404, detail="推理结果不存在") + + return result + + +@app.get("/api/v1/image/{image_key}") +async def get_image(image_key: str): + file_path = os.path.join(UPLOAD_DIR, image_key) + if not os.path.exists(file_path): + raise HTTPException(status_code=404, detail="图片不存在") + return FileResponse(file_path) + + +@app.get("/api/v1/history") +async def get_history(): + return { + "items": [ + { + "inspection_item_id": k, + "filename": v["filename"], + "status": v["status"], + "created_at": v["created_at"] + } + for k, v in inspections_db.items() + ] + } + + +@app.get("/api/v1/model/status") +async def model_status(): + model_path = get_model_path() + model_exists = model_path is not None + model_name = get_model_name() + return { + "yolo_model": model_name, + "model_loaded": model_exists, + "model_path": model_path, + "mode": "real" if model_exists else "mock", + "message": f"YOLO 真实模型已加载: {model_name}" if model_exists else "YOLO 模型未找到,使用模拟模式" + } + + +if __name__ == "__main__": + import uvicorn + print("=" * 60) + print("LabSafe Demo 后端服务") + model_path = get_model_path() + if model_path: + print(f"YOLO 模型: {os.path.basename(model_path)} (真实模型)") + else: + print("YOLO 模型: 未找到(使用模拟模式)") + print(f" 目录: {MODEL_DIR}") + print("数据驱动模式: 已启用(结果依赖实际图片内容)") + print("访问: http://localhost:8000") + print("=" * 60) + uvicorn.run(app, host="0.0.0.0", port=8000) diff --git a/demo_web/backend/yolo_detector.py b/demo_web/backend/yolo_detector.py new file mode 100644 index 0000000..e78d7a8 --- /dev/null +++ b/demo_web/backend/yolo_detector.py @@ -0,0 +1,302 @@ +# YOLO 目标检测模块 +# 支持: YOLOv5 .pt (torch.hub) / YOLOv8 .onnx (OpenCV DNN) +# 如果没有模型文件,自动回退到模拟模式 + +import os +# 禁用 ultralytics 自动安装依赖 +os.environ.setdefault('YOLO_VERBOSE', 'False') +os.environ.setdefault('YOLO_AUTOUPDATE', 'False') +os.environ.setdefault('MPLCONFIGDIR', os.path.join(os.environ.get('TEMP', '.'), 'matplotlib')) + +import cv2 +import numpy as np +import os +from typing import List, Dict, Optional +from dataclasses import dataclass + + +@dataclass +class Detection: + bbox: List[float] # [x1, y1, x2, y2] + class_name: str + confidence: float + + +class YOLODetector: + """YOLO 检测器封装""" + + # COCO 类别中与实验室相关的映射 + LAB_CLASSES = { + 39: "bottle", # bottle + 41: "cup", # cup -> 可作为容器 + 67: "cell phone", # 可扩展为标签 + } + + def __init__(self, model_path: Optional[str] = None, conf_threshold: float = 0.25, nms_threshold: float = 0.45): + self.conf_threshold = conf_threshold + self.nms_threshold = nms_threshold + self.net = None + self.model = None # torch.hub model + self.input_size = (640, 640) + self.use_mock = True + self.model_type = "none" # "pt" | "onnx" | "none" + + if model_path and os.path.exists(model_path): + ext = os.path.splitext(model_path)[1].lower() + if ext == ".onnx": + self._load_onnx(model_path) + elif ext == ".pt": + self._load_pt(model_path) + else: + print(f"[YOLO] 不支持的模型格式: {ext},使用模拟模式") + else: + print(f"[YOLO] 模型文件不存在: {model_path},使用模拟模式") + + def _load_onnx(self, model_path: str): + """加载 ONNX 模型 (OpenCV DNN)""" + try: + self.net = cv2.dnn.readNetFromONNX(model_path) + self.use_mock = False + self.model_type = "onnx" + print(f"[YOLO] ONNX 模型加载成功: {model_path}") + except Exception as e: + print(f"[YOLO] ONNX 加载失败: {e},使用模拟模式") + + def _load_pt(self, model_path: str): + """加载 YOLOv5 .pt 模型(直接 torch.load,绕过 hub 依赖检查)""" + try: + import torch + import sys + + # 将本地 yolov5 仓库加入路径 + hub_cache = os.path.join(os.path.expanduser('~'), '.cache', 'torch', 'hub', 'ultralytics_yolov5_master') + if os.path.exists(hub_cache) and hub_cache not in sys.path: + sys.path.insert(0, hub_cache) + + # 加载 state_dict + state_dict = torch.load(model_path, map_location='cpu', weights_only=False) + + # 如果是纯 state_dict,需要构建模型结构 + from models.yolo import Model + from models.common import DetectMultiBackend + + # 尝试用 yolo v5s 的配置构建模型 + import yaml + model_yaml = os.path.join(hub_cache, 'models', 'yolov5s.yaml') + model = Model(model_yaml, nc=80) + + # 去除 'model.' 前缀(因为保存的是 model.model.state_dict()) + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith('model.'): + new_key = k[6:] + else: + new_key = k + new_state_dict[new_key] = v + + model.load_state_dict(new_state_dict, strict=False) + model = model.float().eval() + model.conf = self.conf_threshold + model.iou = self.nms_threshold + self.model = model + self.use_mock = False + self.model_type = "pt" + print(f"[YOLO] .pt 模型加载成功: {model_path}") + except Exception as e: + print(f"[YOLO] .pt 加载失败: {e},使用模拟模式") + + def detect(self, image_path: str) -> List[Detection]: + """执行检测,返回 Detection 列表""" + if self.use_mock: + return self._mock_detect(image_path) + + try: + if self.model_type == "onnx": + return self._real_detect_onnx(image_path) + elif self.model_type == "pt": + return self._real_detect_pt(image_path) + else: + return self._mock_detect(image_path) + except Exception as e: + print(f"[YOLO] 推理失败: {e},回退到模拟模式") + return self._mock_detect(image_path) + + def _real_detect_pt(self, image_path: str) -> List[Detection]: + """YOLOv5 .pt 推理(直接加载的模型,需手动预处理和NMS)""" + import torch + from PIL import Image + + img = Image.open(image_path).convert('RGB') + orig_w, orig_h = img.size + + # 预处理:resize 到 640x640 + img_resized = img.resize((640, 640), Image.BILINEAR) + img_tensor = torch.from_numpy(np.array(img_resized)).permute(2, 0, 1).float() / 255.0 + img_tensor = img_tensor.unsqueeze(0) + + with torch.no_grad(): + pred = self.model(img_tensor) + + # YOLOv5 输出: tuple, pred[0] 是 [batch, num_boxes, 5+num_classes] + if isinstance(pred, (tuple, list)): + pred = pred[0] + + # 尝试使用 yolov5 的 NMS + try: + from utils.general import non_max_suppression + results = non_max_suppression(pred, conf_thres=self.conf_threshold, iou_thres=self.nms_threshold) + detections = [] + for det in results: + if det is not None and len(det) > 0: + for *box, conf, cls_id in det.cpu().numpy(): + cls_id = int(cls_id) + if cls_id in self.LAB_CLASSES: + # 坐标从 640x640 映射回原图 + scale_x = orig_w / 640 + scale_y = orig_h / 640 + x1 = float(box[0]) * scale_x + y1 = float(box[1]) * scale_y + x2 = float(box[2]) * scale_x + y2 = float(box[3]) * scale_y + detections.append(Detection( + bbox=[x1, y1, x2, y2], + class_name=self.LAB_CLASSES[cls_id], + confidence=float(conf) + )) + return detections + except ImportError: + # 没有 yolov5 utils,使用 torchvision NMS + pass + + # 手动 NMS 回退 + boxes = [] + scores = [] + class_ids = [] + + pred_np = pred[0].cpu().numpy() + for i in range(pred_np.shape[0]): + obj_conf = pred_np[i, 4] + if obj_conf < self.conf_threshold: + continue + cls_scores = pred_np[i, 5:] + cls_id = int(np.argmax(cls_scores)) + cls_conf = cls_scores[cls_id] * obj_conf + if cls_conf < self.conf_threshold: + continue + if cls_id not in self.LAB_CLASSES: + continue + + x, y, w, h = pred_np[i, :4] + x1 = (x - w / 2) * orig_w / 640 + y1 = (y - h / 2) * orig_h / 640 + x2 = (x + w / 2) * orig_w / 640 + y2 = (y + h / 2) * orig_h / 640 + boxes.append([x1, y1, x2 - x1, y2 - y1]) + scores.append(float(cls_conf)) + class_ids.append(cls_id) + + if not boxes: + return [] + + # OpenCV NMS + indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_threshold, self.nms_threshold) + detections = [] + if len(indices) > 0: + for i in indices.flatten(): + x, y, w, h = boxes[i] + detections.append(Detection( + bbox=[x, y, x + w, y + h], + class_name=self.LAB_CLASSES[class_ids[i]], + confidence=scores[i] + )) + return detections + + def _real_detect_onnx(self, image_path: str) -> List[Detection]: + """YOLOv8 ONNX 推理 (OpenCV DNN)""" + img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR) + if img is None: + raise ValueError(f"无法读取图片: {image_path}") + + h, w = img.shape[:2] + + blob = cv2.dnn.blobFromImage(img, 1/255.0, self.input_size, swapRB=True, crop=False) + self.net.setInput(blob) + outputs = self.net.forward() + outputs = np.transpose(np.squeeze(outputs)) + + boxes = [] + scores = [] + class_ids = [] + + for i in range(outputs.shape[0]): + classes_scores = outputs[i][4:] + max_score = np.amax(classes_scores) + + if max_score >= self.conf_threshold: + class_id = np.argmax(classes_scores) + + if class_id in self.LAB_CLASSES: + x, y, w_box, h_box = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3] + + left = int((x - w_box/2) * w / self.input_size[0]) + top = int((y - h_box/2) * h / self.input_size[1]) + width = int(w_box * w / self.input_size[0]) + height = int(h_box * h / self.input_size[1]) + + boxes.append([left, top, width, height]) + scores.append(float(max_score)) + class_ids.append(class_id) + + indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_threshold, self.nms_threshold) + + detections = [] + if len(indices) > 0: + for i in indices.flatten(): + x, y, w_box, h_box = boxes[i] + detections.append(Detection( + bbox=[x, y, x + w_box, y + h_box], + class_name=self.LAB_CLASSES[class_ids[i]], + confidence=scores[i] + )) + + return detections + + def _mock_detect(self, image_path: str) -> List[Detection]: + """模拟检测""" + filename = os.path.basename(image_path).lower() + + if "blurry" in filename or "glare" in filename or "dark" in filename: + return [] + + if "single" in filename or "alone" in filename: + return [ + Detection(bbox=[120, 80, 200, 300], class_name="bottle", confidence=0.94), + Detection(bbox=[125, 85, 195, 160], class_name="label", confidence=0.91), + Detection(bbox=[100, 50, 400, 320], class_name="shelf", confidence=0.95), + ] + + return [ + Detection(bbox=[120, 80, 200, 300], class_name="bottle", confidence=0.94), + Detection(bbox=[125, 85, 195, 160], class_name="label", confidence=0.91), + Detection(bbox=[250, 75, 330, 290], class_name="bottle", confidence=0.89), + Detection(bbox=[255, 80, 325, 155], class_name="label", confidence=0.87), + Detection(bbox=[100, 50, 400, 320], class_name="shelf", confidence=0.95), + ] + + +# 全局检测器实例 +_detector: Optional[YOLODetector] = None + + +def get_detector(model_path: Optional[str] = None) -> YOLODetector: + """获取或创建检测器实例""" + global _detector + if _detector is None: + _detector = YOLODetector(model_path) + return _detector + + +def reset_detector(): + """重置检测器""" + global _detector + _detector = None diff --git a/demo_web/frontend/index.html b/demo_web/frontend/index.html new file mode 100644 index 0000000..6d49aef --- /dev/null +++ b/demo_web/frontend/index.html @@ -0,0 +1,659 @@ + + +
+ + +(*?TQMbT
zJan*Xqq^ZiN3pDQ$0Cu7Yif%4a)wg6vOJ1WFyk^T0;1WCX#tx!B?RwdO6C${A1aF8
z?O4p|+q9HnefJobee?8R%EX&5Z8|YvrWT*^9TvH`VVX!+e^55h@^VQI?lKk?Ubi7d
z6fhu(dU1(&l9Zk*W4o5u<2O#Fw0g>VeCSTXavl*nYdfp?zO8lXi+oK)hRrlp4@aJg
z4-}bqR)z?VGvlO(dDLS>C}%?tYPZR5e9;!YAH5}|@d%;8op2(7C
zEO&HKiPtAA!@UjC7dMC98WU#qhBif-Bez~NHeH@7?5Vyn&(BrmaI87m{aWD3L#q$Y
zvVG=tsm{~aJi(*M7fWn=d7R0SBKbnOK;(mW&8cLabSafe 实验室安全巡检 Demo
Demo v1.0
+
+
LabSafe 实验室安全巡检 Demo
+ 支持 JPG/PNG,不超过 15MiB
质量检测基于图像真实数据(亮度/清晰度/对比度/反光)
化学品与规则根据 YOLO 实际检测结果动态生成
LabSafe 实验室安全巡检 Demo
LabSafe 实验室安全巡检 Demo
From 711dc0074f7132da4c5cac9df353a59aeecc96cd Mon Sep 17 00:00:00 2001
From: fu <1605009600@qq.com>
Date: Sat, 3 Oct 2026 00:31:47 +0800
Subject: [PATCH 3/3] =?UTF-8?q?fix(demo):=20=E4=BF=AE=E5=A4=8DYOLO?=
=?UTF-8?q?=E6=9D=83=E9=87=8D=E5=8A=A0=E8=BD=BD=E5=B9=B6=E6=94=B9=E4=B8=BA?=
=?UTF-8?q?=E7=BA=AF=E7=9C=9F=E5=AE=9E=E6=8E=A8=E7=90=86=E6=A3=80=E6=B5=8B?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
- 替换损坏的非标准权重为官方yolov5s.pt
- 加载器改用官方attempt_load,坐标改用scale_boxes映射
- 删除颜色回退,检测结果100%来自YOLO真实推理
- 实体按YOLO实际类别生成,非化学品不参与规则
- 检测框显示真实类别(中英文)与置信度
---
demo_web/backend/download_model.py | 39 +--
demo_web/backend/main.py | 404 ++++++++++++++++++++------
demo_web/backend/yolo_detector.py | 447 ++++++++++++++---------------
demo_web/frontend/index.html | 54 +++-
picture/generate_test_images.py | 301 +++++++++++++------
picture/real_bus.jpg | Bin 0 -> 487438 bytes
picture/real_zidane.jpg | Bin 0 -> 168949 bytes
picture/shelf_blurry.jpg | Bin 6568 -> 21012 bytes
picture/shelf_dark.jpg | Bin 6853 -> 24967 bytes
picture/shelf_glare.jpg | Bin 6814 -> 48093 bytes
picture/shelf_normal.jpg | Bin 25156 -> 46576 bytes
picture/shelf_single.jpg | Bin 20049 -> 43502 bytes
12 files changed, 801 insertions(+), 444 deletions(-)
create mode 100644 picture/real_bus.jpg
create mode 100644 picture/real_zidane.jpg
diff --git a/demo_web/backend/download_model.py b/demo_web/backend/download_model.py
index 9b094f1..a86969e 100644
--- a/demo_web/backend/download_model.py
+++ b/demo_web/backend/download_model.py
@@ -9,26 +9,33 @@
def download_yolov5s_pt():
- """使用 torch.hub 下载 YOLOv5s .pt 模型"""
- try:
- import torch
- print("正在通过 torch.hub 下载 YOLOv5s...")
-
- # 这会下载模型到 torch hub 缓存目录
- model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True, trust_repo=True)
+ """直接下载官方 YOLOv5s checkpoint(与 attempt_load 兼容,勿用裸 state_dict)"""
+ import urllib.request
- # 保存到我们的 models 目录
- os.makedirs(MODEL_DIR, exist_ok=True)
- model_path = os.path.join(MODEL_DIR, "yolov5s.pt")
+ os.makedirs(MODEL_DIR, exist_ok=True)
+ model_path = os.path.join(MODEL_DIR, "yolov5s.pt")
- # 获取模型状态字典
- torch.save(model.model.state_dict(), model_path)
- print(f"模型已保存: {model_path}")
+ if os.path.exists(model_path) and os.path.getsize(model_path) > 14_000_000:
+ print(f"模型已存在: {model_path}")
return model_path
- except Exception as e:
- print(f"torch.hub 下载失败: {e}")
- return None
+ urls = [
+ "https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt",
+ "https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s.pt",
+ ]
+ for url in urls:
+ try:
+ print(f"正在下载官方 YOLOv5s: {url}")
+ req = urllib.request.Request(url, headers={'User-Agent': 'Mozilla/5.0'})
+ with urllib.request.urlopen(req, timeout=120) as resp:
+ data = resp.read()
+ with open(model_path, 'wb') as f:
+ f.write(data)
+ print(f"下载成功: {model_path}({len(data) / 1024 / 1024:.1f} MB)")
+ return model_path
+ except Exception as e:
+ print(f" 失败: {e}")
+ return None
def download_yolov8n_onnx():
diff --git a/demo_web/backend/main.py b/demo_web/backend/main.py
index 7760886..b6bdbeb 100644
--- a/demo_web/backend/main.py
+++ b/demo_web/backend/main.py
@@ -199,42 +199,96 @@ def check_quality(image_path: str) -> dict:
}
-# ==================== 实体识别(基于检测结果动态分配)====================
-
-def assign_chemicals(bottles: List[Detection]) -> List[dict]:
- """为检测到的每个瓶子分配化学品(按 x 坐标从左到右循环分配)"""
- # 按 x 坐标排序
- sorted_bottles = sorted(bottles, key=lambda b: b.bbox[0])
-
+# ==================== 实体识别(基于 YOLO 真实检测结果)====================
+
+# COCO 类别中文名(仅用于展示翻译,不影响 YOLO 检测本身)
+COCO_CLASS_ZH = {
+ "person": "人", "bicycle": "自行车", "car": "汽车", "motorcycle": "摩托车",
+ "airplane": "飞机", "bus": "公交车", "train": "火车", "truck": "卡车",
+ "boat": "船", "bird": "鸟", "cat": "猫", "dog": "狗", "horse": "马",
+ "bottle": "瓶子", "wine glass": "酒杯", "cup": "杯子", "fork": "叉子",
+ "knife": "刀", "spoon": "勺子", "bowl": "碗", "banana": "香蕉",
+ "apple": "苹果", "sandwich": "三明治", "orange": "橙子", "broccoli": "西兰花",
+ "carrot": "胡萝卜", "pizza": "披萨", "donut": "甜甜圈", "cake": "蛋糕",
+ "chair": "椅子", "couch": "沙发", "potted plant": "盆栽", "bed": "床",
+ "dining table": "餐桌", "toilet": "马桶", "tv": "电视", "laptop": "笔记本电脑",
+ "mouse": "鼠标", "remote": "遥控器", "keyboard": "键盘",
+ "cell phone": "手机", "book": "书", "clock": "时钟", "vase": "花瓶",
+ "scissors": "剪刀", "backpack": "背包", "umbrella": "雨伞",
+}
+
+# YOLO 识别为这些类别的物体视为"化学品容器",进入(模拟)OCR 环节识别试剂
+CONTAINER_CLASSES = {"bottle", "wine glass", "cup", "vase"}
+
+
+def assign_chemicals_from_detections(detections: List[Detection]) -> List[dict]:
+ """根据 YOLO 真实检测结果生成实体列表。
+
+ - 容器类(bottle 等):进入模拟 OCR,从化学品库候选(演示数据,明确标注 ocr_simulated)
+ - 其他类别(person/cup/...):按 YOLO 实际类别生成"非化学品物品"实体,不参与规则
+ """
+ sorted_dets = sorted(detections, key=lambda d: d.bbox[0])
entities = []
- for i, bottle in enumerate(sorted_bottles):
- chem = CHEMICAL_LIBRARY[i % len(CHEMICAL_LIBRARY)]
- # 置信度基于检测置信度
- ocr_conf = round(bottle.confidence * 0.98, 2)
- entities.append({
- "raw_text": chem["name"],
- "entity_id": chem["entity_id"],
- "standard_name": chem["name"],
- "aliases": chem["aliases"],
- "cas_no": chem["cas_no"],
- "hazard_class": chem["hazard_class"],
- "confidence": ocr_conf,
- "match_type": "exact",
- "bbox": bottle.bbox
- })
+ chem_idx = 0 # 仅对容器类滚动取用化学品候选
+
+ for i, det in enumerate(sorted_dets):
+ cls = det.class_name
+ zh_name = COCO_CLASS_ZH.get(cls, cls)
+
+ if cls in CONTAINER_CLASSES:
+ # 容器类:模拟 OCR 识别标签(真实系统中由 OCR + 词典标准化完成)
+ chem = CHEMICAL_LIBRARY[chem_idx % len(CHEMICAL_LIBRARY)]
+ chem_idx += 1
+ entities.append({
+ "raw_text": chem["name"],
+ "entity_id": chem["entity_id"],
+ "standard_name": chem["name"],
+ "aliases": chem["aliases"],
+ "cas_no": chem["cas_no"],
+ "hazard_class": chem["hazard_class"],
+ "confidence": det.confidence,
+ "match_type": "ocr_simulated",
+ "is_chemical": True,
+ "bbox": det.bbox,
+ "detected_class": cls,
+ "detected_class_zh": zh_name,
+ })
+ else:
+ # 非化学品物品:名称、类别、置信度全部来自 YOLO 真实输出
+ entities.append({
+ "raw_text": zh_name,
+ "entity_id": f"obj-{cls}-{i:03d}",
+ "standard_name": zh_name,
+ "aliases": [cls],
+ "cas_no": None,
+ "hazard_class": "非化学品",
+ "confidence": det.confidence,
+ "match_type": "yolo_detected",
+ "is_chemical": False,
+ "bbox": det.bbox,
+ "detected_class": cls,
+ "detected_class_zh": zh_name,
+ })
return entities
+def assign_chemicals(bottles: List[Detection]) -> List[dict]:
+ """兼容旧调用"""
+ return assign_chemicals_from_detections(bottles)
+
+
# ==================== 相邻关系(基于真实 bbox 计算)====================
def calc_proximity(entities: List[dict]) -> List[dict]:
- """基于实体 bbox 计算相邻关系"""
+ """基于实体 bbox 计算相邻关系(仅化学品实体参与,非化学品物品不参与规则)"""
+ # 只保留化学品(YOLO 识别为容器 + OCR 识别出试剂的实体)
+ chem_entities = [e for e in entities if e.get("is_chemical")]
pairs = []
- for i in range(len(entities)):
- for j in range(i + 1, len(entities)):
- a = entities[i]
- b = entities[j]
+ for i in range(len(chem_entities)):
+ for j in range(i + 1, len(chem_entities)):
+ a = chem_entities[i]
+ b = chem_entities[j]
# 计算中心点距离
ax = (a["bbox"][0] + a["bbox"][2]) / 2
@@ -324,90 +378,246 @@ def calc_severity(rule_hits: List[dict], entity_count: int) -> dict:
# ==================== 目标检测(YOLO + 轮廓分析回退)====================
def detect_objects(image_path: str) -> list:
- """目标检测:优先 YOLO,若无结果则用轮廓分析回退"""
+ """目标检测:100% YOLO 真实推理结果,不做任何写死补充"""
detector = get_detector(get_model_path())
detections = detector.detect(image_path)
- bottles = [d for d in detections if d.class_name == "bottle"]
+ # 仅过滤无效输出(类别名缺失或占位符 "0"),保留 YOLO 识别出的全部真实类别与置信度
+ result = [d for d in detections
+ if d.confidence >= 0.25 and d.class_name and d.class_name != "0"]
+
+ return result
- # 如果 YOLO 没检测到瓶子,用图像轮廓分析回退
+
+def postprocess_bottles(bottles: list, image_path: str) -> list:
+ """后处理 YOLO 瓶子检测结果:合并窄框 + 扩展宽度"""
if not bottles:
- bottles = contour_bottle_detection(image_path)
- detections = bottles
+ return bottles
+
+ img = cv2_imread_unicode(image_path)
+ if img is not None:
+ h, w = img.shape[:2]
+ else:
+ h, w = 480, 640
+
+ processed = []
+ for b in bottles:
+ x1, y1, x2, y2 = b.bbox
+ bw = x2 - x1
+ bh = y2 - y1
+
+ # 如果框太窄(宽 < 高的 0.3),按高度扩展到合理宽度
+ if bw < bh * 0.3 and bh > 50:
+ target_w = min(bh * 0.45, 100) # 瓶子宽高比约 0.4-0.5
+ cx = (x1 + x2) / 2
+ x1_new = max(0, cx - target_w / 2)
+ x2_new = min(w, cx + target_w / 2)
+ b.bbox = [x1_new, y1, x2_new, y2]
+
+ processed.append(b)
+
+ # 合并高度重叠的相近框
+ merged = merge_nearby_boxes(processed)
+ return merged
- return detections
+def merge_nearby_boxes(boxes: list, max_gap: int = 25, iou_threshold: float = 0.2) -> list:
+ """合并水平距离近且高度重叠的框"""
+ if not boxes:
+ return boxes
+
+ # 按 x 坐标排序
+ boxes = sorted(boxes, key=lambda b: b.bbox[0])
+ merged = []
-def contour_bottle_detection(image_path: str) -> list:
- """基于轮廓分析的瓶子检测(用于合成图片等 YOLO 无法识别的场景)"""
+ for box in boxes:
+ bx1, by1, bx2, by2 = box.bbox
+ bw = bx2 - bx1
+
+ should_merge = False
+ for m in merged:
+ mx1, my1, mx2, my2 = m.bbox
+
+ # 水平距离
+ gap = bx1 - mx2
+ if gap < 0: # 有重叠
+ # 计算 IoU
+ ix1 = max(bx1, mx1)
+ iy1 = max(by1, my1)
+ ix2 = min(bx2, mx2)
+ iy2 = min(by2, my2)
+ inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
+ area1 = (bx2 - bx1) * (by2 - by1)
+ area2 = (mx2 - mx1) * (my2 - my1)
+ iou = inter / min(area1, area2) if min(area1, area2) > 0 else 0
+ if iou > iou_threshold:
+ should_merge = True
+ # 合并
+ m.bbox = [min(mx1, bx1), min(my1, by1), max(mx2, bx2), max(my2, by2)]
+ m.confidence = max(m.confidence, box.confidence)
+ break
+ elif gap < max_gap:
+ # 水平距离近,检查垂直重叠
+ v_overlap = min(by2, my2) - max(by1, my1)
+ min_h = min(by2 - by1, my2 - my1)
+ if min_h > 0 and v_overlap / min_h > 0.4:
+ should_merge = True
+ m.bbox = [min(mx1, bx1), min(my1, by1), max(mx2, bx2), max(my2, by2)]
+ m.confidence = max(m.confidence, box.confidence)
+ break
+
+ if not should_merge:
+ merged.append(box)
+
+ return merged
+
+
+def color_cap_detection(image_path: str) -> list:
+ """颜色+轮廓检测:优先检测红色瓶盖,从瓶盖向下扩展瓶身"""
from yolo_detector import Detection
img = cv2_imread_unicode(image_path)
if img is None:
return []
+ h, w = img.shape[:2]
+ hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
- h, w = gray.shape
- bottles = []
+ # 策略1:检测红色瓶盖(HSV 红色范围)
+ red_mask1 = cv2.inRange(hsv, np.array([0, 80, 80]), np.array([10, 255, 255]))
+ red_mask2 = cv2.inRange(hsv, np.array([160, 80, 80]), np.array([180, 255, 255]))
+ red_mask = cv2.bitwise_or(red_mask1, red_mask2)
- # 策略:检测瓶子/标签的边框颜色(灰度 130-165 的灰色边框)
- border_mask = cv2.inRange(gray, 130, 165)
- kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
- border_closed = cv2.morphologyEx(border_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+ kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
+ red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+ red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel, iterations=1)
- contours, _ = cv2.findContours(border_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
- for cnt in contours:
- area = cv2.contourArea(cnt)
- if area < (h * w * 0.003):
- continue
+ red_contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
- x, y, bw, bh = cv2.boundingRect(cnt)
- aspect_ratio = bh / bw if bw > 0 else 0
- # 瓶子高大于宽
- if aspect_ratio < 1.2 or aspect_ratio > 10.0:
- continue
- if bw > w * 0.8 or bh > h * 0.8:
+ bottles = []
+ cap_centers = []
+ for c in red_contours:
+ area = cv2.contourArea(c)
+ if area < 100:
continue
-
- confidence = min(0.5 + area / (h * w) * 15, 0.85)
- bottles.append(Detection(
- bbox=[float(x), float(y), float(x + bw), float(y + bh)],
- class_name="bottle",
- confidence=round(confidence, 2)
- ))
-
- # 策略2:如果没找到,尝试检测整体矩形物体(边缘检测 + 霍夫变换思路)
- if not bottles:
- blurred = cv2.GaussianBlur(gray, (3, 3), 0)
- edges = cv2.Canny(blurred, 30, 100)
- # 找垂直和水平直线
- contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
- for cnt in contours:
- area = cv2.contourArea(cnt)
- if area < (h * w * 0.01):
- continue
- x, y, bw, bh = cv2.boundingRect(cnt)
- aspect_ratio = bh / bw if bw > 0 else 0
- if aspect_ratio < 1.3 or aspect_ratio > 8.0:
- continue
- if bw > w * 0.5 or bh > h * 0.7:
- continue
- confidence = min(0.4 + area / (h * w) * 10, 0.7)
+ x, y, cw, ch = cv2.boundingRect(c)
+ ratio = ch / cw if cw > 0 else 0
+ if 0.3 < ratio < 3.0:
+ cap_centers.append((x + cw // 2, y + ch // 2, cw, ch, area))
+
+ # 从瓶盖向下扫描,动态确定瓶身左右和上下边界
+ for cx, cy, cap_w, cap_h, cap_area in cap_centers:
+ body_w_est = int(cap_w * 2.8) # 预估瓶身宽度
+ x1_est = max(0, cx - body_w_est // 2)
+ x2_est = min(w, cx + body_w_est // 2)
+ y_start = cy + cap_h // 2 # 从瓶盖下沿开始
+
+ # 货架背景参考色(取瓶盖上方区域的中位数 BGR)
+ bg_y1 = max(0, cy - cap_h - 30)
+ bg_y2 = max(1, cy - cap_h - 5)
+ bg_strip = img[bg_y1:bg_y2, x1_est:x2_est]
+ if bg_strip.size > 0:
+ bg_color = np.median(bg_strip.reshape(-1, 3), axis=0)
+ else:
+ bg_color = np.array([45, 85, 130], dtype=np.float32) # 默认棕色货架 BGR
+
+ # 逐行向下扫描:瓶身像素与背景色差异明显
+ strip_half = max(8, body_w_est // 3)
+ last_body_row = y_start
+ bg_gap = 0
+ for row in range(y_start, min(h, y_start + int(h * 0.6))):
+ rx1 = max(0, cx - strip_half)
+ rx2 = min(w, cx + strip_half)
+ row_pixels = img[row, rx1:rx2].astype(np.float32)
+ # 与背景色的距离
+ dist = np.sqrt(((row_pixels - bg_color) ** 2).sum(axis=1))
+ body_ratio = float((dist > 45).mean())
+ if body_ratio > 0.35:
+ last_body_row = row
+ bg_gap = 0
+ else:
+ bg_gap += 1
+ # 连续 25 行都是背景,认为瓶身结束
+ if bg_gap > 25 and last_body_row > y_start:
+ break
+
+ # 左右边界:在瓶身中段扫描,找到与背景不同的连续区域
+ mid_y = (y_start + last_body_row) // 2
+ col_pixels = img[mid_y, x1_est:x2_est].astype(np.float32)
+ col_dist = np.sqrt(((col_pixels - bg_color) ** 2).sum(axis=1))
+ body_cols = np.where(col_dist > 40)[0]
+ if len(body_cols) > 5:
+ bx1 = float(x1_est + int(body_cols[0]) - 3)
+ bx2 = float(x1_est + int(body_cols[-1]) + 3)
+ else:
+ bx1, bx2 = float(x1_est), float(x2_est)
+
+ by1 = float(max(0, cy - cap_h // 2 - 2)) # 框顶部包含瓶盖
+ by2 = float(min(h, last_body_row + 4)) # 框底部到瓶身下沿
+
+ # 验证:高度合理(至少瓶盖高的 6 倍)
+ if by2 - by1 > cap_h * 6:
bottles.append(Detection(
- bbox=[float(x), float(y), float(x + bw), float(y + bh)],
+ bbox=[bx1, by1, bx2, by2],
class_name="bottle",
- confidence=round(confidence, 2)
+ confidence=0.78
))
- # 去重和合并重叠框
+ # 策略2:如果没有瓶盖,用浅色矩形检测(瓶身/标签)
+ if not bottles:
+ # 检测浅色区域(标签/瓶身)
+ light_mask = cv2.inRange(gray, 170, 255)
+ light_mask = cv2.morphologyEx(light_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+ contours, _ = cv2.findContours(light_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+ candidates = []
+ for c in contours:
+ area = cv2.contourArea(c)
+ if area < 1000:
+ continue
+ x, y, bw, bh = cv2.boundingRect(c)
+ aspect = bh / bw if bw > 0 else 0
+ if 1.0 < aspect < 6.0 and bw > 25 and bh > 60:
+ candidates.append(Detection(
+ bbox=[float(x), float(y), float(x + bw), float(y + bh)],
+ class_name="bottle",
+ confidence=0.68
+ ))
+ bottles = candidates
+
+ # 策略3:如果仍无结果,用边缘检测
+ if not bottles:
+ blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+ edges = cv2.Canny(blurred, 30, 100)
+ contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+ candidates = []
+ for c in contours:
+ area = cv2.contourArea(c)
+ if area < 500:
+ continue
+ x, y, bw, bh = cv2.boundingRect(c)
+ aspect = bh / bw if bw > 0 else 0
+ if 1.2 < aspect < 8 and bw > 25 and ch > 50:
+ candidates.append(Detection(
+ bbox=[float(x), float(y), float(x + bw), float(y + bh)],
+ class_name="bottle",
+ confidence=0.60
+ ))
+ bottles = candidates
+
+ # 去重
bottles = merge_overlapping_boxes(bottles)
-
- # 按 x 坐标排序
bottles.sort(key=lambda b: b.bbox[0])
+
+ print(f"[颜色检测] 检测到 {len(bottles)} 个瓶子, 瓶盖数: {len(cap_centers)}")
return bottles
+def contour_bottle_detection(image_path: str) -> list:
+ """旧版轮廓检测(保留作为兼容接口)"""
+ return color_cap_detection(image_path)
+
+
def merge_overlapping_boxes(boxes: list, iou_threshold: float = 0.3) -> list:
"""合并重叠的检测框"""
if not boxes:
@@ -448,14 +658,18 @@ def run_inference(image_path: str, filename: str, image_key: str = "") -> dict:
if not quality_result["passed"]:
return build_quality_fail_result(quality_result, image_key)
- # 2. 目标检测(YOLO + 轮廓分析回退)
+ # 2. 目标检测(100% YOLO 真实推理结果)
detections = detect_objects(image_path)
- det_dicts = [{"bbox": d.bbox, "class_name": d.class_name, "confidence": d.confidence} for d in detections]
-
- bottles = [d for d in detections if d.class_name == "bottle"]
-
- # 3. 如果没检测到瓶子(零候选也进入待人工复核,由人工完成结论)
- if not bottles:
+ det_dicts = [{
+ "bbox": d.bbox,
+ "class_name": d.class_name,
+ "class_zh": COCO_CLASS_ZH.get(d.class_name, d.class_name),
+ "confidence": d.confidence
+ } for d in detections]
+
+ # 3. 根据检测到的类别分配化学品(所有检测到的物体都参与)
+ # YOLO 真实识别的物体类别和置信度直接显示在检测框中
+ if not detections:
return {
"inference_run_id": str(uuid.uuid4()),
"pipeline_version": "1.0.0",
@@ -470,17 +684,17 @@ def run_inference(image_path: str, filename: str, image_key: str = "") -> dict:
"error_code": None,
"image_key": image_key,
"quality": quality_result,
- "detections": det_dicts,
+ "detections": [],
"ocr_fields": [],
"entities": [],
"proximity_pairs": [],
"rule_hits": [],
- "severity_summary": {"overall": "low", "confidence_weighted_score": 0.5, "uncertainty_factors": ["no_bottle_detected"]},
+ "severity_summary": {"overall": "low", "confidence_weighted_score": 0.5, "uncertainty_factors": ["no_detection"]},
"timing_ms": {"total": 1500}
}
- # 4. 实体识别(按瓶子位置动态分配化学品)
- entities = assign_chemicals(bottles)
+ # 4. 实体识别(按检测框位置动态分配化学品,类别名映射到化学品库)
+ entities = assign_chemicals_from_detections(detections)
# 5. 相邻关系(基于真实 bbox 距离)
proximity_pairs = calc_proximity(entities)
@@ -491,9 +705,11 @@ def run_inference(image_path: str, filename: str, image_key: str = "") -> dict:
# 7. 严重度评估
severity = calc_severity(rule_hits, len(entities))
- # 8. 构造 OCR 字段(基于分配的化学品)
+ # 8. 构造 OCR 字段(仅化学品容器有模拟 OCR 文本,非化学品物品无 OCR)
ocr_fields = []
for e in entities:
+ if not e.get("is_chemical"):
+ continue
cx = int((e["bbox"][0] + e["bbox"][2]) / 2)
cy = int((e["bbox"][1] + e["bbox"][3]) / 2)
ocr_fields.append({
diff --git a/demo_web/backend/yolo_detector.py b/demo_web/backend/yolo_detector.py
index e78d7a8..04e3505 100644
--- a/demo_web/backend/yolo_detector.py
+++ b/demo_web/backend/yolo_detector.py
@@ -1,302 +1,289 @@
# YOLO 目标检测模块
-# 支持: YOLOv5 .pt (torch.hub) / YOLOv8 .onnx (OpenCV DNN)
-# 如果没有模型文件,自动回退到模拟模式
+# 直接通过 torch 加载 YOLOv5 .pt(不依赖 ultralytics 的下载逻辑)
+# 无模型时回退到颜色+轮廓分析
import os
-# 禁用 ultralytics 自动安装依赖
-os.environ.setdefault('YOLO_VERBOSE', 'False')
-os.environ.setdefault('YOLO_AUTOUPDATE', 'False')
-os.environ.setdefault('MPLCONFIGDIR', os.path.join(os.environ.get('TEMP', '.'), 'matplotlib'))
-
+import sys
import cv2
import numpy as np
-import os
-from typing import List, Dict, Optional
+import torch
+from typing import List, Optional
from dataclasses import dataclass
@dataclass
class Detection:
- bbox: List[float] # [x1, y1, x2, y2]
+ bbox: List[float] # [x1, y1, x2, y2] 原图像素坐标
class_name: str
confidence: float
class YOLODetector:
- """YOLO 检测器封装"""
-
- # COCO 类别中与实验室相关的映射
- LAB_CLASSES = {
- 39: "bottle", # bottle
- 41: "cup", # cup -> 可作为容器
- 67: "cell phone", # 可扩展为标签
- }
+ """YOLOv5 检测器(直接 torch 加载)"""
+
+ COCO_NAMES = [
+ "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train",
+ "truck", "boat", "traffic light", "fire hydrant", "stop sign",
+ "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep",
+ "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
+ "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
+ "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
+ "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork",
+ "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
+ "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
+ "couch", "potted plant", "bed", "dining table", "toilet", "tv",
+ "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave",
+ "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
+ "scissors", "teddy bear", "hair drier", "toothbrush"
+ ]
def __init__(self, model_path: Optional[str] = None, conf_threshold: float = 0.25, nms_threshold: float = 0.45):
self.conf_threshold = conf_threshold
self.nms_threshold = nms_threshold
- self.net = None
- self.model = None # torch.hub model
- self.input_size = (640, 640)
+ self.model = None
self.use_mock = True
- self.model_type = "none" # "pt" | "onnx" | "none"
+ self.model_type = "none"
+ self._yolov5_path = None
if model_path and os.path.exists(model_path):
- ext = os.path.splitext(model_path)[1].lower()
- if ext == ".onnx":
- self._load_onnx(model_path)
- elif ext == ".pt":
- self._load_pt(model_path)
- else:
- print(f"[YOLO] 不支持的模型格式: {ext},使用模拟模式")
+ try:
+ self._load_yolov5_direct(model_path)
+ except Exception as e:
+ print(f"[YOLO] 直接加载失败: {e}")
else:
- print(f"[YOLO] 模型文件不存在: {model_path},使用模拟模式")
+ print(f"[YOLO] 模型文件不存在: {model_path}")
- def _load_onnx(self, model_path: str):
- """加载 ONNX 模型 (OpenCV DNN)"""
- try:
- self.net = cv2.dnn.readNetFromONNX(model_path)
- self.use_mock = False
- self.model_type = "onnx"
- print(f"[YOLO] ONNX 模型加载成功: {model_path}")
- except Exception as e:
- print(f"[YOLO] ONNX 加载失败: {e},使用模拟模式")
+ def _load_yolov5_direct(self, model_path: str):
+ """通过本地 yolov5 源码加载权重(优先官方 attempt_load,不走网络)"""
+ hub = os.path.join(os.path.expanduser('~'), '.cache', 'torch', 'hub', 'ultralytics_yolov5_master')
+ if not os.path.exists(hub):
+ print("[YOLO] 本地无 yolov5 源码缓存,无法加载")
+ return
- def _load_pt(self, model_path: str):
- """加载 YOLOv5 .pt 模型(直接 torch.load,绕过 hub 依赖检查)"""
+ sys.path.insert(0, hub)
+ self._yolov5_path = hub
+
+ # 方式1:官方加载器(标准 checkpoint,含完整 BN/结构校验)
try:
- import torch
- import sys
-
- # 将本地 yolov5 仓库加入路径
- hub_cache = os.path.join(os.path.expanduser('~'), '.cache', 'torch', 'hub', 'ultralytics_yolov5_master')
- if os.path.exists(hub_cache) and hub_cache not in sys.path:
- sys.path.insert(0, hub_cache)
-
- # 加载 state_dict
- state_dict = torch.load(model_path, map_location='cpu', weights_only=False)
-
- # 如果是纯 state_dict,需要构建模型结构
- from models.yolo import Model
- from models.common import DetectMultiBackend
-
- # 尝试用 yolo v5s 的配置构建模型
- import yaml
- model_yaml = os.path.join(hub_cache, 'models', 'yolov5s.yaml')
- model = Model(model_yaml, nc=80)
-
- # 去除 'model.' 前缀(因为保存的是 model.model.state_dict())
- new_state_dict = {}
- for k, v in state_dict.items():
- if k.startswith('model.'):
- new_key = k[6:]
- else:
- new_key = k
- new_state_dict[new_key] = v
-
- model.load_state_dict(new_state_dict, strict=False)
+ from models.experimental import attempt_load
+ model = attempt_load(model_path, device=torch.device('cpu'))
model = model.float().eval()
- model.conf = self.conf_threshold
- model.iou = self.nms_threshold
self.model = model
self.use_mock = False
- self.model_type = "pt"
- print(f"[YOLO] .pt 模型加载成功: {model_path}")
+ self.model_type = "attempt_load"
+ print(f"[YOLO] 官方加载成功: {os.path.basename(model_path)}")
+ return
except Exception as e:
- print(f"[YOLO] .pt 加载失败: {e},使用模拟模式")
+ print(f"[YOLO] attempt_load 失败({e}),尝试裸 state_dict 加载")
+
+ # 方式2:裸 state_dict 兜底(key 可能带 model. 前缀)
+ from models.yolo import Model
+ yaml_path = os.path.join(hub, 'models', 'yolov5s.yaml')
+ model = Model(yaml_path, nc=80)
+ ckpt = torch.load(model_path, map_location='cpu', weights_only=False)
+ result = model.load_state_dict(ckpt, strict=False)
+ if len(result.missing_keys) > 50:
+ # 尝试剥离一层前缀
+ ckpt = {k[6:] if k.startswith('model.') else k: v for k, v in ckpt.items()}
+ model.load_state_dict(ckpt, strict=False)
+ model = model.float().eval()
+ if isinstance(model.names, list) and (model.names[0] == 0 or model.names[0] == '0'):
+ model.names = self.COCO_NAMES
+ self.model = model
+ self.use_mock = False
+ self.model_type = "state_dict"
+ print(f"[YOLO] state_dict 加载成功: {os.path.basename(model_path)}")
def detect(self, image_path: str) -> List[Detection]:
- """执行检测,返回 Detection 列表"""
- if self.use_mock:
- return self._mock_detect(image_path)
+ """执行检测:仅返回 YOLO 真实推理结果,模型不可用时返回空列表(不伪造检测)"""
+ if self.use_mock or self.model is None:
+ print("[YOLO] 模型未加载,跳过检测(无真实推理结果)")
+ return []
try:
- if self.model_type == "onnx":
- return self._real_detect_onnx(image_path)
- elif self.model_type == "pt":
- return self._real_detect_pt(image_path)
- else:
- return self._mock_detect(image_path)
+ return self._real_detect(image_path)
except Exception as e:
- print(f"[YOLO] 推理失败: {e},回退到模拟模式")
- return self._mock_detect(image_path)
+ print(f"[YOLO] 推理失败: {e}")
+ return []
+
+ def _real_detect(self, image_path: str) -> List[Detection]:
+ """YOLOv5 推理:标准 letterbox 预处理 + scale_boxes 坐标映射"""
+ from utils.general import non_max_suppression, scale_boxes
+ try:
+ from utils.augmentations import letterbox
+ except ImportError:
+ from utils.general import letterbox
- def _real_detect_pt(self, image_path: str) -> List[Detection]:
- """YOLOv5 .pt 推理(直接加载的模型,需手动预处理和NMS)"""
- import torch
- from PIL import Image
+ img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
+ if img is None:
+ return []
- img = Image.open(image_path).convert('RGB')
- orig_w, orig_h = img.size
+ orig_h, orig_w = img.shape[:2]
- # 预处理:resize 到 640x640
- img_resized = img.resize((640, 640), Image.BILINEAR)
- img_tensor = torch.from_numpy(np.array(img_resized)).permute(2, 0, 1).float() / 255.0
- img_tensor = img_tensor.unsqueeze(0)
+ # letterbox:等比缩放 + 填充到 640x640
+ im = letterbox(img, 640, auto=False)[0]
+ im = im.transpose((2, 0, 1))[::-1] # HWC BGR -> CHW RGB
+ im = np.ascontiguousarray(im)
+ t = torch.from_numpy(im).float() / 255.0
+ t = t.unsqueeze(0)
with torch.no_grad():
- pred = self.model(img_tensor)
+ pred = self.model(t)
- # YOLOv5 输出: tuple, pred[0] 是 [batch, num_boxes, 5+num_classes]
- if isinstance(pred, (tuple, list)):
- pred = pred[0]
+ pred = non_max_suppression(pred, self.conf_threshold, self.nms_threshold, max_det=50)
- # 尝试使用 yolov5 的 NMS
- try:
- from utils.general import non_max_suppression
- results = non_max_suppression(pred, conf_thres=self.conf_threshold, iou_thres=self.nms_threshold)
- detections = []
- for det in results:
- if det is not None and len(det) > 0:
- for *box, conf, cls_id in det.cpu().numpy():
- cls_id = int(cls_id)
- if cls_id in self.LAB_CLASSES:
- # 坐标从 640x640 映射回原图
- scale_x = orig_w / 640
- scale_y = orig_h / 640
- x1 = float(box[0]) * scale_x
- y1 = float(box[1]) * scale_y
- x2 = float(box[2]) * scale_x
- y2 = float(box[3]) * scale_y
- detections.append(Detection(
- bbox=[x1, y1, x2, y2],
- class_name=self.LAB_CLASSES[cls_id],
- confidence=float(conf)
- ))
- return detections
- except ImportError:
- # 没有 yolov5 utils,使用 torchvision NMS
- pass
-
- # 手动 NMS 回退
- boxes = []
- scores = []
- class_ids = []
-
- pred_np = pred[0].cpu().numpy()
- for i in range(pred_np.shape[0]):
- obj_conf = pred_np[i, 4]
- if obj_conf < self.conf_threshold:
- continue
- cls_scores = pred_np[i, 5:]
- cls_id = int(np.argmax(cls_scores))
- cls_conf = cls_scores[cls_id] * obj_conf
- if cls_conf < self.conf_threshold:
- continue
- if cls_id not in self.LAB_CLASSES:
- continue
-
- x, y, w, h = pred_np[i, :4]
- x1 = (x - w / 2) * orig_w / 640
- y1 = (y - h / 2) * orig_h / 640
- x2 = (x + w / 2) * orig_w / 640
- y2 = (y + h / 2) * orig_h / 640
- boxes.append([x1, y1, x2 - x1, y2 - y1])
- scores.append(float(cls_conf))
- class_ids.append(cls_id)
+ detections = []
+ names = self.model.names if hasattr(self.model, 'names') else {i: n for i, n in enumerate(self.COCO_NAMES)}
- if not boxes:
- return []
+ if pred and len(pred) > 0 and pred[0] is not None and len(pred[0]) > 0:
+ det = pred[0]
+ # 坐标从 640x640 letterbox 空间映射回原图(官方函数,自动扣除填充)
+ det[:, :4] = scale_boxes(im.shape[1:], det[:, :4], (orig_h, orig_w)).round()
+ for d in det:
+ x1, y1, x2, y2, conf, cls_id = d.cpu().numpy()
+ cls_id = int(cls_id)
+ class_name = names.get(cls_id, f"class_{cls_id}") if isinstance(names, dict) else (names[cls_id] if cls_id < len(names) else f"class_{cls_id}")
- # OpenCV NMS
- indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_threshold, self.nms_threshold)
- detections = []
- if len(indices) > 0:
- for i in indices.flatten():
- x, y, w, h = boxes[i]
detections.append(Detection(
- bbox=[x, y, x + w, y + h],
- class_name=self.LAB_CLASSES[class_ids[i]],
- confidence=scores[i]
+ bbox=[float(max(0, x1)), float(max(0, y1)),
+ float(min(orig_w, x2)), float(min(orig_h, y2))],
+ class_name=class_name,
+ confidence=float(conf)
))
+
+ detections.sort(key=lambda d: d.confidence, reverse=True)
+ print(f"[YOLO] 检测到 {len(detections)} 个目标: {[(d.class_name, round(d.confidence, 2)) for d in detections[:10]]}")
return detections
- def _real_detect_onnx(self, image_path: str) -> List[Detection]:
- """YOLOv8 ONNX 推理 (OpenCV DNN)"""
+ def _color_contour_detect(self, image_path: str) -> List[Detection]:
+ """颜色 + 轮廓分析(无模型时或合成图片回退)"""
img = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
if img is None:
- raise ValueError(f"无法读取图片: {image_path}")
+ return []
h, w = img.shape[:2]
+ hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+ gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
- blob = cv2.dnn.blobFromImage(img, 1/255.0, self.input_size, swapRB=True, crop=False)
- self.net.setInput(blob)
- outputs = self.net.forward()
- outputs = np.transpose(np.squeeze(outputs))
-
- boxes = []
- scores = []
- class_ids = []
-
- for i in range(outputs.shape[0]):
- classes_scores = outputs[i][4:]
- max_score = np.amax(classes_scores)
-
- if max_score >= self.conf_threshold:
- class_id = np.argmax(classes_scores)
-
- if class_id in self.LAB_CLASSES:
- x, y, w_box, h_box = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
-
- left = int((x - w_box/2) * w / self.input_size[0])
- top = int((y - h_box/2) * h / self.input_size[1])
- width = int(w_box * w / self.input_size[0])
- height = int(h_box * h / self.input_size[1])
+ detections = []
- boxes.append([left, top, width, height])
- scores.append(float(max_score))
- class_ids.append(class_id)
+ # 策略1:检测红色瓶盖(HSV 红色范围)
+ red_mask1 = cv2.inRange(hsv, np.array([0, 80, 80]), np.array([10, 255, 255]))
+ red_mask2 = cv2.inRange(hsv, np.array([160, 80, 80]), np.array([180, 255, 255]))
+ red_mask = cv2.bitwise_or(red_mask1, red_mask2)
- indices = cv2.dnn.NMSBoxes(boxes, scores, self.conf_threshold, self.nms_threshold)
+ kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
+ red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_CLOSE, kernel, iterations=2)
+ red_mask = cv2.morphologyEx(red_mask, cv2.MORPH_OPEN, kernel, iterations=1)
- detections = []
- if len(indices) > 0:
- for i in indices.flatten():
- x, y, w_box, h_box = boxes[i]
+ red_contours, _ = cv2.findContours(red_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+ cap_centers = []
+ for c in red_contours:
+ area = cv2.contourArea(c)
+ if area < 100:
+ continue
+ x, y, cw, ch = cv2.boundingRect(c)
+ if 0.5 < ch / cw if cw > 0 else 0 < 3.0:
+ cap_centers.append((x + cw // 2, y + ch // 2, cw, ch, area))
+
+ # 策略2:基于瓶盖位置向下扩展瓶身
+ for cx, cy, cap_w, cap_h, cap_area in cap_centers:
+ # 瓶身宽度约为瓶盖的 2-3 倍
+ body_w = int(cap_w * 2.5)
+ body_h = int(cap_h * 12) # 瓶身高约瓶盖的 12 倍
+ bx1 = max(0, cx - body_w // 2)
+ by1 = max(0, cy - cap_h // 2)
+ bx2 = min(w, cx + body_w // 2)
+ by2 = min(h, by1 + body_h)
+
+ # 验证:瓶身区域应该有浅色像素(标签/瓶身)
+ roi = img[by1:by2, bx1:bx2]
+ if roi.size > 0:
+ mean_brightness = np.mean(cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY))
+ if mean_brightness > 100: # 瓶身区域足够亮
+ detections.append(Detection(
+ bbox=[float(bx1), float(by1), float(bx2), float(by2)],
+ class_name="bottle",
+ confidence=0.75
+ ))
+
+ # 策略3:如果没有检测到瓶盖,用灰度轮廓
+ if not detections:
+ blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+ edges = cv2.Canny(blurred, 30, 100)
+ contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+ candidates = []
+ for c in contours:
+ area = cv2.contourArea(c)
+ if area < 500:
+ continue
+ x, y, cw, ch = cv2.boundingRect(c)
+ aspect = ch / cw if cw > 0 else 0
+ if 1.2 < aspect < 8 and cw > 25 and ch > 50:
+ candidates.append((x, y, cw, ch, area))
+
+ candidates.sort(key=lambda c: c[4], reverse=True)
+ for x, y, cw, ch, area in candidates[:10]:
detections.append(Detection(
- bbox=[x, y, x + w_box, y + h_box],
- class_name=self.LAB_CLASSES[class_ids[i]],
- confidence=scores[i]
+ bbox=[float(x), float(y), float(x + cw), float(y + ch)],
+ class_name="bottle",
+ confidence=0.65
))
+ # 去重
+ detections = self._merge_nearby(detections)
+ print(f"[YOLO-颜色轮廓] 检测到 {len(detections)} 个目标")
return detections
- def _mock_detect(self, image_path: str) -> List[Detection]:
- """模拟检测"""
- filename = os.path.basename(image_path).lower()
-
- if "blurry" in filename or "glare" in filename or "dark" in filename:
- return []
-
- if "single" in filename or "alone" in filename:
- return [
- Detection(bbox=[120, 80, 200, 300], class_name="bottle", confidence=0.94),
- Detection(bbox=[125, 85, 195, 160], class_name="label", confidence=0.91),
- Detection(bbox=[100, 50, 400, 320], class_name="shelf", confidence=0.95),
- ]
-
- return [
- Detection(bbox=[120, 80, 200, 300], class_name="bottle", confidence=0.94),
- Detection(bbox=[125, 85, 195, 160], class_name="label", confidence=0.91),
- Detection(bbox=[250, 75, 330, 290], class_name="bottle", confidence=0.89),
- Detection(bbox=[255, 80, 325, 155], class_name="label", confidence=0.87),
- Detection(bbox=[100, 50, 400, 320], class_name="shelf", confidence=0.95),
- ]
-
-
-# 全局检测器实例
+ def _merge_nearby(self, boxes: List[Detection], max_gap: int = 30) -> List[Detection]:
+ """合并水平距离近的框"""
+ if not boxes:
+ return boxes
+ boxes = sorted(boxes, key=lambda b: b.bbox[0])
+ merged = []
+ for box in boxes:
+ bx1, by1, bx2, by2 = box.bbox
+ should_merge = False
+ for m in merged:
+ mx1, my1, mx2, my2 = m.bbox
+ gap = bx1 - mx2
+ if gap < 0:
+ ix1 = max(bx1, mx1)
+ iy1 = max(by1, my1)
+ ix2 = min(bx2, mx2)
+ iy2 = min(by2, my2)
+ inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
+ a1 = (bx2 - bx1) * (by2 - by1)
+ a2 = (mx2 - mx1) * (my2 - my1)
+ iou = inter / min(a1, a2) if min(a1, a2) > 0 else 0
+ if iou > 0.2:
+ m.bbox = [min(mx1, bx1), min(my1, by1), max(mx2, bx2), max(my2, by2)]
+ m.confidence = max(m.confidence, box.confidence)
+ should_merge = True
+ break
+ elif gap < max_gap:
+ v_overlap = min(by2, my2) - max(by1, my1)
+ min_h = min(by2 - by1, my2 - my1)
+ if min_h > 0 and v_overlap / min_h > 0.4:
+ m.bbox = [min(mx1, bx1), min(my1, by1), max(mx2, bx2), max(my2, by2)]
+ m.confidence = max(m.confidence, box.confidence)
+ should_merge = True
+ break
+ if not should_merge:
+ merged.append(box)
+ return merged
+
+
+# 全局检测器
_detector: Optional[YOLODetector] = None
-
def get_detector(model_path: Optional[str] = None) -> YOLODetector:
- """获取或创建检测器实例"""
global _detector
if _detector is None:
_detector = YOLODetector(model_path)
return _detector
-
def reset_detector():
- """重置检测器"""
global _detector
_detector = None
diff --git a/demo_web/frontend/index.html b/demo_web/frontend/index.html
index 47e1021..2cf6b28 100644
--- a/demo_web/frontend/index.html
+++ b/demo_web/frontend/index.html
@@ -785,20 +785,37 @@ 图像质量指标
`;
}
- // 实体卡片:OCR 原文 + 置信度 + 词典候选
+ // 实体卡片:YOLO 真实检测结果,化学品容器附模拟 OCR 结果
if (data.entities && data.entities.length > 0) {
html += `
识别到的化学品 (${data.entities.length})
- ${data.entities.map(e => `
- 识别到的物体 (${data.entities.length})
+ ${data.entities.map(e => {
+ const clsZh = e.detected_class_zh || e.detected_class || '';
+ if (e.is_chemical) {
+ return `
+ 人工复核(AI 结论仅供参考,由人工完成确认)
box.style.borderColor = d.class_name === 'bottle' ? '#e94560' : '#1890ff';
const label = document.createElement('div');
+ // 检测框颜色按类别区分
+ const colors = {
+ bottle: '#e94560', cup: '#1890ff', person: '#52c41a',
+ book: '#722ed1', scissors: '#fa8c16', 'cell phone': '#13c2c2',
+ keyboard: '#a0d911', laptop: '#2f54eb', tv: '#f5222d'
+ };
+ const color = colors[d.class_name] || '#666';
+ box.style.borderColor = color;
+
label.className = 'detection-label';
- label.style.background = d.class_name === 'bottle' ? '#e94560' : '#1890ff';
- label.innerHTML = `#${i + 1} ${d.class_name} ${(d.confidence * 100).toFixed(0)}%`;
+ label.style.background = color;
+ // 检测框显示 YOLO 实际识别类别(中文+英文)与真实置信度
+ const zh = d.class_zh || d.class_name;
+ label.innerHTML = `#${i + 1} ${zh} ${(d.confidence * 100).toFixed(0)}%`;
box.appendChild(label);
container.appendChild(box);
});
diff --git a/picture/generate_test_images.py b/picture/generate_test_images.py
index 48eaf5c..f8afde6 100644
--- a/picture/generate_test_images.py
+++ b/picture/generate_test_images.py
@@ -1,129 +1,254 @@
-# 生成测试图片用于 LabSafe Demo
+# 生成更真实的测试图片用于 LabSafe Demo
# 运行: python generate_test_images.py
from PIL import Image, ImageDraw, ImageFont, ImageFilter
import os
+import math
+
+
+def cv2_imread_unicode(path):
+ import cv2
+ import numpy as np
+ return cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)
+
+
+def draw_realistic_bottle(draw, img, cx, cy, width, height, liquid_color, label_text, font, small_font):
+ """绘制更真实的试剂瓶(圆柱体 + 渐变 + 阴影 + 标签 + 瓶盖)"""
+ x = cx - width // 2
+ top = cy - height // 2
+ bottom = cy + height // 2
+
+ # 1. 阴影(投射在货架上)
+ shadow_offset = 8
+ for i in range(width // 2):
+ alpha = int(80 * (1 - i / (width / 2)))
+ draw.ellipse(
+ [x + i, bottom + shadow_offset, x + width - i, bottom + shadow_offset + 6],
+ fill=(max(0, 139 - alpha), max(0, 90 - alpha), max(0, 43 - alpha))
+ )
+
+ # 2. 瓶身渐变(左暗右亮模拟圆柱光照)
+ body_left = x + 4
+ body_top = top + 25 # 瓶盖下方
+ body_right = x + width - 4
+ body_bottom = bottom - 5
+
+ # 用多条竖线模拟渐变
+ for px in range(body_left, body_right):
+ t = (px - body_left) / (body_right - body_left)
+ # 中间最亮,两侧暗(圆柱光照模型)
+ brightness = 1.0 - abs(t - 0.4) * 1.2
+ brightness = max(0.3, min(1.0, brightness))
+ r = int(liquid_color[0] * brightness + 40)
+ g = int(liquid_color[1] * brightness + 40)
+ b = int(liquid_color[2] * brightness + 40)
+ r, g, b = min(r, 255), min(g, 255), min(b, 255)
+ draw.line([(px, body_top), (px, body_bottom)], fill=(r, g, b))
+
+ # 瓶身边框
+ draw.rectangle([body_left, body_top, body_right, body_bottom],
+ outline=(60, 60, 60), width=2)
+
+ # 3. 瓶肩(圆弧过渡)
+ shoulder_h = 15
+ for i in range(shoulder_h):
+ t = i / shoulder_h
+ offset = int((1 - t) * width * 0.15)
+ draw.line([(x + 4 + offset, body_top - shoulder_h + i),
+ (x + width - 4 - offset, body_top - shoulder_h + i)],
+ fill=(int(liquid_color[0] * 0.6 + 30), int(liquid_color[1] * 0.6 + 30), int(liquid_color[2] * 0.6 + 30)))
+
+ # 4. 瓶颈
+ neck_w = width // 3
+ neck_x = x + (width - neck_w) // 2
+ neck_top = top + 8
+ draw.rectangle([neck_x, neck_top, neck_x + neck_w, body_top - shoulder_h + 2],
+ fill=(int(liquid_color[0] * 0.7), int(liquid_color[1] * 0.7), int(liquid_color[2] * 0.7)),
+ outline=(60, 60, 60), width=1)
+
+ # 5. 瓶盖
+ cap_w = neck_w + 6
+ cap_x = x + (width - cap_w) // 2
+ cap_top = top
+ cap_bottom = top + 12
+ # 瓶盖渐变
+ for px in range(cap_x, cap_x + cap_w):
+ t = (px - cap_x) / cap_w
+ brightness = 1.0 - abs(t - 0.5) * 0.8
+ r = int(200 * brightness)
+ g = int(40 * brightness)
+ b = int(40 * brightness)
+ draw.line([(px, cap_top), (px, cap_bottom)], fill=(r, g, b))
+ draw.rectangle([cap_x, cap_top, cap_x + cap_w, cap_bottom], outline=(120, 20, 20), width=2)
+
+ # 6. 标签(白色矩形 + 文字 + 液位线)
+ label_w = width - 12
+ label_h = 55
+ label_x = x + 6
+ label_y = body_top + 10
+
+ # 标签背景(微泛黄)
+ for py in range(label_y, label_y + label_h):
+ for px in range(label_x, label_x + label_w):
+ pass # 已用 rectangle 填充
+
+ draw.rectangle([label_x, label_y, label_x + label_w, label_y + label_h],
+ fill=(252, 248, 235), outline=(100, 100, 100), width=1)
+
+ # 标签文字
+ draw.text((label_x + 8, label_y + 6), label_text, fill=(20, 20, 20), font=font)
+ draw.text((label_x + 8, label_y + 32), "2025-01", fill=(80, 80, 80), font=small_font)
+
+ # 7. 高光(左上角白色弧线模拟反光)
+ highlight_x = body_left + 4
+ for py in range(body_top + 5, body_bottom - 5, 2):
+ for px in range(highlight_x, highlight_x + 6, 2):
+ offset_y = (py - body_top) / (body_bottom - body_top)
+ fade = 1.0 - abs(offset_y - 0.3) * 3
+ if fade > 0:
+ alpha = int(60 * fade)
+ pixel = img.getpixel((px, py))
+ new_r = min(255, pixel[0] + alpha)
+ new_g = min(255, pixel[1] + alpha)
+ new_b = min(255, pixel[2] + alpha)
+ img.putpixel((px, py), (new_r, new_g, new_b))
+
+
+def draw_shelf(draw, w, h):
+ """绘制货架(多层木板 + 阴影)"""
+ # 外框
+ draw.rectangle([30, 30, w - 30, h - 30], fill=(100, 65, 30), outline=(60, 40, 15), width=3)
+ # 内壁
+ draw.rectangle([40, 40, w - 40, h - 40], fill=(130, 85, 45), outline=(80, 55, 25), width=2)
+ # 隔板
+ shelf_y = h // 2 + 20
+ draw.rectangle([40, shelf_y, w - 40, shelf_y + 8], fill=(90, 60, 28), outline=(60, 40, 18), width=1)
+ # 背景纹理(横纹)
+ for y in range(42, h - 42, 3):
+ c = 128 + (y % 6)
+ draw.line([(42, y), (w - 42, y)], fill=(c, c - 40, c - 60))
-def draw_label(draw, x, y, width, height, text, color, font):
- """绘制模拟试剂瓶标签"""
- # 瓶身
- draw.rectangle([x, y, x+width, y+height], fill=(220, 220, 220), outline=(150, 150, 150), width=2)
- # 标签区域
- label_y = y + 20
- label_height = 80
- draw.rectangle([x+5, label_y, x+width-5, label_y+label_height], fill=(255, 255, 240), outline=(100, 100, 100), width=1)
- # 文字
- draw.text((x+10, label_y+10), text, fill=(0, 0, 0), font=font)
- # 日期
- draw.text((x+10, label_y+40), "2025-01-15", fill=(80, 80, 80), font=font)
- # 瓶盖
- draw.rectangle([x+20, y-10, x+width-20, y], fill=color, outline=(100, 100, 100), width=1)
def create_normal_scene():
- """场景1: 乙醇与丙酮相邻存放"""
- img = Image.new('RGB', (640, 480), (245, 245, 240))
+ """场景1: 乙醇与丙酮相邻存放(更真实的试剂瓶)"""
+ w, h = 640, 480
+ img = Image.new('RGB', (w, h), (240, 235, 225))
draw = ImageDraw.Draw(img)
-
- # 货架背景
- draw.rectangle([50, 50, 590, 430], fill=(139, 90, 43), outline=(100, 70, 30), width=3)
- draw.rectangle([60, 60, 580, 420], fill=(160, 110, 60), outline=(120, 80, 40), width=2)
-
+
try:
- font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 24)
- small_font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 16)
+ font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 20)
+ small_font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 14)
except:
font = ImageFont.load_default()
- small_font = ImageFont.load_default()
-
- # 乙醇瓶
- draw_label(draw, 120, 150, 80, 160, "乙醇", (200, 50, 50), font)
- # 丙酮瓶
- draw_label(draw, 250, 150, 80, 160, "丙酮", (200, 50, 50), font)
-
- # 说明文字
- draw.text((180, 360), "乙醇 + 丙酮 同柜存放", fill=(200, 0, 0), font=font)
-
+ small_font = font
+
+ draw_shelf(draw, w, h)
+
+ # 两个瓶子(位置错开,模拟真实摆放)
+ draw_realistic_bottle(draw, img, 200, 260, 90, 220, (180, 180, 185), "乙醇", font, small_font)
+ draw_realistic_bottle(draw, img, 340, 260, 90, 220, (170, 170, 175), "丙酮", font, small_font)
+
+ img = img.filter(ImageFilter.GaussianBlur(radius=0.3))
return img
+
def create_blurry_image():
- """场景2: 模糊图片(触发质量拦截)"""
- img = Image.new('RGB', (640, 480), (245, 245, 240))
+ """场景2: 模糊图片"""
+ w, h = 640, 480
+ img = Image.new('RGB', (w, h), (240, 235, 225))
draw = ImageDraw.Draw(img)
-
+
try:
- font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 24)
+ font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 20)
+ small_font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 14)
except:
font = ImageFont.load_default()
-
- # 模糊的内容
- draw.text((200, 200), "此图片故意模糊", fill=(100, 100, 100), font=font)
- draw.text((180, 240), "测试质量检测拦截", fill=(100, 100, 100), font=font)
-
- # 应用高斯模糊
- img = img.filter(ImageFilter.GaussianBlur(radius=8))
+ small_font = font
+
+ draw_shelf(draw, w, h)
+ draw_realistic_bottle(draw, img, 200, 260, 90, 220, (180, 180, 185), "乙醇", font, small_font)
+ draw_realistic_bottle(draw, img, 340, 260, 90, 220, (170, 170, 175), "丙酮", font, small_font)
+
+ img = img.filter(ImageFilter.GaussianBlur(radius=6))
return img
+
def create_single_bottle():
"""场景3: 单独存放的硫酸"""
- img = Image.new('RGB', (640, 480), (245, 245, 240))
+ w, h = 640, 480
+ img = Image.new('RGB', (w, h), (240, 235, 225))
draw = ImageDraw.Draw(img)
-
- # 货架背景
- draw.rectangle([50, 50, 590, 430], fill=(139, 90, 43), outline=(100, 70, 30), width=3)
- draw.rectangle([60, 60, 580, 420], fill=(160, 110, 60), outline=(120, 80, 40), width=2)
-
+
try:
- font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 24)
+ font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 20)
+ small_font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 14)
except:
font = ImageFont.load_default()
-
- # 单个硫酸瓶
- draw_label(draw, 250, 150, 80, 160, "硫酸", (200, 100, 50), font)
-
- draw.text((200, 360), "硫酸 单独存放", fill=(0, 100, 0), font=font)
-
+ small_font = font
+
+ draw_shelf(draw, w, h)
+ draw_realistic_bottle(draw, img, 320, 260, 95, 230, (200, 180, 100), "硫酸", font, small_font)
+
+ img = img.filter(ImageFilter.GaussianBlur(radius=0.3))
return img
+
def create_glare_image():
- """场景4: 反光图片(触发质量拦截)"""
- img = Image.new('RGB', (640, 480), (245, 245, 240))
+ """场景4: 反光图片"""
+ w, h = 640, 480
+ img = Image.new('RGB', (w, h), (240, 235, 225))
draw = ImageDraw.Draw(img)
-
+
try:
- font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 24)
+ font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 20)
+ small_font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 14)
except:
font = ImageFont.load_default()
-
- # 标签区域有强光反射
- draw.text((200, 180), "强光反射区域", fill=(255, 255, 255), font=font)
-
- # 添加白色高光区域模拟反光
- for i in range(100):
- alpha = int(255 * (1 - i/100))
- x = 150 + i * 3
- y = 100 + i * 2
- draw.ellipse([x, y, x+200-i, y+150-i], fill=(255, 255, 255))
-
+ small_font = font
+
+ draw_shelf(draw, w, h)
+ draw_realistic_bottle(draw, img, 200, 260, 90, 220, (180, 180, 185), "乙醇", font, small_font)
+
+ # 大面积白色高光
+ for i in range(200):
+ x = 150 + i
+ y = 100 + i // 2
+ r = 255
+ g = 255
+ b = 255
+ draw.ellipse([x - 60, y - 40, x + 60, y + 40], fill=(r, g, b, 180))
+
return img
+
def create_dark_image():
- """场景5: 光线不足(触发质量拦截)"""
- img = Image.new('RGB', (640, 480), (30, 30, 35))
+ """场景5: 光线不足"""
+ w, h = 640, 480
+ img = Image.new('RGB', (w, h), (240, 235, 225))
draw = ImageDraw.Draw(img)
-
+
try:
- font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 24)
+ font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 20)
+ small_font = ImageFont.truetype("C:/Windows/Fonts/simhei.ttf", 14)
except:
font = ImageFont.load_default()
-
- draw.text((200, 220), "光线不足场景", fill=(80, 80, 80), font=font)
-
+ small_font = font
+
+ draw_shelf(draw, w, h)
+ draw_realistic_bottle(draw, img, 200, 260, 90, 220, (180, 180, 185), "乙醇", font, small_font)
+ draw_realistic_bottle(draw, img, 340, 260, 90, 220, (170, 170, 175), "丙酮", font, small_font)
+
+ # 整体变暗
+ import numpy as np
+ arr = np.array(img)
+ arr = (arr * 0.25).astype(np.uint8)
+ img = Image.fromarray(arr)
return img
+
if __name__ == "__main__":
output_dir = os.path.dirname(os.path.abspath(__file__))
-
- # 生成各场景图片
+
images = {
"shelf_normal.jpg": create_normal_scene(),
"shelf_blurry.jpg": create_blurry_image(),
@@ -131,16 +256,10 @@ def create_dark_image():
"shelf_glare.jpg": create_glare_image(),
"shelf_dark.jpg": create_dark_image(),
}
-
+
for filename, img in images.items():
filepath = os.path.join(output_dir, filename)
- img.save(filepath, "JPEG", quality=90)
+ img.save(filepath, "JPEG", quality=92)
print(f"生成: {filepath}")
-
+
print(f"\n共生成 {len(images)} 张测试图片")
- print("使用说明:")
- print(" shelf_normal.jpg - 乙醇+丙酮相邻(高风险场景)")
- print(" shelf_blurry.jpg - 模糊图片(质量拦截)")
- print(" shelf_single.jpg - 单独硫酸(低风险场景)")
- print(" shelf_glare.jpg - 反光图片(质量拦截)")
- print(" shelf_dark.jpg - 光线不足(质量拦截)")
diff --git a/picture/real_bus.jpg b/picture/real_bus.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..b43e311165c785f000eb7493ff8fb662d06a3f83
GIT binary patch
literal 487438
zcmeFYbyODL*Ec*DA)V4KCDNT2AR-_jqKI^Ncc+4McY{cYlr%_p2}pN$ch@`USHHjY
ztovT;S?_xPdw=KR%zS3gKKtywPs~1NhP&~*c>q&NTv8l>Kmgzc_yg`15Z%O_j12%l
zMh2h<000p{ga7@
zWU5*P1