初始化

algo/__init__.py

@@ -0,0 +1 @@
+from .model import YoloModel

algo/model.py

@@ -0,0 +1,92 @@
+from ultralytics import YOLO
+from pathlib import Path
+import cv2, subprocess, numpy as np, time
+
+
+# 开始训练回调
+def on_train_start(trainer):
+    full_path = trainer.save_dir
+    p = Path(full_path)
+    folder_name = p.parent.name
+    print('开始模型训练', folder_name)
+
+
+# 结束训练回调
+def on_train_end(trainer):
+    full_path = trainer.save_dir
+    p = Path(full_path)
+    folder_name = p.parent.name
+    print('模型训练结束', folder_name)
+
+
+# 每轮训练结束回调函数
+def on_train_epoch_end(trainer):
+    full_path = trainer.save_dir
+    p = Path(full_path)
+    folder_name = p.parent.name
+    print('当前训练轮数', trainer.epoch, '当前项目编号', folder_name)
+
+
+class YoloModel:
+
+    def __init__(self, pt_url):
+        self.model = YOLO(pt_url)
+        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)
+
+    def train(self, data, epochs, project, name):
+        train_result = self.model.train(
+            data=data,
+            epochs=epochs,
+            imgsz=640,
+            device=0,
+            project=project,
+            name=name,
+            verbose=False
+        )
+        return train_result
+
+    def detect(self, input_rtsp, output_rtsp, fps: int):
+        w, h = 1920, 1080
+        # 1. 拉流
+        pull = [
+            'ffmpeg', '-hide_banner', '-loglevel', 'error',
+            '-i', input_rtsp,
+            '-an', '-f', 'rawvideo', '-pix_fmt', 'bgr24',
+            '-vf', f'fps={fps},scale={w}:{h}', '-'
+        ]
+        pull_proc = subprocess.Popen(pull, stdout=subprocess.PIPE, bufsize=w * h * 3 * 2)
+
+        # 2. 推流（mediamtx 监听 8554）
+        push = [
+            'ffmpeg',
+            '-hide_banner',
+            '-loglevel',
+            'error',
+            '-y',
+            '-f',
+            'rawvideo',
+            '-pix_fmt', 'yuv420p',
+            '-s', f'{w}x{h}', '-r', str(fps),
+            '-i', '-', '-c:v', 'libx264', '-preset', 'ultrafast', '-tune', 'zerolatency',
+            '-f', 'flv', output_rtsp
+        ]
+        push_proc = subprocess.Popen(push, stdin=subprocess.PIPE)
+
+        while True:
+            raw = pull_proc.stdout.read(w * h * 3)
+            if len(raw) != w * h * 3:  # 网络丢包，直接跳
+                time.sleep(0.01)
+                continue
+            frame = np.frombuffer(raw, np.uint8).reshape((h, w, 3))
+
+            # 3. 推理（缩图 + stream 模式）
+            small = cv2.resize(frame, (w, h))
+            results = self.model(small, stream=True, verbose=False)
+            for r in results:
+                # 4. 画框（返回 RGB → 转回 BGR）
+                img = r.plot()
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2YUV_I420)
+                # 5. 写回
+                push_proc.stdin.write(img.tobytes())