sun15684712291/ultralytics_demo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

初始化

Browse Source
This commit is contained in:
sunyg
2025-09-15 14:37:32 +08:00
commit 58ab12eaba
7 changed files with 155 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
admin/index.html Normal file
View File

@@ -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>

10
admin/js/flv.js Normal file
View File

File diff suppressed because one or more lines are too long

1
algo/__init__.py Normal file
View File

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

92
algo/model.py Normal file
View File

@@ -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())

6
main.py Normal file
View File

@@ -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)

2
requirement.txt Normal file
View File

@@ -0,0 +1,2 @@
ultralytics==8.3.151
ffmpeg-python

BIN
yolo11n.pt Normal file
View File

Binary file not shown.