项目初始化

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/chache_rtsp.iml

@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$">
+      <excludeFolder url="file://$MODULE_DIR$/venv" />
+    </content>
+    <orderEntry type="jdk" jdkName="Python 3.10 (chache-ult)" jdkType="Python SDK" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,41 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="DuplicatedCode" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <Languages>
+        <language minSize="250" name="Python" />
+      </Languages>
+    </inspection_tool>
+    <inspection_tool class="Eslint" enabled="true" level="WARNING" enabled_by_default="true" />
+    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
+      <option name="ignoredPackages">
+        <value>
+          <list size="8">
+            <item index="0" class="java.lang.String" itemvalue="protobuf" />
+            <item index="1" class="java.lang.String" itemvalue="mkl-random" />
+            <item index="2" class="java.lang.String" itemvalue="numpy" />
+            <item index="3" class="java.lang.String" itemvalue="starlette" />
+            <item index="4" class="java.lang.String" itemvalue="charset-normalizer" />
+            <item index="5" class="java.lang.String" itemvalue="h11" />
+            <item index="6" class="java.lang.String" itemvalue="obs" />
+            <item index="7" class="java.lang.String" itemvalue="torch" />
+          </list>
+        </value>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyPep8Inspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="W292" />
+        </list>
+      </option>
+    </inspection_tool>
+    <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
+      <option name="ignoredErrors">
+        <list>
+          <option value="N803" />
+        </list>
+      </option>
+    </inspection_tool>
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/misc.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.10 (chache-ult)" project-jdk-type="Python SDK" />
+</project>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/chache_rtsp.iml" filepath="$PROJECT_DIR$/.idea/chache_rtsp.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="$PROJECT_DIR$" vcs="Git" />
+  </component>
+</project>

algo/__init__.py

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

algo/yolov11.py

@@ -0,0 +1,17 @@
+from ultralytics import YOLO
+
+
+class YoloModel:
+
+    def __init__(self, pt_url=None):
+        if pt_url:
+            self.model = YOLO(pt_url)
+        else:
+            self.model = YOLO('yolo11n.pt')
+
+    def predict(self, source):
+        """
+        对rtsp视频流进行预测
+        """
+        results = self.model.predict(source=source, conf=0.6, device='cpu')
+        return results

chache/chache.yaml

@@ -0,0 +1,3 @@
+names:
+  0: forklift
+  1: person

config.ini

@@ -0,0 +1,4 @@
+[config]
+urls=rtsp://admin:xkrs0425@172.16.20.15:554/D6/sub/av_stream,rtsp://admin:qd123456@172.16.20.13:554/D4/sub/av_stream
+wss=ws://172.16.10.227:8099/websocket/message/rtsp/d6,ws://172.16.10.227:8099/websocket/message/rtsp/d5
+api=http://172.16.10.227:8099/business/t_monitor_ai/add

cut.py

@@ -0,0 +1,40 @@
+import cv2
+import time
+
+# RTSP流地址
+rtsp_url = "rtsp://admin:xkrs0425@172.16.20.15:554/D6/main/av_stream"
+
+# 通过opencv创建一个视频捕获对象
+cap = cv2.VideoCapture(rtsp_url)
+
+# 检查是否成功打开
+if not cap.isOpened():
+    print("Error: Could not open video stream from provided RTSP URL.")
+else:
+    print("Successfully opened RTSP stream, starting to capture frames...")
+
+    # 等待几秒让摄像头预热（可选）
+    time.sleep(2)
+
+    # 循环读取帧
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            print("Failed to grab frame")
+            break
+
+        # 显示当前帧（可选）
+        cv2.imshow('RTSP Stream', frame)
+
+        # 保存截图到文件
+        img_name = f"D:\\syg\\jk2\\frame_{int(time.time())}.png"
+        cv2.imwrite(img_name, frame)
+        print(f"{img_name} saved!")
+
+        # 按键检测，按下 'q' 键退出循环
+        if cv2.waitKey(1) & 0xFF == ord('q'):
+            break
+
+    # 结束后释放资源
+    cap.release()
+    cv2.destroyAllWindows()

ffmpeg_demo.py

@@ -0,0 +1,38 @@
+import subprocess
+import numpy as np
+from pathlib import Path
+from algo import YoloModel
+
+
+def run():
+    FILE = Path(__file__).resolve()
+    ROOT = FILE.parents[0]  # root directory
+    model = YoloModel(ROOT / "chache/best.pt")
+    width, height = 640, 640
+    fps = 0.2  # 每 5 秒一帧
+    command = [
+        'ffmpeg',
+        '-i', 'rtsp://admin:xkrs0425@172.16.20.15:554/D6/sub/av_stream',
+        '-an',  # 不处理音频
+        '-f', 'image2pipe',  # 输出为图像流
+        '-pix_fmt', 'bgr24',  # OpenCV 默认是 BGR 格式
+        '-r', str(fps),  # 每秒 0.2 帧 => 每 5 秒一帧
+        '-vcodec', 'rawvideo',
+        '-'
+    ]
+    # 启动 FFmpeg 子进程
+    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=10 ** 8)
+    while True:
+        # 读取一帧原始数据（大小 = width * height * 3）0
+        raw_frame = pipe.stdout.read(width * height * 3)
+        if not raw_frame:
+            break
+        # 转换为 NumPy 数组并 reshape 成图像
+        frame = np.frombuffer(raw_frame, dtype=np.uint8).reshape((height, width, 3))
+        print(f"Frame shape: {frame.shape}, dtype: {frame.dtype}, min/max: {frame.min(), frame.max()}")
+        # # 使用 YOLO 进行推理
+        result = model.predict(frame)
+
+
+if __name__ == '__main__':
+    run()

jk/jk.yaml

@@ -0,0 +1,3 @@
+names:
+  0: car
+  1: person

main.py

@@ -0,0 +1,96 @@
+import cv2
+import requests
+import threading
+import subprocess
+import numpy as np
+from pathlib import Path
+from collections import Counter
+from websocket import WebSocketApp
+from configparser import ConfigParser
+
+from algo import YoloModel
+
+
+FILE = Path(__file__).resolve()
+ROOT = FILE.parents[0]
+model = YoloModel(ROOT / 'chache/best.pt')
+
+
+def run_client(client: WebSocketApp):
+    client.run_forever()
+
+
+def get_config():
+    # 获取websocket的配置连接
+    config_path = ROOT / "config.ini"
+    config = ConfigParser()
+    config.read(config_path)
+    rtsp_url = config.get('config', 'urls')
+    wss = config.get('config', 'wss')
+    api_url = config.get('config', 'api')
+    return wss.split(','), rtsp_url.split(','), api_url
+
+
+def process_camera(rtsp_url, ws_url, api_url=None):
+    websocket = WebSocketApp(url=ws_url, on_open=on_open)
+    threading.Thread(target=run_client, args=(websocket,), daemon=True).start()
+    width, height = 1920, 1080
+    command = [
+        'ffmpeg',
+        '-i', rtsp_url,
+        '-an',  # 不处理音频
+        '-f', 'image2pipe',  # 输出为图像流
+        '-pix_fmt', 'bgr24',  # OpenCV 默认是 BGR
+        '-vf', f'fps=1,scale={width}:{height}',  # 每秒 1 帧 + 缩放
+        '-vcodec', 'rawvideo',
+        '-'
+    ]
+    # 启动 FFmpeg 子进程
+    pipe = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=10 ** 8)
+
+    while True:
+        # 读取一帧原始数据（大小 = width * height * 3）0
+        raw_frame = pipe.stdout.read(width * height * 3)
+        if not raw_frame:
+            break
+        # 转换为 NumPy 数组并 reshape 成图像
+        frame = np.frombuffer(raw_frame, dtype=np.uint8).reshape((height, width, 3))
+        # # 使用 YOLO 进行推理
+        results = model.predict(frame)
+        img = results[0].plot()
+        ret, jpeg = cv2.imencode('.jpg', img)
+        if ret:
+            frame_data = jpeg.tobytes()
+            websocket.send_bytes(frame_data)
+        if api_url is not None:
+            # 获取类别索引
+            classes = results[0].boxes.cls.cpu().numpy().astype(int)
+            # 统计每个类别的数量
+            class_counts = Counter(classes)
+            forklift_count = class_counts.get(0, 0)
+            person_count = class_counts.get(1, 0)
+            try:
+                data = {
+                    "result": f"0, {person_count}, {forklift_count}"
+                }
+                response = requests.post(api_url, json=data)
+                if response.status_code == 200:
+                    print("接口调用success")
+                else:
+                    print("接口error")
+            except requests.exceptions.RequestException as exc:
+                print("接口except")
+
+
+def on_open(ws):
+    print(ws.url, '已连接')
+
+
+if __name__ == '__main__':
+    websocket_urls, rtsp_urls, api_url = get_config()
+    for i in range(len(websocket_urls)):
+        if i == 0:
+            thread = threading.Thread(target=process_camera, args=(rtsp_urls[i], websocket_urls[i], api_url))
+        else:
+            thread = threading.Thread(target=process_camera, args=(rtsp_urls[i], websocket_urls[i], None))
+        thread.start()

requirement.txt

@@ -0,0 +1,4 @@
+# yolo
+ultralytics==8.3.151
+
+websocket-client

start.bat

@@ -0,0 +1,12 @@
+@echo off
+:: 激活 chache 虚拟环境
+call conda activate chache
+
+:: 切换到 .bat 文件所在目录（main.py 所在目录）
+cd /d %~dp0%
+
+:: 运行主程序
+python main.py
+
+:: 防止窗口关闭，方便查看输出结果
+pause