zfbox/web/API/viedo.py

import cv2
import asyncio
import av
import torch
import json
import time
import os
import sys
import numpy as np
from . import api
from quart import jsonify, request,render_template, Response,websocket
from aiortc import RTCPeerConnection, RTCSessionDescription, VideoStreamTrack
from aiortc.contrib.media import MediaBlackhole, MediaPlayer, MediaRecorder
from fractions import Fraction
from shapely.geometry import Point, Polygon
import threading
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
rtsp_url = 'rtsp://127.0.0.1/live1'
#-------------------基于WEBRTC实现拉流
pcs = set()     #创建一个空的集合，去重复且无序

'''
--------------基础信息--后续封装在函数里---------------
'''
#print(f"Current working directory (video.py): {os.getcwd()}")
# 自定义模型文件的路径
model_path = 'model/mode_test/yolov5s.pt'  # 假设模型文件名为 yolov5s.pt 并与执行文件在同一个目录
# 本地YOLOv5仓库路径
repo_path = 'model/base_model/yolov5'

# 加载自定义模型
model = torch.hub.load(repo_path, 'custom', path=model_path, source='local')
# 定义监控区域（多边形顶点）
region_points = [(100, 100), (500, 100), (500, 400), (100, 400)]
polygon = Polygon(region_points)

# 打开摄像头
def get_camera():
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise IOError("Cannot open webcam")
    return cap

# cap = get_camera()
'''
---------------------传输--------------------------
'''
class VideoTransformTrack(VideoStreamTrack):
    kind = "video"
    def __init__(self,strUrl,itype=1):  #0-usb,1-RTSP,2-海康SDK
        super().__init__()
        self.cap = cv2.VideoCapture(strUrl)
        if not self.cap.isOpened():
            raise Exception("Unable to open video source")
        # Set desired resolution and frame rate
        # self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)  # Set width to 1280
        # self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)  # Set height to 720
        # self.cap.set(cv2.CAP_PROP_FPS, 30)  # Set frame rate to 30 FPS
        self.frame_rate = int(self.cap.get(cv2.CAP_PROP_FPS)) or 15  # Default to 30 if unable to get FPS
        self.time_base = Fraction(1, self.frame_rate)

    async def recv(self):
        start_time = time.time()  # 开始时间
        ret, frame = self.cap.read()
        if not ret:
            raise Exception("Unable to read frame")
        # # Convert frame to RGB
        # frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        #
        # # Convert to VideoFrame
        # video_frame = av.VideoFrame.from_ndarray(frame_rgb, format="rgb24")
        # # Set timestamp
        # video_frame.pts = int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))
        # video_frame.time_base = self.time_base

        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        #img = frame.to_ndarray(format="bgr24")
        # 使用 Yolov5 进行目标检测
        results = model(img)
        detections = results.pandas().xyxy[0].to_dict(orient="records")
        # 绘制检测结果
        for det in detections:
            if det['name'] == 'person':  # 只检测人员
                x1, y1, x2, y2 = int(det['xmin']), int(det['ymin']), int(det['xmax']), int(det['ymax'])
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(img, f"{det['name']} {det['confidence']:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9,
                            (0, 255, 0), 2)
        # 将检测结果图像转换为帧
        new_frame = av.VideoFrame.from_ndarray(img, format="rgb24")
        new_frame.pts = int(self.cap.get(cv2.CAP_PROP_POS_FRAMES))
        new_frame.time_base = self.time_base
        end_time = time.time()  # 结束时间
        print(f"Processing time: {end_time - start_time} seconds")
        return new_frame

@api.route('/offer', methods=['POST'])
async def offer():
    #接收客户端的连接请求
    params = await request.json
    url_type = params["url_type"]
    url = params["url"]
    #提取客户端发来的SDP，生成服务器端的SDP
    offer = RTCSessionDescription(sdp=params["sdp"], type=params["type"])
    # # 配置 STUN 服务器
    # ice_servers = [{"urls": ["stun:stun.voipbuster.com"]}]
    # pc = RTCPeerConnection(configuration={"iceServers": ice_servers})
    pc = RTCPeerConnection()       #实例化一个rtcp对象
    pcs.add(pc)         #集合中添加一个对象，若已存在则不添加

    @pc.on("datachannel")
    def on_datachannel(channel):
        @channel.on("message")
        def on_message(message):
            print(f'Message received: {message}')
            # 处理接收到的数据
            if message == 'pause':
                # 执行暂停操作
                pass
            elif message == 'resume':
                # 执行继续操作
                pass
    #监听RTC连接状态
    @pc.on("iconnectionstatechange")      #当ice连接状态发生变化时
    async def iconnectionstatechange():
        if pc.iceConnectionState == "failed":
            await pc.close()
            pcs.discard(pc)     #移除对象

    # 添加视频轨道
    video_track = VideoTransformTrack(url,url_type)
    pc.addTrack(video_track)

    # @pc.on('track')  --- stream.getTracks().forEach(track => pc.addTrack(track, stream)); 猜测是这里触发的添加了摄像头通道
    # def on_track(track):
    #     if track.kind == 'video':
    #         local_video = VideoTransformTrack(track)
    #         pc.addTrack(local_video)

    # 记录客户端 SDP
    await pc.setRemoteDescription(offer)
    # 生成本地 SDP
    answer = await pc.createAnswer()
    # 记录本地 SDP
    await pc.setLocalDescription(answer)

    print("返回sdp")
    return jsonify({
        "sdp": pc.localDescription.sdp,
        "type": pc.localDescription.type
    })

@api.route('/shutdown', methods=['POST'])
async def shutdown():
    coros = [pc.close() for pc in pcs]
    await asyncio.gather(*coros)
    pcs.clear()
    return 'Server shutting down...'

'''
----------------------实现方式二，网页显示分析画面--基础版
'''

def is_point_in_polygon(point, polygon):
    return polygon.contains(Point(point))

frame_buffer = []
def camera_thread():
    global cap, frame_buffer
    while True:
        try:
            ret, frame = cap.read()
            if not ret:
                logger.warning("Camera disconnected. Reconnecting...")
                cap.release()
                cap = get_camera()
                continue
            frame = cv2.resize(frame, (640, 480))  # 降低视频分辨率
            frame_buffer.append(frame)
            if len(frame_buffer) > 10:  # 保持缓冲区大小不超过10帧
                frame_buffer.pop(0)
                logger.debug("drop one frame!")
        except Exception as e:
            logger.error(f"Error in camera thread: {e}")
            continue

#threading.Thread(target=camera_thread, daemon=True).start()

async def generate_frames():
    global frame_buffer
    while True:
        frame = None
        buffer = None
        try:
            if frame_buffer:
                frame = frame_buffer.pop(0)

                # 进行推理
                results = model(frame)
                detections = results.pandas().xyxy[0]

                # 绘制监控区域
                cv2.polylines(frame, [np.array(region_points, np.int32)], isClosed=True, color=(0, 255, 0), thickness=2)

                for _, row in detections.iterrows():
                    if row['name'] == 'person':
                        # 获取人员检测框的中心点
                        x_center = (row['xmin'] + row['xmax']) / 2
                        y_center = (row['ymin'] + row['ymax']) / 2

                        if is_point_in_polygon((x_center, y_center), polygon):
                            # 触发报警
                            cv2.putText(frame, 'Alert: Intrusion Detected', (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                        (0, 0, 255), 2)
                            #logger.info('Alert: Intrusion Detected')

                        # 绘制检测框
                        cv2.rectangle(frame, (int(row['xmin']), int(row['ymin'])), (int(row['xmax']), int(row['ymax'])),
                                      (255, 0, 0), 2)
                        cv2.circle(frame, (int(x_center), int(y_center)), 5, (0, 0, 255), -1)

                ret, buffer = cv2.imencode('.jpg', frame)
                if not ret:
                    continue
                frame_bytes = buffer.tobytes()
                yield (b'--frame\r\n'
                       b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
            await asyncio.sleep(0.2)
        except Exception as e:
            logger.error(f"Error in generate_frames: {e}")
        finally:
            if frame is not None:
                del frame
            if buffer is not None:
                del buffer

@api.route('/video_feed')
async def video_feed():
    #return jsonify(1)
    return Response(generate_frames(), mimetype='multipart/x-mixed-replace; boundary=frame')

'''
----------------------实现方式三，网页显示分析画面--WEBRTC
'''


class VideoTransformTrack_new(VideoStreamTrack):
    def __init__(self, track):
        super().__init__()
        self.track = track

    async def recv(self):
        frame = await self.track.recv()
        img = frame.to_ndarray(format="bgr24")

        # 使用 Yolov5 进行目标检测
        results = model(img)
        detections = results.pandas().xyxy[0].to_dict(orient="records")

        # 绘制检测结果
        for det in detections:
            if det['name'] == 'person':  # 只检测人员
                x1, y1, x2, y2 = int(det['xmin']), int(det['ymin']), int(det['xmax']), int(det['ymax'])
                cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(img, f"{det['name']} {det['confidence']:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9,
                            (0, 255, 0), 2)

        # 将检测结果图像转换为帧
        new_frame = frame.from_ndarray(img, format="bgr24")
        return new_frame

pcs = set()


@api.websocket('/ws')
async def ws():
    params = await websocket.receive_json()
    offer = RTCSessionDescription(sdp=params['sdp'], type=params['type'])

    pc = RTCPeerConnection()
    pcs.add(pc)

    @pc.on('iceconnectionstatechange')
    async def on_iceconnectionstatechange():
        if pc.iceConnectionState == 'failed':
            await pc.close()
            pcs.discard(pc)

    @pc.on('track')
    def on_track(track):
        if track.kind == 'video':
            local_video = VideoTransformTrack(track)
            pc.addTrack(local_video)

    await pc.setRemoteDescription(offer)
    answer = await pc.createAnswer()
    await pc.setLocalDescription(answer)

    await websocket.send(json.dumps({
        'sdp': pc.localDescription.sdp,
        'type': pc.localDescription.type
    }))

    try:
        while True:
            await asyncio.sleep(1)
    finally:
        await pc.close()
        pcs.discard(pc)