保姆级教程:车流量统计AI系统,从0到1落地实战(上)

defpreprocess_warpAffine(image, dst_width=640, dst_height=640):"""图像预处理函数，用于将输入图像变形并缩放到指定大小。该函数首先计算缩放因子，以确保图像按比例缩放到目标尺寸，然后计算偏移量，以确保缩放后的图像居中。使用计算得到的缩放和偏移参数对图像进行仿射变换，并对变换后的图像进行颜色反转、归一化和维度转换，以适应深度学习模型的输入要求。参数:image: 输入的图像，通常是一个二维或三维的numpy数组。dst_width: 目标图像的宽度，默认值为640像素。dst_height: 目标图像的高度，默认值为640像素。返回:img_pre: 预处理后的图像，是一个四维的PyTorch张量。IM: 逆仿射变换矩阵，用于将变形后的图像恢复到原始状态。"""    # 计算缩放因子，确保图像按比例缩放到目标尺寸scale = min((dst_width / image.shape[1], dst_height / image.shape[0]))    # 计算偏移量，确保缩放后的图像居中ox = int((dst_width  - scale * image.shape[1]) / 2)oy = int((dst_height - scale * image.shape[0]) / 2)    # 构建仿射变换矩阵M = np.array([[scale,0, ox],[0,scale, oy]],dtype=np.float32)    # 对图像进行仿射变换，设置目标尺寸、插值方法和边界填充方式img_pre = cv2.warpAffine(image, M, (dst_width, dst_height), flags=cv2.INTER_LINEAR,borderMode=cv2.BORDER_CONSTANT, borderValue=(114, 114, 114))    # 计算逆仿射变换矩阵，用于将变形后的图像恢复到原始状态IM = cv2.invertAffineTransform(M)    # 对变换后的图像进行颜色反转和归一化处理img_pre = (img_pre[...,::-1] / 255.0).astype(np.float32)    # 调整图像数据的维度顺序，适应深度学习模型的输入要求img_pre = img_pre.transpose(2, 0, 1)[None]    # 将处理后的图像数据转换为PyTorch张量img_pre = torch.from_numpy(img_pre)    # 设置PyTorch打印选项，提高输出精度torch.set_printoptions(precision=6)    # 返回预处理后的图像和逆仿射变换矩阵returnimg_pre, IM

# 计算两个边界框的交并比（Intersection over Union，IoU）defiou(box1, box2):# 内部函数，用于计算边界框的面积defarea_box(box):return (box[2] - box[0]) * (box[3] - box[1])# 计算交集的左上角和右下角坐标    left   = max(box1[0], box2[0])    top    = max(box1[1], box2[1])    right  = min(box1[2], box2[2])    bottom = min(box1[3], box2[3])# 计算交集的面积    cross  = max((right-left), 0) * max((bottom-top), 0)# 计算并集的面积    union  = area_box(box1) + area_box(box2) - cross# 如果交集或并集的面积为0，则返回0if cross == 0or union == 0:return0# 返回交并比return cross / union# 非极大值抑制（Non-Maximum Suppression，NMS）算法，用于去除多余的边界框defNMS(boxes, iou_thres):# 初始化一个标记列表，用于标记是否移除某个边界框    remove_flags = [False] * len(boxes)# 初始化一个列表，用于保存保留下来的边界框    keep_boxes = []# 遍历所有边界框for i, ibox in enumerate(boxes):# 如果当前边界框已被标记为移除，则跳过if remove_flags[i]:continue# 将当前边界框添加到保留列表中        keep_boxes.append(ibox)# 继续遍历其余的边界框for j in range(i + 1, len(boxes)):# 如果当前边界框已被标记为移除，则跳过if remove_flags[j]:continue# 获取下一个边界框            jbox = boxes[j]# 如果两个边界框的类别标签不同，则跳过if(ibox[5] != jbox[5]):continue# 如果两个边界框的交并比大于设定的阈值，则标记第二个边界框为移除if iou(ibox, jbox) > iou_thres:                remove_flags[j] = True# 返回保留下来的边界框列表return keep_boxes# 后处理函数，用于处理模型的预测结果defpostprocess(pred, IM=[], conf_thres=0.25, iou_thres=0.45):# 初始化一个列表，用于保存处理后的边界框    boxes = []# 遍历模型的预测结果for item in pred[0]:# 提取边界框的中心坐标、宽度、高度        cx, cy, w, h = item[:4]# 提取类别标签和置信度        label = item[4:].argmax()        confidence = item[4 + label]# 如果置信度低于设定的阈值，则跳过if confidence < conf_thres:continue# 计算边界框的左上角和右下角坐标        left    = cx - w * 0.5        top     = cy - h * 0.5        right   = cx + w * 0.5        bottom  = cy + h * 0.5# 将边界框的坐标、置信度和类别标签添加到列表中        boxes.append([left, top, right, bottom, confidence, label])# 将边界框列表转换为NumPy数组    boxes = np.array(boxes)# 如果没有边界框，则返回空列表if0 == len(boxes):return []# 应用图像的缩放比例和偏移量，将边界框坐标转换回原始图像的坐标系    lr = boxes[:,[0, 2]]    tb = boxes[:,[1, 3]]    boxes[:,[0,2]] = IM[0][0] * lr + IM[0][2]    boxes[:,[1,3]] = IM[1][1] * tb + IM[1][2]# 根据置信度降序排序边界框    boxes = sorted(boxes.tolist(), key=lambda x:x[4], reverse=True)# 应用非极大值抑制，去除多余的边界框return NMS(boxes, iou_thres)

model = YOLO(model='yolo11x.pt')  # load a pretrained model (recommended for training)model.export(format="onnx", opset=16, imgsz=640, simplify=True, save_dir='.')

defbuild_engine(onnx_model_path, engine_path, precision_set="fp16"):"""    从ONNX模型生成TensorRT引擎。    参数:    onnx_model_path: str，ONNX模型的路径。    engine_path: str，生成的TensorRT引擎的保存路径。    precision_set: str，精度设置，默认为"fp16"。    返回:    无    """# 创建TensorRT日志记录器    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)# 创建TensorRT构建器    builder = trt.Builder(TRT_LOGGER)# 创建网络定义，并设置显式批处理标志    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))# 创建ONNX解析器    parser = trt.OnnxParser(network, TRT_LOGGER)# 打开ONNX模型文件并解析with open(onnx_model_path, "rb") as model:ifnot parser.parse(model.read()):# 如果解析失败，打印错误信息for error in range(parser.num_errors):                print(parser.get_error(error))# 创建构建配置    config = builder.create_builder_config()# 设置最大工作空间大小为1GB    config.max_workspace_size = 1 << 30# 1GB# 如果精度设置为fp16，启用FP16模式if precision_set == "fp16":        config.set_flag(trt.BuilderFlag.FP16)# 如果引擎文件不存在，构建并保存引擎ifnot os.path.exists(engine_path):        engine = builder.build_engine(network, config)with open(engine_path, "wb") as f:# 序列化引擎并写入文件            f.write(engine.serialize())

__global__ voidletterbox(const uchar* srcData, constint srcH, constint srcW, uchar* tgtData, constint tgtH, constint tgtW, constint rszH, constint rszW, constint startY, constint startX){int ix = threadIdx.x + blockDim.x * blockIdx.x;int iy = threadIdx.y + blockDim.y * blockIdx.y;int idx = ix + iy * tgtW;int idx3 = idx * 3;if ( ix > tgtW || iy > tgtH ) return;  // thread out of target range// gray region on target imageif ( iy < startY || iy > (startY + rszH - 1) ) {        tgtData[idx3] = 114;        tgtData[idx3 + 1] = 114;        tgtData[idx3 + 2] = 114;return;    }if ( ix < startX || ix > (startX + rszW - 1) ){        tgtData[idx3] = 114;        tgtData[idx3 + 1] = 114;        tgtData[idx3 + 2] = 114;return;    }float scaleY = (float)rszH / (float)srcH;float scaleX = (float)rszW / (float)srcW;// (ix,iy)为目标图像坐标// (before_x,before_y)原图坐标float beforeX = float(ix - startX + 0.5) / scaleX - 0.5;float beforeY = float(iy - startY + 0.5) / scaleY - 0.5;// 原图像坐标四个相邻点// 获得变换前最近的四个顶点,取整int topY = static_cast<int>(beforeY);int bottomY = topY + 1;int leftX = static_cast<int>(beforeX);int rightX = leftX + 1;//计算变换前坐标的小数部分float u = beforeX - leftX;float v = beforeY - topY;if (topY >= srcH - 1 && leftX >= srcW - 1)  //右下角    {for (int k = 0; k < 3; k++)        {            tgtData[idx3 + k] = (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k];        }    }elseif (topY >= srcH - 1)  // 最后一行    {for (int k = 0; k < 3; k++)        {            tgtData[idx3 + k]            = (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k]            + (u) * (1. - v) * srcData[(rightX + topY * srcW) * 3 + k];        }    }elseif (leftX >= srcW - 1)  // 最后一列    {for (int k = 0; k < 3; k++)        {            tgtData[idx3 + k]            = (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k]            + (1. - u) * (v) * srcData[(leftX + bottomY * srcW) * 3 + k];        }    }else// 非最后一行或最后一列情况    {for (int k = 0; k < 3; k++)        {            tgtData[idx3 + k]            = (1. - u) * (1. - v) * srcData[(leftX + topY * srcW) * 3 + k]            + (u) * (1. - v) * srcData[(rightX + topY * srcW) * 3 + k]            + (1. - u) * (v) * srcData[(leftX + bottomY * srcW) * 3 + k]            + u * v * srcData[(rightX + bottomY * srcW) * 3 + k];        }    }}

__global__ voidprocess(const uchar* srcData, float* tgtData, constint h, constint w){int ix = threadIdx.x + blockIdx.x * blockDim.x;int iy = threadIdx.y + blockIdx.y * blockDim.y;int idx = ix + iy * w;int idx3 = idx * 3;if (ix < w && iy < h)    {        tgtData[idx] = (float)srcData[idx3 + 2] / 255.0;  // R pixel        tgtData[idx + h * w] = (float)srcData[idx3 + 1] / 255.0;  // G pixel        tgtData[idx + h * w * 2] = (float)srcData[idx3] / 255.0;  // B pixel    }}

// ------------------ decode ( get class and conf ) --------------------__global__ voiddecode_kernel(float* src, float* dst, int numBboxes, int numClasses, float confThresh, int maxObjects, int numBoxElement){int position = blockDim.x * blockIdx.x + threadIdx.x;if (position >= numBboxes) return;float* pitem = src + (4 + numClasses) * position;float* classConf = pitem + 4;float confidence = 0;int label = 0;for (int i = 0; i < numClasses; i++){if (classConf[i] > confidence){            confidence = classConf[i];            label = i;        }    }if (confidence < confThresh) return;int index = (int)atomicAdd(dst, 1);if (index >= maxObjects) return;float cx     = pitem[0];float cy     = pitem[1];float width  = pitem[2];float height = pitem[3];float left   = cx - width * 0.5f;float top    = cy - height * 0.5f;float right  = cx + width * 0.5f;float bottom = cy + height * 0.5f;float* pout_item = dst + 1 + index * numBoxElement;    pout_item[0] = left;    pout_item[1] = top;    pout_item[2] = right;    pout_item[3] = bottom;    pout_item[4] = confidence;    pout_item[5] = label;    pout_item[6] = 1;  // 1 = keep, 0 = ignore}voiddecode(float* src, float* dst, int numBboxes, int numClasses, float confThresh, int maxObjects, int numBoxElement, cudaStream_t stream){    cudaMemsetAsync(dst, 0, sizeof(int), stream);int blockSize = 256;int gridSize = (numBboxes + blockSize - 1) / blockSize;    decode_kernel<<<gridSize, blockSize, 0, stream>>>(src, dst, numBboxes, numClasses, confThresh, maxObjects, numBoxElement);}

__device__ floatbox_iou(float aleft, float atop, float aright, float abottom, float bleft, float btop, float bright, float bbottom){float cleft = max(aleft, bleft);float ctop = max(atop, btop);float cright = min(aright, bright);float cbottom = min(abottom, bbottom);float c_area = max(cright - cleft, 0.0f) * max(cbottom - ctop, 0.0f);if (c_area == 0.0f) return0.0f;float a_area = max(0.0f, aright - aleft) * max(0.0f, abottom - atop);float b_area = max(0.0f, bright - bleft) * max(0.0f, bbottom - btop);return c_area / (a_area + b_area - c_area);}__global__ voidnms_kernel(float* data, float kNmsThresh, int maxObjects, int numBoxElement){int position = blockDim.x * blockIdx.x + threadIdx.x;int count = min((int)data[0], maxObjects);if (position >= count) return;// left, top, right, bottom, confidence, class, keepflagfloat* pcurrent = data + 1 + position * numBoxElement;float* pitem;for (int i = 0; i < count; i++){        pitem = data + 1 + i * numBoxElement;if (i == position || pcurrent[5] != pitem[5]) continue;if (pitem[4] >= pcurrent[4]){if (pitem[4] == pcurrent[4] && i < position) continue;float iou = box_iou(                pcurrent[0], pcurrent[1], pcurrent[2], pcurrent[3],                pitem[0], pitem[1], pitem[2], pitem[3]            );if (iou > kNmsThresh){                pcurrent[6] = 0;  // 1 = keep, 0 = ignorereturn;            }        }    }}voidnms(float* data, float kNmsThresh, int maxObjects, int numBoxElement, cudaStream_t stream){int blockSize = maxObjects < 256?maxObjects:256;int gridSize = (maxObjects + blockSize - 1) / blockSize;    nms_kernel<<<gridSize, blockSize, 0, stream>>>(data, kNmsThresh, maxObjects, numBoxElement);}

ffmpeg -re -i "${INPUT_MP4}" -c:v libx264 -preset veryfast -maxrate 3000k -bufsize 6000k -pix_fmt yuv420p -g 50 -c:a aac -b:a 160k -ac 2 -ar 44100 -f rtsp "${RTSP_URL}"

#!/bin/bash# 检查参数数量if [ "$#" -ne 2 ]; then    echo "Usage: $0 <input_folder> <rtsp_base_url>"    exit 1fi# ffmpeg路径export PATH=../../3rd/ffmpeg/bin:$PATHexport LD_LIBRARY_PATH=../../3rd/ffmpeg/lib# 输入的文件夹路径INPUT_FOLDER="$1"# 输出的RTSP URL基础RTSP_BASE_URL="$2"# 检查输入文件夹是否存在if [ ! -d "$INPUT_FOLDER" ]; then    echo "Error: Input folder $INPUT_FOLDER not found!"    exit 1fi# 遍历输入文件夹中的所有MP4文件for INPUT_MP4 in "$INPUT_FOLDER"/*.mp4; do    # 检查文件是否存在    if [ ! -f "$INPUT_MP4" ]; then        echo "No MP4 files found in $INPUT_FOLDER!"        continue    fi    # 获取文件名（不带路径）    FILENAME=$(basename "$INPUT_MP4")    # 去掉文件扩展名    BASENAME="${FILENAME%.*}"    # 构建RTSP URL    RTSP_URL="${RTSP_BASE_URL}/${BASENAME}"    # 使用 FFmpeg 将 mp4 流转换为 RTSP 流    function mp4_to_rtsp {        ffmpeg -re -i "${INPUT_MP4}" -c:v libx264 -preset veryfast -maxrate 3000k -bufsize 6000k -pix_fmt yuv420p -g 50 -c:a aac -b:a 160k -ac 2 -ar 44100 -f rtsp "${RTSP_URL}"    }    # 启动推送MP4文件的FFmpeg进程    echo "Starting to push MP4 file $INPUT_MP4 to RTSP server at $RTSP_URL..."    mp4_to_rtsp &done# 等待所有后台任务完成waitecho "All processes have finished."

#ifndef RTSPSTREAMPROCESSOR_H#define RTSPSTREAMPROCESSOR_H#include"common.h"class RTSPStreamProcessor {public:    // 构造函数，初始化线程池    RTSPStreamProcessor(int numThreads, BlockingQueue<sSingleFrameInfo>& queue);    // 析构函数，停止线程池并等待所有线程结束    ~RTSPStreamProcessor();    // 将RTSP URL添加到任务队列中    voidenqueue(const std::string& rtspUrl);private:    // 处理单个RTSP流    voidprocessRTSPStream(const std::string& rtspUrl);    // 工作线程池    std::vector<std::thread> workers;    // 任务队列    std::queue<std::function<void()>> tasks;    // 保护任务队列的互斥锁    std::mutex queue_mutex;    // 条件变量，用于线程间的通信    std::condition_variable condition;    // 标志位，指示是否停止处理任务    bool stop;    uint32_t mGlobalFrameNum = 0;    BlockingQueue<sSingleFrameInfo> &frameQueue;};#endif// RTSPSTREAMPROCESSOR_H

#include"RTSPStreamProcessor.h"// 构造函数，初始化线程池RTSPStreamProcessor::RTSPStreamProcessor(int numThreads, BlockingQueue<sSingleFrameInfo>& queue) :    stop(false), frameQueue(queue) {    for (int i = 0; i < numThreads; ++i) {        workers.emplace_back([this] {            for (;;) {                std::function<void()> task;                {                    std::unique_lock<std::mutex> lock(this->queue_mutex);                    // 等待任务或停止信号                    this->condition.wait(lock, [this] { return this->stop || !this->tasks.empty(); });                    if (this->stop && this->tasks.empty()) {                        return;                    }                    // 获取任务并移除                    task = std::move(this->tasks.front());                    this->tasks.pop();                }                // 执行任务                task();            }        });    }}// 析构函数，停止线程池并等待所有线程结束RTSPStreamProcessor::~RTSPStreamProcessor() {    {        std::unique_lock<std::mutex> lock(queue_mutex);        stop = true;    }    // 通知所有线程停止    condition.notify_all();    for (std::thread& worker : workers) {        worker.join();    }}// 将RTSP URL添加到任务队列中voidRTSPStreamProcessor::enqueue(const std::string& rtspUrl){    {        std::unique_lock<std::mutex> lock(queue_mutex);        // 将处理RTSP流的任务添加到队列        tasks.emplace([this, rtspUrl] { processRTSPStream(rtspUrl); });    }    // 通知一个线程有新任务    condition.notify_one();}// 处理单个RTSP流voidRTSPStreamProcessor::processRTSPStream(const std::string& rtspUrl){    AVFormatContext *pFormatCtx = nullptr;    AVCodecContext *pCodecCtx = nullptr;    AVFrame *pFrame = nullptr, *pFrameBGR = nullptr;    AVPacket *packet = nullptr;    uint8_t *outBuffer  = nullptr;    SwsContext *imgConvertCtx = nullptr;    int videoStream;    bool bConnected = false;    while (true){        if (!bConnected) {            avformat_network_init();            AVDictionary *avdic=NULL;            av_dict_set(&avdic, "buffer_size", "655360", 0);            av_dict_set(&avdic,"rtsp_transport","tcp",0);            av_dict_set(&avdic, "max_delay", "500000", 0);            av_dict_set(&avdic, "stimeout", "5000000", 0);            if (avformat_open_input(&pFormatCtx, rtspUrl.c_str(), nullptr, &avdic) != 0)            {                std::cout << "can't open the file." << std::endl;                bConnected = false;                continue;            }            if (avformat_find_stream_info(pFormatCtx, nullptr) < 0)            {                std::cout << "can't find stream infomation" << std::endl;                bConnected = false;                continue ;            }            videoStream = -1;            for (unsigned int i = 0; i < pFormatCtx->nb_streams; i++)            {                if (pFormatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)                {                    videoStream = i;                }            }            if (videoStream == -1)            {                std::cout << "can't find a video stream" << std::endl;                bConnected = false;                continue ;            }            const AVCodec *pCodec = avcodec_find_decoder(pFormatCtx->streams[videoStream]->codecpar->codec_id);            if (pCodec == NULL)            {                std::cout << "can't find a codec" << std::endl;                bConnected = false;                continue ;            }            pCodecCtx = avcodec_alloc_context3(pCodec);            if (avcodec_parameters_to_context(pCodecCtx, pFormatCtx->streams[videoStream]->codecpar) < 0) {                std::cerr << "Failed to allocate codec context." << std::endl;            }            pCodecCtx->thread_count = 4;            pCodecCtx->thread_type = FF_THREAD_SLICE;            if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0)            {                std::cout << "can't open a codec" << std::endl;                bConnected = false;                continue ;            }            pFrame = av_frame_alloc();            pFrameBGR = av_frame_alloc();            imgConvertCtx = sws_getContext(pCodecCtx->width, pCodecCtx->height,                                           pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height,                                           AV_PIX_FMT_BGR24, SWS_BICUBIC, NULL, NULL, NULL);            int numBytes = av_image_get_buffer_size(AV_PIX_FMT_BGR24, pCodecCtx->width, pCodecCtx->height, 1);            outBuffer = (uint8_t *) av_malloc(numBytes * sizeof(uint8_t));            av_image_fill_arrays(pFrameBGR->data, pFrameBGR->linesize, outBuffer, AV_PIX_FMT_BGR24,                                 pCodecCtx->width, pCodecCtx->height, 1);            packet = (AVPacket *) malloc(sizeof(AVPacket));            av_new_packet(packet, pCodecCtx->width * pCodecCtx->height);            bConnected = true;        }        uint32_t frameCount = 0;        while (av_read_frame(pFormatCtx, packet) >= 0)        {            if (packet->stream_index == videoStream)            {                int ret = avcodec_send_packet(pCodecCtx, packet);                if (ret < 0) {//                    fprintf(stderr, "Error sending a packet for decoding: %s\n", av_err2str(ret));                    continue;                }else{                    ret = avcodec_receive_frame(pCodecCtx, pFrame);                    if (ret < 0) {//                        fprintf(stderr, "Error during decoding: %s\n", av_err2str(ret));                    }else{                        sws_scale(imgConvertCtx,                                  (uint8_t const * const *) pFrame->data,                                  pFrame->linesize, 0, pCodecCtx->height, pFrameBGR->data,                                  pFrameBGR->linesize);                        cv::Mat img(pFrame->height,                                    pFrame->width,                                    CV_8UC3,                                    pFrameBGR->data[0]);//                        sSingleFrameInfo frameInfo;//                        frameInfo.frame = img.clone();//                        frameInfo.rtspUrl = rtspUrl;//                        frameInfo.localFrameNum = frameCount;                        queue_mutex.lock();                        cv::imshow("rtsp:" + rtspUrl, img);                        if (cv::waitKey(1) == 27) break;                        queue_mutex.unlock();                        std::cout << "rtsp:" << rtspUrl << "frameCount: " << frameCount << std::endl;//                        std::vector<std::string> rtspUrls = splitString(rtspUrl, '/');//                        const std::string savePath = "/media/dling/data_1/ubuntu/everyDay/20241230/" + rtspUrls.back() + "/";//                        createDirectories(savePath);//                        std::string saveName = savePath + std::to_string(frameCount) + ".jpg";//                        cv::imwrite(saveName, img);                        frameCount++;                    }                }            }            av_packet_unref(packet);        }        bConnected = false;        av_free(outBuffer);        outBuffer = nullptr;        av_freep(&pFrameBGR);        av_freep(&pFrame);        sws_freeContext(imgConvertCtx);        imgConvertCtx = nullptr;        avcodec_free_context(&pCodecCtx);        pCodecCtx = nullptr;        avformat_close_input(&pFormatCtx);        pFormatCtx = nullptr;        av_packet_free(&packet);        packet = nullptr;        usleep(1000000);    }}

深耕企业安全管理+AI领域，通过“技术+商业+内容”的融合视角，深度参与AI产业化落地。

全网20W+粉丝AI知识博主，人工智能技术文章超1000W+阅读，《30天入门人工智能》课程，全网2000+名学员。

主导构建的AI知识平台www.jiangdabai.com累计访问已超800万次；

大家一起加油！