diff --git a/ai/yolo/app_test.py b/ai/yolo/app_test.py new file mode 100644 index 0000000..93202e4 --- /dev/null +++ b/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() diff --git a/ai/yolo/export.py b/ai/yolo/export.py new file mode 100644 index 0000000..9bfa1ee --- /dev/null +++ b/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 diff --git a/ai/yolo/train.py b/ai/yolo/train.py new file mode 100644 index 0000000..04630aa --- /dev/null +++ b/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 \ No newline at end of file diff --git a/ai/yolo/valid.py b/ai/yolo/valid.py new file mode 100644 index 0000000..1071642 --- /dev/null +++ b/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}")