yolo训练部署相关代码

ai/yolo/app_test.py

@@ -0,0 +1,220 @@
+import logging
+
+import cv2
+import numpy as np
+import torch
+from ultralytics import YOLO
+import os
+
+from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw, ImageFont
+import numpy as np
+import cv2
+import io
+
+# 设置工作目录
+os.chdir("/home/bbit/mine/yolo/")
+from PIL import Image, ImageDraw, ImageFont
+import cv2
+import numpy as np
+
+def draw_annotations(img_bgr, boxes, labels, confidences=None,
+                     font_path="/usr/share/fonts/truetype/wqy/wqy-zenhei.ttc"):
+    """
+    绘制带中文、置信度、不同颜色的标注框
+    :param img_bgr: np.ndarray, BGR 图像
+    :param boxes: list of xyxy
+    :param labels: list of str
+    :param confidences: list of float, 可选，用于显示置信度
+    :param font_path: 字体路径
+    :return: np.ndarray, BGR 标注图
+    """
+    h, w, _ = img_bgr.shape
+    line_width = max(2, int(w / 400))
+    font_size = max(20, int(w / 40))
+
+    img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+    pil_img = Image.fromarray(img)
+    draw = ImageDraw.Draw(pil_img)
+    font = ImageFont.truetype(font_path, font_size)
+
+    # 定义类别颜色映射
+    color_map = {
+        "正茧": "green",
+        "双宫茧": "grey",
+        "黄斑茧": "red",
+        "毛茧": "white",
+        "蛆壳茧": "purple"
+    }
+
+    for idx, (b, label) in enumerate(zip(boxes, labels), start=1):
+        conf = confidences[idx-1] if confidences else None
+        display_text = f"{label}#{idx}"
+        if conf is not None:
+            display_text += f" {conf:.2f}"
+
+        x1, y1, x2, y2 = map(int, b)
+        box_color = color_map.get(label, "red")
+
+        # 画框
+        draw.rectangle([x1, y1, x2, y2], outline=box_color, width=line_width)
+
+        # 文本边界
+        bbox = draw.textbbox((0,0), display_text, font=font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+
+        # 文本背景
+        draw.rectangle([x1, y1 - text_height - 4, x1 + text_width + 4, y1], fill="black")
+        draw.text((x1 + 2, y1 - text_height - 2), display_text, fill="white", font=font)
+
+    annotated = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
+    return annotated
+
+
+
+
+logger = logging.getLogger("yolo_service")
+
+
+class YOLOSingleton:
+    _model = None
+    _ready = False
+
+    @classmethod
+    def instance(cls):
+        """获取模型单例"""
+        if cls._model is None:
+            logger.warning("模型尚未初始化")
+        return cls._model
+
+    @classmethod
+    def is_ready(cls):
+        return cls._ready
+
+    @classmethod
+    def init_model(cls):
+        """初始化模型，失败时不抛出异常"""
+        try:
+            cls._model = YOLO("runs/detect/train13/weights/best.pt")
+            cls._model.to("cuda" if torch.cuda.is_available() else "cpu")
+            cls._ready = True
+            logger.info("✅ YOLO 模型加载完成")
+        except Exception as e:
+            cls._model = None
+            cls._ready = False
+            logger.error(f"❌ YOLO 模型加载失败: {e}")
+
+
+    @classmethod
+    def detect(cls, img_bytes: bytes):
+        if not cls._ready or cls._model is None:
+            raise RuntimeError("模型未加载或不可用")
+
+        label_map = {
+            "normal": "正茧",
+            "double_pupa": "双宫茧",
+            "spot": "黄斑茧",
+            "hairy": "毛茧",
+            "maggot_shell": "蛆壳茧"
+        }
+
+        img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
+        results = cls._model(img, conf=0.45)
+        r = results[0]
+        boxes = r.boxes
+        names = cls._model.names
+
+        total = len(boxes)
+        class_counts = {}
+        confidences = []
+        box_list = []
+        label_list = []
+
+        for idx, b in enumerate(boxes):
+            cls_id = int(b.cls)
+            conf = float(b.conf)
+            label_en = names.get(cls_id, str(cls_id))
+            label_cn = label_map.get(label_en, label_en)
+            class_counts[label_cn] = class_counts.get(label_cn, 0) + 1
+            confidences.append(conf)
+            box_list.append(b.xyxy[0])
+            label_list.append(label_cn)
+
+        # 用 PIL 绘制中文
+        annotated = draw_annotations(img, box_list, label_list, confidences)
+
+        _, buffer = cv2.imencode(".jpg", annotated)
+        img_bytes_out = buffer.tobytes()
+
+        result_json = {
+            "total_objects": total,
+            "class_counts": class_counts,
+            "min_confidence": round(min(confidences), 4) if confidences else 0.0,
+            "max_confidence": round(max(confidences), 4) if confidences else 0.0,
+            "avg_confidence": (
+                round(sum(confidences) / len(confidences), 4) if confidences else 0.0
+            ),
+            "speed_ms": r.speed,  # 直接来自 YOLO
+        }
+
+        return img_bytes_out, result_json
+
+
+
+import logging
+from pathlib import Path
+
+# 配置日志
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s"
+)
+logger = logging.getLogger("yolo_test")
+import os
+
+@staticmethod
+def save_image(img_bytes: bytes, save_path: str):
+    """
+    将检测后的字节流保存为图片文件
+
+    :param img_bytes: 图片的二进制字节（通常来自 detect() 返回）
+    :param save_path: 保存路径（例如 '/home/bbit/mine/yolo/output/result.jpg'）
+    """
+    # 确保目录存在
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+
+    # 写入文件
+    with open(save_path, "wb") as f:
+        f.write(img_bytes)
+
+def main():
+    # 初始化模型
+    YOLOSingleton.init_model()
+
+    if not YOLOSingleton.is_ready():
+        logger.error("模型未初始化成功，退出")
+        return
+
+    # 测试图片路径
+    img_path = Path("/home/bbit/mine/yolo/valid/raw_1.jpg")
+    if not img_path.exists():
+        logger.error(f"图片不存在: {img_path}")
+        return
+
+    # 读取图片为字节
+    with open(img_path, "rb") as f:
+        img_bytes = f.read()
+
+    # 执行检测
+    try:
+        img_bytes_out,result = YOLOSingleton.detect(img_bytes)
+        logger.info("检测结果:")
+        import json
+        print(json.dumps(result, indent=4))
+        save_image(img_bytes_out,'/home/bbit/mine/yolo/output/result.jpg')
+    except Exception as e:
+        logger.error(f"检测失败: {e}")
+
+if __name__ == "__main__":
+    main()

ai/yolo/export.py

@@ -0,0 +1,10 @@
+from ultralytics import YOLO
+import os
+
+# 设置工作目录
+os.chdir("/home/bbit/mine/yolo/")
+
+# 加载模型
+model = YOLO("runs/detect/train13/weights/best.pt")
+
+model.export(format="onnx")  # 默认导出 FP32 ONNX

ai/yolo/train.py

@@ -0,0 +1,34 @@
+from ultralytics import YOLO
+import os
+
+# 设置工作目录为你的项目路径
+os.chdir("/home/bbit/mine/yolo/")
+
+# Load a pretrained YOLO11n model
+model = YOLO("models/yolo11x.pt")
+
+# Train the model on the COCO8 dataset for 100 epochs
+train_results = model.train(
+    data="SilkwormCocoonDataset/data.yaml",
+    epochs=100,
+    imgsz=768,
+    batch=8,
+    device="cuda:0",
+    amp=True,
+    half=True,
+    optimizer="auto",
+    warmup_epochs=3,
+    mosaic=1.0,
+    mixup=0.1
+)
+
+
+# Evaluate the model's performance on the validation set
+metrics = model.val()
+
+# # Perform object detection on an image
+# results = model("path/to/image.jpg")  # Predict on an image
+# results[0].show()  # Display results
+
+# Export the model to ONNX format for deployment
+path = model.export(format="models/yolo11x_silkworm_cocoon")  # Returns the path to the exported model

ai/yolo/valid.py

@@ -0,0 +1,22 @@
+from ultralytics import YOLO
+import os
+
+# 设置工作目录
+os.chdir("/home/bbit/mine/yolo/")
+
+# 加载模型
+model = YOLO("runs/detect/train13/weights/best.pt")
+
+# 输入图片路径
+img_path = "valid/raw_1.jpg"
+
+# 推理
+results = model(img_path)
+
+# 生成输出路径（同目录下）
+save_path = os.path.splitext(img_path)[0] + "_detected.jpg"
+
+# 保存检测结果
+results[0].save(filename=save_path)
+
+print(f"检测结果已保存至: {save_path}")