提高版本到yolov11

algo/yolov11.py

@@ -0,0 +1,113 @@
+import time
+import asyncio
+from pathlib import Path
+from ultralytics import YOLO
+
+from utils.websocket_server import room_manager
+
+
+# 开始训练回调
+def on_train_start(trainer):
+    full_path = trainer.save_dir
+    p = Path(full_path)
+    room_name = 'train_' + p.parent.name
+    asyncio.run(room_manager.send_to_room(room_name, 'start'))
+
+
+# 结束训练回调
+def on_train_end(trainer):
+    full_path = trainer.save_dir
+    p = Path(full_path)
+    room_name = 'train_' + p.parent.name
+    asyncio.run(room_manager.send_to_room(room_name, 'end'))
+
+
+# 每轮训练结束回调函数
+def on_train_epoch_end(trainer):
+    full_path = trainer.save_dir
+    p = Path(full_path)
+    room_name = 'train_' + p.parent.name
+    asyncio.run(room_manager.send_to_room(room_name, trainer.epoch))
+
+
+# 预测开始回调
+def on_predict_start(predictor):
+    full_path = predictor.save_dir
+    p = Path(full_path)
+    room_name = 'detect_' + p.parent.parent.name
+    asyncio.run(room_manager.send_to_room(room_name, 'start'))
+
+
+# 每个批次预测结束回调
+def on_predict_batch_end(predictor):
+    full_path = predictor.save_dir
+    p = Path(full_path)
+    room_name = 'detect_' + p.parent.parent.name
+    asyncio.run(room_manager.send_to_room(room_name, predictor.seen))
+
+
+# 预测结束回调
+def on_predict_end(predictor):
+    full_path = predictor.save_dir
+    p = Path(full_path)
+    room_name = 'detect_' + p.parent.parent.name
+    asyncio.run(room_manager.send_to_room(room_name, 'end'))
+
+
+class YoloModel:
+
+    def __init__(self, pt_url=None):
+        if pt_url:
+            self.model = YOLO(pt_url)
+        else:
+            self.model = YOLO('yolo11n.pt')
+
+    def train(self, data, epochs, project, name, patience):
+        """
+        模型训练
+        """
+        self.model.add_callback('on_train_start', on_train_start)
+        self.model.add_callback('on_train_end', on_train_end)
+        self.model.add_callback('on_train_epoch_end', on_train_epoch_end)
+        self.model.train(
+            data=data,
+            epochs=epochs,
+            imgsz=640,
+            device=0,
+            project=project,
+            name=name,
+            patience=patience,
+            verbose=False
+        )
+
+    def predict_folder(self, source, name, project):
+        """
+        对文件夹中的内容（图片或者视频）进行预测
+        """
+        self.model.add_callback('on_predict_start', on_predict_start)
+        self.model.add_callback('on_predict_batch_end', on_predict_batch_end)
+        self.model.add_callback('on_predict_end', on_predict_end)
+        self.model.predict(
+            source=source,
+            name=name,
+            project=project,
+            save=True,
+            save_txt=True,
+            device='0',
+            conf=0.6
+        )
+
+    def predict_rtsp(self, source, detect_id):
+        """
+        对rtsp视频流进行预测
+        """
+        room_name = 'detect_rtsp_' + str(detect_id)
+        start_time = time.time()
+        for result in self.model.predict(source=source, stream=True, device='0', conf=0.6):
+            # 检查是否已经超过10分钟（600秒）
+            elapsed_time = time.time() - start_time
+            if elapsed_time > 600:  # 600 seconds = 10 minutes
+                break
+            frame = result.plot()
+            frame_data = frame.tobytes()
+            asyncio.run(room_manager.send_stream_to_room(room_name, frame_data))