commit eb17d1f7f5400cf67b15c7a61a72e8bd7f470d68 Author: sunyugang <sun_shine1229@163.com> Date: Thu Jun 26 16:48:49 2025 +0800 项目初始化 diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..13566b8 --- /dev/null +++ b/.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 diff --git a/.idea/chache_rtsp.iml b/.idea/chache_rtsp.iml new file mode 100644 index 0000000..39816ea --- /dev/null +++ b/.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> \ No newline at end of file diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml new file mode 100644 index 0000000..22230de --- /dev/null +++ b/.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> \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml new file mode 100644 index 0000000..105ce2d --- /dev/null +++ b/.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> \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100644 index 0000000..32aca06 --- /dev/null +++ b/.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> \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..9b3d31b --- /dev/null +++ b/.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> \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.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> \ No newline at end of file diff --git a/algo/__init__.py b/algo/__init__.py new file mode 100644 index 0000000..944d0ee --- /dev/null +++ b/algo/__init__.py @@ -0,0 +1 @@ +from .yolov11 import YoloModel \ No newline at end of file diff --git a/algo/yolov11.py b/algo/yolov11.py new file mode 100644 index 0000000..109ca1e --- /dev/null +++ b/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 \ No newline at end of file diff --git a/chache/best.pt b/chache/best.pt new file mode 100644 index 0000000..fca823f Binary files /dev/null and b/chache/best.pt differ diff --git a/chache/chache.yaml b/chache/chache.yaml new file mode 100644 index 0000000..afb46ff --- /dev/null +++ b/chache/chache.yaml @@ -0,0 +1,3 @@ +names: + 0: forklift + 1: person \ No newline at end of file diff --git a/config.ini b/config.ini new file mode 100644 index 0000000..525c3d4 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/cut.py b/cut.py new file mode 100644 index 0000000..28b8136 --- /dev/null +++ b/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() diff --git a/ffmpeg_demo.py b/ffmpeg_demo.py new file mode 100644 index 0000000..31d7334 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/jk/best.pt b/jk/best.pt new file mode 100644 index 0000000..81aa5a4 Binary files /dev/null and b/jk/best.pt differ diff --git a/jk/jk.yaml b/jk/jk.yaml new file mode 100644 index 0000000..c5d1add --- /dev/null +++ b/jk/jk.yaml @@ -0,0 +1,3 @@ +names: + 0: car + 1: person diff --git a/main.py b/main.py new file mode 100644 index 0000000..0b979ac --- /dev/null +++ b/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() \ No newline at end of file diff --git a/requirement.txt b/requirement.txt new file mode 100644 index 0000000..a99d823 --- /dev/null +++ b/requirement.txt @@ -0,0 +1,4 @@ +# yolo +ultralytics==8.3.151 + +websocket-client \ No newline at end of file diff --git a/start.bat b/start.bat new file mode 100644 index 0000000..0ead7ec --- /dev/null +++ b/start.bat @@ -0,0 +1,12 @@ +@echo off +:: 激活 chache 虚拟环境 +call conda activate chache + +:: 切换到 .bat 文件所在目录(main.py 所在目录) +cd /d %~dp0% + +:: 运行主程序 +python main.py + +:: 防止窗口关闭,方便查看输出结果 +pause \ No newline at end of file