Upload 5 files (#2)

- Upload 5 files (5aac0025bbbd31131f86cdaf86c6c545b0e5c2ab)

test_api/algorithm.py

@@ -0,0 +1,173 @@
+from ultralytics import YOLO
+import streamlit as st
+import cv2
+#import pafy
+import numpy as np
+import settings
+import requests
+
+
+
+class YoloV8Detection:
+
+    def __init__(self):
+        pass
+
+    def display_tracker_options(self):
+        col1, col2 = st.columns(2)
+        with col1:
+            display_tracker = st.radio("显示追踪", ('是', '否'))
+            is_display_tracker = True if display_tracker == '是' else False
+            if is_display_tracker:
+                with col2:
+                    tracker_type = st.radio("追踪器选择", ("bytetrack.yaml", "botsort.yaml"))
+                    return is_display_tracker, tracker_type
+        return is_display_tracker, None
+
+
+
+    def _display_detected_frames(self,conf, model, st_frame, image, is_display_tracking=None, tracker=None):
+
+
+        image = cv2.resize(image, (720, int(720*(9/16))))
+
+        if is_display_tracking:
+            res = model.track(image, conf=conf, persist=True, tracker=tracker)
+        else:
+            res = model.predict(image, conf=conf)
+
+        res_plotted = res[0].plot()
+        try:
+            st_frame.image(res_plotted,
+                        caption='实时检测',
+                        channels="BGR",
+                        use_column_width=True
+                        )
+        except requests.exceptions.RequestException as e:
+            st.write("Unable to get image, using placeholder")
+            st.image("placeholder.png")
+        
+
+    def play_rtsp_stream(self,conf, model,display_video=False):
+        
+        source_rtsp = st.sidebar.text_input("rtsp stream url",value="http://127.0.0.1:8999/live/test.live.flv")
+        is_display_tracker, tracker = self.display_tracker_options()
+
+        start = st.sidebar.button('检测目标')
+        stop = st.sidebar.button('停止') 
+
+        if start:
+            try:
+                vid_cap = cv2.VideoCapture(source_rtsp)
+                st_frame = st.empty()
+
+                while (vid_cap.isOpened()):
+                    success, image = vid_cap.read()
+                    if success:
+                        #print("Success")
+                        if stop:
+                            break
+                        
+                        elif display_video:
+                            self._display_detected_frames(conf,
+                                                    model,
+                                                    st_frame,
+                                                    image,
+                                                    is_display_tracker,
+                                                    tracker
+                                                    )
+
+                    else:
+                        vid_cap.release()
+                        print("Error")
+                        break
+            except Exception as e:
+                st.sidebar.error("Error loading RTSP stream: " + str(e))
+        
+            
+
+class BoundaryDetection(YoloV8Detection):
+    def __init__(self,intrusion_area= [(100, 200), (400, 200), (500, 300), (100, 300)]):
+        self.intrusion_area = intrusion_area
+
+    def _display_detected_frames(self, conf, model, st_frame, image, is_display_tracking=None, tracker=None):
+        
+        image = cv2.resize(image, (720, int(720*(9/16))))
+
+        if is_display_tracking:
+            res = model.track(image, conf=conf, persist=True, tracker=tracker)
+        else:
+            res = model.predict(image, conf=conf)
+
+        cv2.polylines(image, [np.array(self.intrusion_area)], isClosed=True, color=(0, 0, 255), thickness=2)
+        intrusion_detected = False
+
+        for result in res:
+            # for xyxy_box in result.boxes.xyxy:
+            print(f"boxes是:{result.boxes}\n")
+
+            box = result.boxes.xyxy 
+             
+            print(f"box是:{box}\n")
+            if box.size(0) != 0:
+                for (xyxy,obj) in zip(box,result.boxes):
+                    x, y, x2, y2 = int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3])
+                    cx, cy = (x + x2) / 2, (y + y2) / 2
+
+                    print(f"中点坐标是:{cx, cy}\n")
+                    if cv2.pointPolygonTest(np.array(self.intrusion_area), (cx, cy), False) >= 0:
+                    # 有人进入特定区域，进行报警
+                        print("有人入侵!\n")
+                        intrusion_detected = True
+
+                    # 添加框出入侵者的代码
+                        cv2.rectangle(image, (x, y), (x2, y2), (0, 255, 0), 2)
+                        cv2.putText(image, f"{result.names[int(obj.cls[0])]} {obj.conf[0]:.2f}", (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
+                        
+            else:
+                continue
+
+        if intrusion_detected:
+            cv2.putText(image, "Intrusion detected!", (10, 30), cv2.FONT_HERSHEY_TRIPLEX, 1, (0, 0, 255), 2)
+
+        # res_plotted = res[0].plot()
+        # res_plotted=cv2.polylines(res_plotted, [np.array(self.intrusion_area)], isClosed=True, color=(0, 0, 255), thickness=2)
+
+
+        try:
+            st_frame.image(image,
+                        caption='实时检测',
+                        channels="BGR",
+                        use_column_width=True
+                        )
+        except requests.exceptions.RequestException as e:
+            st.write("Unable to get image, using placeholder")
+            st.image("placeholder.png")
+
+
+    def play_rtsp_stream(self, conf, model,display_video=False):
+        source_rtsp = st.sidebar.text_input("rtsp stream url")
+        is_display_tracker, tracker = self.display_tracker_options()
+        if st.sidebar.button('检测目标'):
+            try:
+                vid_cap = cv2.VideoCapture(source_rtsp)
+                st_frame = st.empty()
+                while (vid_cap.isOpened()):
+                    success, image = vid_cap.read()
+                    if success:
+                        
+                        if display_video:
+                            self._display_detected_frames(conf,
+                                                    model,
+                                                    st_frame,
+                                                    image,
+                                                    is_display_tracker,
+                                                    tracker
+                                                    )
+                    else:
+                        vid_cap.release()
+                        break
+            except Exception as e:
+                st.sidebar.error("Error loading RTSP stream: " + str(e))
+
+

test_api/app.py

@@ -0,0 +1,154 @@
+
+from pathlib import Path
+import PIL
+
+# External packages
+import streamlit as st
+
+# Local Modules
+import settings
+import helper
+import test_api.algorithm as algorithm
+import multiprocessing 
+import time
+import requests
+import socketio
+import base64
+import io
+from PIL import Image
+
+
+st.set_page_config(
+    page_title="YOLOv8 目标检测",
+    page_icon="🤖",
+    layout="wide",
+    initial_sidebar_state="expanded"  # 或者 "collapsed"
+)
+
+
+# # Main page heading
+st.title("YOLOv8 目标检测")
+
+# Sidebar
+st.sidebar.header("模型配置")
+
+# Model Options
+model_type = st.sidebar.selectbox(
+    "任务选择", ['检测', '分割',"越界检测","行为检测"])
+
+confidence = float(st.sidebar.slider(
+    "选择模型Confidence", 25, 100, 40)) / 100
+
+# Selecting Detection Or Segmentation
+if model_type == '检测':
+    model_path = Path(settings.DETECTION_MODEL)
+elif model_type == '分割':
+    model_path = Path(settings.SEGMENTATION_MODEL)
+elif model_type == "越界检测":
+    model_path = Path(settings.SEGMENTATION_MODEL)
+elif model_type == "行为检测":
+    model_path = Path(settings.SEGMENTATION_MODEL)
+
+# Load Pre-trained ML Model
+try:
+    model = helper.load_model(model_path)
+except Exception as ex:
+    st.error(f"Unable to load model. Check the specified path: {model_path}")
+    st.error(ex)
+
+st.sidebar.header("图像/视频 配置")
+source_radio = st.sidebar.radio(
+    "选择来源", settings.SOURCES_LIST)
+
+source_img = None
+# If image is selected
+if source_radio == settings.IMAGE:
+    source_img = st.sidebar.file_uploader(
+        "选择一张图像...", type=("jpg", "jpeg", "png", 'bmp', 'webp'))
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        try:
+            if source_img is None:
+                default_image_path = str(settings.DEFAULT_IMAGE)
+                default_image = PIL.Image.open(default_image_path)
+                st.image(default_image_path, caption="默认图像",
+                         use_column_width=True)
+            else:
+                uploaded_image = PIL.Image.open(source_img)
+                st.image(source_img, caption="Uploaded Image",
+                         use_column_width=True)
+        except Exception as ex:
+            st.error("Error occurred while opening the image.")
+            st.error(ex)
+
+    with col2:
+        if source_img is None:
+            default_detected_image_path = str(settings.DEFAULT_DETECT_IMAGE)
+            default_detected_image = PIL.Image.open(
+                default_detected_image_path)
+            st.image(default_detected_image_path, caption='检测图像',
+                     use_column_width=True)
+        else:
+            if st.sidebar.button('检测目标'):
+                res = model.predict(uploaded_image,
+                                    conf=confidence
+                                    )
+                boxes = res[0].boxes
+                res_plotted = res[0].plot()[:, :, ::-1]
+                st.image(res_plotted, caption='Detected Image',
+                         use_column_width=True)
+                try:
+                    with st.expander("Detection Results"):
+                        for box in boxes:
+                            st.write(box.data)
+                except Exception as ex:
+                    # st.write(ex)
+                    st.write("No image is uploaded yet!")
+
+
+
+
+elif source_radio == settings.RTSP:
+    if model_type == '检测':
+        src ={
+	"video_url": "http://127.0.0.1:8999/live/test.live.flv"
+}
+        start = st.sidebar.button('检测目标')
+        stop = st.sidebar.button('停止') 
+
+        if start:
+            response = requests.post('http://192.168.110.232:7555/analyzerControlAdd', json=src)
+            print(response.text)
+
+    elif model_type == '分割':
+        yolov8=algorithm.YoloV8Detection()
+        yolov8.play_rtsp_stream(confidence,model,display_video=True)
+
+    elif model_type == "越界检测":
+        yolov8_b=algorithm.BoundaryDetection()
+        yolov8_b.play_rtsp_stream(confidence,model,display_video=True)
+
+
+    #helper.play_rtsp_stream(confidence, model)
+
+else:
+    st.error("Please select a valid source type!")
+
+
+# 创建 SocketIO 实例
+sio = socketio.Client()
+
+# 连接到 WebSocket 服务器
+sio.connect('http://192.168.110.232:7555')
+# 接收帧数据并显示
+@sio.on('frame_data')
+def handle_frame_data(data):
+    image_bytes = base64.b64decode(data['image'])
+    image = Image.open(io.BytesIO(image_bytes))
+    st.image(image, caption='Real-time Detection', channels='BGR')
+
+# 等待连接断开
+st.text('Waiting for real-time updates...')
+sio.wait()

test_api/helper.py

@@ -0,0 +1,80 @@
+from ultralytics import YOLO
+import streamlit as st
+import cv2
+#import pafy
+import requests
+import test_api.settings as settings
+import numpy as np
+from test_api.algorithm import *
+
+def load_model(model_path):
+
+    model = YOLO(model_path)
+    return model
+
+    
+def display_tracker_options():
+    col1, col2 = st.columns(2)
+    with col1:
+        display_tracker = st.radio("显示追踪", ('Yes', 'No'))
+        is_display_tracker = True if display_tracker == 'Yes' else False
+        if is_display_tracker:
+            with col2:
+                tracker_type = st.radio("追踪器", ("bytetrack.yaml", "botsort.yaml"))
+                return is_display_tracker, tracker_type
+    return is_display_tracker, None
+
+
+
+def _display_detected_frames(conf, model, st_frame, image, is_display_tracking=None, tracker=None):
+
+
+    # Resize the image to a standard size
+    image = cv2.resize(image, (720, int(720*(9/16))))
+
+    # Display object tracking, if specified
+    if is_display_tracking:
+        res = model.track(image, conf=conf, persist=True, tracker=tracker)
+    else:
+        # Predict the objects in the image using the YOLOv8 model
+        res = model.predict(image, conf=conf)
+
+    # # Plot the detected objects on the video frame
+    res_plotted = res[0].plot()
+    try:
+       st_frame.image(res_plotted,
+                   caption='实时检测',
+                   channels="BGR",
+                   use_column_width=True
+                   )
+    except requests.exceptions.RequestException as e:
+        st.write("Unable to get image, using placeholder")
+        st.image("placeholder.png")
+        
+
+def play_rtsp_stream(conf, model):
+    source_rtsp = st.sidebar.text_input("rtsp stream url")
+    is_display_tracker, tracker = display_tracker_options()
+    if st.sidebar.button('检测目标'):
+        try:
+            vid_cap = cv2.VideoCapture(source_rtsp)
+            st_frame = st.empty()
+            while (vid_cap.isOpened()):
+                success, image = vid_cap.read()
+                if success:
+                    _display_detected_frames(conf,
+                                             model,
+                                             st_frame,
+                                             image,
+                                             is_display_tracker,
+                                             tracker
+                                             )
+                else:
+                    vid_cap.release()
+                    break
+        except Exception as e:
+            st.sidebar.error("Error loading RTSP stream: " + str(e))
+
+
+
+

test_api/yoloflaskapi.py

@@ -0,0 +1,96 @@
+import argparse
+import sys
+import base64
+import json
+import cv2
+import numpy as np
+import torch
+from pathlib import Path
+from ultralytics import YOLO
+import json
+from flask import Flask,request_finished
+from flask import Flask, request, jsonify
+import threading
+import time
+from flask import Flask
+from flask_cors import CORS
+from queue import Queue
+from flask_socketio import SocketIO, emit
+import pickle
+
+app = Flask(__name__)
+CORS(app)
+socketio = SocketIO(app)
+
+algorithm_running = False
+
+data = {
+    "code": 0,
+    "msg": "unknown error",
+}
+
+@socketio.on('connect')
+def handle_connect():
+    print('Client connected')
+
+def algorithm_thread(source_rtsp):
+    global algorithm_running
+    global data
+
+
+    
+    vid_cap = cv2.VideoCapture(source_rtsp)
+
+    while algorithm_running:
+        success, image = vid_cap.read()
+        
+        if success:
+            
+            res = model(image)
+            res_plotted = res[0].plot()
+            # 将帧转换为 base64 编码的字符串
+            _, buffer = cv2.imencode('.jpg', res_plotted)
+            frame_bytes = buffer.tobytes()
+            frame_base64 = base64.b64encode(frame_bytes).decode('utf-8')
+
+            # 发送帧数据到客户端
+            socketio.emit('frame_data', {'image': frame_base64})
+
+    vid_cap.release()
+
+@app.route("/analyzerControlAdd",methods=['POST'])
+def start_algorithm():
+    global algorithm_running
+    try:
+        params = request.get_json()
+    except:
+        params = request.form
+    source_rtsp = params.get("video_url")
+    print(f"video_url",source_rtsp)
+    if not algorithm_running :
+        algorithm_running = True
+        print("start algorithm")
+        threading.Thread(target=algorithm_thread(source_rtsp)).start()
+
+    
+
+
+
+@app.route("/analyzerControlCancel", methods=['POST'])
+def stop_algorithm():
+    global algorithm_running
+    global data
+    algorithm_running = False
+
+
+    return json.dumps({"code": 1000,"msg": "算法停止！"}, ensure_ascii=False)
+
+if __name__ == "__main__":
+
+
+    model = YOLO('../weights/yolov8n.pt')
+
+    app.run(host="0.0.0.0",port=7555)
+
+
+