1.配置好正常的Ultralytics源码环境
来源YOLOv8和YOLOv11的源码环境均可用,更换训练好的best.pt文件,然后后面图片的类别更改部分代码即可。
2.从训练结果找到自己的对应检测类型
本代码初始目的是检测快递盒的状态,一共五种状态, 0: 'JAP-front', 1: 'JAP-side', 2: 'USA-front', 3: 'USA-side', 4: 'other',按这种序号方式对应,下面有用。
3.需要更改的代码部位
代码需要更改的部分如下的self.class_names,序号不变,从0开始,一一对照更改成自己的检测类型名字:
class BoxDetector: def __init__(self): self.model = None self.current_annotated_image = None self.class_names = { 0: 'JAP-front', 1: 'JAP-side', 2: 'USA-front', 3: 'USA-side', 4: 'other' } def load_model(self, model_path): """加载YOLO目标检测模型""" try: self.model = YOLO(model_path) print(f"模型加载成功: {model_path}") return True except Exception as e: print(f"模型加载失败: {e}") return False def detect_objects(self, frame): """对输入图像进行目标检测""" if self.model is None: print("错误: 请先加载模型") return None, [] # 使用YOLO模型进行预测 results = self.model.predict(frame) if results: # 保存带标注的图像 self.current_annotated_image = results[0].plot() # 提取检测到的类别名称 detected_classes = [] for box in results[0].boxes: class_id = int(box.cls[0]) if class_id in self.class_names: detected_classes.append(self.class_names[class_id]) return self.current_annotated_image, detected_classes return None, []