yolov8 tracking编码为web 和 rtsp流输出_yolov8 rtsp

作者：IT小白 | 2024-04-11 05:36:10

踩

yolov8 rtsp

1 基础工作

打开cmd
输入 conda env list
输入 conda activate py38
查看 nvidia-smi
查看 nvcc，如下图所示 cuda为11.7 ，为确认可以查看program files 下面的cuda 安装，看到11.7 就行了，读者可以自行确认自己的版本。
在这里插入图片描述
查看nvidia-smi

打开yolov8 track以后，发现速度超级慢，打开cpu占用率奇高，gpu为零，打开cmd
import torch
torch.cuda.is_available()
为false
确认pytorch安装的为cpu版本，重新安装pytorch cuda 版本，去pytorch网站，点开后复制安装命令，把cuda版本改成自己本机的版本。

在这里插入图片描述
输入命令
python yolo\v8\detect\detect_and_trk.py model=yolov8s.pt source=“d:\test.mp4” show=True

gpu占用率上升到40% 左右。

2 输出到web

接下来我们要做的工作为rtsp server 编码上传，为了让其他机器得到结果，我们必须编程并且得到图像结果，让后使用实时传输协议直接转成rtp 流，让本机成为rtsp server，使用vlc 等工具，可以直接拉取到编码后的流。
修改代码，上传结果，先把命令修改成为
python yolo\v8\detect\detect_and_trk.py model=yolov8s.pt source=“d:\test.mp4”
然后在代码里面自己用opencv 来show，也就是增加两行代码

        cv2.imshow("qbshow",im0)
        cv2.waitKey(3)
1
2

没问题show 出来了，也就是im0 是我们在增加了选框的地方，在show 代码的基础上把图片发送出去就行，我们选择使用websocket 发送，或者socket 直接发送都行，这样我们需要一个客户端还是服务端的选择，我们选择让python成为服务端，目的是为了不让python需要断线重连，python 成为服务端有很多好处，
1 是可以直接给web传送图片
2 是可以直接给

    def write_results(self, idx, preds, batch):

        p, im, im0 = batch
        log_string = ""
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim
        self.seen += 1
        im0 = im0.copy()
        if self.webcam:  # batch_size >= 1
            log_string += f'{idx}: '
            frame = self.dataset.count
        else:
            frame = getattr(self.dataset, 'frame', 0)
        # tracker
        self.data_path = p
    
        save_path = str(self.save_dir / p.name)  # im.jpg
        self.txt_path = str(self.save_dir / 'labels' / p.stem) + ('' if self.dataset.mode == 'image' else f'_{frame}')
        log_string += '%gx%g ' % im.shape[2:]  # print string
        self.annotator = self.get_annotator(im0)
        
        det = preds[idx]
        self.all_outputs.append(det)
        if len(det) == 0:
            return log_string
        for c in det[:, 5].unique():
            n = (det[:, 5] == c).sum()  # detections per class
            log_string += f"{n} {self.model.names[int(c)]}{'s' * (n > 1)}, "
    
    
        # #..................USE TRACK FUNCTION....................
        dets_to_sort = np.empty((0,6))
        
        for x1,y1,x2,y2,conf,detclass in det.cpu().detach().numpy():
            dets_to_sort = np.vstack((dets_to_sort, 
                        np.array([x1, y1, x2, y2, conf, detclass])))
        
        tracked_dets = tracker.update(dets_to_sort)
        tracks =tracker.getTrackers()
        
        for track in tracks:
            [cv2.line(im0, (int(track.centroidarr[i][0]),
                        int(track.centroidarr[i][1])), 
                        (int(track.centroidarr[i+1][0]),
                        int(track.centroidarr[i+1][1])),
                        rand_color_list[track.id], thickness=3) 
                        for i,_ in  enumerate(track.centroidarr) 
                            if i < len(track.centroidarr)-1 ] 
        

        if len(tracked_dets)>0:
            bbox_xyxy = tracked_dets[:,:4]
            identities = tracked_dets[:, 8]
            categories = tracked_dets[:, 4]
            draw_boxes(im0, bbox_xyxy, identities, categories, self.model.names)
           
        gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
        cv2.imshow("qbshow",im0)
        cv2.waitKey(3)
        return log_string
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

2.1 web中显示

在python中加上websocket server的代码，以下是一个示例代码，读者可以自己测试，但是要应用到python 的图像处理server中，我们要做一些修改

import asyncio
import threading
import websockets
import time
CONNECTIONS = set()


async def server_recv(websocket):
    while True:
        try:
            recv_text = await websocket.recv()
            print("recv:", recv_text)
        except websockets.ConnectionClosed as e:
            # 客户端关闭连接，跳出对客户端的读取，结束函数
            print(e.code)
            await asyncio.sleep(0.01)
            break

async def server_hands(websocket):
  
    CONNECTIONS.add(websocket)
    print(CONNECTIONS)
    try:
        await websocket.wait_closed()
    finally:
        CONNECTIONS.remove(websocket) 



def message_all(message):
    websockets.broadcast(CONNECTIONS, message)   

async def handler(websocket, path):
    # 处理新的 WebSocket 连接
    print("New WebSocket route is ",path)

    try:
        await server_hands(websocket) # 握手加入队列
        await server_recv(websocket)  # 接收客户端消息并处理
     
    except websockets.exceptions.ConnectionClosedError as e:
        print(f"Connection closed unexpectedly: {e}")
    finally:
        pass
        # 处理完毕，关闭 WebSocket 连接
    print("WebSocket connection closed")


async def sockrun():
    async with websockets.serve(handler, "", 9090):
        await asyncio.Future()  # run forever

def main_thread(): 
    print("main")
    asyncio.run(sockrun())
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

在这里插入图片描述

2.2 修改

修改的地方为我们不能让websocket 阻塞主线程，虽然websocket为协程处理，但是依然要放到线程里面

def predict(cfg):
    init_tracker()
    random_color_list()
        
    cfg.model = cfg.model or "yolov8n.pt"
    cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2)  # check image size
    cfg.source = cfg.source if cfg.source is not None else ROOT / "assets"
    predictor = DetectionPredictor(cfg)
    predictor()


if __name__ == "__main__":
    thread = threading.Thread(target=main_thread)
    thread.start()
    predict()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

值得注意的地方是：
很多人喜欢把编码做成base64 ，这样没有必要，服务程序里面把二进制编码成base64，消耗了cpu，然后发送还大了很多，实际上我们直接发送二进制就行了，发送时

   cv2.waitKey(3)
   _, encimg = cv2.imencode('.jpg', im0)
   bys = np.array(encimg).tobytes()
   message_all(bys)
1
2
3
4

web端测试代码

<!DOCTYPE html>
<html lang="en">

<head>
    <meta charset="UTF-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>websocket</title>
    
</head>
<body>
    <div>
        <button onclick="connecteClient()">开始接收</button>
    </div>

    <br>
    <div>
        <img src="" id="python_img" alt="websock" style="text-align:left; width: 640px; height: 360px;"> 
    </div>
    <script>
        function connecteClient() {
            var ws = new WebSocket("ws://127.0.0.1:9090");

            ws.onopen = function () {
                console.log("WebSocket 连接成功");
            };

            ws.onmessage = function (evt) {
                var received_msg = evt.data;
                blobToDataURI(received_msg, function (result) {
                    document.getElementById("python_img").src = result;
                })
            };

            ws.onclose = function () {
                console.log("连接关闭...");
            };
        }

       
        function blobToDataURI(blob, callback) {
            var reader = new FileReader();
            reader.readAsDataURL(blob);
            reader.onload = function (e) {
                callback(e.target.result);
            }
        }
    </script>
</body>
</html>
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

其中接收到二进制完成以后直接显示，不用服务端编码成base64，下图表明web端可以正常拿到jpg图片显示

在这里插入图片描述

其中message_all 是广播，只要是websocket client 我们就给他发一份，这样做，不但可以发送给web，也可以发送给rtsp server，我们让rtsp server 链接上来，取到jpg后，解码成为rgb24，然后编码成为h264，让vlc 可以链接rtsp server取到rtsp 流

3 rtsp server 编码h264 输出

我们的rtsp server 端使用c++ 编写,为了加快速度，使用了内存共享，而没有使用websocket client，当然是可以的，考虑到这里要解码，又要编码输出，还有可能需要使用硬件编码，为了提高效率，就这么做了，而且还简单，缺点就是这样多了一个问题，rtspserve只能启动在本机上。先编写一个h264Frame， h265 是一样的道理，后面再更改，暂时运行在windows上，等完成了再修改。注意以下代码只能运行在windows上

class H264Frame :public TThreadRunable
{
	const wchar_t*  sharedMemoryMutex = L"shared_memory_mutex";
	const wchar_t*  sharedMemoryName = L"shared_memory";
	LPVOID v_lpBase = NULL;
	HANDLE v_hMapFile = NULL;
	HANDLE 	m_shareMutex = NULL;
	std::shared_ptr<xop::RtspServer> v_server;
	std::string v_livename;
	std::wstring v_pathname;
	std::wstring v_pathmutex_name;

	int v_fps = 10;

	Encoder v_encoder;
	xop::MediaSessionId v_id;
	int v_init = 0;
	int v_quality = 28;
public:
	H264Frame(std::shared_ptr<xop::RtspServer> s, const char* livename, const wchar_t* pathname, xop::MediaSessionId session_id)
	{
		v_server = s;
		v_livename = livename;
		v_pathname = pathname;
		if (pathname != NULL)
		{
			v_pathmutex_name = pathname;
			v_pathmutex_name += L"_mutex";
		}
		v_id = session_id;
	}
	~H264Frame() {}

	bool Open()
	{
		if (v_pathname.empty())
			v_hMapFile = OpenFileMappingW(FILE_MAP_ALL_ACCESS, NULL, sharedMemoryName);
		else
			v_hMapFile = OpenFileMappingW(FILE_MAP_ALL_ACCESS, NULL, v_pathname.c_str());

		if (v_hMapFile == NULL)
		{
			//printf(" Waiting shared memory creation......\n");
			return false;
		}
		return true;
	}
	void Close()
	{
		if (m_shareMutex != NULL)
			ReleaseMutex(m_shareMutex);
		if (v_hMapFile != NULL)
			CloseHandle(v_hMapFile);
	}

	bool IsOpened() const
	{
		return (v_hMapFile != NULL);
	}

	void sendFrame(uint8_t* data ,int size, int sessionid)
	{
		xop::AVFrame videoFrame = { 0 };
		videoFrame.type = 0;
		videoFrame.size = size;// sizeofdata;// frame_size;
		videoFrame.timestamp = xop::H264Source::GetTimestamp();
		//videoFrame.notcopy = data;
		videoFrame.buffer.reset(new uint8_t[videoFrame.size]);
		std::memcpy(videoFrame.buffer.get(), data, videoFrame.size);
		v_server->PushFrame(sessionid, xop::channel_0, videoFrame);
	}
	int ReadFrame(int sessionid)
	{
		if (v_hMapFile)
		{

			v_lpBase = MapViewOfFile(v_hMapFile, FILE_MAP_ALL_ACCESS, 0, 0, 0);
			unsigned char* pMemory = (unsigned char*)v_lpBase;

			int w = 0, h = 0;
			int frame_size = 0;
			std::memcpy(&w, pMemory, 4);
			std::memcpy(&h, pMemory + 4, 4);
			std::memcpy(&frame_size, pMemory + 8, 4);
			//printf("read the width %d height %d framesize %d\n", imageWidth,imageHeight, frame_size);
			uint8_t* rgb = pMemory + 12;
			if (v_init == 0)
			{
				v_encoder.Encoder_Open(v_fps, w, h, w, h, v_quality);
				v_init = 1;
			}
			
			v_encoder.RGB2YUV(rgb, w, h);
			AVPacket * pkt = v_encoder.EncodeVideo();
			//while (!*framebuf++);
			if (pkt!=NULL) {
				uint8_t* framebuf = pkt->data;
				frame_size = pkt->size;
				uint8_t* spsStart = NULL;
				int spsLen = 0;
				uint8_t* ppsStart = NULL;
				int ppsLen = 0;
				uint8_t* seStart = NULL;
				int seLen = 0;
				uint8_t* frameStart = 0;//非关键帧地址或者关键帧地址
				int frameLen = 0;
				AnalyseNalu_no0x((const uint8_t*)framebuf, frame_size, &spsStart, spsLen, &ppsStart, ppsLen, &seStart, seLen, &frameStart, frameLen);
				if (spsStart != NULL && ppsStart!=NULL)
				{
					sendFrame(spsStart, spsLen, sessionid);
					sendFrame(ppsStart, ppsLen, sessionid);
				}
				if (frameStart != NULL)
				{
					sendFrame(frameStart, frameLen, sessionid);
				}


				av_packet_free(&pkt);
			}
			UnmapViewOfFile(pMemory);
		}
		return -1;
	}

	void Run()
	{
		m_shareMutex = CreateMutexW(NULL, false, sharedMemoryMutex/*v_pathmutex_name.c_str()*/);		v_encoder.Encoder_Open(10, 4096, 1080, 4096, 1080, 27);
		while (1)
		{
			if (Open() == true)
			{
				break;
			}
			else
			{
				printf("can not open share mem error\n");
				Sleep(3000);
			}
		}
		printf("wait ok\n");
		//检查是否改变参数
		//重新开始
		int64_t m_start_clock = 0;
		float delay = 1000000.0f / (float)(v_fps);//微妙

		//合并的rgb
		uint8_t * rgb = NULL;
		uint8_t * cam = NULL;

		int64_t start_timestamp = av_gettime_relative();//  GetTimestamp32();
		while (1)
		{
			if (IsStop())
				break;
			//获取一帧
			//rgb = 
			if (v_hMapFile)
			{
				//printf("start to lock sharemutex\n");
				WaitForSingleObject(m_shareMutex, INFINITE);

				//printf("read frame\n");
				ReadFrame(v_id);

				ReleaseMutex(m_shareMutex);
			}
			else
			{
				printf("Shared memory handle error\n");
				break;
			}
			
			
			int64_t tnow = av_gettime_relative();// GetTimestamp32();
			int64_t total = tnow - start_timestamp; //總共花費的時間
			int64_t fix_consume = ++m_start_clock * delay;
			if (total < fix_consume) {
				int64_t diff = fix_consume - total;
				if (diff > 0) { //相差5000微妙以上
					std::this_thread::sleep_for(std::chrono::microseconds(diff));
				}
			}
		}
		if (v_init == 0)
		{
			v_encoder.Encoder_Close();
			v_init = 1;
		}
	}

};


int main(int argc, char **argv)
{	
	
	std::string suffix = "live";
	std::string ip = "127.0.0.1";
	std::string port = "8554";
	std::string rtsp_url = "rtsp://" + ip + ":" + port + "/" + suffix;
	
	std::shared_ptr<xop::EventLoop> event_loop(new xop::EventLoop());
	std::shared_ptr<xop::RtspServer> server = xop::RtspServer::Create(event_loop.get());

	if (!server->Start("0.0.0.0", atoi(port.c_str()))) {
		printf("RTSP Server listen on %s failed.\n", port.c_str());
		return 0;
	}

#ifdef AUTH_CONFIG
	server->SetAuthConfig("-_-", "admin", "12345");
#endif

	xop::MediaSession *session = xop::MediaSession::CreateNew("live"); 
	session->AddSource(xop::channel_0, xop::H264Source::CreateNew()); 
	//session->StartMulticast(); 
	session->AddNotifyConnectedCallback([] (xop::MediaSessionId sessionId, std::string peer_ip, uint16_t peer_port){
		printf("RTSP client connect, ip=%s, port=%hu \n", peer_ip.c_str(), peer_port);
	});
   
	session->AddNotifyDisconnectedCallback([](xop::MediaSessionId sessionId, std::string peer_ip, uint16_t peer_port) {
		printf("RTSP client disconnect, ip=%s, port=%hu \n", peer_ip.c_str(), peer_port);
	});

	xop::MediaSessionId session_id = server->AddSession(session);
      
	H264Frame h264frame(server, suffix.c_str(), L"shared_memory",session_id);

	//std::thread t1(SendFrameThread, server.get(), session_id, &h264_file);
	//t1.detach(); 

	std::cout << "Play URL: " << rtsp_url << std::endl;
	h264frame.Start();
	while (1) {
		xop::Timer::Sleep(100);
	}
	h264frame.Join();
	getchar();
	return 0;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241

完成以后，打开python 服务端，打开rtsp服务端，最后打开vlc查看
在这里插入图片描述
这样就完成了python端使用pytorch ，yolo 等等共享给c++ 的内存

其中python端的内存共享示例使用如下

import mmap
import contextlib
import time

while True:
    with contextlib.closing(mmap.mmap(-1, 25, tagname='test', access=mmap.ACCESS_READ)) as m:
        s = m.read(1024)#.replace('\x00', '')
        print(s)
    time.sleep(1)
1
2
3
4
5
6
7
8
9

读者可自行完成剩余的代码

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/IT小白/article/detail/403155