初始化

2025-09-15 14:37:32 +08:00
commit 58ab12eaba
7 changed files with 155 additions and 0 deletions
--- a/admin/index.html
+++ b/admin/index.html
@@ -0,0 +1,44 @@
+<!DOCTYPE html>
+<html lang="zh">
+<head>
+  <meta charset="UTF-8">
+  <title>FLV 流播放</title>
+  <script src="/js/flv.js"></script>
+  <style>
+    body {
+      background: #000;
+      display: flex;
+      justify-content: center;
+      align-items: center;
+      height: 100vh;
+      margin: 0;
+    }
+    video {
+      width: 80%;
+      max-width: 960px;
+      background: #000;
+    }
+  </style>
+</head>
+<body>
+  <video id="videoElement" controls autoplay></video>
+
+  <script>
+    if (flvjs.isSupported()) {
+      const videoElement = document.getElementById('videoElement');
+      const flvPlayer = flvjs.createPlayer({
+        type: 'flv',
+        url: 'http://172.16.10.97:8000/live/stream.flv', // 你的流地址
+        isLive: true,   // 表示是直播流
+        cors: true,     // 如果跨域要设置
+        hasAudio: false // 如果你推流里没有音频，可以加上这个
+      });
+      flvPlayer.attachMediaElement(videoElement);
+      flvPlayer.load();
+      flvPlayer.play();
+    } else {
+      alert("当前浏览器不支持 flv.js 播放 FLV 流");
+    }
+  </script>
+</body>
+</html>
--- a/admin/js/flv.js
+++ b/admin/js/flv.js
--- a/algo/init.py
+++ b/algo/init.py
@@ -0,0 +1 @@
+from .model import YoloModel
--- a/algo/model.py
+++ b/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())
--- a/main.py
+++ b/main.py
@@ -0,0 +1,6 @@
+from algo import YoloModel
+
+
+if __name__ == '__main__':
+    model = YoloModel('yolo11n.pt')
+    model.detect('rtsp://admin:xkrs0425@172.16.20.15:554/D6/main/av_stream', 'rtmp://localhost/live/stream', 15)
--- a/requirement.txt
+++ b/requirement.txt
@@ -0,0 +1,2 @@
+ultralytics==8.3.151
+ffmpeg-python
--- a/yolo11n.pt
+++ b/yolo11n.pt